PMAT-3477: Qwen3.5 4B/9B/27B on CPU and CUDA — block_count loader fix + Gated DeltaNet GQA (tiled head mapping)

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<f32>,
         state: &GpuBuffer<f32>,
         output: &GpuBuffer<f32>,
+        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();

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()

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,

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)

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<Option<(GpuBuffer<f32>, GpuBuffer<f32>)>>,
     /// `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<Self> {
         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<Qwen35CudaState> {
         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,
         )?;