feat(qwen4): add the Qwen3.8-Flash-Next GDN golden and its prefill gate

docs/models/qwen4/bring-up.md

@@ -93,7 +93,20 @@ The 36 `linear_attention` layers will run Qwen3.5's Gated DeltaNet decode kernel
 
 The three kernel limits (`GDN_AOT_KEY_HEAD_DIM`, `GDN_AOT_VALUE_HEAD_DIM`, `LINEAR_CONV_MAX_KERNEL_DIM`) are defined next to the wrappers. This line's `src/config.rs` keeps a CPU-only mirror of them, so its config tests need no device, until slice B adds the device dependency.
 
-Still to verify is the reuse on this line's activations. The tensor contract derives every GDN shape from the frozen config and validates all 1658 headers, but nothing has run this line's GDN block against a reference yet. That is slice B's M0 exit (#1127), against the `transformers` recurrence, and it has to cover several decode steps from a nonzero state (so `(v - k·s)` is live), finite checks before any reduction, and the gated norm plus `out_proj` that complete the block (the norm's sigmoid entry point is #1129's).
+Still to verify is the **decode** half of that reuse on this line's activations. The tensor contract derives every GDN shape from the frozen config and validates all 1658 headers, and the prefill half is now gated (next section), but no decode step has run against the `transformers` recurrence. That is slice B's M0 exit (#1127), and it has to cover several decode steps from a nonzero state (so `(v - k·s)` is live), finite checks before any reduction, and `out_proj`, which completes the block.
+
+## GDN prefill gate
+
+The prefill half runs against a committed component golden. `test_data/qwen38fn-gdn-golden.safetensors` (written by `tools/accuracy/dump_qwen38fn_gdn_golden.py`, a pure recorder of `Qwen4ExpTextGatedDeltaNet` under transformers 5.16.1 at the frozen geometry) stores, per case S=16/64/80/128, the pre-conv `mixed_qkv`, the reference's post-conv pre-repeat `query`/`key`/`value`, the output gate `z`, the derived pair `g`/`beta`, and the reference's gated-norm output `out` — plus the two weights a GDN kernel applies to them (`conv1d.weight`, `norm.weight`). S=80/128 cross the reference's 64-token chunk boundary, so the carried recurrent state is covered.
+
+`pegainfer-qwen35/tests/qwen38fn_gdn_prefill_parity.rs` replays those activations through the in-tree stack — `conv1d_prefill_batch_into` → the Triton AOT chunk-wise delta rule → the shared gated RMSNorm with the **sigmoid** activation this checkpoint's `output_gate_type` selects — and compares `out`. The conv leg is bit-exact (max abs 0.0 over 2.95M elements); the gated-norm output lands within 6.25e-2 (two bf16 ulps of values up to ~7, 30% of the tolerance budget) on every case, multi-chunk included, so the carried state adds no drift of its own. Three negative controls keep the gate from passing vacuously: silu instead of sigmoid (88–130x past the bound, 24.8–25.8% of elements out), a skipped conv leg (29–34x, 52.7–53.3%), and S=128 replayed as two single-chunk runs (the first half still matches, the tail misses by 15.4x). The fixture stores the gate pair *after* derivation, so the `a`/`b` → `g`/`beta` path, `out_proj`, decode, and real checkpoint weights are **not** covered here.
+
+```
+cargo test --release --locked -p pegainfer-qwen35 --features qwen35 \
+  --test qwen38fn_gdn_prefill_parity -- --nocapture
+```
+
+The gate replays recorded activations, so it needs no checkpoint. Like every `qwen35` gate it is not executed in CI — there is no GPU runner; the Qwen3.5 Clippy row compiles it only. To let it drive the chunk stages with a recorded gate pair, the Qwen3.5 chunk-wise prefill entry point is split into a prepare stage and a `from_gates` stage; production still runs prepare → the chunk stages in the same order over the same kernels.
 
 ## Traps
 
@@ -110,7 +123,7 @@ Still to verify is the reuse on this line's activations. The tensor contract der
 | Criterion | Blocked on |
 | --- | --- |
 | Tensor manifest enforced before the first H2D | slice B's loader, which must classify the headers it maps before any upload |
-| GDN decode reuse, executed | slice B: the wrappers are shared in `pegainfer-kernels::ops` and the contract pins their geometry and tensor names, but nothing has run this line's GDN block against the reference. The model-graph fixture that does so is slice B's M0 exit (#1127); the GDN prefill path and the gated norm composition land with it. |
+| GDN decode reuse, executed | slice B: the wrappers are shared in `pegainfer-kernels::ops` and the contract pins their geometry and tensor names; the prefill half (conv → chunk delta rule → sigmoid gated norm) is gated above. Decode steps from a nonzero state, `out_proj`, and the model-graph fixture are slice B's M0 exit (#1127). |
 | Bookend gate (embedding → mixer → lm_head → argmax) | 1 GPU, three shards (1 + 130 + 131, 5.29 GB, ~2.6 GB of weights on device) and a CPU f64 reference for the mixer. No layer weights needed. |
 | HF golden fixtures (logits + intermediate probes) | the fleet — 360 GB, plus ≥102.4 GB host RAM if the n-gram table is materialized |
 | Tokenizer / chat-template ID fixtures | transformers **5.16.0+** — the checkpoint's own tokenizer is `Qwen2Tokenizer`, which is older, but fixtures are compared against the reference implementation; nothing in the tree commits token-ID fixtures today |

pegainfer-qwen35/Cargo.toml

@@ -63,6 +63,10 @@ required-features = ["qwen35"]
 name = "serving_tp2"
 required-features = ["qwen35"]
 
+[[test]]
+name = "qwen38fn_gdn_prefill_parity"
+required-features = ["qwen35"]
+
 [[bench]]
 harness = false
 name = "qwen35_ops"

pegainfer-qwen35/src/lib.rs

@@ -72,8 +72,11 @@ pub mod runtime {
     pub use crate::weights::Qwen35Model;
 }
 
-/// Public operator surface used by Qwen3.5-local benches.
+/// Public operator surface used by Qwen3.5-local benches and gates.
 pub mod runtime_ops {
+    pub use crate::ops::conv1d_prefill_batch_into;
+    pub use crate::ops::gated_delta_rule_prefill_chunk_prepare_into;
+    pub use crate::ops::gated_delta_rule_prefill_chunkwise_from_gates_into;
     pub use crate::ops::gated_delta_rule_prefill_chunkwise_into;
     pub use crate::ops::rms_norm_batch_offset_into;
     pub use crate::ops::rms_norm_offset_into;

pegainfer-qwen35/src/ops.rs

@@ -23,7 +23,9 @@ pub(crate) use pegainfer_kernels::ops::paged_attention_batch_decode_hd256_into;
 pub(crate) use pegainfer_kernels::ops::paged_attention_batch_decode_split_hd256_into;
 pub(crate) use pegainfer_kernels::ops::paged_attention_batch_decode_via_prefill_hd256_into;
 pub(crate) use pegainfer_kernels::ops::qk_norm_partial_rope_batched_decode_hd256_into;
-pub(crate) use recurrent::conv1d_prefill_batch_into;
+pub use recurrent::conv1d_prefill_batch_into;
+pub use recurrent::gated_delta_rule_prefill_chunk_prepare_into;
+pub use recurrent::gated_delta_rule_prefill_chunkwise_from_gates_into;
 pub use recurrent::gated_delta_rule_prefill_chunkwise_into;
 
 use crate::recurrent;

pegainfer-qwen35/src/prefill_buffers.rs

@@ -20,10 +20,17 @@ use super::config::LocalGeometry;
 /// Buffers are explicit because the chunk-wise path is naturally a multi-stage
 /// pipeline rather than one opaque kernel launch.
 pub struct GdrChunkwiseScratch35 {
-    /// Chunk-local cumulative gate, fp32: [seq_len, num_value_heads]
-    pub(crate) g_cumsum: CudaSlice<f32>,
-    /// Beta values, fp32: [seq_len, num_value_heads]
-    pub(crate) beta: CudaSlice<f32>,
+    /// Raw per-token gate on entry to the chunk stages, folded in place into the
+    /// chunk-local cumulative gate, fp32: [seq_len, num_value_heads].
+    ///
+    /// `pub` because the prepare stage is what normally fills it: a caller that
+    /// drives the chunk stages with its own gate pair (see
+    /// `gated_delta_rule_prefill_chunkwise_from_gates_into`) writes the raw gate
+    /// here.
+    pub g_cumsum: CudaSlice<f32>,
+    /// Beta values, fp32: [seq_len, num_value_heads]. `pub` for the same reason
+    /// as `g_cumsum`.
+    pub beta: CudaSlice<f32>,
 
     /// Expanded + normalized q in token-major layout: [seq_len, num_value_heads * key_dim]
     pub(crate) q_expanded: HiddenStates,

pegainfer-qwen35/src/recurrent.rs

@@ -5,6 +5,8 @@ use cudarc::driver::DevicePtrMut;
 use pegainfer_core::tensor::DeviceContext;
 use pegainfer_core::tensor::DeviceVec;
 use pegainfer_core::tensor::HiddenStates;
+use pegainfer_kernels::ops::GDN_AOT_KEY_HEAD_DIM;
+use pegainfer_kernels::ops::GDN_AOT_VALUE_HEAD_DIM;
 
 use crate::ffi;
 use crate::prefill_buffers::GdrChunkwiseScratch35;
@@ -52,8 +54,11 @@ pub(crate) fn gated_delta_rule_decode_vec_into(
 }
 
 /// Causal depthwise conv1d prefill over a HiddenStates batch.
+///
+/// `pub` so the crate's operator surface (`runtime_ops`) can hand it to local
+/// gates and benches; the module itself stays crate-private.
 #[allow(clippy::too_many_arguments)]
-pub(crate) fn conv1d_prefill_batch_into(
+pub fn conv1d_prefill_batch_into(
     ctx: &DeviceContext,
     x_seq: &HiddenStates,
     conv_weight: &DeviceVec,
@@ -86,32 +91,61 @@ pub(crate) fn conv1d_prefill_batch_into(
     }
 }
 
+/// Prepare stage of the chunk-wise GDR prefill: head-expand q/k to the value-head
+/// count, L2-normalize them, split v, and derive the raw gate pair `g`/`beta`
+/// from the `a`/`b` projections, `dt_bias` and `A_log`.
+///
+/// This is the Qwen3.5-specific half of the operator — the gate derivation is
+/// what a line reusing the GDN stack replaces. On return `scratch.g_cumsum`
+/// still holds the *raw* per-token gate; the chunk-local cumsum folds it in
+/// [`gated_delta_rule_prefill_chunkwise_from_gates_into`].
 #[allow(clippy::too_many_arguments)]
-fn gated_delta_rule_prefill_chunk_prepare_into(
+pub fn gated_delta_rule_prefill_chunk_prepare_into(
     ctx: &DeviceContext,
     qkv: &HiddenStates,
     b_proj: &HiddenStates,
     a_proj: &HiddenStates,
     dt_bias: &DeviceVec,
     a_log: &CudaSlice<f32>,
-    q_out: &mut HiddenStates,
-    k_out: &mut HiddenStates,
-    v_out: &mut HiddenStates,
-    g_out: &mut CudaSlice<f32>,
-    beta_out: &mut CudaSlice<f32>,
+    scratch: &mut GdrChunkwiseScratch35,
     num_key_heads: usize,
     num_value_heads: usize,
 ) -> Result<()> {
+    let seq_len = qkv.seq_len;
+    assert_eq!(b_proj.seq_len, seq_len);
+    assert_eq!(a_proj.seq_len, seq_len);
+    assert_eq!(b_proj.hidden_dim, num_value_heads);
+    assert_eq!(a_proj.hidden_dim, num_value_heads);
+    assert_eq!(dt_bias.len, num_value_heads);
+    assert_eq!(a_log.len(), num_value_heads);
+    assert_eq!(scratch.q_expanded.seq_len, seq_len);
+    assert_eq!(scratch.k_expanded.seq_len, seq_len);
+    assert_eq!(scratch.v_raw.seq_len, seq_len);
+    assert_eq!(
+        scratch.q_expanded.hidden_dim,
+        num_value_heads * GDN_AOT_KEY_HEAD_DIM
+    );
+    assert_eq!(
+        scratch.k_expanded.hidden_dim,
+        num_value_heads * GDN_AOT_KEY_HEAD_DIM
+    );
+    assert_eq!(
+        scratch.v_raw.hidden_dim,
+        num_value_heads * GDN_AOT_VALUE_HEAD_DIM
+    );
+    assert_eq!(scratch.g_cumsum.len(), seq_len * num_value_heads);
+    assert_eq!(scratch.beta.len(), seq_len * num_value_heads);
+
     let (qkv_ptr, _gqkv) = qkv.data.device_ptr(&ctx.stream);
     let (b_ptr, _gb) = b_proj.data.device_ptr(&ctx.stream);
     let (a_ptr, _ga) = a_proj.data.device_ptr(&ctx.stream);
     let (dt_ptr, _gdt) = dt_bias.data.device_ptr(&ctx.stream);
     let (alog_ptr, _gal) = a_log.device_ptr(&ctx.stream);
-    let (q_out_ptr, _gqo) = q_out.data.device_ptr_mut(&ctx.stream);
-    let (k_out_ptr, _gko) = k_out.data.device_ptr_mut(&ctx.stream);
-    let (v_out_ptr, _gvo) = v_out.data.device_ptr_mut(&ctx.stream);
-    let (g_out_ptr, _ggo) = g_out.device_ptr_mut(&ctx.stream);
-    let (beta_out_ptr, _gbetao) = beta_out.device_ptr_mut(&ctx.stream);
+    let (q_out_ptr, _gqo) = scratch.q_expanded.data.device_ptr_mut(&ctx.stream);
+    let (k_out_ptr, _gko) = scratch.k_expanded.data.device_ptr_mut(&ctx.stream);
+    let (v_out_ptr, _gvo) = scratch.v_raw.data.device_ptr_mut(&ctx.stream);
+    let (g_out_ptr, _ggo) = scratch.g_cumsum.device_ptr_mut(&ctx.stream);
+    let (beta_out_ptr, _gbetao) = scratch.beta.device_ptr_mut(&ctx.stream);
 
     let result = unsafe {
         ffi::gated_delta_rule_prefill_chunk_prepare_cuda(
@@ -336,7 +370,9 @@ fn gated_delta_rule_prefill_chunk_o_stage_into(
 /// Chunk-wise GDR prefill operator contract for Qwen3.5.
 ///
 /// The chunk-wise path is an explicit multi-stage operator with pre-allocated
-/// scratch instead of one opaque kernel launch.
+/// scratch instead of one opaque kernel launch. This entry point is the prepare
+/// stage followed by the chunk stages; a caller that already holds the gate pair
+/// drives [`gated_delta_rule_prefill_chunkwise_from_gates_into`] directly.
 #[allow(clippy::too_many_arguments)]
 pub fn gated_delta_rule_prefill_chunkwise_into(
     ctx: &DeviceContext,
@@ -353,50 +389,88 @@ pub fn gated_delta_rule_prefill_chunkwise_into(
     key_dim: usize,
     val_dim: usize,
 ) -> Result<()> {
-    assert_eq!(scratch.q_expanded.seq_len, qkv.seq_len);
-    assert_eq!(scratch.k_expanded.seq_len, qkv.seq_len);
-    assert_eq!(scratch.v_raw.seq_len, qkv.seq_len);
-    assert_eq!(scratch.w.seq_len, qkv.seq_len);
-    assert_eq!(scratch.u.seq_len, qkv.seq_len);
-    assert_eq!(scratch.v_new.seq_len, qkv.seq_len);
+    assert_eq!(output.seq_len, qkv.seq_len);
+    assert_eq!(output.hidden_dim, num_value_heads * val_dim);
+
+    gated_delta_rule_prefill_chunk_prepare_into(
+        ctx,
+        qkv,
+        b_proj,
+        a_proj,
+        dt_bias,
+        a_log,
+        scratch,
+        num_key_heads,
+        num_value_heads,
+    )
+    .map_err(|e| anyhow::anyhow!("GDR prefill prepare failed: {e}"))?;
+    gated_delta_rule_prefill_chunkwise_from_gates_into(
+        ctx,
+        scratch,
+        state,
+        output,
+        num_value_heads,
+        key_dim,
+        val_dim,
+    )
+}
+
+/// The chunk stages of the chunk-wise GDR prefill, over a scratch that already
+/// carries the head-expanded q/k, the raw v and the raw gate pair: `g_cumsum`
+/// still holds the per-token `g` this function folds into its chunk-local prefix
+/// sum, and `beta` the per-token beta.
+///
+/// This is the model-agnostic half of the operator — everything downstream of
+/// the gate derivation, and the entry point a line reusing the GDN stack drives
+/// with its own `g`/`beta`. `state` is the recurrent state, read as the initial
+/// state and overwritten with the final one, so consecutive chunks carry it.
+pub fn gated_delta_rule_prefill_chunkwise_from_gates_into(
+    ctx: &DeviceContext,
+    scratch: &mut GdrChunkwiseScratch35,
+    state: &mut CudaSlice<f32>,
+    output: &mut HiddenStates,
+    num_value_heads: usize,
+    key_dim: usize,
+    val_dim: usize,
+) -> Result<()> {
+    let seq_len = scratch.q_expanded.seq_len;
+    assert_eq!(output.seq_len, seq_len);
+    assert_eq!(output.hidden_dim, num_value_heads * val_dim);
+    assert_eq!(scratch.k_expanded.seq_len, seq_len);
+    assert_eq!(scratch.v_raw.seq_len, seq_len);
+    assert_eq!(scratch.w.seq_len, seq_len);
+    assert_eq!(scratch.u.seq_len, seq_len);
+    assert_eq!(scratch.v_new.seq_len, seq_len);
     assert_eq!(scratch.q_expanded.hidden_dim, num_value_heads * key_dim);
     assert_eq!(scratch.k_expanded.hidden_dim, num_value_heads * key_dim);
     assert_eq!(scratch.v_raw.hidden_dim, num_value_heads * val_dim);
     assert_eq!(scratch.w.hidden_dim, num_value_heads * key_dim);
     assert_eq!(scratch.u.hidden_dim, num_value_heads * val_dim);
     assert_eq!(scratch.v_new.hidden_dim, num_value_heads * val_dim);
 
-    let expected_gate_len = qkv.seq_len * num_value_heads;
-    let expected_chunk_a_len = qkv.seq_len * num_value_heads * GdrChunkwiseScratch35::CHUNK_SIZE;
+    let expected_gate_len = seq_len * num_value_heads;
+    let expected_chunk_a_len = seq_len * num_value_heads * GdrChunkwiseScratch35::CHUNK_SIZE;
     let expected_chunk_ai_len = expected_chunk_a_len;
     let expected_chunk_state_len =
-        GdrChunkwiseScratch35::num_chunks(qkv.seq_len) * num_value_heads * val_dim * key_dim;
+        GdrChunkwiseScratch35::num_chunks(seq_len) * num_value_heads * val_dim * key_dim;
+    // The head dims are compiled into every chunk kernel (the Triton AOT
+    // specializations bake 128/128); the checks above only tie them to the
+    // caller's own scratch layout, so pin them against the compiled pair, and
+    // the recurrent state — which this function carries across chunks — against
+    // the one state slot a chunk fills.
+    assert_eq!(key_dim, GDN_AOT_KEY_HEAD_DIM);
+    assert_eq!(val_dim, GDN_AOT_VALUE_HEAD_DIM);
+    assert_eq!(state.len(), num_value_heads * val_dim * key_dim);
     assert_eq!(scratch.g_cumsum.len(), expected_gate_len);
     assert_eq!(scratch.beta.len(), expected_gate_len);
     assert_eq!(scratch.a_tril.len(), expected_chunk_a_len);
     assert_eq!(scratch.a_inv.len(), expected_chunk_ai_len);
     assert_eq!(scratch.chunk_state.len(), expected_chunk_state_len);
 
-    gated_delta_rule_prefill_chunk_prepare_into(
-        ctx,
-        qkv,
-        b_proj,
-        a_proj,
-        dt_bias,
-        a_log,
-        &mut scratch.q_expanded,
-        &mut scratch.k_expanded,
-        &mut scratch.v_raw,
-        &mut scratch.g_cumsum,
-        &mut scratch.beta,
-        num_key_heads,
-        num_value_heads,
-    )
-    .map_err(|e| anyhow::anyhow!("GDR prefill prepare failed: {e}"))?;
     gated_delta_rule_prefill_chunk_cumsum_inplace(
         ctx,
         &mut scratch.g_cumsum,
-        qkv.seq_len,
+        seq_len,
         num_value_heads,
     )
     .map_err(|e| anyhow::anyhow!("GDR prefill cumsum failed: {e}"))?;
@@ -413,7 +487,7 @@ pub fn gated_delta_rule_prefill_chunkwise_into(
         ctx,
         &scratch.a_tril,
         &mut scratch.a_inv,
-        qkv.seq_len,
+        seq_len,
         num_value_heads,
     )
     .map_err(|e| anyhow::anyhow!("GDR prefill solve failed: {e}"))?;