From 258a311fc8a25065fb4cd917f950c06fc527c9e8 Mon Sep 17 00:00:00 2001 From: Noah Gift Date: Sat, 19 Sep 2026 11:18:36 +0200 Subject: [PATCH 1/4] =?UTF-8?q?PMAT-3477:=20Gated=20DeltaNet=20with=20num?= =?UTF-8?q?=5Fv=5Fheads=20>=20num=5Fk=5Fheads=20(GQA=20in=20the=20recurren?= =?UTF-8?q?ce)=20on=20the=20CPU=20forward=20=E2=80=94=20Qwen3.5=204B/9B/27?= =?UTF-8?q?B=20load=20and=20run=20(#3346,=20#3510)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pmat-Ticket: PMAT-3477 Co-Authored-By: Claude Fable 5.1 --- .../gguf/inference/forward/forward_qwen35.rs | 150 +++++++- .../forward/forward_qwen35_gqa_tests.rs | 349 ++++++++++++++++++ 2 files changed, 481 insertions(+), 18 deletions(-) create mode 100644 crates/aprender-serve/src/gguf/inference/forward/forward_qwen35_gqa_tests.rs diff --git a/crates/aprender-serve/src/gguf/inference/forward/forward_qwen35.rs b/crates/aprender-serve/src/gguf/inference/forward/forward_qwen35.rs index c4f9ad77da..873ceef04c 100644 --- a/crates/aprender-serve/src/gguf/inference/forward/forward_qwen35.rs +++ b/crates/aprender-serve/src/gguf/inference/forward/forward_qwen35.rs @@ -117,7 +117,12 @@ pub fn causal_conv1d( } } -/// The gated delta-rule recurrence for a single token exactly as delta-net-base.cpp computes it. +/// The gated delta-rule recurrence for a single token exactly as delta-net-base.cpp computes it, +/// for a file whose Gated `DeltaNet` has one key/query head per value head +/// (`num_k_heads == num_v_heads`, `head_k_dim == head_v_dim`) — Qwen3.5-0.8B and -2B. +/// +/// This is [`delta_rule_recurrence_gqa`] at ratio 1; it is kept as its own entry point so the +/// 0.8B call sites and the GPU parity tests read unchanged. pub fn delta_rule_recurrence( q: &[f32], k: &[f32], @@ -129,25 +134,119 @@ pub fn delta_rule_recurrence( num_v_heads: usize, head_v_dim: usize, ) { - assert_eq!(q.len(), num_v_heads * head_v_dim); - assert_eq!(k.len(), num_v_heads * head_v_dim); + delta_rule_recurrence_gqa( + q, + k, + v, + beta, + gate, + state, + output, + num_v_heads, + head_v_dim, + num_v_heads, + head_v_dim, + ); +} + +/// The gated delta-rule recurrence for a single token, with **grouped** key/query heads +/// (PMAT-3477, #3346/#3510). +/// +/// Qwen3.5 4B and 9B carry `linear_num_value_heads = 32` against `linear_num_key_heads = 16`, +/// and 27B carries 48 against 16, so `q` and `k` are narrower than `v`. The reference +/// (HF `modeling_qwen3_next.py::GatedDeltaNet`, llama.cpp `delta-net-base.cpp`) expands q and +/// k with `repeat_interleave(num_v_heads / num_k_heads, dim = heads)`: +/// +/// > **value head `h` reads key/query head `h / (num_v_heads / num_k_heads)`** — grouped, not +/// > strided. With `num_k_heads = 16` and `num_v_heads = 32`, value heads 0 and 1 share key +/// > head 0, value heads 2 and 3 share key head 1, and so on. +/// +/// OPEN, MEASURED (PMAT-3477, 2026-09-19): this mapping makes 4B/9B/27B **load and run** — +/// the shape panic is gone and every logit is finite — but the 4B completion is not yet +/// coherent, and a three-prompt A/B against the other candidate mapping (`h % num_k_heads`, +/// which is what a plain `ggml_repeat` of the head axis would give) favours THAT one: +/// +/// | prompt (4B Q4\_K\_M, `--no-gpu`, 12 tokens) | `h / ratio` (here) | `h % num_k_heads` | +/// |---|---|---| +/// | `The capital of France is` | `META自来看似inglyigitar好莱坞erer` | `总部位于法国本土中央kaçinglyweisehood…` | +/// | `1, 2, 3, 4, 5, 6,` | `抱歉erweiseinglyinglyinglyingly…` | `第七下一个 **.**新加坡和新舊舊…` | +/// | `Water freezes at a temperature of` | `думанlyurahusaticallyALLYyronpan…` | `通常情况下所说的普通普通普通…` | +/// +/// The strided column is on-topic (and "the seventh, the next one" is the right continuation +/// of the counting prompt) and then degenerates; this column is off-topic from token 0. +/// Neither is coherent, so at least one further 4B-specific defect remains and the mapping is +/// NOT settled by that reading alone. What is kept here is the reference semantics — HF +/// `repeat_interleave(dim = heads)` and llama.cpp's `convert_hf_to_gguf.py::Qwen3NextModel`, +/// whose `in_proj_qkvz` de-interleave emits v group-major — because the local llama.cpp +/// (39173bcac, 2026-01-15) has no `qwen35` converter to read, so the file's own v ordering +/// could not be confirmed. Settle it against a reference implementation of `qwen35`, not +/// against another sample of generated text. +/// +/// Shapes (this is the contract the CUDA `DeltaRuleRecurrenceKernel` / `Qwen35CudaModel` +/// mirror — the kernel's `num_v_heads`/`head_v_dim` pair gains `num_k_heads`/`head_k_dim` and +/// the same `h / ratio` read): +/// +/// | argument | length | +/// |---|---| +/// | `q`, `k` | `num_k_heads * head_k_dim` | +/// | `v`, `output` | `num_v_heads * head_v_dim` | +/// | `beta`, `gate` (dt) | `num_v_heads` — every gate is per VALUE head | +/// | `state` | `num_v_heads * head_v_dim * head_k_dim` | +/// +/// The recurrent state of value head `h` is `S ∈ R^(head_k_dim × head_v_dim)` laid out so that +/// `S[i][j] = state[h * head_v_dim * head_k_dim + j * head_k_dim + i]` — memory row `j` is +/// column `j` of `S`, `i` runs over the key dim and `j` over the value dim. Every Qwen3.5 size +/// ships `head_k_dim == head_v_dim == 128`, so the block is square in practice; the two dims +/// are kept distinct here so a future file with a rectangular state is a shape, not a rewrite. +/// +/// The scale is `1/sqrt(head_k_dim)` (the query/key width, as `fla`'s +/// `chunk_gated_delta_rule` defaults it) — identical to the previous +/// `1/sqrt(head_v_dim)` on every file that exists today. +/// +/// # Panics +/// If any slice length disagrees with the table above, or if `num_v_heads` is not a positive +/// multiple of `num_k_heads`. +pub fn delta_rule_recurrence_gqa( + q: &[f32], + k: &[f32], + v: &[f32], + beta: &[f32], + gate: &[f32], + state: &mut [f32], + output: &mut [f32], + num_k_heads: usize, + head_k_dim: usize, + num_v_heads: usize, + head_v_dim: usize, +) { + assert!( + num_k_heads > 0 && num_v_heads % num_k_heads == 0, + "Gated DeltaNet: num_v_heads ({num_v_heads}) must be a positive multiple of num_k_heads \ + ({num_k_heads})" + ); + assert_eq!(q.len(), num_k_heads * head_k_dim); + assert_eq!(k.len(), num_k_heads * head_k_dim); assert_eq!(v.len(), num_v_heads * head_v_dim); assert_eq!(beta.len(), num_v_heads); assert_eq!(gate.len(), num_v_heads); - assert_eq!(state.len(), num_v_heads * head_v_dim * head_v_dim); + assert_eq!(state.len(), num_v_heads * head_v_dim * head_k_dim); assert_eq!(output.len(), num_v_heads * head_v_dim); - let scale = 1.0 / (head_v_dim as f32).sqrt(); + let scale = 1.0 / (head_k_dim as f32).sqrt(); + let heads_per_k = num_v_heads / num_k_heads; for h in 0..num_v_heads { - let q_h = &q[h * head_v_dim..(h + 1) * head_v_dim]; - let k_h = &k[h * head_v_dim..(h + 1) * head_v_dim]; + // repeat_interleave: value head h reads key/query head h / heads_per_k. + let kh = h / heads_per_k; + let q_h = &q[kh * head_k_dim..(kh + 1) * head_k_dim]; + let k_h = &k[kh * head_k_dim..(kh + 1) * head_k_dim]; let v_h = &v[h * head_v_dim..(h + 1) * head_v_dim]; let beta_val = beta[h]; let gate_val = gate[h]; - let state_offset = h * head_v_dim * head_v_dim; - let s_h = &mut state[state_offset..state_offset + head_v_dim * head_v_dim]; + let state_stride = head_v_dim * head_k_dim; + let state_offset = h * state_stride; + let s_h = &mut state[state_offset..state_offset + state_stride]; // 1. S_h *= exp(gate_val) let exp_gate = gate_val.exp(); @@ -156,14 +255,14 @@ pub fn delta_rule_recurrence( } // 2. delta = (v_h - S_h^T * k_h) * beta_val - // Note: s_h[j * head_v_dim + i] is S[i][j]. + // Note: s_h[j * head_k_dim + i] is S[i][j]. // So row j of s_h in memory is column j of S. // sum = dot(row j of s_h, k_h) let mut delta = vec![0.0; head_v_dim]; for j in 0..head_v_dim { - let row_j = &s_h[j * head_v_dim..(j + 1) * head_v_dim]; + let row_j = &s_h[j * head_k_dim..(j + 1) * head_k_dim]; let mut sum = 0.0; - for i in 0..head_v_dim { + for i in 0..head_k_dim { sum += row_j[i] * k_h[i]; } delta[j] = (v_h[j] - sum) * beta_val; @@ -171,18 +270,18 @@ pub fn delta_rule_recurrence( // 3. S_h += k_h * delta^T for j in 0..head_v_dim { - let row_j = &mut s_h[j * head_v_dim..(j + 1) * head_v_dim]; + let row_j = &mut s_h[j * head_k_dim..(j + 1) * head_k_dim]; let d_j = delta[j]; - for i in 0..head_v_dim { + for i in 0..head_k_dim { row_j[i] += k_h[i] * d_j; } } // 4. out_h = S_h^T * q_h * scale for j in 0..head_v_dim { - let row_j = &s_h[j * head_v_dim..(j + 1) * head_v_dim]; + let row_j = &s_h[j * head_k_dim..(j + 1) * head_k_dim]; let mut sum = 0.0; - for i in 0..head_v_dim { + for i in 0..head_k_dim { sum += row_j[i] * q_h[i]; } output[h * head_v_dim + j] = sum * scale; @@ -341,7 +440,10 @@ impl Qwen35State { let convalue_dim = head_k_dim * num_k_heads * 2 + head_value_dim * num_v_heads; Self { conv_states: vec![vec![0.0; convalue_dim * 3]; num_layers], - ssm_states: vec![vec![0.0; num_v_heads * head_value_dim * head_value_dim]; num_layers], + // One [head_k_dim x head_value_dim] recurrent state per VALUE head — the key width + // is the state's row length, which is only the same as the value width because + // every Qwen3.5 size ships head_k_dim == head_v_dim (PMAT-3477). + ssm_states: vec![vec![0.0; num_v_heads * head_value_dim * head_k_dim]; num_layers], kv_cache: crate::gguf::OwnedQuantizedKVCache::new( num_layers, num_kv_heads * head_dim, @@ -926,8 +1028,12 @@ impl<'a> Qwen35Model<'a> { self.base .fused_matmul_into(normed, &d.attn_gate, &mut gate)?; + // q and k are num_k_heads wide, v is num_v_heads wide: on 4B/9B (32 value heads to 16 + // key heads) and 27B (48 to 16) the recurrence shares each key/query head across + // num_v_heads / num_k_heads value heads (PMAT-3477, #3346/#3510). On 0.8B and 2B the + // ratio is 1 and this is the pre-GQA arithmetic, unchanged. let mut out_h = vec![0.0; v_dim]; - delta_rule_recurrence( + delta_rule_recurrence_gqa( &q, &k, &v, @@ -935,6 +1041,8 @@ impl<'a> Qwen35Model<'a> { &dt, &mut cache.ssm_states[il][..], &mut out_h, + self.num_k_heads, + self.head_k_dim, self.num_v_heads, self.head_v_dim, ); @@ -1648,3 +1756,9 @@ mod qwen35_math_tests { #[cfg(test)] #[path = "forward_qwen35_contract_tests.rs"] mod qhf_contract_tests; + +/// The Gated `DeltaNet` head mapping when `num_v_heads > num_k_heads` +/// (Qwen3.5 4B/9B/27B) — PMAT-3477, #3346/#3510. +#[cfg(test)] +#[path = "forward_qwen35_gqa_tests.rs"] +mod qwen35_gqa_tests; diff --git a/crates/aprender-serve/src/gguf/inference/forward/forward_qwen35_gqa_tests.rs b/crates/aprender-serve/src/gguf/inference/forward/forward_qwen35_gqa_tests.rs new file mode 100644 index 0000000000..1b6e880708 --- /dev/null +++ b/crates/aprender-serve/src/gguf/inference/forward/forward_qwen35_gqa_tests.rs @@ -0,0 +1,349 @@ +//! PMAT-3477 / aprender#3346, #3510: the Gated `DeltaNet` recurrence when +//! `num_v_heads > num_k_heads` (grouped-query attention *inside* the recurrence). +//! +//! Qwen3.5-0.8B and -2B have `num_k_heads == num_v_heads`, so the whole hybrid +//! stack was written as if q, k and v always had the same head count. 4B and 9B +//! carry `linear_num_value_heads = 32` against `linear_num_key_heads = 16`, and +//! 27B carries 48 against 16, so `q`/`k` are half (or a third) as wide as `v` and +//! the recurrence panicked on its own `assert_eq!(q.len(), num_v_heads * +//! head_v_dim)` before a single token was produced. +//! +//! The reference semantics (HF `modeling_qwen3_next.py::GatedDeltaNet`, llama.cpp +//! `delta-net-base.cpp`) are a `repeat_interleave` on the head axis: value head +//! `h` reads key/query head `h / (num_v_heads / num_k_heads)`. The first test +//! below pins exactly that, against a host reference that literally expands q and +//! k and then calls the old `num_k_heads == num_v_heads` code path — so the +//! mapping is fixed independently of any model file, and independently of the +//! generalised implementation's own arithmetic. + +use super::{delta_rule_recurrence, delta_rule_recurrence_gqa}; + +/// A deterministic, dependency-free value stream: an LCG mapped into [-1, 1). +/// Fixtures must be reproducible byte for byte — two runs of the same test have +/// to compare the same numbers, and the GPU worker mirroring this mapping needs +/// to be able to regenerate them. +struct Lcg(u64); + +impl Lcg { + fn new(seed: u64) -> Self { + Self(seed) + } + fn next_f32(&mut self) -> f32 { + self.0 = self + .0 + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + let bits = (self.0 >> 33) as u32; + (f32::from(u16::try_from(bits & 0xFFFF).unwrap_or(0)) / 32768.0) - 1.0 + } + fn vec(&mut self, n: usize) -> Vec { + (0..n).map(|_| self.next_f32()).collect() + } +} + +/// `repeat_interleave(x, ratio)` on the head axis: head `h` of the output is head +/// `h / ratio` of the input. This is the host reference for the mapping — it is +/// the definition, not a re-derivation of the implementation. +fn repeat_interleave_heads(x: &[f32], head_dim: usize, ratio: usize) -> Vec { + let num_heads = x.len() / head_dim; + let mut out = Vec::with_capacity(num_heads * ratio * head_dim); + for h in 0..num_heads * ratio { + let src = h / ratio; + out.extend_from_slice(&x[src * head_dim..(src + 1) * head_dim]); + } + out +} + +/// The same expansion under the WRONG (strided / `h % num_k_heads`) mapping. Used +/// only to prove the fixture discriminates: if the two expansions produced the +/// same recurrence output, the test above would pass for a broken implementation. +fn strided_heads(x: &[f32], head_dim: usize, num_k_heads: usize, num_v_heads: usize) -> Vec { + let mut out = Vec::with_capacity(num_v_heads * head_dim); + for h in 0..num_v_heads { + let src = h % num_k_heads; + out.extend_from_slice(&x[src * head_dim..(src + 1) * head_dim]); + } + out +} + +/// GQA in the recurrence: with `num_k_heads = 2`, `num_v_heads = 4`, `D = 8`, the +/// generalised recurrence must equal the old `num_k_heads == num_v_heads` code +/// path fed `repeat_interleave`d q and k — on the OUTPUT and on the recurrent +/// STATE the token leaves behind, which is what the next token reads. +/// +/// Held at exact equality: the generalised loop performs the same additions in +/// the same order on the same values, so any difference is a semantic one. +#[test] +fn qwen35_gqa_recurrence_is_repeat_interleave_of_q_and_k() { + let (num_k_heads, num_v_heads, d) = (2usize, 4usize, 8usize); + let ratio = num_v_heads / num_k_heads; + + let mut rng = Lcg::new(0x3477); + let q = rng.vec(num_k_heads * d); + let k = rng.vec(num_k_heads * d); + let v = rng.vec(num_v_heads * d); + let beta = rng.vec(num_v_heads); + let gate: Vec = rng.vec(num_v_heads).iter().map(|g| g * 0.1).collect(); + let state0 = rng.vec(num_v_heads * d * d); + + let mut state_gqa = state0.clone(); + let mut out_gqa = vec![0.0f32; num_v_heads * d]; + delta_rule_recurrence_gqa( + &q, + &k, + &v, + &beta, + &gate, + &mut state_gqa, + &mut out_gqa, + num_k_heads, + d, + num_v_heads, + d, + ); + + let q_exp = repeat_interleave_heads(&q, d, ratio); + let k_exp = repeat_interleave_heads(&k, d, ratio); + let mut state_ref = state0.clone(); + let mut out_ref = vec![0.0f32; num_v_heads * d]; + delta_rule_recurrence( + &q_exp, + &k_exp, + &v, + &beta, + &gate, + &mut state_ref, + &mut out_ref, + num_v_heads, + d, + ); + + assert_eq!( + out_gqa, out_ref, + "the GQA recurrence output is not the repeat_interleave reference" + ); + assert_eq!( + state_gqa, state_ref, + "the GQA recurrence left a different recurrent state than the repeat_interleave reference" + ); + + // Anti-vacuity: the strided mapping (h % num_k_heads), the other obvious + // reading of "share the key heads", must give a DIFFERENT answer — otherwise + // the assertions above would hold for an implementation that got the mapping + // backwards. + let q_strided = strided_heads(&q, d, num_k_heads, num_v_heads); + let k_strided = strided_heads(&k, d, num_k_heads, num_v_heads); + let mut state_strided = state0; + let mut out_strided = vec![0.0f32; num_v_heads * d]; + delta_rule_recurrence( + &q_strided, + &k_strided, + &v, + &beta, + &gate, + &mut state_strided, + &mut out_strided, + num_v_heads, + d, + ); + assert_ne!( + out_strided, out_ref, + "the fixture does not discriminate: the strided head mapping gives the same output as \ + repeat_interleave, so this test would pass for a wrong mapping" + ); +} + +/// `ratio == 1` (0.8B and 2B: `num_k_heads == num_v_heads == 16`) must be BYTE +/// IDENTICAL to the pre-GQA recurrence. The GPU parity tests on the real 0.8B +/// file are the end-to-end regression guard; this is the unit-level one. +#[test] +fn qwen35_gqa_ratio_one_is_byte_identical_to_the_legacy_recurrence() { + let (num_heads, d) = (3usize, 4usize); + let mut rng = Lcg::new(0x0808); + let q = rng.vec(num_heads * d); + let k = rng.vec(num_heads * d); + let v = rng.vec(num_heads * d); + let beta = rng.vec(num_heads); + let gate: Vec = rng.vec(num_heads).iter().map(|g| g * 0.1).collect(); + let state0 = rng.vec(num_heads * d * d); + + let mut state_new = state0.clone(); + let mut out_new = vec![0.0f32; num_heads * d]; + delta_rule_recurrence_gqa( + &q, + &k, + &v, + &beta, + &gate, + &mut state_new, + &mut out_new, + num_heads, + d, + num_heads, + d, + ); + + let mut state_old = state0; + let mut out_old = vec![0.0f32; num_heads * d]; + delta_rule_recurrence( + &q, + &k, + &v, + &beta, + &gate, + &mut state_old, + &mut out_old, + num_heads, + d, + ); + + assert_eq!(out_new, out_old, "ratio 1 changed the output bit pattern"); + assert_eq!(state_new, state_old, "ratio 1 changed the recurrent state"); +} + +/// A rectangular state (`head_k_dim != head_v_dim`) must be indexed as +/// `S[i][j] = state[h * head_v_dim * head_k_dim + j * head_k_dim + i]`, `i` over +/// the key dim and `j` over the value dim. Every Qwen3.5 size ships +/// `head_k_dim == head_v_dim == 128`, so nothing in the fleet exercises the two +/// dims being distinct — which is exactly why it is pinned here rather than left +/// to a future file to discover. The oracle is arithmetic done by hand on a +/// one-head, `k_dim = 2`, `v_dim = 3` case with an identity-shaped state. +#[test] +fn qwen35_gqa_state_is_rectangular_k_dim_by_v_dim() { + let (head_k_dim, head_v_dim) = (2usize, 3usize); + let q = [1.0f32, 2.0]; + let k = [0.5f32, 0.5]; + let v = [1.0f32, -1.0, 2.0]; + let beta = [0.5f32]; + let gate = [0.0f32]; // exp(0) = 1 + // S is [2 x 3] (i over k_dim, j over v_dim); memory row j is column j of S. + let mut state = [ + 1.0f32, 0.0, // j = 0: S[0][0], S[1][0] + 0.0, 1.0, // j = 1 + 1.0, 1.0, // j = 2 + ]; + let mut output = [0.0f32; 3]; + delta_rule_recurrence_gqa( + &q, + &k, + &v, + &beta, + &gate, + &mut state, + &mut output, + 1, + head_k_dim, + 1, + head_v_dim, + ); + + // S^T k = [0.5, 0.5, 1.0]; delta = (v - S^T k) * 0.5 = [0.25, -0.75, 0.5] + // S += k delta^T -> row j gets k * delta[j]: + // j=0: [1.125, 0.125] j=1: [-0.375, 0.625] j=2: [1.25, 1.25] + // out[j] = dot(row j, q) / sqrt(head_k_dim) + let s2 = 2.0f32.sqrt(); + let want = [ + (1.125 + 2.0 * 0.125) / s2, + (-0.375 + 2.0 * 0.625) / s2, + (1.25 + 2.0 * 1.25) / s2, + ]; + for (j, w) in want.iter().enumerate() { + assert!( + (output[j] - w).abs() < 1e-6, + "output[{j}] = {}, want {w} (state: {state:?})", + output[j] + ); + } + assert!((state[0] - 1.125).abs() < 1e-6, "{state:?}"); + assert!((state[1] - 0.125).abs() < 1e-6, "{state:?}"); + assert!((state[4] - 1.25).abs() < 1e-6, "{state:?}"); + assert!((state[5] - 1.25).abs() < 1e-6, "{state:?}"); +} + +/// The real 4B file (`linear_num_value_heads = 32`, `linear_num_key_heads = 16`): +/// load it through the production path and run three tokens of a fixed prompt +/// through `forward_single_qwen35`. Before this ticket this panicked in the +/// recurrence's shape assertion at the first `DeltaNet` layer of the first token. +/// +/// Asserts the loader's own head shape (a `num_v_heads == num_k_heads` file would +/// make the test vacuous), that every logit is finite, and that the argmax is a +/// real vocabulary entry. The argmax is PRINTED, not asserted against a golden: +/// pinning a golden token belongs with a tokenizer-level e2e test, and the +/// acceptance command (`apr run --no-gpu`) is what reads the text. +#[test] +fn qwen35_gqa_4b_real_file_runs_three_tokens() { + const MODEL_PATH: &str = "/home/noah/models/Qwen3.5-4B-Q4_K_M.gguf"; + if !std::path::Path::new(MODEL_PATH).exists() { + eprintln!("SKIP: {MODEL_PATH} is absent"); + return; + } + let mapped = crate::gguf::MappedGGUFModel::from_path(MODEL_PATH).expect("map the 4B GGUF"); + let base = + super::Qwen35Model::create_base_model(&mapped.model, mapped.data()).expect("4B base model"); + let qwen = super::Qwen35Model::from_model_and_layers(&base, &mapped.model, mapped.data()) + .expect("4B hybrid layers"); + + assert!( + qwen.num_v_heads > qwen.num_k_heads, + "this file is not a GQA DeltaNet file (num_v_heads {} vs num_k_heads {}), so it cannot \ + falsify the head mapping", + qwen.num_v_heads, + qwen.num_k_heads + ); + assert_eq!( + qwen.num_v_heads % qwen.num_k_heads, + 0, + "num_v_heads {} is not a multiple of num_k_heads {}", + qwen.num_v_heads, + qwen.num_k_heads + ); + + // The per-head value width the gated RMSNorm weight actually carries must be + // the head_v_dim the recurrence and the state are sized with: a mismatch here + // is a silently mis-shaped forward, not a crash. + if let Some(super::Qwen35OwnedLayer::DeltaNet(d)) = qwen + .layers + .iter() + .find(|l| matches!(l, super::Qwen35OwnedLayer::DeltaNet(_))) + { + assert_eq!( + d.ssm_norm_weight.len(), + qwen.head_v_dim, + "ssm_norm_weight is {} wide but head_v_dim is {}", + d.ssm_norm_weight.len(), + qwen.head_v_dim + ); + assert_eq!( + d.ssm_a.len(), + qwen.num_v_heads, + "ssm_a (A) is per value head" + ); + } + + // "The capital of France is" under the Qwen BPE vocabulary. No golden is + // asserted on the output, so a drifted id costs nothing but a less + // interesting print. + let prompt = [785u32, 6722, 315, 9625, 374]; + let mut state = qwen.new_state(prompt.len() + 4); + let mut logits = Vec::new(); + for (pos, &token) in prompt.iter().take(3).enumerate() { + logits = qwen + .forward_single_qwen35(token, &mut state, pos) + .expect("4B forward"); + assert_eq!(logits.len(), qwen.base.config.vocab_size); + assert!( + logits.iter().all(|l| l.is_finite()), + "position {pos}: the 4B forward produced a non-finite logit" + ); + } + let argmax = crate::gguf::ops::argmax(&logits); + assert!( + (argmax as usize) < qwen.base.config.vocab_size, + "argmax {argmax} is outside the vocabulary" + ); + println!( + "qwen35 4B (num_k_heads {}, num_v_heads {}, head_k_dim {}, head_v_dim {}): argmax after \ + 3 tokens = {argmax}", + qwen.num_k_heads, qwen.num_v_heads, qwen.head_k_dim, qwen.head_v_dim + ); +} From 01236e3e78275dd73bbdc4ba5df298209a9ed02e Mon Sep 17 00:00:00 2001 From: Noah Gift Date: Sat, 19 Sep 2026 12:57:06 +0200 Subject: [PATCH 2/4] =?UTF-8?q?PMAT-3477:=20Qwen3.5=204B/9B=20coherent=20o?= =?UTF-8?q?n=20the=20CPU=20=E2=80=94=20load=5Fqwen35=5Flayers=20reads=20{a?= =?UTF-8?q?rch}.block=5Fcount=20(a=20bare=20key=20never=20matched,=20every?= =?UTF-8?q?=20file=20got=2024=20layers)=20and=20the=20Gated=20DeltaNet=20G?= =?UTF-8?q?QA=20mapping=20is=20the=20ggml=20tiled=20order=20llama.cpp=20co?= =?UTF-8?q?nverts=20to;=204B=20greedy=20tokens=20pinned=20to=20llama.cpp?= =?UTF-8?q?=20(#3346,=20#3510)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pmat-Ticket: PMAT-3477 Co-Authored-By: Claude Fable 5.1 --- .../gguf/inference/forward/forward_qwen35.rs | 64 +++-- .../forward/forward_qwen35_gqa_tests.rs | 260 +++++++++++++----- crates/aprender-serve/src/gguf/qwen35_load.rs | 84 +++++- 3 files changed, 303 insertions(+), 105 deletions(-) diff --git a/crates/aprender-serve/src/gguf/inference/forward/forward_qwen35.rs b/crates/aprender-serve/src/gguf/inference/forward/forward_qwen35.rs index 873ceef04c..f46e6eb5b0 100644 --- a/crates/aprender-serve/src/gguf/inference/forward/forward_qwen35.rs +++ b/crates/aprender-serve/src/gguf/inference/forward/forward_qwen35.rs @@ -153,38 +153,46 @@ pub fn delta_rule_recurrence( /// (PMAT-3477, #3346/#3510). /// /// Qwen3.5 4B and 9B carry `linear_num_value_heads = 32` against `linear_num_key_heads = 16`, -/// and 27B carries 48 against 16, so `q` and `k` are narrower than `v`. The reference -/// (HF `modeling_qwen3_next.py::GatedDeltaNet`, llama.cpp `delta-net-base.cpp`) expands q and -/// k with `repeat_interleave(num_v_heads / num_k_heads, dim = heads)`: +/// and 27B carries 48 against 16, so `q` and `k` are narrower than `v`. A **GGUF** file shares +/// them TILED: /// -/// > **value head `h` reads key/query head `h / (num_v_heads / num_k_heads)`** — grouped, not -/// > strided. With `num_k_heads = 16` and `num_v_heads = 32`, value heads 0 and 1 share key -/// > head 0, value heads 2 and 3 share key head 1, and so on. +/// > **value head `h` reads key/query head `h % num_k_heads`.** With `num_k_heads = 16` and +/// > `num_v_heads = 32`, value heads 0 and 16 share key head 0, value heads 1 and 17 share key +/// > head 1, and so on. /// -/// OPEN, MEASURED (PMAT-3477, 2026-09-19): this mapping makes 4B/9B/27B **load and run** — -/// the shape panic is gone and every logit is finite — but the 4B completion is not yet -/// coherent, and a three-prompt A/B against the other candidate mapping (`h % num_k_heads`, -/// which is what a plain `ggml_repeat` of the head axis would give) favours THAT one: +/// This is not the HF checkpoint's own order, and that is the trap. HF +/// (`modeling_qwen3_next.py::GatedDeltaNet`) stores the value heads grouped by key head and +/// expands q/k with `repeat_interleave` (`h / ratio`); the **conversion permutes the value +/// heads out of that order**, and everything downstream is written for the permuted file: /// -/// | prompt (4B Q4\_K\_M, `--no-gpu`, 12 tokens) | `h / ratio` (here) | `h % num_k_heads` | -/// |---|---|---| -/// | `The capital of France is` | `META自来看似inglyigitar好莱坞erer` | `总部位于法国本土中央kaçinglyweisehood…` | -/// | `1, 2, 3, 4, 5, 6,` | `抱歉erweiseinglyinglyinglyingly…` | `第七下一个 **.**新加坡和新舊舊…` | -/// | `Water freezes at a temperature of` | `думанlyurahusaticallyALLYyronpan…` | `通常情况下所说的普通普通普通…` | +/// * `llama.cpp/conversion/qwen.py:455-464` — `_LinearAttentionVReorderBase`, which +/// `Qwen3_5TextModel` (line 639, `MODEL_ARCH.QWEN35`) is built from: *"reorders V heads from +/// grouped to tiled order for ggml broadcast … The HF weights store V heads grouped by K +/// head … ggml binary ops use tiled broadcast … We reorder V heads to tiled order so +/// `ggml_repeat` can replace the expensive interleaved repeat"*. `modify_tensors` +/// (lines 568-613) permutes **every** v-head-indexed tensor together — the v rows of +/// `in_proj_qkv`, `in_proj_z` (the gate), `in_proj_a`/`in_proj_b` (dt and beta), +/// `A_log`/`dt_bias`, the v channels of `conv1d`, and the v columns of `out_proj` — so the +/// loader needs no compensating permutation anywhere: only this index changes. +/// * `llama.cpp/src/models/qwen35.cpp:436-441` — the graph expands q and k with +/// `ggml_repeat_4d`, which tiles. +/// * `llama.cpp/ggml/src/ggml-cpu/ops.cpp:10976-10977` — the fused kernel that skips that +/// repeat reads `iq1 = iv1 % neq1; ik1 = iv1 % nek1;`. /// -/// The strided column is on-topic (and "the seventh, the next one" is the right continuation -/// of the counting prompt) and then degenerates; this column is off-topic from token 0. -/// Neither is coherent, so at least one further 4B-specific defect remains and the mapping is -/// NOT settled by that reading alone. What is kept here is the reference semantics — HF -/// `repeat_interleave(dim = heads)` and llama.cpp's `convert_hf_to_gguf.py::Qwen3NextModel`, -/// whose `in_proj_qkvz` de-interleave emits v group-major — because the local llama.cpp -/// (39173bcac, 2026-01-15) has no `qwen35` converter to read, so the file's own v ordering -/// could not be confirmed. Settle it against a reference implementation of `qwen35`, not -/// against another sample of generated text. +/// MEASURED, so that nobody re-opens this from a text sample: with this mapping the 4B CPU +/// forward agrees with `llama-eval-callback` (`-ub 1`, autoregressive graph, same token +/// stream) on **every** `DeltaNet` intermediate of layer 0 at position 0 — `attn_norm` +/// -26.4532 vs -26.4532, `q_conv_predelta` 1.4485 vs 1.4614, `k_conv_predelta` 13.3249 vs +/// 13.3100, `gate` -18.2186 vs -18.2173, `beta_sigmoid` 21.0785 vs 21.0794, `attn_output` +/// 0.9502 vs 0.9535, `new_state` 38.7674 vs 38.7736, `l_out` 0.2917 vs 0.3056 (tensor sums; +/// the residual is the Q4_K/Q5_K activation quantisation, which is ~0.3% on a 8192-wide +/// sum). The *incoherent* 4B/9B text that survived this mapping was never a DeltaNet defect +/// at all — [`crate::gguf::qwen35_load::load_qwen35_layers`] was reading a bare +/// `block_count` key that never matched, so it built 24 layers for a 32-block file. /// /// Shapes (this is the contract the CUDA `DeltaRuleRecurrenceKernel` / `Qwen35CudaModel` /// mirror — the kernel's `num_v_heads`/`head_v_dim` pair gains `num_k_heads`/`head_k_dim` and -/// the same `h / ratio` read): +/// the same `h % num_k_heads` read): /// /// | argument | length | /// |---|---| @@ -233,11 +241,11 @@ pub fn delta_rule_recurrence_gqa( assert_eq!(output.len(), num_v_heads * head_v_dim); let scale = 1.0 / (head_k_dim as f32).sqrt(); - let heads_per_k = num_v_heads / num_k_heads; for h in 0..num_v_heads { - // repeat_interleave: value head h reads key/query head h / heads_per_k. - let kh = h / heads_per_k; + // ggml_repeat (tiled): value head h reads key/query head h % num_k_heads. The GGUF + // conversion already permuted the value heads into this order — see the doc comment. + let kh = h % num_k_heads; let q_h = &q[kh * head_k_dim..(kh + 1) * head_k_dim]; let k_h = &k[kh * head_k_dim..(kh + 1) * head_k_dim]; let v_h = &v[h * head_v_dim..(h + 1) * head_v_dim]; diff --git a/crates/aprender-serve/src/gguf/inference/forward/forward_qwen35_gqa_tests.rs b/crates/aprender-serve/src/gguf/inference/forward/forward_qwen35_gqa_tests.rs index 1b6e880708..04b96b2a22 100644 --- a/crates/aprender-serve/src/gguf/inference/forward/forward_qwen35_gqa_tests.rs +++ b/crates/aprender-serve/src/gguf/inference/forward/forward_qwen35_gqa_tests.rs @@ -8,13 +8,31 @@ //! the recurrence panicked on its own `assert_eq!(q.len(), num_v_heads * //! head_v_dim)` before a single token was produced. //! -//! The reference semantics (HF `modeling_qwen3_next.py::GatedDeltaNet`, llama.cpp -//! `delta-net-base.cpp`) are a `repeat_interleave` on the head axis: value head -//! `h` reads key/query head `h / (num_v_heads / num_k_heads)`. The first test -//! below pins exactly that, against a host reference that literally expands q and -//! k and then calls the old `num_k_heads == num_v_heads` code path — so the -//! mapping is fixed independently of any model file, and independently of the -//! generalised implementation's own arithmetic. +//! The mapping a **GGUF** file carries is `ggml`'s TILED broadcast: value head +//! `h` reads key/query head `h % num_k_heads`. Source, in the order it settles +//! the question: +//! +//! * `llama.cpp/conversion/qwen.py:455-464` — `_LinearAttentionVReorderBase`, +//! the class `Qwen3_5TextModel` (line 639) is built from: *"reorders V heads +//! from grouped to tiled order for ggml broadcast … The HF weights store V +//! heads grouped by K head: `[G0_v0..v{r-1}, G1_v0..v{r-1}, …]`. ggml binary +//! ops use tiled broadcast … We reorder V heads to tiled order"*. Every +//! v-head-indexed tensor is permuted together (the v rows of `in_proj_qkv`, +//! `in_proj_z`, `in_proj_a`/`_b`, `A_log`/`dt_bias`, the v channels of +//! `conv1d`, the v columns of `out_proj` — `modify_tensors`, lines 568-613). +//! * `llama.cpp/src/models/qwen35.cpp:436-441` — when `num_k_heads != +//! num_v_heads` the graph expands q and k with `ggml_repeat_4d`, which tiles +//! (`h % num_k_heads`), not `repeat_interleave`. +//! * `llama.cpp/ggml/src/ggml-cpu/ops.cpp:10976-10977` — the fused kernel that +//! skips that repeat indexes `iq1 = iv1 % neq1; ik1 = iv1 % nek1;`. +//! +//! So HF's `repeat_interleave` (`h / ratio`) is the right reading of the HF +//! *checkpoint*, and the wrong reading of the *file we load*: the conversion has +//! already permuted the value heads into tiled order. The first test below pins +//! the tiled mapping against a host reference that literally expands q and k and +//! then calls the `num_k_heads == num_v_heads` code path — so the mapping is +//! fixed independently of any model file, and independently of the generalised +//! implementation's own arithmetic. use super::{delta_rule_recurrence, delta_rule_recurrence_gqa}; @@ -41,40 +59,44 @@ impl Lcg { } } -/// `repeat_interleave(x, ratio)` on the head axis: head `h` of the output is head -/// `h / ratio` of the input. This is the host reference for the mapping — it is -/// the definition, not a re-derivation of the implementation. -fn repeat_interleave_heads(x: &[f32], head_dim: usize, ratio: usize) -> Vec { - let num_heads = x.len() / head_dim; - let mut out = Vec::with_capacity(num_heads * ratio * head_dim); - for h in 0..num_heads * ratio { - let src = h / ratio; +/// `ggml_repeat` on the head axis: head `h` of the output is head +/// `h % num_k_heads` of the input. This is the host reference for the mapping — +/// it is the definition (`conversion/qwen.py` writes the value heads in exactly +/// this order so that `ggml_repeat_4d` is correct), not a re-derivation of the +/// implementation. +fn tiled_heads(x: &[f32], head_dim: usize, num_k_heads: usize, num_v_heads: usize) -> Vec { + let mut out = Vec::with_capacity(num_v_heads * head_dim); + for h in 0..num_v_heads { + let src = h % num_k_heads; out.extend_from_slice(&x[src * head_dim..(src + 1) * head_dim]); } out } -/// The same expansion under the WRONG (strided / `h % num_k_heads`) mapping. Used -/// only to prove the fixture discriminates: if the two expansions produced the -/// same recurrence output, the test above would pass for a broken implementation. -fn strided_heads(x: &[f32], head_dim: usize, num_k_heads: usize, num_v_heads: usize) -> Vec { - let mut out = Vec::with_capacity(num_v_heads * head_dim); - for h in 0..num_v_heads { - let src = h % num_k_heads; +/// The same expansion under the HF-checkpoint (`repeat_interleave` / `h / ratio`) +/// mapping, which a GGUF file does NOT carry. Used only to prove the fixture +/// discriminates: if the two expansions produced the same recurrence output, the +/// test above would pass for a broken implementation. +fn repeat_interleave_heads(x: &[f32], head_dim: usize, ratio: usize) -> Vec { + let num_heads = x.len() / head_dim; + let mut out = Vec::with_capacity(num_heads * ratio * head_dim); + for h in 0..num_heads * ratio { + let src = h / ratio; out.extend_from_slice(&x[src * head_dim..(src + 1) * head_dim]); } out } /// GQA in the recurrence: with `num_k_heads = 2`, `num_v_heads = 4`, `D = 8`, the -/// generalised recurrence must equal the old `num_k_heads == num_v_heads` code -/// path fed `repeat_interleave`d q and k — on the OUTPUT and on the recurrent -/// STATE the token leaves behind, which is what the next token reads. +/// generalised recurrence must equal the `num_k_heads == num_v_heads` code path +/// fed TILED q and k (`ggml_repeat`, `h % num_k_heads` — see the module header +/// for the three source lines) — on the OUTPUT and on the recurrent STATE the +/// token leaves behind, which is what the next token reads. /// /// Held at exact equality: the generalised loop performs the same additions in /// the same order on the same values, so any difference is a semantic one. #[test] -fn qwen35_gqa_recurrence_is_repeat_interleave_of_q_and_k() { +fn qwen35_gqa_recurrence_tiles_q_and_k_over_the_value_heads() { let (num_k_heads, num_v_heads, d) = (2usize, 4usize, 8usize); let ratio = num_v_heads / num_k_heads; @@ -102,8 +124,8 @@ fn qwen35_gqa_recurrence_is_repeat_interleave_of_q_and_k() { d, ); - let q_exp = repeat_interleave_heads(&q, d, ratio); - let k_exp = repeat_interleave_heads(&k, d, ratio); + let q_exp = tiled_heads(&q, d, num_k_heads, num_v_heads); + let k_exp = tiled_heads(&k, d, num_k_heads, num_v_heads); let mut state_ref = state0.clone(); let mut out_ref = vec![0.0f32; num_v_heads * d]; delta_rule_recurrence( @@ -120,36 +142,37 @@ fn qwen35_gqa_recurrence_is_repeat_interleave_of_q_and_k() { assert_eq!( out_gqa, out_ref, - "the GQA recurrence output is not the repeat_interleave reference" + "the GQA recurrence output is not the tiled (ggml_repeat) reference" ); assert_eq!( state_gqa, state_ref, - "the GQA recurrence left a different recurrent state than the repeat_interleave reference" + "the GQA recurrence left a different recurrent state than the tiled reference" ); - // Anti-vacuity: the strided mapping (h % num_k_heads), the other obvious - // reading of "share the key heads", must give a DIFFERENT answer — otherwise - // the assertions above would hold for an implementation that got the mapping + // Anti-vacuity: the HF-checkpoint mapping (repeat_interleave, h / ratio) — + // the other obvious reading of "share the key heads", and the one the + // conversion permutes AWAY — must give a DIFFERENT answer, otherwise the + // assertions above would hold for an implementation that got the mapping // backwards. - let q_strided = strided_heads(&q, d, num_k_heads, num_v_heads); - let k_strided = strided_heads(&k, d, num_k_heads, num_v_heads); - let mut state_strided = state0; - let mut out_strided = vec![0.0f32; num_v_heads * d]; + let q_interleaved = repeat_interleave_heads(&q, d, ratio); + let k_interleaved = repeat_interleave_heads(&k, d, ratio); + let mut state_interleaved = state0; + let mut out_interleaved = vec![0.0f32; num_v_heads * d]; delta_rule_recurrence( - &q_strided, - &k_strided, + &q_interleaved, + &k_interleaved, &v, &beta, &gate, - &mut state_strided, - &mut out_strided, + &mut state_interleaved, + &mut out_interleaved, num_v_heads, d, ); assert_ne!( - out_strided, out_ref, - "the fixture does not discriminate: the strided head mapping gives the same output as \ - repeat_interleave, so this test would pass for a wrong mapping" + out_interleaved, out_ref, + "the fixture does not discriminate: repeat_interleave gives the same output as the tiled \ + mapping, so this test would pass for a wrong mapping" ); } @@ -261,17 +284,31 @@ fn qwen35_gqa_state_is_rectangular_k_dim_by_v_dim() { } /// The real 4B file (`linear_num_value_heads = 32`, `linear_num_key_heads = 16`): -/// load it through the production path and run three tokens of a fixed prompt -/// through `forward_single_qwen35`. Before this ticket this panicked in the -/// recurrence's shape assertion at the first `DeltaNet` layer of the first token. +/// load it through the production path and take FOUR greedy tokens of a fixed +/// prompt through `forward_single_qwen35`, against the tokens llama.cpp produces +/// for the same raw prompt. Before this ticket this panicked in the recurrence's +/// shape assertion at the first `DeltaNet` layer of the first token. /// -/// Asserts the loader's own head shape (a `num_v_heads == num_k_heads` file would -/// make the test vacuous), that every logit is finite, and that the argmax is a -/// real vocabulary entry. The argmax is PRINTED, not asserted against a golden: -/// pinning a golden token belongs with a tokenizer-level e2e test, and the -/// acceptance command (`apr run --no-gpu`) is what reads the text. +/// The reference (llama.cpp master `60b06ab9a`, CPU build, this box): +/// +/// ```text +/// llama-completion -m /home/noah/models/Qwen3.5-4B-Q4_K_M.gguf \ +/// -p "The capital of France is" -n 8 --temp 0 -ngl 0 -no-cnv --seed 0 -t 16 +/// The capital of France is Paris. +/// A. True +/// B +/// ``` +/// +/// `-no-cnv` is load-bearing: without it llama.cpp applies the model's chat +/// template, and the comparison is then against a different prompt. (So does +/// `apr run`, which is why a CLI-level A/B against this reference compares two +/// different prompts and cannot settle anything.) The eight continuation ids +/// under this file's own vocabulary (248320 entries — NOT the Qwen2/Qwen3 +/// vocabulary; `The` is 760 here and 785 there) are `[11751 " Paris", 13 ".", +/// 198 "\n", 32 "A", 13 ".", 2912 " True", 198 "\n", 33 "B"]`, and all eight are +/// asserted. #[test] -fn qwen35_gqa_4b_real_file_runs_three_tokens() { +fn qwen35_gqa_4b_matches_llama_cpp_greedy_tokens() { const MODEL_PATH: &str = "/home/noah/models/Qwen3.5-4B-Q4_K_M.gguf"; if !std::path::Path::new(MODEL_PATH).exists() { eprintln!("SKIP: {MODEL_PATH} is absent"); @@ -320,13 +357,26 @@ fn qwen35_gqa_4b_real_file_runs_three_tokens() { ); } - // "The capital of France is" under the Qwen BPE vocabulary. No golden is - // asserted on the output, so a drifted id costs nothing but a less - // interesting print. - let prompt = [785u32, 6722, 315, 9625, 374]; - let mut state = qwen.new_state(prompt.len() + 4); + // The loader must take the depth from the file, not from a default: 4B is 32 blocks, + // and a silently-24-deep stack is exactly what produced finite logits and drifting text. + assert_eq!( + qwen.layers.len(), + mapped + .model + .num_layers() + .expect("4B GGUF carries qwen35.block_count"), + "the hybrid loader built {} layers for a {:?}-block file", + qwen.layers.len(), + mapped.model.num_layers() + ); + + // "The capital of France is" under THIS file's vocabulary. + let prompt = [760u32, 6511, 314, 9338, 369]; + const WANT: [u32; 8] = [11751, 13, 198, 32, 13, 2912, 198, 33]; + + let mut state = qwen.new_state(prompt.len() + WANT.len()); let mut logits = Vec::new(); - for (pos, &token) in prompt.iter().take(3).enumerate() { + for (pos, &token) in prompt.iter().enumerate() { logits = qwen .forward_single_qwen35(token, &mut state, pos) .expect("4B forward"); @@ -336,14 +386,90 @@ fn qwen35_gqa_4b_real_file_runs_three_tokens() { "position {pos}: the 4B forward produced a non-finite logit" ); } - let argmax = crate::gguf::ops::argmax(&logits); - assert!( - (argmax as usize) < qwen.base.config.vocab_size, - "argmax {argmax} is outside the vocabulary" + + let mut got = Vec::with_capacity(WANT.len()); + for step in 0..WANT.len() { + let next = crate::gguf::ops::argmax(&logits); + got.push(next); + logits = qwen + .forward_single_qwen35(next, &mut state, prompt.len() + step) + .expect("4B forward"); + } + + assert_eq!( + got, + WANT.to_vec(), + "the 4B CPU forward does not follow llama.cpp greedily (head mapping / Gated DeltaNet \ + arithmetic): got {got:?}, want {:?} — num_k_heads {}, num_v_heads {}, head_k_dim {}, \ + head_v_dim {}", + WANT, + qwen.num_k_heads, + qwen.num_v_heads, + qwen.head_k_dim, + qwen.head_v_dim ); - println!( - "qwen35 4B (num_k_heads {}, num_v_heads {}, head_k_dim {}, head_v_dim {}): argmax after \ - 3 tokens = {argmax}", - qwen.num_k_heads, qwen.num_v_heads, qwen.head_k_dim, qwen.head_v_dim +} + +/// CONTROL: the same comparison on 0.8B (`num_k_heads == num_v_heads == 16`), a +/// file the CPU forward already served before this ticket. It shares every line +/// of the hybrid stack with 4B except the head mapping, so a RED here would mean +/// the 4B failure is not a GQA defect at all. +/// +/// Reference (llama.cpp master `60b06ab9a`, CPU build, this box): +/// +/// ```text +/// llama-completion -m /home/noah/models/Qwen3.5-0.8B-Q4_K_M.gguf \ +/// -p "The capital of France is" -n 8 --temp 0 -ngl 0 -no-cnv --seed 0 -t 16 +/// The capital of France is the capital of the country. +/// The +/// ``` +/// +/// Only the first FOUR ids are asserted, and deliberately so: this prompt is a +/// near-tie on 0.8B and the reference itself moved. A January llama.cpp build on +/// another box answers `[7172 " located", 303 " in", 279 " the", 9514 " south"]` +/// for the same file and the same prompt, where master answers `[279 " the", +/// 6511 " capital", 314 " of", 279 " the"]`. A RED here is therefore only +/// evidence of a defect when it is reproduced against the llama.cpp build named +/// above; the load-bearing 0.8B guard is the CUDA `forward_single` parity test, +/// which compares apr to apr and cannot drift with an upstream build. +#[test] +fn qwen35_0_8b_matches_llama_cpp_greedy_tokens() { + const MODEL_PATH: &str = "/home/noah/models/Qwen3.5-0.8B-Q4_K_M.gguf"; + if !std::path::Path::new(MODEL_PATH).exists() { + eprintln!("SKIP: {MODEL_PATH} is absent"); + return; + } + let mapped = crate::gguf::MappedGGUFModel::from_path(MODEL_PATH).expect("map the 0.8B GGUF"); + let base = super::Qwen35Model::create_base_model(&mapped.model, mapped.data()) + .expect("0.8B base model"); + let qwen = super::Qwen35Model::from_model_and_layers(&base, &mapped.model, mapped.data()) + .expect("0.8B hybrid layers"); + assert_eq!( + qwen.num_v_heads, qwen.num_k_heads, + "0.8B is supposed to be the ratio-1 control" + ); + + let prompt = [760u32, 6511, 314, 9338, 369]; + const WANT: [u32; 4] = [279, 6511, 314, 279]; // " the" " capital" " of" " the" + + let mut state = qwen.new_state(prompt.len() + WANT.len()); + let mut logits = Vec::new(); + for (pos, &token) in prompt.iter().enumerate() { + logits = qwen + .forward_single_qwen35(token, &mut state, pos) + .expect("0.8B forward"); + } + let mut got = Vec::with_capacity(WANT.len()); + for step in 0..WANT.len() { + let next = crate::gguf::ops::argmax(&logits); + got.push(next); + logits = qwen + .forward_single_qwen35(next, &mut state, prompt.len() + step) + .expect("0.8B forward"); + } + assert_eq!( + got, + WANT.to_vec(), + "the 0.8B CPU forward does not follow llama.cpp greedily: got {got:?}" ); } diff --git a/crates/aprender-serve/src/gguf/qwen35_load.rs b/crates/aprender-serve/src/gguf/qwen35_load.rs index 671ffc9f33..8f724d7e80 100644 --- a/crates/aprender-serve/src/gguf/qwen35_load.rs +++ b/crates/aprender-serve/src/gguf/qwen35_load.rs @@ -51,19 +51,40 @@ fn as_u32(v: &GGUFValue) -> Option { } } -pub fn load_qwen35_layers(model: &GGUFModel, data: &[u8]) -> Result> { - let mut layers = Vec::new(); - let num_layers = model +/// Read `{arch}.{suffix}` — GGUF metadata keys are architecture-prefixed +/// (`qwen35.block_count`, never a bare `block_count`). +/// +/// PMAT-3477 (#3346/#3510): this loader used to read the BARE keys and fall back to +/// `24` layers / interval `4`. Both lookups always missed, so every Qwen3.5 file was +/// built with exactly 24 layers. 0.8B and 2B *have* 24 blocks, so they were right by +/// coincidence and the default looked correct; 4B and 9B have 32 and 27B has 48, so +/// their last 8 (resp. 24) blocks were silently dropped and `lm_head` ran on a +/// three-quarters-deep residual stream. That is the whole "near-coherent drift" +/// symptom — the shapes are all consistent, the logits are all finite, and the model +/// is simply not finished. The head mapping was never the second defect. +fn arch_u32(model: &GGUFModel, suffix: &str) -> Option { + let arch = model.architecture()?; + model .metadata - .get("block_count") + .get(&crate::gguf::keys::arch_key(arch, suffix)) .and_then(as_u32) - .unwrap_or(24) as usize; +} - let interval = model - .metadata - .get("full_attention_interval") - .and_then(as_u32) - .unwrap_or(4) as usize; +/// # Errors +/// A missing `{arch}.block_count`, or a layer tensor the file does not carry. +pub fn load_qwen35_layers(model: &GGUFModel, data: &[u8]) -> Result> { + let mut layers = Vec::new(); + // Fail closed: a guessed depth is a silently truncated model, not a degraded one. + let num_layers = + arch_u32(model, "block_count").ok_or_else(|| crate::error::RealizarError::InvalidShape { + reason: format!( + "Qwen3.5 loader: {}.block_count is missing from the GGUF metadata; refusing to \ + guess the layer count", + model.architecture().unwrap_or("") + ), + })? as usize; + + let interval = arch_u32(model, "full_attention_interval").unwrap_or(4) as usize; for i in 0..num_layers { if (i + 1) % interval == 0 { @@ -250,4 +271,47 @@ mod tests { assert_eq!(num_delta_net, 18); assert_eq!(num_full_attn, 6); } + + /// PMAT-3477 (#3346/#3510): the depth must come from `{arch}.block_count`, not from a + /// default. + /// + /// `test_qwen35_load_fixture` above cannot see this defect: 0.8B (and 2B) have exactly + /// 24 blocks, which is the number the old bare-key lookup fell back to, so it passed + /// while every deeper file was silently truncated. 4B declares 32 and 9B/27B more, so + /// this is the first case where the fallback and the file disagree. Mutate + /// `arch_u32(model, "block_count")` back to `model.metadata.get("block_count")` and this + /// goes RED at 24 != 32; the 0.8B test stays green. + #[test] + fn qwen35_layer_count_is_read_from_the_arch_prefixed_block_count() { + const MODEL_PATH: &str = "/home/noah/models/Qwen3.5-4B-Q4_K_M.gguf"; + if !std::path::Path::new(MODEL_PATH).exists() { + eprintln!("SKIP: {MODEL_PATH} is absent"); + return; + } + let data = std::fs::read(MODEL_PATH).expect("read the 4B GGUF"); + let model = GGUFModel::from_bytes(&data).expect("parse the 4B GGUF"); + + let declared = model + .num_layers() + .expect("qwen35.block_count is present in the 4B file"); + assert_ne!( + declared, 24, + "this file no longer falsifies the 24-layer fallback" + ); + + let layers = load_qwen35_layers(&model, &data).expect("extract the 4B layers"); + assert_eq!( + layers.len(), + declared, + "the loader built {} layers for a {declared}-block file", + layers.len() + ); + + // The hybrid schedule is (i + 1) % interval, so the last block is a full-attention + // one: a truncated stack would also end on the wrong kind. + assert!( + matches!(layers.last(), Some(Qwen35Layer::Attention(_))), + "the last block of a 32-block, interval-4 file is a full-attention block" + ); + } } From c58062517b7cf73f4ddabfca7158de11568341c8 Mon Sep 17 00:00:00 2001 From: Noah Gift Date: Sat, 19 Sep 2026 13:19:39 +0200 Subject: [PATCH 3/4] =?UTF-8?q?PMAT-3477:=20Gated=20DeltaNet=20GQA=20on=20?= =?UTF-8?q?CUDA=20=E2=80=94=20value=20heads=20read=20key=20head=20h=20%=20?= =?UTF-8?q?num=5Fk=5Fheads;=20Qwen3.5-4B/9B/27B=20on=20the=20GPU=20(#3090,?= =?UTF-8?q?=20#3346,=20#3510)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pmat-Ticket: PMAT-3477 Co-Authored-By: Claude Fable 5.1 --- .../src/kernels/gdn/delta_rule.rs | 324 ++++++++++---- .../src/cuda/executor/gdn_ops.rs | 24 +- crates/aprender-serve/src/cuda/generate.rs | 5 +- crates/aprender-serve/src/cuda/kernel_type.rs | 11 +- .../src/cuda/kernels_generate_gemm_cuda.rs | 5 +- .../src/gguf/cuda/forward_qwen35_cuda.rs | 40 +- .../gguf/cuda/forward_qwen35_cuda_tests.rs | 415 +++++++++++++++++- 7 files changed, 725 insertions(+), 99 deletions(-) diff --git a/crates/aprender-gpu/src/kernels/gdn/delta_rule.rs b/crates/aprender-gpu/src/kernels/gdn/delta_rule.rs index 10ebae3c0c..96022dc215 100644 --- a/crates/aprender-gpu/src/kernels/gdn/delta_rule.rs +++ b/crates/aprender-gpu/src/kernels/gdn/delta_rule.rs @@ -1,22 +1,32 @@ //! PMAT-3477: the gated delta-rule recurrence, one token. //! -//! `delta_rule_recurrence` in the CPU reference, per value head `h` with the state -//! `S_h` stored transposed (`s_h[j * D + i] == S[i][j]`, so memory row `j` is column -//! `j` of `S`): +//! `delta_rule_recurrence_gqa` in the CPU reference, per **value** head `h` with the +//! state `S_h` stored transposed (`s_h[j * Dk + i] == S[i][j]`, so memory row `j` is +//! column `j` of `S`, `i` runs over `head_k_dim` and `j` over `head_v_dim`): //! //! ```text -//! 1. s_h *= exp(gate[h]) -//! 2. delta[j] = (v[j] - sum_i s_h[j*D+i] * k[i]) * beta[h] -//! 3. s_h[j*D+i] += k[i] * delta[j] -//! 4. out[j] = (sum_i s_h[j*D+i] * q[i]) * D^-0.5 +//! kh = h % num_k_heads (the key/query head this value head reads) +//! 1. s_h *= exp(gate[h]) +//! 2. delta[j] = (v[j] - sum_i s_h[j*Dk+i] * k[kh*Dk+i]) * beta[h] +//! 3. s_h[j*Dk+i] += k[kh*Dk+i] * delta[j] +//! 4. out[j] = (sum_i s_h[j*Dk+i] * q[kh*Dk+i]) * Dk^-0.5 //! ``` //! +//! `num_v_heads` may exceed `num_k_heads` (Qwen3.5-4B/9B: 32 against 16; 27B: 48 +//! against 16). The mapping is `h % num_k_heads`, **tiled**, not `h / ratio`: the GGUF +//! conversion (`_LinearAttentionVReorderBase`) already permutes the value heads out of +//! HF's grouped order so that llama.cpp can expand q/k with `ggml_repeat_4d`. The CPU +//! reference (`gguf/inference/forward/forward_qwen35.rs`) carries the citations and the +//! measurement; this kernel only mirrors it. +//! //! Grid: `(num_v_heads, 1, 1)`, Block: `(head_v_dim, 1, 1)`. Thread `j` owns memory //! row `j` of `S_h` **exclusively**: every one of steps 1–4 touches only row `j` for //! output `j`, so the whole recurrence runs with no barrier and no cross-thread //! dependency. Steps 1 and 2 are fused into one ascending pass over `i` and steps 3 //! and 4 into a second, which keeps the fp32 accumulation order identical to the CPU -//! loop (`i = 0..D`) — the per-layer L∞ ≤ 1e-3 parity contract depends on it. +//! loop (`i = 0..head_k_dim`) — the per-layer L∞ ≤ 1e-3 parity contract depends on it. +//! Two value heads that share a key head still touch disjoint state and output, so the +//! grouped case needs no more synchronisation than the symmetric one. use crate::kernels::gdn::emit_exp_f32; use crate::kernels::Kernel; @@ -25,20 +35,32 @@ use crate::ptx::{PtxKernel, PtxReg, PtxType}; /// Gated delta-rule recurrence for a single token. /// -/// For Qwen3.5-0.8B: `num_v_heads = 16`, `head_v_dim = 128`. +/// For Qwen3.5-0.8B and -2B: `num_k_heads = num_v_heads = 16`, +/// `head_k_dim = head_v_dim = 128`. For -4B and -9B: `num_k_heads = 16`, +/// `num_v_heads = 32`. For -27B: `num_k_heads = 16`, `num_v_heads = 48`. #[derive(Debug, Clone, Copy)] pub struct DeltaRuleRecurrenceKernel { - /// Number of value heads. + /// Number of key/query heads — `q` and `k` are `num_k_heads * head_k_dim` long. + pub num_k_heads: u32, + /// Key/query head width `Dk`, the state's row length and the recurrence's scale. + pub head_k_dim: u32, + /// Number of value heads — one block each. pub num_v_heads: u32, - /// Value head width `D` (also the state's row and column count). + /// Value head width `Dv`, the state's row count and the output width per head. pub head_v_dim: u32, } impl DeltaRuleRecurrenceKernel { /// Create the kernel. + /// + /// `num_v_heads` is expected to be a positive multiple of `num_k_heads`, exactly as + /// the CPU reference asserts; a `num_k_heads` of zero is clamped to one so the PTX + /// never emits a `rem.u32` by zero. #[must_use] - pub const fn new(num_v_heads: u32, head_v_dim: u32) -> Self { + pub const fn new(num_k_heads: u32, head_k_dim: u32, num_v_heads: u32, head_v_dim: u32) -> Self { Self { + num_k_heads: if num_k_heads == 0 { 1 } else { num_k_heads }, + head_k_dim, num_v_heads, head_v_dim, } @@ -63,26 +85,29 @@ impl Kernel for DeltaRuleRecurrenceKernel { } fn build_ptx(&self) -> PtxKernel { - let d = self.head_v_dim; - // The CPU computes `1.0 / (head_v_dim as f32).sqrt()` once; the same value is + let dk = self.head_k_dim; + let dv = self.head_v_dim; + let nk = self.num_k_heads; + let nv = self.num_v_heads; + // The CPU computes `1.0 / (head_k_dim as f32).sqrt()` once; the same value is // folded in here as an immediate so no rsqrt approximation enters the output. - let scale = 1.0 / (d as f32).sqrt(); + let scale = 1.0 / (dk as f32).sqrt(); PtxKernel::new(self.name()) - .param(PtxType::U64, "q_ptr") // [num_v_heads * D] - .param(PtxType::U64, "k_ptr") // [num_v_heads * D] - .param(PtxType::U64, "v_ptr") // [num_v_heads * D] + .param(PtxType::U64, "q_ptr") // [num_k_heads * Dk] + .param(PtxType::U64, "k_ptr") // [num_k_heads * Dk] + .param(PtxType::U64, "v_ptr") // [num_v_heads * Dv] .param(PtxType::U64, "beta_ptr") // [num_v_heads] .param(PtxType::U64, "gate_ptr") // [num_v_heads] (dt) - .param(PtxType::U64, "state_ptr") // [num_v_heads * D * D], updated in place - .param(PtxType::U64, "output_ptr") // [num_v_heads * D] + .param(PtxType::U64, "state_ptr") // [num_v_heads * Dv * Dk], updated in place + .param(PtxType::U64, "output_ptr") // [num_v_heads * Dv] .shared_memory(0) .build(|ctx| { let j = ctx.special_reg(PtxReg::TidX); let h = ctx.special_reg(PtxReg::CtaIdX); - let d_r = ctx.mov_u32_imm(d); - let row_in_bounds = ctx.setp_lt_u32(j, d_r); + let dv_r = ctx.mov_u32_imm(dv); + let row_in_bounds = ctx.setp_lt_u32(j, dv_r); ctx.branch_if_not(row_in_bounds, "gdn_dr_exit"); let q_ptr = ctx.load_param_u64("q_ptr"); @@ -94,20 +119,28 @@ impl Kernel for DeltaRuleRecurrenceKernel { let output_ptr = ctx.load_param_u64("output_ptr"); let four = ctx.mov_u32_imm(4); - let d_bytes = ctx.mov_u32_imm(d * 4); - - // q/k/v/out head base: h * D * 4 - let head_off = ctx.mul_wide_u32_reg(h, d_bytes); - let q_base = ctx.add_u64(q_ptr, head_off); - let k_base = ctx.add_u64(k_ptr, head_off); - let v_base = ctx.add_u64(v_ptr, head_off); - let out_base = ctx.add_u64(output_ptr, head_off); - - // state row base: (h * D * D + j * D) * 4 - let state_head_bytes = ctx.mov_u32_imm(d * d * 4); + let dk_r = ctx.mov_u32_imm(dk); + let dk_bytes = ctx.mov_u32_imm(dk * 4); + let dv_bytes = ctx.mov_u32_imm(dv * 4); + + // q/k head base: (h % num_k_heads) * Dk * 4. With nk == nv the block + // index IS the key head, so the `rem` is not emitted at all and the + // symmetric PTX is unchanged. + let kh = if nk == nv { h } else { ctx.rem_u32(h, nk) }; + let qk_head_off = ctx.mul_wide_u32_reg(kh, dk_bytes); + let q_base = ctx.add_u64(q_ptr, qk_head_off); + let k_base = ctx.add_u64(k_ptr, qk_head_off); + + // v/out head base: h * Dv * 4 + let v_head_off = ctx.mul_wide_u32_reg(h, dv_bytes); + let v_base = ctx.add_u64(v_ptr, v_head_off); + let out_base = ctx.add_u64(output_ptr, v_head_off); + + // state row base: (h * Dv * Dk + j * Dk) * 4 + let state_head_bytes = ctx.mov_u32_imm(dv * dk * 4); let state_head_off = ctx.mul_wide_u32_reg(h, state_head_bytes); let state_head = ctx.add_u64(state_ptr, state_head_off); - let row_off = ctx.mul_wide_u32_reg(j, d_bytes); + let row_off = ctx.mul_wide_u32_reg(j, dk_bytes); let s_row = ctx.add_u64(state_head, row_off); // Per-head scalars, read once. @@ -122,7 +155,7 @@ impl Kernel for DeltaRuleRecurrenceKernel { let sum = ctx.mov_f32_imm(0.0); let i = ctx.mov_u32_imm(0); ctx.label("gdn_dr_decay_loop"); - let go = ctx.setp_lt_u32(i, d_r); + let go = ctx.setp_lt_u32(i, dk_r); ctx.branch_if_not(go, "gdn_dr_decay_end"); let off = ctx.mul_wide_u32_reg(i, four); let s_addr = ctx.add_u64(s_row, off); @@ -148,7 +181,7 @@ impl Kernel for DeltaRuleRecurrenceKernel { let out_sum = ctx.mov_f32_imm(0.0); let i2 = ctx.mov_u32_imm(0); ctx.label("gdn_dr_update_loop"); - let go2 = ctx.setp_lt_u32(i2, d_r); + let go2 = ctx.setp_lt_u32(i2, dk_r); ctx.branch_if_not(go2, "gdn_dr_update_end"); let off2 = ctx.mul_wide_u32_reg(i2, four); let s_addr2 = ctx.add_u64(s_row, off2); @@ -183,18 +216,59 @@ mod ptx_tests { #[test] fn gdn_delta_rule_ptx_shape() { - let kernel = DeltaRuleRecurrenceKernel::new(16, 128); + let kernel = DeltaRuleRecurrenceKernel::new(16, 128, 16, 128); let ptx = kernel.emit_ptx(); assert!(ptx.contains(".entry gdn_delta_rule_recurrence"), "{ptx}"); // exp(gate) is the only transcendental in the recurrence. assert_eq!(ptx.matches("ex2.approx.f32").count(), 1, "{ptx}"); // No barrier: thread j owns state row j exclusively. assert!(!ptx.contains("bar.sync"), "{ptx}"); - // D^-0.5 is a host-computed immediate, not an rsqrt. + // Dk^-0.5 is a host-computed immediate, not an rsqrt. assert!(!ptx.contains("rsqrt"), "{ptx}"); + // Symmetric heads: the block index IS the key head, so no modulo is emitted. + assert!(!ptx.contains("rem.u32"), "{ptx}"); assert_eq!(kernel.grid(), (16, 1, 1)); assert_eq!(kernel.block(), (128, 1, 1)); } + + /// Qwen3.5-4B/9B (`nk = 16`, `nv = 32`) and -27B (`nk = 16`, `nv = 48`): the block + /// count follows the VALUE heads and the key head is `h % num_k_heads`. + #[test] + fn gdn_delta_rule_ptx_grouped_heads_emit_the_tiled_modulo() { + for (nk, nv) in [(16u32, 32u32), (16, 48)] { + let kernel = DeltaRuleRecurrenceKernel::new(nk, 128, nv, 128); + let ptx = kernel.emit_ptx(); + assert_eq!(kernel.grid(), (nv, 1, 1), "one block per VALUE head"); + assert_eq!(kernel.block(), (128, 1, 1)); + assert_eq!( + ptx.matches("rem.u32").count(), + 1, + "the key head must be h % {nk}: {ptx}" + ); + assert!( + ptx.contains(&format!("{nk};")) || ptx.contains(&format!("{nk} ")), + "{ptx}" + ); + } + } + + /// A rectangular state (`Dk != Dv`) sizes the row length from the KEY dim and the + /// row count from the VALUE dim. + #[test] + fn gdn_delta_rule_rectangular_state_uses_dk_for_the_row() { + let kernel = DeltaRuleRecurrenceKernel::new(2, 8, 4, 6); + assert_eq!(kernel.grid(), (4, 1, 1)); + assert_eq!(kernel.block(), (6, 1, 1), "one thread per state ROW (Dv)"); + let ptx = kernel.emit_ptx(); + // The per-head state block is Dv * Dk * 4 = 6 * 8 * 4 = 192 bytes. + assert!(ptx.contains("192"), "{ptx}"); + } + + /// `num_k_heads = 0` would be a `rem.u32` by zero; it is clamped instead. + #[test] + fn gdn_delta_rule_zero_key_heads_is_clamped() { + assert_eq!(DeltaRuleRecurrenceKernel::new(0, 8, 4, 8).num_k_heads, 1); + } } /// Device parity against a verbatim port of `delta_rule_recurrence`. @@ -205,10 +279,11 @@ mod gdn_delta_rule_device_tests { use crate::driver::{CudaContext, CudaStream, GpuBuffer}; use crate::kernels::gdn::test_support::{assert_close, run_kernel, Lcg, HEAD_DIM, NUM_V_HEADS}; - /// Verbatim port of `aprender-serve`'s `delta_rule_recurrence` - /// (forward_qwen35.rs:121). + /// Verbatim port of `aprender-serve`'s `delta_rule_recurrence_gqa` + /// (forward_qwen35.rs:216) — the specification for this kernel. `h % num_k_heads` + /// is the reference's own head mapping, not a restatement of the kernel's. #[allow(clippy::too_many_arguments)] - fn delta_rule_recurrence( + fn delta_rule_recurrence_gqa( q: &[f32], k: &[f32], v: &[f32], @@ -216,20 +291,29 @@ mod gdn_delta_rule_device_tests { gate: &[f32], state: &mut [f32], output: &mut [f32], + num_k_heads: usize, + head_k_dim: usize, num_v_heads: usize, head_v_dim: usize, ) { - let scale = 1.0 / (head_v_dim as f32).sqrt(); + assert_eq!(q.len(), num_k_heads * head_k_dim); + assert_eq!(k.len(), num_k_heads * head_k_dim); + assert_eq!(v.len(), num_v_heads * head_v_dim); + assert_eq!(state.len(), num_v_heads * head_v_dim * head_k_dim); + + let scale = 1.0 / (head_k_dim as f32).sqrt(); for h in 0..num_v_heads { - let q_h = &q[h * head_v_dim..(h + 1) * head_v_dim]; - let k_h = &k[h * head_v_dim..(h + 1) * head_v_dim]; + let kh = h % num_k_heads; + let q_h = &q[kh * head_k_dim..(kh + 1) * head_k_dim]; + let k_h = &k[kh * head_k_dim..(kh + 1) * head_k_dim]; let v_h = &v[h * head_v_dim..(h + 1) * head_v_dim]; let beta_val = beta[h]; let gate_val = gate[h]; - let state_offset = h * head_v_dim * head_v_dim; - let s_h = &mut state[state_offset..state_offset + head_v_dim * head_v_dim]; + let state_stride = head_v_dim * head_k_dim; + let state_offset = h * state_stride; + let s_h = &mut state[state_offset..state_offset + state_stride]; let exp_gate = gate_val.exp(); for s in s_h.iter_mut() { @@ -238,26 +322,26 @@ mod gdn_delta_rule_device_tests { let mut delta = vec![0.0; head_v_dim]; for j in 0..head_v_dim { - let row_j = &s_h[j * head_v_dim..(j + 1) * head_v_dim]; + let row_j = &s_h[j * head_k_dim..(j + 1) * head_k_dim]; let mut sum = 0.0; - for i in 0..head_v_dim { + for i in 0..head_k_dim { sum += row_j[i] * k_h[i]; } delta[j] = (v_h[j] - sum) * beta_val; } for j in 0..head_v_dim { - let row_j = &mut s_h[j * head_v_dim..(j + 1) * head_v_dim]; + let row_j = &mut s_h[j * head_k_dim..(j + 1) * head_k_dim]; let d_j = delta[j]; - for i in 0..head_v_dim { + for i in 0..head_k_dim { row_j[i] += k_h[i] * d_j; } } for j in 0..head_v_dim { - let row_j = &s_h[j * head_v_dim..(j + 1) * head_v_dim]; + let row_j = &s_h[j * head_k_dim..(j + 1) * head_k_dim]; let mut sum = 0.0; - for i in 0..head_v_dim { + for i in 0..head_k_dim { sum += row_j[i] * q_h[i]; } output[h * head_v_dim + j] = sum * scale; @@ -276,35 +360,60 @@ mod gdn_delta_rule_device_tests { } } - #[test] - fn gdn_delta_rule_recurrence_matches_cpu_reference_three_steps() { + /// Per-element budget, **relative to the reference's own scale** (that is what + /// [`assert_close`] applies it as — an absolute tolerance on outputs that live + /// near 5e-3 is vacuous, PMAT-3477). + /// + /// MEASURED on this box (RTX 4090, sm_89), worst over all four shapes below and + /// all three steps, output and state alike: **2.2e-7** relative + /// (`nk 2 dk 8 nv 4 dv 8`, 1.49e-8 against a 6.84e-2 scale). The kernel keeps the + /// CPU's fp32 accumulation order, so the whole residual is one ulp of the + /// accumulator; 1e-5 is ~45x that and still two thousand times tighter than the + /// 1e-2 a mis-indexed head mapping produces. + const TOL: f32 = 1e-5; + + /// Three decode steps of the kernel against the host port, at any head shape. + /// + /// Returns `false` when there is no CUDA device (so the caller can say SKIPPED). + fn three_steps_match_the_reference( + num_k_heads: usize, + head_k_dim: usize, + num_v_heads: usize, + head_v_dim: usize, + seed: u32, + ) -> bool { let Ok(ctx) = CudaContext::new(0) else { - println!("gdn_delta_rule_recurrence: no CUDA device — SKIPPED."); - return; + return false; }; let stream = CudaStream::new(&ctx).expect("stream"); - let n = NUM_V_HEADS * HEAD_DIM; - let mut rng = Lcg::new(0x3477_0004); - let mut host_state = rng.vec(NUM_V_HEADS * HEAD_DIM * HEAD_DIM, 0.05); + let qk_n = num_k_heads * head_k_dim; + let n = num_v_heads * head_v_dim; + let mut rng = Lcg::new(seed); + let mut host_state = rng.vec(num_v_heads * head_v_dim * head_k_dim, 0.05); let state_buf = GpuBuffer::from_host(&ctx, &host_state).expect("state"); let out_buf = GpuBuffer::::new(&ctx, n).expect("out"); - let kernel = DeltaRuleRecurrenceKernel::new(NUM_V_HEADS as u32, HEAD_DIM as u32); + let kernel = DeltaRuleRecurrenceKernel::new( + num_k_heads as u32, + head_k_dim as u32, + num_v_heads as u32, + head_v_dim as u32, + ); for step in 0..3 { - let mut q = rng.vec(n, 1.0); - let mut k = rng.vec(n, 1.0); - l2_norm_per_head(&mut q, HEAD_DIM); - l2_norm_per_head(&mut k, HEAD_DIM); + let mut q = rng.vec(qk_n, 1.0); + let mut k = rng.vec(qk_n, 1.0); + l2_norm_per_head(&mut q, head_k_dim); + l2_norm_per_head(&mut k, head_k_dim); let v = rng.vec(n, 1.0); - let beta: Vec = (0..NUM_V_HEADS).map(|_| 0.5 + 0.25 * rng.next()).collect(); + let beta: Vec = (0..num_v_heads).map(|_| 0.5 + 0.25 * rng.next()).collect(); // dt = softplus(...) * ssm_a, with ssm_a negative: the state decays. - let gate: Vec = (0..NUM_V_HEADS) + let gate: Vec = (0..num_v_heads) .map(|_| -0.1 - 0.3 * rng.next().abs()) .collect(); let mut want = vec![0.0f32; n]; - delta_rule_recurrence( + delta_rule_recurrence_gqa( &q, &k, &v, @@ -312,8 +421,10 @@ mod gdn_delta_rule_device_tests { &gate, &mut host_state, &mut want, - NUM_V_HEADS, - HEAD_DIM, + num_k_heads, + head_k_dim, + num_v_heads, + head_v_dim, ); let q_buf = GpuBuffer::from_host(&ctx, &q).expect("q"); @@ -339,13 +450,15 @@ mod gdn_delta_rule_device_tests { &mut args, ); + let shape = + format!("nk {num_k_heads} dk {head_k_dim} nv {num_v_heads} dv {head_v_dim}"); let mut got = vec![0.0f32; n]; out_buf.copy_to_host(&mut got).expect("download output"); assert_close( &got, &want, - 1e-3, - &format!("delta-rule output, step {step}"), + TOL, + &format!("delta-rule output, step {step} ({shape})"), ); let mut got_state = vec![0.0f32; host_state.len()]; @@ -355,13 +468,76 @@ mod gdn_delta_rule_device_tests { assert_close( &got_state, &host_state, - 1e-3, - &format!("delta-rule state after step {step}"), + TOL, + &format!("delta-rule state after step {step} ({shape})"), ); // The fixture must actually move the state, or the carry-over is untested. let moved = got_state.iter().map(|s| s.abs()).fold(0.0f32, f32::max); - assert!(moved > 1e-3, "state is inert at step {step}"); + assert!(moved > 1e-3, "state is inert at step {step} ({shape})"); + } + true + } + + /// Qwen3.5-0.8B / -2B: one key head per value head, square state. Unchanged. + #[test] + fn gdn_delta_rule_recurrence_matches_cpu_reference_three_steps() { + if !three_steps_match_the_reference( + NUM_V_HEADS, + HEAD_DIM, + NUM_V_HEADS, + HEAD_DIM, + 0x3477_0004, + ) { + println!("gdn_delta_rule_recurrence: no CUDA device — SKIPPED."); + } + } + + /// PMAT-3477 (#3346/#3510): two value heads per key head, as Qwen3.5-4B/9B have. + /// + /// This is the case the kernel refused before: `q` and `k` are HALF as long as `v`, + /// and value heads 0/2 and 1/3 must read key heads 0 and 1 — `h % num_k_heads`. + /// A kernel that kept the old one-stride-for-everything indexing reads `q`/`k` past + /// their allocation for `h >= num_k_heads`, so this fails loudly rather than + /// silently. + #[test] + fn gdn_delta_rule_recurrence_grouped_key_heads_match_cpu_reference() { + if !three_steps_match_the_reference(2, 8, 4, 8, 0x3477_0032) { + println!("gdn_delta_rule_recurrence (grouped): no CUDA device — SKIPPED."); + } + } + + /// The same, with three value heads per key head — `num_v_heads / num_k_heads` is + /// not hard-coded to 2 anywhere. + #[test] + fn gdn_delta_rule_recurrence_ratio_three_matches_cpu_reference() { + if !three_steps_match_the_reference(2, 8, 6, 8, 0x3477_0033) { + println!("gdn_delta_rule_recurrence (ratio 3): no CUDA device — SKIPPED."); + } + } + + /// Qwen3.5-27B's own head counts — `nk = 16`, `nv = 48`, ratio 3 — at a small head + /// width so the test costs nothing. + /// + /// The 27B file is 16 GB and there is no host here that can hold it alongside a + /// CPU reference run, so this is the only place its head shape is exercised at all. + /// The two smaller ratio tests above use `nk = 2`, which cannot tell `h % nk` from + /// `h & (nk - 1)` or from a ratio hard-coded to 2; 16 and 48 pin both. + #[test] + fn gdn_delta_rule_recurrence_27b_head_counts_match_cpu_reference() { + if !three_steps_match_the_reference(16, 8, 48, 8, 0x3477_0048) { + println!("gdn_delta_rule_recurrence (27B heads): no CUDA device — SKIPPED."); + } + } + + /// A rectangular state (`Dk = 8`, `Dv = 6`): the state row length follows the KEY + /// dim, the row count and the output follow the VALUE dim, and the scale is + /// `Dk^-0.5`. No Qwen3.5 file is rectangular today; this is what keeps the two dims + /// from silently collapsing back into one. + #[test] + fn gdn_delta_rule_recurrence_rectangular_state_matches_cpu_reference() { + if !three_steps_match_the_reference(2, 8, 4, 6, 0x3477_0034) { + println!("gdn_delta_rule_recurrence (rectangular): no CUDA device — SKIPPED."); } } } diff --git a/crates/aprender-serve/src/cuda/executor/gdn_ops.rs b/crates/aprender-serve/src/cuda/executor/gdn_ops.rs index 1409e583cd..5446ac9d6b 100644 --- a/crates/aprender-serve/src/cuda/executor/gdn_ops.rs +++ b/crates/aprender-serve/src/cuda/executor/gdn_ops.rs @@ -375,8 +375,13 @@ impl CudaExecutor { } /// The gated delta-rule recurrence for one token - /// (`delta_rule_recurrence`). `state` is `[num_v_heads * D * D]` and is - /// updated in place; `output` is `[num_v_heads * D]`. + /// (`delta_rule_recurrence_gqa`). + /// + /// `q` and `k` are `[num_k_heads * head_k_dim]`, `v` and `output` are + /// `[num_v_heads * head_v_dim]`, `beta` and `gate` are per VALUE head, and + /// `state` is `[num_v_heads * head_v_dim * head_k_dim]`, updated in place. + /// Value head `h` reads key/query head `h % num_k_heads` — the tiled order + /// the GGUF conversion writes (PMAT-3477, #3346/#3510). /// /// # Errors /// PTX compilation or kernel launch failure, or a null device pointer. @@ -390,16 +395,25 @@ impl CudaExecutor { gate: &GpuBuffer, state: &GpuBuffer, output: &GpuBuffer, + num_k_heads: u32, + head_k_dim: u32, num_v_heads: u32, head_v_dim: u32, ) -> Result<(), GpuError> { - let kernel = - trueno_gpu::kernels::gdn::DeltaRuleRecurrenceKernel::new(num_v_heads, head_v_dim); + let kernel = trueno_gpu::kernels::gdn::DeltaRuleRecurrenceKernel::new( + num_k_heads, + head_k_dim, + num_v_heads, + head_v_dim, + ); let kernel_type = KernelType::GdnDeltaRule { num_v_heads, head_v_dim, + num_k_heads, + head_k_dim, }; - let cache_key = format!("gdn_delta_rule_{num_v_heads}_{head_v_dim}"); + let cache_key = + format!("gdn_delta_rule_{num_v_heads}_{head_v_dim}_{num_k_heads}_{head_k_dim}"); let kernel_name = self.gdn_prepare(&kernel_type, &cache_key)?; let (gx, _, _) = kernel.grid(); let (bx, _, _) = kernel.block(); diff --git a/crates/aprender-serve/src/cuda/generate.rs b/crates/aprender-serve/src/cuda/generate.rs index 9aeaa792e4..24f8638e0b 100644 --- a/crates/aprender-serve/src/cuda/generate.rs +++ b/crates/aprender-serve/src/cuda/generate.rs @@ -269,8 +269,9 @@ impl CudaKernels { PerHeadL2NormKernel::new(*head_dim, *num_heads, *epsilon).emit_ptx() }, KernelType::GdnGates { num_heads } => GdnGatesKernel::new(*num_heads).emit_ptx(), - KernelType::GdnDeltaRule { num_v_heads, head_v_dim } => { - DeltaRuleRecurrenceKernel::new(*num_v_heads, *head_v_dim).emit_ptx() + KernelType::GdnDeltaRule { num_v_heads, head_v_dim, num_k_heads, head_k_dim } => { + DeltaRuleRecurrenceKernel::new(*num_k_heads, *head_k_dim, *num_v_heads, *head_v_dim) + .emit_ptx() }, KernelType::GdnGatedRmsNorm { head_dim, num_heads, epsilon } => { GatedRmsNormKernel::new(*head_dim, *num_heads, *epsilon).emit_ptx() diff --git a/crates/aprender-serve/src/cuda/kernel_type.rs b/crates/aprender-serve/src/cuda/kernel_type.rs index 5deecac9e6..ef25c3f597 100644 --- a/crates/aprender-serve/src/cuda/kernel_type.rs +++ b/crates/aprender-serve/src/cuda/kernel_type.rs @@ -583,8 +583,15 @@ pub enum KernelType { }, /// PMAT-3477 (#3090): the Gated `DeltaNet` per-head `dt`/`beta` gates. GdnGates { num_heads: u32 }, - /// PMAT-3477 (#3090): the gated delta-rule recurrence, one token. - GdnDeltaRule { num_v_heads: u32, head_v_dim: u32 }, + /// PMAT-3477 (#3090, #3346/#3510): the gated delta-rule recurrence, one token. + /// `num_k_heads`/`head_k_dim` are APPENDED (never reordered — a line-keyed guard + /// baseline points into this file); value head `h` reads key head `h % num_k_heads`. + GdnDeltaRule { + num_v_heads: u32, + head_v_dim: u32, + num_k_heads: u32, + head_k_dim: u32, + }, /// PMAT-3477 (#3090): gated RMSNorm — the Gated `DeltaNet` output norm. GdnGatedRmsNorm { head_dim: u32, diff --git a/crates/aprender-serve/src/cuda/kernels_generate_gemm_cuda.rs b/crates/aprender-serve/src/cuda/kernels_generate_gemm_cuda.rs index 1b74153ad3..f210cc9740 100644 --- a/crates/aprender-serve/src/cuda/kernels_generate_gemm_cuda.rs +++ b/crates/aprender-serve/src/cuda/kernels_generate_gemm_cuda.rs @@ -311,8 +311,9 @@ impl CudaKernels { KernelType::GdnGates { num_heads } => { GdnGatesKernel::new(*num_heads).emit_ptx_for_target(target) }, - KernelType::GdnDeltaRule { num_v_heads, head_v_dim } => { - DeltaRuleRecurrenceKernel::new(*num_v_heads, *head_v_dim).emit_ptx_for_target(target) + KernelType::GdnDeltaRule { num_v_heads, head_v_dim, num_k_heads, head_k_dim } => { + DeltaRuleRecurrenceKernel::new(*num_k_heads, *head_k_dim, *num_v_heads, *head_v_dim) + .emit_ptx_for_target(target) }, KernelType::GdnGatedRmsNorm { head_dim, num_heads, epsilon } => { GatedRmsNormKernel::new(*head_dim, *num_heads, *epsilon).emit_ptx_for_target(target) diff --git a/crates/aprender-serve/src/gguf/cuda/forward_qwen35_cuda.rs b/crates/aprender-serve/src/gguf/cuda/forward_qwen35_cuda.rs index b56ebc0e52..8ff1f85732 100644 --- a/crates/aprender-serve/src/gguf/cuda/forward_qwen35_cuda.rs +++ b/crates/aprender-serve/src/gguf/cuda/forward_qwen35_cuda.rs @@ -30,7 +30,7 @@ //! kinds, the output norm and the `lm_head`. //! //! Every state buffer is sized from the config -//! (`num_v_heads * head_v_dim * head_v_dim`, `conv_dim * (conv_kernel - 1)`, +//! (`num_v_heads * head_v_dim * head_k_dim`, `conv_dim * (conv_kernel - 1)`, //! `max_seq_len * num_kv_heads * head_dim`) — never from a constant. use super::{OwnedQuantizedTensor, RealizarError, Result}; @@ -108,7 +108,7 @@ pub struct Qwen35CudaState { kv: Vec, GpuBuffer)>>, /// `conv_dim * (conv_kernel - 1)`. conv_len: usize, - /// `num_v_heads * head_v_dim * head_v_dim`. + /// `num_v_heads * head_v_dim * head_k_dim`. ssm_len: usize, /// `num_kv_heads * head_dim` — one KV cache row. kv_row: usize, @@ -358,19 +358,29 @@ impl<'a> Qwen35CudaModel<'a> { } } - /// The delta rule reads q, k and v with one head stride, so a file whose key - /// heads differ from its value heads cannot run this kernel — refuse rather - /// than index past a head. - fn check_head_symmetry(d: Qwen35CudaDims) -> Result<()> { - if d.head_k_dim == d.head_v_dim && d.num_k_heads == d.num_v_heads { + /// The delta rule maps value head `h` onto key head `h % num_k_heads`, so the + /// ONLY shape it cannot serve is one whose value heads are not a whole number + /// of key-head groups (PMAT-3477, #3346/#3510). + /// + /// This used to refuse every `num_k_heads != num_v_heads` file outright, which + /// is what kept Qwen3.5-4B/9B (32 value heads against 16 key heads) and -27B + /// (48 against 16) off the GPU. `head_k_dim != head_v_dim` is likewise no + /// longer a refusal: the kernel sizes the state row from the key dim and the + /// row count from the value dim. + /// Renamed from `check_head_symmetry`: it no longer asks for symmetry, and a + /// predicate whose name outlives what it tests is the next reader's wrong + /// diagnosis. + fn check_head_grouping(d: Qwen35CudaDims) -> Result<()> { + if d.num_k_heads > 0 && d.num_v_heads % d.num_k_heads == 0 { return Ok(()); } Err(RealizarError::UnsupportedOperation { operation: "qwen35_cuda_deltanet".to_string(), reason: format!( - "the delta-rule kernel indexes q/k/v with one head stride: \ - head_k_dim {} != head_v_dim {} or num_k_heads {} != num_v_heads {}", - d.head_k_dim, d.head_v_dim, d.num_k_heads, d.num_v_heads + "the delta rule maps value head h onto key head h % num_k_heads, so \ + num_v_heads must be a positive multiple of num_k_heads: num_k_heads {} \ + does not divide num_v_heads {}", + d.num_k_heads, d.num_v_heads ), }) } @@ -500,7 +510,7 @@ impl<'a> Qwen35CudaModel<'a> { max_seq_len: usize, ) -> Result { let dims = Self::dims_of(model); - Self::check_head_symmetry(dims)?; + Self::check_head_grouping(dims)?; if max_seq_len == 0 { return Err(RealizarError::InvalidShape { reason: "qwen35_cuda: max_seq_len must be at least 1".to_string(), @@ -605,7 +615,11 @@ impl<'a> Qwen35CudaModel<'a> { max_seq_len: usize, ) -> Result { let conv_len = (dims.conv_dim * (dims.conv_kernel - 1)) as usize; - let ssm_len = (dims.num_v_heads * dims.head_v_dim * dims.head_v_dim) as usize; + // The recurrent state of one value head is [head_v_dim rows x head_k_dim], + // laid out `s[j * head_k_dim + i] == S[i][j]` — the CPU reference's own + // layout. Sizing it from head_v_dim twice was only right because every + // file so far ships head_k_dim == head_v_dim (PMAT-3477). + let ssm_len = (dims.num_v_heads * dims.head_v_dim * dims.head_k_dim) as usize; let kv_row = (dims.num_kv_heads * dims.attn_head_dim) as usize; let mut conv = Vec::with_capacity(layers.len()); let mut ssm = Vec::with_capacity(layers.len()); @@ -1288,6 +1302,8 @@ impl<'a> Qwen35CudaModel<'a> { &s.dt, &state.ssm[il], &s.out_h, + d.num_k_heads, + d.head_k_dim, d.num_v_heads, d.head_v_dim, )?; diff --git a/crates/aprender-serve/src/gguf/cuda/forward_qwen35_cuda_tests.rs b/crates/aprender-serve/src/gguf/cuda/forward_qwen35_cuda_tests.rs index 116359c2d5..2d1ea33116 100644 --- a/crates/aprender-serve/src/gguf/cuda/forward_qwen35_cuda_tests.rs +++ b/crates/aprender-serve/src/gguf/cuda/forward_qwen35_cuda_tests.rs @@ -1088,8 +1088,8 @@ fn qwen35_cuda_state_is_sized_from_the_config() { ); assert_eq!( gpu.state().ssm_len(), - qwen.num_v_heads * qwen.head_v_dim * qwen.head_v_dim, - "recurrent state is num_v_heads * head_v_dim^2" + qwen.num_v_heads * qwen.head_v_dim * qwen.head_k_dim, + "recurrent state is num_v_heads * head_v_dim * head_k_dim (equal dims on this file)" ); } @@ -1300,3 +1300,414 @@ fn qwen35_cuda_dp4a_gemv_is_catastrophic_through_the_recurrence() { /// Keep the layer type in the compiled surface: the parity tests bind it, and a /// rename of the CPU struct must break here, not silently at phase 2. const _: Option<&Qwen35OwnedDeltaNetLayer> = None; + +// ============================================================================ +// PMAT-3477 (#3346/#3510): the GQA `DeltaNet` file — num_v_heads > num_k_heads. +// +// Everything above runs Qwen3.5-0.8B, whose 16 key heads and 16 value heads make +// `h % num_k_heads` the identity: NOT ONE assertion above can tell the tiled head +// mapping from the old one-stride-for-everything indexing, and the CUDA model +// refused this shape outright until this ticket. These two tests are the only +// place the grouped mapping is exercised end to end on a real file. +// ============================================================================ + +/// The real GQA file: 32 value heads against 16 key heads, 32 blocks. +const MODEL_PATH_4B: &str = "/home/noah/models/Qwen3.5-4B-Q4_K_M.gguf"; + +/// "The capital of France is" under the 4B file's own 248320-entry vocabulary +/// (`The` is 760 here, 785 in the Qwen2/Qwen3 vocabulary) — the same prompt the +/// CPU test `qwen35_gqa_4b_matches_llama_cpp_greedy_tokens` pins against +/// llama.cpp's `[11751 " Paris", 13 ".", 198 "\n", 32 "A"]`. +const PROMPT_4B: [u32; 5] = [760, 6511, 314, 9338, 369]; + +/// Budget for one whole wired 4B `DeltaNet` layer, hidden state and recurrent +/// state alike. +/// +/// MEASURED on this file, 24 `DeltaNet` layers x 2 positions = 48 comparisons: +/// worst hidden **2.250e-2** (pos 1, layer 2), worst recurrent state +/// **1.165e-2**. It is looser than the 0.8B [`LAYER_TOL`] for the reason that +/// file's own budget is looser than [`TOL`] — the wired layer runs the Q4_K/Q8_0 +/// projection GEMVs, whose CPU side quantizes its activation to Q8_K, and 4B is +/// wider and deeper than 0.8B so that error accumulates over more of it. It is a +/// budget, not a bar: the same 2.2x headroom over the measurement that +/// [`LAYER_TOL`] carries. +/// +/// It still has teeth. Mutating the kernel's head mapping from `h % num_k_heads` +/// to the grouped `h / (num_v_heads / num_k_heads)` — one line in +/// `aprender-gpu`'s `delta_rule.rs`, reverted — turns this RED immediately (see +/// the test's own doc comment for the reading). +const LAYER_BUDGET_4B: f32 = 5e-2; + +/// Skip unless both the device and a named model file are here. +macro_rules! qwen35_cuda_file_or_skip { + ($path:expr) => {{ + if !std::path::Path::new($path).exists() { + eprintln!("SKIP: {} is absent", $path); + return; + } + crate::cuda_executor_or_skip!(0) + }}; +} + +/// Load a Qwen3.5 file and assert it really is a grouped-head one, so a file +/// swap cannot quietly turn a GQA test into another 0.8B run. +fn assert_is_gqa(qwen: &Qwen35Model<'_>) { + assert!( + qwen.num_v_heads > qwen.num_k_heads, + "{MODEL_PATH_4B} is not a GQA DeltaNet file (num_v_heads {} vs num_k_heads {}); this \ + test cannot falsify the head mapping against it", + qwen.num_v_heads, + qwen.num_k_heads + ); + assert_eq!( + qwen.num_v_heads % qwen.num_k_heads, + 0, + "num_v_heads {} is not a multiple of num_k_heads {}", + qwen.num_v_heads, + qwen.num_k_heads + ); +} + +/// Teacher-forced per-layer parity of the GPU Gated `DeltaNet` block against the +/// CPU on the 4B file, over the first two tokens and every `DeltaNet` layer — +/// the layer output hidden state, the causal-conv window and the recurrent +/// state, at the same budgets the 0.8B suite uses. +/// +/// Two tokens, not three: the 4B file is 32 blocks of ~2.7 GB of weights and +/// this walks every `DeltaNet` layer at every position with a host round-trip +/// per layer. +/// +/// What this adds over the 0.8B layer test: `q` and `k` here are HALF as long as +/// `v` (`num_k_heads * head_k_dim = 2048` against `num_v_heads * head_v_dim = +/// 4096`), so a kernel that indexed q/k with the value-head stride would read +/// past their allocation for every value head `h >= 16`. +/// +/// MEASURED: worst hidden 2.250e-2 (pos 1, layer 2), worst recurrent state +/// 1.165e-2, over 48 comparisons — both inside [`LAYER_BUDGET_4B`]. With the +/// head mapping mutated to `h / ratio` they are 6.319e-1 and 2.445e0. +#[test] +#[serial_test::serial] +fn qwen35_cuda_4b_deltanet_layers_match_cpu_on_the_real_file() { + let executor = qwen35_cuda_file_or_skip!(MODEL_PATH_4B); + let mapped = crate::gguf::MappedGGUFModel::from_path(MODEL_PATH_4B).expect("map the 4B GGUF"); + let base = load_cpu_model(&mapped); + let qwen = Qwen35Model::from_model_and_layers(&base, &mapped.model, mapped.data()) + .expect("4B hybrid layers"); + assert_is_gqa(&qwen); + + let deltanet_layers = qwen + .layers + .iter() + .filter(|l| matches!(l, Qwen35OwnedLayer::DeltaNet(_))) + .count(); + assert!( + deltanet_layers > 0, + "the 4B file must carry DeltaNet layers" + ); + + let mut gpu = Qwen35CudaModel::new(&qwen, executor).expect("build the 4B CUDA model"); + gpu.pin_reference_gemv(); + + let hidden_dim = qwen.base.config.hidden_dim; + let tokens = &PROMPT_4B[..2]; + let mut cpu_state = qwen.new_state(tokens.len() + 1); + let mut normed = vec![0.0f32; hidden_dim]; + let mut post_normed = vec![0.0f32; hidden_dim]; + let mut compared = 0usize; + let mut worst_hidden = 0.0f32; + let mut worst_ssm = 0.0f32; + let mut worst_where = String::new(); + + for (pos, &token) in tokens.iter().enumerate() { + let mut hidden = qwen.base.token_embedding() + [(token as usize) * hidden_dim..(token as usize + 1) * hidden_dim] + .to_vec(); + + for (il, layer) in qwen.layers.iter().enumerate() { + match layer { + Qwen35OwnedLayer::DeltaNet(d) => { + gpu.upload_layer(il, &cpu_state.conv_states[il], &cpu_state.ssm_states[il]) + .expect("seed the device state"); + let dev = GpuBuffer::from_host(gpu.executor_mut().context(), &hidden) + .expect("upload hidden"); + + qwen.forward_deltanet( + d, + &mut hidden, + &mut cpu_state, + il, + pos, + &mut normed, + &mut post_normed, + ) + .expect("cpu deltanet"); + + gpu.forward_deltanet_layer(il, &dev).expect("gpu deltanet"); + gpu.executor_mut().sync_stream().expect("sync"); + let mut got = vec![0.0f32; hidden_dim]; + dev.copy_to_host(&mut got).expect("download hidden"); + + let l = rel_linf(&got, &hidden); + if l > worst_hidden { + worst_hidden = l; + worst_where = format!("pos {pos} layer {il}"); + } + + let (conv, ssm) = gpu.download_layer(il).expect("download state"); + assert_rel_linf( + &conv, + &cpu_state.conv_states[il], + TOL, + &format!("4B pos {pos} layer {il} conv window"), + ); + let ls = rel_linf(&ssm, &cpu_state.ssm_states[il]); + if ls > worst_ssm { + worst_ssm = ls; + } + compared += 1; + }, + Qwen35OwnedLayer::Attention(a) => { + qwen.forward_attention( + a, + &mut hidden, + &mut cpu_state, + il, + pos, + &mut normed, + &mut post_normed, + ) + .expect("cpu attention"); + }, + } + } + cpu_state.kv_cache.advance(); + } + + assert_eq!( + compared, + deltanet_layers * tokens.len(), + "every 4B DeltaNet layer must be compared at every position" + ); + eprintln!( + "[4b-layer] worst hidden {worst_hidden:.3e} at {worst_where}, worst ssm {worst_ssm:.3e}, \ + over {compared} comparisons" + ); + assert!( + worst_hidden <= LAYER_BUDGET_4B, + "4B layer hidden: worst relative L-inf {worst_hidden:.3e} at {worst_where} exceeds the \ + measured budget {LAYER_BUDGET_4B:.0e}" + ); + assert!( + worst_ssm <= LAYER_BUDGET_4B, + "4B recurrent state: worst relative L-inf {worst_ssm:.3e} exceeds the measured budget \ + {LAYER_BUDGET_4B:.0e}" + ); +} + +/// End to end on the 4B file: the five prompt positions plus FOUR greedy +/// continuation tokens, GPU against CPU, **argmax-equal at every position**. +/// +/// The continuation is driven by the GPU's own argmax and fed to both sides, so +/// a single divergent token ends the agreement — this is the assertion that says +/// the GPU runs the same model, not merely a similar one. The CPU side is itself +/// pinned to llama.cpp's `[11751, 13, 198, 32]` by +/// `qwen35_gqa_4b_matches_llama_cpp_greedy_tokens`, so those four ids are +/// asserted here too: without that the pair could agree on garbage. +/// +/// FALSIFIED: with the kernel's head mapping mutated to the grouped +/// `h / (num_v_heads / num_k_heads)` this fails at the FIRST prompt position on +/// the argmax assertion, not on a tolerance. +#[test] +#[serial_test::serial] +fn qwen35_cuda_4b_forward_single_matches_cpu_argmax_end_to_end() { + let executor = qwen35_cuda_file_or_skip!(MODEL_PATH_4B); + let mapped = crate::gguf::MappedGGUFModel::from_path(MODEL_PATH_4B).expect("map the 4B GGUF"); + let base = load_cpu_model(&mapped); + let qwen = Qwen35Model::from_model_and_layers(&base, &mapped.model, mapped.data()) + .expect("4B hybrid layers"); + assert_is_gqa(&qwen); + + /// llama.cpp master `60b06ab9a`, CPU, `-no-cnv --temp 0`, same raw prompt. + const WANT: [u32; 4] = [11751, 13, 198, 32]; + + let mut gpu = Qwen35CudaModel::new(&qwen, executor).expect("build the 4B CUDA model"); + gpu.pin_reference_gemv(); + let mut gpu_state = gpu.new_state().expect("device state"); + let mut cpu_state = qwen.new_state(PROMPT_4B.len() + WANT.len()); + + let mut cpu_logits = Vec::new(); + let mut gpu_logits = Vec::new(); + for (pos, &token) in PROMPT_4B.iter().enumerate() { + cpu_logits = qwen + .forward_single_qwen35(token, &mut cpu_state, pos) + .expect("cpu 4B forward"); + gpu_logits = gpu + .forward_single(token, &mut gpu_state, pos) + .expect("gpu 4B forward"); + assert!( + gpu_logits.iter().all(|l| l.is_finite()), + "4B pos {pos}: the GPU forward produced a non-finite logit" + ); + let (cos, linf) = assert_forward_parity( + &gpu_logits, + &cpu_logits, + LOGITS_BUDGET, + &format!("4B prompt pos {pos} logits"), + ); + eprintln!( + "[4b-e2e] prompt pos {pos}: argmax {} cosine {cos:.6} relative L-inf {linf:.3e}", + crate::gguf::ops::argmax(&cpu_logits), + ); + } + + let mut got = Vec::with_capacity(WANT.len()); + for step in 0..WANT.len() { + let pos = PROMPT_4B.len() + step; + // The GPU's own choice drives both sides: an argmax the GPU alone picks + // takes the two forwards apart immediately instead of being re-seeded. + let next = crate::gguf::ops::argmax(&gpu_logits); + assert_eq!( + next, + crate::gguf::ops::argmax(&cpu_logits), + "4B step {step}: GPU and CPU disagree on the next token" + ); + got.push(next); + let token = next; + cpu_logits = qwen + .forward_single_qwen35(token, &mut cpu_state, pos) + .expect("cpu 4B forward"); + gpu_logits = gpu + .forward_single(token, &mut gpu_state, pos) + .expect("gpu 4B forward"); + let (cos, linf) = assert_forward_parity( + &gpu_logits, + &cpu_logits, + LOGITS_BUDGET, + &format!("4B decode pos {pos} logits"), + ); + eprintln!("[4b-e2e] decode pos {pos}: token {next} cosine {cos:.6} L-inf {linf:.3e}"); + } + + assert_eq!( + got, + WANT.to_vec(), + "the 4B GPU forward does not follow llama.cpp greedily: got {got:?}, want {WANT:?} — \ + num_k_heads {}, num_v_heads {}, head_k_dim {}, head_v_dim {}", + qwen.num_k_heads, + qwen.num_v_heads, + qwen.head_k_dim, + qwen.head_v_dim + ); + assert_eq!( + gpu_state.kv_len(), + PROMPT_4B.len() + WANT.len(), + "the device KV cache must have advanced once per position" + ); +} + +/// The device state of a GQA file is `num_v_heads * head_v_dim * head_k_dim`, +/// and its conv window is `(2 * k_dim + v_dim) * (conv_kernel - 1)` — the two +/// sizes that were wrong (or unreachable) while `head_k_dim` and `head_v_dim` +/// were assumed equal. +#[test] +#[serial_test::serial] +fn qwen35_cuda_4b_state_is_sized_from_the_grouped_config() { + let executor = qwen35_cuda_file_or_skip!(MODEL_PATH_4B); + let mapped = crate::gguf::MappedGGUFModel::from_path(MODEL_PATH_4B).expect("map the 4B GGUF"); + let base = load_cpu_model(&mapped); + let qwen = Qwen35Model::from_model_and_layers(&base, &mapped.model, mapped.data()) + .expect("4B hybrid layers"); + assert_is_gqa(&qwen); + let conv_dim = qwen.head_k_dim * qwen.num_k_heads * 2 + qwen.head_v_dim * qwen.num_v_heads; + + let gpu = Qwen35CudaModel::new(&qwen, executor).expect("build the 4B CUDA model"); + assert_eq!( + gpu.state().conv_len(), + conv_dim * (qwen.conv_kernel - 1), + "conv window is (2 * k_dim + v_dim) * (conv_kernel - 1)" + ); + assert_eq!( + gpu.state().ssm_len(), + qwen.num_v_heads * qwen.head_v_dim * qwen.head_k_dim, + "the recurrent state is num_v_heads * head_v_dim * head_k_dim" + ); + assert_eq!( + gpu.state().ssm_len(), + qwen.new_state(1).ssm_states[0].len(), + "the device state must be the same length as the CPU's own" + ); +} + +/// The one head shape the CUDA `DeltaNet` still refuses, and the one it must not. +/// +/// `check_head_grouping` is the whole guard between an indivisible file and a +/// kernel that would read `q`/`k` at `h % num_k_heads` for a value head with no +/// key head to land on. Nothing else exercised it: the refusal it replaced was +/// `num_k_heads != num_v_heads`, so every test on this file (0.8B, 16/16) and +/// every 4B test above (16/32) takes the Ok arm. This runs with no device and no +/// model file, so it is the one assertion here that CI actually reaches. +#[test] +fn qwen35_cuda_refuses_only_value_heads_that_do_not_group() { + /// `Qwen35CudaDims` with only the four head fields that decide the grouping. + fn dims( + num_k_heads: u32, + num_v_heads: u32, + head_k_dim: u32, + head_v_dim: u32, + ) -> super::Qwen35CudaDims { + super::Qwen35CudaDims { + hidden_dim: 1024, + intermediate_dim: 3072, + conv_dim: num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim, + k_dim: num_k_heads * head_k_dim, + v_dim: num_v_heads * head_v_dim, + head_k_dim, + num_k_heads, + head_v_dim, + num_v_heads, + conv_kernel: 4, + eps: 1e-6, + num_heads: 16, + num_kv_heads: 2, + attn_head_dim: 128, + n_rot: 128, + theta_scale: 0.5, + vocab_size: 32, + } + } + + // Every real Qwen3.5 head shape must be accepted — these are the files this + // ticket exists to admit, and a guard that refused one would be found only by + // a 16 GB download. + for (nk, nv, label) in [ + (16u32, 16u32, "0.8B / 2B"), + (16, 32, "4B / 9B"), + (16, 48, "27B"), + ] { + Qwen35CudaModel::check_head_grouping(dims(nk, nv, 128, 128)) + .unwrap_or_else(|e| panic!("{label} ({nk} key heads, {nv} value heads) refused: {e}")); + } + + // A rectangular state is no longer a refusal: the kernel sizes the state row + // from the key dim and the row count from the value dim. + Qwen35CudaModel::check_head_grouping(dims(16, 32, 128, 64)) + .expect("head_k_dim != head_v_dim must be accepted"); + + // 5 value heads onto 2 key heads: head 4 would wrap onto key head 0 and read a + // group of one where every other group has two. + let err = Qwen35CudaModel::check_head_grouping(dims(2, 5, 8, 8)) + .expect_err("5 value heads do not group onto 2 key heads"); + let text = format!("{err}"); + assert!( + text.contains("qwen35_cuda_deltanet"), + "the refusal must name the operation: {text}" + ); + assert!( + text.contains('2') && text.contains('5'), + "the refusal must quote both head counts so the file is identifiable: {text}" + ); + + // Zero key heads would be a division by zero, not merely a bad grouping. + Qwen35CudaModel::check_head_grouping(dims(0, 16, 128, 128)) + .expect_err("zero key heads must be refused, not divided by"); +} From f36aa1bea21fd369c186801fbdba0bb8291ea104 Mon Sep 17 00:00:00 2001 From: Noah Gift Date: Sat, 19 Sep 2026 14:14:31 +0200 Subject: [PATCH 4/4] =?UTF-8?q?docs(audits):=20PMAT-3477=20review=20quorum?= =?UTF-8?q?=20receipt=20for=20#3527=20=E2=80=94=203/3=20PASS=20on=20the=20?= =?UTF-8?q?Qwen3.5=20sizes=20diff?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pmat-Ticket: PMAT-3477 Co-Authored-By: Claude Fable 5.1 --- docs/audits/quorum-PMAT-3477-pr3527.json | 144 +++++++++++++++++++++++ 1 file changed, 144 insertions(+) create mode 100644 docs/audits/quorum-PMAT-3477-pr3527.json diff --git a/docs/audits/quorum-PMAT-3477-pr3527.json b/docs/audits/quorum-PMAT-3477-pr3527.json new file mode 100644 index 0000000000..b78ee4eea0 --- /dev/null +++ b/docs/audits/quorum-PMAT-3477-pr3527.json @@ -0,0 +1,144 @@ +{ + "ticket": "PMAT-3477", + "pr": 3527, + "head": "c58062517", + "width": 3, + "executor": "agy", + "mode": "plan (read-only review)", + "dispatched_by": "orchestrator directly via agy-lane.sh (paiml-agy-delegate hit the opus session rate limit 429 at 11:5xZ)", + "agreed": true, + "lanes": [ + { + "lane": 1, + "conversation_id": "bc6a9c8e-842a-4d0c-b486-a64165ac5559", + "status": "SUCCESS", + "verdict": "PASS", + "findings": [ + { + "file": "crates/aprender-serve/src/gguf/inference/forward/forward_qwen35.rs", + "line": 1277, + "claim": "GQA head mapping applied consistently to CPU recurrence (h % num_k_heads), PTX (ctx.rem_u32), state sizing (nv * Dv * Dk), and conv1d/l2norm/gate widths." + }, + { + "file": "crates/aprender-serve/src/gguf/cuda/forward_qwen35_cuda_tests.rs", + "line": 942, + "claim": "Tests are not vacuous. The new 4B test fails on first argmax diverge if h/ratio is used. Layer test verifies the count strictly (fails on 24). The old 0.8B tests pass on 24 and h/ratio only because 0.8B config has exactly 24 layers and ratio=1." + }, + { + "file": "crates/aprender-serve/src/gguf/qwen35_load.rs", + "line": 61, + "claim": "Silent fallback for block_count removed and replaced with a clean error (fail closed). Refusal site check_head_grouping properly opened for divisible heads but kept closed for indivisible configs (and num_k_heads=0)." + }, + { + "file": "crates/aprender-serve/src/cuda/kernel_type.rs", + "line": 589, + "claim": "GdnDeltaRule appends new fields rather than reordering, protecting the line-keyed baseline. Complexity and literal ratchets passed cleanly in CI emulation." + } + ] + }, + { + "lane": 2, + "conversation_id": "2acc7232-ef48-49bb-a5e3-f7ac6f597793", + "status": "SUCCESS", + "verdict": "PASS", + "findings": [ + { + "file": "crates/aprender-serve/src/gguf/inference/forward/forward_qwen35.rs", + "line": 1277, + "claim": "CPU Gated DeltaNet GQA recurrence correctly implements tiled head mapping h % num_k_heads and transposed state layout [num_v_heads * head_v_dim * head_k_dim]." + }, + { + "file": "crates/aprender-serve/src/gguf/inference/forward/forward_qwen35.rs", + "line": 451, + "claim": "Qwen35State recurrent state allocation correctly sizes ssm_states with head_k_dim row length per value head." + }, + { + "file": "crates/aprender-serve/src/gguf/inference/forward/forward_qwen35.rs", + "line": 987, + "claim": "Causal conv1d, per-head L2 normalization, and gate projections correctly use q,k = nk*Dk and v,gate,out = nv*Dv." + }, + { + "file": "crates/aprender-gpu/src/kernels/gdn/delta_rule.rs", + "line": 145, + "claim": "PTX delta-rule kernel emits rem.u32(h, nk) when nk != nv, optimizes to identity when nk == nv, and indexes state with Dk stride." + }, + { + "file": "crates/aprender-serve/src/gguf/cuda/forward_qwen35_cuda.rs", + "line": 373, + "claim": "check_head_grouping replaces check_head_symmetry, failing closed on indivisible head groupings while admitting 4B/9B/27B and rectangular states." + }, + { + "file": "crates/aprender-serve/src/gguf/cuda/forward_qwen35_cuda.rs", + "line": 622, + "claim": "CUDA recurrent state buffer is correctly sized from num_v_heads * head_v_dim * head_k_dim instead of squaring head_v_dim." + }, + { + "file": "crates/aprender-serve/src/gguf/qwen35_load.rs", + "line": 1901, + "claim": "load_qwen35_layers reads arch-prefixed block_count and fails closed with InvalidShape error instead of silent fallback to 24 layers." + }, + { + "file": "crates/aprender-serve/src/cuda/kernel_type.rs", + "line": 592, + "claim": "KernelType::GdnDeltaRule appends new fields to preserve line 236 referenced in claim_literal_baseline.txt." + }, + { + "file": "crates/aprender-serve/src/gguf/inference/forward/forward_qwen35_gqa_tests.rs", + "line": 1555, + "claim": "Test qwen35_gqa_recurrence_tiles_q_and_k_over_the_value_heads explicitly checks assert_ne against repeat_interleave (h / ratio), proving non-vacuity." + }, + { + "file": "crates/aprender-serve/src/gguf/qwen35_load.rs", + "line": 1941, + "claim": "Test qwen35_layer_count_is_read_from_the_arch_prefixed_block_count asserts declared != 24 on the 4B model, proving non-vacuity against the 24-layer fallback." + } + ] + }, + { + "lane": 3, + "conversation_id": "6471c971-a62d-4f70-b0bd-f3e15d8597bb", + "status": "SUCCESS", + "verdict": "PASS", + "findings": [ + { + "file": "crates/aprender-serve/src/gguf/inference/forward/forward_qwen35.rs", + "line": 254, + "claim": "GQA head mapping (h % num_k_heads) is correctly and consistently implemented across CPU recurrence, PTX kernel, state sizing [num_v_heads * Dk * Dv], and tensor widths." + }, + { + "file": "crates/aprender-gpu/src/kernels/gdn/delta_rule.rs", + "line": 129, + "claim": "DeltaRuleRecurrenceKernel in PTX emits rem.u32(h, nk) for grouped heads and maps key/value head offsets and rectangular state sizing correctly." + }, + { + "file": "crates/aprender-serve/src/gguf/qwen35_load.rs", + "line": 78, + "claim": "load_qwen35_layers fails closed on missing {arch}.block_count and eliminates silent 24-layer fallback." + }, + { + "file": "crates/aprender-serve/src/gguf/inference/forward/forward_qwen35_gqa_tests.rs", + "line": 172, + "claim": "Non-vacuous anti-vacuity assertions verify that repeat_interleave (h / ratio) and 24-layer fallback fail tests." + }, + { + "file": "crates/aprender-serve/src/gguf/cuda/forward_qwen35_cuda.rs", + "line": 373, + "claim": "check_head_grouping accurately admits divisible GQA head ratios (16/16, 16/32, 16/48) and rectangular state while properly refusing indivisible head counts." + } + ] + } + ], + "void": [ + { + "lane": 3, + "conversation_id": "7acae5cd-87be-434c-8e7c-cffc02c9491b", + "reason": "BLIND: tried to run the cuda test suite and timed out; relaunched read-only" + } + ], + "sibling_read_check": "lanes ran in separate self-contained clones under .claude/worktrees; no shared $WORK verdict file", + "orchestrator_rerun": { + "cargo test -p aprender-gpu --lib --features cuda delta_rule": "9 passed", + "cargo test -p aprender-serve --lib --features cuda qwen35_cuda": "16 passed", + "apr run --gpu 4B/9B": "\"**Paris**.\" on the CUDA backend line" + } +} \ No newline at end of file