Fix sparse MLA cache layout validation and eight-head decode

[ormandj]

📌 Description

Follow-up to flashinfer-ai/flashinfer#5075, targeting its branch. Preserve its canonical payload and runtime-stride design while tightening cache-view validation and bringing over the remaining GLM NoPE fix and regression coverage from #5022.

Flat 2D cache parsing currently substitutes the packed width for the actual block stride. A view with gaps between pages therefore reads the wrong page. The bindings also accept misaligned cache origins and block strides, and prefill accepts padded token strides for footer-scale caches whose kernels require packed rows. Preserve the actual block stride and reject these unsupported layouts before launching kernels.

GLM NoPE with eight query heads currently selects a kernel that uses sixteen scratch rows, while the caller allocates eight. Add the dedicated eight-head instantiation. Carry over compact/padded, masked-row, and CUDA graph replay tests; extend coverage to aligned storage offsets, additional padded row widths, RoPE prefill routes, and scratch guard regions.

These binding and scratch gaps also exist before flashinfer-ai#5075. This contribution retains Gabriel Wu's broader masking and layout implementation.

Invalidate the earlier calibration-cache schema: it can retain the former 656-byte GLM footprint and measured tuning choices after the payload becomes 528 bytes, as reported by @qsang-nv. This retires version-1 tuning files for all sparse MLA families. Normal untuned dispatch remains available until they are regenerated. Remember a stale file’s timestamp so repeated misses do not reread it; load replacement tuning when the file changes.

Add explicit coverage distinguishing GLM's 512 FP8 values plus four FP32 scales from DeepSeek-V4.1's 512 FP8 values plus sixteen UE8M0 scales. Both have a 528-byte payload, so byte size alone must not select the format. A deterministic per-group test checks all sixteen DeepSeek-V4.1 scale groups across decode chunk boundaries.

🔍 Related Issues

• Parent: flashinfer-ai/flashinfer#5075.
• Original eight-head fix and GLM regression coverage: flashinfer-ai/flashinfer#5022.
• Companion integration: sgl-project/sglang#38430.

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

Configured hooks pass on the changed files. The full-repository hook run is not claimed.

Author-run CPU validation: 64 passed, 17 GPU cases skipped across test_sparse_mla_sm120_dispatch.py, test_sparse_mla_sm120_cpb_model.py, and test_sparse_mla_nvfp4_sm120_plan.py. The new calibration regression fails when the module uses schema 1 and passes with schema 2, covering old constants, measured overrides and crossover entries, plus a current-schema save/reload.

Author-run GPU validation on RTX PRO 6000 Blackwell Max-Q (SM120):

Scope	Result
tests/attention/test_sparse_mla_sm120.py	696 passed
Compute Sanitizer memcheck, layout/masking/scratch/scale-group subset	171 passed, 0 errors
SGLang native adapter, compact and padded rows, holes and empty rows	6 passed

These GPU runs used the final changed files from this branch over an integration containing flashinfer-ai#5075, including a fresh run of the SGLang adapter tests. Source hashes were checked before testing. This is author-run validation of that integration; the complete repository suite and other GPU architectures are not claimed.

python -m pytest -q tests/attention/test_sparse_mla_sm120_dispatch.py tests/attention/test_sparse_mla_sm120_cpb_model.py tests/attention/test_sparse_mla_nvfp4_sm120_plan.py
python -m pytest -q tests/attention/test_sparse_mla_sm120.py

The regression matrix covers compact and padded GLM rows, physical page gaps, aligned storage offsets, rejected misalignment/footer-row gaps, eight-head scratch guards, poisoned unused cache entries, and CUDA graph replay with changed inputs. The standalone decode tests also cover its existing packed-row checks for both main and extra caches. Whole-model serving checks passed with the combined integration, but do not isolate this PR's effect; no serving-speedup or exact-output-equivalence claim is made.

🔬 Experimental Track

• This PR is experimental: it adds or changes code under flashinfer/experimental/ and/or an @flashinfer_experimental_api. Tracking issue: #
  • The tracking issue names an owner, the reason for the experimental path, and a graduation plan with a target release.
  • Core changes are limited to a thin entry point (signature, shared validation, feature-gate check, backend selection, handoff).
  • Tests live in tests/experimental/ and were validated on the intended hardware; a runnable example is included.
  • Nothing is registered in flashinfer/aot.py, and no experimental backend is reachable from backend="auto" without FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1. (Calling an @flashinfer_experimental_api or naming a backend explicitly is itself the opt-in and needs no environment variable.)
  • Test scope declared below. The experimental CI lane runs exactly these targets, so keep them as narrow as the change allows.

# One target per line: a directory or a file. (A pytest ::selector is not
# supported -- the sharding runner cannot consume one.) Must be under
# tests/experimental/ and must exist. Delete these comment lines and add yours, e.g.
#
#   tests/experimental/test_my_backend.py
#   tests/experimental/my_backend/
#
# Declaring the whole tree (tests/experimental/) is allowed but means every
# experimental PR pays for every other feature's tests, in every matrix cell.

Reviewer Notes

Review physical block strides, packed footer rows, the eight-head scratch contract, and calibration invalidation. Flat 2D GLM caches use packed 528-byte rows; existing padded 656-byte pools need an explicit token axis in a 3D/4D view.

The separate NVFP4 slot-zero masking issue reported on flashinfer-ai#5075 is outside this FP8 change. DeepSeek-V4.1 dual-cache prefill above 64 query tokens remains unsupported by the parent.

Please preserve the contributed commits' authorship when incorporating them, including co-author credit if squash-merging.

Developed with AI assistance.

[lucifer1004]

Thanks @ormandj — this is a great catch set, and the timing needs one note: flashinfer-ai#5075 was force-pushed back today to its pre-DSV4.1 state (the DSV4.1 cache-format support moved to a follow-up PR so the NaN/row-stride fixes can land first). Your branch was based on the pre-rewind head, so I integrated your changes directly rather than merging this PR as-is.

Now in flashinfer-ai#5075 (single squashed commit, Co-authored-by preserved per your request):

• the cache-view validation hardening (physical block stride for flat 2D caches, origin/block-stride 16B alignment, packed footer rows) with its tests;
• the GLM NoPE eight-head decode instantiation and its scratch-guard coverage (retaining the Support compact GLM NoPE FP8 rows and fix masked reads flashinfer-ai/flashinfer#5022 implementation);
• the calibration schema invalidation — I had just pushed the same v2 bump in parallel; I kept the merged comment and, on top, your stale-mtime memoization so stale files are not reparsed on every lookup;
• the flat page-gap decode tests.

Deliberately not taken (they belong with the DSV4.1 follow-up PR, since the rewound branch has no DSV4_1 type): the two 528B-format disambiguation / DSV4.1 scale-group commits. If you are willing, rebase just those two onto the follow-up once it is open — they are exactly the coverage it needs. During integration I also had to adapt test_sparse_mla_sm120_decode_footer_row_gap_rejected: it parametrized _MODEL_TYPE_DSV4_1 and passed the explicit model_type FFI arg, neither of which exists on the rewound branch; the test now covers DSV4 only, and the DSV4.1 entry can return with the follow-up.

Two replies to your notes:

• The NVFP4 slot-zero masking gap reported by qsang-nv is now fixed directly in feat(sm120): runtime KV row stride + canonical 528B GLM53_NOPE rows + NaN-safe masked gathers in sparse MLA flashinfer-ai/flashinfer#5075 (both NVFP4 kernels gather the shared zero row; the streaming one also serves the NVFP4 prefill route), with a poisoned-slot regression test.
• "DeepSeek-V4.1 dual-cache prefill above 64 query tokens remains unsupported" — acknowledged; that is the follow-up PR's scope.

Closing this PR since the content is incorporated — thank you again. (Developed-with-AI disclosure noted; same here per repository policy.)