[Bugfix][DSv4.1] Fix FlashInfer DSpark non-causal attention (#57432)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Codex <noreply@openai.com>
Woosuk Kwon committed
Sep 17, 2026 at 15:14 UTC
80447d27655918da6bfbccd0d3a40e975bda220a
3 files changed
+214
-6
tests/v1/attention/test_dspark_noncausal_sparse_mla.py
+150
@@ -604,3 +604,153 @@ def test_dspark_noncausal_differs_from_causal(
604
f"non-causal backend output matches the causal reference "
605
f"(max abs diff={causal_err}); future-pointing indices are not attended to"
606
)
607
+
608
+
609
+@pytest.mark.parametrize("context_len", [20, 128, 900])
610
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn])
611
+@pytest.mark.parametrize("num_heads", [16, 64])
612
+def test_dsv41_flashinfer_dspark_window_matches_reference(
613
+ context_len, dtype, num_heads, monkeypatch
614
+):
615
+ """Draft queries see the full block without treating padded slots as keys."""
616
+ if not current_platform.is_device_capability_family(100):
617
+ pytest.skip("DSV4 TRTLLM sparse attention requires SM100")
618
+ from vllm.models.deepseek_v41.nvidia.flashinfer_sparse import (
619
+ DeepseekSparseSWAFlashInferMetadataBuilder,
620
+ DeepseekV4FlashInferMLAAttention,
621
+ )
622
+ from vllm.models.deepseek_v41.sparse_mla import (
623
+ DeepseekV41SparseSWAMetadataBuilder,
624
+ )
625
+
626
+ torch.manual_seed(123)
627
+ device = "cuda"
628
+ query_lens = [5, 3]
629
+ context_lens = [context_len, context_len + 137]
630
+ num_real_tokens, num_tokens, head_dim = 8, 9, 512
631
+ cache = torch.randn(32, 128, head_dim, device=device, dtype=torch.bfloat16).to(
632
+ dtype
633
+ )
634
+ query = torch.randn(
635
+ num_tokens, num_heads, head_dim, device=device, dtype=torch.bfloat16
636
+ ).to(dtype)
637
+ indices = torch.full((num_tokens, 256), -1, device=device, dtype=torch.int32)
638
+ visible_indices = []
639
+ visible_lens = []
640
+ for req, (context, query_len) in enumerate(zip(context_lens, query_lens)):
641
+ visible = (
642
+ torch.arange(max(context - 128, 0), context + query_len, device=device)
643
+ + req * 16 * 128
644
+ )
645
+ visible_indices.extend([visible] * query_len)
646
+ visible_lens.extend([visible.numel()] * query_len)
647
+ for token, visible in enumerate(visible_indices):
648
+ indices[token, : visible.numel()] = visible.to(torch.int32)
649
+ query_start_loc = torch.tensor([0, 5, 8, 9], dtype=torch.int32)
650
+ metadata = SimpleNamespace(
651
+ num_decodes=3,
652
+ num_prefills=0,
653
+ num_decode_tokens=num_tokens,
654
+ num_prefill_tokens=0,
655
+ seq_lens=torch.tensor(
656
+ [context_lens[0] + 5, context_lens[1] + 3, 1],
657
+ device=device,
658
+ dtype=torch.int32,
659
+ ),
660
+ query_start_loc=query_start_loc.to(device),
661
+ query_start_loc_cpu=query_start_loc,
662
+ token_to_req_indices=torch.tensor(
663
+ [0] * 5 + [1] * 3 + [2], device=device, dtype=torch.int32
664
+ ),
665
+ decode_swa_indices=indices,
666
+ decode_swa_width=256,
667
+ decode_swa_lens=torch.tensor(
668
+ visible_lens + [0], device=device, dtype=torch.int32
669
+ ),
670
+ block_table=torch.arange(48, device=device, dtype=torch.int32).view(3, -1),
671
+ block_size=128,
672
+ flashinfer_sparse_index_cache={},
673
+ max_decode_query_len=5,
674
+ )
675
+
676
+ # Exercise FlashInfer preparation without constructing a model/config.
677
+ def init_parent(builder):
678
+ builder._max_tokens = num_tokens
679
+ builder.device = device
680
+ builder.window_size = 128
681
+
682
+ monkeypatch.setattr(DeepseekV41SparseSWAMetadataBuilder, "__init__", init_parent)
683
+ monkeypatch.setattr(
684
+ DeepseekV41SparseSWAMetadataBuilder,
685
+ "build",
686
+ lambda *args: metadata,
687
+ )
688
+ builder = DeepseekSparseSWAFlashInferMetadataBuilder()
689
+ common_metadata = SimpleNamespace(causal=False)
690
+ builder.build(0, common_metadata)
691
+ prepared = (
692
+ metadata.flashinfer_decode_topk_lens,
693
+ metadata.flashinfer_decode_seq_lens,
694
+ )
695
+ attention = SimpleNamespace(
696
+ kv_cache_torch_dtype=dtype,
697
+ window_size=128,
698
+ compress_ratio=0,
699
+ topk_indices_buffer=torch.empty(
700
+ num_tokens, 0, device=device, dtype=torch.int32
701
+ ),
702
+ scale=1 / math.sqrt(head_dim),
703
+ _flashinfer_fp8_bmm1_scale=1 / math.sqrt(head_dim),
704
+ _flashinfer_fp8_bmm2_scale=1.0,
705
+ attn_sink=None,
706
+ )
707
+ attention._build_sparse_index_metadata = MethodType(
708
+ DeepseekV4FlashInferMLAAttention._build_sparse_index_metadata, attention
709
+ )
710
+ output = torch.empty_like(query, dtype=torch.bfloat16)
711
+
712
+ def forward():
713
+ metadata.flashinfer_sparse_index_cache.clear()
714
+ DeepseekV4FlashInferMLAAttention._forward(
715
+ attention, query, None, cache, metadata, None, True, output
716
+ )
717
+
718
+ def check_output():
719
+ references = []
720
+ for token, visible in enumerate(visible_indices):
721
+ keys = cache.flatten(0, 1)[visible].float()
722
+ weights = torch.softmax(
723
+ query[token].float() @ keys.T / math.sqrt(head_dim), -1
724
+ )
725
+ references.append(weights @ keys)
726
+ atol = 0.01 if dtype == torch.bfloat16 else 0.05
727
+ torch.testing.assert_close(
728
+ output[:num_real_tokens].float(),
729
+ torch.stack(references),
730
+ atol=atol,
731
+ rtol=0.05,
732
+ )
733
+
734
+ forward()
735
+ check_output()
736
+ graph = torch.cuda.CUDAGraph()
737
+ with torch.cuda.graph(graph):
738
+ forward()
739
+ # Replay must consume the current inputs, including KV written since capture.
740
+ query.copy_(torch.randn_like(query, dtype=torch.bfloat16).to(dtype))
741
+ cache.copy_(torch.randn_like(cache, dtype=torch.bfloat16).to(dtype))
742
+ # A new step changes visibility as well as Q/KV, without recapturing.
743
+ metadata.seq_lens.sub_(1)
744
+ metadata.decode_swa_lens[:num_real_tokens].sub_(1)
745
+ for token, visible in enumerate(visible_indices):
746
+ indices[token, visible.numel() - 1] = -1
747
+ visible_indices[token] = visible[:-1]
748
+ builder.build(0, common_metadata)
749
+ assert metadata.flashinfer_decode_topk_lens.data_ptr() == prepared[0].data_ptr()
750
+ assert metadata.flashinfer_decode_seq_lens.data_ptr() == prepared[1].data_ptr()
751
+ torch.testing.assert_close(prepared[0], metadata.decode_swa_lens.clamp_min(128))
752
+ torch.testing.assert_close(
753
+ prepared[1], metadata.seq_lens[metadata.token_to_req_indices.long()]
754
+ )
755
+ graph.replay()
756
+ check_output()
vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py
+62
-6
@@ -26,7 +26,11 @@ from vllm.models.deepseek_v41.sparse_mla import (
26
)
27
from vllm.platforms.interface import DeviceCapability
28
from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4
29
-from vllm.v1.attention.backend import AttentionCGSupport, MultipleOf
29
+from vllm.v1.attention.backend import (
30
+ AttentionCGSupport,
31
+ CommonAttentionMetadata,
32
+ MultipleOf,
33
+)
34
from vllm.v1.attention.backends.mla.compressor_utils import (
35
get_dspark_swa_index_width,
36
)
@@ -183,6 +187,39 @@ class DeepseekSparseSWAFlashInferMetadataBuilder(DeepseekV41SparseSWAMetadataBui
187
188
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
189
190
+ def __init__(self, *args, **kwargs):
191
+ super().__init__(*args, **kwargs)
192
+ # Graphs retain these addresses while each build refreshes their contents.
193
+ self._decode_topk_lens = torch.empty(
194
+ self._max_tokens, dtype=torch.int32, device=self.device
195
+ )
196
+ self._decode_seq_lens = torch.empty_like(self._decode_topk_lens)
197
+
198
+ def build(
199
+ self,
200
+ common_prefix_len: int,
201
+ common_attn_metadata: CommonAttentionMetadata,
202
+ fast_build: bool = False,
203
+ ) -> "DeepseekSparseSWAMetadata":
204
+ metadata = super().build(common_prefix_len, common_attn_metadata, fast_build)
205
+ num_tokens = metadata.num_decode_tokens
206
+ if not common_attn_metadata.causal and num_tokens > 0:
207
+ assert metadata.decode_swa_lens is not None
208
+ assert metadata.seq_lens is not None
209
+ assert metadata.token_to_req_indices is not None
210
+ topk_lens = self._decode_topk_lens[:num_tokens]
211
+ seq_lens = self._decode_seq_lens[:num_tokens]
212
+ torch.clamp(metadata.decode_swa_lens, min=self.window_size, out=topk_lens)
213
+ torch.index_select(
214
+ metadata.seq_lens,
215
+ 0,
216
+ metadata.token_to_req_indices[:num_tokens],
217
+ out=seq_lens,
218
+ )
219
+ metadata.flashinfer_decode_topk_lens = topk_lens
220
+ metadata.flashinfer_decode_seq_lens = seq_lens
221
+ return metadata
222
+
223
224
class DeepseekSparseSWAFlashInferBackend(DeepseekSparseSWABackend):
225
@staticmethod
@@ -493,21 +530,40 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
530
# Keep the TRTLLM-gen decode/prefill split: the launcher is tuned for
531
# uniform-q batches, and this avoids flattening mixed batches into one call.
532
if num_decode_tokens > 0:
533
+ decode_query = query[:num_decode_tokens]
534
+ decode_output = output[:num_decode_tokens]
535
decode_cu = query_start_loc[: num_decodes + 1]
536
+ decode_seq_lens = seq_lens[:num_decodes]
537
+ decode_topk_lens = sparse_topk_lens[:num_decode_tokens]
538
+ max_decode_query_len = swa_metadata.max_decode_query_len
539
+ if swa_metadata.decode_swa_width > self.window_size:
540
+ # DSpark's non-causal window extends past the fixed 128 SWA
541
+ # columns into the aliased compressed pool. Exclude padding,
542
+ # and expose the full block to each query instead of letting
543
+ # TRTLLM derive a causal SWA length from its query position.
544
+ assert swa_only
545
+ assert swa_metadata.flashinfer_decode_topk_lens is not None
546
+ assert swa_metadata.flashinfer_decode_seq_lens is not None
547
+ decode_topk_lens = swa_metadata.flashinfer_decode_topk_lens
548
+ decode_seq_lens = swa_metadata.flashinfer_decode_seq_lens
549
+ decode_query = decode_query.unsqueeze(1)
550
+ decode_output = decode_output.unsqueeze(1)
551
+ decode_cu = None
552
+ max_decode_query_len = 1
553
flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
498
- query=query[:num_decode_tokens],
554
+ query=decode_query,
555
swa_kv_cache=swa_k_cache,
556
workspace_buffer=workspace,
557
sparse_indices=sparse_indices[:num_decode_tokens],
558
compressed_kv_cache=compressed_kv_cache,
503
- sparse_topk_lens=sparse_topk_lens[:num_decode_tokens],
504
- seq_lens=seq_lens[:num_decodes],
505
- out=output[:num_decode_tokens],
559
+ sparse_topk_lens=decode_topk_lens,
560
+ seq_lens=decode_seq_lens,
561
+ out=decode_output,
562
bmm1_scale=bmm1_scale,
563
bmm2_scale=bmm2_scale,
564
sinks=self.attn_sink,
565
cum_seq_lens_q=decode_cu,
510
- max_q_len=swa_metadata.max_decode_query_len,
566
+ max_q_len=max_decode_query_len,
567
)
568
569
if num_prefill_tokens > 0:
vllm/v1/attention/backends/mla/sparse_swa.py
+2
@@ -205,6 +205,8 @@ class DeepseekSparseSWAMetadata:
205
num_decode_tokens: int = 0
206
num_prefill_tokens: int = 0
207
max_decode_query_len: int = 1
208
+ flashinfer_decode_topk_lens: torch.Tensor | None = None
209
+ flashinfer_decode_seq_lens: torch.Tensor | None = None
210
211
# Pre-computed prefill metadata shared across all DeepseekV4 attention layers.
212
prefill_seq_lens: torch.Tensor | None = None