diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/inference/ggml_graph.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/inference/ggml_graph.rs index 43e2263fc..d1ef7a824 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/inference/ggml_graph.rs +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/inference/ggml_graph.rs @@ -630,6 +630,11 @@ pub struct DecodeBuffers { pub kv_v_cache: Vec, /// 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> = 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::sync::OnceLock::new(); let _wc = WEIGHT_CACHE.get_or_init(std::collections::HashMap::new); diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/inference/llm.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/inference/llm.rs index 778685d22..ba80c0767 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/inference/llm.rs +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/inference/llm.rs @@ -189,7 +189,7 @@ impl LLMPipeline { static GRAPH: std::sync::OnceLock = 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"); diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/tests/ggml_kv_cache.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/tests/ggml_kv_cache.rs new file mode 100644 index 000000000..728acf141 --- /dev/null +++ b/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); +}