[unified-memory] Stride the KV translate kernel and route every translate through it (6/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.

6/7 — stride the KV translate kernel and route every translate through it

Base: #40329 (5/7). 5 commits, 11 files, +361/−159.

Why

main computes v2p[t // ps] * ps + t % ps in two places: the fused Triton kernel
(_translate_loc_fused, used by the read table, the SWA read and the write translate) and a
torch body (translate_kv_loc, used by the composites and the PD / HiCache transfer
translates). They disagree on the ids nobody looks at: the torch body faults on a negative or
past-the-table id and on 2-D input, while the kernel sends them to the page-0 sink. The
kernel could not take every caller, because it flat-addresses its input and so needs it
contiguous, while the SWA read path hands down page_table[:bs, :max_seq_len], a column
slice of the capture-stable buffer.

Commit 1 — walk the kernel by its strides

loc and out are walked as a (rows, cols) view with row and column strides passed per
launch. The decomposition sits behind a STRIDED constexpr, so a contiguous pair still takes
the flat path with no div/mod; that is what the per-forward write translate under cuda-graph
capture uses.

• The strides are runtime arguments, not tl.constexpr. A page table's live width
  changes every few decode steps, and a constexpr stride would recompile with it.
• Offsets are computed in int64. A row stride times a row count can exceed 2^31, and an
  int32 overflow here would mis-address silently.
• out_width stays on the flat path, by assert. That mode writes a packed lane range
  [0, W) with a zero tail; a non-contiguous 1-D loc would send the tail lanes over the
  live prefix, so the assert is load-bearing.
• The out=None allocation no longer mirrors a strided loc, so the caller gets a
  contiguous result.

Commit 2 — one implementation

translate_kv_loc calls the fused kernel, and the torch body and translate_kv_loc_for_kernel
are deleted. For the callers that were on the torch path — the composites' translate_kv_loc,
the PD transfer translates, retraction's get_cpu_copy / load_cpu_copy and
_materialize_inflight_write_set (HiCache's transfer translate was already on the kernel) — a negative or past-the-table id now resolves to the
page-0 sink instead of faulting; 2-D, int32 and strided inputs are accepted; and each call is
one launch instead of up to six. Trading a fault for the sink is deliberate: these calls can
run inside a captured graph, where a device-side assert is not a diagnostic. Paths already on
the fused kernel — the read table, the SWA read, the pool's FULL→SWA — do not change, and
every id the torch body accepted translates to the same value.

Tests

• test_matches_reference_on_a_strided_view and test_translate_accepts_a_strided_page_table
  fail on 5/7 (loc must be contiguous) and pass here.
• test_a_loc_past_the_sentinel_lands_on_the_sink, test_an_out_of_range_loc_lands_on_the_sink
  and test_padding_loc_lands_on_the_sink pin the sink semantics.
• test_out_width_needs_a_packed_loc pins the branch that would otherwise regress silently.
• kernels/ops/memory/test_write_loc_to_kernel_ids_strided.py (GPU) is the only test that
  enters the kernel, so it is the one that can catch a wrong row * rs + col * cs; its oracle
  is the kernel on a contiguous copy, plus the CPU reference. Its lanes cover non-zero
  in-page offsets, the -1 sink, a tombstoned page, a loc below the sentinel and a page past
  the table, in int32 and int64. Dropping the past-the-table guard or the in-page offset from
  the kernel fails it.

Validation

A GPU differential of translate_kv_loc against the closed-form definition matches in all 114
cases: int32 / int64; page sizes 1, 4, 64; 1-D, 2-D, stride-2, column-slice, transposed and
empty inputs; tombstoned pages; −1, −2, −(ps + 1) and past-the-table locs; out=None, packed,
column-slice (neighbours untouched) and in-place out=. The 5/7 implementation agrees on all 50
inputs it accepts. UnifiedSWAKVPool.translate_loc_from_full_to_swa is one CUDA launch at 5/7,
at this PR and at 7/7.

Unit tests

At this head on the H200, the strided kernel test and test_unified_mla_views.py pass (24
tests), and the kernel test fails when the past-the-table guard or the in-page offset is
removed from the kernel. The stack tip (#40331) was run over the wider unit scope on the H200 and
end to end on B200; see its description.

Follow-ups this PR does not do

• Lift the fa3/fa4 local-attention refusal under --enable-unified-memory. Commit 1 removes
  one technical reason for it, but validating the lift needs a local-attention model.
• Fold the Mamba id translate. UnifiedMambaPool.translate is still a plain torch gather; at
  page size 1 it computes the same value but does not clamp, so a tombstoned row returns -1
  rather than the sink.

---

CI States

Latest PR Test (Base): ⏳ Run #36944767639
Latest PR Test (Extra): ⏳ Run #36944767374
Latest PR Test (AMD ROCm 10): ⏳ Run #36944767640