From 950634f2003cb6bc1d6939ed96a3f07568ab4497 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 16:32:06 -0400 Subject: [PATCH] ggml decode: the KV cache was sized by two literals the write site ignored `DecodeBuffers::new` built the KV cache with let max_seq = 4096usize; let num_attn_layers = 6; // Nemotron: 6 attention layers while `encode_decode_step` indexes it with the model's real attention-layer ordinal and the real position: let attn_idx = layers.iter().take(layer.idx + 1) .filter(|l| l.layer_type == LayerType::Attention).count() - 1; let dst_offset = _position as u64 * kv_row_f16; Two consequences, neither of which reported anything: - A model with a seventh attention layer got `attn_idx == 6`, the `if attn_idx < bufs.kv_k_cache.len()` guard was false, and that layer's K and V were simply never written. Attention then read a zero-filled cache: wrong output, no error. - `_position` is never bounded anywhere in the function, so a position at or past 4096 gave `set_buffer` an offset past the end of the buffer. Both lengths now come from the model. `init` resolves the layer types and passes the attention-layer count; `max_seq` comes from `config.max_seq_len`, falling back only when it is 0, which would otherwise allocate a zero-length cache that fails every write. The capacity is recorded on `DecodeBuffers` so the decode path can check rather than assume, and `encode_decode_step` now checks both once, up front, where an error can propagate -- the per-site `if idx < len` guards are deep inside nested closures where a `return Err` cannot. tests/ggml_kv_cache.rs pins the two lengths: one buffer per attention layer for counts of 1, 6, 7 and 11 (6 being the literal that hid this), buffer length following max_seq_len for 512, 4096 and 8192, and a zero-length config still producing a usable cache. Scope: this path is behind the GGML_GRAPH environment variable, so it is opt-in rather than the default decode path. Still assumed, and worth separate attention: `GRAPH` and `LAYERS_CACHE` are process-global `OnceLock`s, so a second model loaded in the same process reuses the first one's buffers and layer pattern. Co-Authored-By: Claude Opus 5 --- .../src/inference/ggml_graph.rs | 74 +++++++++++++--- .../jouleclaw-omni/src/inference/llm.rs | 2 +- .../jouleclaw-omni/tests/ggml_kv_cache.rs | 86 +++++++++++++++++++ 3 files changed, 151 insertions(+), 11 deletions(-) create mode 100644 jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/tests/ggml_kv_cache.rs 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); +} -- 2.54.0 (Apple Git-157)