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

//! 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);
}