[SM120] Split the DSv4 extra KV cache to pbs=64 for sparse-MLA prefill

[kkr16]

Motivation

On SM120/SM121 (RTX PRO 6000 Blackwell, GB10/DGX Spark), serving DeepSeek-V4 with the default FlashInfer sparse-MLA backend aborts on any prompt longer than 64 tokens:

tvm.error.InternalError: Check failed: (ok) is false: Unsupported sparse-MLA prefill
configuration: model=DSV4 num_heads=8 topk=128 page_block_size=64 topk_extra=512
extra_page_block_size=128

flash_mla_with_kvcache_sm120() routes a forward to _flash_mla_sm120_prefill whenever q.shape[0] > SM120_DECODE_MAX_TOKENS (64). That function splits the main KV cache from the SWA pool's page size down to 64 via _split_kv_pages_to_64, because the FlashInfer SM120 kernels only instantiate page_block_size=64. It does not do the same for the DSv4 secondary (C4A/C128A) cache, which stays at 128.

On the kernel side, dispatch_dsv4_dual in sparse_mla_sm120_prefill.cu accepts extra_page_block_size of 64 or 2 only — in both the fulltile fast path and the general DISPATCH_BY_NH_PBSX branch. It returns false, and the TVM_FFI_ICHECK(ok) in sparse_mla_sm120.cu fires.

Every other parameter already matched: num_heads=8 ✅, topk=128 ✅, main page_block_size=64 ✅. Only the extra cache was wrong.

The decode path is unaffected and is deliberately left alone — it accepts extra_page_block_size=128 and has always worked. That is why short prompts succeed and only prefill fails.

This is distinct from the topk-bucketing reports in sgl-project#33134 and flashinfer-ai/flashinfer#3828, which reach the same error string via a different parameter (topk=192/256 vs the instantiated 128/512/1024/2048 buckets). This PR is specifically about extra_page_block_size.

Continues the page-split work in sgl-project#32320, sgl-project#35116 and sgl-project#29927.

Modifications

_split_kv_pages_to_64 was already generic in src_pbs (assert src_pbs % 64 == 0) and the extra cache uses the identical 584-byte-per-token footer layout, so the fix is small:

1. Call the existing helper for the extra cache in _flash_mla_sm120_prefill, deriving the source page size from the tensor shape (handling the 3D, NHD and HND forms the binding documents) and skipping the work when it is already 64.

2. Add a key_suffix parameter to _split_kv_pages_to_64. The helper's persistent output and touched-page mask buffers are keyed per device (flash_mla_sm120_split:{dev}, flash_mla_sm120_mask:{dev}) and are shared, so a second call within one forward aliases the first.

   This second part is not cosmetic. With a shared key the extra-cache split reuses the main cache's buffer and silently overwrites it — measured at 512 of 1024 main-cache tokens corrupted in the added test's configuration. key_suffix defaults to "", so the existing main-cache call site and its buffer key names are byte-identical to today.

Token indices are invariant under the split (dst_page = src_page*ratio + sub ⇒ dst_page*64 + off == src_page*src_pbs + sub*64 + off), so no index remapping is needed — the same property the main-cache split already relies on.

Why split in Python rather than instantiate extra_page_block_size=128 in FlashInfer: splitting to 64 is valid against every released FlashInfer wheel and fixes users today. A new C++ template instantiation would require a FlashInfer release before it helps anyone and would leave all current wheels broken. Happy to follow up on the FlashInfer side as well.

Accuracy Tests

All runs on 8× RTX PRO 6000 Blackwell (SM120), TP=8, DeepSeek-V4.1-Flash.

1. The page split is bitwise neutral. The same logical extra-cache contents were materialised twice — as 16 pages of pbs=128 (takes the new split path) and as 32 pages of pbs=64 (no split) — and fed to the same kernel:

layout sanity: sampled token byte mismatches = 0
extra cache pbs=128 (SPLIT to 64) vs pbs=64 (NATIVE, no split)
  bitwise identical : True
  max |delta|       : 0.00000000
  mismatching elems : 0 / 524288

So the split contributes exactly zero numerical difference.

2. Dual-cache prefill agrees with the in-tree reference. _flash_mla_sm120_prefill vs the pure-PyTorch _sm120_sparse_decode_fwd, 128 tokens / 8 heads / topk=128 / extra_topk=512:

pair	max abs delta	mean abs delta	cosine
flashinfer vs torch	2.563e-3	3.83e-4	0.99975300
triton vs torch	4.88e-4	2.7e-5	0.99999636
flashinfer vs triton	2.563e-3	3.86e-4	0.99975163

FlashInfer sits further from the torch reference than Triton does, but per (1) that gap is a pre-existing property of the kernel and not of this change — identical inputs produce bitwise-identical output whether or not the split runs.

3. The unsplit path still reproduces the bug. Bypassing the extra-cache split in the same harness reproduces the original error exactly:

Check failed: (ok) is false: Unsupported sparse-MLA prefill configuration:
model=DSV4 num_heads=8 topk=128 page_block_size=64 topk_extra=512 extra_page_block_size=128

4. End-to-end. Sweeping prompt length across the 64-token boundary against a live server (server-reported prompt_tokens):

reps=1   prompt_tokens=29    decode  OK
reps=3   prompt_tokens=57    decode  OK
reps=4   prompt_tokens=71    PREFILL OK
reps=8   prompt_tokens=127   PREFILL OK
reps=16  prompt_tokens=239   PREFILL OK
reps=32  prompt_tokens=463   PREFILL OK
reps=64  prompt_tokens=911   PREFILL OK
failures=0

Before this patch every prompt_tokens > 64 row aborted the server. A full serving sweep (1k/1k, 1k/8k, 8k/1k at concurrency 32/64/128) also runs clean, including 8k prompts that exercise multi-page chunked prefill.

5. New unit test. test/registered/kernels/ops/attention/test_flash_mla_sm120_page_split.py — this file previously had no test coverage. It covers:

• byte-exact footer round-trip for pbs=256 (regression guard for the existing main-cache path) and pbs=128 (the extra cache)
• key_suffix isolating the output/mask buffers, so the extra split cannot clobber the main split
• default key_suffix="" preserving the historic buffer key names
• touched_indices masking still restricting copies to referenced pages

It exercises only the Triton page-copy path, so it runs on any CUDA GPU and needs no SM120 hardware.

Ran 5 tests in 0.672s
OK

Speed Tests and Profiling

There is no meaningful before/after throughput comparison for this path: before the patch, every DSv4 prefill of more than 64 tokens aborted the server, so the baseline does not run.

On cost: the extra-cache split reuses the same touched-page masking introduced in sgl-project#32320, so only source pages actually referenced by extra_indices are copied rather than the whole pool. The additional persistent buffer is num_dst_pages * 37440 bytes, grow-only, allocated once per device.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• 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.