[https://nvbugs/6565412][fix] Size trtllm-gen and thop decode buffers for beam search

cpp/tensorrt_llm/nanobind/thop/bindings.cpp

@@ -180,8 +180,8 @@ void initBindings(nb::module_& m)
         nb::arg("relative_attention_bias") = std::nullopt, nb::arg("relative_attention_max_distance") = 0,
         nb::arg("spec_decoding_target_max_draft_tokens") = std::nullopt, nb::arg("quant_scale_qkv") = std::nullopt,
         nb::arg("dsv4_inv_rope_cos_sin_cache") = std::nullopt, nb::arg("enable_dsv4_epilogue_fusion") = false,
-        nb::arg("force_prepare_spec_dec_tree_mask") = false, "Multi-head attention operation",
-        nb::call_guard<nb::gil_scoped_release>());
+        nb::arg("force_prepare_spec_dec_tree_mask") = false, nb::arg("max_num_sequences") = std::nullopt,
+        "Multi-head attention operation", nb::call_guard<nb::gil_scoped_release>());
 
     m.def(
         "get_helix_workspace_size_per_rank",

cpp/tensorrt_llm/thop/attentionOp.cpp

@@ -335,11 +335,12 @@ class RunnerBase
 public:
     int32_t beam_width;
     int32_t max_num_requests;
+    int32_t max_num_sequences;
     int32_t attention_window_size;
 
     auto data() const
     {
-        return std::make_tuple(beam_width, max_num_requests, attention_window_size);
+        return std::make_tuple(beam_width, max_num_requests, max_num_sequences, attention_window_size);
     };
 
     virtual ~RunnerBase() = default;
@@ -406,7 +407,7 @@ class Runner : public RunnerBase
         // is not enough.
         // The attention kernel might split the heads into multiple blocks, so we might need to reserve more semaphores.
         // Use mMultiProcessorCount as the lower-bound to make sure we reserve enough semaphores.
-        op.reserveSemaphoreArray(std::max(op.mNumHeads * max_num_requests, op.getMultiProcessorCount()));
+        op.reserveSemaphoreArray(std::max(op.mNumHeads * max_num_sequences, op.getMultiProcessorCount()));
     }
 
     int64_t getWorkspaceSize(AttentionOp const& op, int const num_tokens, int const max_attention_window_size,
@@ -415,7 +416,7 @@ class Runner : public RunnerBase
         size_t const context_workspace_size = op.getWorkspaceSizeForContext(
             op.mType, max_num_requests, op.mMaxContextLength, 0, num_tokens, ctx_total_kv_len);
         size_t const generation_workspace_size = op.getWorkspaceSizeForGeneration(
-            op.mType, max_num_requests, max_attention_window_size, num_gen_tokens, max_blocks_per_sequence);
+            op.mType, max_num_sequences, max_attention_window_size, num_gen_tokens, max_blocks_per_sequence);
 
         return std::max(context_workspace_size, generation_workspace_size);
     }
@@ -1113,7 +1114,7 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
     std::optional<torch::Tensor> relative_attention_bias, int64_t relative_attention_max_distance,
     std::optional<int64_t> spec_decoding_target_max_draft_tokens, std::optional<torch::Tensor> quant_scale_qkv,
     std::optional<torch::Tensor> dsv4_inv_rope_cos_sin_cache, bool enable_dsv4_epilogue_fusion,
-    bool const force_prepare_spec_dec_tree_mask)
+    bool const force_prepare_spec_dec_tree_mask, std::optional<int64_t> const max_num_sequences)
 {
     TLLM_LOG_TRACE("Attention op starts at layer %d", local_layer_idx);
     // Use these tensors to infer if the attention is using KV cache
@@ -1193,6 +1194,7 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
 #endif
     runner->beam_width = beam_width;
     runner->max_num_requests = max_num_requests;
+    runner->max_num_sequences = max_num_sequences.value_or(max_num_requests);
     runner->attention_window_size = attention_window_size;
 
     auto op = std::make_shared<AttentionOp>();

cpp/tensorrt_llm/thop/attentionOp.h

@@ -95,7 +95,7 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
     std::optional<int64_t> spec_decoding_target_max_draft_tokens = std::nullopt,
     std::optional<torch::Tensor> quant_scale_qkv = std::nullopt,
     std::optional<torch::Tensor> dsv4_inv_rope_cos_sin_cache = std::nullopt, bool enable_dsv4_epilogue_fusion = false,
-    bool const force_prepare_spec_dec_tree_mask = false);
+    bool const force_prepare_spec_dec_tree_mask = false, std::optional<int64_t> const max_num_sequences = std::nullopt);
 
 struct KvCachePoolPointers
 {

tensorrt_llm/_torch/attention_backend/fmha/fallback.py

@@ -85,6 +85,7 @@ def forward(
             block_ids_per_seq=metadata.block_ids_per_seq,
             tokens_per_block=metadata.tokens_per_block,
             max_num_requests=metadata.max_num_requests,
+            max_num_sequences=metadata.max_num_sequences,
             beam_width=metadata.effective_beam_width,
             use_paged_context_fmha=metadata.use_paged_context_fmha,
             helix_position_offsets=metadata.helix_position_offsets,

tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py

@@ -212,10 +212,10 @@ def _cached_build(
 
 def _get_multi_ctas_kv_counter_size(
     num_heads: int,
-    max_num_requests: int,
+    max_num_sequences: int,
     multi_processor_count: int,
 ) -> int:
-    num_counters = max(num_heads * max_num_requests, multi_processor_count)
+    num_counters = max(num_heads * max_num_sequences, multi_processor_count)
     aligned_num_counters = (
         (num_counters + _MULTI_CTAS_KV_COUNTER_ALIGNMENT - 1)
         // _MULTI_CTAS_KV_COUNTER_ALIGNMENT
@@ -943,9 +943,11 @@ def prepare_workspace(
         if self._multi_processor_count is None:
             self._multi_processor_count = self._get_multi_processor_count(q.device)
 
+        # One counter per head per decoder sequence; beam search expands each
+        # request into ``beam_width`` sequences.
         required_counter_size = _get_multi_ctas_kv_counter_size(
             attn.num_heads,
-            metadata.max_num_requests,
+            metadata.max_num_sequences or metadata.max_num_requests,
             self._multi_processor_count,
         )
         counter_buffer = self._multi_ctas_kv_counter_buffer

tests/unittest/_torch/attention/test_fmha_page_index.py

@@ -1,6 +1,15 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 from types import SimpleNamespace
 
-from tensorrt_llm._torch.attention_backend.fmha.flashinfer_trtllm_gen import FlashInferTrtllmGenFmha
+import pytest
+import torch
+
+from tensorrt_llm._torch.attention_backend.fmha.flashinfer_trtllm_gen import (
+    FlashInferTrtllmGenFmha,
+    _get_multi_ctas_kv_counter_size,
+)
 
 
 def _get_total_num_blocks(manager: SimpleNamespace, kv_factor: int = 2) -> int:
@@ -25,3 +34,65 @@ def test_flashinfer_preserves_legacy_pool_scaling() -> None:
         num_local_layers=36,
     )
     assert _get_total_num_blocks(manager, kv_factor=2) == 1024 * 36 * 2
+
+
+def test_multi_ctas_kv_counter_size_covers_beam_expanded_batch() -> None:
+    # The kernel keeps one counter per head per decoder sequence. Sizing off the
+    # request count alone under-allocates under beam search, but only once the
+    # product clears the multi-processor floor, so pick a case that does.
+    num_heads, batch, beam, sm_count = 6, 16, 2, 148
+    needed = num_heads * batch * beam * torch.int32.itemsize
+    assert _get_multi_ctas_kv_counter_size(num_heads, batch, sm_count) < needed
+    assert _get_multi_ctas_kv_counter_size(num_heads, batch * beam, sm_count) >= needed
+
+
+def test_multi_ctas_kv_counter_size_keeps_multi_processor_floor() -> None:
+    num_heads, batch, sm_count = 6, 1, 148
+    assert _get_multi_ctas_kv_counter_size(num_heads, batch, sm_count) >= (
+        sm_count * torch.int32.itemsize
+    )
+
+
+def test_prepare_workspace_sizes_counter_for_max_num_sequences(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    num_heads, max_num_requests, beam_width, sm_count = 6, 16, 2, 148
+    max_num_sequences = max_num_requests * beam_width
+
+    def check_counter_size_args(
+        actual_num_heads: int,
+        actual_max_num_sequences: int,
+        actual_sm_count: int,
+    ) -> int:
+        assert (actual_num_heads, actual_max_num_sequences, actual_sm_count) == (
+            num_heads,
+            max_num_sequences,
+            sm_count,
+        )
+        raise RuntimeError("counter size arguments observed")
+
+    monkeypatch.setattr(
+        "tensorrt_llm._torch.attention_backend.fmha.flashinfer_trtllm_gen."
+        "_get_multi_ctas_kv_counter_size",
+        check_counter_size_args,
+    )
+
+    fmha = SimpleNamespace(
+        attn=SimpleNamespace(num_heads=num_heads),
+        _multi_processor_count=sm_count,
+    )
+    metadata = SimpleNamespace(
+        max_num_requests=max_num_requests,
+        beam_width=beam_width,
+        max_num_sequences=max_num_sequences,
+    )
+    with pytest.raises(RuntimeError, match="counter size arguments observed"):
+        FlashInferTrtllmGenFmha.prepare_workspace(
+            fmha,
+            q=SimpleNamespace(),
+            k=None,
+            v=None,
+            metadata=metadata,
+            forward_args=SimpleNamespace(),
+            workspace=SimpleNamespace(),
+        )