From 629b9ef601c83d59398460252844b17b9fe784b7 Mon Sep 17 00:00:00 2001
From: ixlmar <206748156+ixlmar@users.noreply.github.com>
Date: Tue, 29 Sep 2026 15:10:23 +0000
Subject: [PATCH 1/4] perf: do not materialize embedding bias rows on host

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/llm_request.py |  8 +-
 tensorrt_llm/_torch/pyexecutor/py_executor.py | 11 +++
 .../_torch/pyexecutor/sampler/sampler.py      |  2 +-
 .../pyexecutor/sampler/sampler_features.py    | 49 ++++++----
 tensorrt_llm/executor/base_worker.py          | 30 +++++-
 tensorrt_llm/sampling_params.py               | 27 ++++--
 tensorrt_llm/serve/openai_protocol.py         | 17 ++--
 tests/unittest/executor/test_base_worker.py   | 18 ++++
 tests/unittest/grpc/smg/test_smg.py           |  2 +-
 tests/unittest/llmapi/test_executor.py        |  3 +-
 tests/unittest/llmapi/test_llm_pytorch.py     | 91 ++++++++++++-------
 tests/unittest/llmapi/test_sampling_params.py | 18 ++++
 12 files changed, 202 insertions(+), 74 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py
index 299163310c24..5b6eb2063be0 100644
--- a/tensorrt_llm/_torch/pyexecutor/llm_request.py
+++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py
@@ -18,7 +18,7 @@
 from tensorrt_llm.executor.result import SimpleTokenLogprobs, TokenLogprobs
 from tensorrt_llm.inputs.multimodal import strip_mm_encoder_inputs
 from tensorrt_llm.inputs.registry import get_multimodal_encoder_item_metadata
-from tensorrt_llm.sampling_params import LogprobMode
+from tensorrt_llm.sampling_params import EmbeddingBias, LogprobMode
 
 SamplingConfig = tllm_executor.SamplingConfig
 
@@ -907,6 +907,7 @@ def __init__(
             logits_chunk_size: int = 8,
             logprobs_mode: LogprobMode = LogprobMode.RAW,
             logprobs_simple_format: bool = False,
+            embedding_bias: Optional[EmbeddingBias] = None,
             return_routed_experts: bool = False,
             **kwargs):
         self.py_sampling_strategy: "Strategy | None" = None
@@ -915,8 +916,7 @@ def __init__(
                                                     None)
         # Owned entirely on the Python side: the TorchSampler is the only
         # consumer, so the C++ request no longer carries a copy.
-        self.py_embedding_bias: Optional[torch.Tensor] = kwargs.pop(
-            "embedding_bias", None)
+        self.py_embedding_bias: EmbeddingBias | None = embedding_bias
         self.py_lora_path: str | None = kwargs.pop("py_lora_path", None)
         # Multimodal data
         self.py_multimodal_data = kwargs.pop("py_multimodal_data", None)
@@ -1577,7 +1577,7 @@ def executor_request_to_llm_request(
         sampling_config=sampling_config,
         is_streaming=executor_request.streaming,
         end_id=executor_request.end_id,
-        embedding_bias=executor_request.embedding_bias,
+        embedding_bias=getattr(executor_request, "py_embedding_bias", None),
         stop_words_list=stop_words_list,
         position_ids=position_ids,
         prompt_embedding_table=None if executor_request.prompt_tuning_config
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index fc4c549f63ef..37242890b936 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -480,6 +480,7 @@ def __init__(
         # related modules
         self.resource_manager = resource_manager
         self.model_engine = model_engine
+        self.vocab_size: int = self.model_engine.model.model_config.pretrained_config.vocab_size
         self._enable_adp_dummy_fixes = getattr(model_engine,
                                                "_enable_adp_dummy_fixes", False)
         self._enable_scheduler_aware_adp_dummy = getattr(
@@ -5625,6 +5626,14 @@ def _validate_request_budget(self, request: LlmRequest) -> None:
                     f"max_new_tokens={request.max_new_tokens}, "
                     f"beam_width={request.py_beam_width}).")
 
+    def _validate_embedding_bias(self, request: LlmRequest) -> None:
+        embedding_bias = request.py_embedding_bias
+        if embedding_bias is not None:
+            for idx, val in embedding_bias:
+                if idx < 0 or idx >= self.vocab_size:
+                    raise ValueError(
+                        "Embedding bias index must be in [0, vocab_size - 1]")
+
     def _validate_request(self, request: LlmRequest):
         # Validate context-side pipelined-transfer constraints.
         disagg_params = request.py_disaggregated_params
@@ -5706,6 +5715,8 @@ def _validate_request(self, request: LlmRequest):
         # Check token ID ranges
         self._validate_token_id_range(request)
 
+        self._validate_embedding_bias(request)
+
         # Perform sampler-specific validation
         self.sampler.validate_request(request)
 
diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
index cca1cbd12a6c..29f9324ce266 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
@@ -778,7 +778,7 @@ def _request_rewrites_logits(request: LlmRequest) -> bool:
         in-graph, so a request using one cannot be sampled there at all.
         """
         return (
-            request._py_embedding_bias_1d is not None
+            request.py_embedding_bias is not None
             or bool(getattr(request, "py_bad_words", None))
             or bool(getattr(request, "py_no_repeat_ngram_size", None))
             or has_occurrence_penalty(request)
diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_features.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_features.py
index 6d7a50008166..b7f680c9bb13 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_features.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_features.py
@@ -47,6 +47,7 @@
 
 from tensorrt_llm._utils import nvtx_range, prefer_pinned
 from tensorrt_llm.bindings.executor import FinishReason
+from tensorrt_llm.sampling_params import EmbeddingBias
 
 from ..llm_request import LlmRequest
 from .sampler_common import DEFAULT_BEAM_IDX
@@ -496,11 +497,8 @@ def provision_bias_index() -> int:
         _next_bias_index += 1
         return bias_index
 
-    # Indices of unique bias tensors
-    #
-    # NB: hash(torch.Tensor) is equivalent to id(torch.Tensor), and does not
-    #     depend on tensor contents, cf. https://github.com/pytorch/pytorch/issues/2569
-    bias_to_index: dict[torch.Tensor, int] = defaultdict(provision_bias_index)
+    # Indices of unique bias rows
+    bias_to_index: dict[EmbeddingBias, int] = defaultdict(provision_bias_index)
 
     # Source indices for bias application
     bias_gather_indices: list[int] = []
@@ -523,10 +521,6 @@ def provision_bias_index() -> int:
 
     if not bias_to_index:
         return
-    # NB: take the reference shape from the collected biases rather than from the
-    #     loop variable: that holds the *last* request's bias, which is None
-    #     whenever a biased request is followed by an unbiased one.
-    bias_tensors = tuple(bias_to_index)
 
     bias_gather_indices_cuda = torch.tensor(
         bias_gather_indices, pin_memory=prefer_pinned(), dtype=torch.int32
@@ -534,16 +528,39 @@ def provision_bias_index() -> int:
     logits_bias_mask_cuda = torch.tensor(
         logits_bias_masks, pin_memory=prefer_pinned(), dtype=torch.bool
     ).to(logits.device, non_blocking=True)
-    biases_tensor = torch.empty(
-        (len(bias_tensors), *bias_tensors[0].shape), pin_memory=prefer_pinned()
+
+    vocab_size = logits.size(-1)
+    unique_biases_tensor_cuda = torch.zeros(
+        (len(bias_to_index), vocab_size),
+        dtype=logits.dtype,
+        device=logits.device,
     )
-    biases_tensor = torch.stack(
-        bias_tensors,
-        out=biases_tensor,
+    unique_biases_flat_idx_cuda = torch.tensor(
+        [
+            vocab_size * unique_bias_idx + logit_idx
+            for sparse_bias, unique_bias_idx in bias_to_index.items()
+            for (logit_idx, bias_value) in sparse_bias
+        ],
+        pin_memory=prefer_pinned(),
+        dtype=torch.int32,
+    ).to(logits.device, non_blocking=True)
+    unique_biases_flat_values_cuda = torch.tensor(
+        [
+            bias_value
+            for sparse_bias, unique_bias_idx in bias_to_index.items()
+            for (logit_idx, bias_value) in sparse_bias
+        ],
+        pin_memory=prefer_pinned(),
+        dtype=logits.dtype,
+    ).to(logits.device, non_blocking=True)
+    unique_biases_tensor_cuda.view(-1).scatter_(
+        dim=0,
+        index=unique_biases_flat_idx_cuda,
+        src=unique_biases_flat_values_cuda,
     )
-    biases_tensor_cuda = biases_tensor.to(logits.device, non_blocking=True)
 
-    biases_tensor_cuda = torch.index_select(biases_tensor_cuda, 0, bias_gather_indices_cuda)
+    biases_tensor_cuda = torch.index_select(unique_biases_tensor_cuda, 0, bias_gather_indices_cuda)
+
     # NB: Avoiding logits[bias_scatter_indices] += biases_tensor (and torch.Tensor.scatter_add_), because it
     #     is unclear if this allows for repeated indices, cf.
     #         https://docs.pytorch.org/docs/2.8/generated/torch.Tensor.index_put_.html#torch-tensor-index-put
diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py
index c4dfa7552244..c4a408a6d462 100644
--- a/tensorrt_llm/executor/base_worker.py
+++ b/tensorrt_llm/executor/base_worker.py
@@ -40,7 +40,8 @@
 from ..llmapi.tracer import global_tracer
 from ..llmapi.utils import _SyncQueue, configure_cpu_affinity, logger_debug
 from ..runtime import ModelConfig
-from ..sampling_params import BatchedLogitsProcessor, SamplingParams
+from ..sampling_params import (BatchedLogitsProcessor, EmbeddingBias,
+                               SamplingParams)
 from .executor import GenerationExecutor, IterationResultQueue
 from .ipc import FusedIpcQueue, IpcQueue
 from .postproc_worker import (PostprocParams, PostprocWorker,
@@ -338,6 +339,24 @@ def _load_prompt_adapter(self,
             model_config=self._runtime_model_config,
             uids=[str(prompt_adapter_request.adapter_id)])
 
+    @staticmethod
+    def _embedding_bias_sparse_to_dense(embedding_bias: EmbeddingBias, *,
+                                        vocab_size: int) -> torch.Tensor:
+        return torch.zeros(
+            vocab_size,
+            dtype=torch.float32,
+        ).scatter_(
+            dim=0,
+            index=torch.tensor(
+                [idx for (idx, val) in embedding_bias],
+                dtype=torch.int32,
+            ),
+            src=torch.tensor(
+                [val for (idx, val) in embedding_bias],
+                dtype=torch.float32,
+            ),
+        )
+
     def _enqueue_request(self,
                          request: GenerationRequest,
                          result_wait_queue=None) -> int:
@@ -490,7 +509,13 @@ def _deduce_max_tokens(request: GenerationRequest) -> int:
                 bad_words=request.sampling_params._get_bad_words(),
                 stop_words=[] if request.sampling_params.ignore_eos else
                 request.sampling_params._get_stop_words(),
-                embedding_bias=request.sampling_params.embedding_bias,
+                embedding_bias=(
+                    None if request.sampling_params._embedding_bias_sparse
+                    is None or self._is_pytorch_backend else
+                    self._embedding_bias_sparse_to_dense(
+                        request.sampling_params._embedding_bias_sparse,
+                        vocab_size=self._runtime_model_config.vocab_size,
+                    )),
                 lora_config=lora_config,
                 prompt_tuning_config=prompt_tuning_config,
                 multimodal_input=multimodal_input,
@@ -512,6 +537,7 @@ def _deduce_max_tokens(request: GenerationRequest) -> int:
                 request.sampling_params.logprobs_simple_format)
             executor_request.py_return_routed_experts = (
                 request.sampling_params.return_routed_experts)
+            executor_request.py_embedding_bias = request.sampling_params._embedding_bias_sparse
 
             # here we add executor_request.py_disaggregated_params= request.disaggregated_params for python cache transceiver
             if self._is_pytorch_backend and request.disaggregated_params is not None:
diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py
index bfa3d909a429..cb280775d4bb 100644
--- a/tensorrt_llm/sampling_params.py
+++ b/tensorrt_llm/sampling_params.py
@@ -17,7 +17,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Iterable, Mapping
 from dataclasses import dataclass, field, fields
-from typing import Any, List, NamedTuple, Optional, Tuple, Union
+from typing import Any, List, NamedTuple, Optional, Tuple, TypeAlias, Union
 
 import torch
 from pydantic import BaseModel
@@ -182,6 +182,9 @@ def __call__(
         pass  # noqa
 
 
+EmbeddingBias: TypeAlias = tuple[tuple[int, float], ...]
+
+
 @dataclass(slots=True, kw_only=True)
 class SamplingParams:
     """Sampling parameters for text generation.
@@ -204,7 +207,10 @@ class SamplingParams:
         stop (str, List[str], optional): A string or a list of strings that stop the generation when they are generated. The returned output will not contain the stop strings unless include_stop_str_in_output is True. Defaults to None.
         stop_token_ids (List[int], optional): A list of token ids that stop the generation when they are generated. Defaults to None.
         include_stop_str_in_output (bool): Whether to include the stop strings in output text. Defaults to False.
-        embedding_bias (torch.Tensor, optional): The embedding bias tensor. Expected type is kFP32 and shape is [vocab_size]. Defaults to None.
+        embedding_bias (EmbeddingBias, torch.Tensor, List[float], optional): The embedding bias tensor in dense or sparse representation.
+            Dense representations can be torch.Tensor (float32 dtype, shape (vocab_size,)) or list of length vocab_size.
+            Sparse representation is a tuple of (vocabulary index, bias value) pairs. The behavior for repeated indices is undefined.
+            Defaults to None.
         logits_processor (tensorrt_llm.sampling_params.LogitsProcessor, List[tensorrt_llm.sampling_params.LogitsProcessor], optional): The logits postprocessor callback(s). Defaults to None.
             If a list, each processor is applied in order during generation (supported in PyTorch backend only).
         apply_batched_logits_processor (bool): Whether to apply batched logits postprocessor callback. Defaults to False.
@@ -291,7 +297,8 @@ class SamplingParams:
     include_stop_str_in_output: bool = False
     _stop_word_ids: Optional[List[List[int]]] = field(default=None, init=False, repr=False)
 
-    embedding_bias: Optional[torch.Tensor] = None
+    embedding_bias: Optional[Union[torch.Tensor, list[float], EmbeddingBias]] = None
+    _embedding_bias_sparse: Optional[EmbeddingBias] = field(default=None, init=False, repr=False)
     logits_processor: Optional[Union[LogitsProcessor, List[LogitsProcessor]]] = None
     apply_batched_logits_processor: bool = False
 
@@ -373,10 +380,18 @@ def __post_init__(self):
         self.best_of = self.best_of or self.n
 
         if self.embedding_bias is not None:
-            if isinstance(self.embedding_bias, torch.Tensor):
-                self.embedding_bias = self.embedding_bias.detach().clone()
+            if isinstance(self.embedding_bias, (torch.Tensor, list)):
+                if isinstance(self.embedding_bias, torch.Tensor):
+                    embedding_bias_list = self.embedding_bias.tolist()
+                else:
+                    embedding_bias_list = self.embedding_bias
+                self._embedding_bias_sparse = tuple(
+                    [(idx, val) for idx, val in enumerate(embedding_bias_list) if val != 0]
+                )
             else:
-                self.embedding_bias = torch.tensor(self.embedding_bias, dtype=torch.float32)
+                self._embedding_bias_sparse = self.embedding_bias
+        else:
+            self._embedding_bias_sparse = None
 
         self._validate()
 
diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py
index 0d85fd554229..3ff7c7b83ca3 100644
--- a/tensorrt_llm/serve/openai_protocol.py
+++ b/tensorrt_llm/serve/openai_protocol.py
@@ -56,10 +56,10 @@
 from tensorrt_llm.inputs.media_io import MediaModality
 from tensorrt_llm.llmapi import ConversationParams as LlmConversationParams
 from tensorrt_llm.llmapi import DisaggregatedParams as LlmDisaggregatedParams
-from tensorrt_llm.llmapi import (DisaggScheduleStyle, GuidedDecodingParams,
-                                 SamplingParams)
+from tensorrt_llm.llmapi import DisaggScheduleStyle, GuidedDecodingParams
 from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory
-from tensorrt_llm.sampling_params import (check_logprobs_limit,
+from tensorrt_llm.sampling_params import (EmbeddingBias, SamplingParams,
+                                          check_logprobs_limit,
                                           validate_thinking_token_budget)
 from tensorrt_llm.scheduling_params import AgentHierarchy
 from tensorrt_llm.visual_gen.params import MediaRole
@@ -98,8 +98,8 @@ def ensure_request_chat_template_allowed(request: Any,
 
 def _logit_bias_to_embedding_bias(
         logit_bias: Optional[Dict[str, float]],
-        vocab_size: Optional[int]) -> Optional[torch.Tensor]:
-    """Convert OpenAI logit_bias dict to embedding_bias tensor for sampling."""
+        vocab_size: Optional[int]) -> Optional[EmbeddingBias]:
+    """Convert OpenAI logit_bias dict to sparse embedding_bias for sampling."""
     if logit_bias is None:
         return None
     if vocab_size is None:
@@ -112,7 +112,7 @@ def _logit_bias_to_embedding_bias(
         raise ValueError("vocab_size must be positive when logit_bias is used")
 
     # Create 1D zeros tensor as expected by executor API (will be unsqueezed to [1, vocab_size] internally)
-    embedding_bias = torch.zeros(vocab_size, dtype=torch.float32)
+    embedding_bias_list = []
 
     # Apply biases for specified token IDs
     for token_str, bias in logit_bias.items():
@@ -141,9 +141,10 @@ def _logit_bias_to_embedding_bias(
             raise ValueError(
                 f"logit_bias value for token ID {token_id} must be in "
                 f"[{_LOGIT_BIAS_MIN:g}, {_LOGIT_BIAS_MAX:g}]")
-        embedding_bias[token_id] = bias_value
+        embedding_bias_list.append((token_id, bias_value))
 
-    return embedding_bias
+    return tuple(
+        sorted(embedding_bias_list, key=lambda idx_and_val: idx_and_val[0]))
 
 
 class OpenAIBaseModel(BaseModel):
diff --git a/tests/unittest/executor/test_base_worker.py b/tests/unittest/executor/test_base_worker.py
index bc9a814e76c7..0590ba2dadf8 100644
--- a/tests/unittest/executor/test_base_worker.py
+++ b/tests/unittest/executor/test_base_worker.py
@@ -21,6 +21,24 @@
 model_path = llm_models_root() / default_model_name
 
 
+@pytest.mark.cpu_only
+def test_embedding_bias_sparse_to_dense():
+    vocab_size = 123
+    embedding_bias = ((1, -0.5), (42, -2.0))
+    dense_bias = BaseWorker._embedding_bias_sparse_to_dense(
+        embedding_bias,
+        vocab_size=vocab_size,
+    )
+    assert dense_bias.shape == (vocab_size, )
+    for idx, val in embedding_bias:
+        assert dense_bias[idx] == val
+    idx_nz = set(idx for idx, _ in embedding_bias)
+    for idx in range(vocab_size):
+        if idx in idx_nz:
+            continue
+        assert dense_bias[idx] == 0
+
+
 @pytest.mark.cpu_only
 def test_enqueue_request_wraps_lora_load_error():
 
diff --git a/tests/unittest/grpc/smg/test_smg.py b/tests/unittest/grpc/smg/test_smg.py
index bbf6f5a31717..4df4dfbc7e4b 100644
--- a/tests/unittest/grpc/smg/test_smg.py
+++ b/tests/unittest/grpc/smg/test_smg.py
@@ -482,7 +482,7 @@ def test_all_sampling_config_fields(self):
 
         # Embedding bias converted to torch.Tensor
         assert params.embedding_bias is not None
-        assert len(params.embedding_bias) == 12
+        assert params.embedding_bias == ((10, 1.5), (11, -1.5))
 
     def test_embedding_bias_must_match_vocab_size(self):
         """Embedding bias is rejected when it cannot match model vocabulary."""
diff --git a/tests/unittest/llmapi/test_executor.py b/tests/unittest/llmapi/test_executor.py
index 5aaa3433794a..c9d4d8b4fe91 100644
--- a/tests/unittest/llmapi/test_executor.py
+++ b/tests/unittest/llmapi/test_executor.py
@@ -8,7 +8,6 @@
 from queue import Empty
 
 import pytest
-import torch
 import zmq
 
 from tensorrt_llm._utils import mpi_world_size
@@ -361,7 +360,7 @@ def test_ZeroMqQueue_serialization_complicated_dataclass():
         80, None, tllm.KvCacheTransferMode.DRAM, "test_dir")
 
     sampling_params = SamplingParams(max_tokens=4,
-                                     embedding_bias=torch.randn(2, 2))
+                                     embedding_bias=((42, -2.0), (123, -1.0)))
 
     for i in range(iterations):
         request = GenerationRequest(prompt_token_ids=[i],
diff --git a/tests/unittest/llmapi/test_llm_pytorch.py b/tests/unittest/llmapi/test_llm_pytorch.py
index 5a5d80c99382..0369117382a3 100644
--- a/tests/unittest/llmapi/test_llm_pytorch.py
+++ b/tests/unittest/llmapi/test_llm_pytorch.py
@@ -7,7 +7,7 @@
 import random
 import time
 from contextlib import contextmanager, nullcontext
-from typing import Optional
+from typing import Generator, Optional
 
 import pytest
 
@@ -19,7 +19,7 @@
 from tensorrt_llm.llmapi.llm_args import (NGramDecodingConfig, SchedulerConfig,
                                           WaitingQueuePolicy)
 from tensorrt_llm.metrics import MetricNames
-from tensorrt_llm.sampling_params import SamplingParams
+from tensorrt_llm.sampling_params import EmbeddingBias, SamplingParams
 
 # isort: off
 from .lora_test_utils import (create_mock_nemo_lora_checkpoint,
@@ -139,41 +139,36 @@ def test_llm_abort_request(sampling_params):
     run_llm_abort_request(llm=llm, sampling_params=sampling_params)
 
 
-@contextmanager
-def _validate_invalid_token_error_scope():
-    with pytest.raises(RuntimeError) as exc_info:
-        yield
-    assert "Token ID out of range" in str(exc_info.value)
-
-
 @force_ampere
 @pytest.mark.part1
-def test_llm_invalid_input_token():
-    llm = LLM(model=llama_model_path, kv_cache_config=global_kvcache_config)
-    prompts = [
-        [-1],
-    ]
-    # NB: exc_info in _validate_invalid_token_error_scope creates a reference
-    #     to a traceback which outlives the scope of 'exc_info' and prevents
-    #     deletion of 'llm'. However, using the context manager protocol is
-    #     anyways more robust than delegating cleanup to __del__.
-    with llm:
-        with _validate_invalid_token_error_scope():
+class TestLlmInputValidation:
+
+    @pytest.fixture(scope="class")
+    def llm(self) -> Generator[LLM, None, None]:
+        with LLM(model=llama_model_path,
+                 kv_cache_config=global_kvcache_config) as llm:
+            yield llm
+
+    @staticmethod
+    @contextmanager
+    def _validate_invalid_token_error_scope():
+        with pytest.raises(RuntimeError) as exc_info:
+            yield
+        assert "Token ID out of range" in str(exc_info.value)
+
+    @pytest.mark.threadleak(enabled=False)
+    def test_llm_invalid_input_token(self, llm: LLM):
+        prompts = [
+            [-1],
+        ]
+        with self._validate_invalid_token_error_scope():
             llm.generate(
                 prompts,
                 sampling_params=SamplingParams(max_tokens=5),
             )
 
-
-@force_ampere
-@pytest.mark.part0
-def test_llm_invalid_input_token_async():
-    llm = LLM(model=llama_model_path, kv_cache_config=global_kvcache_config)
-    # NB: exc_info in _validate_invalid_token_error_scope creates a reference
-    #     to a traceback which outlives the scope of 'exc_info' and prevents
-    #     deletion of 'llm'. However, using the context manager protocol is
-    #     anyways more robust than delegating cleanup to __del__.
-    with llm:
+    @pytest.mark.threadleak(enabled=False)
+    def test_llm_invalid_input_token_async(self, llm: LLM):
         prompts = [
             [-1],
             [42],
@@ -189,7 +184,7 @@ def test_llm_invalid_input_token_async():
                     ) for submit_idx in submit_order
                 ]
                 for collect_idx in collect_order:
-                    with _validate_invalid_token_error_scope(
+                    with self._validate_invalid_token_error_scope(
                     ) if submit_order[collect_idx] in fail_idx else nullcontext(
                     ):
                         print(
@@ -197,6 +192,36 @@ def test_llm_invalid_input_token_async():
                         )
                         futures[collect_idx].result()
 
+    @pytest.mark.threadleak(enabled=False)
+    @pytest.mark.parametrize(
+        "embedding_bias_in",
+        [
+            # NB: Tensors with fewer elements than vocab_size are currently implicitly zero-padded
+            torch.tensor([1] * 32001, dtype=torch.float32),  # invalid length
+            [1] * 32001,  # invalid length
+            ((42, -0.5), (-1, -2.5)),  # out of vocabulary
+            ((42, -0.5), (999999, -2.5)),  # out of vocabulary
+        ],
+    )
+    def test_llm_invalid_embedding_bias(
+        self,
+        embedding_bias_in: torch.Tensor | list[float] | EmbeddingBias,
+        llm: LLM,
+    ):
+        prompts = [
+            [0],
+        ]
+        with pytest.raises(
+                RuntimeError,
+                match=
+                ".*Embedding bias index must be in \\[0, vocab_size - 1\\].*",
+        ):
+            llm.generate(
+                prompts,
+                sampling_params=SamplingParams(
+                    max_tokens=5, embedding_bias=embedding_bias_in),
+            )
+
 
 @skip_ray
 @pytest.mark.part3
@@ -296,9 +321,7 @@ def test_embedding_bias_with_torch_sampler_strategies():
     """Test embedding bias application in TorchSampler."""
     tokenizer = AutoTokenizer.from_pretrained(llama_model_path)
     biased_word_id = tokenizer.encode("Z", add_special_tokens=False)[-1]
-    vocab_size_padded = 32000
-    embedding_bias = torch.zeros(vocab_size_padded)
-    embedding_bias[biased_word_id] = torch.finfo(torch.float32).max
+    embedding_bias = ((biased_word_id, torch.finfo(torch.float32).max), )
 
     sampling_kwargs = {
         "max_tokens": 6,
diff --git a/tests/unittest/llmapi/test_sampling_params.py b/tests/unittest/llmapi/test_sampling_params.py
index 3c42fc097cf5..19ae94beda4b 100644
--- a/tests/unittest/llmapi/test_sampling_params.py
+++ b/tests/unittest/llmapi/test_sampling_params.py
@@ -269,6 +269,24 @@ def test_sampling_params_rejects_nan(field):
         SamplingParams(**{field: float("nan")})
 
 
+@pytest.mark.parametrize(
+    "embedding_bias_in",
+    [
+        None,
+        torch.tensor([0, 0, -1.0, 0, -2.5], dtype=torch.float32),
+        [0, 0, -1.0, 0, -2.5],
+        ((2, -1.0), (4, -2.5)),
+    ],
+)
+def test_sampling_params_sets_embedding_bias_sparse(embedding_bias_in):
+    params = SamplingParams(embedding_bias=embedding_bias_in)
+    assert params.embedding_bias is embedding_bias_in
+    if embedding_bias_in is None:
+        assert params._embedding_bias_sparse is None
+    else:
+        assert params._embedding_bias_sparse == ((2, -1.0), (4, -2.5))
+
+
 @pytest.mark.parametrize(
     ("field", "value"),
     [

From ab9f0ed36cd485c5b778006cf19fb2a17ee9cbf0 Mon Sep 17 00:00:00 2001
From: ixlmar <206748156+ixlmar@users.noreply.github.com>
Date: Wed, 30 Sep 2026 17:26:18 +0200
Subject: [PATCH 2/4] fix SMG test

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>
---
 tests/unittest/grpc/smg/test_smg.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/unittest/grpc/smg/test_smg.py b/tests/unittest/grpc/smg/test_smg.py
index 4df4dfbc7e4b..295bbff1aaab 100644
--- a/tests/unittest/grpc/smg/test_smg.py
+++ b/tests/unittest/grpc/smg/test_smg.py
@@ -482,7 +482,7 @@ def test_all_sampling_config_fields(self):
 
         # Embedding bias converted to torch.Tensor
         assert params.embedding_bias is not None
-        assert params.embedding_bias == ((10, 1.5), (11, -1.5))
+        assert params.embedding_bias == embedding_bias
 
     def test_embedding_bias_must_match_vocab_size(self):
         """Embedding bias is rejected when it cannot match model vocabulary."""

From abc60aad82b3a7a7c1b1bbc3f0106d2b9aae49ba Mon Sep 17 00:00:00 2001
From: ixlmar <206748156+ixlmar@users.noreply.github.com>
Date: Wed, 30 Sep 2026 17:11:44 +0000
Subject: [PATCH 3/4] fix PyExecutor mock in a unit test

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>
---
 tests/unittest/_torch/sampler/test_beam_search.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py
index 0114ac5ba624..7e243fab7190 100644
--- a/tests/unittest/_torch/sampler/test_beam_search.py
+++ b/tests/unittest/_torch/sampler/test_beam_search.py
@@ -1984,6 +1984,7 @@ def test_vbws_rejects_decreasing_beam_width_array(beam_width_array: list[int],
         _validate_token_id_range=lambda _request: None,
         sampler=types.SimpleNamespace(validate_request=lambda _request: None),
         _validate_request_budget=lambda _request: None,
+        _validate_embedding_bias=lambda _request: None,
     )
     validate = functools.partial(
         PyExecutor._validate_request,

From ee7cf2a84a6aa0ba644b409ff38f971ff57dd41f Mon Sep 17 00:00:00 2001
From: ixlmar <206748156+ixlmar@users.noreply.github.com>
Date: Mon, 5 Oct 2026 09:58:41 +0000
Subject: [PATCH 4/4] fix fallout in another test

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>
---
 tests/unittest/_torch/sampler/test_in_graph_sampling.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/unittest/_torch/sampler/test_in_graph_sampling.py b/tests/unittest/_torch/sampler/test_in_graph_sampling.py
index 4660aa29641c..9dc36de81a39 100644
--- a/tests/unittest/_torch/sampler/test_in_graph_sampling.py
+++ b/tests/unittest/_torch/sampler/test_in_graph_sampling.py
@@ -69,7 +69,7 @@ class _Request:
     py_is_draft: bool = False
     py_logits_post_processors: Any = None
     guided_decoding_params: Any = None
-    _py_embedding_bias_1d: Any = None
+    py_embedding_bias: Any = None
 
     def get_beam_width_by_iter(self, for_next_iteration: bool = False) -> int:
         return self.sampling_config.beam_width
@@ -131,7 +131,7 @@ def test_sampling_params_outside_the_tier_are_full(self, sampler, sampling):
             pytest.param(dict(py_min_length=50), id="min_length"),
             pytest.param(dict(py_bad_words=[[7]]), id="bad_words"),
             pytest.param(dict(py_no_repeat_ngram_size=3), id="no_repeat_ngram"),
-            pytest.param(dict(_py_embedding_bias_1d=object()), id="embedding_bias"),
+            pytest.param(dict(py_embedding_bias=object()), id="embedding_bias"),
             # Consume the logits after the forward returns.
             pytest.param(dict(guided_decoding_params=object()), id="guided_decoding"),
             pytest.param(dict(py_logits_post_processors=[object()]), id="logits_processor"),