Implement Qwen text core

This commit is contained in:
Georg Bauer
2026-09-03 20:31:53 +02:00
parent 3773cfda2e
commit c414640050
9 changed files with 2766 additions and 52 deletions

408
metal/qwen38.metal Normal file
View File

@@ -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<float>((uint)value << 16);
}
static inline ushort qwen_to_bf16(float value) {
uint bits = as_type<uint>(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]));
}

View File

@@ -128,6 +128,30 @@ int ds4_gpu_set_aux_model_map_range(const void *model_map,
uint64_t map_offset, uint64_t map_offset,
uint64_t map_size); 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); 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_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_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); int ds4_gpu_q8_cache_suppressed(void);

View File

@@ -4359,6 +4359,7 @@ static NSString *ds4_gpu_full_source(void) {
@[@"DS4_METAL_GLM53_BF16_SOURCE", @"metal/glm53_bf16.metal"], @[@"DS4_METAL_GLM53_BF16_SOURCE", @"metal/glm53_bf16.metal"],
@[@"DS4_METAL_GLM53_VISION_SOURCE", @"metal/glm53_vision.metal"], @[@"DS4_METAL_GLM53_VISION_SOURCE", @"metal/glm53_vision.metal"],
@[@"DS4_METAL_GLM53_KDA_SOURCE", @"metal/glm53_kda.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_MOE_SOURCE", @"metal/moe.metal"],
@[@"DS4_METAL_DSV4_HC_SOURCE", @"metal/dsv4_hc.metal"], @[@"DS4_METAL_DSV4_HC_SOURCE", @"metal/dsv4_hc.metal"],
@[@"DS4_METAL_UNARY_SOURCE", @"metal/unary.metal"], @[@"DS4_METAL_UNARY_SOURCE", @"metal/unary.metal"],
@@ -11737,6 +11738,68 @@ static id<MTLBuffer> ds4_gpu_wrap_model_exact_range_owned(
DS4_GPU_EXACT_VIEW_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<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(kernel);
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> 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<MTLBuffer> 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 ds4_gpu_stream_expert_cache_configured_count(void) {
uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget();
if (budget > DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) { if (budget > DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) {

View File

@@ -53,6 +53,7 @@ pub(crate) fn checkpoint_model(path: &Path) -> Option<ModelChoice> {
let model_size_offset = match &magic { let model_size_offset = match &magic {
b"DS4RKV01" => 40, b"DS4RKV01" => 40,
b"DS4GLM01" => 44, b"DS4GLM01" => 44,
b"DS4QWN01" => return Some(ModelChoice::Qwen38FlashNext),
_ => return None, _ => return None,
}; };
file.seek(SeekFrom::Start(model_size_offset)).ok()?; file.seek(SeekFrom::Start(model_size_offset)).ok()?;
@@ -301,6 +302,115 @@ pub(crate) struct Model {
tokenizer: Tokenizer, tokenizer: Tokenizer,
} }
enum LoadedModel {
Gguf(Box<Model>),
Qwen(Box<qwen::QwenModel>),
}
impl From<Model> 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<Self, String> {
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<i32> {
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<i32> {
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<i32> {
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<Vec<u8>> {
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)] #[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) struct ModelSummary { pub(crate) struct ModelSummary {
pub(crate) model: ModelChoice, pub(crate) model: ModelChoice,
@@ -320,18 +430,7 @@ impl Model {
&settings.artifacts, &settings.artifacts,
)?; )?;
if settings.model.is_qwen38() { if settings.model.is_qwen38() {
let context = u32::try_from(settings.context_tokens) return Err("Qwen must be opened through its dedicated safetensors loader".into());
.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,
));
} }
let mut model = Self::open_main(&settings.artifacts.model, settings.model)?; let mut model = Self::open_main(&settings.artifacts.model, settings.model)?;
if settings.execution.warm_weights { if settings.execution.warm_weights {
@@ -622,7 +721,7 @@ impl Generator {
pub(crate) fn open(settings: &EngineSettings, metrics: Arc<Metrics>) -> Result<Self, String> { pub(crate) fn open(settings: &EngineSettings, metrics: Arc<Metrics>) -> Result<Self, String> {
let simulated_memory = let simulated_memory =
SimulatedMemory::acquire(settings.diagnostics.simulated_used_memory_bytes)?; SimulatedMemory::acquire(settings.diagnostics.simulated_used_memory_bytes)?;
let model = Model::open(settings)?; let model = LoadedModel::open(settings)?;
let executor = metal::Executor::open_configured( let executor = metal::Executor::open_configured(
model, model,
settings.context_tokens.max(1) as u32, settings.context_tokens.max(1) as u32,

View File

@@ -3,16 +3,18 @@ mod glm;
mod gpu; mod gpu;
mod hotlist; mod hotlist;
mod profile; mod profile;
mod qwen;
mod vision; mod vision;
pub(super) use vision::VisionEmbedding; pub(super) use vision::VisionEmbedding;
use glm::GlmExecutor; use glm::GlmExecutor;
use gpu::*; use gpu::*;
use profile::ExpertProfile; 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::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::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::model::ModelChoice;
use crate::settings::{ use crate::settings::{
EngineSpeculativeSettings, EngineSsdSettings, EngineSteeringSettings, ReasoningMode, EngineSpeculativeSettings, EngineSsdSettings, EngineSteeringSettings, ReasoningMode,
@@ -44,7 +46,7 @@ fn environment_present(name: &CStr) -> bool {
!unsafe { getenv(name.as_ptr()) }.is_null() !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_FLASH_ATTN_SOURCE", "flash_attn.metal"),
("DS4_METAL_DENSE_SOURCE", "dense.metal"), ("DS4_METAL_DENSE_SOURCE", "dense.metal"),
("DS4_METAL_MOE_SOURCE", "moe.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_BF16_SOURCE", "glm53_bf16.metal"),
("DS4_METAL_GLM53_VISION_SOURCE", "glm53_vision.metal"), ("DS4_METAL_GLM53_VISION_SOURCE", "glm53_vision.metal"),
("DS4_METAL_GLM53_KDA_SOURCE", "glm53_kda.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 // 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 { pub(super) enum Executor {
DeepSeek(Box<DeepSeekExecutor>), DeepSeek(Box<DeepSeekExecutor>),
Glm(Box<GlmExecutor>), Glm(Box<GlmExecutor>),
Qwen(Box<QwenExecutor>),
} }
pub(super) enum ResidentState { pub(super) enum ResidentState {
DeepSeek(Box<DeepSeekResidentState>), DeepSeek(Box<DeepSeekResidentState>),
Glm(Box<glm::GlmResidentState>), Glm(Box<glm::GlmResidentState>),
Qwen(Box<qwen::QwenResidentState>),
} }
impl Executor { impl Executor {
@@ -4380,7 +4385,7 @@ impl Executor {
prefill_chunk: u32, prefill_chunk: u32,
) -> Result<Self, String> { ) -> Result<Self, String> {
Self::open_configured( Self::open_configured(
model, LoadedModel::from(model),
context, context,
quality, quality,
prefill_chunk, prefill_chunk,
@@ -4415,7 +4420,7 @@ impl Executor {
#[allow(clippy::too_many_arguments)] #[allow(clippy::too_many_arguments)]
pub(super) fn open_configured( pub(super) fn open_configured(
model: Model, model: impl Into<LoadedModel>,
context: u32, context: u32,
quality: bool, quality: bool,
prefill_chunk: u32, prefill_chunk: u32,
@@ -4425,32 +4430,39 @@ impl Executor {
steering: EngineSteeringSettings, steering: EngineSteeringSettings,
expert_profile_path: Option<&str>, expert_profile_path: Option<&str>,
) -> Result<Self, String> { ) -> Result<Self, String> {
match model.shape.family { match model.into() {
ModelFamily::DeepSeek => DeepSeekExecutor::open_profile( LoadedModel::Gguf(model) if model.shape.family == ModelFamily::DeepSeek => {
model, DeepSeekExecutor::open_profile(
context, *model,
quality, context,
prefill_chunk, quality,
power_percent, prefill_chunk,
speculative, power_percent,
ssd, speculative,
steering, ssd,
expert_profile_path, steering,
) expert_profile_path,
.map(Box::new) )
.map(Self::DeepSeek), .map(Box::new)
ModelFamily::Glm => GlmExecutor::open_profile( .map(Self::DeepSeek)
model, }
context, LoadedModel::Gguf(model) if model.shape.family == ModelFamily::Glm => {
quality, GlmExecutor::open_profile(
ssd, *model,
speculative, context,
steering, quality,
expert_profile_path, ssd,
) speculative,
.map(Box::new) steering,
.map(Self::Glm), expert_profile_path,
ModelFamily::Qwen => unreachable!("Qwen uses its dedicated executor"), )
.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)?; executor.eval(token)?;
Ok(vec![token]) Ok(vec![token])
} }
Self::Qwen(executor) => {
executor.eval(token)?;
Ok(vec![token])
}
} }
} }
@@ -4490,6 +4506,7 @@ impl Executor {
match self { match self {
Self::DeepSeek(executor) => executor.eval(token), Self::DeepSeek(executor) => executor.eval(token),
Self::Glm(executor) => executor.eval(token), Self::Glm(executor) => executor.eval(token),
Self::Qwen(executor) => executor.eval(token),
} }
} }
@@ -4499,7 +4516,7 @@ impl Executor {
) -> Result<Vec<vision::VisionEmbedding>, String> { ) -> Result<Vec<vision::VisionEmbedding>, String> {
match self { match self {
Self::Glm(executor) => executor.encode_visions(encoded), 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 { match self {
Self::Glm(executor) => executor.set_vision_overlays(overlays), Self::Glm(executor) => executor.set_vision_overlays(overlays),
Self::DeepSeek(_) if overlays.is_empty() => Ok(()), 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; let _ = reasoning;
executor.eval_speculative_greedy(token, max_tokens, cancelled) 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 { match self {
Self::DeepSeek(executor) => executor.prefill(tokens, progress), Self::DeepSeek(executor) => executor.prefill(tokens, progress),
Self::Glm(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 { match self {
Self::DeepSeek(executor) => executor.logits(), Self::DeepSeek(executor) => executor.logits(),
Self::Glm(executor) => executor.logits(), Self::Glm(executor) => executor.logits(),
Self::Qwen(executor) => executor.logits(),
} }
} }
@@ -4554,13 +4579,15 @@ impl Executor {
match self { match self {
Self::DeepSeek(executor) => executor.execution_stats(), Self::DeepSeek(executor) => executor.execution_stats(),
Self::Glm(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 { match self {
Self::DeepSeek(executor) => executor.model(), Self::DeepSeek(executor) => ModelRef::Gguf(executor.model()),
Self::Glm(executor) => executor.model(), Self::Glm(executor) => ModelRef::Gguf(executor.model()),
Self::Qwen(executor) => ModelRef::Qwen(executor.model()),
} }
} }
@@ -4568,6 +4595,7 @@ impl Executor {
match self { match self {
Self::DeepSeek(executor) => executor.context(), Self::DeepSeek(executor) => executor.context(),
Self::Glm(executor) => executor.context(), Self::Glm(executor) => executor.context(),
Self::Qwen(executor) => executor.context(),
} }
} }
@@ -4575,6 +4603,7 @@ impl Executor {
match self { match self {
Self::DeepSeek(executor) => executor.position(), Self::DeepSeek(executor) => executor.position(),
Self::Glm(executor) => executor.position(), Self::Glm(executor) => executor.position(),
Self::Qwen(executor) => executor.position(),
} }
} }
@@ -4582,6 +4611,7 @@ impl Executor {
match self { match self {
Self::DeepSeek(executor) => executor.reset(), Self::DeepSeek(executor) => executor.reset(),
Self::Glm(executor) => executor.reset(), Self::Glm(executor) => executor.reset(),
Self::Qwen(executor) => executor.reset(),
} }
} }
@@ -4596,6 +4626,9 @@ impl Executor {
Some(ResidentState::Glm(_)) => { Some(ResidentState::Glm(_)) => {
return Err("resident session belongs to a different model family".into()); 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, None => None,
}; };
executor.swap_resident_state(&mut inner)?; executor.swap_resident_state(&mut inner)?;
@@ -4607,11 +4640,25 @@ impl Executor {
Some(ResidentState::DeepSeek(_)) => { Some(ResidentState::DeepSeek(_)) => {
return Err("resident session belongs to a different model family".into()); 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, None => None,
}; };
executor.swap_resident_state(&mut inner)?; executor.swap_resident_state(&mut inner)?;
*state = inner.map(|state| ResidentState::Glm(Box::new(state))); *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(()) Ok(())
} }
@@ -4620,6 +4667,7 @@ impl Executor {
match self { match self {
Self::DeepSeek(executor) => executor.align_prompt(tokens), Self::DeepSeek(executor) => executor.align_prompt(tokens),
Self::Glm(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 { match self {
Self::DeepSeek(executor) => executor.tokens(), Self::DeepSeek(executor) => executor.tokens(),
Self::Glm(executor) => executor.tokens(), Self::Glm(executor) => executor.tokens(),
Self::Qwen(executor) => executor.tokens(),
} }
} }
@@ -4634,6 +4683,7 @@ impl Executor {
match self { match self {
Self::DeepSeek(executor) => executor.checkpoint_tag(), Self::DeepSeek(executor) => executor.checkpoint_tag(),
Self::Glm(executor) => executor.checkpoint_tag(), Self::Glm(executor) => executor.checkpoint_tag(),
Self::Qwen(executor) => executor.checkpoint_tag(),
} }
} }
@@ -4641,6 +4691,7 @@ impl Executor {
match self { match self {
Self::DeepSeek(executor) => executor.note_checkpoint_tag(tag), Self::DeepSeek(executor) => executor.note_checkpoint_tag(tag),
Self::Glm(executor) => executor.note_checkpoint_tag(tag), Self::Glm(executor) => executor.note_checkpoint_tag(tag),
Self::Qwen(executor) => executor.note_checkpoint_tag(tag),
} }
} }
} }

View File

@@ -320,6 +320,7 @@ impl Executor {
match self { match self {
Self::DeepSeek(executor) => executor.save_checkpoint(path, tag, progress), Self::DeepSeek(executor) => executor.save_checkpoint(path, tag, progress),
Self::Glm(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 { match self {
Self::DeepSeek(executor) => executor.load_checkpoint(path, progress), Self::DeepSeek(executor) => executor.load_checkpoint(path, progress),
Self::Glm(executor) => executor.load_checkpoint(path, progress), Self::Glm(executor) => executor.load_checkpoint(path, progress),
Self::Qwen(executor) => executor.load_checkpoint(path, progress),
} }
} }
} }

View File

@@ -5,6 +5,22 @@ pub(super) struct GpuTensor {
_private: [u8; 0], _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)] #[derive(Clone, Copy, Default)]
#[repr(C)] #[repr(C)]
pub(super) struct Glm53VisionLayerWeights { pub(super) struct Glm53VisionLayerWeights {
@@ -81,6 +97,18 @@ unsafe extern "C" {
map_size: u64, map_size: u64,
max_tensor_bytes: u64, max_tensor_bytes: u64,
) -> i32; ) -> 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( pub(super) fn ds4_gpu_set_transient_model_map_range(
model_map: *const c_void, model_map: *const c_void,
model_size: u64, model_size: u64,
@@ -1804,7 +1832,7 @@ unsafe extern "C" {
} }
pub(super) struct Context { pub(super) struct Context {
_model_file: File, _model_file: Option<File>,
} }
impl Context { impl Context {
@@ -1908,9 +1936,29 @@ impl Context {
} }
} }
Ok(Self { Ok(Self {
_model_file: model_file, _model_file: Some(model_file),
}) })
} }
pub(super) fn open_qwen(admission_bytes: u64) -> Result<Self, String> {
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 { impl Drop for Context {

1862
src/engine/metal/qwen.rs Normal file

File diff suppressed because it is too large Load Diff

View File

@@ -1,8 +1,13 @@
use super::tokenizer::Tokenizer; 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::de::{MapAccess, Visitor};
use serde::{Deserialize, Deserializer}; use serde::{Deserialize, Deserializer};
use serde_json::Value; use serde_json::Value;
use std::collections::{BTreeMap, BTreeSet}; use sha2::{Digest, Sha256};
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::fmt; use std::fmt;
use std::fs::{self, File}; use std::fs::{self, File};
use std::io::Read; use std::io::Read;
@@ -88,6 +93,158 @@ pub(super) struct ArtifactBindings {
pub(super) mtp: Vec<TensorBinding>, pub(super) mtp: Vec<TensorBinding>,
} }
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<u64>,
pub(super) quant_bits: Option<u32>,
pub(super) group_size: Option<u64>,
pub(super) range: std::ops::Range<u64>,
}
pub(super) struct QwenModel {
tokenizer: Tokenizer,
memory: MemoryPlan,
maps: Vec<QwenMap>,
tensors: HashMap<String, QwenTensor>,
identity: [u8; 32],
}
impl QwenModel {
pub(super) fn open(root: &Path, context: u32) -> Result<Self, String> {
let loaded = load(root, context, false)?;
let mut paths = loaded
.bindings
.core
.iter()
.map(|binding| binding.file.clone())
.collect::<Vec<_>>();
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::<HashMap<_, _>>();
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<i32> {
self.tokenizer
.encode_conversation(system, messages, reasoning)
}
pub(super) fn render_history(
&self,
system: &str,
messages: &[ChatTurn],
reasoning: ReasoningMode,
) -> Vec<i32> {
self.tokenizer.encode_history(system, messages, reasoning)
}
pub(super) fn render_continuation(
&self,
prompt: &str,
reasoning: ReasoningMode,
skip_previous_eos: bool,
) -> Vec<i32> {
self.tokenizer
.encode_continuation(prompt, reasoning, skip_previous_eos)
}
pub(super) fn token_bytes(&self, token: i32) -> Option<Vec<u8>> {
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> { pub(crate) fn validate_artifacts(root: &Path) -> Result<(), String> {
load(root, 262_144, true).map(|_| ()) load(root, 262_144, true).map(|_| ())
} }