[Bugfix][DSv4.1] Fix FlashInfer DSpark non-causal attention

[WoosukKwon]

DSpark's non-causal draft window has 133 valid keys at draft length 5, padded to 256 indices. The DSV4.1 FlashInfer integration passed the padded width as the active sparse length, so padding entered softmax normalization. A constant-value reference expecting 1 returned 0.51953125 (133/256). Short contexts also inherited causal visibility from the launcher's per-request query positions.

Use the valid SWA length, retaining FlashInfer's required 128-column minimum, and present each draft query with the full non-causal sequence length. Prepare the per-token active lengths and sequence lengths once in the metadata builder, using persistent buffers refreshed before FULL CUDA graph replay and shared across layers. Singleton query/output views remove the need to construct query offsets in each layer. This changes vLLM metadata only; no FlashInfer source changes or new Triton kernel are needed. The two int32 buffers use 8 bytes per configured maximum batched token per builder (64 KiB at capacity 8192).

Validation on 4× GB200 with FlashInfer 0.6.18.post1, DeepSeek-V4.1-Flash (dba1be0a40aa45a94ad051997016db3960a90277), DSpark length 5 and full decode graphs. DEP4 uses MegaMoE/adaptive verification; TP4 uses auto MoE/fixed verification.

Full GSM8K: 1,319 identical 5-shot chat prompts per pair, greedy, seed 123, thinking off, 1,024-token output cap. Accuracy below uses flexible extraction.

Configuration	Acceptance before → fixed (including bonus)	GSM8K before → fixed
DEP4	3.1823 → 3.8248 (+20.2%)	1268 → 1271 / 1319 (96.13% → 96.36%)
TP4	3.2075 → 3.8879 (+21.2%)	1266 → 1263 / 1319 (95.98% → 95.75%)

All 1,319 requests completed in each fixed run, with no errors or empty responses. DEP had no capped outputs; TP had one (versus zero before). Strict extraction: DEP 1042 → 1042; TP 1060 → 1049. These evaluations establish acceptance recovery, not exact greedy-output equivalence or an accuracy improvement. MegaAttention measured 3.8437 acceptance on the DEP cohort. These are short-context results, not 100K-context throughput measurements.

.venv/bin/python -m pytest tests/v1/attention/test_dspark_noncausal_sparse_mla.py \
  -k test_dsv41_flashinfer_dspark_window_matches_reference -q -rs
.venv/bin/pre-commit run --files \
  vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py \
  vllm/v1/attention/backends/mla/sparse_swa.py \
  tests/v1/attention/test_dspark_noncausal_sparse_mla.py

The original 12 regression cases failed before the correctness fix. All 12 expanded cases pass with persistent metadata: BF16/FP8, 16/64 heads, short/long windows, mixed query lengths, padding, and CUDA graph replay after Q/KV and visible-length updates, with unchanged buffer addresses. Attention outputs are bitwise identical to the initial correctness fix across eight FP8 shapes (16/64 heads, batch 1/8/32/128). Nine full draft-forward runs have zero observed eager/graph, changed-input, restored-input and post-replay error. Pre-commit checks pass. Validation ran on the optimization worktree; all three changed files are identical in this isolated PR branch. The GSM8K results above evaluate the correctness fix; the metadata refactor was validated by parity and replay checks, without repeating the full accuracy runs.

Metadata-refactor performance on GB200 ×4, relative to the initial correctness fix: synthetic weights and 102,400 KV tokens/request, draft length 5, TP auto MoE / DEP MegaMoE. Medians of five distributed-MAX samples, 20 graph replays/sample after 3 seconds of warmup; each comparison uses the same runner and repeats one endpoint to expose order effects.

Workload	Initial fix (µs)	Prepared metadata (µs)	Draft-forward result
TP4 B1	403 / 425	366	9–14% lower latency
DEP4 B8	703 / 691	669	3–5% lower latency
DEP4 B128	1370	1431 / 1327	Inconclusive

The forward boundary excludes metadata preparation, LM head, scheduler and sampling; these are not serving-throughput gains. Preparation costs 3.9–8.5 µs per build in a separate CUPTI/CUDA-graph/cold-L2 microbenchmark and is shared across layers. Exact commands and raw results are archived locally under /home/woosuk/.cache/v41-dspark-opt/metadata-opt/metadata-paired-002-*; the microbenchmark results are in prepared-results.json.

Duplicate check: searched open DSpark/FlashInfer, acceptance, noncausal sparse, and padding PRs. #52499 and #57292 address SM120 routing/page geometry; #54834 addresses dense FlashInfer cache-layout resolution; #55692 addresses block-table bounds. None fixes the SM100 DSV4.1 non-causal active lengths and visibility corrected here.

AI assistance: implemented and validated with OpenAI Codex.

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[WoosukKwon]

/ci run

[github-actions]

❌ This PR is 7 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

[WoosukKwon]

/ci run

[github-actions]

✅ Triggered Buildkite CI #89709 for commit 6554023f06ce.