[SM120] flash_mla: allocate the page-split buffer outside inference mode

[AliceChenyy]

Purpose

_split_kv_pages_to_64 keeps two lazily allocated persistent buffers in the same
buffers dict, with the same lifetime: allocated once on first use, reused across
autotune, CUDA graph capture and steady-state serving.

On current main only one of them is protected:

    buf = buffers.get(key)                      # page-split destination
    if buf is None or buf.shape[0] < num_dst_pages:
        buf = torch.empty(...)                  # <- no guard
        buffers[key] = buf
...
        mbuf = buffers.get(mkey)                # dirty mask
        if mbuf is None or mbuf.shape[0] < N:
            # The first allocation can happen under inference mode (autotune),
            # but the buffer is zeroed again later during CUDA graph capture
            # outside inference mode -- an inference tensor cannot be mutated
            # there, so force a normal tensor.
            with torch.inference_mode(False):   # <- guarded
                mbuf = torch.empty(N, dtype=torch.int8, device=dev)
            buffers[mkey] = mbuf

The reasoning in that comment applies verbatim to buf: it is written during CUDA graph
capture too. This PR closes the asymmetry.

The mbuf guard was added after @moxcat reported the failure while running DSV4 on
SM120; the fix at the time covered both lazy buffers in this function, and the buf
half was lost somewhere between that backport and main. The same pattern, with the same
justification, also exists in
sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py.

When it does fire, the symptom is:

RuntimeError: Inplace update to inference tensor outside InferenceMode is not allowed

What I could and could not verify

I could not construct a configuration on 4x RTX 6000D (SM120) where the unguarded
allocation actually lands under inference mode. Instrumenting the allocation site:

config	autotune ran	mode at allocation	result
--moe-runner-backend deep_gemm	no (0 invocations)	inference_mode=False	boots, no error
--moe-runner-backend flashinfer_mxfp4, DeepSeek-V4-Flash-0731	yes (4 invocations)	inference_mode=False	boots, no error

In both, buf is first allocated during Capturing batches, so the ordering never
arises. So I am not attaching a before/after reproducer — the case for this change is
the asymmetry with mbuf, not a failure I can demonstrate on demand.

@moxcat — if you still have the configuration that hit this originally, it would turn the
table above into a real before/after.

Context

Split out of #29927 at @Fridge003's request, since it is not SM120 enablement work.
#29927 now depends on it.

Scope

One allocation site. Buffer contents, shape, caching key and lifetime are unchanged;
only the mode it is allocated under. No-op wherever the allocation already happens
outside inference mode.

[moxcat]

The original configuration still exists, and I re-ran it today on both the
v0.5.16 backport stack where I hit the failure and on current main (2e7c85d), with
both allocation sites instrumented to print inference-mode state, stream-capture
state, and Tensor.is_inference() at each lazy allocation. All runs on 2x RTX PRO
6000 Blackwell (sm 120), TP2. Short version: your table is correct on main; the
arming ordering is real but branch-specific; and the missing split-buffer crash is
writer type, not allocation mode.

On the v0.5.16 stack (autotune enabled, DSPARK, CUDA graphs; exact flags in the fold
below), first touch of _split_kv_pages_to_64 was during autotune under
inference_mode=True, both ranks, every boot:

variant	resulting tensors	outcome
both allocations wrapped (as this stack served from July)	ordinary	boots, serves
split bare, mask wrapped (this PR's before state)	split buffer is_inference=True, and it stays in the persistent buffers dict	boots and serves anyway; no inference-tensor error in the log
both bare (the July pre-fix shape)	both is_inference=True	capture aborts at mask.zero_(): "Capture cuda graph failed: Inplace update to inference tensor outside InferenceMode is not allowed"

On current main with the same decisive flags (TP2, DSPARK, autotune enabled), first
touch still lands at autotune start, but with inference_mode=False; both buffers
come up ordinary. The difference is in the code: v0.5.16 runs the autotune warmup
under with torch.inference_mode(), autotune(...)
(srt/model_executor/runner/flashinfer_autotune.py); main removed that wrap in
#33788, closing #33470. So this path cannot arm the trap on main today, which
confirms your table from the machine that hit the original failure.

Why row two serves despite holding a live inference tensor: the split buffer is
written only by the Triton _page_split_kernel, and raw kernel writes never reach
the dispatcher check that raises this error. The one every-call torch-level in-place
write in the function is the mask's .zero_(), hence the observed crashes are
mask-site. The unguarded buf is a latent state: it becomes an inference
tensor whenever a caller first touches this function under inference mode, and
whether that aborts depends on which writer touches it next. A future torch-level
write, or a warmup path running under inference mode again, turns it into a boot
failure. That supports this PR as written: hygiene that closes the asymmetry, not a
fix for a crash reproducible on today's main.

Exact configurations and July history

v0.5.16 stack: sglang v0.5.16 with the #32320 backport, torch 2.11, DeepSeek-V4-Flash
DSpark checkpoint, TP2, EP2, --speculative-algorithm DSPARK,
--speculative-num-draft-tokens 6 (checkpoint block size 5), --kv-cache-dtype fp8_e4m3, --chunked-prefill-size 4096, --cuda-graph-max-bs 4,
--mem-fraction-static 0.93, context 65536, FlashInfer autotune enabled.

Current main (2e7c85d): DeepSeek-V4-Flash-0731, TP2, --moe-runner-backend flashinfer_mxfp4, --speculative-algorithm DSPARK, --swa-full-tokens-ratio 0.1,
--kv-cache-dtype fp8_e4m3, context 32768, autotune enabled (not disabled as in the
#29927 eval command), --mem-fraction-static 0.96 (TP2 needs the headroom that the
TP4 command gets from halving per-rank weights).

July history, stated precisely: what I observed in July was this failure class on the
v0.5.16 backport, and I wrapped both lazy allocations in one edit at the time,
shortly before the same mask failure was reported and fixed on #32320. The
split-buffer half of that edit was the sibling allocation wrapped by the same
reasoning, not a second observed traceback; I do not have, and am not claiming, a
split-site crash. Re-running the pre-fix shape today, the abort is at mask.zero_()
during capture, as in the table. The closed PR #30759 from July is consistent with
this era: it worked around lazy allocation during FlashInfer autotune inside
torch.inference_mode() by allocating eagerly.

Row three of the table doubles as a live boot-time reproduction of the failure class,
if one is ever wanted for the record. Happy to share the raw probe logs if useful.