feat(attention): add PagedAttention v1 audit and MLA oracle

.gitignore

@@ -104,3 +104,4 @@ roy_validate.log
 # hostlock self-test scratch (scripts/hostlock_test.sh writes these into the repo
 # root deliberately -- /tmp is off-limits here -- and an aborted run leaves them)
 .hostlock-selftest*
+.scratch-paged/

.squad/decisions/inbox/leon-paged-attention-audit-rev2.md

@@ -0,0 +1,91 @@
+# Decision: Microsoft PagedAttention — corrected audit (rev2) + typed validator/oracle
+
+- **Author:** Leon (Engine Dev, KV & Buffers)
+- **Date:** 2026-08-24
+- **Supersedes:** Sapper's rejected audit (`squad/paged-attention-audit` @ `0a4726200`,
+  REJECTED by Gaff). Rewritten independently from authoritative ORT v1.29.0 source
+  under the reviewer-protocol lockout (no Sapper input).
+- **Branch:** `squad/paged-attention-audit-rev2` (based on `main`).
+- **Reviewer:** independent code review — **APPROVE** (both findings fixed; 16/16 tests, clippy clean).
+- **Status:** NOT merged (task requires explicit approval).
+
+## Core correction
+
+The rejected audit STOPed on a **stale/wrong schema**. The authoritative ORT
+**1.29.0** `com.microsoft::PagedAttention` v1 schema (`bert_defs.cc:1385-1754`,
+`paged_attention_helper.h`) is **15 attributes / 17 inputs / 3 outputs** and
+supports **`kv_cache_layout="LATENT"` absorbed MLA**. That flips the verdict:
+dense MLA (GLM-5.2 `--glm-full-attention`, DeepSeek V2/V3) **is expressible** as
+PagedAttention. (The brief's "16 attrs" is off by one; source is authoritative.)
+
+Key schema facts (all cited in the audit and ported in `validate.rs`):
+- `block_size` power-of-two **≥ 16** (NOT `%256`).
+- `v_head_size < head_size` allowed **only** in LATENT (DSV3 576/512); explicit
+  `scale` **required** when `v_head_size != head_size` (the "scale trap").
+- `rotary_offset` for a partial-RoPE suffix; `head_sink`; `slot_mapping` (`-1`
+  suppresses a write); quantized cache storage `int8`/`float8e4m3fn` (`int4`/
+  `float4` reserved but rejected by shipped kernels).
+- Providers: **CUDA** (6 typed kernels) + **WebGPU** (fp16, SEPARATE-only, typed
+  `NOT_IMPLEMENTED` for every other mode). **No CPU kernel.**
+
+## One-authority invariant (the constraint that matters to KV & Buffers)
+
+> `onnx-genai-kv` remains the **sole** allocator/owner of KV pages, block tables,
+> slot mapping, sequence lengths and lifetimes. PagedAttention v1 **allocates no
+> KV cache** — it aliases the caller's buffers and mutates them in place
+> (`key_cache_out` must alias `key_cache`). Integration is therefore **additive**
+> on `onnx-genai-kv`: add a **token-major** page layout
+> `[block_size, kv_num_heads, head_size]` (and a **LATENT single-cache** variant
+> of width `head_size`), set `page_size == block_size` (pow2 ≥ 16), and emit
+> device `block_table` / `past_seqlens` / `cumulative_sequence_length` /
+> `slot_mapping` from the page table. **No second KV manager is created.**
+
+A hard blocker would exist only if the op required its **own** cache
+allocation/resize — v1 does not. (CUDA-EP transient compute scratch is
+EP-managed workspace, not a KV-cache authority.)
+
+## Applicability (first-principles, not model-name allowlist)
+
+- **Expressible now:** GLM-5.2 dense MLA (`--glm-full-attention`), DeepSeek V2/V3
+  MLA — via LATENT + narrower `v_head_size` + partial-RoPE `rotary_offset`.
+- **NOT expressible:** GLM DSA / IndexShare and DeepSeek-V4 CSA/HCA
+  sparse-selection — **no query-dependent sparse-index input** exists in the v1
+  schema; `block_table` is per-sequence and the kernel reads the whole prefix.
+  (Custom `pkg.nxrt::{IndexShare,CompressedSparseAttention,SparseKvGather}` ops
+  do exist in this tree, but they are a different, non-PagedAttention path.)
+
+## Delivered (this branch)
+
+1. **Corrected audit** — `docs/models/PAGED_ATTENTION_AUDIT.md` (commit `f4c26c53b`).
+2. **`onnx-genai-paged-attention`** crate (commits `3107beadc`, `3bd5f4b53`),
+   test/reference-only (upstream has no CPU kernel — this is **not** upstream CPU support):
+   - `validate.rs` — faithful `CheckInputs` port → typed `INVALID_ARGUMENT`.
+   - `backend.rs` — capability gate → typed `NOT_IMPLEMENTED` (mirrors the WebGPU
+     EP precedent); presets `webgpu_separate()`, `glm_dense_mla_latent()`.
+   - `oracle.rs` — f32 dense paged attention (SEPARATE GQA + LATENT absorbed MLA).
+   - `tests/equivalence.rs` — 16 tests incl. **absorbed-LATENT == decomposed-MLA**
+     knockout at tiny **and** DeepSeek-V3 576/512/64 dims.
+
+## Proven
+
+- Absorbed MLA folds `W_UK` into the query (`W_UK^T q_nope`) and `W_UV`/`W_O`
+  after the op; `scale = 1/sqrt(qk_nope+qk_rope)`; latent row `[c_kv ; k_rope]`,
+  V = leading `v_head_size` channels. The knockout compares two structurally
+  different code paths and matches — the applicability claim is executable, not
+  asserted.
+
+## Remaining gates (deferred, each an independently reviewed follow-on slice)
+
+- **Slice 3 — CUDA LATENT kernel:** capture/replay-safe, stable buffers, no host
+  sync/alloc in capture; validated against this CPU oracle at GLM dims
+  (qk=192, v=128, partial RoPE), prefill + decode.
+- **Slice 4 — Mobius `--paged-attention` export:** opt-in, GLM-5.2 dense MLA only,
+  default stays DSA/decomposed; gate on **MLA layout properties** (latent width,
+  `v_head_size`, rotary suffix, cache dtype, pow2≥16 block), not a model-name
+  allowlist; keep a decomposed oracle path.
+- **onnx-genai-kv additive change:** token-major (+LATENT) page layout +
+  block_table/slot_mapping view. **No second KV manager.**
+- **Full-size + measurement:** coordinate with GLM GGUF/safetensors work; **no
+  full-size perf claim** before official checkpoint runs on idle A100 (CUDA
+  events, capture/replay, page/VRAM accounting, n≥3). Tiny correctness is only a gate.
+- **Merge:** blocked pending explicit approval.

Cargo.toml

@@ -9,6 +9,7 @@ members = [
     "crates/onnx-genai-comfyui-config",
     "crates/onnx-genai-runtime-config",
     "crates/onnx-genai-kv",
+    "crates/onnx-genai-paged-attention",
     "crates/onnx-genai-scheduler",
     "crates/onnx-genai-engine",
     "crates/onnx-genai-ort",

crates/onnx-genai-paged-attention/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "onnx-genai-paged-attention"
+description = "Typed schema validator and CPU reference oracle for com.microsoft::PagedAttention v1 (ORT 1.29.0)"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+keywords = ["llm", "inference", "attention", "paged", "mla"]
+categories = ["science"]
+
+[dependencies]
+thiserror = { workspace = true }
+
+[lints]
+workspace = true

crates/onnx-genai-paged-attention/src/backend.rs

@@ -0,0 +1,172 @@
+//! Backend capability gate.
+//!
+//! [`check_inputs`](crate::validate::check_inputs) accepts everything the schema
+//! accepts. A concrete kernel implements only a *subset* of those schema-valid
+//! modes, so it must reject the rest with a typed
+//! [`PagedAttentionError::NotImplemented`] — never silently miscompute. This is
+//! exactly what the upstream WebGPU kernel does
+//! (`contrib_ops/webgpu/bert/paged_attention.cc:500-560`), and what our native
+//! CUDA `LATENT` subset does in the other direction.
+//!
+//! Keeping this separate from the schema validator is the `design-discipline`
+//! rule: one question (is this *legal*?) and one mechanism; a second question
+//! (can *this kernel* run it?) and a second mechanism.
+
+use crate::params::{PagedAttentionInputs, PagedAttentionParameters};
+use crate::types::{KvCacheLayout, PagedAttentionError};
+
+/// Declares which schema-valid optional modes a particular kernel implements.
+/// Anything set to `false` (or a layout not in `layouts`) is rejected by
+/// [`check_backend_support`] with a typed `NOT_IMPLEMENTED`.
+#[derive(Debug, Clone)]
+pub struct NativeSubset {
+    pub name: &'static str,
+    pub separate: bool,
+    pub latent: bool,
+    pub softcap: bool,
+    pub local_window: bool,
+    pub narrower_v_head_size: bool,
+    pub rotary_offset: bool,
+    pub head_sink: bool,
+    pub qk_norm: bool,
+    pub quantized_cache: bool,
+    pub slot_mapping: bool,
+    pub attention_metadata: bool,
+    pub packed_qkv: bool,
+}
+
+impl NativeSubset {
+    /// The exact subset the upstream **WebGPU** kernel implements: `SEPARATE`
+    /// only, every optional feature rejected. Packed-QKV *is* supported (the
+    /// kernel splits it via `RunSplitPackedQKV`, `webgpu/bert/paged_attention.cc`),
+    /// so it is allowed here. Encoded so the native side can be diffed against a
+    /// known-good precedent.
+    #[must_use]
+    pub fn webgpu_separate() -> Self {
+        Self {
+            name: "webgpu-separate",
+            separate: true,
+            latent: false,
+            softcap: false,
+            local_window: false,
+            narrower_v_head_size: false,
+            rotary_offset: false,
+            head_sink: false,
+            qk_norm: false,
+            quantized_cache: false,
+            slot_mapping: false,
+            attention_metadata: false,
+            packed_qkv: true,
+        }
+    }
+
+    /// The native first-slice target: GLM-5.2 `--glm-full-attention` dense MLA,
+    /// i.e. the `LATENT` absorbed path with a narrower `v_head_size`, partial
+    /// RoPE (`rotary_offset`), scheduler `slot_mapping`, and host
+    /// `attention_metadata` for CUDA-graph capture. Every other optional mode is
+    /// rejected until a later slice claims it.
+    #[must_use]
+    pub fn glm_dense_mla_latent() -> Self {
+        Self {
+            name: "glm-dense-mla-latent",
+            separate: false,
+            latent: true,
+            softcap: false,
+            local_window: false,
+            narrower_v_head_size: true,
+            rotary_offset: true,
+            head_sink: false,
+            qk_norm: false,
+            quantized_cache: false,
+            slot_mapping: true,
+            attention_metadata: true,
+            packed_qkv: false,
+        }
+    }
+}
+
+/// Reject any schema-valid mode this `subset` does not implement.
+///
+/// `params` must be the output of [`check_inputs`](crate::validate::check_inputs)
+/// (i.e. the node is already schema-valid).
+///
+/// # Errors
+/// Returns [`PagedAttentionError::NotImplemented`] for the first unsupported mode.
+pub fn check_backend_support(
+    params: &PagedAttentionParameters,
+    inputs: &PagedAttentionInputs,
+    subset: &NativeSubset,
+) -> Result<(), PagedAttentionError> {
+    let who = subset.name;
+    let layout = if params.is_latent_kv {
+        KvCacheLayout::Latent
+    } else {
+        KvCacheLayout::Separate
+    };
+    let layout_ok = match layout {
+        KvCacheLayout::Separate => subset.separate,
+        KvCacheLayout::Latent => subset.latent,
+    };
+    if !layout_ok {
+        return Err(PagedAttentionError::unimplemented(format!(
+            "PagedAttention ({who}): kv_cache_layout={} is not supported.",
+            if params.is_latent_kv {
+                "LATENT"
+            } else {
+                "SEPARATE"
+            }
+        )));
+    }
+    if params.is_packed_qkv && !subset.packed_qkv {
+        return Err(PagedAttentionError::unimplemented(format!(
+            "PagedAttention ({who}): packed QKV is not supported."
+        )));
+    }
+    if params.softcap != 0.0 && !subset.softcap {
+        return Err(PagedAttentionError::unimplemented(format!(
+            "PagedAttention ({who}): non-zero softcap is not supported."
+        )));
+    }
+    if params.local_window_size != -1 && !subset.local_window {
+        return Err(PagedAttentionError::unimplemented(format!(
+            "PagedAttention ({who}): local_window_size != -1 is not supported."
+        )));
+    }
+    if params.v_head_size != params.head_size && !subset.narrower_v_head_size {
+        return Err(PagedAttentionError::unimplemented(format!(
+            "PagedAttention ({who}): v_head_size != head_size is not supported."
+        )));
+    }
+    if params.rotary_offset != 0 && !subset.rotary_offset {
+        return Err(PagedAttentionError::unimplemented(format!(
+            "PagedAttention ({who}): rotary_offset != 0 is not supported."
+        )));
+    }
+    if params.use_head_sink && !subset.head_sink {
+        return Err(PagedAttentionError::unimplemented(format!(
+            "PagedAttention ({who}): head_sink input is not supported."
+        )));
+    }
+    if params.use_qk_norm && !subset.qk_norm {
+        return Err(PagedAttentionError::unimplemented(format!(
+            "PagedAttention ({who}): q_norm_weight/k_norm_weight inputs are not supported."
+        )));
+    }
+    let quantized = !params.k_quant_type.is_none() || !params.v_quant_type.is_none();
+    if quantized && !subset.quantized_cache {
+        return Err(PagedAttentionError::unimplemented(format!(
+            "PagedAttention ({who}): quantized KV cache is not supported."
+        )));
+    }
+    if inputs.slot_mapping.is_some() && !subset.slot_mapping {
+        return Err(PagedAttentionError::unimplemented(format!(
+            "PagedAttention ({who}): slot_mapping input is not supported."
+        )));
+    }
+    if inputs.attention_metadata.is_some() && !subset.attention_metadata {
+        return Err(PagedAttentionError::unimplemented(format!(
+            "PagedAttention ({who}): attention_metadata input is not supported."
+        )));
+    }
+    Ok(())
+}

crates/onnx-genai-paged-attention/src/lib.rs

@@ -0,0 +1,43 @@
+//! Typed schema validator and CPU reference oracle for
+//! `com.microsoft::PagedAttention` version 1, as shipped in ONNX Runtime
+//! **1.29.0** (the version this workspace pins — see
+//! `crates/onnx-genai-ort/ort-sys/build.rs`).
+//!
+//! This crate is **test/reference infrastructure**, not a runtime kernel. It
+//! exists to:
+//!
+//! 1. Reproduce, in Rust and from first principles, the exact acceptance rules
+//!    of the authoritative helper
+//!    (`onnxruntime/contrib_ops/cpu/bert/paged_attention_helper.h`) so that a
+//!    native execution provider can validate a `PagedAttention` node and
+//!    **reject every unsupported optional mode with a typed reason** rather than
+//!    silently miscompute (mirroring the upstream WebGPU EP, which implements
+//!    only the `SEPARATE` subset and returns `NOT_IMPLEMENTED` for the rest).
+//! 2. Provide a numerically exact CPU **oracle** for the two dense modes the op
+//!    supports — `SEPARATE` (dense GQA/MHA) and `LATENT` (absorbed Multi-head
+//!    Latent Attention, `v_head_size < head_size`, partial RoPE) — for use as
+//!    the correctness gate of a future CUDA `LATENT` kernel.
+//!
+//! Upstream ORT has **no CPU kernel** for this op; nothing here should be
+//! presented as upstream CPU support.
+//!
+//! ## The two validation layers
+//!
+//! [`validate::check_inputs`] is the *schema* layer: it accepts exactly what the
+//! ORT v1.29.0 helper accepts and returns [`PagedAttentionError::InvalidArgument`]
+//! for the same violations. [`backend::check_backend_support`] is the *backend*
+//! layer: given a [`backend::NativeSubset`] describing which optional features a
+//! particular kernel implements, it returns [`PagedAttentionError::NotImplemented`]
+//! for any schema-valid mode outside that subset. Keeping the two separate is the
+//! `design-discipline` rule — one question, one mechanism.
+
+pub mod backend;
+pub mod oracle;
+pub mod params;
+pub mod types;
+pub mod validate;
+
+pub use backend::{NativeSubset, check_backend_support};
+pub use oracle::{PagedAttentionData, paged_attention_reference};
+pub use params::{PagedAttentionAttributes, PagedAttentionInputs, PagedAttentionParameters, Shape};
+pub use types::{KvCacheDtype, KvCacheLayout, KvQuantType, PagedAttentionError};