234 lines
7.6 KiB
Rust
234 lines
7.6 KiB
Rust
use super::*;
|
|
use crate::{MlxRopeAttentionParameters, MlxTextRopeParameters};
|
|
|
|
#[test]
|
|
fn layout_reuses_last_concrete_cache_of_the_same_attention_kind() -> Result<()> {
|
|
let config = sample_text_config(
|
|
&[
|
|
"full_attention",
|
|
"sliding_attention",
|
|
"full_attention",
|
|
"sliding_attention",
|
|
"full_attention",
|
|
"sliding_attention",
|
|
],
|
|
2,
|
|
);
|
|
let layout = GemmaKvCacheLayout::from_text_config(&config, 1)?;
|
|
|
|
assert_eq!(layout.first_kv_shared_layer_idx, 4);
|
|
assert_eq!(layout.layer_idx_to_cache_idx, vec![0, 1, 2, 3, 2, 3]);
|
|
assert_eq!(layout.first_full_cache_idx, Some(0));
|
|
assert_eq!(layout.first_sliding_cache_idx, Some(1));
|
|
assert_eq!(layout.cache_specs.len(), 4);
|
|
assert_eq!(layout.cache_specs[0].attention, GemmaAttentionKind::Full);
|
|
assert_eq!(layout.cache_specs[1].attention, GemmaAttentionKind::Sliding);
|
|
assert!(layout.is_kv_shared_layer(4)?);
|
|
assert!(!layout.is_kv_shared_layer(3)?);
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn layout_uses_global_kv_shape_for_full_attention_layers() -> Result<()> {
|
|
let mut config = sample_text_config(
|
|
&["sliding_attention", "full_attention", "sliding_attention"],
|
|
0,
|
|
);
|
|
config.attention_k_eq_v = true;
|
|
config.head_dim = 4;
|
|
config.global_head_dim = 8;
|
|
config.num_key_value_heads = 2;
|
|
config.num_global_key_value_heads = Some(1);
|
|
config.max_position_embeddings = 32;
|
|
config.sliding_window = 6;
|
|
|
|
let layout = GemmaKvCacheLayout::from_text_config(&config, 1)?;
|
|
|
|
assert_eq!(layout.cache_specs.len(), 3);
|
|
assert_eq!(
|
|
layout.cache_specs[0],
|
|
GemmaKvCacheSpec::new(GemmaAttentionKind::Sliding, 1, 2, 4, 6)?
|
|
);
|
|
assert_eq!(
|
|
layout.cache_specs[1],
|
|
GemmaKvCacheSpec::new(GemmaAttentionKind::Full, 1, 1, 8, 32)?
|
|
);
|
|
assert_eq!(
|
|
layout.cache_specs[2],
|
|
GemmaKvCacheSpec::new(GemmaAttentionKind::Sliding, 1, 2, 4, 6)?
|
|
);
|
|
assert_eq!(layout.cache_spec_for_layer(1)?, &layout.cache_specs[1]);
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn full_cache_prefill_then_decode_preserves_head_layout() -> Result<()> {
|
|
let spec = GemmaKvCacheSpec::new(GemmaAttentionKind::Full, 1, 2, 3, 8)?;
|
|
let mut cache = GemmaKvCache::<i32>::new(spec)?;
|
|
|
|
let prefill_keys = test_tensor(
|
|
KvTensorShape {
|
|
batch_size: 1,
|
|
kv_head_count: 2,
|
|
seq_len: 2,
|
|
head_dim: 3,
|
|
},
|
|
0,
|
|
)?;
|
|
let prefill_values = test_tensor(prefill_keys.shape(), 1000)?;
|
|
cache.update_and_fetch(prefill_keys.view(), prefill_values.view())?;
|
|
|
|
let decode_keys = test_tensor(
|
|
KvTensorShape {
|
|
batch_size: 1,
|
|
kv_head_count: 2,
|
|
seq_len: 1,
|
|
head_dim: 3,
|
|
},
|
|
200,
|
|
)?;
|
|
let decode_values = test_tensor(decode_keys.shape(), 1200)?;
|
|
let state = cache.update_and_fetch(decode_keys.view(), decode_values.view())?;
|
|
|
|
assert_eq!(state.stored_tokens(), 3);
|
|
assert_eq!(state.start_position(), 0);
|
|
assert_eq!(state.offset(), 3);
|
|
assert_eq!(state.keys.row(0, 0, 0)?, &[0, 1, 2]);
|
|
assert_eq!(state.keys.row(0, 0, 1)?, &[10, 11, 12]);
|
|
assert_eq!(state.keys.row(0, 0, 2)?, &[200, 201, 202]);
|
|
assert_eq!(state.keys.row(0, 1, 0)?, &[100, 101, 102]);
|
|
assert_eq!(state.keys.row(0, 1, 2)?, &[300, 301, 302]);
|
|
assert_eq!(state.values.row(0, 1, 2)?, &[1300, 1301, 1302]);
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn sliding_cache_keeps_only_the_last_window_and_tracks_absolute_positions() -> Result<()> {
|
|
let spec = GemmaKvCacheSpec::new(GemmaAttentionKind::Sliding, 1, 1, 2, 3)?;
|
|
let mut cache = GemmaKvCache::<i32>::new(spec)?;
|
|
|
|
let first = test_tensor(
|
|
KvTensorShape {
|
|
batch_size: 1,
|
|
kv_head_count: 1,
|
|
seq_len: 2,
|
|
head_dim: 2,
|
|
},
|
|
0,
|
|
)?;
|
|
cache.update_and_fetch(first.view(), first.view())?;
|
|
|
|
let second = test_tensor(
|
|
KvTensorShape {
|
|
batch_size: 1,
|
|
kv_head_count: 1,
|
|
seq_len: 2,
|
|
head_dim: 2,
|
|
},
|
|
100,
|
|
)?;
|
|
let state = cache.update_and_fetch(second.view(), second.view())?;
|
|
|
|
assert_eq!(state.stored_tokens(), 3);
|
|
assert_eq!(state.start_position(), 1);
|
|
assert_eq!(state.offset(), 4);
|
|
assert_eq!(state.keys.row(0, 0, 0)?, &[10, 11]);
|
|
assert_eq!(state.keys.row(0, 0, 1)?, &[100, 101]);
|
|
assert_eq!(state.keys.row(0, 0, 2)?, &[110, 111]);
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn decode_step_reads_match_expected_batch_head_layout() -> Result<()> {
|
|
let spec = GemmaKvCacheSpec::new(GemmaAttentionKind::Full, 2, 2, 2, 4)?;
|
|
let mut cache = GemmaKvCache::<i32>::new(spec)?;
|
|
|
|
let decode = test_tensor(
|
|
KvTensorShape {
|
|
batch_size: 2,
|
|
kv_head_count: 2,
|
|
seq_len: 1,
|
|
head_dim: 2,
|
|
},
|
|
500,
|
|
)?;
|
|
let state = cache.update_and_fetch(decode.view(), decode.view())?;
|
|
|
|
assert_eq!(state.keys.row(0, 0, 0)?, &[500, 501]);
|
|
assert_eq!(state.keys.row(0, 1, 0)?, &[600, 601]);
|
|
assert_eq!(state.keys.row(1, 0, 0)?, &[1500, 1501]);
|
|
assert_eq!(state.keys.row(1, 1, 0)?, &[1600, 1601]);
|
|
Ok(())
|
|
}
|
|
|
|
fn sample_text_config(layer_types: &[&str], num_kv_shared_layers: u32) -> MlxTextConfig {
|
|
MlxTextConfig {
|
|
attention_bias: false,
|
|
attention_dropout: 0.0,
|
|
attention_k_eq_v: false,
|
|
bos_token_id: 2,
|
|
dtype: "bfloat16".to_owned(),
|
|
enable_moe_block: true,
|
|
eos_token_id: 1,
|
|
final_logit_softcapping: 0.0,
|
|
global_head_dim: 4,
|
|
head_dim: 4,
|
|
hidden_activation: "gelu_pytorch_tanh".to_owned(),
|
|
hidden_size: 16,
|
|
hidden_size_per_layer_input: 0,
|
|
initializer_range: 0.02,
|
|
intermediate_size: 32,
|
|
layer_types: layer_types.iter().map(|item| (*item).to_owned()).collect(),
|
|
max_position_embeddings: 16,
|
|
model_type: "gemma4".to_owned(),
|
|
moe_intermediate_size: Some(32),
|
|
expert_intermediate_size: None,
|
|
num_attention_heads: 4,
|
|
num_experts: Some(8),
|
|
num_global_key_value_heads: Some(1),
|
|
num_hidden_layers: layer_types.len() as u32,
|
|
num_key_value_heads: 2,
|
|
num_kv_shared_layers,
|
|
pad_token_id: 0,
|
|
rms_norm_eps: 1e-6,
|
|
rope_parameters: MlxTextRopeParameters {
|
|
full_attention: MlxRopeAttentionParameters {
|
|
partial_rotary_factor: None,
|
|
rope_theta: 10_000.0,
|
|
rope_type: "default".to_owned(),
|
|
},
|
|
sliding_attention: MlxRopeAttentionParameters {
|
|
partial_rotary_factor: None,
|
|
rope_theta: 10_000.0,
|
|
rope_type: "default".to_owned(),
|
|
},
|
|
},
|
|
sliding_window: 4,
|
|
tie_word_embeddings: true,
|
|
top_k_experts: Some(2),
|
|
use_bidirectional_attention: Some("never".to_owned()),
|
|
use_cache: true,
|
|
use_double_wide_mlp: false,
|
|
vocab_size: 256,
|
|
vocab_size_per_layer_input: 0,
|
|
}
|
|
}
|
|
|
|
fn test_tensor(shape: KvTensorShape, base: i32) -> Result<KvTensor<i32>> {
|
|
let mut data = Vec::with_capacity(shape.element_count()?);
|
|
for batch in 0..shape.batch_size {
|
|
for head in 0..shape.kv_head_count {
|
|
for token in 0..shape.seq_len {
|
|
for dim in 0..shape.head_dim {
|
|
data.push(
|
|
base + (batch as i32) * 1000
|
|
+ (head as i32) * 100
|
|
+ (token as i32) * 10
|
|
+ dim as i32,
|
|
);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
KvTensor::from_vec(shape, data)
|
|
}
|