diff --git a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp index 75073f276c92..429110a40496 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 b7d7f716acb6..e472894a22ea 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); } @@ -1113,7 +1114,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 @@ -1193,6 +1194,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 2b5b164c5574..d9e92b42f5cd 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.h +++ b/cpp/tensorrt_llm/thop/attentionOp.h @@ -95,7 +95,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 c78cc155b071..4cfcda50bac2 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 e68d77a19fe7..02b52af0e585 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_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 diff --git a/tests/unittest/_torch/attention/test_fmha_page_index.py b/tests/unittest/_torch/attention/test_fmha_page_index.py index 50bb4b2c1349..9ffabeb9ede1 100644 --- a/tests/unittest/_torch/attention/test_fmha_page_index.py +++ b/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(), + )