diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 4633850204b6..aac79746e665 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -43,6 +43,80 @@ _FORCE_RAGGED_FA2 = False """Used for testing.""" +_MAX_CUDA_THREADS_PER_BLOCK = 1024 + + +def _slice_paged_kv_cache_heads( + paged_kv_cache: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + start: int, + end: int, + kv_layout: str, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if kv_layout == "HND": + head_axis = 2 + elif kv_layout == "NHD": + head_axis = 3 + else: + raise ValueError(f"Unsupported kv_layout: {kv_layout}") + + if isinstance(paged_kv_cache, tuple): + head_axis -= 1 + index = [slice(None)] * 4 + index[head_axis] = slice(start, end) + return tuple(cache[tuple(index)] for cache in paged_kv_cache) + + index = [slice(None)] * 5 + index[head_axis] = slice(start, end) + return paged_kv_cache[tuple(index)] + + +def _append_paged_kv_cache( + append_key: torch.Tensor, + append_value: torch.Tensor, + batch_indices: torch.Tensor, + positions: torch.Tensor, + paged_kv_cache: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + kv_indices: torch.Tensor, + kv_indptr: torch.Tensor, + kv_last_page_len: torch.Tensor, + kv_layout: str = "NHD", +) -> None: + """Split FlashInfer paged-KV appends that exceed CUDA's CTA limit.""" + head_dim = append_key.shape[-1] + vec_size = max(16 // append_key.element_size(), head_dim // 32) + threads_per_head = head_dim // vec_size + max_heads_per_launch = _MAX_CUDA_THREADS_PER_BLOCK // threads_per_head + + num_kv_heads = append_key.shape[1] + if num_kv_heads <= max_heads_per_launch: + flashinfer.page.append_paged_kv_cache( + append_key=append_key, + append_value=append_value, + batch_indices=batch_indices, + positions=positions, + paged_kv_cache=paged_kv_cache, + kv_indices=kv_indices, + kv_indptr=kv_indptr, + kv_last_page_len=kv_last_page_len, + kv_layout=kv_layout, + ) + return + + for start in range(0, num_kv_heads, max_heads_per_launch): + end = min(start + max_heads_per_launch, num_kv_heads) + flashinfer.page.append_paged_kv_cache( + append_key=append_key[:, start:end], + append_value=append_value[:, start:end], + batch_indices=batch_indices, + positions=positions, + paged_kv_cache=_slice_paged_kv_cache_heads(paged_kv_cache, start, + end, kv_layout), + kv_indices=kv_indices, + kv_indptr=kv_indptr, + kv_last_page_len=kv_last_page_len, + kv_layout=kv_layout, + ) + @dataclass(kw_only=True, frozen=True) class FlashInferMultiItemParams: @@ -1748,7 +1822,7 @@ def forward_impl( f"KV cache dtype {kv_cache.dtype} does not match k/v dtype {k.dtype}/{v.dtype}" ) - flashinfer.page.append_paged_kv_cache( + _append_paged_kv_cache( append_key=k, append_value=v, batch_indices=metadata.batch_indices, diff --git a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py index 0a0d89dd5adf..d39ecaf61865 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py @@ -36,6 +36,7 @@ "attention_mask_data", # custom-mask code path "out_scale_sf", # promoted into ``out_scale`` in ``TrtllmAttention.forward`` for NVFP4 path "skip_mla_rope_generation", # handled in ``TrtllmAttention.forward`` for the test-only MLA path + "timestep", # used to populate skip-softmax params in ``TrtllmAttention.forward`` } ) diff --git a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py index 795e1766edc1..fe4fa839ad8d 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py @@ -408,9 +408,6 @@ class FlashInferTrtllmGenFmha(PhasedFmha): (320, 256), (576, 512), } - MISSING_MLA_GENERATION_KERNELS = { - (576, 512, 32), - } def __init__(self, attn: "TrtllmAttention"): super().__init__(attn) @@ -529,7 +526,6 @@ def _get_attention_chunk_size(attn: "TrtllmAttention") -> int: def _check_mla_generation_support( cls, head_size: int, - tokens_per_block: int, kv_lora_rank: Optional[int], qk_rope_head_dim: Optional[int], ) -> Tuple[bool, str]: @@ -567,14 +563,6 @@ def _check_mla_generation_support( f"headDimQk={head_dim_qk}, headDimV={head_dim_v}. Supported: {supported}.", ) - if (head_dim_qk, head_dim_v, tokens_per_block) in cls.MISSING_MLA_GENERATION_KERNELS: - return ( - False, - f"[Generation][MLA] Missing TRTLLM-GEN decode kernel for " - f"headDimQk={head_dim_qk}, headDimV={head_dim_v}, " - f"tokens_per_block={tokens_per_block}.", - ) - return True, "" def is_supported( @@ -765,7 +753,6 @@ def _is_supported_with_reason( if is_mla_enable: supported, reason = self._check_mla_generation_support( head_size=attn.head_dim, - tokens_per_block=tokens_per_block, kv_lora_rank=attn.kv_lora_rank, qk_rope_head_dim=attn.qk_rope_head_dim, ) @@ -1192,7 +1179,7 @@ def run_mla_generation( batch_beam = params.num_requests * meta.beam_width if params.attention_input is None: raise RuntimeError("MLA generation requires attention_input.") - kv_cache, block_tables = thop.build_trtllm_gen_kv_cache_metadata( + kv_cache, block_tables, _kv_scale_pool = thop.build_trtllm_gen_kv_cache_metadata( meta.host_kv_cache_pool_pointers, # host_kv_cache_pool_pointers meta.host_kv_cache_pool_mapping, # host_kv_cache_pool_mapping meta.kv_cache_block_offsets, # kv_cache_block_offsets @@ -1222,9 +1209,43 @@ def run_mla_generation( mla_head_dim_qk = kv_lora_rank + qk_rope_head_dim q_len_per_req = params.num_tokens // batch_beam if batch_beam > 0 else 1 - query = params.qkv_input.view(batch_beam, q_len_per_req, attn.num_heads, mla_head_dim_qk) + if QuantMode(attn.quant_mode).has_fp8_kv_cache(): + quant_q_buffer = fwd.quant_q_buffer + bmm1_scale_buffer = fwd.mla_bmm1_scale + bmm2_scale_buffer = fwd.mla_bmm2_scale + if quant_q_buffer is None or bmm1_scale_buffer is None or bmm2_scale_buffer is None: + raise RuntimeError( + "FP8 MLA generation requires quant_q_buffer, " + "mla_bmm1_scale, and mla_bmm2_scale." + ) - bmm1_scale = 1.0 / (attn.q_scaling * math.sqrt(qk_nope_head_dim + qk_rope_head_dim)) + expected_q_elements = params.num_tokens * attn.num_heads * mla_head_dim_qk + if quant_q_buffer.numel() < expected_q_elements: + raise RuntimeError( + f"FP8 MLA quant_q_buffer has {quant_q_buffer.numel()} elements; " + f"expected at least {expected_q_elements}." + ) + if bmm1_scale_buffer.dtype != torch.float32 or bmm1_scale_buffer.numel() < 1: + raise RuntimeError("FP8 MLA bmm1 scale must contain a float32 value.") + if bmm2_scale_buffer.dtype != torch.float32 or bmm2_scale_buffer.numel() < 1: + raise RuntimeError("FP8 MLA bmm2 scale must contain a float32 value.") + + query = ( + quant_q_buffer.view(torch.uint8) + .flatten()[:expected_q_elements] + .view(torch.float8_e4m3fn) + .view(batch_beam, q_len_per_req, attn.num_heads, mla_head_dim_qk) + ) + # FlashInfer converts tensor BMM1 scales to log2 internally. The + # producer stores the regular scale at index 0 and log2 at index 1. + bmm1_scale = bmm1_scale_buffer.flatten()[:1] + bmm2_scale = bmm2_scale_buffer.flatten()[:1] + else: + query = params.qkv_input.view( + batch_beam, q_len_per_req, attn.num_heads, mla_head_dim_qk + ) + bmm1_scale = 1.0 / (attn.q_scaling * math.sqrt(qk_nope_head_dim + qk_rope_head_dim)) + bmm2_scale = 1.0 workspace_buffer = params.workspace.view(-1, 4) _clear_multi_ctas_kv_counter_workspace( workspace_buffer, attn.num_heads, meta.max_num_requests, self._multi_processor_count @@ -1243,7 +1264,7 @@ def run_mla_generation( 0, # sparse_mla_top_k params.context_buf.view(batch_beam, q_len_per_req, attn.num_heads, kv_lora_rank), # out bmm1_scale, # bmm1_scale - 1.0, # bmm2_scale + bmm2_scale, # bmm2_scale fwd.attention_sinks, # sinks None, # skip_softmax_threshold_scale_factor self._enable_pdl, # enable_pdl diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 420fd14a1fe0..bb3c8859b9fe 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -1312,11 +1312,6 @@ def use_nvfp4_output( or metadata.runtime_features.has_speculative_draft_tokens ) if metadata.runtime_features else False - # This is a workaround for https://nvbugs/5624818 - # Paged context FMHA is forced on SM90 for correctness - if get_sm_version() == 90: - use_paged_context_fmha = True - return self._is_nvfp4_output_kernel_available( tokens_per_block=metadata.tokens_per_block, attention_mask=attention_mask, @@ -1487,17 +1482,6 @@ def forward( # Cross-attention uses the THOP path; the trtllm-gen backend API does # not carry encoder K/V tensors yet. - # cpp/tensorrt_llm/thop/attentionOp.cpp enables mFP8ContextFMHA for an - # FP8 KV cache only when use_paged_context_fmha is true. Force paged - # context so QKV preprocessing and context FMHA use the FP8 path. - if self.has_fp8_kv_cache: - metadata.use_paged_context_fmha = True - - # SM90 forces `use_paged_context_fmha` on for correctness - # (https://nvbugs/5624818). - if get_sm_version() == 90: - metadata.use_paged_context_fmha = True - # Sparse mqa/gqa attention uses generation kernel which reads Q from qPtr (separate buffer). # Force paged context FMHA so QKV preprocessing writes Q to q_buf_2_. if (self.sparse_params is not None and getattr( diff --git a/tests/unittest/_torch/attention/backend_capability.py b/tests/unittest/_torch/attention/backend_capability.py index ab117f9ddebb..91d11be455df 100644 --- a/tests/unittest/_torch/attention/backend_capability.py +++ b/tests/unittest/_torch/attention/backend_capability.py @@ -68,20 +68,9 @@ } _FLASHINFER_PAGED_UNSUPPORTED_HEAD_DIMS = (96, 512) -_FLASHINFER_PAGED_APPEND_OVER_1024_THREADS = ( - # bf16/fp16 append launches 16 threads/head for head_dim=128. With 128 KV - # heads, that exceeds CUDA's 1024 threads/block launch limit. - (128, 128, "bfloat16"), - (128, 128, "float16"), -) _TRTLLM_PAGED_UNSUPPORTED_HEAD_DIMS = (512,) _TRTLLM_BLACKWELL_PAGED_UNSUPPORTED_HEAD_DIMS = (96,) -_TRTLLM_BLACKWELL_SLIDING_DECODE_UNSTABLE = ( - # Gemma3-27B local layers and Gemma4-31B sliding layers. - (32, 16, 128, 1024), - (32, 16, 256, 1024), -) def required_features(case) -> set: @@ -133,27 +122,16 @@ def unsupported_reason(backend: str, case) -> Optional[str]: if kv_layout is not None and kv_layout not in caps.get("kv_layouts", ()): return f"{backend} does not support kv_layout '{kv_layout}'" - # FlashInfer's standard paged path has shape-specific kernel limits in this - # suite. Keep the skip list explicit so newly supported shapes are not - # hidden by a broad dispatch-set predicate. + # FlashInfer's standard paged kernels reject these head dimensions. Keep + # the skip list explicit so newly supported shapes are not hidden by a + # broad dispatch-set predicate. if ( backend == "FLASHINFER" and getattr(case, "cache", "paged") != "none" and not getattr(case, "is_mla", False) + and case.head_dim in _FLASHINFER_PAGED_UNSUPPORTED_HEAD_DIMS ): - kv_dtype = getattr(case, "kv_dtype", None) or getattr(case, "dtype", None) - if case.head_dim in _FLASHINFER_PAGED_UNSUPPORTED_HEAD_DIMS: - return f"FLASHINFER paged attention is unstable for head_dim {case.head_dim}" - if ( - case.head_dim, - case.num_kv_heads, - kv_dtype, - ) in _FLASHINFER_PAGED_APPEND_OVER_1024_THREADS: - return ( - "FLASHINFER paged KV append exceeds CUDA's 1024 threads/block " - f"limit for head_dim={case.head_dim}, " - f"num_kv_heads={case.num_kv_heads}, kv_dtype={kv_dtype}" - ) + return f"FLASHINFER paged kernels do not support head_dim {case.head_dim}" # TRTLLM's standard paged FMHA/MMHA kernels in this build do not cover the # Gemma4 head_dim 512 path. @@ -163,7 +141,7 @@ def unsupported_reason(backend: str, case) -> Optional[str]: and not getattr(case, "is_mla", False) and case.head_dim in _TRTLLM_PAGED_UNSUPPORTED_HEAD_DIMS ): - return f"TRTLLM paged attention is unstable for head_dim {case.head_dim}" + return f"TRTLLM paged attention does not support head_dim {case.head_dim}" # TRTLLM's Blackwell paged fallback path aborts for the Phi-3 head_dim 96 # shape. Hopper covers that context config, but Blackwell must skip it @@ -175,31 +153,7 @@ def unsupported_reason(backend: str, case) -> Optional[str]: and not getattr(case, "is_mla", False) and case.head_dim in _TRTLLM_BLACKWELL_PAGED_UNSUPPORTED_HEAD_DIMS ): - return f"TRTLLM Blackwell paged attention is unstable for head_dim {case.head_dim}" - - # TRTLLM's Blackwell pure-decode sliding-window path is numerically unstable - # for the listed Gemma GQA shapes. Mixed batches still use a different path - # and match the golden. - if ( - backend == "TRTLLM" - and sm >= 100 - and getattr(case, "sliding_window", None) is not None - and getattr(case, "cache", "paged") != "none" - and not getattr(case, "is_mla", False) - and case.num_contexts == 0 - and ( - case.num_heads, - case.num_kv_heads, - case.head_dim, - case.sliding_window, - ) - in _TRTLLM_BLACKWELL_SLIDING_DECODE_UNSTABLE - ): - return ( - "TRTLLM Blackwell sliding-window pure decode is unstable for " - f"num_heads={case.num_heads}, num_kv_heads={case.num_kv_heads}, " - f"head_dim={case.head_dim}, window={case.sliding_window}" - ) + return f"TRTLLM Blackwell paged attention does not support head_dim {case.head_dim}" # TRTLLM's Blackwell no-cache fallback mismatches the Vanilla golden for the # Qwen2-VL vision tower's head_dim 80 workload; other no-cache head dims in @@ -213,74 +167,4 @@ def unsupported_reason(backend: str, case) -> Optional[str]: ): return "TRTLLM Blackwell no-cache fallback is unstable for head_dim 80" - # TRTLLM MLA context is validated by test_attention_mla.py through the - # production MLA module path. This standalone backend harness feeds already - # up-projected random K/V tensors, which does not match the TRTLLM MLA - # context op contract and produces invalid output, while Vanilla/FlashInfer - # can still validate the up-projected context math and cache append. - if ( - backend == "TRTLLM" - and getattr(case, "is_mla", False) - and getattr(case, "num_contexts", 0) == len(getattr(case, "seq_lens", ())) - ): - return ( - "TRTLLM MLA context is covered by test_attention_mla.py via the " - "production MLA module path; standalone up-projected backend " - "inputs are validated with Vanilla/FlashInfer" - ) - - # On Blackwell, TRTLLM-Gen is the supported MLA generation path. The current - # kernel set explicitly lacks the DeepSeek-family decode shape at page_size - # 32, and the generic fallback then tries to JIT an unsupported generated - # kernel. FlashInfer/Vanilla still validate these standalone MLA gen cases. - if ( - backend == "TRTLLM" - and sm >= 100 - and getattr(case, "is_mla", False) - and getattr(case, "num_contexts", 0) == 0 - ): - head_dim_qk = getattr(case, "kv_lora_rank", 0) + getattr(case, "qk_rope_head_dim", 0) - head_dim_v = getattr(case, "kv_lora_rank", 0) - if (head_dim_qk, head_dim_v, getattr(case, "page_size", None)) == ( - 576, - 512, - 32, - ): - return ( - "TRTLLM-Gen Blackwell MLA generation is missing the decode " - "kernel for headDimQk=576, headDimV=512, page_size=32" - ) - - # TRTLLM's fp8 generation (XQA) path computes in fp8 and needs the model's - # real KV-dequant + output scale state (fed by the Attention module's - # projection layers). A bare standalone backend lacks it: supplying a unit - # output scale lets the kernel run, but the MHA path then produces garbage - # (~1e5 abs error) while only GQA happens to tolerate it. fp8 *context* - # (pure prefill) is exercised; FlashInfer validates fp8 decode on every arch. - if backend == "TRTLLM" and getattr(case, "kv_dtype", None) == "fp8": - has_generation = case.num_contexts < len(case.seq_lens) - if has_generation: - return ( - "TRTLLM fp8 KV generation (XQA) needs the model's fp8 scale " - "state, absent in the standalone backend; fp8 context is covered " - "and FlashInfer validates fp8 decode" - ) - - # TRTLLM cross-attention works through the standard plumbing (prepare() - # derives the cross kv_lens from seq_lens_kv; the backend builds cross_kv from - # k/v). The aligned q_len == kv_len case is exercised on TRTLLM. The - # q_len != kv_len case is numerically correct in isolation (verified vs the - # golden to ~2e-3) but exhibits a cross-case state-dependent mismatch when run - # after other cross cases in the same process (a TRTLLM cross-path global not - # reset between fresh backend instances), so it is validated via - # FlashInfer/Vanilla in the suite. - if ( - backend == "TRTLLM" - and getattr(case, "is_cross", False) - and list(case.seq_lens) != list(case.seq_lens_kv) - ): - return ( - "TRTLLM cross q_len != kv_len is correct standalone (~2e-3) but " - "flaky across cross cases in-process; validated via FlashInfer/Vanilla" - ) return None diff --git a/tests/unittest/_torch/attention/backend_case.py b/tests/unittest/_torch/attention/backend_case.py index f1a9583ab864..cc30f1e051e0 100644 --- a/tests/unittest/_torch/attention/backend_case.py +++ b/tests/unittest/_torch/attention/backend_case.py @@ -80,8 +80,7 @@ class BackendCase: # Cross-attention: new KV (encoder) tokens per request. None => self-attention # (KV tokens == seq_lens). When set, the case is cross-attention (must be - # non-causal). Same-length cross runs on TRTLLM/FlashInfer/Vanilla; unequal - # q/kv lengths are gated for TRTLLM in the capability matrix. + # non-causal). seq_lens_kv: Optional[List[int]] = None dtype: str = "float16" @@ -113,9 +112,9 @@ class BackendCase: # mla_rope_generation (feeding a pre-formed fused_q and explicit q_pe), so # all three backends run RoPE-free and stay aligned; TRTLLM additionally # pre-writes the new latent into the cache and Python-initializes the - # trtllm-gen scheduler buffers. MLA *context* (up-projected K/V) runs on - # Vanilla/FlashInfer; the TRTLLM context FMHA fuses RoPE and is validated in - # test_attention_mla.py. + # trtllm-gen scheduler buffers. MLA context uses coherent up-projected K/V + # and latent-cache inputs: TRTLLM fuses RoPE while Vanilla/FlashInfer receive + # the equivalent pre-rotated tensors. v_head_dim: Optional[int] = None q_lora_rank: Optional[int] = None kv_lora_rank: Optional[int] = None @@ -562,18 +561,104 @@ def _create_metadata(AttentionCls, case, mgr): mgr.shutdown() +def _mla_context_pos_embd_params(case: BackendCase) -> PositionalEmbeddingParams: + """Build the GPT-J-style RoPE configuration required by TRTLLM MLA context.""" + if case.rope is None: + raise ValueError("TRTLLM MLA context requires RoPE parameters.") + + rope_config = dict(case.rope) + rope_config.update(dim=case.qk_rope_head_dim, duplicate_data=True) + return PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gptj, + rope=_rope_params_from_dict(rope_config), + is_neox=False, + ) + + def generate_mla_context_inputs(case: BackendCase, seed: int = 0) -> Dict: - """Random up-projected MLA context inputs (asymmetric K/V).""" + """Random production-layout MLA context inputs before RoPE.""" gen = torch.Generator(device="cuda").manual_seed(seed) cdt = case.compute_dtype H, Hkv = case.num_heads, case.num_kv_heads qk_head = case.qk_nope_head_dim + case.qk_rope_head_dim - d_latent = case.kv_lora_rank + case.qk_rope_head_dim + compressed_kv = _randn(gen, cdt, case.nnz_q, case.kv_lora_rank) + k_pe = _randn(gen, cdt, case.nnz_q, case.qk_rope_head_dim) + packed_kv = _randn( + gen, + cdt, + case.nnz_q, + Hkv * (case.qk_nope_head_dim + case.v_head_dim), + ) + k_nope, v = packed_kv.split([Hkv * case.qk_nope_head_dim, Hkv * case.v_head_dim], dim=-1) + k = torch.cat( + [ + k_nope.view(-1, Hkv, case.qk_nope_head_dim), + k_pe.view(-1, 1, case.qk_rope_head_dim).expand(-1, Hkv, -1), + ], + dim=-1, + ).view(-1, Hkv * qk_head) return dict( q=_randn(gen, cdt, case.nnz_q, H * qk_head), - k=_randn(gen, cdt, case.nnz_q, Hkv * qk_head), - v=_randn(gen, cdt, case.nnz_q, Hkv * case.v_head_dim), - latent_cache=_randn(gen, cdt, case.nnz_q, d_latent), # appended for later gen + k=k, + # Keep the split view: TRTLLM MLA context expects token stride to include + # the packed k_nope portion, and Vanilla/FlashInfer support that layout. + v=v, + compressed_kv=compressed_kv, + k_pe=k_pe, + latent_cache=torch.cat([compressed_kv, k_pe], dim=-1), + ) + + +def _prepare_mla_context_inputs( + case: BackendCase, + inputs: Dict, + pos_embd_params: PositionalEmbeddingParams, + *, + fuse_rope: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Prepare backend inputs and the expected post-RoPE latent cache.""" + position_ids = make_position_ids(case.seq_lens, case.num_cached_tokens) + rope_params = pos_embd_params.rope + assert rope_params is not None + + rotated_k_pe = apply_rope( + inputs["k_pe"], + position_ids, + rope_params, + case.qk_rope_head_dim, + is_neox=pos_embd_params.is_neox, + ) + expected_latent_cache = torch.cat([inputs["compressed_kv"], rotated_k_pe], dim=-1) + + if fuse_rope: + return ( + inputs["q"].clone(), + inputs["k"].clone(), + inputs["v"], + inputs["latent_cache"], + expected_latent_cache, + ) + + q = inputs["q"].clone().view(-1, case.num_heads, case.head_dim) + q_pe = q[..., case.qk_nope_head_dim :].reshape( + case.nnz_q, case.num_heads * case.qk_rope_head_dim + ) + q[..., case.qk_nope_head_dim :] = apply_rope( + q_pe, + position_ids, + rope_params, + case.qk_rope_head_dim, + is_neox=pos_embd_params.is_neox, + ).view(case.nnz_q, case.num_heads, case.qk_rope_head_dim) + + k = inputs["k"].clone().view(-1, case.num_kv_heads, case.head_dim) + k[..., case.qk_nope_head_dim :] = rotated_k_pe.view(case.nnz_q, 1, case.qk_rope_head_dim) + return ( + q.view(case.nnz_q, -1), + k.view(case.nnz_q, -1), + inputs["v"], + expected_latent_cache, + expected_latent_cache, ) @@ -582,6 +667,7 @@ def _run_mla_context_backend(case, backend, inputs, *, kv_layout: str) -> torch. AttentionCls = get_attention_backend(backend) qk_head = case.qk_nope_head_dim + case.qk_rope_head_dim request_ids = list(range(case.num_seqs)) + pos_embd_params = _mla_context_pos_embd_params(case) attn = create_attention( backend, layer_idx=0, @@ -589,6 +675,7 @@ def _run_mla_context_backend(case, backend, inputs, *, kv_layout: str) -> torch. head_dim=qk_head, num_kv_heads=case.num_kv_heads, q_scaling=case.q_scaling, + pos_embd_params=pos_embd_params, is_mla_enable=True, q_lora_rank=case.q_lora_rank, kv_lora_rank=case.kv_lora_rank, @@ -598,7 +685,15 @@ def _run_mla_context_backend(case, backend, inputs, *, kv_layout: str) -> torch. ) mgr = _build_mla_kv_cache_manager(case, backend) mgr.add_dummy_requests(request_ids, case.token_nums) - expected_latents = _split_packed_tokens(inputs["latent_cache"], case.seq_lens) + fuse_rope = AttentionCls.support_fused_rope() + q, k, v, latent_cache, expected_latent_cache = _prepare_mla_context_inputs( + case, + inputs, + pos_embd_params, + fuse_rope=fuse_rope, + ) + expected_latents = _split_packed_tokens(expected_latent_cache, case.seq_lens) + cache_atol, cache_rtol = _tolerances(case, case.compute_dtype) if fuse_rope else (0.0, 0.0) metadata = AttentionCls.Metadata( num_contexts=case.num_contexts, kv_cache_params=KVCacheParams( @@ -615,12 +710,12 @@ def _run_mla_context_backend(case, backend, inputs, *, kv_layout: str) -> torch. metadata.prepare() try: out = attn.forward( - inputs["q"], - inputs["k"], - inputs["v"], + q, + k, + v, metadata, forward_args=AttentionForwardArgs( - latent_cache=inputs["latent_cache"], + latent_cache=latent_cache, attention_input_type=AttentionInputType.context_only, ), ) @@ -635,6 +730,8 @@ def _run_mla_context_backend(case, backend, inputs, *, kv_layout: str) -> torch. expected_latents, kv_layout=metadata.kv_layout, cache_kind="mla", + atol=cache_atol, + rtol=cache_rtol, ) return out[: case.nnz_q].contiguous() finally: @@ -865,7 +962,7 @@ def create_metadata(AttentionCls, case, mgr, *, num_contexts: int = 0): max_num_tokens=case.max_num_tokens, kv_cache_manager=mgr, request_ids=request_ids, - prompt_lens=case.token_nums, + prompt_lens=case.seq_lens if case.is_cross else case.token_nums, kv_layout=kv_layout, ) diff --git a/tests/unittest/_torch/attention/test_attention_op_sync.py b/tests/unittest/_torch/attention/test_attention_op_sync.py index d1beb769f16d..cfd56978b848 100644 --- a/tests/unittest/_torch/attention/test_attention_op_sync.py +++ b/tests/unittest/_torch/attention/test_attention_op_sync.py @@ -59,20 +59,13 @@ } _THOP_KWARG_SOURCE_ALIASES: dict[str, tuple[str, tuple[str, ...]]] = { + "beam_width": ("metadata", ("effective_beam_width",)), "context_lengths": ("metadata", ("prompt_lens_cuda_runtime",)), "head_size": ("attn", ("head_dim",)), "host_context_lengths": ("metadata", ("prompt_lens_cpu_runtime",)), "host_past_key_value_lengths": ("metadata", ("kv_lens_runtime",)), "host_request_types": ("metadata", ("host_request_types_runtime",)), "sequence_length": ("metadata", ("kv_lens_cuda_runtime",)), - "skip_softmax_threshold_scale_factor_decode": ( - "skip_softmax_kernel_params", - ("threshold_scale_factor_decode",), - ), - "skip_softmax_threshold_scale_factor_prefill": ( - "skip_softmax_kernel_params", - ("threshold_scale_factor_prefill",), - ), "spec_decoding_target_max_draft_tokens": ( "metadata", ("max_total_draft_tokens",),