120 lines
4.8 KiB
Rust
120 lines
4.8 KiB
Rust
use super::array::Dtype;
|
|
use super::{ops, qsa_attention};
|
|
|
|
#[test]
|
|
#[ignore = "Metal and exhaustive pinned gate fixtures; external supervisor required"]
|
|
fn issue109_gate_phase_and_dtype_laws() {
|
|
super::configure_sources().unwrap();
|
|
let _context = crate::engine::metal::gpu::Context::open_qwen(0).unwrap();
|
|
let streams = super::stream::Streams::new(Some(super::allocator::Allocator::new().unwrap()));
|
|
let gpu = streams.default_stream(super::stream::Device::Gpu).unwrap();
|
|
let root = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE109_GATE").unwrap());
|
|
for (name, dtype) in [
|
|
("bf16", Dtype::BF16),
|
|
("f16", Dtype::F16),
|
|
("f32", Dtype::F32),
|
|
] {
|
|
let load = |part: &str| {
|
|
let bytes = std::fs::read(root.join(format!("{name}-{part}.f32"))).unwrap();
|
|
let array = super::array::Array::new(
|
|
&[(bytes.len() / 4) as i32],
|
|
Dtype::F32,
|
|
super::scalar_buffer(&bytes).unwrap(),
|
|
)
|
|
.unwrap();
|
|
ops::astype(&array, dtype, false, gpu).unwrap()
|
|
};
|
|
let gate = load("gate");
|
|
let output = load("output");
|
|
for phase in ["prefill", "ar_decode", "decode_verify"] {
|
|
let actual =
|
|
qsa_attention::attention_gate(&output, &gate, phase == "decode_verify", gpu)
|
|
.unwrap();
|
|
let actual = ops::read_f32_as_f64(
|
|
&ops::astype(&actual, Dtype::F32, false, gpu).unwrap(),
|
|
&streams,
|
|
gpu,
|
|
)
|
|
.unwrap()
|
|
.into_iter()
|
|
.flat_map(|v| (v as f32).to_le_bytes())
|
|
.collect::<Vec<_>>();
|
|
let expected = std::fs::read(root.join(format!("{name}-{phase}.f32"))).unwrap();
|
|
assert_eq!(actual.len(), expected.len());
|
|
let different = actual
|
|
.chunks_exact(4)
|
|
.zip(expected.chunks_exact(4))
|
|
.enumerate()
|
|
.filter_map(|(i, (a, b))| (a != b).then_some(i))
|
|
.collect::<Vec<_>>();
|
|
assert!(
|
|
different.is_empty(),
|
|
"{name}/{phase}: {} different rows; first {:?}",
|
|
different.len(),
|
|
different.first()
|
|
);
|
|
println!(
|
|
"{}",
|
|
serde_json::json!({"case":"gate","dtype":name,"phase":phase,"values":actual.len()/4,"exact":true})
|
|
);
|
|
if dtype == Dtype::BF16 {
|
|
let inject =
|
|
super::hyper::inject_gate(&gate, phase == "decode_verify", gpu).unwrap();
|
|
let actual = ops::read_f32_as_f64(
|
|
&ops::astype(&inject, Dtype::F32, false, gpu).unwrap(),
|
|
&streams,
|
|
gpu,
|
|
)
|
|
.unwrap()
|
|
.into_iter()
|
|
.flat_map(|v| (v as f32).to_le_bytes())
|
|
.collect::<Vec<_>>();
|
|
assert!(
|
|
actual == std::fs::read(root.join(format!("inject-{phase}.f32"))).unwrap(),
|
|
"inject/{phase}"
|
|
);
|
|
println!(
|
|
"{}",
|
|
serde_json::json!({"case":"inject","phase":phase,"values":actual.len()/4,"exact":true})
|
|
);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
thread_local! {
|
|
static RECEIPTS: std::cell::RefCell<Option<(std::path::PathBuf, usize)>> = const { std::cell::RefCell::new(None) };
|
|
}
|
|
|
|
pub(super) fn capture(directory: Option<std::path::PathBuf>) {
|
|
RECEIPTS.with_borrow_mut(|slot| {
|
|
*slot = directory.map(|d| {
|
|
std::fs::create_dir_all(&d).unwrap();
|
|
(d, 0)
|
|
})
|
|
});
|
|
}
|
|
|
|
pub(super) fn record(
|
|
execution: &super::execution::Execution,
|
|
ids: &super::array::Array,
|
|
output: &super::model::Output,
|
|
verify: bool,
|
|
) -> Result<(), String> {
|
|
RECEIPTS.with_borrow_mut(|slot| {
|
|
let Some((directory, index)) = slot else { return Ok(()); };
|
|
let ids = ops::read_i64(ids, &execution.streams, execution.stream)?;
|
|
let mut record=serde_json::json!({"index":*index,"ids":ids,"verify":verify});
|
|
for (part, a) in [("logits",&output.logits),("hidden",&output.hidden)] {
|
|
let a=a.as_ref().ok_or("missing #109 output")?;
|
|
let bytes=execution.values_f32(a)?.into_iter().flat_map(f32::to_le_bytes).collect::<Vec<_>>();
|
|
std::fs::write(directory.join(format!("{index}-{part}.f32")),&bytes).map_err(|e|e.to_string())?;
|
|
record[part]=serde_json::json!({"shape":a.layout().shape(),"sha256":crate::engine::metal::qwen_mtplx_tests::hex_digest(&bytes)});
|
|
}
|
|
std::fs::write(directory.join(format!("{index}.json")),serde_json::to_vec(&record).unwrap()).map_err(|e|e.to_string())?;
|
|
println!("{}",serde_json::json!({"case":"verify_forward","record":record}));
|
|
*index+=1;
|
|
Ok(())
|
|
})
|
|
}
|