feat(cuda-ep): native PagedAttention LATENT kernel for GLM-5.2 dense MLA (slice 3A.2)

[justinchuby]

Slice 3A.2 — native onnx-genai CUDA support for com.microsoft::PagedAttention LATENT (GLM-5.2 dense MLA)

Builds on the merged audit/validator/oracle (#1940) and the token-major/LATENT KV index emission (#1955). This slice adds the native CUDA kernel for the exact ORT v1 com.microsoft::PagedAttention LATENT (absorbed-MLA) subset that GLM-5.2 dense MLA needs, with onnx-genai-kv remaining the sole page/cache authority.

Depends on / stacks after #1955 (already merged into main).

What this adds

• crates/onnx-runtime-ep-cuda/src/kernels/paged_attention.rs — NVRTC f16/bf16 write + attention kernels implemented from the oracle equations (not copied upstream source): partial-RoPE suffix write into the paged latent cache, online softmax over the latent cache honoring local_window_size/softcap, V taken from the leading v_head_size channels of the same latent row. Plus PagedAttentionFactory, PagedAttentionLatentKernel, and unsupported_reason().
• Five-place op registration in kernels/mod.rs + the unsupported_reason arm in provider.rs::supports_op.

Invariants held

• Default-off / typed subset only. The op is claimed only when the typed geometry+dtype validator proves the exact supported subset (fp16/bf16 LATENT, single latent KV head, GLM qk=192/v=128/partial RoPE, block pow2 ≥ 16). Every unsupported optional mode returns a typed NotImplemented reason rather than silently miscomputing: non-LATENT layout, quantized cache (k/v quant type + int4/float4 cache dtype), head_sink, q/k-norm, k/v scales, present value/value_cache, non-f16/bf16.
• One-authority. No op-side allocation and no second KV manager — the kernel consumes/mutates the caller's page buffers in place. In-place key_cache_out alias contract enforced (non-aliased cache output is rejected).
• Capture/replay safe. Warmed kernel signature; no host sync/allocation during capture. Proven by a capture + 3×replay test that is bit-equal to eager.
• No Mobius changes / no export claims in this slice. No full-size performance claim.

Tests (native CUDA vs onnx-genai-paged-attention oracle, verified on an idle A100)

• Parity: tiny prefill/decode; GLM dims (qk=192/v=128, partial RoPE) first-token/prefill/decode; block pow2≥16 boundary + multi-request; slot -1 skip; no-rotary; eager equivalence (fp16 err ~2e-4, bf16 ~2e-3 — well under tol).
• Capture + 3 replays bit-equal to eager (check_capture_error()==0).
• Rejections: non-aliased key_cache_out and missing required input (block_table).
• Measurement (tiny-shape correctness gate, not a perf claim): CUDA-event timing, n=5, prefill + decode, page/VRAM accounting, op-side alloc = 0. Prefill med ~0.050 ms, decode med ~0.073 ms on tiny GLM shapes.
• 5 pure unit tests cover the typed rejections without a GPU.
• Existing onnx-runtime-ep-cuda lib suite green (559 passed) — no regressions; KV geometry validation (feat(kv): token-major/LATENT PagedAttention index emission (slice 3A.1) #1955 + Gaff's sibling constraints) unchanged.

Reviewer / gates

• Draft — for independent review by Gaff or Roy (final approval required; reviewer excludes Leon/Sapper). Do not merge without explicit approval.
• Remaining gates: full-size GLM checkpoint runs (correctness + measurement) coordinated with the ongoing GLM GGUF/safetensors work; 3B Mobius opt-in --paged-attention export only starts after 3A approval.

GPU tests are gated behind the gpu-tests cargo feature and the CUDA env; run on A100 with --features gpu-tests.

[codecov]

Codecov Report

❌ Patch coverage is 93.93939% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.23%. Comparing base (011fbb2) to head (a52a90a).
⚠️ Report is 21 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/onnx-genai-kv/src/paged_index.rs	93.93%	0 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1978      +/-   ##
==========================================
+ Coverage   80.21%   80.23%   +0.02%     
==========================================
  Files         421      422       +1     
  Lines      203480   203638     +158     
  Branches   203480   203638     +158     
==========================================
+ Hits       163231   163399     +168     
+ Misses      34675    34670       -5     
+ Partials     5574     5569       -5     

Flag	Coverage Δ
cli-ort-linux	72.51% <ø> (?)
cli-ort-windows	72.10% <ø> (+0.09%)	⬆️
offline	80.44% <93.93%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
crates/onnx-genai-kv/src/paged_index.rs	90.06% <93.93%> (+0.25%)	⬆️

... and 8 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[justinchuby]

Independent review round (pre-request, reviewer excluding Leon/Sapper)

An independent code-review pass (read + build, traced every math path against onnx-genai-paged-attention::oracle) returned no high-severity issues and confirmed:

• RoPE (interleaved + non-interleaved), online softmax, write-scatter, V = leading v_head_size, scale/softcap order, token→(batch,pos) mapping, output layout, and shared-mem sizing all match the oracle.
• In-place key_cache_out alias check compares the correct pointers and skips when output[1] is absent.
• Capture-safe: no synchronize/host-alloc/H2D/NVRTC-compile on the execute() path after warm; identical idiom to rotary_embedding.rs.
• Five-place registration complete; rejection set consistent between unsupported_reason() and execute().

Three defensive findings (1 Medium, 1 Low, 1 nit) were fixed in a52a90aa1:

1. (Medium) execute() now cross-checks the attribute-derived geometry (num_heads, head_size, batch) against the actual query/key/cumulative_sequence_length extents — a schema-typed but geometry-inconsistent node now returns a typed KernelFailed instead of an out-of-bounds device read that would tear down the CUDA context.
2. (Low) unsupported_reason() now rejects kv_num_heads != 1 at the supports_op gate (was only enforced at create()/execute()), so a multi-head-KV node is routed elsewhere rather than claimed-then-hard-failed.
3. (nit) sin_cache dtype validated symmetrically with cos_cache.

New tests: unit rejects_multi_head_kv; GPU rejection test extended with a num_heads geometry-mismatch case. Full GPU suite (10) + unit tests (6) green on an idle A100; parity/capture/measurement numbers unchanged.

Still a draft — final approval required from Gaff or Roy (not Leon/Sapper). Remaining gate: full-size GLM checkpoint runs (coordinated with the GLM GGUF/safetensors work); 3B Mobius opt-in export starts only after 3A approval.

[justinchuby]

Independent review verdict — PR #1978 (native CUDA com.microsoft::PagedAttention v1, LATENT subset / slice 3A.2)

APPROVE 3A.2 FOR MERGE

Reviewed read-only at exact HEAD a52a90aa1311de33327604d99e896c0982430b20 (author Leon), against the merged typed validator/CPU oracle (#1940), #1955 PageId plan, and the authoritative ORT v1.29 schema. Verified in an isolated worktree; independently built and ran tests on an idle A100. Requested by Justin Chu.

Gates verified

• Equation-by-equation parity vs the CPU oracle (LATENT path): cache write + partial-RoPE suffix, interleaved & non-interleaved RoPE (cos/sin indexing), online-softmax max/sum ordering (running m/denom/acc with corr = exp(m_old − m_new), corr = 0 when m_old = −inf) — algebraically identical to the oracle's two-pass; scale→softcap order; causal 0..=pos and window (pos − j) > window semantics; slot -1 skip; V = leading v_head_size channels of the latent row; output layout/dtype. Kernel matches oracle exactly.
• supports_op / create / execute agreement: five-place registration consistent (mod.rs pub-mod / CUDA_COVERED_OPS / dtype arm / factory OpKey(...,1); provider.rs unsupported_reason arm). Unsupported routes (quantized cache, separate value cache, head_sink, q/k norm, k/v scales, wrong kv-heads/layout/dtype/options) are rejected before any support claim; no model-name allowlist.
• Extent cross-checks before device access: query hidden = num_heads·head_size, key hidden = head_size, cumulative_seqlens = batch+1, 4-D key_cache, block_size pow2 ≥ 16, v_head_size ∈ 1..=head_size, rotary_offset + rotary_dim ≤ head_size, int32 index dtypes, output dtype, in-place key_cache_out alias contract. No OOB / context poison.
• One allocator/page authority: onnx-genai-kv remains sole allocator/page authority; op-side device alloc = 0 (execute issues only two launches over caller pointers; no H2D/sync/compile in the hot path).
• Capture-safety: NVRTC module cached by key; only the first non-capturing warm compiles/loads under a synchronizing gate; a warmed-signature check rejects capture unless the exact shape/dtype signature was warmed; capture + 3 replays are bit-equal to eager; check_capture_error() == 0.
• ORT v1.29 schema: input order query(0) key(1) value(2) key_cache(3) value_cache(4) cumulative_sequence_length(5) past_seqlens(6) block_table(7) cos_cache(8) sin_cache(9) slot_mapping(10) head_sink(11) q_norm_weight(12) k_norm_weight(13) k_scale(14) v_scale(15) attention_metadata(16); outputs output(0) key_cache_out(1) value_cache_out(2) — matches the kernel IN_* constants exactly.
• Native-cuda build contains the new code: builds under --features cuda,cuda-13000,gpu-tests; no ORT-CUDA confusion.
• Oracle parity numbers (far under tolerance; tol fp16 = 3e-2, bf16 = 1.2e-1): fp16 output err ~2e-4 (GLM first-token exactly 0), bf16 ~2e-3; block-boundary, multi-request, slot-skip, and no-rotary cases all pass. Measurement uses CUDA events on an idle A100, n=5 (prefill med ~0.041 ms / decode ~0.068 ms), op-side alloc = 0, with an explicit "NO full-size performance claim."

Independent test / lint results

• GPU parity/capture/rejection/measurement suite: 10/10 pass (A100, serial, --test-threads=1).
• Unit tests: 6/6 paged_attention:: rejection tests pass.
• cargo fmt --check: clean.
• cargo clippy … --tests: the new/changed files (paged_attention.rs, paged_index.rs) produce zero diagnostics. The one blocking clippy error is the pre-existing PI-approximation literal at optimizer.rs:4811, which this PR does not touch (confirmed: PR changes 7 files; optimizer.rs is not among them). Out of scope.

Merge sequencing

• Merge MAY precede full-size GLM. This is a default-off, typed LATENT subset that is byte-identical when not selected and makes no full-size / performance / export claim; it is a safe, honest seam.
• Next 3A follow-up gate: full-size GLM-5.2 checkpoint correctness + measurement on real artifacts (still owed within 3A).
• 3B gate: Mobius opt-in --paged-attention export — starts only after 3A approval; not part of this slice.

Non-blocking observations

• PR base is 18 commits behind current main (5081ddb2e); none of those 18 commits touch any file in this PR, so it is non-conflicting — recommend a rebase onto latest main before final merge for hygiene.
• block_table per-element value validity (e.g., ring-buffer -1 entries) is trusted from onnx-genai-kv, consistent with the oracle and out of scope for this non-evicting GLM slice.
• Geometry cross-checks fail closed with typed errors in execute rather than in supports_op; the enumerated routing rejections are correctly in supports_op, so this is acceptable defense-in-depth.

Draft PR — not merged. Independent read-only review; requested by Justin Chu.