[Perf][DSA] Pass topk_length to flash_mla_sparse_fwd in the sparse attention path

[zkyue]

Motivation

sgl_kernel.flash_mla_sparse_fwd accepts an optional per-row topk_length tensor, and the vendored FlashMLA sparse prefill kernels (both sm90 and sm100) early-exit their top-k loop after ceil_div(topk_length[row], B_TOPK) blocks. The DSA backend (dsa_backend.py::_forward_flashmla_sparse, the DeepSeek-V3.2 / GLM-5 path) never passed it, so every row whose context is shorter than index_topk scanned the full -1-padded top-k width for nothing. With index_topk = 2048 that is:

• the first 2048 tokens of every prefill sequence (roughly topk / (2L) of the sparse-prefill kernel work: ~25% for a 4k prompt, ~12% at 8k),
• short-context decode / MTP verify / draft-extend rows.

This is pure padding tax: the kernel masks -1 indices either way, so the skipped tail contributes nothing to the output. deepseek_v4_backend.py already passes topk_length on its sparse-prefill path; this PR brings the DSA backend in line. Prefill performance on the V3.2/DSA path has been a reported concern (e.g. #14498) — this recovers part of the sparse-prefill kernel time for short/medium contexts for free. The mechanism (trailing fully-invalid top-k blocks) is analyzed in deepseek-ai/FlashMLA#196.

Modifications

16-line diff in python/sglang/srt/layers/attention/dsa_backend.py:

• _forward_flashmla_sparse gains an optional topk_length parameter, forwarded to flash_mla_sparse_fwd (with a defensive row-count guard that falls back to the old full-width behavior on mismatch).
• Both call sites (extend and decode) pass metadata.dsa_cache_seqlens_int32 — the existing "seqlens clipped to topk" metadata, which is exactly the per-row count of valid indices. No new tensors or allocations in the hot path:
  • int32, 1-D, contiguous, and row-aligned with q (pad_dsa_cache_seqlens / _pad_topk_indices pad to the same row count; pad rows get length 0, matching their all--1 indices);
  • already updated in place under CUDA graph replay (fused_dsa_*_metadata / .copy_()); the flashmla_kv decode path already consumes it inside captured graphs today.

Correctness relies on top-k emitters writing all valid indices in the first min(seqlen, topk) slots with -1 padding at the tail, which holds for all backends: sgl-kernel topk.cu (short rows: indice[i] = (i < length) ? i : -1; long rows fill all slots), the torch fallback (torch.topk sorted output puts masked -inf entries last), and the flashinfer / fused-v2 transforms (documented "-1 padded" contract). The DSA test fixture codifies the same contract (_make_dsa_sparse_topk_rows tail-pads every pattern, including the non-trailing strided / head_tail selections).

Accuracy Tests

Kernel-level A/B on B200 (sm100, torch 2.11.0, sglang-kernel 0.4.4), emulating the exact call-site layout (q [s_q, 128, 576] bf16, kv bf16, d_v=512, topk=2048, indices tail-padded with -1, topk_length[i] = min(i+1, 2048)): out, lse and max_logits are bitwise identical with vs without topk_length, for the ramp case, the fully-valid case (topk_length == topk everywhere — no behavior change when there is nothing to skip), and empty rows (topk_length == 0 gives the same output as an all--1 row).

test/registered/attention/unittests/dsa/test_dsa.py (B200): results with this patch are identical to an unpatched upstream/main baseline run in the same environment — all sparse-path tests pass (test_sparse_topk_cases, test_sparse_non_trailing_index_cases (strided/head_tail), test_sparse_prefill_impl_variants/test_sparse_decode_impl_variants/test_sparse_cuda_graph_decode_impl_variants (flashmla_sparse + flashmla_kv), fp8 prefill/decode, speculative forward modes, sparse layout robustness). The only failures in both runs are pre-existing environment issues in the dense MHA fallback and the trtllm variant (failure sets byte-identical between baseline and patch).

sgl-kernel/tests/test_flashmla.py -k prefill: 20 passed.

Speed Tests and Profiling

Paired interleaved CUDA-event timing (median of 50 pairs) of flash_mla_sparse_fwd on B200, same layout as above. Note: shared node with ~85-90% ambient neighbor load, so treat these as a noise band; the controlled-environment number for the same kernel mechanism is in deepseek-ai/FlashMLA#196.

case	no topk_length	with topk_length	speedup
s_q=4096, ramp	1.695 ms	1.414 ms	1.20x
s_q=8192, ramp	5.555 ms	4.818 ms	1.15x
s_q=8192, fully-valid	6.244 ms	6.263 ms	1.00x (parity)

A second order-balanced run at s_q=8192 under heavier ambient load gave 1.05-1.07x (ramp) and 0.997x (fully-valid), i.e. the ramp win tracks the expected topk/(2L) work reduction minus fixed overheads, and the fully-valid case is unchanged.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests (covered by the existing DSA attention unit matrix; results verified identical to baseline).
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ⏳ Run #30599366957
Latest PR Test (Extra): ❌ Run #30599366700

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[DarkSharpness]

LGTM. Actually I don't know why this is missing. Perhaps FlashMLA updated its interface 2 months ago and we didn't follow up correctly?

[zkyue]

Thanks for the review! Your guess matches what we found when tracing this: topk_length was added to the kernel interface upstream and deepseek_v4_backend.py picked it up at that point, but these two dsa_backend.py call sites predate the interface change and were never revisited — the metadata tensor (dsa_cache_seqlens_int32) was already computed and graph-safe, just not plumbed through. So it is exactly a missed follow-up rather than anything deeper, which is also why the fix is a pure two-call-site diff with no kernel or metadata changes.

[zkyue]

Friendly bump 🙏 This is a small (+16/-0) perf-only change on the DSA sparse-attention path (passing topk_length through to flash_mla_sparse_fwd), open since Jul 14 and approved by @DarkSharpness. Since dsa_backend.py is owned by the attention CODEOWNERS, could one of @Fridge003 @hebiao064 @HaiShaw take a look and help kick off CI when you get a chance? Happy to address any feedback — thanks!

[Fridge003]

/rerun-test test/registered/models_e2e/test_dsa_glm52_tp_mtp.py test/registered/models_e2e/test_dsa_glm52_dp_mtp.py

[github-actions]

Results for /rerun-test test/registered/models_e2e/test_dsa_glm52_tp_mtp.py test/registered/models_e2e/test_dsa_glm52_dp_mtp.py:

🚀 8-gpu-h200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/models_e2e/test_dsa_glm52_tp_mtp.py
cd test/ && python3 registered/models_e2e/test_dsa_glm52_dp_mtp.py