From 7dbb471fe02f05819079060a8ad7e9d06de99aad Mon Sep 17 00:00:00 2001 From: yihwang-nv Date: Thu, 16 Jul 2026 23:47:56 -0700 Subject: [PATCH 1/3] [None][chore] Update flashinfer-python from 0.6.14 to 0.6.15 Bump flashinfer-python dependency to the latest stable release. Updated version pins in requirements.txt, security_scanning/pyproject.toml, and ATTRIBUTIONS-Python.md. Signed-off-by: yihwang-nv --- ATTRIBUTIONS-Python.md | 2 +- requirements.txt | 2 +- security_scanning/pyproject.toml | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/ATTRIBUTIONS-Python.md b/ATTRIBUTIONS-Python.md index e43a78cafda0..7e74b846a85f 100644 --- a/ATTRIBUTIONS-Python.md +++ b/ATTRIBUTIONS-Python.md @@ -5261,7 +5261,7 @@ For more information, please refer to - `Tracker`: https://github.com/tox-dev/py-filelock/issues -## flashinfer-python (0.6.14) +## flashinfer-python (0.6.15) ### Licenses License: `Apache-2.0` diff --git a/requirements.txt b/requirements.txt index dab577614f57..a7692309fb4a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -56,7 +56,7 @@ ordered-set peft>=0.18.1,<0.19.0 patchelf einops -flashinfer-python==0.6.14 +flashinfer-python==0.6.15 xgrammar==0.1.32 llguidance==0.7.29 jsonschema diff --git a/security_scanning/pyproject.toml b/security_scanning/pyproject.toml index ec9b4f532e06..a1729b6c267a 100644 --- a/security_scanning/pyproject.toml +++ b/security_scanning/pyproject.toml @@ -56,7 +56,7 @@ dependencies = [ "peft (>=0.18.1,<0.19.0)", "patchelf (>=0.17.2.4,<0.18.0.0)", "einops (>=0.8.2,<0.9.0)", - "flashinfer-python (==0.6.14)", + "flashinfer-python (==0.6.15)", "xgrammar (==0.1.32)", "llguidance (==0.7.29)", "jsonschema (>=4.26.0,<5.0.0)", From cbcde2c0d835fb6ace909192f10c7423b1dad99d Mon Sep 17 00:00:00 2001 From: yihwang-nv Date: Fri, 17 Jul 2026 00:01:17 -0700 Subject: [PATCH 2/3] [None][chore] Remove multi-CTAs counter workaround fixed in flashinfer 0.6.15 flashinfer 0.6.15 includes the fix for the missing multi-CTAs counter buffer initialization (flashinfer-ai/flashinfer#3433). Remove the _clear_multi_ctas_kv_counter_workspace workaround calls from run_generation() and run_mla_generation(), and drop the now-unused helper functions _clear_multi_ctas_kv_counter_workspace and _get_multi_ctas_kv_counter_size. Signed-off-by: yihwang-nv --- .../fmha/flashinfer_trtllm_gen.py | 41 ------------------- 1 file changed, 41 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 21e67006fc95..6d6196451348 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py @@ -67,28 +67,6 @@ ) -def _clear_multi_ctas_kv_counter_workspace( - fmha_workspace: torch.Tensor, - num_heads: int, - max_num_requests: int, - multi_processor_count: Optional[int], -) -> None: - counter_size = _get_multi_ctas_kv_counter_size( - num_heads, - max_num_requests, - multi_processor_count, - ) - fmha_workspace.flatten().narrow(0, 0, counter_size).zero_() - - -def _get_multi_ctas_kv_counter_size( - num_heads: int, - max_num_requests: int, - multi_processor_count: Optional[int], -) -> int: - return max(num_heads * max_num_requests, multi_processor_count or 0) * torch.int32.itemsize - - def _get_bmm1_scale_log2(bmm1_scale: torch.Tensor) -> torch.Tensor: if bmm1_scale.numel() < 2: raise RuntimeError("trtllm-gen bmm1_scale workspace must contain raw and log2 scales.") @@ -1092,22 +1070,6 @@ def run_generation( params.is_cross, # is_cross ) - # FIXME: Flashinfer trtllm-gen API doesn't support a separate - # multi CTAs counter buffer. We have to clear a small buffer - # before trtllm_gen_batch_decode_with_kv_cache. - # - # We must also avoid clearing the workspace only when it is - # resized. The warmup phase may have already cached the workspace - # pointer; if the capture phase skips the zeroing step, the - # CUDA graph will not include the counter initialization. We - # have already verified—specifically in the context of the GPTOSS-20B - # test graph replay scenario—that this skipping logic is unsafe. - # - # https://github.com/flashinfer-ai/flashinfer/issues/3433 - _clear_multi_ctas_kv_counter_workspace( - fmha_workspace, attn.num_heads, meta.max_num_requests, self._multi_processor_count - ) - q_len_per_req = None if is_multi_token_gen else params.input_seq_length decode_max_q_len = max_q_len if is_multi_token_gen else None decode_cu_seqlens = cu_seqlens if is_multi_token_gen else None @@ -1233,9 +1195,6 @@ def run_mla_generation( bmm1_scale = 1.0 / (attn.q_scaling * math.sqrt(qk_nope_head_dim + qk_rope_head_dim)) bmm2_scale = 1.0 workspace_buffer = params.workspace.view(-1, 4) - _clear_multi_ctas_kv_counter_workspace( - workspace_buffer, attn.num_heads, meta.max_num_requests, self._multi_processor_count - ) flashinfer.mla.trtllm_batch_decode_with_kv_cache_mla( query, # query From 42f72b0b03419a3517ba373fef77e5d13c7a40c7 Mon Sep 17 00:00:00 2001 From: yihwang-nv Date: Sun, 19 Jul 2026 22:50:30 -0700 Subject: [PATCH 3/3] [None][fix] Pass multi-CTA counter buffer to FlashInfer Signed-off-by: yihwang-nv --- .../fmha/flashinfer_trtllm_gen.py | 53 +++++++++++++++++++ 1 file changed, 53 insertions(+) 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 6d6196451348..ff2002c36d57 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py @@ -67,6 +67,23 @@ ) +_MULTI_CTAS_KV_COUNTER_ALIGNMENT = 8 + + +def _get_multi_ctas_kv_counter_size( + num_heads: int, + max_num_requests: int, + multi_processor_count: int, +) -> int: + num_counters = max(num_heads * max_num_requests, multi_processor_count) + aligned_num_counters = ( + (num_counters + _MULTI_CTAS_KV_COUNTER_ALIGNMENT - 1) + // _MULTI_CTAS_KV_COUNTER_ALIGNMENT + * _MULTI_CTAS_KV_COUNTER_ALIGNMENT + ) + return aligned_num_counters * torch.int32.itemsize + + def _get_bmm1_scale_log2(bmm1_scale: torch.Tensor) -> torch.Tensor: if bmm1_scale.numel() < 2: raise RuntimeError("trtllm-gen bmm1_scale workspace must contain raw and log2 scales.") @@ -77,6 +94,7 @@ def _trtllm_gen_batch_decode_with_kv_cache( query: torch.Tensor, kv_pool: torch.Tensor, workspace_buffer: torch.Tensor, + multi_ctas_kv_counter_buffer: torch.Tensor, block_tables: torch.Tensor, seq_lens: torch.Tensor, max_seq_len: int, @@ -116,6 +134,7 @@ def _trtllm_gen_batch_decode_with_kv_cache( kv_pool, kv_pool, workspace_buffer, + multi_ctas_kv_counter_buffer, block_tables, seq_lens, decode_max_q_len, @@ -147,6 +166,7 @@ def _trtllm_gen_batch_context_with_kv_cache( query: torch.Tensor, kv_pool: torch.Tensor, workspace_buffer: torch.Tensor, + multi_ctas_kv_counter_buffer: torch.Tensor, block_tables: torch.Tensor, seq_lens: torch.Tensor, max_q_len: int, @@ -177,6 +197,7 @@ def _trtllm_gen_batch_context_with_kv_cache( kv_pool, kv_pool, workspace_buffer, + multi_ctas_kv_counter_buffer, block_tables, seq_lens, max_q_len, @@ -398,6 +419,7 @@ def __init__(self, attn: "TrtllmAttention"): # Lazily set on the first forward() call from the query device. self._multi_processor_count: Optional[int] = None + self._multi_ctas_kv_counter_buffer: Optional[torch.Tensor] = None def _get_total_num_blocks(self, meta: "TrtllmAttentionMetadata") -> int: kv_cache_manager = meta.kv_cache_manager @@ -780,6 +802,28 @@ def prepare_workspace( if self._multi_processor_count is None: self._multi_processor_count = self._get_multi_processor_count(q.device) + required_counter_size = _get_multi_ctas_kv_counter_size( + attn.num_heads, + metadata.max_num_requests, + self._multi_processor_count, + ) + counter_buffer = self._multi_ctas_kv_counter_buffer + if ( + counter_buffer is None + or counter_buffer.device != q.device + or counter_buffer.numel() * counter_buffer.element_size() < required_counter_size + ): + if metadata.is_cuda_graph and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "The trtllm-gen multi-CTA KV counter buffer must be allocated " + "before CUDA graph capture." + ) + self._multi_ctas_kv_counter_buffer = torch.zeros( + required_counter_size, + dtype=torch.uint8, + device=q.device, + ) + num_tokens = q.size(0) attention_input_type = forward_args.attention_input_type is_gen_only = attention_input_type == AttentionInputType.generation_only @@ -813,6 +857,12 @@ def prepare_workspace( required_workspace_numel = math.ceil(required_workspace_size / workspace.element_size()) workspace.resize_((required_workspace_numel,)) + def _get_multi_ctas_kv_counter_buffer(self) -> torch.Tensor: + counter_buffer = self._multi_ctas_kv_counter_buffer + if counter_buffer is None: + raise RuntimeError("The trtllm-gen multi-CTA KV counter buffer is not initialized.") + return counter_buffer + @staticmethod def _compute_window_left( cyclic_attention_window_size: int, @@ -939,6 +989,7 @@ def run_context( q_processed, # query kv_pool, # kv_pool fmha_workspace, # workspace_buffer + self._get_multi_ctas_kv_counter_buffer(), # multi_ctas_kv_counter_buffer block_tables, # block_tables params.sequence_lengths, # seq_lens max_q_len, # max_q_len @@ -1093,6 +1144,7 @@ def run_generation( q_processed, # query kv_pool, # kv_pool fmha_workspace, # workspace_buffer + self._get_multi_ctas_kv_counter_buffer(), # multi_ctas_kv_counter_buffer block_tables, # block_tables params.sequence_lengths, # seq_lens max_kv_len, # max_seq_len @@ -1216,4 +1268,5 @@ def run_mla_generation( "trtllm-gen", # backend True, # is_var_seq self.USE_SHARED_PAGED_KV_IDX, # uses_shared_paged_kv_idx + multi_ctas_kv_counter_buffer=self._get_multi_ctas_kv_counter_buffer(), )