1798 lines
64 KiB
Rust
1798 lines
64 KiB
Rust
use super::{
|
|
bf16_round_to_f32, bf16_word_to_f32, bf16_words_from_f32_bits, bytes_from_bf16_words,
|
|
compile_default_pipeline, default_model_path, exact_qproj_layout, print_cached_artifacts,
|
|
read_bf16_buffer_bits, read_exact_kv_cache_tensor_bits, read_f32_file_as_bf16_words,
|
|
run_layer_plan, run_layer_plan_from_sequence, run_layer_plan_with_session,
|
|
run_layer_plan_with_session_from_sequence, run_layer_sequence, run_layer_sequence_from_inputs,
|
|
write_bf16_words_as_f32_file, CachedLayerInputs, CachedLayerSequenceInputs,
|
|
ExactMetalQprojLayout, ExactMetalTextRuntimeSession, Layer0CachedArtifacts, Layer0CachedPlan,
|
|
Layer0CachedStage, LayerExecutionSession, LayerTensorNames, MlxAffineQprojRowArgs,
|
|
MlxIndexedSafetensors, DECODE_ROPE_OFFSET, NORM_LEN, PREFILL_ROPE_OFFSET,
|
|
};
|
|
use crate::fnv1a64_u32_words;
|
|
use makepad_ggml::backend::metal::{BufferStorageMode, MetalBufferBindingRef, MetalSize};
|
|
use std::env::temp_dir;
|
|
use std::path::PathBuf;
|
|
use std::{fs, io::Write};
|
|
|
|
fn u4_group_terms_row_totals(
|
|
x_bf16_words: &[u16],
|
|
weights: &[u32],
|
|
scales: &[u16],
|
|
biases: &[u16],
|
|
row: usize,
|
|
weight_stride_words: usize,
|
|
qparams_per_row: usize,
|
|
) -> (u32, u32) {
|
|
const VALUES_PER_THREAD: usize = 8;
|
|
const BLOCK_SIZE: usize = VALUES_PER_THREAD * 32;
|
|
const GROUPS_PER_BLOCK: usize = BLOCK_SIZE / 64;
|
|
|
|
let row_weight_base = row * weight_stride_words;
|
|
let row_qparam_base = row * qparams_per_row;
|
|
let mut total_plain = 0.0f32;
|
|
let mut total_groupbf16 = 0.0f32;
|
|
|
|
for block in 0..(x_bf16_words.len() / BLOCK_SIZE) {
|
|
let block_weight_base = block * (BLOCK_SIZE / 8);
|
|
let block_x_base = block * BLOCK_SIZE;
|
|
for group in 0..GROUPS_PER_BLOCK {
|
|
let scale =
|
|
bf16_word_to_f32(scales[row_qparam_base + block * GROUPS_PER_BLOCK + group]);
|
|
let bias = bf16_word_to_f32(biases[row_qparam_base + block * GROUPS_PER_BLOCK + group]);
|
|
let mut group_sum = 0.0f32;
|
|
let mut group_accum = 0.0f32;
|
|
|
|
for lane in 0..8 {
|
|
let lane_in_block = group * 8 + lane;
|
|
let lane_x_base = block_x_base + lane_in_block * VALUES_PER_THREAD;
|
|
let w = weights[row_weight_base + block_weight_base + lane_in_block];
|
|
|
|
let mut x_thread = [0.0f32; VALUES_PER_THREAD];
|
|
let mut lane_sum = 0.0f32;
|
|
for i in (0..VALUES_PER_THREAD).step_by(4) {
|
|
let x0 = bf16_word_to_f32(x_bf16_words[lane_x_base + i]);
|
|
let x1 = bf16_word_to_f32(x_bf16_words[lane_x_base + i + 1]);
|
|
let x2 = bf16_word_to_f32(x_bf16_words[lane_x_base + i + 2]);
|
|
let x3 = bf16_word_to_f32(x_bf16_words[lane_x_base + i + 3]);
|
|
lane_sum += x0 + x1 + x2 + x3;
|
|
x_thread[i] = x0;
|
|
x_thread[i + 1] = x1 / 16.0;
|
|
x_thread[i + 2] = x2 / 256.0;
|
|
x_thread[i + 3] = x3 / 4096.0;
|
|
}
|
|
|
|
let ws0 = (w & 0xFFFF) as u16;
|
|
let ws1 = (w >> 16) as u16;
|
|
let lane_accum = x_thread[0] * ((ws0 & 0x000F) as f32)
|
|
+ x_thread[1] * ((ws0 & 0x00F0) as f32)
|
|
+ x_thread[2] * ((ws0 & 0x0F00) as f32)
|
|
+ x_thread[3] * ((ws0 & 0xF000) as f32)
|
|
+ x_thread[4] * ((ws1 & 0x000F) as f32)
|
|
+ x_thread[5] * ((ws1 & 0x00F0) as f32)
|
|
+ x_thread[6] * ((ws1 & 0x0F00) as f32)
|
|
+ x_thread[7] * ((ws1 & 0xF000) as f32);
|
|
|
|
group_sum += lane_sum;
|
|
group_accum += lane_accum;
|
|
}
|
|
|
|
total_plain += scale * group_accum + bias * group_sum;
|
|
total_groupbf16 +=
|
|
bf16_round_to_f32(scale * group_accum) + bf16_round_to_f32(bias * group_sum);
|
|
}
|
|
}
|
|
|
|
(
|
|
bf16_round_to_f32(total_plain).to_bits(),
|
|
bf16_round_to_f32(total_groupbf16).to_bits(),
|
|
)
|
|
}
|
|
|
|
fn empty_artifacts() -> Layer0CachedArtifacts {
|
|
Layer0CachedArtifacts {
|
|
backend_name: "test".to_string(),
|
|
model_path: PathBuf::from("test.safetensors"),
|
|
layer_idx: 0,
|
|
selected_stage: None,
|
|
prefill_rope_offset: PREFILL_ROPE_OFFSET,
|
|
decode_rope_offset: DECODE_ROPE_OFFSET,
|
|
q_head_count: 0,
|
|
k_head_count: 0,
|
|
v_head_count: 0,
|
|
q_heads_per_kv: 0,
|
|
head_dim: 0,
|
|
prefill_input_norm_bits: Vec::new(),
|
|
prefill_v_proj_bits: Vec::new(),
|
|
prefill_q_bits: Vec::new(),
|
|
prefill_k_bits: Vec::new(),
|
|
prefill_v_bits: Vec::new(),
|
|
decode_input_norm_bits: Vec::new(),
|
|
decode_v_proj_bits: Vec::new(),
|
|
decode_q_bits: Vec::new(),
|
|
decode_k_bits: Vec::new(),
|
|
decode_v_bits: Vec::new(),
|
|
full_k_bits: Vec::new(),
|
|
full_v_bits: Vec::new(),
|
|
attention_score_bits: Vec::new(),
|
|
attention_prob_bits: Vec::new(),
|
|
attention_out_bits: Vec::new(),
|
|
attention_oproj_bits: None,
|
|
post_attention_norm_bits: None,
|
|
post_attention_residual_bits: None,
|
|
pre_feedforward_norm_bits: None,
|
|
dense_gate_bits: None,
|
|
dense_up_bits: None,
|
|
dense_geglu_bits: None,
|
|
dense_down_bits: None,
|
|
router_output: None,
|
|
moe_expert_gate_bits: None,
|
|
moe_expert_up_bits: None,
|
|
moe_expert_geglu_bits: None,
|
|
moe_expert_down_bits: None,
|
|
post_ffn_norm1_bits: None,
|
|
moe_expert_out_bits: None,
|
|
moe_post_ffn_norm2_bits: None,
|
|
moe_merge_bits: None,
|
|
prefill_post_ffn_residual_bits: None,
|
|
post_ffn_residual_bits: None,
|
|
}
|
|
}
|
|
|
|
fn device_qproj_row_bits(
|
|
session: &mut LayerExecutionSession,
|
|
input_words: &[u16],
|
|
layout: ExactMetalQprojLayout,
|
|
weight_name: &str,
|
|
scales_name: &str,
|
|
biases_name: &str,
|
|
) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
|
|
let runtime = session.runtime.clone();
|
|
let x_buf = runtime.create_buffer_with_bytes(
|
|
&bytes_from_bf16_words(input_words),
|
|
BufferStorageMode::Private,
|
|
)?;
|
|
let out_buf = runtime.create_buffer(layout.out_len() * 2, BufferStorageMode::Private)?;
|
|
let pipeline = compile_default_pipeline(&runtime, "kernel_mlx_affine_qproj_row_bf16")?;
|
|
let args = MlxAffineQprojRowArgs {
|
|
n_in: u32::try_from(input_words.len())?,
|
|
weight_words_per_row: layout.weight_words_per_row,
|
|
qparams_per_row: layout.qparams_per_row,
|
|
out_rows: layout.out_rows,
|
|
};
|
|
let bindings = [
|
|
MetalBufferBindingRef {
|
|
index: 1,
|
|
buffer: &x_buf,
|
|
offset_bytes: 0,
|
|
},
|
|
MetalBufferBindingRef {
|
|
index: 2,
|
|
buffer: &session.private_weight_buffer(weight_name)?,
|
|
offset_bytes: 0,
|
|
},
|
|
MetalBufferBindingRef {
|
|
index: 3,
|
|
buffer: &session.private_weight_buffer(scales_name)?,
|
|
offset_bytes: 0,
|
|
},
|
|
MetalBufferBindingRef {
|
|
index: 4,
|
|
buffer: &session.private_weight_buffer(biases_name)?,
|
|
offset_bytes: 0,
|
|
},
|
|
MetalBufferBindingRef {
|
|
index: 5,
|
|
buffer: &out_buf,
|
|
offset_bytes: 0,
|
|
},
|
|
];
|
|
let threads_per_threadgroup = MetalSize {
|
|
width: 256,
|
|
height: 1,
|
|
depth: 1,
|
|
};
|
|
let threadgroups = MetalSize {
|
|
width: (layout.out_len() as u64).div_ceil(threads_per_threadgroup.width),
|
|
height: 1,
|
|
depth: 1,
|
|
};
|
|
runtime.begin_command_batch()?;
|
|
runtime.dispatch_compute(
|
|
&pipeline,
|
|
super::bytes_of(&args),
|
|
&bindings,
|
|
&[],
|
|
threadgroups,
|
|
threads_per_threadgroup,
|
|
)?;
|
|
runtime.end_command_batch()?;
|
|
runtime.wait_idle()?;
|
|
read_bf16_buffer_bits(&runtime, &out_buf, layout.out_len())
|
|
}
|
|
|
|
fn write_f32_bits_file(path: &PathBuf, bits: &[u32]) {
|
|
let mut bytes = Vec::with_capacity(bits.len() * size_of::<u32>());
|
|
for word in bits {
|
|
bytes.extend_from_slice(&word.to_le_bytes());
|
|
}
|
|
fs::write(path, bytes).unwrap();
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct FormattedSayHiLayerOutputs {
|
|
layer_idx: usize,
|
|
step1_prefill_hash: u64,
|
|
step1_decode_hash: u64,
|
|
step2_decode_hash: u64,
|
|
step1_prefill_f32_path: PathBuf,
|
|
step1_decode_f32_path: PathBuf,
|
|
step2_decode_f32_path: PathBuf,
|
|
}
|
|
|
|
fn post_ffn_only_plan() -> Layer0CachedPlan {
|
|
let mut plan = Layer0CachedPlan::new();
|
|
plan.require_stage(Layer0CachedStage::PostFfnResidual);
|
|
plan
|
|
}
|
|
|
|
fn write_formatted_say_hi_outputs_through_layer(
|
|
last_layer_idx: usize,
|
|
) -> Vec<FormattedSayHiLayerOutputs> {
|
|
let model_path = default_model_path();
|
|
let model_root = super::model_root_dir(&model_path).unwrap();
|
|
let weights = MlxIndexedSafetensors::load(&model_root).unwrap();
|
|
let mut session = LayerExecutionSession::load(model_path).unwrap();
|
|
let mut token2_words = weights.embed_token_bf16_words(2).unwrap();
|
|
let mut token105_words = weights.embed_token_bf16_words(105).unwrap();
|
|
let mut token2364_words = weights.embed_token_bf16_words(2_364).unwrap();
|
|
let tmp_dir = temp_dir();
|
|
let mut outputs = Vec::with_capacity(last_layer_idx + 1);
|
|
|
|
for layer_idx in 0..=last_layer_idx {
|
|
let step1 = run_layer_plan_with_session(
|
|
&mut session,
|
|
layer_idx,
|
|
CachedLayerInputs {
|
|
prefill_input_words: token2_words.clone(),
|
|
decode_input_words: token105_words.clone(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: 1,
|
|
validate_against_oracle: false,
|
|
},
|
|
post_ffn_only_plan(),
|
|
)
|
|
.unwrap();
|
|
let step2 = run_layer_plan_with_session_from_sequence(
|
|
&mut session,
|
|
layer_idx,
|
|
CachedLayerSequenceInputs {
|
|
prefill_input_words_list: vec![token2_words.clone(), token105_words.clone()],
|
|
decode_input_words: token2364_words.clone(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: 2,
|
|
validate_against_oracle: false,
|
|
},
|
|
post_ffn_only_plan(),
|
|
)
|
|
.unwrap();
|
|
|
|
let step1_prefill_words = step1
|
|
.prefill_layer_output_bf16_words()
|
|
.expect("missing step1 prefill output");
|
|
let step1_decode_words = step1
|
|
.bf16_words_for_stage(Layer0CachedStage::PostFfnResidual)
|
|
.expect("missing step1 decode output");
|
|
let step2_decode_words = step2
|
|
.bf16_words_for_stage(Layer0CachedStage::PostFfnResidual)
|
|
.expect("missing step2 decode output");
|
|
|
|
let step1_prefill_f32_path =
|
|
tmp_dir.join(format!("gemma_say_hi_layer{layer_idx}_prefill_2_f32.bin"));
|
|
let step1_decode_f32_path =
|
|
tmp_dir.join(format!("gemma_say_hi_layer{layer_idx}_decode_105_f32.bin"));
|
|
let step2_decode_f32_path =
|
|
tmp_dir.join(format!("gemma_say_hi_layer{layer_idx}_decode_2364_f32.bin"));
|
|
write_bf16_words_as_f32_file(&step1_prefill_f32_path, &step1_prefill_words).unwrap();
|
|
write_bf16_words_as_f32_file(&step1_decode_f32_path, &step1_decode_words).unwrap();
|
|
write_bf16_words_as_f32_file(&step2_decode_f32_path, &step2_decode_words).unwrap();
|
|
|
|
let entry = FormattedSayHiLayerOutputs {
|
|
layer_idx,
|
|
step1_prefill_hash: fnv1a64_u32_words(
|
|
step1
|
|
.prefill_layer_output_bits()
|
|
.expect("missing step1 prefill bits"),
|
|
),
|
|
step1_decode_hash: fnv1a64_u32_words(
|
|
step1
|
|
.layer_output_bits()
|
|
.expect("missing step1 decode bits"),
|
|
),
|
|
step2_decode_hash: fnv1a64_u32_words(
|
|
step2
|
|
.layer_output_bits()
|
|
.expect("missing step2 decode bits"),
|
|
),
|
|
step1_prefill_f32_path,
|
|
step1_decode_f32_path,
|
|
step2_decode_f32_path,
|
|
};
|
|
outputs.push(entry);
|
|
|
|
token2_words = step1_prefill_words;
|
|
token105_words = step1_decode_words;
|
|
token2364_words = step2_decode_words;
|
|
}
|
|
|
|
outputs
|
|
}
|
|
|
|
fn rms_norm_unweighted_rows_f32(x: &[f32], head_dim: usize, eps: f32) -> Vec<f32> {
|
|
assert!(head_dim != 0);
|
|
assert_eq!(x.len() % head_dim, 0);
|
|
let mut out = Vec::with_capacity(x.len());
|
|
for row in x.chunks_exact(head_dim) {
|
|
let mut mean_square = 0.0f32;
|
|
for value in row {
|
|
mean_square += value * value;
|
|
}
|
|
mean_square /= head_dim as f32;
|
|
let inv_rms = 1.0f32 / (mean_square + eps).sqrt();
|
|
for value in row {
|
|
out.push(bf16_round_to_f32(*value * inv_rms));
|
|
}
|
|
}
|
|
out
|
|
}
|
|
|
|
fn formatted_say_hi_prompt_token_ids() -> Vec<u32> {
|
|
vec![2, 105, 2364, 107, 30_468, 5_631, 106, 107, 105, 4_368, 107]
|
|
}
|
|
|
|
fn teacher_forced_prompt_hidden_words_through_layer(last_layer_idx: usize) -> Vec<Vec<u16>> {
|
|
let prompt_token_ids = formatted_say_hi_prompt_token_ids();
|
|
let model_path = default_model_path();
|
|
let mut session = LayerExecutionSession::load(model_path).unwrap();
|
|
let weights = session.weights.clone();
|
|
let mut token_hidden_words = prompt_token_ids
|
|
.iter()
|
|
.map(|token_id| weights.embed_token_bf16_words(*token_id).unwrap())
|
|
.collect::<Vec<_>>();
|
|
|
|
for layer_idx in 0..=last_layer_idx {
|
|
let mut next_token_hidden_words = vec![Vec::new(); prompt_token_ids.len()];
|
|
for pos in 0..(prompt_token_ids.len() - 1) {
|
|
let artifacts = run_layer_plan_with_session_from_sequence(
|
|
&mut session,
|
|
layer_idx,
|
|
CachedLayerSequenceInputs {
|
|
prefill_input_words_list: token_hidden_words[..=pos].to_vec(),
|
|
decode_input_words: token_hidden_words[pos + 1].clone(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: i32::try_from(pos + 1).unwrap(),
|
|
validate_against_oracle: false,
|
|
},
|
|
post_ffn_only_plan(),
|
|
)
|
|
.unwrap();
|
|
|
|
let prefill_output_words = artifacts
|
|
.prefill_layer_output_bf16_words()
|
|
.expect("missing teacher-forced prefill output");
|
|
let decode_output_words = artifacts
|
|
.bf16_words_for_stage(Layer0CachedStage::PostFfnResidual)
|
|
.expect("missing teacher-forced decode output");
|
|
|
|
if pos == 0 {
|
|
next_token_hidden_words[pos] = prefill_output_words;
|
|
} else {
|
|
assert_eq!(
|
|
next_token_hidden_words[pos], prefill_output_words,
|
|
"teacher-forced layer {} token {} prefill output disagreed with prior decode path",
|
|
layer_idx, pos
|
|
);
|
|
}
|
|
next_token_hidden_words[pos + 1] = decode_output_words;
|
|
}
|
|
token_hidden_words = next_token_hidden_words;
|
|
}
|
|
|
|
token_hidden_words
|
|
}
|
|
|
|
fn write_teacher_forced_hidden_inputs_for_oracle(
|
|
input_layer_idx: usize,
|
|
decode_token_position: usize,
|
|
) {
|
|
let token_hidden_words = teacher_forced_prompt_hidden_words_through_layer(input_layer_idx);
|
|
let tmp_dir = temp_dir();
|
|
let mut prefill_paths = Vec::new();
|
|
for token_position in 0..decode_token_position {
|
|
let path = tmp_dir.join(format!(
|
|
"gemma_say_hi_layer{input_layer_idx}_token{token_position}_f32.bin"
|
|
));
|
|
write_bf16_words_as_f32_file(&path, &token_hidden_words[token_position]).unwrap();
|
|
println!(
|
|
"token_position={} hidden_fnv1a64=0x{:016X} f32_path={}",
|
|
token_position,
|
|
fnv1a64_u32_words(
|
|
&token_hidden_words[token_position]
|
|
.iter()
|
|
.copied()
|
|
.map(bf16_word_to_f32)
|
|
.map(f32::to_bits)
|
|
.collect::<Vec<_>>()
|
|
),
|
|
path.display()
|
|
);
|
|
prefill_paths.push(path);
|
|
}
|
|
|
|
let decode_path = tmp_dir.join(format!(
|
|
"gemma_say_hi_layer{input_layer_idx}_token{decode_token_position}_f32.bin"
|
|
));
|
|
write_bf16_words_as_f32_file(&decode_path, &token_hidden_words[decode_token_position]).unwrap();
|
|
println!(
|
|
"token_position={} hidden_fnv1a64=0x{:016X} f32_path={}",
|
|
decode_token_position,
|
|
fnv1a64_u32_words(
|
|
&token_hidden_words[decode_token_position]
|
|
.iter()
|
|
.copied()
|
|
.map(bf16_word_to_f32)
|
|
.map(f32::to_bits)
|
|
.collect::<Vec<_>>()
|
|
),
|
|
decode_path.display()
|
|
);
|
|
println!(
|
|
"prefill_input_f32_files={}",
|
|
prefill_paths
|
|
.iter()
|
|
.map(|path| path.display().to_string())
|
|
.collect::<Vec<_>>()
|
|
.join(",")
|
|
);
|
|
println!("decode_input_f32_file={}", decode_path.display());
|
|
}
|
|
|
|
fn teacher_forced_prompt_step_artifacts(
|
|
layer_idx: usize,
|
|
input_layer_idx: usize,
|
|
decode_token_position: usize,
|
|
) -> Layer0CachedArtifacts {
|
|
let token_hidden_words = teacher_forced_prompt_hidden_words_through_layer(input_layer_idx);
|
|
let model_path = default_model_path();
|
|
let mut session = LayerExecutionSession::load(model_path).unwrap();
|
|
run_layer_plan_with_session_from_sequence(
|
|
&mut session,
|
|
layer_idx,
|
|
CachedLayerSequenceInputs {
|
|
prefill_input_words_list: token_hidden_words[..decode_token_position].to_vec(),
|
|
decode_input_words: token_hidden_words[decode_token_position].clone(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: i32::try_from(decode_token_position).unwrap(),
|
|
validate_against_oracle: false,
|
|
},
|
|
post_ffn_only_plan(),
|
|
)
|
|
.unwrap()
|
|
}
|
|
|
|
fn teacher_env_usize(name: &str) -> usize {
|
|
std::env::var(name)
|
|
.unwrap_or_else(|_| panic!("missing {name}"))
|
|
.parse()
|
|
.unwrap_or_else(|_| panic!("invalid {name}"))
|
|
}
|
|
|
|
#[test]
|
|
fn post_ffn_residual_plan_pulls_in_full_cached_layer_path() {
|
|
let mut plan = Layer0CachedPlan::new();
|
|
plan.require_stage(Layer0CachedStage::PostFfnResidual);
|
|
|
|
for stage in [
|
|
Layer0CachedStage::AttentionOproj,
|
|
Layer0CachedStage::PostAttentionResidual,
|
|
Layer0CachedStage::PreFeedforwardNorm,
|
|
Layer0CachedStage::DenseGate,
|
|
Layer0CachedStage::DenseUp,
|
|
Layer0CachedStage::DenseGeGlu,
|
|
Layer0CachedStage::DenseDown,
|
|
Layer0CachedStage::PostFfnNorm1,
|
|
Layer0CachedStage::Router,
|
|
Layer0CachedStage::MoeExpertGate,
|
|
Layer0CachedStage::MoeExpertUp,
|
|
Layer0CachedStage::MoeExpertGeGlu,
|
|
Layer0CachedStage::MoeExpertDown,
|
|
Layer0CachedStage::MoeExpertOut,
|
|
Layer0CachedStage::MoePostFfnNorm2,
|
|
Layer0CachedStage::MoeMerge,
|
|
Layer0CachedStage::PostFfnResidual,
|
|
] {
|
|
assert!(plan.requires(stage), "missing dependency for {stage:?}");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn display_stage_matches_previous_priority_order() {
|
|
let mut plan = Layer0CachedPlan::new();
|
|
plan.require_stage(Layer0CachedStage::Router);
|
|
plan.require_stage(Layer0CachedStage::DenseUp);
|
|
|
|
assert_eq!(plan.display_stage(), Some(Layer0CachedStage::DenseUp));
|
|
}
|
|
|
|
#[test]
|
|
fn evaluation_order_is_dependency_first() {
|
|
let mut plan = Layer0CachedPlan::new();
|
|
plan.require_stage(Layer0CachedStage::MoeMerge);
|
|
|
|
assert_eq!(
|
|
plan.evaluation_order().last().copied(),
|
|
Some(Layer0CachedStage::MoeMerge)
|
|
);
|
|
assert!(plan.requires(Layer0CachedStage::PostFfnNorm1));
|
|
assert!(plan.requires(Layer0CachedStage::MoePostFfnNorm2));
|
|
}
|
|
|
|
#[test]
|
|
fn artifacts_expose_stage_bits_and_bf16_words() {
|
|
let mut artifacts = empty_artifacts();
|
|
artifacts.post_ffn_residual_bits = Some(vec![0x3F80_0000, 0xC020_0000]);
|
|
|
|
assert_eq!(
|
|
artifacts.layer_output_bits(),
|
|
Some([0x3F80_0000, 0xC020_0000].as_slice())
|
|
);
|
|
assert_eq!(
|
|
artifacts.bf16_words_for_stage(Layer0CachedStage::PostFfnResidual),
|
|
Some(vec![0x3F80, 0xC020])
|
|
);
|
|
assert_eq!(
|
|
artifacts.tensor_bits_for_stage(Layer0CachedStage::Router),
|
|
None
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_layer28_step1_v_path_matches_local_mlx_math() {
|
|
let outputs = write_formatted_say_hi_outputs_through_layer(27);
|
|
let layer27 = outputs.last().expect("missing layer 27 outputs");
|
|
let model_path = default_model_path();
|
|
let model_root = super::model_root_dir(&model_path).unwrap();
|
|
let weights = MlxIndexedSafetensors::load(&model_root).unwrap();
|
|
let text_config = &weights.snapshot.config.text_config;
|
|
let layer_idx = 28usize;
|
|
let layer_type = text_config.layer_types.get(layer_idx).unwrap();
|
|
let attention_k_eq_v = text_config.attention_k_eq_v && layer_type == "full_attention";
|
|
let layer_names = LayerTensorNames::for_layer(layer_idx, attention_k_eq_v);
|
|
|
|
let mut session = LayerExecutionSession::load(model_path).unwrap();
|
|
let artifacts = run_layer_plan_with_session_from_sequence(
|
|
&mut session,
|
|
layer_idx,
|
|
CachedLayerSequenceInputs {
|
|
prefill_input_words_list: vec![read_f32_file_as_bf16_words(
|
|
&layer27.step1_prefill_f32_path,
|
|
)
|
|
.unwrap()],
|
|
decode_input_words: read_f32_file_as_bf16_words(&layer27.step1_decode_f32_path)
|
|
.unwrap(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: 1,
|
|
validate_against_oracle: false,
|
|
},
|
|
post_ffn_only_plan(),
|
|
)
|
|
.unwrap();
|
|
|
|
let expected_input_norm = weights
|
|
.header_for_tensor(&layer_names.input_norm_weight_name)
|
|
.unwrap()
|
|
.rms_norm_weighted_f32(
|
|
&read_f32_file_as_bf16_words(&layer27.step1_decode_f32_path).unwrap(),
|
|
&layer_names.input_norm_weight_name,
|
|
weights.snapshot.config.text_config.rms_norm_eps,
|
|
)
|
|
.unwrap();
|
|
let expected_input_norm_bits = expected_input_norm
|
|
.iter()
|
|
.copied()
|
|
.map(f32::to_bits)
|
|
.collect::<Vec<_>>();
|
|
let decode_h_words = bf16_words_from_f32_bits(&artifacts.decode_input_norm_bits);
|
|
let expected_v_proj = weights
|
|
.header_for_tensor(&layer_names.v.weight_name)
|
|
.unwrap()
|
|
.affine_quantized_matmul_t_f32(
|
|
&decode_h_words,
|
|
&layer_names.v.weight_name,
|
|
&layer_names.v.scales_name,
|
|
&layer_names.v.biases_name,
|
|
weights.snapshot.config.quantization.group_size as u64,
|
|
weights.snapshot.config.quantization.bits,
|
|
)
|
|
.unwrap();
|
|
let expected_k_proj = weights
|
|
.header_for_tensor(&layer_names.k.weight_name)
|
|
.unwrap()
|
|
.affine_quantized_matmul_t_f32(
|
|
&decode_h_words,
|
|
&layer_names.k.weight_name,
|
|
&layer_names.k.scales_name,
|
|
&layer_names.k.biases_name,
|
|
weights.snapshot.config.quantization.group_size as u64,
|
|
weights.snapshot.config.quantization.bits,
|
|
)
|
|
.unwrap();
|
|
let expected_v_proj_bits = expected_v_proj
|
|
.iter()
|
|
.copied()
|
|
.map(f32::to_bits)
|
|
.collect::<Vec<_>>();
|
|
let expected_k_proj_bits = expected_k_proj
|
|
.iter()
|
|
.copied()
|
|
.map(f32::to_bits)
|
|
.collect::<Vec<_>>();
|
|
let expected_v_norm_bits = rms_norm_unweighted_rows_f32(
|
|
&expected_v_proj,
|
|
artifacts.head_dim,
|
|
weights.snapshot.config.text_config.rms_norm_eps,
|
|
)
|
|
.into_iter()
|
|
.map(f32::to_bits)
|
|
.collect::<Vec<_>>();
|
|
eprintln!(
|
|
"layer28_step1 decode_input_norm expected=0x{:016X} actual=0x{:016X}",
|
|
fnv1a64_u32_words(&expected_input_norm_bits),
|
|
fnv1a64_u32_words(&artifacts.decode_input_norm_bits),
|
|
);
|
|
eprintln!(
|
|
"layer28_step1 decode_k_proj expected=0x{:016X}",
|
|
fnv1a64_u32_words(&expected_k_proj_bits),
|
|
);
|
|
eprintln!(
|
|
"layer28_step1 decode_v_proj expected=0x{:016X} actual=0x{:016X}",
|
|
fnv1a64_u32_words(&expected_v_proj_bits),
|
|
fnv1a64_u32_words(&artifacts.decode_v_proj_bits),
|
|
);
|
|
eprintln!(
|
|
"layer28_step1 decode_v_proj expected_first16={}",
|
|
expected_v_proj_bits
|
|
.iter()
|
|
.take(16)
|
|
.map(|bits| format!("0x{bits:08X}"))
|
|
.collect::<Vec<_>>()
|
|
.join(",")
|
|
);
|
|
eprintln!(
|
|
"layer28_step1 decode_v_proj actual_first16={}",
|
|
artifacts
|
|
.decode_v_proj_bits
|
|
.iter()
|
|
.take(16)
|
|
.map(|bits| format!("0x{bits:08X}"))
|
|
.collect::<Vec<_>>()
|
|
.join(",")
|
|
);
|
|
|
|
assert_eq!(
|
|
fnv1a64_u32_words(&expected_input_norm_bits),
|
|
fnv1a64_u32_words(&artifacts.decode_input_norm_bits),
|
|
"layer 28 step1 decode_input_norm diverged from local MLX math"
|
|
);
|
|
assert_eq!(
|
|
fnv1a64_u32_words(&expected_v_proj_bits),
|
|
fnv1a64_u32_words(&artifacts.decode_v_proj_bits),
|
|
"layer 28 step1 decode_v_proj diverged from local MLX math"
|
|
);
|
|
assert_eq!(
|
|
fnv1a64_u32_words(&expected_v_norm_bits),
|
|
fnv1a64_u32_words(&artifacts.decode_v_bits),
|
|
"layer 28 step1 decode_v_norm diverged from local MLX math"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_layer28_step1_v_row665_term_models() {
|
|
let outputs = write_formatted_say_hi_outputs_through_layer(27);
|
|
let layer27 = outputs.last().expect("missing layer 27 outputs");
|
|
let model_path = default_model_path();
|
|
let model_root = super::model_root_dir(&model_path).unwrap();
|
|
let weights = MlxIndexedSafetensors::load(&model_root).unwrap();
|
|
let text_config = &weights.snapshot.config.text_config;
|
|
let layer_idx = 28usize;
|
|
let layer_type = text_config.layer_types.get(layer_idx).unwrap();
|
|
let attention_k_eq_v = text_config.attention_k_eq_v && layer_type == "full_attention";
|
|
let layer_names = LayerTensorNames::for_layer(layer_idx, attention_k_eq_v);
|
|
|
|
let mut session = LayerExecutionSession::load(model_path).unwrap();
|
|
let artifacts = run_layer_plan_with_session_from_sequence(
|
|
&mut session,
|
|
layer_idx,
|
|
CachedLayerSequenceInputs {
|
|
prefill_input_words_list: vec![read_f32_file_as_bf16_words(
|
|
&layer27.step1_prefill_f32_path,
|
|
)
|
|
.unwrap()],
|
|
decode_input_words: read_f32_file_as_bf16_words(&layer27.step1_decode_f32_path)
|
|
.unwrap(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: 1,
|
|
validate_against_oracle: false,
|
|
},
|
|
post_ffn_only_plan(),
|
|
)
|
|
.unwrap();
|
|
|
|
let v_weight_entry = weights.tensor(&layer_names.v.weight_name).unwrap();
|
|
let v_scales_entry = weights.tensor(&layer_names.v.scales_name).unwrap();
|
|
let v_weights = weights
|
|
.header_for_tensor(&layer_names.v.weight_name)
|
|
.unwrap()
|
|
.read_u32_tensor_words(&layer_names.v.weight_name)
|
|
.unwrap();
|
|
let v_scales = weights
|
|
.read_bf16_tensor_words(&layer_names.v.scales_name)
|
|
.unwrap();
|
|
let v_biases = weights
|
|
.read_bf16_tensor_words(&layer_names.v.biases_name)
|
|
.unwrap();
|
|
let decode_h_words = bf16_words_from_f32_bits(&artifacts.decode_input_norm_bits);
|
|
let row = 665usize;
|
|
let (plain_bits, groupbf16_bits) = u4_group_terms_row_totals(
|
|
&decode_h_words,
|
|
&v_weights,
|
|
&v_scales,
|
|
&v_biases,
|
|
row,
|
|
v_weight_entry.shape[1] as usize,
|
|
v_scales_entry.shape[1] as usize,
|
|
);
|
|
|
|
eprintln!(
|
|
"layer28_step1 row665 actual=0x{:08X} plain_seq=0x{:08X} groupbf16_seq=0x{:08X}",
|
|
artifacts.decode_v_proj_bits[row], plain_bits, groupbf16_bits,
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_layer28_step1_device_qproj_matches_device_qmv() {
|
|
let outputs = write_formatted_say_hi_outputs_through_layer(27);
|
|
let layer27 = outputs.last().expect("missing layer 27 outputs");
|
|
let model_path = default_model_path();
|
|
let model_root = super::model_root_dir(&model_path).unwrap();
|
|
let weights = MlxIndexedSafetensors::load(&model_root).unwrap();
|
|
let text_config = &weights.snapshot.config.text_config;
|
|
let layer_idx = 28usize;
|
|
let layer_type = text_config.layer_types.get(layer_idx).unwrap();
|
|
let attention_k_eq_v = text_config.attention_k_eq_v && layer_type == "full_attention";
|
|
let layer_names = LayerTensorNames::for_layer(layer_idx, attention_k_eq_v);
|
|
let v_weight_entry = weights.tensor(&layer_names.v.weight_name).unwrap();
|
|
let v_scales_entry = weights.tensor(&layer_names.v.scales_name).unwrap();
|
|
let v_layout = exact_qproj_layout(
|
|
v_weight_entry.shape[1] as u32,
|
|
v_scales_entry.shape[1] as u32,
|
|
u32::try_from(v_weight_entry.shape[0]).unwrap(),
|
|
weights.snapshot.config.quantization.bits,
|
|
);
|
|
|
|
let mut session = LayerExecutionSession::load(model_path).unwrap();
|
|
let artifacts = run_layer_plan_with_session_from_sequence(
|
|
&mut session,
|
|
layer_idx,
|
|
CachedLayerSequenceInputs {
|
|
prefill_input_words_list: vec![read_f32_file_as_bf16_words(
|
|
&layer27.step1_prefill_f32_path,
|
|
)
|
|
.unwrap()],
|
|
decode_input_words: read_f32_file_as_bf16_words(&layer27.step1_decode_f32_path)
|
|
.unwrap(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: 1,
|
|
validate_against_oracle: false,
|
|
},
|
|
post_ffn_only_plan(),
|
|
)
|
|
.unwrap();
|
|
|
|
let decode_h_words = bf16_words_from_f32_bits(&artifacts.decode_input_norm_bits);
|
|
let qproj_bits = device_qproj_row_bits(
|
|
&mut session,
|
|
&decode_h_words,
|
|
v_layout,
|
|
&layer_names.v.weight_name,
|
|
&layer_names.v.scales_name,
|
|
&layer_names.v.biases_name,
|
|
)
|
|
.unwrap();
|
|
|
|
eprintln!(
|
|
"layer28_step1 device_qproj expected=0x{:016X} qmv=0x{:016X}",
|
|
fnv1a64_u32_words(&qproj_bits),
|
|
fnv1a64_u32_words(&artifacts.decode_v_proj_bits),
|
|
);
|
|
eprintln!(
|
|
"layer28_step1 device_qproj first16={}",
|
|
qproj_bits
|
|
.iter()
|
|
.take(16)
|
|
.map(|bits| format!("0x{bits:08X}"))
|
|
.collect::<Vec<_>>()
|
|
.join(",")
|
|
);
|
|
|
|
assert_eq!(
|
|
fnv1a64_u32_words(&qproj_bits),
|
|
fnv1a64_u32_words(&artifacts.decode_v_proj_bits),
|
|
"layer 28 step1 device qproj row kernel diverged from qmv output"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_layer28_step1_writes_rust_decode_v_proj_file() {
|
|
let outputs = write_formatted_say_hi_outputs_through_layer(27);
|
|
let layer27 = outputs.last().expect("missing layer 27 outputs");
|
|
let model_path = default_model_path();
|
|
let mut session = LayerExecutionSession::load(model_path).unwrap();
|
|
let artifacts = run_layer_plan_with_session_from_sequence(
|
|
&mut session,
|
|
28,
|
|
CachedLayerSequenceInputs {
|
|
prefill_input_words_list: vec![read_f32_file_as_bf16_words(
|
|
&layer27.step1_prefill_f32_path,
|
|
)
|
|
.unwrap()],
|
|
decode_input_words: read_f32_file_as_bf16_words(&layer27.step1_decode_f32_path)
|
|
.unwrap(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: 1,
|
|
validate_against_oracle: false,
|
|
},
|
|
post_ffn_only_plan(),
|
|
)
|
|
.unwrap();
|
|
|
|
let out_path = temp_dir().join("rust_layer28_step1_decode_v_proj_f32.bin");
|
|
write_f32_bits_file(&out_path, &artifacts.decode_v_proj_bits);
|
|
println!("rust_decode_v_proj_f32_path={}", out_path.display());
|
|
println!(
|
|
"rust_decode_v_proj_fnv1a64=0x{:016X}",
|
|
fnv1a64_u32_words(&artifacts.decode_v_proj_bits)
|
|
);
|
|
assert_eq!(artifacts.decode_v_proj_bits.len(), 2048);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_layer28_step1_writes_rust_decode_v_norm_file() {
|
|
let outputs = write_formatted_say_hi_outputs_through_layer(27);
|
|
let layer27 = outputs.last().expect("missing layer 27 outputs");
|
|
let model_path = default_model_path();
|
|
let mut session = LayerExecutionSession::load(model_path).unwrap();
|
|
let artifacts = run_layer_plan_with_session_from_sequence(
|
|
&mut session,
|
|
28,
|
|
CachedLayerSequenceInputs {
|
|
prefill_input_words_list: vec![read_f32_file_as_bf16_words(
|
|
&layer27.step1_prefill_f32_path,
|
|
)
|
|
.unwrap()],
|
|
decode_input_words: read_f32_file_as_bf16_words(&layer27.step1_decode_f32_path)
|
|
.unwrap(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: 1,
|
|
validate_against_oracle: false,
|
|
},
|
|
post_ffn_only_plan(),
|
|
)
|
|
.unwrap();
|
|
|
|
let out_path = temp_dir().join("rust_layer28_step1_decode_v_norm_f32.bin");
|
|
write_f32_bits_file(&out_path, &artifacts.decode_v_bits);
|
|
println!("rust_decode_v_norm_f32_path={}", out_path.display());
|
|
println!(
|
|
"rust_decode_v_norm_fnv1a64=0x{:016X}",
|
|
fnv1a64_u32_words(&artifacts.decode_v_bits)
|
|
);
|
|
assert_eq!(artifacts.decode_v_bits.len(), 2048);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_layer28_step1_writes_rust_post_ffn_residual_file() {
|
|
let outputs = write_formatted_say_hi_outputs_through_layer(27);
|
|
let layer27 = outputs.last().expect("missing layer 27 outputs");
|
|
let model_path = default_model_path();
|
|
let mut session = LayerExecutionSession::load(model_path).unwrap();
|
|
let artifacts = run_layer_plan_with_session_from_sequence(
|
|
&mut session,
|
|
28,
|
|
CachedLayerSequenceInputs {
|
|
prefill_input_words_list: vec![read_f32_file_as_bf16_words(
|
|
&layer27.step1_prefill_f32_path,
|
|
)
|
|
.unwrap()],
|
|
decode_input_words: read_f32_file_as_bf16_words(&layer27.step1_decode_f32_path)
|
|
.unwrap(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: 1,
|
|
validate_against_oracle: false,
|
|
},
|
|
post_ffn_only_plan(),
|
|
)
|
|
.unwrap();
|
|
|
|
let out_bits = artifacts
|
|
.post_ffn_residual_bits
|
|
.as_ref()
|
|
.expect("missing post-ffn residual bits");
|
|
let out_path = temp_dir().join("rust_layer28_step1_post_ffn_residual_f32.bin");
|
|
write_f32_bits_file(&out_path, out_bits);
|
|
println!("rust_post_ffn_residual_f32_path={}", out_path.display());
|
|
println!(
|
|
"rust_post_ffn_residual_fnv1a64=0x{:016X}",
|
|
fnv1a64_u32_words(out_bits)
|
|
);
|
|
assert_eq!(out_bits.len(), 2816);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_layer28_step1_writes_rust_attention_output_file() {
|
|
let outputs = write_formatted_say_hi_outputs_through_layer(27);
|
|
let layer27 = outputs.last().expect("missing layer 27 outputs");
|
|
let model_path = default_model_path();
|
|
let mut session = LayerExecutionSession::load(model_path).unwrap();
|
|
let artifacts = run_layer_plan_with_session_from_sequence(
|
|
&mut session,
|
|
28,
|
|
CachedLayerSequenceInputs {
|
|
prefill_input_words_list: vec![read_f32_file_as_bf16_words(
|
|
&layer27.step1_prefill_f32_path,
|
|
)
|
|
.unwrap()],
|
|
decode_input_words: read_f32_file_as_bf16_words(&layer27.step1_decode_f32_path)
|
|
.unwrap(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: 1,
|
|
validate_against_oracle: false,
|
|
},
|
|
post_ffn_only_plan(),
|
|
)
|
|
.unwrap();
|
|
|
|
let out_path = temp_dir().join("rust_layer28_step1_attention_output_f32.bin");
|
|
write_f32_bits_file(&out_path, &artifacts.attention_out_bits);
|
|
println!("rust_attention_output_f32_path={}", out_path.display());
|
|
println!(
|
|
"rust_attention_output_fnv1a64=0x{:016X}",
|
|
fnv1a64_u32_words(&artifacts.attention_out_bits)
|
|
);
|
|
assert_eq!(artifacts.attention_out_bits.len(), 8192);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_layer28_step1_writes_rust_full_v_cache_file() {
|
|
let outputs = write_formatted_say_hi_outputs_through_layer(27);
|
|
let layer27 = outputs.last().expect("missing layer 27 outputs");
|
|
let model_path = default_model_path();
|
|
let mut session = LayerExecutionSession::load(model_path).unwrap();
|
|
let artifacts = run_layer_plan_with_session_from_sequence(
|
|
&mut session,
|
|
28,
|
|
CachedLayerSequenceInputs {
|
|
prefill_input_words_list: vec![read_f32_file_as_bf16_words(
|
|
&layer27.step1_prefill_f32_path,
|
|
)
|
|
.unwrap()],
|
|
decode_input_words: read_f32_file_as_bf16_words(&layer27.step1_decode_f32_path)
|
|
.unwrap(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: 1,
|
|
validate_against_oracle: false,
|
|
},
|
|
post_ffn_only_plan(),
|
|
)
|
|
.unwrap();
|
|
|
|
let out_path = temp_dir().join("rust_layer28_step1_full_v_cache_f32.bin");
|
|
write_f32_bits_file(&out_path, &artifacts.full_v_bits);
|
|
println!("rust_full_v_cache_f32_path={}", out_path.display());
|
|
println!(
|
|
"rust_full_v_cache_fnv1a64=0x{:016X}",
|
|
fnv1a64_u32_words(&artifacts.full_v_bits)
|
|
);
|
|
assert_eq!(artifacts.full_v_bits.len(), 4096);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_layer28_step1_writes_rust_attention_scores_file() {
|
|
let outputs = write_formatted_say_hi_outputs_through_layer(27);
|
|
let layer27 = outputs.last().expect("missing layer 27 outputs");
|
|
let model_path = default_model_path();
|
|
let mut session = LayerExecutionSession::load(model_path).unwrap();
|
|
let artifacts = run_layer_plan_with_session_from_sequence(
|
|
&mut session,
|
|
28,
|
|
CachedLayerSequenceInputs {
|
|
prefill_input_words_list: vec![read_f32_file_as_bf16_words(
|
|
&layer27.step1_prefill_f32_path,
|
|
)
|
|
.unwrap()],
|
|
decode_input_words: read_f32_file_as_bf16_words(&layer27.step1_decode_f32_path)
|
|
.unwrap(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: 1,
|
|
validate_against_oracle: false,
|
|
},
|
|
post_ffn_only_plan(),
|
|
)
|
|
.unwrap();
|
|
|
|
let out_path = temp_dir().join("rust_layer28_step1_attention_scores_f32.bin");
|
|
write_f32_bits_file(&out_path, &artifacts.attention_score_bits);
|
|
println!("rust_attention_scores_f32_path={}", out_path.display());
|
|
println!(
|
|
"rust_attention_scores_fnv1a64=0x{:016X}",
|
|
fnv1a64_u32_words(&artifacts.attention_score_bits)
|
|
);
|
|
assert_eq!(artifacts.attention_score_bits.len(), 32);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_layer28_step1_writes_rust_attention_probs_file() {
|
|
let outputs = write_formatted_say_hi_outputs_through_layer(27);
|
|
let layer27 = outputs.last().expect("missing layer 27 outputs");
|
|
let model_path = default_model_path();
|
|
let mut session = LayerExecutionSession::load(model_path).unwrap();
|
|
let artifacts = run_layer_plan_with_session_from_sequence(
|
|
&mut session,
|
|
28,
|
|
CachedLayerSequenceInputs {
|
|
prefill_input_words_list: vec![read_f32_file_as_bf16_words(
|
|
&layer27.step1_prefill_f32_path,
|
|
)
|
|
.unwrap()],
|
|
decode_input_words: read_f32_file_as_bf16_words(&layer27.step1_decode_f32_path)
|
|
.unwrap(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: 1,
|
|
validate_against_oracle: false,
|
|
},
|
|
post_ffn_only_plan(),
|
|
)
|
|
.unwrap();
|
|
|
|
let out_path = temp_dir().join("rust_layer28_step1_attention_probs_f32.bin");
|
|
write_f32_bits_file(&out_path, &artifacts.attention_prob_bits);
|
|
println!("rust_attention_probs_f32_path={}", out_path.display());
|
|
println!(
|
|
"rust_attention_probs_fnv1a64=0x{:016X}",
|
|
fnv1a64_u32_words(&artifacts.attention_prob_bits)
|
|
);
|
|
assert_eq!(artifacts.attention_prob_bits.len(), 32);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_layer28_step1_writes_rust_attention_oproj_file() {
|
|
let outputs = write_formatted_say_hi_outputs_through_layer(27);
|
|
let layer27 = outputs.last().expect("missing layer 27 outputs");
|
|
let model_path = default_model_path();
|
|
let mut session = LayerExecutionSession::load(model_path).unwrap();
|
|
let artifacts = run_layer_plan_with_session_from_sequence(
|
|
&mut session,
|
|
28,
|
|
CachedLayerSequenceInputs {
|
|
prefill_input_words_list: vec![read_f32_file_as_bf16_words(
|
|
&layer27.step1_prefill_f32_path,
|
|
)
|
|
.unwrap()],
|
|
decode_input_words: read_f32_file_as_bf16_words(&layer27.step1_decode_f32_path)
|
|
.unwrap(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: 1,
|
|
validate_against_oracle: false,
|
|
},
|
|
post_ffn_only_plan(),
|
|
)
|
|
.unwrap();
|
|
|
|
let out_bits = artifacts
|
|
.attention_oproj_bits
|
|
.as_ref()
|
|
.expect("missing attention oproj bits");
|
|
let out_path = temp_dir().join("rust_layer28_step1_attention_oproj_f32.bin");
|
|
write_f32_bits_file(&out_path, out_bits);
|
|
println!("rust_attention_oproj_f32_path={}", out_path.display());
|
|
println!(
|
|
"rust_attention_oproj_fnv1a64=0x{:016X}",
|
|
fnv1a64_u32_words(out_bits)
|
|
);
|
|
assert_eq!(out_bits.len(), 2816);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn layer0_to_layer1_hidden_state_handoff_executes() {
|
|
let mut plan = Layer0CachedPlan::new();
|
|
plan.require_stage(Layer0CachedStage::PostFfnResidual);
|
|
let outputs = run_layer_sequence(default_model_path(), &[0, 1], plan).unwrap();
|
|
assert_eq!(outputs.len(), 2);
|
|
assert!(outputs[0].prefill_layer_output_bits().is_some());
|
|
assert!(outputs[1].prefill_layer_output_bits().is_some());
|
|
assert!(outputs[1].layer_output_bits().is_some());
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn all_30_text_layers_execute_from_synthetic_hidden_state_handoff() {
|
|
let mut plan = Layer0CachedPlan::new();
|
|
plan.require_stage(Layer0CachedStage::PostFfnResidual);
|
|
let layer_indices = (0usize..30).collect::<Vec<_>>();
|
|
let outputs = run_layer_sequence(default_model_path(), &layer_indices, plan).unwrap();
|
|
assert_eq!(outputs.len(), 30);
|
|
assert!(outputs[29].prefill_layer_output_bits().is_some());
|
|
assert!(outputs[29].layer_output_bits().is_some());
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn layer0_real_two_token_inputs_report_exact_cached_stage_hashes() {
|
|
let model_path = default_model_path();
|
|
let model_root = super::model_root_dir(&model_path).unwrap();
|
|
let weights = MlxIndexedSafetensors::load(&model_root).unwrap();
|
|
let mut plan = Layer0CachedPlan::new();
|
|
plan.require_stage(Layer0CachedStage::PostFfnResidual);
|
|
let outputs = run_layer_sequence_from_inputs(
|
|
model_path,
|
|
&[0],
|
|
CachedLayerInputs {
|
|
prefill_input_words: weights.embed_token_bf16_words(30_468).unwrap(),
|
|
decode_input_words: weights.embed_token_bf16_words(5_631).unwrap(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: 1,
|
|
validate_against_oracle: false,
|
|
},
|
|
plan,
|
|
)
|
|
.unwrap();
|
|
assert_eq!(outputs.len(), 1);
|
|
print_cached_artifacts(&outputs[0]);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn single_prefill_sequence_path_matches_single_prefill_plan() {
|
|
let model_path = default_model_path();
|
|
let model_root = super::model_root_dir(&model_path).unwrap();
|
|
let weights = MlxIndexedSafetensors::load(&model_root).unwrap();
|
|
let mut plan = Layer0CachedPlan::new();
|
|
plan.require_stage(Layer0CachedStage::PostFfnResidual);
|
|
let single = run_layer_plan(
|
|
model_path.clone(),
|
|
0,
|
|
CachedLayerInputs {
|
|
prefill_input_words: weights.embed_token_bf16_words(30_468).unwrap(),
|
|
decode_input_words: weights.embed_token_bf16_words(5_631).unwrap(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: 1,
|
|
validate_against_oracle: false,
|
|
},
|
|
plan,
|
|
)
|
|
.unwrap();
|
|
let sequence = run_layer_plan_from_sequence(
|
|
model_path,
|
|
0,
|
|
CachedLayerSequenceInputs {
|
|
prefill_input_words_list: vec![weights.embed_token_bf16_words(30_468).unwrap()],
|
|
decode_input_words: weights.embed_token_bf16_words(5_631).unwrap(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: 1,
|
|
validate_against_oracle: false,
|
|
},
|
|
plan,
|
|
)
|
|
.unwrap();
|
|
assert_eq!(sequence.prefill_k_bits, single.prefill_k_bits);
|
|
assert_eq!(sequence.full_k_bits, single.full_k_bits);
|
|
assert_eq!(sequence.full_v_bits, single.full_v_bits);
|
|
assert_eq!(sequence.layer_output_bits(), single.layer_output_bits());
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn layer0_formatted_say_hi_prefix_step_matches_local_mlx_hash() {
|
|
let model_path = default_model_path();
|
|
let model_root = super::model_root_dir(&model_path).unwrap();
|
|
let weights = MlxIndexedSafetensors::load(&model_root).unwrap();
|
|
let mut plan = Layer0CachedPlan::new();
|
|
plan.require_stage(Layer0CachedStage::PostFfnResidual);
|
|
let artifacts = run_layer_plan_from_sequence(
|
|
model_path,
|
|
0,
|
|
CachedLayerSequenceInputs {
|
|
prefill_input_words_list: vec![
|
|
weights.embed_token_bf16_words(2).unwrap(),
|
|
weights.embed_token_bf16_words(105).unwrap(),
|
|
],
|
|
decode_input_words: weights.embed_token_bf16_words(2_364).unwrap(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: 2,
|
|
validate_against_oracle: false,
|
|
},
|
|
plan,
|
|
)
|
|
.unwrap();
|
|
assert_eq!(
|
|
fnv1a64_u32_words(&artifacts.decode_q_bits),
|
|
0x9B1CAF70FB269479
|
|
);
|
|
assert_eq!(
|
|
fnv1a64_u32_words(&artifacts.full_k_bits),
|
|
0xF4944920E989FCF5
|
|
);
|
|
assert_eq!(
|
|
fnv1a64_u32_words(
|
|
artifacts
|
|
.layer_output_bits()
|
|
.expect("missing post-ffn residual bits for formatted prompt step"),
|
|
),
|
|
0xA062311D5B7C20A4
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_first_prefix_step_writes_layer0_outputs_for_layer1_oracle() {
|
|
let model_path = default_model_path();
|
|
let model_root = super::model_root_dir(&model_path).unwrap();
|
|
let weights = MlxIndexedSafetensors::load(&model_root).unwrap();
|
|
let mut plan = Layer0CachedPlan::new();
|
|
plan.require_stage(Layer0CachedStage::PostFfnResidual);
|
|
let artifacts = run_layer_plan(
|
|
model_path,
|
|
0,
|
|
CachedLayerInputs {
|
|
prefill_input_words: weights.embed_token_bf16_words(2).unwrap(),
|
|
decode_input_words: weights.embed_token_bf16_words(105).unwrap(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: 1,
|
|
validate_against_oracle: false,
|
|
},
|
|
plan,
|
|
)
|
|
.unwrap();
|
|
let prefill_words = artifacts
|
|
.prefill_layer_output_bf16_words()
|
|
.expect("missing layer-0 prefill output");
|
|
let decode_words = artifacts
|
|
.bf16_words_for_stage(Layer0CachedStage::PostFfnResidual)
|
|
.expect("missing layer-0 decode output");
|
|
let prefill_path = temp_dir().join("gemma_say_hi_layer0_prefill_2_f32.bin");
|
|
let decode_path = temp_dir().join("gemma_say_hi_layer0_decode_105_f32.bin");
|
|
write_bf16_words_as_f32_file(&prefill_path, &prefill_words).unwrap();
|
|
write_bf16_words_as_f32_file(&decode_path, &decode_words).unwrap();
|
|
println!("prefill_f32_path={}", prefill_path.display());
|
|
println!("decode_f32_path={}", decode_path.display());
|
|
println!(
|
|
"prefill_fnv1a64=0x{:016X}",
|
|
fnv1a64_u32_words(
|
|
artifacts
|
|
.prefill_layer_output_bits()
|
|
.expect("missing layer-0 prefill bits"),
|
|
)
|
|
);
|
|
println!(
|
|
"decode_fnv1a64=0x{:016X}",
|
|
fnv1a64_u32_words(
|
|
artifacts
|
|
.layer_output_bits()
|
|
.expect("missing layer-0 decode bits"),
|
|
)
|
|
);
|
|
assert_eq!(prefill_words.len(), NORM_LEN);
|
|
assert_eq!(decode_words.len(), NORM_LEN);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_second_prefix_step_writes_layer0_decode_output_for_layer1_oracle() {
|
|
let model_path = default_model_path();
|
|
let model_root = super::model_root_dir(&model_path).unwrap();
|
|
let weights = MlxIndexedSafetensors::load(&model_root).unwrap();
|
|
let mut plan = Layer0CachedPlan::new();
|
|
plan.require_stage(Layer0CachedStage::PostFfnResidual);
|
|
let artifacts = run_layer_plan_from_sequence(
|
|
model_path,
|
|
0,
|
|
CachedLayerSequenceInputs {
|
|
prefill_input_words_list: vec![
|
|
weights.embed_token_bf16_words(2).unwrap(),
|
|
weights.embed_token_bf16_words(105).unwrap(),
|
|
],
|
|
decode_input_words: weights.embed_token_bf16_words(2_364).unwrap(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: 2,
|
|
validate_against_oracle: false,
|
|
},
|
|
plan,
|
|
)
|
|
.unwrap();
|
|
let decode_words = artifacts
|
|
.bf16_words_for_stage(Layer0CachedStage::PostFfnResidual)
|
|
.expect("missing layer-0 decode output for token 2364");
|
|
let decode_path = temp_dir().join("gemma_say_hi_layer0_decode_2364_f32.bin");
|
|
write_bf16_words_as_f32_file(&decode_path, &decode_words).unwrap();
|
|
println!("decode_f32_path={}", decode_path.display());
|
|
println!(
|
|
"decode_fnv1a64=0x{:016X}",
|
|
fnv1a64_u32_words(
|
|
artifacts
|
|
.layer_output_bits()
|
|
.expect("missing layer-0 decode bits for token 2364"),
|
|
)
|
|
);
|
|
assert_eq!(decode_words.len(), NORM_LEN);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_prefix_step_writes_outputs_through_text_tower() {
|
|
let outputs = write_formatted_say_hi_outputs_through_layer(29);
|
|
let manifest_path = temp_dir().join("gemma_say_hi_layer_hashes.txt");
|
|
let mut manifest = fs::File::create(&manifest_path).unwrap();
|
|
for entry in &outputs {
|
|
writeln!(
|
|
manifest,
|
|
"layer={} step1_prefill_fnv1a64=0x{:016X} step1_decode105_fnv1a64=0x{:016X} step2_decode2364_fnv1a64=0x{:016X}",
|
|
entry.layer_idx,
|
|
entry.step1_prefill_hash,
|
|
entry.step1_decode_hash,
|
|
entry.step2_decode_hash,
|
|
)
|
|
.unwrap();
|
|
println!(
|
|
"layer={} step1_prefill_fnv1a64=0x{:016X} step1_decode105_fnv1a64=0x{:016X} step2_decode2364_fnv1a64=0x{:016X}",
|
|
entry.layer_idx,
|
|
entry.step1_prefill_hash,
|
|
entry.step1_decode_hash,
|
|
entry.step2_decode_hash,
|
|
);
|
|
println!(
|
|
"layer={} step1_prefill_f32_path={}",
|
|
entry.layer_idx,
|
|
entry.step1_prefill_f32_path.display()
|
|
);
|
|
println!(
|
|
"layer={} step1_decode105_f32_path={}",
|
|
entry.layer_idx,
|
|
entry.step1_decode_f32_path.display()
|
|
);
|
|
println!(
|
|
"layer={} step2_decode2364_f32_path={}",
|
|
entry.layer_idx,
|
|
entry.step2_decode_f32_path.display()
|
|
);
|
|
}
|
|
println!("manifest_path={}", manifest_path.display());
|
|
assert_eq!(outputs.len(), 30);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_full_prompt_teacher_forced_hash_manifest() {
|
|
let model_path = default_model_path();
|
|
let model_root = super::model_root_dir(&model_path).unwrap();
|
|
let weights = MlxIndexedSafetensors::load(&model_root).unwrap();
|
|
let num_layers = weights.snapshot.config.text_config.num_hidden_layers as usize;
|
|
|
|
for layer_idx in 0..num_layers {
|
|
let token_hidden_words = teacher_forced_prompt_hidden_words_through_layer(layer_idx);
|
|
for (position, hidden_words) in token_hidden_words.iter().enumerate() {
|
|
let hidden_bits = hidden_words
|
|
.iter()
|
|
.copied()
|
|
.map(bf16_word_to_f32)
|
|
.map(f32::to_bits)
|
|
.collect::<Vec<_>>();
|
|
println!(
|
|
"token_position={} layer_idx={} hidden_fnv1a64=0x{:016X}",
|
|
position,
|
|
layer_idx,
|
|
fnv1a64_u32_words(&hidden_bits)
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_layer6_token8_writes_layer5_hidden_inputs_for_oracle() {
|
|
write_teacher_forced_hidden_inputs_for_oracle(5, 8);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_layer17_token8_writes_layer16_hidden_inputs_for_oracle() {
|
|
write_teacher_forced_hidden_inputs_for_oracle(16, 8);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_layer29_token7_writes_layer28_hidden_inputs_for_oracle() {
|
|
write_teacher_forced_hidden_inputs_for_oracle(28, 7);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_layer17_token8_writes_rust_attention_output_file() {
|
|
let artifacts = teacher_forced_prompt_step_artifacts(17, 16, 8);
|
|
let out_path = temp_dir().join("rust_layer17_token8_attention_output_f32.bin");
|
|
write_f32_bits_file(&out_path, &artifacts.attention_out_bits);
|
|
println!("rust_attention_output_f32_path={}", out_path.display());
|
|
println!(
|
|
"rust_attention_output_fnv1a64=0x{:016X}",
|
|
fnv1a64_u32_words(&artifacts.attention_out_bits)
|
|
);
|
|
assert_eq!(artifacts.attention_out_bits.len(), 8192);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn teacher_forced_prompt_hidden_inputs_for_oracle_from_env() {
|
|
let input_layer_idx = teacher_env_usize("MAKEPAD_TEACHER_INPUT_LAYER");
|
|
let decode_token_position = teacher_env_usize("MAKEPAD_TEACHER_DECODE_POS");
|
|
write_teacher_forced_hidden_inputs_for_oracle(input_layer_idx, decode_token_position);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn teacher_forced_prompt_attention_output_file_from_env() {
|
|
let layer_idx = teacher_env_usize("MAKEPAD_TEACHER_LAYER");
|
|
let input_layer_idx = teacher_env_usize("MAKEPAD_TEACHER_INPUT_LAYER");
|
|
let decode_token_position = teacher_env_usize("MAKEPAD_TEACHER_DECODE_POS");
|
|
let artifacts =
|
|
teacher_forced_prompt_step_artifacts(layer_idx, input_layer_idx, decode_token_position);
|
|
let out_path = temp_dir().join(format!(
|
|
"rust_layer{layer_idx}_token{decode_token_position}_attention_output_f32.bin"
|
|
));
|
|
write_f32_bits_file(&out_path, &artifacts.attention_out_bits);
|
|
println!("rust_attention_output_f32_path={}", out_path.display());
|
|
println!(
|
|
"rust_attention_output_fnv1a64=0x{:016X}",
|
|
fnv1a64_u32_words(&artifacts.attention_out_bits)
|
|
);
|
|
assert_eq!(artifacts.attention_out_bits.len(), 8192);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_full_prompt_teacher_forced_plan_matches_exact_backend() {
|
|
let prompt_token_ids = formatted_say_hi_prompt_token_ids();
|
|
let model_path = default_model_path();
|
|
let mut session = LayerExecutionSession::load(model_path.clone()).unwrap();
|
|
let weights = session.weights.clone();
|
|
let num_layers = weights.snapshot.config.text_config.num_hidden_layers as usize;
|
|
let mut token_hidden_words = prompt_token_ids
|
|
.iter()
|
|
.map(|token_id| weights.embed_token_bf16_words(*token_id).unwrap())
|
|
.collect::<Vec<_>>();
|
|
|
|
for layer_idx in 0..num_layers {
|
|
let mut next_token_hidden_words = vec![Vec::new(); prompt_token_ids.len()];
|
|
for pos in 0..(prompt_token_ids.len() - 1) {
|
|
let artifacts = run_layer_plan_with_session_from_sequence(
|
|
&mut session,
|
|
layer_idx,
|
|
CachedLayerSequenceInputs {
|
|
prefill_input_words_list: token_hidden_words[..=pos].to_vec(),
|
|
decode_input_words: token_hidden_words[pos + 1].clone(),
|
|
prefill_rope_offset: 0,
|
|
decode_rope_offset: i32::try_from(pos + 1).unwrap(),
|
|
validate_against_oracle: false,
|
|
},
|
|
post_ffn_only_plan(),
|
|
)
|
|
.unwrap();
|
|
|
|
let prefill_output_words = artifacts
|
|
.prefill_layer_output_bf16_words()
|
|
.expect("missing teacher-forced prefill output");
|
|
let decode_output_words = artifacts
|
|
.bf16_words_for_stage(Layer0CachedStage::PostFfnResidual)
|
|
.expect("missing teacher-forced decode output");
|
|
|
|
if pos == 0 {
|
|
next_token_hidden_words[pos] = prefill_output_words;
|
|
} else {
|
|
assert_eq!(
|
|
next_token_hidden_words[pos], prefill_output_words,
|
|
"teacher-forced layer {} token {} prefill output disagreed with prior decode path",
|
|
layer_idx, pos
|
|
);
|
|
}
|
|
next_token_hidden_words[pos + 1] = decode_output_words;
|
|
}
|
|
token_hidden_words = next_token_hidden_words;
|
|
}
|
|
|
|
let explicit_final_hidden_words = token_hidden_words
|
|
.last()
|
|
.expect("missing explicit final hidden")
|
|
.clone();
|
|
let explicit_final_hidden_bits = read_bf16_buffer_bits(
|
|
&session.runtime,
|
|
&session
|
|
.runtime
|
|
.create_buffer_with_bytes(
|
|
&bytes_from_bf16_words(&explicit_final_hidden_words),
|
|
BufferStorageMode::Shared,
|
|
)
|
|
.unwrap(),
|
|
explicit_final_hidden_words.len(),
|
|
)
|
|
.unwrap();
|
|
let explicit_final_norm_words = weights
|
|
.final_text_norm_bf16_words(&explicit_final_hidden_words)
|
|
.unwrap();
|
|
let explicit_next = weights
|
|
.tied_text_logits_top1_f32(&explicit_final_norm_words)
|
|
.unwrap();
|
|
|
|
let mut backend = ExactMetalTextRuntimeSession::load(model_path).unwrap();
|
|
backend.reset_kv_caches().unwrap();
|
|
let mut backend_final_hidden_words = Vec::new();
|
|
for (position, token_id) in prompt_token_ids.iter().copied().enumerate() {
|
|
backend_final_hidden_words = backend
|
|
.eval_token_hidden_state_from_token_id(token_id, position)
|
|
.unwrap();
|
|
}
|
|
let backend_next = backend
|
|
.greedy_token_from_hidden_words(&backend_final_hidden_words)
|
|
.unwrap();
|
|
|
|
println!(
|
|
"explicit_final_hidden_fnv1a64=0x{:016X}",
|
|
fnv1a64_u32_words(&explicit_final_hidden_bits)
|
|
);
|
|
println!(
|
|
"backend_final_hidden_fnv1a64=0x{:016X}",
|
|
fnv1a64_u32_words(
|
|
&backend_final_hidden_words
|
|
.iter()
|
|
.copied()
|
|
.map(|word| (bf16_word_to_f32(word)).to_bits())
|
|
.collect::<Vec<_>>()
|
|
)
|
|
);
|
|
println!(
|
|
"explicit_next_token_id={} backend_next_token_id={}",
|
|
explicit_next.token_id, backend_next.token_id
|
|
);
|
|
|
|
assert_eq!(explicit_final_hidden_words, backend_final_hidden_words);
|
|
assert_eq!(explicit_next.token_id, backend_next.token_id);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn exact_runtime_reuses_cached_layer_workspace() {
|
|
let model_path = default_model_path();
|
|
let model_root = super::model_root_dir(&model_path).unwrap();
|
|
let weights = MlxIndexedSafetensors::load(&model_root).unwrap();
|
|
let mut runtime = ExactMetalTextRuntimeSession::load(model_path).unwrap();
|
|
let token0 = weights.embed_token_bf16_words(30_468).unwrap();
|
|
let token1 = weights.embed_token_bf16_words(5_631).unwrap();
|
|
|
|
runtime.eval_layer_hidden_state(0, &token0, 0).unwrap();
|
|
assert_eq!(runtime.layer_workspaces.len(), 1);
|
|
let workspace_ptr = runtime
|
|
.layer_workspaces
|
|
.get(&0)
|
|
.map(|workspace| workspace as *const _)
|
|
.unwrap();
|
|
|
|
runtime.eval_layer_hidden_state(0, &token1, 1).unwrap();
|
|
assert_eq!(runtime.layer_workspaces.len(), 1);
|
|
let workspace_ptr_after = runtime
|
|
.layer_workspaces
|
|
.get(&0)
|
|
.map(|workspace| workspace as *const _)
|
|
.unwrap();
|
|
assert_eq!(workspace_ptr, workspace_ptr_after);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn formatted_say_hi_token0_layer0_active_stage_hashes() {
|
|
let model_path = default_model_path();
|
|
let mut runtime = ExactMetalTextRuntimeSession::load(model_path).unwrap();
|
|
runtime.reset_kv_caches().unwrap();
|
|
|
|
let workspace = runtime.layer_workspace(0).unwrap();
|
|
let input_buffer = runtime.token_input_buffer().unwrap();
|
|
let output_buffer = workspace.buffers.post_ffn_residual_out.clone();
|
|
runtime
|
|
.dequantize_token_embedding_into_buffer(2, &input_buffer)
|
|
.unwrap();
|
|
runtime
|
|
.eval_layer_hidden_state_core(0, None, Some(&input_buffer), Some(&output_buffer), 0, false)
|
|
.unwrap();
|
|
let workspace = runtime.layer_workspace(0).unwrap();
|
|
let metal = runtime.session.runtime.clone();
|
|
|
|
let print_bits = |stage: &str, bits: &[u32]| {
|
|
println!(
|
|
"stage_hash token_position=0 layer_idx=0 stage={} fnv1a64=0x{:016X}",
|
|
stage,
|
|
fnv1a64_u32_words(bits)
|
|
);
|
|
};
|
|
let print_buffer =
|
|
|stage: &str, buffer: &makepad_ggml::backend::metal::MetalBuffer, len: usize| {
|
|
let bits = read_bf16_buffer_bits(&metal, buffer, len).unwrap();
|
|
print_bits(stage, &bits);
|
|
};
|
|
|
|
let input_bits = read_bf16_buffer_bits(&metal, &input_buffer, NORM_LEN).unwrap();
|
|
print_bits("input", &input_bits);
|
|
print_buffer("input_norm", &workspace.buffers.h, NORM_LEN);
|
|
print_buffer("q", &workspace.buffers.q_rope, workspace.q_proj.out_len());
|
|
print_buffer("k", &workspace.buffers.k_rope, workspace.k_proj.out_len());
|
|
print_buffer("v", &workspace.buffers.v_norm, workspace.k_proj.out_len());
|
|
|
|
{
|
|
let kv_cache = runtime.kv_cache_for_layer(0).unwrap();
|
|
let full_k_bits = read_exact_kv_cache_tensor_bits(
|
|
&metal,
|
|
&kv_cache,
|
|
kv_cache
|
|
.key_buffer()
|
|
.expect("stage hash dump expects bf16 key cache storage"),
|
|
)
|
|
.unwrap();
|
|
let full_v_bits =
|
|
read_exact_kv_cache_tensor_bits(&metal, &kv_cache, &kv_cache.value_buffer).unwrap();
|
|
print_bits("full_k", &full_k_bits);
|
|
print_bits("full_v", &full_v_bits);
|
|
}
|
|
|
|
print_buffer(
|
|
"attention_output",
|
|
&workspace.buffers.attn_out,
|
|
workspace.q_proj.out_len(),
|
|
);
|
|
print_buffer(
|
|
"attention_oproj",
|
|
&workspace.buffers.o_proj_out,
|
|
workspace.o_proj.out_len(),
|
|
);
|
|
print_buffer(
|
|
"post_attention_residual",
|
|
&workspace.buffers.residual_out,
|
|
workspace.post_attention_norm_len,
|
|
);
|
|
print_buffer(
|
|
"pre_feedforward_norm",
|
|
&workspace.buffers.pre_feedforward_norm_out,
|
|
workspace.pre_feedforward_norm_len,
|
|
);
|
|
print_buffer(
|
|
"dense_down",
|
|
&workspace.buffers.mlp_down_out,
|
|
workspace.mlp_down.out_len(),
|
|
);
|
|
print_buffer(
|
|
"post_ffn_norm1",
|
|
&workspace.buffers.post_feedforward_norm1_out,
|
|
workspace.post_feedforward_norm1_len,
|
|
);
|
|
print_buffer(
|
|
"router_scaled",
|
|
&workspace.buffers.router_scaled_out,
|
|
workspace.post_attention_norm_len,
|
|
);
|
|
print_buffer(
|
|
"moe_pre_ffn_norm2",
|
|
&workspace.buffers.pre_feedforward_norm2_out,
|
|
workspace.pre_feedforward_norm2_len,
|
|
);
|
|
print_buffer(
|
|
"expert_scores",
|
|
&workspace.buffers.router_proj_out,
|
|
workspace.router_proj.out_len(),
|
|
);
|
|
print_buffer(
|
|
"router_probs",
|
|
&workspace.buffers.router_probs_out,
|
|
workspace.router_proj.out_len(),
|
|
);
|
|
print_buffer(
|
|
"moe_expert_out",
|
|
&workspace.buffers.moe_weighted_out,
|
|
workspace.post_feedforward_norm2_len,
|
|
);
|
|
print_buffer(
|
|
"moe_post_ffn_norm2",
|
|
&workspace.buffers.moe_post_ffn_norm2_out,
|
|
workspace.post_feedforward_norm2_len,
|
|
);
|
|
print_buffer(
|
|
"moe_merge",
|
|
&workspace.buffers.moe_merge_out,
|
|
workspace.post_feedforward_norm1_len,
|
|
);
|
|
print_buffer(
|
|
"post_feedforward_norm",
|
|
&workspace.buffers.post_feedforward_norm_out,
|
|
workspace.post_feedforward_norm_len,
|
|
);
|
|
print_buffer(
|
|
"layer_output",
|
|
&workspace.buffers.post_ffn_residual_out,
|
|
workspace.post_feedforward_norm1_len,
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn exact_token_embedding_matches_cpu_dequant() {
|
|
let model_path = default_model_path();
|
|
let model_root = super::model_root_dir(&model_path).unwrap();
|
|
let weights = MlxIndexedSafetensors::load(&model_root).unwrap();
|
|
let mut runtime = ExactMetalTextRuntimeSession::load(model_path).unwrap();
|
|
|
|
let token_id = 2u32;
|
|
let expected_words = weights.embed_token_bf16_words(token_id).unwrap();
|
|
let expected_bits = expected_words
|
|
.iter()
|
|
.copied()
|
|
.map(|word| (word as u32) << 16)
|
|
.collect::<Vec<_>>();
|
|
|
|
let input_buffer = runtime.token_input_buffer().unwrap();
|
|
runtime
|
|
.dequantize_token_embedding_into_buffer(token_id, &input_buffer)
|
|
.unwrap();
|
|
let got_bits = read_bf16_buffer_bits(
|
|
&runtime.session.runtime,
|
|
&input_buffer,
|
|
expected_words.len(),
|
|
)
|
|
.unwrap();
|
|
|
|
println!(
|
|
"token_id={} expected_fnv1a64=0x{:016X} got_fnv1a64=0x{:016X}",
|
|
token_id,
|
|
fnv1a64_u32_words(&expected_bits),
|
|
fnv1a64_u32_words(&got_bits)
|
|
);
|
|
if got_bits != expected_bits {
|
|
for idx in 0..16.min(expected_words.len()) {
|
|
let expected_word = expected_words[idx];
|
|
let got_word = (got_bits[idx] >> 16) as u16;
|
|
println!(
|
|
"idx={} expected=0x{:04X} ({:.7}) got=0x{:04X} ({:.7})",
|
|
idx,
|
|
expected_word,
|
|
super::bf16_word_to_f32(expected_word),
|
|
got_word,
|
|
super::bf16_word_to_f32(got_word),
|
|
);
|
|
}
|
|
}
|
|
assert_eq!(got_bits, expected_bits);
|
|
}
|