From 87ccf67d0c2fb0b152696b8c178c3822a68d9ddd Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Thu, 3 Sep 2026 21:12:20 +0200 Subject: [PATCH] Stream Qwen PLE embeddings --- metal/qwen38.metal | 180 ++++++++-- native/metal/ds4_metal.m | 9 +- src/engine/metal/qwen.rs | 711 ++++++++++++++++++++++++++++++++++++--- src/engine/qwen.rs | 71 +++- 4 files changed, 893 insertions(+), 78 deletions(-) diff --git a/metal/qwen38.metal b/metal/qwen38.metal index b558d81..9a148f6 100644 --- a/metal/qwen38.metal +++ b/metal/qwen38.metal @@ -16,6 +16,52 @@ static inline ushort qwen_to_bf16(float value) { return (ushort)(bits >> 16); } +static inline ushort qwen_weight_u16( + device const uchar *data, + uint byte_offset, + ulong index) { + const ulong byte = (ulong)byte_offset + index * 2u; + if ((byte & 1u) == 0u) return *((device const ushort *)(data + byte)); + return (ushort)data[byte] | ((ushort)data[byte + 1u] << 8u); +} + +static inline uint qwen_weight_u32( + device const uchar *data, + uint byte_offset, + ulong index) { + const ulong byte = (ulong)byte_offset + index * 4u; + if ((byte & 3u) == 0u) return *((device const uint *)(data + byte)); + return (uint)data[byte] | + ((uint)data[byte + 1u] << 8u) | + ((uint)data[byte + 2u] << 16u) | + ((uint)data[byte + 3u] << 24u); +} + +static inline float qwen_quant_weight( + device const uchar *packed, + device const uchar *scales, + device const uchar *biases, + uint packed_offset, + uint scales_offset, + uint biases_offset, + uint row, + uint column, + uint in_dim, + uint bits, + uint group_size) { + const uint per_word = 32u / bits; + const uint packed_columns = in_dim / per_word; + const uint groups = in_dim / group_size; + const uint word = qwen_weight_u32(packed, packed_offset, + (ulong)row * packed_columns + column / per_word); + const uint mask = (1u << bits) - 1u; + const uint quant = (word >> ((column % per_word) * bits)) & mask; + const uint group = row * groups + column / group_size; + return fma((float)quant, + qwen_bf16(qwen_weight_u16(scales, scales_offset, group)), + qwen_bf16(qwen_weight_u16(biases, biases_offset, group))); +} + static inline float qwen_quant_value( device const uint *packed, device const ushort *scales, @@ -39,17 +85,18 @@ kernel void kernel_qwen_affine_mv( constant qwen_kernel_args &args [[buffer(0)]], device float *out [[buffer(1)]], device const float *x [[buffer(2)]], - device const uint *packed [[buffer(5)]], - device const ushort *scales [[buffer(6)]], - device const ushort *biases [[buffer(7)]], + device const uchar *packed [[buffer(5)]], + device const uchar *scales [[buffer(6)]], + device const uchar *biases [[buffer(7)]], uint row [[thread_position_in_grid]]) { const uint in_dim = args.u[0]; const uint out_dim = args.u[1]; if (row >= out_dim) return; float sum = 0.0f; for (uint column = 0; column < in_dim; column++) { - sum = fma(qwen_quant_value(packed, scales, biases, row, column, - in_dim, args.u[2], args.u[3]), x[column], sum); + sum = fma(qwen_quant_weight(packed, scales, biases, + args.u[13], args.u[14], args.u[15], row, column, + in_dim, args.u[2], args.u[3]), x[column], sum); } out[row] = sum; } @@ -57,25 +104,43 @@ kernel void kernel_qwen_affine_mv( kernel void kernel_qwen_affine_embedding( constant qwen_kernel_args &args [[buffer(0)]], device float *out [[buffer(1)]], - device const uint *packed [[buffer(5)]], - device const ushort *scales [[buffer(6)]], - device const ushort *biases [[buffer(7)]], + device const uchar *packed [[buffer(5)]], + device const uchar *scales [[buffer(6)]], + device const uchar *biases [[buffer(7)]], uint column [[thread_position_in_grid]]) { if (column >= args.u[0]) return; - out[column] = qwen_quant_value(packed, scales, biases, args.u[4], column, - args.u[0], args.u[2], args.u[3]); + out[column] = qwen_quant_weight(packed, scales, biases, + args.u[13], args.u[14], args.u[15], + args.u[4], column, args.u[0], args.u[2], args.u[3]); +} + +kernel void kernel_qwen_ple_dequant( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const uint *packed [[buffer(2)]], + device const ushort *scales [[buffer(3)]], + device const ushort *biases [[buffer(4)]], + uint index [[thread_position_in_grid]]) { + const uint dim = args.u[0]; + if (index >= dim * args.u[1]) return; + const uint row = index / dim; + const uint column = index % dim; + out[index] = qwen_quant_value(packed, scales, biases, row, column, + dim, args.u[2], args.u[3]); } kernel void kernel_qwen_bf16_mv( constant qwen_kernel_args &args [[buffer(0)]], device float *out [[buffer(1)]], device const float *x [[buffer(2)]], - device const ushort *weights [[buffer(5)]], + device const uchar *weights [[buffer(5)]], uint row [[thread_position_in_grid]]) { if (row >= args.u[1]) return; float sum = 0.0f; for (uint column = 0; column < args.u[0]; column++) { - sum = fma(qwen_bf16(weights[(ulong)row * args.u[0] + column]), x[column], sum); + sum = fma(qwen_bf16(qwen_weight_u16(weights, args.u[13], + (ulong)row * args.u[0] + column)), + x[column], sum); } out[row] = sum; } @@ -92,7 +157,7 @@ kernel void kernel_qwen_zero_rms( constant qwen_kernel_args &args [[buffer(0)]], device float *out [[buffer(1)]], device const float *x [[buffer(2)]], - device const ushort *weight [[buffer(5)]], + device const uchar *weight [[buffer(5)]], uint group [[thread_position_in_grid]]) { const uint width = args.u[0]; const uint group_size = args.u[1]; @@ -103,7 +168,8 @@ kernel void kernel_qwen_zero_rms( const float scale = rsqrt(variance / (float)group_size + args.f[0]); for (uint i = 0; i < group_size; i++) { const uint index = start + i; - out[index] = x[index] * scale * (1.0f + qwen_bf16(weight[index])); + out[index] = x[index] * scale * + (1.0f + qwen_bf16(qwen_weight_u16(weight, args.u[13], index))); } } @@ -160,19 +226,73 @@ kernel void kernel_qwen_hyper_inject( out[index] = residual[index] + block[index % hidden] * gate[index / hidden]; } +kernel void kernel_qwen_ple_gate( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *key [[buffer(2)]], + device const float *query [[buffer(3)]], + device const float *value [[buffer(4)]], + uint stream [[thread_position_in_grid]]) { + const uint hidden = args.u[0]; + if (stream >= 4u) return; + const ulong base = (ulong)stream * hidden; + float score = 0.0f; + for (uint i = 0; i < hidden; i++) score = fma(key[base + i], query[base + i], score); + score *= rsqrt((float)hidden); + const float transformed = copysign(sqrt(max(abs(score), 1.0e-6f)), score); + const float gate = 1.0f / (1.0f + exp(-transformed)); + for (uint i = 0; i < hidden; i++) out[base + i] = value[i] * gate; +} + +kernel void kernel_qwen_ple_conv( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *gated [[buffer(2)]], + device const float *normalized [[buffer(3)]], + device ushort *state [[buffer(4)]], + device const uchar *weight [[buffer(5)]], + uint channel [[thread_position_in_grid]]) { + if (channel >= args.u[0]) return; + device ushort *history = state + (ulong)channel * 9u; + float value = fma(qwen_bf16(history[0]), + qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u)), + fma(qwen_bf16(history[3]), + qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 1u)), + fma(qwen_bf16(history[6]), + qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 2u)), + normalized[channel] * + qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 3u))))); + for (uint i = 0; i < 8u; i++) history[i] = history[i + 1u]; + history[8] = qwen_to_bf16(normalized[channel]); + out[channel] = gated[channel] + value / (1.0f + exp(-value)); +} + +kernel void kernel_qwen_add( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *a [[buffer(2)]], + device const float *b [[buffer(3)]], + uint index [[thread_position_in_grid]]) { + if (index < args.u[0]) out[index] = a[index] + b[index]; +} + kernel void kernel_qwen_conv_silu( constant qwen_kernel_args &args [[buffer(0)]], device float *out [[buffer(1)]], device const float *x [[buffer(2)]], device ushort *state [[buffer(3)]], - device const ushort *weight [[buffer(5)]], + device const uchar *weight [[buffer(5)]], uint channel [[thread_position_in_grid]]) { if (channel >= args.u[0]) return; device ushort *history = state + (ulong)channel * 3u; - float value = fma(qwen_bf16(history[0]), qwen_bf16(weight[(ulong)channel * 4u]), - fma(qwen_bf16(history[1]), qwen_bf16(weight[(ulong)channel * 4u + 1u]), - fma(qwen_bf16(history[2]), qwen_bf16(weight[(ulong)channel * 4u + 2u]), - x[channel] * qwen_bf16(weight[(ulong)channel * 4u + 3u])))); + float value = fma(qwen_bf16(history[0]), + qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u)), + fma(qwen_bf16(history[1]), + qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 1u)), + fma(qwen_bf16(history[2]), + qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 2u)), + x[channel] * + qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 3u))))); history[0] = history[1]; history[1] = history[2]; history[2] = qwen_to_bf16(x[channel]); @@ -185,8 +305,8 @@ kernel void kernel_qwen_gdn_step( device const float *qkv [[buffer(2)]], device const float *controls [[buffer(3)]], device float *state [[buffer(4)]], - device const ushort *a_log [[buffer(5)]], - device const ushort *dt_bias [[buffer(6)]], + device const uchar *a_log [[buffer(5)]], + device const uchar *dt_bias [[buffer(6)]], uint2 gid [[thread_position_in_grid]]) { const uint value_index = gid.x; const uint head = gid.y; @@ -207,9 +327,10 @@ kernel void kernel_qwen_gdn_step( const float qscale = rsqrt(qsum + args.f[0]) * rsqrt((float)dim); const float kscale = rsqrt(ksum + args.f[0]); const float beta = 1.0f / (1.0f + exp(-controls[args.u[3] + head])); - const float step = controls[args.u[4] + head] + qwen_bf16(dt_bias[head]); + const float step = controls[args.u[4] + head] + + qwen_bf16(qwen_weight_u16(dt_bias, args.u[14], head)); const float softplus = max(step, 0.0f) + log(1.0f + exp(-abs(step))); - const float decay = exp(-exp(qwen_bf16(a_log[head])) * softplus); + const float decay = exp(-exp(qwen_bf16(qwen_weight_u16(a_log, args.u[13], head))) * softplus); device float *column = state + ((ulong)head * dim * dim) + value_index; float prediction = 0.0f; for (uint i = 0; i < dim; i++) { @@ -232,7 +353,7 @@ kernel void kernel_qwen_gdn_norm_gate( device float *out [[buffer(1)]], device const float *x [[buffer(2)]], device const float *controls [[buffer(3)]], - device const ushort *weight [[buffer(5)]], + device const uchar *weight [[buffer(5)]], uint head [[thread_position_in_grid]]) { const uint dim = args.u[0]; if (head >= args.u[1]) return; @@ -242,7 +363,8 @@ kernel void kernel_qwen_gdn_norm_gate( const float scale = rsqrt(variance / (float)dim + args.f[0]); for (uint i = 0; i < dim; i++) { const ulong index = (ulong)head * dim + i; - out[index] = row[i] * scale * qwen_bf16(weight[i]) / + out[index] = row[i] * scale * + qwen_bf16(qwen_weight_u16(weight, args.u[13], i)) / (1.0f + exp(-controls[index])); } } @@ -324,7 +446,7 @@ kernel void kernel_qwen_head_norm_rope( constant qwen_kernel_args &args [[buffer(0)]], device float *out [[buffer(1)]], device const float *x [[buffer(2)]], - device const ushort *weight [[buffer(5)]], + device const uchar *weight [[buffer(5)]], uint2 gid [[thread_position_in_grid]]) { const uint column = gid.x; const uint head = gid.y; @@ -335,11 +457,13 @@ kernel void kernel_qwen_head_norm_rope( float variance = 0.0f; for (uint i = 0; i < dim; i++) variance = fma(row[i], row[i], variance); const float scale = rsqrt(variance / (float)dim + args.f[0]); - float value = row[column] * scale * (1.0f + qwen_bf16(weight[column])); + float value = row[column] * scale * + (1.0f + qwen_bf16(qwen_weight_u16(weight, args.u[13], column))); if (column < rotary) { const uint rotary_half = rotary / 2u; const uint pair = column < rotary_half ? column + rotary_half : column - rotary_half; - const float paired = row[pair] * scale * (1.0f + qwen_bf16(weight[pair])); + const float paired = row[pair] * scale * + (1.0f + qwen_bf16(qwen_weight_u16(weight, args.u[13], pair))); const float theta = (float)args.u[3] * pow(args.f[1], -2.0f * (float)(column % rotary_half) / (float)rotary); value = value * cos(theta) + (column < rotary_half ? -paired : paired) * sin(theta); } diff --git a/native/metal/ds4_metal.m b/native/metal/ds4_metal.m index 45a737c..36ce2ad 100644 --- a/native/metal/ds4_metal.m +++ b/native/metal/ds4_metal.m @@ -11762,7 +11762,7 @@ int ds4_gpu_qwen_dispatch( id enc = ds4_gpu_compute_encoder(cb); if (!enc) return 0; [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; + ds4_gpu_qwen_kernel_args bound_args = *args; const DS4MetalTensor *tensors[4] = { ds4_gpu_tensor_const_obj(out), @@ -11788,8 +11788,13 @@ int ds4_gpu_qwen_dispatch( if (owned) [cb commit]; return 0; } - [enc setBuffer:weight offset:(NSUInteger)inner atIndex:5 + i]; + /* Metal resource offsets are four-byte aligned. Preserve an arbitrary + * safetensors data offset for the Qwen kernels to decode explicitly. */ + const uint64_t aligned_inner = inner & ~3ull; + bound_args.u[13 + i] = (uint32_t)(inner - aligned_inner); + [enc setBuffer:weight offset:(NSUInteger)aligned_inner atIndex:5 + i]; } + [enc setBytes:&bound_args length:sizeof(bound_args) atIndex:0]; const NSUInteger width = pipeline.threadExecutionWidth; const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup; diff --git a/src/engine/metal/qwen.rs b/src/engine/metal/qwen.rs index e2a56ab..0ddf00d 100644 --- a/src/engine/metal/qwen.rs +++ b/src/engine/metal/qwen.rs @@ -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, + ple_contract: PleContract, + ple_state: PleState, scratch: Scratch, logits: Vec, tokens: Vec, @@ -127,6 +164,7 @@ pub(in crate::engine) struct QwenExecutor { pub(in crate::engine) struct QwenResidentState { states: Vec, + ple_state: PleState, logits: Vec, tokens: Vec, position: u32, @@ -135,11 +173,14 @@ pub(in crate::engine) struct QwenResidentState { impl QwenExecutor { pub(super) fn open(model: QwenModel, context: u32) -> Result { + 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 { 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, String> { .collect() } +fn allocate_ple_state() -> Result { + 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 { + let multipliers = read_i64_array::<3>( + model, + "language_model.model.layers.1.ple.ple_embedding.layer_multipliers", + )?; + let sizes = read_i64_array::( + model, + "language_model.model.layers.1.ple.ple_embedding.ngram_heads_vocab_sizes", + )?; + let offsets = read_i64_array::( + 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(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::>() + }) + .collect::>() + }; + 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::>(); + let scales = (0..30).map(|value| (value + 17) as u8).collect::>(); + let biases = (0..30).map(|value| (value + 47) as u8).collect::>(); + 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(); } diff --git a/src/engine/qwen.rs b/src/engine/qwen.rs index 9ecb43b..7a4494c 100644 --- a/src/engine/qwen.rs +++ b/src/engine/qwen.rs @@ -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 { 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::>(); @@ -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());