[TRTLLM-14865][feat] Support occurrence penalties with beam search for TorchSampler

docs/source/features/sampling.md

@@ -196,8 +196,10 @@ modes.
     `repetition_penalty`. Values `<= 0` have no effect, and values larger than the prompt
     are clamped to the prompt length.
 
-  * Occurrence penalties are not supported in combination with beam search; such requests
-    are rejected.
+  * With beam search the occurrence history is kept per beam rather than per request:
+    each beam is penalized against the tokens on its own path, and whenever a beam
+    continues another one it inherits that beam's history. The prompt seeds every beam
+    alike, so `prompt_ignore_length` applies to all of them equally.
 
 * If `no_repeat_ngram_size = n` is specified, any token that would recreate an `n`-gram already
   present in the sequence (prompt included) is excluded from sampling. `None` or `0` disables

tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py

@@ -379,8 +379,148 @@ def update_occurrence_workspace(
             # Marking a dense bool mask is idempotent, so duplicate tokens are safe.
             presence_prefix_cuda[prefix_slots, prefix_tokens] = True
 
-    # fullgraph=True is safe here: served model has fixed shapes and compiles ~2 graphs,
-    # well under the default limit (8)
+    # --- Beam-search occurrence counts --------------------------------------
+    # Counterpart of the per-beam workspace handling in the C++ ``batchApplyPenalty``
+    # kernel (``penaltyKernels.cu:151-171``): a beam does not re-walk its history,
+    # it inherits its parent beam's counts and appends the single token it just
+    # emitted. ``counts_cuda`` is flat ``[num_slots * max_beam_width, vocab_size]``;
+    # beam ``b`` of slot ``s`` owns row ``s * max_beam_width + b``, which collapses to
+    # plain slot indexing at ``max_beam_width == 1``. A single-beam engine never calls
+    # this op -- it folds inside the penalty graph instead.
+
+    @staticmethod
+    @torch.compile(fullgraph=True, mode="max-autotune-no-cudagraphs")
+    def _update_beam_occurrence_counts_impl(
+        counts_cuda: torch.Tensor,
+        active_cuda: torch.Tensor,
+        has_previous_token_cuda: torch.Tensor,
+        beam_slot_cuda: torch.Tensor,
+        new_tokens: torch.Tensor,
+        predecessor_beams: torch.Tensor,
+        seq_slots: torch.Tensor,
+        request_num_beams: torch.Tensor,
+        max_beam_width: int,
+    ) -> None:
+        """Re-parent every beam onto the beam it continues, then fold in its token.
+
+        Beam ``b`` of slot ``s`` takes over the counts of ``predecessor_beams[s, b]``
+        before this step's token is folded in. Slots that did not sample last step read
+        the identity permutation instead, gated by ``has_previous_token``; single-beam
+        slots on a beam engine are gated out by ``beam_slot_cuda``, since their
+        ``predecessor_beams`` row is never written.
+
+        ``armed`` gates per *slot* (``.unsqueeze(1)``), so beams past the current width are
+        re-parented too, from a clamped and possibly stale parent -- safe because nothing
+        reads such a row until the width grows to cover it, and growth re-parents it from a
+        beam that was valid last step. The gate is device-side, so this runs
+        unconditionally rather than pay a D2H sync.
+
+        NB: ``fullgraph=True``, mutates ``counts_cuda`` in place, and every batch-varying
+        dim-0 argument must be marked dynamic by the caller -- an unmarked peer forces the
+        marked dims to specialize. Add such an argument only together with its
+        ``mark_dynamic`` in ``update_beam_occurrence_counts``.
+        """
+        vocab = counts_cuda.size(-1)
+        beam_ids = torch.arange(max_beam_width, device=counts_cuda.device)
+
+        armed = (
+            active_cuda[seq_slots] & has_previous_token_cuda[seq_slots] & beam_slot_cuda[seq_slots]
+        ).unsqueeze(1)
+        # Beams past the current width hold stale parents; clamping keeps the
+        # gather in range. Their rows are never read while masked out, and a later
+        # beam-width growth only ever re-gathers from a valid parent row.
+        parent = predecessor_beams[seq_slots].to(torch.int64).clamp(0, max_beam_width - 1)
+        src_beam = torch.where(armed, parent, beam_ids.expand_as(parent))
+        base = seq_slots.unsqueeze(1) * max_beam_width
+        counts_cuda.index_copy_(
+            0,
+            (base + beam_ids).reshape(-1),
+            counts_cuda.index_select(0, (base + src_beam).reshape(-1)),
+        )
+
+        # Same masked flat scatter as the single-beam fold, fanned out over the
+        # beam axis: masked entries add 0 at counts[row, 0], so inactive, unarmed,
+        # out-of-layout and padded-token beams are no-ops. A beam the previous step
+        # did not produce carries BEAM_SEARCH_PAD_TOKEN (-1), which the range check
+        # rejects -- that is what confines the fold to the beams that actually
+        # sampled, since ``request_num_beams`` is the (wider) row-layout width.
+        previous_token = new_tokens[0, seq_slots, :].to(torch.int64)  # [R, max_beam_width]
+        fold_ok = (
+            (active_cuda[seq_slots] & has_previous_token_cuda[seq_slots]).unsqueeze(1)
+            & (beam_ids.unsqueeze(0) < request_num_beams.unsqueeze(1))
+            & (previous_token >= 0)
+            & (previous_token < vocab)
+        )
+        rows = base + beam_ids  # [R, max_beam_width]
+        flat_index = rows * vocab + torch.where(
+            fold_ok, previous_token, previous_token.new_zeros(())
+        )
+        counts_cuda.view(-1).scatter_add_(
+            0, flat_index.reshape(-1), fold_ok.reshape(-1).to(counts_cuda.dtype)
+        )
+
+    @staticmethod
+    def update_beam_occurrence_counts(
+        counts_cuda: torch.Tensor,
+        active_cuda: torch.Tensor,
+        has_previous_token_cuda: torch.Tensor,
+        beam_slot_cuda: torch.Tensor,
+        new_tokens: torch.Tensor,
+        predecessor_beams: torch.Tensor,
+        seq_slots: torch.Tensor,
+        request_num_beams: torch.Tensor,
+        max_beam_width: int,
+    ) -> None:
+        """Advance the per-beam occurrence counts by one step, in place.
+
+        Re-parents every beam onto the beam it continues, then folds in the token
+        each beam sampled last step. Must run before anything reads the counts for
+        this step, and before this step's sampling overwrites ``predecessor_beams``.
+        This wrapper only marks the batch-varying dims dynamic; the work is in
+        ``_update_beam_occurrence_counts_impl``.
+
+        Also covers single-beam requests sharing a beam engine: ``beam_slot_cuda``
+        turns their re-parent into the identity, and ``request_num_beams == 1``
+        confines their fold to beam 0. Context requests are likewise laid out at
+        one beam and read as such.
+
+        Args:
+            counts_cuda: ``int32[num_slots * max_beam_width, vocab_size]`` workspace.
+            active_cuda / has_previous_token_cuda: per-slot gates, length ``num_slots``.
+            new_tokens: ``[max_tokens, num_slots, max_beam_width]`` device buffer
+                holding the previous step's sampled token per beam.
+            predecessor_beams: ``int32[num_slots, max_beam_width]``, the parent beam
+                of each beam as written by the previous step's beam search.
+            seq_slots: ``int64[R]`` slot per request.
+            request_num_beams: ``[R]`` row-layout beam width per request, i.e. the
+                static admission width the logits rows are laid out at (1 for a
+                context request), not the per-iteration width. Beams at or past it
+                are skipped; under a growing ``beam_width_array`` the beams between
+                the per-iteration and the layout width are skipped instead by their
+                BEAM_SEARCH_PAD_TOKEN.
+        """
+        if seq_slots.numel() == 0:
+            return
+        # Batch-varying dim-0 tensors; mark every one, or an unmarked peer forces the
+        # marked dims to specialize (cf. apply_batched_occurrence_penalties). The other
+        # arguments keep dim 0 == num_slots (fixed), so they must NOT be marked.
+        torch._dynamo.mark_dynamic(seq_slots, 0)
+        torch._dynamo.mark_dynamic(request_num_beams, 0)
+        Fusions._update_beam_occurrence_counts_impl(
+            counts_cuda,
+            active_cuda,
+            has_previous_token_cuda,
+            beam_slot_cuda,
+            new_tokens,
+            predecessor_beams,
+            seq_slots,
+            request_num_beams,
+            max_beam_width,
+        )
+
+    # fullgraph=True is safe here: 4 cache entries per process against a default limit of 8
+    # -- the size-1 specialization of the dynamic dim, times the `prefix_seen_cuda` Optional
+    # flipping on first use. The beam parameters are per-engine constants and cost nothing.
     @staticmethod
     @torch.compile(fullgraph=True, mode="max-autotune-no-cudagraphs")
     def _apply_occurrence_penalties_impl(
@@ -396,24 +536,32 @@ def _apply_occurrence_penalties_impl(
         repetition_cuda: torch.Tensor,
         presence_cuda: torch.Tensor,
         frequency_cuda: torch.Tensor,
+        request_num_beams: Optional[torch.Tensor],
+        max_beam_width: int,
+        fold_pending: bool,
     ) -> None:
         vocab = logits.size(-1)
 
         # Fold the device-pending sampled token into the persistent counts, once per armed
         # active slot, before the gather reads them, via one flat scatter_add. Masked entries
         # add 0 at counts[slot, 0], so inactive/unarmed/out-of-range slots are no-ops.
-        previous_token = new_tokens[0, seq_slots, 0].to(torch.int64)
-        fold_ok = (
-            active_cuda[seq_slots]
-            & has_previous_token_cuda[seq_slots]
-            & (request_num_steps > 0)
-            & (previous_token >= 0)
-            & (previous_token < vocab)
-        )
-        flat_index = seq_slots * vocab + torch.where(
-            fold_ok, previous_token, previous_token.new_zeros(())
-        )
-        counts_cuda.view(-1).scatter_add_(0, flat_index, fold_ok.to(counts_cuda.dtype))
+        # Only reached on a single-beam engine, so ``max_beam_width`` is 1 and the row
+        # scaling folds away. A beam engine has to re-parent before folding, which it
+        # cannot do here, so it folds in ``update_beam_occurrence_counts`` instead.
+        if fold_pending:
+            slot_rows = seq_slots * max_beam_width
+            previous_token = new_tokens[0, seq_slots, 0].to(torch.int64)
+            fold_ok = (
+                active_cuda[seq_slots]
+                & has_previous_token_cuda[seq_slots]
+                & (request_num_steps > 0)
+                & (previous_token >= 0)
+                & (previous_token < vocab)
+            )
+            flat_index = slot_rows * vocab + torch.where(
+                fold_ok, previous_token, previous_token.new_zeros(())
+            )
+            counts_cuda.view(-1).scatter_add_(0, flat_index, fold_ok.to(counts_cuda.dtype))
 
         # Map each logits row to its owning request with a broadcasted range comparison.
         # This is O(T * R), but T and R are both small (rows per step x requests) and the
@@ -424,12 +572,29 @@ def _apply_occurrence_penalties_impl(
         # reads back from the device, and that per-step D2H sync destroys the overlap
         # between the sampler's host work and the model forward (measured ~20x slower).
         rows = torch.arange(logits.size(0), device=logits.device).unsqueeze(1)  # [T, 1]
-        owned = (rows >= request_offsets) & (rows < request_offsets + request_num_steps)  # [T, R]
+        # A request owns num_steps * num_beams rows, laid out beam-major / step-minor.
+        # request_num_beams is None on a single-beam engine, collapsing the span to
+        # num_steps and the counts row to plain slot indexing.
+        span = (
+            request_num_steps
+            if request_num_beams is None
+            else request_num_steps * request_num_beams
+        )
+        owned = (rows >= request_offsets) & (rows < request_offsets + span)  # [T, R]
         row_owned = owned.any(dim=1)  # [T]
         row_slot = (owned * seq_slots).sum(dim=1)  # [T]; slot per row, 0 for unowned
         row_active = row_owned & active_cuda[row_slot]
 
-        count = counts_cuda[row_slot]
+        if request_num_beams is None:
+            row_counts = row_slot
+        else:
+            # Recover the beam from the row's offset within its request: beam-major means
+            # beam = local_index // num_steps.
+            local = (owned * (rows - request_offsets)).sum(dim=1)  # [T]
+            steps = (owned * request_num_steps).sum(dim=1).clamp(min=1)  # [T]; avoid //0
+            row_counts = row_slot * max_beam_width + local // steps
+
+        count = counts_cuda[row_counts]
         rep = repetition_cuda[row_slot].unsqueeze(1)
         pre = presence_cuda[row_slot].unsqueeze(1)
         freq = frequency_cuda[row_slot].unsqueeze(1)
@@ -466,14 +631,17 @@ def apply_batched_occurrence_penalties(
         repetition_cuda: torch.Tensor,
         presence_cuda: torch.Tensor,
         frequency_cuda: torch.Tensor,
+        request_num_beams: Optional[torch.Tensor] = None,
+        max_beam_width: int = 1,
+        fold_pending: bool = True,
     ) -> None:
         """Apply occurrence penalties to ``logits`` in place, before temperature handling.
 
         Args:
             logits: ``[T, vocab_size]`` packed generated-token logits, where
                 ``T == sum(num_steps * num_beams)``. Request ``r`` owns the rows
-                ``request_offsets[r] + step`` for ``step in [0, request_num_steps[r])``;
-                rows no request owns are left bit-identical. Modified in place.
+                ``request_offsets[r] + beam * num_steps[r] + step``, i.e. beam-major /
+                step-minor; rows no request owns are left bit-identical. Modified in place.
             counts_cuda / presence_prefix_cuda: the occurrence workspace; see
                 ``PenaltyHandler.PenaltyStore`` for their semantics.
             active_cuda / has_previous_token_cuda / repetition_cuda / presence_cuda /
@@ -484,6 +652,18 @@ def apply_batched_occurrence_penalties(
             request_offsets / request_num_steps: ``[R]`` device tensors, already
                 staged by the caller. The owned spans must not overlap, but they need
                 not be ordered, and rows they skip are left bit-identical.
+            request_num_beams: ``[R]`` row-layout beam width per request -- the static
+                admission width ModelEngine lays the rows out at, which under a growing
+                ``beam_width_array`` exceeds the per-iteration width -- so a row can be
+                mapped back to the beam that owns it. None on a single-beam engine.
+            max_beam_width: the engine's beam width, the stride of the counts rows.
+            fold_pending: whether to fold the device-pending token here. False on a beam
+                engine, where ``update_beam_occurrence_counts`` has already re-parented
+                and folded every slot.
+
+        The last three are compile-time constants to Dynamo (an Optional and two Python
+        scalars), so each engine specializes to its own graph and neither carries the
+        other's branches.
 
         All heavy lifting is fused into the single compiled ``_apply_occurrence_penalties_impl``
         graph; this wrapper only marks the batch-varying dims dynamic.
@@ -498,6 +678,8 @@ def apply_batched_occurrence_penalties(
         torch._dynamo.mark_dynamic(seq_slots, 0)
         torch._dynamo.mark_dynamic(request_offsets, 0)
         torch._dynamo.mark_dynamic(request_num_steps, 0)
+        if request_num_beams is not None:
+            torch._dynamo.mark_dynamic(request_num_beams, 0)
         Fusions._apply_occurrence_penalties_impl(
             logits,
             counts_cuda,
@@ -511,4 +693,7 @@ def apply_batched_occurrence_penalties(
             repetition_cuda,
             presence_cuda,
             frequency_cuda,
+            request_num_beams,
+            max_beam_width,
+            fold_pending,
         )