ghostlink-core 2.4.0

//! KV cache primitives for a single sequence within a single transformer layer.
//!
//! Stores one contiguous, pre-allocated buffer for all tokens' keys and one
//! for all values (shape: `max_tokens_per_seq * num_heads * hidden_dim_per_head`
//! floats each), instead of one small heap allocation per token. The previous
//! design called `resize_with` on a `Vec<KVCacheEntry>` where every entry
//! independently allocated its own `keys`/`values` `Vec<f32>` — for the
//! default 8192-token max sequence length, a full-length `initialize()`
//! performed up to 16,384 separate small heap allocations. A single upfront
//! allocation eliminates that, keeps a sequence's KV data in one contiguous
//! cache-friendly region, and enables `read_range` — a single batched read
//! over a token span, the shape attention actually needs, instead of N
//! separate per-token reads.
//!
//! Reads and writes are guarded by an `RwLock` rather than a `Mutex`: real
//! attention reads happen far more often than writes (every position reads
//! all previous positions' KV every decode step; only one write happens per
//! step), so letting reads proceed concurrently with each other matters.
//!
//! Zero-copy read paths (`with_read_kv`, `with_read_range`, `read_range_into`)
//! allow inspecting or copying cached keys and values directly without owned
//! heap allocations.
//!
//! ## Read API: prefer the zero-copy closure forms
//!
//! `with_read_kv` / `with_read_range` are the intended read API — they hand
//! out borrowed slices under the read lock and allocate nothing. The owned
//! forms `read_kv` / `read_range` exist only as a reference implementation and
//! as the equivalence oracle for the zero-copy paths' tests; they are
//! `#[doc(hidden)]` and deliberately **not** part of the public surface.
//!
//! Measured on the reference host (`cargo bench -p ghostlink-core --bench
//! kv_cache`, `wide` config), the difference is not a micro-optimization:
//!
//! | operation | owned | zero-copy |
//! |---|---|---|
//! | `read_range_1024` | 4.57 ms | 12.8 ns |
//! | `decode_step_256` | 141.9 ms | 1.26 ms |
//! | `concurrent_4_readers_256` | 1.63 ms | 221 ns |
//!
//! A naive wiring of the owned paths into a decode loop costs 100x-350,000x.
//! Use the closure forms.
//!
//! This module has no current caller in `runtime.rs` — Ghostlink delegates
//! actual model execution to an external inference engine (llama-server /
//! Ollama, see `ghost-link::native_engine`) rather than computing attention
//! itself. It stays self-contained as a ready-to-use primitive for a future
//! local execution path, so it can evolve without affecting the stable
//! public API exported from `lib.rs`.

use std::sync::{Arc, RwLock};

pub const DEFAULT_MAX_TOKENS_PER_SEQ: usize = 8192;
pub const DEFAULT_NUM_HEADS: usize = 4;
pub const DEFAULT_HIDDEN_DIM_PER_HEAD: usize = 64;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct KVCacheConfig {
    pub max_tokens_per_seq: usize,
    pub num_heads: usize,
    pub hidden_dim_per_head: usize,
}

impl KVCacheConfig {
    /// Number of `f32` elements one token's keys (or values) occupy.
    pub const fn token_width(&self) -> usize {
        self.num_heads * self.hidden_dim_per_head
    }
}

impl Default for KVCacheConfig {
    fn default() -> Self {
        Self {
            max_tokens_per_seq: DEFAULT_MAX_TOKENS_PER_SEQ,
            num_heads: DEFAULT_NUM_HEADS,
            hidden_dim_per_head: DEFAULT_HIDDEN_DIM_PER_HEAD,
        }
    }
}

/// Owned copy of one token's cached key/value vectors, returned by `read_kv`.
/// Small and fixed-size (`config.token_width()` floats each) — not a copy of
/// the whole cache.
#[derive(Clone, Debug, PartialEq)]
pub struct KVCacheEntry {
    pub keys: Vec<f32>,
    pub values: Vec<f32>,
}

/// Borrowed slice view over a contiguous token range's cached keys and values.
#[derive(Debug, PartialEq)]
pub struct KVSpan<'a> {
    pub keys: &'a [f32],
    pub values: &'a [f32],
}

struct KvCacheState {
    /// Empty until the first `initialize()` call; allocated exactly once
    /// after that, sized for `config.max_tokens_per_seq` tokens.
    keys: Vec<f32>,
    values: Vec<f32>,
    /// Number of tokens actually written so far (<= config.max_tokens_per_seq).
    current_len: usize,
}

/// Cheap to clone (an `Arc` around the shared, lock-guarded buffer) — matches
/// the shared-state pattern used elsewhere in this crate (e.g. `ClusterState`).
#[derive(Clone)]
pub struct LayerKvCache {
    pub config: KVCacheConfig,
    state: Arc<RwLock<KvCacheState>>,
}

impl std::fmt::Debug for LayerKvCache {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        // Deliberately does not print the backing buffers — they can be
        // megabytes (max_tokens_per_seq * token_width * 2 floats).
        f.debug_struct("LayerKvCache")
            .field("config", &self.config)
            .field("current_len", &self.len())
            .finish()
    }
}

impl Default for LayerKvCache {
    fn default() -> Self {
        Self::new(KVCacheConfig::default())
    }
}

impl LayerKvCache {
    /// Creates an empty cache. Does not allocate the backing buffer yet —
    /// `initialize()` performs the single upfront allocation, sized once for
    /// `config.max_tokens_per_seq`.
    pub fn new(config: KVCacheConfig) -> Self {
        Self {
            config,
            state: Arc::new(RwLock::new(KvCacheState {
                keys: Vec::new(),
                values: Vec::new(),
                current_len: 0,
            })),
        }
    }

    /// Ensures the backing buffer exists (allocating it on the first call
    /// only) and marks the first `seq_len` tokens (clamped to
    /// `config.max_tokens_per_seq`) as valid.
    ///
    /// Safe to call again on the same cache — e.g. to grow `current_len` as
    /// more tokens are generated — without losing previously written data:
    /// the buffer is allocated once and never reallocated by this method.
    pub fn initialize(&self, seq_len: usize) -> Result<(), String> {
        if seq_len == 0 {
            return Err("sequence length must be greater than zero".to_string());
        }

        let bounded_len = seq_len.min(self.config.max_tokens_per_seq);
        let width = self.config.token_width();
        let mut state = self
            .state
            .write()
            .map_err(|_| "kv cache lock poisoned".to_string())?;

        if state.keys.is_empty() {
            let capacity = self.config.max_tokens_per_seq * width;
            state.keys = vec![0.0; capacity];
            state.values = vec![0.0; capacity];
        }

        state.current_len = state.current_len.max(bounded_len);
        Ok(())
    }

    /// Writes one token's key/value vectors into the cache at `token_idx`.
    /// Copies directly into the pre-allocated buffer slice — no allocation.
    pub fn write_kv(&self, token_idx: usize, keys: &[f32], values: &[f32]) -> Result<(), String> {
        if keys.len() != values.len() {
            return Err("keys/values length mismatch".to_string());
        }
        let width = self.config.token_width();
        if keys.len() != width {
            return Err(format!(
                "expected {width} elements per token (num_heads * hidden_dim_per_head), got {}",
                keys.len()
            ));
        }
        if token_idx >= self.config.max_tokens_per_seq {
            return Err(format!(
                "token index {token_idx} out of bounds (max {})",
                self.config.max_tokens_per_seq
            ));
        }

        // Pre-compute slice range bounds once upfront to reduce arithmetic
        // calculations inside the lock-holding section.
        let start = token_idx * width;
        let end = start + width;

        let mut state = self
            .state
            .write()
            .map_err(|_| "kv cache lock poisoned".to_string())?;
        if state.keys.is_empty() {
            return Err("cache not initialized — call initialize() first".to_string());
        }

        state.keys[start..end].copy_from_slice(keys);
        state.values[start..end].copy_from_slice(values);
        state.current_len = state.current_len.max(token_idx + 1);
        Ok(())
    }

    /// Writes multiple tokens' key/value vectors under a single write-lock
    /// acquisition, validating every entry upfront (before taking the lock)
    /// so a bad entry anywhere in the batch fails the whole call without
    /// partially writing it. Same per-token semantics as calling `write_kv`
    /// in a loop — this only changes how many times the lock is acquired.
    pub fn write_kv_batch(&self, entries: &[(usize, &[f32], &[f32])]) -> Result<(), String> {
        if entries.is_empty() {
            return Ok(());
        }

        let width = self.config.token_width();

        // Validate all entries upfront before taking write lock
        for (token_idx, keys, values) in entries {
            if keys.len() != values.len() {
                return Err("keys/values length mismatch".to_string());
            }
            if keys.len() != width {
                return Err(format!("expected {width} elements, got {}", keys.len()));
            }
            if *token_idx >= self.config.max_tokens_per_seq {
                return Err(format!("token index {token_idx} out of bounds"));
            }
        }

        // Single write lock acquisition
        let mut state = self
            .state
            .write()
            .map_err(|_| "kv cache lock poisoned".to_string())?;
        if state.keys.is_empty() {
            return Err("cache not initialized".to_string());
        }

        let mut max_token_idx = 0;
        // Direct reference destructuring and upfront slice range calculation
        // avoid redundant arithmetic and pointer dereferences in the hot loop.
        for &(token_idx, keys, values) in entries {
            let start = token_idx * width;
            let end = start + width;
            state.keys[start..end].copy_from_slice(keys);
            state.values[start..end].copy_from_slice(values);
            if token_idx > max_token_idx {
                max_token_idx = token_idx;
            }
        }

        state.current_len = state.current_len.max(max_token_idx + 1);
        Ok(())
    }

    /// Inspects one token's cached key/value slices directly under a read lock via a closure.
    /// Zero heap allocation on this path.
    pub fn with_read_kv<F, R>(&self, token_idx: usize, f: F) -> Result<R, String>
    where
        F: FnOnce(&[f32], &[f32]) -> R,
    {
        let width = self.config.token_width();
        let state = self
            .state
            .read()
            .map_err(|_| "kv cache lock poisoned".to_string())?;

        if token_idx >= state.current_len {
            return Err(format!(
                "token index {token_idx} out of bounds (current length {})",
                state.current_len
            ));
        }

        let start = token_idx * width;
        let end = start + width;
        Ok(f(&state.keys[start..end], &state.values[start..end]))
    }

    /// Reads one token's cached key/value vectors as an owned, fixed-size
    /// copy (`config.token_width()` floats each — not the whole cache).
    ///
    /// **Prefer [`Self::with_read_kv`]** — this allocates two `Vec`s per call
    /// and is kept only as a reference implementation / test oracle (see the
    /// module docs for the measured cost). Not part of the public API.
    #[doc(hidden)]
    pub fn read_kv(&self, token_idx: usize) -> Result<KVCacheEntry, String> {
        self.with_read_kv(token_idx, |keys, values| KVCacheEntry {
            keys: keys.to_vec(),
            values: values.to_vec(),
        })
    }

    /// Inspects keys/values for a contiguous token range (`[start_token, end_token)`)
    /// directly under a read lock via a closure that accepts a `KVSpan`.
    /// Zero heap allocation on this path.
    pub fn with_read_range<F, R>(
        &self,
        start_token: usize,
        end_token: usize,
        f: F,
    ) -> Result<R, String>
    where
        F: FnOnce(KVSpan<'_>) -> R,
    {
        if start_token > end_token {
            return Err("start_token must be <= end_token".to_string());
        }
        let width = self.config.token_width();
        let state = self
            .state
            .read()
            .map_err(|_| "kv cache lock poisoned".to_string())?;

        if end_token > state.current_len {
            return Err(format!(
                "end_token {end_token} exceeds current length {}",
                state.current_len
            ));
        }

        let start = start_token * width;
        let end = end_token * width;
        let span = KVSpan {
            keys: &state.keys[start..end],
            values: &state.values[start..end],
        };
        Ok(f(span))
    }

    /// Reads keys/values for a contiguous token range (`[start_token,
    /// end_token)`) in one call — the shape attention actually wants (all
    /// prior positions at once), instead of `end_token - start_token`
    /// separate lock acquisitions and small allocations.
    ///
    /// **Prefer [`Self::with_read_range`] or [`Self::read_range_into`]** —
    /// this allocates two `Vec`s sized to the whole span on every call, which
    /// in a decode loop is 100x-350,000x slower than the zero-copy forms (see
    /// the module docs). Kept only as a reference implementation / test
    /// oracle. Not part of the public API.
    #[doc(hidden)]
    pub fn read_range(
        &self,
        start_token: usize,
        end_token: usize,
    ) -> Result<(Vec<f32>, Vec<f32>), String> {
        self.with_read_range(start_token, end_token, |span| {
            (span.keys.to_vec(), span.values.to_vec())
        })
    }

    /// Copies keys and values for a contiguous token range into pre-allocated
    /// caller slices `keys_out` and `values_out`.
    /// Rejects output buffers if they are too small for the requested token range.
    pub fn read_range_into(
        &self,
        start_token: usize,
        end_token: usize,
        keys_out: &mut [f32],
        values_out: &mut [f32],
    ) -> Result<(), String> {
        if start_token > end_token {
            return Err("start_token must be <= end_token".to_string());
        }
        let num_tokens = end_token - start_token;
        let width = self.config.token_width();
        let required_len = num_tokens * width;

        if keys_out.len() < required_len {
            return Err(format!(
                "keys_out buffer too short: expected at least {required_len} elements, got {}",
                keys_out.len()
            ));
        }
        if values_out.len() < required_len {
            return Err(format!(
                "values_out buffer too short: expected at least {required_len} elements, got {}",
                values_out.len()
            ));
        }

        self.with_read_range(start_token, end_token, |span| {
            keys_out[..required_len].copy_from_slice(span.keys);
            values_out[..required_len].copy_from_slice(span.values);
        })
    }

    pub fn len(&self) -> usize {
        self.state.read().map(|s| s.current_len).unwrap_or(0)
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use proptest::prelude::*;

    fn small_config() -> KVCacheConfig {
        // Small width keeps test assertions easy to read; 2 heads * 4 dims = 8.
        KVCacheConfig {
            max_tokens_per_seq: 16,
            num_heads: 2,
            hidden_dim_per_head: 4,
        }
    }

    #[test]
    fn initialize_rejects_zero_len() {
        let cache = LayerKvCache::default();
        assert!(cache.initialize(0).is_err());
    }

    #[test]
    fn initialize_and_round_trip() {
        let cache = LayerKvCache::new(small_config());
        cache.initialize(8).expect("cache should initialize");
        assert_eq!(cache.len(), 8);

        let width = cache.config.token_width();
        let keys = vec![1.0_f32; width];
        let values = vec![2.0_f32; width];
        cache
            .write_kv(2, &keys, &values)
            .expect("write should succeed");

        let entry = cache.read_kv(2).expect("entry should exist");
        assert_eq!(entry.keys, keys);
        assert_eq!(entry.values, values);
    }

    #[test]
    fn growing_seq_len_preserves_prior_writes() {
        let cache = LayerKvCache::new(small_config());
        cache.initialize(4).unwrap();

        let width = cache.config.token_width();
        let keys = vec![7.0_f32; width];
        let values = vec![9.0_f32; width];
        cache.write_kv(1, &keys, &values).unwrap();

        // Growing current_len must not touch the already-allocated buffer,
        // so token 1's data must survive.
        cache.initialize(10).unwrap();
        assert_eq!(cache.len(), 10);

        let entry = cache.read_kv(1).unwrap();
        assert_eq!(entry.keys, keys);
        assert_eq!(entry.values, values);
    }

    #[test]
    fn write_kv_rejects_wrong_width() {
        let cache = LayerKvCache::new(small_config());
        cache.initialize(4).unwrap();
        let wrong_width = vec![0.0_f32; cache.config.token_width() + 1];
        assert!(cache.write_kv(0, &wrong_width, &wrong_width).is_err());
    }

    #[test]
    fn write_kv_rejects_out_of_bounds_token() {
        let cache = LayerKvCache::new(small_config());
        cache.initialize(4).unwrap();
        let width = cache.config.token_width();
        let data = vec![0.0_f32; width];
        assert!(cache.write_kv(999, &data, &data).is_err());
    }

    #[test]
    fn write_kv_before_initialize_is_rejected() {
        let cache = LayerKvCache::new(small_config());
        let width = cache.config.token_width();
        let data = vec![0.0_f32; width];
        assert!(cache.write_kv(0, &data, &data).is_err());
    }

    #[test]
    fn read_range_returns_contiguous_data() {
        let cache = LayerKvCache::new(small_config());
        cache.initialize(4).unwrap();
        let width = cache.config.token_width();

        for token_idx in 0..4 {
            let keys = vec![token_idx as f32; width];
            let values = vec![(token_idx * 10) as f32; width];
            cache.write_kv(token_idx, &keys, &values).unwrap();
        }

        let (keys, values) = cache.read_range(1, 3).unwrap();
        assert_eq!(keys.len(), 2 * width);
        assert_eq!(values.len(), 2 * width);
        // First token in the range is token_idx 1.
        assert_eq!(keys[0], 1.0);
        assert_eq!(values[0], 10.0);
        // Second token in the range is token_idx 2.
        assert_eq!(keys[width], 2.0);
        assert_eq!(values[width], 20.0);
    }

    #[test]
    fn read_range_rejects_end_beyond_current_len() {
        let cache = LayerKvCache::new(small_config());
        cache.initialize(4).unwrap();
        assert!(cache.read_range(0, 100).is_err());
    }

    #[test]
    fn is_empty_and_len_track_initialization() {
        let cache = LayerKvCache::new(small_config());
        assert!(cache.is_empty());
        cache.initialize(5).unwrap();
        assert!(!cache.is_empty());
        assert_eq!(cache.len(), 5);
    }

    #[test]
    fn concurrent_reads_do_not_corrupt_or_deadlock() {
        let cache = LayerKvCache::new(small_config());
        cache.initialize(8).unwrap();
        let width = cache.config.token_width();
        for token_idx in 0..8 {
            let v = vec![token_idx as f32; width];
            cache.write_kv(token_idx, &v, &v).unwrap();
        }

        let mut handles = Vec::new();
        for _ in 0..8 {
            let cache = cache.clone();
            handles.push(std::thread::spawn(move || {
                for token_idx in 0..8 {
                    let entry = cache.read_kv(token_idx).unwrap();
                    assert_eq!(entry.keys[0], token_idx as f32);
                }
            }));
        }
        for h in handles {
            h.join().unwrap();
        }
    }

    #[test]
    fn clone_shares_the_same_underlying_cache() {
        let cache = LayerKvCache::new(small_config());
        cache.initialize(4).unwrap();
        let width = cache.config.token_width();
        let data = vec![3.0_f32; width];

        let cache_clone = cache.clone();
        cache_clone.write_kv(0, &data, &data).unwrap();

        // Written through the clone, visible through the original — same
        // underlying Arc<RwLock<..>>, not a deep copy.
        assert_eq!(cache.read_kv(0).unwrap().keys, data);
    }

    #[test]
    fn write_kv_batch_multiple() {
        let cache = LayerKvCache::new(small_config());
        cache.initialize(8).unwrap();
        let width = cache.config.token_width();

        let k0 = vec![0.0_f32; width];
        let v0 = vec![0.0_f32; width];
        let k1 = vec![1.0_f32; width];
        let v1 = vec![10.0_f32; width];

        let entries = [
            (0, k0.as_slice(), v0.as_slice()),
            (1, k1.as_slice(), v1.as_slice()),
        ];

        cache.write_kv_batch(&entries).unwrap();
        assert_eq!(cache.read_kv(0).unwrap().keys[0], 0.0);
        assert_eq!(cache.read_kv(1).unwrap().values[0], 10.0);
    }

    #[test]
    fn write_kv_batch_empty() {
        let cache = LayerKvCache::new(small_config());
        cache.initialize(4).unwrap();
        assert!(cache.write_kv_batch(&[]).is_ok());
    }

    #[test]
    fn write_kv_batch_validation_fails() {
        let cache = LayerKvCache::new(small_config());
        cache.initialize(4).unwrap();
        let width = cache.config.token_width();

        let k0 = vec![1.0_f32; width];
        let v0 = vec![2.0_f32; width];
        let bad_k = vec![0.0_f32; width + 1];

        let entries = [
            (0, k0.as_slice(), v0.as_slice()),
            (1, bad_k.as_slice(), v0.as_slice()),
        ];

        assert!(cache.write_kv_batch(&entries).is_err());
    }

    #[test]
    fn write_kv_batch_is_atomic_on_validation_failure() {
        let cache = LayerKvCache::new(small_config());
        cache.initialize(4).unwrap();
        let width = cache.config.token_width();

        let good_keys = vec![5.0_f32; width];
        let bad_keys = vec![0.0_f32; width + 1];

        let entries = [
            (0, good_keys.as_slice(), good_keys.as_slice()),
            (1, bad_keys.as_slice(), good_keys.as_slice()),
        ];

        assert!(cache.write_kv_batch(&entries).is_err());
        // Upfront validation must reject the whole batch before the write
        // lock is even taken, so entry 0 must still hold its zero-initialized
        // default rather than the batch's (never-applied) value.
        assert_eq!(cache.read_kv(0).unwrap().keys[0], 0.0);
    }

    #[test]
    fn with_read_kv_matches_read_kv() {
        let cache = LayerKvCache::new(small_config());
        cache.initialize(4).unwrap();
        let width = cache.config.token_width();
        let keys = vec![1.5_f32; width];
        let values = vec![2.5_f32; width];
        cache.write_kv(1, &keys, &values).unwrap();

        let owned = cache.read_kv(1).unwrap();
        cache
            .with_read_kv(1, |k, v| {
                assert_eq!(k, owned.keys.as_slice());
                assert_eq!(v, owned.values.as_slice());
            })
            .unwrap();
    }

    #[test]
    fn with_read_range_matches_read_range() {
        let cache = LayerKvCache::new(small_config());
        cache.initialize(4).unwrap();
        let width = cache.config.token_width();
        for t in 0..4 {
            let k = vec![t as f32 + 0.1; width];
            let v = vec![t as f32 + 0.2; width];
            cache.write_kv(t, &k, &v).unwrap();
        }

        let (owned_keys, owned_values) = cache.read_range(1, 3).unwrap();
        cache
            .with_read_range(1, 3, |span| {
                assert_eq!(span.keys, owned_keys.as_slice());
                assert_eq!(span.values, owned_values.as_slice());
            })
            .unwrap();
    }

    #[test]
    fn read_range_into_matches_read_range() {
        let cache = LayerKvCache::new(small_config());
        cache.initialize(4).unwrap();
        let width = cache.config.token_width();
        for t in 0..4 {
            let k = vec![t as f32 + 0.1; width];
            let v = vec![t as f32 + 0.2; width];
            cache.write_kv(t, &k, &v).unwrap();
        }

        let (owned_keys, owned_values) = cache.read_range(0, 4).unwrap();
        let mut keys_out = vec![0.0_f32; 4 * width];
        let mut values_out = vec![0.0_f32; 4 * width];

        cache
            .read_range_into(0, 4, &mut keys_out, &mut values_out)
            .unwrap();

        assert_eq!(keys_out, owned_keys);
        assert_eq!(values_out, owned_values);
    }

    #[test]
    fn read_range_into_rejects_short_output_buffers() {
        let cache = LayerKvCache::new(small_config());
        cache.initialize(4).unwrap();
        let width = cache.config.token_width();
        let mut short_buf = vec![0.0_f32; width - 1];
        let mut ok_buf = vec![0.0_f32; width];

        assert!(cache
            .read_range_into(0, 1, &mut short_buf, &mut ok_buf)
            .is_err());
        assert!(cache
            .read_range_into(0, 1, &mut ok_buf, &mut short_buf)
            .is_err());
    }

    #[test]
    fn concurrent_mixed_readers_and_one_writer_no_corruption() {
        let config = KVCacheConfig {
            max_tokens_per_seq: 256,
            num_heads: 2,
            hidden_dim_per_head: 4,
        };
        let cache = LayerKvCache::new(config);
        let width = config.token_width();
        cache.initialize(1).unwrap();

        let initial_keys = vec![1.0_f32; width];
        let initial_vals = vec![2.0_f32; width];
        cache.write_kv(0, &initial_keys, &initial_vals).unwrap();

        let stop_flag = Arc::new(std::sync::atomic::AtomicBool::new(false));
        let mut handles = Vec::new();

        // Writer thread appends tokens
        let writer_cache = cache.clone();
        let writer_stop = stop_flag.clone();
        handles.push(std::thread::spawn(move || {
            let mut t = 1;
            while !writer_stop.load(std::sync::atomic::Ordering::Relaxed) && t < 256 {
                let k = vec![t as f32; width];
                let v = vec![(t * 10) as f32; width];
                if writer_cache.write_kv(t, &k, &v).is_ok() {
                    t += 1;
                }
                std::thread::yield_now();
            }
        }));

        // 4 Reader threads
        for _ in 0..4 {
            let reader_cache = cache.clone();
            let reader_stop = stop_flag.clone();
            handles.push(std::thread::spawn(move || {
                while !reader_stop.load(std::sync::atomic::Ordering::Relaxed) {
                    let len = reader_cache.len();
                    if len > 0 {
                        reader_cache
                            .with_read_range(0, len, |span| {
                                for i in 0..len {
                                    let token_k = &span.keys[i * width..(i + 1) * width];
                                    let token_v = &span.values[i * width..(i + 1) * width];
                                    if i == 0 {
                                        assert_eq!(token_k[0], 1.0);
                                        assert_eq!(token_v[0], 2.0);
                                    } else {
                                        assert_eq!(token_k[0], i as f32);
                                        assert_eq!(token_v[0], (i * 10) as f32);
                                    }
                                }
                            })
                            .unwrap();
                    }
                    std::thread::yield_now();
                }
            }));
        }

        std::thread::sleep(std::time::Duration::from_millis(50));
        stop_flag.store(true, std::sync::atomic::Ordering::Relaxed);

        for h in handles {
            h.join().unwrap();
        }
    }

    proptest! {
        #[test]
        fn proptest_random_writes_and_read_range_equals_concatenation(
            token_count in 1usize..16,
            seed in 0u64..1000
        ) {
            let config = KVCacheConfig {
                max_tokens_per_seq: 16,
                num_heads: 2,
                hidden_dim_per_head: 4,
            };
            let cache = LayerKvCache::new(config);
            let width = config.token_width();
            cache.initialize(token_count).unwrap();

            let mut expected_keys = Vec::new();
            let mut expected_values = Vec::new();

            for t in 0..token_count {
                let k: Vec<f32> = (0..width).map(|i| (seed + t as u64 * 10 + i as u64) as f32).collect();
                let v: Vec<f32> = (0..width).map(|i| (seed + t as u64 * 100 + i as u64) as f32).collect();
                cache.write_kv(t, &k, &v).unwrap();
                expected_keys.extend_from_slice(&k);
                expected_values.extend_from_slice(&v);
            }

            let (read_k, read_v) = cache.read_range(0, token_count).unwrap();
            prop_assert_eq!(read_k, expected_keys.clone());
            prop_assert_eq!(read_v, expected_values.clone());

            cache.with_read_range(0, token_count, |span| {
                assert_eq!(span.keys, expected_keys.as_slice());
                assert_eq!(span.values, expected_values.as_slice());
            }).unwrap();
        }
    }
}