[SM120] Only split touched SWA pages in FlashMLA page-split kernel

[TTThanos]

On SM120 the FlashMLA decode path splits the entire SWA KV pool (pbs=256 -> pbs=64) on every attention layer every decode step, but only ~2*batch pages are actually referenced by the sparse indices. Add a mark + masked-copy pass so only touched pages are copied; the persistent dst buffer keeps stale (unreferenced) data for the rest.

Fully fixed-grid / no D2H sync, so CUDA graph capture/replay is unaffected. Numerically equivalent to the full split on all touched pages (verified max_abs_diff=0.0).

Profile (DSv4-Flash, TP4+DP2, cc=8, 12 decode steps):
_page_split_kernel 16.8% -> 2.2% of GPU time (77.9ms -> 8.1ms)
median ITL -20%, mean TPOT -18%

Motivation

Modifications

Only copy pages that are actually referenced, via a fixed-grid mark + masked-copy pass (no D2H sync, CUDA-graph friendly):

_page_mark_kernel (new): one program per index element; for each valid token index t (skip -1), sets mask[t // src_pbs] = 1. Concurrent stores of the same value 1 are safe — no atomic needed.
_page_split_kernel: added mask_ptr + HAS_MASK: tl.constexpr. When HAS_MASK is set, each program loads mask[page_idx] and returns early if 0, so untouched pages incur only a single byte load instead
of a full data+scale copy. The grid stays (N * ratio,) (fixed) so launch shape is unchanged under CUDA graph capture.
split_kv_pages_to_64: new optional touched_indices arg. When provided: zero() the persistent int8 mask → run _page_mark_kernel → run the masked page_split_kernel. When None, behaves exactly as
before (full split, HAS_MASK=False). The mask buffer is lazily allocated and reused across steps; zero() is a memset that captures cleanly into the graph.
_flash_mla_flashinfer: computes idx (already needed downstream) before the split and passes it as touched_indices.

Accuracy Tests

DSv4-Flash, 8x RTX PRO 5000 72GB, TP4+DP2+dp-attention, `--moe-runner-backend marlin`, `--cuda-graph-max-bs 32`, random 4096-in / 150-out.                                                

torch profiler (12 decode steps, TP0):
| kernel | before | after |
|---|---|---|
| `_page_split_kernel` | 77.87ms (16.8%) | 8.11ms (2.2%) |
| `sparse_mla_decode_dsv4_kernel` | 12.13ms (2.6%) | 10.87ms (2.9%) |
| total GPU kernel time | 462.99ms | 373.47ms |

End-to-end (cc=8, 24 prompts), verified in two independent containers:
| metric | before | after |
|---|---|---|
| median ITL | 34.65 ms | 27.77 ms (-20%) |
| mean TPOT | 41.67 ms | 34.12 ms (-18%) |

Speed Tests and Profiling

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.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #30241155013
Latest PR Test (Extra): ❌ Run #30241154893

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[ormandj]

@TTThanos I opened TTThanos#1 against your branch. It fixes the current isort failure and adds an SM120-gated regression through the production page-split wrapper. The targeted test passes on RTX PRO 6000 Blackwell, and the full test_flash_mla_backends.py file is green (14 tests, one existing skip). The only production-file change is import ordering.

[ormandj]

I found one current-main integration issue while running a matched TP2 A/B on 2x RTX PRO 6000 Blackwell. The current head creates the persistent page mask during inference-mode autotune; target-verify CUDA graph capture then aborts at mask.zero_() with Inplace update to inference tensor outside InferenceMode is not allowed.

Focused fix and lifecycle regression: TTThanos#2

The unfixed module fails the strengthened test and full server startup. With that fix, all 14 registered FlashMLA backend tests pass on SM120 (1 skipped), target-verify/draft graph capture completes, and the server becomes ready.

Canonical matched A/B (same image ancestry/config/overlays; runtime delta is #32320 plus the mask-allocation fix; 30 s warmup/settle and 30 s sustained cells, client inside pod over localhost):

Concurrency	Control	Candidate	Delta
1	141.16	167.83	+18.89%
2	240.70	285.16	+18.47%
4	387.29	454.20	+17.28%
8	595.19	662.67	+11.34%
16	865.75	940.36	+8.62%
32	1299.89	1410.79	+8.53%

Coding median improved 146.68 -> 174.78 tok/s (+19.15%). Cold prefill (not this decode path) was -1.6% to -2.5% in this single pair, so I am not attributing a prefill effect without repeated evidence.

[ormandj]

Quality follow-up: I ran the full pinned GSM8K test split (1,319 questions, identical five-shot prompts, temperature 0, seed 0, parallel 8) against the matched control and fixed candidate.

• control: 1249/1319 (94.693%)
• candidate: 1248/1319 (94.617%)
• paired table: candidate-only 14, control-only 15, both correct 1234, both wrong 56
• McNemar exact p=1.0
• paired bootstrap delta CI95: [-0.910, +0.758] percentage points
• zero invalid extractions/finish reasons; structural gates passed

The harness strict raw-count gate is false because the candidate is lower by one answer; I am not relabeling that as a pass. The 14-vs-15 discordance is symmetric and exact response text agreed on only 31.3% of rows, so this pair provides no evidence of a quality change under normal batching variability. The kernel regression separately verifies every referenced data/scale byte is copied exactly and unreferenced pages are not read.

[moxcat]

Independent confirmation of the inference-tensor trap that TTThanos#2 just fixed for the page mask, plus the same fix for the sibling buffer.

I hit this failure class on a v0.5.16 backport of this PR (2x RTX PRO 6000 Blackwell, sm_120, TP2, torch 2.11) before that fix landed, and wrapped both lazy allocations at the time. At the current head the page-split destination buffer (key flash_mla_sm120_split, allocated just above the mask in the same path) still has the bare lazy allocation, subject to the same trap the mask fix addresses: created under inference mode, it comes up as an inference tensor that later in-place updates cannot touch ("Inplace update to inference tensor outside InferenceMode is not allowed"). Same two-line shape as the mask fix:

--- a/python/sglang/kernels/ops/attention/flash_mla_sm120.py
+++ b/python/sglang/kernels/ops/attention/flash_mla_sm120.py
@@ -393,12 +393,17 @@
     key = f"flash_mla_sm120_split:{dev}"
     buf = buffers.get(key)
     if buf is None or buf.shape[0] < num_dst_pages:
-        buf = torch.empty(
-            num_dst_pages,
-            _BYTES_PER_DST_PAGE_PADDED,
-            dtype=torch.uint8,
-            device=dev,
-        )
+        # The first allocation can happen under inference mode (autotune),
+        # but the buffer is written again later during CUDA graph capture
+        # outside inference mode, where an inference tensor cannot be
+        # mutated, so force a normal tensor.
+        with torch.inference_mode(False):
+            buf = torch.empty(
+                num_dst_pages,
+                _BYTES_PER_DST_PAGE_PADDED,
+                dtype=torch.uint8,
+                device=dev,
+            )
         buffers[key] = buf
     out = buf[:num_dst_pages]
 

The same constraint is already documented in-tree for the persistent metadata buffer in srt/layers/attention/deepseek_v4_backend.py.

With both allocations wrapped, the backport has been serving DeepSeek-V4-Flash with the DSPARK drafter for about 24 hours (across two brief config restarts) with correctness gates green (GSM8K-50, needle retrieval, tool calls). The lines are in TTThanos#3 for one-click absorption, or fold them in directly, whichever is easier.

[Fridge003]

/rerun-test test/registered/kernels/ops/attention/test_flash_mla_backends.py

[github-actions]

Results for /rerun-test test/registered/kernels/ops/attention/test_flash_mla_backends.py:

🚀 1-gpu-h100 (1 test): ❌ View workflow run

cd test/ && python3 registered/kernels/ops/attention/test_flash_mla_backends.py