[https://nvbugs/6577736][perf] cache embedding biases (PR to be merged after #19412)

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/ops/vanilla.py

@@ -710,3 +710,38 @@ def apply_batched_occurrence_penalties(
             max_beam_width,
             fold_pending,
         )
+
+    @staticmethod
+    @torch.compile(
+        dynamic=None,
+        fullgraph=True,
+        options=dict(
+            split_reductions=False,
+        ),
+    )
+    def _gather_masked_inplace_add_impl(
+        summands_cuda: torch.Tensor,
+        summand_indices_cuda: torch.Tensor,
+        output_cuda: torch.Tensor,
+        output_mask_cuda: torch.Tensor,
+    ) -> None:
+        output_cuda[output_mask_cuda, ...] += torch.index_select(
+            summands_cuda, 0, summand_indices_cuda
+        )
+
+    @staticmethod
+    def gather_masked_inplace_add(
+        summands_cuda: torch.Tensor,
+        summand_indices_cuda: torch.Tensor,
+        output_cuda: torch.Tensor,
+        output_mask_cuda: torch.Tensor,
+    ) -> None:
+        # NB: helper function for apply_embedding_bias, torch.compile is expected to reduce
+        #     data movement.
+        torch._dynamo.mark_dynamic(summands_cuda, 0)
+        torch._dynamo.mark_dynamic(summand_indices_cuda, 0)
+        torch._dynamo.mark_dynamic(output_cuda, 0)
+        torch._dynamo.mark_dynamic(output_mask_cuda, 0)
+        Fusions._gather_masked_inplace_add_impl(
+            summands_cuda, summand_indices_cuda, output_cuda, output_mask_cuda
+        )

tensorrt_llm/_torch/pyexecutor/sampler/sampler.py

@@ -76,6 +76,7 @@
 )
 from .sampler_features import (
     AsyncWorkerMixin,
+    EmbeddingBiasCache,
     SamplerEvent,
     _PackedStepIndexer,
     _UnpackedStepIndexer,
@@ -590,6 +591,8 @@ def __init__(self, args: Args):
         # constructed before _create_store, which reads the shape.
         self._log_probs = LogProbsHandler(self)
 
+        self._embedding_bias_cache: EmbeddingBiasCache | None = None
+
         # The Torch sampler hard-depends on flashinfer. Enforce it once here, at
         # construction, so the check stays out of the CUDA-graph-captured
         # sampling loop.
@@ -778,7 +781,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)
@@ -2688,8 +2691,12 @@ def _process_requests(
         seq_lens_cuda = seq_lens_host.to(device="cuda", non_blocking=True)
 
         # Handle embedding bias
-        apply_embedding_bias(
-            logits_cuda, sampling_requests, sampling_requests_metadata.req_num_steps
+        self._embedding_bias_cache = apply_embedding_bias(
+            logits_cuda,
+            sampling_requests,
+            sampling_requests_metadata.req_num_steps,
+            max_num_sequences=self.max_num_sequences,
+            maybe_embedding_bias_cache=self._embedding_bias_cache,
         )
 
         # Apply repetition/presence/frequency penalties in place, before the greedy fast

tensorrt_llm/_torch/pyexecutor/sampler/sampler_features.py

@@ -41,14 +41,16 @@
 from contextlib import contextmanager
 from dataclasses import dataclass
 from itertools import repeat
-from typing import Any, Iterator, Optional, cast
+from typing import Any, Final, Iterator, Optional, cast
 
 import torch
 
 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 .ops.vanilla import Fusions
 from .sampler_common import DEFAULT_BEAM_IDX
 
 if sys.version_info[:2] >= (3, 12):
@@ -458,11 +460,111 @@ def fast_greedy_sample_kernel(
     return next_tokens
 
 
+class EmbeddingBiasCache:
+    def __init__(
+        self, *, logit_dtype: torch.dtype, vocab_size: int, max_seq_slots: int, device: torch.device
+    ):
+        self._flat_cache = torch.empty(
+            (max_seq_slots, vocab_size), dtype=logit_dtype, device=device
+        )
+        self._device = device
+        self._logit_dtype = logit_dtype
+        self._vocab_size = vocab_size
+        self.CACHE_SENTINEL_ATTR_NAME: Final[str] = "_py_has_cached_embedding_bias"
+
+    def update(self, reqs: list[LlmRequest]) -> None:
+        # NB: This could be extended to allow reuse of equivalent rows in _flat_cache
+        #     (potentially relevant for output sequence length <~ batch size).
+
+        reqs_needing_update = [
+            req
+            for req in reqs
+            if req.py_embedding_bias and not getattr(req, self.CACHE_SENTINEL_ATTR_NAME, None)
+        ]
+
+        if not reqs_needing_update:
+            return
+
+        _next_bias_index = 0
+
+        def provision_bias_index() -> int:
+            nonlocal _next_bias_index
+            bias_index = _next_bias_index
+            _next_bias_index += 1
+            return bias_index
+
+        # Indices of unique bias rows
+        bias_to_index: dict[EmbeddingBias, int] = defaultdict(provision_bias_index)
+
+        # Source indices for bias cache update
+        bias_gather_indices: list[int] = []
+        # Destination indices for bias cache update
+        bias_scatter_indices: list[int] = []
+
+        # Collect bias information
+        for req in reqs_needing_update:
+            req_bias = req.py_embedding_bias
+            assert req_bias is not None
+            bias_gather_indices.append(bias_to_index[req_bias])
+            seq_slot = req.py_seq_slot
+            assert seq_slot is not None
+            bias_scatter_indices.append(seq_slot)
+
+        bias_gather_indices_cuda = torch.tensor(
+            bias_gather_indices, pin_memory=prefer_pinned(), dtype=torch.int32
+        ).to(self._device, non_blocking=True)
+        bias_scatter_indices_cuda = torch.tensor(
+            bias_scatter_indices, pin_memory=prefer_pinned(), dtype=torch.int32
+        ).to(self._device, non_blocking=True)
+
+        unique_biases_tensor_cuda = torch.zeros(
+            (len(bias_to_index), self._vocab_size),
+            dtype=self._logit_dtype,
+            device=self._device,
+        )
+        unique_biases_flat_idx_cuda = torch.tensor(
+            [
+                self._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(self._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=self._logit_dtype,
+        ).to(self._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,
+        )
+
+        self._flat_cache[bias_scatter_indices_cuda, :] = unique_biases_tensor_cuda[
+            bias_gather_indices_cuda, :
+        ]
+
+        for req in reqs_needing_update:
+            setattr(req, self.CACHE_SENTINEL_ATTR_NAME, True)
+
+    def get(self) -> torch.Tensor:
+        return self._flat_cache
+
+
 def apply_embedding_bias(
     logits: torch.Tensor,
     requests: list[LlmRequest],
     request_steps: torch.Tensor,
-) -> None:
+    max_num_sequences: int,
+    maybe_embedding_bias_cache: EmbeddingBiasCache | None = None,
+) -> EmbeddingBiasCache | None:
     """Apply embedding bias (aka logit bias) to logits.
 
     Arguments:
@@ -485,22 +587,23 @@ def apply_embedding_bias(
     #              and low reuse).
     #     Since read-caching is expected to help in typical cases, option (ii) is implemented here.
 
+    if all(not req.py_embedding_bias for req in requests):
+        return maybe_embedding_bias_cache
+
+    vocab_size = logits.size(-1)
+    if maybe_embedding_bias_cache is None:
+        embedding_bias_cache = EmbeddingBiasCache(
+            logit_dtype=logits.dtype,
+            vocab_size=vocab_size,
+            max_seq_slots=max_num_sequences,
+            device=logits.device,
+        )
+    else:
+        embedding_bias_cache = maybe_embedding_bias_cache
+
     # Track which logits require logit bias application
     request_steps_list = request_steps.tolist()
     logits_bias_masks = [False] * logits.size(0)
-    _next_bias_index = 0
-
-    def provision_bias_index() -> int:
-        nonlocal _next_bias_index
-        bias_index = _next_bias_index
-        _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)
 
     # Source indices for bias application
     bias_gather_indices: list[int] = []
@@ -509,47 +612,38 @@ def provision_bias_index() -> int:
     #
     # NB: the mask indexes rows of 'logits', not requests: a request contributes
     #     'steps' consecutive rows (> 1 under speculative decoding), so the offset
-    #     has to accumulate 'steps' rather than track the request index. It must
-    #     stay in step with 'bias_gather_indices', which is built in row order.
+    #     has to accumulate 'steps' rather than track the request index.
     row_offset = 0
+    reqs_using_bias = []
     for req, steps in zip(requests, request_steps_list):
         req_bias = req.py_embedding_bias
-        if req_bias is not None:
+        if req_bias:
+            reqs_using_bias.append(req)
             for j in range(row_offset, row_offset + steps):
                 logits_bias_masks[j] = True
-            req_bias_index = bias_to_index[req_bias]
-            bias_gather_indices.extend(repeat(req_bias_index, steps))
+            seq_slot = req.py_seq_slot
+            assert seq_slot is not None
+            bias_gather_indices.extend(repeat(seq_slot, steps))
         row_offset += steps
 
-    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)
+    embedding_bias_cache.update(reqs_using_bias)
+    biases_by_slot_cuda = embedding_bias_cache.get()  # indexed by seq_slot
 
     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()
-    )
-    biases_tensor = torch.stack(
-        bias_tensors,
-        out=biases_tensor,
+
+    Fusions.gather_masked_inplace_add(
+        summands_cuda=biases_by_slot_cuda,
+        summand_indices_cuda=bias_gather_indices_cuda,
+        output_cuda=logits,
+        output_mask_cuda=logits_bias_mask_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)
-    # 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
-    #     and thus introduces read-after-write dependencies (including possible false
-    #     sharing).
-    logits[logits_bias_mask_cuda] += biases_tensor_cuda
+
+    return embedding_bias_cache
 
 
 # --------------------------------------------------------------------------