From 2caa41575b33a83c934e575fc3f545abaa4db4cb Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Mon, 13 Jul 2026 00:28:16 -0700 Subject: [PATCH 01/17] [None][feat] Add DSA sparse attention to VanillaAttention Signed-off-by: Yihan Wang --- .../_torch/attention_backend/vanilla.py | 294 ++++++++- .../sparse/dsa/test_dsa_attention_backends.py | 52 ++ .../sparse/dsa/test_dsa_sparse_mla.py | 619 +++++------------- .../sparse/dsa/test_dsa_vanilla_attention.py | 260 ++++++++ 4 files changed, 778 insertions(+), 447 deletions(-) create mode 100644 tests/unittest/_torch/attention/sparse/dsa/test_dsa_attention_backends.py create mode 100644 tests/unittest/_torch/attention/sparse/dsa/test_dsa_vanilla_attention.py diff --git a/tensorrt_llm/_torch/attention_backend/vanilla.py b/tensorrt_llm/_torch/attention_backend/vanilla.py index eef2fc520906..2c12d0a551e5 100644 --- a/tensorrt_llm/_torch/attention_backend/vanilla.py +++ b/tensorrt_llm/_torch/attention_backend/vanilla.py @@ -16,7 +16,8 @@ from .interface import (AttentionBackend, AttentionForwardArgs, AttentionInputType, AttentionMask, AttentionMetadata, - PredefinedAttentionMask, merge_attention_forward_args) + PositionalEmbeddingParams, PredefinedAttentionMask, + merge_attention_forward_args) from .sparse.kernel import triton_index_gather from .sparse.params import SparseParams @@ -94,6 +95,7 @@ def __init__( num_kv_heads: Optional[int] = None, quant_config: Optional[QuantConfig] = None, q_scaling: Optional[float] = None, + pos_embd_params: Optional[PositionalEmbeddingParams] = None, sparse_params: Optional[SparseParams] = None, **kwargs, ): @@ -114,10 +116,96 @@ def __init__( self.qk_nope_head_dim = mla_params.qk_nope_head_dim self.v_head_dim = mla_params.v_head_dim + self.dsa_rope_cos_sin = None + self.dsa_rope_is_neox = True + if (self.is_mla_enable + and getattr(self.sparse_params, "algorithm", None) == "dsa" + and pos_embd_params is not None + and pos_embd_params.rope is not None): + self.dsa_rope_cos_sin = pos_embd_params.rope.create_rope_const_params( + interleave=False)[1].reshape(pos_embd_params.rope.max_positions, + 2, -1) + self.dsa_rope_is_neox = pos_embd_params.is_neox + @classmethod def support_mla(cls) -> bool: return True + @staticmethod + def _apply_rotary_embedding(x: torch.Tensor, cos: torch.Tensor, + sin: torch.Tensor, + is_neox: bool) -> torch.Tensor: + """Apply RoPE to ``x`` using one cos/sin row per packed token.""" + cos = cos.to(device=x.device, dtype=x.dtype).unsqueeze(1) + sin = sin.to(device=x.device, dtype=x.dtype).unsqueeze(1) + rotary_dim = cos.shape[-1] * 2 + x_rotary, x_pass = x[..., :rotary_dim], x[..., rotary_dim:] + if is_neox: + x1, x2 = x_rotary.chunk(2, dim=-1) + else: + x1, x2 = x_rotary[..., ::2], x_rotary[..., 1::2] + out1 = x1 * cos - x2 * sin + out2 = x2 * cos + x1 * sin + if is_neox: + rotated = torch.cat((out1, out2), dim=-1) + else: + rotated = torch.stack((out1, out2), dim=-1).flatten(-2) + return torch.cat((rotated, x_pass), dim=-1) + + def _prepare_dsa_mla_inputs( + self, + fused_q: torch.Tensor, + latent_cache: torch.Tensor, + q_pe: Optional[torch.Tensor], + positions: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Apply DSA MLA RoPE to raw packed query and latent-cache inputs. + + As with the other attention paths, omitting positional-embedding + parameters means the caller already applied RoPE. + """ + if self.dsa_rope_cos_sin is None: + return fused_q, latent_cache + + if positions.numel() == 0: + return fused_q, latent_cache + max_position = int(positions.max().item()) + if max_position >= self.dsa_rope_cos_sin.shape[0]: + raise ValueError( + f"DSA position {max_position} exceeds the configured RoPE table " + f"size {self.dsa_rope_cos_sin.shape[0]}") + + num_tokens = fused_q.shape[0] + fused_head_dim = self.kv_lora_rank + self.qk_rope_head_dim + query = fused_q.view(num_tokens, self.num_heads, fused_head_dim).clone() + if q_pe is None: + raise ValueError( + "Vanilla DSA requires raw q_pe when RoPE parameters are configured" + ) + expected_numel = num_tokens * self.num_heads * self.qk_rope_head_dim + if q_pe.numel() != expected_numel: + raise ValueError( + f"DSA q_pe has {q_pe.numel()} elements, expected {expected_numel}" + ) + query_rope = q_pe.reshape(num_tokens, self.num_heads, + self.qk_rope_head_dim) + + if latent_cache.shape[1] != fused_head_dim: + raise ValueError( + f"DSA latent cache width must be {fused_head_dim}, got " + f"{latent_cache.shape[1]}") + latent_cache = latent_cache.clone() + key_rope = latent_cache[:, self.kv_lora_rank:].unsqueeze(1) + cos_sin = self.dsa_rope_cos_sin.index_select( + 0, + positions.to(device=self.dsa_rope_cos_sin.device, dtype=torch.long)) + cos, sin = cos_sin.unbind(dim=1) + query[..., -self.qk_rope_head_dim:] = self._apply_rotary_embedding( + query_rope, cos, sin, self.dsa_rope_is_neox) + latent_cache[:, self.kv_lora_rank:] = self._apply_rotary_embedding( + key_rope, cos, sin, self.dsa_rope_is_neox).squeeze(1) + return query.view(num_tokens, -1), latent_cache + def _single_request_sparse_attn_predict( self, q: torch.Tensor, k: Optional[torch.Tensor], v: Optional[torch.Tensor], kv_cache_tensor: torch.Tensor, @@ -571,6 +659,191 @@ 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_dsa( + self, + fused_q: torch.Tensor, + metadata: VanillaAttentionMetadata, + latent_cache: torch.Tensor, + q_pe: Optional[torch.Tensor], + topk_indices: torch.Tensor, + attention_input_type: AttentionInputType, + ) -> torch.Tensor: + """Run DSA selected attention from caller-provided local top-k rows. + + DSA's indexer owns selection. This golden consumes its request-local + token positions, gathers the selected latent K/V, and performs the + absorbed MLA attention directly in PyTorch. + """ + 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 + + phase_token_start = sum(seq_lens[:seq_start]) + if metadata.position_ids is not None: + positions = metadata.position_ids.reshape( + -1)[phase_token_start:phase_token_start + num_phase_tokens].to( + device=fused_q.device, dtype=torch.long) + if positions.numel() != num_phase_tokens: + raise ValueError( + "DSA metadata does not provide one position ID per phase token" + ) + else: + positions = torch.cat([ + torch.arange(int(past), + int(past) + q_len, + device=fused_q.device, + dtype=torch.long) for past, q_len in zip( + phase_past_tokens, phase_seq_lens, strict=True) + ]) + fused_q, latent_cache = self._prepare_dsa_mla_inputs( + fused_q, latent_cache, q_pe, positions) + + 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, @@ -635,6 +908,25 @@ 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.") + sparse_algorithm = getattr(self.sparse_params, "algorithm", None) + if forward_args.topk_indices is not None: + if sparse_algorithm != "dsa": + raise ValueError( + "Vanilla selected MLA currently supports only DSA") + if k is not None or v is not None: + raise ValueError( + "Vanilla DSA expects absorbed queries and latent cache, " + "not explicit K/V tensors") + return self._mla_forward_dsa( + q, + metadata, + forward_args.latent_cache, + forward_args.q_pe, + forward_args.topk_indices, + forward_args.attention_input_type, + ) + if sparse_algorithm == "dsa": + raise ValueError("Vanilla DSA requires topk_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/tests/unittest/_torch/attention/sparse/dsa/test_dsa_attention_backends.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_attention_backends.py new file mode 100644 index 000000000000..6ada2a406dcd --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_attention_backends.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Differential DSA backend tests using VanillaAttention as the golden.""" + +import pytest +from utils.util import skip_pre_blackwell + +from .test_dsa_sparse_mla import _assert_matches_vanilla, _test_sparse_attention_mla, scenarios + +BACKENDS_UNDER_TEST = ("TRTLLM",) + +DSA_CASES = { + "bf16-paged-context-generation": dict( + scenario=scenarios[0], + context_sequence_lengths=[160], + generation_seq_len_q=1, + num_generation_steps=2, + sparse_topk=128, + seed=123, + topk_seed=456, + ), +} + + +@skip_pre_blackwell +@pytest.mark.parametrize("name", list(DSA_CASES), ids=lambda name: name) +def test_dsa_attention_backend(name: str) -> None: + """Run Vanilla first, then compare every production DSA phase to it.""" + case = DSA_CASES[name] + golden = _test_sparse_attention_mla("VANILLA", **case) + + for backend in BACKENDS_UNDER_TEST: + actual = _test_sparse_attention_mla(backend, **case) + _assert_matches_vanilla( + actual, + golden, + case["scenario"].kv_cache_dtype, + backend, + ) 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 efbe2e2e56d2..f4a0820b75a6 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 @@ -35,279 +35,18 @@ ) from tensorrt_llm._torch.attention_backend.sparse.dsa import 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 from tensorrt_llm.bindings.executor import KvCacheConfig -from tensorrt_llm.functional import PositionEmbeddingType, RopeEmbeddingUtils +from tensorrt_llm.functional import PositionEmbeddingType from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig 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 @@ -356,22 +95,11 @@ 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) - """ - 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 - ) - 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) @@ -380,7 +108,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 @@ -388,7 +116,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) @@ -397,7 +129,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 @@ -443,6 +175,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) @@ -464,7 +216,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 @@ -499,14 +278,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, @@ -524,6 +303,9 @@ def test_sparse_attention_mla( context_sequence_lengths, generation_seq_len_q, num_generation_steps, + sparse_topk, + seed, + topk_seed, ) @@ -546,20 +328,27 @@ 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 # 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 = [] @@ -681,26 +470,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, @@ -731,34 +500,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 @@ -802,9 +565,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) @@ -821,12 +586,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( @@ -853,7 +617,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): @@ -865,7 +629,10 @@ 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 + context_sequence_lengths, + sparse_topk, + device, + generator=topk_generator, ) result = ctx_layers[layer_idx].forward( fused_q.clone(), @@ -873,153 +640,113 @@ def yarn_get_mscale(scale=1, mscale=1): None, attn_metadata, attention_input_type=AttentionInputType.context_only, - latent_cache=latent_cache, + latent_cache=latent_cache.clone(), q_pe=q_pe, topk_indices=topk_indices, ) - 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, - ) - 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 + cached_lens, + generation_seq_len_q, + sparse_topk, + device, + generator=topk_generator, ) + if is_vanilla: + backend_fused_q = fused_q.clone() + backend_latent_cache = latent_cache.clone() + generation_kwargs = {} + 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, + ) + 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, + } 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, - topk_indices=topk_indices, - ) - 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, + **generation_kwargs, ) - 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/sparse/dsa/test_dsa_vanilla_attention.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_vanilla_attention.py new file mode 100644 index 000000000000..e5b2e5bf82e8 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_vanilla_attention.py @@ -0,0 +1,260 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +import tensorrt_llm +from tensorrt_llm._torch.attention_backend.interface import AttentionInputType, MLAParams +from tensorrt_llm._torch.attention_backend.sparse.dsa import DSAParams +from tensorrt_llm._torch.attention_backend.vanilla import VanillaAttention, VanillaAttentionMetadata +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.bindings.executor import KvCacheConfig +from tensorrt_llm.mapping import Mapping + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _make_metadata( + manager: KVCacheManager, + request_ids: list[int], + seq_lens: list[int], + cached_lens: list[int], + num_contexts: int, +) -> VanillaAttentionMetadata: + metadata = VanillaAttentionMetadata( + seq_lens=torch.tensor(seq_lens, dtype=torch.int), + request_ids=request_ids, + max_num_requests=len(request_ids), + num_contexts=num_contexts, + max_num_tokens=sum(seq_lens), + kv_cache_manager=manager, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=cached_lens, + ), + ) + metadata.prepare() + return metadata + + +def _make_singleton_topk(selections: list[int], device: torch.device) -> torch.Tensor: + topk = torch.full( + (len(selections), 4), + -1, + dtype=torch.int32, + device=device, + ) + topk[:, 0] = torch.tensor(selections, dtype=torch.int32, device=device) + return topk + + +def _repeat_for_query_heads(values: torch.Tensor, num_heads: int) -> torch.Tensor: + return values.unsqueeze(1).expand(-1, num_heads, -1).reshape(values.shape[0], -1) + + +def test_dsa_selected_mla_context_generation_and_mixed_phases() -> None: + """Exercise ragged, paged DSA cache access with an analytic singleton oracle.""" + torch.manual_seed(123) + device = torch.device("cuda") + dtype = torch.bfloat16 + num_heads = 4 + kv_lora_rank = 8 + qk_nope_head_dim = 8 + qk_rope_head_dim = 4 + fused_head_dim = kv_lora_rank + qk_rope_head_dim + context_lens = [5, 3] + generation_len = 2 + request_ids = [0, 1] + final_lens = [length + generation_len for length in context_lens] + allocated_lens = [final_lens[0], final_lens[1] + 1] + + manager = KVCacheManager( + KvCacheConfig(max_tokens=32, enable_block_reuse=False), + tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY, + num_layers=1, + num_kv_heads=1, + head_dim=fused_head_dim, + tokens_per_block=4, + max_seq_len=max(allocated_lens), + max_batch_size=len(request_ids), + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=tensorrt_llm.bindings.DataType.BF16, + ) + manager.add_dummy_requests(request_ids, allocated_lens) + + attention = VanillaAttention( + layer_idx=0, + num_heads=num_heads, + head_dim=fused_head_dim, + num_kv_heads=1, + q_scaling=1.25, + mla_params=MLAParams( + q_lora_rank=8, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + qk_nope_head_dim=qk_nope_head_dim, + v_head_dim=kv_lora_rank, + ), + sparse_params=DSAParams(), + ) + + try: + num_context_tokens = sum(context_lens) + context_q = torch.randn( + num_context_tokens, + num_heads * fused_head_dim, + device=device, + dtype=dtype, + ) + context_latent = torch.randn( + num_context_tokens, + fused_head_dim, + device=device, + dtype=dtype, + ) + context_metadata = _make_metadata( + manager, + request_ids, + context_lens, + cached_lens=[0, 0], + num_contexts=len(request_ids), + ) + context_output = attention.forward( + context_q, + None, + None, + context_metadata, + attention_input_type=AttentionInputType.context_only, + latent_cache=context_latent, + topk_indices=_make_singleton_topk( + [0, 1, 2, 3, 4, 0, 1, 2], + device, + ), + ) + context_expected = _repeat_for_query_heads( + context_latent[:, :kv_lora_rank], + num_heads, + ) + torch.testing.assert_close(context_output, context_expected) + + num_generation_tokens = len(request_ids) * generation_len + generation_q = torch.randn( + num_generation_tokens, + num_heads * fused_head_dim, + device=device, + dtype=dtype, + ) + generation_latent = torch.randn( + num_generation_tokens, + fused_head_dim, + device=device, + dtype=dtype, + ) + generation_metadata = _make_metadata( + manager, + request_ids, + [generation_len] * len(request_ids), + cached_lens=context_lens, + num_contexts=0, + ) + generation_output = attention.forward( + generation_q, + None, + None, + generation_metadata, + attention_input_type=AttentionInputType.generation_only, + latent_cache=generation_latent, + topk_indices=_make_singleton_topk([4, 6, 0, 4], device), + ) + request_0_cache = torch.cat((context_latent[:5], generation_latent[:2])) + request_1_cache = torch.cat((context_latent[5:], generation_latent[2:])) + generation_selected_values = torch.stack( + ( + request_0_cache[4, :kv_lora_rank], + request_0_cache[6, :kv_lora_rank], + request_1_cache[0, :kv_lora_rank], + request_1_cache[4, :kv_lora_rank], + ) + ) + generation_expected = _repeat_for_query_heads( + generation_selected_values, + num_heads, + ) + torch.testing.assert_close(generation_output, generation_expected) + + mixed_metadata = _make_metadata( + manager, + request_ids, + [2, 1], + cached_lens=[0, final_lens[1]], + num_contexts=1, + ) + mixed_context_q = torch.randn( + 2, + num_heads * fused_head_dim, + device=device, + dtype=dtype, + ) + mixed_context_latent = torch.randn( + 2, + fused_head_dim, + device=device, + dtype=dtype, + ) + mixed_context_output = attention.forward( + mixed_context_q, + None, + None, + mixed_metadata, + attention_input_type=AttentionInputType.context_only, + latent_cache=mixed_context_latent, + topk_indices=_make_singleton_topk([0, 1], device), + ) + mixed_context_expected = _repeat_for_query_heads( + mixed_context_latent[:, :kv_lora_rank], + num_heads, + ) + torch.testing.assert_close(mixed_context_output, mixed_context_expected) + + mixed_generation_q = torch.randn( + 1, + num_heads * fused_head_dim, + device=device, + dtype=dtype, + ) + mixed_generation_latent = torch.randn( + 1, + fused_head_dim, + device=device, + dtype=dtype, + ) + mixed_generation_output = attention.forward( + mixed_generation_q, + None, + None, + mixed_metadata, + attention_input_type=AttentionInputType.generation_only, + latent_cache=mixed_generation_latent, + topk_indices=_make_singleton_topk([final_lens[1]], device), + ) + mixed_generation_expected = _repeat_for_query_heads( + mixed_generation_latent[:, :kv_lora_rank], + num_heads, + ) + torch.testing.assert_close(mixed_generation_output, mixed_generation_expected) + finally: + manager.shutdown() From b3238f5e0634b7a8b0c6438f336b4a37a8b6f56b Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Tue, 14 Jul 2026 23:12:05 -0700 Subject: [PATCH 02/17] [None][refactor] Refine the implementation Signed-off-by: Yihan Wang --- .../_torch/attention_backend/vanilla.py | 199 +++++- .../_torch/attention/backend_capability.py | 12 +- .../unittest/_torch/attention/backend_case.py | 573 +++++++++++++++++- .../_torch/attention/model_attn_config.py | 70 ++- .../sparse/dsa/test_dsa_attention_backends.py | 52 -- .../sparse/dsa/test_dsa_vanilla_attention.py | 260 -------- .../sparse/test_sparse_mla_forward.py | 69 +-- .../attention/test_attention_backends.py | 71 +++ 8 files changed, 902 insertions(+), 404 deletions(-) delete mode 100644 tests/unittest/_torch/attention/sparse/dsa/test_dsa_attention_backends.py delete mode 100644 tests/unittest/_torch/attention/sparse/dsa/test_dsa_vanilla_attention.py diff --git a/tensorrt_llm/_torch/attention_backend/vanilla.py b/tensorrt_llm/_torch/attention_backend/vanilla.py index 2c12d0a551e5..29f65d2ac3f7 100644 --- a/tensorrt_llm/_torch/attention_backend/vanilla.py +++ b/tensorrt_llm/_torch/attention_backend/vanilla.py @@ -2,7 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 import math -from typing import Optional +from dataclasses import replace +from typing import Callable, Optional import torch import torch.nn.functional as F @@ -116,21 +117,41 @@ def __init__( self.qk_nope_head_dim = mla_params.qk_nope_head_dim self.v_head_dim = mla_params.v_head_dim - self.dsa_rope_cos_sin = None - self.dsa_rope_is_neox = True + self.sparse_mla_rope_cos_sin = None + self.sparse_mla_rope_is_neox = True if (self.is_mla_enable and getattr(self.sparse_params, "algorithm", None) == "dsa" and pos_embd_params is not None and pos_embd_params.rope is not None): - self.dsa_rope_cos_sin = pos_embd_params.rope.create_rope_const_params( + self.sparse_mla_rope_cos_sin = pos_embd_params.rope.create_rope_const_params( interleave=False)[1].reshape(pos_embd_params.rope.max_positions, 2, -1) - self.dsa_rope_is_neox = pos_embd_params.is_neox + self.sparse_mla_rope_is_neox = pos_embd_params.is_neox @classmethod def support_mla(cls) -> bool: return True + def sparse_kv_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + metadata: VanillaAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """Lower backend-neutral sparse KV selection for Vanilla attention.""" + return None, None + + 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.""" + return forward_args.topk_indices, None + @staticmethod def _apply_rotary_embedding(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, @@ -152,58 +173,59 @@ def _apply_rotary_embedding(x: torch.Tensor, cos: torch.Tensor, rotated = torch.stack((out1, out2), dim=-1).flatten(-2) return torch.cat((rotated, x_pass), dim=-1) - def _prepare_dsa_mla_inputs( + def _prepare_sparse_mla_inputs( self, fused_q: torch.Tensor, latent_cache: torch.Tensor, q_pe: Optional[torch.Tensor], positions: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - """Apply DSA MLA RoPE to raw packed query and latent-cache inputs. + """Apply sparse MLA RoPE to raw packed query and latent-cache inputs. As with the other attention paths, omitting positional-embedding parameters means the caller already applied RoPE. """ - if self.dsa_rope_cos_sin is None: + if self.sparse_mla_rope_cos_sin is None: return fused_q, latent_cache if positions.numel() == 0: return fused_q, latent_cache max_position = int(positions.max().item()) - if max_position >= self.dsa_rope_cos_sin.shape[0]: + if max_position >= self.sparse_mla_rope_cos_sin.shape[0]: raise ValueError( - f"DSA position {max_position} exceeds the configured RoPE table " - f"size {self.dsa_rope_cos_sin.shape[0]}") + f"Sparse MLA position {max_position} exceeds the configured RoPE table " + f"size {self.sparse_mla_rope_cos_sin.shape[0]}") num_tokens = fused_q.shape[0] fused_head_dim = self.kv_lora_rank + self.qk_rope_head_dim query = fused_q.view(num_tokens, self.num_heads, fused_head_dim).clone() if q_pe is None: raise ValueError( - "Vanilla DSA requires raw q_pe when RoPE parameters are configured" + "Vanilla sparse MLA requires raw q_pe when RoPE parameters are configured" ) expected_numel = num_tokens * self.num_heads * self.qk_rope_head_dim if q_pe.numel() != expected_numel: raise ValueError( - f"DSA q_pe has {q_pe.numel()} elements, expected {expected_numel}" + f"Sparse MLA q_pe has {q_pe.numel()} elements, expected {expected_numel}" ) query_rope = q_pe.reshape(num_tokens, self.num_heads, self.qk_rope_head_dim) if latent_cache.shape[1] != fused_head_dim: raise ValueError( - f"DSA latent cache width must be {fused_head_dim}, got " + f"Sparse MLA latent cache width must be {fused_head_dim}, got " f"{latent_cache.shape[1]}") latent_cache = latent_cache.clone() key_rope = latent_cache[:, self.kv_lora_rank:].unsqueeze(1) - cos_sin = self.dsa_rope_cos_sin.index_select( + cos_sin = self.sparse_mla_rope_cos_sin.index_select( 0, - positions.to(device=self.dsa_rope_cos_sin.device, dtype=torch.long)) + positions.to(device=self.sparse_mla_rope_cos_sin.device, + dtype=torch.long)) cos, sin = cos_sin.unbind(dim=1) query[..., -self.qk_rope_head_dim:] = self._apply_rotary_embedding( - query_rope, cos, sin, self.dsa_rope_is_neox) + query_rope, cos, sin, self.sparse_mla_rope_is_neox) latent_cache[:, self.kv_lora_rank:] = self._apply_rotary_embedding( - key_rope, cos, sin, self.dsa_rope_is_neox).squeeze(1) + key_rope, cos, sin, self.sparse_mla_rope_is_neox).squeeze(1) return query.view(num_tokens, -1), latent_cache def _single_request_sparse_attn_predict( @@ -691,7 +713,7 @@ def _load_mla_latent_cache(kv_cache: torch.Tensor, block_ids: list[int], remaining -= num_tokens return torch.cat(chunks, dim=0) - def _mla_forward_dsa( + def _mla_forward_sparse( self, fused_q: torch.Tensor, metadata: VanillaAttentionMetadata, @@ -700,11 +722,11 @@ def _mla_forward_dsa( topk_indices: torch.Tensor, attention_input_type: AttentionInputType, ) -> torch.Tensor: - """Run DSA selected attention from caller-provided local top-k rows. + """Run selected sparse MLA from caller-provided local top-k rows. - DSA's indexer owns selection. This golden consumes its request-local - token positions, gathers the selected latent K/V, and performs the - absorbed MLA attention directly in PyTorch. + The sparse algorithm owns selection. This golden consumes its + request-local token positions, gathers the selected latent K/V, and + performs the absorbed MLA attention directly in PyTorch. """ if attention_input_type == AttentionInputType.context_only: seq_start, seq_end = 0, metadata.num_contexts @@ -796,7 +818,7 @@ def _mla_forward_dsa( dtype=torch.long) for past, q_len in zip( phase_past_tokens, phase_seq_lens, strict=True) ]) - fused_q, latent_cache = self._prepare_dsa_mla_inputs( + fused_q, latent_cache = self._prepare_sparse_mla_inputs( fused_q, latent_cache, q_pe, positions) from .utils import append_mla_latent_cache @@ -908,25 +930,42 @@ 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.") - sparse_algorithm = getattr(self.sparse_params, "algorithm", None) - if forward_args.topk_indices is not None: + 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") + kv_idx, kv_off = self.sparse_kv_predict(q, k, metadata, + forward_args) + at_idx, at_off = self.sparse_attn_predict( + q, k, metadata, forward_args) + forward_args.sparse_prediction = replace( + forward_args.sparse_prediction, + sparse_kv_indices=kv_idx, + sparse_kv_offsets=kv_off, + sparse_attn_indices=at_idx, + sparse_attn_offsets=at_off, + sparse_attn_indices_block_size=getattr( + self.sparse_params, "indices_block_size"), + ) + sparse_attn_indices = ( + forward_args.sparse_prediction.sparse_attn_indices) + if sparse_attn_indices is not None: if k is not None or v is not None: raise ValueError( - "Vanilla DSA expects absorbed queries and latent cache, " + "Vanilla sparse MLA expects absorbed queries and latent cache, " "not explicit K/V tensors") - return self._mla_forward_dsa( + return self._mla_forward_sparse( q, metadata, forward_args.latent_cache, forward_args.q_pe, - forward_args.topk_indices, + sparse_attn_indices, forward_args.attention_input_type, ) - if sparse_algorithm == "dsa": - raise ValueError("Vanilla DSA requires topk_indices") + 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, @@ -997,3 +1036,101 @@ def forward(self, attn_output = attn_output.view(q_len, -1) return attn_output + + +class VanillaIndexer: + """fp32 reference for the production DSA / DeepSeek-V4 sparse ``Indexer``. + + The production indexer (``sparse/dsa.py``) selects the top-k KV a query may + attend to. ``VanillaIndexer`` mirrors that selection math in plain fp32 so it + can serve as the indexer golden, analogous to how ``VanillaAttention`` is the + attention golden. + + It is a standalone reference, deliberately **not** wired into + ``VanillaAttention.forward``: + + * The indexer consumes index-space inputs (``qr`` / ``hidden_states`` / an + index-space K cache) that ``forward(q, k, v, metadata)`` never receives; + folding it into the forward path would widen the generic backend interface + with algorithm-specific tensors. + * A meaningful indexer golden must run against the *production indexer's own + weights*, so this class **wraps a production ``Indexer`` instance** and + reads its ``wq_b`` / ``weights_proj`` / ``softmax_scale`` / ... rather than + owning independent (non-comparable) parameters. + + It owns the parts common to DSA and DeepSeek-V4 -- the index-space query + projection, the per-head token weights, the logit scoring, and the top-k. + Algorithm-specific K comes from a compressor / ``wk`` projection, so the + caller supplies the reference K and this class handles the rest. + + Selection is discrete: top-k over near-tied logits, and an fp32 reference vs + an fp8/fp4 kernel can pick different borderline tokens. Compare the selected + index *set*, never exact attention outputs. + """ + + def __init__(self, indexer): + self.indexer = indexer + self.n_heads = indexer.n_heads + self.head_dim = indexer.head_dim + self.rope_dim = indexer.rope_dim + self.softmax_scale = indexer.softmax_scale + self.indexer_k_dtype = getattr(indexer, "indexer_k_dtype", None) + + @property + def uses_fp4(self) -> bool: + return self.indexer_k_dtype == "fp4" + + def project_query( + self, + qr: torch.Tensor, + position_ids: torch.Tensor, + freqs_cis: torch.Tensor, + rope_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], + *, + fp4_prep: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, + ) -> torch.Tensor: + """Index-space query: ``wq_b`` GEMM + RoPE on the rope slice. + + ``rope_fn(slice, freqs)`` applies RoPE in place (the caller injects the + algorithm's rotary helper). ``fp4_prep`` optionally applies the fp4 + indexer quant/dequant. Returns ``[num_tokens, n_heads, head_dim]``. + """ + num_tokens = qr.shape[0] + q = F.linear(qr, self.indexer.wq_b.weight) + q = q.view(num_tokens, self.n_heads, self.head_dim).unsqueeze(0) + rope_fn(q[..., -self.rope_dim:], freqs_cis[position_ids.long()]) + q = q.squeeze(0) + if fp4_prep is not None: + q = fp4_prep(q) + return q + + def token_weights(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Per-head token weights: ``weights_proj`` GEMM scaled by ``n_heads**-0.5``.""" + weights = F.linear(hidden_states, self.indexer.weights_proj.weight) + return weights.float() * (self.n_heads**-0.5) + + def scores(self, q_row: torch.Tensor, k: torch.Tensor, + weights_row: torch.Tensor) -> torch.Tensor: + """Per-KV index logits for one query token (fp32). + + ``q_row`` is ``[n_heads, head_dim]``, ``k`` is ``[num_kv, head_dim]``, + ``weights_row`` is ``[n_heads]``. Mirrors the production indexer: per-head + ReLU(q·k) scaled by ``softmax_scale``, combined with the per-head weights. + Returns ``[num_kv]`` logits. + """ + head_scores = torch.einsum("hd,kd->hk", q_row.float(), k.float()) + head_scores = F.relu(head_scores) * self.softmax_scale + return (head_scores * weights_row.float().unsqueeze(-1)).sum(dim=0) + + def topk_from_scores(self, scores: torch.Tensor, + topk_tokens: int) -> torch.Tensor: + """Top-k KV positions for one query token, ``-1``-padded to ``topk_tokens``.""" + row = torch.full((topk_tokens, ), + -1, + dtype=torch.int32, + device=scores.device) + if scores.numel() == 0: + return row + k = min(topk_tokens, scores.numel()) + row[:k] = torch.topk(scores.float(), k, dim=-1).indices.to(torch.int32) + return row diff --git a/tests/unittest/_torch/attention/backend_capability.py b/tests/unittest/_torch/attention/backend_capability.py index 91d11be455df..2f73bcdbaa4a 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,12 @@ 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.get("algorithm") + if backend == "TRTLLM" and algorithm == "dsa" and sm < 100: + return f"TRTLLM DSA selected attention requires sm>=100 (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..8d6f8d549296 100644 --- a/tests/unittest/_torch/attention/backend_case.py +++ b/tests/unittest/_torch/attention/backend_case.py @@ -14,28 +14,34 @@ import math from dataclasses import asdict, dataclass +from types import SimpleNamespace from typing import Dict, List, Optional import torch from backend_capability import BACKEND_CAPS, unsupported_reason from kv_cache_utils import apply_rope, fill_kv_cache_logical, make_position_ids +from pydantic import TypeAdapter import tensorrt_llm from tensorrt_llm._torch.attention_backend.interface import ( AttentionForwardArgs, AttentionInputType, + MLAParams, PositionalEmbeddingParams, 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.attention_backend.vanilla import VanillaAttention from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 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 @@ -55,6 +61,7 @@ # included when available, so callers can iterate this list unconditionally. BACKENDS_UNDER_TEST = ("TRTLLM",) + (("FLASHINFER",) if IS_FLASHINFER_AVAILABLE else ()) DEFAULT_MAX_NUM_TOKENS = 8192 +_SPARSE_CONFIG_ADAPTER = TypeAdapter(SparseAttentionConfig) def _dtype_to_torch(dtype: str): @@ -92,7 +99,19 @@ class BackendCase: q_scaling: float = 1.0 page_size: int = 64 cache: str = "paged" # "paged" | "none" - sparse: str = "off" # "off" | "degenerate" + # Serialized user-facing sparse config. It is lowered independently into + # backend params, metadata params, and the sparse KV-cache manager, exactly + # as in production. + sparse_attention_config: Optional[dict] = None + # Backend-neutral sparse execution contract. This keeps the runner from + # inferring DSA/MLA semantics from an algorithm name. + sparse_attention_family: Optional[str] = None # "standard" | "mla" + sparse_selection_unit: Optional[str] = None # "none" | "token" | "block" + sparse_topk: Optional[int] = None + # Fake backend-neutral selection policy used in backend-level tests. The + # real algorithm/indexer is covered separately. + sparse_selection: str = "random" # "random" | "singleton" + sparse_generation_tokens_per_seq: int = 1 # 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,11 +135,15 @@ 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 qk_rope_head_dim: Optional[int] = None use_kv_cache_manager_v2: bool = False + hf_config_overrides: Optional[dict] = None + atol: Optional[float] = None + rtol: Optional[float] = None @property def num_seqs(self) -> int: @@ -134,6 +157,20 @@ 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 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 +238,54 @@ def _rope_params_from_dict(d: dict) -> RopeParams: return RopeParams(**kwargs) +def _sparse_config(case: BackendCase): + """Restore the production sparse config from its serialized case form.""" + if case.sparse_attention_config is None: + return None + return _SPARSE_CONFIG_ADAPTER.validate_python(case.sparse_attention_config) + + +def _validate_sparse_case(case: BackendCase) -> None: + """Reject incomplete or unsupported sparse harness contracts explicitly.""" + harness_fields = ( + case.sparse_attention_family, + case.sparse_selection_unit, + case.sparse_topk, + ) + if not case.is_sparse: + if any(value is not None for value in harness_fields): + raise ValueError("Sparse harness fields require sparse_attention_config") + return + if any(value is None for value in harness_fields): + raise ValueError("Sparse backend cases require family, selection unit, and top-k") + if case.sparse_attention_family == "mla" and not case.is_mla: + raise ValueError("Sparse MLA harness requires is_mla=True") + if case.sparse_selection_unit not in ("token", "block"): + raise ValueError( + f"Unsupported sparse selection unit: {case.sparse_selection_unit!r} " + "(expected 'token' or 'block')" + ) + if case.sparse_topk <= 0: + raise ValueError("Sparse backend cases require sparse_topk > 0") + + +def _sparse_pretrained_config(case: BackendCase) -> SimpleNamespace: + """Build one HF-like config shared by every sparse lowering boundary.""" + values = dict( + rms_norm_eps=1e-6, + hidden_size=case.hidden_size, + num_attention_heads=case.num_heads, + num_key_value_heads=case.num_kv_heads, + 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, + v_head_dim=case.v_head_dim, + ) + values.update(case.hf_config_overrides or {}) + return SimpleNamespace(**values) + + 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 +360,18 @@ 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, + model_config=None, + pretrained_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 +382,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 +398,20 @@ def _build_mla_kv_cache_manager(case: BackendCase, backend: str): dtype=torch_dtype_to_binding(case.compute_dtype), ) + if sparse_config is not None and backend != "VANILLA": + cls = get_sparse_attn_kv_cache_manager(sparse_config) + if model_config is None or pretrained_config is None: + raise ValueError("Sparse cache manager requires model and pretrained configs") + kwargs.update( + sparse_attention_config=sparse_config, + model_config=model_config, + pretrained_config=pretrained_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 +440,230 @@ def generate_mla_gen_inputs(case: BackendCase, seed: int = 0) -> Dict: ) +def _build_sparse_topk_indices( + case: BackendCase, + generator: torch.Generator, +) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Build causal request-local selections for a sparse backend case.""" + if case.sparse_selection_unit != "token": + raise ValueError( + f"Token selection builder cannot handle {case.sparse_selection_unit!r} selections" + ) + topk = case.sparse_topk + if topk is None or topk <= 0: + raise ValueError("Sparse token-selection cases require a positive top-k") + if case.sparse_selection not in ("random", "singleton"): + raise ValueError(f"Unsupported sparse selection policy: {case.sparse_selection}") + + 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") + if case.sparse_selection == "singleton" + else None + ) + row = 0 + 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 case.sparse_selection == "singleton": + if max_index + 1 <= topk: + # DSA is effectively dense while the visible prefix fits + # within top-k; a real indexer returns every causal token. + indices[row, : max_index + 1] = torch.arange( + max_index + 1, + dtype=torch.int32, + device="cuda", + ) + else: + # Repeating one logical token across all top-k slots keeps + # the physical shape production-valid while yielding an + # analytic singleton-value oracle for truly sparse rows. + selector = row % 3 + selected = (0, max_index, max_index // 2)[selector] + indices[row].fill_(selected) + singleton_oracle_mask[row] = True + else: + valid = min(max_index + 1, topk) + selected = torch.randperm( + max_index + 1, + generator=generator, + device="cuda", + )[:valid] + indices[row, :valid] = torch.sort(selected).values.to(torch.int32) + row += 1 + return indices, singleton_oracle_mask + + +def _build_sparse_block_indices( + case: BackendCase, + generator: torch.Generator, + block_size: int, +) -> torch.Tensor: + """Build causal request-local *block* selections for a block-sparse case. + + Skeleton for the next block-granular algorithm family (e.g. RocketKV, whose + ``indices_block_size == page_size``). Mirrors ``_build_sparse_topk_indices`` + but selects logical KV *blocks* instead of tokens: for a query at causal + position ``p`` the selectable blocks are ``0 .. p // block_size`` (the block + holding ``p`` is inclusive), and each row is padded to ``sparse_topk`` with + ``-1``. Returns ``[nnz_q, sparse_topk]`` int32 block indices. + """ + if case.sparse_selection_unit != "block": + raise ValueError( + f"Block selection builder cannot handle {case.sparse_selection_unit!r} selections" + ) + if block_size <= 0: + raise ValueError("Block-selection cases require a positive block size") + topk = case.sparse_topk + if topk is None or topk <= 0: + raise ValueError("Block-selection cases require a positive top-k") + if case.sparse_selection not in ("random", "singleton"): + raise ValueError(f"Unsupported sparse selection policy: {case.sparse_selection}") + + indices = torch.full((case.nnz_q, topk), -1, dtype=torch.int32, device="cuda") + row = 0 + for cached_len, q_len in zip(case.num_cached_tokens, case.seq_lens, strict=True): + for token_idx in range(q_len): + max_block = (cached_len + token_idx) // block_size + num_blocks = max_block + 1 + if case.sparse_selection == "singleton": + indices[row].fill_(row % num_blocks) + else: + valid = min(num_blocks, topk) + selected = torch.randperm(num_blocks, generator=generator, device="cuda")[:valid] + indices[row, :valid] = torch.sort(selected).values.to(torch.int32) + row += 1 + return indices + + +def generate_sparse_block_inputs(case: BackendCase, seed: int = 0) -> Dict: + """Skeleton input generator for block-selected sparse attention. + + TODO(sparse-block): flesh out once a block-granular algorithm lands. The + reusable ``_build_sparse_block_indices`` builder is wired here; still missing + are (a) the standard/MLA input tensors for the block family and (b) the + matching Vanilla block reference forward in ``VanillaAttention`` that + ``_run_sparse_block_backend`` will call. Until both exist this raises so the + contract is explicit rather than silently mis-run. + """ + sparse_config = _sparse_config(case) + assert sparse_config is not None + sparse_params = sparse_config.to_sparse_params( + layer_idx=0, pretrained_config=_sparse_pretrained_config(case) + ) + block_size = getattr(sparse_params, "indices_block_size") + selection_gen = torch.Generator(device="cuda").manual_seed(seed + 1) + _ = _build_sparse_block_indices(case, selection_gen, block_size) + raise NotImplementedError( + "Block-selected sparse attention is a skeleton: add the block input " + "tensors and a VanillaAttention block reference forward, plus a block " + "ModelAttnConfig (e.g. RocketKV), before enabling this path." + ) + + +def _run_sparse_block_backend( + case: BackendCase, + backend: str, + inputs: Dict, + *, + kv_layout: str, +) -> torch.Tensor: + """Skeleton runner for block-selected sparse attention. + + TODO(sparse-block): implement the block-family forward once a Vanilla block + reference exists, following ``_run_sparse_mla_backend`` (lower the production + ``sparse_attention_config`` via ``to_sparse_params`` / the sparse KV-cache + manager, inject ``inputs`` block selections, and compare against the Vanilla + golden). + """ + raise NotImplementedError( + "Block-selected sparse backend runner is a skeleton; see " + "_run_sparse_mla_backend for the token-selected MLA reference." + ) + + +def generate_sparse_mla_inputs(case: BackendCase, seed: int = 0) -> Dict: + """Generate raw absorbed-MLA inputs plus backend-neutral sparse selections.""" + if case.sparse_attention_family != "mla" or case.sparse_selection_unit != "token": + raise ValueError("This generator supports token-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) + fused_q = torch.cat((q_nope, q_pe), dim=-1).reshape(case.nnz_q, num_heads * d_latent) + compressed_kv = _randn(gen, cdt, case.nnz_q, kv_lora_rank) + k_pe = _randn(gen, cdt, case.nnz_q, qk_rope_head_dim) + latent_cache = torch.cat((compressed_kv, k_pe), dim=-1) + + 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) + 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) + + 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) + expected_output = None + if case.sparse_selection == "singleton": + 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, + q_pe=q_pe, + latent_cache=latent_cache, + 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 +788,199 @@ 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 = _sparse_config(case) + assert sparse_config is not None + pretrained_config = _sparse_pretrained_config(case) + sparse_params = sparse_config.to_sparse_params( + layer_idx=0, + pretrained_config=pretrained_config, + ) + sparse_metadata_params = sparse_config.to_sparse_metadata_params( + pretrained_config=pretrained_config + ) + AttentionCls = ( + VanillaAttention + if backend == "VANILLA" + else 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) + model_config = ModelConfig( + mapping=mapping, + sparse_attention_config=sparse_config, + pretrained_config=pretrained_config, + ) + if backend == "VANILLA": + attn = VanillaAttention( + 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, + sparse_params=sparse_params, + mla_params=MLAParams( + q_lora_rank=case.q_lora_rank, + kv_lora_rank=case.kv_lora_rank, + qk_rope_head_dim=case.qk_rope_head_dim, + qk_nope_head_dim=case.qk_nope_head_dim, + # Selected sparse MLA returns latent values; the model's V + # projection lives outside the standalone backend. + v_head_dim=case.kv_lora_rank, + predicted_tokens_per_seq=case.sparse_generation_tokens_per_seq, + hidden_size=case.hidden_size, + ), + dtype=case.compute_dtype, + skip_create_weights_in_init=True, + ) + else: + 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, + v_head_dim=case.kv_lora_rank, + hidden_size=case.hidden_size, + predicted_tokens_per_seq=case.sparse_generation_tokens_per_seq, + sparse_params=sparse_params, + dtype=case.compute_dtype, + skip_create_weights_in_init=True, + ) + if backend == "TRTLLM": + # Weight creation is intentionally skipped because the harness injects + # selections instead of running the indexer. Quant/FMHA state is still + # required by the backend forward path. + attn.update_quant_config(None) + mgr = _build_mla_kv_cache_manager( + case, + backend, + sparse_config, + model_config, + pretrained_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: + fused_q = inputs["fused_q"][token_slice].clone() + q_pe = inputs["q_pe"][token_slice] + latent_cache = inputs["latent_cache"][token_slice].clone() + forward_kwargs = {} + if backend == "TRTLLM" and attention_input_type == AttentionInputType.generation_only: + cu_q_seqlens = torch.empty(case.num_seqs + 1, dtype=torch.int32, device="cuda") + cu_kv_seqlens = torch.empty(case.num_seqs + 1, dtype=torch.int32, device="cuda") + fmha_scheduler_counter = torch.empty(1, dtype=torch.uint32, device="cuda") + attn.mla_rope_generation( + fused_q, + q_pe, + latent_cache, + metadata, + cu_q_seqlens, + cu_kv_seqlens, + fmha_scheduler_counter, + None, + None, + None, + ) + forward_kwargs.update( + cu_q_seqlens=cu_q_seqlens, + cu_kv_seqlens=cu_kv_seqlens, + fmha_scheduler_counter=fmha_scheduler_counter, + ) + + forward_args = AttentionForwardArgs( + latent_cache=latent_cache, + q_pe=q_pe, + topk_indices=inputs["topk_indices"][token_slice], + attention_input_type=attention_input_type, + **forward_kwargs, + ) + out = attn.forward( + fused_q, + None, + None, + metadata, + forward_args=forward_args, + ) + assert forward_args.sparse_prediction.sparse_attn_indices is not None + assert ( + forward_args.sparse_prediction.sparse_attn_indices_block_size + == sparse_params.indices_block_size + ) + 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 +1281,11 @@ 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.atol is not None or case.rtol is not None: + if case.atol is None or case.rtol is None: + raise ValueError("BackendCase.atol and rtol must be set together") + return case.atol, case.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 +1406,14 @@ 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: + sparse_kind = (case.sparse_attention_family, case.sparse_selection_unit) + if sparse_kind == ("mla", "token"): + return _run_sparse_mla_backend(case, backend, inputs, kv_layout=kv_layout) + if case.sparse_selection_unit == "block": + return _run_sparse_block_backend(case, backend, inputs, kv_layout=kv_layout) + raise ValueError(f"Unsupported sparse harness contract: {sparse_kind}") + if case.is_mla: if case.is_context_only: return _run_mla_context_backend(case, backend, inputs, kv_layout=kv_layout) @@ -1063,8 +1602,17 @@ def run_case(case: BackendCase, *, seed: int = 0) -> Dict[str, torch.Tensor]: Returns the per-backend outputs (including ``"VANILLA"`` golden) for callers that want the raw tensors (e.g. the minimizer). """ + _validate_sparse_case(case) is_mla = case.is_mla - if is_mla: + if case.is_sparse: + sparse_kind = (case.sparse_attention_family, case.sparse_selection_unit) + if sparse_kind == ("mla", "token"): + inputs = generate_sparse_mla_inputs(case, seed) + elif case.sparse_selection_unit == "block": + inputs = generate_sparse_block_inputs(case, seed) + else: + raise ValueError(f"Unsupported sparse harness contract: {sparse_kind}") + elif is_mla: if case.is_context_only: inputs = generate_mla_context_inputs(case, seed) else: @@ -1074,6 +1622,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 case.sparse_selection == "singleton" and 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 = [] @@ -1098,7 +1651,7 @@ 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: + 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 29f292a89c12..e8571ec4606d 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,13 +48,36 @@ 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 List, Optional +from typing import List, Literal, Optional, Tuple + +from tensorrt_llm.llmapi.llm_args import BaseSparseAttentionConfig, DeepSeekSparseAttentionConfig + + +@dataclass(frozen=True) +class SparseHarnessConfig: + """Backend-neutral semantics and dimensions for a sparse test case. + + Algorithms that produce the same selection representation reuse one + generator/runner. A new representation extends this common harness once, + rather than adding per-backend tests. + """ + + attention_family: Literal["standard", "mla"] + selection_unit: Literal["none", "token", "block"] + topk: int + page_size: int = 64 + kv_layout: Literal["NHD", "HND"] = "HND" + use_kv_cache_manager_v2: bool = False + compute_dtypes: Tuple[str, ...] = ("bfloat16",) + selection_patterns: Tuple[str, ...] = ("random", "singleton") + generation_tokens_per_seq: int = 2 @dataclass(frozen=True) @@ -78,6 +103,12 @@ 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 + sparse_attention_config: Optional[BaseSparseAttentionConfig] = None + sparse_harness_config: Optional[SparseHarnessConfig] = None + hf_config_overrides: Optional[dict] = None + atol: Optional[float] = None + rtol: Optional[float] = None # --------------------------------------------------------------------------- @@ -569,6 +600,37 @@ 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, + ), + sparse_harness_config=SparseHarnessConfig( + attention_family="mla", + selection_unit="token", + topk=128, + ), + atol=0.1, + rtol=0.01, + ), # 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_attention_backends.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_attention_backends.py deleted file mode 100644 index 6ada2a406dcd..000000000000 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_attention_backends.py +++ /dev/null @@ -1,52 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Differential DSA backend tests using VanillaAttention as the golden.""" - -import pytest -from utils.util import skip_pre_blackwell - -from .test_dsa_sparse_mla import _assert_matches_vanilla, _test_sparse_attention_mla, scenarios - -BACKENDS_UNDER_TEST = ("TRTLLM",) - -DSA_CASES = { - "bf16-paged-context-generation": dict( - scenario=scenarios[0], - context_sequence_lengths=[160], - generation_seq_len_q=1, - num_generation_steps=2, - sparse_topk=128, - seed=123, - topk_seed=456, - ), -} - - -@skip_pre_blackwell -@pytest.mark.parametrize("name", list(DSA_CASES), ids=lambda name: name) -def test_dsa_attention_backend(name: str) -> None: - """Run Vanilla first, then compare every production DSA phase to it.""" - case = DSA_CASES[name] - golden = _test_sparse_attention_mla("VANILLA", **case) - - for backend in BACKENDS_UNDER_TEST: - actual = _test_sparse_attention_mla(backend, **case) - _assert_matches_vanilla( - actual, - golden, - case["scenario"].kv_cache_dtype, - backend, - ) diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_vanilla_attention.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_vanilla_attention.py deleted file mode 100644 index e5b2e5bf82e8..000000000000 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_vanilla_attention.py +++ /dev/null @@ -1,260 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import pytest -import torch - -import tensorrt_llm -from tensorrt_llm._torch.attention_backend.interface import AttentionInputType, MLAParams -from tensorrt_llm._torch.attention_backend.sparse.dsa import DSAParams -from tensorrt_llm._torch.attention_backend.vanilla import VanillaAttention, VanillaAttentionMetadata -from tensorrt_llm._torch.metadata import KVCacheParams -from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager -from tensorrt_llm.bindings.executor import KvCacheConfig -from tensorrt_llm.mapping import Mapping - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") - - -def _make_metadata( - manager: KVCacheManager, - request_ids: list[int], - seq_lens: list[int], - cached_lens: list[int], - num_contexts: int, -) -> VanillaAttentionMetadata: - metadata = VanillaAttentionMetadata( - seq_lens=torch.tensor(seq_lens, dtype=torch.int), - request_ids=request_ids, - max_num_requests=len(request_ids), - num_contexts=num_contexts, - max_num_tokens=sum(seq_lens), - kv_cache_manager=manager, - kv_cache_params=KVCacheParams( - use_cache=True, - num_cached_tokens_per_seq=cached_lens, - ), - ) - metadata.prepare() - return metadata - - -def _make_singleton_topk(selections: list[int], device: torch.device) -> torch.Tensor: - topk = torch.full( - (len(selections), 4), - -1, - dtype=torch.int32, - device=device, - ) - topk[:, 0] = torch.tensor(selections, dtype=torch.int32, device=device) - return topk - - -def _repeat_for_query_heads(values: torch.Tensor, num_heads: int) -> torch.Tensor: - return values.unsqueeze(1).expand(-1, num_heads, -1).reshape(values.shape[0], -1) - - -def test_dsa_selected_mla_context_generation_and_mixed_phases() -> None: - """Exercise ragged, paged DSA cache access with an analytic singleton oracle.""" - torch.manual_seed(123) - device = torch.device("cuda") - dtype = torch.bfloat16 - num_heads = 4 - kv_lora_rank = 8 - qk_nope_head_dim = 8 - qk_rope_head_dim = 4 - fused_head_dim = kv_lora_rank + qk_rope_head_dim - context_lens = [5, 3] - generation_len = 2 - request_ids = [0, 1] - final_lens = [length + generation_len for length in context_lens] - allocated_lens = [final_lens[0], final_lens[1] + 1] - - manager = KVCacheManager( - KvCacheConfig(max_tokens=32, enable_block_reuse=False), - tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY, - num_layers=1, - num_kv_heads=1, - head_dim=fused_head_dim, - tokens_per_block=4, - max_seq_len=max(allocated_lens), - max_batch_size=len(request_ids), - mapping=Mapping(world_size=1, tp_size=1, rank=0), - dtype=tensorrt_llm.bindings.DataType.BF16, - ) - manager.add_dummy_requests(request_ids, allocated_lens) - - attention = VanillaAttention( - layer_idx=0, - num_heads=num_heads, - head_dim=fused_head_dim, - num_kv_heads=1, - q_scaling=1.25, - mla_params=MLAParams( - q_lora_rank=8, - kv_lora_rank=kv_lora_rank, - qk_rope_head_dim=qk_rope_head_dim, - qk_nope_head_dim=qk_nope_head_dim, - v_head_dim=kv_lora_rank, - ), - sparse_params=DSAParams(), - ) - - try: - num_context_tokens = sum(context_lens) - context_q = torch.randn( - num_context_tokens, - num_heads * fused_head_dim, - device=device, - dtype=dtype, - ) - context_latent = torch.randn( - num_context_tokens, - fused_head_dim, - device=device, - dtype=dtype, - ) - context_metadata = _make_metadata( - manager, - request_ids, - context_lens, - cached_lens=[0, 0], - num_contexts=len(request_ids), - ) - context_output = attention.forward( - context_q, - None, - None, - context_metadata, - attention_input_type=AttentionInputType.context_only, - latent_cache=context_latent, - topk_indices=_make_singleton_topk( - [0, 1, 2, 3, 4, 0, 1, 2], - device, - ), - ) - context_expected = _repeat_for_query_heads( - context_latent[:, :kv_lora_rank], - num_heads, - ) - torch.testing.assert_close(context_output, context_expected) - - num_generation_tokens = len(request_ids) * generation_len - generation_q = torch.randn( - num_generation_tokens, - num_heads * fused_head_dim, - device=device, - dtype=dtype, - ) - generation_latent = torch.randn( - num_generation_tokens, - fused_head_dim, - device=device, - dtype=dtype, - ) - generation_metadata = _make_metadata( - manager, - request_ids, - [generation_len] * len(request_ids), - cached_lens=context_lens, - num_contexts=0, - ) - generation_output = attention.forward( - generation_q, - None, - None, - generation_metadata, - attention_input_type=AttentionInputType.generation_only, - latent_cache=generation_latent, - topk_indices=_make_singleton_topk([4, 6, 0, 4], device), - ) - request_0_cache = torch.cat((context_latent[:5], generation_latent[:2])) - request_1_cache = torch.cat((context_latent[5:], generation_latent[2:])) - generation_selected_values = torch.stack( - ( - request_0_cache[4, :kv_lora_rank], - request_0_cache[6, :kv_lora_rank], - request_1_cache[0, :kv_lora_rank], - request_1_cache[4, :kv_lora_rank], - ) - ) - generation_expected = _repeat_for_query_heads( - generation_selected_values, - num_heads, - ) - torch.testing.assert_close(generation_output, generation_expected) - - mixed_metadata = _make_metadata( - manager, - request_ids, - [2, 1], - cached_lens=[0, final_lens[1]], - num_contexts=1, - ) - mixed_context_q = torch.randn( - 2, - num_heads * fused_head_dim, - device=device, - dtype=dtype, - ) - mixed_context_latent = torch.randn( - 2, - fused_head_dim, - device=device, - dtype=dtype, - ) - mixed_context_output = attention.forward( - mixed_context_q, - None, - None, - mixed_metadata, - attention_input_type=AttentionInputType.context_only, - latent_cache=mixed_context_latent, - topk_indices=_make_singleton_topk([0, 1], device), - ) - mixed_context_expected = _repeat_for_query_heads( - mixed_context_latent[:, :kv_lora_rank], - num_heads, - ) - torch.testing.assert_close(mixed_context_output, mixed_context_expected) - - mixed_generation_q = torch.randn( - 1, - num_heads * fused_head_dim, - device=device, - dtype=dtype, - ) - mixed_generation_latent = torch.randn( - 1, - fused_head_dim, - device=device, - dtype=dtype, - ) - mixed_generation_output = attention.forward( - mixed_generation_q, - None, - None, - mixed_metadata, - attention_input_type=AttentionInputType.generation_only, - latent_cache=mixed_generation_latent, - topk_indices=_make_singleton_topk([final_lens[1]], device), - ) - mixed_generation_expected = _repeat_for_query_heads( - mixed_generation_latent[:, :kv_lora_rank], - num_heads, - ) - torch.testing.assert_close(mixed_generation_output, mixed_generation_expected) - finally: - manager.shutdown() diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py b/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py index 72454df6fe86..104742a01eb0 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py @@ -22,7 +22,6 @@ import pytest import torch -import torch.nn.functional as F import tensorrt_llm import tensorrt_llm.bindings @@ -33,6 +32,7 @@ from tensorrt_llm._torch.attention_backend.sparse.dsa import (HAS_FAST_HADAMARD, DSACacheManager) from tensorrt_llm._torch.attention_backend.utils import get_attention_backend +from tensorrt_llm._torch.attention_backend.vanilla import VanillaIndexer from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.modules.mla import MLA @@ -271,18 +271,6 @@ def _copy_ref_compressor_weights(ref_compressor: RefCompressor, ref_compressor.norm.weight.data.copy_(compressor.norm.weight.data) -def _topk_from_scores(scores: torch.Tensor, topk_tokens: int) -> torch.Tensor: - row = torch.full((topk_tokens, ), - -1, - dtype=torch.int32, - device=scores.device) - if scores.numel() == 0: - return row - k = min(topk_tokens, scores.numel()) - row[:k] = torch.topk(scores.float(), k, dim=-1).indices.to(torch.int32) - return row - - def _ceil_pow2_scale(amax: torch.Tensor, max_value_inv: float, min_amax: float) -> torch.Tensor: scaled = torch.clamp(amax.float(), min=min_amax) * max_value_inv @@ -355,14 +343,6 @@ def _prepare_fp4_indexer_k_reference(k: torch.Tensor) -> torch.Tensor: min_amax=fp4_min_amax) -def _reference_indexer_scores(q: torch.Tensor, k: torch.Tensor, - weights: torch.Tensor, - softmax_scale: float) -> torch.Tensor: - head_scores = torch.einsum("hd,kd->hk", q.float(), k.float()) - head_scores = F.relu(head_scores) * softmax_scale - return (head_scores * weights.float().unsqueeze(-1)).sum(dim=0) - - def calculate_reference_deepseek_v4_topk_indices( mla, ref_indexer_compressor: RefCompressor, @@ -378,26 +358,29 @@ def calculate_reference_deepseek_v4_topk_indices( compress_ratio: int, device: torch.device, ) -> torch.Tensor: - """Independent PyTorch reference for HF DS4 ratio-4 indexer top-k indices.""" + """Independent PyTorch reference for HF DS4 ratio-4 indexer top-k indices. + + The index-space query projection, per-head token weights, logit scoring, and + top-k are shared with the DSA path via :class:`VanillaIndexer`; this function + adds the DeepSeek-V4-specific compressor that produces the reference K. + """ indexer = mla.mqa.indexer + vindexer = VanillaIndexer(indexer) num_tokens = hidden_states.shape[0] topk_indices = torch.full((num_tokens, topk_tokens), -1, dtype=torch.int32, device=device) - q_ref = F.linear(qr, indexer.wq_b.weight) - q_ref = q_ref.view(num_tokens, indexer.n_heads, indexer.head_dim) - q_ref = q_ref.unsqueeze(0) - apply_rotary_emb(q_ref[..., -indexer.rope_dim:], - freqs_cis[position_ids.long()]) - q_ref = q_ref.squeeze(0) - use_fp4_indexer = indexer.indexer_k_dtype == "fp4" - if use_fp4_indexer: - q_ref = _prepare_fp4_indexer_q_reference(q_ref) - - weights = F.linear(hidden_states, indexer.weights_proj.weight) - weights = weights.float() * (indexer.n_heads**-0.5) + use_fp4_indexer = vindexer.uses_fp4 + q_ref = vindexer.project_query( + qr, + position_ids, + freqs_cis, + apply_rotary_emb, + fp4_prep=_prepare_fp4_indexer_q_reference if use_fp4_indexer else None, + ) + weights = vindexer.token_weights(hidden_states) offset = 0 for req_idx in ctx_indices: @@ -421,11 +404,10 @@ def calculate_reference_deepseek_v4_topk_indices( for token_idx in range(seq_len): valid_len = (token_idx + 1) // compress_ratio k_for_scores = compressed_kv_for_scores[:valid_len] - scores = _reference_indexer_scores(q_ref[offset + token_idx], - k_for_scores, - weights[offset + token_idx], - indexer.softmax_scale) - topk_indices[offset + token_idx] = _topk_from_scores( + scores = vindexer.scores(q_ref[offset + token_idx], + k_for_scores, + weights[offset + token_idx]) + topk_indices[offset + token_idx] = vindexer.topk_from_scores( scores, topk_tokens) offset += seq_len @@ -467,11 +449,10 @@ def calculate_reference_deepseek_v4_topk_indices( compressed_kv_for_scores = ( _prepare_fp4_indexer_k_reference(compressed_kv) if use_fp4_indexer else compressed_kv) - scores = _reference_indexer_scores(q_ref[token_idx], - compressed_kv_for_scores[:valid_len], - weights[token_idx], - indexer.softmax_scale) - topk_indices[token_idx] = _topk_from_scores(scores, topk_tokens) + scores = vindexer.scores(q_ref[token_idx], + compressed_kv_for_scores[:valid_len], + weights[token_idx]) + topk_indices[token_idx] = vindexer.topk_from_scores(scores, topk_tokens) return topk_indices diff --git a/tests/unittest/_torch/attention/test_attention_backends.py b/tests/unittest/_torch/attention/test_attention_backends.py index 8bae3401ba7a..857912d8bb17 100644 --- a/tests/unittest/_torch/attention/test_attention_backends.py +++ b/tests/unittest/_torch/attention/test_attention_backends.py @@ -78,6 +78,33 @@ def _phases_from_window(window: int) -> dict: def _phases_for(cfg: ModelAttnConfig) -> dict: + if cfg.sparse_attention_config is not None: + harness = cfg.sparse_harness_config + if harness is None: + raise ValueError(f"Sparse model config {cfg.id} must define a sparse harness") + if (harness.attention_family == "mla") != cfg.is_mla: + raise ValueError( + f"Sparse harness family {harness.attention_family!r} does not match {cfg.id}" + ) + long_len = harness.topk + 32 + generation_len = harness.generation_tokens_per_seq + return { + "ctx": dict( + seq_lens=[long_len, 73, 41], + num_cached_tokens=[0, 0, 0], + num_contexts=3, + ), + "gen": dict( + seq_lens=[generation_len] * 3, + num_cached_tokens=[long_len, 73, 41], + num_contexts=0, + ), + "mix": dict( + seq_lens=[long_len, generation_len, generation_len], + num_cached_tokens=[0, long_len, 73], + num_contexts=1, + ), + } if cfg.mask != "sliding": return _PHASES @@ -113,6 +140,20 @@ 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: + harness = cfg.sparse_harness_config + assert harness is not None + common.update( + sparse_attention_config=cfg.sparse_attention_config.model_dump(mode="json"), + sparse_attention_family=harness.attention_family, + sparse_selection_unit=harness.selection_unit, + sparse_topk=harness.topk, + sparse_generation_tokens_per_seq=harness.generation_tokens_per_seq, + hf_config_overrides=cfg.hf_config_overrides, + atol=cfg.atol, + rtol=cfg.rtol, ) return common @@ -136,6 +177,36 @@ def _expand(cfg: ModelAttnConfig, precisions, kv_layouts, page_sizes): common = _common(cfg) phases = _phases_for(cfg) + # The model's sparse harness declares the common input/selection semantics + # and cache constraints. Fake request-local selections isolate backend + # execution from the separately-tested algorithm/indexer. + if cfg.sparse_attention_config is not None: + harness = cfg.sparse_harness_config + assert harness is not None + for dtype, kvd in precisions: + if kvd is not None or dtype not in harness.compute_dtypes: + continue + for phase_name in ("ctx", "gen", "mix"): + for selection in harness.selection_patterns: + manager = "v2" if harness.use_kv_cache_manager_v2 else "v1" + tag = ( + f"{_prec_tag(dtype, kvd)}-{harness.kv_layout}" + f"-p{harness.page_size}-{manager}-{selection}" + ) + yield ( + f"{cfg.id}-{phase_name}-{tag}", + BackendCase( + page_size=harness.page_size, + kv_layout=harness.kv_layout, + dtype=dtype, + sparse_selection=selection, + use_kv_cache_manager_v2=harness.use_kv_cache_manager_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: From 108df3fa0ff3699c2edc07f8e8aebf446adf566d Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Wed, 15 Jul 2026 02:33:32 -0700 Subject: [PATCH 03/17] [None][refactor] Address review on sparse attention test harness - Register DSA in the vanilla sparse-backend factory so VanillaAttention is built through create_attention like every other backend (no special-casing in the sparse runner); update_quant_config is now called unconditionally. - Derive sparse family / selection unit / top-k from is_mla and sparse_attention_config, and move the model-agnostic sweep dimensions to shared constants, so ModelAttnConfig no longer carries a SparseHarnessConfig or hf_config_overrides/atol/rtol. Sparse tolerance is derived in _tolerances. - Clarify why the backend-only sparse oracle skips the captured-CUDA-graph path. Signed-off-by: Yihan Wang --- .../_torch/attention_backend/sparse/utils.py | 6 + .../unittest/_torch/attention/backend_case.py | 105 ++++++++---------- .../_torch/attention/model_attn_config.py | 37 +----- .../attention/test_attention_backends.py | 105 ++++++++++-------- 4 files changed, 117 insertions(+), 136 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/utils.py b/tensorrt_llm/_torch/attention_backend/sparse/utils.py index f438d998939f..b9ba6edda535 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/utils.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/utils.py @@ -48,6 +48,12 @@ def get_vanilla_sparse_attn_attention_backend( from .rocket import RocketVanillaAttention if sparse_params.algorithm == "rocket": return RocketVanillaAttention + elif sparse_params.algorithm == "dsa": + # DSA selected-attention is implemented directly in the base + # VanillaAttention (its `_mla_forward_sparse` golden), so the vanilla + # slot for DSA is VanillaAttention itself. + from ..vanilla import VanillaAttention + return VanillaAttention elif sparse_params.algorithm == "minimax_m3": return get_minimax_m3_attention_backend_cls() else: diff --git a/tests/unittest/_torch/attention/backend_case.py b/tests/unittest/_torch/attention/backend_case.py index 8d6f8d549296..f866580364ba 100644 --- a/tests/unittest/_torch/attention/backend_case.py +++ b/tests/unittest/_torch/attention/backend_case.py @@ -26,14 +26,12 @@ from tensorrt_llm._torch.attention_backend.interface import ( AttentionForwardArgs, AttentionInputType, - MLAParams, PositionalEmbeddingParams, 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.attention_backend.vanilla import VanillaAttention from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.model_config import ModelConfig @@ -56,6 +54,11 @@ # 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 +# Selected-sparse-MLA golden-vs-backend tolerance. Looser than dense bf16: the +# selected-attention path accumulates over gathered latent rows in bf16, so the +# TRTLLM kernel and the Vanilla golden diverge more than dense attention does. +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. @@ -141,7 +144,9 @@ class BackendCase: qk_nope_head_dim: Optional[int] = None qk_rope_head_dim: Optional[int] = None use_kv_cache_manager_v2: bool = False - hf_config_overrides: Optional[dict] = None + # Low-level (atol, rtol) override for a single case, e.g. a captured/replayed + # case. Left unset by the model sweep, which derives tolerances by dtype (and + # a sparse default) in ``_tolerances``. atol: Optional[float] = None rtol: Optional[float] = None @@ -282,7 +287,6 @@ def _sparse_pretrained_config(case: BackendCase) -> SimpleNamespace: qk_rope_head_dim=case.qk_rope_head_dim, v_head_dim=case.v_head_dim, ) - values.update(case.hf_config_overrides or {}) return SimpleNamespace(**values) @@ -806,11 +810,7 @@ def _run_sparse_mla_backend( sparse_metadata_params = sparse_config.to_sparse_metadata_params( pretrained_config=pretrained_config ) - AttentionCls = ( - VanillaAttention - if backend == "VANILLA" - else get_attention_backend(backend, sparse_params=sparse_params) - ) + 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) @@ -820,55 +820,35 @@ def _run_sparse_mla_backend( sparse_attention_config=sparse_config, pretrained_config=pretrained_config, ) - if backend == "VANILLA": - attn = VanillaAttention( - 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, - sparse_params=sparse_params, - mla_params=MLAParams( - q_lora_rank=case.q_lora_rank, - kv_lora_rank=case.kv_lora_rank, - qk_rope_head_dim=case.qk_rope_head_dim, - qk_nope_head_dim=case.qk_nope_head_dim, - # Selected sparse MLA returns latent values; the model's V - # projection lives outside the standalone backend. - v_head_dim=case.kv_lora_rank, - predicted_tokens_per_seq=case.sparse_generation_tokens_per_seq, - hidden_size=case.hidden_size, - ), - dtype=case.compute_dtype, - skip_create_weights_in_init=True, - ) - else: - 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, - v_head_dim=case.kv_lora_rank, - hidden_size=case.hidden_size, - predicted_tokens_per_seq=case.sparse_generation_tokens_per_seq, - sparse_params=sparse_params, - dtype=case.compute_dtype, - skip_create_weights_in_init=True, - ) - if backend == "TRTLLM": - # Weight creation is intentionally skipped because the harness injects - # selections instead of running the indexer. Quant/FMHA state is still - # required by the backend forward path. - attn.update_quant_config(None) + # Every backend (Vanilla included) is built through the production factory so + # the harness does not special-case backend construction. + 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=case.sparse_generation_tokens_per_seq, + sparse_params=sparse_params, + dtype=case.compute_dtype, + skip_create_weights_in_init=True, + ) + # Weight creation is skipped (the harness injects selections instead of + # running the indexer); update_quant_config still initializes the quant/FMHA + # state the fused backends need before forward, and is a safe no-op for + # Vanilla (the base implementation only stores quant_config). + attn.update_quant_config(None) mgr = _build_mla_kv_cache_manager( case, backend, @@ -1286,6 +1266,9 @@ def _tolerances(case: "BackendCase", kv_dtype) -> tuple: raise ValueError("BackendCase.atol and rtol must be set together") return case.atol, case.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) @@ -1650,7 +1633,11 @@ 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. + # eager golden. Sparse cases are excluded: the standalone sparse runner + # is an eager correctness oracle that validates injected selections on the + # host and rebuilds each request's logical cache with Python loops, none + # of which is graph-capturable. Graph coverage of the DSA decode kernel is + # exercised at the model level, not by this backend-only oracle. if case.is_gen_only and not case.is_sparse: cg_out = run_backend( case, diff --git a/tests/unittest/_torch/attention/model_attn_config.py b/tests/unittest/_torch/attention/model_attn_config.py index e8571ec4606d..00c53d30d14e 100644 --- a/tests/unittest/_torch/attention/model_attn_config.py +++ b/tests/unittest/_torch/attention/model_attn_config.py @@ -55,31 +55,11 @@ """ from dataclasses import dataclass -from typing import List, Literal, Optional, Tuple +from typing import List, Optional from tensorrt_llm.llmapi.llm_args import BaseSparseAttentionConfig, DeepSeekSparseAttentionConfig -@dataclass(frozen=True) -class SparseHarnessConfig: - """Backend-neutral semantics and dimensions for a sparse test case. - - Algorithms that produce the same selection representation reuse one - generator/runner. A new representation extends this common harness once, - rather than adding per-backend tests. - """ - - attention_family: Literal["standard", "mla"] - selection_unit: Literal["none", "token", "block"] - topk: int - page_size: int = 64 - kv_layout: Literal["NHD", "HND"] = "HND" - use_kv_cache_manager_v2: bool = False - compute_dtypes: Tuple[str, ...] = ("bfloat16",) - selection_patterns: Tuple[str, ...] = ("random", "singleton") - generation_tokens_per_seq: int = 2 - - @dataclass(frozen=True) class ModelAttnConfig: id: str # short, stable test id @@ -104,11 +84,11 @@ class ModelAttnConfig: 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` etc. + # Everything the sparse test sweep needs (attention family, selection unit, + # top-k, tolerances) is derived from this + the fields above, so no separate + # harness/override/tolerance knobs live on the model config. sparse_attention_config: Optional[BaseSparseAttentionConfig] = None - sparse_harness_config: Optional[SparseHarnessConfig] = None - hf_config_overrides: Optional[dict] = None - atol: Optional[float] = None - rtol: Optional[float] = None # --------------------------------------------------------------------------- @@ -623,13 +603,6 @@ class ModelAttnConfig: index_topk=128, skip_indexer_for_short_seqs=False, ), - sparse_harness_config=SparseHarnessConfig( - attention_family="mla", - selection_unit="token", - topk=128, - ), - atol=0.1, - rtol=0.01, ), # DeepSeek-V3: 128 Q heads, qk_nope=128, qk_rope=64, kv_lora=512, v=128. ModelAttnConfig( diff --git a/tests/unittest/_torch/attention/test_attention_backends.py b/tests/unittest/_torch/attention/test_attention_backends.py index 857912d8bb17..32ef7a61ea10 100644 --- a/tests/unittest/_torch/attention/test_attention_backends.py +++ b/tests/unittest/_torch/attention/test_attention_backends.py @@ -77,17 +77,42 @@ def _phases_from_window(window: int) -> dict: _PHASES = _phases_from_window(_NON_SLIDING_PHASE_WINDOW) +# Sparse test sweep dimensions are model-agnostic: one shared set for every +# sparse config rather than per-model knobs. Everything else (attention family, +# selection unit, top-k) is derived from the model's is_mla flag and its +# sparse_attention_config. +_SPARSE_COMPUTE_DTYPE = "bfloat16" +_SPARSE_KV_LAYOUT = "HND" +_SPARSE_PAGE_SIZE = 64 +_SPARSE_USE_KVM_V2 = False +_SPARSE_SELECTION_PATTERNS = ("random", "singleton") +_SPARSE_GENERATION_TOKENS_PER_SEQ = 2 +# Selection representation per algorithm (token-selecting vs block-selecting). +_SPARSE_SELECTION_UNIT = {"dsa": "token", "deepseek_v4": "token", "rocket": "block"} + + +def _sparse_selection_unit(cfg: ModelAttnConfig) -> str: + algo = cfg.sparse_attention_config.algorithm + if algo not in _SPARSE_SELECTION_UNIT: + raise ValueError(f"Unknown selection unit for sparse algorithm {algo!r}") + return _SPARSE_SELECTION_UNIT[algo] + + +def _sparse_topk(cfg: ModelAttnConfig) -> int: + """The selection budget: the indexer's top-k (index_topk) or the algo's topk.""" + sac = cfg.sparse_attention_config + topk = getattr(sac, "index_topk", None) + if topk is None: + topk = getattr(sac, "topk", None) + if topk is None: + raise ValueError(f"Cannot derive sparse top-k from {sac.algorithm!r} config") + return topk + + def _phases_for(cfg: ModelAttnConfig) -> dict: if cfg.sparse_attention_config is not None: - harness = cfg.sparse_harness_config - if harness is None: - raise ValueError(f"Sparse model config {cfg.id} must define a sparse harness") - if (harness.attention_family == "mla") != cfg.is_mla: - raise ValueError( - f"Sparse harness family {harness.attention_family!r} does not match {cfg.id}" - ) - long_len = harness.topk + 32 - generation_len = harness.generation_tokens_per_seq + long_len = _sparse_topk(cfg) + 32 + generation_len = _SPARSE_GENERATION_TOKENS_PER_SEQ return { "ctx": dict( seq_lens=[long_len, 73, 41], @@ -143,17 +168,12 @@ def _common(cfg: ModelAttnConfig) -> dict: hidden_size=cfg.hidden_size, ) if cfg.sparse_attention_config is not None: - harness = cfg.sparse_harness_config - assert harness is not None common.update( sparse_attention_config=cfg.sparse_attention_config.model_dump(mode="json"), - sparse_attention_family=harness.attention_family, - sparse_selection_unit=harness.selection_unit, - sparse_topk=harness.topk, - sparse_generation_tokens_per_seq=harness.generation_tokens_per_seq, - hf_config_overrides=cfg.hf_config_overrides, - atol=cfg.atol, - rtol=cfg.rtol, + sparse_attention_family="mla" if cfg.is_mla else "standard", + sparse_selection_unit=_sparse_selection_unit(cfg), + sparse_topk=_sparse_topk(cfg), + sparse_generation_tokens_per_seq=_SPARSE_GENERATION_TOKENS_PER_SEQ, ) return common @@ -177,34 +197,29 @@ def _expand(cfg: ModelAttnConfig, precisions, kv_layouts, page_sizes): common = _common(cfg) phases = _phases_for(cfg) - # The model's sparse harness declares the common input/selection semantics - # and cache constraints. Fake request-local selections isolate backend - # execution from the separately-tested algorithm/indexer. + # Sparse cases use one model-agnostic sweep (bf16 latent cache, fixed layout/ + # page/manager). Fake request-local selections isolate backend execution from + # the separately-tested algorithm/indexer. if cfg.sparse_attention_config is not None: - harness = cfg.sparse_harness_config - assert harness is not None - for dtype, kvd in precisions: - if kvd is not None or dtype not in harness.compute_dtypes: - continue - for phase_name in ("ctx", "gen", "mix"): - for selection in harness.selection_patterns: - manager = "v2" if harness.use_kv_cache_manager_v2 else "v1" - tag = ( - f"{_prec_tag(dtype, kvd)}-{harness.kv_layout}" - f"-p{harness.page_size}-{manager}-{selection}" - ) - yield ( - f"{cfg.id}-{phase_name}-{tag}", - BackendCase( - page_size=harness.page_size, - kv_layout=harness.kv_layout, - dtype=dtype, - sparse_selection=selection, - use_kv_cache_manager_v2=harness.use_kv_cache_manager_v2, - **phases[phase_name], - **common, - ), - ) + manager = "v2" if _SPARSE_USE_KVM_V2 else "v1" + for phase_name in ("ctx", "gen", "mix"): + for selection in _SPARSE_SELECTION_PATTERNS: + tag = ( + f"{_prec_tag(_SPARSE_COMPUTE_DTYPE, None)}-{_SPARSE_KV_LAYOUT}" + f"-p{_SPARSE_PAGE_SIZE}-{manager}-{selection}" + ) + yield ( + f"{cfg.id}-{phase_name}-{tag}", + BackendCase( + page_size=_SPARSE_PAGE_SIZE, + kv_layout=_SPARSE_KV_LAYOUT, + dtype=_SPARSE_COMPUTE_DTYPE, + sparse_selection=selection, + use_kv_cache_manager_v2=_SPARSE_USE_KVM_V2, + **phases[phase_name], + **common, + ), + ) return # Bidirectional, KV-cache-free DiT / encoder workloads: only compute dtype. From 3d8c7c1ff050a073ccbdf01fe58e12866aad7053 Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Wed, 15 Jul 2026 22:01:46 -0700 Subject: [PATCH 04/17] [None][refactor] Address second review round on sparse test harness - Add BaseSparseAttentionConfig.sparse_topk property so ModelAttnConfig and BackendCase resolve the selection budget from the config instead of guessing. - Make sparse_selection_unit and sparse_topk derived properties of BackendCase; drop the sparse_attention_family field (use is_mla). - Share one _phases(long_len, gen_len) helper between the dense and sparse sweeps. - DSA capability gate: allow sm>=90 (FlashMLA); only FP8-KV DSA needs sm>=100. - Remove code comments that referenced the removed design. Signed-off-by: Yihan Wang --- .../_torch/attention_backend/sparse/utils.py | 4 +- tensorrt_llm/llmapi/llm_args.py | 14 +++ .../_torch/attention/backend_capability.py | 9 +- .../unittest/_torch/attention/backend_case.py | 90 +++++++++---------- .../_torch/attention/model_attn_config.py | 6 +- .../attention/test_attention_backends.py | 62 +++---------- 6 files changed, 78 insertions(+), 107 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/utils.py b/tensorrt_llm/_torch/attention_backend/sparse/utils.py index b9ba6edda535..a6611460ed4a 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/utils.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/utils.py @@ -49,9 +49,7 @@ def get_vanilla_sparse_attn_attention_backend( if sparse_params.algorithm == "rocket": return RocketVanillaAttention elif sparse_params.algorithm == "dsa": - # DSA selected-attention is implemented directly in the base - # VanillaAttention (its `_mla_forward_sparse` golden), so the vanilla - # slot for DSA is VanillaAttention itself. + # DSA's selected-attention golden lives in the base VanillaAttention. from ..vanilla import VanillaAttention return VanillaAttention elif sparse_params.algorithm == "minimax_m3": diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 5c6c89e88b3c..8c33c6645988 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -586,6 +586,20 @@ class BaseSparseAttentionConfig(StrictBaseModel): """Configuration for sparse attention.""" algorithm: str + @property + def sparse_topk(self) -> int: + """Per-token selection budget (top-k). + + The field name is not yet unified across algorithms (DSA / DeepSeek-V4 + use ``index_topk``, RocketKV uses ``topk``), so resolve it here. + """ + for name in ("index_topk", "topk"): + value = getattr(self, name, None) + if value is not None: + return value + raise NotImplementedError( + f"{type(self).__name__} does not expose a selection top-k.") + def supports_backend(self, backend: str) -> bool: """Override if the sparse attention algorithm does not support a subset of the possible backends. diff --git a/tests/unittest/_torch/attention/backend_capability.py b/tests/unittest/_torch/attention/backend_capability.py index 2f73bcdbaa4a..75ea258c6afe 100644 --- a/tests/unittest/_torch/attention/backend_capability.py +++ b/tests/unittest/_torch/attention/backend_capability.py @@ -117,8 +117,13 @@ def unsupported_reason(backend: str, case) -> Optional[str]: sparse_config = getattr(case, "sparse_attention_config", None) if sparse_config is not None: algorithm = sparse_config.get("algorithm") - if backend == "TRTLLM" and algorithm == "dsa" and sm < 100: - return f"TRTLLM DSA selected attention requires sm>=100 (have sm{sm})" + if backend == "TRTLLM" and algorithm == "dsa": + # DSA runs on Hopper (FlashMLA) and Blackwell; only the FP8 KV-cache + # DSA path is Blackwell-only. + if sm < 90: + return f"TRTLLM DSA requires sm>=90 (have sm{sm})" + if kv_dtype == "fp8" and sm < 100: + return f"TRTLLM DSA with FP8 KV cache requires sm>=100 (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 diff --git a/tests/unittest/_torch/attention/backend_case.py b/tests/unittest/_torch/attention/backend_case.py index f866580364ba..244d9d466aa8 100644 --- a/tests/unittest/_torch/attention/backend_case.py +++ b/tests/unittest/_torch/attention/backend_case.py @@ -54,9 +54,8 @@ # 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 -# Selected-sparse-MLA golden-vs-backend tolerance. Looser than dense bf16: the -# selected-attention path accumulates over gathered latent rows in bf16, so the -# TRTLLM kernel and the Vanilla golden diverge more than dense attention does. +# Golden-vs-backend tolerance for selected sparse MLA (bf16 latent-gather +# accumulation). SPARSE_ATOL = 1e-1 SPARSE_RTOL = 1e-2 @@ -65,6 +64,8 @@ BACKENDS_UNDER_TEST = ("TRTLLM",) + (("FLASHINFER",) if IS_FLASHINFER_AVAILABLE else ()) DEFAULT_MAX_NUM_TOKENS = 8192 _SPARSE_CONFIG_ADAPTER = TypeAdapter(SparseAttentionConfig) +# Selection representation per sparse algorithm (token- vs block-selecting). +_SELECTION_UNIT_BY_ALGORITHM = {"dsa": "token", "deepseek_v4": "token", "rocket": "block"} def _dtype_to_torch(dtype: str): @@ -102,17 +103,13 @@ class BackendCase: q_scaling: float = 1.0 page_size: int = 64 cache: str = "paged" # "paged" | "none" - # Serialized user-facing sparse config. It is lowered independently into - # backend params, metadata params, and the sparse KV-cache manager, exactly - # as in production. + # Serialized 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[dict] = None - # Backend-neutral sparse execution contract. This keeps the runner from - # inferring DSA/MLA semantics from an algorithm name. - sparse_attention_family: Optional[str] = None # "standard" | "mla" - sparse_selection_unit: Optional[str] = None # "none" | "token" | "block" - sparse_topk: Optional[int] = None - # Fake backend-neutral selection policy used in backend-level tests. The - # real algorithm/indexer is covered separately. + # Injected selection policy for backend-level tests; the real algorithm / + # indexer is covered separately. sparse_selection: str = "random" # "random" | "singleton" sparse_generation_tokens_per_seq: int = 1 # RoPE config: RopeParams kwargs (+ optional "is_neox"), or None to disable. @@ -144,9 +141,7 @@ class BackendCase: qk_nope_head_dim: Optional[int] = None qk_rope_head_dim: Optional[int] = None use_kv_cache_manager_v2: bool = False - # Low-level (atol, rtol) override for a single case, e.g. a captured/replayed - # case. Left unset by the model sweep, which derives tolerances by dtype (and - # a sparse default) in ``_tolerances``. + # Optional per-case (atol, rtol) override, e.g. for a captured/replayed case. atol: Optional[float] = None rtol: Optional[float] = None @@ -166,6 +161,19 @@ def is_cross(self) -> bool: def is_sparse(self) -> bool: return self.sparse_attention_config is not None + @property + def sparse_selection_unit(self) -> Optional[str]: + """Selection representation ("token" / "block") for this algorithm.""" + if self.sparse_attention_config is None: + return None + return _SELECTION_UNIT_BY_ALGORITHM.get(self.sparse_attention_config["algorithm"]) + + @property + def sparse_topk(self) -> Optional[int]: + """Per-token selection budget, resolved from the sparse config.""" + config = _sparse_config(self) + return None if config is None else config.sparse_topk + @property def prompt_lens(self) -> List[int]: """Original prompt lengths expected by fused generation kernels.""" @@ -251,27 +259,16 @@ def _sparse_config(case: BackendCase): def _validate_sparse_case(case: BackendCase) -> None: - """Reject incomplete or unsupported sparse harness contracts explicitly.""" - harness_fields = ( - case.sparse_attention_family, - case.sparse_selection_unit, - case.sparse_topk, - ) + """Reject incomplete or unsupported sparse contracts explicitly.""" if not case.is_sparse: - if any(value is not None for value in harness_fields): - raise ValueError("Sparse harness fields require sparse_attention_config") return - if any(value is None for value in harness_fields): - raise ValueError("Sparse backend cases require family, selection unit, and top-k") - if case.sparse_attention_family == "mla" and not case.is_mla: - raise ValueError("Sparse MLA harness requires is_mla=True") if case.sparse_selection_unit not in ("token", "block"): raise ValueError( f"Unsupported sparse selection unit: {case.sparse_selection_unit!r} " "(expected 'token' or 'block')" ) - if case.sparse_topk <= 0: - raise ValueError("Sparse backend cases require sparse_topk > 0") + if case.sparse_topk is None or case.sparse_topk <= 0: + raise ValueError("Sparse backend cases require a positive top-k") def _sparse_pretrained_config(case: BackendCase) -> SimpleNamespace: @@ -593,7 +590,7 @@ def _run_sparse_block_backend( def generate_sparse_mla_inputs(case: BackendCase, seed: int = 0) -> Dict: """Generate raw absorbed-MLA inputs plus backend-neutral sparse selections.""" - if case.sparse_attention_family != "mla" or case.sparse_selection_unit != "token": + if not case.is_mla or case.sparse_selection_unit != "token": raise ValueError("This generator supports token-selected sparse MLA only") gen = torch.Generator(device="cuda").manual_seed(seed) selection_gen = torch.Generator(device="cuda").manual_seed(seed + 1) @@ -820,8 +817,6 @@ def _run_sparse_mla_backend( sparse_attention_config=sparse_config, pretrained_config=pretrained_config, ) - # Every backend (Vanilla included) is built through the production factory so - # the harness does not special-case backend construction. attn = create_attention( backend, layer_idx=0, @@ -844,10 +839,8 @@ def _run_sparse_mla_backend( dtype=case.compute_dtype, skip_create_weights_in_init=True, ) - # Weight creation is skipped (the harness injects selections instead of - # running the indexer); update_quant_config still initializes the quant/FMHA - # state the fused backends need before forward, and is a safe no-op for - # Vanilla (the base implementation only stores quant_config). + # 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, @@ -1390,12 +1383,14 @@ def run_backend( matrix). MLA cases are dispatched to the absorbed-generation path. """ if case.is_sparse: - sparse_kind = (case.sparse_attention_family, case.sparse_selection_unit) - if sparse_kind == ("mla", "token"): + if case.is_mla and case.sparse_selection_unit == "token": return _run_sparse_mla_backend(case, backend, inputs, kv_layout=kv_layout) if case.sparse_selection_unit == "block": return _run_sparse_block_backend(case, backend, inputs, kv_layout=kv_layout) - raise ValueError(f"Unsupported sparse harness contract: {sparse_kind}") + raise ValueError( + f"Unsupported sparse contract: is_mla={case.is_mla}, " + f"selection_unit={case.sparse_selection_unit!r}" + ) if case.is_mla: if case.is_context_only: @@ -1588,13 +1583,15 @@ def run_case(case: BackendCase, *, seed: int = 0) -> Dict[str, torch.Tensor]: _validate_sparse_case(case) is_mla = case.is_mla if case.is_sparse: - sparse_kind = (case.sparse_attention_family, case.sparse_selection_unit) - if sparse_kind == ("mla", "token"): + if case.is_mla and case.sparse_selection_unit == "token": inputs = generate_sparse_mla_inputs(case, seed) elif case.sparse_selection_unit == "block": inputs = generate_sparse_block_inputs(case, seed) else: - raise ValueError(f"Unsupported sparse harness contract: {sparse_kind}") + raise ValueError( + f"Unsupported sparse contract: is_mla={case.is_mla}, " + f"selection_unit={case.sparse_selection_unit!r}" + ) elif is_mla: if case.is_context_only: inputs = generate_mla_context_inputs(case, seed) @@ -1633,11 +1630,8 @@ 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. Sparse cases are excluded: the standalone sparse runner - # is an eager correctness oracle that validates injected selections on the - # host and rebuilds each request's logical cache with Python loops, none - # of which is graph-capturable. Graph coverage of the DSA decode kernel is - # exercised at the model level, not by this backend-only oracle. + # 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, diff --git a/tests/unittest/_torch/attention/model_attn_config.py b/tests/unittest/_torch/attention/model_attn_config.py index 00c53d30d14e..0bce41be381d 100644 --- a/tests/unittest/_torch/attention/model_attn_config.py +++ b/tests/unittest/_torch/attention/model_attn_config.py @@ -84,10 +84,8 @@ class ModelAttnConfig: 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` etc. - # Everything the sparse test sweep needs (attention family, selection unit, - # top-k, tolerances) is derived from this + the fields above, so no separate - # harness/override/tolerance knobs live on the model config. + # 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[BaseSparseAttentionConfig] = None diff --git a/tests/unittest/_torch/attention/test_attention_backends.py b/tests/unittest/_torch/attention/test_attention_backends.py index 32ef7a61ea10..289a72da23c7 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,77 +59,42 @@ 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) -# Sparse test sweep dimensions are model-agnostic: one shared set for every -# sparse config rather than per-model knobs. Everything else (attention family, -# selection unit, top-k) is derived from the model's is_mla flag and its -# sparse_attention_config. +# 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 _SPARSE_SELECTION_PATTERNS = ("random", "singleton") _SPARSE_GENERATION_TOKENS_PER_SEQ = 2 -# Selection representation per algorithm (token-selecting vs block-selecting). -_SPARSE_SELECTION_UNIT = {"dsa": "token", "deepseek_v4": "token", "rocket": "block"} - - -def _sparse_selection_unit(cfg: ModelAttnConfig) -> str: - algo = cfg.sparse_attention_config.algorithm - if algo not in _SPARSE_SELECTION_UNIT: - raise ValueError(f"Unknown selection unit for sparse algorithm {algo!r}") - return _SPARSE_SELECTION_UNIT[algo] - - -def _sparse_topk(cfg: ModelAttnConfig) -> int: - """The selection budget: the indexer's top-k (index_topk) or the algo's topk.""" - sac = cfg.sparse_attention_config - topk = getattr(sac, "index_topk", None) - if topk is None: - topk = getattr(sac, "topk", None) - if topk is None: - raise ValueError(f"Cannot derive sparse top-k from {sac.algorithm!r} config") - return topk def _phases_for(cfg: ModelAttnConfig) -> dict: if cfg.sparse_attention_config is not None: - long_len = _sparse_topk(cfg) + 32 - generation_len = _SPARSE_GENERATION_TOKENS_PER_SEQ - return { - "ctx": dict( - seq_lens=[long_len, 73, 41], - num_cached_tokens=[0, 0, 0], - num_contexts=3, - ), - "gen": dict( - seq_lens=[generation_len] * 3, - num_cached_tokens=[long_len, 73, 41], - num_contexts=0, - ), - "mix": dict( - seq_lens=[long_len, generation_len, generation_len], - num_cached_tokens=[0, long_len, 73], - num_contexts=1, - ), - } + return _phases( + cfg.sparse_attention_config.sparse_topk + 32, _SPARSE_GENERATION_TOKENS_PER_SEQ + ) if cfg.mask != "sliding": return _PHASES @@ -170,9 +135,6 @@ def _common(cfg: ModelAttnConfig) -> dict: if cfg.sparse_attention_config is not None: common.update( sparse_attention_config=cfg.sparse_attention_config.model_dump(mode="json"), - sparse_attention_family="mla" if cfg.is_mla else "standard", - sparse_selection_unit=_sparse_selection_unit(cfg), - sparse_topk=_sparse_topk(cfg), sparse_generation_tokens_per_seq=_SPARSE_GENERATION_TOKENS_PER_SEQ, ) return common From 419e2feb6029a802a28cc2c3f1fbc17b54561eb8 Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Wed, 15 Jul 2026 22:10:49 -0700 Subject: [PATCH 05/17] [None][refactor] Drop redundant comment in vanilla sparse factory Signed-off-by: Yihan Wang --- tensorrt_llm/_torch/attention_backend/sparse/utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/utils.py b/tensorrt_llm/_torch/attention_backend/sparse/utils.py index a6611460ed4a..c1f718d4a418 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/utils.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/utils.py @@ -49,7 +49,6 @@ def get_vanilla_sparse_attn_attention_backend( if sparse_params.algorithm == "rocket": return RocketVanillaAttention elif sparse_params.algorithm == "dsa": - # DSA's selected-attention golden lives in the base VanillaAttention. from ..vanilla import VanillaAttention return VanillaAttention elif sparse_params.algorithm == "minimax_m3": From b5a82704c1bcde0672ddf3927c1c60e8b5804cf1 Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Wed, 15 Jul 2026 23:35:18 -0700 Subject: [PATCH 06/17] [None][refactor] Third review round: simplify sparse test harness - Store the SparseAttentionConfig object on ModelAttnConfig/BackendCase (not a JSON dict); drop _sparse_config and the TypeAdapter. - Drop _sparse_pretrained_config: DSA lowering and the sparse KV-cache manager work with pretrained_config=None. - Remove BackendCase.atol/rtol (unused) and sparse_generation_tokens_per_seq (no spec-decoding support; gen_len is 1). - One mixed selection per phase (alternating singleton/random rows); remove sparse_selection, the pattern loop, and the id suffix. - Move the _validate_sparse_case call inside the is_sparse block. Signed-off-by: Yihan Wang --- .../_torch/attention/backend_capability.py | 2 +- .../unittest/_torch/attention/backend_case.py | 177 ++++++------------ .../_torch/attention/model_attn_config.py | 4 +- .../attention/test_attention_backends.py | 46 ++--- 4 files changed, 77 insertions(+), 152 deletions(-) diff --git a/tests/unittest/_torch/attention/backend_capability.py b/tests/unittest/_torch/attention/backend_capability.py index 75ea258c6afe..f34374f98007 100644 --- a/tests/unittest/_torch/attention/backend_capability.py +++ b/tests/unittest/_torch/attention/backend_capability.py @@ -116,7 +116,7 @@ def unsupported_reason(backend: str, case) -> Optional[str]: sparse_config = getattr(case, "sparse_attention_config", None) if sparse_config is not None: - algorithm = sparse_config.get("algorithm") + algorithm = sparse_config.algorithm if backend == "TRTLLM" and algorithm == "dsa": # DSA runs on Hopper (FlashMLA) and Blackwell; only the FP8 KV-cache # DSA path is Blackwell-only. diff --git a/tests/unittest/_torch/attention/backend_case.py b/tests/unittest/_torch/attention/backend_case.py index 244d9d466aa8..bea540182d1b 100644 --- a/tests/unittest/_torch/attention/backend_case.py +++ b/tests/unittest/_torch/attention/backend_case.py @@ -14,13 +14,11 @@ import math from dataclasses import asdict, dataclass -from types import SimpleNamespace from typing import Dict, List, Optional import torch from backend_capability import BACKEND_CAPS, unsupported_reason from kv_cache_utils import apply_rope, fill_kv_cache_logical, make_position_ids -from pydantic import TypeAdapter import tensorrt_llm from tensorrt_llm._torch.attention_backend.interface import ( @@ -63,7 +61,6 @@ # included when available, so callers can iterate this list unconditionally. BACKENDS_UNDER_TEST = ("TRTLLM",) + (("FLASHINFER",) if IS_FLASHINFER_AVAILABLE else ()) DEFAULT_MAX_NUM_TOKENS = 8192 -_SPARSE_CONFIG_ADAPTER = TypeAdapter(SparseAttentionConfig) # Selection representation per sparse algorithm (token- vs block-selecting). _SELECTION_UNIT_BY_ALGORITHM = {"dsa": "token", "deepseek_v4": "token", "rocket": "block"} @@ -103,15 +100,11 @@ class BackendCase: q_scaling: float = 1.0 page_size: int = 64 cache: str = "paged" # "paged" | "none" - # Serialized 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[dict] = None - # Injected selection policy for backend-level tests; the real algorithm / - # indexer is covered separately. - sparse_selection: str = "random" # "random" | "singleton" - sparse_generation_tokens_per_seq: int = 1 + # 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 @@ -141,9 +134,6 @@ class BackendCase: qk_nope_head_dim: Optional[int] = None qk_rope_head_dim: Optional[int] = None use_kv_cache_manager_v2: bool = False - # Optional per-case (atol, rtol) override, e.g. for a captured/replayed case. - atol: Optional[float] = None - rtol: Optional[float] = None @property def num_seqs(self) -> int: @@ -166,13 +156,14 @@ def sparse_selection_unit(self) -> Optional[str]: """Selection representation ("token" / "block") for this algorithm.""" if self.sparse_attention_config is None: return None - return _SELECTION_UNIT_BY_ALGORITHM.get(self.sparse_attention_config["algorithm"]) + return _SELECTION_UNIT_BY_ALGORITHM.get(self.sparse_attention_config.algorithm) @property def sparse_topk(self) -> Optional[int]: """Per-token selection budget, resolved from the sparse config.""" - config = _sparse_config(self) - return None if config is None else config.sparse_topk + if self.sparse_attention_config is None: + return None + return self.sparse_attention_config.sparse_topk @property def prompt_lens(self) -> List[int]: @@ -251,17 +242,8 @@ def _rope_params_from_dict(d: dict) -> RopeParams: return RopeParams(**kwargs) -def _sparse_config(case: BackendCase): - """Restore the production sparse config from its serialized case form.""" - if case.sparse_attention_config is None: - return None - return _SPARSE_CONFIG_ADAPTER.validate_python(case.sparse_attention_config) - - def _validate_sparse_case(case: BackendCase) -> None: - """Reject incomplete or unsupported sparse contracts explicitly.""" - if not case.is_sparse: - return + """Reject unsupported sparse contracts (called only for sparse cases).""" if case.sparse_selection_unit not in ("token", "block"): raise ValueError( f"Unsupported sparse selection unit: {case.sparse_selection_unit!r} " @@ -271,22 +253,6 @@ def _validate_sparse_case(case: BackendCase) -> None: raise ValueError("Sparse backend cases require a positive top-k") -def _sparse_pretrained_config(case: BackendCase) -> SimpleNamespace: - """Build one HF-like config shared by every sparse lowering boundary.""" - values = dict( - rms_norm_eps=1e-6, - hidden_size=case.hidden_size, - num_attention_heads=case.num_heads, - num_key_value_heads=case.num_kv_heads, - 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, - v_head_dim=case.v_head_dim, - ) - return SimpleNamespace(**values) - - 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) @@ -366,7 +332,6 @@ def _build_mla_kv_cache_manager( backend: str, sparse_config=None, model_config=None, - pretrained_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 @@ -401,12 +366,11 @@ def _build_mla_kv_cache_manager( if sparse_config is not None and backend != "VANILLA": cls = get_sparse_attn_kv_cache_manager(sparse_config) - if model_config is None or pretrained_config is None: - raise ValueError("Sparse cache manager requires model and pretrained configs") + if model_config is None: + raise ValueError("Sparse cache manager requires a model config") kwargs.update( sparse_attention_config=sparse_config, model_config=model_config, - pretrained_config=pretrained_config, ) else: cls = KVCacheManagerV2 if case.use_kv_cache_manager_v2 else KVCacheManager @@ -444,8 +408,15 @@ def generate_mla_gen_inputs(case: BackendCase, seed: int = 0) -> Dict: def _build_sparse_topk_indices( case: BackendCase, generator: torch.Generator, -) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - """Build causal request-local selections for a sparse backend case.""" +) -> 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. + """ if case.sparse_selection_unit != "token": raise ValueError( f"Token selection builder cannot handle {case.sparse_selection_unit!r} selections" @@ -453,49 +424,28 @@ def _build_sparse_topk_indices( topk = case.sparse_topk if topk is None or topk <= 0: raise ValueError("Sparse token-selection cases require a positive top-k") - if case.sparse_selection not in ("random", "singleton"): - raise ValueError(f"Unsupported sparse selection policy: {case.sparse_selection}") - - 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") - if case.sparse_selection == "singleton" - else None - ) + + 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 case.sparse_selection == "singleton": - if max_index + 1 <= topk: - # DSA is effectively dense while the visible prefix fits - # within top-k; a real indexer returns every causal token. - indices[row, : max_index + 1] = torch.arange( - max_index + 1, - dtype=torch.int32, - device="cuda", - ) - else: - # Repeating one logical token across all top-k slots keeps - # the physical shape production-valid while yielding an - # analytic singleton-value oracle for truly sparse rows. - selector = row % 3 - selected = (0, max_index, max_index // 2)[selector] - indices[row].fill_(selected) - singleton_oracle_mask[row] = True + 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: - valid = min(max_index + 1, topk) - selected = torch.randperm( - max_index + 1, - generator=generator, - device="cuda", - )[:valid] - indices[row, :valid] = torch.sort(selected).values.to(torch.int32) + 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 @@ -523,21 +473,15 @@ def _build_sparse_block_indices( topk = case.sparse_topk if topk is None or topk <= 0: raise ValueError("Block-selection cases require a positive top-k") - if case.sparse_selection not in ("random", "singleton"): - raise ValueError(f"Unsupported sparse selection policy: {case.sparse_selection}") indices = torch.full((case.nnz_q, topk), -1, dtype=torch.int32, device="cuda") row = 0 for cached_len, q_len in zip(case.num_cached_tokens, case.seq_lens, strict=True): for token_idx in range(q_len): - max_block = (cached_len + token_idx) // block_size - num_blocks = max_block + 1 - if case.sparse_selection == "singleton": - indices[row].fill_(row % num_blocks) - else: - valid = min(num_blocks, topk) - selected = torch.randperm(num_blocks, generator=generator, device="cuda")[:valid] - indices[row, :valid] = torch.sort(selected).values.to(torch.int32) + num_blocks = (cached_len + token_idx) // block_size + 1 + valid = min(num_blocks, topk) + selected = torch.randperm(num_blocks, generator=generator, device="cuda")[:valid] + indices[row, :valid] = torch.sort(selected).values.to(torch.int32) row += 1 return indices @@ -552,11 +496,9 @@ def generate_sparse_block_inputs(case: BackendCase, seed: int = 0) -> Dict: ``_run_sparse_block_backend`` will call. Until both exist this raises so the contract is explicit rather than silently mis-run. """ - sparse_config = _sparse_config(case) + sparse_config = case.sparse_attention_config assert sparse_config is not None - sparse_params = sparse_config.to_sparse_params( - layer_idx=0, pretrained_config=_sparse_pretrained_config(case) - ) + sparse_params = sparse_config.to_sparse_params(layer_idx=None, pretrained_config=None) block_size = getattr(sparse_params, "indices_block_size") selection_gen = torch.Generator(device="cuda").manual_seed(seed + 1) _ = _build_sparse_block_indices(case, selection_gen, block_size) @@ -636,8 +578,10 @@ def generate_sparse_mla_inputs(case: BackendCase, seed: int = 0) -> Dict: 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 case.sparse_selection == "singleton": + if bool(singleton_oracle_mask.any()): new_per_seq = _split_packed_tokens(expected_new_latent, case.seq_lens) selected_values = [] row = 0 @@ -797,16 +741,12 @@ def _run_sparse_mla_backend( kv_layout: str, ) -> torch.Tensor: """Run selected sparse MLA through production backend/config lowering.""" - sparse_config = _sparse_config(case) + sparse_config = case.sparse_attention_config assert sparse_config is not None - pretrained_config = _sparse_pretrained_config(case) - sparse_params = sparse_config.to_sparse_params( - layer_idx=0, - pretrained_config=pretrained_config, - ) - sparse_metadata_params = sparse_config.to_sparse_metadata_params( - pretrained_config=pretrained_config - ) + # 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 @@ -815,7 +755,6 @@ def _run_sparse_mla_backend( model_config = ModelConfig( mapping=mapping, sparse_attention_config=sparse_config, - pretrained_config=pretrained_config, ) attn = create_attention( backend, @@ -834,7 +773,7 @@ def _run_sparse_mla_backend( # 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=case.sparse_generation_tokens_per_seq, + predicted_tokens_per_seq=1, sparse_params=sparse_params, dtype=case.compute_dtype, skip_create_weights_in_init=True, @@ -847,7 +786,6 @@ def _run_sparse_mla_backend( backend, sparse_config, model_config, - pretrained_config, ) try: @@ -1254,11 +1192,6 @@ 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.atol is not None or case.rtol is not None: - if case.atol is None or case.rtol is None: - raise ValueError("BackendCase.atol and rtol must be set together") - return case.atol, case.rtol - if case.is_sparse: return SPARSE_ATOL, SPARSE_RTOL @@ -1580,9 +1513,9 @@ def run_case(case: BackendCase, *, seed: int = 0) -> Dict[str, torch.Tensor]: Returns the per-backend outputs (including ``"VANILLA"`` golden) for callers that want the raw tensors (e.g. the minimizer). """ - _validate_sparse_case(case) is_mla = case.is_mla if case.is_sparse: + _validate_sparse_case(case) if case.is_mla and case.sparse_selection_unit == "token": inputs = generate_sparse_mla_inputs(case, seed) elif case.sparse_selection_unit == "block": @@ -1602,7 +1535,7 @@ 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 case.sparse_selection == "singleton" and inputs.get("expected_output") is not None: + 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]) diff --git a/tests/unittest/_torch/attention/model_attn_config.py b/tests/unittest/_torch/attention/model_attn_config.py index 0bce41be381d..076f739e5e13 100644 --- a/tests/unittest/_torch/attention/model_attn_config.py +++ b/tests/unittest/_torch/attention/model_attn_config.py @@ -57,7 +57,7 @@ from dataclasses import dataclass from typing import List, Optional -from tensorrt_llm.llmapi.llm_args import BaseSparseAttentionConfig, DeepSeekSparseAttentionConfig +from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig, SparseAttentionConfig @dataclass(frozen=True) @@ -86,7 +86,7 @@ class ModelAttnConfig: 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[BaseSparseAttentionConfig] = None + sparse_attention_config: Optional[SparseAttentionConfig] = None # --------------------------------------------------------------------------- diff --git a/tests/unittest/_torch/attention/test_attention_backends.py b/tests/unittest/_torch/attention/test_attention_backends.py index 289a72da23c7..07bdca0e6fc7 100644 --- a/tests/unittest/_torch/attention/test_attention_backends.py +++ b/tests/unittest/_torch/attention/test_attention_backends.py @@ -86,15 +86,11 @@ def _phases_from_window(window: int) -> dict: _SPARSE_KV_LAYOUT = "HND" _SPARSE_PAGE_SIZE = 64 _SPARSE_USE_KVM_V2 = False -_SPARSE_SELECTION_PATTERNS = ("random", "singleton") -_SPARSE_GENERATION_TOKENS_PER_SEQ = 2 def _phases_for(cfg: ModelAttnConfig) -> dict: if cfg.sparse_attention_config is not None: - return _phases( - cfg.sparse_attention_config.sparse_topk + 32, _SPARSE_GENERATION_TOKENS_PER_SEQ - ) + return _phases(cfg.sparse_attention_config.sparse_topk + 32, gen_len=1) if cfg.mask != "sliding": return _PHASES @@ -133,10 +129,7 @@ def _common(cfg: ModelAttnConfig) -> dict: hidden_size=cfg.hidden_size, ) if cfg.sparse_attention_config is not None: - common.update( - sparse_attention_config=cfg.sparse_attention_config.model_dump(mode="json"), - sparse_generation_tokens_per_seq=_SPARSE_GENERATION_TOKENS_PER_SEQ, - ) + common.update(sparse_attention_config=cfg.sparse_attention_config) return common @@ -160,28 +153,27 @@ def _expand(cfg: ModelAttnConfig, precisions, kv_layouts, page_sizes): phases = _phases_for(cfg) # Sparse cases use one model-agnostic sweep (bf16 latent cache, fixed layout/ - # page/manager). Fake request-local selections isolate backend execution from + # 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"): - for selection in _SPARSE_SELECTION_PATTERNS: - tag = ( - f"{_prec_tag(_SPARSE_COMPUTE_DTYPE, None)}-{_SPARSE_KV_LAYOUT}" - f"-p{_SPARSE_PAGE_SIZE}-{manager}-{selection}" - ) - yield ( - f"{cfg.id}-{phase_name}-{tag}", - BackendCase( - page_size=_SPARSE_PAGE_SIZE, - kv_layout=_SPARSE_KV_LAYOUT, - dtype=_SPARSE_COMPUTE_DTYPE, - sparse_selection=selection, - use_kv_cache_manager_v2=_SPARSE_USE_KVM_V2, - **phases[phase_name], - **common, - ), - ) + 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. From ad273667a13bd75447ba4d8b3d639a9f21f2be09 Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Thu, 16 Jul 2026 01:23:52 -0700 Subject: [PATCH 07/17] [None][refactor] Unify sparse selection: token is block_size 1 A token selection is a block selection with block_size == 1, so the harness no longer distinguishes token vs block selection units. Selection granularity is derived from the config (get_indices_block_size(); 1 for DSA) and the single selected-MLA path is exercised by DSA. Removes the unverified block-selection scaffolding (builder + runner/generator stubs) and the selection-unit plumbing; VanillaAttention._mla_forward_sparse guards block_size == 1. block_size > 1 is deferred to a follow-up PR (RocketKV / MiniMax-M3). Signed-off-by: Yihan Wang --- .../_torch/attention_backend/vanilla.py | 11 +- .../unittest/_torch/attention/backend_case.py | 132 +++--------------- 2 files changed, 29 insertions(+), 114 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/vanilla.py b/tensorrt_llm/_torch/attention_backend/vanilla.py index 29f65d2ac3f7..3ce3ea5a1419 100644 --- a/tensorrt_llm/_torch/attention_backend/vanilla.py +++ b/tensorrt_llm/_torch/attention_backend/vanilla.py @@ -725,9 +725,16 @@ def _mla_forward_sparse( """Run selected sparse MLA from caller-provided local top-k rows. The sparse algorithm owns selection. This golden consumes its - request-local token positions, gathers the selected latent K/V, and - performs the absorbed MLA attention directly in PyTorch. + request-local selections, gathers the selected latent K/V, and performs + the absorbed MLA attention directly in PyTorch. Selections are block + indices; this reference implements ``block_size == 1`` (token selection), + which is what the MLA sparse algorithms (DSA / DeepSeek-V4) use. """ + block_size = getattr(self.sparse_params, "indices_block_size", 1) + if block_size != 1: + raise NotImplementedError( + "Vanilla selected MLA supports block_size 1 (token selection); " + f"got block_size {block_size}") if attention_input_type == AttentionInputType.context_only: seq_start, seq_end = 0, metadata.num_contexts elif attention_input_type == AttentionInputType.generation_only: diff --git a/tests/unittest/_torch/attention/backend_case.py b/tests/unittest/_torch/attention/backend_case.py index bea540182d1b..1de44af04944 100644 --- a/tests/unittest/_torch/attention/backend_case.py +++ b/tests/unittest/_torch/attention/backend_case.py @@ -61,8 +61,6 @@ # included when available, so callers can iterate this list unconditionally. BACKENDS_UNDER_TEST = ("TRTLLM",) + (("FLASHINFER",) if IS_FLASHINFER_AVAILABLE else ()) DEFAULT_MAX_NUM_TOKENS = 8192 -# Selection representation per sparse algorithm (token- vs block-selecting). -_SELECTION_UNIT_BY_ALGORITHM = {"dsa": "token", "deepseek_v4": "token", "rocket": "block"} def _dtype_to_torch(dtype: str): @@ -152,18 +150,18 @@ def is_sparse(self) -> bool: return self.sparse_attention_config is not None @property - def sparse_selection_unit(self) -> Optional[str]: - """Selection representation ("token" / "block") for this algorithm.""" + def sparse_topk(self) -> Optional[int]: + """Per-token selection budget, resolved from the sparse config.""" if self.sparse_attention_config is None: return None - return _SELECTION_UNIT_BY_ALGORITHM.get(self.sparse_attention_config.algorithm) + return self.sparse_attention_config.sparse_topk @property - def sparse_topk(self) -> Optional[int]: - """Per-token selection budget, resolved from the sparse config.""" + def sparse_block_size(self) -> Optional[int]: + """Selection granularity: token selection is a block of size 1.""" if self.sparse_attention_config is None: return None - return self.sparse_attention_config.sparse_topk + return self.sparse_attention_config.get_indices_block_size() @property def prompt_lens(self) -> List[int]: @@ -244,13 +242,10 @@ def _rope_params_from_dict(d: dict) -> RopeParams: def _validate_sparse_case(case: BackendCase) -> None: """Reject unsupported sparse contracts (called only for sparse cases).""" - if case.sparse_selection_unit not in ("token", "block"): - raise ValueError( - f"Unsupported sparse selection unit: {case.sparse_selection_unit!r} " - "(expected 'token' or 'block')" - ) if case.sparse_topk is None or case.sparse_topk <= 0: raise ValueError("Sparse backend cases require a positive top-k") + if case.sparse_block_size is None or case.sparse_block_size < 1: + raise ValueError("Sparse backend cases require a positive selection block size") def _randn(gen: torch.Generator, dtype: torch.dtype, *shape) -> torch.Tensor: @@ -416,14 +411,17 @@ def _build_sparse_topk_indices( ``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 block indices; with ``block_size == 1`` (DSA) a block is a + single token, so these are token selections. """ - if case.sparse_selection_unit != "token": + if case.sparse_block_size != 1: raise ValueError( - f"Token selection builder cannot handle {case.sparse_selection_unit!r} selections" + f"Sparse MLA selection builder only supports block_size 1, got {case.sparse_block_size}" ) topk = case.sparse_topk if topk is None or topk <= 0: - raise ValueError("Sparse token-selection cases require a positive top-k") + 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") @@ -450,90 +448,10 @@ def _build_sparse_topk_indices( return indices, singleton_oracle_mask -def _build_sparse_block_indices( - case: BackendCase, - generator: torch.Generator, - block_size: int, -) -> torch.Tensor: - """Build causal request-local *block* selections for a block-sparse case. - - Skeleton for the next block-granular algorithm family (e.g. RocketKV, whose - ``indices_block_size == page_size``). Mirrors ``_build_sparse_topk_indices`` - but selects logical KV *blocks* instead of tokens: for a query at causal - position ``p`` the selectable blocks are ``0 .. p // block_size`` (the block - holding ``p`` is inclusive), and each row is padded to ``sparse_topk`` with - ``-1``. Returns ``[nnz_q, sparse_topk]`` int32 block indices. - """ - if case.sparse_selection_unit != "block": - raise ValueError( - f"Block selection builder cannot handle {case.sparse_selection_unit!r} selections" - ) - if block_size <= 0: - raise ValueError("Block-selection cases require a positive block size") - topk = case.sparse_topk - if topk is None or topk <= 0: - raise ValueError("Block-selection cases require a positive top-k") - - indices = torch.full((case.nnz_q, topk), -1, dtype=torch.int32, device="cuda") - row = 0 - for cached_len, q_len in zip(case.num_cached_tokens, case.seq_lens, strict=True): - for token_idx in range(q_len): - num_blocks = (cached_len + token_idx) // block_size + 1 - valid = min(num_blocks, topk) - selected = torch.randperm(num_blocks, generator=generator, device="cuda")[:valid] - indices[row, :valid] = torch.sort(selected).values.to(torch.int32) - row += 1 - return indices - - -def generate_sparse_block_inputs(case: BackendCase, seed: int = 0) -> Dict: - """Skeleton input generator for block-selected sparse attention. - - TODO(sparse-block): flesh out once a block-granular algorithm lands. The - reusable ``_build_sparse_block_indices`` builder is wired here; still missing - are (a) the standard/MLA input tensors for the block family and (b) the - matching Vanilla block reference forward in ``VanillaAttention`` that - ``_run_sparse_block_backend`` will call. Until both exist this raises so the - contract is explicit rather than silently mis-run. - """ - sparse_config = case.sparse_attention_config - assert sparse_config is not None - sparse_params = sparse_config.to_sparse_params(layer_idx=None, pretrained_config=None) - block_size = getattr(sparse_params, "indices_block_size") - selection_gen = torch.Generator(device="cuda").manual_seed(seed + 1) - _ = _build_sparse_block_indices(case, selection_gen, block_size) - raise NotImplementedError( - "Block-selected sparse attention is a skeleton: add the block input " - "tensors and a VanillaAttention block reference forward, plus a block " - "ModelAttnConfig (e.g. RocketKV), before enabling this path." - ) - - -def _run_sparse_block_backend( - case: BackendCase, - backend: str, - inputs: Dict, - *, - kv_layout: str, -) -> torch.Tensor: - """Skeleton runner for block-selected sparse attention. - - TODO(sparse-block): implement the block-family forward once a Vanilla block - reference exists, following ``_run_sparse_mla_backend`` (lower the production - ``sparse_attention_config`` via ``to_sparse_params`` / the sparse KV-cache - manager, inject ``inputs`` block selections, and compare against the Vanilla - golden). - """ - raise NotImplementedError( - "Block-selected sparse backend runner is a skeleton; see " - "_run_sparse_mla_backend for the token-selected MLA reference." - ) - - 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 or case.sparse_selection_unit != "token": - raise ValueError("This generator supports token-selected sparse MLA only") + if not case.is_mla or case.sparse_block_size != 1: + raise ValueError("This generator supports selected sparse MLA with block_size 1 only") gen = torch.Generator(device="cuda").manual_seed(seed) selection_gen = torch.Generator(device="cuda").manual_seed(seed + 1) cdt = case.compute_dtype @@ -1316,14 +1234,9 @@ def run_backend( matrix). MLA cases are dispatched to the absorbed-generation path. """ if case.is_sparse: - if case.is_mla and case.sparse_selection_unit == "token": + if case.is_mla: return _run_sparse_mla_backend(case, backend, inputs, kv_layout=kv_layout) - if case.sparse_selection_unit == "block": - return _run_sparse_block_backend(case, backend, inputs, kv_layout=kv_layout) - raise ValueError( - f"Unsupported sparse contract: is_mla={case.is_mla}, " - f"selection_unit={case.sparse_selection_unit!r}" - ) + raise ValueError(f"Unsupported sparse contract: is_mla={case.is_mla}") if case.is_mla: if case.is_context_only: @@ -1516,15 +1429,10 @@ def run_case(case: BackendCase, *, seed: int = 0) -> Dict[str, torch.Tensor]: is_mla = case.is_mla if case.is_sparse: _validate_sparse_case(case) - if case.is_mla and case.sparse_selection_unit == "token": + if case.is_mla: inputs = generate_sparse_mla_inputs(case, seed) - elif case.sparse_selection_unit == "block": - inputs = generate_sparse_block_inputs(case, seed) else: - raise ValueError( - f"Unsupported sparse contract: is_mla={case.is_mla}, " - f"selection_unit={case.sparse_selection_unit!r}" - ) + 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) From ec910124515551b3535d6134fdfbc3929e65cb83 Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Thu, 16 Jul 2026 01:32:35 -0700 Subject: [PATCH 08/17] [None][refactor] Use full names for sparse prediction locals in vanilla forward Signed-off-by: Yihan Wang --- tensorrt_llm/_torch/attention_backend/vanilla.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/vanilla.py b/tensorrt_llm/_torch/attention_backend/vanilla.py index 3ce3ea5a1419..e02dc8491202 100644 --- a/tensorrt_llm/_torch/attention_backend/vanilla.py +++ b/tensorrt_llm/_torch/attention_backend/vanilla.py @@ -942,16 +942,16 @@ def forward(self, if sparse_algorithm != "dsa": raise ValueError( "Vanilla selected MLA currently supports only DSA") - kv_idx, kv_off = self.sparse_kv_predict(q, k, metadata, - forward_args) - at_idx, at_off = self.sparse_attn_predict( + sparse_kv_indices, sparse_kv_offsets = self.sparse_kv_predict( + q, k, metadata, forward_args) + sparse_attn_indices, sparse_attn_offsets = self.sparse_attn_predict( q, k, metadata, forward_args) forward_args.sparse_prediction = replace( forward_args.sparse_prediction, - sparse_kv_indices=kv_idx, - sparse_kv_offsets=kv_off, - sparse_attn_indices=at_idx, - sparse_attn_offsets=at_off, + sparse_kv_indices=sparse_kv_indices, + sparse_kv_offsets=sparse_kv_offsets, + sparse_attn_indices=sparse_attn_indices, + sparse_attn_offsets=sparse_attn_offsets, sparse_attn_indices_block_size=getattr( self.sparse_params, "indices_block_size"), ) From e4ebfcb88a4b3d4078e4859fd940677843deba35 Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Thu, 16 Jul 2026 01:54:35 -0700 Subject: [PATCH 09/17] [None][refactor] Move VanillaIndexer out of production vanilla.py into its test VanillaIndexer is a test-only fp32 reference used only by test_sparse_mla_forward.py and has a different API from the production DSA Indexer, so it does not belong in the production attention backend. Signed-off-by: Yihan Wang --- .../_torch/attention_backend/vanilla.py | 100 +----------------- .../sparse/test_sparse_mla_forward.py | 80 +++++++++++++- 2 files changed, 79 insertions(+), 101 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/vanilla.py b/tensorrt_llm/_torch/attention_backend/vanilla.py index e02dc8491202..ea193348dab8 100644 --- a/tensorrt_llm/_torch/attention_backend/vanilla.py +++ b/tensorrt_llm/_torch/attention_backend/vanilla.py @@ -3,7 +3,7 @@ import math from dataclasses import replace -from typing import Callable, Optional +from typing import Optional import torch import torch.nn.functional as F @@ -1043,101 +1043,3 @@ def forward(self, attn_output = attn_output.view(q_len, -1) return attn_output - - -class VanillaIndexer: - """fp32 reference for the production DSA / DeepSeek-V4 sparse ``Indexer``. - - The production indexer (``sparse/dsa.py``) selects the top-k KV a query may - attend to. ``VanillaIndexer`` mirrors that selection math in plain fp32 so it - can serve as the indexer golden, analogous to how ``VanillaAttention`` is the - attention golden. - - It is a standalone reference, deliberately **not** wired into - ``VanillaAttention.forward``: - - * The indexer consumes index-space inputs (``qr`` / ``hidden_states`` / an - index-space K cache) that ``forward(q, k, v, metadata)`` never receives; - folding it into the forward path would widen the generic backend interface - with algorithm-specific tensors. - * A meaningful indexer golden must run against the *production indexer's own - weights*, so this class **wraps a production ``Indexer`` instance** and - reads its ``wq_b`` / ``weights_proj`` / ``softmax_scale`` / ... rather than - owning independent (non-comparable) parameters. - - It owns the parts common to DSA and DeepSeek-V4 -- the index-space query - projection, the per-head token weights, the logit scoring, and the top-k. - Algorithm-specific K comes from a compressor / ``wk`` projection, so the - caller supplies the reference K and this class handles the rest. - - Selection is discrete: top-k over near-tied logits, and an fp32 reference vs - an fp8/fp4 kernel can pick different borderline tokens. Compare the selected - index *set*, never exact attention outputs. - """ - - def __init__(self, indexer): - self.indexer = indexer - self.n_heads = indexer.n_heads - self.head_dim = indexer.head_dim - self.rope_dim = indexer.rope_dim - self.softmax_scale = indexer.softmax_scale - self.indexer_k_dtype = getattr(indexer, "indexer_k_dtype", None) - - @property - def uses_fp4(self) -> bool: - return self.indexer_k_dtype == "fp4" - - def project_query( - self, - qr: torch.Tensor, - position_ids: torch.Tensor, - freqs_cis: torch.Tensor, - rope_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], - *, - fp4_prep: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, - ) -> torch.Tensor: - """Index-space query: ``wq_b`` GEMM + RoPE on the rope slice. - - ``rope_fn(slice, freqs)`` applies RoPE in place (the caller injects the - algorithm's rotary helper). ``fp4_prep`` optionally applies the fp4 - indexer quant/dequant. Returns ``[num_tokens, n_heads, head_dim]``. - """ - num_tokens = qr.shape[0] - q = F.linear(qr, self.indexer.wq_b.weight) - q = q.view(num_tokens, self.n_heads, self.head_dim).unsqueeze(0) - rope_fn(q[..., -self.rope_dim:], freqs_cis[position_ids.long()]) - q = q.squeeze(0) - if fp4_prep is not None: - q = fp4_prep(q) - return q - - def token_weights(self, hidden_states: torch.Tensor) -> torch.Tensor: - """Per-head token weights: ``weights_proj`` GEMM scaled by ``n_heads**-0.5``.""" - weights = F.linear(hidden_states, self.indexer.weights_proj.weight) - return weights.float() * (self.n_heads**-0.5) - - def scores(self, q_row: torch.Tensor, k: torch.Tensor, - weights_row: torch.Tensor) -> torch.Tensor: - """Per-KV index logits for one query token (fp32). - - ``q_row`` is ``[n_heads, head_dim]``, ``k`` is ``[num_kv, head_dim]``, - ``weights_row`` is ``[n_heads]``. Mirrors the production indexer: per-head - ReLU(q·k) scaled by ``softmax_scale``, combined with the per-head weights. - Returns ``[num_kv]`` logits. - """ - head_scores = torch.einsum("hd,kd->hk", q_row.float(), k.float()) - head_scores = F.relu(head_scores) * self.softmax_scale - return (head_scores * weights_row.float().unsqueeze(-1)).sum(dim=0) - - def topk_from_scores(self, scores: torch.Tensor, - topk_tokens: int) -> torch.Tensor: - """Top-k KV positions for one query token, ``-1``-padded to ``topk_tokens``.""" - row = torch.full((topk_tokens, ), - -1, - dtype=torch.int32, - device=scores.device) - if scores.numel() == 0: - return row - k = min(topk_tokens, scores.numel()) - row[:k] = torch.topk(scores.float(), k, dim=-1).indices.to(torch.int32) - return row diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py b/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py index 104742a01eb0..260a6dfe0e8f 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py @@ -18,10 +18,11 @@ import os from dataclasses import dataclass from types import SimpleNamespace -from typing import List, Optional +from typing import Callable, List, Optional import pytest import torch +import torch.nn.functional as F import tensorrt_llm import tensorrt_llm.bindings @@ -32,7 +33,6 @@ from tensorrt_llm._torch.attention_backend.sparse.dsa import (HAS_FAST_HADAMARD, DSACacheManager) from tensorrt_llm._torch.attention_backend.utils import get_attention_backend -from tensorrt_llm._torch.attention_backend.vanilla import VanillaIndexer from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.modules.mla import MLA @@ -59,6 +59,82 @@ HAS_FLASH_MLA = False +class VanillaIndexer: + """fp32 reference for the production DSA / DeepSeek-V4 sparse ``Indexer``. + + Wraps a production ``Indexer`` instance and reproduces its selection math in + plain fp32 (index-space query projection, per-head token weights, logit + scoring, top-k) so the trtllm indexer's selection can be compared against a + same-weights reference. Algorithm-specific K comes from a compressor / ``wk`` + projection, so the caller supplies the reference K. + + Selection is discrete: top-k over near-tied logits, and an fp32 reference vs + an fp8/fp4 kernel can pick different borderline tokens. Compare the selected + index *set*, never exact attention outputs. + """ + + def __init__(self, indexer): + self.indexer = indexer + self.n_heads = indexer.n_heads + self.head_dim = indexer.head_dim + self.rope_dim = indexer.rope_dim + self.softmax_scale = indexer.softmax_scale + self.indexer_k_dtype = getattr(indexer, "indexer_k_dtype", None) + + @property + def uses_fp4(self) -> bool: + return self.indexer_k_dtype == "fp4" + + def project_query( + self, + qr: torch.Tensor, + position_ids: torch.Tensor, + freqs_cis: torch.Tensor, + rope_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], + *, + fp4_prep: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, + ) -> torch.Tensor: + """Index-space query: ``wq_b`` GEMM + RoPE on the rope slice. + + ``rope_fn(slice, freqs)`` applies RoPE in place (the caller injects the + algorithm's rotary helper). ``fp4_prep`` optionally applies the fp4 + indexer quant/dequant. Returns ``[num_tokens, n_heads, head_dim]``. + """ + num_tokens = qr.shape[0] + q = F.linear(qr, self.indexer.wq_b.weight) + q = q.view(num_tokens, self.n_heads, self.head_dim).unsqueeze(0) + rope_fn(q[..., -self.rope_dim:], freqs_cis[position_ids.long()]) + q = q.squeeze(0) + if fp4_prep is not None: + q = fp4_prep(q) + return q + + def token_weights(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Per-head token weights: ``weights_proj`` GEMM scaled by ``n_heads**-0.5``.""" + weights = F.linear(hidden_states, self.indexer.weights_proj.weight) + return weights.float() * (self.n_heads**-0.5) + + def scores(self, q_row: torch.Tensor, k: torch.Tensor, + weights_row: torch.Tensor) -> torch.Tensor: + """Per-KV index logits for one query token (fp32).""" + head_scores = torch.einsum("hd,kd->hk", q_row.float(), k.float()) + head_scores = F.relu(head_scores) * self.softmax_scale + return (head_scores * weights_row.float().unsqueeze(-1)).sum(dim=0) + + def topk_from_scores(self, scores: torch.Tensor, + topk_tokens: int) -> torch.Tensor: + """Top-k KV positions for one query token, ``-1``-padded to ``topk_tokens``.""" + row = torch.full((topk_tokens, ), + -1, + dtype=torch.int32, + device=scores.device) + if scores.numel() == 0: + return row + k = min(topk_tokens, scores.numel()) + row[:k] = torch.topk(scores.float(), k, dim=-1).indices.to(torch.int32) + return row + + @dataclass class BatchSpec: """Batch specification for testing, following vLLM's pattern.""" From 04732cfdfc05efc83ffbbfcce5948c50197f7a45 Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Thu, 16 Jul 2026 02:02:01 -0700 Subject: [PATCH 10/17] [None][refactor] Vanilla uses the sparse KV-cache manager; drop model_config The sparse KV-cache manager takes sparse_attention_config directly, so model_config/pretrained_config are not needed, and Vanilla can use the same sparse cache manager as the other backends (no VANILLA special-case). Signed-off-by: Yihan Wang --- .../unittest/_torch/attention/backend_case.py | 22 +++---------------- 1 file changed, 3 insertions(+), 19 deletions(-) diff --git a/tests/unittest/_torch/attention/backend_case.py b/tests/unittest/_torch/attention/backend_case.py index 1de44af04944..58ea1f4594b8 100644 --- a/tests/unittest/_torch/attention/backend_case.py +++ b/tests/unittest/_torch/attention/backend_case.py @@ -32,7 +32,6 @@ 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 -from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import str_dtype_to_torch, torch_dtype_to_binding @@ -326,7 +325,6 @@ def _build_mla_kv_cache_manager( case: BackendCase, backend: str, sparse_config=None, - model_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 @@ -359,14 +357,9 @@ def _build_mla_kv_cache_manager( dtype=torch_dtype_to_binding(case.compute_dtype), ) - if sparse_config is not None and backend != "VANILLA": + if sparse_config is not None: cls = get_sparse_attn_kv_cache_manager(sparse_config) - if model_config is None: - raise ValueError("Sparse cache manager requires a model config") - kwargs.update( - sparse_attention_config=sparse_config, - model_config=model_config, - ) + kwargs.update(sparse_attention_config=sparse_config) else: cls = KVCacheManagerV2 if case.use_kv_cache_manager_v2 else KVCacheManager @@ -670,10 +663,6 @@ def _run_sparse_mla_backend( 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) - model_config = ModelConfig( - mapping=mapping, - sparse_attention_config=sparse_config, - ) attn = create_attention( backend, layer_idx=0, @@ -699,12 +688,7 @@ def _run_sparse_mla_backend( # 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, - model_config, - ) + mgr = _build_mla_kv_cache_manager(case, backend, sparse_config) try: mgr.add_dummy_requests(request_ids, case.token_nums) From 1592fb94f5244f9c156e5fb3ea7edb0b587b23e1 Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Fri, 17 Jul 2026 00:17:49 -0700 Subject: [PATCH 11/17] [None][refactor] Vanilla DSA: host RoPE, sm100 gate, token-only selection - Remove VanillaAttention's internal sparse-MLA RoPE. The harness now applies RoPE on the host and feeds pre-formed inputs, using skip_mla_rope_generation for the absorbed generation path (matches production: the MLA module runs RoPE, the attention backend does not). - Context phase feeds raw inputs so the TRTLLM MLA context kernel ropes once (avoids double RoPE); Vanilla and the generation path use the pre-RoPE'd inputs. - Gate TRTLLM DSA to sm>=100: the trtllm-gen DynamicTokenSparse FMHA kernels are Blackwell-only; on Hopper MLA generation falls back to dense FlashMLA, which has no per-token sparse path. - Drop the block-unit (sparse_block_size) abstraction; the suite is token-selection only (DSA / DeepSeek-V4). RocketKV block sparse is out of scope for this suite. Signed-off-by: Yihan Wang --- .../_torch/attention_backend/vanilla.py | 124 +----------------- .../_torch/attention/backend_capability.py | 12 +- .../unittest/_torch/attention/backend_case.py | 99 +++++++------- 3 files changed, 60 insertions(+), 175 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/vanilla.py b/tensorrt_llm/_torch/attention_backend/vanilla.py index ea193348dab8..20399cbee3a9 100644 --- a/tensorrt_llm/_torch/attention_backend/vanilla.py +++ b/tensorrt_llm/_torch/attention_backend/vanilla.py @@ -117,17 +117,6 @@ def __init__( self.qk_nope_head_dim = mla_params.qk_nope_head_dim self.v_head_dim = mla_params.v_head_dim - self.sparse_mla_rope_cos_sin = None - self.sparse_mla_rope_is_neox = True - if (self.is_mla_enable - and getattr(self.sparse_params, "algorithm", None) == "dsa" - and pos_embd_params is not None - and pos_embd_params.rope is not None): - self.sparse_mla_rope_cos_sin = pos_embd_params.rope.create_rope_const_params( - interleave=False)[1].reshape(pos_embd_params.rope.max_positions, - 2, -1) - self.sparse_mla_rope_is_neox = pos_embd_params.is_neox - @classmethod def support_mla(cls) -> bool: return True @@ -152,82 +141,6 @@ def sparse_attn_predict( """Use request-local caller-provided selections in the Vanilla backend.""" return forward_args.topk_indices, None - @staticmethod - def _apply_rotary_embedding(x: torch.Tensor, cos: torch.Tensor, - sin: torch.Tensor, - is_neox: bool) -> torch.Tensor: - """Apply RoPE to ``x`` using one cos/sin row per packed token.""" - cos = cos.to(device=x.device, dtype=x.dtype).unsqueeze(1) - sin = sin.to(device=x.device, dtype=x.dtype).unsqueeze(1) - rotary_dim = cos.shape[-1] * 2 - x_rotary, x_pass = x[..., :rotary_dim], x[..., rotary_dim:] - if is_neox: - x1, x2 = x_rotary.chunk(2, dim=-1) - else: - x1, x2 = x_rotary[..., ::2], x_rotary[..., 1::2] - out1 = x1 * cos - x2 * sin - out2 = x2 * cos + x1 * sin - if is_neox: - rotated = torch.cat((out1, out2), dim=-1) - else: - rotated = torch.stack((out1, out2), dim=-1).flatten(-2) - return torch.cat((rotated, x_pass), dim=-1) - - def _prepare_sparse_mla_inputs( - self, - fused_q: torch.Tensor, - latent_cache: torch.Tensor, - q_pe: Optional[torch.Tensor], - positions: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Apply sparse MLA RoPE to raw packed query and latent-cache inputs. - - As with the other attention paths, omitting positional-embedding - parameters means the caller already applied RoPE. - """ - if self.sparse_mla_rope_cos_sin is None: - return fused_q, latent_cache - - if positions.numel() == 0: - return fused_q, latent_cache - max_position = int(positions.max().item()) - if max_position >= self.sparse_mla_rope_cos_sin.shape[0]: - raise ValueError( - f"Sparse MLA position {max_position} exceeds the configured RoPE table " - f"size {self.sparse_mla_rope_cos_sin.shape[0]}") - - num_tokens = fused_q.shape[0] - fused_head_dim = self.kv_lora_rank + self.qk_rope_head_dim - query = fused_q.view(num_tokens, self.num_heads, fused_head_dim).clone() - if q_pe is None: - raise ValueError( - "Vanilla sparse MLA requires raw q_pe when RoPE parameters are configured" - ) - expected_numel = num_tokens * self.num_heads * self.qk_rope_head_dim - if q_pe.numel() != expected_numel: - raise ValueError( - f"Sparse MLA q_pe has {q_pe.numel()} elements, expected {expected_numel}" - ) - query_rope = q_pe.reshape(num_tokens, self.num_heads, - self.qk_rope_head_dim) - - if latent_cache.shape[1] != fused_head_dim: - raise ValueError( - f"Sparse MLA latent cache width must be {fused_head_dim}, got " - f"{latent_cache.shape[1]}") - latent_cache = latent_cache.clone() - key_rope = latent_cache[:, self.kv_lora_rank:].unsqueeze(1) - cos_sin = self.sparse_mla_rope_cos_sin.index_select( - 0, - positions.to(device=self.sparse_mla_rope_cos_sin.device, - dtype=torch.long)) - cos, sin = cos_sin.unbind(dim=1) - query[..., -self.qk_rope_head_dim:] = self._apply_rotary_embedding( - query_rope, cos, sin, self.sparse_mla_rope_is_neox) - latent_cache[:, self.kv_lora_rank:] = self._apply_rotary_embedding( - key_rope, cos, sin, self.sparse_mla_rope_is_neox).squeeze(1) - return query.view(num_tokens, -1), latent_cache - def _single_request_sparse_attn_predict( self, q: torch.Tensor, k: Optional[torch.Tensor], v: Optional[torch.Tensor], kv_cache_tensor: torch.Tensor, @@ -718,23 +631,19 @@ def _mla_forward_sparse( fused_q: torch.Tensor, metadata: VanillaAttentionMetadata, latent_cache: torch.Tensor, - q_pe: Optional[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 selections, gathers the selected latent K/V, and performs - the absorbed MLA attention directly in PyTorch. Selections are block - indices; this reference implements ``block_size == 1`` (token selection), - which is what the MLA sparse algorithms (DSA / DeepSeek-V4) use. + 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. """ - block_size = getattr(self.sparse_params, "indices_block_size", 1) - if block_size != 1: - raise NotImplementedError( - "Vanilla selected MLA supports block_size 1 (token selection); " - f"got block_size {block_size}") if attention_input_type == AttentionInputType.context_only: seq_start, seq_end = 0, metadata.num_contexts elif attention_input_type == AttentionInputType.generation_only: @@ -808,26 +717,6 @@ def _mla_forward_sparse( raise ValueError("DSA top-k index selects a future token") del valid_mask, kv_lengths, causal_limits - phase_token_start = sum(seq_lens[:seq_start]) - if metadata.position_ids is not None: - positions = metadata.position_ids.reshape( - -1)[phase_token_start:phase_token_start + num_phase_tokens].to( - device=fused_q.device, dtype=torch.long) - if positions.numel() != num_phase_tokens: - raise ValueError( - "DSA metadata does not provide one position ID per phase token" - ) - else: - positions = torch.cat([ - torch.arange(int(past), - int(past) + q_len, - device=fused_q.device, - dtype=torch.long) for past, q_len in zip( - phase_past_tokens, phase_seq_lens, strict=True) - ]) - fused_q, latent_cache = self._prepare_sparse_mla_inputs( - fused_q, latent_cache, q_pe, positions) - from .utils import append_mla_latent_cache kv_cache = append_mla_latent_cache( metadata.kv_cache_manager, @@ -966,7 +855,6 @@ def forward(self, q, metadata, forward_args.latent_cache, - forward_args.q_pe, sparse_attn_indices, forward_args.attention_input_type, ) diff --git a/tests/unittest/_torch/attention/backend_capability.py b/tests/unittest/_torch/attention/backend_capability.py index f34374f98007..6031f5cba72e 100644 --- a/tests/unittest/_torch/attention/backend_capability.py +++ b/tests/unittest/_torch/attention/backend_capability.py @@ -118,12 +118,12 @@ def unsupported_reason(backend: str, case) -> Optional[str]: if sparse_config is not None: algorithm = sparse_config.algorithm if backend == "TRTLLM" and algorithm == "dsa": - # DSA runs on Hopper (FlashMLA) and Blackwell; only the FP8 KV-cache - # DSA path is Blackwell-only. - if sm < 90: - return f"TRTLLM DSA requires sm>=90 (have sm{sm})" - if kv_dtype == "fp8" and sm < 100: - return f"TRTLLM DSA with FP8 KV cache requires sm>=100 (have sm{sm})" + # 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 diff --git a/tests/unittest/_torch/attention/backend_case.py b/tests/unittest/_torch/attention/backend_case.py index 58ea1f4594b8..88a6f4631203 100644 --- a/tests/unittest/_torch/attention/backend_case.py +++ b/tests/unittest/_torch/attention/backend_case.py @@ -155,13 +155,6 @@ def sparse_topk(self) -> Optional[int]: return None return self.sparse_attention_config.sparse_topk - @property - def sparse_block_size(self) -> Optional[int]: - """Selection granularity: token selection is a block of size 1.""" - if self.sparse_attention_config is None: - return None - return self.sparse_attention_config.get_indices_block_size() - @property def prompt_lens(self) -> List[int]: """Original prompt lengths expected by fused generation kernels.""" @@ -243,8 +236,6 @@ 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") - if case.sparse_block_size is None or case.sparse_block_size < 1: - raise ValueError("Sparse backend cases require a positive selection block size") def _randn(gen: torch.Generator, dtype: torch.dtype, *shape) -> torch.Tensor: @@ -405,13 +396,8 @@ def _build_sparse_topk_indices( 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 block indices; with ``block_size == 1`` (DSA) a block is a - single token, so these are token selections. + Selections are per-token (DSA / DeepSeek-V4 select individual tokens). """ - if case.sparse_block_size != 1: - raise ValueError( - f"Sparse MLA selection builder only supports block_size 1, got {case.sparse_block_size}" - ) topk = case.sparse_topk if topk is None or topk <= 0: raise ValueError("Sparse selection cases require a positive top-k") @@ -443,8 +429,8 @@ def _build_sparse_topk_indices( 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 or case.sparse_block_size != 1: - raise ValueError("This generator supports selected sparse MLA with block_size 1 only") + 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 @@ -455,15 +441,29 @@ def generate_sparse_mla_inputs(case: BackendCase, seed: int = 0) -> Dict: 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) - fused_q = torch.cat((q_nope, q_pe), dim=-1).reshape(case.nnz_q, num_heads * d_latent) compressed_kv = _randn(gen, cdt, case.nnz_q, kv_lora_rank) k_pe = _randn(gen, cdt, case.nnz_q, qk_rope_head_dim) - latent_cache = torch.cat((compressed_kv, k_pe), dim=-1) 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, @@ -472,6 +472,10 @@ def generate_sparse_mla_inputs(case: BackendCase, seed: int = 0) -> Dict: 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: @@ -510,8 +514,14 @@ def generate_sparse_mla_inputs(case: BackendCase, seed: int = 0) -> Dict: return dict( fused_q=fused_q, - q_pe=q_pe, + # 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, @@ -728,38 +738,29 @@ def _run_sparse_mla_backend( outputs = [] for attention_input_type, token_slice in phases: - fused_q = inputs["fused_q"][token_slice].clone() - q_pe = inputs["q_pe"][token_slice] - latent_cache = inputs["latent_cache"][token_slice].clone() - forward_kwargs = {} - if backend == "TRTLLM" and attention_input_type == AttentionInputType.generation_only: - cu_q_seqlens = torch.empty(case.num_seqs + 1, dtype=torch.int32, device="cuda") - cu_kv_seqlens = torch.empty(case.num_seqs + 1, dtype=torch.int32, device="cuda") - fmha_scheduler_counter = torch.empty(1, dtype=torch.uint32, device="cuda") - attn.mla_rope_generation( - fused_q, - q_pe, - latent_cache, - metadata, - cu_q_seqlens, - cu_kv_seqlens, - fmha_scheduler_counter, - None, - None, - None, - ) - forward_kwargs.update( - cu_q_seqlens=cu_q_seqlens, - cu_kv_seqlens=cu_kv_seqlens, - fmha_scheduler_counter=fmha_scheduler_counter, - ) - + # 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, - **forward_kwargs, + skip_mla_rope_generation=not kernel_ropes, ) out = attn.forward( fused_q, @@ -769,10 +770,6 @@ def _run_sparse_mla_backend( forward_args=forward_args, ) assert forward_args.sparse_prediction.sparse_attn_indices is not None - assert ( - forward_args.sparse_prediction.sparse_attn_indices_block_size - == sparse_params.indices_block_size - ) outputs.append(out[0] if isinstance(out, tuple) else out) expected_latents = _split_packed_tokens(inputs["expected_new_latent"], case.seq_lens) From cd819c4130010059c67e3b499e0c9da8a56a3c41 Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Tue, 21 Jul 2026 01:47:36 -0700 Subject: [PATCH 12/17] [None][refactor] Address review: drop unused pos_embd_params, revert indexer test - Remove the now-unused pos_embd_params from VanillaAttention.__init__ (and its import); it only fed the sparse-MLA RoPE that was removed. create_attention passes it as a kwarg, which the base backend swallows, so this is a no-op. - Fully revert tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py: VanillaIndexer is not used by the model-driven suite and has a different API than the production DSA Indexer, so restore the original file. Signed-off-by: Yihan Wang --- .../_torch/attention_backend/vanilla.py | 4 +- .../sparse/test_sparse_mla_forward.py | 145 ++++++------------ 2 files changed, 45 insertions(+), 104 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/vanilla.py b/tensorrt_llm/_torch/attention_backend/vanilla.py index 20399cbee3a9..4f0071f1a13c 100644 --- a/tensorrt_llm/_torch/attention_backend/vanilla.py +++ b/tensorrt_llm/_torch/attention_backend/vanilla.py @@ -17,8 +17,7 @@ from .interface import (AttentionBackend, AttentionForwardArgs, AttentionInputType, AttentionMask, AttentionMetadata, - PositionalEmbeddingParams, PredefinedAttentionMask, - merge_attention_forward_args) + PredefinedAttentionMask, merge_attention_forward_args) from .sparse.kernel import triton_index_gather from .sparse.params import SparseParams @@ -96,7 +95,6 @@ def __init__( num_kv_heads: Optional[int] = None, quant_config: Optional[QuantConfig] = None, q_scaling: Optional[float] = None, - pos_embd_params: Optional[PositionalEmbeddingParams] = None, sparse_params: Optional[SparseParams] = None, **kwargs, ): diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py b/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py index 260a6dfe0e8f..72454df6fe86 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py @@ -18,7 +18,7 @@ import os from dataclasses import dataclass from types import SimpleNamespace -from typing import Callable, List, Optional +from typing import List, Optional import pytest import torch @@ -59,82 +59,6 @@ HAS_FLASH_MLA = False -class VanillaIndexer: - """fp32 reference for the production DSA / DeepSeek-V4 sparse ``Indexer``. - - Wraps a production ``Indexer`` instance and reproduces its selection math in - plain fp32 (index-space query projection, per-head token weights, logit - scoring, top-k) so the trtllm indexer's selection can be compared against a - same-weights reference. Algorithm-specific K comes from a compressor / ``wk`` - projection, so the caller supplies the reference K. - - Selection is discrete: top-k over near-tied logits, and an fp32 reference vs - an fp8/fp4 kernel can pick different borderline tokens. Compare the selected - index *set*, never exact attention outputs. - """ - - def __init__(self, indexer): - self.indexer = indexer - self.n_heads = indexer.n_heads - self.head_dim = indexer.head_dim - self.rope_dim = indexer.rope_dim - self.softmax_scale = indexer.softmax_scale - self.indexer_k_dtype = getattr(indexer, "indexer_k_dtype", None) - - @property - def uses_fp4(self) -> bool: - return self.indexer_k_dtype == "fp4" - - def project_query( - self, - qr: torch.Tensor, - position_ids: torch.Tensor, - freqs_cis: torch.Tensor, - rope_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], - *, - fp4_prep: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, - ) -> torch.Tensor: - """Index-space query: ``wq_b`` GEMM + RoPE on the rope slice. - - ``rope_fn(slice, freqs)`` applies RoPE in place (the caller injects the - algorithm's rotary helper). ``fp4_prep`` optionally applies the fp4 - indexer quant/dequant. Returns ``[num_tokens, n_heads, head_dim]``. - """ - num_tokens = qr.shape[0] - q = F.linear(qr, self.indexer.wq_b.weight) - q = q.view(num_tokens, self.n_heads, self.head_dim).unsqueeze(0) - rope_fn(q[..., -self.rope_dim:], freqs_cis[position_ids.long()]) - q = q.squeeze(0) - if fp4_prep is not None: - q = fp4_prep(q) - return q - - def token_weights(self, hidden_states: torch.Tensor) -> torch.Tensor: - """Per-head token weights: ``weights_proj`` GEMM scaled by ``n_heads**-0.5``.""" - weights = F.linear(hidden_states, self.indexer.weights_proj.weight) - return weights.float() * (self.n_heads**-0.5) - - def scores(self, q_row: torch.Tensor, k: torch.Tensor, - weights_row: torch.Tensor) -> torch.Tensor: - """Per-KV index logits for one query token (fp32).""" - head_scores = torch.einsum("hd,kd->hk", q_row.float(), k.float()) - head_scores = F.relu(head_scores) * self.softmax_scale - return (head_scores * weights_row.float().unsqueeze(-1)).sum(dim=0) - - def topk_from_scores(self, scores: torch.Tensor, - topk_tokens: int) -> torch.Tensor: - """Top-k KV positions for one query token, ``-1``-padded to ``topk_tokens``.""" - row = torch.full((topk_tokens, ), - -1, - dtype=torch.int32, - device=scores.device) - if scores.numel() == 0: - return row - k = min(topk_tokens, scores.numel()) - row[:k] = torch.topk(scores.float(), k, dim=-1).indices.to(torch.int32) - return row - - @dataclass class BatchSpec: """Batch specification for testing, following vLLM's pattern.""" @@ -347,6 +271,18 @@ def _copy_ref_compressor_weights(ref_compressor: RefCompressor, ref_compressor.norm.weight.data.copy_(compressor.norm.weight.data) +def _topk_from_scores(scores: torch.Tensor, topk_tokens: int) -> torch.Tensor: + row = torch.full((topk_tokens, ), + -1, + dtype=torch.int32, + device=scores.device) + if scores.numel() == 0: + return row + k = min(topk_tokens, scores.numel()) + row[:k] = torch.topk(scores.float(), k, dim=-1).indices.to(torch.int32) + return row + + def _ceil_pow2_scale(amax: torch.Tensor, max_value_inv: float, min_amax: float) -> torch.Tensor: scaled = torch.clamp(amax.float(), min=min_amax) * max_value_inv @@ -419,6 +355,14 @@ def _prepare_fp4_indexer_k_reference(k: torch.Tensor) -> torch.Tensor: min_amax=fp4_min_amax) +def _reference_indexer_scores(q: torch.Tensor, k: torch.Tensor, + weights: torch.Tensor, + softmax_scale: float) -> torch.Tensor: + head_scores = torch.einsum("hd,kd->hk", q.float(), k.float()) + head_scores = F.relu(head_scores) * softmax_scale + return (head_scores * weights.float().unsqueeze(-1)).sum(dim=0) + + def calculate_reference_deepseek_v4_topk_indices( mla, ref_indexer_compressor: RefCompressor, @@ -434,29 +378,26 @@ def calculate_reference_deepseek_v4_topk_indices( compress_ratio: int, device: torch.device, ) -> torch.Tensor: - """Independent PyTorch reference for HF DS4 ratio-4 indexer top-k indices. - - The index-space query projection, per-head token weights, logit scoring, and - top-k are shared with the DSA path via :class:`VanillaIndexer`; this function - adds the DeepSeek-V4-specific compressor that produces the reference K. - """ + """Independent PyTorch reference for HF DS4 ratio-4 indexer top-k indices.""" indexer = mla.mqa.indexer - vindexer = VanillaIndexer(indexer) num_tokens = hidden_states.shape[0] topk_indices = torch.full((num_tokens, topk_tokens), -1, dtype=torch.int32, device=device) - use_fp4_indexer = vindexer.uses_fp4 - q_ref = vindexer.project_query( - qr, - position_ids, - freqs_cis, - apply_rotary_emb, - fp4_prep=_prepare_fp4_indexer_q_reference if use_fp4_indexer else None, - ) - weights = vindexer.token_weights(hidden_states) + q_ref = F.linear(qr, indexer.wq_b.weight) + q_ref = q_ref.view(num_tokens, indexer.n_heads, indexer.head_dim) + q_ref = q_ref.unsqueeze(0) + apply_rotary_emb(q_ref[..., -indexer.rope_dim:], + freqs_cis[position_ids.long()]) + q_ref = q_ref.squeeze(0) + use_fp4_indexer = indexer.indexer_k_dtype == "fp4" + if use_fp4_indexer: + q_ref = _prepare_fp4_indexer_q_reference(q_ref) + + weights = F.linear(hidden_states, indexer.weights_proj.weight) + weights = weights.float() * (indexer.n_heads**-0.5) offset = 0 for req_idx in ctx_indices: @@ -480,10 +421,11 @@ def calculate_reference_deepseek_v4_topk_indices( for token_idx in range(seq_len): valid_len = (token_idx + 1) // compress_ratio k_for_scores = compressed_kv_for_scores[:valid_len] - scores = vindexer.scores(q_ref[offset + token_idx], - k_for_scores, - weights[offset + token_idx]) - topk_indices[offset + token_idx] = vindexer.topk_from_scores( + scores = _reference_indexer_scores(q_ref[offset + token_idx], + k_for_scores, + weights[offset + token_idx], + indexer.softmax_scale) + topk_indices[offset + token_idx] = _topk_from_scores( scores, topk_tokens) offset += seq_len @@ -525,10 +467,11 @@ def calculate_reference_deepseek_v4_topk_indices( compressed_kv_for_scores = ( _prepare_fp4_indexer_k_reference(compressed_kv) if use_fp4_indexer else compressed_kv) - scores = vindexer.scores(q_ref[token_idx], - compressed_kv_for_scores[:valid_len], - weights[token_idx]) - topk_indices[token_idx] = vindexer.topk_from_scores(scores, topk_tokens) + scores = _reference_indexer_scores(q_ref[token_idx], + compressed_kv_for_scores[:valid_len], + weights[token_idx], + indexer.softmax_scale) + topk_indices[token_idx] = _topk_from_scores(scores, topk_tokens) return topk_indices From 034f841f774d7657e229aa1dc99d1b984aa2b7f6 Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Mon, 3 Aug 2026 01:14:17 -0700 Subject: [PATCH 13/17] [None][chore] Address the review comments in https://github.com/NVIDIA/TensorRT-LLM/pull/16714 Signed-off-by: Yihan Wang --- .../_torch/attention_backend/vanilla.py | 20 ++++++++++++------- .../_torch/pyexecutor/py_executor_creator.py | 8 ++------ 2 files changed, 15 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/vanilla.py b/tensorrt_llm/_torch/attention_backend/vanilla.py index 5d13b9d77aea..4151528dcf0e 100644 --- a/tensorrt_llm/_torch/attention_backend/vanilla.py +++ b/tensorrt_llm/_torch/attention_backend/vanilla.py @@ -237,13 +237,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: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 9f6c0e5cbef2..faf3ae08b68f 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"): From 3e35f63a0a711d0444b00dcbf8632c06a57ffe4a Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Mon, 3 Aug 2026 06:03:44 -0700 Subject: [PATCH 14/17] Revert LLM API changes Signed-off-by: Yihan Wang --- tensorrt_llm/llmapi/llm_args.py | 14 -------------- tests/unittest/_torch/attention/backend_case.py | 7 ++++--- .../unittest/_torch/attention/model_attn_config.py | 12 ++++++++++-- .../_torch/attention/test_attention_backends.py | 2 +- 4 files changed, 15 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index b9135dc58843..d325b69f65d0 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -632,20 +632,6 @@ class BaseSparseAttentionConfig(StrictBaseModel): """Configuration for sparse attention.""" algorithm: str - @property - def sparse_topk(self) -> int: - """Per-token selection budget (top-k). - - The field name is not yet unified across algorithms (DSA / DeepSeek-V4 - use ``index_topk``, RocketKV uses ``topk``), so resolve it here. - """ - for name in ("index_topk", "topk"): - value = getattr(self, name, None) - if value is not None: - return value - raise NotImplementedError( - f"{type(self).__name__} does not expose a selection top-k.") - def supports_backend(self, backend: str) -> bool: """Override if the sparse attention algorithm does not support a subset of the possible backends. diff --git a/tests/unittest/_torch/attention/backend_case.py b/tests/unittest/_torch/attention/backend_case.py index 88a6f4631203..8834e923746b 100644 --- a/tests/unittest/_torch/attention/backend_case.py +++ b/tests/unittest/_torch/attention/backend_case.py @@ -150,10 +150,11 @@ def is_sparse(self) -> bool: @property def sparse_topk(self) -> Optional[int]: - """Per-token selection budget, resolved from the sparse config.""" - if self.sparse_attention_config is None: + """Per-token selection budget (``index_topk``) from the sparse config.""" + cfg = self.sparse_attention_config + if cfg is None: return None - return self.sparse_attention_config.sparse_topk + return cfg.index_topk @property def prompt_lens(self) -> List[int]: diff --git a/tests/unittest/_torch/attention/model_attn_config.py b/tests/unittest/_torch/attention/model_attn_config.py index 543d8a8b8e52..41d656eb3bcd 100644 --- a/tests/unittest/_torch/attention/model_attn_config.py +++ b/tests/unittest/_torch/attention/model_attn_config.py @@ -57,10 +57,10 @@ from dataclasses import dataclass from typing import Literal, Optional -AttentionPhase = Literal["ctx", "gen"] - from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig, SparseAttentionConfig +AttentionPhase = Literal["ctx", "gen"] + @dataclass(frozen=True) class ModelAttnConfig: @@ -90,6 +90,14 @@ class ModelAttnConfig: # 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 + # --------------------------------------------------------------------------- # Standard self-attention (decoder LLMs). The dominant configurations. diff --git a/tests/unittest/_torch/attention/test_attention_backends.py b/tests/unittest/_torch/attention/test_attention_backends.py index f1433b31e2d2..a752276a39d0 100644 --- a/tests/unittest/_torch/attention/test_attention_backends.py +++ b/tests/unittest/_torch/attention/test_attention_backends.py @@ -90,7 +90,7 @@ def _phases_from_window(window: int) -> dict: def _phases_for(cfg: ModelAttnConfig) -> dict: if cfg.sparse_attention_config is not None: - return _phases(cfg.sparse_attention_config.sparse_topk + 32, gen_len=1) + return _phases(cfg.sparse_topk + 32, gen_len=1) if cfg.mask != "sliding": return _PHASES From cc2acdc390816f9acea56291b42989ea79816be8 Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Mon, 10 Aug 2026 04:27:13 -0700 Subject: [PATCH 15/17] [None][test] Fix vanilla DSA sparse MLA golden after main merge The origin merge refactored AttentionForwardArgs and replaced the DSA backend, breaking the vanilla-vs-TRTLLM sparse MLA comparison: - Route caller-provided top-k through sparse_backend_args.topk_indices (VanillaAttention.sparse_attn_predict) and rename the dropped forward_args.sparse_prediction to sparse_runtime_params in vanilla.py and backend_case.py. - Restore the is_vanilla branching in the test (direct selection for vanilla, indexer mock + sparse_backend_args for TRTLLM) and drop the duplicated generation forward kwargs left by the merge. - VanillaAttention applies no RoPE itself while the kernel does via rope_append, so pre-apply the matching yarn RoPE to the vanilla q_pe/k_pe at absolute positions before the golden forward. All 12 parametrized cases pass. Signed-off-by: Yihan Wang --- .../_torch/attention_backend/vanilla.py | 11 +- .../unittest/_torch/attention/backend_case.py | 2 +- .../sparse/dsa/test_dsa_sparse_mla.py | 181 +++++++++++++++--- 3 files changed, 161 insertions(+), 33 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/vanilla.py b/tensorrt_llm/_torch/attention_backend/vanilla.py index e072fd67b7f7..17decfb40b51 100644 --- a/tensorrt_llm/_torch/attention_backend/vanilla.py +++ b/tensorrt_llm/_torch/attention_backend/vanilla.py @@ -137,7 +137,10 @@ def sparse_attn_predict( forward_args: AttentionForwardArgs, ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: """Use request-local caller-provided selections in the Vanilla backend.""" - return forward_args.topk_indices, None + 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], @@ -914,8 +917,8 @@ def forward(self, q, k, metadata, forward_args) sparse_attn_indices, sparse_attn_offsets = self.sparse_attn_predict( q, k, metadata, forward_args) - forward_args.sparse_prediction = replace( - forward_args.sparse_prediction, + forward_args.sparse_runtime_params = replace( + forward_args.sparse_runtime_params, sparse_kv_indices=sparse_kv_indices, sparse_kv_offsets=sparse_kv_offsets, sparse_attn_indices=sparse_attn_indices, @@ -924,7 +927,7 @@ def forward(self, self.sparse_params, "indices_block_size"), ) sparse_attn_indices = ( - forward_args.sparse_prediction.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( diff --git a/tests/unittest/_torch/attention/backend_case.py b/tests/unittest/_torch/attention/backend_case.py index 8834e923746b..e097f5d41ff3 100644 --- a/tests/unittest/_torch/attention/backend_case.py +++ b/tests/unittest/_torch/attention/backend_case.py @@ -770,7 +770,7 @@ def _run_sparse_mla_backend( metadata, forward_args=forward_args, ) - assert forward_args.sparse_prediction.sparse_attn_indices is not None + 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) 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 81278f24b302..b8323c42b917 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 @@ -41,7 +41,7 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str from tensorrt_llm.bindings.executor import KvCacheConfig -from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.functional import PositionEmbeddingType, RopeEmbeddingUtils from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig @@ -96,6 +96,88 @@ class RopeConfig: model_type: str = "deepseek_v3" +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_q_pe_for_vanilla( + fused_q: torch.Tensor, + rope_cos_sin: torch.Tensor, + seq_lens: List[int], + past_lens: List[int], + num_heads: int, + kv_lora_rank: int, + qk_rope_head_dim: int, +) -> torch.Tensor: + """Apply the kernel's yarn RoPE to the q_pe portion of ``fused_q`` per request. + + ``VanillaAttention`` expects RoPE-pre-applied inputs (it performs no positional + embedding itself), whereas the TRTLLM kernel applies RoPE internally via + ``rope_append``. This rotates each request's query rope dims at absolute positions + ``[past_len, past_len + seq_len)`` so the Vanilla golden matches the kernel. + """ + fused_q = fused_q.clone() + fused_head_dim = kv_lora_rank + qk_rope_head_dim + off = 0 + for seq_len, past in zip(seq_lens, past_lens): + seg = fused_q[off : off + seq_len].view(seq_len, num_heads, fused_head_dim) + q_rope = seg[..., -qk_rope_head_dim:] + cos, sin = rope_cos_sin[past : past + 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) + seg[..., -qk_rope_head_dim:] = q_rope + fused_q[off : off + seq_len] = seg.view(seq_len, -1) + off += seq_len + return fused_q + + +def _rotate_k_pe_for_vanilla( + k_pe: torch.Tensor, + rope_cos_sin: torch.Tensor, + seq_lens: List[int], + past_lens: List[int], +) -> torch.Tensor: + """Apply the kernel's yarn RoPE to ``k_pe`` per request (positions ``[past, past+seq)``).""" + out = [] + off = 0 + for seq_len, past in zip(seq_lens, past_lens): + seg = k_pe[off : off + seq_len].unsqueeze(-2) + cos, sin = rope_cos_sin[past : past + seq_len].chunk(2, dim=-2) + seg = seg.unflatten(-1, [-1, 2]).transpose(-2, -1).flatten(start_dim=-2) + seg = ((seg * cos) + (rotate_half(seg) * sin)).to(dtype=seg.dtype) + seg = seg.unflatten(-1, [2, -1]).transpose(-2, -1).flatten(start_dim=-2) + out.append(seg) + off += seq_len + return torch.cat(out).squeeze(-2) + + +def _yarn_rope_cos_sin(rope_config, device: torch.device) -> torch.Tensor: + """Precompute the yarn RoPE cos/sin table used by the TRTLLM kernel.""" + 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) + ) + + def _build_sparse_topk_indices_context( seq_lens: List[int], topk: int, @@ -346,6 +428,9 @@ def _run_test_for_backend( # 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 + # Vanilla applies no positional embedding itself (the kernel does via rope_append), so the + # golden must receive RoPE-pre-applied q_pe/k_pe. Precompute the matching yarn cos/sin table. + rope_cos_sin = _yarn_rope_cos_sin(rope_config, device) if is_vanilla else None # Set seed for reproducibility. torch.manual_seed(seed) @@ -635,19 +720,49 @@ def create_layer(layer_idx: int, num_kv_heads: int): device, generator=topk_generator, ) - ctx_layers[layer_idx].indexer.forward_from_projected = Mock( - return_value=topk_indices - ) - result = ctx_layers[layer_idx].forward( - fused_q.clone(), - None, - None, - attn_metadata, - attention_input_type=AttentionInputType.context_only, - latent_cache=latent_cache.clone(), - q_pe=q_pe, - sparse_backend_args=DSABackendForwardArgs(indexer_intermediates=[]), - ) + if is_vanilla: + # VanillaAttention has no indexer and applies no RoPE itself; feed it the + # sparse selection plus RoPE-pre-applied q_pe/k_pe (context positions 0..len). + ctx_past = [0] * len(context_sequence_lengths) + fused_q_v = _rotate_q_pe_for_vanilla( + fused_q, + rope_cos_sin, + context_sequence_lengths, + ctx_past, + num_heads, + kv_lora_rank, + qk_rope_head_dim, + ) + k_pe_v = _rotate_k_pe_for_vanilla( + k_pe, rope_cos_sin, context_sequence_lengths, ctx_past + ) + latent_cache_v = torch.cat([compressed_kv, k_pe_v], dim=-1) + result = ctx_layers[layer_idx].forward( + fused_q_v.clone(), + None, + None, + attn_metadata, + attention_input_type=AttentionInputType.context_only, + latent_cache=latent_cache_v.clone(), + q_pe=q_pe, + sparse_backend_args=DSABackendForwardArgs( + indexer_intermediates=[], topk_indices=topk_indices + ), + ) + else: + ctx_layers[layer_idx].indexer.forward_from_projected = Mock( + return_value=topk_indices + ) + result = ctx_layers[layer_idx].forward( + fused_q.clone(), + None, + None, + attn_metadata, + attention_input_type=AttentionInputType.context_only, + latent_cache=latent_cache.clone(), + q_pe=q_pe, + sparse_backend_args=DSABackendForwardArgs(indexer_intermediates=[]), + ) 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] @@ -665,13 +780,27 @@ def create_layer(layer_idx: int, num_kv_heads: int): device, generator=topk_generator, ) - gen_layers[layer_idx].indexer.forward_from_projected = Mock( - return_value=topk_indices - ) if is_vanilla: - backend_fused_q = fused_q.clone() - backend_latent_cache = latent_cache.clone() - generation_kwargs = {} + # VanillaAttention has no indexer and applies no RoPE itself; RoPE-pre-apply + # q_pe/k_pe at this step's absolute positions [cached_len, cached_len+seq_q). + seq_q = [generation_seq_len_q] * len(context_sequence_lengths) + fused_q_v = _rotate_q_pe_for_vanilla( + fused_q, + rope_cos_sin, + seq_q, + cached_lens, + num_heads, + kv_lora_rank, + qk_rope_head_dim, + ) + k_pe_v = _rotate_k_pe_for_vanilla(k_pe, rope_cos_sin, seq_q, cached_lens) + backend_fused_q = fused_q_v.clone() + backend_latent_cache = torch.cat([compressed_kv, k_pe_v], dim=-1).clone() + 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) @@ -720,6 +849,9 @@ def create_layer(layer_idx: int, num_kv_heads: int): 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 = { @@ -729,6 +861,7 @@ def create_layer(layer_idx: int, num_kv_heads: int): "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( backend_fused_q, @@ -738,15 +871,7 @@ def create_layer(layer_idx: int, num_kv_heads: int): attention_input_type=AttentionInputType.generation_only, latent_cache=backend_latent_cache, q_pe=q_pe, - topk_indices=topk_indices, **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=[]), ) # Record results for the Vanilla-golden comparison. print( From 434b838e0ab28cd9a8defbac6a32eb2792c4c2f1 Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Mon, 10 Aug 2026 04:35:24 -0700 Subject: [PATCH 16/17] [None][chore] Remove unused VanillaAttention.sparse_kv_predict VanillaAttention.sparse_kv_predict only ever returned (None, None) and its sole caller wrote those Nones back as the already-default sparse_runtime_params values. Vanilla handles sparse selection inline via sparse_attn_predict (not through the sparse hooks), so the method is dead. Drop it and the corresponding kv-index assignments. Signed-off-by: Yihan Wang --- tensorrt_llm/_torch/attention_backend/vanilla.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/vanilla.py b/tensorrt_llm/_torch/attention_backend/vanilla.py index 17decfb40b51..05e64f1f8a7f 100644 --- a/tensorrt_llm/_torch/attention_backend/vanilla.py +++ b/tensorrt_llm/_torch/attention_backend/vanilla.py @@ -119,16 +119,6 @@ def __init__( def support_mla(cls) -> bool: return True - def sparse_kv_predict( - self, - q: torch.Tensor, - k: Optional[torch.Tensor], - metadata: VanillaAttentionMetadata, - forward_args: AttentionForwardArgs, - ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: - """Lower backend-neutral sparse KV selection for Vanilla attention.""" - return None, None - def sparse_attn_predict( self, q: torch.Tensor, @@ -913,14 +903,10 @@ def forward(self, if sparse_algorithm != "dsa": raise ValueError( "Vanilla selected MLA currently supports only DSA") - sparse_kv_indices, sparse_kv_offsets = self.sparse_kv_predict( - q, k, metadata, forward_args) 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_kv_indices=sparse_kv_indices, - sparse_kv_offsets=sparse_kv_offsets, sparse_attn_indices=sparse_attn_indices, sparse_attn_offsets=sparse_attn_offsets, sparse_attn_indices_block_size=getattr( From 4f7ba697f29e72e760cd16891a57f8d59dd86bb0 Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Mon, 10 Aug 2026 04:45:54 -0700 Subject: [PATCH 17/17] [None][test] Simplify vanilla DSA RoPE helpers and dedup forward call Merge the q_pe/k_pe rotate helpers into one _rope_qk_for_vanilla over a shared _apply_yarn_rope core, unify the context vanilla/non-vanilla forward() into a single call, and trim redundant comments. Signed-off-by: Yihan Wang --- .../sparse/dsa/test_dsa_sparse_mla.py | 125 +++++++----------- 1 file changed, 46 insertions(+), 79 deletions(-) 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 b8323c42b917..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 @@ -97,66 +97,51 @@ class RopeConfig: 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_q_pe_for_vanilla( +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, - kv_lora_rank: int, qk_rope_head_dim: int, -) -> torch.Tensor: - """Apply the kernel's yarn RoPE to the q_pe portion of ``fused_q`` per request. +) -> tuple[torch.Tensor, torch.Tensor]: + """RoPE-pre-apply the yarn positional embedding to the Vanilla golden's q_pe/k_pe. - ``VanillaAttention`` expects RoPE-pre-applied inputs (it performs no positional - embedding itself), whereas the TRTLLM kernel applies RoPE internally via - ``rope_append``. This rotates each request's query rope dims at absolute positions - ``[past_len, past_len + seq_len)`` so the Vanilla golden matches the kernel. + 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. """ fused_q = fused_q.clone() - fused_head_dim = kv_lora_rank + qk_rope_head_dim - off = 0 - for seq_len, past in zip(seq_lens, past_lens): - seg = fused_q[off : off + seq_len].view(seq_len, num_heads, fused_head_dim) - q_rope = seg[..., -qk_rope_head_dim:] - cos, sin = rope_cos_sin[past : past + 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) - seg[..., -qk_rope_head_dim:] = q_rope - fused_q[off : off + seq_len] = seg.view(seq_len, -1) - off += seq_len - return fused_q - - -def _rotate_k_pe_for_vanilla( - k_pe: torch.Tensor, - rope_cos_sin: torch.Tensor, - seq_lens: List[int], - past_lens: List[int], -) -> torch.Tensor: - """Apply the kernel's yarn RoPE to ``k_pe`` per request (positions ``[past, past+seq)``).""" - out = [] + latent_cache = latent_cache.clone() + fused_head_dim = latent_cache.shape[-1] off = 0 for seq_len, past in zip(seq_lens, past_lens): - seg = k_pe[off : off + seq_len].unsqueeze(-2) cos, sin = rope_cos_sin[past : past + seq_len].chunk(2, dim=-2) - seg = seg.unflatten(-1, [-1, 2]).transpose(-2, -1).flatten(start_dim=-2) - seg = ((seg * cos) + (rotate_half(seg) * sin)).to(dtype=seg.dtype) - seg = seg.unflatten(-1, [2, -1]).transpose(-2, -1).flatten(start_dim=-2) - out.append(seg) + 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 torch.cat(out).squeeze(-2) + return fused_q, latent_cache def _yarn_rope_cos_sin(rope_config, device: torch.device) -> torch.Tensor: - """Precompute the yarn RoPE cos/sin table used by the TRTLLM kernel.""" return ( torch.tensor( RopeEmbeddingUtils.create_sinusoidal_positions_yarn( @@ -428,8 +413,6 @@ def _run_test_for_backend( # 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 - # Vanilla applies no positional embedding itself (the kernel does via rope_append), so the - # golden must receive RoPE-pre-applied q_pe/k_pe. Precompute the matching yarn cos/sin table. rope_cos_sin = _yarn_rope_cos_sin(rope_config, device) if is_vanilla else None # Set seed for reproducibility. @@ -721,48 +704,36 @@ def create_layer(layer_idx: int, num_kv_heads: int): generator=topk_generator, ) if is_vanilla: - # VanillaAttention has no indexer and applies no RoPE itself; feed it the - # sparse selection plus RoPE-pre-applied q_pe/k_pe (context positions 0..len). - ctx_past = [0] * len(context_sequence_lengths) - fused_q_v = _rotate_q_pe_for_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, - ctx_past, + [0] * len(context_sequence_lengths), num_heads, - kv_lora_rank, qk_rope_head_dim, ) - k_pe_v = _rotate_k_pe_for_vanilla( - k_pe, rope_cos_sin, context_sequence_lengths, ctx_past - ) - latent_cache_v = torch.cat([compressed_kv, k_pe_v], dim=-1) - result = ctx_layers[layer_idx].forward( - fused_q_v.clone(), - None, - None, - attn_metadata, - attention_input_type=AttentionInputType.context_only, - latent_cache=latent_cache_v.clone(), - q_pe=q_pe, - sparse_backend_args=DSABackendForwardArgs( - indexer_intermediates=[], topk_indices=topk_indices - ), + ctx_sba = DSABackendForwardArgs( + indexer_intermediates=[], topk_indices=topk_indices ) else: ctx_layers[layer_idx].indexer.forward_from_projected = Mock( return_value=topk_indices ) - result = ctx_layers[layer_idx].forward( - fused_q.clone(), - None, - None, - attn_metadata, - attention_input_type=AttentionInputType.context_only, - latent_cache=latent_cache.clone(), - q_pe=q_pe, - sparse_backend_args=DSABackendForwardArgs(indexer_intermediates=[]), - ) + ctx_fused_q, ctx_latent = fused_q, latent_cache + ctx_sba = DSABackendForwardArgs(indexer_intermediates=[]) + result = ctx_layers[layer_idx].forward( + ctx_fused_q.clone(), + None, + None, + attn_metadata, + attention_input_type=AttentionInputType.context_only, + latent_cache=ctx_latent.clone(), + q_pe=q_pe, + sparse_backend_args=ctx_sba, + ) 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] @@ -781,21 +752,17 @@ def create_layer(layer_idx: int, num_kv_heads: int): generator=topk_generator, ) if is_vanilla: - # VanillaAttention has no indexer and applies no RoPE itself; RoPE-pre-apply - # q_pe/k_pe at this step's absolute positions [cached_len, cached_len+seq_q). + # Pre-apply RoPE at absolute positions [cached_len, cached_len+seq_q). seq_q = [generation_seq_len_q] * len(context_sequence_lengths) - fused_q_v = _rotate_q_pe_for_vanilla( + backend_fused_q, backend_latent_cache = _rope_qk_for_vanilla( fused_q, + latent_cache, rope_cos_sin, seq_q, cached_lens, num_heads, - kv_lora_rank, qk_rope_head_dim, ) - k_pe_v = _rotate_k_pe_for_vanilla(k_pe, rope_cos_sin, seq_q, cached_lens) - backend_fused_q = fused_q_v.clone() - backend_latent_cache = torch.cat([compressed_kv, k_pe_v], dim=-1).clone() generation_kwargs = { "sparse_backend_args": DSABackendForwardArgs( indexer_intermediates=[], topk_indices=topk_indices