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

Add GLM 5.3 Flash support

Browse Source
This commit is contained in:
Georg Bauer committed 2026-09-01 19:06:20 +02:00
1 parent 9a33c61ea6
commit 46d6a976a5
31 files changed
+7806 -447
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

AGENTS.md
+1
View File
@@ -4,6 +4,7 @@
- Keep changes focused, handle errors explicitly, and add the smallest useful test for non-trivial behavior. - Keep changes focused, handle errors explicitly, and add the smallest useful test for non-trivial behavior.
- Preserve `rustfmt` output and keep Clippy warning-free. - Preserve `rustfmt` output and keep Clippy warning-free.
- Treat DS4 as the behavioral oracle. Model execution, token processing, context accounting, and KV-cache behavior must remain identical to DS4. Differences are bugs unless they are unavoidable Rust/platform requirements and the user explicitly approves them before implementation; when uncertain, preserve DS4 behavior. - Treat DS4 as the behavioral oracle. Model execution, token processing, context accounting, and KV-cache behavior must remain identical to DS4. Differences are bugs unless they are unavoidable Rust/platform requirements and the user explicitly approves them before implementation; when uncertain, preserve DS4 behavior.
- Keep all application, model-loading, inference-graph, session, and orchestration code in Rust. Never copy, vendor, compile, or call DS4's C implementation; use it only as a behavioral and performance reference. Native code is limited to the Objective-C Metal bridge and `.metal` kernels required to execute GPU work.
- This is not a GitHub project. Use direct `git` commands for version control and the `tea` CLI for forge operations; do not use GitHub tools or workflows. - This is not a GitHub project. Use direct `git` commands for version control and the `tea` CLI for forge operations; do not use GitHub tools or workflows.
- Issues are managed with the command "tea" run from the project directory. - Issues are managed with the command "tea" run from the project directory.
docs/METAL_PARITY.md
+64
View File
@@ -1,5 +1,69 @@
# Apple Metal parity for issue #78 # Apple Metal parity for issue #78
## GLM 5.3 Flash parity target
- Previous DS4 oracle: `8db89fe083ae4d17c9a2428ccd29803d3ae8f577`
- Current DS4 oracle: `ec7642cdd9ec81d01ad4b1fd8f8a3d1511533748`
- Main artifact: `GLM-5.3-Flash-Q2.gguf`, SHA-256
`e81fd6241c6e55a64e1e14e47a3eab61a173fa8d7e4b5c1d1848827119705b32`
- Vision artifact: `GLM-5.3-Flash-Vision-Encoder.gguf`, SHA-256
`ae23e14c6979e889051b2e4a39351abcdafb161e18e606fae4d8c40095a4bf3a`
- Machine: MacBook Pro `Mac17,7`, Apple M5 Max, 128 GiB unified memory,
resident Q2 weights
The applicable commits after the previous oracle are:
| DS4 commit(s) | Scope carried into DS4Server |
| --- | --- |
| `147109a` | GLM 5.3 shape, GGUF binding, KDA/DSA attention, hyperconnections, routed MoE, BF16 tensors, output head, integrated MTP, and Metal kernels |
| `4771329`, `9f95d9f` | Vision encoder, image preprocessing, prompt overlays, and compact-prefill correctness |
| `b0c31af` | Pooled indexer cache, bounded attention memory, and batched prefill |
| `747eafc`, `1df8ff0` | Fixed-image encoder oracle, replay/logit checks, and multimodal quality gates |
| `8adb221` | Directional attention and FFN steering at the DS4 residual points |
| `639d4eb` | Continued-prefill cache and position behavior |
| `6cf658a` | BF16 decode projection schedule and Metal kernel specialization |
| `01f3316`, `81c8c39`, `5da3206` | Current Metal compact-cache ABI, Q2/full-checkpoint binding, and GLM streaming compatibility |
CUDA- and ROCm-only commits in the range are excluded. The Rust release gates
compare three-run medians against this exact DS4 build, require identical
greedy token hashes, and fail only when Rust is more than 2% slower than DS4;
there is no upper bound when Rust is faster. Text measures an
8192-token frontier plus steady decode, both with and without integrated MTP.
Vision compares the fixed encoder embedding (cosine at least 0.995, mean
absolute error at most 0.001, maximum error at most 0.06), encoder latency,
multimodal prefill, steady decode, image-conditioned logits, and replay.
### GLM 5.3 Flash acceptance results
The text fixture is `ds4.c`, rendered as one no-system chat turn (`-sys ''` in
DS4), truncated to exactly 8192 tokens, with a 32768-token allocation and 64
greedy output tokens. Same-process three-run medians are used because both DS4
and DS4Server exhibit a load-order/thermal artifact with this 96 GB model.
| Capability | Metric | DS4 | DS4Server | Result |
| --- | --- | ---: | ---: | --- |
| Text | Prefill tok/s | 288.91 | 295.21 | Rust 2.18% faster |
| Text | Decode tok/s | 22.58 | 22.78 | Rust 0.89% faster |
| Integrated MTP | Decode tok/s | 26.55 | 26.80 | Rust 0.93% faster |
| Vision encoder | Fixed-image seconds | 2.6255 | 2.4354 | Rust 7.24% faster |
| Vision prompt | Prefill tok/s | 185.81 | 201.51 | Rust 8.45% faster |
| Vision prompt | Decode tok/s | 26.40 | 27.21 | Rust 3.09% faster |
Text and MTP produced the same token SHA-256,
`0b61e613708d8367fc87f5a32fbfd4ba098b40ee920645b1e38381e867131108`.
The 1094-token fixed-image prompt produced the same vision token SHA-256,
`424fdadc10c0a2bd95b405a3985151e84aa844a209d2a2ca72cd01c39dfa53db`.
The Rust vision embedding compared with the DS4 oracle at cosine
`0.999958726`, mean absolute error `0.000036702`, and maximum error
`0.006153620`.
The checkpoint advertises a 1048576-token maximum. On this 128 GiB machine,
32768 is the default manageable context and 50000 is the extended option: the
50000-token run used 3.51 GiB of live runtime tensors in addition to 90.61 GiB
of mapped model wrappers and retained the exact output hash. Larger allocations
leave too little operating-system and application headroom for a reliable
default.
## Frozen comparison ## Frozen comparison
- DS4 oracle: `8db89fe083ae4d17c9a2428ccd29803d3ae8f577` - DS4 oracle: `8db89fe083ae4d17c9a2428ccd29803d3ae8f577`
metal/dsv4_misc.metal
+227 -47
View File
@@ -195,6 +195,17 @@ struct ds4_metal_args_glm_store_indexer_k {
float pad1; float pad1;
}; };
struct ds4_metal_args_glm53_indexer_pool_update {
uint32_t pos0;
uint32_t n_tokens;
uint32_t cache_cap;
uint32_t head_dim;
uint32_t pool_size;
uint32_t cache_f16;
float eps;
uint32_t pad0;
};
struct ds4_metal_args_glm_attention_full { struct ds4_metal_args_glm_attention_full {
uint32_t pos0; uint32_t pos0;
uint32_t n_tokens; uint32_t n_tokens;
@@ -221,6 +232,15 @@ struct ds4_metal_args_glm_fill_selected_range_batch {
uint32_t pad_row; uint32_t pad_row;
}; };
struct ds4_metal_args_glm53_expand_pool_selection {
uint32_t n_tokens;
uint32_t pos0;
uint32_t selected_pools;
uint32_t index_topk;
uint32_t pool_size;
uint32_t output_width;
};
struct ds4_metal_args_glm_indexer_rope_tail { struct ds4_metal_args_glm_indexer_rope_tail {
uint32_t n_tokens; uint32_t n_tokens;
uint32_t n_head; uint32_t n_head;
@@ -252,6 +272,8 @@ struct ds4_metal_args_glm_indexer_scores_batch {
uint32_t head_dim; uint32_t head_dim;
uint32_t pos0; uint32_t pos0;
uint32_t cache_f16; uint32_t cache_f16;
uint32_t row_group_size;
uint32_t pad0;
uint64_t q_token_stride; uint64_t q_token_stride;
uint64_t q_head_stride; uint64_t q_head_stride;
uint64_t weights_token_stride; uint64_t weights_token_stride;
@@ -259,6 +281,13 @@ struct ds4_metal_args_glm_indexer_scores_batch {
float scale; float scale;
}; };
static inline uint glm_indexer_batch_visible_rows(
constant ds4_metal_args_glm_indexer_scores_batch &args,
uint token) {
const uint group = max(args.row_group_size, 1u);
return min((args.pos0 + token + 1u) / group, args.n_rows);
}
struct ds4_metal_args_glm_qk_lowrank { struct ds4_metal_args_glm_qk_lowrank {
uint32_t n_head; uint32_t n_head;
uint32_t kv_lora_dim; uint32_t kv_lora_dim;
@@ -933,6 +962,117 @@ kernel void kernel_glm_store_indexer_k(
} }
} }
static inline float glm53_pool_bf16_to_f32(ushort value) {
return as_type<float>((uint)value << 16);
}
kernel void kernel_glm53_indexer_pool_update(
constant ds4_metal_args_glm53_indexer_pool_update &args,
device const char *raw_k,
device const char *gate,
device const float *norm_weight,
device const float *norm_bias,
device const ushort *ape,
device char *pool_cache,
device float *tail_k,
device float *tail_gate,
threadgroup float *shared [[threadgroup(0)]],
uint tid [[thread_index_in_threadgroup]],
uint3 tgpig [[threadgroup_position_in_grid]]) {
if (args.head_dim == 0u || args.pool_size == 0u ||
tid >= args.head_dim || args.n_tokens == 0u) return;
const uint pool = args.pos0 / args.pool_size + tgpig.x;
const uint pool_start = pool * args.pool_size;
const uint input_end = args.pos0 + args.n_tokens;
if (pool_start >= input_end || pool_start + args.pool_size <= args.pos0) return;
threadgroup float *rows = shared;
threadgroup float *mean = rows + args.pool_size * args.head_dim;
threadgroup float *inv = mean + args.pool_size;
const bool complete = pool_start + args.pool_size <= input_end;
for (uint r = 0; r < args.pool_size; r++) {
const uint pos = pool_start + r;
float k_value = 0.0f;
float gate_value = 0.0f;
if (pos >= args.pos0 && pos < input_end) {
const uint src_row = pos - args.pos0;
k_value = ((device const float *)raw_k)[
(uint64_t)src_row * args.head_dim + tid];
gate_value = ((device const float *)gate)[
(uint64_t)src_row * args.head_dim + tid];
if (!complete) {
tail_k[(uint64_t)r * args.head_dim + tid] = k_value;
tail_gate[(uint64_t)r * args.head_dim + tid] = gate_value;
}
} else {
k_value = tail_k[(uint64_t)r * args.head_dim + tid];
gate_value = tail_gate[(uint64_t)r * args.head_dim + tid];
}
rows[(uint64_t)r * args.head_dim + tid] = k_value;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (!complete || pool >= (args.cache_cap + args.pool_size - 1u) / args.pool_size) {
return;
}
if (tid < args.pool_size) {
const uint r = tid;
float sum = 0.0f;
for (uint d = 0; d < args.head_dim; d++) {
sum += rows[(uint64_t)r * args.head_dim + d];
}
const float m = sum / (float)args.head_dim;
float ss = 0.0f;
for (uint d = 0; d < args.head_dim; d++) {
const float delta = rows[(uint64_t)r * args.head_dim + d] - m;
ss += delta * delta;
}
mean[r] = m;
inv[r] = rsqrt(ss / (float)args.head_dim + args.eps);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
float max_logit = -INFINITY;
float logits[4];
for (uint r = 0; r < args.pool_size; r++) {
const uint pos = pool_start + r;
float gate_value;
if (pos >= args.pos0) {
const uint src_row = pos - args.pos0;
gate_value = ((device const float *)gate)[
(uint64_t)src_row * args.head_dim + tid];
} else {
gate_value = tail_gate[(uint64_t)r * args.head_dim + tid];
}
logits[r] = gate_value +
glm53_pool_bf16_to_f32(ape[(uint64_t)r * args.head_dim + tid]);
max_logit = max(max_logit, logits[r]);
}
float denom = 0.0f;
for (uint r = 0; r < args.pool_size; r++) {
logits[r] = exp(logits[r] - max_logit);
denom += logits[r];
}
float pooled = 0.0f;
for (uint r = 0; r < args.pool_size; r++) {
const float normalized =
(rows[(uint64_t)r * args.head_dim + tid] - mean[r]) * inv[r] *
norm_weight[tid] + norm_bias[tid];
pooled += (logits[r] / denom) * normalized;
}
const uint64_t dst_index = (uint64_t)pool * args.head_dim + tid;
if (args.cache_f16 != 0u) {
((device half *)pool_cache)[dst_index] = (half)pooled;
} else {
((device float *)pool_cache)[dst_index] = pooled;
}
}
static inline void glm_dense_cache_store_f32_or_f16( static inline void glm_dense_cache_store_f32_or_f16(
device char *base, device char *base,
uint64_t index, uint64_t index,
@@ -1011,7 +1151,8 @@ kernel void kernel_glm_build_kv_cache(
corr_dims); corr_dims);
} }
const float theta_base = (float)pos; const float theta_base = (float)pos;
const float inv_ndims = -1.0f / (float)args.qk_rope; const float inv_ndims = args.qk_rope != 0u ?
-1.0f / (float)args.qk_rope : 0.0f;
for (uint r = tid * 2u; r < args.qk_rope; r += nth * 2u) { for (uint r = tid * 2u; r < args.qk_rope; r += nth * 2u) {
#ifdef DS4_METAL_ROPE_EXP2_LOG2 #ifdef DS4_METAL_ROPE_EXP2_LOG2
const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base)); const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
@@ -1096,7 +1237,8 @@ kernel void kernel_glm_build_kv_cache_decode_group4(
corr_dims); corr_dims);
} }
const float theta_base = (float)pos; const float theta_base = (float)pos;
const float inv_ndims = -1.0f / (float)args.qk_rope; const float inv_ndims = args.qk_rope != 0u ?
-1.0f / (float)args.qk_rope : 0.0f;
for (uint r = tid * 2u; r < args.qk_rope; r += 512u) { for (uint r = tid * 2u; r < args.qk_rope; r += 512u) {
#ifdef DS4_METAL_ROPE_EXP2_LOG2 #ifdef DS4_METAL_ROPE_EXP2_LOG2
const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base)); const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
@@ -1179,7 +1321,8 @@ kernel void kernel_glm_build_kv_cache_flash(
corr_dims); corr_dims);
} }
const float theta_base = (float)pos; const float theta_base = (float)pos;
const float inv_ndims = -1.0f / (float)args.qk_rope; const float inv_ndims = args.qk_rope != 0u ?
-1.0f / (float)args.qk_rope : 0.0f;
for (uint r = tid * 2u; r < args.qk_rope; r += nth * 2u) { for (uint r = tid * 2u; r < args.qk_rope; r += nth * 2u) {
#ifdef DS4_METAL_ROPE_EXP2_LOG2 #ifdef DS4_METAL_ROPE_EXP2_LOG2
const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base)); const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
@@ -1400,6 +1543,35 @@ kernel void kernel_glm_fill_selected_range_batch(
selected[gid] = slot < visible ? slot : args.pad_row; selected[gid] = slot < visible ? slot : args.pad_row;
} }
kernel void kernel_glm53_expand_pool_selection(
constant ds4_metal_args_glm53_expand_pool_selection &args,
device const uint32_t *pool_selected,
device uint32_t *raw_selected,
uint gid [[thread_position_in_grid]]) {
const uint total = args.n_tokens * args.output_width;
if (gid >= total || args.output_width == 0u || args.pool_size == 0u) return;
const uint token = gid / args.output_width;
const uint slot = gid - token * args.output_width;
uint value = 0xffffffffu;
if (slot < args.index_topk) {
const uint pool_slot = slot / args.pool_size;
if (pool_slot < args.selected_pools) {
const uint pool = pool_selected[
(uint64_t)token * args.selected_pools + pool_slot];
value = pool * args.pool_size + slot % args.pool_size;
}
} else {
const uint tail_slot = slot - args.index_topk;
const uint visible = args.pos0 + token + 1u;
const uint tail_count = visible % args.pool_size;
if (tail_slot < tail_count) {
value = visible - tail_count + tail_slot;
}
}
raw_selected[gid] = value;
}
kernel void kernel_glm_indexer_rope_tail_f32( kernel void kernel_glm_indexer_rope_tail_f32(
constant ds4_metal_args_glm_indexer_rope_tail & args, constant ds4_metal_args_glm_indexer_rope_tail & args,
device char *x, device char *x,
@@ -1849,7 +2021,7 @@ kernel void kernel_glm_indexer_scores_batch(
device float *dst = (device float *)(scores + device float *dst = (device float *)(scores +
(uint64_t)token * args.score_token_stride) + row; (uint64_t)token * args.score_token_stride) + row;
const uint visible = min(args.pos0 + token + 1u, args.n_rows); const uint visible = glm_indexer_batch_visible_rows(args, token);
if (row >= visible) { if (row >= visible) {
if (tid == 0) *dst = -INFINITY; if (tid == 0) *dst = -INFINITY;
return; return;
@@ -1910,7 +2082,7 @@ kernel void kernel_glm_indexer_scores_tiled_f32(
const uint last_token = min(token_base + TM, args.n_tokens); const uint last_token = min(token_base + TM, args.n_tokens);
const uint max_visible = last_token > token_base ? const uint max_visible = last_token > token_base ?
min(args.pos0 + last_token, args.n_rows) : 0u; glm_indexer_batch_visible_rows(args, last_token - 1u) : 0u;
if (row_base >= max_visible) { if (row_base >= max_visible) {
for (uint i = tid; i < TM*TN; i += 128) { for (uint i = tid; i < TM*TN; i += 128) {
@@ -2005,13 +2177,13 @@ kernel void kernel_glm_indexer_scores_tiled_f32(
} }
if (token0 < args.n_tokens && row0 < args.n_rows) { if (token0 < args.n_tokens && row0 < args.n_rows) {
const uint visible = min(args.pos0 + token0 + 1u, args.n_rows); const uint visible = glm_indexer_batch_visible_rows(args, token0);
device float *dst = (device float *)(scores + device float *dst = (device float *)(scores +
(uint64_t)token0 * args.score_token_stride) + row0; (uint64_t)token0 * args.score_token_stride) + row0;
*dst = row0 < visible ? acc0 : -INFINITY; *dst = row0 < visible ? acc0 : -INFINITY;
} }
if (token1 < args.n_tokens && row1 < args.n_rows) { if (token1 < args.n_tokens && row1 < args.n_rows) {
const uint visible = min(args.pos0 + token1 + 1u, args.n_rows); const uint visible = glm_indexer_batch_visible_rows(args, token1);
device float *dst = (device float *)(scores + device float *dst = (device float *)(scores +
(uint64_t)token1 * args.score_token_stride) + row1; (uint64_t)token1 * args.score_token_stride) + row1;
*dst = row1 < visible ? acc1 : -INFINITY; *dst = row1 < visible ? acc1 : -INFINITY;
@@ -2043,7 +2215,7 @@ kernel void kernel_glm_indexer_scores_tiled(
const uint last_token = min(token_base + TM, args.n_tokens); const uint last_token = min(token_base + TM, args.n_tokens);
const uint max_visible = last_token > token_base ? const uint max_visible = last_token > token_base ?
min(args.pos0 + last_token, args.n_rows) : 0u; glm_indexer_batch_visible_rows(args, last_token - 1u) : 0u;
if (row_base >= max_visible) { if (row_base >= max_visible) {
for (uint i = tid; i < TM*TN; i += 128) { for (uint i = tid; i < TM*TN; i += 128) {
@@ -2138,13 +2310,13 @@ kernel void kernel_glm_indexer_scores_tiled(
} }
if (token0 < args.n_tokens && row0 < args.n_rows) { if (token0 < args.n_tokens && row0 < args.n_rows) {
const uint visible = min(args.pos0 + token0 + 1u, args.n_rows); const uint visible = glm_indexer_batch_visible_rows(args, token0);
device float *dst = (device float *)(scores + device float *dst = (device float *)(scores +
(uint64_t)token0 * args.score_token_stride) + row0; (uint64_t)token0 * args.score_token_stride) + row0;
*dst = row0 < visible ? acc0 : -INFINITY; *dst = row0 < visible ? acc0 : -INFINITY;
} }
if (token1 < args.n_tokens && row1 < args.n_rows) { if (token1 < args.n_tokens && row1 < args.n_rows) {
const uint visible = min(args.pos0 + token1 + 1u, args.n_rows); const uint visible = glm_indexer_batch_visible_rows(args, token1);
device float *dst = (device float *)(scores + device float *dst = (device float *)(scores +
(uint64_t)token1 * args.score_token_stride) + row1; (uint64_t)token1 * args.score_token_stride) + row1;
*dst = row1 < visible ? acc1 : -INFINITY; *dst = row1 < visible ? acc1 : -INFINITY;
@@ -2241,7 +2413,6 @@ kernel void kernel_glm_qk_lowrank_q8_0_glm52_sg(
ushort3 ntg_u [[threads_per_threadgroup]], ushort3 ntg_u [[threads_per_threadgroup]],
ushort tiisg [[thread_index_in_simdgroup]], ushort tiisg [[thread_index_in_simdgroup]],
ushort sgitg [[simdgroup_index_in_threadgroup]]) { ushort sgitg [[simdgroup_index_in_threadgroup]]) {
constexpr uint n_head = 64u;
constexpr uint kv_lora_dim = 512u; constexpr uint kv_lora_dim = 512u;
constexpr uint qk_nope = 192u; constexpr uint qk_nope = 192u;
constexpr uint qk_dim = 256u; constexpr uint qk_dim = 256u;
@@ -2249,8 +2420,8 @@ kernel void kernel_glm_qk_lowrank_q8_0_glm52_sg(
const uint head = tgpig.x; const uint head = tgpig.x;
const uint wt = args.weight_type; const uint wt = args.weight_type;
if (head >= n_head || if (head >= args.n_head ||
args.n_head != n_head || (args.n_head != 32u && args.n_head != 64u) ||
args.kv_lora_dim != kv_lora_dim || args.kv_lora_dim != kv_lora_dim ||
args.qk_nope != qk_nope || args.qk_nope != qk_nope ||
args.qk_dim != qk_dim || args.qk_dim != qk_dim ||
@@ -3624,7 +3795,7 @@ kernel void kernel_glm_attention_indexed_batch_lora_group8_vec_impl(
args.n_selected == 0u || args.n_selected == 0u ||
args.cache_f16 == 0u || args.cache_f16 == 0u ||
args.kv_lora_dim != 512u || args.kv_lora_dim != 512u ||
args.qk_rope != 64u) { (args.qk_rope != 0u && args.qk_rope != 64u)) {
return; return;
} }
@@ -3668,7 +3839,7 @@ kernel void kernel_glm_attention_indexed_batch_lora_group8_vec_impl(
} }
float corr_dims[2] = {0.0f, 0.0f}; float corr_dims[2] = {0.0f, 0.0f};
if (args.ext_factor != 0.0f) { if (args.qk_rope != 0u && args.ext_factor != 0.0f) {
glm_rope_yarn_corr_dims((int)args.qk_rope, glm_rope_yarn_corr_dims((int)args.qk_rope,
(int)args.n_ctx_orig, (int)args.n_ctx_orig,
args.freq_base, args.freq_base,
@@ -3822,9 +3993,8 @@ kernel void kernel_glm_attention_indexed_batch_lora_group8_vec_causal_impl(
const uint head = tgpig.x * group_heads + head_in_group + args.head_base; const uint head = tgpig.x * group_heads + head_in_group + args.head_base;
if (token >= args.n_tokens || if (token >= args.n_tokens ||
args.n_selected == 0u || args.n_selected == 0u ||
args.cache_f16 == 0u ||
args.kv_lora_dim != 512u || args.kv_lora_dim != 512u ||
args.qk_rope != 64u) { (args.qk_rope != 0u && args.qk_rope != 64u)) {
return; return;
} }
@@ -3869,7 +4039,7 @@ kernel void kernel_glm_attention_indexed_batch_lora_group8_vec_causal_impl(
} }
float corr_dims[2] = {0.0f, 0.0f}; float corr_dims[2] = {0.0f, 0.0f};
if (args.ext_factor != 0.0f) { if (args.qk_rope != 0u && args.ext_factor != 0.0f) {
glm_rope_yarn_corr_dims((int)args.qk_rope, glm_rope_yarn_corr_dims((int)args.qk_rope,
(int)args.n_ctx_orig, (int)args.n_ctx_orig,
args.freq_base, args.freq_base,
@@ -3891,10 +4061,17 @@ kernel void kernel_glm_attention_indexed_batch_lora_group8_vec_causal_impl(
const uint rr = off / kv_vecs; const uint rr = off / kv_vecs;
const uint vv = off - rr * kv_vecs; const uint vv = off - rr * kv_vecs;
const uint row = base + rr; const uint row = base + rr;
device const half4 *src = if (args.cache_f16 != 0u) {
(device const half4 *)((device const half *)kv_lora_cache + device const half4 *src =
(uint64_t)row * args.kv_lora_dim); (device const half4 *)((device const half *)kv_lora_cache +
kv_shared[off] = src[vv]; (uint64_t)row * args.kv_lora_dim);
kv_shared[off] = src[vv];
} else {
device const float4 *src =
(device const float4 *)((device const float *)kv_lora_cache +
(uint64_t)row * args.kv_lora_dim);
kv_shared[off] = (half4)src[vv];
}
} }
for (uint off = tid; off < rows * rope_vecs; off += 256u) { for (uint off = tid; off < rows * rope_vecs; off += 256u) {
const uint rr = off / rope_vecs; const uint rr = off / rope_vecs;
@@ -3903,29 +4080,31 @@ kernel void kernel_glm_attention_indexed_batch_lora_group8_vec_causal_impl(
const uint row = base + rr; const uint row = base + rr;
const uint64_t rope_base = (uint64_t)row * args.qk_rope; const uint64_t rope_base = (uint64_t)row * args.qk_rope;
const float2 y0 = const float2 y0 =
glm_cache_load_rotated_rope_pair_f16_only(k_rope_cache, glm_cache_load_rotated_rope_pair(k_rope_cache,
rope_base, rope_base,
r, r,
row, row,
args.qk_rope, args.qk_rope,
args.freq_base, args.cache_f16,
args.freq_scale, args.freq_base,
args.ext_factor, args.freq_scale,
args.attn_factor, args.ext_factor,
corr_dims[0], args.attn_factor,
corr_dims[1]); corr_dims[0],
corr_dims[1]);
const float2 y1 = const float2 y1 =
glm_cache_load_rotated_rope_pair_f16_only(k_rope_cache, glm_cache_load_rotated_rope_pair(k_rope_cache,
rope_base, rope_base,
r + 2u, r + 2u,
row, row,
args.qk_rope, args.qk_rope,
args.freq_base, args.cache_f16,
args.freq_scale, args.freq_base,
args.ext_factor, args.freq_scale,
args.attn_factor, args.ext_factor,
corr_dims[0], args.attn_factor,
corr_dims[1]); corr_dims[0],
corr_dims[1]);
rope_shared[off] = float4(y0.x, y0.y, y1.x, y1.y); rope_shared[off] = float4(y0.x, y0.y, y1.x, y1.y);
} }
threadgroup_barrier(mem_flags::mem_threadgroup); threadgroup_barrier(mem_flags::mem_threadgroup);
@@ -4586,14 +4765,15 @@ kernel void kernel_glm_router_select_one(
threadgroup float *scratch [[threadgroup(0)]], threadgroup float *scratch [[threadgroup(0)]],
uint token [[threadgroup_position_in_grid]], uint token [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]]) { uint tid [[thread_position_in_threadgroup]]) {
const uint sort_width = args.n_expert > 256u ? 512u : 256u;
threadgroup float *sel_scores = scratch; threadgroup float *sel_scores = scratch;
threadgroup int32_t *idx = (threadgroup int32_t *)(scratch + 256); threadgroup int32_t *idx = (threadgroup int32_t *)(scratch + sort_width);
device const float *token_logits = logits + (uint64_t)token * args.n_expert; device const float *token_logits = logits + (uint64_t)token * args.n_expert;
device int32_t *token_selected = selected + (uint64_t)token * args.n_expert_used; device int32_t *token_selected = selected + (uint64_t)token * args.n_expert_used;
device float *token_weights = weights + (uint64_t)token * args.n_expert_used; device float *token_weights = weights + (uint64_t)token * args.n_expert_used;
device float *token_probs = probs + (uint64_t)token * args.n_expert; device float *token_probs = probs + (uint64_t)token * args.n_expert;
const uint n_expert = min(args.n_expert, 256u); const uint n_expert = min(args.n_expert, 512u);
const bool active = tid < n_expert; const bool active = tid < n_expert;
const float p = active ? ds4_glm_router_sigmoid(token_logits[tid]) : 0.0f; const float p = active ? ds4_glm_router_sigmoid(token_logits[tid]) : 0.0f;
if (active) token_probs[tid] = p; if (active) token_probs[tid] = p;
@@ -4601,7 +4781,7 @@ kernel void kernel_glm_router_select_one(
idx[tid] = (int32_t)tid; idx[tid] = (int32_t)tid;
threadgroup_barrier(mem_flags::mem_threadgroup); threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint k = 2; k <= 256; k <<= 1) { for (uint k = 2; k <= sort_width; k <<= 1) {
for (uint j = k >> 1; j > 0; j >>= 1) { for (uint j = k >> 1; j > 0; j >>= 1) {
const uint other = tid ^ j; const uint other = tid ^ j;
if (other > tid) { if (other > tid) {
metal/glm53_bf16.metal
+143
View File
@@ -0,0 +1,143 @@
// BF16 model-weight kernels used by GLM-5.3 Flash.
static inline float glm53_bf16_to_f32(ushort value) {
return as_type<float>((uint)value << 16);
}
struct glm53_bf16_matmul_args {
uint in_dim;
uint out_dim;
uint n_rows;
};
kernel void kernel_glm53_embedding_bf16(
constant glm53_bf16_matmul_args &args,
device const ushort *weights,
device const int *tokens,
device float *out,
uint2 gid [[thread_position_in_grid]]) {
const uint d = gid.x;
const uint row = gid.y;
if (d >= args.in_dim || row >= args.n_rows) return;
const int token = tokens[row];
out[(ulong)row * args.in_dim + d] =
token >= 0 && (uint)token < args.out_dim
? glm53_bf16_to_f32(weights[(ulong)(uint)token * args.in_dim + d])
: 0.0f;
}
static inline void glm53_mul_mv_bf16_f32_row(
constant glm53_bf16_matmul_args &args,
device const ushort *weights,
device const float *x,
device float *out,
uint2 tgpig,
ushort lane,
ushort sg,
ushort nsg) {
const uint out_row = tgpig.x * (uint)nsg + sg;
const uint token = tgpig.y;
if (out_row >= args.out_dim || token >= args.n_rows) return;
device const ushort *w = weights + (ulong)out_row * args.in_dim;
device const float *xr = x + (ulong)token * args.in_dim;
float sum = 0.0f;
uint k = lane;
for (; k + 224u < args.in_dim; k += 256u) {
const ushort w0 = w[k];
const ushort w1 = w[k + 32u];
const ushort w2 = w[k + 64u];
const ushort w3 = w[k + 96u];
const ushort w4 = w[k + 128u];
const ushort w5 = w[k + 160u];
const ushort w6 = w[k + 192u];
const ushort w7 = w[k + 224u];
const float x0 = xr[k];
const float x1 = xr[k + 32u];
const float x2 = xr[k + 64u];
const float x3 = xr[k + 96u];
const float x4 = xr[k + 128u];
const float x5 = xr[k + 160u];
const float x6 = xr[k + 192u];
const float x7 = xr[k + 224u];
sum = fma(glm53_bf16_to_f32(w0), x0, sum);
sum = fma(glm53_bf16_to_f32(w1), x1, sum);
sum = fma(glm53_bf16_to_f32(w2), x2, sum);
sum = fma(glm53_bf16_to_f32(w3), x3, sum);
sum = fma(glm53_bf16_to_f32(w4), x4, sum);
sum = fma(glm53_bf16_to_f32(w5), x5, sum);
sum = fma(glm53_bf16_to_f32(w6), x6, sum);
sum = fma(glm53_bf16_to_f32(w7), x7, sum);
}
for (; k < args.in_dim; k += 32u) {
sum = fma(glm53_bf16_to_f32(w[k]), xr[k], sum);
}
sum = simd_sum(sum);
if (lane == 0u) out[(ulong)token * args.out_dim + out_row] = sum;
}
/* One simdgroup owns one output row. Eight independent loads expose enough
* memory-level parallelism for decode without changing the reduction tree. */
kernel void kernel_glm53_mul_mv_bf16_f32(
constant glm53_bf16_matmul_args &args,
device const ushort *weights,
device const float *x,
device float *out,
uint2 tgpig [[threadgroup_position_in_grid]],
ushort lane [[thread_index_in_simdgroup]],
ushort sg [[simdgroup_index_in_threadgroup]],
ushort nsg [[simdgroups_per_threadgroup]]) {
glm53_mul_mv_bf16_f32_row(args, weights, x, out,
tgpig, lane, sg, nsg);
}
kernel void kernel_glm53_mul_mv_bf16_f32_qkv(
constant glm53_bf16_matmul_args &args,
device const ushort *weights_q,
device const ushort *weights_k,
device const ushort *weights_v,
device const float *x,
device float *out_q,
device float *out_k,
device float *out_v,
uint3 tgpig [[threadgroup_position_in_grid]],
ushort lane [[thread_index_in_simdgroup]],
ushort sg [[simdgroup_index_in_threadgroup]],
ushort nsg [[simdgroups_per_threadgroup]]) {
device const ushort *weights = tgpig.z == 0u ? weights_q :
(tgpig.z == 1u ? weights_k : weights_v);
device float *out = tgpig.z == 0u ? out_q :
(tgpig.z == 1u ? out_k : out_v);
glm53_mul_mv_bf16_f32_row(args, weights, x, out,
tgpig.xy, lane, sg, nsg);
}
struct glm53_bf16_block16 {
ushort v[16];
};
template <typename type4x4>
void glm53_dequantize_bf16(
device const glm53_bf16_block16 *src,
short il,
thread type4x4 &reg) {
(void)il;
float4x4 values;
for (short i = 0; i < 16; i++) {
values[i / 4][i % 4] = glm53_bf16_to_f32(src->v[i]);
}
reg = (type4x4)values;
}
typedef decltype(kernel_mul_mm<
half, half4x4, simdgroup_half8x8,
half, half2x4, simdgroup_half8x8,
glm53_bf16_block16, 1, glm53_dequantize_bf16,
float, float4x4, float, float2x4>) glm53_mul_mm_bf16_t;
template [[host_name("kernel_glm53_mul_mm_bf16_f32")]]
kernel glm53_mul_mm_bf16_t kernel_mul_mm<
half, half4x4, simdgroup_half8x8,
half, half2x4, simdgroup_half8x8,
glm53_bf16_block16, 1, glm53_dequantize_bf16,
half, half4x4, float, float2x4>;
metal/glm53_kda.metal
+314
View File
@@ -0,0 +1,314 @@
// Kimi Delta Attention kernels, adapted from the kimi-k3 branch.
struct glm53_kda_args {
uint n_heads;
uint n_rows;
float lower_bound;
float norm_eps;
};
/*
* One threadgroup owns one (sequence, head). Four simdgroups update four
* value rows concurrently; every lane owns four adjacent key columns.
*/
kernel void kernel_glm53_kda_decode(
constant glm53_kda_args &args,
device const float *q_in,
device const float *k_in,
device const float *v_in,
device const float *raw_gate,
device const float *raw_beta,
device const float *output_gate,
device const float *q_conv,
device const float *k_conv,
device const float *v_conv,
device const float *a_log,
device const float *dt_bias,
device const float *output_norm,
device float *conv_state,
device float *state,
device float *out,
threadgroup float *scratch [[threadgroup(0)]],
uint2 tgpig [[threadgroup_position_in_grid]],
ushort tid [[thread_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]],
ushort sg [[simdgroup_index_in_threadgroup]]) {
constexpr uint D = 128u;
constexpr uint HISTORY = 3u;
const uint row = tgpig.x;
const uint head = tgpig.y;
if (row >= args.n_rows || head >= args.n_heads) return;
threadgroup float *sq = scratch;
threadgroup float *sk = sq + D;
threadgroup float *sd = sk + D;
threadgroup float *sv = sd + D;
threadgroup float *so = sv + D;
threadgroup float *reduce_q = so + D;
threadgroup float *reduce_k = reduce_q + 4u;
threadgroup float *reduce_o = reduce_k + 4u;
threadgroup float *beta_shared = reduce_o + 4u;
const uint projection = args.n_heads * D;
const uint channel = head * D + tid;
const ulong input_base = (ulong)row * projection + head * D;
const ulong conv_row_stride = 3ul * HISTORY * projection;
if (tid < D) {
float q_acc = 0.0f;
float k_acc = 0.0f;
float v_acc = 0.0f;
device float *q_state = conv_state +
(ulong)row * conv_row_stride;
device float *k_state = q_state + HISTORY * projection;
device float *v_state = k_state + HISTORY * projection;
for (uint w = 0; w < HISTORY; w++) {
q_acc = fma(q_state[(ulong)w * projection + channel],
q_conv[(ulong)channel * 4u + w], q_acc);
k_acc = fma(k_state[(ulong)w * projection + channel],
k_conv[(ulong)channel * 4u + w], k_acc);
v_acc = fma(v_state[(ulong)w * projection + channel],
v_conv[(ulong)channel * 4u + w], v_acc);
}
const float q_new = q_in[input_base + tid];
const float k_new = k_in[input_base + tid];
const float v_new = v_in[input_base + tid];
q_acc = fma(q_new, q_conv[(ulong)channel * 4u + 3u], q_acc);
k_acc = fma(k_new, k_conv[(ulong)channel * 4u + 3u], k_acc);
v_acc = fma(v_new, v_conv[(ulong)channel * 4u + 3u], v_acc);
q_state[channel] = q_state[projection + channel];
q_state[projection + channel] = q_state[2ul * projection + channel];
q_state[2ul * projection + channel] = q_new;
k_state[channel] = k_state[projection + channel];
k_state[projection + channel] = k_state[2ul * projection + channel];
k_state[2ul * projection + channel] = k_new;
v_state[channel] = v_state[projection + channel];
v_state[projection + channel] = v_state[2ul * projection + channel];
v_state[2ul * projection + channel] = v_new;
sq[tid] = q_acc / (1.0f + exp(-q_acc));
sk[tid] = k_acc / (1.0f + exp(-k_acc));
sv[tid] = v_acc / (1.0f + exp(-v_acc));
const float gate = raw_gate[input_base + tid] + dt_bias[channel];
sd[tid] = exp(args.lower_bound *
(1.0f / (1.0f + exp(-exp(a_log[head]) * gate))));
}
if (tid == 0u) {
beta_shared[0] =
1.0f / (1.0f + exp(-raw_beta[(ulong)row * args.n_heads + head]));
}
threadgroup_barrier(mem_flags::mem_threadgroup |
mem_flags::mem_device);
float q_sumsq = sq[tid] * sq[tid];
float k_sumsq = sk[tid] * sk[tid];
q_sumsq = simd_sum(q_sumsq);
k_sumsq = simd_sum(k_sumsq);
if (lane == 0u) {
reduce_q[sg] = q_sumsq;
reduce_k[sg] = k_sumsq;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
float q_total = lane < 4u ? reduce_q[lane] : 0.0f;
float k_total = lane < 4u ? reduce_k[lane] : 0.0f;
q_total = simd_sum(q_total);
k_total = simd_sum(k_total);
const float q_scale = rsqrt(q_total + 1.0e-6f) * 0x1.6a09e6p-4f;
const float k_scale = rsqrt(k_total + 1.0e-6f);
if (tid < D) {
sq[tid] *= q_scale;
sk[tid] *= k_scale;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
const uint k0 = lane * 4u;
const float4 q4 = *((threadgroup float4 *)(sq + k0));
const float4 k4 = *((threadgroup float4 *)(sk + k0));
const float4 decay4 = *((threadgroup float4 *)(sd + k0));
const ulong state_head =
((ulong)row * args.n_heads + head) * D * D;
for (uint value = sg; value < D; value += 4u) {
device float4 *hptr =
(device float4 *)(state + state_head + (ulong)value * D + k0);
float4 h = *hptr * decay4;
float hk = dot(h, k4);
hk = simd_sum(hk);
const float delta_v = (sv[value] - hk) * beta_shared[0];
h = fma(k4, float4(delta_v), h);
*hptr = h;
float hq = simd_sum(dot(h, q4));
if (lane == 0u) so[value] = hq;
}
threadgroup_barrier(mem_flags::mem_threadgroup |
mem_flags::mem_device);
float o_sumsq = so[tid] * so[tid];
o_sumsq = simd_sum(o_sumsq);
if (lane == 0u) reduce_o[sg] = o_sumsq;
threadgroup_barrier(mem_flags::mem_threadgroup);
float o_total = lane < 4u ? reduce_o[lane] : 0.0f;
o_total = simd_sum(o_total);
const float o_scale = rsqrt(o_total / (float)D + args.norm_eps);
if (tid < D) {
const ulong index = input_base + tid;
const float gate =
1.0f / (1.0f + exp(-output_gate[index]));
out[index] = so[tid] * o_scale * output_norm[tid] * gate;
}
}
kernel void kernel_glm53_kda_prefill_prepare(
constant glm53_kda_args &args,
device float *q,
device float *k,
device float *v,
device float *raw_gate,
device const float *q_conv,
device const float *k_conv,
device const float *v_conv,
device const float *a_log,
device const float *dt_bias,
device float *conv_state,
threadgroup float *scratch [[threadgroup(0)]],
uint head [[threadgroup_position_in_grid]],
ushort tid [[thread_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]],
ushort sg [[simdgroup_index_in_threadgroup]]) {
constexpr uint D = 128u;
constexpr uint HISTORY = 3u;
if (head >= args.n_heads) return;
threadgroup float *sq = scratch;
threadgroup float *sk = sq + D;
threadgroup float *reduce_q = sk + D;
threadgroup float *reduce_k = reduce_q + 4u;
const uint projection = args.n_heads * D;
const uint channel = head * D + tid;
device float *q_state = conv_state;
device float *k_state = q_state + HISTORY * projection;
device float *v_state = k_state + HISTORY * projection;
for (uint token = 0; token < args.n_rows; token++) {
const ulong index = (ulong)token * projection + channel;
float q_acc = 0.0f;
float k_acc = 0.0f;
float v_acc = 0.0f;
for (uint w = 0; w < HISTORY; w++) {
q_acc = fma(q_state[(ulong)w * projection + channel],
q_conv[(ulong)channel * 4u + w], q_acc);
k_acc = fma(k_state[(ulong)w * projection + channel],
k_conv[(ulong)channel * 4u + w], k_acc);
v_acc = fma(v_state[(ulong)w * projection + channel],
v_conv[(ulong)channel * 4u + w], v_acc);
}
const float q_new = q[index];
const float k_new = k[index];
const float v_new = v[index];
q_acc = fma(q_new, q_conv[(ulong)channel * 4u + 3u], q_acc);
k_acc = fma(k_new, k_conv[(ulong)channel * 4u + 3u], k_acc);
v_acc = fma(v_new, v_conv[(ulong)channel * 4u + 3u], v_acc);
q_state[channel] = q_state[projection + channel];
q_state[projection + channel] = q_state[2ul * projection + channel];
q_state[2ul * projection + channel] = q_new;
k_state[channel] = k_state[projection + channel];
k_state[projection + channel] = k_state[2ul * projection + channel];
k_state[2ul * projection + channel] = k_new;
v_state[channel] = v_state[projection + channel];
v_state[projection + channel] = v_state[2ul * projection + channel];
v_state[2ul * projection + channel] = v_new;
sq[tid] = q_acc / (1.0f + exp(-q_acc));
sk[tid] = k_acc / (1.0f + exp(-k_acc));
v[index] = v_acc / (1.0f + exp(-v_acc));
const float gate = raw_gate[index] + dt_bias[channel];
raw_gate[index] = exp(args.lower_bound *
(1.0f / (1.0f + exp(-exp(a_log[head]) * gate))));
threadgroup_barrier(mem_flags::mem_threadgroup |
mem_flags::mem_device);
float q_sumsq = simd_sum(sq[tid] * sq[tid]);
float k_sumsq = simd_sum(sk[tid] * sk[tid]);
if (lane == 0u) {
reduce_q[sg] = q_sumsq;
reduce_k[sg] = k_sumsq;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
float q_total = lane < 4u ? reduce_q[lane] : 0.0f;
float k_total = lane < 4u ? reduce_k[lane] : 0.0f;
q_total = simd_sum(q_total);
k_total = simd_sum(k_total);
q[index] = sq[tid] * rsqrt(q_total + 1.0e-6f) *
0x1.6a09e6p-4f;
k[index] = sk[tid] * rsqrt(k_total + 1.0e-6f);
threadgroup_barrier(mem_flags::mem_threadgroup |
mem_flags::mem_device);
}
}
kernel void kernel_glm53_kda_prefill_recurrence(
constant glm53_kda_args &args,
device const float *q,
device const float *k,
device const float *v,
device const float *decay,
device const float *raw_beta,
device float *state,
device float *out,
uint2 tgpig [[threadgroup_position_in_grid]],
ushort lane [[thread_index_in_simdgroup]],
ushort sg [[simdgroup_index_in_threadgroup]]) {
constexpr uint D = 128u;
const uint head = tgpig.x;
const uint value = tgpig.y * 4u + sg;
if (head >= args.n_heads || value >= D) return;
const uint projection = args.n_heads * D;
const uint k0 = lane * 4u;
device float4 *state_ptr = (device float4 *)(
state + ((ulong)head * D + value) * D + k0);
float4 h = *state_ptr;
for (uint token = 0; token < args.n_rows; token++) {
const ulong base = (ulong)token * projection + head * D;
const float4 q4 = *((device const float4 *)(q + base + k0));
const float4 k4 = *((device const float4 *)(k + base + k0));
const float4 decay4 =
*((device const float4 *)(decay + base + k0));
h *= decay4;
const float hk = simd_sum(dot(h, k4));
const float beta = 1.0f /
(1.0f + exp(-raw_beta[(ulong)token * args.n_heads + head]));
const float delta_v = (v[base + value] - hk) * beta;
h = fma(k4, float4(delta_v), h);
const float result = simd_sum(dot(h, q4));
if (lane == 0u) out[base + value] = result;
}
*state_ptr = h;
}
kernel void kernel_glm53_kda_prefill_output(
constant glm53_kda_args &args,
device float *out,
device const float *output_gate,
device const float *output_norm,
threadgroup float *partial [[threadgroup(0)]],
uint2 tgpig [[threadgroup_position_in_grid]],
ushort tid [[thread_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]],
ushort sg [[simdgroup_index_in_threadgroup]]) {
constexpr uint D = 128u;
const uint token = tgpig.x;
const uint head = tgpig.y;
if (token >= args.n_rows || head >= args.n_heads) return;
const uint projection = args.n_heads * D;
const ulong base = (ulong)token * projection + head * D;
const float raw = out[base + tid];
float sumsq = simd_sum(raw * raw);
if (lane == 0u) partial[sg] = sumsq;
threadgroup_barrier(mem_flags::mem_threadgroup);
float total = lane < 4u ? partial[lane] : 0.0f;
total = simd_sum(total);
const float scale = rsqrt(total / (float)D + args.norm_eps);
out[base + tid] = raw * scale * output_norm[tid] /
(1.0f + exp(-output_gate[base + tid]));
}
metal/glm53_vision.metal
+271
View File
@@ -0,0 +1,271 @@
// GLM-5.3 Flash vision operations not covered by the shared BF16 matmuls.
struct glm53_vision_rows_args {
uint width;
uint rows;
float eps;
};
struct glm53_vision_qkv_args {
uint rows;
uint grid_h;
uint grid_w;
float eps;
};
struct glm53_vision_attention_args {
uint rows;
float scale;
};
struct glm53_vision_scatter_args {
uint dst_row;
uint image_row;
uint rows;
uint total_rows;
uint width;
uint hc;
};
static inline float glm53_vision_erf(float x) {
const float sign = x < 0.0f ? -1.0f : 1.0f;
const float a = abs(x);
const float t = 1.0f / (1.0f + 0.3275911f * a);
const float p = (((((1.061405429f * t - 1.453152027f) * t) +
1.421413741f) * t - 0.284496736f) * t +
0.254829592f) * t;
return sign * (1.0f - p * exp(-a * a));
}
kernel void kernel_glm53_vision_add_bias(
constant glm53_vision_rows_args &args,
device float *x,
device const ushort *bias,
uint2 gid [[thread_position_in_grid]]) {
if (gid.x >= args.width || gid.y >= args.rows) return;
x[(ulong)gid.y * args.width + gid.x] += glm53_bf16_to_f32(bias[gid.x]);
}
kernel void kernel_glm53_vision_rms_bf16(
constant glm53_vision_rows_args &args,
device const float *x,
device const ushort *weight,
device float *out,
threadgroup float *partial,
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]],
ushort sg [[simdgroup_index_in_threadgroup]],
ushort nsg [[simdgroups_per_threadgroup]]) {
if (row >= args.rows) return;
device const float *xr = x + (ulong)row * args.width;
device float *yr = out + (ulong)row * args.width;
float sum = 0.0f;
for (uint d = tid; d < args.width; d += 256u) sum = fma(xr[d], xr[d], sum);
sum = simd_sum(sum);
if (lane == 0u) partial[sg] = sum;
threadgroup_barrier(mem_flags::mem_threadgroup);
if (sg == 0u) {
float v = lane < nsg ? partial[lane] : 0.0f;
v = simd_sum(v);
if (lane == 0u) partial[0] = rsqrt(v / (float)args.width + args.eps);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
const float inv = partial[0];
for (uint d = tid; d < args.width; d += 256u) {
yr[d] = xr[d] * inv * glm53_bf16_to_f32(weight[d]);
}
}
kernel void kernel_glm53_vision_qkv_rope(
constant glm53_vision_qkv_args &args,
device const float *qkv,
device const ushort *bias,
device const ushort *q_weight,
device const ushort *k_weight,
device float *q,
device float *k,
device float *v,
uint2 group [[threadgroup_position_in_grid]],
ushort lane [[thread_index_in_simdgroup]]) {
const uint row = group.x;
const uint head = group.y;
if (row >= args.rows || head >= 16u) return;
const ulong qkv_base = (ulong)row * 3072u + (ulong)head * 64u;
const ulong out_base = (ulong)row * 1024u + (ulong)head * 64u;
float q0 = qkv[qkv_base + lane] + glm53_bf16_to_f32(bias[(ulong)head * 64u + lane]);
float q1 = qkv[qkv_base + lane + 32u] +
glm53_bf16_to_f32(bias[(ulong)head * 64u + lane + 32u]);
float k0 = qkv[qkv_base + 1024u + lane] +
glm53_bf16_to_f32(bias[1024u + (ulong)head * 64u + lane]);
float k1 = qkv[qkv_base + 1024u + lane + 32u] +
glm53_bf16_to_f32(bias[1024u + (ulong)head * 64u + lane + 32u]);
const float qsum = simd_sum(fma(q0, q0, q1 * q1));
const float ksum = simd_sum(fma(k0, k0, k1 * k1));
const float qinv = rsqrt(qsum / 64.0f + args.eps);
const float kinv = rsqrt(ksum / 64.0f + args.eps);
q0 *= qinv * glm53_bf16_to_f32(q_weight[lane]);
q1 *= qinv * glm53_bf16_to_f32(q_weight[lane + 32u]);
k0 *= kinv * glm53_bf16_to_f32(k_weight[lane]);
k1 *= kinv * glm53_bf16_to_f32(k_weight[lane + 32u]);
const uint merge_w = args.grid_w / 2u;
const uint group_index = row / 4u;
const uint within = row & 3u;
const uint py = (group_index / merge_w) * 2u + within / 2u;
const uint px = (group_index % merge_w) * 2u + within % 2u;
const uint freq_index = lane & 15u;
const uint pos = lane < 16u ? py : px;
const float inv_freq = powr(10000.0f, -(float)freq_index / 16.0f);
const float angle = (float)pos * inv_freq;
const float cs = cos(angle);
const float sn = sin(angle);
q[out_base + lane] = q0 * cs - q1 * sn;
q[out_base + lane + 32u] = q1 * cs + q0 * sn;
k[out_base + lane] = k0 * cs - k1 * sn;
k[out_base + lane + 32u] = k1 * cs + k0 * sn;
v[out_base + lane] = qkv[qkv_base + 2048u + lane] +
glm53_bf16_to_f32(bias[2048u + (ulong)head * 64u + lane]);
v[out_base + lane + 32u] = qkv[qkv_base + 2048u + lane + 32u] +
glm53_bf16_to_f32(bias[2048u + (ulong)head * 64u + lane + 32u]);
}
/* A simdgroup owns one query/head and keeps its 64 output values in registers.
* This is quadratic in compute, as the model graph requires, but linear in
* memory and never materializes the attention matrix. */
kernel void kernel_glm53_vision_attention(
constant glm53_vision_attention_args &args,
device const float *q,
device const float *k,
device const float *v,
device float *out,
uint2 group [[threadgroup_position_in_grid]],
ushort lane [[thread_index_in_simdgroup]]) {
const uint row = group.x;
const uint head = group.y;
if (row >= args.rows || head >= 16u) return;
const ulong base = (ulong)row * 1024u + (ulong)head * 64u;
const float q0 = q[base + lane];
const float q1 = q[base + lane + 32u];
float acc0 = 0.0f, acc1 = 0.0f;
float max_score = -INFINITY;
float denom = 0.0f;
for (uint key_row = 0; key_row < args.rows; key_row++) {
const ulong kb = (ulong)key_row * 1024u + (ulong)head * 64u;
float score = simd_sum(q0 * k[kb + lane] + q1 * k[kb + lane + 32u]);
score *= args.scale;
const float next_max = max(max_score, score);
const float old_scale = max_score == -INFINITY ? 0.0f : exp(max_score - next_max);
const float new_scale = exp(score - next_max);
denom = denom * old_scale + new_scale;
acc0 = acc0 * old_scale + new_scale * v[kb + lane];
acc1 = acc1 * old_scale + new_scale * v[kb + lane + 32u];
max_score = next_max;
}
out[base + lane] = acc0 / denom;
out[base + lane + 32u] = acc1 / denom;
}
kernel void kernel_glm53_vision_bias_residual(
constant glm53_vision_rows_args &args,
device float *x,
device const ushort *bias,
device const float *residual,
uint2 gid [[thread_position_in_grid]]) {
if (gid.x >= args.width || gid.y >= args.rows) return;
const ulong off = (ulong)gid.y * args.width + gid.x;
x[off] += glm53_bf16_to_f32(bias[gid.x]) + residual[off];
}
kernel void kernel_glm53_vision_swiglu_bias(
constant glm53_vision_rows_args &args,
device const float *gate,
device const ushort *gate_bias,
device const float *up,
device const ushort *up_bias,
device float *out,
uint2 gid [[thread_position_in_grid]]) {
if (gid.x >= args.width || gid.y >= args.rows) return;
const ulong off = (ulong)gid.y * args.width + gid.x;
const float g = min(gate[off] + glm53_bf16_to_f32(gate_bias[gid.x]), 10.0f);
const float u = clamp(up[off] + glm53_bf16_to_f32(up_bias[gid.x]), -10.0f, 10.0f);
out[off] = (g / (1.0f + exp(-g))) * u;
}
kernel void kernel_glm53_vision_downsample_reorder(
device const float *x,
device float *out,
uint2 gid [[thread_position_in_grid]]) {
const uint d = gid.x;
const uint row = gid.y;
if (d >= 4096u) return;
const uint channel = d / 4u;
const uint within = d & 3u;
out[(ulong)row * 4096u + d] = x[((ulong)row * 4u + within) * 1024u + channel];
}
kernel void kernel_glm53_vision_layernorm_gelu(
constant glm53_vision_rows_args &args,
device const float *x,
device const ushort *weight,
device const ushort *bias,
device float *out,
threadgroup float *partial,
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]],
ushort sg [[simdgroup_index_in_threadgroup]],
ushort nsg [[simdgroups_per_threadgroup]]) {
if (row >= args.rows) return;
device const float *xr = x + (ulong)row * args.width;
device float *yr = out + (ulong)row * args.width;
float sum = 0.0f;
for (uint d = tid; d < args.width; d += 256u) sum += xr[d];
sum = simd_sum(sum);
if (lane == 0u) partial[sg] = sum;
threadgroup_barrier(mem_flags::mem_threadgroup);
if (sg == 0u) {
float v = lane < nsg ? partial[lane] : 0.0f;
v = simd_sum(v);
if (lane == 0u) partial[0] = v / (float)args.width;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
const float mean = partial[0];
float var = 0.0f;
for (uint d = tid; d < args.width; d += 256u) {
const float centered = xr[d] - mean;
var = fma(centered, centered, var);
}
var = simd_sum(var);
if (lane == 0u) partial[sg] = var;
threadgroup_barrier(mem_flags::mem_threadgroup);
if (sg == 0u) {
float v = lane < nsg ? partial[lane] : 0.0f;
v = simd_sum(v);
if (lane == 0u) partial[0] = rsqrt(v / (float)args.width + args.eps);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
const float inv = partial[0];
const float inv_sqrt2 = 0.7071067811865475f;
for (uint d = tid; d < args.width; d += 256u) {
float v = (xr[d] - mean) * inv * glm53_bf16_to_f32(weight[d]) +
glm53_bf16_to_f32(bias[d]);
yr[d] = 0.5f * v * (1.0f + glm53_vision_erf(v * inv_sqrt2));
}
}
kernel void kernel_glm53_vision_scatter_hc(
constant glm53_vision_scatter_args &args,
device float *hc,
device const float *image,
uint2 gid [[thread_position_in_grid]]) {
const uint d = gid.x;
const uint linear_row = gid.y;
if (d >= args.width || linear_row >= args.rows * args.hc) return;
const uint image_delta = linear_row / args.hc;
const uint hc_index = linear_row % args.hc;
const ulong dst = ((ulong)(args.dst_row + image_delta) * args.hc + hc_index) *
args.width + d;
const ulong src = (ulong)(args.image_row + image_delta) * args.width + d;
if (args.dst_row + image_delta < args.total_rows) hc[dst] = image[src];
}
metal/moe.metal
+64 -25
View File
@@ -423,6 +423,7 @@ struct ds4_metal_glm_routed_moe_args {
uint32_t n_tokens; uint32_t n_tokens;
uint32_t mid_token_stride; uint32_t mid_token_stride;
uint32_t down_type; uint32_t down_type;
float swiglu_clamp;
/* Expert ownership under tensor parallelism: tp_world 0/1 = full /* Expert ownership under tensor parallelism: tp_world 0/1 = full
* compute; otherwise each rank owns a contiguous expert range. */ * compute; otherwise each rank owns a contiguous expert range. */
int32_t tp_rank; int32_t tp_rank;
@@ -436,6 +437,14 @@ struct ds4_metal_glm_routed_moe_args {
uint64_t down_row_bytes; uint64_t down_row_bytes;
}; };
static inline float ds4_glm_swiglu(float gate, float up, float limit) {
if (limit > 1.0e-6f) {
gate = min(gate, limit);
up = clamp(up, -limit, limit);
}
return (gate / (1.0f + exp(-gate))) * up;
}
static inline bool ds4_tp_owns_expert(int expert, int n_total, static inline bool ds4_tp_owns_expert(int expert, int n_total,
int tp_rank, int tp_world) { int tp_rank, int tp_world) {
@@ -712,18 +721,20 @@ kernel void kernel_glm_q4_K_pair_swiglu_f32(
const uint64_t mid_off = (uint64_t)token * args.mid_token_stride + const uint64_t mid_off = (uint64_t)token * args.mid_token_stride +
(uint64_t)slot * args.mid_dim + row; (uint64_t)slot * args.mid_dim + row;
const int expert = selected[selected_off]; const int expert = selected[selected_off];
if (expert < 0 || (uint)expert >= args.n_total_expert) { if (!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) {
if (tid == 0u) mid[mid_off] = 0.0f; if (tid == 0u) mid[mid_off] = 0.0f;
return; return;
} }
const int local_expert = expert - args.tp_expert_base;
device const block_q4_K *gate_row = device const block_q4_K *gate_row =
(device const block_q4_K *)(gate + (device const block_q4_K *)(gate +
(uint64_t)(uint)expert * args.gate_expert_bytes + (uint64_t)(uint)local_expert * args.gate_expert_bytes +
(uint64_t)row * args.gate_row_bytes); (uint64_t)row * args.gate_row_bytes);
device const block_q4_K *up_row = device const block_q4_K *up_row =
(device const block_q4_K *)(up + (device const block_q4_K *)(up +
(uint64_t)(uint)expert * args.up_expert_bytes + (uint64_t)(uint)local_expert * args.up_expert_bytes +
(uint64_t)row * args.up_row_bytes); (uint64_t)row * args.up_row_bytes);
float acc_gate = 0.0f; float acc_gate = 0.0f;
@@ -747,10 +758,8 @@ kernel void kernel_glm_q4_K_pair_swiglu_f32(
} }
if (tid == 0u) { if (tid == 0u) {
const float g = scratch[0]; mid[mid_off] = ds4_glm_swiglu(scratch[0], scratch[ntg],
const float u = scratch[ntg]; args.swiglu_clamp) * weights[selected_off];
const float sw = g / (1.0f + exp(-g));
mid[mid_off] = sw * u * weights[selected_off];
} }
} }
@@ -887,8 +896,8 @@ static inline void glm_q2_K_pair_swiglu_simd_f32_impl(
const float g = simd_sum(sumg[row]); const float g = simd_sum(sumg[row]);
const float u = simd_sum(sumu[row]); const float u = simd_sum(sumu[row]);
if (tiisg == 0u) { if (tiisg == 0u) {
const float sw = g / (1.0f + exp(-g)); mid[mid_base + row0 + (uint)row] =
mid[mid_base + row0 + (uint)row] = sw * u * weights[selected_off]; ds4_glm_swiglu(g, u, args.swiglu_clamp) * weights[selected_off];
} }
} }
@@ -944,7 +953,9 @@ kernel void kernel_glm_q2_K_addr_pair_swiglu2_f32(
(uint64_t)slot * args.mid_dim; (uint64_t)slot * args.mid_dim;
if (row0 >= args.mid_dim) return; if (row0 >= args.mid_dim) return;
if (expert < 0 || (uint)expert >= args.n_total_expert) { if (expert < 0 || (uint)expert >= args.n_total_expert ||
!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) {
if (tiisg == 0u) { if (tiisg == 0u) {
for (short row = 0; for (short row = 0;
row < N_R0_GLM_Q2_PAIR2_K && row0 + (uint)row < args.mid_dim; row < N_R0_GLM_Q2_PAIR2_K && row0 + (uint)row < args.mid_dim;
@@ -1005,7 +1016,9 @@ kernel void kernel_glm_q2_K_addr_pair_swiglu2_f32_masked(
(uint64_t)slot * args.mid_dim; (uint64_t)slot * args.mid_dim;
if (row0 >= args.mid_dim) return; if (row0 >= args.mid_dim) return;
if (expert < 0 || (uint)expert >= args.n_total_expert) { if (expert < 0 || (uint)expert >= args.n_total_expert ||
!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) {
if (tiisg == 0u) { if (tiisg == 0u) {
for (short row = 0; for (short row = 0;
row < N_R0_GLM_Q2_PAIR2_K && row0 + (uint)row < args.mid_dim; row < N_R0_GLM_Q2_PAIR2_K && row0 + (uint)row < args.mid_dim;
@@ -1191,8 +1204,8 @@ static inline void glm_q4_K_pair_swiglu_simd_f32_impl(
const float g = simd_sum(sumg[row]); const float g = simd_sum(sumg[row]);
const float u = simd_sum(sumu[row]); const float u = simd_sum(sumu[row]);
if (tiisg == 0u) { if (tiisg == 0u) {
const float sw = g / (1.0f + exp(-g)); mid[mid_base + row0 + (uint)row] =
mid[mid_base + row0 + (uint)row] = sw * u * weights[selected_off]; ds4_glm_swiglu(g, u, args.swiglu_clamp) * weights[selected_off];
} }
} }
@@ -1216,9 +1229,12 @@ kernel void kernel_glm_q4_K_pair_swiglu2_f32(
if (slot >= args.n_expert_used || token >= args.n_tokens) return; if (slot >= args.n_expert_used || token >= args.n_tokens) return;
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
const int expert = selected[selected_off]; const int expert = selected[selected_off];
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) return;
glm_q4_K_pair_swiglu_simd_f32_impl<N_R0_GLM_Q4_PAIR2_K>( glm_q4_K_pair_swiglu_simd_f32_impl<N_R0_GLM_Q4_PAIR2_K>(
args, gate, up, x, weights, mid, scratch, args, gate, up, x, weights, mid, scratch,
tgpig, slot, token, selected_off, expert, tiisg, sgitg); tgpig, slot, token, selected_off,
expert - args.tp_expert_base, tiisg, sgitg);
} }
kernel void kernel_glm_q4_K_addr_pair_swiglu_f32( kernel void kernel_glm_q4_K_addr_pair_swiglu_f32(
@@ -1245,7 +1261,9 @@ kernel void kernel_glm_q4_K_addr_pair_swiglu_f32(
(uint64_t)slot * args.mid_dim; (uint64_t)slot * args.mid_dim;
if (row0 >= args.mid_dim) return; if (row0 >= args.mid_dim) return;
if (expert < 0 || (uint)expert >= args.n_total_expert) { if (expert < 0 || (uint)expert >= args.n_total_expert ||
!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) {
if (tiisg == 0u) { if (tiisg == 0u) {
for (short row = 0; for (short row = 0;
row < N_R0_Q4_K && row0 + (uint)row < args.mid_dim; row < N_R0_Q4_K && row0 + (uint)row < args.mid_dim;
@@ -1306,7 +1324,9 @@ kernel void kernel_glm_q4_K_addr_pair_swiglu_f32_masked(
(uint64_t)slot * args.mid_dim; (uint64_t)slot * args.mid_dim;
if (row0 >= args.mid_dim) return; if (row0 >= args.mid_dim) return;
if (expert < 0 || (uint)expert >= args.n_total_expert) { if (expert < 0 || (uint)expert >= args.n_total_expert ||
!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) {
if (tiisg == 0u) { if (tiisg == 0u) {
for (short row = 0; for (short row = 0;
row < N_R0_Q4_K && row0 + (uint)row < args.mid_dim; row < N_R0_Q4_K && row0 + (uint)row < args.mid_dim;
@@ -1358,9 +1378,12 @@ kernel void kernel_glm_q4_K_pair_swiglu4_f32(
if (slot >= args.n_expert_used || token >= args.n_tokens) return; if (slot >= args.n_expert_used || token >= args.n_tokens) return;
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
const int expert = selected[selected_off]; const int expert = selected[selected_off];
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) return;
glm_q4_K_pair_swiglu_simd_f32_impl<N_R0_GLM_Q4_PAIR_K>( glm_q4_K_pair_swiglu_simd_f32_impl<N_R0_GLM_Q4_PAIR_K>(
args, gate, up, x, weights, mid, scratch, args, gate, up, x, weights, mid, scratch,
tgpig, slot, token, selected_off, expert, tiisg, sgitg); tgpig, slot, token, selected_off,
expert - args.tp_expert_base, tiisg, sgitg);
} }
kernel void kernel_glm_q4_K_pair_swiglu2_mapped_f32( kernel void kernel_glm_q4_K_pair_swiglu2_mapped_f32(
@@ -1378,6 +1401,8 @@ kernel void kernel_glm_q4_K_pair_swiglu2_mapped_f32(
ushort sgitg [[simdgroup_index_in_threadgroup]]) { ushort sgitg [[simdgroup_index_in_threadgroup]]) {
const uint expert = tgpig.z; const uint expert = tgpig.z;
if (expert >= args.n_total_expert) return; if (expert >= args.n_total_expert) return;
if (!ds4_tp_owns_expert((int)expert, args.n_total_expert,
args.tp_rank, args.tp_world)) return;
const uint count = htpe[expert]; const uint count = htpe[expert];
const uint map_base = tgpig.y * 32u; const uint map_base = tgpig.y * 32u;
for (uint i = 0; i < 32u; i++) { for (uint i = 0; i < 32u; i++) {
@@ -1391,7 +1416,8 @@ kernel void kernel_glm_q4_K_pair_swiglu2_mapped_f32(
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
glm_q4_K_pair_swiglu_simd_f32_impl<N_R0_Q4_K>( glm_q4_K_pair_swiglu_simd_f32_impl<N_R0_Q4_K>(
args, gate, up, x, weights, mid, scratch, args, gate, up, x, weights, mid, scratch,
tgpig, slot, token, selected_off, (int)expert, tiisg, sgitg); tgpig, slot, token, selected_off,
(int)expert - args.tp_expert_base, tiisg, sgitg);
} }
} }
@@ -1411,6 +1437,8 @@ kernel void kernel_glm_q4_K_pair_swiglu2_mapped_row_f32(
const uint expert = tgpig.z; const uint expert = tgpig.z;
const uint map_row = tgpig.y; const uint map_row = tgpig.y;
if (expert >= args.n_total_expert || map_row >= htpe[expert]) return; if (expert >= args.n_total_expert || map_row >= htpe[expert]) return;
if (!ds4_tp_owns_expert((int)expert, args.n_total_expert,
args.tp_rank, args.tp_world)) return;
const int id = hids[(uint64_t)expert * args.n_tokens + map_row]; const int id = hids[(uint64_t)expert * args.n_tokens + map_row];
if (id < 0) return; if (id < 0) return;
const uint token = (uint)id / args.n_expert_used; const uint token = (uint)id / args.n_expert_used;
@@ -1419,7 +1447,8 @@ kernel void kernel_glm_q4_K_pair_swiglu2_mapped_row_f32(
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
glm_q4_K_pair_swiglu_simd_f32_impl<N_R0_Q4_K>( glm_q4_K_pair_swiglu_simd_f32_impl<N_R0_Q4_K>(
args, gate, up, x, weights, mid, scratch, args, gate, up, x, weights, mid, scratch,
tgpig, slot, token, selected_off, (int)expert, tiisg, sgitg); tgpig, slot, token, selected_off,
(int)expert - args.tp_expert_base, tiisg, sgitg);
} }
static inline void glm_q5_K_pair_swiglu_f32_impl( static inline void glm_q5_K_pair_swiglu_f32_impl(
@@ -1591,8 +1620,8 @@ static inline void glm_q5_K_pair_swiglu_f32_impl(
const float g = simd_sum(sumg[row]); const float g = simd_sum(sumg[row]);
const float u = simd_sum(sumu[row]); const float u = simd_sum(sumu[row]);
if (tiisg == 0u) { if (tiisg == 0u) {
const float sw = g / (1.0f + exp(-g)); mid[mid_base + row0 + (uint)row] =
mid[mid_base + row0 + (uint)row] = sw * u * weights[selected_off]; ds4_glm_swiglu(g, u, args.swiglu_clamp) * weights[selected_off];
} }
} }
@@ -1913,6 +1942,8 @@ kernel void kernel_glm_q2_K_addr_down_f32(
for (uint slot = 0; slot < args.n_expert_used; slot++) { for (uint slot = 0; slot < args.n_expert_used; slot++) {
const int expert = selected[selected_base + slot]; const int expert = selected[selected_base + slot];
if (expert < 0 || (uint)expert >= args.n_total_expert) continue; if (expert < 0 || (uint)expert >= args.n_total_expert) continue;
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) continue;
const uint64_t down_addr = down_addrs[(uint)expert]; const uint64_t down_addr = down_addrs[(uint)expert];
if (down_addr == 0) continue; if (down_addr == 0) continue;
device const block_q2_K *x = device const block_q2_K *x =
@@ -1996,10 +2027,12 @@ kernel void kernel_glm_q4_K_down_f32(
const uint64_t mid_base = (uint64_t)token * args.mid_token_stride; const uint64_t mid_base = (uint64_t)token * args.mid_token_stride;
for (uint slot = 0; slot < args.n_expert_used; slot++) { for (uint slot = 0; slot < args.n_expert_used; slot++) {
const int expert = selected[selected_base + slot]; const int expert = selected[selected_base + slot];
if (expert < 0 || (uint)expert >= args.n_total_expert) continue; if (!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) continue;
device const block_q4_K *down_row = device const block_q4_K *down_row =
(device const block_q4_K *)(down + (device const block_q4_K *)(down +
(uint64_t)(uint)expert * args.down_expert_bytes + (uint64_t)(uint)(expert - args.tp_expert_base) *
args.down_expert_bytes +
(uint64_t)row * args.down_row_bytes); (uint64_t)row * args.down_row_bytes);
device const float *slot_mid = mid + mid_base + (uint64_t)slot * args.mid_dim; device const float *slot_mid = mid + mid_base + (uint64_t)slot * args.mid_dim;
for (uint k = tid; k < args.mid_dim; k += ntg) { for (uint k = tid; k < args.mid_dim; k += ntg) {
@@ -2038,6 +2071,8 @@ kernel void kernel_glm_q4_K_addr_down_f32(
for (uint slot = 0; slot < args.n_expert_used; slot++) { for (uint slot = 0; slot < args.n_expert_used; slot++) {
const int expert = selected[selected_base + slot]; const int expert = selected[selected_base + slot];
if (expert < 0 || (uint)expert >= args.n_total_expert) continue; if (expert < 0 || (uint)expert >= args.n_total_expert) continue;
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) continue;
const uint64_t down_addr = down_addrs[(uint)expert]; const uint64_t down_addr = down_addrs[(uint)expert];
if (down_addr == 0) continue; if (down_addr == 0) continue;
device const block_q4_K *down_row = device const block_q4_K *down_row =
@@ -2093,11 +2128,13 @@ kernel void kernel_glm_q4_K_down_simd_f32(
const uint64_t mid_base = (uint64_t)token * args.mid_token_stride; const uint64_t mid_base = (uint64_t)token * args.mid_token_stride;
for (uint slot = 0; slot < args.n_expert_used; slot++) { for (uint slot = 0; slot < args.n_expert_used; slot++) {
const int expert = selected[selected_base + slot]; const int expert = selected[selected_base + slot];
if (expert < 0 || (uint)expert >= args.n_total_expert) continue; if (!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) continue;
device const block_q4_K *x = device const block_q4_K *x =
(device const block_q4_K *)(down + (device const block_q4_K *)(down +
(uint64_t)(uint)expert * args.down_expert_bytes + (uint64_t)(uint)(expert - args.tp_expert_base) *
args.down_expert_bytes +
(uint64_t)row0 * args.down_row_bytes); (uint64_t)row0 * args.down_row_bytes);
device const float *y = mid + mid_base + (uint64_t)slot * args.mid_dim; device const float *y = mid + mid_base + (uint64_t)slot * args.mid_dim;
device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir;
@@ -2197,6 +2234,8 @@ kernel void kernel_glm_q4_K_addr_down_simd_f32(
for (uint slot = 0; slot < args.n_expert_used; slot++) { for (uint slot = 0; slot < args.n_expert_used; slot++) {
const int expert = selected[selected_base + slot]; const int expert = selected[selected_base + slot];
if (expert < 0 || (uint)expert >= args.n_total_expert) continue; if (expert < 0 || (uint)expert >= args.n_total_expert) continue;
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) continue;
const uint64_t down_addr = down_addrs[(uint)expert]; const uint64_t down_addr = down_addrs[(uint)expert];
if (down_addr == 0) continue; if (down_addr == 0) continue;
native/metal/README.md
+5 -6
View File
@@ -1,15 +1,14 @@
# Vendored DS4 Metal boundary # Vendored DS4 Metal boundary
These files track the DS4 Metal boundary audited at commit The Metal bridge and GPU declarations track the DS4 boundary audited at commit
`8db89fe083ae4d17c9a2428ccd29803d3ae8f577` (2026-08-28): `ec7642cdd9ec81d01ad4b1fd8f8a3d1511533748` (2026-08-31):
- `ds4_metal.m` - `ds4_metal.m`
- `ds4.h`
- `ds4_gpu.h` - `ds4_gpu.h`
- `ds4_ssd.h`
- `LICENSE`
The matching Metal kernels live in the repository-level `metal/` directory. The matching Metal kernels, including `glm53_*.metal`, live in the
repository-level `metal/` directory. The other headers remain DS4Server-owned
interfaces and are not copies of DS4 application code.
DS4Server builds and bundles this local snapshot; it does not read a sibling DS4Server builds and bundles this local snapshot; it does not read a sibling
DS4 checkout. Rust owns the model, graph, session, sampling, and lifecycle. DS4 checkout. Rust owns the model, graph, session, sampling, and lifecycle.
Objective-C remains only at the platform Metal boundary. Objective-C remains only at the platform Metal boundary.
native/metal/ds4_gpu.h
+223
View File
@@ -118,6 +118,11 @@ int ds4_gpu_build_derived_artifacts(const void *model_map, uint64_t model_size,
int ds4_gpu_model_range_replaced(const void *model_map, uint64_t offset, int ds4_gpu_model_range_replaced(const void *model_map, uint64_t offset,
uint64_t bytes); uint64_t bytes);
int ds4_gpu_set_model_map_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size, uint64_t max_tensor_bytes); int ds4_gpu_set_model_map_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size, uint64_t max_tensor_bytes);
/* Add a secondary GGUF mapping without replacing the primary model mapping. */
int ds4_gpu_set_aux_model_map_range(const void *model_map,
uint64_t model_size,
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); 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_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);
@@ -668,6 +673,17 @@ int ds4_gpu_matmul_q8_0_pair_tensor(
const ds4_gpu_tensor *x, const ds4_gpu_tensor *x,
uint64_t n_tok); uint64_t n_tok);
int ds4_gpu_matmul_q4_K_pair_decode_tensor(
ds4_gpu_tensor *out0,
ds4_gpu_tensor *out1,
const void *model_map,
uint64_t model_size,
uint64_t weight0_offset,
uint64_t weight1_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x);
/* Multi-row decode projections that preserve the one-row reduction order. */ /* Multi-row decode projections that preserve the one-row reduction order. */
int ds4_gpu_matmul_q8_0_decode_rows_exact_tensor( int ds4_gpu_matmul_q8_0_decode_rows_exact_tensor(
ds4_gpu_tensor *out, ds4_gpu_tensor *out,
@@ -1276,6 +1292,37 @@ int ds4_gpu_glm_store_indexer_k_tensor(
float beta_slow, float beta_slow,
bool cache_f16); bool cache_f16);
/* GLM-5.3 pools four normalized indexer keys with a learned, per-channel
* softmax. Partial pools are retained in tail_k/tail_gate across calls. */
int ds4_gpu_glm53_indexer_pool_update_tensor(
ds4_gpu_tensor *pool_cache,
ds4_gpu_tensor *tail_k,
ds4_gpu_tensor *tail_gate,
const ds4_gpu_tensor *raw_k,
const ds4_gpu_tensor *gate,
const void *model_map,
uint64_t model_size,
uint64_t norm_weight_offset,
uint64_t norm_bias_offset,
uint64_t ape_offset,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t head_dim,
uint32_t pool_size,
float eps,
bool cache_f16);
int ds4_gpu_glm53_expand_pool_selection_tensor(
ds4_gpu_tensor *raw_selected,
const ds4_gpu_tensor *pool_selected,
uint32_t n_tokens,
uint32_t pos0,
uint32_t selected_pools,
uint32_t index_topk,
uint32_t pool_size,
uint32_t output_width);
int ds4_gpu_glm_build_kv_cache_tensor( int ds4_gpu_glm_build_kv_cache_tensor(
ds4_gpu_tensor *key_cache, ds4_gpu_tensor *key_cache,
ds4_gpu_tensor *value_cache, ds4_gpu_tensor *value_cache,
@@ -1388,6 +1435,20 @@ int ds4_gpu_glm_indexer_scores_batch_tensor(
float scale, float scale,
bool cache_f16); bool cache_f16);
int ds4_gpu_glm53_indexer_scores_batch_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *indexer_key_cache,
uint32_t n_rows,
uint32_t n_tokens,
uint32_t pos0,
uint32_t pool_size,
uint32_t n_head,
uint32_t head_dim,
float scale,
bool cache_f16);
int ds4_gpu_glm_qk_lowrank_q8_0_tensor( int ds4_gpu_glm_qk_lowrank_q8_0_tensor(
ds4_gpu_tensor *qk_low, ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *q, const ds4_gpu_tensor *q,
@@ -1698,6 +1759,21 @@ int ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor(
float beta_fast, float beta_fast,
float beta_slow); float beta_slow);
/* Dense causal MLA over the shared compact latent cache. qk_low and lora_out
* are [token, head, kv_lora_dim]; the F16 cache is shared by all heads. */
int ds4_gpu_glm_attention_dense_compact_lora_causal_tensor(
ds4_gpu_tensor *lora_out,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
uint32_t q_row0,
uint32_t n_q,
uint32_t n_kv,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_dim);
int ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor( int ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor(
ds4_gpu_tensor *lora_out, ds4_gpu_tensor *lora_out,
const ds4_gpu_tensor *q, const ds4_gpu_tensor *q,
@@ -2366,6 +2442,7 @@ int ds4_gpu_glm_routed_moe_one_tensor(
const ds4_gpu_tensor *weights, const ds4_gpu_tensor *weights,
uint32_t n_total_expert, uint32_t n_total_expert,
uint32_t n_expert, uint32_t n_expert,
float swiglu_clamp,
uint32_t layer_index, uint32_t layer_index,
const ds4_gpu_tensor *x, const ds4_gpu_tensor *x,
bool force_resident); bool force_resident);
@@ -2394,6 +2471,7 @@ int ds4_gpu_glm_routed_moe_batch_tensor(
const ds4_gpu_tensor *weights, const ds4_gpu_tensor *weights,
uint32_t n_total_expert, uint32_t n_total_expert,
uint32_t n_expert, uint32_t n_expert,
float swiglu_clamp,
uint32_t layer_index, uint32_t layer_index,
const ds4_gpu_tensor *x, const ds4_gpu_tensor *x,
uint32_t n_tokens, uint32_t n_tokens,
@@ -2424,6 +2502,7 @@ int ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor(
const ds4_gpu_tensor *weights, const ds4_gpu_tensor *weights,
uint32_t n_total_expert, uint32_t n_total_expert,
uint32_t n_expert, uint32_t n_expert,
float swiglu_clamp,
uint32_t layer_index, uint32_t layer_index,
const ds4_gpu_tensor *x, const ds4_gpu_tensor *x,
uint32_t n_tokens, uint32_t n_tokens,
@@ -2431,6 +2510,10 @@ int ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor(
int ds4_gpu_routed_moe_set_selected_override(const int32_t *selected, uint32_t n_selected); int ds4_gpu_routed_moe_set_selected_override(const int32_t *selected, uint32_t n_selected);
void ds4_gpu_set_glm_mtp_verify_mode(bool enabled); void ds4_gpu_set_glm_mtp_verify_mode(bool enabled);
#ifdef DS4_ROCM_BUILD
int ds4_gpu_dspark_gfx1151_fast_path(void);
void ds4_gpu_set_dspark_verify_mode(bool enabled);
#endif
int ds4_gpu_matmul_q8_0_kslice_hc_expand_add_tensor( int ds4_gpu_matmul_q8_0_kslice_hc_expand_add_tensor(
ds4_gpu_tensor *out_hc, ds4_gpu_tensor *out_hc,
@@ -2850,6 +2933,146 @@ int ds4_gpu_matmul_q8_0_hc_expand_tensor(
uint32_t n_embd, uint32_t n_embd,
uint32_t n_hc); uint32_t n_hc);
int ds4_gpu_glm53_embedding_bf16(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
const ds4_gpu_tensor *token_ids,
uint32_t n_tokens,
uint32_t n_embd,
uint32_t n_vocab);
int ds4_gpu_glm53_matmul_bf16(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t in_dim,
uint32_t out_dim,
const ds4_gpu_tensor *x,
uint32_t n_rows);
int ds4_gpu_glm53_matmul_bf16_qkv(
ds4_gpu_tensor *out_q,
ds4_gpu_tensor *out_k,
ds4_gpu_tensor *out_v,
const void *model_map,
uint64_t model_size,
uint64_t weight_q_offset,
uint64_t weight_k_offset,
uint64_t weight_v_offset,
uint32_t in_dim,
uint32_t out_dim,
const ds4_gpu_tensor *x);
#ifndef DS4_GLM53_VISION_TYPES_DEFINED
#define DS4_GLM53_VISION_TYPES_DEFINED
#define DS4_GLM53_VISION_LAYERS 24u
typedef struct {
uint64_t norm1;
uint64_t qkv_weight;
uint64_t qkv_bias;
uint64_t q_norm;
uint64_t k_norm;
uint64_t attn_proj_weight;
uint64_t attn_proj_bias;
uint64_t norm2;
uint64_t gate_weight;
uint64_t gate_bias;
uint64_t up_weight;
uint64_t up_bias;
uint64_t down_weight;
uint64_t down_bias;
} ds4_glm53_vision_layer_weights;
typedef struct {
uint64_t patch_weight;
uint64_t patch_bias;
uint64_t post_norm;
uint64_t downsample_weight;
uint64_t downsample_bias;
uint64_t merger_proj;
uint64_t merger_norm;
uint64_t merger_norm_bias;
uint64_t merger_gate;
uint64_t merger_up;
uint64_t merger_down;
ds4_glm53_vision_layer_weights layer[DS4_GLM53_VISION_LAYERS];
} ds4_glm53_vision_weights;
#endif
/* Encode normalized, block-major image patches into 4096-wide language-model
* embeddings. GPU implementations keep every intermediate on device. */
int ds4_gpu_glm53_vision_encode(
float *out,
const float *patches,
uint32_t grid_h,
uint32_t grid_w,
const void *model_map,
uint64_t model_size,
const ds4_glm53_vision_weights *weights);
/* Replace token rows with projected image embeddings and repeat each row into
* every GLM hyperconnection stream. Must be called in an active command batch. */
int ds4_gpu_glm53_scatter_image_hc(
ds4_gpu_tensor *hc,
const ds4_gpu_tensor *image,
uint32_t dst_row,
uint32_t image_row,
uint32_t rows,
uint32_t total_rows,
uint32_t n_embd,
uint32_t n_hc);
/* GLM-5.3 Kimi Delta Attention. Recurrent and convolution state stay FP32. */
int ds4_gpu_glm53_kda_decode(
ds4_gpu_tensor *out,
ds4_gpu_tensor *conv_state,
ds4_gpu_tensor *recurrent_state,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *k,
const ds4_gpu_tensor *v,
const ds4_gpu_tensor *raw_gate,
const ds4_gpu_tensor *raw_beta,
const ds4_gpu_tensor *output_gate,
const void *model_map,
uint64_t model_size,
uint64_t q_conv_offset,
uint64_t k_conv_offset,
uint64_t v_conv_offset,
uint64_t a_log_offset,
uint64_t dt_bias_offset,
uint64_t output_norm_offset,
uint32_t n_heads,
uint32_t n_rows,
float gate_lower_bound,
float norm_eps);
int ds4_gpu_glm53_kda_prefill(
ds4_gpu_tensor *out,
ds4_gpu_tensor *conv_state,
ds4_gpu_tensor *recurrent_state,
ds4_gpu_tensor *q,
ds4_gpu_tensor *k,
ds4_gpu_tensor *v,
ds4_gpu_tensor *raw_gate,
const ds4_gpu_tensor *raw_beta,
const ds4_gpu_tensor *output_gate,
const void *model_map,
uint64_t model_size,
uint64_t q_conv_offset,
uint64_t k_conv_offset,
uint64_t v_conv_offset,
uint64_t a_log_offset,
uint64_t dt_bias_offset,
uint64_t output_norm_offset,
uint32_t n_heads,
uint32_t n_tokens,
float gate_lower_bound,
float norm_eps);
/* Decode-island CUDA graph capture (CUDA backend; Metal/ROCm/CPU stub it /* Decode-island CUDA graph capture (CUDA backend; Metal/ROCm/CPU stub it
* out and stay eager). Design ported from the Entrpi/ds4 batched-serving * out and stay eager). Design ported from the Entrpi/ds4 batched-serving
* fork's per-layer decode graph capture. The key identifies a captured * fork's per-layer decode graph capture. The key identifies a captured
native/metal/ds4_metal.m
+1573 -91
View File
File diff suppressed because it is too large. Load diff
src/agent.rs
+2 -2
View File
@@ -3612,7 +3612,7 @@ pub(crate) fn parse_tool_calls(
model: ModelChoice, model: ModelChoice,
text: &str, text: &str,
) -> Result<(String, Vec<ToolCall>), String> { ) -> Result<(String, Vec<ToolCall>), String> {
let (content, calls) = if model == ModelChoice::Glm52 { let (content, calls) = if model.is_glm() {
parse_glm_calls(text)? parse_glm_calls(text)?
} else { } else {
crate::dsml::parse_tool_calls(text)? crate::dsml::parse_tool_calls(text)?
@@ -3643,7 +3643,7 @@ fn system_prompt_with_tools(
ralph_child: bool, ralph_child: bool,
) -> String { ) -> String {
let schemas = tool_schemas(dev_brain, ralph_child); let schemas = tool_schemas(dev_brain, ralph_child);
let tools = if model == ModelChoice::Glm52 { let tools = if model.is_glm() {
format!( format!(
"You are a coding agent running in a local workspace. Use tools for local file and system work. Avoid printing large file contents or code blocks as answers; edit files with tools, then summarize briefly.\n\n# Tools\n\n<tools>\n{schemas}\n</tools>\n\nFor a function call, output exactly: <tool_call>function-name<arg_key>key</arg_key><arg_value>value</arg_value></tool_call>\nTool calls are not allowed inside <think></think>. Pass numbers and booleans as JSON primitives, not quoted strings. When a tool fails validation or execution, use its code, field, expected, and received feedback to correct the next call. Preserve the current system configuration unless the user explicitly asks otherwise." "You are a coding agent running in a local workspace. Use tools for local file and system work. Avoid printing large file contents or code blocks as answers; edit files with tools, then summarize briefly.\n\n# Tools\n\n<tools>\n{schemas}\n</tools>\n\nFor a function call, output exactly: <tool_call>function-name<arg_key>key</arg_key><arg_value>value</arg_value></tool_call>\nTool calls are not allowed inside <think></think>. Pass numbers and booleans as JSON primitives, not quoted strings. When a tool fails validation or execution, use its code, field, expected, and received feedback to correct the next call. Preserve the current system configuration unless the user explicitly asks otherwise."
) )
src/app.rs
+64
View File
@@ -35,6 +35,7 @@ use crate::settings::{
ReasoningMode, RuntimePreferences, SpeculativePreferences, SsdPreferences, SteeringPreferences, ReasoningMode, RuntimePreferences, SpeculativePreferences, SsdPreferences, SteeringPreferences,
StreamingCacheBudget, StreamingCacheBudget,
}; };
use base64::Engine as _;
use iced::widget::{markdown, scrollable, text_editor}; use iced::widget::{markdown, scrollable, text_editor};
use iced::{Size, Subscription, Task, keyboard, mouse, window}; use iced::{Size, Subscription, Task, keyboard, mouse, window};
use rfd::AsyncFileDialog; use rfd::AsyncFileDialog;
@@ -115,6 +116,7 @@ pub(crate) struct App {
project_name_input: String, project_name_input: String,
model_download: ModelDownload, model_download: ModelDownload,
pub(super) composer: text_editor::Content, pub(super) composer: text_editor::Content,
pub(super) pending_vision_image: Option<PendingVisionImage>,
pub(super) queued_inputs: VecDeque<String>, pub(super) queued_inputs: VecDeque<String>,
pub(super) conversation: Vec<ChatMessage>, pub(super) conversation: Vec<ChatMessage>,
/// Follow appended chat content until the user scrolls away from the tail. /// Follow appended chat content until the user scrolls away from the tail.
@@ -191,6 +193,7 @@ struct ChatSnapshot {
selected_session: Option<i32>, selected_session: Option<i32>,
permission_mode: PermissionMode, permission_mode: PermissionMode,
composer: text_editor::Content, composer: text_editor::Content,
pending_vision_image: Option<PendingVisionImage>,
queued_inputs: VecDeque<String>, queued_inputs: VecDeque<String>,
conversation: Vec<ChatMessage>, conversation: Vec<ChatMessage>,
chat_follow_tail: bool, chat_follow_tail: bool,
@@ -231,6 +234,11 @@ struct ChatSnapshot {
skip_compaction_once: bool, skip_compaction_once: bool,
} }
pub(super) struct PendingVisionImage {
pub(super) name: String,
marker: String,
}
#[derive(Clone, Debug, Eq, PartialEq)] #[derive(Clone, Debug, Eq, PartialEq)]
pub(super) struct ProjectChoice { pub(super) struct ProjectChoice {
id: i32, id: i32,
@@ -480,6 +488,9 @@ pub(crate) enum Message {
StopModelDownload, StopModelDownload,
DownloadProgressTick, DownloadProgressTick,
ComposerAction(text_editor::Action), ComposerAction(text_editor::Action),
ChooseVisionImage,
VisionImagePicked(Result<Option<(String, Vec<u8>)>, String>),
RemoveVisionImage,
TranscriptAction(usize, text_editor::Action), TranscriptAction(usize, text_editor::Action),
ToggleReasoning(usize), ToggleReasoning(usize),
OpenLink(markdown::Uri), OpenLink(markdown::Uri),
@@ -641,6 +652,7 @@ impl App {
project_name_input: String::new(), project_name_input: String::new(),
model_download: ModelDownload::Idle, model_download: ModelDownload::Idle,
composer: text_editor::Content::new(), composer: text_editor::Content::new(),
pending_vision_image: None,
queued_inputs: VecDeque::new(), queued_inputs: VecDeque::new(),
conversation: Vec::new(), conversation: Vec::new(),
chat_follow_tail: true, chat_follow_tail: true,
@@ -805,6 +817,7 @@ impl App {
project_name_input: String::new(), project_name_input: String::new(),
model_download: ModelDownload::Idle, model_download: ModelDownload::Idle,
composer: text_editor::Content::new(), composer: text_editor::Content::new(),
pending_vision_image: None,
queued_inputs: VecDeque::new(), queued_inputs: VecDeque::new(),
conversation: Vec::new(), conversation: Vec::new(),
chat_follow_tail: true, chat_follow_tail: true,
@@ -882,6 +895,7 @@ impl App {
selected_session: self.selected_session.take(), selected_session: self.selected_session.take(),
permission_mode: self.permission_mode, permission_mode: self.permission_mode,
composer: std::mem::take(&mut self.composer), composer: std::mem::take(&mut self.composer),
pending_vision_image: self.pending_vision_image.take(),
queued_inputs: std::mem::take(&mut self.queued_inputs), queued_inputs: std::mem::take(&mut self.queued_inputs),
conversation: std::mem::take(&mut self.conversation), conversation: std::mem::take(&mut self.conversation),
chat_follow_tail: std::mem::replace(&mut self.chat_follow_tail, true), chat_follow_tail: std::mem::replace(&mut self.chat_follow_tail, true),
@@ -932,6 +946,7 @@ impl App {
self.selected_session = snapshot.selected_session; self.selected_session = snapshot.selected_session;
self.permission_mode = snapshot.permission_mode; self.permission_mode = snapshot.permission_mode;
self.composer = snapshot.composer; self.composer = snapshot.composer;
self.pending_vision_image = snapshot.pending_vision_image;
self.queued_inputs = snapshot.queued_inputs; self.queued_inputs = snapshot.queued_inputs;
self.conversation = snapshot.conversation; self.conversation = snapshot.conversation;
self.chat_follow_tail = snapshot.chat_follow_tail; self.chat_follow_tail = snapshot.chat_follow_tail;
@@ -1290,6 +1305,55 @@ impl App {
} }
Message::MetricsTick => self.sample_metrics(), Message::MetricsTick => self.sample_metrics(),
Message::ComposerAction(action) => self.composer.perform(action), Message::ComposerAction(action) => self.composer.perform(action),
Message::ChooseVisionImage => {
if self.config.model != ModelChoice::Glm53Flash || self.generating {
return Task::none();
}
return Task::perform(
async {
let file = AsyncFileDialog::new()
.set_title("Attach an image")
.add_filter("Image", &["png", "jpg", "jpeg"])
.pick_file()
.await;
let Some(file) = file else {
return Ok::<_, String>(None);
};
let name = file.file_name();
let bytes = file.read().await;
Ok::<_, String>(Some((name, bytes)))
},
Message::VisionImagePicked,
);
}
Message::VisionImagePicked(result) => match result {
Ok(Some((name, bytes))) => {
if bytes.is_empty() || bytes.len() > 64 * 1024 * 1024 {
self.error = Some("Image must be between 1 byte and 64 MiB.".into());
} else {
let mime = match image::guess_format(&bytes) {
Ok(image::ImageFormat::Png) => Some("image/png"),
Ok(image::ImageFormat::Jpeg) => Some("image/jpeg"),
_ => None,
};
if let Some(mime) = mime {
let encoded = base64::engine::general_purpose::STANDARD.encode(bytes);
self.pending_vision_image = Some(PendingVisionImage {
name,
marker: crate::engine::vision_data_marker(&format!(
"data:{mime};base64,{encoded}"
)),
});
self.error = None;
} else {
self.error = Some("Only PNG and JPEG images are supported.".into());
}
}
}
Ok(None) => {}
Err(error) => self.error = Some(error),
},
Message::RemoveVisionImage => self.pending_vision_image = None,
Message::TranscriptAction(index, action) => { Message::TranscriptAction(index, action) => {
if !action.is_edit() if !action.is_edit()
&& let Some(message) = self.conversation.get_mut(index) && let Some(message) = self.conversation.get_mut(index)
src/app/generation.rs
+43 -6
View File
@@ -517,7 +517,10 @@ impl App {
if self.selected_project.is_none() { if self.selected_project.is_none() {
return; return;
} }
let prompt = self.composer.text().trim().to_owned(); let mut prompt = self.composer.text().trim().to_owned();
if prompt.is_empty() && self.pending_vision_image.is_some() {
prompt = "Describe this image.".into();
}
if prompt.is_empty() { if prompt.is_empty() {
return; return;
} }
@@ -659,13 +662,20 @@ impl App {
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
let assistant_reasoning = effective.turn.reasoning_mode != ReasoningMode::Direct; let assistant_reasoning = effective.turn.reasoning_mode != ReasoningMode::Direct;
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
let (visible_prompt, base_model_prompt) = vision_prompts(
&prompt,
self.pending_vision_image
.as_ref()
.map(|image| (image.name.as_str(), image.marker.as_str())),
);
#[cfg(target_os = "macos")]
let model_prompt = if self.config.a2ui_enabled { let model_prompt = if self.config.a2ui_enabled {
format!( format!(
"{prompt}\n\nA2UI client metadata:\n{}", "{base_model_prompt}\n\nA2UI client metadata:\n{}",
self.a2ui.client_metadata() self.a2ui.client_metadata()
) )
} else { } else {
prompt.clone() base_model_prompt
}; };
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
let mut injected_system = self.workspace_instruction_messages(&[], opening_turn); let mut injected_system = self.workspace_instruction_messages(&[], opening_turn);
@@ -750,8 +760,8 @@ impl App {
}; };
let mut saved = match database.start_chat_turn( let mut saved = match database.start_chat_turn(
session_id, session_id,
&prompt, &visible_prompt,
(model_prompt != prompt).then_some(model_prompt.as_str()), (model_prompt != visible_prompt).then_some(model_prompt.as_str()),
&injected_system, &injected_system,
assistant_reasoning, assistant_reasoning,
) { ) {
@@ -811,6 +821,7 @@ impl App {
} }
assistant.reasoning_open = assistant_reasoning; assistant.reasoning_open = assistant_reasoning;
self.composer = text_editor::Content::new(); self.composer = text_editor::Content::new();
self.pending_vision_image = None;
self.conversation.push(user); self.conversation.push(user);
self.conversation.push(assistant); self.conversation.push(assistant);
self.generating = true; self.generating = true;
@@ -2459,6 +2470,18 @@ impl App {
} }
} }
fn vision_prompts(prompt: &str, image: Option<(&str, &str)>) -> (String, String) {
image.map_or_else(
|| (prompt.to_owned(), prompt.to_owned()),
|(name, marker)| {
(
format!("[Image: {name}]\n{prompt}"),
format!("{marker}\n{prompt}"),
)
},
)
}
/// Reduces a model reply to a single sidebar-sized line, or `None` if nothing /// Reduces a model reply to a single sidebar-sized line, or `None` if nothing
/// usable came back. /// usable came back.
pub(super) fn session_title(reply: &str) -> Option<String> { pub(super) fn session_title(reply: &str) -> Option<String> {
@@ -2485,7 +2508,7 @@ mod tests {
ChatMessage, TOOL_PROTOCOL_CORRECTION, TurnSummary, chat_turn, compacted_context_start, ChatMessage, TOOL_PROTOCOL_CORRECTION, TurnSummary, chat_turn, compacted_context_start,
correction_already_sent, extension_context_visibility, has_chat_after_last_compaction, correction_already_sent, extension_context_visibility, has_chat_after_last_compaction,
has_misplaced_tool_call, is_empty_response, promote_legacy_turn_summaries, queued_prompt, has_misplaced_tool_call, is_empty_response, promote_legacy_turn_summaries, queued_prompt,
sync_a2ui_message, title_context, sync_a2ui_message, title_context, vision_prompts,
}; };
use crate::engine::ChatTurn; use crate::engine::ChatTurn;
use crate::model::ModelChoice; use crate::model::ModelChoice;
@@ -2553,6 +2576,20 @@ mod tests {
assert_eq!(user.content, "visible question"); assert_eq!(user.content, "visible question");
} }
#[test]
fn attached_image_marker_is_model_only() {
let (visible, model) = vision_prompts(
"Describe this.",
Some(("diagram.png", "<|vision_start|>payload<|vision_end|>")),
);
assert_eq!(visible, "[Image: diagram.png]\nDescribe this.");
assert_eq!(
model,
"<|vision_start|>payload<|vision_end|>\nDescribe this."
);
}
#[test] #[test]
fn legacy_generation_rows_collapse_into_their_user_turn() { fn legacy_generation_rows_collapse_into_their_user_turn() {
let mut user = assistant(None, "question"); let mut user = assistant(None, "question");
src/app/preferences.rs
+9 -4
View File
@@ -593,6 +593,9 @@ impl App {
return; return;
} }
self.config = config; self.config = config;
if self.config.model != ModelChoice::Glm53Flash {
self.pending_vision_image = None;
}
self.context_limit = self.config.active_generation().context_tokens.max(0) as u32; self.context_limit = self.config.active_generation().context_tokens.max(0) as u32;
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
{ {
@@ -653,9 +656,11 @@ impl App {
self.preference_draft.default_reasoning_mode = mode; self.preference_draft.default_reasoning_mode = mode;
self.preference_draft.load_generation(model, mode); self.preference_draft.load_generation(model, mode);
self.preference_draft.load_acceleration(model); self.preference_draft.load_acceleration(model);
if model == ModelChoice::Glm52 { if model.is_glm() {
self.preference_draft.power_percent.clear(); self.preference_draft.power_percent.clear();
self.preference_draft.prefill_chunk.clear(); self.preference_draft.prefill_chunk.clear();
}
if model == ModelChoice::Glm52 {
self.preference_draft.directional_steering_file.clear(); self.preference_draft.directional_steering_file.clear();
self.preference_draft.directional_steering_ffn.clear(); self.preference_draft.directional_steering_ffn.clear();
self.preference_draft.directional_steering_attn.clear(); self.preference_draft.directional_steering_attn.clear();
@@ -851,7 +856,7 @@ impl App {
} }
Message::PreferenceGlmMtpChanged(value) => { Message::PreferenceGlmMtpChanged(value) => {
self.preference_draft.glm_mtp = self.preference_draft.glm_mtp =
self.preference_draft.acceleration_model == ModelChoice::Glm52 && value; self.preference_draft.acceleration_model.supports_glm_mtp() && value;
if !self.preference_draft.glm_mtp { if !self.preference_draft.glm_mtp {
self.preference_draft.glm_mtp_timing = false; self.preference_draft.glm_mtp_timing = false;
} }
@@ -859,7 +864,7 @@ impl App {
} }
Message::PreferenceGlmMtpTimingChanged(value) => { Message::PreferenceGlmMtpTimingChanged(value) => {
self.preference_draft.glm_mtp_timing = self.preference_draft.glm_mtp_timing =
self.preference_draft.acceleration_model == ModelChoice::Glm52 && value; self.preference_draft.acceleration_model.supports_glm_mtp() && value;
if self.preference_draft.glm_mtp_timing { if self.preference_draft.glm_mtp_timing {
self.preference_draft.glm_mtp = true; self.preference_draft.glm_mtp = true;
} }
@@ -907,7 +912,7 @@ impl App {
self.preference_error = None; self.preference_error = None;
} }
Message::PreferenceSsdFullLayersChanged(value) => { Message::PreferenceSsdFullLayersChanged(value) => {
if self.preference_draft.acceleration_model == ModelChoice::Glm52 { if self.preference_draft.acceleration_model.is_glm() {
self.preference_draft.ssd_full_layers = value; self.preference_draft.ssd_full_layers = value;
} }
self.preference_error = None; self.preference_error = None;
src/app/view/chat.rs
+39 -2
View File
@@ -250,7 +250,8 @@ impl App {
.padding(8) .padding(8)
.style(stop_button_style) .style(stop_button_style)
.on_press(Message::StopGeneration) .on_press(Message::StopGeneration)
} else if self.composer.text().trim().is_empty() { } else if self.composer.text().trim().is_empty() && self.pending_vision_image.is_none()
{
action_button(icon(ICON_SEND, 18)).padding(8) action_button(icon(ICON_SEND, 18)).padding(8)
} else { } else {
action_button(icon(ICON_SEND, 18)) action_button(icon(ICON_SEND, 18))
@@ -263,6 +264,20 @@ impl App {
self.context_used.min(self.context_limit) as f32 / self.context_limit as f32 self.context_used.min(self.context_limit) as f32 / self.context_limit as f32
}; };
let mut composer_content = column![composer].spacing(6); let mut composer_content = column![composer].spacing(6);
if let Some(image) = &self.pending_vision_image {
composer_content = composer_content.push(
row![
icon(ICON_PAPERCLIP, 14),
text(format!("Image · {}", image.name)).size(12),
button(text("Remove").size(11))
.padding(0)
.style(button::text)
.on_press(Message::RemoveVisionImage),
]
.spacing(7)
.align_y(Alignment::Center),
);
}
for queued in &self.queued_inputs { for queued in &self.queued_inputs {
let mut queued = queued.replace('\n', " "); let mut queued = queued.replace('\n', " ");
if queued.chars().count() > 120 { if queued.chars().count() > 120 {
@@ -322,10 +337,32 @@ impl App {
.padding([2, 6]) .padding([2, 6])
.into() .into()
}; };
let vision_ready = self.config.model == ModelChoice::Glm53Flash
&& model::engine_artifacts(ModelChoice::Glm53Flash, false, &models_path())
.vision
.is_some();
let attach = if vision_ready && !self.generating && self.pending_vision_image.is_none()
{
action_button(icon(ICON_PAPERCLIP, 19))
.padding(4)
.on_press(Message::ChooseVisionImage)
} else {
action_button(icon(ICON_PAPERCLIP, 19)).padding(4)
};
composer_content = composer_content =
composer_content.push( composer_content.push(
row![ row![
icon(ICON_PAPERCLIP, 19), tooltip(
attach,
container(text(if vision_ready {
"Attach PNG or JPEG"
} else {
"GLM 5.3 Flash vision sidecar is not ready"
}))
.padding(8)
.style(preference_group_style),
tooltip::Position::Top,
),
tooltip( tooltip(
context_pie(context_fraction, 19), context_pie(context_fraction, 19),
container( container(
src/app/view/preferences.rs
+24 -12
View File
@@ -14,12 +14,16 @@ impl App {
.on_toggle_maybe(dspark_toggle), .on_toggle_maybe(dspark_toggle),
"Speculative decoding with the managed DSpark draft artifact: a small model proposes tokens that the main model verifies in one pass. Usually a large speedup; the target model may also stream routed experts from SSD.", "Speculative decoding with the managed DSpark draft artifact: a small model proposes tokens that the main model verifies in one pass. Usually a large speedup; the target model may also stream routed experts from SSD.",
); );
let glm_mtp_toggle: Option<fn(bool) -> Message> = let glm_mtp_toggle: Option<fn(bool) -> Message> = self
(self.preference_draft.acceleration_model == ModelChoice::Glm52) .preference_draft
.then_some(Message::PreferenceGlmMtpChanged); .acceleration_model
let glm_mtp_timing_toggle: Option<fn(bool) -> Message> = .supports_glm_mtp()
(self.preference_draft.acceleration_model == ModelChoice::Glm52) .then_some(Message::PreferenceGlmMtpChanged);
.then_some(Message::PreferenceGlmMtpTimingChanged); let glm_mtp_timing_toggle: Option<fn(bool) -> Message> = self
.preference_draft
.acceleration_model
.supports_glm_mtp()
.then_some(Message::PreferenceGlmMtpTimingChanged);
let dspark_strict_toggle: Option<fn(bool) -> Message> = self let dspark_strict_toggle: Option<fn(bool) -> Message> = self
.preference_draft .preference_draft
.acceleration_model .acceleration_model
@@ -71,14 +75,16 @@ impl App {
"0.8 (DeepSeek V4 Flash default)", "0.8 (DeepSeek V4 Flash default)",
&self.preference_draft.dspark_confidence_threshold, &self.preference_draft.dspark_confidence_threshold,
); );
if self.preference_draft.model != ModelChoice::Glm52 { if !self.preference_draft.model.is_glm() {
power = power.on_input(Message::PreferencePowerChanged); power = power.on_input(Message::PreferencePowerChanged);
prefill = prefill.on_input(Message::PreferencePrefillChunkChanged); prefill = prefill.on_input(Message::PreferencePrefillChunkChanged);
}
if self.preference_draft.model != ModelChoice::Glm52 {
steering_file = steering_file.on_input(Message::PreferenceSteeringFileChanged); steering_file = steering_file.on_input(Message::PreferenceSteeringFileChanged);
steering_ffn = steering_ffn.on_input(Message::PreferenceSteeringFfnChanged); steering_ffn = steering_ffn.on_input(Message::PreferenceSteeringFfnChanged);
steering_attn = steering_attn.on_input(Message::PreferenceSteeringAttnChanged); steering_attn = steering_attn.on_input(Message::PreferenceSteeringAttnChanged);
} }
if self.preference_draft.acceleration_model == ModelChoice::Glm52 { if self.preference_draft.acceleration_model.is_glm() {
ssd_full_layers = ssd_full_layers.on_input(Message::PreferenceSsdFullLayersChanged); ssd_full_layers = ssd_full_layers.on_input(Message::PreferenceSsdFullLayersChanged);
} }
if self.preference_draft.acceleration_model.supports_dspark() { if self.preference_draft.acceleration_model.supports_dspark() {
@@ -467,6 +473,12 @@ impl App {
text_input("32768", &self.preference_draft.context_tokens) text_input("32768", &self.preference_draft.context_tokens)
.on_input(Message::PreferenceContextChanged), .on_input(Message::PreferenceContextChanged),
), ),
text(if self.preference_draft.generation_model == ModelChoice::Glm53Flash {
"GLM 5.3 Flash: 32768 is the recommended default on this 128 GB machine; 50000 is the validated extended-session target."
} else {
""
})
.size(12),
preference_input_row( preference_input_row(
"Maximum generated tokens", "Maximum generated tokens",
"Hard stop for a single reply, counted from the first generated token. It bounds runaway answers and reasoning loops; it does not reserve memory.", "Hard stop for a single reply, counted from the first generated token. It bounds runaway answers and reasoning loops; it does not reserve memory.",
@@ -550,8 +562,8 @@ impl App {
.on_toggle(Message::PreferenceWarmWeightsChanged), .on_toggle(Message::PreferenceWarmWeightsChanged),
"Reads every mapped weight page once at load, so the first reply is not interrupted by page faults from disk. Loading takes longer and memory pressure rises immediately.", "Reads every mapped weight page once at load, so the first reply is not interrupted by page faults from disk. Loading takes longer and memory pressure rises immediately.",
), ),
text(if self.preference_draft.model == ModelChoice::Glm52 { text(if self.preference_draft.model.is_glm() {
"GLM 5.2 uses full GPU power and selects prefill chunks automatically." "GLM uses full GPU power and selects prefill chunks automatically."
} else { } else {
"Blank numeric values preserve DS4's automatic engine behavior." "Blank numeric values preserve DS4's automatic engine behavior."
}) })
@@ -595,7 +607,7 @@ impl App {
toggle(self.preference_draft.glm_mtp) toggle(self.preference_draft.glm_mtp)
.label("Enable integrated GLM MTP") .label("Enable integrated GLM MTP")
.on_toggle_maybe(glm_mtp_toggle), .on_toggle_maybe(glm_mtp_toggle),
"Uses the prediction head built into GLM 5.2 for speculative decoding, so no separate draft model is loaded. Available for GLM 5.2 only.", "Uses the prediction head built into GLM for speculative decoding, so no separate draft model is loaded.",
), ),
hint( hint(
toggle(self.preference_draft.glm_mtp_timing) toggle(self.preference_draft.glm_mtp_timing)
@@ -623,7 +635,7 @@ impl App {
), ),
text(if self.preference_draft.acceleration_model.supports_dspark() { text(if self.preference_draft.acceleration_model.supports_dspark() {
"DeepSeek V4 Flash 0731 uses its managed DSpark support artifact." "DeepSeek V4 Flash 0731 uses its managed DSpark support artifact."
} else if self.preference_draft.acceleration_model == ModelChoice::Glm52 { } else if self.preference_draft.acceleration_model.supports_glm_mtp() {
"GLM MTP is integrated; DSpark is unavailable for this model." "GLM MTP is integrated; DSpark is unavailable for this model."
} else { } else {
"No speculative-decoding support is available for this model." "No speculative-decoding support is available for this model."
src/engine.rs
+202 -28
View File
@@ -12,8 +12,10 @@ use crate::model::{ModelChoice, validate_engine_artifacts};
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
use crate::settings::TurnSettings; use crate::settings::TurnSettings;
use crate::settings::{EngineSettings, ReasoningMode}; use crate::settings::{EngineSettings, ReasoningMode};
#[cfg(target_os = "macos")]
use base64::Engine as _;
use gguf::{ use gguf::{
F16, F32, Gguf, I32, IQ2_XXS, MXFP4, Q2_K, Q4_0, Q4_K, Q5_K, Q6_K, Q8_0, Tensor, Value, BF16, F16, F32, Gguf, I32, IQ2_XXS, MXFP4, Q2_K, Q4_0, Q4_K, Q5_K, Q6_K, Q8_0, Tensor, Value,
}; };
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
use kvstore::{KvStore, StoreReason}; use kvstore::{KvStore, StoreReason};
@@ -32,6 +34,19 @@ use std::time::Instant;
use tokenizer::Tokenizer; use tokenizer::Tokenizer;
use validation::{SupportKind, validate_main, validate_support}; use validation::{SupportKind, validate_main, validate_support};
const VISION_DATA_START: &str = "<|ds4server_image_data|>";
const VISION_DATA_END: &str = "<|/ds4server_image_data|>";
const VISION_TOKEN_START: &str = "\u{fdd0}ds4-image:";
const VISION_TOKEN_END: &str = "\u{fdd1}";
const VISION_IMAGE_TOKEN: i32 = 154_854;
const VISION_START_TOKEN: i32 = 154_830;
const VISION_END_TOKEN: i32 = 154_831;
type VisionOverlays = Vec<(u32, metal::VisionEmbedding)>;
pub(crate) fn vision_data_marker(uri: &str) -> String {
format!("{VISION_DATA_START}{uri}{VISION_DATA_END}")
}
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
unsafe extern "C" { unsafe extern "C" {
fn mmap( fn mmap(
@@ -47,14 +62,14 @@ unsafe extern "C" {
fn munmap(address: *mut std::ffi::c_void, length: usize) -> i32; fn munmap(address: *mut std::ffi::c_void, length: usize) -> i32;
} }
pub(crate) use validation::validate_model_artifact; pub(crate) use validation::{validate_model_artifact, validate_vision_artifact};
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
pub(crate) use kvstore::sweep_unreachable as sweep_transient_cache; pub(crate) use kvstore::sweep_unreachable as sweep_transient_cache;
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
pub(crate) use metal::configure_sources as configure_metal_sources; pub(crate) use metal::configure_sources as configure_metal_sources;
const DENSE: &[u32] = &[Q8_0, Q4_K, Q4_0]; const DENSE: &[u32] = &[BF16, Q8_0, Q4_K, Q4_0];
const ROUTED: &[u32] = &[Q8_0, IQ2_XXS, Q2_K, Q4_K, Q5_K, Q6_K, MXFP4]; const ROUTED: &[u32] = &[Q8_0, IQ2_XXS, Q2_K, Q4_K, Q5_K, Q6_K, MXFP4];
const PLAIN: &[u32] = &[F16, F32]; const PLAIN: &[u32] = &[F16, F32];
const DSPARK_DENSE: &[u32] = &[F16, F32, Q8_0]; const DSPARK_DENSE: &[u32] = &[F16, F32, Q8_0];
@@ -211,9 +226,53 @@ const GLM: Shape = Shape {
original_context: 1_048_576, original_context: 1_048_576,
}; };
const GLM53_FLASH: Shape = Shape {
model: ModelChoice::Glm53Flash,
family: ModelFamily::Glm,
layers: 46,
embd: 4096,
vocab: 154_880,
heads: 64,
head_kv: 1,
head_dim: 512,
value_dim: 256,
rot: 0,
out_groups: 0,
lora_q: 1536,
lora_o: 0,
experts: 288,
experts_used: 8,
expert_shared: 1,
ff_expert: 2048,
ff_dense: 12_288,
hash_layers: 0,
sliding_window: 0,
indexer_heads: 32,
indexer_head_dim: 128,
indexer_top_k: 2048,
hc: 4,
hc_sinkhorn: 20,
nextn: 1,
leading_dense: 3,
kv_lora: 512,
key_mla: 256,
value_mla: 256,
rms_epsilon: 1.0e-5,
hc_epsilon: 1.0e-6,
expert_weight_scale: 2.5,
swiglu_clamp: 10.0,
rope_base: 0.0,
rope_scale: 0.0,
rope_beta_fast: 0.0,
rope_beta_slow: 0.0,
compress_rope_base: 0.0,
original_context: 1_048_576,
};
pub(crate) struct Model { pub(crate) struct Model {
main: Gguf, main: Gguf,
support: Option<Gguf>, support: Option<Gguf>,
vision: Option<Gguf>,
support_kind: Option<SupportKind>, support_kind: Option<SupportKind>,
shape: Shape, shape: Shape,
tokenizer: Tokenizer, tokenizer: Tokenizer,
@@ -226,6 +285,7 @@ pub(crate) struct ModelSummary {
pub(crate) tensor_count: usize, pub(crate) tensor_count: usize,
pub(crate) vocabulary_size: usize, pub(crate) vocabulary_size: usize,
pub(crate) support_loaded: bool, pub(crate) support_loaded: bool,
pub(crate) vision_loaded: bool,
} }
impl Model { impl Model {
@@ -249,6 +309,14 @@ impl Model {
model.support = Some(support); model.support = Some(support);
model.support_kind = Some(kind); model.support_kind = Some(kind);
} }
if let Some(path) = &settings.artifacts.vision {
validate_vision_artifact(path)?;
let vision = Gguf::open(path)?;
if settings.execution.warm_weights {
vision.warm()?;
}
model.vision = Some(vision);
}
Ok(model) Ok(model)
} }
@@ -266,6 +334,7 @@ impl Model {
Ok(Self { Ok(Self {
main, main,
support: None, support: None,
vision: None,
support_kind: None, support_kind: None,
shape, shape,
tokenizer, tokenizer,
@@ -275,14 +344,21 @@ impl Model {
pub(crate) fn summary(&self) -> ModelSummary { pub(crate) fn summary(&self) -> ModelSummary {
ModelSummary { ModelSummary {
model: self.shape.model, model: self.shape.model,
mapped_bytes: self.main.len() + self.support.as_ref().map_or(0, Gguf::len), mapped_bytes: self.main.len()
+ self.support.as_ref().map_or(0, Gguf::len)
+ self.vision.as_ref().map_or(0, Gguf::len),
tensor_count: self.main.tensors.len() tensor_count: self.main.tensors.len()
+ self + self
.support .support
.as_ref() .as_ref()
.map_or(0, |support| support.tensors.len()), .map_or(0, |support| support.tensors.len())
+ self
.vision
.as_ref()
.map_or(0, |vision| vision.tensors.len()),
vocabulary_size: self.tokenizer.vocab_size(), vocabulary_size: self.tokenizer.vocab_size(),
support_loaded: self.support.is_some(), support_loaded: self.support.is_some(),
vision_loaded: self.vision.is_some(),
} }
} }
@@ -293,6 +369,9 @@ impl Model {
if let Some(support) = &self.support { if let Some(support) = &self.support {
hash.update(support.checkpoint_identity()); hash.update(support.checkpoint_identity());
} }
if let Some(vision) = &self.vision {
hash.update(vision.checkpoint_identity());
}
hash.finalize().into() hash.finalize().into()
} }
@@ -732,7 +811,7 @@ impl Generator {
} }
phase("Updating system prompt cache…"); phase("Updating system prompt cache…");
let completed = self.prefill_suffix(&tokens, reused, cancelled, progress)?; let completed = self.prefill_suffix(&tokens, reused, false, cancelled, progress)?;
if completed != tokens.len() - reused { if completed != tokens.len() - reused {
return Err("generation cancelled while updating the system prompt cache".into()); return Err("generation cancelled while updating the system prompt cache".into());
} }
@@ -1055,11 +1134,12 @@ impl Generator {
&mut self, &mut self,
tokens: &[i32], tokens: &[i32],
reused: usize, reused: usize,
has_vision: bool,
cancelled: &AtomicBool, cancelled: &AtomicBool,
progress: &mut impl FnMut(u32, u32, Option<f32>), progress: &mut impl FnMut(u32, u32, Option<f32>),
) -> Result<usize, String> { ) -> Result<usize, String> {
let suffix = &tokens[reused..]; let suffix = &tokens[reused..];
if (reused == 0 && tokens.len() > 1) || suffix.len() >= 4 { if has_vision || (reused == 0 && tokens.len() > 1) || suffix.len() >= 4 {
let context = self.executor.context(); let context = self.executor.context();
self.executor.prefill(suffix, |used| { self.executor.prefill(suffix, |used| {
progress(used, context, None); progress(used, context, None);
@@ -1079,6 +1159,87 @@ impl Generator {
} }
} }
fn render_multimodal_conversation(
&mut self,
system: &str,
messages: &[ChatTurn],
reasoning: ReasoningMode,
) -> Result<(Vec<i32>, VisionOverlays), String> {
let mut rendered = messages.to_vec();
let mut embeddings = Vec::new();
let mut total_images = 0_usize;
let mut total_bytes = 0_usize;
for message in &mut rendered {
if !message.content.contains(VISION_DATA_START) {
continue;
}
if !message.user && !message.tool {
return Err("vision input is allowed only in user or tool messages".into());
}
let mut content = String::with_capacity(message.content.len());
let mut rest = message.content.as_str();
while let Some(start) = rest.find(VISION_DATA_START) {
content.push_str(&rest[..start]);
let encoded = &rest[start + VISION_DATA_START.len()..];
let end = encoded
.find(VISION_DATA_END)
.ok_or("unterminated image input")?;
let uri = &encoded[..end];
let payload = uri
.strip_prefix("data:image/png;base64,")
.or_else(|| uri.strip_prefix("data:image/jpeg;base64,"))
.ok_or("image input must be an inline PNG or JPEG data URI")?;
total_images += 1;
if total_images > 16 {
return Err("a request may contain at most 16 images".into());
}
let bytes = base64::engine::general_purpose::STANDARD
.decode(payload)
.map_err(|_| "image data URI contains invalid base64")?;
total_bytes = total_bytes
.checked_add(bytes.len())
.ok_or("image input size overflow")?;
if total_bytes > 64 * 1024 * 1024 {
return Err("image inputs exceed the 64 MiB request limit".into());
}
let embedding = self.executor.encode_vision(&bytes)?;
content.push_str(VISION_TOKEN_START);
content.push_str(&embedding.tokens.to_string());
content.push_str(VISION_TOKEN_END);
embeddings.push(embedding);
rest = &encoded[end + VISION_DATA_END.len()..];
}
content.push_str(rest);
message.content = content;
}
let tokens = self
.executor
.model()
.render_conversation(system, &rendered, reasoning);
let mut overlays = Vec::with_capacity(embeddings.len());
let mut cursor = 0_usize;
for embedding in embeddings {
let count = embedding.tokens as usize;
let relative = tokens[cursor..]
.windows(count + 2)
.position(|window| {
window[0] == VISION_START_TOKEN
&& window[count + 1] == VISION_END_TOKEN
&& window[1..count + 1]
.iter()
.all(|token| *token == VISION_IMAGE_TOKEN)
})
.ok_or("rendered prompt lost an image placeholder")?;
let start = cursor + relative + 1;
overlays.push((
u32::try_from(start).map_err(|_| "image prompt position overflow")?,
embedding,
));
cursor = start + count + 1;
}
Ok((tokens, overlays))
}
fn generate_inner( fn generate_inner(
&mut self, &mut self,
messages: &[ChatTurn], messages: &[ChatTurn],
@@ -1087,29 +1248,41 @@ impl Generator {
emit: &mut impl FnMut(bool, String), emit: &mut impl FnMut(bool, String),
progress: &mut impl FnMut(u32, u32, Option<f32>), progress: &mut impl FnMut(u32, u32, Option<f32>),
) -> Result<(GenerationOutput, bool), String> { ) -> Result<(GenerationOutput, bool), String> {
let tokens = match messages.split_last() { let has_vision = messages
Some((latest, history)) .iter()
if latest.user .any(|message| message.content.contains(VISION_DATA_START));
&& self.executor.checkpoint_tag() let (tokens, overlays) = if has_vision {
== conversation_tag( self.render_multimodal_conversation(
&settings.system_prompt,
settings.reasoning_mode,
history,
) =>
{
let mut tokens = self.executor.tokens().to_vec();
tokens.extend(self.executor.model().render_continuation(
&latest.content,
settings.reasoning_mode,
latest.skip_previous_eos,
));
tokens
}
_ => self.executor.model().render_conversation(
&settings.system_prompt, &settings.system_prompt,
messages, messages,
settings.reasoning_mode, settings.reasoning_mode,
), )?
} else {
let tokens = match messages.split_last() {
Some((latest, history))
if latest.user
&& self.executor.checkpoint_tag()
== conversation_tag(
&settings.system_prompt,
settings.reasoning_mode,
history,
) =>
{
let mut tokens = self.executor.tokens().to_vec();
tokens.extend(self.executor.model().render_continuation(
&latest.content,
settings.reasoning_mode,
latest.skip_previous_eos,
));
tokens
}
_ => self.executor.model().render_conversation(
&settings.system_prompt,
messages,
settings.reasoning_mode,
),
};
(tokens, Vec::new())
}; };
if tokens.is_empty() { if tokens.is_empty() {
return Err("the rendered prompt is empty".into()); return Err("the rendered prompt is empty".into());
@@ -1122,6 +1295,7 @@ impl Generator {
)); ));
} }
let reused = self.executor.align_prompt(&tokens)?; let reused = self.executor.align_prompt(&tokens)?;
self.executor.set_vision_overlays(overlays)?;
self.metrics.kv_prefix_reused(reused); self.metrics.kv_prefix_reused(reused);
progress(self.executor.position(), self.executor.context(), None); progress(self.executor.position(), self.executor.context(), None);
let mut rng = Rng::new(settings.seed.unwrap_or(0x4453_3453_4552_5645)); let mut rng = Rng::new(settings.seed.unwrap_or(0x4453_3453_4552_5645));
@@ -1140,7 +1314,7 @@ impl Generator {
let mut pending_utf8 = Vec::new(); let mut pending_utf8 = Vec::new();
let prompt_tokens = tokens.len(); let prompt_tokens = tokens.len();
let suffix = &tokens[reused..]; let suffix = &tokens[reused..];
let completed = self.prefill_suffix(&tokens, reused, cancelled, progress)?; let completed = self.prefill_suffix(&tokens, reused, has_vision, cancelled, progress)?;
self.publish_execution_stats(); self.publish_execution_stats();
if completed != suffix.len() { if completed != suffix.len() {
return Ok(( return Ok((
src/engine/gguf.rs
+2
View File
@@ -19,6 +19,7 @@ pub(super) const Q5_K: u32 = 13;
pub(super) const Q6_K: u32 = 14; pub(super) const Q6_K: u32 = 14;
pub(super) const IQ2_XXS: u32 = 16; pub(super) const IQ2_XXS: u32 = 16;
pub(super) const I32: u32 = 26; pub(super) const I32: u32 = 26;
pub(super) const BF16: u32 = 30;
pub(super) const MXFP4: u32 = 39; pub(super) const MXFP4: u32 = 39;
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
@@ -488,6 +489,7 @@ impl<'a> Cursor<'a> {
| "deepseek4.dspark.target_layer_ids" | "deepseek4.dspark.target_layer_ids"
| "deepseek4.dspark_target_layer_ids" | "deepseek4.dspark_target_layer_ids"
| "dspark.target_layer_ids" | "dspark.target_layer_ids"
| "glm5-next.layer_types"
); );
if !keep { if !keep {
for _ in 0..len { for _ in 0..len {
src/engine/metal.rs
+32 -5
View File
@@ -3,12 +3,14 @@ mod glm;
mod gpu; mod gpu;
mod hotlist; mod hotlist;
mod profile; mod profile;
mod vision;
pub(super) use vision::VisionEmbedding;
use glm::GlmExecutor; use glm::GlmExecutor;
use gpu::*; use gpu::*;
use profile::ExpertProfile; use profile::ExpertProfile;
use super::gguf::{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::{Model, ModelFamily, Rng, exact_delta_sample};
use crate::model::ModelChoice; use crate::model::ModelChoice;
@@ -42,7 +44,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); 19] = [ const SOURCES: [(&str, &str); 22] = [
("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"),
@@ -62,6 +64,9 @@ const SOURCES: [(&str, &str); 19] = [
("DS4_METAL_NORM_SOURCE", "norm.metal"), ("DS4_METAL_NORM_SOURCE", "norm.metal"),
("DS4_METAL_BIN_SOURCE", "bin.metal"), ("DS4_METAL_BIN_SOURCE", "bin.metal"),
("DS4_METAL_SET_ROWS_SOURCE", "set_rows.metal"), ("DS4_METAL_SET_ROWS_SOURCE", "set_rows.metal"),
("DS4_METAL_GLM53_BF16_SOURCE", "glm53_bf16.metal"),
("DS4_METAL_GLM53_VISION_SOURCE", "glm53_vision.metal"),
("DS4_METAL_GLM53_KDA_SOURCE", "glm53_kda.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
@@ -1593,7 +1598,7 @@ impl Dspark {
} }
} }
struct Steering { pub(super) struct Steering {
directions: Buffer, directions: Buffer,
attention_scale: f32, attention_scale: f32,
ffn_scale: f32, ffn_scale: f32,
@@ -1615,6 +1620,7 @@ impl Steering {
let expected = model let expected = model
.shape .shape
.layers .layers
.saturating_sub(model.shape.nextn)
.checked_mul(model.shape.embd as u32) .checked_mul(model.shape.embd as u32)
.and_then(|values| values.checked_mul(4)) .and_then(|values| values.checked_mul(4))
.ok_or("directional steering size overflow")? as usize; .ok_or("directional steering size overflow")? as usize;
@@ -2168,7 +2174,9 @@ impl SsdPlan {
let hotlist = match model.shape.model { let hotlist = match model.shape.model {
ModelChoice::DeepSeekV4Flash0731 => hotlist::FLASH, ModelChoice::DeepSeekV4Flash0731 => hotlist::FLASH,
ModelChoice::DeepSeekV4Pro => hotlist::PRO, ModelChoice::DeepSeekV4Pro => hotlist::PRO,
ModelChoice::Glm52 => unreachable!("GLM uses its dedicated executor"), ModelChoice::Glm52 | ModelChoice::Glm53Flash => {
unreachable!("GLM uses its dedicated executor")
}
}; };
for &(layer, expert) in hotlist { for &(layer, expert) in hotlist {
if loaded == self.preload_experts { if loaded == self.preload_experts {
@@ -4435,6 +4443,7 @@ impl Executor {
quality, quality,
ssd, ssd,
speculative, speculative,
steering,
expert_profile_path, expert_profile_path,
) )
.map(Box::new) .map(Box::new)
@@ -4481,6 +4490,24 @@ impl Executor {
} }
} }
pub(super) fn encode_vision(&self, encoded: &[u8]) -> Result<vision::VisionEmbedding, String> {
match self {
Self::Glm(executor) => executor.encode_vision(encoded),
Self::DeepSeek(_) => Err("vision input requires GLM 5.3 Flash".into()),
}
}
pub(super) fn set_vision_overlays(
&mut self,
overlays: Vec<(u32, vision::VisionEmbedding)>,
) -> Result<(), String> {
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()),
}
}
pub(super) fn eval_speculative_greedy( pub(super) fn eval_speculative_greedy(
&mut self, &mut self,
token: i32, token: i32,
@@ -7648,7 +7675,7 @@ fn compression_ratio(shape: super::Shape, layer: u32) -> u32 {
} }
crate::model::ModelChoice::DeepSeekV4Flash0731 crate::model::ModelChoice::DeepSeekV4Flash0731
| crate::model::ModelChoice::DeepSeekV4Pro => 128, | crate::model::ModelChoice::DeepSeekV4Pro => 128,
crate::model::ModelChoice::Glm52 => 0, crate::model::ModelChoice::Glm52 | crate::model::ModelChoice::Glm53Flash => 0,
} }
} }
src/engine/metal/glm.rs
+3530 -191
View File
File diff suppressed because it is too large. Load diff
src/engine/metal/gpu.rs
+279
View File
@@ -5,6 +5,42 @@ pub(super) struct GpuTensor {
_private: [u8; 0], _private: [u8; 0],
} }
#[derive(Clone, Copy, Default)]
#[repr(C)]
pub(super) struct Glm53VisionLayerWeights {
pub(super) norm1: u64,
pub(super) qkv_weight: u64,
pub(super) qkv_bias: u64,
pub(super) q_norm: u64,
pub(super) k_norm: u64,
pub(super) attn_proj_weight: u64,
pub(super) attn_proj_bias: u64,
pub(super) norm2: u64,
pub(super) gate_weight: u64,
pub(super) gate_bias: u64,
pub(super) up_weight: u64,
pub(super) up_bias: u64,
pub(super) down_weight: u64,
pub(super) down_bias: u64,
}
#[derive(Clone, Copy, Default)]
#[repr(C)]
pub(super) struct Glm53VisionWeights {
pub(super) patch_weight: u64,
pub(super) patch_bias: u64,
pub(super) post_norm: u64,
pub(super) downsample_weight: u64,
pub(super) downsample_bias: u64,
pub(super) merger_proj: u64,
pub(super) merger_norm: u64,
pub(super) merger_norm_bias: u64,
pub(super) merger_gate: u64,
pub(super) merger_up: u64,
pub(super) merger_down: u64,
pub(super) layer: [Glm53VisionLayerWeights; 24],
}
#[repr(C)] #[repr(C)]
pub(super) struct StreamExpertTable { pub(super) struct StreamExpertTable {
pub(super) model_map: *const c_void, pub(super) model_map: *const c_void,
@@ -93,6 +129,12 @@ unsafe extern "C" {
count: u32, count: u32,
) -> i32; ) -> i32;
pub(super) fn ds4_gpu_flush_commands() -> i32; pub(super) fn ds4_gpu_flush_commands() -> i32;
pub(super) fn ds4_gpu_flush_encoder() -> i32;
pub(super) fn ds4_gpu_argmax_tensor(
out: *mut GpuTensor,
logits: *const GpuTensor,
vocab: u32,
) -> i32;
pub(super) fn ds4_gpu_device_is_pre_m5_apple_silicon() -> i32; pub(super) fn ds4_gpu_device_is_pre_m5_apple_silicon() -> i32;
pub(super) fn ds4_gpu_device_is_m5_apple_silicon() -> i32; pub(super) fn ds4_gpu_device_is_m5_apple_silicon() -> i32;
#[cfg(test)] #[cfg(test)]
@@ -227,6 +269,65 @@ unsafe extern "C" {
embd: u32, embd: u32,
hc: u32, hc: u32,
) -> i32; ) -> i32;
pub(super) fn ds4_gpu_repeat_hc_rows_tensor(
out: *mut GpuTensor,
x: *const GpuTensor,
rows: u32,
embd: u32,
hc: u32,
) -> i32;
pub(super) fn ds4_gpu_glm53_embedding_bf16(
out: *mut GpuTensor,
map: *const c_void,
size: u64,
weight: u64,
tokens: *const GpuTensor,
rows: u32,
embd: u32,
vocab: u32,
) -> i32;
pub(super) fn ds4_gpu_glm53_vision_encode(
out: *mut f32,
patches: *const f32,
grid_h: u32,
grid_w: u32,
map: *const c_void,
size: u64,
weights: *const Glm53VisionWeights,
) -> i32;
pub(super) fn ds4_gpu_glm53_scatter_image_hc(
hc: *mut GpuTensor,
image: *const GpuTensor,
dst_row: u32,
image_row: u32,
rows: u32,
total_rows: u32,
embd: u32,
hc_count: u32,
) -> i32;
pub(super) fn ds4_gpu_glm53_matmul_bf16(
out: *mut GpuTensor,
map: *const c_void,
size: u64,
weight: u64,
input: u32,
output: u32,
x: *const GpuTensor,
rows: u32,
) -> i32;
pub(super) fn ds4_gpu_glm53_matmul_bf16_qkv(
q: *mut GpuTensor,
k: *mut GpuTensor,
v: *mut GpuTensor,
map: *const c_void,
size: u64,
q_weight: u64,
k_weight: u64,
v_weight: u64,
input: u32,
output: u32,
x: *const GpuTensor,
) -> i32;
pub(super) fn ds4_gpu_attention_noncausal_raw_batch_heads_tensor( pub(super) fn ds4_gpu_attention_noncausal_raw_batch_heads_tensor(
out: *mut GpuTensor, out: *mut GpuTensor,
map: *const c_void, map: *const c_void,
@@ -417,6 +518,35 @@ unsafe extern "C" {
beta_slow: f32, beta_slow: f32,
cache_f16: bool, cache_f16: bool,
) -> i32; ) -> i32;
pub(super) fn ds4_gpu_glm53_indexer_pool_update_tensor(
cache: *mut GpuTensor,
tail_k: *mut GpuTensor,
tail_gate: *mut GpuTensor,
raw_k: *const GpuTensor,
gate: *const GpuTensor,
map: *const c_void,
size: u64,
norm_weight: u64,
norm_bias: u64,
ape: u64,
pos: u32,
rows: u32,
cache_cap: u32,
head_dim: u32,
pool_size: u32,
eps: f32,
cache_f16: bool,
) -> i32;
pub(super) fn ds4_gpu_glm53_expand_pool_selection_tensor(
selected: *mut GpuTensor,
pools: *const GpuTensor,
rows: u32,
pos: u32,
selected_pools: u32,
top_k: u32,
pool_size: u32,
width: u32,
) -> i32;
pub(super) fn ds4_gpu_glm_fill_selected_range_tensor( pub(super) fn ds4_gpu_glm_fill_selected_range_tensor(
selected: *mut GpuTensor, selected: *mut GpuTensor,
count: u32, count: u32,
@@ -467,6 +597,66 @@ unsafe extern "C" {
scale: f32, scale: f32,
cache_f16: bool, cache_f16: bool,
) -> i32; ) -> i32;
pub(super) fn ds4_gpu_glm53_indexer_scores_batch_tensor(
scores: *mut GpuTensor,
q: *const GpuTensor,
weights: *const GpuTensor,
cache: *const GpuTensor,
visible: u32,
rows: u32,
pos: u32,
pool_size: u32,
heads: u32,
head_dim: u32,
scale: f32,
cache_f16: bool,
) -> i32;
pub(super) fn ds4_gpu_glm53_kda_decode(
out: *mut GpuTensor,
conv: *mut GpuTensor,
recurrent: *mut GpuTensor,
q: *const GpuTensor,
k: *const GpuTensor,
v: *const GpuTensor,
gate: *const GpuTensor,
beta: *const GpuTensor,
output_gate: *const GpuTensor,
map: *const c_void,
size: u64,
q_conv: u64,
k_conv: u64,
v_conv: u64,
a_log: u64,
dt_bias: u64,
output_norm: u64,
heads: u32,
rows: u32,
gate_lower_bound: f32,
eps: f32,
) -> i32;
pub(super) fn ds4_gpu_glm53_kda_prefill(
out: *mut GpuTensor,
conv: *mut GpuTensor,
recurrent: *mut GpuTensor,
q: *mut GpuTensor,
k: *mut GpuTensor,
v: *mut GpuTensor,
gate: *mut GpuTensor,
beta: *const GpuTensor,
output_gate: *const GpuTensor,
map: *const c_void,
size: u64,
q_conv: u64,
k_conv: u64,
v_conv: u64,
a_log: u64,
dt_bias: u64,
output_norm: u64,
heads: u32,
rows: u32,
gate_lower_bound: f32,
eps: f32,
) -> i32;
pub(super) fn ds4_gpu_glm_qk_lowrank_typed_tensor( pub(super) fn ds4_gpu_glm_qk_lowrank_typed_tensor(
out: *mut GpuTensor, out: *mut GpuTensor,
q: *const GpuTensor, q: *const GpuTensor,
@@ -554,6 +744,19 @@ unsafe extern "C" {
beta_fast: f32, beta_fast: f32,
beta_slow: f32, beta_slow: f32,
) -> i32; ) -> i32;
pub(super) fn ds4_gpu_glm_attention_dense_compact_lora_causal_tensor(
out: *mut GpuTensor,
qk_low: *const GpuTensor,
kv_cache: *const GpuTensor,
q_row0: u32,
rows: u32,
selected: u32,
cache_cap: u32,
cache_f16: bool,
heads: u32,
kv_lora: u32,
q_nope: u32,
) -> i32;
pub(super) fn ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor( pub(super) fn ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor(
out: *mut GpuTensor, out: *mut GpuTensor,
q: *const GpuTensor, q: *const GpuTensor,
@@ -777,6 +980,20 @@ unsafe extern "C" {
eps: f32, eps: f32,
norm_eps: f32, norm_eps: f32,
) -> i32; ) -> i32;
pub(super) fn ds4_gpu_hc_split_weighted_sum_tensor(
out: *mut GpuTensor,
split: *mut GpuTensor,
mix: *const GpuTensor,
residual: *const GpuTensor,
map: *const c_void,
size: u64,
scale: u64,
base: u64,
embd: u32,
hc: u32,
iterations: u32,
eps: f32,
) -> i32;
pub(super) fn ds4_gpu_hc_rms_norm_mix_f16_available() -> i32; pub(super) fn ds4_gpu_hc_rms_norm_mix_f16_available() -> i32;
pub(super) fn ds4_gpu_hc_rms_norm_mix_f16_tensor( pub(super) fn ds4_gpu_hc_rms_norm_mix_f16_tensor(
out: *mut GpuTensor, out: *mut GpuTensor,
@@ -1359,6 +1576,15 @@ unsafe extern "C" {
embd: u32, embd: u32,
hc: u32, hc: u32,
) -> i32; ) -> i32;
pub(super) fn ds4_gpu_hc_expand_tensor(
out: *mut GpuTensor,
block: *const GpuTensor,
residual: *const GpuTensor,
post: *const GpuTensor,
combine: *const GpuTensor,
embd: u32,
hc: u32,
) -> i32;
pub(super) fn ds4_gpu_hc_expand_add_split_tensor( pub(super) fn ds4_gpu_hc_expand_add_split_tensor(
out: *mut GpuTensor, out: *mut GpuTensor,
block: *const GpuTensor, block: *const GpuTensor,
@@ -1463,6 +1689,44 @@ unsafe extern "C" {
x: *const GpuTensor, x: *const GpuTensor,
clamp: f32, clamp: f32,
) -> i32; ) -> i32;
pub(super) fn ds4_gpu_shared_mid_swiglu_q8_0_tensor(
mid: *mut GpuTensor,
map: *const c_void,
size: u64,
gate_weight: u64,
up_weight: u64,
input: u64,
output: u64,
x: *const GpuTensor,
clamp: f32,
) -> i32;
pub(super) fn ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor(
gate: *mut GpuTensor,
up: *mut GpuTensor,
mid: *mut GpuTensor,
map: *const c_void,
size: u64,
gate_weight: u64,
up_weight: u64,
input: u64,
output: u64,
x: *const GpuTensor,
clamp: f32,
) -> i32;
pub(super) fn ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor(
gate: *mut GpuTensor,
up: *mut GpuTensor,
mid: *mut GpuTensor,
map: *const c_void,
size: u64,
gate_weight: u64,
up_weight: u64,
input: u64,
output: u64,
x: *const GpuTensor,
rows: u64,
clamp: f32,
) -> i32;
pub(super) fn ds4_gpu_router_shared_gate_up_q8_0_tensor( pub(super) fn ds4_gpu_router_shared_gate_up_q8_0_tensor(
router_logits: *mut GpuTensor, router_logits: *mut GpuTensor,
gate: *mut GpuTensor, gate: *mut GpuTensor,
@@ -1598,6 +1862,21 @@ impl Context {
return Err(error); return Err(error);
} }
} }
if let Some(vision) = &model.vision {
let mapped = unsafe {
ds4_gpu_set_model_map_range(
vision.map_ptr().cast(),
vision.len(),
vision.data_offset(),
vision.len() - vision.data_offset(),
vision.max_tensor_bytes(),
)
};
if let Err(error) = check(mapped, "vision-model mapping") {
unsafe { ds4_gpu_cleanup() };
return Err(error);
}
}
unsafe { ds4_gpu_set_quality(quality) }; unsafe { ds4_gpu_set_quality(quality) };
let model_file = File::open(model.main.path()).map_err(|error| { let model_file = File::open(model.main.path()).map_err(|error| {
unsafe { ds4_gpu_cleanup() }; unsafe { ds4_gpu_cleanup() };
src/engine/metal/vision.rs
+311
View File
@@ -0,0 +1,311 @@
use super::gpu::{Glm53VisionLayerWeights, Glm53VisionWeights, ds4_gpu_glm53_vision_encode};
use super::{Gguf, call};
use image::{DynamicImage, ImageDecoder, ImageReader};
use std::io::Cursor;
const EMBEDDING: usize = 4096;
const PATCH: u32 = 14;
const MERGE: u32 = 2;
const MIN_TOKENS: u32 = 16;
const MAX_TOKENS: u32 = 8000;
pub(super) struct VisionEncoder {
weights: Glm53VisionWeights,
}
pub(in crate::engine) struct VisionEmbedding {
pub(in crate::engine) values: Vec<f32>,
pub(in crate::engine) tokens: u32,
pub(in crate::engine) width: u32,
pub(in crate::engine) height: u32,
pub(in crate::engine) content_width: u32,
pub(in crate::engine) content_height: u32,
}
impl VisionEncoder {
pub(super) fn bind(model: &Gguf) -> Result<Self, String> {
let offset = |name: &str| model.tensor(name).map(|tensor| tensor.offset);
let mut weights = Glm53VisionWeights {
patch_weight: offset("model.visual.patch_embed.proj.weight")?,
patch_bias: offset("model.visual.patch_embed.proj.bias")?,
post_norm: offset("model.visual.post_layernorm.weight")?,
downsample_weight: offset("model.visual.downsample.weight")?,
downsample_bias: offset("model.visual.downsample.bias")?,
merger_proj: offset("model.visual.merger.proj.weight")?,
merger_norm: offset("model.visual.merger.post_projection_norm.weight")?,
merger_norm_bias: offset("model.visual.merger.post_projection_norm.bias")?,
merger_gate: offset("model.visual.merger.gate_proj.weight")?,
merger_up: offset("model.visual.merger.up_proj.weight")?,
merger_down: offset("model.visual.merger.down_proj.weight")?,
..Glm53VisionWeights::default()
};
for (layer, target) in weights.layer.iter_mut().enumerate() {
let name = |suffix: &str| format!("model.visual.blocks.{layer}.{suffix}");
*target = Glm53VisionLayerWeights {
norm1: offset(&name("norm1.weight"))?,
qkv_weight: offset(&name("attn.qkv.weight"))?,
qkv_bias: offset(&name("attn.qkv.bias"))?,
q_norm: offset(&name("attn.q_norm.weight"))?,
k_norm: offset(&name("attn.k_norm.weight"))?,
attn_proj_weight: offset(&name("attn.proj.weight"))?,
attn_proj_bias: offset(&name("attn.proj.bias"))?,
norm2: offset(&name("norm2.weight"))?,
gate_weight: offset(&name("mlp.gate_proj.weight"))?,
gate_bias: offset(&name("mlp.gate_proj.bias"))?,
up_weight: offset(&name("mlp.up_proj.weight"))?,
up_bias: offset(&name("mlp.up_proj.bias"))?,
down_weight: offset(&name("mlp.down_proj.weight"))?,
down_bias: offset(&name("mlp.down_proj.bias"))?,
};
}
Ok(Self { weights })
}
pub(super) fn encode(&self, model: &Gguf, encoded: &[u8]) -> Result<VisionEmbedding, String> {
if encoded.is_empty() || encoded.len() > 64 * 1024 * 1024 {
return Err("image is empty or exceeds the 64 MiB encoded limit".into());
}
let reader = ImageReader::new(Cursor::new(encoded))
.with_guessed_format()
.map_err(|error| error.to_string())?;
let mut decoder = reader.into_decoder().map_err(|error| error.to_string())?;
let orientation = decoder.orientation().map_err(|error| error.to_string())?;
let mut image = DynamicImage::from_decoder(decoder).map_err(|error| error.to_string())?;
image.apply_orientation(orientation);
let rgb = image.into_rgb8();
let (width, height) = rgb.dimensions();
if width == 0
|| height == 0
|| width > 16_384
|| height > 16_384
|| u64::from(width) * u64::from(height) > 64 * 1024 * 1024
{
return Err("image dimensions exceed the GLM 5.3 vision limits".into());
}
let patches = preprocess(rgb.as_raw(), width, height)?;
let tokens = patches.grid_height * patches.grid_width / 4;
let mut values = vec![0.0_f32; tokens as usize * EMBEDDING];
call(
unsafe {
ds4_gpu_glm53_vision_encode(
values.as_mut_ptr(),
patches.values.as_ptr(),
patches.grid_height,
patches.grid_width,
model.map_ptr().cast(),
model.len(),
&self.weights,
)
},
"encoding a GLM 5.3 image",
)?;
Ok(VisionEmbedding {
values,
tokens,
width,
height,
content_width: patches.content_width,
content_height: patches.content_height,
})
}
}
struct Patches {
values: Vec<f32>,
content_width: u32,
content_height: u32,
grid_width: u32,
grid_height: u32,
}
fn preprocess(rgb: &[u8], width: u32, height: u32) -> Result<Patches, String> {
const MEAN: [f32; 3] = [0.48145466, 0.4578275, 0.40821073];
const STDDEV: [f32; 3] = [0.26862954, 0.261_302_6, 0.275_777_1];
let (target_height, target_width) = smart_resize(height, width)?;
let mut scale = (target_height as f64 / height as f64).min(target_width as f64 / width as f64);
if 2 * u64::from(height) * u64::from(width) >= 2 * 28 * 28 * MIN_TOKENS as u64 && scale > 1.0 {
scale = 1.0;
}
let content_height = ((height as f64 * scale).floor() as u32).clamp(1, target_height);
let content_width = ((width as f64 * scale).floor() as u32).clamp(1, target_width);
let mut canvas = vec![0.0_f32; target_height as usize * target_width as usize * 3];
if content_width == width && content_height == height {
for y in 0..height {
for x in 0..width {
let source = (y as usize * width as usize + x as usize) * 3;
let target = (y as usize * target_width as usize + x as usize) * 3;
for channel in 0..3 {
canvas[target + channel] = f32::from(rgb[source + channel]);
}
}
}
} else {
resize_bicubic(
rgb,
width,
height,
&mut canvas,
content_width,
content_height,
target_width,
);
}
for y in 0..target_height {
for x in 0..target_width {
let pixel = (y as usize * target_width as usize + x as usize) * 3;
for channel in 0..3 {
let value = if x < content_width && y < content_height {
canvas[pixel + channel]
} else {
0.0
};
canvas[pixel + channel] = (value / 255.0 - MEAN[channel]) / STDDEV[channel];
}
}
}
let grid_height = target_height / PATCH;
let grid_width = target_width / PATCH;
let patch_values = grid_height as usize * grid_width as usize * 3 * 2 * 14 * 14;
let mut values = Vec::with_capacity(patch_values);
for block_y in 0..grid_height / MERGE {
for block_x in 0..grid_width / MERGE {
for merge_y in 0..MERGE {
for merge_x in 0..MERGE {
let patch_y = block_y * MERGE + merge_y;
let patch_x = block_x * MERGE + merge_x;
for channel in 0..3 {
for _ in 0..2 {
for y in 0..PATCH {
for x in 0..PATCH {
let pixel = ((patch_y * PATCH + y) as usize
* target_width as usize
+ (patch_x * PATCH + x) as usize)
* 3;
values.push(canvas[pixel + channel]);
}
}
}
}
}
}
}
}
if values.len() != patch_values {
return Err("internal GLM 5.3 vision patch layout mismatch".into());
}
Ok(Patches {
values,
content_width,
content_height,
grid_width,
grid_height,
})
}
fn smart_resize(height: u32, width: u32) -> Result<(u32, u32), String> {
let factor = 28_u32;
let align = |value: u32| value.div_ceil(factor) * factor;
let pixels_per_token = 2_u64 * factor as u64 * factor as u64;
let min_pixels = MIN_TOKENS as u64 * pixels_per_token;
let max_pixels = MAX_TOKENS as u64 * pixels_per_token;
let mut aligned_height = align(height);
let mut aligned_width = align(width);
let mut budget = 2_u64 * aligned_height as u64 * aligned_width as u64;
if budget < min_pixels {
let scale = (min_pixels as f64 / (2.0 * height as f64 * width as f64)).sqrt();
aligned_height = align((height as f64 * scale).ceil() as u32);
aligned_width = align((width as f64 * scale).ceil() as u32);
budget = 2_u64 * aligned_height as u64 * aligned_width as u64;
}
if budget > max_pixels {
let (mut low, mut high) = (1_u32, height);
aligned_height = factor;
aligned_width = factor;
while low <= high {
let content_height = low + (high - low) / 2;
let content_width =
((width as f64 * content_height as f64 / height as f64).floor() as u32).max(1);
let candidate_height = align(content_height);
let candidate_width = align(content_width);
if 2_u64 * candidate_height as u64 * candidate_width as u64 <= max_pixels {
aligned_height = candidate_height;
aligned_width = candidate_width;
low = content_height + 1;
} else {
high = content_height - 1;
}
}
}
Ok((aligned_height, aligned_width))
}
#[allow(clippy::too_many_arguments)]
fn resize_bicubic(
source: &[u8],
source_width: u32,
source_height: u32,
target: &mut [f32],
target_width: u32,
target_height: u32,
target_stride: u32,
) {
let scale_x = source_width as f64 / target_width as f64;
let scale_y = source_height as f64 / target_height as f64;
let filter_x = if scale_x >= 1.0 { 1.0 / scale_x } else { 1.0 };
let filter_y = if scale_y >= 1.0 { 1.0 / scale_y } else { 1.0 };
let support_x = if scale_x >= 1.0 { 2.0 * scale_x } else { 2.0 };
let support_y = if scale_y >= 1.0 { 2.0 * scale_y } else { 2.0 };
for dy in 0..target_height {
let center_y = scale_y * (dy as f64 + 0.5);
let y0 = (center_y - support_y + 0.5).max(0.0) as u32;
let y1 = (center_y + support_y + 0.5).min(source_height as f64) as u32;
for dx in 0..target_width {
let center_x = scale_x * (dx as f64 + 0.5);
let x0 = (center_x - support_x + 0.5).max(0.0) as u32;
let x1 = (center_x + support_x + 0.5).min(source_width as f64) as u32;
let mut sum = [0.0; 3];
let mut weight_sum = 0.0;
for iy in y0..y1 {
let wy = cubic((iy as f64 + 0.5 - center_y) * filter_y);
for ix in x0..x1 {
let weight = wy * cubic((ix as f64 + 0.5 - center_x) * filter_x);
let pixel = (iy as usize * source_width as usize + ix as usize) * 3;
for channel in 0..3 {
sum[channel] += source[pixel + channel] as f64 * weight;
}
weight_sum += weight;
}
}
let pixel = (dy as usize * target_stride as usize + dx as usize) * 3;
for channel in 0..3 {
target[pixel + channel] =
(sum[channel] / weight_sum).round().clamp(0.0, 255.0) as f32;
}
}
}
}
fn cubic(mut x: f64) -> f64 {
const A: f64 = -0.5;
x = x.abs();
if x < 1.0 {
((A + 2.0) * x - (A + 3.0)) * x * x + 1.0
} else if x < 2.0 {
((A * x - 5.0 * A) * x + 8.0 * A) * x - 4.0 * A
} else {
0.0
}
}
#[cfg(test)]
mod tests {
use super::smart_resize;
#[test]
fn glm53_resize_matches_reference_token_grids() {
assert_eq!(smart_resize(28, 28), Ok((112, 112)));
assert_eq!(smart_resize(1024, 1024), Ok((1036, 1036)));
assert_eq!(smart_resize(1080, 1920), Ok((1092, 1932)));
}
}
src/engine/tokenizer.rs
+20 -1
View File
@@ -1,5 +1,8 @@
use super::gguf::Gguf; use super::gguf::Gguf;
use super::{ChatTurn, ModelFamily}; use super::{
ChatTurn, ModelFamily, VISION_END_TOKEN, VISION_IMAGE_TOKEN, VISION_START_TOKEN,
VISION_TOKEN_END, VISION_TOKEN_START,
};
use crate::settings::ReasoningMode; use crate::settings::ReasoningMode;
use std::collections::HashMap; use std::collections::HashMap;
@@ -160,6 +163,22 @@ impl Tokenizer {
let mut span = 0; let mut span = 0;
let mut position = 0; let mut position = 0;
while position < bytes.len() { while position < bytes.len() {
if bytes[position..].starts_with(VISION_TOKEN_START.as_bytes())
&& let Some(relative_end) =
text[position + VISION_TOKEN_START.len()..].find(VISION_TOKEN_END)
{
let count_start = position + VISION_TOKEN_START.len();
let count_end = count_start + relative_end;
if let Ok(count) = text[count_start..count_end].parse::<usize>() {
self.tokenize_plain(&text[span..position], &mut output);
output.push(VISION_START_TOKEN);
output.extend(std::iter::repeat_n(VISION_IMAGE_TOKEN, count));
output.push(VISION_END_TOKEN);
position = count_end + VISION_TOKEN_END.len();
span = position;
continue;
}
}
let special = self let special = self
.rendered_specials .rendered_specials
.iter() .iter()
src/engine/validation.rs
+186 -10
View File
@@ -9,7 +9,7 @@ pub(crate) fn validate_model_artifact(
let model = Gguf::open(path)?; let model = Gguf::open(path)?;
let shape = match expected { let shape = match expected {
ModelChoice::DeepSeekV4Flash0731 => FLASH_0731, ModelChoice::DeepSeekV4Flash0731 => FLASH_0731,
ModelChoice::DeepSeekV4Pro | ModelChoice::Glm52 => { ModelChoice::DeepSeekV4Pro | ModelChoice::Glm52 | ModelChoice::Glm53Flash => {
return Err(format!("{expected} does not use an external support GGUF")); return Err(format!("{expected} does not use an external support GGUF"));
} }
}; };
@@ -35,6 +35,79 @@ pub(crate) fn validate_model_artifact(
} }
} }
pub(crate) fn validate_vision_artifact(path: &Path) -> Result<(), String> {
let model = Gguf::open(path)?;
if model.bytes("general.architecture")? != b"glm5-next-vision" {
return Err("vision GGUF architecture is not glm5-next-vision".into());
}
if model.tensors.len() != 347 {
return Err(format!(
"vision GGUF has {} tensors, expected 347",
model.tensors.len()
));
}
for (key, expected) in [
("block_count", 24),
("embedding_length", 1024),
("feed_forward_length", 4096),
("attention.head_count", 16),
("projection_length", 4096),
("projection.feed_forward_length", 10_240),
("patch_size", 14),
("temporal_patch_size", 2),
("spatial_merge_size", 2),
("image_token_id", VISION_IMAGE_TOKEN as u64),
("image_start_token_id", VISION_START_TOKEN as u64),
("image_end_token_id", VISION_END_TOKEN as u64),
] {
expect_u64(&model, &format!("glm5-next-vision.{key}"), expected)?;
}
let bf16 = &[BF16];
for (name, dims) in [
(
"model.visual.patch_embed.proj.weight",
vec![14, 14, 2, 3, 1024],
),
("model.visual.patch_embed.proj.bias", vec![1024]),
("model.visual.post_layernorm.weight", vec![1024]),
("model.visual.downsample.weight", vec![2, 2, 1024, 4096]),
("model.visual.downsample.bias", vec![4096]),
("model.visual.merger.proj.weight", vec![4096, 4096]),
(
"model.visual.merger.post_projection_norm.weight",
vec![4096],
),
("model.visual.merger.post_projection_norm.bias", vec![4096]),
("model.visual.merger.gate_proj.weight", vec![4096, 10_240]),
("model.visual.merger.up_proj.weight", vec![4096, 10_240]),
("model.visual.merger.down_proj.weight", vec![10_240, 4096]),
] {
expect(&model, name, bf16, &dims)?;
}
for layer in 0..24 {
let name = |suffix: &str| format!("model.visual.blocks.{layer}.{suffix}");
for (suffix, dims) in [
("norm1.weight", vec![1024]),
("attn.qkv.weight", vec![1024, 3072]),
("attn.qkv.bias", vec![3072]),
("attn.q_norm.weight", vec![64]),
("attn.k_norm.weight", vec![64]),
("attn.proj.weight", vec![1024, 1024]),
("attn.proj.bias", vec![1024]),
("norm2.weight", vec![1024]),
("mlp.gate_proj.weight", vec![1024, 4096]),
("mlp.gate_proj.bias", vec![4096]),
("mlp.up_proj.weight", vec![1024, 4096]),
("mlp.up_proj.bias", vec![4096]),
("mlp.down_proj.weight", vec![4096, 1024]),
("mlp.down_proj.bias", vec![1024]),
] {
expect(&model, &name(suffix), bf16, &dims)?;
}
}
Ok(())
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)] #[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(super) enum SupportKind { pub(super) enum SupportKind {
DSpark, DSpark,
@@ -116,14 +189,11 @@ pub(super) fn validate_support(model: &Gguf, shape: &Shape) -> Result<SupportKin
} }
pub(super) fn validate_main(model: &Gguf, expected: ModelChoice) -> Result<Shape, String> { pub(super) fn validate_main(model: &Gguf, expected: ModelChoice) -> Result<Shape, String> {
let family = if model.bytes("general.architecture").ok() == Some(b"glm-dsa") { let architecture = model.bytes("general.architecture")?;
ModelFamily::Glm let shape = match architecture {
} else { b"glm-dsa" => GLM,
ModelFamily::DeepSeek b"glm5-next" => GLM53_FLASH,
}; _ => match (model.u32("deepseek4.block_count")?, expected) {
let shape = match family {
ModelFamily::Glm => GLM,
ModelFamily::DeepSeek => match (model.u32("deepseek4.block_count")?, expected) {
(43, ModelChoice::DeepSeekV4Flash0731) => FLASH_0731, (43, ModelChoice::DeepSeekV4Flash0731) => FLASH_0731,
(43, _) => FLASH_0731, (43, _) => FLASH_0731,
(61, _) => PRO, (61, _) => PRO,
@@ -143,6 +213,9 @@ pub(super) fn validate_main(model: &Gguf, expected: ModelChoice) -> Result<Shape
} }
fn validate_metadata(model: &Gguf, shape: &Shape) -> Result<(), String> { fn validate_metadata(model: &Gguf, shape: &Shape) -> Result<(), String> {
if shape.model == ModelChoice::Glm53Flash {
return validate_glm53_metadata(model, shape);
}
let prefix = if shape.family == ModelFamily::Glm { let prefix = if shape.family == ModelFamily::Glm {
"glm-dsa" "glm-dsa"
} else { } else {
@@ -266,13 +339,102 @@ fn validate_metadata(model: &Gguf, shape: &Shape) -> Result<(), String> {
Ok(()) Ok(())
} }
fn validate_glm53_metadata(model: &Gguf, shape: &Shape) -> Result<(), String> {
let prefix = "glm5-next";
for (key, expected) in [
("block_count", u64::from(shape.layers)),
("trunk_block_count", u64::from(shape.layers - shape.nextn)),
("nextn_predict_layers", u64::from(shape.nextn)),
("context_length", shape.original_context),
("embedding_length", shape.embd),
("vocab_size", shape.vocab),
("feed_forward_length", shape.ff_dense),
("expert_feed_forward_length", shape.ff_expert),
("expert_count", shape.experts),
("expert_used_count", shape.experts_used),
("expert_shared_count", shape.expert_shared),
("leading_dense_block_count", u64::from(shape.leading_dense)),
("attention.head_count", shape.heads),
("attention.key_length", shape.key_mla),
("attention.value_length", shape.value_mla),
("attention.q_lora_rank", shape.lora_q),
("attention.kv_lora_rank", shape.kv_lora),
("attention.rope_dimension_count", shape.rot),
("attention.indexer.head_count", shape.indexer_heads),
("attention.indexer.key_length", shape.indexer_head_dim),
("attention.indexer.top_k", shape.indexer_top_k),
("attention.indexer.pool_size", 4),
("linear_attention.head_count", 64),
("linear_attention.head_dimension", 128),
("linear_attention.conv_kernel", 4),
("hyper_connection.count", shape.hc),
("hyper_connection.sinkhorn_iterations", shape.hc_sinkhorn),
] {
expect_u64(model, &format!("{prefix}.{key}"), expected)?;
}
for (key, expected) in [
("expert_weights_scale", shape.expert_weight_scale),
("swiglu_limit", shape.swiglu_clamp),
("attention.layer_norm_rms_epsilon", shape.rms_epsilon),
("linear_attention.gate_lower_bound", -5.0),
("hyper_connection.epsilon", shape.hc_epsilon),
] {
expect_float(model, &format!("{prefix}.{key}"), expected)?;
}
if !model.boolean("glm5-next.expert_weights_norm")? {
return Err("glm5-next.expert_weights_norm must be true".into());
}
let layer_types = model.u32s("glm5-next.layer_types")?;
if layer_types.len() != shape.layers as usize {
return Err("glm5-next.layer_types must contain one entry per layer".into());
}
for (layer, &kind) in layer_types.iter().enumerate() {
let expected =
u32::from(layer + shape.nextn as usize >= shape.layers as usize || layer % 4 == 3);
if kind != expected {
return Err(format!(
"unexpected GLM 5.3 attention type at layer {layer}"
));
}
}
Ok(())
}
fn validate_tensors(model: &Gguf, shape: &Shape) -> Result<(), String> { fn validate_tensors(model: &Gguf, shape: &Shape) -> Result<(), String> {
match shape.family { match shape.family {
ModelFamily::DeepSeek => validate_deepseek_tensors(model, shape), ModelFamily::DeepSeek => validate_deepseek_tensors(model, shape),
ModelFamily::Glm if shape.model == ModelChoice::Glm53Flash => {
validate_glm53_tensors(model, shape)
}
ModelFamily::Glm => validate_glm_tensors(model, shape), ModelFamily::Glm => validate_glm_tensors(model, shape),
} }
} }
fn validate_glm53_tensors(model: &Gguf, shape: &Shape) -> Result<(), String> {
expect(
model,
"token_embd.weight",
DENSE,
&[shape.embd, shape.vocab],
)?;
expect(model, "output_norm.weight", &[F32], &[shape.embd])?;
expect(model, "output.weight", DENSE, &[shape.embd, shape.vocab])?;
expect(model, "blk.0.kda_q.weight", DENSE, &[shape.embd, 8192])?;
expect(
model,
"blk.3.attn_q_a.weight",
DENSE,
&[shape.embd, shape.lora_q],
)?;
expect(
model,
"blk.45.nextn.eh_proj.weight",
DENSE,
&[2 * shape.embd, shape.embd],
)?;
Ok(())
}
fn validate_deepseek_tensors(model: &Gguf, shape: &Shape) -> Result<(), String> { fn validate_deepseek_tensors(model: &Gguf, shape: &Shape) -> Result<(), String> {
let hc_dim = shape.embd * shape.hc; let hc_dim = shape.embd * shape.hc;
let hc_mix = 2 * shape.hc + shape.hc * shape.hc; let hc_mix = 2 * shape.hc + shape.hc * shape.hc;
@@ -902,7 +1064,7 @@ fn compression_ratio(shape: &Shape, layer: u32) -> u32 {
4 4
} }
ModelChoice::DeepSeekV4Flash0731 | ModelChoice::DeepSeekV4Pro => 128, ModelChoice::DeepSeekV4Flash0731 | ModelChoice::DeepSeekV4Pro => 128,
ModelChoice::Glm52 => 0, ModelChoice::Glm52 | ModelChoice::Glm53Flash => 0,
} }
} }
@@ -1157,4 +1319,18 @@ mod tests {
validate_model_artifact(&path, ModelChoice::DeepSeekV4Flash0731, true).unwrap(); validate_model_artifact(&path, ModelChoice::DeepSeekV4Flash0731, true).unwrap();
} }
} }
#[test]
fn configured_glm53_vision_fixture_passes_the_exact_layout() {
if let Some(path) = std::env::var_os("DS4SERVER_GLM53_VISION") {
validate_vision_artifact(Path::new(&path)).unwrap();
}
}
#[test]
fn configured_glm53_main_fixture_passes_the_exact_layout() {
if let Some(path) = std::env::var_os("DS4SERVER_GLM53_MODEL") {
validate_model_artifact(Path::new(&path), ModelChoice::Glm53Flash, false).unwrap();
}
}
} }
src/metrics.rs
+2
View File
@@ -847,6 +847,7 @@ fn model_code(model: ModelChoice) -> u8 {
ModelChoice::DeepSeekV4Pro => 2, ModelChoice::DeepSeekV4Pro => 2,
ModelChoice::Glm52 => 3, ModelChoice::Glm52 => 3,
ModelChoice::DeepSeekV4Flash0731 => 4, ModelChoice::DeepSeekV4Flash0731 => 4,
ModelChoice::Glm53Flash => 5,
} }
} }
@@ -856,6 +857,7 @@ fn model_name(value: u8) -> &'static str {
2 => "DeepSeek V4 Pro", 2 => "DeepSeek V4 Pro",
3 => "GLM 5.2", 3 => "GLM 5.2",
4 => "DeepSeek V4 Flash 0731", 4 => "DeepSeek V4 Flash 0731",
5 => "GLM 5.3 Flash Q2",
_ => "No model loaded", _ => "No model loaded",
} }
} }
src/model.rs
+66 -3
View File
@@ -9,20 +9,24 @@ use std::fmt;
use std::fs; use std::fs;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
pub(crate) const MODEL_CHOICES: [ModelChoice; 3] = [ pub(crate) const MODEL_CHOICES: [ModelChoice; 4] = [
ModelChoice::DeepSeekV4Flash0731, ModelChoice::DeepSeekV4Flash0731,
ModelChoice::DeepSeekV4Pro, ModelChoice::DeepSeekV4Pro,
ModelChoice::Glm52, ModelChoice::Glm52,
ModelChoice::Glm53Flash,
]; ];
pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 4] = [ pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 6] = [
ManagedArtifactId::DeepSeekV4Flash0731, ManagedArtifactId::DeepSeekV4Flash0731,
ManagedArtifactId::DeepSeekV4Flash0731Dspark, ManagedArtifactId::DeepSeekV4Flash0731Dspark,
ManagedArtifactId::DeepSeekV4Pro, ManagedArtifactId::DeepSeekV4Pro,
ManagedArtifactId::Glm52, ManagedArtifactId::Glm52,
ManagedArtifactId::Glm53Flash,
ManagedArtifactId::Glm53FlashVision,
]; ];
const DEEPSEEK_REPOSITORY: &str = "antirez/deepseek-v4-gguf"; const DEEPSEEK_REPOSITORY: &str = "antirez/deepseek-v4-gguf";
const GLM_REPOSITORY: &str = "antirez/glm-5.2-gguf"; const GLM_REPOSITORY: &str = "antirez/glm-5.2-gguf";
const GLM53_REPOSITORY: &str = "antirez/glm-5.3-flash-gguf";
const FLASH_0731: Artifact = Artifact { const FLASH_0731: Artifact = Artifact {
label: "DeepSeek V4 Flash 0731 model", label: "DeepSeek V4 Flash 0731 model",
@@ -31,6 +35,7 @@ const FLASH_0731: Artifact = Artifact {
size: 86_720_111_488, size: 86_720_111_488,
sha256: "ca22ae2f838e14077c22bc1c1417b71b45b5e5a3687bd96c2ac6e17fdb6261c0", sha256: "ca22ae2f838e14077c22bc1c1417b71b45b5e5a3687bd96c2ac6e17fdb6261c0",
support: Some(false), support: Some(false),
vision: false,
}; };
const FLASH_0731_DSPARK: Artifact = Artifact { const FLASH_0731_DSPARK: Artifact = Artifact {
label: "DeepSeek V4 Flash 0731 DSpark support", label: "DeepSeek V4 Flash 0731 DSpark support",
@@ -39,6 +44,7 @@ const FLASH_0731_DSPARK: Artifact = Artifact {
size: 5_989_114_272, size: 5_989_114_272,
sha256: "7e319924541db3f7a163ed7e11d7532a70d48228ab59d36cb81e1d4511885360", sha256: "7e319924541db3f7a163ed7e11d7532a70d48228ab59d36cb81e1d4511885360",
support: Some(true), support: Some(true),
vision: false,
}; };
const PRO: Artifact = Artifact { const PRO: Artifact = Artifact {
label: "DeepSeek V4 Pro 0813 model", label: "DeepSeek V4 Pro 0813 model",
@@ -47,6 +53,7 @@ const PRO: Artifact = Artifact {
size: 464_627_334_560, size: 464_627_334_560,
sha256: "c4d997ab9894b6c78b759f7869fe1726b6314b6515f6ff82607df3797c5eb193", sha256: "c4d997ab9894b6c78b759f7869fe1726b6314b6515f6ff82607df3797c5eb193",
support: Some(false), support: Some(false),
vision: false,
}; };
const GLM: Artifact = Artifact { const GLM: Artifact = Artifact {
label: "GLM 5.2 model", label: "GLM 5.2 model",
@@ -55,6 +62,25 @@ const GLM: Artifact = Artifact {
size: 211_075_856_448, size: 211_075_856_448,
sha256: "a49de64c5020432bdae23de36a423a9660a5621bc0db8d12b66bd8814b07fea0", sha256: "a49de64c5020432bdae23de36a423a9660a5621bc0db8d12b66bd8814b07fea0",
support: Some(false), support: Some(false),
vision: false,
};
const GLM53_FLASH: Artifact = Artifact {
label: "GLM 5.3 Flash Q2 model",
file_name: "GLM-5.3-Flash-Q2.gguf",
repository: GLM53_REPOSITORY,
size: 96_505_816_384,
sha256: "e81fd6241c6e55a64e1e14e47a3eab61a173fa8d7e4b5c1d1848827119705b32",
support: Some(false),
vision: false,
};
const GLM53_FLASH_VISION: Artifact = Artifact {
label: "GLM 5.3 Flash vision encoder",
file_name: "GLM-5.3-Flash-Vision-Encoder.gguf",
repository: GLM53_REPOSITORY,
size: 1_127_280_960,
sha256: "ae23e14c6979e889051b2e4a39351abcdafb161e18e606fae4d8c40095a4bf3a",
support: None,
vision: true,
}; };
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)] #[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
@@ -66,6 +92,8 @@ pub(crate) enum ModelChoice {
DeepSeekV4Pro, DeepSeekV4Pro,
#[serde(rename = "glm-5.2")] #[serde(rename = "glm-5.2")]
Glm52, Glm52,
#[serde(rename = "glm-5.3-flash")]
Glm53Flash,
} }
impl ModelChoice { impl ModelChoice {
@@ -74,6 +102,7 @@ impl ModelChoice {
Self::DeepSeekV4Flash0731 => "deepseek-v4-flash-0731", Self::DeepSeekV4Flash0731 => "deepseek-v4-flash-0731",
Self::DeepSeekV4Pro => "deepseek-v4-pro", Self::DeepSeekV4Pro => "deepseek-v4-pro",
Self::Glm52 => "glm-5.2", Self::Glm52 => "glm-5.2",
Self::Glm53Flash => "glm-5.3-flash",
} }
} }
@@ -85,26 +114,40 @@ impl ModelChoice {
self == Self::DeepSeekV4Flash0731 self == Self::DeepSeekV4Flash0731
} }
pub(crate) fn is_glm(self) -> bool {
matches!(self, Self::Glm52 | Self::Glm53Flash)
}
pub(crate) fn supports_glm_mtp(self) -> bool {
self.is_glm()
}
fn main_artifact(self) -> &'static Artifact { fn main_artifact(self) -> &'static Artifact {
match self { match self {
Self::DeepSeekV4Flash0731 => &FLASH_0731, Self::DeepSeekV4Flash0731 => &FLASH_0731,
Self::DeepSeekV4Pro => &PRO, Self::DeepSeekV4Pro => &PRO,
Self::Glm52 => &GLM, Self::Glm52 => &GLM,
Self::Glm53Flash => &GLM53_FLASH,
} }
} }
fn dspark_artifact(self) -> Option<&'static Artifact> { fn dspark_artifact(self) -> Option<&'static Artifact> {
match self { match self {
Self::DeepSeekV4Flash0731 => Some(&FLASH_0731_DSPARK), Self::DeepSeekV4Flash0731 => Some(&FLASH_0731_DSPARK),
Self::DeepSeekV4Pro | Self::Glm52 => None, Self::DeepSeekV4Pro | Self::Glm52 | Self::Glm53Flash => None,
} }
} }
fn vision_artifact(self) -> Option<&'static Artifact> {
(self == Self::Glm53Flash).then_some(&GLM53_FLASH_VISION)
}
#[cfg(test)] #[cfg(test)]
fn artifacts(self, dspark_enabled: bool) -> impl Iterator<Item = &'static Artifact> { fn artifacts(self, dspark_enabled: bool) -> impl Iterator<Item = &'static Artifact> {
[ [
Some(self.main_artifact()), Some(self.main_artifact()),
dspark_enabled.then(|| self.dspark_artifact()).flatten(), dspark_enabled.then(|| self.dspark_artifact()).flatten(),
self.vision_artifact(),
] ]
.into_iter() .into_iter()
.flatten() .flatten()
@@ -115,6 +158,7 @@ impl ModelChoice {
pub(crate) struct EngineArtifacts { pub(crate) struct EngineArtifacts {
pub(crate) model: PathBuf, pub(crate) model: PathBuf,
pub(crate) support: Option<PathBuf>, pub(crate) support: Option<PathBuf>,
pub(crate) vision: Option<PathBuf>,
} }
pub(crate) fn engine_artifacts( pub(crate) fn engine_artifacts(
@@ -131,6 +175,11 @@ pub(crate) fn engine_artifacts(
} else { } else {
None None
}, },
vision: model.vision_artifact().and_then(|artifact| {
artifact
.is_installed(model, models_path)
.then(|| artifact.path(model, models_path))
}),
} }
} }
@@ -158,7 +207,14 @@ pub(crate) fn validate_engine_artifacts(
"{} is not compatible with the selected {model} checkpoint", "{} is not compatible with the selected {model} checkpoint",
path.display() path.display()
)), )),
}?;
if let Some(path) = artifacts.vision.as_deref() {
model
.vision_artifact()
.ok_or_else(|| format!("{} is not compatible with {model}", path.display()))?
.validate_installed_path(path)?;
} }
Ok(())
} }
#[derive(Clone, Copy, Debug, Eq, PartialEq)] #[derive(Clone, Copy, Debug, Eq, PartialEq)]
@@ -167,6 +223,8 @@ pub(crate) enum ManagedArtifactId {
DeepSeekV4Flash0731Dspark, DeepSeekV4Flash0731Dspark,
DeepSeekV4Pro, DeepSeekV4Pro,
Glm52, Glm52,
Glm53Flash,
Glm53FlashVision,
} }
impl ManagedArtifactId { impl ManagedArtifactId {
@@ -177,6 +235,7 @@ impl ManagedArtifactId {
} }
Self::DeepSeekV4Pro => ModelChoice::DeepSeekV4Pro, Self::DeepSeekV4Pro => ModelChoice::DeepSeekV4Pro,
Self::Glm52 => ModelChoice::Glm52, Self::Glm52 => ModelChoice::Glm52,
Self::Glm53Flash | Self::Glm53FlashVision => ModelChoice::Glm53Flash,
} }
} }
@@ -186,6 +245,8 @@ impl ManagedArtifactId {
Self::DeepSeekV4Flash0731Dspark => &FLASH_0731_DSPARK, Self::DeepSeekV4Flash0731Dspark => &FLASH_0731_DSPARK,
Self::DeepSeekV4Pro => &PRO, Self::DeepSeekV4Pro => &PRO,
Self::Glm52 => &GLM, Self::Glm52 => &GLM,
Self::Glm53Flash => &GLM53_FLASH,
Self::Glm53FlashVision => &GLM53_FLASH_VISION,
} }
} }
} }
@@ -280,6 +341,7 @@ impl fmt::Display for ModelChoice {
Self::DeepSeekV4Flash0731 => "DeepSeek V4 Flash 0731", Self::DeepSeekV4Flash0731 => "DeepSeek V4 Flash 0731",
Self::DeepSeekV4Pro => "DeepSeek V4 Pro 0813", Self::DeepSeekV4Pro => "DeepSeek V4 Pro 0813",
Self::Glm52 => "GLM 5.2", Self::Glm52 => "GLM 5.2",
Self::Glm53Flash => "GLM 5.3 Flash Q2",
}) })
} }
} }
@@ -291,6 +353,7 @@ struct Artifact {
size: u64, size: u64,
sha256: &'static str, sha256: &'static str,
support: Option<bool>, support: Option<bool>,
vision: bool,
} }
impl Artifact { impl Artifact {
src/model/transfer.rs
+11 -3
View File
@@ -205,7 +205,9 @@ fn verify(
path.display() path.display()
)); ));
} }
if let Some(support) = artifact.support { if artifact.vision {
crate::engine::validate_vision_artifact(path)?;
} else if let Some(support) = artifact.support {
crate::engine::validate_model_artifact(path, model, support)?; crate::engine::validate_model_artifact(path, model, support)?;
} }
Ok(DownloadOutcome::Complete) Ok(DownloadOutcome::Complete)
@@ -330,8 +332,8 @@ mod tests {
Some(ModelChoice::DeepSeekV4Flash0731) Some(ModelChoice::DeepSeekV4Flash0731)
); );
assert!(ModelChoice::from_id("unknown").is_none()); assert!(ModelChoice::from_id("unknown").is_none());
assert_eq!(MODEL_CHOICES.len(), 3); assert_eq!(MODEL_CHOICES.len(), 4);
assert_eq!(MANAGED_ARTIFACTS.len(), 4); assert_eq!(MANAGED_ARTIFACTS.len(), 6);
assert_eq!(ModelChoice::Glm52.main_artifact().size, 211_075_856_448); assert_eq!(ModelChoice::Glm52.main_artifact().size, 211_075_856_448);
assert_eq!( assert_eq!(
ModelChoice::DeepSeekV4Flash0731.main_artifact().size, ModelChoice::DeepSeekV4Flash0731.main_artifact().size,
@@ -339,6 +341,7 @@ mod tests {
); );
assert_eq!(ModelChoice::DeepSeekV4Flash0731.artifacts(true).count(), 2); assert_eq!(ModelChoice::DeepSeekV4Flash0731.artifacts(true).count(), 2);
assert_eq!(ModelChoice::Glm52.artifacts(true).count(), 1); assert_eq!(ModelChoice::Glm52.artifacts(true).count(), 1);
assert_eq!(ModelChoice::Glm53Flash.artifacts(true).count(), 2);
let id = SystemTime::now() let id = SystemTime::now()
.duration_since(UNIX_EPOCH) .duration_since(UNIX_EPOCH)
@@ -366,6 +369,7 @@ mod tests {
size: 0, size: 0,
sha256: "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", sha256: "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
support: None, support: None,
vision: false,
}; };
let partial = empty.partial_path(ModelChoice::DeepSeekV4Flash0731, &models_path); let partial = empty.partial_path(ModelChoice::DeepSeekV4Flash0731, &models_path);
fs::create_dir_all(partial.parent().unwrap()).unwrap(); fs::create_dir_all(partial.parent().unwrap()).unwrap();
@@ -393,6 +397,7 @@ mod tests {
&EngineArtifacts { &EngineArtifacts {
model: installed.clone(), model: installed.clone(),
support: None, support: None,
vision: None,
}, },
) )
.is_err() .is_err()
@@ -404,6 +409,7 @@ mod tests {
&EngineArtifacts { &EngineArtifacts {
model: installed, model: installed,
support: None, support: None,
vision: None,
}, },
) )
.is_err() .is_err()
@@ -428,6 +434,7 @@ mod tests {
size: 3, size: 3,
sha256: "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad", sha256: "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad",
support: None, support: None,
vision: false,
}; };
let verified_bytes = AtomicU64::new(999); let verified_bytes = AtomicU64::new(999);
@@ -593,6 +600,7 @@ mod tests {
size: 10, size: 10,
sha256: "unused", sha256: "unused",
support: None, support: None,
vision: false,
}; };
let partial = artifact.partial_path(ModelChoice::DeepSeekV4Flash0731, &directory); let partial = artifact.partial_path(ModelChoice::DeepSeekV4Flash0731, &directory);
fs::create_dir_all(partial.parent().unwrap()).unwrap(); fs::create_dir_all(partial.parent().unwrap()).unwrap();
src/server.rs
+13
View File
@@ -554,6 +554,12 @@ fn model_alias(id: &str) -> Option<ModelChoice> {
| "zai/glm-5.2" | "zai/glm-5.2"
| "zai/glm-5.2-chat" | "zai/glm-5.2-chat"
| "zai/glm-5.2-reasoner" => Some(ModelChoice::Glm52), | "zai/glm-5.2-reasoner" => Some(ModelChoice::Glm52),
"glm-5.3-flash"
| "glm-5.3-flash-chat"
| "glm-5.3-flash-reasoner"
| "zai/glm-5.3-flash"
| "zai/glm-5.3-flash-chat"
| "zai/glm-5.3-flash-reasoner" => Some(ModelChoice::Glm53Flash),
_ => ModelChoice::from_id(id), _ => ModelChoice::from_id(id),
} }
} }
@@ -900,6 +906,13 @@ mod tests {
content_text(&json!([{"type": "text", "text": "one"}, " two"])), content_text(&json!([{"type": "text", "text": "one"}, " two"])),
"one two" "one two"
); );
assert_eq!(
content_text(&json!([{
"type": "image_url",
"image_url": {"url": "data:image/png;base64,YQ=="}
}])),
"<|ds4server_image_data|>data:image/png;base64,YQ==<|/ds4server_image_data|>"
);
} }
#[test] #[test]
src/server/request.rs
+41 -1
View File
@@ -54,6 +54,23 @@ pub(super) fn anthropic_request(value: Value) -> Result<ChatRequest, (u16, Strin
for block in blocks { for block in blocks {
match block.get("type").and_then(Value::as_str).unwrap_or("text") { match block.get("type").and_then(Value::as_str).unwrap_or("text") {
"text" => text.push_str(block.get("text").and_then(Value::as_str).unwrap_or("")), "text" => text.push_str(block.get("text").and_then(Value::as_str).unwrap_or("")),
"image" => {
let source = block.get("source").and_then(Value::as_object);
let media = source
.and_then(|source| source.get("media_type"))
.and_then(Value::as_str)
.unwrap_or("");
let data = source
.and_then(|source| source.get("data"))
.and_then(Value::as_str)
.unwrap_or("");
if !matches!(media, "image/png" | "image/jpeg") || data.is_empty() {
return Err((400, "image input must be inline PNG or JPEG base64".into()));
}
text.push_str(&crate::engine::vision_data_marker(&format!(
"data:{media};base64,{data}"
)));
}
"thinking" | "redacted_thinking" => reasoning.push_str( "thinking" | "redacted_thinking" => reasoning.push_str(
block block
.get("thinking") .get("thinking")
@@ -382,6 +399,24 @@ fn responses_content_text(value: &Value) -> Result<String, (u16, String)> {
text.push_str(value); text.push_str(value);
} }
} }
Value::Object(part)
if part.get("type").and_then(Value::as_str) == Some("input_image") =>
{
let url = part
.get("image_url")
.or_else(|| part.get("url"))
.and_then(Value::as_str)
.ok_or_else(|| (400, "invalid image input".into()))?;
if !url.starts_with("data:image/png;base64,")
&& !url.starts_with("data:image/jpeg;base64,")
{
return Err((
400,
"image input must be an inline PNG or JPEG data URI".into(),
));
}
text.push_str(&crate::engine::vision_data_marker(url));
}
_ => return Err((400, "invalid JSON request".into())), _ => return Err((400, "invalid JSON request".into())),
} }
} }
@@ -574,7 +609,12 @@ fn request_reasoning(
|| (explicit_thinking.is_none() || (explicit_thinking.is_none()
&& matches!( && matches!(
model_id, model_id,
"deepseek-chat" | "glm-5.2-chat" | "glm-5.2-no-think" | "glm-5.2-nothink" "deepseek-chat"
| "glm-5.2-chat"
| "glm-5.2-no-think"
| "glm-5.2-nothink"
| "glm-5.3-flash-chat"
| "zai/glm-5.3-flash-chat"
)) ))
{ {
reasoning = ReasoningMode::Direct; reasoning = ReasoningMode::Direct;
src/server/tools.rs
+37 -2
View File
@@ -365,6 +365,7 @@ pub(super) fn render_messages(
let mut turns = Vec::<ChatTurn>::new(); let mut turns = Vec::<ChatTurn>::new();
for message in messages { for message in messages {
validate_inline_images(&message.content)?;
let content = content_text(&message.content); let content = content_text(&message.content);
match message.role.as_str() { match message.role.as_str() {
"system" | "developer" => { "system" | "developer" => {
@@ -426,6 +427,30 @@ pub(super) fn render_messages(
Ok((system, turns)) Ok((system, turns))
} }
fn validate_inline_images(value: &Value) -> Result<(), (u16, String)> {
let Some(parts) = value.as_array() else {
return Ok(());
};
for part in parts {
let Some(image) = part.get("image_url") else {
continue;
};
let url = image
.get("url")
.or(Some(image))
.and_then(Value::as_str)
.ok_or_else(|| (400, "invalid image input".into()))?;
if !url.starts_with("data:image/png;base64,") && !url.starts_with("data:image/jpeg;base64,")
{
return Err((
400,
"image input must be an inline PNG or JPEG data URI".into(),
));
}
}
Ok(())
}
pub(super) fn validate_tool_results( pub(super) fn validate_tool_results(
state: &State, state: &State,
messages: &[ApiMessage], messages: &[ApiMessage],
@@ -779,8 +804,18 @@ pub(super) fn content_text(value: &Value) -> String {
Value::Array(parts) => parts Value::Array(parts) => parts
.iter() .iter()
.filter_map(|part| match part { .filter_map(|part| match part {
Value::String(text) => Some(text.as_str()), Value::String(text) => Some(text.clone()),
Value::Object(object) => object.get("text").and_then(Value::as_str), Value::Object(object) => object
.get("text")
.and_then(Value::as_str)
.map(str::to_owned)
.or_else(|| {
let url = object
.get("image_url")
.and_then(|value| value.get("url").or(Some(value)))
.and_then(Value::as_str)?;
Some(crate::engine::vision_data_marker(url))
}),
_ => None, _ => None,
}) })
.collect(), .collect(),
src/settings.rs
+8 -8
View File
@@ -34,8 +34,8 @@ impl SpeculativePreferences {
if self.glm_mtp_timing && !self.glm_mtp { if self.glm_mtp_timing && !self.glm_mtp {
return Err("GLM MTP timing requires GLM MTP.".into()); return Err("GLM MTP timing requires GLM MTP.".into());
} }
if model != ModelChoice::Glm52 && (self.glm_mtp || self.glm_mtp_timing) { if !model.supports_glm_mtp() && (self.glm_mtp || self.glm_mtp_timing) {
return Err("GLM MTP is available only for GLM 5.2.".into()); return Err("GLM MTP is available only for GLM models.".into());
} }
if self.dspark_enabled && !model.supports_dspark() { if self.dspark_enabled && !model.supports_dspark() {
return Err("DSpark is not available for the selected model.".into()); return Err("DSpark is not available for the selected model.".into());
@@ -166,8 +166,8 @@ impl SsdPreferences {
{ {
return Err("SSD full-layer count is too large.".into()); return Err("SSD full-layer count is too large.".into());
} }
if self.full_layers.is_some_and(|layers| layers > 0) && model != ModelChoice::Glm52 { if self.full_layers.is_some_and(|layers| layers > 0) && !model.is_glm() {
return Err("Fully resident SSD layers are available only for GLM 5.2.".into()); return Err("Fully resident SSD layers are available only for GLM models.".into());
} }
if let Some(StreamingCacheBudget::Gib(gib)) = self.cache { if let Some(StreamingCacheBudget::Gib(gib)) = self.cache {
validate_gib("SSD cache budget", gib)?; validate_gib("SSD cache budget", gib)?;
@@ -443,12 +443,12 @@ impl ExecutionPreferences {
{ {
return Err("Prefill chunk is too large.".into()); return Err("Prefill chunk is too large.".into());
} }
if model == ModelChoice::Glm52 { if model.is_glm() {
if self.power_percent.is_some_and(|power| power != 100) { if self.power_percent.is_some_and(|power| power != 100) {
return Err("GLM 5.2 currently requires 100% GPU power.".into()); return Err("GLM currently requires 100% GPU power.".into());
} }
if self.prefill_chunk.is_some() { if self.prefill_chunk.is_some() {
return Err("GLM 5.2 selects its prefill chunk automatically.".into()); return Err("GLM selects its prefill chunk automatically.".into());
} }
} }
Ok(()) Ok(())
@@ -569,7 +569,7 @@ impl GenerationPreferences {
model: ModelChoice, model: ModelChoice,
kv_cache: KvCacheSettings, kv_cache: KvCacheSettings,
) -> TurnSettings { ) -> TurnSettings {
let glm = model == ModelChoice::Glm52; let glm = model.is_glm();
TurnSettings { TurnSettings {
kv_cache, kv_cache,
context_tokens: self.context_tokens, context_tokens: self.context_tokens,