[Perf][DSv4] Fuse inverse RoPE + FP8 quant into FlashInfer sparse MLA (#58621)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

tests/kernels/attention/test_flashmla_sparse.py

@@ -511,6 +511,86 @@ def test_flashinfer_mixed_sparse_indices_separates_window_and_padded_width():
     assert sparse_lens.cpu().tolist() == [padded_width]
 
 
+def test_flashinfer_mixed_sparse_indices_noncausal_rows_have_no_active_gaps():
+    """A DSpark non-causal row keeps -1 out of the kernel's active ranges.
+
+    Each draft token's causal window holds its first visible entries, the rest
+    of the block starts at column `window`, and the compressed entries follow
+    directly; `sparse_topk_lens` ends right after them.
+    """
+    from vllm.models.deepseek_v4.common.ops.cache_utils import (
+        build_flashinfer_mixed_sparse_indices,
+    )
+
+    device = torch.device("cuda")
+    window, width = 4, 8
+    # One request: 1 context token, a 2-token draft block, slots 0..2 visible.
+    decode_swa = torch.full((2, width), -1, dtype=torch.int32, device=device)
+    decode_swa[:, :3] = torch.arange(3, dtype=torch.int32, device=device)
+    compressed = torch.tensor(
+        [[100, 101, -1, -1]] * 2, dtype=torch.int32, device=device
+    )
+    sparse_indices, sparse_lens = build_flashinfer_mixed_sparse_indices(
+        decode_swa_indices=decode_swa,
+        decode_compressed_indices=compressed,
+        decode_compressed_topk_lens=torch.tensor(
+            [2, 2], dtype=torch.int32, device=device
+        ),
+        prefill_topk_indices=torch.empty((0, 4), dtype=torch.int32, device=device),
+        query_start_loc=torch.tensor([0, 2], dtype=torch.int32, device=device),
+        seq_lens=torch.tensor([3], dtype=torch.int32, device=device),
+        token_to_req_indices=torch.tensor([0, 0], dtype=torch.int32, device=device),
+        swa_block_table=torch.tensor([[0]], dtype=torch.int32, device=device),
+        swa_block_size=64,
+        compressed_block_table=torch.tensor([[0]], dtype=torch.int32, device=device),
+        compressed_block_size=64,
+        window_size=window,
+        compress_ratio=128,
+        topk=0,
+    )
+
+    assert sparse_indices.cpu().tolist() == [
+        # pos 1: window [0, 1]; slot 2 spills to column 4; compressed after it.
+        [0, 1, -1, -1, 2, 100, 101, -1, -1, -1, -1, -1],
+        # pos 2: window [0, 1, 2]; nothing spills; compressed at column 4.
+        [0, 1, 2, -1, 100, 101, -1, -1, -1, -1, -1, -1],
+    ]
+    assert sparse_lens.cpu().tolist() == [7, 6]
+
+
+def test_flashinfer_mixed_sparse_indices_noncausal_pad_rows_keep_min_length():
+    """A CUDA-graph pad token of the DSpark graph (a zero-length request) gets
+    a causal pad row's shape instead of negative spill arithmetic."""
+    from vllm.models.deepseek_v4.common.ops.cache_utils import (
+        build_flashinfer_mixed_sparse_indices,
+    )
+
+    device = torch.device("cuda")
+    window, width = 128, 192
+    decode_swa = torch.full((3, width), -1, dtype=torch.int32, device=device)
+    decode_swa[:2, :3] = torch.arange(3, dtype=torch.int32, device=device)
+    sparse_indices, sparse_lens = build_flashinfer_mixed_sparse_indices(
+        decode_swa_indices=decode_swa,
+        decode_compressed_indices=None,
+        decode_compressed_topk_lens=None,
+        prefill_topk_indices=torch.empty((0, 0), dtype=torch.int32, device=device),
+        # Request 1 is the padding: no query tokens, seq_len 0.
+        query_start_loc=torch.tensor([0, 2, 2], dtype=torch.int32, device=device),
+        seq_lens=torch.tensor([3, 0], dtype=torch.int32, device=device),
+        token_to_req_indices=torch.tensor([0, 0, 1], dtype=torch.int32, device=device),
+        swa_block_table=torch.zeros((2, 1), dtype=torch.int32, device=device),
+        swa_block_size=64,
+        compressed_block_table=None,
+        compressed_block_size=64,
+        window_size=window,
+        compress_ratio=1,
+        topk=0,
+    )
+
+    assert sparse_lens.cpu().tolist() == [129, 128, 128]
+    assert (sparse_indices[2] == -1).all()
+
+
 @pytest.mark.parametrize("model", ["deepseek_v4", "deepseek_v41"])
 @pytest.mark.parametrize("page_padding", [0, 64])
 def test_flashinfer_sparse_forward_reads_packed_kv_rows(model, page_padding):
@@ -565,3 +645,198 @@ def test_flashinfer_sparse_forward_reads_packed_kv_rows(model, page_padding):
     attn._forward(query, None, cache, metadata, None, True, output)
     expected = torch.tensor([2, 3, 3], dtype=output.dtype, device=device)
     torch.testing.assert_close(output, expected[:, None, None].expand_as(output))
+
+
+def _make_rope_quant_attn(mod, num_rows: int, device: str):
+    """A bare DSv4 FlashInfer layer set up for RopeQuant: 128 heads, fp8 KV."""
+    from types import SimpleNamespace
+
+    heads, groups, rank = 128, 16, 256
+    angles = torch.arange(2048, device=device, dtype=torch.float32)[:, None] * 0.05
+    angles = angles + torch.arange(32, device=device)[None] * 0.1
+    attn = object.__new__(mod.DeepseekV4FlashInferMLAAttention)
+    attn.compress_ratio = 1
+    attn.window_size = 128
+    attn.scale = 512**-0.5
+    attn.kv_cache_torch_dtype = torch.float8_e4m3fn
+    attn._flashinfer_fp8_bmm1_scale = attn.scale
+    attn._flashinfer_fp8_bmm2_scale = 1.0
+    attn.attn_sink = torch.linspace(-1, 1, heads, device=device)
+    attn.padded_heads = heads
+    attn.topk_indices_buffer = torch.empty(
+        (num_rows, 0), dtype=torch.int32, device=device
+    )
+    attn.rotary_emb = SimpleNamespace(
+        cos_sin_cache=torch.cat((angles.cos(), angles.sin()), -1).contiguous()
+    )
+    attn.n_local_heads, attn.n_local_groups, attn.o_lora_rank = heads, groups, rank
+    attn.nope_head_dim, attn.rope_head_dim = 448, 64
+    attn._einsum_recipe, attn._tma_aligned_scales = (1, 1, 128), True
+    attn.wo_a = SimpleNamespace(
+        weight=(torch.randn((groups, rank, 4096), device=device) * 0.05).to(
+            torch.float8_e4m3fn
+        ),
+        weight_scale=torch.ones((groups, rank, 32), device=device),
+    )
+    attn.wo_b = lambda z: z
+    return attn
+
+
+def _rope_quant_vs_unfused(attn, q, cache, metadata, positions, num_tokens):
+    """Run `_forward` unfused and fused; return (bf16 attention, z_unfused, z_fused)."""
+    from vllm.models.deepseek_v4.nvidia.ops.o_proj import rope_quant_attn_out
+
+    heads, groups = attn.n_local_heads, attn.n_local_groups
+    unfused = torch.empty(q.shape, dtype=torch.bfloat16, device=q.device)
+    attn._forward(q, None, cache, metadata, None, True, unfused)
+    fused = rope_quant_attn_out(q.shape[0], q.device)
+    metadata.flashinfer_sparse_index_cache.clear()
+    attn._forward(q, None, cache, metadata, None, True, fused)
+    z_unfused = attn._o_proj(unfused[:num_tokens, :heads], positions)
+    z_fused = attn._o_proj(fused, positions)[:num_tokens]
+    return unfused[:num_tokens], z_unfused, z_fused.view(num_tokens, groups, -1)
+
+
+def _assert_projects_alike(z_unfused, z_fused):
+    z_unfused = z_unfused.float().view_as(z_fused)
+    rel = (z_fused.float() - z_unfused).norm() / z_unfused.norm()
+    assert rel < 2e-2, rel
+
+
+def _skip_unless_rope_quant():
+    from vllm.platforms import current_platform
+
+    if not current_platform.is_device_capability_family(100):
+        pytest.skip("Requires FlashInfer TRTLLM sparse MLA on SM100/SM103")
+
+
+def test_flashinfer_rope_quant_matches_unfused_o_proj():
+    """RopeQuant's fused output must project like the bf16 + inv-RoPE/quant path.
+
+    One mixed step (two decodes, a three-token prefill) plus CUDA-graph pad rows,
+    so both calls run over the full buffer with un-rebased cum_seq_lens_q.
+    """
+    from vllm.models.deepseek_v4.nvidia import flashinfer_sparse as mod
+    from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
+
+    _skip_unless_rope_quant()
+    torch.manual_seed(0)
+    device = "cuda"
+    block_size = 64
+    q_lens, seq_lens = [1, 1, 3], [5, 9, 70]
+    num_tokens, num_rows = sum(q_lens), sum(q_lens) + 3
+    cache = torch.randn((3, block_size, 512), device=device).to(torch.float8_e4m3fn)
+    attn = _make_rope_quant_attn(mod, num_rows, device)
+
+    decode_indices = torch.full((2, 128), -1, dtype=torch.int32, device=device)
+    for i, seq_len in enumerate(seq_lens[:2]):
+        decode_indices[i, :seq_len] = torch.arange(
+            i * block_size, i * block_size + seq_len, device=device
+        )
+    query_start = torch.tensor([0, 1, 2, 5], dtype=torch.int32)
+    metadata = DeepseekSparseSWAMetadata(
+        block_table=torch.tensor([[0], [1], [2]], dtype=torch.int32, device=device),
+        slot_mapping=torch.zeros(num_tokens, dtype=torch.int64, device=device),
+        block_size=block_size,
+        seq_lens=torch.tensor(seq_lens, dtype=torch.int32, device=device),
+        query_start_loc=query_start.to(device),
+        query_start_loc_cpu=query_start,
+        token_to_req_indices=torch.tensor(
+            [0, 1, 2, 2, 2], dtype=torch.int32, device=device
+        ),
+        decode_swa_indices=decode_indices,
+        decode_swa_lens=torch.tensor(seq_lens[:2], dtype=torch.int32, device=device),
+        decode_swa_width=128,
+        is_valid_token=torch.ones(num_tokens, dtype=torch.bool, device=device),
+        num_decodes=2,
+        num_prefills=1,
+        num_decode_tokens=2,
+        num_prefill_tokens=3,
+        max_decode_query_len=1,
+    )
+    positions = torch.tensor([4, 8, 67, 68, 69], device=device)
+    q = torch.randn((num_rows, 128, 512), device=device).to(torch.float8_e4m3fn)
+    _, z_unfused, z_fused = _rope_quant_vs_unfused(
+        attn, q, cache, metadata, positions, num_tokens
+    )
+    _assert_projects_alike(z_unfused, z_fused)
+
+
+@pytest.mark.parametrize("context_len", [0, 20, 127, 900])
+def test_flashinfer_dspark_noncausal_block_sees_future_tokens(context_len):
+    """Every DSpark draft token attends to the whole block, fused or not.
+
+    The kernel counts every entry of its active index ranges as a key, so the
+    builder must place the block around each query's causal window without -1
+    gaps, down to an empty context.
+    """
+    from vllm.models.deepseek_v4.nvidia import flashinfer_sparse as mod
+    from vllm.v1.attention.backends.mla.compressor_utils import (
+        get_dspark_swa_index_width,
+    )
+    from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
+
+    _skip_unless_rope_quant()
+    torch.manual_seed(0)
+    device = "cuda"
+    block_size, window, draft = 64, 128, 4
+    width = get_dspark_swa_index_width(window, draft)
+    contexts = [context_len, context_len + 37]
+    num_reqs, num_tokens = len(contexts), len(contexts) * draft
+    num_rows = num_tokens + 2
+    pages = -(-(max(contexts) + draft) // block_size)
+    cache = torch.randn((num_reqs * pages, block_size, 512), device=device)
+    cache = cache.clamp_(-1, 1).to(torch.float8_e4m3fn)
+    attn = _make_rope_quant_attn(mod, num_rows, device)
+    attn.attn_sink = torch.full((128,), -float("inf"), device=device)
+
+    # As the DSpark SWA builder lays them out: the trailing window of context
+    # and the whole block, contiguous from column 0.
+    decode_indices = torch.full((num_tokens, width), -1, dtype=torch.int32)
+    visible, positions = [], []
+    for r, ctx in enumerate(contexts):
+        slots = torch.arange(max(ctx - window, 0), ctx + draft) + r * pages * block_size
+        for k in range(draft):
+            decode_indices[r * draft + k, : slots.numel()] = slots.to(torch.int32)
+            visible.append(slots)
+            positions.append(ctx + k)
+    query_start = torch.arange(0, num_tokens + 1, draft, dtype=torch.int32)
+    seq_lens = torch.tensor([c + draft for c in contexts], dtype=torch.int32)
+    metadata = DeepseekSparseSWAMetadata(
+        block_table=torch.arange(
+            num_reqs * pages, dtype=torch.int32, device=device
+        ).view(num_reqs, pages),
+        slot_mapping=torch.zeros(num_tokens, dtype=torch.int64, device=device),
+        block_size=block_size,
+        seq_lens=seq_lens.to(device),
+        query_start_loc=query_start.to(device),
+        query_start_loc_cpu=query_start,
+        token_to_req_indices=torch.arange(
+            num_reqs, dtype=torch.int32, device=device
+        ).repeat_interleave(draft),
+        decode_swa_indices=decode_indices.to(device),
+        decode_swa_lens=torch.tensor(
+            [v.numel() for v in visible], dtype=torch.int32, device=device
+        ),
+        decode_swa_width=width,
+        is_valid_token=torch.ones(num_tokens, dtype=torch.bool, device=device),
+        num_decodes=num_reqs,
+        num_prefills=0,
+        num_decode_tokens=num_tokens,
+        num_prefill_tokens=0,
+        max_decode_query_len=draft,
+    )
+    q = torch.randn((num_rows, 128, 512), device=device).clamp_(-1, 1)
+    q = q.to(torch.float8_e4m3fn)
+    attention, z_unfused, z_fused = _rope_quant_vs_unfused(
+        attn, q, cache, metadata, torch.tensor(positions, device=device), num_tokens
+    )
+
+    keys_all = cache.flatten(0, 1).float()
+    for token, slots in enumerate(visible):
+        keys = keys_all[slots.to(device)]
+        weights = torch.softmax(q[token].float() @ keys.T * attn.scale, -1)
+        torch.testing.assert_close(
+            attention[token].float(), weights @ keys, atol=0.05, rtol=0.05
+        )
+    _assert_projects_alike(z_unfused, z_fused)

vllm/models/deepseek_v4/attention.py

@@ -43,6 +43,7 @@ from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.models.utils import extract_layer_index
@@ -171,14 +172,20 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
         q: torch.Tensor,
         kv: torch.Tensor,
         positions: torch.Tensor,
-        output: torch.Tensor,
+        output: "torch.Tensor | QuantizedActivation",
     ) -> None:
         """Platform-specific sparse MLA forward; writes attention into ``output``."""
         raise NotImplementedError
 
     @abstractmethod
-    def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
-        """Inverse-RoPE + wo_a + wo_b output projection (platform-specific)."""
+    def _o_proj(
+        self, o: "torch.Tensor | QuantizedActivation", positions: torch.Tensor
+    ) -> torch.Tensor:
+        """Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
+
+        ``o`` is the live heads of the bf16 attention output, or the
+        QuantizedActivation of a layer whose ``_alloc_attn_out`` returns one.
+        """
         raise NotImplementedError
 
     def _uses_fp8_ds_mla_layout(self) -> bool:
@@ -467,14 +474,9 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
         hidden_states: torch.Tensor,
         llama_4_scaling: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        # Pre-allocate attention output with FlashMLA-padded head count.
-        # The op writes into `o_padded`; we slice to n_local_heads after.
-        num_tokens = hidden_states.shape[0]
-        o_padded = torch.empty(
-            (num_tokens, self.padded_heads, self.head_dim),
-            dtype=hidden_states.dtype,
-            device=hidden_states.device,
-        )
+        # The eager attention region writes into a caller-owned buffer
+        # (breakable_cudagraph needs in-place outputs).
+        attn_out = self._alloc_attn_out(hidden_states.shape[0], hidden_states)
 
         # Keep the attention input preparation in the captured graph. Only the
         # sparse indexer and MLA attention run in the eager break below.
@@ -492,12 +494,27 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
             indexer_kv_score,
             indexer_weights,
             positions,
-            o_padded,
+            attn_out,
         )
-        o = o_padded[:, : self.n_local_heads, :]
+        if isinstance(attn_out, torch.Tensor):
+            attn_out = attn_out[:, : self.n_local_heads, :]
+        return self._o_proj(attn_out, positions)
 
-        # Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
-        return self._o_proj(o, positions)
+    def _alloc_attn_out(
+        self, num_tokens: int, hidden_states: torch.Tensor
+    ) -> "torch.Tensor | QuantizedActivation":
+        """The buffer ``forward_mqa`` fills.
+
+        A bf16 ``[num_tokens, padded_heads, head_dim]`` buffer by default, whose
+        padding heads are sliced off before ``_o_proj``. A layer whose kernel
+        also does the inverse RoPE and the FP8 cast returns a
+        QuantizedActivation instead and gets it back in ``_o_proj`` whole.
+        """
+        return torch.empty(
+            (num_tokens, self.padded_heads, self.head_dim),
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
 
     def _split_qkv_and_norm(
         self, qr_kv: torch.Tensor
@@ -529,7 +546,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
         indexer_kv_score: torch.Tensor,
         indexer_weights: torch.Tensor,
         positions: torch.Tensor,
-        o_padded: torch.Tensor,
+        o_padded: "torch.Tensor | QuantizedActivation",
     ) -> None:
         """Wide eager region: the whole of ``_prepare_and_attn`` runs eagerly.
 
@@ -558,7 +575,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
         indexer_kv_score: torch.Tensor,
         indexer_weights: torch.Tensor,
         positions: torch.Tensor,
-        o_padded: torch.Tensor,
+        o_padded: "torch.Tensor | QuantizedActivation",
     ) -> None:
         """Attention input preparation followed by the sparse indexer and MLA.
 
@@ -723,7 +740,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
         q: torch.Tensor,
         kv: torch.Tensor,
         positions: torch.Tensor,
-        out: torch.Tensor,
+        out: "torch.Tensor | QuantizedActivation",
     ) -> None:
         if self.indexer is not None and index_q is not None:
             assert index_weights is not None
@@ -736,7 +753,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
             )
 
         # MLA attention writes into the pre-allocated `out` buffer
-        # ([num_tokens, padded_heads, head_dim]).
+        # (see _alloc_attn_out).
         self.forward_mqa(q, kv, positions, out)
 
     def _fused_qnorm_rope_kv_insert(

vllm/models/deepseek_v4/common/ops/cache_utils.py

@@ -1099,6 +1099,7 @@ def build_flashinfer_mixed_sparse_indices(
     prefill_left_visible: torch.Tensor | None = None,
     prefill_right_visible: torch.Tensor | None = None,
     max_image_tokens: int = 0,
+    num_rows: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches.
 
@@ -1109,6 +1110,9 @@ def build_flashinfer_mixed_sparse_indices(
     tokens derive their SWA window from the position and translate local
     compressed indices to global slots via the block tables.
 
+    ``num_rows`` (>= ``num_tokens``) sizes both outputs for a kernel call that
+    spans a CUDA-graph-padded batch; the rows past ``num_tokens`` stay unset.
+
     When ``prefill_left_visible``/``prefill_right_visible`` are given (vision
     variant), the SWA column region widens by ``max_image_tokens`` and prefill
     tokens inside an image span get a bidirectionally widened window; decode
@@ -1170,13 +1174,15 @@ def build_flashinfer_mixed_sparse_indices(
     # by ``sparse_topk_lens``, so padding never changes the attention result.
     padded_topk = max(topk, decode_compressed_topk)
     padded_topk = (padded_topk + 3) // 4 * 4
+    num_rows = num_tokens if num_rows is None else num_rows
+    assert num_rows >= num_tokens
     sparse_indices = torch.empty(
-        (num_tokens, swa_total_width + padded_topk),
+        (num_rows, swa_total_width + padded_topk),
         dtype=torch.int32,
         device=decode_swa_indices.device,
     )
     sparse_topk_lens = torch.empty(
-        num_tokens, dtype=torch.int32, device=decode_swa_indices.device
+        num_rows, dtype=torch.int32, device=decode_swa_indices.device
     )
     if num_tokens == 0:
         return sparse_indices, sparse_topk_lens
@@ -1193,9 +1199,10 @@ def build_flashinfer_mixed_sparse_indices(
         if compressed_block_span is None
         else compressed_block_span
     )
+    # The launch grid follows the output rows, so hand over only the real ones.
     _BUILD_FLASHINFER_MIXED_SPARSE_INDICES_KERNEL(
-        sparse_indices,
-        sparse_topk_lens,
+        sparse_indices[:num_tokens],
+        sparse_topk_lens[:num_tokens],
         decode_swa_indices,
         decode_compressed_indices,
         decode_compressed_topk_lens,
@@ -1322,12 +1329,52 @@ class BuildFlashinferMixedSparseIndicesKernel(
         SWA_TOTAL_WIDTH: tl.constexpr = SWA_INDEX_WIDTH + IMAGE_WIDTH
 
         if token_idx < NUM_DECODE_TOKENS:
+            # A decode row wider than the window is DSpark's non-causal block:
+            # the block-anchored window of context followed by the whole block,
+            # contiguous from column 0. The kernel treats every entry of its
+            # active ranges as a key -- the first min(WINDOW_SIZE, pos + 1)
+            # columns and [WINDOW_SIZE, sparse_topk_lens) -- so -1 must not
+            # appear in either. Fill the causal window with the first visible
+            # entries and put the rest from column WINDOW_SIZE on, followed
+            # directly by the compressed entries.
+            NONCAUSAL: tl.constexpr = SWA_INDEX_WIDTH > WINDOW_SIZE
+            compressed_start = SWA_TOTAL_WIDTH
+            if NONCAUSAL:
+                req_idx = tl.load(token_to_req_indices_ptr + token_idx)
+                query_start = tl.load(query_start_loc_ptr + req_idx)
+                query_len = tl.load(query_start_loc_ptr + req_idx + 1) - query_start
+                context_len = tl.load(seq_lens_ptr + req_idx) - query_len
+                token_in_query = token_idx - query_start
+                window_len = tl.minimum(context_len + token_in_query + 1, WINDOW_SIZE)
+                spill_len = (
+                    tl.minimum(context_len, WINDOW_SIZE) + query_len - window_len
+                )
+                # A CUDA-graph pad token sits outside its empty request, so
+                # the kernel never reads its row; keep it the shape of a
+                # causal pad row (all -1, WINDOW_SIZE long) rather than let
+                # the formulas above go negative.
+                is_block_token = (token_in_query >= 0) & (token_in_query < query_len)
+                window_len = tl.where(is_block_token, window_len, 0)
+                spill_len = tl.where(is_block_token, spill_len, 0)
+                compressed_start = WINDOW_SIZE + spill_len
             for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE):
                 offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
                 mask = offset < SWA_TOTAL_WIDTH
+                src = offset
+                src_mask = offset < SWA_INDEX_WIDTH
+                if NONCAUSAL:
+                    # Columns from compressed_start on belong to the loops below.
+                    mask = mask & (offset < compressed_start)
+                    spill = offset >= WINDOW_SIZE
+                    src = tl.where(spill, offset - WINDOW_SIZE + window_len, offset)
+                    src_mask = tl.where(
+                        spill,
+                        offset < WINDOW_SIZE + spill_len,
+                        offset < window_len,
+                    )
                 values = tl.load(
-                    decode_swa_indices_ptr + token_idx * decode_swa_stride + offset,
-                    mask=offset < SWA_INDEX_WIDTH,
+                    decode_swa_indices_ptr + token_idx * decode_swa_stride + src,
+                    mask=src_mask,
                     other=-1,
                 )
                 values = _remap_flashinfer_index(values, swa_block_size, swa_block_span)
@@ -1372,12 +1419,27 @@ class BuildFlashinferMixedSparseIndicesKernel(
                 tl.store(
                     sparse_indices_ptr
                     + token_idx * sparse_indices_stride
-                    + SWA_TOTAL_WIDTH
+                    + compressed_start
                     + offset,
                     values,
                     mask=mask,
                 )
 
+            if NONCAUSAL:
+                # The row's remaining columns, left over by the shifted
+                # compressed block.
+                for i in range(0, SWA_TOTAL_WIDTH - WINDOW_SIZE, WINDOW_BLOCK_SIZE):
+                    offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
+                    tl.store(
+                        sparse_indices_ptr
+                        + token_idx * sparse_indices_stride
+                        + compressed_start
+                        + PADDED_TOP_K
+                        + offset,
+                        tl.full((WINDOW_BLOCK_SIZE,), -1, dtype=tl.int32),
+                        mask=offset < SWA_TOTAL_WIDTH - compressed_start,
+                    )
+
             if DECODE_COMPRESSED_TOPK == 0:
                 compressed_len = tl.zeros((), dtype=tl.int32)
             elif not DECODE_COMPRESSED_INDICES_ARE_LOCAL:
@@ -1388,7 +1450,9 @@ class BuildFlashinferMixedSparseIndicesKernel(
                 else:
                     compressed_len = tl.full((), DECODE_COMPRESSED_TOPK, dtype=tl.int32)
 
-            tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + compressed_len)
+            tl.store(
+                sparse_topk_lens_ptr + token_idx, compressed_start + compressed_len
+            )
             return
 
         prefill_idx = token_idx - NUM_DECODE_TOKENS

vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py

@@ -9,6 +9,8 @@ import torch
 from vllm.config import VllmConfig
 from vllm.config.cache import CacheDType
 from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
 from vllm.models.deepseek_v4.attention import DeepseekV4Attention
 from vllm.models.deepseek_v4.common.ops import (
     build_flashinfer_mixed_sparse_indices,
@@ -17,6 +19,8 @@ from vllm.models.deepseek_v4.common.ops import (
 from vllm.models.deepseek_v4.nvidia.ops.o_proj import (
     compute_fp8_einsum_recipe,
     deep_gemm_fp8_o_proj,
+    rope_quant_attn_out,
+    rope_quant_unsupported_reason,
 )
 from vllm.models.deepseek_v4.sparse_mla import (
     DeepseekV4FlashMLAMetadata,
@@ -37,6 +41,8 @@ from vllm.v1.attention.backends.mla.sparse_swa import (
 if TYPE_CHECKING:
     from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
 
+logger = init_logger(__name__)
+
 _FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
 _flashinfer_dsv4_workspace_by_device: dict[torch.device, torch.Tensor] = {}
 
@@ -203,7 +209,9 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
     def get_padded_num_q_heads(cls, num_heads: int) -> int:
         return _pad_to_supported_q_heads(num_heads)
 
-    def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+    def _o_proj(
+        self, o: torch.Tensor | QuantizedActivation, positions: torch.Tensor
+    ) -> torch.Tensor:
         return deep_gemm_fp8_o_proj(
             o,
             positions,
@@ -219,9 +227,30 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
             tma_aligned_scales=self._tma_aligned_scales,
         )
 
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
+    def _alloc_attn_out(
+        self, num_tokens: int, hidden_states: torch.Tensor
+    ) -> torch.Tensor | QuantizedActivation:
+        if self._fuse_rope_quant:
+            return rope_quant_attn_out(num_tokens, hidden_states.device)
+        return super()._alloc_attn_out(num_tokens, hidden_states)
+
+    def __init__(self, vllm_config: VllmConfig, *args, **kwargs) -> None:
+        super().__init__(vllm_config, *args, **kwargs)
         self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe()
+        # RopeQuant derives each query's RoPE position as seq_len - q_len + i,
+        # which matches `positions` for DSpark's non-causal draft block too,
+        # since every call passes the real per-request seq_lens. The one
+        # exception is a draft block running past max_model_len, whose clamped
+        # positions it does not reproduce; that only costs those drafts.
+        reason = rope_quant_unsupported_reason(self)
+        self._fuse_rope_quant = reason is None
+        if reason is None:
+            logger.info_once(
+                "FLASHINFER_MLA_SPARSE_DSV4 fuses the inverse RoPE and FP8 "
+                "quant of the attention output."
+            )
+        else:
+            logger.debug_once("FLASHINFER_MLA_SPARSE_DSV4 RopeQuant off: %s", reason)
         # Per-tensor FP8 scale buffers + precomputed scalar BMM scales. Only the
         # per-tensor FP8 cache path consumes these; bf16 reads ``self.scale``.
         if self.kv_cache_torch_dtype != torch.float8_e4m3fn:
@@ -253,32 +282,39 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
         q: torch.Tensor,
         kv: torch.Tensor,
         positions: torch.Tensor,
-        output: torch.Tensor,
+        output: torch.Tensor | QuantizedActivation,
     ) -> None:
-        # The TRTLLM-gen kernel requires h_q in {64, 128}, so the output buffer
-        # is allocated at the padded head count while q arrives at the local
-        # head count; _forward pads q to match before the launcher.
-        assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], (
-            f"output buffer shape {output.shape} incompatible with q shape {q.shape}"
-        )
-        assert output.shape[1] >= q.shape[1], (
-            f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}"
-        )
-        # Per-tensor FP8 q produces a bf16 attention output.
-        expected_output_dtype = (
-            torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype
-        )
-        assert output.dtype == expected_output_dtype, (
-            f"output dtype {output.dtype} must match expected {expected_output_dtype} "
-            f"for q dtype {q.dtype}"
-        )
+        if isinstance(output, QuantizedActivation):
+            assert output.data.shape[0] == q.shape[0]
+        else:
+            # The TRTLLM-gen kernel requires h_q in {64, 128}, so the output
+            # buffer is allocated at the padded head count while q arrives at
+            # the local head count; _forward pads q to match before the launcher.
+            assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], (
+                f"output buffer shape {output.shape} incompatible with q {q.shape}"
+            )
+            assert output.shape[1] >= q.shape[1], (
+                f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}"
+            )
+            # Per-tensor FP8 q produces a bf16 attention output.
+            expected_output_dtype = (
+                torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype
+            )
+            assert output.dtype == expected_output_dtype, (
+                f"output dtype {output.dtype} must match expected "
+                f"{expected_output_dtype} for q dtype {q.dtype}"
+            )
 
         forward_context = get_forward_context()
         attn_metadata = forward_context.attn_metadata
         if attn_metadata is None:
             # Warmup dummy run: FlashInfer reads the cache directly and lazily
             # allocates its workspace, so nothing to reserve here.
-            output.zero_()
+            if isinstance(output, QuantizedActivation):
+                output.data.zero_()
+                output.scale.zero_()
+            else:
+                output.zero_()
             return
 
         assert isinstance(attn_metadata, dict)
@@ -313,11 +349,12 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
         swa_metadata: "DeepseekSparseSWAMetadata",
         attn_metadata: DeepseekV4FlashMLAMetadata | None,
         swa_only: bool,
+        num_rows: int | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         """Build the combined sparse-index tensors for the mixed batch.
 
         Returns ``(compressed_kv_cache, seq_lens, sparse_indices,
-        sparse_topk_lens)``.
+        sparse_topk_lens)``, sized ``num_rows`` (default: the real tokens).
         """
         num_decodes = swa_metadata.num_decodes
         num_prefills = swa_metadata.num_prefills
@@ -412,6 +449,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
             if swa_only
             else ("c128a" if self.compress_ratio == 128 else "c4a")
         )
+        if num_rows is not None:
+            cache_key = f"{cache_key}:{num_rows}"
         cached_sparse = swa_metadata.flashinfer_sparse_index_cache.get(cache_key, None)
         if cached_sparse is None:
             swa_block_span = _packed_block_span(swa_k_cache)
@@ -439,8 +478,9 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
                 prefill_right_visible=swa_metadata.prefill_right_visible,
                 # getattr for tests that bypass __init__ via object.__new__.
                 max_image_tokens=getattr(self, "max_image_tokens", 0),
+                num_rows=num_rows,
             )
-            if cache_key != "c4a":
+            if not cache_key.startswith("c4a"):
                 swa_metadata.flashinfer_sparse_index_cache[cache_key] = (
                     sparse_indices,
                     sparse_topk_lens,
@@ -457,7 +497,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
         swa_metadata: "DeepseekSparseSWAMetadata",
         attn_metadata: DeepseekV4FlashMLAMetadata | None,
         swa_only: bool,
-        output: torch.Tensor,
+        output: torch.Tensor | QuantizedActivation,
     ) -> None:
         assert self.kv_cache_torch_dtype in (torch.bfloat16, torch.float8_e4m3fn)
         num_decodes = swa_metadata.num_decodes
@@ -468,6 +508,11 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
         num_tokens = num_decode_tokens + num_prefill_tokens
         if num_tokens == 0:
             return
+        # RopeQuant's output is group-major with a group stride of exactly
+        # sum_q tokens, so both calls span every row of the (CUDA-graph
+        # padded) buffers and select their tokens through cum_seq_lens_q
+        # alone; the kernel never touches rows outside it.
+        rope_quant = isinstance(output, QuantizedActivation)
 
         (
             compressed_kv_cache,
@@ -480,12 +525,12 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
             swa_metadata=swa_metadata,
             attn_metadata=attn_metadata,
             swa_only=swa_only,
+            num_rows=q.shape[0] if rope_quant else None,
         )
 
         # CUDA graph execution can pad q/output past the scheduled token count;
         # restrict to the real tokens (the launcher validates sparse indices).
-        query = q[:num_tokens]
-        output = output[:num_tokens]
+        query = q if rope_quant else q[:num_tokens]
         bmm1_scale: float | torch.Tensor = self.scale
         bmm2_scale: float | torch.Tensor = 1.0
         if self.kv_cache_torch_dtype == torch.float8_e4m3fn:
@@ -499,7 +544,9 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
         # The TRTLLM-gen sparse-MLA kernel requires h_q in {64, 128}; zero-pad
         # the query heads to the allocated output head count. Padded heads attend
         # to the shared KV and are sliced off downstream (output is padded too).
-        padded_heads = output.shape[1]
+        padded_heads = (
+            output.shape[1] if isinstance(output, torch.Tensor) else self.padded_heads
+        )
         if query.shape[1] < padded_heads:
             padded_query = query.new_zeros(
                 (query.shape[0], padded_heads, query.shape[2])
@@ -512,51 +559,60 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
         query_start_loc_cpu = swa_metadata.query_start_loc_cpu
         assert query_start_loc is not None and query_start_loc_cpu is not None
 
+        def run(
+            rows: slice, reqs: slice, cum_seq_lens_q: torch.Tensor, max_q_len: int
+        ) -> None:
+            if rope_quant:
+                assert isinstance(output, QuantizedActivation)
+                rows = slice(None)
+                out_kwargs = dict(
+                    out=output.data,
+                    dsv4_output_scale=output.scale,
+                    dsv4_inv_rope_cos_sin_cache=self.rotary_emb.cos_sin_cache,
+                )
+            else:
+                assert isinstance(output, torch.Tensor)
+                out_kwargs = dict(out=output[rows])
+            flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
+                query=query[rows],
+                swa_kv_cache=swa_k_cache,
+                workspace_buffer=workspace,
+                sparse_indices=sparse_indices[rows],
+                sparse_indices_are_storage_offsets=True,
+                compressed_kv_cache=compressed_kv_cache,
+                sparse_topk_lens=sparse_topk_lens[rows],
+                seq_lens=seq_lens[reqs],
+                bmm1_scale=bmm1_scale,
+                bmm2_scale=bmm2_scale,
+                sinks=self.attn_sink,
+                cum_seq_lens_q=cum_seq_lens_q,
+                max_q_len=max_q_len,
+                **out_kwargs,
+            )
+
         # Keep the TRTLLM-gen decode/prefill split: the launcher is tuned for
         # uniform-q batches, and this avoids flattening mixed batches into one call.
         if num_decode_tokens > 0:
-            decode_cu = query_start_loc[: num_decodes + 1]
-            flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
-                query=query[:num_decode_tokens],
-                swa_kv_cache=swa_k_cache,
-                workspace_buffer=workspace,
-                sparse_indices=sparse_indices[:num_decode_tokens],
-                sparse_indices_are_storage_offsets=True,
-                compressed_kv_cache=compressed_kv_cache,
-                sparse_topk_lens=sparse_topk_lens[:num_decode_tokens],
-                seq_lens=seq_lens[:num_decodes],
-                out=output[:num_decode_tokens],
-                bmm1_scale=bmm1_scale,
-                bmm2_scale=bmm2_scale,
-                sinks=self.attn_sink,
-                cum_seq_lens_q=decode_cu,
-                max_q_len=swa_metadata.max_decode_query_len,
+            run(
+                slice(0, num_decode_tokens),
+                slice(0, num_decodes),
+                query_start_loc[: num_decodes + 1],
+                swa_metadata.max_decode_query_len,
             )
 
         if num_prefill_tokens > 0:
-            # The prefill query view re-anchors at offset 0, so rebase the
-            # cumulative query offsets to start at 0.
-            prefill_cu = (
-                query_start_loc[num_decodes : num_reqs + 1]
-                - query_start_loc[num_decodes]
-            )
+            prefill_cu = query_start_loc[num_decodes : num_reqs + 1]
+            if not rope_quant:
+                # The prefill query view re-anchors at offset 0, so rebase the
+                # cumulative query offsets to start at 0.
+                prefill_cu = prefill_cu - query_start_loc[num_decodes]
             prefill_cu_cpu = query_start_loc_cpu[num_decodes : num_reqs + 1]
             prefill_lens_cpu = prefill_cu_cpu[1:] - prefill_cu_cpu[:-1]
-            flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
-                query=query[num_decode_tokens:num_tokens],
-                swa_kv_cache=swa_k_cache,
-                workspace_buffer=workspace,
-                sparse_indices=sparse_indices[num_decode_tokens:num_tokens],
-                sparse_indices_are_storage_offsets=True,
-                compressed_kv_cache=compressed_kv_cache,
-                sparse_topk_lens=sparse_topk_lens[num_decode_tokens:num_tokens],
-                seq_lens=seq_lens[num_decodes:num_reqs],
-                out=output[num_decode_tokens:num_tokens],
-                bmm1_scale=bmm1_scale,
-                bmm2_scale=bmm2_scale,
-                sinks=self.attn_sink,
-                cum_seq_lens_q=prefill_cu,
-                max_q_len=int(prefill_lens_cpu.max().item()),
+            run(
+                slice(num_decode_tokens, num_tokens),
+                slice(num_decodes, num_reqs),
+                prefill_cu,
+                int(prefill_lens_cpu.max().item()),
             )
 
 

vllm/models/deepseek_v4/nvidia/ops/o_proj.py

@@ -5,11 +5,15 @@ from collections.abc import Callable
 import torch
 import torch.nn as nn
 
+from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    kFp8Dynamic128Sym,
+)
 from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import (
     fused_inv_rope_fp8_quant,
 )
 from vllm.platforms import current_platform
-from vllm.utils.deep_gemm import fp8_einsum
+from vllm.utils.deep_gemm import fp8_einsum, get_tma_aligned_size
 
 
 def compute_fp8_einsum_recipe(
@@ -29,7 +33,7 @@ def compute_fp8_einsum_recipe(
 
 
 def deep_gemm_fp8_o_proj(
-    o: torch.Tensor,
+    o: torch.Tensor | QuantizedActivation,
     positions: torch.Tensor,
     cos_sin_cache: torch.Tensor,
     wo_a: nn.Module,
@@ -49,23 +53,31 @@ def deep_gemm_fp8_o_proj(
     layer selects the recipe at initialization. ``wo_b`` is any callable over
     the flattened ``z``: the projection module itself, or a wrapper that also
     reduce-scatters its output (DeepSeek-V4.1 GEMM-RS).
+
+    A QuantizedActivation ``o`` was already rotated and cast by the attention
+    kernel, one ``wo_a`` group per slot with the live groups first, so only
+    the padding groups are sliced off.
     """
     use_fp8 = wo_a.weight.dtype == torch.float8_e4m3fn
-    o_proj_input, o_scale = fused_inv_rope_fp8_quant(
-        o,
-        positions,
-        cos_sin_cache,
-        n_groups=n_groups,
-        heads_per_group=heads_per_group,
-        nope_dim=nope_dim,
-        rope_dim=rope_dim,
-        quant_group_size=einsum_recipe[2],
-        tma_aligned_scales=tma_aligned_scales,
-        quantize=use_fp8,
-    )
+    if isinstance(o, QuantizedActivation):
+        assert use_fp8
+        o_proj_input, o_scale = o.data[:, :n_groups], o.scale[:, :n_groups]
+    else:
+        o_proj_input, o_scale = fused_inv_rope_fp8_quant(
+            o,
+            positions,
+            cos_sin_cache,
+            n_groups=n_groups,
+            heads_per_group=heads_per_group,
+            nope_dim=nope_dim,
+            rope_dim=rope_dim,
+            quant_group_size=einsum_recipe[2],
+            tma_aligned_scales=tma_aligned_scales,
+            quantize=use_fp8,
+        )
     z = torch.empty(
-        (o.shape[0], n_groups, o_lora_rank),
-        device=o.device,
+        (o_proj_input.shape[0], n_groups, o_lora_rank),
+        device=o_proj_input.device,
         dtype=torch.bfloat16,
     )
     if use_fp8:
@@ -89,3 +101,58 @@ def deep_gemm_fp8_o_proj(
             out=z.transpose(0, 1),
         )
     return wo_b(z.flatten(1))
+
+
+# TRTLLM-gen's DSv4 RopeQuant epilogue (flashinfer-ai/flashinfer#4918) has one
+# fixed schedule: 128 query heads in groups of 8 512-wide heads, and FP8 with
+# one packed UE8M0 scale per 128 elements.
+_ROPE_QUANT_NUM_HEADS = 128
+_ROPE_QUANT_HEADS_PER_GROUP = 8
+_ROPE_QUANT_HEAD_DIM = 512
+_ROPE_QUANT_BLOCK = 128
+
+
+def rope_quant_unsupported_reason(layer: torch.nn.Module) -> str | None:
+    """Why ``layer`` cannot take FlashInfer's fused inverse RoPE + FP8 output.
+
+    DSv4 always has 8 heads per ``wo_a`` group at any TP size, and the rest of
+    the kernel's fixed shape (head/RoPE dims, per-128 einsum scales, FP32 RoPE
+    cache) is what the SM100 layer uses anyway. What varies is the local head
+    count and the configs the layer also serves unquantized.
+    """
+    if layer.n_local_heads != _ROPE_QUANT_NUM_HEADS:
+        # Padding fewer heads up to the kernel's fixed 128 costs more
+        # attention than the fused epilogue saves.
+        return f"RopeQuant needs {_ROPE_QUANT_NUM_HEADS} local query heads"
+    if layer.kv_cache_torch_dtype != torch.float8_e4m3fn:
+        return "RopeQuant needs the per-tensor FP8 KV cache"
+    if layer.wo_a.weight.dtype != torch.float8_e4m3fn:
+        return "RopeQuant needs an FP8 wo_a"
+    return None
+
+
+def rope_quant_attn_out(num_tokens: int, device: torch.device) -> QuantizedActivation:
+    """Allocate the output pair the RopeQuant kernel writes.
+
+    Values are group-major, ``[num_tokens, 16, 8 * 512]`` over physical
+    ``[group, token, K]``. Scales are MN-major packed UE8M0, one int32 per
+    head over its four 128-wide blocks. Both are what ``fp8_einsum`` consumes.
+    """
+    n_groups = _ROPE_QUANT_NUM_HEADS // _ROPE_QUANT_HEADS_PER_GROUP
+    group_width = _ROPE_QUANT_HEADS_PER_GROUP * _ROPE_QUANT_HEAD_DIM
+    data = torch.empty(
+        (n_groups, num_tokens, group_width), dtype=torch.float8_e4m3fn, device=device
+    ).transpose(0, 1)
+    aligned = get_tma_aligned_size(num_tokens, torch.int32.itemsize)
+    scale = torch.empty(
+        (n_groups, group_width // (_ROPE_QUANT_BLOCK * 4), aligned),
+        dtype=torch.int32,
+        device=device,
+    ).permute(2, 0, 1)[:num_tokens]
+    return QuantizedActivation(
+        data=data,
+        scale=scale,
+        orig_dtype=torch.bfloat16,
+        orig_shape=data.shape,
+        quant_key=kFp8Dynamic128Sym,
+    )