[unified-memory] Mark the write loc physical and check it at every write door (4/7)

[caihuali95]

This is one of seven PRs split out of #38592. That PR was +2105/−1366 across 41
files, twelve of them attention backends on the serving hot path, and it interleaved
hardening, a layout change and dead-code deletion in a single diff — a reviewer could not
tell which hunk was safe-by-construction and which changed behaviour. #38592 is retained as the token-major
layout PR (3/7), which is the core of the series.

Stack — each PR is based on the one above it, and every one is green on its own:

#	PR	what it does	size
1	#40326	derive KV row addresses from strides, not shapes	21 files, +733/−118
2	#40327	build paged KV views through one helper	13 files, +704/−47
3	#38592	token-major layout	36 files, +1215/−1108
4	#40328	mark the write loc physical and check it at every write door	45 files, +935/−132
5	#40329	remove the kernel-page multiplier plumbing	17 files, +117/−230
6	#40330	stride the KV translate kernel and route every translate through it	11 files, +361/−159
7	#40331	name the fused KV translate for what it computes	28 files, +119/−121

Review order is the table order. 1 and 2 are no-ops on today's contiguous pools (2 also
fixes how HND pools with more than one KV head are read); 3 is where the behaviour changes
and deserves the most attention; 4 adds the physical write-loc check; 5 is deletion; 6
consolidates the translate; 7 is naming only.

4/7 — mark the write loc physical, and check the mark at every write door

Base: #38592 (3/7). 9 commits, 45 files, +935/−132.

Why

A unified pool's write door needs physical token ids. The allocator hands out virtual ids,
and KVIndexTranslator.rebind_write_loc translates the batch's write loc at ForwardBatch
construction; every other pool allocates physically, so the two coincide. Under 3/7's
token-major views a virtual id is in range and, by value, the same kind of integer as a
physical one, so a loc that skipped the rebind writes to the wrong slots without any error.
3/7 deletes the only probe that could tell them apart, because the new layout makes it
vacuous. This PR replaces it with an explicit mark, checked at every write door.

The contract

piece	role
ForwardBatch.out_cache_loc_is_physical	the batch's write loc holds physical ids. rebind_write_loc sets it on every pool, translating or not; the capture-time and draft-runner constructors that build a batch outside init_new set it themselves
KVWriteLoc.physical	carries the mark to the pool; False unless stated
KVWriteLoc.for_batch(forward_batch, *, swa_loc, full_loc)	the batch's out_cache_loc with the batch's mark. It wraps nothing else, because the mark describes that tensor only
KVWriteLoc.for_layer(forward_batch, layer, …)	the same, except that a cross-attention layer gets encoder_out_cache_loc, which nothing translates and so is never marked
KVCache.requires_physical_write_loc	set by the unified pools. Their write doors raise ValueError on a loc not marked physical. The check is a flag read with no device sync, so it runs on every write-door call. Fused writers that scatter through get_kv_buffer, and a replayed cuda graph, do not pass through a door

A loc produced separately states its own mark with KVWriteLoc(loc, physical=…):
trtllm_mla's capture-stable decode buffer (the runner translates into it), triton's
rank-local DCP slot of a physical loc, and the dspark / dflash draft writes (draft pools are
always static). HybridLinearKVPool and the SWA composites forward the KVWriteLoc to their
sub-pools instead of unwrapping it, so the mark reaches the door that checks it.

Completeness is enforced

test_physical_write_loc.py parses every file that can hand a write loc to a pool under
--enable-unified-memory — the attention backends on the unified allowlists (tied to
kv_cache_hook's allowed sets, so a newly admitted backend joins the scan) plus models/,
layers/cp/, speculative/, model_executor/ and disaggregation/ — and requires the loc
argument of every write-door call to be KVWriteLoc.for_batch / for_layer or a
KVWriteLoc(...) with an explicit physical=. It follows local names, parameters and helper
returns, and a self-test shows it flags a bare tensor and an unmarked KVWriteLoc. The
adoption covers 44 write-door calls, including fa3's two MLA writes, triton's DCP writes and
the draft-model writes.

The pool side is enforced the same way. Producers now hand a KVWriteLoc to every pool,
static ones included, so the same file also requires every set_kv_buffer /
set_mla_kv_buffer override under srt/ to unwrap its loc argument, forward it untouched to
another door, take only *args, or only raise. It covers 30 doors.

When a new backend must wrap

If the backend can receive a unified pool, its write doors pass
KVWriteLoc.for_batch(forward_batch, …), or for_layer when it serves cross-attention. If
it translates the ids itself, it passes KVWriteLoc(loc, physical=True). The ratchet above
fails otherwise.

Static-pool paths the new loc type reaches

Wrapping the loc changes what static pools and backends receive, not just unified ones:

• triton with --dcp-size > 1 compares positions.numel() with the rank-local loc's length;
  it now reads the wrapped tensor (loc.loc.numel()). test_dcp_layout_unit.py drives the
  DCP write with a stub pool.
• MLATokenToKVPoolFP4.set_mla_kv_buffer, HiSparseDSATokenToKVPool's two doors,
  LayerSplitDSATokenToKVPool.set_mla_kv_buffer and
  PageInterleaveMLATokenToKVPool.set_mla_kv_buffer take loc_info and unwrap it, like
  their siblings.
• Hand-built batch stubs that drive a write path carry out_cache_loc_is_physical
  (test_cp_strategy_unit.py, test_fused_fp8_kv_write.py).

On the H200, before the stack was rebased onto current main, each against the base it then had
(0318a8d) (GSM8K 5-shot, n=200):

model	configuration	base	without these fixes	this PR
Qwen3.5-9B	tp2, dcp2, triton	0.855	AttributeError at the first KV write	0.835
DeepSeek-V2-Lite	fp4_mx_block16 KV cache, flashinfer, --max-total-tokens 262144	0.175	AttributeError in graph capture	0.180

DeepSeek-V2-Lite is a base model: with a bf16 KV cache on the same setup it scores
0.325, so the FP4 rows are to be read against each other. The token cap keeps
flashinfer's FP4 decode, which dequantizes the whole pool, inside memory at both commits.

Validation

Model families exercised — a hybrid of every pool shape the unified pool supports, plus a
non-hybrid control. Perf covers the gated-deltanet hybrid (Qwen3.5-9B), the MLA hybrid
(Kimi-Linear-48B-A3B), the sliding-window hybrid (gpt-oss-20b), the tri-pool hybrid
(Inkling and Inkling-Small), and dense Llama-3.1-8B as a control. Accuracy covers the same
set and additionally the Mamba hybrid, Falcon-H1-7B. Each model is run unified against a
baseline on the same model, backend and configuration, so every delta isolates the pool.

Median inter-token latency, unified vs baseline: +1.03% (measured run-to-run noise floor
for this metric: 1.9%). Median GSM8K delta at n=200: +0.50 pt (binomial delta-sigma ≈ 3.0 pt).

Re-validated on the previous base (0318a8d) after review feedback: the stack tip (7/7) gives +0.95% median inter-token latency over 40 paired configurations (median per-configuration shift against the previous run of the same 40: +0.12 pp) and −0.25 pt median GSM8K over 26, with the full unit scope green.

Against the preceding PR's run on the same matrix the median per-model shift is −0.11 pp —
this PR moves nothing.

End to end with the always-on check

Each run below boots with --enable-unified-memory and writes through the checked doors;
none refuses a write (GSM8K 5-shot, n=100). The unified-vs-baseline comparison for these
backends is in #38592.

box	model	backend	GSM8K
H200	Kimi-Linear-48B-A3B, tp2	fa3	0.890
B200	Kimi-Linear-48B-A3B, tp2	trtllm_mla	0.900
B200	Kimi-Linear-48B-A3B, tp2	cutedsl_mla	0.890
B200	Qwen3.5-9B, tp1, cuda graphs off	trtllm_mha	0.900

aiter's write doors are marked, but aiter is not admitted under --enable-unified-memory, so
the ratchet does not scan it, and its doors have not been run on ROCm.

---

CI States

Latest PR Test (Base): ⏳ Run #36944620308
Latest PR Test (Extra): ⏳ Run #36944620104
Latest PR Test (AMD ROCm 10): ⏳ Run #36944620296