[unified-memory] Let Kimi-Linear use the paged MLA attention backends

[ch-wan]

Stacked on #32971 (base is feature/unified-memory-kimi-linear-triton). Review that one first; this PR's diff is just the two commits on top.

Some refactor work is needed after Kimi K3 is merged to main.

What this does

With the dense per-layer MLA views from #32971 in place, the stock paged MLA kernels can read the unified pool directly — only their kv_indices / block tables need remapping to dense ids. No change to the physical layout, the allocator, or compaction.

So this opens the page-major full-attention allowlist from {triton} to the paged MLA backends when --enable-unified-memory is set on an MLA model:

backend	role	how
trtllm_mla	prefill + decode	block table via create_flashmla_kv_indices_triton, now fused with v2p_ptr + PAGE_MULT
flashinfer	prefill + decode	token-level kv_indices remapped to dense, written back in place (see below)
cutedsl_mla	decode	subclasses TRTLLMMLABackend, inherits the dense path with zero extra code
tokenspeed_mla	decode (fp8)	same inheritance

flashmla / cutlass_mla share the create_flashmla block-table path and can be added the same way once exercised.

Cuda-graph safety

trtllm_mla decode gets a capture-stable dense write-loc buffer (cuda_graph_out_cache_loc_dense) filled out of graph, so the in-graph set_mla_kv_buffer captures no translate. The replay prep zeroes the buffer's tail beyond the raw batch size — without that, pad rows retain stale dense locs from earlier larger replays and scatter garbage KV into live pages (a bug that only appears once a server has load history, and which bf16 can mask through batch saturation).

Validation

Kimi-Linear-48B-A3B-Instruct (27 layers = 20 KDA + 7 MLA, kv_lora_rank 512 + qk_rope_head_dim 64), single node 8×B300 SXM6, TP1 per cell, GSM8K 200q. Every cell has a same-commit, same-GPU static-pool twin, so the unified column is a delta and not an absolute:

prefill	decode	KV	static	unified
triton	triton	bf16	0.910	0.910
trtllm_mla	trtllm_mla	bf16	0.905	0.895
trtllm_mla	cutedsl_mla	bf16	0.915	0.915
trtllm_mla	tokenspeed_mla	fp8_e4m3, ps64	0.905	0.905
flashinfer	flashinfer	bf16	0.910	0.925 (eager)

The 0.895–0.925 band is noise: the static arms span the same 0.905–0.915, and each unified arm matches its own twin inside it.

Unit tests: test_unified_mla_dense_block_table.py (new) plus the #32971 suite and the allowlist test — 43 passed on B300.

Correction: flashinfer decode is admitted, not excluded

This section replaced an earlier one titled "The allowlist is per role: flashinfer decode is excluded." That described an interim version of the third commit that was superseded before merge; the description was not updated at the time. What follows is what actually merged. The diagnosis and the numbers below are unchanged and still accurate — only the remedy differs.

The second commit fixes a bug found while validating the first. Admitting flashinfer for decode under unified served silently wrong output:

config	GSM8K
unified + flashinfer decode + cuda graph	0.000 (Invalid 1.000, reproduced twice, no crash, no warning)
unified + flashinfer decode + eager	0.925
static + flashinfer decode + cuda graph	0.910

Cause: FlashInferMLAIndicesUpdaterDecode.call_begin_forward did its virtual→dense remap as a per-forward gather returning a fresh tensor, rebinding a local. On replay, the capture-stable buffer the captured kernel actually reads kept the raw virtual ids — fast_mla_decode_plan ignores its kv_indices argument entirely, so the rebind was dead. trtllm_mla escaped this only because it was given the out-of-graph capture-stable buffer described above.

The remedy that merged writes the translation back in place rather than rejecting the backend:

valid = kv_indices[:paged_kernel_lens_sum]
valid.copy_(self._translate_kv_loc_dense(valid))

bounded to the prefix the index kernel just filled, so the stale tail can never index the v2p table out of bounds. This is why a per-role split was not taken: --attention-backend flashinfer resolves to flashinfer for both roles, so a decode-only rejection would have rejected the only invocation anyone actually uses, leaving only a never-exercised flashinfer prefill + triton decode mix. _handle_page_major_kv_layout therefore keeps a single allowed_full set, with flashinfer admitted for both roles.

The in-place write is correct because of a flashinfer-side contract: under use_cuda_graph, BatchMLAPagedAttentionWrapper keeps the buffer passed at construction as _kv_indices_buf and plan() copies into it rather than rebinding, so the kernel reads exactly the buffer we wrote. No flashinfer patch is needed.

Post-merge validation (this was measured after merge; every flashinfer figure above was eager, so the captured path had not actually been run). Kimi-Linear-48B, 2×H200 TP2, --attention-backend flashinfer, cuda graph enabled and confirmed replaying (cuda graph: True on every decode batch line):

examples	unified	static twin
200	0.905	0.925
400	0.917	0.902

Pooled over 600: 0.913 unified vs 0.910 static. The gap changes sign between the two sample sizes — noise, not a regression.

Fusing the remap into create_flashinfer_kv_indices_triton (v2p_ptr + PAGE_MULT, as create_flashmla_kv_indices_triton already does) remains a worthwhile follow-up — it saves a launch and a temporary, and makes the remap impossible to skip — but it is an optimization, not a correctness prerequisite.

Two pre-existing issues found, NOT introduced here

Both reproduce on the static pool at this commit, so they belong to main, not to this PR. Filing separately; noting them here because they shape what a reviewer can reproduce on Kimi-Linear:

1. flashinfer prefill in a mixed pair crashes. ValueError: q.shape[0] (2586) does not match qo_indptr[-1] (664) from the chunked-MHA path (flashinfer_mla_backend.py mha_chunk_kv_cache.forward ← deepseek_common/.../forward_mha.py forward_normal_core ← hybrid_attn_backend.py). Fires whenever flashinfer prefill is paired with a different decode backend; a uniform flashinfer pair never hits it.
2. fp8 × flashinfer prefill materialises the whole KV pool as bf16. k_buf = self.token_to_kv_pool.get_key_buffer(layer_id).to(q.dtype) → OutOfMemoryError: Tried to allocate 270.37 GiB. Routing tokenspeed_mla through a trtllm_mla prefill instead gives the 0.905 == 0.905 row above, so the fp8 decode backend itself is fine.

Limitations

• TP1 only. TP1 is what made an 8-cell parallel sweep cheap (48B bf16 = 98 GB fits one 275 GB B300) and it exercises the same dense-view code, but it does not cover TP-sharded behavior. Worth a TP8 confirmation before this is called a platform result.
• Correctness only — no perf A/B, no retraction or long-context stress in this round.
• Based on 9f5655340; wants a rebase (main is ~24 commits ahead).

---

CI States

Latest PR Test (Base): ❌ Run #30614286898
Latest PR Test (Extra): ❌ Run #30614286592

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: cf06c6f7e4

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[ch-wan]

/tag-and-rerun-ci

[ch-wan]

CI test: https://github.com/sgl-project/sglang/actions/runs/30590695606/job/91077312520?pr=32972