diff --git a/tensorrt_llm/_torch/attention_backend/sparse/registry.py b/tensorrt_llm/_torch/attention_backend/sparse/registry.py index cd2cea5d2db2..51340a14581e 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/registry.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/registry.py @@ -67,6 +67,10 @@ def get_vanilla_sparse_attn_attention_backend( if sparse_params.algorithm == "rocket": return RocketVanillaAttention + elif sparse_params.algorithm == "dsa": + from ..vanilla import VanillaAttention + + return VanillaAttention elif sparse_params.algorithm == "minimax_m3": return _resolve_minimax_m3_backend_cls(sparse_params) else: diff --git a/tensorrt_llm/_torch/attention_backend/vanilla.py b/tensorrt_llm/_torch/attention_backend/vanilla.py index ac1ef7aa6d0f..05e64f1f8a7f 100644 --- a/tensorrt_llm/_torch/attention_backend/vanilla.py +++ b/tensorrt_llm/_torch/attention_backend/vanilla.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import math +from dataclasses import replace from typing import Optional import torch @@ -118,6 +119,19 @@ def __init__( def support_mla(cls) -> bool: return True + def sparse_attn_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + metadata: VanillaAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """Use request-local caller-provided selections in the Vanilla backend.""" + sparse_backend_args = forward_args.sparse_backend_args + topk_indices = (sparse_backend_args.topk_indices + if sparse_backend_args is not None else None) + return topk_indices, None + def _single_request_sparse_attn_predict( self, q: torch.Tensor, k: Optional[torch.Tensor], v: Optional[torch.Tensor], kv_cache_tensor: torch.Tensor, @@ -217,13 +231,19 @@ def _single_request_update_kv_cache(self, assert blk != BAD_PAGE_INDEX, ( f"Writing new KV into an evicted/invalid page (pos {pos}); " "block_ids/metadata are inconsistent.") - dst = torch.arange(off, off + n, device=kv_cache_tensor.device) - kv_cache_tensor[blk, 0].view(dtype=access_type).index_copy_( - 0, dst, - k_selected[0, written:written + n].view(dtype=access_type)) - kv_cache_tensor[blk, 1].view(dtype=access_type).index_copy_( - 0, dst, - v_selected[0, written:written + n].view(dtype=access_type)) + # Slicing the outermost (token) dim keeps the destination + # contiguous, so a plain copy_ avoids the per-iteration arange + # and scatter. view(access_type) reinterprets to an int of the + # same width so copy_ works for dtypes (e.g. fp8) it otherwise + # rejects. + kv_cache_tensor[blk, 0, + off:off + n].view(dtype=access_type).copy_( + k_selected[0, written:written + + n].view(dtype=access_type)) + kv_cache_tensor[blk, 1, + off:off + n].view(dtype=access_type).copy_( + v_selected[0, written:written + + n].view(dtype=access_type)) written += n if sparse_kv_indices is not None: @@ -646,6 +666,174 @@ def _mla_forward_generation(self, fused_q: torch.Tensor, return torch.cat(outputs, dim=0) + @staticmethod + def _load_mla_latent_cache(kv_cache: torch.Tensor, block_ids: list[int], + kv_len: int, kv_layout: str) -> torch.Tensor: + """Materialize one request's logical MLA cache from its pages.""" + if kv_len <= 0: + raise ValueError(f"MLA KV length must be positive, got {kv_len}") + if kv_layout == "NHD": + tokens_per_block = kv_cache.shape[2] + elif kv_layout == "HND": + tokens_per_block = kv_cache.shape[3] + else: + raise ValueError(f"Unsupported KV cache layout: {kv_layout}") + + valid_block_ids = [block_id for block_id in block_ids if block_id != -1] + num_required_blocks = math.ceil(kv_len / tokens_per_block) + if len(valid_block_ids) < num_required_blocks: + raise ValueError( + f"MLA cache has {len(valid_block_ids)} blocks, but " + f"{num_required_blocks} are required for {kv_len} tokens") + + chunks = [] + remaining = kv_len + for block_id in valid_block_ids[:num_required_blocks]: + num_tokens = min(tokens_per_block, remaining) + if kv_layout == "NHD": + chunk = kv_cache[block_id, 0, :num_tokens, 0, :] + else: + chunk = kv_cache[block_id, 0, 0, :num_tokens, :] + chunks.append(chunk) + remaining -= num_tokens + return torch.cat(chunks, dim=0) + + def _mla_forward_sparse( + self, + fused_q: torch.Tensor, + metadata: VanillaAttentionMetadata, + latent_cache: torch.Tensor, + topk_indices: torch.Tensor, + attention_input_type: AttentionInputType, + ) -> torch.Tensor: + """Run selected sparse MLA from caller-provided local top-k rows. + + The sparse algorithm owns selection. This golden consumes its + request-local per-token selections, gathers the selected latent K/V, and + performs the absorbed MLA attention directly in PyTorch. The MLA sparse + algorithms (DSA / DeepSeek-V4) select individual tokens. + + ``fused_q`` and ``latent_cache`` arrive with RoPE already applied (like + every other Vanilla attention path); the caller owns positional embedding. + """ + if attention_input_type == AttentionInputType.context_only: + seq_start, seq_end = 0, metadata.num_contexts + elif attention_input_type == AttentionInputType.generation_only: + seq_start, seq_end = metadata.num_contexts, metadata.num_seqs + else: + raise ValueError( + "Vanilla DSA requires a context-only or generation-only input") + + seq_lens = metadata.seq_lens.tolist() + phase_seq_lens = seq_lens[seq_start:seq_end] + num_phase_tokens = sum(phase_seq_lens) + fused_head_dim = self.kv_lora_rank + self.qk_rope_head_dim + if fused_q.shape[0] != num_phase_tokens: + raise ValueError( + f"DSA query has {fused_q.shape[0]} tokens, but metadata " + f"describes {num_phase_tokens} tokens for this phase") + if fused_q.ndim != 2 or fused_q.shape[ + 1] != self.num_heads * fused_head_dim: + raise ValueError( + "DSA query must have shape " + f"[{num_phase_tokens}, {self.num_heads * fused_head_dim}]; " + f"got {tuple(fused_q.shape)}") + if (latent_cache.ndim != 2 + or latent_cache.shape != (num_phase_tokens, fused_head_dim)): + raise ValueError("DSA latent cache must have shape " + f"[{num_phase_tokens}, {fused_head_dim}]; " + f"got {tuple(latent_cache.shape)}") + if topk_indices.ndim != 2 or topk_indices.shape[0] != num_phase_tokens: + raise ValueError( + "DSA top-k indices must have shape [num_phase_tokens, top_k]; " + f"got {tuple(topk_indices.shape)}") + if topk_indices.dtype != torch.int32: + raise ValueError( + f"DSA top-k indices must have dtype int32, got {topk_indices.dtype}" + ) + + request_ids = metadata.request_ids[seq_start:seq_end] + past_tokens = metadata.kv_cache_params.num_cached_tokens_per_seq + phase_past_tokens = past_tokens[seq_start:seq_end] + valid_mask = topk_indices >= 0 + if torch.any(topk_indices < -1): + raise ValueError("DSA top-k indices may only use -1 as padding") + if torch.any(~valid_mask.any(dim=1)): + raise ValueError( + "Every DSA query token must select at least one KV token") + + kv_lengths = torch.cat([ + torch.full( + (q_len, ), + int(past) + q_len, + dtype=topk_indices.dtype, + device=topk_indices.device, + ) for past, q_len in zip( + phase_past_tokens, phase_seq_lens, strict=True) + ]) + if torch.any(valid_mask & (topk_indices >= kv_lengths.unsqueeze(1))): + raise ValueError( + "DSA top-k index is out of bounds for its request-local KV length" + ) + + causal_limits = torch.cat([ + torch.arange( + int(past), + int(past) + q_len, + dtype=topk_indices.dtype, + device=topk_indices.device, + ) for past, q_len in zip( + phase_past_tokens, phase_seq_lens, strict=True) + ]) + if torch.any(valid_mask & (topk_indices > causal_limits.unsqueeze(1))): + raise ValueError("DSA top-k index selects a future token") + del valid_mask, kv_lengths, causal_limits + + from .utils import append_mla_latent_cache + kv_cache = append_mla_latent_cache( + metadata.kv_cache_manager, + self.layer_idx, + request_ids, + phase_seq_lens, + phase_past_tokens, + latent_cache, + kv_layout=metadata.kv_layout, + ) + + q = fused_q.view(num_phase_tokens, self.num_heads, fused_head_dim) + qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + scale = 1.0 / (math.sqrt(qk_head_dim) * + (self.q_scaling if self.q_scaling is not None else 1.0)) + + outputs = [] + token_offset = 0 + for phase_idx, q_len in enumerate(phase_seq_lens): + seq_idx = seq_start + phase_idx + kv_len = int(phase_past_tokens[phase_idx]) + q_len + latent = self._load_mla_latent_cache( + kv_cache, metadata.block_ids_per_seq[seq_idx], kv_len, + metadata.kv_layout).to(q.dtype) + + per_token_outputs = [] + for token_idx in range(q_len): + row = topk_indices[token_offset + token_idx] + selected = row[row >= 0].to(device=q.device, dtype=torch.long) + selected_latent = latent.index_select(0, selected) + query = q[token_offset + token_idx] + scores = torch.matmul(query, selected_latent.transpose( + 0, 1)) * scale + probabilities = F.softmax(scores, dim=-1, + dtype=torch.float32).to(q.dtype) + values = selected_latent[:, :self.kv_lora_rank] + per_token_outputs.append(torch.matmul(probabilities, values)) + + outputs.append( + torch.stack(per_token_outputs).reshape( + q_len, self.num_heads * self.kv_lora_rank)) + token_offset += q_len + + return torch.cat(outputs, dim=0) + def _mla_forward_context(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, metadata: VanillaAttentionMetadata, @@ -710,6 +898,37 @@ def forward(self, raise ValueError("Vanilla MLA requires a KV cache manager.") if forward_args.latent_cache is None: raise ValueError("Vanilla MLA requires latent_cache.") + if self.sparse_params is not None: + sparse_algorithm = self.sparse_params.algorithm + if sparse_algorithm != "dsa": + raise ValueError( + "Vanilla selected MLA currently supports only DSA") + sparse_attn_indices, sparse_attn_offsets = self.sparse_attn_predict( + q, k, metadata, forward_args) + forward_args.sparse_runtime_params = replace( + forward_args.sparse_runtime_params, + sparse_attn_indices=sparse_attn_indices, + sparse_attn_offsets=sparse_attn_offsets, + sparse_attn_indices_block_size=getattr( + self.sparse_params, "indices_block_size"), + ) + sparse_attn_indices = ( + forward_args.sparse_runtime_params.sparse_attn_indices) + if sparse_attn_indices is not None: + if k is not None or v is not None: + raise ValueError( + "Vanilla sparse MLA expects absorbed queries and latent cache, " + "not explicit K/V tensors") + return self._mla_forward_sparse( + q, + metadata, + forward_args.latent_cache, + sparse_attn_indices, + forward_args.attention_input_type, + ) + if self.sparse_params is not None: + raise ValueError( + "Vanilla sparse MLA requires sparse attention indices") if forward_args.attention_input_type == AttentionInputType.context_only: assert k is not None and v is not None return self._mla_forward_context(q, k, v, metadata, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 680d607efeb5..dbb7b5138c28 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -418,12 +418,8 @@ def create_py_executor( tokens_per_block = kv_cache_config.tokens_per_block - # RocketKV's Vanilla path keeps its landmark (KT) cache in a single block per - # sequence: RocketVanillaAttention writes the whole sequence into - # kt_cache_block_offsets[0], and kt_tokens_per_block is derived from - # tokens_per_block. It does not support a paged KT cache, so force one block - # per sequence for it. Plain Vanilla attention supports paged KV cache and is - # left untouched. + # RocketKV's Vanilla path does not support a paged KT cache, so force one + # block per sequence for it. See RocketVanillaAttention for detail. sparse_config = llm_args.sparse_attention_config if (llm_args.attn_backend == "VANILLA" and sparse_config is not None and getattr(sparse_config, "algorithm", None) == "rocket"): diff --git a/tests/unittest/_torch/attention/backend_capability.py b/tests/unittest/_torch/attention/backend_capability.py index 19d9d7230dee..029f26584546 100644 --- a/tests/unittest/_torch/attention/backend_capability.py +++ b/tests/unittest/_torch/attention/backend_capability.py @@ -20,7 +20,7 @@ # fp4_kv - NVFP4 KV cache (Blackwell only) # sliding_window - sliding-window attention via attention_window_size # no_cache - ragged/prefill forward with kv_cache_manager=None -# sparse - sparse-attention forward plumbing (degenerate regime here) +# sparse - sparse-attention forward plumbing # mla - multi-head latent attention # cross_attn - cross-attention (encoder-decoder) # kv_layouts - supported paged-cache block layouts ("NHD" / "HND") @@ -60,7 +60,7 @@ fp4_kv=False, sliding_window=True, no_cache=True, - sparse=False, + sparse=True, mla=True, cross_attn=True, kv_layouts=("NHD",), # reads the NHD get_buffers view @@ -86,7 +86,7 @@ def required_features(case) -> set: feats.add("sliding_window") if getattr(case, "cache", "paged") == "none": feats.add("no_cache") - if getattr(case, "sparse", "off") != "off": + if getattr(case, "sparse_attention_config", None) is not None: feats.add("sparse") if getattr(case, "is_mla", False): feats.add("mla") @@ -114,6 +114,17 @@ def unsupported_reason(backend: str, case) -> Optional[str]: if not caps.get(feat, False): return f"{backend} does not support feature '{feat}'" + sparse_config = getattr(case, "sparse_attention_config", None) + if sparse_config is not None: + algorithm = sparse_config.algorithm + if backend == "TRTLLM" and algorithm == "dsa": + # DSA selected-attention runs the trtllm-gen DynamicTokenSparse FMHA + # kernels, which only ship for Blackwell (sm_100+). On Hopper (sm90) + # MLA generation falls back to the dense FlashMLA kernel, which has no + # sparse path, so top-k selection is silently ignored. + if sm < 100: + return f"TRTLLM DSA requires sm>=100/Blackwell (have sm{sm})" + # KV-cache block layout: a case may request a specific layout (NHD/HND). A # backend that cannot store the cache that way is skipped (e.g. TRTLLM is # head-major HND only). The Vanilla golden always runs in its native NHD and diff --git a/tests/unittest/_torch/attention/backend_case.py b/tests/unittest/_torch/attention/backend_case.py index cc30f1e051e0..e097f5d41ff3 100644 --- a/tests/unittest/_torch/attention/backend_case.py +++ b/tests/unittest/_torch/attention/backend_case.py @@ -28,6 +28,7 @@ PredefinedAttentionMask, RopeParams, ) +from tensorrt_llm._torch.attention_backend.sparse import get_sparse_attn_kv_cache_manager from tensorrt_llm._torch.attention_backend.utils import create_attention, get_attention_backend from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE from tensorrt_llm._torch.metadata import KVCacheParams @@ -35,7 +36,7 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import str_dtype_to_torch, torch_dtype_to_binding from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType -from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.llmapi.llm_args import KvCacheConfig, SparseAttentionConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo @@ -50,6 +51,10 @@ # differ from the fp16 golden by ~0.1-0.4. Matches test_attention_mla.py (fp8=4e-1). FP8_ATOL = 4e-1 FP4_ATOL = 6e-1 +# Golden-vs-backend tolerance for selected sparse MLA (bf16 latent-gather +# accumulation). +SPARSE_ATOL = 1e-1 +SPARSE_RTOL = 1e-2 # Backends compared against the VanillaAttention golden. FlashInfer is only # included when available, so callers can iterate this list unconditionally. @@ -92,7 +97,11 @@ class BackendCase: q_scaling: float = 1.0 page_size: int = 64 cache: str = "paged" # "paged" | "none" - sparse: str = "off" # "off" | "degenerate" + # User-facing sparse config, lowered into backend params, metadata params, and + # the sparse KV-cache manager exactly as in production. The selection unit and + # top-k are derived from it (properties below); the attention family is + # ``is_mla``. + sparse_attention_config: Optional[SparseAttentionConfig] = None # RoPE config: RopeParams kwargs (+ optional "is_neox"), or None to disable. rope: Optional[dict] = None # When True (and rope set), exercise TRTLLM's in-kernel fused RoPE: TRTLLM @@ -116,6 +125,7 @@ class BackendCase: # and latent-cache inputs: TRTLLM fuses RoPE while Vanilla/FlashInfer receive # the equivalent pre-rotated tensors. v_head_dim: Optional[int] = None + hidden_size: Optional[int] = None q_lora_rank: Optional[int] = None kv_lora_rank: Optional[int] = None qk_nope_head_dim: Optional[int] = None @@ -134,6 +144,28 @@ def nnz_q(self) -> int: def is_cross(self) -> bool: return self.seq_lens_kv is not None + @property + def is_sparse(self) -> bool: + return self.sparse_attention_config is not None + + @property + def sparse_topk(self) -> Optional[int]: + """Per-token selection budget (``index_topk``) from the sparse config.""" + cfg = self.sparse_attention_config + if cfg is None: + return None + return cfg.index_topk + + @property + def prompt_lens(self) -> List[int]: + """Original prompt lengths expected by fused generation kernels.""" + return [ + seq_len if i < self.num_contexts else cached_len + for i, (seq_len, cached_len) in enumerate( + zip(self.seq_lens, self.num_cached_tokens, strict=True) + ) + ] + @property def is_gen_only(self) -> bool: """A uniform pure-decode batch eligible for a captured CUDA graph. @@ -201,6 +233,12 @@ def _rope_params_from_dict(d: dict) -> RopeParams: return RopeParams(**kwargs) +def _validate_sparse_case(case: BackendCase) -> None: + """Reject unsupported sparse contracts (called only for sparse cases).""" + if case.sparse_topk is None or case.sparse_topk <= 0: + raise ValueError("Sparse backend cases require a positive top-k") + + def _randn(gen: torch.Generator, dtype: torch.dtype, *shape) -> torch.Tensor: """Seeded random tensor on cuda in ``dtype`` (shared by all input builders).""" return torch.randn(*shape, generator=gen, device="cuda").to(dtype) @@ -275,10 +313,16 @@ def _build_kv_cache_manager(case: BackendCase, backend: str, kv_dtype: torch.dty # FlashInfer, the comparison validates the absorbed-MQA math regardless of the # RoPE values (RoPE correctness is covered by test_attention_mla.py). # --------------------------------------------------------------------------- -def _build_mla_kv_cache_manager(case: BackendCase, backend: str): +def _build_mla_kv_cache_manager( + case: BackendCase, + backend: str, + sparse_config=None, +): """A SELFKONLY KV cache for MLA: one latent head, head_dim kv_lora+qk_rope.""" d_latent = case.kv_lora_rank + case.qk_rope_head_dim - paged = BACKEND_CAPS[backend]["paged"] + # Sparse selected-attention tests deliberately exercise multiple pages in + # Vanilla too. Dense Vanilla keeps its historical single-block setup. + paged = case.is_sparse or BACKEND_CAPS[backend]["paged"] max_total = max(case.token_nums) if paged: tokens_per_block = case.page_size @@ -289,10 +333,12 @@ def _build_mla_kv_cache_manager(case: BackendCase, backend: str): num_blocks = case.num_seqs * pages_per_seq mapping = Mapping(world_size=1, tp_size=1, rank=0) cache_types = tensorrt_llm.bindings.internal.batch_manager.CacheType - cls = KVCacheManagerV2 if case.use_kv_cache_manager_v2 else KVCacheManager - return cls( - KvCacheConfig(max_tokens=num_blocks * tokens_per_block, enable_block_reuse=False), - cache_types.SELFKONLY, + kwargs = dict( + kv_cache_config=KvCacheConfig( + max_tokens=num_blocks * tokens_per_block, + enable_block_reuse=False, + ), + kv_cache_type=cache_types.SELFKONLY, num_layers=1, num_kv_heads=1, head_dim=d_latent, @@ -303,6 +349,14 @@ def _build_mla_kv_cache_manager(case: BackendCase, backend: str): dtype=torch_dtype_to_binding(case.compute_dtype), ) + if sparse_config is not None: + cls = get_sparse_attn_kv_cache_manager(sparse_config) + kwargs.update(sparse_attention_config=sparse_config) + else: + cls = KVCacheManagerV2 if case.use_kv_cache_manager_v2 else KVCacheManager + + return cls(**kwargs) + def generate_mla_gen_inputs(case: BackendCase, seed: int = 0) -> Dict: """Random absorbed-MLA generation inputs (shared by all backends). @@ -331,6 +385,152 @@ def generate_mla_gen_inputs(case: BackendCase, seed: int = 0) -> Dict: ) +def _build_sparse_topk_indices( + case: BackendCase, + generator: torch.Generator, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build causal request-local selections, mixing singleton and random rows. + + A row whose visible prefix fits within top-k selects every causal token + (dense). A truly sparse row (visible > top-k) is alternately made a + ``singleton`` row -- one logical token repeated across all top-k slots, which + yields an analytic value oracle -- or a ``random`` row. ``singleton_oracle_mask`` + marks the singleton rows so the caller applies the analytic check only there. + + Selections are per-token (DSA / DeepSeek-V4 select individual tokens). + """ + topk = case.sparse_topk + if topk is None or topk <= 0: + raise ValueError("Sparse selection cases require a positive top-k") + + indices = torch.full((case.nnz_q, topk), -1, dtype=torch.int32, device="cuda") + singleton_oracle_mask = torch.zeros(case.nnz_q, dtype=torch.bool, device="cuda") + row = 0 + sparse_row = 0 # counts truly-sparse rows, to alternate singleton / random + for cached_len, q_len in zip(case.num_cached_tokens, case.seq_lens, strict=True): + for token_idx in range(q_len): + max_index = cached_len + token_idx + if max_index + 1 <= topk: + indices[row, : max_index + 1] = torch.arange( + max_index + 1, dtype=torch.int32, device="cuda" + ) + elif sparse_row % 2 == 0: + # Singleton: repeat one logical token across all top-k slots. + selected = (0, max_index, max_index // 2)[row % 3] + indices[row].fill_(selected) + singleton_oracle_mask[row] = True + sparse_row += 1 + else: + selected = torch.randperm(max_index + 1, generator=generator, device="cuda")[:topk] + indices[row] = torch.sort(selected).values.to(torch.int32) + sparse_row += 1 + row += 1 + return indices, singleton_oracle_mask + + +def generate_sparse_mla_inputs(case: BackendCase, seed: int = 0) -> Dict: + """Generate raw absorbed-MLA inputs plus backend-neutral sparse selections.""" + if not case.is_mla: + raise ValueError("This generator supports selected sparse MLA only") + gen = torch.Generator(device="cuda").manual_seed(seed) + selection_gen = torch.Generator(device="cuda").manual_seed(seed + 1) + cdt = case.compute_dtype + num_heads = case.num_heads + kv_lora_rank = case.kv_lora_rank + qk_rope_head_dim = case.qk_rope_head_dim + d_latent = kv_lora_rank + qk_rope_head_dim + + q_nope = _randn(gen, cdt, case.nnz_q, num_heads, kv_lora_rank) + q_pe = _randn(gen, cdt, case.nnz_q, num_heads, qk_rope_head_dim) + compressed_kv = _randn(gen, cdt, case.nnz_q, kv_lora_rank) + k_pe = _randn(gen, cdt, case.nnz_q, qk_rope_head_dim) + + pos_embd_params = _mla_context_pos_embd_params(case) + rope_params = pos_embd_params.rope + assert rope_params is not None + new_positions = make_position_ids(case.seq_lens, case.num_cached_tokens) + # Two input flavors, because RoPE happens in different places per phase: + # * generation: the absorbed path runs with skip_mla_rope_generation, so no + # backend ropes -- feed the RoPE'd (pre-formed) inputs to every backend. + # * context: the TRTLLM MLA context kernel ropes internally (no skip exists + # for context), so it must get RAW inputs and rope them once itself; Vanilla + # (which never ropes) still gets the RoPE'd inputs. + # q_pe rotates per head; k_pe is shared across heads. + rotated_q_pe = apply_rope( + q_pe.reshape(case.nnz_q, num_heads * qk_rope_head_dim), + new_positions, + rope_params, + qk_rope_head_dim, + is_neox=pos_embd_params.is_neox, + ).reshape(case.nnz_q, num_heads, qk_rope_head_dim) + fused_q = torch.cat((q_nope, rotated_q_pe), dim=-1).reshape(case.nnz_q, num_heads * d_latent) + fused_q_raw = torch.cat((q_nope, q_pe), dim=-1).reshape(case.nnz_q, num_heads * d_latent) + rotated_new_k_pe = apply_rope( + k_pe, + new_positions, + rope_params, + qk_rope_head_dim, + is_neox=pos_embd_params.is_neox, + ) + expected_new_latent = torch.cat((compressed_kv, rotated_new_k_pe), dim=-1) + # RoPE'd new-token latent: with skip_mla_rope_generation the backend appends it + # verbatim. The raw variant is roped in-kernel by the TRTLLM context path. + latent_cache = expected_new_latent + latent_cache_raw = torch.cat((compressed_kv, k_pe), dim=-1) + + cached_latent = [] + for cached_len in case.num_cached_tokens: + cached_compressed = _randn(gen, cdt, cached_len, kv_lora_rank) + cached_k_pe = _randn(gen, cdt, cached_len, qk_rope_head_dim) + if cached_len: + cached_positions = torch.arange(cached_len, dtype=torch.int32, device="cuda") + cached_k_pe = apply_rope( + cached_k_pe, + cached_positions, + rope_params, + qk_rope_head_dim, + is_neox=pos_embd_params.is_neox, + ) + cached_latent.append(torch.cat((cached_compressed, cached_k_pe), dim=-1)) + + topk_indices, singleton_oracle_mask = _build_sparse_topk_indices(case, selection_gen) + # Analytic value oracle for singleton rows: a query that selects a single + # logical token must return that token's latent-V slice on every head. + expected_output = None + if bool(singleton_oracle_mask.any()): + new_per_seq = _split_packed_tokens(expected_new_latent, case.seq_lens) + selected_values = [] + row = 0 + for cached, new_tokens in zip(cached_latent, new_per_seq, strict=True): + logical_cache = torch.cat((cached, new_tokens), dim=0) + for _ in range(new_tokens.shape[0]): + selected_values.append(logical_cache[int(topk_indices[row, 0]), :kv_lora_rank]) + row += 1 + expected_output = ( + torch.stack(selected_values) + .unsqueeze(1) + .expand(-1, num_heads, -1) + .reshape(case.nnz_q, num_heads * kv_lora_rank) + ) + + return dict( + fused_q=fused_q, + # The RoPE'd q_pe view is passed explicitly since the MLA RoPE step is + # skipped (skip_mla_rope_generation); it must match the fused_q pe slot. + q_pe=fused_q.view(case.nnz_q, num_heads, d_latent)[..., kv_lora_rank:], + latent_cache=latent_cache, + # Raw (un-RoPE'd) variants for the TRTLLM context path, which ropes in-kernel. + fused_q_raw=fused_q_raw, + q_pe_raw=q_pe, + latent_cache_raw=latent_cache_raw, + cached_latent=cached_latent, + expected_new_latent=expected_new_latent, + topk_indices=topk_indices, + expected_output=expected_output, + expected_output_mask=singleton_oracle_mask, + ) + + def _fill_mla_cache(mgr, layer_idx, request_ids, cached_latent, *, kv_layout="NHD"): """Write the per-request cached latent prefix into the MLA cache pool.""" if all(c.shape[0] == 0 for c in cached_latent): @@ -455,6 +655,143 @@ def _assert_cache_contains_new_tokens( torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol) +def _run_sparse_mla_backend( + case: BackendCase, + backend: str, + inputs: Dict, + *, + kv_layout: str, +) -> torch.Tensor: + """Run selected sparse MLA through production backend/config lowering.""" + sparse_config = case.sparse_attention_config + assert sparse_config is not None + # The DSA config carries every field the lowering needs, so no pretrained + # config is required. + sparse_params = sparse_config.to_sparse_params(layer_idx=None, pretrained_config=None) + sparse_metadata_params = sparse_config.to_sparse_metadata_params(pretrained_config=None) + AttentionCls = get_attention_backend(backend, sparse_params=sparse_params) + request_ids = list(range(case.num_seqs)) + d_latent = case.kv_lora_rank + case.qk_rope_head_dim + pos_embd_params = _mla_context_pos_embd_params(case) + mapping = Mapping(world_size=1, tp_size=1, rank=0) + attn = create_attention( + backend, + layer_idx=0, + num_heads=case.num_heads, + head_dim=d_latent, + num_kv_heads=1, + 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, + qk_nope_head_dim=case.qk_nope_head_dim, + qk_rope_head_dim=case.qk_rope_head_dim, + # Selected sparse MLA returns latent values; the model's V projection + # lives outside the standalone backend, so v_head_dim is the latent width. + v_head_dim=case.kv_lora_rank, + hidden_size=case.hidden_size, + predicted_tokens_per_seq=1, + sparse_params=sparse_params, + dtype=case.compute_dtype, + skip_create_weights_in_init=True, + ) + # Weights are skipped (selections are injected, not produced by an indexer); + # update_quant_config initializes the quant/FMHA state needed before forward. + attn.update_quant_config(None) + mgr = _build_mla_kv_cache_manager(case, backend, sparse_config) + + try: + mgr.add_dummy_requests(request_ids, case.token_nums) + _fill_mla_cache( + mgr, + 0, + request_ids, + inputs["cached_latent"], + kv_layout=kv_layout, + ) + metadata = AttentionCls.Metadata( + num_contexts=case.num_contexts, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=case.num_cached_tokens, + ), + seq_lens=torch.tensor(case.seq_lens, dtype=torch.int), + max_num_requests=case.num_seqs, + max_num_tokens=case.max_num_tokens, + kv_cache_manager=mgr, + request_ids=request_ids, + prompt_lens=case.prompt_lens, + kv_layout=kv_layout, + mapping=mapping, + sparse_metadata_params=sparse_metadata_params, + ) + metadata.prepare() + + num_context_tokens = sum(case.seq_lens[: case.num_contexts]) + phases = [] + if case.num_contexts: + phases.append((AttentionInputType.context_only, slice(0, num_context_tokens))) + if case.num_contexts < case.num_seqs: + phases.append( + (AttentionInputType.generation_only, slice(num_context_tokens, case.nnz_q)) + ) + + outputs = [] + for attention_input_type, token_slice in phases: + # RoPE placement differs by phase (see generate_sparse_mla_inputs): + # * TRTLLM context ropes in-kernel -> feed RAW inputs, no skip. + # * generation (both backends) + Vanilla context consume the + # pre-RoPE'd inputs; TRTLLM generation runs with + # skip_mla_rope_generation (it still appends the new latent + inits + # the trtllm-gen scheduler buffers inside forward). + kernel_ropes = ( + backend == "TRTLLM" and attention_input_type == AttentionInputType.context_only + ) + if kernel_ropes: + fused_q = inputs["fused_q_raw"][token_slice].clone() + q_pe = inputs["q_pe_raw"][token_slice] + latent_cache = inputs["latent_cache_raw"][token_slice].clone() + else: + fused_q = inputs["fused_q"][token_slice].clone() + q_pe = inputs["q_pe"][token_slice] + latent_cache = inputs["latent_cache"][token_slice].clone() + forward_args = AttentionForwardArgs( + latent_cache=latent_cache, + q_pe=q_pe, + topk_indices=inputs["topk_indices"][token_slice], + attention_input_type=attention_input_type, + skip_mla_rope_generation=not kernel_ropes, + ) + out = attn.forward( + fused_q, + None, + None, + metadata, + forward_args=forward_args, + ) + assert forward_args.sparse_runtime_params.sparse_attn_indices is not None + outputs.append(out[0] if isinstance(out, tuple) else out) + + expected_latents = _split_packed_tokens(inputs["expected_new_latent"], case.seq_lens) + cache_atol, cache_rtol = _tolerances(case, case.compute_dtype) + _assert_cache_contains_new_tokens( + mgr, + 0, + request_ids, + case.seq_lens, + case.num_cached_tokens, + expected_latents, + kv_layout=metadata.kv_layout, + cache_kind="mla", + atol=cache_atol, + rtol=cache_rtol, + ) + return torch.cat(outputs, dim=0)[: case.nnz_q].contiguous() + finally: + mgr.shutdown() + + def _run_mla_gen_backend( case, backend, inputs, *, kv_layout: str, cuda_graph=False ) -> torch.Tensor: @@ -755,6 +1092,9 @@ def _tolerances(case: "BackendCase", kv_dtype) -> tuple: dtype is bf16 its coarser mantissa compounds with the quant error, so the quantized atol gets extra headroom and the rtol relaxes to the bf16 rtol. """ + if case.is_sparse: + return SPARSE_ATOL, SPARSE_RTOL + bf16 = case.compute_dtype == torch.bfloat16 if kv_dtype == torch.float8_e4m3fn: return (FP8_ATOL + BF16_ATOL, BF16_RTOL) if bf16 else (FP8_ATOL, RTOL) @@ -875,6 +1215,11 @@ def run_backend( the caller passes a layout the backend supports (gated by the capability matrix). MLA cases are dispatched to the absorbed-generation path. """ + if case.is_sparse: + if case.is_mla: + return _run_sparse_mla_backend(case, backend, inputs, kv_layout=kv_layout) + raise ValueError(f"Unsupported sparse contract: is_mla={case.is_mla}") + if case.is_mla: if case.is_context_only: return _run_mla_context_backend(case, backend, inputs, kv_layout=kv_layout) @@ -1064,7 +1409,13 @@ def run_case(case: BackendCase, *, seed: int = 0) -> Dict[str, torch.Tensor]: that want the raw tensors (e.g. the minimizer). """ is_mla = case.is_mla - if is_mla: + if case.is_sparse: + _validate_sparse_case(case) + if case.is_mla: + inputs = generate_sparse_mla_inputs(case, seed) + else: + raise ValueError(f"Unsupported sparse contract: is_mla={case.is_mla}") + elif is_mla: if case.is_context_only: inputs = generate_mla_context_inputs(case, seed) else: @@ -1074,6 +1425,11 @@ def run_case(case: BackendCase, *, seed: int = 0) -> Dict[str, torch.Tensor]: golden = run_backend(case, "VANILLA", inputs, kv_dtype=case.compute_dtype, kv_layout="NHD") results = {"VANILLA": golden} + if inputs.get("expected_output") is not None: + oracle_mask = inputs["expected_output_mask"] + assert oracle_mask is not None and oracle_mask.any() + torch.testing.assert_close(golden[oracle_mask], inputs["expected_output"][oracle_mask]) + # Evaluate every supported backend before asserting, so one backend's # mismatch does not mask another's. failures = [] @@ -1097,8 +1453,9 @@ def run_case(case: BackendCase, *, seed: int = 0) -> Dict[str, torch.Tensor]: # A gen-only batch also exercises the captured-CUDA-graph path # (production replays a captured decode graph); it must still match the - # eager golden. - if case.is_gen_only: + # eager golden. The sparse runner rebuilds each request's logical cache on + # the host, which is not graph-capturable, so sparse cases are skipped. + if case.is_gen_only and not case.is_sparse: cg_out = run_backend( case, backend, diff --git a/tests/unittest/_torch/attention/model_attn_config.py b/tests/unittest/_torch/attention/model_attn_config.py index b6f32e2fbbe3..41d656eb3bcd 100644 --- a/tests/unittest/_torch/attention/model_attn_config.py +++ b/tests/unittest/_torch/attention/model_attn_config.py @@ -19,7 +19,9 @@ ``rope`` is one of ``None`` / ``"neox"`` / ``"gptj"``. ``mask`` is ``"causal"`` / ``"full"`` / ``"sliding"``. ``no_cache=True`` marks the -bidirectional, KV-cache-free DiT / encoder workloads. +bidirectional, KV-cache-free DiT / encoder workloads. Sparse models carry the +same user-facing sparse-attention config that production lowers independently +for the backend, metadata, and KV-cache manager. ID naming rule: - Use lowercase snake_case: @@ -46,14 +48,17 @@ Qwen2.5-VL with ``use_sliding_window=False``). - Vision/text encoders (SigLip/Radio/CLIP/Parakeet, MiniMax-VL tower) collapse onto the bidirectional MHA tuples already listed (e.g. 16x64 / 12x64 full). -- Sparse/DSA indexer attention (GLM-DSA, NSA, RocketKV) is a separate paradigm - validated under sparse/; there is no dense Vanilla golden for it. +- Sparse indexer correctness is validated under ``sparse/``. This model sweep + injects deterministic backend-neutral selections and validates the selected + attention path against the Vanilla golden. - Multimodal cross variants (Llama4-vision, Gemma4-MM) reuse the cross tuples. """ from dataclasses import dataclass from typing import Literal, Optional +from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig, SparseAttentionConfig + AttentionPhase = Literal["ctx", "gen"] @@ -80,6 +85,18 @@ class ModelAttnConfig: qk_nope_head_dim: Optional[int] = None qk_rope_head_dim: Optional[int] = None v_head_dim: Optional[int] = None + hidden_size: Optional[int] = None + # User-facing sparse config, lowered by production `to_sparse_params()`. The + # sparse sweep derives its other parameters from this and `is_mla`. + sparse_attention_config: Optional[SparseAttentionConfig] = None + + @property + def sparse_topk(self) -> Optional[int]: + """Per-token selection budget (``index_topk``) from the sparse config.""" + cfg = self.sparse_attention_config + if cfg is None: + return None + return cfg.index_topk # --------------------------------------------------------------------------- @@ -571,6 +588,30 @@ class ModelAttnConfig: # MLA (DeepSeek-style absorbed latent attention). num_kv_heads == 1 latent head. # --------------------------------------------------------------------------- _MLA = [ + # DeepSeek-V3.2 DSA uses absorbed MLA for both context and generation. + # The model's V projection is outside the standalone backend; the sparse + # backend itself returns kv_lora_rank-wide latent values per query head. + ModelAttnConfig( + "deepseekv3_2_dsa_mla", + "DeepSeek-V3.2", + num_heads=128, + num_kv_heads=1, + head_dim=192, + rope="gptj", + is_mla=True, + kv_lora_rank=512, + q_lora_rank=1536, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + hidden_size=7168, + sparse_attention_config=DeepSeekSparseAttentionConfig( + index_n_heads=64, + index_head_dim=128, + index_topk=128, + skip_indexer_for_short_seqs=False, + ), + ), # DeepSeek-V3: 128 Q heads, qk_nope=128, qk_rope=64, kv_lora=512, v=128. ModelAttnConfig( "deepseekv3_mla", diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py index fec8746fcc48..94f5b1fbab7e 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py @@ -36,6 +36,7 @@ ) from tensorrt_llm._torch.attention_backend.sparse.dsa import DSABackendForwardArgs, DSACacheManager from tensorrt_llm._torch.attention_backend.utils import get_attention_backend +from tensorrt_llm._torch.attention_backend.vanilla import VanillaAttention from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str @@ -47,268 +48,6 @@ from tensorrt_llm.quantization.mode import QuantAlgo -# Copied from transformers.models.llama.modeling_llama.rotate_half -def rotate_half(x): - """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -def _rotate_k_pe_for_ctx( - k_pe: torch.Tensor, rope_cos_sin: torch.Tensor, sequence_lengths: List[int] -) -> torch.Tensor: - k_pe_ref_list = [] - total_tokens = 0 - for seq_len in sequence_lengths: - k_pe_seq = k_pe[total_tokens : total_tokens + seq_len].unsqueeze(-2) - cos, sin = rope_cos_sin[:seq_len].chunk(2, dim=-2) - k_pe_seq = k_pe_seq.unflatten(-1, [-1, 2]).transpose(-2, -1).flatten(start_dim=-2) - k_pe_seq = ((k_pe_seq * cos) + (rotate_half(k_pe_seq) * sin)).to(dtype=k_pe_seq.dtype) - k_pe_seq = k_pe_seq.unflatten(-1, [2, -1]).transpose(-2, -1).flatten(start_dim=-2) - k_pe_ref_list.append(k_pe_seq) - total_tokens += seq_len - return torch.cat(k_pe_ref_list).squeeze(-2) - - -def _rotate_fused_q_for_ctx( - fused_q: torch.Tensor, - rope_cos_sin: torch.Tensor, - sequence_lengths: List[int], - num_heads: int, - kv_lora_rank: int, - qk_rope_head_dim: int, -) -> torch.Tensor: - fused_q = fused_q.clone() - fused_head_dim = kv_lora_rank + qk_rope_head_dim - total_tokens = 0 - for seq_len in sequence_lengths: - fused_q_seq = fused_q[total_tokens : total_tokens + seq_len].view( - seq_len, num_heads, fused_head_dim - ) - q_rope = fused_q_seq[..., -qk_rope_head_dim:] - cos, sin = rope_cos_sin[:seq_len].chunk(2, dim=-2) - q_rope = q_rope.unflatten(-1, [-1, 2]).transpose(-2, -1).flatten(start_dim=-2) - q_rope = ((q_rope * cos) + (rotate_half(q_rope) * sin)).to(dtype=fused_q.dtype) - q_rope = q_rope.unflatten(-1, [2, -1]).transpose(-2, -1).flatten(start_dim=-2) - fused_q_seq[..., -qk_rope_head_dim:] = q_rope - fused_q[total_tokens : total_tokens + seq_len] = fused_q_seq.view(seq_len, -1) - total_tokens += seq_len - return fused_q - - -def calculate_ref_result_ctx_sparse( - fused_q: torch.Tensor, - latent_cache: torch.Tensor, - sequence_lengths: List[int], - num_heads: int, - kv_lora_rank: int, - v_head_dim: int, - qk_nope_head_dim: int, - qk_rope_head_dim: int, - q_scaling: float, - topk_indices: Optional[torch.Tensor] = None, -): - """ - Reference for sparse MLA context using fused Q and latent cache. - fused_q shape: (total_tokens, num_heads * (kv_lora_rank + qk_rope_head_dim)) - latent_cache shape: (total_tokens, kv_lora_rank + qk_rope_head_dim) - """ - qk_head_dim = qk_nope_head_dim + qk_rope_head_dim - bmm1_scale = 1 / (math.sqrt(qk_head_dim) * q_scaling) - fused_head_dim = kv_lora_rank + qk_rope_head_dim - ref_results = [] - total_tokens = 0 - for seq_len in sequence_lengths: - fused_q_seq = fused_q[total_tokens : total_tokens + seq_len].unflatten( - -1, [num_heads, fused_head_dim] - ) - fused_q_seq = fused_q_seq.transpose(0, 1) # (num_heads, seq_len, fused_head_dim) - - latent_seq = latent_cache[total_tokens : total_tokens + seq_len] - k_seq = latent_seq.unsqueeze(0) # (1, seq_len, fused_head_dim) - v_seq = latent_seq[..., :v_head_dim].unsqueeze(0) # (1, seq_len, v_head_dim) - - k_seq = repeat_kv(k_seq.unsqueeze(0), num_heads).squeeze(0) - v_seq = repeat_kv(v_seq.unsqueeze(0), num_heads).squeeze(0) - - if topk_indices is None: - attn_weights = torch.matmul(fused_q_seq, k_seq.transpose(1, 2)) * bmm1_scale - causal_mask = torch.triu( - torch.ones(seq_len, seq_len, device=fused_q.device, dtype=torch.bool), diagonal=1 - ) - attn_weights = attn_weights.masked_fill(causal_mask, float("-inf")) - attn_weights = torch.nn.functional.softmax( - attn_weights, dim=-1, dtype=torch.float32 - ).to(fused_q.dtype) - attn_output = torch.matmul(attn_weights, v_seq) # (num_heads, seq_len, v_head_dim) - attn_output = ( - attn_output.transpose(0, 1).contiguous().view(seq_len, num_heads * v_head_dim) - ) - ref_results.append(attn_output) - else: - per_token_outputs = [] - token_rows = topk_indices[total_tokens : total_tokens + seq_len] - for token_idx in range(seq_len): - token_indices = token_rows[token_idx] - token_indices = token_indices[token_indices >= 0] - q_tok = fused_q_seq[:, token_idx, :] - k_sel = k_seq[:, token_indices, :] - v_sel = v_seq[:, token_indices, :] - attn_weights = ( - torch.matmul( - q_tok.unsqueeze(1), - k_sel.transpose(1, 2), - ) - * bmm1_scale - ) - attn_weights = torch.nn.functional.softmax( - attn_weights, dim=-1, dtype=torch.float32 - ).to(fused_q.dtype) - attn_output = torch.matmul(attn_weights, v_sel) # (num_heads, 1, v_head_dim) - per_token_outputs.append( - attn_output.transpose(0, 1).contiguous().view(1, num_heads * v_head_dim) - ) - ref_results.append(torch.cat(per_token_outputs, dim=0)) - total_tokens += seq_len - return torch.cat(ref_results) - - -def calculate_ref_result_gen( - fused_q: torch.Tensor, - q_pe: torch.Tensor, - compressed_kv: torch.Tensor, - k_pe: torch.Tensor, - latent_cache: torch.Tensor, - rope_cos_sin: torch.Tensor, - num_heads: int, - kv_lora_rank: int, - v_head_dim: int, - qk_nope_head_dim: int, - qk_rope_head_dim: int, - sequence_lengths: List[int], - q_scaling: float, - topk_indices: Optional[torch.Tensor] = None, -): - """ - use standard attention to calculate the reference result by iterating over each request - 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) - compressed_kv shape: (num_requests, kv_lora_rank) - k_pe shape: (num_requests, qk_rope_head_dim) - latent_cache shape: (total_tokens, kv_lora_rank + qk_rope_head_dim) - rope_cos_sin shape: (max_position_embeddings, 2, qk_rope_head_dim) - """ - num_requests = len(sequence_lengths) - seq_len_q = fused_q.shape[0] // num_requests - - # Reshape inputs for reference calculation - q_reshaped = [] - k_reshaped = [] - v_reshaped = [] - latent_cache_list = [] - total_tokens = 0 - for i in range(num_requests): - fused_q_seq = fused_q[i * seq_len_q : (i + 1) * seq_len_q].unflatten( - -1, [num_heads, kv_lora_rank + qk_rope_head_dim] - ) - q_pe_seq = q_pe[i * seq_len_q : (i + 1) * seq_len_q] - compressed_kv_seq = compressed_kv[i * seq_len_q : (i + 1) * seq_len_q].unsqueeze(-2) - k_pe_seq = k_pe[i * seq_len_q : (i + 1) * seq_len_q].unsqueeze(-2) - latent_cache_seq = latent_cache[ - total_tokens : total_tokens + sequence_lengths[i] - ].unsqueeze(-2) - - cos, sin = rope_cos_sin[sequence_lengths[i] : sequence_lengths[i] + seq_len_q].chunk( - 2, dim=-2 - ) - q_pe_seq = q_pe_seq.unflatten(-1, [-1, 2]).transpose(-2, -1).flatten(start_dim=-2) - k_pe_seq = k_pe_seq.unflatten(-1, [-1, 2]).transpose(-2, -1).flatten(start_dim=-2) - q_pe_seq = ((q_pe_seq * cos) + (rotate_half(q_pe_seq) * sin)).to(dtype=q_pe_seq.dtype) - k_pe_seq = ((k_pe_seq * cos) + (rotate_half(k_pe_seq) * sin)).to(dtype=k_pe_seq.dtype) - q_pe_seq = q_pe_seq.unflatten(-1, [2, -1]).transpose(-2, -1).flatten(start_dim=-2) - k_pe_seq = k_pe_seq.unflatten(-1, [2, -1]).transpose(-2, -1).flatten(start_dim=-2) - fused_q_seq[..., -qk_rope_head_dim:] = q_pe_seq - latent_cache_seq = torch.cat( - [latent_cache_seq, torch.cat([compressed_kv_seq, k_pe_seq], dim=-1)], dim=0 - ) - latent_cache_list.append(latent_cache_seq) - - q_reshaped.append( - fused_q_seq.transpose(0, 1) - ) # (num_heads, seq_len_q, kv_lora_rank + qk_rope_head_dim) - k_reshaped.append( - latent_cache_seq.transpose(0, 1) - ) # (1, seq_len_kv, kv_lora_rank + qk_rope_head_dim) - v_reshaped.append( - latent_cache_seq[..., :v_head_dim].transpose(0, 1) - ) # (1, seq_len_kv, v_head_dim) - - total_tokens += sequence_lengths[i] - - # Calculate reference result batch by batch - ref_results = [] - for i in range(num_requests): - q = q_reshaped[i] # (num_heads, seq_len_q, kv_lora_rank + qk_rope_head_dim) - k = k_reshaped[i] # (1, seq_len_kv, kv_lora_rank + qk_rope_head_dim) - v = v_reshaped[i] # (1, seq_len_kv, v_head_dim) - - # Handle grouped-query attention - k = repeat_kv(k.unsqueeze(0), num_heads).squeeze(0) - v = repeat_kv(v.unsqueeze(0), num_heads).squeeze(0) - - seq_len_q = q.shape[1] - seq_len_kv = k.shape[1] - if topk_indices is None: - # Compute attention scores - attn_weights = torch.matmul(q, k.transpose(1, 2)) / ( - q_scaling * math.sqrt(qk_nope_head_dim + qk_rope_head_dim) - ) - - # Use MTP mask by default if seqlen_q > 1. - mask = torch.zeros(seq_len_q, seq_len_kv, device=q.device, dtype=torch.bool) - for qi in range(seq_len_q): - for ki in range(seq_len_kv - seq_len_q + 1 + qi, seq_len_kv): - mask[qi, ki] = 1 - attn_weights = attn_weights.masked_fill(mask, float("-inf")) - # Apply softmax to get attention probabilities - attn_weights = torch.nn.functional.softmax( - attn_weights, dim=-1, dtype=torch.float32 - ).to(q.dtype) - - # Apply attention weights to values - attn_output = torch.matmul(attn_weights, v) # (num_heads, 1, v_head_dim) - - # Reshape back to (seq_len_q, num_heads*v_head_dim) - attn_output = attn_output.transpose(0, 1).contiguous().view(-1, num_heads * v_head_dim) - ref_results.append(attn_output) - else: - per_token_outputs = [] - for qi in range(seq_len_q): - row = i * seq_len_q + qi - token_indices = topk_indices[row] - token_indices = token_indices[token_indices >= 0] - q_tok = q[:, qi, :] - k_sel = k[:, token_indices, :] - v_sel = v[:, token_indices, :] - attn_weights = torch.matmul( - q_tok.unsqueeze(1), - k_sel.transpose(1, 2), - ) / (q_scaling * math.sqrt(qk_nope_head_dim + qk_rope_head_dim)) - attn_weights = torch.nn.functional.softmax( - attn_weights, dim=-1, dtype=torch.float32 - ).to(q.dtype) - attn_output = torch.matmul(attn_weights, v_sel) # (num_heads, 1, v_head_dim) - per_token_outputs.append( - attn_output.transpose(0, 1).contiguous().view(1, num_heads * v_head_dim) - ) - ref_results.append(torch.cat(per_token_outputs, dim=0)) - - ref_result = torch.cat(ref_results) - latent_cache = torch.cat(latent_cache_list).squeeze(-2) - return ref_result, latent_cache - - @dataclass(kw_only=True, frozen=True) class Scenario: dtype: torch.dtype = torch.bfloat16 @@ -357,22 +96,78 @@ class RopeConfig: model_type: str = "deepseek_v3" -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, - num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) +def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def _apply_yarn_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + x = x.unflatten(-1, [-1, 2]).transpose(-2, -1).flatten(start_dim=-2) + x = ((x * cos) + (rotate_half(x) * sin)).to(dtype) + return x.unflatten(-1, [2, -1]).transpose(-2, -1).flatten(start_dim=-2) + + +def _rope_qk_for_vanilla( + fused_q: torch.Tensor, + latent_cache: torch.Tensor, + rope_cos_sin: torch.Tensor, + seq_lens: List[int], + past_lens: List[int], + num_heads: int, + qk_rope_head_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """RoPE-pre-apply the yarn positional embedding to the Vanilla golden's q_pe/k_pe. + + Vanilla applies no positional embedding itself while the kernel does via ``rope_append``; + the query rope dims (in ``fused_q``) and the k_pe tail (in ``latent_cache``) are rotated at + absolute positions ``[past, past + seq_len)`` per request. """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand( - batch, num_key_value_heads, n_rep, slen, head_dim + fused_q = fused_q.clone() + latent_cache = latent_cache.clone() + fused_head_dim = latent_cache.shape[-1] + off = 0 + for seq_len, past in zip(seq_lens, past_lens): + cos, sin = rope_cos_sin[past : past + seq_len].chunk(2, dim=-2) + q = fused_q[off : off + seq_len].view(seq_len, num_heads, fused_head_dim) + q[..., -qk_rope_head_dim:] = _apply_yarn_rope(q[..., -qk_rope_head_dim:], cos, sin) + fused_q[off : off + seq_len] = q.view(seq_len, -1) + kv = latent_cache[off : off + seq_len] + kv[..., -qk_rope_head_dim:] = _apply_yarn_rope( + kv[..., -qk_rope_head_dim:].unsqueeze(-2), cos, sin + ).squeeze(-2) + off += seq_len + return fused_q, latent_cache + + +def _yarn_rope_cos_sin(rope_config, device: torch.device) -> torch.Tensor: + return ( + torch.tensor( + RopeEmbeddingUtils.create_sinusoidal_positions_yarn( + rope_config.max_position_embeddings, + rope_config.qk_rope_head_dim, + rope_config.rope_theta, + rope_config.rope_scaling["factor"], + rope_config.rope_scaling["original_max_position_embeddings"], + rope_config.rope_scaling["beta_fast"], + rope_config.rope_scaling["beta_slow"], + rope_config.rope_scaling["mscale"], + rope_config.rope_scaling["mscale_all_dim"], + )[1], + dtype=torch.float32, + device=device, + ) + .reshape(rope_config.max_position_embeddings, -1, 2) + .transpose(-2, -1) ) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) def _build_sparse_topk_indices_context( - seq_lens: List[int], topk: int, device: torch.device + seq_lens: List[int], + topk: int, + device: torch.device, + generator: Optional[torch.Generator] = None, ) -> torch.Tensor: total_tokens = sum(seq_lens) topk_indices = torch.full((total_tokens, topk), -1, dtype=torch.int32, device=device) @@ -381,7 +176,7 @@ def _build_sparse_topk_indices_context( for token_idx in range(seq_len): max_index = token_idx valid_len = min(max_index + 1, topk) - indices = torch.randperm(max_index + 1, device=device)[:valid_len] + indices = torch.randperm(max_index + 1, device=device, generator=generator)[:valid_len] indices, _ = torch.sort(indices) topk_indices[token_offset + token_idx, :valid_len] = indices.to(torch.int32) token_offset += seq_len @@ -389,7 +184,11 @@ def _build_sparse_topk_indices_context( def _build_sparse_topk_indices_generation( - cached_lens: List[int], seq_len_q: int, topk: int, device: torch.device + cached_lens: List[int], + seq_len_q: int, + topk: int, + device: torch.device, + generator: Optional[torch.Generator] = None, ) -> torch.Tensor: total_tokens = len(cached_lens) * seq_len_q topk_indices = torch.full((total_tokens, topk), -1, dtype=torch.int32, device=device) @@ -398,7 +197,7 @@ def _build_sparse_topk_indices_generation( for q_idx in range(seq_len_q): max_index = cached_len + q_idx valid_len = min(max_index + 1, topk) - indices = torch.randperm(max_index + 1, device=device)[:valid_len] + indices = torch.randperm(max_index + 1, device=device, generator=generator)[:valid_len] indices, _ = torch.sort(indices) topk_indices[row, :valid_len] = indices.to(torch.int32) row += 1 @@ -444,6 +243,26 @@ def _allocate_kv_cache_for_generation(kv_cache_manager, request_ids, num_tokens: SPARSE_TOPK = 2048 +def _assert_matches_vanilla( + actual: dict[str, torch.Tensor], + golden: dict[str, torch.Tensor], + kv_cache_dtype: torch.dtype, + backend_name: str, +) -> None: + assert actual.keys() == golden.keys(), ( + f"[{backend_name} vs VANILLA golden] phase mismatch: {list(actual)} != {list(golden)}" + ) + atol, rtol = accuracy_dict[kv_cache_dtype] + for phase, golden_output in golden.items(): + torch.testing.assert_close( + actual[phase], + golden_output, + atol=atol, + rtol=rtol, + msg=lambda message: f"[{backend_name} vs VANILLA golden, {phase}]\n{message}", + ) + + # Convert parameterized tests to pytest parametrize @skip_pre_blackwell @pytest.mark.skip_less_device_memory(80000) @@ -465,7 +284,34 @@ def test_sparse_attention_mla( generation_seq_len_q: int, num_generation_steps: int, ): - """Test sparse MLA computation for both context and generation phases""" + """Compare TRTLLM sparse MLA against the VanillaAttention golden.""" + golden = _test_sparse_attention_mla( + "VANILLA", + scenario, + context_sequence_lengths, + generation_seq_len_q, + num_generation_steps, + ) + actual = _test_sparse_attention_mla( + "TRTLLM", + scenario, + context_sequence_lengths, + generation_seq_len_q, + num_generation_steps, + ) + _assert_matches_vanilla(actual, golden, scenario.kv_cache_dtype, "TRTLLM") + + +def _test_sparse_attention_mla( + backend_name: str, + scenario: Scenario, + context_sequence_lengths: List[int], + generation_seq_len_q: int, + num_generation_steps: int, + sparse_topk: int = SPARSE_TOPK, + seed: int = 123, + topk_seed: int = 456, +) -> dict[str, torch.Tensor]: num_heads = scenario.num_heads num_kv_heads = scenario.num_kv_heads q_lora_rank = scenario.q_lora_rank @@ -500,14 +346,14 @@ def test_sparse_attention_mla( dtype = scenario.dtype kv_cache_dtype = scenario.kv_cache_dtype - assert SPARSE_TOPK % 128 == 0 + assert sparse_topk % 128 == 0 print( f"--------------------------------Test for scenario: {scenario} start--------------------------------" ) - _run_test_for_backend( - "TRTLLM", + return _run_test_for_backend( + backend_name, num_heads, num_kv_heads, num_layers, @@ -525,6 +371,9 @@ def test_sparse_attention_mla( context_sequence_lengths, generation_seq_len_q, num_generation_steps, + sparse_topk, + seed, + topk_seed, ) @@ -547,20 +396,28 @@ def _run_test_for_backend( context_sequence_lengths, generation_seq_len_q, num_generation_steps, -): + sparse_topk, + seed, + topk_seed, +) -> dict[str, torch.Tensor]: sparse_config = DeepSeekSparseAttentionConfig( index_n_heads=64, index_head_dim=128, - index_topk=SPARSE_TOPK, + index_topk=sparse_topk, skip_indexer_for_short_seqs=False, ) - AttentionCls = get_attention_backend(backend_name, sparse_config) + is_vanilla = backend_name == "VANILLA" + AttentionCls = ( + VanillaAttention if is_vanilla else get_attention_backend(backend_name, sparse_config) + ) # When rope_append is False, [448: 512) are used for qk_rope_head_dim kv_lora_rank = kv_lora_rank - qk_rope_head_dim if not rope_append else kv_lora_rank head_dim = kv_lora_rank + qk_rope_head_dim + rope_cos_sin = _yarn_rope_cos_sin(rope_config, device) if is_vanilla else None # Set seed for reproducibility. - torch.manual_seed(123) + torch.manual_seed(seed) + topk_generator = torch.Generator(device=device).manual_seed(topk_seed) # Create inputs inputs_per_layer = [] @@ -682,26 +539,6 @@ def _run_test_for_backend( else: print(f"{key}: {val.shape}") - rope_cos_sin = ( - torch.tensor( - RopeEmbeddingUtils.create_sinusoidal_positions_yarn( - rope_config.max_position_embeddings, - rope_config.qk_rope_head_dim, - rope_config.rope_theta, - rope_config.rope_scaling["factor"], - rope_config.rope_scaling["original_max_position_embeddings"], - rope_config.rope_scaling["beta_fast"], - rope_config.rope_scaling["beta_slow"], - rope_config.rope_scaling["mscale"], - rope_config.rope_scaling["mscale_all_dim"], - )[1], - dtype=torch.float32, - device=device, - ) - .reshape(rope_config.max_position_embeddings, -1, 2) - .transpose(-2, -1) - ) - # Setup attention module and metadata pos_embd_params = PositionalEmbeddingParams( type=PositionEmbeddingType.yarn, @@ -732,34 +569,28 @@ def yarn_get_mscale(scale=1, mscale=1): if kv_cache_dtype == torch.float8_e4m3fn: quant_config = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8.value) - ctx_layers = [ - AttentionCls( - layer_idx=layer_idx, - num_heads=num_heads, - head_dim=head_dim, - num_kv_heads=num_kv_heads, - quant_config=quant_config, - q_scaling=q_scaling, - pos_embd_params=pos_embd_params, - mla_params=mla_params, - sparse_attention_config=sparse_config, + def create_layer(layer_idx: int, num_kv_heads: int): + sparse_kwargs = ( + {"sparse_params": sparse_config.to_sparse_params(layer_idx=layer_idx)} + if is_vanilla + else {"sparse_attention_config": sparse_config} ) - for layer_idx in range(num_layers) - ] - gen_layers = [ - AttentionCls( + return AttentionCls( layer_idx=layer_idx, num_heads=num_heads, head_dim=head_dim, - num_kv_heads=1, + num_kv_heads=num_kv_heads, quant_config=quant_config, q_scaling=q_scaling, pos_embd_params=pos_embd_params, mla_params=mla_params, - sparse_attention_config=sparse_config, + **sparse_kwargs, ) - for layer_idx in range(num_layers) - ] + + ctx_layers = [create_layer(layer_idx, num_kv_heads) for layer_idx in range(num_layers)] + gen_layers = [create_layer(layer_idx, 1) for layer_idx in range(num_layers)] + if is_vanilla: + assert all(type(layer) is VanillaAttention for layer in ctx_layers + gen_layers) # NOTE: set up metadata, refer to tensorrt_llm/_torch/pyexecutor/model_engine.py # all layers share the same metadata @@ -803,9 +634,11 @@ def yarn_get_mscale(scale=1, mscale=1): sparse_attn_config=sparse_config, model_config=model_config, ) + outputs = {} try: request_ids = list(range(max_num_contexts)) kv_cache_manager.add_dummy_requests(request_ids, context_sequence_lengths) + metadata_sparse_kwargs = {} if is_vanilla else {"sparse_attention_config": sparse_config} ctx_seq_lens = torch.tensor(context_sequence_lengths, dtype=torch.int) total_ctx_tokens = sum(context_sequence_lengths) @@ -822,12 +655,11 @@ def yarn_get_mscale(scale=1, mscale=1): num_cached_tokens_per_seq=[0 for _ in context_sequence_lengths], ), mapping=mapping, - sparse_attention_config=sparse_config, + **metadata_sparse_kwargs, ) attn_metadata.prepare() # run forward for each step and each layer - latent_cache_ref_all_list = [None for _ in range(num_layers)] for step in range(num_generation_steps + 1): if step > 0: _allocate_kv_cache_for_generation( @@ -854,7 +686,7 @@ def yarn_get_mscale(scale=1, mscale=1): ), mapping=mapping, enable_flash_mla=torch.cuda.get_device_capability() == (9, 0), - sparse_attention_config=sparse_config, + **metadata_sparse_kwargs, ) attn_metadata.prepare() for layer_idx in range(num_layers): @@ -866,167 +698,161 @@ def yarn_get_mscale(scale=1, mscale=1): latent_cache = torch.cat([compressed_kv, k_pe], dim=-1) q_pe = inputs_per_layer[layer_idx]["ctx_q_pe"] topk_indices = _build_sparse_topk_indices_context( - context_sequence_lengths, SPARSE_TOPK, device - ) - ctx_layers[layer_idx].indexer.forward_from_projected = Mock( - return_value=topk_indices + context_sequence_lengths, + sparse_topk, + device, + generator=topk_generator, ) + if is_vanilla: + # Vanilla has no indexer and applies no RoPE; feed the selection directly + # and pre-apply RoPE at context positions 0..len. + ctx_fused_q, ctx_latent = _rope_qk_for_vanilla( + fused_q, + latent_cache, + rope_cos_sin, + context_sequence_lengths, + [0] * len(context_sequence_lengths), + num_heads, + qk_rope_head_dim, + ) + ctx_sba = DSABackendForwardArgs( + indexer_intermediates=[], topk_indices=topk_indices + ) + else: + ctx_layers[layer_idx].indexer.forward_from_projected = Mock( + return_value=topk_indices + ) + ctx_fused_q, ctx_latent = fused_q, latent_cache + ctx_sba = DSABackendForwardArgs(indexer_intermediates=[]) result = ctx_layers[layer_idx].forward( - fused_q.clone(), + ctx_fused_q.clone(), None, None, attn_metadata, attention_input_type=AttentionInputType.context_only, - latent_cache=latent_cache, + latent_cache=ctx_latent.clone(), q_pe=q_pe, - sparse_backend_args=DSABackendForwardArgs(indexer_intermediates=[]), - ) - k_pe_ref = _rotate_k_pe_for_ctx(k_pe, rope_cos_sin, context_sequence_lengths) - latent_cache_ref = torch.cat([compressed_kv, k_pe_ref], dim=-1) - fused_q_rot = _rotate_fused_q_for_ctx( - fused_q, - rope_cos_sin, - context_sequence_lengths, - num_heads, - kv_lora_rank, - qk_rope_head_dim, - ) - ref_result = calculate_ref_result_ctx_sparse( - fused_q_rot, - latent_cache_ref, - context_sequence_lengths, - num_heads, - kv_lora_rank, - v_head_dim, - qk_nope_head_dim, - qk_rope_head_dim, - q_scaling, - topk_indices=topk_indices, + sparse_backend_args=ctx_sba, ) - latent_cache_ref_all_list[layer_idx] = latent_cache_ref else: fused_q = inputs_per_layer[layer_idx]["gen_fused_q_list"][step - 1] q_pe = inputs_per_layer[layer_idx]["gen_q_pe_list"][step - 1] compressed_kv = inputs_per_layer[layer_idx]["gen_compressed_kv_list"][step - 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=fused_q.device - ) - cu_kv_seqlens = torch.empty( - num_seqs + 1, dtype=torch.int32, device=fused_q.device - ) - fmha_scheduler_counter = torch.empty( - 1, dtype=torch.uint32, device=fused_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=fused_q.device) - mla_bmm2_scale = torch.empty(1, dtype=torch.float32, device=fused_q.device) - quant_q_buffer = torch.empty( - num_tokens, - num_heads * head_dim, - dtype=torch.uint8, - device=fused_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, - latent_cache, - attn_metadata, - cu_q_seqlens, - cu_kv_seqlens, - fmha_scheduler_counter, - mla_bmm1_scale, - mla_bmm2_scale, - quant_q_buffer, - ) - cached_lens = [ ctx_len + (step - 1) * generation_seq_len_q for ctx_len in context_sequence_lengths ] topk_indices = _build_sparse_topk_indices_generation( - cached_lens, generation_seq_len_q, SPARSE_TOPK, device - ) - gen_layers[layer_idx].indexer.forward_from_projected = Mock( - return_value=topk_indices + cached_lens, + generation_seq_len_q, + sparse_topk, + device, + generator=topk_generator, ) + if is_vanilla: + # Pre-apply RoPE at absolute positions [cached_len, cached_len+seq_q). + seq_q = [generation_seq_len_q] * len(context_sequence_lengths) + backend_fused_q, backend_latent_cache = _rope_qk_for_vanilla( + fused_q, + latent_cache, + rope_cos_sin, + seq_q, + cached_lens, + num_heads, + qk_rope_head_dim, + ) + generation_kwargs = { + "sparse_backend_args": DSABackendForwardArgs( + indexer_intermediates=[], topk_indices=topk_indices + ) + } + else: + 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=fused_q.device + ) + cu_kv_seqlens = torch.empty( + num_seqs + 1, dtype=torch.int32, device=fused_q.device + ) + fmha_scheduler_counter = torch.empty( + 1, dtype=torch.uint32, device=fused_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=fused_q.device + ) + mla_bmm2_scale = torch.empty( + 1, dtype=torch.float32, device=fused_q.device + ) + quant_q_buffer = torch.empty( + num_tokens, + num_heads * head_dim, + dtype=torch.uint8, + device=fused_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, + latent_cache, + attn_metadata, + cu_q_seqlens, + cu_kv_seqlens, + fmha_scheduler_counter, + mla_bmm1_scale, + mla_bmm2_scale, + quant_q_buffer, + ) + gen_layers[layer_idx].indexer.forward_from_projected = Mock( + return_value=topk_indices + ) + backend_fused_q = fused_q + backend_latent_cache = latent_cache + generation_kwargs = { + "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, + "sparse_backend_args": DSABackendForwardArgs(indexer_intermediates=[]), + } result = gen_layers[layer_idx].forward( - fused_q, + backend_fused_q, None, None, attn_metadata, attention_input_type=AttentionInputType.generation_only, - latent_cache=latent_cache, + latent_cache=backend_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, - sparse_backend_args=DSABackendForwardArgs(indexer_intermediates=[]), + **generation_kwargs, ) - ref_result, latent_cache_ref = calculate_ref_result_gen( - fused_q, - q_pe, - compressed_kv, - k_pe, - latent_cache_ref_all_list[layer_idx], - rope_cos_sin, - num_heads, - kv_lora_rank, - v_head_dim, - qk_nope_head_dim, - qk_rope_head_dim, - [ - ctx_len + (step - 1) * generation_seq_len_q - for ctx_len in context_sequence_lengths - ], - q_scaling, - topk_indices=topk_indices, - ) - latent_cache_ref_all_list[layer_idx] = latent_cache_ref - # Compare results + # Record results for the Vanilla-golden comparison. print( f"{backend_name} output mean: {result.abs().mean().item()}, max: {result.abs().max().item()}" ) - print( - f"Reference output mean: {ref_result.abs().mean().item()}, max: {ref_result.abs().max().item()}" - ) - print( - f"Difference mean: {(result - ref_result).abs().mean().item()}, \ - max: {(result - ref_result).abs().max().item()}" - ) - - # Assert results are close - atol, rtol = accuracy_dict[kv_cache_dtype] - assert torch.allclose(result, ref_result, atol=atol, rtol=rtol), ( - f"Results for sparse MLA in {backend_name} backend don't match reference implementation \ - at layer {layer_idx} in step {step}" - ) - print( f"Test for sparse MLA in {backend_name} backend passed at layer {layer_idx} in step {step}" ) print(f"---- step {step} layer {layer_idx} end ----") + phase = "context" if step == 0 else f"generation_{step - 1}" + outputs[f"{phase}_layer_{layer_idx}"] = result.detach().clone() print(f"Test for sparse MLA in {backend_name} backend passed") + return outputs finally: kv_cache_manager.shutdown() diff --git a/tests/unittest/_torch/attention/test_attention_backends.py b/tests/unittest/_torch/attention/test_attention_backends.py index ca00dcd21125..a752276a39d0 100644 --- a/tests/unittest/_torch/attention/test_attention_backends.py +++ b/tests/unittest/_torch/attention/test_attention_backends.py @@ -50,8 +50,8 @@ def get_long_seq_len(window: int) -> int: return window + 17 -def _phases_from_window(window: int) -> dict: - long_len = get_long_seq_len(window) +def _phases(long_len: int, gen_len: int) -> dict: + """The ctx/gen/mix batch shapes, parameterized by context and generation len.""" return { "ctx": dict( seq_lens=[long_len, 73, 41], @@ -59,25 +59,38 @@ def _phases_from_window(window: int) -> dict: num_contexts=3, ), "gen": dict( - seq_lens=[1, 1, 1], + seq_lens=[gen_len] * 3, num_cached_tokens=[long_len, 73, 41], num_contexts=0, ), "mix": dict( - seq_lens=[long_len, 1, 1], + seq_lens=[long_len, gen_len, gen_len], num_cached_tokens=[0, long_len, 73], num_contexts=1, ), } +def _phases_from_window(window: int) -> dict: + return _phases(get_long_seq_len(window), gen_len=1) + + # Standard self-attention batch phases. Non-sliding cases use a nominal window # only to choose non-tiny, non-power-of-two lengths; the backend still receives # sliding_window=None. _PHASES = _phases_from_window(_NON_SLIDING_PHASE_WINDOW) +# Model-agnostic sparse sweep dimensions, shared by every sparse config. +_SPARSE_COMPUTE_DTYPE = "bfloat16" +_SPARSE_KV_LAYOUT = "HND" +_SPARSE_PAGE_SIZE = 64 +_SPARSE_USE_KVM_V2 = False + + def _phases_for(cfg: ModelAttnConfig) -> dict: + if cfg.sparse_attention_config is not None: + return _phases(cfg.sparse_topk + 32, gen_len=1) if cfg.mask != "sliding": return _PHASES @@ -117,7 +130,10 @@ def _common(cfg: ModelAttnConfig) -> dict: qk_nope_head_dim=cfg.qk_nope_head_dim, qk_rope_head_dim=cfg.qk_rope_head_dim, v_head_dim=cfg.v_head_dim, + hidden_size=cfg.hidden_size, ) + if cfg.sparse_attention_config is not None: + common.update(sparse_attention_config=cfg.sparse_attention_config) return common @@ -140,6 +156,30 @@ def _expand(cfg: ModelAttnConfig, precisions, kv_layouts, page_sizes): common = _common(cfg) phases = _phases_for(cfg) + # Sparse cases use one model-agnostic sweep (bf16 latent cache, fixed layout/ + # page/manager). Each phase mixes singleton and random selection rows in a + # single execution; the injected selections isolate backend execution from + # the separately-tested algorithm/indexer. + if cfg.sparse_attention_config is not None: + manager = "v2" if _SPARSE_USE_KVM_V2 else "v1" + tag = ( + f"{_prec_tag(_SPARSE_COMPUTE_DTYPE, None)}-{_SPARSE_KV_LAYOUT}" + f"-p{_SPARSE_PAGE_SIZE}-{manager}" + ) + for phase_name in ("ctx", "gen", "mix"): + yield ( + f"{cfg.id}-{phase_name}-{tag}", + BackendCase( + page_size=_SPARSE_PAGE_SIZE, + kv_layout=_SPARSE_KV_LAYOUT, + dtype=_SPARSE_COMPUTE_DTYPE, + use_kv_cache_manager_v2=_SPARSE_USE_KVM_V2, + **phases[phase_name], + **common, + ), + ) + return + # Bidirectional, KV-cache-free DiT / encoder workloads: only compute dtype. if cfg.no_cache: for dtype, kvd in precisions: