hugo/DS4Server
1
0
Fork 0
Code Issues 2 Pull Requests Actions Packages Projects Releases Wiki Activity

Stream Qwen PLE embeddings

Browse Source
This commit is contained in:
Georg Bauer
2026-09-03 21:12:20 +02:00
parent c414640050
commit 87ccf67d0c
4 changed files with 893 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files

711
src/engine/metal/qwen.rs
View File

@@ -24,8 +24,14 @@ const EXPERTS_USED: usize = 10;
const EXPERT_WIDTH: u32 = 640;
const VOCAB: u32 = 248_320;
const DENSE_BUDGET: u32 = 2_048;
const PLE_HEADS: usize = 16;
const PLE_HEAD_DIM: u32 = 160;
const PLE_HISTORY: usize = 2;
const PLE_CONV_STATE: u32 = 9;
const EOS_TOKEN: i32 = 248_044;
const PLE_ROW_BYTES: usize = 100;
const CHECKPOINT_MAGIC: &[u8; 8] = b"DS4QWN01";
const CHECKPOINT_VERSION: u32 = 1;
const CHECKPOINT_VERSION: u32 = 2;
const CHECKPOINT_CHUNK: usize = 8 * 1024 * 1024;
#[derive(Clone, Copy)]
@@ -75,6 +81,15 @@ struct Scratch {
k_rope: Buffer,
v: Buffer,
attention: Buffer,
ple_packed: Buffer,
ple_scales: Buffer,
ple_biases: Buffer,
ple_embedding: Buffer,
ple_key: Buffer,
ple_value: Buffer,
ple_gated: Buffer,
ple_norm: Buffer,
ple_output: Buffer,
logits: Buffer,
}
@@ -108,14 +123,36 @@ impl Scratch {
k_rope: Buffer::floats(ATTN_KV_WIDTH.into())?,
v: Buffer::floats(ATTN_KV_WIDTH.into())?,
attention: Buffer::floats(ATTN_WIDTH.into())?,
ple_packed: Buffer::bytes((PLE_HEADS as u64) * 80)?,
ple_scales: Buffer::bytes((PLE_HEADS as u64) * 10)?,
ple_biases: Buffer::bytes((PLE_HEADS as u64) * 10)?,
ple_embedding: Buffer::floats(HIDDEN.into())?,
ple_key: Buffer::floats(HC_WIDTH.into())?,
ple_value: Buffer::floats(HIDDEN.into())?,
ple_gated: Buffer::floats(HC_WIDTH.into())?,
ple_norm: Buffer::floats(HC_WIDTH.into())?,
ple_output: Buffer::floats(HC_WIDTH.into())?,
logits: Buffer::floats(VOCAB.into())?,
})
}
}
struct PleContract {
multipliers: [i64; 3],
sizes: [i64; PLE_HEADS],
offsets: [i64; PLE_HEADS],
}
struct PleState {
history: [i32; PLE_HISTORY],
conv: Buffer,
}
pub(in crate::engine) struct QwenExecutor {
model: QwenModel,
states: Vec<LayerState>,
ple_contract: PleContract,
ple_state: PleState,
scratch: Scratch,
logits: Vec<f32>,
tokens: Vec<i32>,
@@ -127,6 +164,7 @@ pub(in crate::engine) struct QwenExecutor {
pub(in crate::engine) struct QwenResidentState {
states: Vec<LayerState>,
ple_state: PleState,
logits: Vec<f32>,
tokens: Vec<i32>,
position: u32,
@@ -135,11 +173,14 @@ pub(in crate::engine) struct QwenResidentState {
impl QwenExecutor {
pub(super) fn open(model: QwenModel, context: u32) -> Result<Self, String> {
let ple_contract = ple_contract(&model)?;
let native = Context::open_qwen(model.memory().admission)?;
let states = allocate_states(context)?;
Ok(Self {
model,
states,
ple_contract,
ple_state: allocate_ple_state()?,
scratch: Scratch::new()?,
logits: vec![0.0; VOCAB as usize],
tokens: Vec::new(),
@@ -160,11 +201,19 @@ impl QwenExecutor {
self.context
));
}
self.require_ple()?;
if self.position + 1 > DENSE_BUDGET {
return Err(
"Qwen sparse QSA selection is required beyond 2048 tokens; native QSA belongs to issue #97"
.into(),
);
}
self.begin_token(token)?;
for layer in 0..LAYERS {
if layer == 1 {
self.ple(token)?;
}
self.encode_layer(layer)?;
}
self.final_output()?;
@@ -211,11 +260,143 @@ impl QwenExecutor {
commands.finish()
}
fn require_ple(&self) -> Result<(), String> {
Err(
"Qwen PLE injection is required at layer 2; native mapped PLE lookup belongs to issue #96"
.into(),
)
fn ple(&mut self, token: i32) -> Result<(), String> {
let rows = ple_rows(&self.ple_contract, self.ple_state.history, token)?;
let packed = self.model.tensor("ngram.weight")?;
let scales = self.model.tensor("ngram.scales")?;
let biases = self.model.tensor("ngram.biases")?;
let mut packed_stage = [0_u8; PLE_HEADS * 80];
let mut scales_stage = [0_u8; PLE_HEADS * 10];
let mut biases_stage = [0_u8; PLE_HEADS * 10];
let packed_bytes = self.model.tensor_bytes(packed)?;
let scales_bytes = self.model.tensor_bytes(scales)?;
let biases_bytes = self.model.tensor_bytes(biases)?;
for (head, &row) in rows.iter().enumerate() {
let bytes = gather_ple_row(row, packed_bytes, scales_bytes, biases_bytes)?;
packed_stage[head * 80..][..80].copy_from_slice(&bytes[..80]);
scales_stage[head * 10..][..10].copy_from_slice(&bytes[80..90]);
biases_stage[head * 10..][..10].copy_from_slice(&bytes[90..]);
}
self.scratch.ple_packed.write(0, &packed_stage)?;
self.scratch.ple_scales.write(0, &scales_stage)?;
self.scratch.ple_biases.write(0, &biases_stage)?;
let commands = Commands::begin()?;
let mut dequant = args();
dequant.u[0] = PLE_HEAD_DIM;
dequant.u[1] = PLE_HEADS as u32;
dequant.u[2] = 4;
dequant.u[3] = 32;
self.dispatch(
c"kernel_qwen_ple_dequant",
&self.scratch.ple_embedding,
Some(&self.scratch.ple_packed),
Some(&self.scratch.ple_scales),
Some(&self.scratch.ple_biases),
&[],
&dequant,
HIDDEN,
1,
)?;
let prefix = "language_model.model.layers.1.ple";
self.bf16_mv(
self.weight(&format!("{prefix}.key_proj.weight"))?,
&self.scratch.ple_embedding,
&self.scratch.ple_key,
HIDDEN,
HC_WIDTH,
)?;
self.bf16_mv(
self.weight(&format!("{prefix}.value_proj.weight"))?,
&self.scratch.ple_embedding,
&self.scratch.ple_value,
HIDDEN,
HIDDEN,
)?;
for (input, output, name) in [
(&self.scratch.ple_key, &self.scratch.ple_key, "norm_key"),
(&self.scratch.hc, &self.scratch.hc_norm, "norm_query"),
] {
let mut norm = args();
norm.u[0] = HC_WIDTH;
norm.u[1] = HIDDEN;
norm.f[0] = 1.0e-6;
self.dispatch(
c"kernel_qwen_zero_rms",
output,
Some(input),
None,
None,
&[self.view(self.weight(&format!("{prefix}.{name}.weight"))?)],
&norm,
HC,
1,
)?;
}
let mut gate = args();
gate.u[0] = HIDDEN;
self.dispatch(
c"kernel_qwen_ple_gate",
&self.scratch.ple_gated,
Some(&self.scratch.ple_key),
Some(&self.scratch.hc_norm),
Some(&self.scratch.ple_value),
&[],
&gate,
HC,
1,
)?;
let mut norm = args();
norm.u[0] = HC_WIDTH;
norm.u[1] = HIDDEN;
norm.f[0] = 1.0e-6;
self.dispatch(
c"kernel_qwen_zero_rms",
&self.scratch.ple_norm,
Some(&self.scratch.ple_gated),
None,
None,
&[self.view(self.weight(&format!("{prefix}.norm_conv.weight"))?)],
&norm,
HC,
1,
)?;
let mut conv = args();
conv.u[0] = HC_WIDTH;
self.dispatch(
c"kernel_qwen_ple_conv",
&self.scratch.ple_output,
Some(&self.scratch.ple_gated),
Some(&self.scratch.ple_norm),
Some(&self.ple_state.conv),
&[self.view(self.weight(&format!("{prefix}.conv_weight"))?)],
&conv,
HC_WIDTH,
1,
)?;
let mut add = args();
add.u[0] = HC_WIDTH;
self.dispatch(
c"kernel_qwen_add",
&self.scratch.hc_norm,
Some(&self.scratch.hc),
Some(&self.scratch.ple_output),
None,
&[],
&add,
HC_WIDTH,
1,
)?;
self.scratch.hc.copy_from(
0,
&self.scratch.hc_norm,
0,
u64::from(HC_WIDTH) * 4,
"committing Qwen PLE injection",
)?;
commands.finish()?;
self.ple_state.history = advance_ple_history(self.ple_state.history, token);
Ok(())
}
fn encode_layer(&mut self, layer: usize) -> Result<(), String> {
@@ -258,7 +439,15 @@ impl QwenExecutor {
let commands = Commands::begin()?;
for (&expert, &weight) in ids.iter().zip(&weights) {
if !(0..EXPERTS as i32).contains(&expert) || !weight.is_finite() || weight < 0.0 {
return Err("Qwen router produced an invalid top-10 selection".into());
let mut router = vec![0.0; EXPERTS as usize];
self.scratch.router.read_f32(&mut router)?;
let non_finite = router.iter().filter(|value| !value.is_finite()).count();
let mut block = vec![0.0; HIDDEN as usize];
self.scratch.block.read_f32(&mut block)?;
let block_non_finite = block.iter().filter(|value| !value.is_finite()).count();
return Err(format!(
"Qwen layer {layer} router produced invalid expert {expert} with weight {weight} ({non_finite} non-finite logits, {block_non_finite} non-finite inputs)"
));
}
self.expert(&prefix, expert as u32, weight)?;
}
@@ -1071,6 +1260,7 @@ impl QwenExecutor {
pub(super) fn reset(&mut self) -> Result<(), String> {
self.states = allocate_states(self.context)?;
self.ple_state = allocate_ple_state()?;
self.logits.fill(0.0);
self.tokens.clear();
self.position = 0;
@@ -1088,6 +1278,7 @@ impl QwenExecutor {
fn blank_resident(&self) -> Result<QwenResidentState, String> {
Ok(QwenResidentState {
states: allocate_states(self.context)?,
ple_state: allocate_ple_state()?,
logits: vec![0.0; VOCAB as usize],
tokens: Vec::new(),
position: 0,
@@ -1101,6 +1292,7 @@ impl QwenExecutor {
) -> Result<(), String> {
let mut incoming = state.take().map_or_else(|| self.blank_resident(), Ok)?;
std::mem::swap(&mut self.states, &mut incoming.states);
std::mem::swap(&mut self.ple_state, &mut incoming.ple_state);
std::mem::swap(&mut self.logits, &mut incoming.logits);
std::mem::swap(&mut self.tokens, &mut incoming.tokens);
std::mem::swap(&mut self.position, &mut incoming.position);
@@ -1134,6 +1326,9 @@ impl QwenExecutor {
file.write_all(&self.model.checkpoint_identity())
.map_err(|error| error.to_string())?;
file.write_all(&tag).map_err(|error| error.to_string())?;
for token in self.ple_state.history {
write_u32(&mut file, token as u32)?;
}
for &token in &self.tokens {
write_u32(&mut file, token as u32)?;
}
@@ -1141,6 +1336,14 @@ impl QwenExecutor {
write_u32(&mut file, logit.to_bits())?;
}
let mut chunk = vec![0; CHECKPOINT_CHUNK];
write_buffer(
&mut file,
&self.ple_state.conv,
0,
u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 * 2,
&mut chunk,
progress,
)?;
for state in &self.states {
match state {
LayerState::Gdn { conv, recurrent } => {
@@ -1216,6 +1419,14 @@ impl QwenExecutor {
let mut tag = [0; 32];
file.read_exact(&mut tag)
.map_err(|error| error.to_string())?;
let mut ple_history = [0; PLE_HISTORY];
for token in &mut ple_history {
let value = read_u32(&mut file)?;
if value >= VOCAB {
return Err("Qwen checkpoint PLE history is invalid".into());
}
*token = value as i32;
}
let mut tokens = Vec::with_capacity(position as usize);
for _ in 0..position {
let token = read_u32(&mut file)?;
@@ -1230,6 +1441,14 @@ impl QwenExecutor {
}
self.reset()?;
let mut chunk = vec![0; CHECKPOINT_CHUNK];
read_buffer(
&mut file,
&self.ple_state.conv,
0,
u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 * 2,
&mut chunk,
progress,
)?;
for state in &self.states {
match state {
LayerState::Gdn { conv, recurrent } => {
@@ -1272,6 +1491,7 @@ impl QwenExecutor {
self.position = position;
self.tokens = tokens;
self.logits = logits;
self.ple_state.history = ple_history;
self.checkpoint_tag = tag;
Ok(true)
}
@@ -1299,6 +1519,150 @@ fn allocate_states(context: u32) -> Result<Vec<LayerState>, String> {
.collect()
}
fn allocate_ple_state() -> Result<PleState, String> {
let conv = Buffer::bytes(u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 * 2)?;
conv.fill(0.0, u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 / 2)?;
Ok(PleState {
history: [EOS_TOKEN; PLE_HISTORY],
conv,
})
}
fn ple_contract(model: &QwenModel) -> Result<PleContract, String> {
let multipliers = read_i64_array::<3>(
model,
"language_model.model.layers.1.ple.ple_embedding.layer_multipliers",
)?;
let sizes = read_i64_array::<PLE_HEADS>(
model,
"language_model.model.layers.1.ple.ple_embedding.ngram_heads_vocab_sizes",
)?;
let offsets = read_i64_array::<PLE_HEADS>(
model,
"language_model.model.layers.1.ple.ple_embedding.ngram_heads_offsets",
)?;
let expected_multipliers = official_ple_multipliers();
let mut expected_sizes = [0; PLE_HEADS];
let mut expected_offsets = [0; PLE_HEADS];
let mut total = 0_i64;
let mut prime = 19_999_999_i64;
for head in 0..PLE_HEADS {
prime = next_prime(prime);
expected_sizes[head] = prime;
expected_offsets[head] = total;
total += prime;
}
if multipliers != expected_multipliers || sizes != expected_sizes || offsets != expected_offsets
{
return Err("Qwen PLE hash parameters do not match the official contract".into());
}
for (name, shape, bits, group) in [
("ngram.weight", [320_001_536, 20], Some(4), Some(32)),
("ngram.scales", [320_001_536, 5], Some(4), Some(32)),
("ngram.biases", [320_001_536, 5], Some(4), Some(32)),
] {
let tensor = model.tensor(name)?;
if tensor.shape != shape || tensor.quant_bits != bits || tensor.group_size != group {
return Err(format!("{name} does not match the Qwen PLE row layout"));
}
}
Ok(PleContract {
multipliers,
sizes,
offsets,
})
}
fn read_i64_array<const N: usize>(model: &QwenModel, name: &str) -> Result<[i64; N], String> {
let tensor = model.tensor(name)?;
if tensor.dtype != "I64" || tensor.shape != [N as u64] {
return Err(format!("{name} does not match the Qwen PLE integer layout"));
}
let bytes = model.tensor_bytes(tensor)?;
if bytes.len() != N * 8 {
return Err(format!("{name} has an invalid byte length"));
}
Ok(std::array::from_fn(|index| {
i64::from_le_bytes(bytes[index * 8..index * 8 + 8].try_into().unwrap())
}))
}
fn official_ple_multipliers() -> [i64; 3] {
const GAMMA: u64 = 0x9e37_79b9_7f4a_7c15;
const M1: u64 = 0xbf58_476d_1ce4_e5b9;
const M2: u64 = 0x94d0_49bb_1331_11eb;
let bound = (i64::MAX / VOCAB as i64 / 2) as u64;
std::array::from_fn(|index| {
let mut value = 1234_u64.wrapping_add(GAMMA.wrapping_mul(index as u64 + 1));
value = value.wrapping_add(GAMMA);
value = (value ^ (value >> 30)).wrapping_mul(M1);
value = (value ^ (value >> 27)).wrapping_mul(M2);
value ^= value >> 31;
(2 * (value % bound) + 1) as i64
})
}
fn next_prime(mut value: i64) -> i64 {
loop {
value += 1;
if value % 2 != 0
&& (3..=((value as f64).sqrt() as i64))
.step_by(2)
.all(|divisor| value % divisor != 0)
{
return value;
}
}
}
fn ple_rows(
contract: &PleContract,
history: [i32; PLE_HISTORY],
token: i32,
) -> Result<[u64; PLE_HEADS], String> {
if token < 0 || token as u32 >= VOCAB {
return Err(format!("token {token} is outside the Qwen vocabulary"));
}
let shifted = [token as i64, history[1] as i64, history[0] as i64];
let two = shifted[0].wrapping_mul(contract.multipliers[0])
^ shifted[1].wrapping_mul(contract.multipliers[1]);
let three = two ^ shifted[2].wrapping_mul(contract.multipliers[2]);
Ok(std::array::from_fn(|head| {
let mixed = if head < 8 { two } else { three };
(contract.offsets[head] + mixed.rem_euclid(contract.sizes[head])) as u64
}))
}
fn advance_ple_history(history: [i32; PLE_HISTORY], token: i32) -> [i32; PLE_HISTORY] {
if token == EOS_TOKEN {
[EOS_TOKEN; PLE_HISTORY]
} else {
[history[1], token]
}
}
fn gather_ple_row(
row: u64,
packed: &[u8],
scales: &[u8],
biases: &[u8],
) -> Result<[u8; PLE_ROW_BYTES], String> {
let row = usize::try_from(row).map_err(|_| "Qwen PLE row exceeds this platform")?;
let mut value = [0; PLE_ROW_BYTES];
value[..80].copy_from_slice(ple_row_slice(packed, row, 80)?);
value[80..90].copy_from_slice(ple_row_slice(scales, row, 10)?);
value[90..].copy_from_slice(ple_row_slice(biases, row, 10)?);
Ok(value)
}
fn ple_row_slice(data: &[u8], row: usize, width: usize) -> Result<&[u8], String> {
let start = row
.checked_mul(width)
.ok_or_else(|| "Qwen PLE row offset overflows".to_owned())?;
data.get(start..start + width)
.ok_or_else(|| format!("Qwen PLE row {row} is truncated"))
}
fn args() -> QwenKernelArgs {
QwenKernelArgs::default()
}
@@ -1338,6 +1702,7 @@ fn dispatch_qwen(
mod tests {
use super::*;
use memmap2::MmapOptions;
use sha2::{Digest, Sha256};
use std::fs;
use std::path::PathBuf;
@@ -1353,6 +1718,145 @@ mod tests {
);
}
#[test]
fn qwen_ple_hash_contract_matches_golden_boundaries() {
assert_eq!(
official_ple_multipliers(),
[23_703_573_157_769, 20_109_073_645_365, 8_052_911_324_071]
);
let mut contract = PleContract {
multipliers: official_ple_multipliers(),
sizes: [0; PLE_HEADS],
offsets: [0; PLE_HEADS],
};
let mut prime = 19_999_999;
let mut offset = 0;
for head in 0..PLE_HEADS {
prime = next_prime(prime);
contract.sizes[head] = prime;
contract.offsets[head] = offset;
offset += prime;
}
let initial = ple_rows(&contract, [EOS_TOKEN; 2], 1).unwrap();
let repeated = ple_rows(&contract, [1, 1], 1).unwrap();
let boundary = ple_rows(&contract, [EOS_TOKEN, 42], 43).unwrap();
assert_eq!(
initial,
[
16_121_432,
28_938_500,
59_087_997,
73_487_090,
81_148_277,
104_500_129,
120_276_032,
149_373_875,
176_283_436,
184_305_849,
216_528_839,
231_080_079,
257_961_536,
266_068_568,
289_043_455,
305_959_965,
]
);
assert_eq!(
repeated,
[
6_868_091,
38_325_817,
54_054_700,
68_075_137,
82_949_816,
101_241_419,
138_678_867,
155_262_032,
176_541_251,
196_154_476,
215_703_237,
234_413_824,
254_220_543,
274_027_268,
293_962_951,
313_640_732,
]
);
assert_eq!(
boundary,
[
18_529_343,
23_547_650,
56_056_978,
73_570_159,
88_581_601,
113_585_506,
121_091_299,
151_099_148,
175_585_266,
184_439_538,
216_587_431,
222_082_137,
250_284_866,
278_847_169,
281_781_121,
317_050_322,
]
);
assert_eq!(contract.offsets[0], 0);
assert_eq!(contract.offsets[8], 160_000_374);
assert_eq!(contract.offsets[15] + contract.sizes[15], 320_001_446);
for (head, &row) in initial.iter().enumerate() {
assert!(
(contract.offsets[head]..contract.offsets[head] + contract.sizes[head])
.contains(&(row as i64))
);
}
assert_eq!(320_001_536 % 128, 0);
let history = advance_ple_history([EOS_TOKEN; 2], 1);
assert_eq!(advance_ple_history(history, 2), [1, 2]);
assert_eq!(advance_ple_history([1, 2], EOS_TOKEN), [EOS_TOKEN; 2]);
let hash_chunks = |chunk: usize| {
let mut history = [EOS_TOKEN; PLE_HISTORY];
[1, 2, EOS_TOKEN, 3, 4]
.chunks(chunk)
.flat_map(|tokens| {
tokens
.iter()
.map(|&token| {
let rows = ple_rows(&contract, history, token).unwrap();
history = advance_ple_history(history, token);
rows
})
.collect::<Vec<_>>()
})
.collect::<Vec<_>>()
};
assert_eq!(hash_chunks(1), hash_chunks(2));
assert_eq!(hash_chunks(1), hash_chunks(5));
}
#[test]
fn qwen_ple_gather_copies_only_the_requested_row() {
let packed = (0..240).map(|value| value as u8).collect::<Vec<_>>();
let scales = (0..30).map(|value| (value + 17) as u8).collect::<Vec<_>>();
let biases = (0..30).map(|value| (value + 47) as u8).collect::<Vec<_>>();
let direct = gather_ple_row(1, &packed, &scales, &biases).unwrap();
assert_eq!(&direct[..80], &packed[80..160]);
assert_eq!(&direct[80..90], &scales[10..20]);
assert_eq!(&direct[90..], &biases[10..20]);
assert_eq!(
gather_ple_row(1, &packed, &scales, &biases).unwrap(),
direct
);
assert!(
gather_ple_row(3, &packed, &scales, &biases)
.unwrap_err()
.contains("truncated")
);
}
#[test]
#[ignore = "requires Apple Metal"]
fn qwen_metal_primitives_match_reference_vectors() {
@@ -1363,11 +1867,12 @@ mod tests {
input.write_f32(&vec![1.0; 64]).unwrap();
let output = Buffer::floats(64).unwrap();
let path = std::env::temp_dir().join(format!(
"ds4-qwen95-weights-{}-{}",
"ds4-qwen96-weights-{}-{}",
std::process::id(),
std::thread::current().name().unwrap_or("test")
));
let mut bytes = Vec::new();
// Safetensors headers are not guaranteed to align the data region.
let mut bytes = vec![0];
for _ in 0..8 {
bytes.extend_from_slice(&0x3333_3333_u32.to_le_bytes());
}
@@ -1398,7 +1903,7 @@ mod tests {
Some(&input),
None,
None,
&[view(0, 32), view(32, 2), view(34, 2)],
&[view(1, 32), view(33, 2), view(35, 2)],
&affine,
1,
1,
@@ -1421,7 +1926,7 @@ mod tests {
Some(&conv_input),
Some(&conv_state),
None,
&[view(40, 8)],
&[view(41, 8)],
&conv,
1,
1,
@@ -1459,7 +1964,7 @@ mod tests {
Some(&qkv),
Some(&controls),
Some(&recurrent),
&[view(36, 2), view(38, 2)],
&[view(37, 2), view(39, 2)],
&step,
2,
1,
@@ -1479,7 +1984,7 @@ mod tests {
Some(&raw),
Some(&controls),
None,
&[view(48, 4)],
&[view(49, 4)],
&norm,
1,
1,
@@ -1533,7 +2038,7 @@ mod tests {
Some(&norm_input),
None,
None,
&[view(52, 8)],
&[view(53, 8)],
&zero_norm,
2,
1,
@@ -1670,7 +2175,7 @@ mod tests {
Some(&rope_input),
None,
None,
&[view(52, 8)],
&[view(53, 8)],
&rope,
4,
1,
@@ -1749,6 +2254,94 @@ mod tests {
probability * 3.0 + (1.0 - probability) * 7.0,
);
let ple_packed = Buffer::bytes(80).unwrap();
ple_packed.write(0, &[0x33; 80]).unwrap();
let ple_scales = Buffer::bytes(10).unwrap();
ple_scales
.write(0, &bf16(0.5).to_le_bytes().repeat(5))
.unwrap();
let ple_biases = Buffer::bytes(10).unwrap();
ple_biases
.write(0, &bf16(-1.0).to_le_bytes().repeat(5))
.unwrap();
let ple_embedding = Buffer::floats(160).unwrap();
let mut dequant = args();
dequant.u[0] = 160;
dequant.u[1] = 1;
dequant.u[2] = 4;
dequant.u[3] = 32;
dispatch_qwen(
c"kernel_qwen_ple_dequant",
&ple_embedding,
Some(&ple_packed),
Some(&ple_scales),
Some(&ple_biases),
&[],
&dequant,
160,
1,
)
.unwrap();
let mut embedding = [0.0; 160];
ple_embedding.read_f32(&mut embedding).unwrap();
assert!(embedding.into_iter().all(|value| value == 0.5));
let ple_key = Buffer::floats(8).unwrap();
ple_key.write_f32(&[1.0; 8]).unwrap();
let ple_query = Buffer::floats(8).unwrap();
ple_query.write_f32(&[1.0; 8]).unwrap();
let ple_value = Buffer::floats(2).unwrap();
ple_value.write_f32(&[2.0, 3.0]).unwrap();
let ple_gated = Buffer::floats(8).unwrap();
let mut gate = args();
gate.u[0] = 2;
dispatch_qwen(
c"kernel_qwen_ple_gate",
&ple_gated,
Some(&ple_key),
Some(&ple_query),
Some(&ple_value),
&[],
&gate,
4,
1,
)
.unwrap();
let expected_gate = 1.0 / (1.0 + (-2.0_f32.sqrt().sqrt()).exp());
let mut gated = [0.0; 8];
ple_gated.read_f32(&mut gated).unwrap();
for stream in 0..4 {
close(gated[stream * 2], 2.0 * expected_gate);
close(gated[stream * 2 + 1], 3.0 * expected_gate);
}
let ple_normalized = Buffer::floats(1).unwrap();
ple_normalized.write_f32(&[2.0]).unwrap();
let ple_gate_value = Buffer::floats(1).unwrap();
ple_gate_value.write_f32(&[0.5]).unwrap();
let ple_state = Buffer::bytes(18).unwrap();
ple_state.write(0, &[0; 18]).unwrap();
let ple_output = Buffer::floats(1).unwrap();
let mut conv = args();
conv.u[0] = 1;
dispatch_qwen(
c"kernel_qwen_ple_conv",
&ple_output,
Some(&ple_gate_value),
Some(&ple_normalized),
Some(&ple_state),
&[view(41, 8)],
&conv,
1,
1,
)
.unwrap();
ple_output.read_f32(&mut scalar).unwrap();
close(scalar[0], 0.5 + 8.0 / (1.0 + (-8.0_f32).exp()));
let mut ple_history = [0; 18];
ple_state.read(0, &mut ple_history).unwrap();
assert_eq!(&ple_history[16..], &bf16(2.0).to_le_bytes());
drop(map);
drop(file);
fs::remove_file(path).unwrap();
@@ -1762,35 +2355,53 @@ mod tests {
.map(PathBuf::from)
.expect("set DS4SERVER_QWEN38_SOURCE to the pinned artifact directory");
let model = QwenModel::open(&root, 4).unwrap();
let residency_before = model.mapped_residency().unwrap();
let mut executor = QwenExecutor::open(model, 4).unwrap();
let direct = executor.eval(1).unwrap_err();
let batch = executor.prefill(&[1], |_| true).unwrap_err();
assert_eq!(direct, batch);
assert_eq!(executor.position, 0);
assert!(executor.tokens.is_empty());
executor.begin_token(1).unwrap();
executor.encode_layer(0).unwrap();
executor.final_output().unwrap();
executor.eval(1).unwrap();
let residency_after = executor.model().mapped_residency().unwrap();
assert!(residency_after.0 <= executor.model().memory().resident_core);
assert!(residency_after.1 <= executor.model().memory().mapped_ple);
eprintln!(
"Qwen mapped residency core/PLE before {:?}, after {:?}",
residency_before, residency_after
);
assert!(executor.logits.iter().all(|value| value.is_finite()));
let mut digest = Sha256::new();
for value in &executor.logits {
digest.update(value.to_bits().to_le_bytes());
}
let digest: [u8; 32] = digest.finalize().into();
assert_eq!(
digest,
[
137, 244, 133, 253, 201, 214, 196, 144, 249, 130, 28, 63, 124, 75, 32, 40, 16, 148,
148, 123, 5, 50, 90, 165, 101, 44, 223, 164, 62, 54, 164, 86,
]
);
let reference = [
executor.logits[0],
executor.logits[1],
executor.logits[1000],
executor.logits[VOCAB as usize - 1],
];
let next = executor
.logits
.iter()
.enumerate()
.max_by(|a, b| a.1.total_cmp(b.1))
.unwrap()
.0 as i32;
for (actual, expected) in
reference
.into_iter()
.zip([1.483_976_1, 0.575_980_66, -0.008_828_48, 0.047_039_207])
.zip([6.406_557, 2.082_818_3, -3.058_045_6, -0.121_010_3])
{
close(actual, expected);
}
assert_eq!(next, 89_648);
executor.position = 1;
executor.tokens = vec![1];
let checkpoint = std::env::temp_dir().join(format!(
"ds4-qwen95-checkpoint-{}-{}",
"ds4-qwen96-checkpoint-{}-{}",
std::process::id(),
std::thread::current().name().unwrap_or("test")
));
@@ -1798,9 +2409,7 @@ mod tests {
.save_checkpoint(&checkpoint, [7; 32], &mut |_| {})
.unwrap();
executor.begin_token(2).unwrap();
executor.encode_layer(0).unwrap();
executor.final_output().unwrap();
executor.eval(next).unwrap();
let continued = [
executor.logits[0],
executor.logits[1],
@@ -1814,10 +2423,13 @@ mod tests {
conv.read(0, &mut continued_conv).unwrap();
let mut continued_recurrent = vec![0; 4 * 1024];
recurrent.read(0, &mut continued_recurrent).unwrap();
let mut continued_ple = vec![0; HC_WIDTH as usize * PLE_CONV_STATE as usize * 2];
executor.ple_state.conv.read(0, &mut continued_ple).unwrap();
assert!(executor.load_checkpoint(&checkpoint, &mut |_| {}).unwrap());
assert_eq!(executor.position, 1);
assert_eq!(executor.tokens, [1]);
assert_eq!(executor.ple_state.history, [EOS_TOKEN, 1]);
assert_eq!(executor.checkpoint_tag, [7; 32]);
assert_eq!(
[
@@ -1828,9 +2440,7 @@ mod tests {
],
reference
);
executor.begin_token(2).unwrap();
executor.encode_layer(0).unwrap();
executor.final_output().unwrap();
executor.eval(next).unwrap();
for (actual, expected) in [
executor.logits[0],
executor.logits[1],
@@ -1851,11 +2461,36 @@ mod tests {
let mut resumed_recurrent = vec![0; continued_recurrent.len()];
recurrent.read(0, &mut resumed_recurrent).unwrap();
assert_eq!(resumed_recurrent, continued_recurrent);
let mut resumed_ple = vec![0; continued_ple.len()];
executor.ple_state.conv.read(0, &mut resumed_ple).unwrap();
assert_eq!(resumed_ple, continued_ple);
executor.reset().unwrap();
assert_eq!(executor.position, 0);
assert_eq!(executor.ple_state.history, [EOS_TOKEN; PLE_HISTORY]);
let mut reset_ple = vec![1; continued_ple.len()];
executor.ple_state.conv.read(0, &mut reset_ple).unwrap();
assert!(reset_ple.into_iter().all(|byte| byte == 0));
executor.eval(1).unwrap();
for (actual, expected) in [
executor.logits[0],
executor.logits[1],
executor.logits[1000],
executor.logits[VOCAB as usize - 1],
]
.into_iter()
.zip(reference)
{
close(actual, expected);
}
executor.reset().unwrap();
executor.eval(EOS_TOKEN).unwrap();
assert_eq!(executor.ple_state.history, [EOS_TOKEN; PLE_HISTORY]);
executor.position = DENSE_BUDGET;
let error = executor
.attention("language_model.model.layers.3", 3)
.unwrap_err();
executor.context = DENSE_BUDGET + 1;
let error = executor.eval(1).unwrap_err();
assert!(error.contains("issue #97"));
fs::remove_file(checkpoint).unwrap();
}

71
src/engine/qwen.rs
View File

@@ -23,6 +23,7 @@ const MTP_BYTES: u64 = 1_672_575_532;
const KV_BYTES_PER_TOKEN: u64 = 24_576;
const GDN_STATE_BYTES: u64 = 113_246_208;
const GDN_CONV_BYTES: u64 = 2_211_840;
const PLE_CONV_BYTES: u64 = 184_320;
#[derive(Deserialize)]
struct Manifest {
@@ -120,9 +121,9 @@ pub(super) struct QwenModel {
impl QwenModel {
pub(super) fn open(root: &Path, context: u32) -> Result<Self, String> {
let loaded = load(root, context, false)?;
let mut paths = loaded
.bindings
.core
let mut bindings = loaded.bindings.core;
bindings.extend(loaded.bindings.ple);
let mut paths = bindings
.iter()
.map(|binding| binding.file.clone())
.collect::<Vec<_>>();
@@ -138,9 +139,7 @@ impl QwenModel {
map_indices.insert(path.clone(), maps.len());
maps.push(QwenMap { path, map });
}
let tensors = loaded
.bindings
.core
let tensors = bindings
.into_iter()
.map(|binding| {
let tensor = QwenTensor {
@@ -178,6 +177,16 @@ impl QwenModel {
(&self.maps[index].map, &self.maps[index].path)
}
pub(super) fn tensor_bytes<'a>(&'a self, tensor: &QwenTensor) -> Result<&'a [u8], String> {
let map = &self.maps[tensor.map].map;
let start = usize::try_from(tensor.range.start)
.map_err(|_| format!("{} starts beyond this platform", tensor.name))?;
let end = usize::try_from(tensor.range.end)
.map_err(|_| format!("{} ends beyond this platform", tensor.name))?;
map.get(start..end)
.ok_or_else(|| format!("{} is outside its mapped artifact", tensor.name))
}
pub(super) fn checkpoint_identity(&self) -> [u8; 32] {
self.identity
}
@@ -243,6 +252,48 @@ impl QwenModel {
pub(super) fn memory(&self) -> &MemoryPlan {
&self.memory
}
#[cfg(test)]
pub(super) fn mapped_residency(&self) -> Result<(u64, u64), String> {
// SAFETY: sysconf is read-only and has no pointer preconditions.
let page = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
if page <= 0 {
return Err("macOS did not report its virtual-memory page size".into());
}
let page = page as usize;
let mut core = 0_u64;
let mut ple = 0_u64;
for item in &self.maps {
let mut pages = vec![0_i8; item.map.len().div_ceil(page)];
// SAFETY: each read-only mmap and residency vector remain valid for this call.
if unsafe {
libc::mincore(
item.map.as_ptr().cast_mut().cast(),
item.map.len(),
pages.as_mut_ptr(),
)
} != 0
{
return Err(format!(
"cannot inspect residency for {}: {}",
item.path.display(),
std::io::Error::last_os_error()
));
}
let bytes = (pages.iter().filter(|value| **value & 1 != 0).count() * page)
.min(item.map.len()) as u64;
if item
.path
.file_name()
.is_some_and(|name| name == "ngram-table.safetensors")
{
ple += bytes;
} else {
core += bytes;
}
}
Ok((core, ple))
}
}
pub(crate) fn validate_artifacts(root: &Path) -> Result<(), String> {
@@ -587,7 +638,7 @@ pub(super) fn memory_plan(
.checked_mul(u64::from(context))
.ok_or_else(|| "Qwen KV memory size overflows".to_owned())?;
let kv_and_recurrent = kv
.checked_add(GDN_STATE_BYTES + GDN_CONV_BYTES)
.checked_add(GDN_STATE_BYTES + GDN_CONV_BYTES + PLE_CONV_BYTES)
.ok_or_else(|| "Qwen recurrent memory size overflows".to_owned())?;
let prefill_transient = u64::from(prefill_chunk)
.checked_mul((4 * 2_560 + 2_048 + 2_048 + 6_144 + 6_144) * 2)
@@ -630,12 +681,12 @@ mod tests {
assert_eq!(plan.resident_core, CORE_BYTES);
assert_eq!(plan.mapped_ple, PLE_BYTES);
assert_eq!(plan.optional_mtp, MTP_BYTES);
assert_eq!(plan.kv_and_recurrent, 6_557_908_992);
assert_eq!(plan.kv_and_recurrent, 6_558_093_312);
assert_eq!(plan.prefill_transient, 27_262_976);
assert_eq!(plan.admission, 80_000_430_099);
assert_eq!(plan.admission, 80_000_614_419);
let without_mtp = memory_plan(262_144, false, 512).unwrap();
assert_eq!(without_mtp.optional_mtp, MTP_BYTES);
assert_eq!(without_mtp.admission, 78_327_854_567);
assert_eq!(without_mtp.admission, 78_328_038_887);
assert!(memory_plan(0, false, 512).is_err());
assert!(memory_plan(262_145, false, 512).is_err());
assert!(memory_plan(1, false, 0).is_err());