diff --git a/metal/qwen38.metal b/metal/qwen38.metal new file mode 100644 index 0000000..b558d81 --- /dev/null +++ b/metal/qwen38.metal @@ -0,0 +1,408 @@ +// Qwen3.8 Flash Next primitives. Rust owns the graph and all state lifetimes; +// this file contains only the data-parallel kernels executed by Metal. + +struct qwen_kernel_args { + uint u[16]; + float f[8]; +}; + +static inline float qwen_bf16(ushort value) { + return as_type((uint)value << 16); +} + +static inline ushort qwen_to_bf16(float value) { + uint bits = as_type(value); + bits += 0x7fffu + ((bits >> 16) & 1u); + return (ushort)(bits >> 16); +} + +static inline float qwen_quant_value( + device const uint *packed, + device const ushort *scales, + device const ushort *biases, + 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 = packed[(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(scales[group]), qwen_bf16(biases[group])); +} + +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)]], + 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); + } + out[row] = sum; +} + +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)]], + 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]); +} + +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)]], + 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); + } + out[row] = sum; +} + +kernel void kernel_qwen_repeat4( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + uint index [[thread_position_in_grid]]) { + if (index < args.u[0] * 4u) out[index] = x[index % args.u[0]]; +} + +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)]], + uint group [[thread_position_in_grid]]) { + const uint width = args.u[0]; + const uint group_size = args.u[1]; + if (group >= width / group_size) return; + const uint start = group * group_size; + float variance = 0.0f; + for (uint i = 0; i < group_size; i++) variance = fma(x[start + i], x[start + i], variance); + 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])); + } +} + +kernel void kernel_qwen_silu_div4( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + uint index [[thread_position_in_grid]]) { + if (index >= args.u[0]) return; + const float value = x[index] * 0.25f; + out[index] = value / (1.0f + exp(-value)); +} + +kernel void kernel_qwen_sigmoid( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + uint index [[thread_position_in_grid]]) { + if (index < args.u[0]) out[index] = 1.0f / (1.0f + exp(-x[index])); +} + +kernel void kernel_qwen_sigmoid2_div4( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + uint index [[thread_position_in_grid]]) { + if (index < args.u[0]) out[index] = 2.0f / (1.0f + exp(-x[index] * 0.25f)); +} + +kernel void kernel_qwen_hyper_mix( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *normalized [[buffer(2)]], + device const float *mix [[buffer(3)]], + uint index [[thread_position_in_grid]]) { + if (index >= args.u[0]) return; + float value = 0.0f; + for (uint stream = 0; stream < 4u; stream++) { + const uint offset = stream * args.u[0] + index; + value = fma(normalized[offset], mix[offset], value); + } + out[index] = value * 0.25f; +} + +kernel void kernel_qwen_hyper_inject( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *residual [[buffer(2)]], + device const float *block [[buffer(3)]], + device const float *gate [[buffer(4)]], + uint index [[thread_position_in_grid]]) { + const uint hidden = args.u[0]; + if (index >= hidden * 4u) return; + out[index] = residual[index] + block[index % hidden] * gate[index / hidden]; +} + +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)]], + 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])))); + history[0] = history[1]; + history[1] = history[2]; + history[2] = qwen_to_bf16(x[channel]); + out[channel] = value / (1.0f + exp(-value)); +} + +kernel void kernel_qwen_gdn_step( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + 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)]], + uint2 gid [[thread_position_in_grid]]) { + const uint value_index = gid.x; + const uint head = gid.y; + const uint dim = args.u[0]; + const uint key_heads = args.u[1]; + const uint value_heads = args.u[2]; + if (value_index >= dim || head >= value_heads) return; + const uint key_head = head / (value_heads / key_heads); + device const float *q_raw = qkv + (ulong)key_head * dim; + device const float *k_raw = qkv + (ulong)key_heads * dim + (ulong)key_head * dim; + device const float *value = qkv + (ulong)key_heads * dim * 2u + (ulong)head * dim; + float qsum = 0.0f; + float ksum = 0.0f; + for (uint i = 0; i < dim; i++) { + qsum = fma(q_raw[i], q_raw[i], qsum); + ksum = fma(k_raw[i], k_raw[i], ksum); + } + 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 softplus = max(step, 0.0f) + log(1.0f + exp(-abs(step))); + const float decay = exp(-exp(qwen_bf16(a_log[head])) * softplus); + device float *column = state + ((ulong)head * dim * dim) + value_index; + float prediction = 0.0f; + for (uint i = 0; i < dim; i++) { + prediction = fma(column[(ulong)i * dim] * decay, k_raw[i] * kscale, prediction); + } + const float delta = (value[value_index] - prediction) * beta; + float result = 0.0f; + for (uint i = 0; i < dim; i++) { + const ulong offset = (ulong)i * dim; + const float updated = column[offset] * decay + k_raw[i] * kscale * delta; + column[offset] = updated; + result = fma(updated, q_raw[i] * qscale, result); + } + device float *head_out = out + (ulong)head * dim; + head_out[value_index] = result; +} + +kernel void kernel_qwen_gdn_norm_gate( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const float *controls [[buffer(3)]], + device const ushort *weight [[buffer(5)]], + uint head [[thread_position_in_grid]]) { + const uint dim = args.u[0]; + if (head >= args.u[1]) return; + device const float *row = x + (ulong)head * dim; + 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]); + for (uint i = 0; i < dim; i++) { + const ulong index = (ulong)head * dim + i; + out[index] = row[i] * scale * qwen_bf16(weight[i]) / + (1.0f + exp(-controls[index])); + } +} + +kernel void kernel_qwen_swiglu( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *gate [[buffer(2)]], + device const float *up [[buffer(3)]], + uint index [[thread_position_in_grid]]) { + if (index >= args.u[0]) return; + out[index] = gate[index] / (1.0f + exp(-gate[index])) * up[index]; +} + +kernel void kernel_qwen_route_top10( + constant qwen_kernel_args &args [[buffer(0)]], + device int *ids [[buffer(1)]], + device float *weights [[buffer(2)]], + device const float *logits [[buffer(3)]], + uint gid [[thread_position_in_grid]]) { + if (gid != 0u) return; + float max_value = -INFINITY; + for (uint i = 0; i < args.u[0]; i++) max_value = max(max_value, logits[i]); + float sum = 0.0f; + for (uint i = 0; i < args.u[0]; i++) sum += exp(logits[i] - max_value); + float selected_sum = 0.0f; + for (uint slot = 0; slot < 10u; slot++) { + float best = -1.0f; + int best_id = -1; + for (uint i = 0; i < args.u[0]; i++) { + bool used = false; + for (uint j = 0; j < slot; j++) used = used || ids[j] == (int)i; + const float probability = exp(logits[i] - max_value) / sum; + if (!used && probability > best) { + best = probability; + best_id = (int)i; + } + } + ids[slot] = best_id; + weights[slot] = best; + selected_sum += best; + } + for (uint slot = 0; slot < 10u; slot++) weights[slot] /= selected_sum; +} + +kernel void kernel_qwen_accumulate( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + uint index [[thread_position_in_grid]]) { + if (index < args.u[0]) out[index] += x[index] * args.f[0]; +} + +kernel void kernel_qwen_accumulate_sigmoid_scalar( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const float *gate [[buffer(3)]], + uint index [[thread_position_in_grid]]) { + if (index < args.u[0]) out[index] += x[index] / (1.0f + exp(-gate[0])); +} + +kernel void kernel_qwen_split_q_gate( + constant qwen_kernel_args &args [[buffer(0)]], + device float *q [[buffer(1)]], + device const float *packed [[buffer(2)]], + device float *gate [[buffer(3)]], + uint index [[thread_position_in_grid]]) { + const uint heads = args.u[0]; + const uint dim = args.u[1]; + if (index >= heads * dim) return; + const uint head = index / dim; + const uint column = index % dim; + q[index] = packed[(ulong)head * dim * 2u + column]; + gate[index] = packed[(ulong)head * dim * 2u + dim + column]; +} + +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)]], + uint2 gid [[thread_position_in_grid]]) { + const uint column = gid.x; + const uint head = gid.y; + const uint dim = args.u[0]; + const uint rotary = args.u[1]; + if (column >= dim || head >= args.u[2]) return; + device const float *row = x + (ulong)head * dim; + 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])); + 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 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); + } + out[(ulong)head * dim + column] = value; +} + +kernel void kernel_qwen_store_kv_bf16( + constant qwen_kernel_args &args [[buffer(0)]], + device ushort *cache [[buffer(1)]], + device const float *key [[buffer(2)]], + device const float *value [[buffer(3)]], + uint index [[thread_position_in_grid]]) { + const uint width = args.u[0]; + if (index >= width) return; + const ulong base = (ulong)args.u[1] * width * 2u; + cache[base + index] = qwen_to_bf16(key[index]); + cache[base + width + index] = qwen_to_bf16(value[index]); +} + +kernel void kernel_qwen_dense_attention( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *query [[buffer(2)]], + device const ushort *cache [[buffer(3)]], + uint head [[thread_position_in_grid]]) { + const uint heads = args.u[0]; + const uint kv_heads = args.u[1]; + const uint dim = args.u[2]; + const uint tokens = args.u[3]; + if (head >= heads) return; + const uint kv_head = head / (heads / kv_heads); + float max_score = -INFINITY; + for (uint token = 0; token < tokens; token++) { + const ulong base = (ulong)token * kv_heads * dim * 2u + (ulong)kv_head * dim; + float score = 0.0f; + for (uint i = 0; i < dim; i++) score = fma(query[(ulong)head * dim + i], qwen_bf16(cache[base + i]), score); + max_score = max(max_score, score * rsqrt((float)dim)); + } + float denominator = 0.0f; + for (uint token = 0; token < tokens; token++) { + const ulong base = (ulong)token * kv_heads * dim * 2u + (ulong)kv_head * dim; + float score = 0.0f; + for (uint i = 0; i < dim; i++) score = fma(query[(ulong)head * dim + i], qwen_bf16(cache[base + i]), score); + denominator += exp(score * rsqrt((float)dim) - max_score); + } + for (uint column = 0; column < dim; column++) { + float value = 0.0f; + for (uint token = 0; token < tokens; token++) { + const ulong base = (ulong)token * kv_heads * dim * 2u + (ulong)kv_head * dim; + float score = 0.0f; + for (uint i = 0; i < dim; i++) score = fma(query[(ulong)head * dim + i], qwen_bf16(cache[base + i]), score); + const float probability = exp(score * rsqrt((float)dim) - max_score) / denominator; + value = fma(probability, qwen_bf16(cache[base + kv_heads * dim + column]), value); + } + out[(ulong)head * dim + column] = value; + } +} + +kernel void kernel_qwen_gate_attention( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *attention [[buffer(2)]], + device const float *gate [[buffer(3)]], + uint index [[thread_position_in_grid]]) { + if (index < args.u[0]) out[index] = attention[index] / (1.0f + exp(-gate[index])); +} diff --git a/native/metal/ds4_gpu.h b/native/metal/ds4_gpu.h index caf7b83..f848b9e 100644 --- a/native/metal/ds4_gpu.h +++ b/native/metal/ds4_gpu.h @@ -128,6 +128,30 @@ int ds4_gpu_set_aux_model_map_range(const void *model_map, uint64_t map_offset, uint64_t map_size); int ds4_gpu_set_model_map_spans(const void *model_map, uint64_t model_size, const uint64_t *offsets, const uint64_t *sizes, uint32_t count, uint64_t max_tensor_bytes); + +typedef struct { + const void *map; + uint64_t size; + uint64_t offset; + uint64_t bytes; +} ds4_gpu_qwen_weight_view; + +typedef struct { + uint32_t u[16]; + float f[8]; +} ds4_gpu_qwen_kernel_args; + +int ds4_gpu_qwen_dispatch( + const char *kernel, + ds4_gpu_tensor *out, + const ds4_gpu_tensor *a, + const ds4_gpu_tensor *b, + const ds4_gpu_tensor *c, + const ds4_gpu_qwen_weight_view *weights, + uint32_t weight_count, + const ds4_gpu_qwen_kernel_args *args, + uint32_t grid_x, + uint32_t grid_y); int ds4_gpu_cache_model_range(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes, const char *label); int ds4_gpu_cache_q8_f16_range(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes, uint64_t in_dim, uint64_t out_dim, const char *label); int ds4_gpu_q8_cache_suppressed(void); diff --git a/native/metal/ds4_metal.m b/native/metal/ds4_metal.m index 82d3746..45a737c 100644 --- a/native/metal/ds4_metal.m +++ b/native/metal/ds4_metal.m @@ -4359,6 +4359,7 @@ static NSString *ds4_gpu_full_source(void) { @[@"DS4_METAL_GLM53_BF16_SOURCE", @"metal/glm53_bf16.metal"], @[@"DS4_METAL_GLM53_VISION_SOURCE", @"metal/glm53_vision.metal"], @[@"DS4_METAL_GLM53_KDA_SOURCE", @"metal/glm53_kda.metal"], + @[@"DS4_METAL_QWEN38_SOURCE", @"metal/qwen38.metal"], @[@"DS4_METAL_MOE_SOURCE", @"metal/moe.metal"], @[@"DS4_METAL_DSV4_HC_SOURCE", @"metal/dsv4_hc.metal"], @[@"DS4_METAL_UNARY_SOURCE", @"metal/unary.metal"], @@ -11737,6 +11738,68 @@ static id ds4_gpu_wrap_model_exact_range_owned( DS4_GPU_EXACT_VIEW_OWNED); } +int ds4_gpu_qwen_dispatch( + const char *kernel, + ds4_gpu_tensor *out, + const ds4_gpu_tensor *a, + const ds4_gpu_tensor *b, + const ds4_gpu_tensor *c, + const ds4_gpu_qwen_weight_view *weights, + uint32_t weight_count, + const ds4_gpu_qwen_kernel_args *args, + uint32_t grid_x, + uint32_t grid_y) { + if (!kernel || !out || !args || grid_x == 0 || grid_y == 0 || + weight_count > 3 || (weight_count != 0 && !weights)) { + return 0; + } + id pipeline = ds4_gpu_get_pipeline(kernel); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + + const DS4MetalTensor *tensors[4] = { + ds4_gpu_tensor_const_obj(out), + a ? ds4_gpu_tensor_const_obj(a) : nil, + b ? ds4_gpu_tensor_const_obj(b) : nil, + c ? ds4_gpu_tensor_const_obj(c) : nil, + }; + for (uint32_t i = 0; i < 4; i++) { + if (tensors[i]) { + [enc setBuffer:tensors[i].buffer offset:(NSUInteger)tensors[i].offset atIndex:1 + i]; + } + } + for (uint32_t i = 0; i < weight_count; i++) { + uint64_t inner = 0; + id weight = ds4_gpu_wrap_model_exact_range( + weights[i].map, + weights[i].size, + weights[i].offset, + weights[i].bytes, + &inner); + if (!weight) { + ds4_gpu_end_compute_encoder(cb, enc); + if (owned) [cb commit]; + return 0; + } + [enc setBuffer:weight offset:(NSUInteger)inner atIndex:5 + i]; + } + + const NSUInteger width = pipeline.threadExecutionWidth; + const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup; + const NSUInteger threads = MIN(MAX(width, 1u), max_threads); + [enc dispatchThreads:MTLSizeMake(grid_x, grid_y, 1) + threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return owned ? ds4_gpu_finish_command_buffer(cb, 1, kernel) : 1; +} + uint32_t ds4_gpu_stream_expert_cache_configured_count(void) { uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); if (budget > DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) { diff --git a/src/engine.rs b/src/engine.rs index 67325f8..41fe7af 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -53,6 +53,7 @@ pub(crate) fn checkpoint_model(path: &Path) -> Option { let model_size_offset = match &magic { b"DS4RKV01" => 40, b"DS4GLM01" => 44, + b"DS4QWN01" => return Some(ModelChoice::Qwen38FlashNext), _ => return None, }; file.seek(SeekFrom::Start(model_size_offset)).ok()?; @@ -301,6 +302,115 @@ pub(crate) struct Model { tokenizer: Tokenizer, } +enum LoadedModel { + Gguf(Box), + Qwen(Box), +} + +impl From for LoadedModel { + fn from(model: Model) -> Self { + Self::Gguf(Box::new(model)) + } +} + +#[derive(Clone, Copy)] +enum ModelRef<'a> { + Gguf(&'a Model), + Qwen(&'a qwen::QwenModel), +} + +impl LoadedModel { + fn open(settings: &EngineSettings) -> Result { + if settings.model.is_qwen38() { + validate_engine_artifacts( + settings.model, + settings.speculative.dspark, + &settings.artifacts, + )?; + let context = u32::try_from(settings.context_tokens) + .map_err(|_| "Qwen context must be a positive whole number")?; + qwen::QwenModel::open(&settings.artifacts.model, context) + .map(Box::new) + .map(Self::Qwen) + } else { + Model::open(settings).map(Box::new).map(Self::Gguf) + } + } +} + +impl ModelRef<'_> { + fn summary(self) -> ModelSummary { + match self { + Self::Gguf(model) => model.summary(), + Self::Qwen(model) => model.summary(), + } + } + + fn render_conversation( + self, + system: &str, + messages: &[ChatTurn], + reasoning: ReasoningMode, + ) -> Vec { + match self { + Self::Gguf(model) => model.render_conversation(system, messages, reasoning), + Self::Qwen(model) => model.render_conversation(system, messages, reasoning), + } + } + + fn render_history( + self, + system: &str, + messages: &[ChatTurn], + reasoning: ReasoningMode, + ) -> Vec { + match self { + Self::Gguf(model) => model.render_history(system, messages, reasoning), + Self::Qwen(model) => model.render_history(system, messages, reasoning), + } + } + + fn render_continuation( + self, + prompt: &str, + reasoning: ReasoningMode, + skip_previous_eos: bool, + ) -> Vec { + match self { + Self::Gguf(model) => model.render_continuation(prompt, reasoning, skip_previous_eos), + Self::Qwen(model) => model.render_continuation(prompt, reasoning, skip_previous_eos), + } + } + + fn token_bytes(self, token: i32) -> Option> { + match self { + Self::Gguf(model) => model.token_bytes(token), + Self::Qwen(model) => model.token_bytes(token), + } + } + + fn is_stop_token_for_reasoning(self, token: i32, reasoning: ReasoningMode) -> bool { + match self { + Self::Gguf(model) => model.is_stop_token_for_reasoning(token, reasoning), + Self::Qwen(model) => model.is_stop_token_for_reasoning(token, reasoning), + } + } + + fn is_think_start_token(self, token: i32) -> bool { + match self { + Self::Gguf(model) => model.is_think_start_token(token), + Self::Qwen(model) => model.is_think_start_token(token), + } + } + + fn is_think_end_token(self, token: i32) -> bool { + match self { + Self::Gguf(model) => model.is_think_end_token(token), + Self::Qwen(model) => model.is_think_end_token(token), + } + } +} + #[derive(Clone, Debug, Eq, PartialEq)] pub(crate) struct ModelSummary { pub(crate) model: ModelChoice, @@ -320,18 +430,7 @@ impl Model { &settings.artifacts, )?; if settings.model.is_qwen38() { - let context = u32::try_from(settings.context_tokens) - .map_err(|_| "Qwen context must be a positive whole number")?; - let loaded = qwen::load(&settings.artifacts.model, context, false)?; - return Err(format!( - "Qwen3.8 artifacts are valid, but its Rust Metal execution backend is not available until issue #95. Memory plan: {} resident core bytes, {} mapped PLE bytes, {} optional MTP bytes, {} KV/recurrent bytes, {} prefill transient bytes, {} admitted bytes.", - loaded.memory.resident_core, - loaded.memory.mapped_ple, - loaded.memory.optional_mtp, - loaded.memory.kv_and_recurrent, - loaded.memory.prefill_transient, - loaded.memory.admission, - )); + return Err("Qwen must be opened through its dedicated safetensors loader".into()); } let mut model = Self::open_main(&settings.artifacts.model, settings.model)?; if settings.execution.warm_weights { @@ -622,7 +721,7 @@ impl Generator { pub(crate) fn open(settings: &EngineSettings, metrics: Arc) -> Result { let simulated_memory = SimulatedMemory::acquire(settings.diagnostics.simulated_used_memory_bytes)?; - let model = Model::open(settings)?; + let model = LoadedModel::open(settings)?; let executor = metal::Executor::open_configured( model, settings.context_tokens.max(1) as u32, diff --git a/src/engine/metal.rs b/src/engine/metal.rs index 1252521..147835a 100644 --- a/src/engine/metal.rs +++ b/src/engine/metal.rs @@ -3,16 +3,18 @@ mod glm; mod gpu; mod hotlist; mod profile; +mod qwen; mod vision; pub(super) use vision::VisionEmbedding; use glm::GlmExecutor; use gpu::*; use profile::ExpertProfile; +use qwen::QwenExecutor; use super::gguf::{BF16, F16, F32, Gguf, IQ2_XXS, MXFP4, Q2_K, Q4_K, Q8_0, Tensor as GgufTensor}; use super::validation::{DsparkConfig, SupportKind, dspark_config}; -use super::{Model, ModelFamily, Rng, exact_delta_sample}; +use super::{LoadedModel, Model, ModelFamily, ModelRef, Rng, exact_delta_sample}; use crate::model::ModelChoice; use crate::settings::{ EngineSpeculativeSettings, EngineSsdSettings, EngineSteeringSettings, ReasoningMode, @@ -44,7 +46,7 @@ fn environment_present(name: &CStr) -> bool { !unsafe { getenv(name.as_ptr()) }.is_null() } -const SOURCES: [(&str, &str); 22] = [ +const SOURCES: [(&str, &str); 23] = [ ("DS4_METAL_FLASH_ATTN_SOURCE", "flash_attn.metal"), ("DS4_METAL_DENSE_SOURCE", "dense.metal"), ("DS4_METAL_MOE_SOURCE", "moe.metal"), @@ -67,6 +69,7 @@ const SOURCES: [(&str, &str); 22] = [ ("DS4_METAL_GLM53_BF16_SOURCE", "glm53_bf16.metal"), ("DS4_METAL_GLM53_VISION_SOURCE", "glm53_vision.metal"), ("DS4_METAL_GLM53_KDA_SOURCE", "glm53_kda.metal"), + ("DS4_METAL_QWEN38_SOURCE", "qwen38.metal"), ]; // The Metal boundary uses this only to decide whether diagnostic logs get ANSI @@ -4360,15 +4363,17 @@ impl DeepSeekExecutor { } } -/// Model-family dispatch over the two Rust-owned Metal graphs. +/// Model-family dispatch over the Rust-owned Metal graphs. pub(super) enum Executor { DeepSeek(Box), Glm(Box), + Qwen(Box), } pub(super) enum ResidentState { DeepSeek(Box), Glm(Box), + Qwen(Box), } impl Executor { @@ -4380,7 +4385,7 @@ impl Executor { prefill_chunk: u32, ) -> Result { Self::open_configured( - model, + LoadedModel::from(model), context, quality, prefill_chunk, @@ -4415,7 +4420,7 @@ impl Executor { #[allow(clippy::too_many_arguments)] pub(super) fn open_configured( - model: Model, + model: impl Into, context: u32, quality: bool, prefill_chunk: u32, @@ -4425,32 +4430,39 @@ impl Executor { steering: EngineSteeringSettings, expert_profile_path: Option<&str>, ) -> Result { - match model.shape.family { - ModelFamily::DeepSeek => DeepSeekExecutor::open_profile( - model, - context, - quality, - prefill_chunk, - power_percent, - speculative, - ssd, - steering, - expert_profile_path, - ) - .map(Box::new) - .map(Self::DeepSeek), - ModelFamily::Glm => GlmExecutor::open_profile( - model, - context, - quality, - ssd, - speculative, - steering, - expert_profile_path, - ) - .map(Box::new) - .map(Self::Glm), - ModelFamily::Qwen => unreachable!("Qwen uses its dedicated executor"), + match model.into() { + LoadedModel::Gguf(model) if model.shape.family == ModelFamily::DeepSeek => { + DeepSeekExecutor::open_profile( + *model, + context, + quality, + prefill_chunk, + power_percent, + speculative, + ssd, + steering, + expert_profile_path, + ) + .map(Box::new) + .map(Self::DeepSeek) + } + LoadedModel::Gguf(model) if model.shape.family == ModelFamily::Glm => { + GlmExecutor::open_profile( + *model, + context, + quality, + ssd, + speculative, + steering, + expert_profile_path, + ) + .map(Box::new) + .map(Self::Glm) + } + LoadedModel::Gguf(_) => unreachable!("Qwen never uses a GGUF model"), + LoadedModel::Qwen(model) => QwenExecutor::open(*model, context) + .map(Box::new) + .map(Self::Qwen), } } @@ -4483,6 +4495,10 @@ impl Executor { executor.eval(token)?; Ok(vec![token]) } + Self::Qwen(executor) => { + executor.eval(token)?; + Ok(vec![token]) + } } } @@ -4490,6 +4506,7 @@ impl Executor { match self { Self::DeepSeek(executor) => executor.eval(token), Self::Glm(executor) => executor.eval(token), + Self::Qwen(executor) => executor.eval(token), } } @@ -4499,7 +4516,7 @@ impl Executor { ) -> Result, String> { match self { Self::Glm(executor) => executor.encode_visions(encoded), - Self::DeepSeek(_) => Err("vision input requires GLM 5.3 Flash".into()), + Self::DeepSeek(_) | Self::Qwen(_) => Err("vision input requires GLM 5.3 Flash".into()), } } @@ -4510,7 +4527,8 @@ impl Executor { match self { Self::Glm(executor) => executor.set_vision_overlays(overlays), Self::DeepSeek(_) if overlays.is_empty() => Ok(()), - Self::DeepSeek(_) => Err("vision input requires GLM 5.3 Flash".into()), + Self::Qwen(_) if overlays.is_empty() => Ok(()), + Self::DeepSeek(_) | Self::Qwen(_) => Err("vision input requires GLM 5.3 Flash".into()), } } @@ -4529,6 +4547,11 @@ impl Executor { let _ = reasoning; executor.eval_speculative_greedy(token, max_tokens, cancelled) } + Self::Qwen(executor) => { + let _ = (max_tokens, reasoning, cancelled); + executor.eval(token)?; + Ok(vec![token]) + } } } @@ -4540,6 +4563,7 @@ impl Executor { match self { Self::DeepSeek(executor) => executor.prefill(tokens, progress), Self::Glm(executor) => executor.prefill(tokens, progress), + Self::Qwen(executor) => executor.prefill(tokens, progress), } } @@ -4547,6 +4571,7 @@ impl Executor { match self { Self::DeepSeek(executor) => executor.logits(), Self::Glm(executor) => executor.logits(), + Self::Qwen(executor) => executor.logits(), } } @@ -4554,13 +4579,15 @@ impl Executor { match self { Self::DeepSeek(executor) => executor.execution_stats(), Self::Glm(executor) => executor.execution_stats(), + Self::Qwen(executor) => executor.execution_stats(), } } - pub(super) fn model(&self) -> &Model { + pub(super) fn model(&self) -> ModelRef<'_> { match self { - Self::DeepSeek(executor) => executor.model(), - Self::Glm(executor) => executor.model(), + Self::DeepSeek(executor) => ModelRef::Gguf(executor.model()), + Self::Glm(executor) => ModelRef::Gguf(executor.model()), + Self::Qwen(executor) => ModelRef::Qwen(executor.model()), } } @@ -4568,6 +4595,7 @@ impl Executor { match self { Self::DeepSeek(executor) => executor.context(), Self::Glm(executor) => executor.context(), + Self::Qwen(executor) => executor.context(), } } @@ -4575,6 +4603,7 @@ impl Executor { match self { Self::DeepSeek(executor) => executor.position(), Self::Glm(executor) => executor.position(), + Self::Qwen(executor) => executor.position(), } } @@ -4582,6 +4611,7 @@ impl Executor { match self { Self::DeepSeek(executor) => executor.reset(), Self::Glm(executor) => executor.reset(), + Self::Qwen(executor) => executor.reset(), } } @@ -4596,6 +4626,9 @@ impl Executor { Some(ResidentState::Glm(_)) => { return Err("resident session belongs to a different model family".into()); } + Some(ResidentState::Qwen(_)) => { + return Err("resident session belongs to a different model family".into()); + } None => None, }; executor.swap_resident_state(&mut inner)?; @@ -4607,11 +4640,25 @@ impl Executor { Some(ResidentState::DeepSeek(_)) => { return Err("resident session belongs to a different model family".into()); } + Some(ResidentState::Qwen(_)) => { + return Err("resident session belongs to a different model family".into()); + } None => None, }; executor.swap_resident_state(&mut inner)?; *state = inner.map(|state| ResidentState::Glm(Box::new(state))); } + Self::Qwen(executor) => { + let mut inner = match state.take() { + Some(ResidentState::Qwen(state)) => Some(*state), + Some(ResidentState::DeepSeek(_) | ResidentState::Glm(_)) => { + return Err("resident session belongs to a different model family".into()); + } + None => None, + }; + executor.swap_resident_state(&mut inner)?; + *state = inner.map(|state| ResidentState::Qwen(Box::new(state))); + } } Ok(()) } @@ -4620,6 +4667,7 @@ impl Executor { match self { Self::DeepSeek(executor) => executor.align_prompt(tokens), Self::Glm(executor) => executor.align_prompt(tokens), + Self::Qwen(executor) => executor.align_prompt(tokens), } } @@ -4627,6 +4675,7 @@ impl Executor { match self { Self::DeepSeek(executor) => executor.tokens(), Self::Glm(executor) => executor.tokens(), + Self::Qwen(executor) => executor.tokens(), } } @@ -4634,6 +4683,7 @@ impl Executor { match self { Self::DeepSeek(executor) => executor.checkpoint_tag(), Self::Glm(executor) => executor.checkpoint_tag(), + Self::Qwen(executor) => executor.checkpoint_tag(), } } @@ -4641,6 +4691,7 @@ impl Executor { match self { Self::DeepSeek(executor) => executor.note_checkpoint_tag(tag), Self::Glm(executor) => executor.note_checkpoint_tag(tag), + Self::Qwen(executor) => executor.note_checkpoint_tag(tag), } } } diff --git a/src/engine/metal/checkpoint.rs b/src/engine/metal/checkpoint.rs index 647a5fe..605d809 100644 --- a/src/engine/metal/checkpoint.rs +++ b/src/engine/metal/checkpoint.rs @@ -320,6 +320,7 @@ impl Executor { match self { Self::DeepSeek(executor) => executor.save_checkpoint(path, tag, progress), Self::Glm(executor) => executor.save_checkpoint(path, tag, progress), + Self::Qwen(executor) => executor.save_checkpoint(path, tag, progress), } } @@ -331,6 +332,7 @@ impl Executor { match self { Self::DeepSeek(executor) => executor.load_checkpoint(path, progress), Self::Glm(executor) => executor.load_checkpoint(path, progress), + Self::Qwen(executor) => executor.load_checkpoint(path, progress), } } } diff --git a/src/engine/metal/gpu.rs b/src/engine/metal/gpu.rs index bafc937..634c50a 100644 --- a/src/engine/metal/gpu.rs +++ b/src/engine/metal/gpu.rs @@ -5,6 +5,22 @@ pub(super) struct GpuTensor { _private: [u8; 0], } +#[derive(Clone, Copy)] +#[repr(C)] +pub(super) struct QwenWeightView { + pub(super) map: *const c_void, + pub(super) size: u64, + pub(super) offset: u64, + pub(super) bytes: u64, +} + +#[derive(Clone, Copy, Default)] +#[repr(C)] +pub(super) struct QwenKernelArgs { + pub(super) u: [u32; 16], + pub(super) f: [f32; 8], +} + #[derive(Clone, Copy, Default)] #[repr(C)] pub(super) struct Glm53VisionLayerWeights { @@ -81,6 +97,18 @@ unsafe extern "C" { map_size: u64, max_tensor_bytes: u64, ) -> i32; + pub(super) fn ds4_gpu_qwen_dispatch( + kernel: *const c_char, + out: *mut GpuTensor, + a: *const GpuTensor, + b: *const GpuTensor, + c: *const GpuTensor, + weights: *const QwenWeightView, + weight_count: u32, + args: *const QwenKernelArgs, + grid_x: u32, + grid_y: u32, + ) -> i32; pub(super) fn ds4_gpu_set_transient_model_map_range( model_map: *const c_void, model_size: u64, @@ -1804,7 +1832,7 @@ unsafe extern "C" { } pub(super) struct Context { - _model_file: File, + _model_file: Option, } impl Context { @@ -1908,9 +1936,29 @@ impl Context { } } Ok(Self { - _model_file: model_file, + _model_file: Some(model_file), }) } + + pub(super) fn open_qwen(admission_bytes: u64) -> Result { + check(unsafe { ds4_gpu_init() }, "Metal initialization")?; + unsafe { + ds4_gpu_set_glm_model(false); + ds4_gpu_set_ssd_streaming(false); + ds4_gpu_set_decode_pipeline_fast_lookup(0); + } + let recommended = unsafe { ds4_gpu_recommended_working_set_size() }; + if admission_bytes != 0 && recommended != 0 && admission_bytes > recommended { + unsafe { ds4_gpu_cleanup() }; + return Err(format!( + "Qwen model load needs {:.1} GiB including context and scratch, but Metal recommends at most {:.1} GiB", + admission_bytes as f64 / 1_073_741_824.0, + recommended as f64 / 1_073_741_824.0, + )); + } + unsafe { ds4_gpu_set_quality(true) }; + Ok(Self { _model_file: None }) + } } impl Drop for Context { diff --git a/src/engine/metal/qwen.rs b/src/engine/metal/qwen.rs new file mode 100644 index 0000000..e2a56ab --- /dev/null +++ b/src/engine/metal/qwen.rs @@ -0,0 +1,1862 @@ +use super::checkpoint::{read_buffer, read_u32, write_buffer, write_u32}; +use super::*; +use crate::engine::qwen::{QwenModel, QwenTensor}; +use std::ffi::CStr; + +const HIDDEN: u32 = 2_560; +const HC: u32 = 4; +const HC_WIDTH: u32 = HIDDEN * HC; +const HC_RANK: u32 = 320; +const LAYERS: usize = 48; +const GDN_HEADS_K: u32 = 16; +const GDN_HEADS_V: u32 = 48; +const HEAD_DIM: u32 = 128; +const GDN_QKV: u32 = (GDN_HEADS_K * 2 + GDN_HEADS_V) * HEAD_DIM; +const GDN_VALUE: u32 = GDN_HEADS_V * HEAD_DIM; +const GDN_CONTROLS: u32 = GDN_VALUE + GDN_HEADS_V * 2; +const ATTN_HEADS: u32 = 24; +const ATTN_KV_HEADS: u32 = 2; +const ATTN_DIM: u32 = 256; +const ATTN_WIDTH: u32 = ATTN_HEADS * ATTN_DIM; +const ATTN_KV_WIDTH: u32 = ATTN_KV_HEADS * ATTN_DIM; +const EXPERTS: u32 = 512; +const EXPERTS_USED: usize = 10; +const EXPERT_WIDTH: u32 = 640; +const VOCAB: u32 = 248_320; +const DENSE_BUDGET: u32 = 2_048; +const CHECKPOINT_MAGIC: &[u8; 8] = b"DS4QWN01"; +const CHECKPOINT_VERSION: u32 = 1; +const CHECKPOINT_CHUNK: usize = 8 * 1024 * 1024; + +#[derive(Clone, Copy)] +struct Weight<'a> { + tensor: &'a QwenTensor, + expert: Option, +} + +struct Affine<'a> { + packed: Weight<'a>, + scales: Weight<'a>, + biases: Weight<'a>, + bits: u32, + group: u32, +} + +enum LayerState { + Gdn { conv: Buffer, recurrent: Buffer }, + Attention { kv: Buffer }, +} + +struct Scratch { + hidden: Buffer, + hc: Buffer, + hc_norm: Buffer, + hc_mix: Buffer, + rank: Buffer, + block: Buffer, + injection: Buffer, + qkv: Buffer, + controls: Buffer, + gdn_raw: Buffer, + gdn_out: Buffer, + router: Buffer, + route_ids: Buffer, + route_weights: Buffer, + gate: Buffer, + up: Buffer, + mid: Buffer, + expert: Buffer, + moe: Buffer, + shared: Buffer, + q_packed: Buffer, + q: Buffer, + q_gate: Buffer, + k: Buffer, + k_rope: Buffer, + v: Buffer, + attention: Buffer, + logits: Buffer, +} + +impl Scratch { + fn new() -> Result { + Ok(Self { + hidden: Buffer::floats(HIDDEN.into())?, + hc: Buffer::floats(HC_WIDTH.into())?, + hc_norm: Buffer::floats(HC_WIDTH.into())?, + hc_mix: Buffer::floats(HC_WIDTH.into())?, + rank: Buffer::floats(HC_RANK.into())?, + block: Buffer::floats(HIDDEN.into())?, + injection: Buffer::floats(HC.into())?, + qkv: Buffer::floats(GDN_QKV.into())?, + controls: Buffer::floats(GDN_CONTROLS.into())?, + gdn_raw: Buffer::floats(GDN_VALUE.into())?, + gdn_out: Buffer::floats(GDN_VALUE.into())?, + router: Buffer::floats(EXPERTS.into())?, + route_ids: Buffer::bytes((EXPERTS_USED * 4) as u64)?, + route_weights: Buffer::floats(EXPERTS_USED as u64)?, + gate: Buffer::floats(EXPERT_WIDTH.into())?, + up: Buffer::floats(EXPERT_WIDTH.into())?, + mid: Buffer::floats(EXPERT_WIDTH.into())?, + expert: Buffer::floats(HIDDEN.into())?, + moe: Buffer::floats(HIDDEN.into())?, + shared: Buffer::floats(HIDDEN.into())?, + q_packed: Buffer::floats((ATTN_WIDTH * 2).into())?, + q: Buffer::floats(ATTN_WIDTH.into())?, + q_gate: Buffer::floats(ATTN_WIDTH.into())?, + k: Buffer::floats(ATTN_KV_WIDTH.into())?, + k_rope: Buffer::floats(ATTN_KV_WIDTH.into())?, + v: Buffer::floats(ATTN_KV_WIDTH.into())?, + attention: Buffer::floats(ATTN_WIDTH.into())?, + logits: Buffer::floats(VOCAB.into())?, + }) + } +} + +pub(in crate::engine) struct QwenExecutor { + model: QwenModel, + states: Vec, + scratch: Scratch, + logits: Vec, + tokens: Vec, + position: u32, + context: u32, + checkpoint_tag: [u8; 32], + _context: Context, +} + +pub(in crate::engine) struct QwenResidentState { + states: Vec, + logits: Vec, + tokens: Vec, + position: u32, + checkpoint_tag: [u8; 32], +} + +impl QwenExecutor { + pub(super) fn open(model: QwenModel, context: u32) -> Result { + let native = Context::open_qwen(model.memory().admission)?; + let states = allocate_states(context)?; + Ok(Self { + model, + states, + scratch: Scratch::new()?, + logits: vec![0.0; VOCAB as usize], + tokens: Vec::new(), + position: 0, + context, + checkpoint_tag: [0; 32], + _context: native, + }) + } + + pub(super) fn eval(&mut self, token: i32) -> Result<(), String> { + if token < 0 || token as u32 >= VOCAB { + return Err(format!("token {token} is outside the Qwen vocabulary")); + } + if self.position >= self.context { + return Err(format!( + "the Qwen executor supports {} tokens per session", + self.context + )); + } + self.require_ple()?; + + self.begin_token(token)?; + + for layer in 0..LAYERS { + self.encode_layer(layer)?; + } + self.final_output()?; + self.tokens.push(token); + self.position += 1; + Ok(()) + } + + fn begin_token(&self, token: i32) -> Result<(), String> { + let embedding = self.affine("language_model.model.embed_tokens", HIDDEN, VOCAB, None)?; + let mut args = args(); + args.u[0] = HIDDEN; + args.u[2] = embedding.bits; + args.u[3] = embedding.group; + args.u[4] = token as u32; + let commands = Commands::begin()?; + self.dispatch( + c"kernel_qwen_affine_embedding", + &self.scratch.hidden, + None, + None, + None, + &[ + self.view(embedding.packed), + self.view(embedding.scales), + self.view(embedding.biases), + ], + &args, + HIDDEN, + 1, + )?; + args.u[0] = HIDDEN; + self.dispatch( + c"kernel_qwen_repeat4", + &self.scratch.hc, + Some(&self.scratch.hidden), + None, + None, + &[], + &args, + HC_WIDTH, + 1, + )?; + 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 encode_layer(&mut self, layer: usize) -> Result<(), String> { + let prefix = format!("language_model.model.layers.{layer}"); + let commands = Commands::begin()?; + self.hyper_read(&format!("{prefix}.attn_hyper_connection"))?; + match &self.states[layer] { + LayerState::Gdn { .. } => self.gdn(&prefix, layer)?, + LayerState::Attention { .. } => self.attention(&prefix, layer)?, + } + self.hyper_write()?; + self.hyper_read(&format!("{prefix}.mlp_hyper_connection"))?; + self.affine_mv_into( + &self.affine(&format!("{prefix}.mlp.gate"), HIDDEN, EXPERTS, None)?, + &self.scratch.block, + &self.scratch.router, + HIDDEN, + EXPERTS, + )?; + let mut route_args = args(); + route_args.u[0] = EXPERTS; + self.dispatch( + c"kernel_qwen_route_top10", + &self.scratch.route_ids, + Some(&self.scratch.route_weights), + Some(&self.scratch.router), + None, + &[], + &route_args, + 1, + 1, + )?; + commands.finish()?; + + let mut ids = [0_i32; EXPERTS_USED]; + let mut weights = [0.0_f32; EXPERTS_USED]; + self.scratch.route_ids.read_i32(&mut ids)?; + self.scratch.route_weights.read_f32(&mut weights)?; + self.scratch.moe.fill(0.0, HIDDEN.into())?; + 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()); + } + self.expert(&prefix, expert as u32, weight)?; + } + self.shared_expert(&prefix)?; + let mut inject_args = args(); + inject_args.u[0] = HIDDEN; + self.dispatch( + c"kernel_qwen_hyper_inject", + &self.scratch.hc_norm, + Some(&self.scratch.hc), + Some(&self.scratch.moe), + Some(&self.scratch.injection), + &[], + &inject_args, + HC_WIDTH, + 1, + )?; + self.scratch.hc.copy_from( + 0, + &self.scratch.hc_norm, + 0, + u64::from(HC_WIDTH) * 4, + "committing Qwen MoE hyper streams", + )?; + commands.finish() + } + + fn hyper_read(&self, prefix: &str) -> Result<(), String> { + let norm = self.weight(&format!("{prefix}.hc_norm.weight"))?; + let down = self.weight(&format!("{prefix}.input_mix_weight_down.weight"))?; + let up = self.weight(&format!("{prefix}.input_mix_weight_up.weight"))?; + let inject = self.weight(&format!("{prefix}.block_inject_weight.weight"))?; + let mut rms = args(); + rms.u[0] = HC_WIDTH; + rms.u[1] = HIDDEN; + rms.f[0] = 1.0e-6; + self.dispatch( + c"kernel_qwen_zero_rms", + &self.scratch.hc_norm, + Some(&self.scratch.hc), + None, + None, + &[self.view(norm)], + &rms, + HC, + 1, + )?; + self.bf16_mv( + down, + &self.scratch.hc_norm, + &self.scratch.rank, + HC_WIDTH, + HC_RANK, + )?; + let mut unary = args(); + unary.u[0] = HC_RANK; + self.dispatch( + c"kernel_qwen_silu_div4", + &self.scratch.rank, + Some(&self.scratch.rank), + None, + None, + &[], + &unary, + HC_RANK, + 1, + )?; + self.bf16_mv( + up, + &self.scratch.rank, + &self.scratch.hc_mix, + HC_RANK, + HC_WIDTH, + )?; + unary.u[0] = HC_WIDTH; + self.dispatch( + c"kernel_qwen_sigmoid", + &self.scratch.hc_mix, + Some(&self.scratch.hc_mix), + None, + None, + &[], + &unary, + HC_WIDTH, + 1, + )?; + let mut mix = args(); + mix.u[0] = HIDDEN; + self.dispatch( + c"kernel_qwen_hyper_mix", + &self.scratch.block, + Some(&self.scratch.hc_norm), + Some(&self.scratch.hc_mix), + None, + &[], + &mix, + HIDDEN, + 1, + )?; + self.bf16_mv( + inject, + &self.scratch.hc_norm, + &self.scratch.injection, + HC_WIDTH, + HC, + )?; + unary.u[0] = HC; + self.dispatch( + c"kernel_qwen_sigmoid2_div4", + &self.scratch.injection, + Some(&self.scratch.injection), + None, + None, + &[], + &unary, + HC, + 1, + ) + } + + fn hyper_write(&self) -> Result<(), String> { + let mut values = args(); + values.u[0] = HIDDEN; + self.dispatch( + c"kernel_qwen_hyper_inject", + &self.scratch.hc_norm, + Some(&self.scratch.hc), + Some(&self.scratch.hidden), + Some(&self.scratch.injection), + &[], + &values, + HC_WIDTH, + 1, + )?; + self.scratch.hc.copy_from( + 0, + &self.scratch.hc_norm, + 0, + u64::from(HC_WIDTH) * 4, + "committing Qwen attention hyper streams", + ) + } + + fn gdn(&self, prefix: &str, layer: usize) -> Result<(), String> { + let LayerState::Gdn { conv, recurrent } = &self.states[layer] else { + return Err("Qwen GDN graph received attention state".into()); + }; + self.affine_mv_into( + &self.affine( + &format!("{prefix}.linear_attn.in_proj_qkv"), + HIDDEN, + GDN_QKV, + None, + )?, + &self.scratch.block, + &self.scratch.qkv, + HIDDEN, + GDN_QKV, + )?; + for (name, offset, width) in [ + ("in_proj_z", 0_u64, GDN_VALUE), + ("in_proj_b", u64::from(GDN_VALUE), GDN_HEADS_V), + ("in_proj_a", u64::from(GDN_VALUE + GDN_HEADS_V), GDN_HEADS_V), + ] { + let target = self + .scratch + .controls + .view(offset * 4, u64::from(width) * 4)?; + self.affine_mv_into( + &self.affine(&format!("{prefix}.linear_attn.{name}"), HIDDEN, width, None)?, + &self.scratch.block, + &target, + HIDDEN, + width, + )?; + } + let mut conv_args = args(); + conv_args.u[0] = GDN_QKV; + self.dispatch( + c"kernel_qwen_conv_silu", + &self.scratch.qkv, + Some(&self.scratch.qkv), + Some(conv), + None, + &[self.view(self.weight(&format!("{prefix}.linear_attn.conv1d.weight"))?)], + &conv_args, + GDN_QKV, + 1, + )?; + let mut step = args(); + step.u[0] = HEAD_DIM; + step.u[1] = GDN_HEADS_K; + step.u[2] = GDN_HEADS_V; + step.u[3] = GDN_VALUE; + step.u[4] = GDN_VALUE + GDN_HEADS_V; + step.f[0] = 1.0e-6; + self.dispatch( + c"kernel_qwen_gdn_step", + &self.scratch.gdn_raw, + Some(&self.scratch.gdn_out), + Some(&self.scratch.controls), + Some(recurrent), + &[ + self.view(self.weight(&format!("{prefix}.linear_attn.A_log"))?), + self.view(self.weight(&format!("{prefix}.linear_attn.dt_bias"))?), + ], + &step, + HEAD_DIM, + GDN_HEADS_V, + )?; + let mut gate = args(); + gate.u[0] = HEAD_DIM; + gate.u[1] = GDN_HEADS_V; + gate.f[0] = 1.0e-6; + self.dispatch( + c"kernel_qwen_gdn_norm_gate", + &self.scratch.gdn_out, + Some(&self.scratch.gdn_raw), + Some(&self.scratch.controls), + None, + &[self.view(self.weight(&format!("{prefix}.linear_attn.norm.weight"))?)], + &gate, + GDN_HEADS_V, + 1, + )?; + self.affine_mv_into( + &self.affine( + &format!("{prefix}.linear_attn.out_proj"), + GDN_VALUE, + HIDDEN, + None, + )?, + &self.scratch.gdn_out, + &self.scratch.hidden, + GDN_VALUE, + HIDDEN, + ) + } + + fn attention(&self, prefix: &str, layer: usize) -> Result<(), String> { + if self.position + 1 > DENSE_BUDGET { + return Err( + "Qwen sparse QSA selection is required beyond 2048 tokens; native QSA belongs to issue #97" + .into(), + ); + } + let LayerState::Attention { kv } = &self.states[layer] else { + return Err("Qwen attention graph received GDN state".into()); + }; + self.affine_mv_into( + &self.affine( + &format!("{prefix}.self_attn.q_proj"), + HIDDEN, + ATTN_WIDTH * 2, + None, + )?, + &self.scratch.block, + &self.scratch.q_packed, + HIDDEN, + ATTN_WIDTH * 2, + )?; + self.affine_mv_into( + &self.affine( + &format!("{prefix}.self_attn.k_proj"), + HIDDEN, + ATTN_KV_WIDTH, + None, + )?, + &self.scratch.block, + &self.scratch.k, + HIDDEN, + ATTN_KV_WIDTH, + )?; + self.affine_mv_into( + &self.affine( + &format!("{prefix}.self_attn.v_proj"), + HIDDEN, + ATTN_KV_WIDTH, + None, + )?, + &self.scratch.block, + &self.scratch.v, + HIDDEN, + ATTN_KV_WIDTH, + )?; + let mut split = args(); + split.u[0] = ATTN_HEADS; + split.u[1] = ATTN_DIM; + self.dispatch( + c"kernel_qwen_split_q_gate", + &self.scratch.q, + Some(&self.scratch.q_packed), + Some(&self.scratch.q_gate), + None, + &[], + &split, + ATTN_WIDTH, + 1, + )?; + self.head_norm_rope( + &self.scratch.q, + &self.scratch.attention, + self.weight(&format!("{prefix}.self_attn.q_norm.weight"))?, + ATTN_HEADS, + )?; + self.head_norm_rope( + &self.scratch.k, + &self.scratch.k_rope, + self.weight(&format!("{prefix}.self_attn.k_norm.weight"))?, + ATTN_KV_HEADS, + )?; + let mut store = args(); + store.u[0] = ATTN_KV_WIDTH; + store.u[1] = self.position; + self.dispatch( + c"kernel_qwen_store_kv_bf16", + kv, + Some(&self.scratch.k_rope), + Some(&self.scratch.v), + None, + &[], + &store, + ATTN_KV_WIDTH, + 1, + )?; + let mut dense = args(); + dense.u[0] = ATTN_HEADS; + dense.u[1] = ATTN_KV_HEADS; + dense.u[2] = ATTN_DIM; + dense.u[3] = self.position + 1; + self.dispatch( + c"kernel_qwen_dense_attention", + &self.scratch.q, + Some(&self.scratch.attention), + Some(kv), + None, + &[], + &dense, + ATTN_HEADS, + 1, + )?; + let mut gate = args(); + gate.u[0] = ATTN_WIDTH; + self.dispatch( + c"kernel_qwen_gate_attention", + &self.scratch.attention, + Some(&self.scratch.q), + Some(&self.scratch.q_gate), + None, + &[], + &gate, + ATTN_WIDTH, + 1, + )?; + self.affine_mv_into( + &self.affine( + &format!("{prefix}.self_attn.o_proj"), + ATTN_WIDTH, + HIDDEN, + None, + )?, + &self.scratch.attention, + &self.scratch.hidden, + ATTN_WIDTH, + HIDDEN, + ) + } + + fn head_norm_rope( + &self, + input: &Buffer, + output: &Buffer, + weight: Weight<'_>, + heads: u32, + ) -> Result<(), String> { + let mut values = args(); + values.u[0] = ATTN_DIM; + values.u[1] = 64; + values.u[2] = heads; + values.u[3] = self.position; + values.f[0] = 1.0e-6; + values.f[1] = 10_000_000.0; + self.dispatch( + c"kernel_qwen_head_norm_rope", + output, + Some(input), + None, + None, + &[self.view(weight)], + &values, + ATTN_DIM, + heads, + ) + } + + fn expert(&self, prefix: &str, expert: u32, weight: f32) -> Result<(), String> { + for (name, target) in [ + ("gate_proj", &self.scratch.gate), + ("up_proj", &self.scratch.up), + ] { + self.affine_mv_into( + &self.affine( + &format!("{prefix}.mlp.switch_mlp.{name}"), + HIDDEN, + EXPERT_WIDTH, + Some(expert), + )?, + &self.scratch.block, + target, + HIDDEN, + EXPERT_WIDTH, + )?; + } + let mut swiglu = args(); + swiglu.u[0] = EXPERT_WIDTH; + self.dispatch( + c"kernel_qwen_swiglu", + &self.scratch.mid, + Some(&self.scratch.gate), + Some(&self.scratch.up), + None, + &[], + &swiglu, + EXPERT_WIDTH, + 1, + )?; + self.affine_mv_into( + &self.affine( + &format!("{prefix}.mlp.switch_mlp.down_proj"), + EXPERT_WIDTH, + HIDDEN, + Some(expert), + )?, + &self.scratch.mid, + &self.scratch.expert, + EXPERT_WIDTH, + HIDDEN, + )?; + let mut accumulate = args(); + accumulate.u[0] = HIDDEN; + accumulate.f[0] = weight; + self.dispatch( + c"kernel_qwen_accumulate", + &self.scratch.moe, + Some(&self.scratch.expert), + None, + None, + &[], + &accumulate, + HIDDEN, + 1, + ) + } + + fn shared_expert(&self, prefix: &str) -> Result<(), String> { + for (name, target) in [ + ("gate_proj", &self.scratch.gate), + ("up_proj", &self.scratch.up), + ] { + self.affine_mv_into( + &self.affine( + &format!("{prefix}.mlp.shared_expert.{name}"), + HIDDEN, + EXPERT_WIDTH, + None, + )?, + &self.scratch.block, + target, + HIDDEN, + EXPERT_WIDTH, + )?; + } + let mut swiglu = args(); + swiglu.u[0] = EXPERT_WIDTH; + self.dispatch( + c"kernel_qwen_swiglu", + &self.scratch.mid, + Some(&self.scratch.gate), + Some(&self.scratch.up), + None, + &[], + &swiglu, + EXPERT_WIDTH, + 1, + )?; + self.affine_mv_into( + &self.affine( + &format!("{prefix}.mlp.shared_expert.down_proj"), + EXPERT_WIDTH, + HIDDEN, + None, + )?, + &self.scratch.mid, + &self.scratch.shared, + EXPERT_WIDTH, + HIDDEN, + )?; + self.affine_mv_into( + &self.affine(&format!("{prefix}.mlp.shared_expert_gate"), HIDDEN, 1, None)?, + &self.scratch.block, + &self.scratch.gate, + HIDDEN, + 1, + )?; + let mut accumulate = args(); + accumulate.u[0] = HIDDEN; + self.dispatch( + c"kernel_qwen_accumulate_sigmoid_scalar", + &self.scratch.moe, + Some(&self.scratch.shared), + Some(&self.scratch.gate), + None, + &[], + &accumulate, + HIDDEN, + 1, + ) + } + + fn final_output(&mut self) -> Result<(), String> { + let commands = Commands::begin()?; + self.final_mix()?; + self.affine_mv_into( + &self.affine("language_model.lm_head", HIDDEN, VOCAB, None)?, + &self.scratch.block, + &self.scratch.logits, + HIDDEN, + VOCAB, + )?; + commands.finish()?; + self.scratch.logits.read_f32(&mut self.logits) + } + + fn final_mix(&self) -> Result<(), String> { + let prefix = "language_model.model.hyper_connection_mixer"; + let norm = self.weight(&format!("{prefix}.hc_norm.weight"))?; + let down = self.weight(&format!("{prefix}.input_mix_weight_down.weight"))?; + let up = self.weight(&format!("{prefix}.input_mix_weight_up.weight"))?; + let mut rms = args(); + rms.u[0] = HC_WIDTH; + rms.u[1] = HIDDEN; + rms.f[0] = 1.0e-6; + self.dispatch( + c"kernel_qwen_zero_rms", + &self.scratch.hc_norm, + Some(&self.scratch.hc), + None, + None, + &[self.view(norm)], + &rms, + HC, + 1, + )?; + self.bf16_mv( + down, + &self.scratch.hc_norm, + &self.scratch.rank, + HC_WIDTH, + HC_RANK, + )?; + let mut unary = args(); + unary.u[0] = HC_RANK; + self.dispatch( + c"kernel_qwen_silu_div4", + &self.scratch.rank, + Some(&self.scratch.rank), + None, + None, + &[], + &unary, + HC_RANK, + 1, + )?; + self.bf16_mv( + up, + &self.scratch.rank, + &self.scratch.hc_mix, + HC_RANK, + HC_WIDTH, + )?; + unary.u[0] = HC_WIDTH; + self.dispatch( + c"kernel_qwen_sigmoid", + &self.scratch.hc_mix, + Some(&self.scratch.hc_mix), + None, + None, + &[], + &unary, + HC_WIDTH, + 1, + )?; + let mut mix = args(); + mix.u[0] = HIDDEN; + self.dispatch( + c"kernel_qwen_hyper_mix", + &self.scratch.block, + Some(&self.scratch.hc_norm), + Some(&self.scratch.hc_mix), + None, + &[], + &mix, + HIDDEN, + 1, + ) + } + + fn affine( + &self, + prefix: &str, + in_dim: u32, + out_dim: u32, + expert: Option, + ) -> Result, String> { + let packed = self.weight(&format!("{prefix}.weight"))?; + let scales = self.weight(&format!("{prefix}.scales"))?; + let biases = self.weight(&format!("{prefix}.biases"))?; + let bits = packed + .tensor + .quant_bits + .ok_or_else(|| format!("{prefix} is not affine quantized"))?; + let group = packed + .tensor + .group_size + .ok_or_else(|| format!("{prefix} has no affine group size"))? + as u32; + if !matches!(bits, 2 | 4 | 8) || !in_dim.is_multiple_of(group) { + return Err(format!("{prefix} has an incompatible affine layout")); + } + let expected = if expert.is_some() { + vec![ + EXPERTS as u64, + out_dim as u64, + (in_dim / (32 / bits)) as u64, + ] + } else { + vec![out_dim as u64, (in_dim / (32 / bits)) as u64] + }; + if packed.tensor.shape != expected { + return Err(format!("{prefix} has an incompatible executor shape")); + } + let expected_parameters = if expert.is_some() { + vec![EXPERTS as u64, out_dim as u64, (in_dim / group) as u64] + } else { + vec![out_dim as u64, (in_dim / group) as u64] + }; + if packed.tensor.dtype != "U32" + || scales.tensor.dtype != "BF16" + || biases.tensor.dtype != "BF16" + || scales.tensor.shape != expected_parameters + || biases.tensor.shape != expected_parameters + { + return Err(format!("{prefix} has incompatible affine parameters")); + } + Ok(Affine { + packed: Weight { + tensor: packed.tensor, + expert, + }, + scales: Weight { + tensor: scales.tensor, + expert, + }, + biases: Weight { + tensor: biases.tensor, + expert, + }, + bits, + group, + }) + } + + fn weight(&self, name: &str) -> Result, String> { + self.model.tensor(name).map(|tensor| Weight { + tensor, + expert: None, + }) + } + + fn view(&self, weight: Weight<'_>) -> QwenWeightView { + let (map, _) = self.model.map(weight.tensor.map); + let experts = weight.expert.map_or(1, |_| EXPERTS as u64); + let bytes = (weight.tensor.range.end - weight.tensor.range.start) / experts; + let offset = weight.tensor.range.start + u64::from(weight.expert.unwrap_or(0)) * bytes; + QwenWeightView { + map: map.as_ptr().cast(), + size: map.len() as u64, + offset, + bytes, + } + } + + fn affine_mv_into( + &self, + weight: &Affine<'_>, + input: &Buffer, + output: &Buffer, + in_dim: u32, + out_dim: u32, + ) -> Result<(), String> { + let mut values = args(); + values.u[0] = in_dim; + values.u[1] = out_dim; + values.u[2] = weight.bits; + values.u[3] = weight.group; + self.dispatch( + c"kernel_qwen_affine_mv", + output, + Some(input), + None, + None, + &[ + self.view(weight.packed), + self.view(weight.scales), + self.view(weight.biases), + ], + &values, + out_dim, + 1, + ) + } + + fn bf16_mv( + &self, + weight: Weight<'_>, + input: &Buffer, + output: &Buffer, + in_dim: u32, + out_dim: u32, + ) -> Result<(), String> { + if weight.tensor.dtype != "BF16" || weight.tensor.shape != [out_dim as u64, in_dim as u64] { + return Err(format!( + "{} has an incompatible BF16 matrix shape", + weight.tensor.name + )); + } + let mut values = args(); + values.u[0] = in_dim; + values.u[1] = out_dim; + self.dispatch( + c"kernel_qwen_bf16_mv", + output, + Some(input), + None, + None, + &[self.view(weight)], + &values, + out_dim, + 1, + ) + } + + #[allow(clippy::too_many_arguments)] + fn dispatch( + &self, + kernel: &CStr, + out: &Buffer, + a: Option<&Buffer>, + b: Option<&Buffer>, + c: Option<&Buffer>, + weights: &[QwenWeightView], + values: &QwenKernelArgs, + grid_x: u32, + grid_y: u32, + ) -> Result<(), String> { + dispatch_qwen(kernel, out, a, b, c, weights, values, grid_x, grid_y) + } + + pub(super) fn prefill( + &mut self, + tokens: &[i32], + mut progress: impl FnMut(u32) -> bool, + ) -> Result { + let mut completed = 0; + for &token in tokens { + if !progress(completed) { + break; + } + self.eval(token)?; + completed += 1; + } + progress(completed); + Ok(completed as usize) + } + + pub(super) fn logits(&self) -> &[f32] { + &self.logits + } + pub(super) fn execution_stats(&self) -> ExecutionStats { + ExecutionStats::default() + } + pub(super) fn model(&self) -> &QwenModel { + &self.model + } + pub(super) fn context(&self) -> u32 { + self.context + } + pub(super) fn position(&self) -> u32 { + self.position + } + pub(super) fn tokens(&self) -> &[i32] { + &self.tokens + } + pub(super) fn checkpoint_tag(&self) -> [u8; 32] { + self.checkpoint_tag + } + pub(super) fn note_checkpoint_tag(&mut self, tag: [u8; 32]) { + self.checkpoint_tag = tag; + } + + pub(super) fn reset(&mut self) -> Result<(), String> { + self.states = allocate_states(self.context)?; + self.logits.fill(0.0); + self.tokens.clear(); + self.position = 0; + self.checkpoint_tag = [0; 32]; + Ok(()) + } + + pub(super) fn align_prompt(&mut self, tokens: &[i32]) -> Result { + if !tokens.starts_with(&self.tokens) { + self.reset()?; + } + Ok(self.tokens.len()) + } + + fn blank_resident(&self) -> Result { + Ok(QwenResidentState { + states: allocate_states(self.context)?, + logits: vec![0.0; VOCAB as usize], + tokens: Vec::new(), + position: 0, + checkpoint_tag: [0; 32], + }) + } + + pub(super) fn swap_resident_state( + &mut self, + state: &mut Option, + ) -> 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.logits, &mut incoming.logits); + std::mem::swap(&mut self.tokens, &mut incoming.tokens); + std::mem::swap(&mut self.position, &mut incoming.position); + std::mem::swap(&mut self.checkpoint_tag, &mut incoming.checkpoint_tag); + *state = Some(incoming); + Ok(()) + } + + pub(in crate::engine) fn save_checkpoint( + &mut self, + path: &Path, + tag: [u8; 32], + progress: &mut impl FnMut(u64), + ) -> Result<(), String> { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent).map_err(|error| error.to_string())?; + } + let temporary = path.with_extension("tmp"); + let mut file = File::create(&temporary).map_err(|error| error.to_string())?; + file.write_all(CHECKPOINT_MAGIC) + .map_err(|error| error.to_string())?; + for value in [ + CHECKPOINT_VERSION, + self.context, + self.position, + VOCAB, + LAYERS as u32, + ] { + write_u32(&mut file, value)?; + } + 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.tokens { + write_u32(&mut file, token as u32)?; + } + for &logit in &self.logits { + write_u32(&mut file, logit.to_bits())?; + } + let mut chunk = vec![0; CHECKPOINT_CHUNK]; + for state in &self.states { + match state { + LayerState::Gdn { conv, recurrent } => { + write_buffer( + &mut file, + conv, + 0, + u64::from(GDN_QKV) * 3 * 2, + &mut chunk, + progress, + )?; + write_buffer( + &mut file, + recurrent, + 0, + u64::from(GDN_HEADS_V) * HEAD_DIM as u64 * HEAD_DIM as u64 * 4, + &mut chunk, + progress, + )?; + } + LayerState::Attention { kv } => { + write_buffer( + &mut file, + kv, + 0, + u64::from(self.position) * ATTN_KV_WIDTH as u64 * 4, + &mut chunk, + progress, + )?; + } + } + } + file.sync_all().map_err(|error| error.to_string())?; + fs::rename(temporary, path).map_err(|error| error.to_string())?; + self.checkpoint_tag = tag; + Ok(()) + } + + pub(in crate::engine) fn load_checkpoint( + &mut self, + path: &Path, + progress: &mut impl FnMut(u64), + ) -> Result { + let mut file = match File::open(path) { + Ok(file) => file, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false), + Err(error) => return Err(error.to_string()), + }; + let mut magic = [0; 8]; + file.read_exact(&mut magic) + .map_err(|error| error.to_string())?; + if &magic != CHECKPOINT_MAGIC { + return Err("Qwen checkpoint has an invalid signature".into()); + } + for expected in [CHECKPOINT_VERSION, self.context] { + if read_u32(&mut file)? != expected { + return Err("Qwen checkpoint does not match the current executor".into()); + } + } + let position = read_u32(&mut file)?; + if position > self.context + || read_u32(&mut file)? != VOCAB + || read_u32(&mut file)? != LAYERS as u32 + { + return Err("Qwen checkpoint shape is invalid".into()); + } + let mut identity = [0; 32]; + file.read_exact(&mut identity) + .map_err(|error| error.to_string())?; + if identity != self.model.checkpoint_identity() { + return Err("Qwen checkpoint model identity changed".into()); + } + let mut tag = [0; 32]; + file.read_exact(&mut tag) + .map_err(|error| error.to_string())?; + let mut tokens = Vec::with_capacity(position as usize); + for _ in 0..position { + let token = read_u32(&mut file)?; + if token >= VOCAB { + return Err("Qwen checkpoint token is invalid".into()); + } + tokens.push(token as i32); + } + let mut logits = Vec::with_capacity(VOCAB as usize); + for _ in 0..VOCAB { + logits.push(f32::from_bits(read_u32(&mut file)?)); + } + self.reset()?; + let mut chunk = vec![0; CHECKPOINT_CHUNK]; + for state in &self.states { + match state { + LayerState::Gdn { conv, recurrent } => { + read_buffer( + &mut file, + conv, + 0, + u64::from(GDN_QKV) * 3 * 2, + &mut chunk, + progress, + )?; + read_buffer( + &mut file, + recurrent, + 0, + u64::from(GDN_HEADS_V) * HEAD_DIM as u64 * HEAD_DIM as u64 * 4, + &mut chunk, + progress, + )?; + } + LayerState::Attention { kv } => read_buffer( + &mut file, + kv, + 0, + u64::from(position) * ATTN_KV_WIDTH as u64 * 4, + &mut chunk, + progress, + )?, + } + } + let mut trailing = [0]; + if file + .read(&mut trailing) + .map_err(|error| error.to_string())? + != 0 + { + self.reset()?; + return Err("Qwen checkpoint has trailing data".into()); + } + self.position = position; + self.tokens = tokens; + self.logits = logits; + self.checkpoint_tag = tag; + Ok(true) + } +} + +fn allocate_states(context: u32) -> Result, String> { + (0..LAYERS) + .map(|layer| { + if layer % 4 == 3 { + Ok(LayerState::Attention { + kv: Buffer::bytes(u64::from(context) * ATTN_KV_WIDTH as u64 * 4)?, + }) + } else { + let conv = Buffer::bytes(u64::from(GDN_QKV) * 3 * 2)?; + let recurrent = + Buffer::floats(u64::from(GDN_HEADS_V) * HEAD_DIM as u64 * HEAD_DIM as u64)?; + conv.fill(0.0, u64::from(GDN_QKV) * 3 / 2)?; + recurrent.fill( + 0.0, + u64::from(GDN_HEADS_V) * HEAD_DIM as u64 * HEAD_DIM as u64, + )?; + Ok(LayerState::Gdn { conv, recurrent }) + } + }) + .collect() +} + +fn args() -> QwenKernelArgs { + QwenKernelArgs::default() +} + +#[allow(clippy::too_many_arguments)] +fn dispatch_qwen( + kernel: &CStr, + out: &Buffer, + a: Option<&Buffer>, + b: Option<&Buffer>, + c: Option<&Buffer>, + weights: &[QwenWeightView], + values: &QwenKernelArgs, + grid_x: u32, + grid_y: u32, +) -> Result<(), String> { + call( + unsafe { + ds4_gpu_qwen_dispatch( + kernel.as_ptr(), + out.raw(), + a.map_or(std::ptr::null(), |buffer| buffer.raw().cast_const()), + b.map_or(std::ptr::null(), |buffer| buffer.raw().cast_const()), + c.map_or(std::ptr::null(), |buffer| buffer.raw().cast_const()), + weights.as_ptr(), + weights.len() as u32, + values, + grid_x, + grid_y, + ) + }, + kernel.to_str().unwrap_or("running a Qwen Metal kernel"), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use memmap2::MmapOptions; + use std::fs; + use std::path::PathBuf; + + fn bf16(value: f32) -> u16 { + let bits = value.to_bits(); + ((bits + 0x7fff + ((bits >> 16) & 1)) >> 16) as u16 + } + + fn close(actual: f32, expected: f32) { + assert!( + (actual - expected).abs() <= 2.0e-4 * expected.abs().max(1.0), + "{actual} != {expected}" + ); + } + + #[test] + #[ignore = "requires Apple Metal"] + fn qwen_metal_primitives_match_reference_vectors() { + configure_sources().unwrap(); + let _context = Context::open_qwen(0).unwrap(); + + let input = Buffer::floats(64).unwrap(); + 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-{}-{}", + std::process::id(), + std::thread::current().name().unwrap_or("test") + )); + let mut bytes = Vec::new(); + for _ in 0..8 { + bytes.extend_from_slice(&0x3333_3333_u32.to_le_bytes()); + } + for value in [ + 0.5, -1.0, 0.0, 0.0, 1.0, 2.0, 3.0, 4.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, + ] { + bytes.extend_from_slice(&bf16(value).to_le_bytes()); + } + fs::write(&path, bytes).unwrap(); + let file = File::open(&path).unwrap(); + // SAFETY: this test owns the read-only file for the lifetime of the mapping. + let map = unsafe { MmapOptions::new().map(&file).unwrap() }; + let view = |offset, size| QwenWeightView { + map: map.as_ptr().cast(), + size: map.len() as u64, + offset, + bytes: size, + }; + + let mut affine = args(); + affine.u[0] = 64; + affine.u[1] = 1; + affine.u[2] = 4; + affine.u[3] = 64; + dispatch_qwen( + c"kernel_qwen_affine_mv", + &output, + Some(&input), + None, + None, + &[view(0, 32), view(32, 2), view(34, 2)], + &affine, + 1, + 1, + ) + .unwrap(); + let mut scalar = [0.0]; + output.read_f32(&mut scalar).unwrap(); + close(scalar[0], 32.0); + + let conv_input = Buffer::floats(1).unwrap(); + conv_input.write_f32(&[2.0]).unwrap(); + let conv_state = Buffer::bytes(6).unwrap(); + conv_state.write(0, &[0; 6]).unwrap(); + let conv_output = Buffer::floats(1).unwrap(); + let mut conv = args(); + conv.u[0] = 1; + dispatch_qwen( + c"kernel_qwen_conv_silu", + &conv_output, + Some(&conv_input), + Some(&conv_state), + None, + &[view(40, 8)], + &conv, + 1, + 1, + ) + .unwrap(); + conv_output.read_f32(&mut scalar).unwrap(); + close(scalar[0], 8.0 / (1.0 + (-8.0_f32).exp())); + assert_eq!( + { + let mut state = [0; 6]; + conv_state.read(0, &mut state).unwrap(); + state + }, + [0, 0, 0, 0, 0, 64] + ); + + let qkv = Buffer::floats(6).unwrap(); + qkv.write_f32(&[3.0, 4.0, 0.0, 2.0, 5.0, 7.0]).unwrap(); + let controls = Buffer::floats(4).unwrap(); + controls.write_f32(&[0.0; 4]).unwrap(); + let recurrent = Buffer::floats(4).unwrap(); + recurrent.write_f32(&[1.0, 2.0, 3.0, 4.0]).unwrap(); + let raw = Buffer::floats(2).unwrap(); + let gated = Buffer::floats(2).unwrap(); + let mut step = args(); + step.u[0] = 2; + step.u[1] = 1; + step.u[2] = 1; + step.u[3] = 2; + step.u[4] = 3; + step.f[0] = 1.0e-6; + dispatch_qwen( + c"kernel_qwen_gdn_step", + &raw, + Some(&qkv), + Some(&controls), + Some(&recurrent), + &[view(36, 2), view(38, 2)], + &step, + 2, + 1, + ) + .unwrap(); + let mut raw_values = [0.0; 2]; + raw.read_f32(&mut raw_values).unwrap(); + close(raw_values[0], 2.0506096); + close(raw_values[1], 2.9698484); + let mut norm = args(); + norm.u[0] = 2; + norm.u[1] = 1; + norm.f[0] = 1.0e-6; + dispatch_qwen( + c"kernel_qwen_gdn_norm_gate", + &gated, + Some(&raw), + Some(&controls), + None, + &[view(48, 4)], + &norm, + 1, + 1, + ) + .unwrap(); + let mut gated_values = [0.0; 2]; + gated.read_f32(&mut gated_values).unwrap(); + let rms = ((raw_values[0].powi(2) + raw_values[1].powi(2)) * 0.5 + 1.0e-6).sqrt(); + close(gated_values[0], raw_values[0] / rms * 0.5); + close(gated_values[1], raw_values[1] / rms * 0.5); + + let router = Buffer::floats(EXPERTS.into()).unwrap(); + let mut logits = vec![-100.0; EXPERTS as usize]; + for (index, value) in logits.iter_mut().take(EXPERTS_USED).enumerate() { + *value = index as f32; + } + router.write_f32(&logits).unwrap(); + let ids = Buffer::bytes((EXPERTS_USED * 4) as u64).unwrap(); + let weights = Buffer::floats(EXPERTS_USED as u64).unwrap(); + let mut route = args(); + route.u[0] = EXPERTS; + dispatch_qwen( + c"kernel_qwen_route_top10", + &ids, + Some(&weights), + Some(&router), + None, + &[], + &route, + 1, + 1, + ) + .unwrap(); + let mut selected = [0; EXPERTS_USED]; + ids.read_i32(&mut selected).unwrap(); + assert_eq!(selected, [9, 8, 7, 6, 5, 4, 3, 2, 1, 0]); + let mut probabilities = [0.0; EXPERTS_USED]; + weights.read_f32(&mut probabilities).unwrap(); + close(probabilities.iter().sum(), 1.0); + + let norm_input = Buffer::floats(4).unwrap(); + norm_input.write_f32(&[1.0, 2.0, 3.0, 4.0]).unwrap(); + let norm_output = Buffer::floats(4).unwrap(); + let mut zero_norm = args(); + zero_norm.u[0] = 4; + zero_norm.u[1] = 2; + zero_norm.f[0] = 1.0e-6; + dispatch_qwen( + c"kernel_qwen_zero_rms", + &norm_output, + Some(&norm_input), + None, + None, + &[view(52, 8)], + &zero_norm, + 2, + 1, + ) + .unwrap(); + let mut normalized = [0.0; 4]; + norm_output.read_f32(&mut normalized).unwrap(); + for (actual, expected) in normalized.into_iter().zip([ + 1.0 / (2.5_f32 + 1.0e-6).sqrt(), + 2.0 / (2.5_f32 + 1.0e-6).sqrt(), + 3.0 / (12.5_f32 + 1.0e-6).sqrt(), + 4.0 / (12.5_f32 + 1.0e-6).sqrt(), + ]) { + close(actual, expected); + } + + let hyper_input = Buffer::floats(8).unwrap(); + hyper_input + .write_f32(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]) + .unwrap(); + let hyper_weights = Buffer::floats(8).unwrap(); + hyper_weights.write_f32(&[1.0; 8]).unwrap(); + let mixed = Buffer::floats(2).unwrap(); + let mut hyper = args(); + hyper.u[0] = 2; + dispatch_qwen( + c"kernel_qwen_hyper_mix", + &mixed, + Some(&hyper_input), + Some(&hyper_weights), + None, + &[], + &hyper, + 2, + 1, + ) + .unwrap(); + let mut mixed_values = [0.0; 2]; + mixed.read_f32(&mut mixed_values).unwrap(); + assert_eq!(mixed_values, [4.0, 5.0]); + let injection = Buffer::floats(4).unwrap(); + injection.write_f32(&[1.0, 2.0, 3.0, 4.0]).unwrap(); + let combined = Buffer::floats(8).unwrap(); + dispatch_qwen( + c"kernel_qwen_hyper_inject", + &combined, + Some(&hyper_input), + Some(&mixed), + Some(&injection), + &[], + &hyper, + 8, + 1, + ) + .unwrap(); + let mut combined_values = [0.0; 8]; + combined.read_f32(&mut combined_values).unwrap(); + assert_eq!( + combined_values, + [5.0, 7.0, 11.0, 14.0, 17.0, 21.0, 23.0, 28.0] + ); + + let expert_gate = Buffer::floats(2).unwrap(); + expert_gate.write_f32(&[0.0, 2.0]).unwrap(); + let expert_up = Buffer::floats(2).unwrap(); + expert_up.write_f32(&[4.0, 3.0]).unwrap(); + let expert = Buffer::floats(2).unwrap(); + let mut moe = args(); + moe.u[0] = 2; + dispatch_qwen( + c"kernel_qwen_swiglu", + &expert, + Some(&expert_gate), + Some(&expert_up), + None, + &[], + &moe, + 2, + 1, + ) + .unwrap(); + let accumulated = Buffer::floats(2).unwrap(); + accumulated.fill(0.0, 2).unwrap(); + moe.f[0] = 0.25; + dispatch_qwen( + c"kernel_qwen_accumulate", + &accumulated, + Some(&expert), + None, + None, + &[], + &moe, + 2, + 1, + ) + .unwrap(); + let shared = Buffer::floats(2).unwrap(); + shared.write_f32(&[2.0, 4.0]).unwrap(); + let shared_gate = Buffer::floats(1).unwrap(); + shared_gate.write_f32(&[0.0]).unwrap(); + dispatch_qwen( + c"kernel_qwen_accumulate_sigmoid_scalar", + &accumulated, + Some(&shared), + Some(&shared_gate), + None, + &[], + &moe, + 2, + 1, + ) + .unwrap(); + let mut moe_values = [0.0; 2]; + accumulated.read_f32(&mut moe_values).unwrap(); + close(moe_values[0], 1.0); + close( + moe_values[1], + 0.25 * (2.0 / (1.0 + (-2.0_f32).exp()) * 3.0) + 2.0, + ); + + let rope_input = Buffer::floats(4).unwrap(); + rope_input.write_f32(&[1.0, 2.0, 3.0, 4.0]).unwrap(); + let rope_output = Buffer::floats(4).unwrap(); + let mut rope = args(); + rope.u[0] = 4; + rope.u[1] = 4; + rope.u[2] = 1; + rope.u[3] = 1; + rope.f[0] = 1.0e-6; + rope.f[1] = 10_000.0; + dispatch_qwen( + c"kernel_qwen_head_norm_rope", + &rope_output, + Some(&rope_input), + None, + None, + &[view(52, 8)], + &rope, + 4, + 1, + ) + .unwrap(); + let rms = (7.5_f32 + 1.0e-6).sqrt(); + let theta = [1.0_f32, 0.01]; + let normalized = [1.0 / rms, 2.0 / rms, 3.0 / rms, 4.0 / rms]; + let expected_rope = [ + normalized[0] * theta[0].cos() - normalized[2] * theta[0].sin(), + normalized[1] * theta[1].cos() - normalized[3] * theta[1].sin(), + normalized[2] * theta[0].cos() + normalized[0] * theta[0].sin(), + normalized[3] * theta[1].cos() + normalized[1] * theta[1].sin(), + ]; + let mut actual_rope = [0.0; 4]; + rope_output.read_f32(&mut actual_rope).unwrap(); + for (actual, expected) in actual_rope.into_iter().zip(expected_rope) { + close(actual, expected); + } + + let cache = Buffer::bytes(16).unwrap(); + let key = Buffer::floats(2).unwrap(); + let value = Buffer::floats(2).unwrap(); + let mut store = args(); + store.u[0] = 2; + for (position, (keys, values)) in [([1.0, 0.0], [2.0, 3.0]), ([0.0, 1.0], [5.0, 7.0])] + .into_iter() + .enumerate() + { + key.write_f32(&keys).unwrap(); + value.write_f32(&values).unwrap(); + store.u[1] = position as u32; + dispatch_qwen( + c"kernel_qwen_store_kv_bf16", + &cache, + Some(&key), + Some(&value), + None, + &[], + &store, + 2, + 1, + ) + .unwrap(); + } + let query = Buffer::floats(2).unwrap(); + query.write_f32(&[1.0, 0.0]).unwrap(); + let attention = Buffer::floats(2).unwrap(); + let mut dense = args(); + dense.u[0] = 1; + dense.u[1] = 1; + dense.u[2] = 2; + dense.u[3] = 2; + dispatch_qwen( + c"kernel_qwen_dense_attention", + &attention, + Some(&query), + Some(&cache), + None, + &[], + &dense, + 1, + 1, + ) + .unwrap(); + let first = (1.0_f32 / 2.0_f32.sqrt()).exp(); + let probability = first / (first + 1.0); + let mut actual_attention = [0.0; 2]; + attention.read_f32(&mut actual_attention).unwrap(); + close( + actual_attention[0], + probability * 2.0 + (1.0 - probability) * 5.0, + ); + close( + actual_attention[1], + probability * 3.0 + (1.0 - probability) * 7.0, + ); + + drop(map); + drop(file); + fs::remove_file(path).unwrap(); + } + + #[test] + #[ignore = "requires the pinned 105 GB Qwen artifact set and Apple Metal"] + fn qwen_core_boundary_and_checkpoint_are_stable() { + configure_sources().unwrap(); + let root = std::env::var_os("DS4SERVER_QWEN38_SOURCE") + .map(PathBuf::from) + .expect("set DS4SERVER_QWEN38_SOURCE to the pinned artifact directory"); + let model = QwenModel::open(&root, 4).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(); + assert!(executor.logits.iter().all(|value| value.is_finite())); + let reference = [ + executor.logits[0], + executor.logits[1], + executor.logits[1000], + executor.logits[VOCAB as usize - 1], + ]; + for (actual, expected) in + reference + .into_iter() + .zip([1.483_976_1, 0.575_980_66, -0.008_828_48, 0.047_039_207]) + { + close(actual, expected); + } + + executor.position = 1; + executor.tokens = vec![1]; + let checkpoint = std::env::temp_dir().join(format!( + "ds4-qwen95-checkpoint-{}-{}", + std::process::id(), + std::thread::current().name().unwrap_or("test") + )); + executor + .save_checkpoint(&checkpoint, [7; 32], &mut |_| {}) + .unwrap(); + + executor.begin_token(2).unwrap(); + executor.encode_layer(0).unwrap(); + executor.final_output().unwrap(); + let continued = [ + executor.logits[0], + executor.logits[1], + executor.logits[1000], + executor.logits[VOCAB as usize - 1], + ]; + let LayerState::Gdn { conv, recurrent } = &executor.states[0] else { + unreachable!() + }; + let mut continued_conv = vec![0; GDN_QKV as usize * 3 * 2]; + conv.read(0, &mut continued_conv).unwrap(); + let mut continued_recurrent = vec![0; 4 * 1024]; + recurrent.read(0, &mut continued_recurrent).unwrap(); + + assert!(executor.load_checkpoint(&checkpoint, &mut |_| {}).unwrap()); + assert_eq!(executor.position, 1); + assert_eq!(executor.tokens, [1]); + assert_eq!(executor.checkpoint_tag, [7; 32]); + assert_eq!( + [ + executor.logits[0], + executor.logits[1], + executor.logits[1000], + executor.logits[VOCAB as usize - 1], + ], + reference + ); + executor.begin_token(2).unwrap(); + executor.encode_layer(0).unwrap(); + executor.final_output().unwrap(); + for (actual, expected) in [ + executor.logits[0], + executor.logits[1], + executor.logits[1000], + executor.logits[VOCAB as usize - 1], + ] + .into_iter() + .zip(continued) + { + close(actual, expected); + } + let LayerState::Gdn { conv, recurrent } = &executor.states[0] else { + unreachable!() + }; + let mut resumed_conv = vec![0; continued_conv.len()]; + conv.read(0, &mut resumed_conv).unwrap(); + assert_eq!(resumed_conv, continued_conv); + let mut resumed_recurrent = vec![0; continued_recurrent.len()]; + recurrent.read(0, &mut resumed_recurrent).unwrap(); + assert_eq!(resumed_recurrent, continued_recurrent); + + executor.position = DENSE_BUDGET; + let error = executor + .attention("language_model.model.layers.3", 3) + .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 2681ab1..9ecb43b 100644 --- a/src/engine/qwen.rs +++ b/src/engine/qwen.rs @@ -1,8 +1,13 @@ use super::tokenizer::Tokenizer; +use super::{ChatTurn, ModelSummary}; +use crate::model::ModelChoice; +use crate::settings::ReasoningMode; +use memmap2::{Mmap, MmapOptions}; use serde::de::{MapAccess, Visitor}; use serde::{Deserialize, Deserializer}; use serde_json::Value; -use std::collections::{BTreeMap, BTreeSet}; +use sha2::{Digest, Sha256}; +use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::fmt; use std::fs::{self, File}; use std::io::Read; @@ -88,6 +93,158 @@ pub(super) struct ArtifactBindings { pub(super) mtp: Vec, } +pub(super) struct QwenMap { + path: PathBuf, + map: Mmap, +} + +#[derive(Clone)] +pub(super) struct QwenTensor { + pub(super) map: usize, + pub(super) name: String, + pub(super) dtype: String, + pub(super) shape: Vec, + pub(super) quant_bits: Option, + pub(super) group_size: Option, + pub(super) range: std::ops::Range, +} + +pub(super) struct QwenModel { + tokenizer: Tokenizer, + memory: MemoryPlan, + maps: Vec, + tensors: HashMap, + identity: [u8; 32], +} + +impl QwenModel { + pub(super) fn open(root: &Path, context: u32) -> Result { + let loaded = load(root, context, false)?; + let mut paths = loaded + .bindings + .core + .iter() + .map(|binding| binding.file.clone()) + .collect::>(); + paths.sort(); + paths.dedup(); + let mut maps = Vec::with_capacity(paths.len()); + let mut map_indices = HashMap::with_capacity(paths.len()); + for path in paths { + let file = File::open(&path).map_err(|error| format!("{}: {error}", path.display()))?; + // SAFETY: verified managed artifacts remain read-only while the model owns each mapping. + let map = unsafe { MmapOptions::new().map(&file) } + .map_err(|error| format!("cannot map {}: {error}", path.display()))?; + map_indices.insert(path.clone(), maps.len()); + maps.push(QwenMap { path, map }); + } + let tensors = loaded + .bindings + .core + .into_iter() + .map(|binding| { + let tensor = QwenTensor { + map: map_indices[&binding.file], + name: binding.name.clone(), + dtype: binding.dtype, + shape: binding.shape, + quant_bits: binding.quant_bits, + group_size: binding.group_size, + range: binding.range, + }; + (binding.name, tensor) + }) + .collect::>(); + let mut hash = Sha256::new(); + hash.update(b"DS4Server Qwen3.8 checkpoint identity v1"); + hash.update(MANIFEST); + let identity = hash.finalize().into(); + Ok(Self { + tokenizer: loaded.tokenizer, + memory: loaded.memory, + maps, + tensors, + identity, + }) + } + + pub(super) fn tensor(&self, name: &str) -> Result<&QwenTensor, String> { + self.tensors + .get(name) + .ok_or_else(|| format!("Qwen core tensor is missing: {name}")) + } + + pub(super) fn map(&self, index: usize) -> (&[u8], &Path) { + (&self.maps[index].map, &self.maps[index].path) + } + + pub(super) fn checkpoint_identity(&self) -> [u8; 32] { + self.identity + } + + pub(super) fn summary(&self) -> ModelSummary { + ModelSummary { + model: ModelChoice::Qwen38FlashNext, + mapped_bytes: self.maps.iter().map(|item| item.map.len() as u64).sum(), + tensor_count: self.tensors.len(), + vocabulary_size: self.tokenizer.vocab_size(), + support_loaded: false, + vision_loaded: false, + } + } + + pub(super) fn render_conversation( + &self, + system: &str, + messages: &[ChatTurn], + reasoning: ReasoningMode, + ) -> Vec { + self.tokenizer + .encode_conversation(system, messages, reasoning) + } + + pub(super) fn render_history( + &self, + system: &str, + messages: &[ChatTurn], + reasoning: ReasoningMode, + ) -> Vec { + self.tokenizer.encode_history(system, messages, reasoning) + } + + pub(super) fn render_continuation( + &self, + prompt: &str, + reasoning: ReasoningMode, + skip_previous_eos: bool, + ) -> Vec { + self.tokenizer + .encode_continuation(prompt, reasoning, skip_previous_eos) + } + + pub(super) fn token_bytes(&self, token: i32) -> Option> { + self.tokenizer.token_bytes(token) + } + + pub(super) fn is_stop_token_for_reasoning(&self, token: i32, reasoning: ReasoningMode) -> bool { + self.tokenizer.is_stop(token) + || (reasoning == ReasoningMode::Direct + && (self.tokenizer.is_think_start(token) || self.tokenizer.is_think_end(token))) + } + + pub(super) fn is_think_start_token(&self, token: i32) -> bool { + self.tokenizer.is_think_start(token) + } + + pub(super) fn is_think_end_token(&self, token: i32) -> bool { + self.tokenizer.is_think_end(token) + } + + pub(super) fn memory(&self) -> &MemoryPlan { + &self.memory + } +} + pub(crate) fn validate_artifacts(root: &Path) -> Result<(), String> { load(root, 262_144, true).map(|_| ()) }