[https://nvbugs/6577736][perf] do not materialize embedding bias rows on host

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

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)
 

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)

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,27 +521,46 @@ 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
     ).to(logits.device, non_blocking=True)
     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

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:

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()
 

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):

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,

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"),

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():
 

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 == embedding_bias
 
     def test_embedding_bias_must_match_vocab_size(self):
         """Embedding bias is rejected when it cannot match model vocabulary."""