diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index f959e08203f2..49964319a18a 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, - PredefinedAttentionMask, merge_attention_forward_args) + AttentionInputType, AttentionMetadata, + CustomAttentionMask, MLAParams, PredefinedAttentionMask, + merge_attention_forward_args) try: check_cuda_arch() @@ -48,6 +49,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 MLAPlanParams: + """ + Parameters for FlashInfer MLA 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 @@ -75,10 +106,31 @@ 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) - _plan_params_to_wrappers: Dict[PlanParams, FlashInferWrappers] = field(init=False) + # 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_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[MLAPlanParams] = field(init=False, + 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) + _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 +162,214 @@ 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", + ) + 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 + + 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() + + 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, + 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, + ) + + def plan_mla_decode( + self, + plan_params: MLAPlanParams, + ) -> 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", + ) + + 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 + + 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.") + + # 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 + + 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: MLAPlanParams, + ) -> 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._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 + + 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() + + 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, + 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, + ) + + 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. + """ + 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] + 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 + + 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 +545,31 @@ 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._ragged_prefill_wrapper = None + 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, max_batch_size: int, @@ -403,6 +688,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 for plan_params in list(self._plan_params_to_wrappers.keys()): if plan_params.attention_mask_data is None: self._plan_params_to_wrappers[ @@ -422,6 +708,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 @@ -566,6 +858,43 @@ 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 + 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] + self._do_plan_ragged(ragged_indptr, ragged_indptr, + self._mla_ragged_plan_params) + self._mla_ragged_planned = True + + 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 + and self._mla_decode_wrapper is not None): + 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): @@ -613,7 +942,6 @@ def prepare(self) -> None: non_blocking=True) if self.num_generations < bs: kv_lens_buf[self.num_generations:bs].zero_() - if self.cross is not None and self.cross is not self: self.cross.prepare() @@ -802,6 +1130,10 @@ class FlashInferAttention(AttentionBackend[FlashInferAttentionMetadata]): Metadata = FlashInferAttentionMetadata + @classmethod + def support_mla(cls) -> bool: + return True + def __init__( self, layer_idx: int, @@ -811,6 +1143,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 +1153,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 +1167,275 @@ 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", + ): + """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) + + 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) + + 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. + 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 = MLAPlanParams( + 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 _mla_forward_paged_context( + self, + q: torch.Tensor, + metadata: FlashInferAttentionMetadata, + output: torch.Tensor, + latent_cache: torch.Tensor, + ) -> None: + """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: + kv_dtype = torch.float8_e4m3fn + + 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:] + 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 = MLAPlanParams( + 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, @@ -837,7 +1446,40 @@ 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, + attention_input_type: AttentionInputType = AttentionInputType.mixed, ) -> 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: + has_cached_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: + # 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, + 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, + 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) @@ -1021,6 +1663,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) @@ -1037,7 +1680,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. @@ -1045,12 +1697,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) + 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 af9662f95670..74ef299532ce 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" @@ -1721,7 +1711,14 @@ 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) + # 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( @@ -1745,7 +1742,14 @@ 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) + # 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( @@ -2399,6 +2403,18 @@ def forward_context( output: torch.Tensor, latent_cache: Optional[torch.Tensor] = None, ) -> torch.Tensor: + if isinstance(attn_metadata, FlashInferAttentionMetadata): + 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, + 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) @@ -2447,7 +2463,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/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index ce917d723cc9..6a0befbbbdd5 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -1755,6 +1755,18 @@ 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) + @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", + 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) + @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/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 5fc18da8b40a..e0e47957c352 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -226,6 +226,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_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] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[ep4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=True-torch_compile=False] diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 490820cc2c1c..66892a382e49 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -32,6 +32,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=False] - 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=False] - 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_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 d0ac7ed986c5..063125788ed9 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -59,3 +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[enable_chunked_prefill=False] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=True] diff --git a/tests/unittest/_torch/attention/test_attention_mla.py b/tests/unittest/_torch/attention/test_attention_mla.py index a634175edc6e..bd01ed363dc4 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,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=max(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, @@ -686,8 +795,11 @@ 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 +814,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 +874,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,