[None][fix] Size the trtllm-gen multi-CTA KV counter buffer for beam search

tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py

@@ -72,10 +72,10 @@
 
 def _get_multi_ctas_kv_counter_size(
     num_heads: int,
-    max_num_requests: int,
+    max_batch_beam: int,
     multi_processor_count: int,
 ) -> int:
-    num_counters = max(num_heads * max_num_requests, multi_processor_count)
+    num_counters = max(num_heads * max_batch_beam, multi_processor_count)
     aligned_num_counters = (
         (num_counters + _MULTI_CTAS_KV_COUNTER_ALIGNMENT - 1)
         // _MULTI_CTAS_KV_COUNTER_ALIGNMENT
@@ -802,9 +802,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_requests * metadata.beam_width,
             self._multi_processor_count,
         )
         counter_buffer = self._multi_ctas_kv_counter_buffer

tests/unittest/_torch/attention/test_fmha_page_index.py

@@ -1,6 +1,11 @@
 from types import SimpleNamespace
 
-from tensorrt_llm._torch.attention_backend.fmha.flashinfer_trtllm_gen import FlashInferTrtllmGenFmha
+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 +30,20 @@ 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
+    )