[Bugfix] Propagate decode SWA index width for DSpark + FLASHINFER_MLA_SPARSE_DSV4

[ilmarkov]

DSpark non-causal draft batches allocate decode_swa_indices with
noncausal_index_width (wider than window_size). FlashInfer DSV4 sparse
MLA still reshaped / passed self.window_size into
build_flashinfer_mixed_sparse_indices, which crashes those drafts, e.g.:

RuntimeError: shape '[7, 128]' is invalid for input of size 1792
(at `flashinfer_sparse.py` `_build_sparse_index_metadata` when reshaping
`decode_swa_indices`).

• Add decode_swa_width to DeepseekSparseSWAMetadata and set it in the SWA
  builder (window_size or noncausal_index_width).
• Use that width in FlashInfer _build_sparse_index_metadata and the ROCM
  ragged conversion path.

Test Plan

• pytest tests/kernels/attention/test_flashmla_sparse.py
• Manual: DeepSeek-V4-Pro-DSpark with attention_backend: FLASHINFER_MLA_SPARSE_DSV4

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[ilmarkov]

@mgoin @benchislett

[lucifer1004]

Hi @ilmarkov — after coordinating with the maintainers we're consolidating the DSV4 DSpark fixes into a single PR, which I'm opening shortly and which will reference this one. Your work is included rather than replaced: the amd/rocm.py hunk and the decode_swa_width field there are byte-identical to yours, so you're credited as Co-authored-by on that commit.

Two deliberate differences from this PR:

1. We keep window_size logical and derive the padded width inside build_flashinfer_mixed_sparse_indices, passing both to the kernel as WINDOW_SIZE and SWA_INDEX_WIDTH. Passing decode_swa_width as window_size is functionally equivalent today — I instrumented the SWA builder and observed 792 non-causal builds under overlapping prefill/decode load, none carrying prefill tokens, so the kernel's prefill branch that derives swa_len/swa_start_pos from WINDOW_SIZE is never reached. We separate the two defensively, not because your version is wrong.

2. We pad the non-causal width to a multiple of 64 (192 for the K=5 shape) rather than 128 (256). Both buckets dispatch after feat(sm120): consolidate DSV4 sparse MLA top-k 192/256 support flashinfer-ai/flashinfer#4380, but 192 matches the kernel's 64-entry tile, and on 8xRTX PRO 6000 it measures 13-16% faster than 256 at >=8 tokens (equal at 1 token, where the kernel is launch-bound).

The consolidated PR carries four further fixes DSpark needs to run end-to-end, with gsm8k n=1319 evaluations across IFB and P/D.

Happy to reshape this if you'd rather land yours first and have us stack on top — just say the word.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ilmarkov.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[ilmarkov]

Close as resolved in #52809