Repository navigation
feat(attention): add PagedAttention v1 audit and MLA oracle - #1940
Merged
Merged
Conversation
…T (rev2) Independent rewrite (Leon) of the rejected rev-1 audit, from authoritative ORT v1.29.0 source. Corrects the schema (15 attrs / 17 inputs / 3 outputs; block size power-of-two >=16, not %256; LATENT absorbed MLA; v_head_size<head_size; rotary_offset; head_sink; quantized cache; CUDA+WebGPU, no CPU kernel) and flips the verdict: dense MLA (GLM-5.2 --glm-full-attention, DeepSeek-V2/V3) IS expressible via kv_cache_layout=LATENT. Sparse DSA/IndexShare and CSA/HCA remain inexpressible for a first-principles reason (no query-dependent sparse index). States the exact one-authority invariant on onnx-genai-kv (additive, no second KV manager). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…acle Slice 2 of the corrected Microsoft PagedAttention audit: a standalone, test/reference-only crate that reproduces ONNX Runtime 1.29.0's com.microsoft::PagedAttention v1 acceptance rules and provides a numerically exact CPU oracle for the two dense modes the op supports. Why: the rejected rev1 audit concluded STOP from a stale schema. rev1's successor audit (committed prior) shows the real v1.29.0 schema supports kv_cache_layout="LATENT" absorbed MLA, which makes dense MLA (GLM-5.2 --glm-full-attention, DeepSeek V2/V3) a valid PagedAttention target. This crate turns that claim into executable, checkable evidence. Contents: - validate.rs: faithful port of paged_attention_helper.h::CheckInputs (15 attrs / 17 inputs), returning typed INVALID_ARGUMENT — block_size power-of-two >=16 (not %256), LATENT constraints (kv_num_heads==1, no value_cache, v_head_size<head_size, rotary suffix), the explicit-scale trap, quantized-cache storage/scale contract, sub-byte-dtype rejection. - backend.rs: capability gate returning typed NOT_IMPLEMENTED for any schema-valid mode a kernel does not implement, mirroring the upstream WebGPU EP precedent. Presets: webgpu_separate(), glm_dense_mla_latent(). - oracle.rs: f32 dense paged attention over a block cache addressed by block_table/slot_mapping, for SEPARATE (GQA, RoPE, softcap, local window, head_sink, q/k RMSNorm) and LATENT (single latent row, partial RoPE via rotary_offset, V = leading v_head_size channels). Test/reference only — upstream ships no CPU kernel; this is not upstream CPU support. Tests (tests/equivalence.rs, 14 cases): schema accept/reject matrix; backend typed-rejection for every unsupported mode; SEPARATE GQA oracle == naive contiguous dense with partial RoPE; paging placement-invariance across a permuted block table; slot_mapping=-1 write suppression; and the core knockout — absorbed LATENT == decomposed multi-head latent attention at both tiny and DeepSeek-V3 576/512/64 dims — plus head_sink smooth-softmax and q/k RMSNorm paths. Reuses onnx-genai-kv page/block/slot authorities conceptually (no second KV manager introduced). Added to workspace members only (not default-members). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…ches) Independent review of the rev2 audit + validator returned REQUEST-CHANGES with one Major and one Minor schema-fidelity gap; both are fixed here, plus the audit wording and a backend-preset nit. - S1 (Major): port helper.h:425-428 `num_heads % kv_num_heads == 0` into check_inputs. Without it the validator accepted invalid GQA ratios and the oracle would read an out-of-range kv head. Added rejection test (num_heads=4, kv_num_heads=3) and positive cases (4/2, 4/4). - S2 (Minor): the rotary block trusted a caller-supplied rotary_dim; now port group_query_attention_helper::CheckRotaryCaches faithfully — derive rotary_dim = cos_cache.dims[1]*2 and validate rank-2, head_size % 16 == 0, cos/sin dim1 <= head_size/2 and a multiple of 8, and cos dim1 == sin dim1. Added invalid_rotary_caches_rejected test (three shape violations). Updated rotary-using tests to schema-valid dims (head_size %16, rotary_dim %16). - backend nit: webgpu_separate() now allows packed_qkv — the upstream WebGPU kernel splits packed QKV via RunSplitPackedQKV (verified in contrib_ops/webgpu/bert/paged_attention.cc), so rejecting it was stricter than the precedent it mirrors. - audit: soften "allocates nothing" to "allocates no KV cache / aliases the caller's KV buffers"; note the CUDA EP's transient compute scratch is EP-managed workspace, not a KV-cache authority (the one-authority verdict is unchanged). Noted the omitted device-specific max_threads_per_block guard. - broadened the WebGPU backend-gate test to also exercise local_window, rotary_offset, q/k-norm and attention_metadata rejections. cargo test -p onnx-genai-paged-attention: 16 passed. clippy --all-targets clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…eviewer APPROVE) Records the corrected authoritative schema (15 attrs / 17 inputs / 3 outputs, LATENT absorbed MLA), the one-authority invariant (onnx-genai-kv stays the sole KV manager; the op aliases caller buffers and allocates no cache), the first-principles applicability matrix (dense MLA yes; sparse-selection no), the delivered slices, and the remaining deferred gates (CUDA LATENT kernel, Mobius feature, onnx-genai-kv additive layout, full-size measurement, merge approval). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1940 +/- ##
==========================================
- Coverage 80.95% 80.66% -0.29%
==========================================
Files 415 421 +6
Lines 204440 206390 +1950
Branches 204440 206390 +1950
==========================================
+ Hits 165500 166494 +994
- Misses 33368 34267 +899
- Partials 5572 5629 +57
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
justinchuby
added a commit
that referenced
this pull request
Aug 24, 2026
…1) (#1955) ## Slice 3A.1 — KV-authority half of native PagedAttention integration First independently-reviewable increment of slice 3A. Follows #1940 (merged corrected audit + typed validator/oracle crate `onnx-genai-paged-attention`). Slice 3A as briefed (native CUDA `com.microsoft::PagedAttention` LATENT kernel + KV APIs + full parity/capture matrix + A100 measurement) is larger than one safe, verifiable increment. Per the "land in order / no dead code / each slice independently reviewed" directive, this PR lands **Requirement 1 only**: additive **token-major / LATENT paged-index emission** in `onnx-genai-kv`. It is a strict prerequisite for the CUDA kernel and is **fully validated on host (no GPU)**. ### What it adds (strictly additive, read-only) `crates/onnx-genai-kv/src/paged_index.rs`: - `PagedIndexPlan::build(&PageTable, &[PagedRequest])` emits the exact ORT v1 int32 index tensors from the existing page authority: - `block_table[num_seqs, max_num_blocks_per_seq]` — physical `PageId` per logical block, padded with `PAGED_BLOCK_TABLE_PAD`; - `slot_mapping[token_count]` — `page_id*block_size + offset` per query token; `PAGED_SLOT_EMPTY == -1` documented skip sentinel; - `cumulative_sequence_length[num_seqs+1]` (`cu_seqlens_q`), `past_seqlens`, derived `context_lens`. - `LatentCacheGeometry` + `validate()`, and canonical `token_major_element_offset` / `latent_element_offset` so a CUDA kernel and the CPU oracle index the cache through one formula. `block_size` = power-of-two ≥ 16 (matches `check_kv_cache`). - `PagedKvCache::emit_paged_index_plan` convenience. ### One-authority invariant (enforced in code) `onnx-genai-kv` stays the sole owner of page allocation/lifetime. `paged_index` is a read-only view — it allocates/frees/mutates nothing, no second manager, no op-side allocation. The physical block id **is** the `PageId` (`KvViewKind::VirtuallyContiguous`); a kernel binds these host-emitted indices as stable device inputs and updates caller-owned cache tensors in place. ### Byte-identity No change to `Page`/`PageTable` storage; existing head-major layout untouched. A read-only/leak-free test asserts `materialize_sequence` and pool `usage()`/ `stats()` are identical before and after emission. ### Typed rejections (never silent miscompute) — all tested non-pow2/`<16` block size, windowed/attention-sink (non-contiguous) sequences, query>context, missing backing pages, i32 block-id/slot overflow. ### Tests 3 unit + 12 integration (prefill/decode slot math, exact pow2 boundaries 16/32/64, multi-request row-major+padding, page reuse after free, read-only/ leak-free, every typed rejection). Full crate suite (147+…) and `cargo clippy -p onnx-genai-kv --tests` clean. ### "No dead code" note First production consumer is the native CUDA kernel (slice 3A.2). This mirrors the crate's existing `kv_capacity_bucket`/`ensure_kv_capacity`/ `KvCapacityGrowthBackend` authority seam (exercised by tests, consumed by backends). Its consumer lands in the next, independently-reviewed slice. ### Remaining gates - **3A.2** native CUDA LATENT kernel + claim wiring (design in the decision note); quantized cache → typed NotImplemented; then A100 measurement (n≥3). - **3B** Mobius `--paged-attention` export, only after 3A approval. **Reviewer:** independent, excluding Leon (author). Final approval: Gaff or Roy. **Do not merge without explicit approval.** Draft. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 24, 2026
…MLA (slice 3A.2) (#1978) ## 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 (#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`._ --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
com.microsoft::PagedAttentionschemaDense MLA is supported; GLM DSA/IndexShare and DeepSeek-V4 HCA remain separate sparse paths. This slice adds no production CUDA kernel or Mobius emission.
Validation
-D warningsFollow-up: native CUDA LATENT kernel, additive
onnx-genai-kvlayout/view, and opt-in Mobius export.