From bd0efcaabdf65b8657572fa4dd9b2a6e3b1d9d08 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Wed, 29 Jul 2026 00:56:27 -0700 Subject: [PATCH 1/2] [None][fix] Size the trtllm-gen multi-CTA KV counter buffer for beam search The trtllm-gen decode kernel keeps one multi-CTA KV counter per attention head per decoder sequence, but the buffer was sized from the request count alone. Beam search expands each request into beam_width sequences, so the kernel rejects the buffer at warmup with TVM_FFI_CHECK(multi_ctas_kv_counter_size >= counter_bytes) The max(..., multi_processor_count) floor hides the missing factor at small batch, so this only reproduces past a threshold: with 6 decoder heads and 148 SMs, beam 2 passes up to max_batch_size 8 and fails from 16 up. Every in-tree beam-search test runs below that threshold, so CI does not cover it. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- .../fmha/flashinfer_trtllm_gen.py | 8 ++++--- .../_torch/attention/test_fmha_page_index.py | 24 ++++++++++++++++++- 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py index ff2002c36d57..0d7e1d672ea3 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py +++ b/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 diff --git a/tests/unittest/_torch/attention/test_fmha_page_index.py b/tests/unittest/_torch/attention/test_fmha_page_index.py index 50bb4b2c1349..45ae4e579dc5 100644 --- a/tests/unittest/_torch/attention/test_fmha_page_index.py +++ b/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 + ) From b9b2b13df560138500a22f8eab59c71976bd525b Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Thu, 6 Aug 2026 13:26:03 -0700 Subject: [PATCH 2/2] [https://nvbugs/6565412][fix] Source beam sizing from max_num_sequences, including thop Take the bound from metadata.max_num_sequences instead of recomputing max_num_requests * beam_width: beam_width is rewritten per step and is 1 during context, so it cannot size a buffer that must outlive CUDA graph capture. The thop attention path had the same gap and now takes the same bound for its semaphore and generation workspace sizing. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- cpp/tensorrt_llm/nanobind/thop/bindings.cpp | 4 +- cpp/tensorrt_llm/thop/attentionOp.cpp | 10 ++-- cpp/tensorrt_llm/thop/attentionOp.h | 2 +- .../_torch/attention_backend/fmha/fallback.py | 1 + .../fmha/flashinfer_trtllm_gen.py | 6 +-- .../_torch/attention/test_fmha_page_index.py | 46 +++++++++++++++++++ 6 files changed, 59 insertions(+), 10 deletions(-) diff --git a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp index 2ac9cc671277..48fe0fb8f87a 100644 --- a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp +++ b/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::arg("force_prepare_spec_dec_tree_mask") = false, nb::arg("max_num_sequences") = std::nullopt, + "Multi-head attention operation", nb::call_guard()); m.def( "get_helix_workspace_size_per_rank", diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 24d869f704dc..0a5d75aeb987 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/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); } @@ -1114,7 +1115,7 @@ void attention(torch::Tensor q, std::optional k, std::optional relative_attention_bias, int64_t relative_attention_max_distance, std::optional spec_decoding_target_max_draft_tokens, std::optional quant_scale_qkv, std::optional 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 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 @@ -1194,6 +1195,7 @@ void attention(torch::Tensor q, std::optional k, std::optionalbeam_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(); diff --git a/cpp/tensorrt_llm/thop/attentionOp.h b/cpp/tensorrt_llm/thop/attentionOp.h index f28209166e0d..a44e51953a18 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.h +++ b/cpp/tensorrt_llm/thop/attentionOp.h @@ -96,7 +96,7 @@ void attention(torch::Tensor q, std::optional k, std::optional spec_decoding_target_max_draft_tokens = std::nullopt, std::optional quant_scale_qkv = std::nullopt, std::optional 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 const max_num_sequences = std::nullopt); struct KvCachePoolPointers { diff --git a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py index 9c081bf19dfc..6898aa8ce84d 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py +++ b/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, diff --git a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py index d1f72e5fcc83..62ed5b1235e3 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py +++ b/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_batch_beam: int, + max_num_sequences: int, multi_processor_count: int, ) -> int: - num_counters = max(num_heads * max_batch_beam, 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 @@ -946,7 +946,7 @@ def prepare_workspace( # request into ``beam_width`` sequences. required_counter_size = _get_multi_ctas_kv_counter_size( attn.num_heads, - metadata.max_num_requests * metadata.beam_width, + metadata.max_num_sequences or metadata.max_num_requests, self._multi_processor_count, ) counter_buffer = self._multi_ctas_kv_counter_buffer diff --git a/tests/unittest/_torch/attention/test_fmha_page_index.py b/tests/unittest/_torch/attention/test_fmha_page_index.py index c09e5bdf80dd..9ffabeb9ede1 100644 --- a/tests/unittest/_torch/attention/test_fmha_page_index.py +++ b/tests/unittest/_torch/attention/test_fmha_page_index.py @@ -3,6 +3,7 @@ from types import SimpleNamespace +import pytest import torch from tensorrt_llm._torch.attention_backend.fmha.flashinfer_trtllm_gen import ( @@ -50,3 +51,48 @@ def test_multi_ctas_kv_counter_size_keeps_multi_processor_floor() -> None: 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(), + )