use mlx_native::{DType, MlxDevice};
use super::cache::{CacheError, CacheKind, Deepseek4Cache, Deepseek4CachePlan};
use super::Deepseek4Config;
fn config(ratios: Vec<u32>) -> Deepseek4Config {
let layers = ratios.len() as u32;
Deepseek4Config {
num_hidden_layers: layers,
hidden_size: 4096,
hidden_size_out: 16384,
max_position_embeddings: 1_048_576,
vocab_size: 129280,
num_attention_heads: 64,
num_key_value_heads: 1,
head_dim: 512,
rope_head_dim: 64,
rope_theta: 10000.0,
rope_factor: 16.0,
original_context_length: 65536,
yarn_beta_fast: 32.0,
yarn_beta_slow: 1.0,
q_lora_rank: 1024,
o_lora_rank: 1024,
output_groups: 8,
sliding_window: 128,
compress_ratios: ratios,
compress_rope_theta: 160000.0,
index_num_heads: 64,
index_head_dim: 128,
index_top_k: 512,
rms_norm_eps: 1e-6,
num_experts: 256,
num_experts_per_tok: 6,
num_shared_experts: 1,
expert_intermediate_size: 2048,
route_scale: 1.5,
normalize_topk: true,
swiglu_clamp_experts: vec![10.0; layers as usize],
swiglu_clamp_shared: vec![10.0; layers as usize],
hyper_connection_count: 4,
hyper_connection_sinkhorn_iterations: 20,
hyper_connection_epsilon: 1e-6,
hash_layer_count: layers.min(3),
}
}
fn official_config() -> Deepseek4Config {
config(
(0..43)
.map(|layer| {
if layer < 2 {
0
} else if layer % 2 == 0 {
4
} else {
128
}
})
.collect(),
)
}
#[test]
fn official_one_million_context_plan_has_exact_shapes_and_bytes() {
let plan = Deepseek4CachePlan::for_context(&official_config(), 1_048_576).unwrap();
assert_eq!(plan.layers.len(), 43);
assert_eq!(plan.resident_bytes, 7_232_045_056);
assert_eq!(plan.layers[0].attention_kv.shape, vec![128, 512]);
assert_eq!(plan.layers[0].window_kv.shape, vec![128, 512]);
assert!(plan.layers[0].compressed_kv.is_none());
assert!(plan.layers[0].indexer_kv.is_none());
let ratio_four = &plan.layers[2];
assert_eq!(ratio_four.compress_ratio, 4);
assert_eq!(
ratio_four.compressed_kv.as_ref().unwrap().shape,
vec![262_144, 512]
);
assert_eq!(ratio_four.attention_kv.shape, vec![262_272, 512]);
assert_eq!(
ratio_four.indexer_kv.as_ref().unwrap().shape,
vec![262_144, 128]
);
assert_eq!(
ratio_four.main_kv_state.as_ref().unwrap().shape,
vec![1, 8, 1024]
);
assert_eq!(
ratio_four.main_score_state.as_ref().unwrap().shape,
vec![1, 8, 1024]
);
assert_eq!(
ratio_four.indexer_kv_state.as_ref().unwrap().shape,
vec![1, 8, 256]
);
assert_eq!(
ratio_four.indexer_score_state.as_ref().unwrap().shape,
vec![1, 8, 256]
);
assert_eq!(ratio_four.main_kv_state.as_ref().unwrap().dtype, DType::F32);
let ratio_128 = &plan.layers[3];
assert_eq!(ratio_128.compress_ratio, 128);
assert_eq!(
ratio_128.compressed_kv.as_ref().unwrap().shape,
vec![8192, 512]
);
assert_eq!(ratio_128.attention_kv.shape, vec![8320, 512]);
assert_eq!(
ratio_128.main_kv_state.as_ref().unwrap().shape,
vec![1, 128, 512]
);
assert!(ratio_128.indexer_kv.is_none());
assert_eq!(
plan.layers
.iter()
.map(|layer| layer.resident_bytes)
.sum::<u64>(),
plan.resident_bytes
);
}
#[test]
fn malformed_context_schedule_and_overflow_fail_closed() {
let cfg = official_config();
assert!(matches!(
Deepseek4CachePlan::for_context(&cfg, 0),
Err(CacheError::EmptyContext)
));
assert!(matches!(
Deepseek4CachePlan::for_context(&cfg, 1_048_577),
Err(CacheError::ContextBound { .. })
));
let mut malformed = cfg.clone();
malformed.sliding_window = 127;
assert!(matches!(
Deepseek4CachePlan::for_context(&malformed, 1024),
Err(CacheError::SlidingWindow { actual: 127 })
));
malformed = cfg.clone();
malformed.compress_ratios.pop();
assert!(matches!(
Deepseek4CachePlan::for_context(&malformed, 1024),
Err(CacheError::LayerCount { .. })
));
malformed = cfg.clone();
malformed.compress_ratios[0] = 8;
assert!(matches!(
Deepseek4CachePlan::for_context(&malformed, 1024),
Err(CacheError::CompressionRatio { layer: 0, ratio: 8 })
));
let mut overflowing = config(vec![4, 4]);
overflowing.max_position_embeddings = u32::MAX;
overflowing.head_dim = u32::MAX;
assert!(matches!(
Deepseek4CachePlan::for_context(&overflowing, u32::MAX as usize),
Err(CacheError::ByteOverflow {
layer: 1,
kind: CacheKind::AttentionKv
})
));
}
#[test]
fn allocator_materializes_the_plan_as_zeroed_bf16_buffers() {
let mut cfg = config(vec![0, 4, 128]);
cfg.max_position_embeddings = 128;
cfg.head_dim = 32;
cfg.index_head_dim = 16;
let plan = Deepseek4CachePlan::for_context(&cfg, 128).unwrap();
assert_eq!(plan.resident_bytes, 66_624);
let _gpu = crate::inference::hf2q_gpu_test_lock();
let cache = Deepseek4Cache::allocate(&plan, MlxDevice::new().unwrap()).unwrap();
assert_eq!(cache.resident_bytes(), plan.resident_bytes);
assert_eq!(cache.layers().len(), 3);
assert_eq!(cache.layers()[1].attention_kv.dtype(), DType::BF16);
assert_eq!(cache.layers()[1].attention_kv.shape(), &[160, 32]);
assert_eq!(cache.layers()[1].window_kv.dtype(), DType::BF16);
assert_eq!(cache.layers()[1].window_kv.shape(), &[128, 32]);
assert_eq!(
cache.layers()[1].compressed_kv.as_ref().unwrap().shape(),
&[32, 32]
);
assert_eq!(
cache.layers()[1].indexer_kv.as_ref().unwrap().shape(),
&[32, 16]
);
assert_eq!(
cache.layers()[1]
.compressed_kv
.as_ref()
.unwrap()
.byte_offset(),
(128 * 32 * DType::BF16.size_of()) as u64
);
assert_eq!(
cache.layers()[1].main_kv_state.as_ref().unwrap().shape(),
&[1, 8, 64]
);
assert_eq!(
cache.layers()[1]
.indexer_score_state
.as_ref()
.unwrap()
.shape(),
&[1, 8, 32]
);
assert!(cache.layers()[1]
.indexer_kv
.as_ref()
.unwrap()
.as_slice::<u16>()
.unwrap()
.iter()
.all(|value| *value == 0));
assert!(cache.layers()[1]
.main_kv_state
.as_ref()
.unwrap()
.as_slice::<f32>()
.unwrap()
.iter()
.all(|value| *value == 0.0));
assert!(cache.layers()[1]
.main_score_state
.as_ref()
.unwrap()
.as_slice::<f32>()
.unwrap()
.iter()
.all(|value| *value == f32::NEG_INFINITY));
}
#[test]
fn cache_steps_publish_only_complete_groups_and_commit_transactionally() {
let cfg = config(vec![4, 128]);
let plan = Deepseek4CachePlan::for_context(&cfg, 128).unwrap();
let _gpu = crate::inference::hf2q_gpu_test_lock();
let mut cache = Deepseek4Cache::allocate(&plan, MlxDevice::new().unwrap()).unwrap();
let first = cache.plan_next_step().unwrap();
assert_eq!(first.position, 0);
assert_eq!(first.layers[0].window_write_slot, 0);
assert_eq!(first.layers[0].compressed_write_slot, None);
assert_eq!(first.layers[0].compressed_valid_after, 0);
assert!(matches!(
cache.commit_step(1),
Err(CacheError::StepOutOfOrder {
expected: 0,
actual: 1
})
));
cache.commit_step(first.position).unwrap();
for expected in 1..4 {
let step = cache.plan_next_step().unwrap();
assert_eq!(step.position, expected);
cache.commit_step(step.position).unwrap();
}
let after_boundary = cache.plan_next_step().unwrap();
assert_eq!(after_boundary.position, 4);
assert_eq!(after_boundary.layers[0].compressed_write_slot, None);
assert_eq!(after_boundary.layers[0].compressed_valid_after, 1);
cache.reset().unwrap();
for expected in 0..128 {
let step = cache.plan_next_step().unwrap();
assert_eq!(step.position, expected);
if expected == 3 {
assert_eq!(step.layers[0].compressed_write_slot, Some(0));
assert_eq!(step.layers[0].indexer_write_slot, Some(0));
}
if expected == 127 {
assert_eq!(step.layers[0].compressed_write_slot, Some(31));
assert_eq!(step.layers[1].compressed_write_slot, Some(0));
assert_eq!(step.layers[1].indexer_write_slot, None);
}
cache.commit_step(step.position).unwrap();
}
assert!(matches!(
cache.plan_next_step(),
Err(CacheError::ContextExhausted { maximum: 128 })
));
}
#[test]
fn partial_token_poison_requires_reset_before_replay() {
let cfg = config(vec![4, 128]);
let plan = Deepseek4CachePlan::for_context(&cfg, 128).unwrap();
let _gpu = crate::inference::hf2q_gpu_test_lock();
let mut cache = Deepseek4Cache::allocate(&plan, MlxDevice::new().unwrap()).unwrap();
cache.poison();
assert!(cache.is_poisoned());
assert!(matches!(cache.plan_next_step(), Err(CacheError::Poisoned)));
assert!(matches!(cache.commit_step(0), Err(CacheError::Poisoned)));
cache.reset().unwrap();
assert!(!cache.is_poisoned());
assert_eq!(cache.plan_next_step().unwrap().position, 0);
}
#[test]
fn start_zero_prefill_span_counts_complete_groups_and_publishes_once() {
let cfg = config(vec![4, 128]);
let plan = Deepseek4CachePlan::for_context(&cfg, 128).unwrap();
let _gpu = crate::inference::hf2q_gpu_test_lock();
let mut cache = Deepseek4Cache::allocate(&plan, MlxDevice::new().unwrap()).unwrap();
for (tokens, ratio4, ratio128) in [(1, 0, 0), (3, 0, 0), (4, 1, 0), (127, 31, 0), (128, 32, 1)]
{
let span = cache.plan_prefill_start0(tokens).unwrap();
assert_eq!(span.start_position, 0);
assert_eq!(span.token_count, tokens);
assert_eq!(span.layers[0].window_valid_after, tokens);
assert_eq!(span.layers[0].compressed_count, ratio4);
assert_eq!(span.layers[0].compressed_valid_after, ratio4);
assert_eq!(span.layers[0].indexer_count, ratio4);
assert_eq!(span.layers[0].indexer_valid_after, ratio4);
assert_eq!(span.layers[1].compressed_count, ratio128);
assert_eq!(span.layers[1].compressed_valid_after, ratio128);
assert_eq!(span.layers[1].indexer_count, 0);
}
assert!(matches!(
cache.plan_prefill_start0(0),
Err(CacheError::EmptyPrefill)
));
assert!(matches!(
cache.plan_prefill_start0(129),
Err(CacheError::ContextBound {
requested: 129,
maximum: 128
})
));
assert!(matches!(
cache.commit_prefill(1, 4),
Err(CacheError::StepOutOfOrder {
expected: 0,
actual: 1
})
));
cache.commit_prefill(0, 4).unwrap();
assert_eq!(cache.position(), 4);
assert!(matches!(
cache.plan_prefill_start0(4),
Err(CacheError::PrefillNotEmpty { position: 4 })
));
cache.poison();
assert!(matches!(
cache.commit_prefill(4, 1),
Err(CacheError::Poisoned)
));
cache.reset().unwrap();
assert_eq!(cache.position(), 0);
}
#[test]
fn prefill_transactions_cross_window_and_compression_boundaries() {
let cfg = config(vec![4, 128]);
let plan = Deepseek4CachePlan::for_context(&cfg, 512).unwrap();
let _gpu = crate::inference::hf2q_gpu_test_lock();
let mut cache = Deepseek4Cache::allocate(&plan, MlxDevice::new().unwrap()).unwrap();
let first = cache.plan_prefill(127).unwrap();
assert_eq!(first.start_position, 0);
assert_eq!(first.layers[0].compressed_write_start, 0);
assert_eq!(first.layers[0].compressed_count, 31);
assert_eq!(first.layers[0].compressed_valid_after, 31);
cache.commit_prefill(0, 127).unwrap();
let boundary = cache.plan_prefill(2).unwrap();
assert_eq!(boundary.start_position, 127);
assert_eq!(boundary.layers[0].window_source_start, 0);
assert_eq!(boundary.layers[0].window_write_start, 127);
assert_eq!(boundary.layers[0].window_write_count, 2);
assert_eq!(boundary.layers[0].window_valid_after, 128);
assert_eq!(boundary.layers[0].compressed_write_start, 31);
assert_eq!(boundary.layers[0].compressed_count, 1);
assert_eq!(boundary.layers[0].compressed_valid_after, 32);
assert_eq!(boundary.layers[1].compressed_write_start, 0);
assert_eq!(boundary.layers[1].compressed_count, 1);
assert_eq!(boundary.layers[1].compressed_valid_after, 1);
cache.commit_prefill(127, 2).unwrap();
let next = cache.plan_prefill(128).unwrap();
assert_eq!(next.start_position, 129);
assert_eq!(next.layers[0].window_source_start, 0);
assert_eq!(next.layers[0].window_write_start, 1);
assert_eq!(next.layers[0].window_write_count, 128);
assert_eq!(next.layers[0].compressed_write_start, 32);
assert_eq!(next.layers[0].compressed_count, 32);
assert_eq!(next.layers[0].compressed_valid_after, 64);
assert_eq!(next.layers[1].compressed_count, 1);
assert_eq!(next.layers[1].compressed_valid_after, 2);
let wider_than_window = cache.plan_prefill(129).unwrap();
assert_eq!(wider_than_window.start_position, 129);
assert_eq!(wider_than_window.layers[0].window_source_start, 1);
assert_eq!(wider_than_window.layers[0].window_write_start, 2);
assert_eq!(wider_than_window.layers[0].window_write_count, 128);
assert_eq!(wider_than_window.layers[0].window_valid_after, 128);
assert_eq!(wider_than_window.layers[0].compressed_count, 32);
assert_eq!(wider_than_window.layers[1].compressed_count, 1);
}
#[test]
fn wide_prefill_persists_only_the_newest_non_overlapping_window() {
let cfg = config(vec![4, 128]);
let plan = Deepseek4CachePlan::for_context(&cfg, 4096).unwrap();
let _gpu = crate::inference::hf2q_gpu_test_lock();
let cache = Deepseek4Cache::allocate(&plan, MlxDevice::new().unwrap()).unwrap();
let span = cache.plan_prefill_start0(2048).unwrap();
for layer in &span.layers {
assert_eq!(layer.window_source_start, 1920);
assert_eq!(layer.window_write_start, 0);
assert_eq!(layer.window_write_count, 128);
assert_eq!(layer.window_valid_after, 128);
}
}
#[test]
fn compact_snapshot_restores_overwritten_window_state_and_position_without_aliasing() {
let mut cfg = config(vec![0, 4, 128]);
cfg.max_position_embeddings = 128;
cfg.head_dim = 8;
cfg.index_head_dim = 4;
let plan = Deepseek4CachePlan::for_context(&cfg, 128).unwrap();
let _gpu = crate::inference::hf2q_gpu_test_lock();
let mut cache = Deepseek4Cache::allocate(&plan, MlxDevice::new().unwrap()).unwrap();
cache.layers_mut()[1]
.window_kv
.as_mut_slice::<u16>()
.unwrap()[3] = 0x1234;
cache.layers_mut()[1]
.indexer_kv
.as_mut()
.unwrap()
.as_mut_slice::<u16>()
.unwrap()[2] = 0x5678;
cache.layers_mut()[1]
.main_kv_state
.as_mut()
.unwrap()
.as_mut_slice::<f32>()
.unwrap()[1] = 9.25;
for expected in 0..7 {
let step = cache.plan_next_step().unwrap();
assert_eq!(step.position, expected);
cache.commit_step(step.position).unwrap();
}
let snapshot = cache.snapshot().unwrap();
assert_eq!(snapshot.position(), 7);
let expected_snapshot_bytes: u64 = plan
.layers
.iter()
.map(|layer| {
std::iter::once(Some(&layer.window_kv))
.chain([
layer.main_kv_state.as_ref(),
layer.main_score_state.as_ref(),
layer.indexer_kv_state.as_ref(),
layer.indexer_score_state.as_ref(),
])
.flatten()
.map(|buffer| buffer.bytes)
.sum::<u64>()
})
.sum();
assert_eq!(snapshot.resident_bytes(), expected_snapshot_bytes);
assert!(snapshot.resident_bytes() < cache.resident_bytes());
cache.layers_mut()[1]
.window_kv
.as_mut_slice::<u16>()
.unwrap()[3] = 0;
cache.layers_mut()[1]
.main_kv_state
.as_mut()
.unwrap()
.as_mut_slice::<f32>()
.unwrap()[1] = -1.0;
cache.commit_step(7).unwrap();
cache.restore(&snapshot).unwrap();
assert_eq!(cache.position(), 7);
assert_eq!(
cache.layers()[1].window_kv.as_slice::<u16>().unwrap()[3],
0x1234
);
assert_eq!(
cache.layers()[1]
.indexer_kv
.as_ref()
.unwrap()
.as_slice::<u16>()
.unwrap()[2],
0x5678
);
assert_eq!(
cache.layers()[1]
.main_kv_state
.as_ref()
.unwrap()
.as_slice::<f32>()
.unwrap()[1],
9.25
);
cache.layers_mut()[1]
.window_kv
.as_mut_slice::<u16>()
.unwrap()[3] = 0xabcd;
cache.restore(&snapshot).unwrap();
assert_eq!(
cache.layers()[1].window_kv.as_slice::<u16>().unwrap()[3],
0x1234
);
}
#[test]
fn compact_snapshot_never_reads_or_copies_rows_beyond_the_cursor() {
let mut cfg = config(vec![4]);
cfg.max_position_embeddings = 128;
cfg.head_dim = 8;
cfg.index_head_dim = 4;
let plan = Deepseek4CachePlan::for_context(&cfg, 128).unwrap();
let _gpu = crate::inference::hf2q_gpu_test_lock();
let device = MlxDevice::new().unwrap();
let mut source = Deepseek4Cache::allocate_logical(&plan, device.clone()).unwrap();
let mut target = Deepseek4Cache::allocate(&plan, device).unwrap();
source.layers_mut()[0]
.window_kv
.as_mut_slice::<u16>()
.unwrap()
.fill(0x5a5a);
source.layers_mut()[0]
.window_kv
.as_mut_slice::<u16>()
.unwrap()[..8]
.fill(0x1234);
for expected in 0..1 {
source.commit_step(expected).unwrap();
}
let snapshot = source.snapshot().unwrap();
target.restore(&snapshot).unwrap();
let restored = target.layers()[0].window_kv.as_slice::<u16>().unwrap();
assert!(restored[..8].iter().all(|&value| value == 0x1234));
assert!(
restored[8..].iter().all(|&value| value == 0),
"cursor-invisible source tail must remain destination zero state"
);
}
#[test]
fn compatible_growth_migrates_live_prefix_and_leaves_new_tail_zeroed() {
let mut cfg = config(vec![0, 4, 128]);
cfg.max_position_embeddings = 256;
cfg.head_dim = 8;
cfg.index_head_dim = 4;
let source_plan = Deepseek4CachePlan::for_context(&cfg, 128).unwrap();
let target_plan = Deepseek4CachePlan::for_context(&cfg, 256).unwrap();
let _gpu = crate::inference::hf2q_gpu_test_lock();
let mut source = Deepseek4Cache::allocate(&source_plan, MlxDevice::new().unwrap()).unwrap();
let mut target = Deepseek4Cache::allocate(&target_plan, MlxDevice::new().unwrap()).unwrap();
source.layers_mut()[1]
.window_kv
.as_mut_slice::<u16>()
.unwrap()[3] = 0x2345;
source.layers_mut()[1]
.compressed_kv
.as_mut()
.unwrap()
.as_mut_slice::<u16>()
.unwrap()[9] = 0x3456;
let compressed_tail = source.layers()[1]
.compressed_kv
.as_ref()
.unwrap()
.as_slice::<u16>()
.unwrap()
.len();
source.layers_mut()[1]
.compressed_kv
.as_mut()
.unwrap()
.as_mut_slice::<u16>()
.unwrap()[compressed_tail - 1] = 0xcafe;
source.layers_mut()[1]
.indexer_kv
.as_mut()
.unwrap()
.as_mut_slice::<u16>()
.unwrap()[5] = 0x6789;
let indexer_tail = source.layers()[1]
.indexer_kv
.as_ref()
.unwrap()
.as_slice::<u16>()
.unwrap()
.len();
source.layers_mut()[1]
.indexer_kv
.as_mut()
.unwrap()
.as_mut_slice::<u16>()
.unwrap()[indexer_tail - 1] = 0xbeef;
source.layers_mut()[1]
.main_kv_state
.as_mut()
.unwrap()
.as_mut_slice::<f32>()
.unwrap()[1] = 7.5;
for expected in 0..11 {
source.commit_step(expected).unwrap();
}
target.migrate_from(&source, None).unwrap();
assert_eq!(source.position(), 11);
assert_eq!(source.capacity(), 128);
assert_eq!(target.position(), 11);
assert_eq!(target.capacity(), 256);
assert_eq!(
target.layers()[1].window_kv.as_slice::<u16>().unwrap()[3],
0x2345
);
assert_eq!(
target.layers()[1]
.compressed_kv
.as_ref()
.unwrap()
.as_slice::<u16>()
.unwrap()[9],
0x3456
);
assert_eq!(
target.layers()[1]
.compressed_kv
.as_ref()
.unwrap()
.as_slice::<u16>()
.unwrap()[compressed_tail - 1],
0,
"uncommitted compressed tail must not migrate"
);
assert_eq!(
target.layers()[1]
.indexer_kv
.as_ref()
.unwrap()
.as_slice::<u16>()
.unwrap()[5],
0x6789
);
assert_eq!(
target.layers()[1]
.indexer_kv
.as_ref()
.unwrap()
.as_slice::<u16>()
.unwrap()[indexer_tail - 1],
0,
"uncommitted indexer tail must not migrate"
);
assert_eq!(
target.layers()[1]
.main_kv_state
.as_ref()
.unwrap()
.as_slice::<f32>()
.unwrap()[1],
7.5
);
}
#[test]
fn compatible_growth_rebinds_and_restores_the_recovery_anchor() {
let mut cfg = config(vec![4, 128]);
cfg.max_position_embeddings = 256;
cfg.head_dim = 8;
cfg.index_head_dim = 4;
let source_plan = Deepseek4CachePlan::for_context(&cfg, 128).unwrap();
let target_plan = Deepseek4CachePlan::for_context(&cfg, 256).unwrap();
let _gpu = crate::inference::hf2q_gpu_test_lock();
let mut source = Deepseek4Cache::allocate(&source_plan, MlxDevice::new().unwrap()).unwrap();
let mut target = Deepseek4Cache::allocate(&target_plan, MlxDevice::new().unwrap()).unwrap();
source.layers_mut()[0]
.window_kv
.as_mut_slice::<u16>()
.unwrap()[3] = 0x1234;
source.layers_mut()[0]
.indexer_kv
.as_mut()
.unwrap()
.as_mut_slice::<u16>()
.unwrap()[2] = 0x5678;
source.layers_mut()[0]
.main_kv_state
.as_mut()
.unwrap()
.as_mut_slice::<f32>()
.unwrap()[1] = 9.25;
for expected in 0..7 {
source.commit_step(expected).unwrap();
}
let mut snapshot = source.snapshot().unwrap();
source.layers_mut()[0]
.window_kv
.as_mut_slice::<u16>()
.unwrap()[3] = 0xabcd;
source.layers_mut()[0]
.main_kv_state
.as_mut()
.unwrap()
.as_mut_slice::<f32>()
.unwrap()[1] = -1.0;
source.commit_step(7).unwrap();
target.migrate_from(&source, Some(&mut snapshot)).unwrap();
assert_eq!(target.position(), 8);
target.restore(&snapshot).unwrap();
assert_eq!(target.position(), 7);
assert_eq!(
target.layers()[0].window_kv.as_slice::<u16>().unwrap()[3],
0x1234
);
assert_eq!(
target.layers()[0]
.indexer_kv
.as_ref()
.unwrap()
.as_slice::<u16>()
.unwrap()[2],
0x5678
);
assert_eq!(
target.layers()[0]
.main_kv_state
.as_ref()
.unwrap()
.as_slice::<f32>()
.unwrap()[1],
9.25
);
}
#[test]
fn incompatible_or_poisoned_growth_leaves_the_destination_unchanged() {
let mut source_cfg = config(vec![4]);
source_cfg.max_position_embeddings = 256;
source_cfg.head_dim = 8;
source_cfg.index_head_dim = 4;
let mut target_cfg = source_cfg.clone();
target_cfg.index_head_dim = 8;
let source_plan = Deepseek4CachePlan::for_context(&source_cfg, 128).unwrap();
let target_plan = Deepseek4CachePlan::for_context(&target_cfg, 256).unwrap();
let _gpu = crate::inference::hf2q_gpu_test_lock();
let mut source = Deepseek4Cache::allocate(&source_plan, MlxDevice::new().unwrap()).unwrap();
let mut target = Deepseek4Cache::allocate(&target_plan, MlxDevice::new().unwrap()).unwrap();
target.layers_mut()[0]
.attention_kv
.as_mut_slice::<u16>()
.unwrap()[0] = 0xbeef;
assert!(matches!(
target.migrate_from(&source, None),
Err(CacheError::MigrationPlanMismatch {
layer: 0,
kind: CacheKind::IndexerKv
})
));
assert_eq!(
target.layers()[0].attention_kv.as_slice::<u16>().unwrap()[0],
0xbeef
);
source.poison();
assert!(matches!(
target.migrate_from(&source, None),
Err(CacheError::Poisoned)
));
assert_eq!(
target.layers()[0].attention_kv.as_slice::<u16>().unwrap()[0],
0xbeef
);
}
#[test]
fn restore_rejects_a_snapshot_from_a_different_cache_plan() {
let mut cfg = config(vec![4]);
cfg.max_position_embeddings = 128;
cfg.head_dim = 8;
cfg.index_head_dim = 4;
let short = Deepseek4CachePlan::for_context(&cfg, 64).unwrap();
let long = Deepseek4CachePlan::for_context(&cfg, 128).unwrap();
let _gpu = crate::inference::hf2q_gpu_test_lock();
let short_cache = Deepseek4Cache::allocate(&short, MlxDevice::new().unwrap()).unwrap();
let snapshot = short_cache.snapshot().unwrap();
let mut long_cache = Deepseek4Cache::allocate(&long, MlxDevice::new().unwrap()).unwrap();
assert!(matches!(
long_cache.restore(&snapshot),
Err(CacheError::SnapshotPlanMismatch)
));
}