diff --git a/src/distributed/kv_cache_serde/deserialize.rs b/src/distributed/kv_cache_serde/deserialize.rs index d9e76130f..44c60a7cd 100644 --- a/src/distributed/kv_cache_serde/deserialize.rs +++ b/src/distributed/kv_cache_serde/deserialize.rs @@ -19,19 +19,44 @@ //! //! Used by: decode node in disaggregated serving pipeline +use std::collections::{HashMap, HashSet}; + use anyhow::{Context, Result, bail}; use super::serialize::CACHE_FORMAT_VERSION; use super::types::{ - CACHE_FORMAT_VERSION_V1, CacheType, RawTensorData, SerializableCacheState, - SerializableSequenceBackend, validate_raw_tensor, + CACHE_FORMAT_VERSION_V1, CacheIngestLimits, CacheType, ExpectedBlockGeometry, RawTensorData, + SerializableCacheState, SerializableSequenceBackend, validate_raw_tensor, }; -/// Deserialize a `SerializableCacheState` from the binary wire format. +/// Deserialize a `SerializableCacheState` from the binary wire format, applying +/// the default [`CacheIngestLimits`] frame cap. /// /// Returns the deserialized state ready for reconstruction into live /// cache objects. pub fn deserialize_cache_state(buf: &[u8]) -> Result { + deserialize_cache_state_with_limits(buf, &CacheIngestLimits::default()) +} + +/// Deserialize a `SerializableCacheState`, rejecting any frame larger than +/// `limits.max_frame_bytes` before parsing. +/// +/// The cap is the first line of defense at the deserialization boundary: an +/// untrusted peer cannot force a multi-gigabyte allocation by sending an +/// oversized buffer or a forged metadata length. The disaggregated decode path +/// passes a config-derived cap; [`deserialize_cache_state`] uses the generous +/// default backstop. +pub fn deserialize_cache_state_with_limits( + buf: &[u8], + limits: &CacheIngestLimits, +) -> Result { + if buf.len() > limits.max_frame_bytes { + bail!( + "cache state frame is {} bytes, exceeds limit {}", + buf.len(), + limits.max_frame_bytes + ); + } if buf.len() < 10 { bail!( "cache state buffer too short: need at least 10 bytes, got {}", @@ -51,6 +76,14 @@ pub fn deserialize_cache_state(buf: &[u8]) -> Result { let num_layers = u32::from_le_bytes(buf[2..6].try_into()?) as usize; let metadata_len = u32::from_le_bytes(buf[6..10].try_into()?) as usize; + if metadata_len > limits.max_frame_bytes { + bail!( + "cache metadata length {} exceeds frame limit {}", + metadata_len, + limits.max_frame_bytes + ); + } + let metadata_end = 10 + metadata_len; if buf.len() < metadata_end { bail!( @@ -188,15 +221,186 @@ pub fn restore_into_sequence_cache_set( Ok(()) } -/// Restore serialized cache state directly into an active `CachePool` slot. +/// Validate a deserialized paged handoff payload against ingest limits and, +/// when supplied, the decode model's real KV geometry, BEFORE any pool block is +/// acquired. This is the structural gate for a cross-node paged restore. /// -/// Used by: decode-side distributed cache ingestion where paged allocator -/// bookkeeping must be rebuilt alongside the sequence state. +/// Always enforced (no live model needed): +/// - the transferred block count and each slab's byte size stay within `limits`; +/// - every block the restored table references has matching transferred +/// contents on the SAME layer, and every transferred block is referenced by +/// the table on its declared layer (no orphan contents, which would leak a +/// pool block; no cross-layer aliasing; no duplicate id disagreeing on layer); +/// - a paged sequence carries a coherent, non-empty content path (an empty +/// block table, or a table that references blocks with no contents, is +/// rejected; a metadata-only paged payload is not a valid cross-node handoff +/// because a fresh decode pool holds no pre-existing block contents). +/// +/// Enforced when `expected` is provided: +/// - every slab's shape is exactly `[block_size, n_kv_heads, head_dim]` and its +/// dtype matches the decode model, anchoring pool geometry before the first +/// write captures it from the (untrusted) slab. +/// +/// Dense payloads (no paged state, no blocks) pass unchanged. +pub fn validate_paged_payload( + state: &SerializableCacheState, + limits: &CacheIngestLimits, + expected: Option<&ExpectedBlockGeometry>, +) -> Result<()> { + let is_paged = state.sequence_backend == SerializableSequenceBackend::PagedKvCache + || state.paged_state.is_some() + || !state.paged_blocks.is_empty(); + if !is_paged { + return Ok(()); + } + + if state.paged_blocks.len() > limits.max_paged_blocks { + bail!( + "paged handoff carries {} blocks, exceeds limit {}", + state.paged_blocks.len(), + limits.max_paged_blocks + ); + } + + let num_layers = expected + .map(|g| g.num_layers) + .or_else(|| state.paged_state.as_ref().map(|s| s.layers.len())); + + // Per-block: byte caps, layer range, optional geometry anchor. Index each + // supplied block id to its declared layer for the table cross-check below. + let mut block_layer: HashMap = HashMap::new(); + for (i, block) in state.paged_blocks.iter().enumerate() { + for (name, slab) in [("keys", &block.keys), ("values", &block.values)] { + if slab.data.len() > limits.max_block_slab_bytes { + bail!( + "paged block {i} {name} slab is {} bytes, exceeds limit {}", + slab.data.len(), + limits.max_block_slab_bytes + ); + } + } + if let Some(n) = num_layers + && block.layer_idx >= n + { + bail!( + "paged block {i} layer_idx {} out of range for {n} layers", + block.layer_idx + ); + } + if let Some(g) = expected { + let want = g.slab_shape(); + for (name, slab) in [("keys", &block.keys), ("values", &block.values)] { + if slab.shape != want { + bail!( + "paged block {i} {name} shape {:?} does not match decode model geometry {:?}", + slab.shape, + want + ); + } + if slab.dtype != g.dtype { + bail!( + "paged block {i} {name} dtype {} does not match decode model dtype {}", + slab.dtype, + g.dtype + ); + } + } + } + if let Some(prev) = block_layer.insert(block.block_id, block.layer_idx) + && prev != block.layer_idx + { + bail!( + "paged block id {} supplied for both layer {prev} and layer {} (block ids are layer-unique)", + block.block_id, + block.layer_idx + ); + } + } + + let Some(paged_state) = state.paged_state.as_ref() else { + bail!("paged sequence payload is missing paged_state metadata"); + }; + + let has_blocks = !state.paged_blocks.is_empty(); + let mut referenced: HashSet = HashSet::new(); + let mut referenced_count = 0usize; + for (layer_idx, layer) in paged_state.layers.iter().enumerate() { + for &block_id in &layer.block_ids { + referenced_count += 1; + referenced.insert(block_id); + if has_blocks { + match block_layer.get(&block_id) { + Some(&l) if l == layer_idx => {} + Some(&l) => bail!( + "block {block_id} is referenced by layer {layer_idx} but its transferred contents declare layer {l}" + ), + None => bail!( + "block {block_id} referenced by layer {layer_idx} has no transferred contents" + ), + } + } + } + } + + if referenced_count == 0 { + bail!( + "paged handoff carries an empty block table; a cross-node decode handoff must carry a prefilled sequence" + ); + } + if !has_blocks { + bail!( + "paged handoff references {referenced_count} blocks but carries no block contents; metadata-only paged restore is invalid on a fresh decode pool" + ); + } + for block in &state.paged_blocks { + if !referenced.contains(&block.block_id) { + bail!( + "paged block id {} has transferred contents but is not referenced by any layer table", + block.block_id + ); + } + } + + Ok(()) +} + +/// Restore serialized cache state directly into an active `CachePool` slot, +/// applying the default [`CacheIngestLimits`] and no model geometry anchor. +/// +/// Used by: in-process round-trips and the decode-side ingestion that does not +/// (yet) supply the decode model's geometry. Prefer +/// [`restore_into_cache_pool_sequence_anchored`] on the live disaggregated path +/// so transferred slabs are pinned to the decode model. pub fn restore_into_cache_pool_sequence( state: &SerializableCacheState, cache_pool: &mut mlxcel_core::cache::CachePool, seq_id: mlxcel_core::cache::SequenceId, ) -> Result<()> { + restore_into_cache_pool_sequence_anchored( + state, + cache_pool, + seq_id, + &CacheIngestLimits::default(), + None, + ) +} + +/// Restore serialized cache state into an active `CachePool` slot, validating +/// the payload against `limits` and (when supplied) the decode model's real KV +/// geometry before any pool block is acquired. +/// +/// Used by: decode-side distributed cache ingestion where paged allocator +/// bookkeeping must be rebuilt alongside the sequence state, and the untrusted +/// wire payload must be anchored to the local model. +pub fn restore_into_cache_pool_sequence_anchored( + state: &SerializableCacheState, + cache_pool: &mut mlxcel_core::cache::CachePool, + seq_id: mlxcel_core::cache::SequenceId, + limits: &CacheIngestLimits, + expected: Option<&ExpectedBlockGeometry>, +) -> Result<()> { + validate_paged_payload(state, limits, expected)?; + // A true pool-backed cross-node handoff carries the pool block CONTENTS in // `paged_blocks` (#125). That path materializes fresh pool blocks rather // than restoring dense per-layer buffers (empty for pool-backed sequences) diff --git a/src/distributed/kv_cache_serde/mod.rs b/src/distributed/kv_cache_serde/mod.rs index 82e372874..53718209b 100644 --- a/src/distributed/kv_cache_serde/mod.rs +++ b/src/distributed/kv_cache_serde/mod.rs @@ -60,8 +60,9 @@ pub mod types; mod tests; pub use deserialize::{ - deserialize_cache_state, reconstruct_mlx_array, restore_into_cache_pool_sequence, - restore_into_kv_caches, restore_into_sequence_cache_set, + deserialize_cache_state, deserialize_cache_state_with_limits, reconstruct_mlx_array, + restore_into_cache_pool_sequence, restore_into_cache_pool_sequence_anchored, + restore_into_kv_caches, restore_into_sequence_cache_set, validate_paged_payload, }; pub use serialize::{ CACHE_FORMAT_VERSION, extract_chunked_cache_entry, extract_kv_cache_entry, @@ -69,8 +70,8 @@ pub use serialize::{ serialize_sequence_cache_set, }; pub use types::{ - CACHE_FORMAT_VERSION_V1, CACHE_FORMAT_VERSION_V2, CacheMetadata, CacheType, RawTensorData, - SerializableCacheEntry, SerializableCacheState, SerializablePagedBlock, - SerializablePagedLayerState, SerializablePagedSequenceState, SerializableSamplingState, - SerializableSequenceBackend, validate_raw_tensor, + CACHE_FORMAT_VERSION_V1, CACHE_FORMAT_VERSION_V2, CacheIngestLimits, CacheMetadata, CacheType, + ExpectedBlockGeometry, RawTensorData, SerializableCacheEntry, SerializableCacheState, + SerializablePagedBlock, SerializablePagedLayerState, SerializablePagedSequenceState, + SerializableSamplingState, SerializableSequenceBackend, validate_raw_tensor, }; diff --git a/src/distributed/kv_cache_serde/tests.rs b/src/distributed/kv_cache_serde/tests.rs index 071f650ca..544b8829d 100644 --- a/src/distributed/kv_cache_serde/tests.rs +++ b/src/distributed/kv_cache_serde/tests.rs @@ -586,9 +586,12 @@ fn mlx_dtype_element_size_known_types() { mlx_dtype_element_size(mlxcel_core::dtype::FLOAT16).unwrap(), 2 ); + // BFLOAT16 is a 16-bit type: 2 bytes. (#125 fixed `mlx_dtype_element_size` + // to return 2 here so `validate_raw_tensor` stops rejecting valid bf16 + // payloads; this assertion was left stale at the old, wrong value 4.) assert_eq!( mlx_dtype_element_size(mlxcel_core::dtype::BFLOAT16).unwrap(), - 4 + 2 ); assert_eq!(mlx_dtype_element_size(0).unwrap(), 1); // BOOL } @@ -762,3 +765,312 @@ fn chunked_cache_extract_round_trip() { assert!(shape[2] >= 3); assert_eq!(shape[3], 1); } + +// --------------------------------------------------------------------------- +// Untrusted-network handoff hardening tests (#126): frame caps, model-geometry +// anchor, table <-> contents layer consistency, and empty-sequence routing. +// These exercise `validate_paged_payload` / `deserialize_cache_state_with_limits` +// over synthetic payloads, so they need no Metal device. +// --------------------------------------------------------------------------- + +const HARDEN_BS: usize = 4; +const HARDEN_H: i32 = 2; +const HARDEN_D: i32 = 8; + +fn harden_slab() -> RawTensorData { + let dt = mlxcel_core::dtype::FLOAT16; + let elements = HARDEN_BS * HARDEN_H as usize * HARDEN_D as usize; + let bytes = elements * types::mlx_dtype_element_size(dt).unwrap(); + RawTensorData { + data: vec![0u8; bytes], + shape: vec![HARDEN_BS as i32, HARDEN_H, HARDEN_D], + dtype: dt, + } +} + +fn harden_block(block_id: u64, layer_idx: usize) -> SerializablePagedBlock { + SerializablePagedBlock { + block_id, + layer_idx, + keys: harden_slab(), + values: harden_slab(), + } +} + +fn harden_expected() -> ExpectedBlockGeometry { + ExpectedBlockGeometry { + num_layers: 2, + block_size: HARDEN_BS, + n_kv_heads: HARDEN_H, + head_dim: HARDEN_D, + dtype: mlxcel_core::dtype::FLOAT16, + } +} + +/// Coherent 2-layer pool-backed handoff: layer 0 references blocks 10 and 11 +/// (len 6 over block_size 4 = 2 blocks), layer 1 references block 20 (len 3 = 1 +/// block). Each referenced block has matching contents on its own layer. +fn coherent_paged_payload() -> SerializableCacheState { + SerializableCacheState { + cache_type: CacheType::Standard, + entries: vec![], + metadata: CacheMetadata { + prompt_len: 6, + current_offset: 6, + num_layers: 2, + layer_offsets: vec![], + max_size: None, + layer_indices: None, + chunk_size: None, + start_positions: None, + }, + sampling_state: None, + token_history: vec![], + sequence_id: 1, + sequence_backend: SerializableSequenceBackend::PagedKvCache, + paged_state: Some(SerializablePagedSequenceState { + block_size: HARDEN_BS, + bytes_per_block: vec![HARDEN_BS, HARDEN_BS], + layers: vec![ + SerializablePagedLayerState { + block_ids: vec![10, 11], + len: 6, + logical_start: 0, + }, + SerializablePagedLayerState { + block_ids: vec![20], + len: 3, + logical_start: 0, + }, + ], + }), + paged_blocks: vec![ + harden_block(10, 0), + harden_block(11, 0), + harden_block(20, 1), + ], + } +} + +#[test] +fn validate_paged_payload_accepts_coherent_handoff() { + let payload = coherent_paged_payload(); + validate_paged_payload(&payload, &CacheIngestLimits::default(), None).unwrap(); + validate_paged_payload( + &payload, + &CacheIngestLimits::default(), + Some(&harden_expected()), + ) + .unwrap(); +} + +#[test] +fn validate_paged_payload_passes_dense_payload_unchanged() { + let mut payload = coherent_paged_payload(); + payload.sequence_backend = SerializableSequenceBackend::DenseKvCache; + payload.paged_state = None; + payload.paged_blocks = Vec::new(); + validate_paged_payload(&payload, &CacheIngestLimits::default(), None).unwrap(); +} + +#[test] +fn validate_paged_payload_rejects_geometry_shape_mismatch() { + let payload = coherent_paged_payload(); + let mut expected = harden_expected(); + expected.head_dim = HARDEN_D + 1; + let err = validate_paged_payload(&payload, &CacheIngestLimits::default(), Some(&expected)) + .unwrap_err() + .to_string(); + assert!( + err.contains("does not match decode model geometry"), + "{err}" + ); +} + +#[test] +fn validate_paged_payload_rejects_geometry_dtype_mismatch() { + let payload = coherent_paged_payload(); + let mut expected = harden_expected(); + expected.dtype = mlxcel_core::dtype::BFLOAT16; + let err = validate_paged_payload(&payload, &CacheIngestLimits::default(), Some(&expected)) + .unwrap_err() + .to_string(); + assert!(err.contains("does not match decode model dtype"), "{err}"); +} + +#[test] +fn validate_paged_payload_rejects_block_count_over_limit() { + let payload = coherent_paged_payload(); + let limits = CacheIngestLimits { + max_paged_blocks: 2, + ..CacheIngestLimits::default() + }; + let err = validate_paged_payload(&payload, &limits, None) + .unwrap_err() + .to_string(); + assert!(err.contains("exceeds limit"), "{err}"); +} + +#[test] +fn validate_paged_payload_rejects_oversized_slab() { + let payload = coherent_paged_payload(); + let limits = CacheIngestLimits { + max_block_slab_bytes: 1, + ..CacheIngestLimits::default() + }; + let err = validate_paged_payload(&payload, &limits, None) + .unwrap_err() + .to_string(); + assert!(err.contains("exceeds limit"), "{err}"); +} + +#[test] +fn validate_paged_payload_rejects_layer_out_of_range() { + let mut payload = coherent_paged_payload(); + // A block declaring a layer beyond the table's layer count. + payload.paged_blocks.push(harden_block(99, 7)); + payload + .paged_state + .as_mut() + .unwrap() + .layers + .push(SerializablePagedLayerState { + block_ids: vec![99], + len: 1, + logical_start: 0, + }); + let err = validate_paged_payload(&payload, &CacheIngestLimits::default(), None) + .unwrap_err() + .to_string(); + assert!(err.contains("out of range"), "{err}"); +} + +#[test] +fn validate_paged_payload_rejects_missing_block_contents() { + let mut payload = coherent_paged_payload(); + // Drop block 20's contents while the layer-1 table still references it. + payload.paged_blocks.retain(|b| b.block_id != 20); + let err = validate_paged_payload(&payload, &CacheIngestLimits::default(), None) + .unwrap_err() + .to_string(); + assert!(err.contains("no transferred contents"), "{err}"); +} + +#[test] +fn validate_paged_payload_rejects_cross_layer_block() { + let mut payload = coherent_paged_payload(); + // Block 10 is referenced by layer 0, but its contents claim layer 1. + for block in &mut payload.paged_blocks { + if block.block_id == 10 { + block.layer_idx = 1; + } + } + let err = validate_paged_payload(&payload, &CacheIngestLimits::default(), None) + .unwrap_err() + .to_string(); + assert!(err.contains("declare layer"), "{err}"); +} + +#[test] +fn validate_paged_payload_rejects_orphan_contents() { + let mut payload = coherent_paged_payload(); + // A block with contents that no layer table references (a pool-block leak). + payload.paged_blocks.push(harden_block(404, 0)); + let err = validate_paged_payload(&payload, &CacheIngestLimits::default(), None) + .unwrap_err() + .to_string(); + assert!(err.contains("not referenced by any layer table"), "{err}"); +} + +#[test] +fn validate_paged_payload_rejects_metadata_only_handoff() { + let mut payload = coherent_paged_payload(); + // Table still references blocks, but no contents are carried. + payload.paged_blocks = Vec::new(); + let err = validate_paged_payload(&payload, &CacheIngestLimits::default(), None) + .unwrap_err() + .to_string(); + assert!( + err.contains("metadata-only paged restore is invalid"), + "{err}" + ); +} + +#[test] +fn validate_paged_payload_rejects_empty_block_table() { + let mut payload = coherent_paged_payload(); + payload.paged_blocks = Vec::new(); + for layer in &mut payload.paged_state.as_mut().unwrap().layers { + layer.block_ids.clear(); + layer.len = 0; + } + let err = validate_paged_payload(&payload, &CacheIngestLimits::default(), None) + .unwrap_err() + .to_string(); + assert!(err.contains("empty block table"), "{err}"); +} + +#[test] +fn validate_paged_payload_rejects_paged_backend_without_metadata() { + let mut payload = coherent_paged_payload(); + payload.paged_state = None; + payload.paged_blocks = Vec::new(); + let err = validate_paged_payload(&payload, &CacheIngestLimits::default(), None) + .unwrap_err() + .to_string(); + assert!(err.contains("missing paged_state metadata"), "{err}"); +} + +#[test] +fn deserialize_cache_state_rejects_oversized_frame() { + // A well-formed dense frame, then a limit below its size. + let state = SerializableCacheState { + cache_type: CacheType::Standard, + entries: vec![], + metadata: CacheMetadata { + prompt_len: 0, + current_offset: 0, + num_layers: 0, + layer_offsets: vec![], + max_size: None, + layer_indices: None, + chunk_size: None, + start_positions: None, + }, + sampling_state: None, + token_history: vec![], + sequence_id: 1, + sequence_backend: SerializableSequenceBackend::DenseKvCache, + paged_state: None, + paged_blocks: Vec::new(), + }; + let bytes = serialize_cache_state(&state).unwrap(); + let limits = CacheIngestLimits { + max_frame_bytes: bytes.len() - 1, + ..CacheIngestLimits::default() + }; + let err = deserialize_cache_state_with_limits(&bytes, &limits) + .unwrap_err() + .to_string(); + assert!(err.contains("exceeds limit"), "{err}"); + // The same bytes deserialize fine under the default cap. + deserialize_cache_state(&bytes).unwrap(); +} + +#[test] +fn deserialize_cache_state_rejects_forged_metadata_length() { + // Header claims a metadata length far beyond the frame cap. + let mut buf = vec![0u8; 10]; + buf[0] = CACHE_FORMAT_VERSION; + buf[1] = CacheType::Standard as u8; + buf[2..6].copy_from_slice(&0u32.to_le_bytes()); + buf[6..10].copy_from_slice(&u32::MAX.to_le_bytes()); + let limits = CacheIngestLimits { + max_frame_bytes: 4096, + ..CacheIngestLimits::default() + }; + let err = deserialize_cache_state_with_limits(&buf, &limits) + .unwrap_err() + .to_string(); + assert!(err.contains("exceeds frame limit"), "{err}"); +} diff --git a/src/distributed/kv_cache_serde/types.rs b/src/distributed/kv_cache_serde/types.rs index 8e8ec2ca7..34c7ec7ce 100644 --- a/src/distributed/kv_cache_serde/types.rs +++ b/src/distributed/kv_cache_serde/types.rs @@ -382,3 +382,69 @@ pub fn validate_raw_tensor(tensor: &RawTensorData) -> anyhow::Result<()> { Ok(()) } + +/// Bounds applied to a cache payload accepted from an (untrusted) peer node, so +/// a malformed or hostile handoff cannot exhaust decode-node memory or smuggle +/// oversized blocks past the frame reader. +/// +/// The live disaggregated decode path derives tight values from the model and +/// serving config; the [`Default`] values are generous backstops sized for +/// in-process and test round-trips, not a security boundary on their own. The +/// frame cap is the JSON-bloated wire size (the tensor bytes ride inside the +/// metadata frame as number arrays, roughly 4x their raw size). +#[derive(Debug, Clone, Copy)] +pub struct CacheIngestLimits { + /// Maximum accepted wire frame size in bytes (whole buffer + metadata len). + pub max_frame_bytes: usize, + /// Maximum number of transferred paged blocks in one handoff. + pub max_paged_blocks: usize, + /// Maximum bytes for a single block slab (one K or one V slab). + pub max_block_slab_bytes: usize, +} + +impl Default for CacheIngestLimits { + fn default() -> Self { + Self { + // 16 GiB: rejects absurd payloads while leaving headroom for a + // large model / long context handoff under the bloated JSON framing. + max_frame_bytes: 16 << 30, + // 1,048,576 per-layer blocks: far above any realistic single-sequence + // block table; the decode pool's block budget is the real bound. + max_paged_blocks: 1 << 20, + // 64 MiB: one [block_size, n_kv_heads, head_dim] slab is small even + // for wide-GQA models (e.g. 32 * 128 * 256 * 4 = 4 MiB). + max_block_slab_bytes: 64 << 20, + } + } +} + +/// The decode model's real per-layer paged KV block geometry, used to anchor +/// every transferred block slab at the deserialization boundary. +/// +/// A fresh decode pool captures its geometry from the FIRST slab it is asked to +/// write ([`mlxcel_core::cache`]'s `PagedBlockPool::write_block` only validates +/// subsequent writes against the captured geometry), so without this anchor a +/// peer could establish a wrong-shaped pool on the first write. Supplying the +/// model's real `(block_size, n_kv_heads, head_dim, dtype)` rejects any slab +/// that does not match before a single pool block is acquired. +#[derive(Debug, Clone, Copy)] +pub struct ExpectedBlockGeometry { + /// Number of layers the decode model exposes. + pub num_layers: usize, + /// Paged pool block size (slots per block). + pub block_size: usize, + /// Key/value head count. + pub n_kv_heads: i32, + /// Per-head dimension. + pub head_dim: i32, + /// MLX dtype code for the pool tensors (e.g. 9 = float16, 12 = bfloat16). + pub dtype: i32, +} + +impl ExpectedBlockGeometry { + /// The canonical bare layout-A slab shape `[block_size, n_kv_heads, head_dim]` + /// that `read_block_contents` produces and `acquire_and_write_block` accepts. + pub fn slab_shape(&self) -> [i32; 3] { + [self.block_size as i32, self.n_kv_heads, self.head_dim] + } +} diff --git a/src/lib/mlxcel-core/src/cache/paged_detach_tests.rs b/src/lib/mlxcel-core/src/cache/paged_detach_tests.rs index 17155f8c3..ab4b68092 100644 --- a/src/lib/mlxcel-core/src/cache/paged_detach_tests.rs +++ b/src/lib/mlxcel-core/src/cache/paged_detach_tests.rs @@ -1494,6 +1494,150 @@ fn pool_backed_sharers_cow_keeps_other_prefix_intact() { assert_eq!(b_after.1, reference.1, "B's prefix V intact after A's COW"); } +#[test] +fn trim_across_shared_blocks_releases_only_the_trimming_sequences_refs() { + // Two LIVE pool-backed sequences share a 2-block prefix (refcount 2 on each + // block, the state two adopters of a cached prefix reach). Trimming one of + // them across the shared tail block must drop ONLY that sequence's ref: the + // physical block survives for the sibling (refcount 2 -> 1, never freed), + // the sibling's gathered prefix stays bit-identical, and the trimmed + // sequence's own window shrinks to the retained prefix. This is the + // "trim/rewind across shared blocks" correctness case for #126. (rewind is + // an alias of trim, so this covers both.) + let n_kv_heads = 2i32; + let head_dim = 3i32; + let block_size = 4usize; + let layout = fp16_pool_layout(block_size, n_kv_heads, head_dim); + let model = DenseNaturalStubModel::new(1); + let mut pool = CachePool::new(4); + + // Sequence A: 6-token prefix => 2 blocks (head block full, tail half-full). + let seq_a = pool + .allocate_with_layout( + &model, + Some(SequenceStateLayout::paged_kv_cache(layout.clone())), + ) + .unwrap(); + let prefix_len = 6i32; + let prefix_k = prefill_block(1000.0, n_kv_heads, prefix_len, head_dim); + let prefix_v = prefill_block(5000.0, n_kv_heads, prefix_len, head_dim); + write_prefill_for(&mut pool, seq_a, 0, &prefix_k, &prefix_v).unwrap(); + let prefix_blocks = pool + .get_paged_state(seq_a) + .unwrap() + .layer(0) + .unwrap() + .block_ids + .clone(); + assert_eq!(prefix_blocks.len(), 2); + let head_block = prefix_blocks[0]; + let tail_block = prefix_blocks[1]; + + let reference = { + let state = pool.get_paged_state(seq_a).unwrap(); + let (gk, gv) = pool + .paged_pool_ref() + .unwrap() + .gather_visible(&state, 0) + .unwrap() + .expect("A gather"); + (flatten_fp32(&gk), flatten_fp32(&gv)) + }; + + // Sequence B aliases A's prefix block table and pins it (refcount 2), + // exactly the shared state two concurrent adopters of a cached prefix reach. + let seq_b = pool + .allocate_with_layout(&model, Some(SequenceStateLayout::paged_kv_cache(layout))) + .unwrap(); + { + let set = pool.get_mut(seq_b).unwrap(); + let mut state = set.paged.as_ref().unwrap().borrow_mut(); + let layer = state.layer_mut(0).unwrap(); + layer.block_ids = prefix_blocks.clone(); + layer.len = prefix_len as usize; + layer.logical_start = 0; + } + for &id in &prefix_blocks { + pool.paged_pool + .as_ref() + .unwrap() + .borrow_mut() + .retain_block(id) + .unwrap(); + } + assert_eq!(pool.paged_pool_ref().unwrap().refcount(tail_block), 2); + assert_eq!(pool.paged_pool_ref().unwrap().refcount(head_block), 2); + assert_eq!(pool.paged_pool_ref().unwrap().live_block_count(), 2); + + // Trim A by 2 tokens: visible 6 -> 4, dropping the shared tail block from + // A's table. `release_block` decrements the tail's refcount but must NOT + // free it (B still owns it). + let trimmed = pool.trim_paged_tokens(seq_a, 0, 2).unwrap(); + assert_eq!(trimmed, 2); + { + let state = pool.get_paged_state(seq_a).unwrap(); + let layer = state.layer(0).unwrap(); + assert_eq!(layer.len, 4); + assert_eq!(layer.block_ids, vec![head_block]); + } + + assert_eq!( + pool.paged_pool_ref().unwrap().refcount(tail_block), + 1, + "trim dropped only A's ref to the shared tail; B still owns it" + ); + assert_eq!( + pool.paged_pool_ref().unwrap().refcount(head_block), + 2, + "the head block is still shared by both sequences" + ); + assert_eq!( + pool.paged_pool_ref().unwrap().live_block_count(), + 2, + "the trimmed-away block stays live for the sibling (no premature free)" + ); + + // B's full 6-token prefix is untouched by A's trim. + { + let state = pool.get_paged_state(seq_b).unwrap(); + let (gk, gv) = pool + .paged_pool_ref() + .unwrap() + .gather_visible(&state, 0) + .unwrap() + .expect("B gather after A trim"); + assert_eq!(flatten_fp32(&gk), reference.0, "B keys intact after A trim"); + assert_eq!( + flatten_fp32(&gv), + reference.1, + "B values intact after A trim" + ); + } + + // A's own window is exactly the retained 4-token prefix (the head block's + // data). `prefill_block` values are position-encoded and length-independent, + // so the first four tokens equal a fresh 4-token block. + { + let state = pool.get_paged_state(seq_a).unwrap(); + let (gk, gv) = pool + .paged_pool_ref() + .unwrap() + .gather_visible(&state, 0) + .unwrap() + .expect("A gather post-trim"); + assert_eq!( + flatten_fp32(&gk), + flatten_fp32(&prefill_block(1000.0, n_kv_heads, 4, head_dim)), + "A keys are the retained 4-token prefix" + ); + assert_eq!( + flatten_fp32(&gv), + flatten_fp32(&prefill_block(5000.0, n_kv_heads, 4, head_dim)), + "A values are the retained 4-token prefix" + ); + } +} + /// Free reference to DetachedPagedCacheSet to keep the import alive. #[allow(dead_code)] fn _type_alive() -> Option {