From f293f92bde75b0454e232ab8d794301d1397583c Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Wed, 1 Apr 2026 20:10:44 -0700 Subject: [PATCH 01/11] [None][feat] Add FlashInfer MLA attention backend support Implement FlashInfer-based Multi-head Latent Attention (MLA) for DeepSeek-style models, enabling use of BatchPrefillWithRaggedKVCacheWrapper for context and BatchMLAPagedAttentionWrapper for generation. Key design decisions: - KV cache: zero-copy ckv/kpe views split from the existing paged KV pool buffer (kv_lora_rank + qk_rope_head_dim), avoiding extra allocation - Context phase: ragged prefill with expanded Q/K/V after appending latent to paged MLA caches via append_paged_mla_kv_cache - Generation phase: paged MLA decode with q_nope/q_pe split from fused_q; latent cache append handled before each decode step - RoPE applied externally in MLA.forward before latent_cache construction; FlashInfer kernels receive pre-RoPE'd inputs - Plans called unconditionally each forward pass (no stale caching) - backend=auto on BatchMLAPagedAttentionWrapper for FA3 on Hopper - mla_rope_generation stub copies RoPE'd q_pe into fused_q to satisfy the forward_absorption_generation call chain Also fix metadata type assertion in attention.py to allow FlashInferAttentionMetadata for MLA, and replace hasattr duck-typing with isinstance check for kv_lens_cuda_runtime. Validated with 84/84 unit tests and end-to-end DeepSeek-V3-Lite NVFP4 inference with --attention_backend FLASHINFER. Support CUDA graph for flashinfer MLA. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> Drop redundant contiguous/reshape copies in flashinfer MLA path. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> [None][test] Fix FlashInfer MLA unit test and add B200/B300 smoke coverage - test_attention_mla_flashinfer: restrict to SM100 and set max_num_tokens to sum(context_sequence_lengths) so FlashInfer's batch_indices/positions buffers are sized for the full batch (was max(...), causing OOB copy). - Add TestDeepSeekV3Lite::test_bfloat16_flashinfer accuracy smoke test and register it in l0_b200.yml and l0_b300.yml. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- .../_torch/attention_backend/flashinfer.py | 419 +++++++++++++++++- tensorrt_llm/_torch/modules/attention.py | 50 ++- .../defs/accuracy/test_llm_api_pytorch.py | 10 + .../test_lists/test-db/l0_b200.yml | 1 + .../test_lists/test-db/l0_b300.yml | 1 + .../_torch/attention/test_attention_mla.py | 288 +++++++++--- 6 files changed, 692 insertions(+), 77 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 2f179d495b8d..aa27caddf13a 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -16,7 +16,7 @@ from ..utils import get_global_attrs, get_model_extra_attrs from .interface import (AttentionBackend, AttentionForwardArgs, AttentionMetadata, CustomAttentionMask, - PredefinedAttentionMask, merge_attention_forward_args) + PredefinedAttentionMask, merge_attention_forward_args, MLAParams) try: check_cuda_arch() @@ -48,6 +48,36 @@ class PlanParams: window_left: Optional[int] = None +@dataclass(kw_only=True, frozen=True) +class RaggedPlanParams: + """ + Parameters for MLA ragged prefill (context phase with expanded K, V). + """ + + num_heads: int + num_kv_heads: int + head_dim: int + head_dim_vo: int + q_dtype: torch.dtype + kv_dtype: torch.dtype + sm_scale: Optional[float] = None + + +@dataclass(kw_only=True, frozen=True) +class MLADecodePlanParams: + """ + Parameters for FlashInfer MLA decode using BatchMLAPagedAttentionWrapper. + """ + + num_heads: int + kv_lora_rank: int + qk_rope_head_dim: int + page_size: int + q_dtype: torch.dtype + kv_dtype: torch.dtype + sm_scale: Optional[float] = None + + @dataclass(kw_only=True) class FlashInferWrappers: is_planned: bool @@ -79,6 +109,22 @@ class FlashInferAttentionMetadata(AttentionMetadata): _plan_params_to_wrappers: Dict[PlanParams, FlashInferWrappers] = field(init=False) + # MLA ragged prefill wrapper (for context phase with expanded K, V) + _ragged_prefill_wrapper: Optional[ + flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper] = field( + init=False, default=None) + + # MLA decode wrapper (BatchMLAPagedAttentionWrapper) and stable buffers. + # Cached plan params + is-planned flag let prepare() refresh the plan + # outside stream capture (flashinfer plan() does device->host syncs). + _mla_decode_wrapper: Optional[object] = field(init=False, default=None) + _mla_decode_plan_params: Optional[MLADecodePlanParams] = field(init=False, + default=None) + _mla_decode_planned: bool = field(init=False, default=False) + _mla_qo_indptr_buf: Optional[torch.Tensor] = field(init=False, default=None) + _mla_kv_len_arr_buf: Optional[torch.Tensor] = field(init=False, + default=None) + def needs_plan(self, plan_params: PlanParams) -> bool: if plan_params not in self._plan_params_to_wrappers: return True @@ -110,6 +156,113 @@ def get_ragged_prefill_wrapper( assert result is not None, "Ragged prefill wrapper was not created in plan()" return result + def plan_ragged( + self, + qo_indptr: torch.Tensor, + kv_indptr: torch.Tensor, + plan_params: RaggedPlanParams, + ) -> flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper: + """Plan MLA ragged prefill with expanded K, V (not paged).""" + if self._ragged_prefill_wrapper is None: + self._ragged_prefill_wrapper = flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper( + self.workspace_buffer, + "NHD", + ) + + self._ragged_prefill_wrapper.plan( + qo_indptr, + kv_indptr, + plan_params.num_heads, + plan_params.num_kv_heads, + plan_params.head_dim, + head_dim_vo=plan_params.head_dim_vo, + use_fp16_qk_reduction=False, + causal=True, + q_data_type=plan_params.q_dtype, + kv_data_type=plan_params.kv_dtype, + sm_scale=plan_params.sm_scale, + ) + + return self._ragged_prefill_wrapper + + def plan_mla_decode( + self, + plan_params: MLADecodePlanParams, + ) -> object: + """Plan MLA decode using BatchMLAPagedAttentionWrapper. + + Caches the wrapper and plan params; the actual plan() call is driven + by prepare() so it runs outside of CUDA graph capture. + """ + if self._mla_decode_wrapper is None: + self._mla_decode_wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper( + self.workspace_buffer, + use_cuda_graph=self.is_cuda_graph, + qo_indptr=self._mla_qo_indptr_buf, + kv_indptr=self.paged_kv_indptr_decode, + kv_indices=self._paged_kv_indices, + kv_len_arr=self._mla_kv_len_arr_buf, + backend="auto", + ) + + # Cache params so prepare() can re-plan on subsequent forward passes. + self._mla_decode_plan_params = plan_params + + if self._mla_decode_planned: + return self._mla_decode_wrapper + + if self.is_cuda_graph and torch.cuda.is_current_stream_capturing(): + raise ValueError( + "Cannot plan() flashinfer MLA decode while the stream is " + "capturing. Make sure prepare() has run at least one warmup " + "forward pass before capture.") + + self._do_plan_mla_decode(plan_params) + self._mla_decode_planned = True + return self._mla_decode_wrapper + + def _do_plan_mla_decode(self, plan_params: MLADecodePlanParams) -> None: + """Compute MLA decode plan inputs and call wrapper.plan(). + + Must run outside of CUDA graph capture. kv_indptr / kv_indices are + cloned because they alias the wrapper's own buffers and + flashinfer.plan() would otherwise do a self-copy. + """ + num_gen = self.num_generations + kv_indptr = self.paged_kv_indptr_decode[:num_gen + 1] + kv_indices = self._paged_kv_indices[self.num_context_blocks:self. + num_context_blocks + + self.num_generation_blocks].clone() + kv_last_page = self._paged_kv_last_page_len[self.num_contexts:self. + num_contexts + num_gen] + + # _qo_indptr is ordered [context_seqs..., generation_seqs...]; rebase + # the generation slice to 0. + num_ctx = self.num_contexts + qo_indptr = self._qo_indptr[num_ctx:num_ctx + num_gen + + 1] - self._qo_indptr[num_ctx] + + num_pages_per_seq = kv_indptr[1:] - kv_indptr[:-1] + kv_len_arr = (num_pages_per_seq - + 1) * plan_params.page_size + kv_last_page + + kv_indptr = kv_indptr.clone() + + self._mla_decode_wrapper.plan( + qo_indptr, + kv_indptr, + kv_indices, + kv_len_arr, + plan_params.num_heads, + plan_params.kv_lora_rank, + plan_params.qk_rope_head_dim, + plan_params.page_size, + causal=True, + q_data_type=plan_params.q_dtype, + kv_data_type=plan_params.kv_dtype, + sm_scale=plan_params.sm_scale, + ) + @property def paged_kv_indices(self) -> torch.Tensor: return self._paged_kv_indices[:self.num_generation_blocks + @@ -285,6 +438,24 @@ def _post_init_with_buffers(self, buffers) -> None: torch.empty(all_pool_pages, dtype=torch.int, device='cuda')) + # Stable buffers for FlashInfer MLA decode; required for CUDA graphs. + self._mla_qo_indptr_buf = self.get_empty( + buffers, + (self.max_num_requests + 1, ), + dtype=torch.int32, + cache_name="_mla_qo_indptr_buf", + capture_graph=capture_graph, + ) + self._mla_kv_len_arr_buf = self.get_empty( + buffers, + (self.max_num_requests, ), + dtype=torch.int32, + cache_name="_mla_kv_len_arr_buf", + capture_graph=capture_graph, + ) + # Rebind the wrapper to the freshly allocated buffers. + self._mla_decode_wrapper = None + self._mla_decode_planned = False def create_cuda_graph_metadata(self, max_batch_size: int, @@ -613,6 +784,14 @@ def prepare(self) -> None: non_blocking=True) if self.num_generations < bs: kv_lens_buf[self.num_generations:bs].zero_() + # Re-plan the MLA decode wrapper outside of any stream capture. + if (self._mla_decode_plan_params is not None + and self._mla_decode_wrapper is not None): + self._mla_decode_planned = False + if self.num_generations > 0: + torch.cuda.current_stream().synchronize() + self._do_plan_mla_decode(self._mla_decode_plan_params) + self._mla_decode_planned = True if self.cross is not None and self.cross is not self: self.cross.prepare() @@ -802,6 +981,10 @@ class FlashInferAttention(AttentionBackend[FlashInferAttentionMetadata]): Metadata = FlashInferAttentionMetadata + @classmethod + def support_mla(cls) -> bool: + return True + def __init__( self, layer_idx: int, @@ -811,6 +994,7 @@ def __init__( quant_config: Optional[QuantConfig] = None, q_scaling: Optional[float] = None, skip_create_weights_in_init: bool = False, + mla_params: Optional[MLAParams] = None, **kwargs, ): self.flashinfer_backend = kwargs.pop('flashinfer_backend', "fa2") @@ -820,6 +1004,13 @@ def __init__( self.update_quant_config(self.quant_config) self.q_scaling = q_scaling + self.is_mla_enable = mla_params is not None + if self.is_mla_enable: + self.kv_lora_rank = mla_params.kv_lora_rank + self.qk_rope_head_dim = mla_params.qk_rope_head_dim + self.qk_nope_head_dim = mla_params.qk_nope_head_dim + self.v_head_dim = mla_params.v_head_dim + def update_quant_config(self, new_quant_config: Optional[QuantConfig]): self.quant_config = new_quant_config self.has_fp8_kv_cache = False @@ -827,6 +1018,205 @@ def update_quant_config(self, new_quant_config: Optional[QuantConfig]): self.has_fp8_kv_cache = self.quant_config.layer_quant_mode.has_fp8_kv_cache( ) + def mla_rope_generation( + self, + fused_q: torch.Tensor, + q_pe: torch.Tensor, + latent_cache: torch.Tensor, + metadata, + cu_q_seqlens: torch.Tensor, + cu_kv_seqlens: torch.Tensor, + fmha_scheduler_counter: torch.Tensor, + mla_bmm1_scale, + mla_bmm2_scale, + quant_q_buffer, + out_scale=None, + ) -> None: + """Stub for MLA generation rope step used when FlashInfer is the mqa backend. + + FlashInferAttention does not fuse RoPE (support_fused_rope returns False), + so RoPE is applied externally in MLA.forward_impl before this point. + q_pe already has RoPE applied; we just copy it into the rope slot of + fused_q so that forward_absorption_generation can pass fused_q directly + to _mla_forward_generation. The latent_cache KV-cache append is handled + inside _mla_forward_generation when forward() is called. + """ + # fused_q shape: [num_tokens, num_heads, kv_lora_rank + qk_rope_head_dim] + # q_pe shape: [num_tokens, num_heads, qk_rope_head_dim] + fused_q[..., self.kv_lora_rank:] = q_pe + + def _get_mla_caches( + self, + metadata: "FlashInferAttentionMetadata", + kv_dtype: torch.dtype, + ): + """Derive per-instance MLA ckv/kpe cache views from the standard KV buffer. + + For MLA models the KV cache manager allocates a single buffer per layer + with kv_factor=1 and head_dim = kv_lora_rank + qk_rope_head_dim. + get_buffers() (NHD layout) returns a tensor with shape + [num_pages, 1, page_size, 1, kv_lora_rank + qk_rope_head_dim]. + + We squeeze out the singleton kv_factor and num_kv_heads dimensions + to obtain [num_pages, page_size, kv_lora_rank + qk_rope_head_dim] + and then create non-allocating views for ckv and kpe. + + Returns: + (ckv_cache, kpe_cache) with shapes + [num_pages, page_size, kv_lora_rank] and + [num_pages, page_size, qk_rope_head_dim]. + """ + # NHD layout: [num_pages, kv_factor=1, page_size, num_kv_heads=1, head_dim] + kv_buf = metadata.kv_cache_manager.get_buffers(self.layer_idx) + # [num_pages, page_size, kv_lora_rank + qk_rope_head_dim] + combined = kv_buf.squeeze(1).squeeze(2) + if self.has_fp8_kv_cache: + combined = combined.view(torch.float8_e4m3fn) + ckv_cache = combined[..., :self.kv_lora_rank] + kpe_cache = combined[..., self.kv_lora_rank:] + return ckv_cache, kpe_cache + + def _mla_forward_context( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + metadata: FlashInferAttentionMetadata, + output: torch.Tensor, + latent_cache: torch.Tensor, + ) -> None: + """MLA context phase: append latent to MLA caches, run ragged prefill.""" + # 1. Append latent_cache to separate ckv/kpe paged caches. + # latent_cache shape: [num_ctx_tokens, kv_lora_rank + qk_rope_head_dim] + num_ctx_tokens = metadata.num_ctx_tokens + append_ckv = latent_cache[:, :self.kv_lora_rank] + append_kpe = latent_cache[:, self.kv_lora_rank:] + + kv_dtype = q.dtype + if self.has_fp8_kv_cache: + kv_dtype = torch.float8_e4m3fn + append_ckv = append_ckv.to(kv_dtype) + append_kpe = append_kpe.to(kv_dtype) + + ckv_cache, kpe_cache = self._get_mla_caches(metadata, kv_dtype) + + ctx_batch_indices = metadata.batch_indices[:num_ctx_tokens] + ctx_positions = metadata.positions[:num_ctx_tokens] + + flashinfer.page.append_paged_mla_kv_cache( + append_ckv, + append_kpe, + ctx_batch_indices, + ctx_positions, + ckv_cache, + kpe_cache, + metadata.paged_kv_indices, + metadata.paged_kv_indptr, + metadata.paged_kv_last_page_len, + ) + + # 2. Run ragged prefill with expanded q, k, v + num_contexts = metadata.num_contexts + qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + + q_ctx = q[:num_ctx_tokens].view(-1, self.num_heads, qk_head_dim) + k_ctx = k[:num_ctx_tokens].view(-1, self.num_kv_heads, qk_head_dim) + v_ctx = v[:num_ctx_tokens].view(-1, self.num_kv_heads, self.v_head_dim) + + sm_scale = None + if self.q_scaling is not None: + sm_scale = 1 / (math.sqrt(qk_head_dim) * self.q_scaling) + + ragged_params = RaggedPlanParams( + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim=qk_head_dim, + head_dim_vo=self.v_head_dim, + q_dtype=q.dtype, + kv_dtype=k.dtype, + sm_scale=sm_scale, + ) + + qo_indptr = metadata.qo_indptr[:num_contexts + 1] + kv_indptr = qo_indptr # self-attention: same as qo + + wrapper = metadata.plan_ragged(qo_indptr, kv_indptr, ragged_params) + + out_view = output[:num_ctx_tokens].view(-1, self.num_heads, + self.v_head_dim) + wrapper.run(q_ctx, k_ctx, v_ctx, out=out_view) + + def _mla_forward_generation( + self, + q: torch.Tensor, + metadata: FlashInferAttentionMetadata, + output: torch.Tensor, + latent_cache: Optional[torch.Tensor] = None, + ) -> None: + """MLA generation phase: append latent to MLA caches, then BatchMLAPagedAttentionWrapper decode.""" + kv_dtype = q.dtype + if self.has_fp8_kv_cache: + kv_dtype = torch.float8_e4m3fn + ckv_cache, kpe_cache = self._get_mla_caches(metadata, kv_dtype) + + # If latent_cache is provided, append it to the paged MLA KV cache first. + # latent_cache shape: [num_tokens, kv_lora_rank + qk_rope_head_dim] + # RoPE must already be applied to the k_pe portion before calling this. + if latent_cache is not None: + append_ckv = latent_cache[:, :self.kv_lora_rank] + append_kpe = latent_cache[:, self.kv_lora_rank:] + if self.has_fp8_kv_cache: + append_ckv = append_ckv.to(kv_dtype) + append_kpe = append_kpe.to(kv_dtype) + num_ctx_tokens = metadata.num_ctx_tokens + gen_batch_indices = metadata.batch_indices[num_ctx_tokens:] + gen_positions = metadata.positions[num_ctx_tokens:] + flashinfer.page.append_paged_mla_kv_cache( + append_ckv, + append_kpe, + gen_batch_indices, + gen_positions, + ckv_cache, + kpe_cache, + metadata.paged_kv_indices, + metadata.paged_kv_indptr, + metadata.paged_kv_last_page_len, + ) + + # fused_q layout: [num_tokens, num_heads * (kv_lora_rank + qk_rope_head_dim)] + # Split into q_nope (absorbed) and q_pe (rope) + num_tokens = q.shape[0] + q_3d = q.view(num_tokens, self.num_heads, self.head_dim) + q_nope = q_3d[..., :self.kv_lora_rank] + q_pe = q_3d[..., self.kv_lora_rank:] + + # sm_scale is based on qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + if self.q_scaling is not None: + sm_scale = 1.0 / (self.q_scaling * math.sqrt(qk_head_dim)) + else: + sm_scale = 1.0 / math.sqrt(qk_head_dim) + + plan_params = MLADecodePlanParams( + num_heads=self.num_heads, + kv_lora_rank=self.kv_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + page_size=metadata.page_size, + q_dtype=q.dtype, + kv_dtype=kv_dtype, + sm_scale=sm_scale, + ) + + wrapper = metadata.plan_mla_decode(plan_params) + + # output: [num_tokens, num_heads, kv_lora_rank] + wrapper.run(q_nope, + q_pe, + ckv_cache, + kpe_cache, + out=output[:num_tokens].view(-1, self.num_heads, + self.kv_lora_rank)) + def forward_impl( self, q: torch.Tensor, @@ -837,7 +1227,20 @@ def forward_impl( output: torch.Tensor, attention_mask_data: Optional[torch.Tensor] = None, attention_window_size: Optional[int] = None, + latent_cache: Optional[torch.Tensor] = None, ) -> None: + # MLA dispatch + if self.is_mla_enable: + if latent_cache is not None and k is not None and v is not None: + # MLA context phase: ragged prefill + cache append + self._mla_forward_context(q, k, v, metadata, output, + latent_cache) + return + elif k is None and v is None: + # MLA generation phase: paged decode + slice + self._mla_forward_generation(q, metadata, output, latent_cache) + return + # Query q = q.view(-1, self.num_heads, self.head_dim) @@ -1031,7 +1434,16 @@ def forward(self, output = forward_args.output if output is None: - output = torch.empty_like(q) + if self.is_mla_enable and latent_cache is not None and k is not None and v is not None: + # MLA context: output has v_head_dim per head, not head_dim + output = q.new_empty( + [q.shape[0], self.num_heads * self.v_head_dim]) + elif self.is_mla_enable and k is None and v is None: + # MLA generation: output has kv_lora_rank per head + output = q.new_empty( + [q.shape[0], self.num_heads * self.kv_lora_rank]) + else: + output = torch.empty_like(q) # FlashInfer's sliding window attention is inclusive, while the attention window size defined in TRTLLM is exclusive. # So we need to subtract 1 from the attention window size for a consistent behavior. @@ -1046,5 +1458,6 @@ def forward(self, attention_mask_type=attention_mask_type, attention_mask_data=attention_mask_data, attention_window_size=attention_window_size, - output=output) + output=output, + latent_cache=latent_cache) return output diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 68ffa2e8f874..fd78968ee15e 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -51,19 +51,9 @@ def extract_extra_attrs(layer_idx: str, attn_type: str): metadata_ref = extra_attrs.get("attention_metadata", None) assert metadata_ref is not None, "Attention metadata is not set" metadata = metadata_ref() - if attn_type == "mla": - assert isinstance( - metadata, - TrtllmAttentionMetadata, - ) - else: - assert isinstance( - metadata, - FlashInferAttentionMetadata, - ) or isinstance( - metadata, - TrtllmAttentionMetadata, - ) + assert isinstance( + metadata, (FlashInferAttentionMetadata, TrtllmAttentionMetadata) + ), "Metadata must be a subclass of FlashInferAttentionMetadata or TrtllmAttentionMetadata" attn_layers = extra_attrs.get(attn_type + "_layers", None) assert attn_layers is not None, "Attention layer is not registered" @@ -1725,7 +1715,19 @@ def forward_impl(self, latent_cache_ctx = latent_cache[:num_ctx_tokens, ...] if self.apply_rotary_emb: assert position_ids is not None - k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, position_ids) + if isinstance(attn_metadata, FlashInferAttentionMetadata): + # position_ids spans [ctx..., gen...] in mixed batches; + # slice to match q_ctx/k_pe_ctx so external RoPE uses ctx + # positions. + ctx_position_ids = position_ids[..., :num_ctx_tokens] + k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, + ctx_position_ids) + # Rebuild latent_cache with RoPE'd k_pe for backends that + # don't handle fused RoPE internally (e.g., FlashInfer). + latent_cache_ctx = torch.cat([compressed_kv_ctx, k_pe_ctx], + dim=-1) + else: + k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, position_ids) if self.llama_4_scaling: q_ctx = self._attention_scaling( @@ -1749,7 +1751,20 @@ def forward_impl(self, latent_cache_gen = latent_cache[num_ctx_tokens:, ...] if self.apply_rotary_emb: assert position_ids is not None - k_pe_gen = self.apply_rope(q_gen, k_pe_gen, position_ids) + if isinstance(attn_metadata, FlashInferAttentionMetadata): + # position_ids spans [ctx..., gen...] in mixed batches; + # gen positions start at num_ctx_tokens. Without this + # slice the external RoPE op applied ctx positions to gen + # k_pe and poisoned the paged MLA cache. + gen_position_ids = position_ids[..., num_ctx_tokens:] + k_pe_gen = self.apply_rope(q_gen, k_pe_gen, + gen_position_ids) + # Rebuild latent_cache with RoPE'd k_pe for backends that + # don't handle fused RoPE internally (e.g., FlashInfer). + latent_cache_gen = torch.cat([compressed_kv_gen, k_pe_gen], + dim=-1) + else: + k_pe_gen = self.apply_rope(q_gen, k_pe_gen, position_ids) if self.llama_4_scaling: q_gen = self._attention_scaling( @@ -2451,7 +2466,10 @@ def forward_absorption_generation( # fused_q contains 1) the result of the following bmm with shape [num_tokens, num_heads, kv_lora_rank] # 2) rope(q_pe) with shape [num_tokens, num_heads, qk_rope_head_dim]. rope is applied inside AttentionOp - num_seqs = attn_metadata.kv_lens_cuda_runtime.size(0) + if isinstance(attn_metadata, FlashInferAttentionMetadata): + num_seqs = attn_metadata.num_generations + else: + num_seqs = attn_metadata.kv_lens_cuda_runtime.size(0) cu_q_seqlens = torch.empty(num_seqs + 1, dtype=torch.int32, diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 31d6a17c85f1..1bd4fa136d9d 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -1705,6 +1705,16 @@ def test_bfloat16(self, mtp_nextn, attention_dp, cuda_graph, task = GSM8K(self.MODEL_NAME) task.evaluate(llm) + @pytest.mark.skip_less_device_memory(60000) + def test_bfloat16_flashinfer(self): + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.75) + with LLM(self.MODEL_PATH, + kv_cache_config=kv_cache_config, + attn_backend="FLASHINFER", + max_num_tokens=8192) as llm: + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + @pytest.mark.skip_less_device_memory(60000) @parametrize_with_ids("enable_chunked_prefill", [False, True]) @parametrize_with_ids("attention_dp,cuda_graph,overlap_scheduler", diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 2750917eb7a7..cd0a8be11620 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -36,6 +36,7 @@ l0_b200: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-enable_chunked_prefill=True] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-enable_chunked_prefill=True] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_2_model_mtp + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=0-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=TRTLLM-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=False] diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index d0ac7ed986c5..24aa39b589ba 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -59,3 +59,4 @@ l0_b300: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=TRTLLM-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=2-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_chunked_prefill[quant_dtype=fp8-kv_cache_reuse=True-fp8kv=True-overlap_scheduler=True] # Cover nvbugs 6084445 + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer diff --git a/tests/unittest/_torch/attention/test_attention_mla.py b/tests/unittest/_torch/attention/test_attention_mla.py index a634175edc6e..13e7362e56dd 100644 --- a/tests/unittest/_torch/attention/test_attention_mla.py +++ b/tests/unittest/_torch/attention/test_attention_mla.py @@ -292,6 +292,36 @@ class RopeConfig: model_type: str = "deepseek_v3" +def apply_mla_rope(tensor: torch.Tensor, positions: list, + rope_cos_sin: torch.Tensor) -> torch.Tensor: + """Apply MLA-style RoPE to the last dimension of tensor. + + Reorders from interleaved to pair format (unflatten/transpose/flatten), + then applies standard (cos, sin) rotation. Used to pre-apply RoPE for + backends that do not fuse RoPE internally (e.g. FlashInfer MLA). + + Args: + tensor: [..., qk_rope_head_dim] + positions: list of integer position indices, length = tensor.shape[0] + rope_cos_sin: [max_pos, 2, qk_rope_head_dim] + + Returns: + tensor with RoPE applied (same shape/dtype) + """ + pos = torch.tensor(positions, dtype=torch.long) + cos_sin = rope_cos_sin[pos] # [num_tokens, 2, qk_rope_head_dim] + cos = cos_sin[:, 0] # [num_tokens, qk_rope_head_dim] + sin = cos_sin[:, 1] # [num_tokens, qk_rope_head_dim] + # Expand to match tensor dimensions + for _ in range(tensor.dim() - 2): + cos = cos.unsqueeze(-2) + sin = sin.unsqueeze(-2) + # Reorder: interleaved -> pair format + t = tensor.unflatten(-1, [-1, 2]).transpose(-2, -1).flatten(start_dim=-2) + rotated = (t * cos + rotate_half(t) * sin).to(tensor.dtype) + return rotated + + def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, @@ -411,6 +441,82 @@ def test_attention_mla(scenario: Scenario, context_sequence_lengths: List[int], v2_kv_cache) +# FlashInfer MLA test: BF16 only, fewer combos since it's slower +flashinfer_scenarios = [ + Scenario(kv_cache_dtype=torch.bfloat16, + num_layers=num_layers, + kv_cache_tokens_per_block=tokens_per_block) for num_layers in [1] +] + +flashinfer_context_sequence_lengths = [ + [10, 12, 5], + [100, 300, 20, 10], +] + + +@pytest.mark.parametrize("scenario", + flashinfer_scenarios, + ids=lambda x: f"scenario: {x}") +@pytest.mark.parametrize("context_sequence_lengths", + flashinfer_context_sequence_lengths, + ids=lambda x: f"context_sequence_lengths: {x}") +@pytest.mark.parametrize("generation_seq_len_q", [1], + ids=lambda x: f"generation_seq_len_q: {x}") +@pytest.mark.parametrize("num_generation_steps", [10], + ids=lambda x: f"num_generation_steps: {x}") +@pytest.mark.parametrize("v2_kv_cache", [True, False], + ids=["v2_kv_cache", "v1_kv_cache"]) +def test_attention_mla_flashinfer(scenario: Scenario, + context_sequence_lengths: List[int], + generation_seq_len_q: int, + num_generation_steps: List[int], + v2_kv_cache: bool): + """Test FlashInfer MLA computation for both context and generation phases""" + pytest.importorskip("flashinfer") + if (not torch.cuda.is_available() + or torch.cuda.get_device_capability() != (10, 0)): + pytest.skip("FlashInfer MLA test only runs on SM100 (Blackwell)") + + num_heads = scenario.num_heads + num_kv_heads = scenario.num_kv_heads + q_lora_rank = scenario.q_lora_rank + kv_lora_rank = scenario.kv_lora_rank + qk_nope_head_dim = scenario.qk_nope_head_dim + qk_rope_head_dim = scenario.qk_rope_head_dim + v_head_dim = scenario.v_head_dim + rope_config = RopeConfig( + hidden_size=scenario.hidden_size, + num_attention_heads=scenario.num_heads, + rope_scaling={ + "beta_fast": scenario.rope_beta_fast, + "beta_slow": scenario.rope_beta_slow, + "factor": scenario.rope_factor, + "mscale": scenario.rope_mscale, + "mscale_all_dim": scenario.rope_mscale_all_dim, + "original_max_position_embeddings": + scenario.rope_original_max_position_embeddings, + "type": scenario.rope_type, + }, + max_position_embeddings=scenario.max_position_embeddings, + rope_theta=scenario.rope_theta, + qk_rope_head_dim=scenario.qk_rope_head_dim, + model_type=scenario.model_type, + ) + kv_cache_tokens_per_block = scenario.kv_cache_tokens_per_block + num_layers = scenario.num_layers + device = torch.device('cuda') + dtype = scenario.dtype + kv_cache_dtype = scenario.kv_cache_dtype + + _run_test_for_backend("FLASHINFER", num_heads, num_kv_heads, num_layers, + q_lora_rank, kv_lora_rank, qk_nope_head_dim, + qk_rope_head_dim, v_head_dim, rope_config, + kv_cache_tokens_per_block, device, dtype, + kv_cache_dtype, context_sequence_lengths, + generation_seq_len_q, num_generation_steps, + v2_kv_cache) + + def _run_test_for_backend(backend_name, num_heads, num_kv_heads, num_layers, q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, rope_config, @@ -647,7 +753,7 @@ def yarn_get_mscale(scale=1, mscale=1): max_num_requests=len(context_sequence_lengths), num_contexts=len(context_sequence_lengths), prompt_lens=context_sequence_lengths, - max_num_tokens=max(context_sequence_lengths), + max_num_tokens=sum(context_sequence_lengths), kv_cache_manager=kv_cache_manager, kv_cache_params=KVCacheParams( use_cache=True, @@ -668,7 +774,7 @@ def yarn_get_mscale(scale=1, mscale=1): kv_cache.capacity += 1 else: kv_cache_manager.impl.add_token(req_id) - attn_metadata = AttentionCls.Metadata( + gen_metadata_kwargs = dict( seq_lens=torch.tensor([generation_seq_len_q] * len(context_sequence_lengths), dtype=torch.int), @@ -676,7 +782,7 @@ def yarn_get_mscale(scale=1, mscale=1): max_num_requests=len(context_sequence_lengths), num_contexts=0, prompt_lens=context_sequence_lengths, - max_num_tokens=max(context_sequence_lengths), + max_num_tokens=sum(context_sequence_lengths), kv_cache_manager=kv_cache_manager, kv_cache_params=KVCacheParams( use_cache=True, @@ -686,8 +792,12 @@ def yarn_get_mscale(scale=1, mscale=1): ], ), mapping=mapping, - enable_flash_mla=torch.cuda.get_device_capability() == (9, 0), ) + if backend_name == "TRTLLM": + gen_metadata_kwargs[ + 'enable_flash_mla'] = torch.cuda.get_device_capability( + ) == (9, 0) + attn_metadata = AttentionCls.Metadata(**gen_metadata_kwargs) attn_metadata.prepare() for layer_idx in range(num_layers): print( @@ -702,13 +812,38 @@ def yarn_get_mscale(scale=1, mscale=1): latent_cache = torch.cat([compressed_kv, k_pe], dim=-1) # q/k will be modified in the forward pass, so we need to clone them # we should not clone v because we need to keep the stride of v + if backend_name == "FLASHINFER": + # FlashInfer MLA does not fuse RoPE; pre-apply it here. + ctx_positions = [ + pos for ctx_len in context_sequence_lengths + for pos in range(ctx_len) + ] + q_fwd = q.clone().view(-1, num_heads, qk_head_dim) + q_fwd[..., qk_nope_head_dim:] = apply_mla_rope( + q_fwd[..., qk_nope_head_dim:], ctx_positions, + rope_cos_sin) + q_fwd = q_fwd.view(-1, num_heads * qk_head_dim) + k_fwd = k.clone().view(-1, num_kv_heads, qk_head_dim) + k_fwd[..., qk_nope_head_dim:] = apply_mla_rope( + k_fwd[..., qk_nope_head_dim:], ctx_positions, + rope_cos_sin) + k_fwd = k_fwd.view(-1, num_kv_heads * qk_head_dim) + lc_kpe = latent_cache[:, kv_lora_rank:].unsqueeze(1) + lc_kpe_rope = apply_mla_rope(lc_kpe, ctx_positions, + rope_cos_sin).squeeze(1) + latent_cache_fwd = torch.cat( + [latent_cache[:, :kv_lora_rank], lc_kpe_rope], dim=-1) + else: + q_fwd = q.clone() + k_fwd = k.clone() + latent_cache_fwd = latent_cache result = ctx_layers[layer_idx].forward( - q.clone(), - k.clone(), + q_fwd, + k_fwd, v, attn_metadata, attention_input_type=AttentionInputType.context_only, - latent_cache=latent_cache, + latent_cache=latent_cache_fwd, ) ref_result, latent_cache_ref = calculate_ref_result_ctx( q, @@ -737,58 +872,95 @@ def yarn_get_mscale(scale=1, mscale=1): k_pe = inputs_per_layer[layer_idx]["gen_k_pe_list"][step - 1] latent_cache = torch.cat([compressed_kv, k_pe], dim=-1) - num_tokens = fused_q.size(0) - num_seqs = attn_metadata.kv_lens_cuda_runtime.size(0) - cu_q_seqlens = torch.empty(num_seqs + 1, - dtype=torch.int32, - device=q.device) - cu_kv_seqlens = torch.empty(num_seqs + 1, - dtype=torch.int32, - device=q.device) - fmha_scheduler_counter = torch.empty(1, - dtype=torch.uint32, - device=q.device) - has_fp8_kv_cache = gen_layers[ - layer_idx].has_fp8_kv_cache if hasattr( - gen_layers[layer_idx], 'has_fp8_kv_cache') else False - - if has_fp8_kv_cache: - mla_bmm1_scale = torch.empty(2, - dtype=torch.float32, - device=q.device) - mla_bmm2_scale = torch.empty(1, - dtype=torch.float32, - device=q.device) - quant_q_buffer = torch.empty( - num_tokens, - num_heads * (kv_lora_rank + qk_rope_head_dim), - dtype=torch.uint8, - device=q.device) + if backend_name == "FLASHINFER": + # FlashInfer MLA does not fuse RoPE; pre-apply before + # appending to the KV cache. Keep original q_pe and + # latent_cache intact for the reference calculation + # (calculate_ref_result_gen applies RoPE internally). + gen_positions = [ + ctx_len + (step - 1) * generation_seq_len_q + i + for ctx_len in context_sequence_lengths + for i in range(generation_seq_len_q) + ] + q_pe_for_gen = apply_mla_rope(q_pe, gen_positions, + rope_cos_sin) + lc_kpe = latent_cache[:, kv_lora_rank:].unsqueeze(1) + lc_kpe_rope = apply_mla_rope(lc_kpe, gen_positions, + rope_cos_sin).squeeze(1) + latent_cache_for_gen = torch.cat( + [latent_cache[:, :kv_lora_rank], lc_kpe_rope], dim=-1) + # Copy RoPE'd q_pe into the rope portion of fused_q + num_tokens = fused_q.size(0) + fused_q_3d = fused_q.view(num_tokens, num_heads, + kv_lora_rank + qk_rope_head_dim) + fused_q_3d[..., kv_lora_rank:].copy_(q_pe_for_gen) + + result = gen_layers[layer_idx].forward( + fused_q, + None, + None, + attn_metadata, + attention_input_type=AttentionInputType.generation_only, + latent_cache=latent_cache_for_gen, + ) else: - mla_bmm1_scale = None - mla_bmm2_scale = None - quant_q_buffer = None - - gen_layers[layer_idx].mla_rope_generation( - fused_q, q_pe, latent_cache, attn_metadata, cu_q_seqlens, - cu_kv_seqlens, fmha_scheduler_counter, mla_bmm1_scale, - mla_bmm2_scale, quant_q_buffer) - - result = gen_layers[layer_idx].forward( - fused_q, - None, - None, - attn_metadata, - attention_input_type=AttentionInputType.generation_only, - latent_cache=latent_cache, - q_pe=q_pe, - cu_q_seqlens=cu_q_seqlens, - cu_kv_seqlens=cu_kv_seqlens, - fmha_scheduler_counter=fmha_scheduler_counter, - mla_bmm1_scale=mla_bmm1_scale, - mla_bmm2_scale=mla_bmm2_scale, - quant_q_buffer=quant_q_buffer, - ) + q_pe_for_gen = q_pe + latent_cache_for_gen = latent_cache + + num_tokens = fused_q.size(0) + num_seqs = len(context_sequence_lengths) + cu_q_seqlens = torch.empty(num_seqs + 1, + dtype=torch.int32, + device=q.device) + cu_kv_seqlens = torch.empty(num_seqs + 1, + dtype=torch.int32, + device=q.device) + fmha_scheduler_counter = torch.empty(1, + dtype=torch.uint32, + device=q.device) + has_fp8_kv_cache = gen_layers[ + layer_idx].has_fp8_kv_cache if hasattr( + gen_layers[layer_idx], + 'has_fp8_kv_cache') else False + + if has_fp8_kv_cache: + mla_bmm1_scale = torch.empty(2, + dtype=torch.float32, + device=q.device) + mla_bmm2_scale = torch.empty(1, + dtype=torch.float32, + device=q.device) + quant_q_buffer = torch.empty( + num_tokens, + num_heads * (kv_lora_rank + qk_rope_head_dim), + dtype=torch.uint8, + device=q.device) + else: + mla_bmm1_scale = None + mla_bmm2_scale = None + quant_q_buffer = None + + gen_layers[layer_idx].mla_rope_generation( + fused_q, q_pe_for_gen, latent_cache_for_gen, + attn_metadata, cu_q_seqlens, cu_kv_seqlens, + fmha_scheduler_counter, mla_bmm1_scale, mla_bmm2_scale, + quant_q_buffer) + + result = gen_layers[layer_idx].forward( + fused_q, + None, + None, + attn_metadata, + attention_input_type=AttentionInputType.generation_only, + latent_cache=latent_cache, + q_pe=q_pe, + cu_q_seqlens=cu_q_seqlens, + cu_kv_seqlens=cu_kv_seqlens, + fmha_scheduler_counter=fmha_scheduler_counter, + mla_bmm1_scale=mla_bmm1_scale, + mla_bmm2_scale=mla_bmm2_scale, + quant_q_buffer=quant_q_buffer, + ) ref_result, latent_cache_ref = calculate_ref_result_gen( fused_q, q_pe, From dee75ec164729ec54299e98274321b6fb1ca01e7 Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Mon, 11 May 2026 19:45:39 -0700 Subject: [PATCH 02/11] Rebase branch. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- requirements.txt | 4 ++-- tensorrt_llm/_torch/attention_backend/flashinfer.py | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/requirements.txt b/requirements.txt index 5f373dec74d8..c652ea17625f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -54,7 +54,7 @@ ordered-set peft>=0.18.1,<0.19.0 patchelf einops -flashinfer-python==0.6.10 +flashinfer-python==0.6.11 opencv-python-headless xgrammar==0.1.32 llguidance==0.7.29 @@ -71,7 +71,7 @@ xdsl>=0.59.0 # Optional: required for MLIR-based elementwise fusion in AutoDeplo tiktoken blobfile openai-harmony==0.0.4 -nvidia-cutlass-dsl==4.4.2; python_version >= "3.10" +nvidia-cutlass-dsl==4.5.0; python_version >= "3.10" plotly numexpr partial_json_parser diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index aa27caddf13a..0f5b8d37ddab 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -15,8 +15,8 @@ from ..metadata import KVCacheParams from ..utils import get_global_attrs, get_model_extra_attrs from .interface import (AttentionBackend, AttentionForwardArgs, - AttentionMetadata, CustomAttentionMask, - PredefinedAttentionMask, merge_attention_forward_args, MLAParams) + AttentionMetadata, CustomAttentionMask, MLAParams, + PredefinedAttentionMask, merge_attention_forward_args) try: check_cuda_arch() @@ -1418,6 +1418,7 @@ def forward(self, forward_args = merge_attention_forward_args(forward_args, kwargs) attention_mask_data = forward_args.attention_mask_data + latent_cache = forward_args.latent_cache if forward_args.attention_mask == CustomAttentionMask.CUSTOM: assert attention_mask_data is not None, "attention_mask_data is required for custom attention mask." attention_mask_type = int(AttentionMaskType.custom_mask) From a00d56b465b2b8fa95bdeb4984cf72aef3a5451c Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Tue, 12 May 2026 01:12:04 -0700 Subject: [PATCH 03/11] [None][feat] Support FlashInfer MLA chunked prefill Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- .../_torch/attention_backend/flashinfer.py | 155 ++++++++++++++++-- tensorrt_llm/_torch/modules/attention.py | 12 ++ .../defs/accuracy/test_llm_api_pytorch.py | 6 +- .../test_lists/test-db/l0_b200.yml | 3 +- .../test_lists/test-db/l0_b300.yml | 3 +- 5 files changed, 162 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 0f5b8d37ddab..c811bb87a5ae 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -15,8 +15,9 @@ from ..metadata import KVCacheParams from ..utils import get_global_attrs, get_model_extra_attrs from .interface import (AttentionBackend, AttentionForwardArgs, - AttentionMetadata, CustomAttentionMask, MLAParams, - PredefinedAttentionMask, merge_attention_forward_args) + AttentionInputType, AttentionMetadata, + CustomAttentionMask, MLAParams, PredefinedAttentionMask, + merge_attention_forward_args) try: check_cuda_arch() @@ -118,6 +119,7 @@ class FlashInferAttentionMetadata(AttentionMetadata): # Cached plan params + is-planned flag let prepare() refresh the plan # outside stream capture (flashinfer plan() does device->host syncs). _mla_decode_wrapper: Optional[object] = field(init=False, default=None) + _mla_context_wrapper: Optional[object] = field(init=False, default=None) _mla_decode_plan_params: Optional[MLADecodePlanParams] = field(init=False, default=None) _mla_decode_planned: bool = field(init=False, default=False) @@ -221,6 +223,52 @@ def plan_mla_decode( self._mla_decode_planned = True return self._mla_decode_wrapper + def plan_mla_context( + self, + qo_indptr: torch.Tensor, + kv_indptr: torch.Tensor, + kv_indices: torch.Tensor, + kv_last_page_len: torch.Tensor, + plan_params: MLADecodePlanParams, + ) -> object: + """Plan MLA context with cached KV using BatchMLAPagedAttentionWrapper.""" + if self._mla_context_wrapper is None: + self._mla_context_wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper( + self.workspace_buffer, + use_cuda_graph=False, + backend="auto", + ) + + if self.is_cuda_graph and torch.cuda.is_current_stream_capturing(): + raise ValueError( + "Cannot plan() flashinfer MLA context while the stream is " + "capturing. Chunked MLA prefill with FlashInfer does not " + "support CUDA graph capture.") + + num_pages_per_seq = kv_indptr[1:] - kv_indptr[:-1] + kv_len_arr = (num_pages_per_seq - + 1) * plan_params.page_size + kv_last_page_len + + # Must sync after append_paged_mla_kv_cache and before plan(). + torch.cuda.current_stream().synchronize() + + self._mla_context_wrapper.plan( + qo_indptr, + kv_indptr, + kv_indices, + kv_len_arr, + plan_params.num_heads, + plan_params.kv_lora_rank, + plan_params.qk_rope_head_dim, + plan_params.page_size, + causal=True, + q_data_type=plan_params.q_dtype, + kv_data_type=plan_params.kv_dtype, + sm_scale=plan_params.sm_scale, + ) + + return self._mla_context_wrapper + def _do_plan_mla_decode(self, plan_params: MLADecodePlanParams) -> None: """Compute MLA decode plan inputs and call wrapper.plan(). @@ -455,6 +503,7 @@ def _post_init_with_buffers(self, buffers) -> None: ) # Rebind the wrapper to the freshly allocated buffers. self._mla_decode_wrapper = None + self._mla_context_wrapper = None self._mla_decode_planned = False def create_cuda_graph_metadata(self, @@ -1217,6 +1266,76 @@ def _mla_forward_generation( out=output[:num_tokens].view(-1, self.num_heads, self.kv_lora_rank)) + def _mla_forward_cached_context( + self, + q: torch.Tensor, + metadata: FlashInferAttentionMetadata, + output: torch.Tensor, + latent_cache: torch.Tensor, + ) -> None: + """MLA context phase with cached KV: append latent and run paged MLA.""" + num_ctx_tokens = metadata.num_ctx_tokens + kv_dtype = q.dtype + if self.has_fp8_kv_cache: + kv_dtype = torch.float8_e4m3fn + + ckv_cache, kpe_cache = self._get_mla_caches(metadata, kv_dtype) + + append_ckv = latent_cache[:, :self.kv_lora_rank] + append_kpe = latent_cache[:, self.kv_lora_rank:] + if self.has_fp8_kv_cache: + append_ckv = append_ckv.to(kv_dtype) + append_kpe = append_kpe.to(kv_dtype) + + flashinfer.page.append_paged_mla_kv_cache( + append_ckv, + append_kpe, + metadata.batch_indices[:num_ctx_tokens], + metadata.positions[:num_ctx_tokens], + ckv_cache, + kpe_cache, + metadata.paged_kv_indices, + metadata.paged_kv_indptr, + metadata.paged_kv_last_page_len, + ) + + num_tokens = q.shape[0] + q_3d = q.view(num_tokens, self.num_heads, self.head_dim) + q_nope = q_3d[..., :self.kv_lora_rank] + q_pe = q_3d[..., self.kv_lora_rank:] + + qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + if self.q_scaling is not None: + sm_scale = 1.0 / (self.q_scaling * math.sqrt(qk_head_dim)) + else: + sm_scale = 1.0 / math.sqrt(qk_head_dim) + + plan_params = MLADecodePlanParams( + num_heads=self.num_heads, + kv_lora_rank=self.kv_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + page_size=metadata.page_size, + q_dtype=q.dtype, + kv_dtype=kv_dtype, + sm_scale=sm_scale, + ) + + num_contexts = metadata.num_contexts + wrapper = metadata.plan_mla_context( + qo_indptr=metadata.qo_indptr[:num_contexts + 1], + kv_indptr=metadata.paged_kv_indptr_prefill[:num_contexts + 1], + kv_indices=metadata.paged_kv_indices[:metadata.num_context_blocks], + kv_last_page_len=metadata.paged_kv_last_page_len[:num_contexts], + plan_params=plan_params, + ) + + wrapper.run(q_nope, + q_pe, + ckv_cache, + kpe_cache, + out=output[:num_tokens].view(-1, self.num_heads, + self.kv_lora_rank)) + def forward_impl( self, q: torch.Tensor, @@ -1228,6 +1347,7 @@ def forward_impl( attention_mask_data: Optional[torch.Tensor] = None, attention_window_size: Optional[int] = None, latent_cache: Optional[torch.Tensor] = None, + attention_input_type: AttentionInputType = AttentionInputType.mixed, ) -> None: # MLA dispatch if self.is_mla_enable: @@ -1237,8 +1357,15 @@ def forward_impl( latent_cache) return elif k is None and v is None: - # MLA generation phase: paged decode + slice - self._mla_forward_generation(q, metadata, output, latent_cache) + if attention_input_type == AttentionInputType.context_only: + assert latent_cache is not None, ( + "FlashInfer MLA cached context requires latent_cache.") + self._mla_forward_cached_context(q, metadata, output, + latent_cache) + else: + # MLA generation phase: paged decode + slice + self._mla_forward_generation(q, metadata, output, + latent_cache) return # Query @@ -1452,13 +1579,15 @@ def forward(self, if attention_window_size is not None: attention_window_size = attention_window_size - 1 - self.forward_impl(q=q, - k=k, - v=v, - metadata=metadata, - attention_mask_type=attention_mask_type, - attention_mask_data=attention_mask_data, - attention_window_size=attention_window_size, - output=output, - latent_cache=latent_cache) + self.forward_impl( + q=q, + k=k, + v=v, + metadata=metadata, + attention_mask_type=attention_mask_type, + attention_mask_data=attention_mask_data, + attention_window_size=attention_window_size, + output=output, + latent_cache=latent_cache, + attention_input_type=forward_args.attention_input_type) return output diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index fd78968ee15e..2cd9f59afd80 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -2418,6 +2418,18 @@ def forward_context( output: torch.Tensor, latent_cache: Optional[torch.Tensor] = None, ) -> torch.Tensor: + if (isinstance(attn_metadata, FlashInferAttentionMetadata) + and attn_metadata.kv_cache_manager is not None and any( + attn_metadata.kv_cache_params. + num_cached_tokens_per_seq[:attn_metadata.num_contexts])): + return self.forward_absorption_context(q, + compressed_kv, + k_pe, + attn_metadata, + output, + position_ids=position_ids, + latent_cache=latent_cache) + if isinstance(self.mha, TrtllmAttention): assert isinstance(attn_metadata, TrtllmAttentionMetadata) trtllm_attention = cast(TrtllmAttention, self.mha) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 1bd4fa136d9d..6947c3d55cde 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -1706,12 +1706,14 @@ def test_bfloat16(self, mtp_nextn, attention_dp, cuda_graph, task.evaluate(llm) @pytest.mark.skip_less_device_memory(60000) - def test_bfloat16_flashinfer(self): + @parametrize_with_ids("enable_chunked_prefill", [False, True]) + def test_bfloat16_flashinfer(self, enable_chunked_prefill): kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.75) with LLM(self.MODEL_PATH, kv_cache_config=kv_cache_config, attn_backend="FLASHINFER", - max_num_tokens=8192) as llm: + enable_chunked_prefill=enable_chunked_prefill, + max_num_tokens=512 if enable_chunked_prefill else 8192) as llm: task = GSM8K(self.MODEL_NAME) task.evaluate(llm) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index cd0a8be11620..54f6fa4d31ba 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -36,7 +36,8 @@ l0_b200: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-enable_chunked_prefill=True] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-enable_chunked_prefill=True] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_2_model_mtp - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=False] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=True] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=0-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=TRTLLM-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=False] diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 24aa39b589ba..063125788ed9 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -59,4 +59,5 @@ l0_b300: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=TRTLLM-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=2-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_chunked_prefill[quant_dtype=fp8-kv_cache_reuse=True-fp8kv=True-overlap_scheduler=True] # Cover nvbugs 6084445 - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=False] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=True] From e8e6bb935656ea0f82691ad00d0167d6c76d36ac Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Tue, 12 May 2026 19:16:33 -0700 Subject: [PATCH 04/11] Add DeepSeek V3 Lite FlashInfer QA tests Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- tests/integration/test_lists/qa/llm_function_core.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index ef49289cb6fb..8ef7ee438dad 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -220,6 +220,8 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2- accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=True-v2_kv_cache=True] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=False] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=True] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=False] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=True] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_2_model_mtp accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[ep4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[ep4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] From 7445f96434b5f8845a7b86b7af1a49670dc69044 Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Sun, 17 May 2026 22:18:56 -0700 Subject: [PATCH 05/11] [None][fix] Address FlashInfer MLA review comments Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- .../_torch/attention_backend/flashinfer.py | 116 ++++++++++++------ tensorrt_llm/_torch/modules/attention.py | 53 +++----- 2 files changed, 99 insertions(+), 70 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index c811bb87a5ae..47c6bd982210 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -106,6 +106,9 @@ class FlashInferAttentionMetadata(AttentionMetadata): _qo_indptr: torch.Tensor = field(init=False) _kv_indptr: torch.Tensor = field(init=False) _cached_token_lens: torch.Tensor = field(init=False) + kv_lens_cuda_runtime: Optional[torch.Tensor] = field(init=False, + default=None, + repr=False) _plan_params_to_wrappers: Dict[PlanParams, FlashInferWrappers] = field(init=False) @@ -120,8 +123,14 @@ class FlashInferAttentionMetadata(AttentionMetadata): # outside stream capture (flashinfer plan() does device->host syncs). _mla_decode_wrapper: Optional[object] = field(init=False, default=None) _mla_context_wrapper: Optional[object] = field(init=False, default=None) + _mla_ragged_plan_params: Optional[RaggedPlanParams] = field(init=False, + default=None) + _mla_context_plan_params: Optional[MLADecodePlanParams] = field( + init=False, default=None) _mla_decode_plan_params: Optional[MLADecodePlanParams] = field(init=False, default=None) + _mla_ragged_planned: bool = field(init=False, default=False) + _mla_context_planned: bool = field(init=False, default=False) _mla_decode_planned: bool = field(init=False, default=False) _mla_qo_indptr_buf: Optional[torch.Tensor] = field(init=False, default=None) _mla_kv_len_arr_buf: Optional[torch.Tensor] = field(init=False, @@ -170,6 +179,16 @@ def plan_ragged( self.workspace_buffer, "NHD", ) + if self._mla_ragged_plan_params != plan_params: + self._mla_ragged_planned = False + self._mla_ragged_plan_params = plan_params + + if self._mla_ragged_planned: + return self._ragged_prefill_wrapper + + # Split append_paged_mla_kv_cache from plan() when this wrapper needs a + # new plan. Reusing a cached plan avoids this sync on later layers. + torch.cuda.current_stream().synchronize() self._ragged_prefill_wrapper.plan( qo_indptr, @@ -184,6 +203,7 @@ def plan_ragged( kv_data_type=plan_params.kv_dtype, sm_scale=plan_params.sm_scale, ) + self._mla_ragged_planned = True return self._ragged_prefill_wrapper @@ -207,6 +227,9 @@ def plan_mla_decode( backend="auto", ) + if self._mla_decode_plan_params != plan_params: + self._mla_decode_planned = False + # Cache params so prepare() can re-plan on subsequent forward passes. self._mla_decode_plan_params = plan_params @@ -219,6 +242,9 @@ def plan_mla_decode( "capturing. Make sure prepare() has run at least one warmup " "forward pass before capture.") + # Split append_paged_mla_kv_cache from plan() on a cache miss. prepare() + # calls _do_plan_mla_decode() directly and does not need this sync. + torch.cuda.current_stream().synchronize() self._do_plan_mla_decode(plan_params) self._mla_decode_planned = True return self._mla_decode_wrapper @@ -244,12 +270,19 @@ def plan_mla_context( "Cannot plan() flashinfer MLA context while the stream is " "capturing. Chunked MLA prefill with FlashInfer does not " "support CUDA graph capture.") + if self._mla_context_plan_params != plan_params: + self._mla_context_planned = False + self._mla_context_plan_params = plan_params + + if self._mla_context_planned: + return self._mla_context_wrapper num_pages_per_seq = kv_indptr[1:] - kv_indptr[:-1] kv_len_arr = (num_pages_per_seq - 1) * plan_params.page_size + kv_last_page_len - # Must sync after append_paged_mla_kv_cache and before plan(). + # Split append_paged_mla_kv_cache from plan() when this wrapper needs a + # new plan. Reusing a cached plan avoids this sync on later layers. torch.cuda.current_stream().synchronize() self._mla_context_wrapper.plan( @@ -266,6 +299,7 @@ def plan_mla_context( kv_data_type=plan_params.kv_dtype, sm_scale=plan_params.sm_scale, ) + self._mla_context_planned = True return self._mla_context_wrapper @@ -504,6 +538,11 @@ def _post_init_with_buffers(self, buffers) -> None: # Rebind the wrapper to the freshly allocated buffers. self._mla_decode_wrapper = None self._mla_context_wrapper = None + self._mla_ragged_plan_params = None + self._mla_context_plan_params = None + self._mla_decode_plan_params = None + self._mla_ragged_planned = False + self._mla_context_planned = False self._mla_decode_planned = False def create_cuda_graph_metadata(self, @@ -623,6 +662,7 @@ def prepare(self) -> None: self.kv_cache_params = KVCacheParams(use_cache=False) n = self.num_seqs self._cached_token_lens[:n].zero_() + self.kv_lens_cuda_runtime = self.seq_lens_cuda[:n] for plan_params in list(self._plan_params_to_wrappers.keys()): if plan_params.attention_mask_data is None: self._plan_params_to_wrappers[ @@ -645,6 +685,7 @@ def prepare(self) -> None: # number of tokens needed in the kv cache for each sequence after the next pass kv_lens = self.cached_token_lens + self.seq_lens_kv_cuda + self.kv_lens_cuda_runtime = kv_lens[:self.num_seqs] # start and end indices of each sequence in the ragged key and value # for self attention it's the same as qo_indptr so avoid computing twice. @@ -786,6 +827,16 @@ def prepare(self) -> None: else: del self._plan_params_to_wrappers[plan_params] + self._mla_ragged_planned = False + self._mla_context_planned = False + # Re-plan the MLA decode wrapper outside of any stream capture. + if (self.num_generations > 0 + and self._mla_decode_plan_params is not None + and self._mla_decode_wrapper is not None): + self._mla_decode_planned = False + self._do_plan_mla_decode(self._mla_decode_plan_params) + self._mla_decode_planned = True + # VSWA: restore primary pool indices as the default. if (self._vswa_layer_to_pool is not None and self._vswa_pool_indices_cache is not None): @@ -833,15 +884,6 @@ def prepare(self) -> None: non_blocking=True) if self.num_generations < bs: kv_lens_buf[self.num_generations:bs].zero_() - # Re-plan the MLA decode wrapper outside of any stream capture. - if (self._mla_decode_plan_params is not None - and self._mla_decode_wrapper is not None): - self._mla_decode_planned = False - if self.num_generations > 0: - torch.cuda.current_stream().synchronize() - self._do_plan_mla_decode(self._mla_decode_plan_params) - self._mla_decode_planned = True - if self.cross is not None and self.cross is not self: self.cross.prepare() @@ -1097,7 +1139,6 @@ def mla_rope_generation( def _get_mla_caches( self, metadata: "FlashInferAttentionMetadata", - kv_dtype: torch.dtype, ): """Derive per-instance MLA ckv/kpe cache views from the standard KV buffer. @@ -1147,7 +1188,7 @@ def _mla_forward_context( append_ckv = append_ckv.to(kv_dtype) append_kpe = append_kpe.to(kv_dtype) - ckv_cache, kpe_cache = self._get_mla_caches(metadata, kv_dtype) + ckv_cache, kpe_cache = self._get_mla_caches(metadata) ctx_batch_indices = metadata.batch_indices[:num_ctx_tokens] ctx_positions = metadata.positions[:num_ctx_tokens] @@ -1206,31 +1247,32 @@ def _mla_forward_generation( kv_dtype = q.dtype if self.has_fp8_kv_cache: kv_dtype = torch.float8_e4m3fn - ckv_cache, kpe_cache = self._get_mla_caches(metadata, kv_dtype) + ckv_cache, kpe_cache = self._get_mla_caches(metadata) - # If latent_cache is provided, append it to the paged MLA KV cache first. + assert latent_cache is not None, ( + "FlashInfer MLA generation requires latent_cache.") + # Append latent_cache to the paged MLA KV cache first. # latent_cache shape: [num_tokens, kv_lora_rank + qk_rope_head_dim] # RoPE must already be applied to the k_pe portion before calling this. - if latent_cache is not None: - append_ckv = latent_cache[:, :self.kv_lora_rank] - append_kpe = latent_cache[:, self.kv_lora_rank:] - if self.has_fp8_kv_cache: - append_ckv = append_ckv.to(kv_dtype) - append_kpe = append_kpe.to(kv_dtype) - num_ctx_tokens = metadata.num_ctx_tokens - gen_batch_indices = metadata.batch_indices[num_ctx_tokens:] - gen_positions = metadata.positions[num_ctx_tokens:] - flashinfer.page.append_paged_mla_kv_cache( - append_ckv, - append_kpe, - gen_batch_indices, - gen_positions, - ckv_cache, - kpe_cache, - metadata.paged_kv_indices, - metadata.paged_kv_indptr, - metadata.paged_kv_last_page_len, - ) + append_ckv = latent_cache[:, :self.kv_lora_rank] + append_kpe = latent_cache[:, self.kv_lora_rank:] + if self.has_fp8_kv_cache: + append_ckv = append_ckv.to(kv_dtype) + append_kpe = append_kpe.to(kv_dtype) + num_ctx_tokens = metadata.num_ctx_tokens + gen_batch_indices = metadata.batch_indices[num_ctx_tokens:] + gen_positions = metadata.positions[num_ctx_tokens:] + flashinfer.page.append_paged_mla_kv_cache( + append_ckv, + append_kpe, + gen_batch_indices, + gen_positions, + ckv_cache, + kpe_cache, + metadata.paged_kv_indices, + metadata.paged_kv_indptr, + metadata.paged_kv_last_page_len, + ) # fused_q layout: [num_tokens, num_heads * (kv_lora_rank + qk_rope_head_dim)] # Split into q_nope (absorbed) and q_pe (rope) @@ -1279,7 +1321,7 @@ def _mla_forward_cached_context( if self.has_fp8_kv_cache: kv_dtype = torch.float8_e4m3fn - ckv_cache, kpe_cache = self._get_mla_caches(metadata, kv_dtype) + ckv_cache, kpe_cache = self._get_mla_caches(metadata) append_ckv = latent_cache[:, :self.kv_lora_rank] append_kpe = latent_cache[:, self.kv_lora_rank:] @@ -1367,6 +1409,10 @@ def forward_impl( self._mla_forward_generation(q, metadata, output, latent_cache) return + raise ValueError( + "FlashInfer MLA received an unsupported input combination: " + f"k is None={k is None}, v is None={v is None}, " + f"latent_cache is None={latent_cache is None}.") # Query q = q.view(-1, self.num_heads, self.head_dim) diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 2cd9f59afd80..fc9b5d5a6043 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -1715,19 +1715,14 @@ def forward_impl(self, latent_cache_ctx = latent_cache[:num_ctx_tokens, ...] if self.apply_rotary_emb: assert position_ids is not None - if isinstance(attn_metadata, FlashInferAttentionMetadata): - # position_ids spans [ctx..., gen...] in mixed batches; - # slice to match q_ctx/k_pe_ctx so external RoPE uses ctx - # positions. - ctx_position_ids = position_ids[..., :num_ctx_tokens] - k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, - ctx_position_ids) - # Rebuild latent_cache with RoPE'd k_pe for backends that - # don't handle fused RoPE internally (e.g., FlashInfer). - latent_cache_ctx = torch.cat([compressed_kv_ctx, k_pe_ctx], - dim=-1) - else: - k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, position_ids) + # position_ids spans [ctx..., gen...] in mixed batches; slice to + # match q_ctx/k_pe_ctx so external RoPE uses ctx positions. + ctx_position_ids = position_ids[..., :num_ctx_tokens] + k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, ctx_position_ids) + # External RoPE is only used by backends that do not handle + # fused RoPE internally, so keep latent_cache in sync. + latent_cache_ctx = torch.cat([compressed_kv_ctx, k_pe_ctx], + dim=-1) if self.llama_4_scaling: q_ctx = self._attention_scaling( @@ -1751,20 +1746,14 @@ def forward_impl(self, latent_cache_gen = latent_cache[num_ctx_tokens:, ...] if self.apply_rotary_emb: assert position_ids is not None - if isinstance(attn_metadata, FlashInferAttentionMetadata): - # position_ids spans [ctx..., gen...] in mixed batches; - # gen positions start at num_ctx_tokens. Without this - # slice the external RoPE op applied ctx positions to gen - # k_pe and poisoned the paged MLA cache. - gen_position_ids = position_ids[..., num_ctx_tokens:] - k_pe_gen = self.apply_rope(q_gen, k_pe_gen, - gen_position_ids) - # Rebuild latent_cache with RoPE'd k_pe for backends that - # don't handle fused RoPE internally (e.g., FlashInfer). - latent_cache_gen = torch.cat([compressed_kv_gen, k_pe_gen], - dim=-1) - else: - k_pe_gen = self.apply_rope(q_gen, k_pe_gen, position_ids) + # position_ids spans [ctx..., gen...] in mixed batches; gen + # positions start at num_ctx_tokens. + gen_position_ids = position_ids[..., num_ctx_tokens:] + k_pe_gen = self.apply_rope(q_gen, k_pe_gen, gen_position_ids) + # External RoPE is only used by backends that do not handle + # fused RoPE internally, so keep latent_cache in sync. + latent_cache_gen = torch.cat([compressed_kv_gen, k_pe_gen], + dim=-1) if self.llama_4_scaling: q_gen = self._attention_scaling( @@ -2418,10 +2407,7 @@ def forward_context( output: torch.Tensor, latent_cache: Optional[torch.Tensor] = None, ) -> torch.Tensor: - if (isinstance(attn_metadata, FlashInferAttentionMetadata) - and attn_metadata.kv_cache_manager is not None and any( - attn_metadata.kv_cache_params. - num_cached_tokens_per_seq[:attn_metadata.num_contexts])): + if isinstance(attn_metadata, FlashInferAttentionMetadata): return self.forward_absorption_context(q, compressed_kv, k_pe, @@ -2478,10 +2464,7 @@ def forward_absorption_generation( # fused_q contains 1) the result of the following bmm with shape [num_tokens, num_heads, kv_lora_rank] # 2) rope(q_pe) with shape [num_tokens, num_heads, qk_rope_head_dim]. rope is applied inside AttentionOp - if isinstance(attn_metadata, FlashInferAttentionMetadata): - num_seqs = attn_metadata.num_generations - else: - num_seqs = attn_metadata.kv_lens_cuda_runtime.size(0) + num_seqs = attn_metadata.kv_lens_cuda_runtime.size(0) cu_q_seqlens = torch.empty(num_seqs + 1, dtype=torch.int32, From 5f5158b03e7577c0250a423e082f0009a51eb8d5 Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Mon, 18 May 2026 01:35:42 -0700 Subject: [PATCH 06/11] [None][fix] Address remaining FlashInfer MLA review comments Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- .../_torch/attention_backend/flashinfer.py | 98 ++++++++++++++----- 1 file changed, 76 insertions(+), 22 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 47c6bd982210..ecb3db9f31f8 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -65,9 +65,9 @@ class RaggedPlanParams: @dataclass(kw_only=True, frozen=True) -class MLADecodePlanParams: +class MLAPlanParams: """ - Parameters for FlashInfer MLA decode using BatchMLAPagedAttentionWrapper. + Parameters for FlashInfer MLA using BatchMLAPagedAttentionWrapper. """ num_heads: int @@ -118,17 +118,17 @@ class FlashInferAttentionMetadata(AttentionMetadata): flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper] = field( init=False, default=None) - # MLA decode wrapper (BatchMLAPagedAttentionWrapper) and stable buffers. + # MLA wrappers (BatchMLAPagedAttentionWrapper) and stable buffers. # Cached plan params + is-planned flag let prepare() refresh the plan # outside stream capture (flashinfer plan() does device->host syncs). _mla_decode_wrapper: Optional[object] = field(init=False, default=None) _mla_context_wrapper: Optional[object] = field(init=False, default=None) _mla_ragged_plan_params: Optional[RaggedPlanParams] = field(init=False, default=None) - _mla_context_plan_params: Optional[MLADecodePlanParams] = field( - init=False, default=None) - _mla_decode_plan_params: Optional[MLADecodePlanParams] = field(init=False, - default=None) + _mla_context_plan_params: Optional[MLAPlanParams] = field(init=False, + default=None) + _mla_decode_plan_params: Optional[MLAPlanParams] = field(init=False, + default=None) _mla_ragged_planned: bool = field(init=False, default=False) _mla_context_planned: bool = field(init=False, default=False) _mla_decode_planned: bool = field(init=False, default=False) @@ -190,6 +190,18 @@ def plan_ragged( # new plan. Reusing a cached plan avoids this sync on later layers. torch.cuda.current_stream().synchronize() + self._do_plan_ragged(qo_indptr, kv_indptr, plan_params) + self._mla_ragged_planned = True + + return self._ragged_prefill_wrapper + + def _do_plan_ragged( + self, + qo_indptr: torch.Tensor, + kv_indptr: torch.Tensor, + plan_params: RaggedPlanParams, + ) -> None: + assert self._ragged_prefill_wrapper is not None self._ragged_prefill_wrapper.plan( qo_indptr, kv_indptr, @@ -203,13 +215,10 @@ def plan_ragged( kv_data_type=plan_params.kv_dtype, sm_scale=plan_params.sm_scale, ) - self._mla_ragged_planned = True - - return self._ragged_prefill_wrapper def plan_mla_decode( self, - plan_params: MLADecodePlanParams, + plan_params: MLAPlanParams, ) -> object: """Plan MLA decode using BatchMLAPagedAttentionWrapper. @@ -255,7 +264,7 @@ def plan_mla_context( kv_indptr: torch.Tensor, kv_indices: torch.Tensor, kv_last_page_len: torch.Tensor, - plan_params: MLADecodePlanParams, + plan_params: MLAPlanParams, ) -> object: """Plan MLA context with cached KV using BatchMLAPagedAttentionWrapper.""" if self._mla_context_wrapper is None: @@ -277,14 +286,35 @@ def plan_mla_context( if self._mla_context_planned: return self._mla_context_wrapper - num_pages_per_seq = kv_indptr[1:] - kv_indptr[:-1] - kv_len_arr = (num_pages_per_seq - - 1) * plan_params.page_size + kv_last_page_len - # Split append_paged_mla_kv_cache from plan() when this wrapper needs a # new plan. Reusing a cached plan avoids this sync on later layers. torch.cuda.current_stream().synchronize() + self._do_plan_mla_context( + qo_indptr, + kv_indptr, + kv_indices, + kv_last_page_len, + plan_params, + ) + self._mla_context_planned = True + + return self._mla_context_wrapper + + def _do_plan_mla_context( + self, + qo_indptr: torch.Tensor, + kv_indptr: torch.Tensor, + kv_indices: torch.Tensor, + kv_last_page_len: torch.Tensor, + plan_params: MLAPlanParams, + ) -> None: + assert self._mla_context_wrapper is not None + + num_pages_per_seq = kv_indptr[1:] - kv_indptr[:-1] + kv_len_arr = (num_pages_per_seq - + 1) * plan_params.page_size + kv_last_page_len + self._mla_context_wrapper.plan( qo_indptr, kv_indptr, @@ -299,11 +329,8 @@ def plan_mla_context( kv_data_type=plan_params.kv_dtype, sm_scale=plan_params.sm_scale, ) - self._mla_context_planned = True - return self._mla_context_wrapper - - def _do_plan_mla_decode(self, plan_params: MLADecodePlanParams) -> None: + def _do_plan_mla_decode(self, plan_params: MLAPlanParams) -> None: """Compute MLA decode plan inputs and call wrapper.plan(). Must run outside of CUDA graph capture. kv_indptr / kv_indices are @@ -827,8 +854,35 @@ def prepare(self) -> None: else: del self._plan_params_to_wrappers[plan_params] + # Re-plan MLA wrappers outside of forward/capture using the params + # cached by prior warmup forwards. Forward still handles first-use or + # dtype/shape changes by syncing only on a plan cache miss. self._mla_ragged_planned = False + if (self.num_contexts > 0 and self._mla_ragged_plan_params is not None + and self._ragged_prefill_wrapper is not None): + ragged_indptr = self.qo_indptr[:self.num_contexts + 1] + self._do_plan_ragged(ragged_indptr, ragged_indptr, + self._mla_ragged_plan_params) + self._mla_ragged_planned = True + self._mla_context_planned = False + if (self.num_contexts > 0 and self._mla_context_plan_params is not None + and self._mla_context_wrapper is not None): + num_contexts = self.num_contexts + num_context_blocks = self.num_context_blocks + context_qo_indptr = self.qo_indptr[:num_contexts + 1] + context_kv_indptr = self.paged_kv_indptr_prefill[:num_contexts + 1] + context_kv_indices = self._paged_kv_indices[:num_context_blocks] + context_last_page_len = self._paged_kv_last_page_len[:num_contexts] + self._do_plan_mla_context( + qo_indptr=context_qo_indptr, + kv_indptr=context_kv_indptr, + kv_indices=context_kv_indices, + kv_last_page_len=context_last_page_len, + plan_params=self._mla_context_plan_params, + ) + self._mla_context_planned = True + # Re-plan the MLA decode wrapper outside of any stream capture. if (self.num_generations > 0 and self._mla_decode_plan_params is not None @@ -1288,7 +1342,7 @@ def _mla_forward_generation( else: sm_scale = 1.0 / math.sqrt(qk_head_dim) - plan_params = MLADecodePlanParams( + plan_params = MLAPlanParams( num_heads=self.num_heads, kv_lora_rank=self.kv_lora_rank, qk_rope_head_dim=self.qk_rope_head_dim, @@ -1352,7 +1406,7 @@ def _mla_forward_cached_context( else: sm_scale = 1.0 / math.sqrt(qk_head_dim) - plan_params = MLADecodePlanParams( + plan_params = MLAPlanParams( num_heads=self.num_heads, kv_lora_rank=self.kv_lora_rank, qk_rope_head_dim=self.qk_rope_head_dim, From 033420f6a1eb7020def03652b885a7c9e1df3134 Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Tue, 19 May 2026 20:07:20 -0700 Subject: [PATCH 07/11] Route FlashInfer MLA normal prefill through default path Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- tensorrt_llm/_torch/modules/attention.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index fc9b5d5a6043..589ba860dd72 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -2408,13 +2408,15 @@ def forward_context( latent_cache: Optional[torch.Tensor] = None, ) -> torch.Tensor: if isinstance(attn_metadata, FlashInferAttentionMetadata): - return self.forward_absorption_context(q, - compressed_kv, - k_pe, - attn_metadata, - output, - position_ids=position_ids, - latent_cache=latent_cache) + if attn_metadata.runtime_features.chunked_prefill: + return self.forward_absorption_context( + q, + compressed_kv, + k_pe, + attn_metadata, + output, + position_ids=position_ids, + latent_cache=latent_cache) if isinstance(self.mha, TrtllmAttention): assert isinstance(attn_metadata, TrtllmAttentionMetadata) From bcc19e5a1bb132b8759ea2ec5f5d9060174c48bf Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Tue, 19 May 2026 21:31:45 -0700 Subject: [PATCH 08/11] Fix FlashInfer MLA cached context dispatch Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- .../_torch/attention_backend/flashinfer.py | 34 +++++++++++++++---- tensorrt_llm/_torch/modules/attention.py | 5 ++- 2 files changed, 32 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index ecb3db9f31f8..78df0a7784c3 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -129,6 +129,7 @@ class FlashInferAttentionMetadata(AttentionMetadata): default=None) _mla_decode_plan_params: Optional[MLAPlanParams] = field(init=False, default=None) + num_ctx_cached_tokens: int = field(init=False, default=0) _mla_ragged_planned: bool = field(init=False, default=False) _mla_context_planned: bool = field(init=False, default=False) _mla_decode_planned: bool = field(init=False, default=False) @@ -689,6 +690,7 @@ def prepare(self) -> None: self.kv_cache_params = KVCacheParams(use_cache=False) n = self.num_seqs self._cached_token_lens[:n].zero_() + self.num_ctx_cached_tokens = 0 self.kv_lens_cuda_runtime = self.seq_lens_cuda[:n] for plan_params in list(self._plan_params_to_wrappers.keys()): if plan_params.attention_mask_data is None: @@ -709,6 +711,12 @@ def prepare(self) -> None: self.kv_cache_params.num_cached_tokens_per_seq, dtype=torch.int) self._cached_token_lens[:cached_token_lens.size(0)].copy_( cached_token_lens, non_blocking=True) + if self.num_contexts > 0: + self.num_ctx_cached_tokens = sum( + self.kv_cache_params.num_cached_tokens_per_seq[:self. + num_contexts]) + else: + self.num_ctx_cached_tokens = 0 # number of tokens needed in the kv cache for each sequence after the next pass kv_lens = self.cached_token_lens + self.seq_lens_kv_cuda @@ -1362,14 +1370,14 @@ def _mla_forward_generation( out=output[:num_tokens].view(-1, self.num_heads, self.kv_lora_rank)) - def _mla_forward_cached_context( + def _mla_forward_paged_context( self, q: torch.Tensor, metadata: FlashInferAttentionMetadata, output: torch.Tensor, latent_cache: torch.Tensor, ) -> None: - """MLA context phase with cached KV: append latent and run paged MLA.""" + """MLA context phase with paged KV: append latent and run paged MLA.""" num_ctx_tokens = metadata.num_ctx_tokens kv_dtype = q.dtype if self.has_fp8_kv_cache: @@ -1453,11 +1461,25 @@ def forward_impl( latent_cache) return elif k is None and v is None: - if attention_input_type == AttentionInputType.context_only: + has_cached_context = ( + attention_input_type == AttentionInputType.context_only + and metadata.enable_context_mla_with_cached_kv + and metadata.num_ctx_cached_tokens > 0) + has_first_chunk_context = ( + attention_input_type == AttentionInputType.context_only + and metadata.enable_context_mla_with_cached_kv + and metadata.num_ctx_cached_tokens == 0) + if has_cached_context or has_first_chunk_context: + # Context MLA with cached KV uses paged MLA. The first + # chunk has no cached tokens yet, but still uses this path. assert latent_cache is not None, ( - "FlashInfer MLA cached context requires latent_cache.") - self._mla_forward_cached_context(q, metadata, output, - latent_cache) + "FlashInfer MLA paged context requires latent_cache.") + self._mla_forward_paged_context(q, metadata, output, + latent_cache) + elif attention_input_type == AttentionInputType.context_only: + raise ValueError( + "FlashInfer MLA context without cached KV " + "requires key/value tensors.") else: # MLA generation phase: paged decode + slice self._mla_forward_generation(q, metadata, output, diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 589ba860dd72..9d9f1929efe8 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -2408,7 +2408,10 @@ def forward_context( latent_cache: Optional[torch.Tensor] = None, ) -> torch.Tensor: if isinstance(attn_metadata, FlashInferAttentionMetadata): - if attn_metadata.runtime_features.chunked_prefill: + if attn_metadata.enable_context_mla_with_cached_kv: + # The zero-cached first chunk still needs FlashInfer's paged + # MLA context path; the backend checks num_ctx_cached_tokens + # when distinguishing cached vs first-chunk context. return self.forward_absorption_context( q, compressed_kv, From 7dcf29ad9d35543c02ea2c20dc593b4539d1442d Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Tue, 19 May 2026 21:40:05 -0700 Subject: [PATCH 09/11] Address FlashInfer MLA review comments Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- .../_torch/attention_backend/flashinfer.py | 44 +++++++++---------- tensorrt_llm/_torch/modules/attention.py | 2 +- .../_torch/attention/test_attention_mla.py | 10 +++-- 3 files changed, 27 insertions(+), 29 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 78df0a7784c3..2d432c18121d 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -106,21 +106,15 @@ class FlashInferAttentionMetadata(AttentionMetadata): _qo_indptr: torch.Tensor = field(init=False) _kv_indptr: torch.Tensor = field(init=False) _cached_token_lens: torch.Tensor = field(init=False) - kv_lens_cuda_runtime: Optional[torch.Tensor] = field(init=False, - default=None, - repr=False) - _plan_params_to_wrappers: Dict[PlanParams, FlashInferWrappers] = field(init=False) - # MLA ragged prefill wrapper (for context phase with expanded K, V) + # MLA wrappers and stable buffers. + # Cached plan params + is-planned flag let prepare() refresh the plan + # outside stream capture (flashinfer plan() does device->host syncs). _ragged_prefill_wrapper: Optional[ flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper] = field( init=False, default=None) - - # MLA wrappers (BatchMLAPagedAttentionWrapper) and stable buffers. - # Cached plan params + is-planned flag let prepare() refresh the plan - # outside stream capture (flashinfer plan() does device->host syncs). _mla_decode_wrapper: Optional[object] = field(init=False, default=None) _mla_context_wrapper: Optional[object] = field(init=False, default=None) _mla_ragged_plan_params: Optional[RaggedPlanParams] = field(init=False, @@ -187,6 +181,12 @@ def plan_ragged( if self._mla_ragged_planned: return self._ragged_prefill_wrapper + if self.is_cuda_graph and torch.cuda.is_current_stream_capturing(): + raise ValueError( + "Cannot plan() flashinfer MLA ragged prefill while the stream " + "is capturing. Make sure prepare() has run at least one " + "warmup forward pass before capture.") + # Split append_paged_mla_kv_cache from plan() when this wrapper needs a # new plan. Reusing a cached plan avoids this sync on later layers. torch.cuda.current_stream().synchronize() @@ -275,11 +275,6 @@ def plan_mla_context( backend="auto", ) - if self.is_cuda_graph and torch.cuda.is_current_stream_capturing(): - raise ValueError( - "Cannot plan() flashinfer MLA context while the stream is " - "capturing. Chunked MLA prefill with FlashInfer does not " - "support CUDA graph capture.") if self._mla_context_plan_params != plan_params: self._mla_context_planned = False self._mla_context_plan_params = plan_params @@ -287,6 +282,12 @@ def plan_mla_context( if self._mla_context_planned: return self._mla_context_wrapper + if self.is_cuda_graph and torch.cuda.is_current_stream_capturing(): + raise ValueError( + "Cannot plan() flashinfer MLA context while the stream is " + "capturing. Make sure prepare() has run at least one warmup " + "forward pass before capture.") + # Split append_paged_mla_kv_cache from plan() when this wrapper needs a # new plan. Reusing a cached plan avoids this sync on later layers. torch.cuda.current_stream().synchronize() @@ -334,15 +335,13 @@ def _do_plan_mla_context( def _do_plan_mla_decode(self, plan_params: MLAPlanParams) -> None: """Compute MLA decode plan inputs and call wrapper.plan(). - Must run outside of CUDA graph capture. kv_indptr / kv_indices are - cloned because they alias the wrapper's own buffers and - flashinfer.plan() would otherwise do a self-copy. + Must run outside of CUDA graph capture. """ num_gen = self.num_generations kv_indptr = self.paged_kv_indptr_decode[:num_gen + 1] kv_indices = self._paged_kv_indices[self.num_context_blocks:self. num_context_blocks + - self.num_generation_blocks].clone() + self.num_generation_blocks] kv_last_page = self._paged_kv_last_page_len[self.num_contexts:self. num_contexts + num_gen] @@ -356,8 +355,6 @@ def _do_plan_mla_decode(self, plan_params: MLAPlanParams) -> None: kv_len_arr = (num_pages_per_seq - 1) * plan_params.page_size + kv_last_page - kv_indptr = kv_indptr.clone() - self._mla_decode_wrapper.plan( qo_indptr, kv_indptr, @@ -564,6 +561,7 @@ def _post_init_with_buffers(self, buffers) -> None: capture_graph=capture_graph, ) # Rebind the wrapper to the freshly allocated buffers. + self._ragged_prefill_wrapper = None self._mla_decode_wrapper = None self._mla_context_wrapper = None self._mla_ragged_plan_params = None @@ -691,7 +689,6 @@ def prepare(self) -> None: n = self.num_seqs self._cached_token_lens[:n].zero_() self.num_ctx_cached_tokens = 0 - self.kv_lens_cuda_runtime = self.seq_lens_cuda[:n] for plan_params in list(self._plan_params_to_wrappers.keys()): if plan_params.attention_mask_data is None: self._plan_params_to_wrappers[ @@ -720,7 +717,6 @@ def prepare(self) -> None: # number of tokens needed in the kv cache for each sequence after the next pass kv_lens = self.cached_token_lens + self.seq_lens_kv_cuda - self.kv_lens_cuda_runtime = kv_lens[:self.num_seqs] # start and end indices of each sequence in the ragged key and value # for self attention it's the same as qo_indptr so avoid computing twice. @@ -866,6 +862,8 @@ def prepare(self) -> None: # cached by prior warmup forwards. Forward still handles first-use or # dtype/shape changes by syncing only on a plan cache miss. self._mla_ragged_planned = False + self._mla_context_planned = False + self._mla_decode_planned = False if (self.num_contexts > 0 and self._mla_ragged_plan_params is not None and self._ragged_prefill_wrapper is not None): ragged_indptr = self.qo_indptr[:self.num_contexts + 1] @@ -873,7 +871,6 @@ def prepare(self) -> None: self._mla_ragged_plan_params) self._mla_ragged_planned = True - self._mla_context_planned = False if (self.num_contexts > 0 and self._mla_context_plan_params is not None and self._mla_context_wrapper is not None): num_contexts = self.num_contexts @@ -895,7 +892,6 @@ def prepare(self) -> None: if (self.num_generations > 0 and self._mla_decode_plan_params is not None and self._mla_decode_wrapper is not None): - self._mla_decode_planned = False self._do_plan_mla_decode(self._mla_decode_plan_params) self._mla_decode_planned = True diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 9d9f1929efe8..9414e0fb2837 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -2469,7 +2469,7 @@ def forward_absorption_generation( # fused_q contains 1) the result of the following bmm with shape [num_tokens, num_heads, kv_lora_rank] # 2) rope(q_pe) with shape [num_tokens, num_heads, qk_rope_head_dim]. rope is applied inside AttentionOp - num_seqs = attn_metadata.kv_lens_cuda_runtime.size(0) + num_seqs = attn_metadata.num_seqs cu_q_seqlens = torch.empty(num_seqs + 1, dtype=torch.int32, diff --git a/tests/unittest/_torch/attention/test_attention_mla.py b/tests/unittest/_torch/attention/test_attention_mla.py index 13e7362e56dd..bd01ed363dc4 100644 --- a/tests/unittest/_torch/attention/test_attention_mla.py +++ b/tests/unittest/_torch/attention/test_attention_mla.py @@ -782,7 +782,10 @@ def yarn_get_mscale(scale=1, mscale=1): max_num_requests=len(context_sequence_lengths), num_contexts=0, prompt_lens=context_sequence_lengths, - max_num_tokens=sum(context_sequence_lengths), + max_num_tokens=max( + sum(context_sequence_lengths), + generation_seq_len_q * len(context_sequence_lengths), + ), kv_cache_manager=kv_cache_manager, kv_cache_params=KVCacheParams( use_cache=True, @@ -794,9 +797,8 @@ def yarn_get_mscale(scale=1, mscale=1): mapping=mapping, ) if backend_name == "TRTLLM": - gen_metadata_kwargs[ - 'enable_flash_mla'] = torch.cuda.get_device_capability( - ) == (9, 0) + gen_metadata_kwargs["enable_flash_mla"] = ( + torch.cuda.get_device_capability() == (9, 0)) attn_metadata = AttentionCls.Metadata(**gen_metadata_kwargs) attn_metadata.prepare() for layer_idx in range(num_layers): From 1545fbcad9451e5a8723aa901f3a1c6ce1d7a127 Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Tue, 19 May 2026 21:56:34 -0700 Subject: [PATCH 10/11] Route first MLA context chunk through normal prefill Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- tensorrt_llm/_torch/attention_backend/flashinfer.py | 9 ++------- tensorrt_llm/_torch/modules/attention.py | 6 ++---- 2 files changed, 4 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 2d432c18121d..6a1936a1c4cb 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -1461,13 +1461,8 @@ def forward_impl( attention_input_type == AttentionInputType.context_only and metadata.enable_context_mla_with_cached_kv and metadata.num_ctx_cached_tokens > 0) - has_first_chunk_context = ( - attention_input_type == AttentionInputType.context_only - and metadata.enable_context_mla_with_cached_kv - and metadata.num_ctx_cached_tokens == 0) - if has_cached_context or has_first_chunk_context: - # Context MLA with cached KV uses paged MLA. The first - # chunk has no cached tokens yet, but still uses this path. + if has_cached_context: + # Context MLA with cached KV uses paged MLA. assert latent_cache is not None, ( "FlashInfer MLA paged context requires latent_cache.") self._mla_forward_paged_context(q, metadata, output, diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 9414e0fb2837..11874b812802 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -2408,10 +2408,8 @@ def forward_context( latent_cache: Optional[torch.Tensor] = None, ) -> torch.Tensor: if isinstance(attn_metadata, FlashInferAttentionMetadata): - if attn_metadata.enable_context_mla_with_cached_kv: - # The zero-cached first chunk still needs FlashInfer's paged - # MLA context path; the backend checks num_ctx_cached_tokens - # when distinguishing cached vs first-chunk context. + if (attn_metadata.enable_context_mla_with_cached_kv + and attn_metadata.num_ctx_cached_tokens > 0): return self.forward_absorption_context( q, compressed_kv, From 0d71ba3dcc01a24b3037e49f867b261815cf5511 Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Tue, 19 May 2026 22:02:31 -0700 Subject: [PATCH 11/11] Fix FlashInfer MLA error formatting Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- tensorrt_llm/_torch/attention_backend/flashinfer.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 6a1936a1c4cb..8fd9d60baa48 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -1468,9 +1468,8 @@ def forward_impl( self._mla_forward_paged_context(q, metadata, output, latent_cache) elif attention_input_type == AttentionInputType.context_only: - raise ValueError( - "FlashInfer MLA context without cached KV " - "requires key/value tensors.") + raise ValueError("FlashInfer MLA context without cached KV " + "requires key/value tensors.") else: # MLA generation phase: paged decode + slice self._mla_forward_generation(q, metadata, output,