jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/inference/ggml_graph.rs

@@ -630,6 +630,11 @@ pub struct DecodeBuffers {
     pub kv_v_cache: Vec<metal::Buffer>,
     /// Temp F16 buffer for F32→F16 conversion: [`num_kv_heads` * `head_dim`] F16
     pub attn_kv_f16: metal::Buffer,
+    /// Positions each KV cache row can hold, and how many attention
+    /// layers were allocated. Kept so the decode path can check a model
+    /// against the capacity instead of assuming it fits.
+    pub kv_max_seq: usize,
+    pub kv_attn_layers: usize,
     /// Causal mask for flash attention: [`max_seq_len`] F16
     /// mask[i] = 0 for i < `seq_len` (attend), -INF for i >= `seq_len` (ignore)
     pub attn_mask: metal::Buffer,
@@ -643,7 +648,15 @@ pub struct DecodeBuffers {
 #[cfg(feature = "metal")]
 impl DecodeBuffers {
     /// Allocate all decode buffers from model config.
-    pub fn new(device: &metal::DeviceRef, config: &crate::inference::model::ModelConfig) -> Self {
+    pub fn new(
+        device: &metal::DeviceRef,
+        config: &crate::inference::model::ModelConfig,
+        num_attn_layers: usize,
+    ) -> Self {
+        // A config that never had max_seq_len filled in would otherwise
+        // allocate a zero-length cache and fail every write.
+        let kv_max_seq = if config.max_seq_len == 0 { 4096 } else { config.max_seq_len };
+        let kv_attn_layers = num_attn_layers;
         let f32_buf = |size: usize| -> metal::Buffer {
             let buf = device.new_buffer((size * 4) as u64, metal::MTLResourceOptions::StorageModeShared);
             // Zero-initialize — critical for SSM state buffers (garbage causes NaN in scan)
@@ -738,12 +751,17 @@ impl DecodeBuffers {
                 let head_dim_ssm = d_inner.checked_div(n_heads_ssm).unwrap_or(64);
                 f32_buf(d_inner + d_state * head_dim_ssm * n_heads_ssm)
             },
-            // KV cache: allocate per attention layer
-            // Nemotron has 6 attention layers. max_seq_len=4096 default.
-            // Zero-fill to prevent NaN from garbage data in unused positions.
+            // KV cache: one buffer per ATTENTION layer, not per layer. The
+            // count used to be the literal 6 ("Nemotron: 6 attention layers")
+            // and the length the literal 4096, while the write site indexes
+            // with the model's real attention-layer ordinal and the real
+            // position. A model with more attention layers had its KV writes
+            // silently skipped by a bounds guard, and a position past 4096
+            // addressed past the end of the buffer. Both now come from the
+            // model, and encode_decode_step checks them.
             kv_k_cache: {
-                let max_seq = 4096usize;
-                let num_attn_layers = 6; // Nemotron: 6 attention layers
+                let max_seq = kv_max_seq;
+                let num_attn_layers = kv_attn_layers;
                 let kv_row_bytes = (num_kv_heads * head_dim * 2) as u64; // F16
                 let total = max_seq as u64 * kv_row_bytes;
                 (0..num_attn_layers).map(|_| {
@@ -753,8 +771,8 @@ impl DecodeBuffers {
                 }).collect()
             },
             kv_v_cache: {
-                let max_seq = 4096usize;
-                let num_attn_layers = 6;
+                let max_seq = kv_max_seq;
+                let num_attn_layers = kv_attn_layers;
                 let kv_row_bytes = (num_kv_heads * head_dim * 2) as u64;
                 let total = max_seq as u64 * kv_row_bytes;
                 (0..num_attn_layers).map(|_| {
@@ -764,6 +782,8 @@ impl DecodeBuffers {
                 }).collect()
             },
             attn_kv_f16: device.new_buffer((num_kv_heads * head_dim * 2) as u64, metal::MTLResourceOptions::StorageModeShared),
+            kv_max_seq,
+            kv_attn_layers,
             // Causal mask: [max_seq_len] F16, pre-filled with -INF (0xFC00 = -inf in F16)
             // Each decode step sets mask[position] = 0 to "unlock" that position.
             attn_mask: {
@@ -876,12 +896,23 @@ impl GgmlGraph {
 
     /// Initialize shader + buffers.
     #[cfg(feature = "metal")]
-    pub fn init(&mut self, device: &metal::DeviceRef, config: &crate::inference::model::ModelConfig) -> Result<()> {
+    pub fn init(
+        &mut self,
+        device: &metal::DeviceRef,
+        config: &crate::inference::model::ModelConfig,
+        model: &crate::inference::model::Model,
+    ) -> Result<()> {
         if self.shader.is_none() {
             self.shader = Some(GgmlShaderLibrary::new(device)?);
         }
         if self.buffers.is_none() {
-            self.buffers = Some(DecodeBuffers::new(device, config));
+            // The KV cache is indexed by attention-layer ordinal, so its
+            // length has to be this model's attention-layer count.
+            let num_attn_layers = resolve_layer_types(model, config.num_layers)
+                .iter()
+                .filter(|l| l.layer_type == LayerType::Attention)
+                .count();
+            self.buffers = Some(DecodeBuffers::new(device, config, num_attn_layers));
         }
         Ok(())
     }
@@ -1363,6 +1394,29 @@ pub fn encode_decode_step(
     static LAYERS_CACHE: std::sync::OnceLock<Vec<LayerInfo>> = std::sync::OnceLock::new();
     let layers = LAYERS_CACHE.get_or_init(|| resolve_layer_types(_model, config.num_layers));
 
+    // The KV cache is indexed by attention-layer ordinal and by position, and
+    // both used to be assumed rather than checked: a model with more attention
+    // layers than the cache had its writes dropped by an `if idx < len` guard,
+    // and a position past the cache length addressed past the end of the
+    // buffer. Check once, here, where it can be reported.
+    let attn_layers = layers
+        .iter()
+        .filter(|l| l.layer_type == LayerType::Attention)
+        .count();
+    if attn_layers > bufs.kv_attn_layers {
+        return Err(crate::core::Error::internal(format!(
+            "ggml decode: this model has {attn_layers} attention layers but the KV cache \
+             was built for {}; the extra layers' writes would be dropped silently",
+            bufs.kv_attn_layers
+        )));
+    }
+    if _position >= bufs.kv_max_seq {
+        return Err(crate::core::Error::internal(format!(
+            "ggml decode: position {_position} is past the KV cache length {}",
+            bufs.kv_max_seq
+        )));
+    }
+
     // Cache weight buffer pointers (stable across tokens — weights don't move)
     static WEIGHT_CACHE: std::sync::OnceLock<std::collections::HashMap<String, usize>> = std::sync::OnceLock::new();
     let _wc = WEIGHT_CACHE.get_or_init(std::collections::HashMap::new);

jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/inference/llm.rs

@@ -189,7 +189,7 @@ impl LLMPipeline {
             static GRAPH: std::sync::OnceLock<ggml_graph::GgmlGraph> = std::sync::OnceLock::new();
             let graph_ref = GRAPH.get_or_init(|| {
                 let mut g = ggml_graph::GgmlGraph::new();
-                if let Err(e) = g.init(self.compute.device().raw(), config) {
+                if let Err(e) = g.init(self.compute.device().raw(), config, &self.model) {
                     eprintln!("[ggml] init failed: {e}");
                 } else {
                     eprintln!("[ggml] graph initialized");

jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/tests/ggml_kv_cache.rs

@@ -0,0 +1,86 @@
+//! The ggml decode path's KV cache is indexed by attention-layer ordinal and
+//! by position. Both were literals -- 6 layers and 4096 positions -- while the
+//! write site used the model's real ordinal and the real position, so a model
+//! with more attention layers had its writes dropped by an `if idx < len`
+//! guard, and a position past 4096 addressed past the end of the buffer.
+//!
+//! These tests pin the two lengths against what was asked for. They need a
+//! Metal device and do not skip without one.
+
+#![cfg(feature = "metal")]
+
+use jouleclaw_omni::hal::metal::MetalDevice;
+use jouleclaw_omni::inference::ggml_graph::DecodeBuffers;
+use jouleclaw_omni::inference::model::ModelConfig;
+
+/// A small but complete hybrid config. Every field below is read by
+/// `DecodeBuffers::new` to size a buffer, and Metal rejects a zero-length
+/// allocation, so none of them may be left at its default of 0.
+fn config(num_layers: usize, max_seq_len: usize) -> ModelConfig {
+    ModelConfig {
+        num_layers,
+        max_seq_len,
+        hidden_size: 64,
+        num_heads: 4,
+        num_kv_heads: 2,
+        intermediate_size: 128,
+        num_experts: 4,
+        num_active_experts: 2,
+        ssm_inner_size: 32,
+        ssm_state_size: 8,
+        ssm_group_count: 2,
+        ssm_time_step_rank: 4,
+        ssm_conv_kernel: 4,
+        vocab_size: Some(1024),
+        ..Default::default()
+    }
+}
+
+#[test]
+fn kv_cache_has_one_buffer_per_attention_layer() {
+    let device = MetalDevice::new().expect("a Metal device is required by these tests");
+
+    // The count is the caller's, not a constant: 6 was the literal that made a
+    // 7th attention layer's writes vanish.
+    for attn_layers in [1usize, 6, 7, 11] {
+        let bufs = DecodeBuffers::new(device.raw(), &config(52, 4096), attn_layers);
+        assert_eq!(
+            bufs.kv_k_cache.len(), attn_layers,
+            "K cache must hold one buffer per attention layer"
+        );
+        assert_eq!(
+            bufs.kv_v_cache.len(), attn_layers,
+            "V cache must hold one buffer per attention layer"
+        );
+        assert_eq!(bufs.kv_attn_layers, attn_layers, "recorded capacity must match");
+    }
+}
+
+#[test]
+fn kv_cache_length_follows_the_config_not_a_literal() {
+    let device = MetalDevice::new().expect("a Metal device is required by these tests");
+
+    // 4096 was the literal. A model configured for more must get more, or its
+    // later positions write past the end of the buffer.
+    for max_seq in [512usize, 4096, 8192] {
+        let bufs = DecodeBuffers::new(device.raw(), &config(52, max_seq), 2);
+        assert_eq!(bufs.kv_max_seq, max_seq, "recorded length must follow the config");
+        // Two bytes per F16 element, num_kv_heads * head_dim elements per row.
+        let head_dim = 64 / 4;
+        let row = 2 * head_dim * 2;
+        assert_eq!(
+            bufs.kv_k_cache[0].length() as usize, max_seq * row,
+            "each buffer must be long enough for every position"
+        );
+    }
+}
+
+#[test]
+fn a_config_with_no_max_seq_len_still_allocates() {
+    let device = MetalDevice::new().expect("a Metal device is required by these tests");
+    // max_seq_len == 0 would allocate a zero-length buffer and fail every
+    // write, so it falls back rather than producing an unusable cache.
+    let bufs = DecodeBuffers::new(device.raw(), &config(52, 0), 2);
+    assert!(bufs.kv_max_seq > 0, "a zero-length KV cache is never usable");
+    assert!(bufs.kv_k_cache[0].length() > 0);
+}