From 46d6a976a543a10981bf6ccafed99d61443371a0 Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Tue, 1 Sep 2026 19:06:20 +0200 Subject: [PATCH] Add GLM 5.3 Flash support --- AGENTS.md | 1 + docs/METAL_PARITY.md | 64 + metal/dsv4_misc.metal | 274 ++- metal/glm53_bf16.metal | 143 ++ metal/glm53_kda.metal | 314 +++ metal/glm53_vision.metal | 271 +++ metal/moe.metal | 89 +- native/metal/README.md | 11 +- native/metal/ds4_gpu.h | 223 +++ native/metal/ds4_metal.m | 1664 +++++++++++++++- src/agent.rs | 4 +- src/app.rs | 64 + src/app/generation.rs | 49 +- src/app/preferences.rs | 13 +- src/app/view/chat.rs | 41 +- src/app/view/preferences.rs | 36 +- src/engine.rs | 230 ++- src/engine/gguf.rs | 2 + src/engine/metal.rs | 37 +- src/engine/metal/glm.rs | 3721 +++++++++++++++++++++++++++++++++-- src/engine/metal/gpu.rs | 279 +++ src/engine/metal/vision.rs | 311 +++ src/engine/tokenizer.rs | 21 +- src/engine/validation.rs | 196 +- src/metrics.rs | 2 + src/model.rs | 69 +- src/model/transfer.rs | 14 +- src/server.rs | 13 + src/server/request.rs | 42 +- src/server/tools.rs | 39 +- src/settings.rs | 16 +- 31 files changed, 7806 insertions(+), 447 deletions(-) create mode 100644 metal/glm53_bf16.metal create mode 100644 metal/glm53_kda.metal create mode 100644 metal/glm53_vision.metal create mode 100644 src/engine/metal/vision.rs diff --git a/AGENTS.md b/AGENTS.md index de94a74..d32dc54 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -4,6 +4,7 @@ - Keep changes focused, handle errors explicitly, and add the smallest useful test for non-trivial behavior. - 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. +- 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. - Issues are managed with the command "tea" run from the project directory. diff --git a/docs/METAL_PARITY.md b/docs/METAL_PARITY.md index 266d809..266cc27 100644 --- a/docs/METAL_PARITY.md +++ b/docs/METAL_PARITY.md @@ -1,5 +1,69 @@ # 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 - DS4 oracle: `8db89fe083ae4d17c9a2428ccd29803d3ae8f577` diff --git a/metal/dsv4_misc.metal b/metal/dsv4_misc.metal index f24ca11..567b4e0 100644 --- a/metal/dsv4_misc.metal +++ b/metal/dsv4_misc.metal @@ -195,6 +195,17 @@ struct ds4_metal_args_glm_store_indexer_k { 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 { uint32_t pos0; uint32_t n_tokens; @@ -221,6 +232,15 @@ struct ds4_metal_args_glm_fill_selected_range_batch { 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 { uint32_t n_tokens; uint32_t n_head; @@ -252,6 +272,8 @@ struct ds4_metal_args_glm_indexer_scores_batch { uint32_t head_dim; uint32_t pos0; uint32_t cache_f16; + uint32_t row_group_size; + uint32_t pad0; uint64_t q_token_stride; uint64_t q_head_stride; uint64_t weights_token_stride; @@ -259,6 +281,13 @@ struct ds4_metal_args_glm_indexer_scores_batch { 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 { uint32_t n_head; 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((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( device char *base, uint64_t index, @@ -1011,7 +1151,8 @@ kernel void kernel_glm_build_kv_cache( corr_dims); } 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) { #ifdef DS4_METAL_ROPE_EXP2_LOG2 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); } 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) { #ifdef DS4_METAL_ROPE_EXP2_LOG2 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); } 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) { #ifdef DS4_METAL_ROPE_EXP2_LOG2 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; } +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( constant ds4_metal_args_glm_indexer_rope_tail & args, device char *x, @@ -1849,7 +2021,7 @@ kernel void kernel_glm_indexer_scores_batch( device float *dst = (device float *)(scores + (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 (tid == 0) *dst = -INFINITY; 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 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) { 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) { - 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 + (uint64_t)token0 * args.score_token_stride) + row0; *dst = row0 < visible ? acc0 : -INFINITY; } 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 + (uint64_t)token1 * args.score_token_stride) + row1; *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 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) { 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) { - 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 + (uint64_t)token0 * args.score_token_stride) + row0; *dst = row0 < visible ? acc0 : -INFINITY; } 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 + (uint64_t)token1 * args.score_token_stride) + row1; *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]], ushort tiisg [[thread_index_in_simdgroup]], ushort sgitg [[simdgroup_index_in_threadgroup]]) { - constexpr uint n_head = 64u; constexpr uint kv_lora_dim = 512u; constexpr uint qk_nope = 192u; 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 wt = args.weight_type; - if (head >= n_head || - args.n_head != n_head || + if (head >= args.n_head || + (args.n_head != 32u && args.n_head != 64u) || args.kv_lora_dim != kv_lora_dim || args.qk_nope != qk_nope || 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.cache_f16 == 0u || args.kv_lora_dim != 512u || - args.qk_rope != 64u) { + (args.qk_rope != 0u && args.qk_rope != 64u)) { return; } @@ -3668,7 +3839,7 @@ kernel void kernel_glm_attention_indexed_batch_lora_group8_vec_impl( } 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, (int)args.n_ctx_orig, 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; if (token >= args.n_tokens || args.n_selected == 0u || - args.cache_f16 == 0u || args.kv_lora_dim != 512u || - args.qk_rope != 64u) { + (args.qk_rope != 0u && args.qk_rope != 64u)) { 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}; - 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, (int)args.n_ctx_orig, 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 vv = off - rr * kv_vecs; const uint row = base + rr; - device const half4 *src = - (device const half4 *)((device const half *)kv_lora_cache + - (uint64_t)row * args.kv_lora_dim); - kv_shared[off] = src[vv]; + if (args.cache_f16 != 0u) { + device const half4 *src = + (device const half4 *)((device const half *)kv_lora_cache + + (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) { 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 uint64_t rope_base = (uint64_t)row * args.qk_rope; const float2 y0 = - glm_cache_load_rotated_rope_pair_f16_only(k_rope_cache, - rope_base, - r, - row, - args.qk_rope, - args.freq_base, - args.freq_scale, - args.ext_factor, - args.attn_factor, - corr_dims[0], - corr_dims[1]); + glm_cache_load_rotated_rope_pair(k_rope_cache, + rope_base, + r, + row, + args.qk_rope, + args.cache_f16, + args.freq_base, + args.freq_scale, + args.ext_factor, + args.attn_factor, + corr_dims[0], + corr_dims[1]); const float2 y1 = - glm_cache_load_rotated_rope_pair_f16_only(k_rope_cache, - rope_base, - r + 2u, - row, - args.qk_rope, - args.freq_base, - args.freq_scale, - args.ext_factor, - args.attn_factor, - corr_dims[0], - corr_dims[1]); + glm_cache_load_rotated_rope_pair(k_rope_cache, + rope_base, + r + 2u, + row, + args.qk_rope, + args.cache_f16, + args.freq_base, + args.freq_scale, + args.ext_factor, + args.attn_factor, + corr_dims[0], + corr_dims[1]); rope_shared[off] = float4(y0.x, y0.y, y1.x, y1.y); } threadgroup_barrier(mem_flags::mem_threadgroup); @@ -4586,14 +4765,15 @@ kernel void kernel_glm_router_select_one( threadgroup float *scratch [[threadgroup(0)]], uint token [[threadgroup_position_in_grid]], uint tid [[thread_position_in_threadgroup]]) { + const uint sort_width = args.n_expert > 256u ? 512u : 256u; 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 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_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 float p = active ? ds4_glm_router_sigmoid(token_logits[tid]) : 0.0f; if (active) token_probs[tid] = p; @@ -4601,7 +4781,7 @@ kernel void kernel_glm_router_select_one( idx[tid] = (int32_t)tid; 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) { const uint other = tid ^ j; if (other > tid) { diff --git a/metal/glm53_bf16.metal b/metal/glm53_bf16.metal new file mode 100644 index 0000000..a46dd81 --- /dev/null +++ b/metal/glm53_bf16.metal @@ -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((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 +void glm53_dequantize_bf16( + device const glm53_bf16_block16 *src, + short il, + thread type4x4 ®) { + (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>; diff --git a/metal/glm53_kda.metal b/metal/glm53_kda.metal new file mode 100644 index 0000000..c21d30e --- /dev/null +++ b/metal/glm53_kda.metal @@ -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])); +} diff --git a/metal/glm53_vision.metal b/metal/glm53_vision.metal new file mode 100644 index 0000000..986304c --- /dev/null +++ b/metal/glm53_vision.metal @@ -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]; +} diff --git a/metal/moe.metal b/metal/moe.metal index 7128d47..7aeb9d9 100644 --- a/metal/moe.metal +++ b/metal/moe.metal @@ -423,6 +423,7 @@ struct ds4_metal_glm_routed_moe_args { uint32_t n_tokens; uint32_t mid_token_stride; uint32_t down_type; + float swiglu_clamp; /* Expert ownership under tensor parallelism: tp_world 0/1 = full * compute; otherwise each rank owns a contiguous expert range. */ int32_t tp_rank; @@ -436,6 +437,14 @@ struct ds4_metal_glm_routed_moe_args { 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, 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 + (uint64_t)slot * args.mid_dim + row; 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; return; } + const int local_expert = expert - args.tp_expert_base; device const block_q4_K *gate_row = (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); device const block_q4_K *up_row = (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); float acc_gate = 0.0f; @@ -747,10 +758,8 @@ kernel void kernel_glm_q4_K_pair_swiglu_f32( } if (tid == 0u) { - const float g = scratch[0]; - const float u = scratch[ntg]; - const float sw = g / (1.0f + exp(-g)); - mid[mid_off] = sw * u * weights[selected_off]; + mid[mid_off] = ds4_glm_swiglu(scratch[0], scratch[ntg], + args.swiglu_clamp) * 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 u = simd_sum(sumu[row]); if (tiisg == 0u) { - const float sw = g / (1.0f + exp(-g)); - mid[mid_base + row0 + (uint)row] = sw * u * weights[selected_off]; + mid[mid_base + row0 + (uint)row] = + 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; 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) { for (short row = 0; 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; 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) { for (short row = 0; 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 u = simd_sum(sumu[row]); if (tiisg == 0u) { - const float sw = g / (1.0f + exp(-g)); - mid[mid_base + row0 + (uint)row] = sw * u * weights[selected_off]; + mid[mid_base + row0 + (uint)row] = + 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; const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; 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( 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( @@ -1245,7 +1261,9 @@ kernel void kernel_glm_q4_K_addr_pair_swiglu_f32( (uint64_t)slot * args.mid_dim; 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) { for (short row = 0; 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; 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) { for (short row = 0; 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; const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; 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( 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( @@ -1378,6 +1401,8 @@ kernel void kernel_glm_q4_K_pair_swiglu2_mapped_f32( ushort sgitg [[simdgroup_index_in_threadgroup]]) { const uint expert = tgpig.z; 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 map_base = tgpig.y * 32u; 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; glm_q4_K_pair_swiglu_simd_f32_impl( 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 map_row = tgpig.y; 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]; if (id < 0) return; 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; glm_q4_K_pair_swiglu_simd_f32_impl( 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( @@ -1591,8 +1620,8 @@ static inline void glm_q5_K_pair_swiglu_f32_impl( const float g = simd_sum(sumg[row]); const float u = simd_sum(sumu[row]); if (tiisg == 0u) { - const float sw = g / (1.0f + exp(-g)); - mid[mid_base + row0 + (uint)row] = sw * u * weights[selected_off]; + mid[mid_base + row0 + (uint)row] = + 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++) { 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; const uint64_t down_addr = down_addrs[(uint)expert]; if (down_addr == 0) continue; 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; for (uint slot = 0; slot < args.n_expert_used; 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 + - (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); device const float *slot_mid = mid + mid_base + (uint64_t)slot * args.mid_dim; 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++) { 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; const uint64_t down_addr = down_addrs[(uint)expert]; if (down_addr == 0) continue; 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; for (uint slot = 0; slot < args.n_expert_used; 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 *)(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); 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; @@ -2197,6 +2234,8 @@ kernel void kernel_glm_q4_K_addr_down_simd_f32( for (uint slot = 0; slot < args.n_expert_used; 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; const uint64_t down_addr = down_addrs[(uint)expert]; if (down_addr == 0) continue; diff --git a/native/metal/README.md b/native/metal/README.md index 9e777cc..643bd0d 100644 --- a/native/metal/README.md +++ b/native/metal/README.md @@ -1,15 +1,14 @@ # Vendored DS4 Metal boundary -These files track the DS4 Metal boundary audited at commit -`8db89fe083ae4d17c9a2428ccd29803d3ae8f577` (2026-08-28): +The Metal bridge and GPU declarations track the DS4 boundary audited at commit +`ec7642cdd9ec81d01ad4b1fd8f8a3d1511533748` (2026-08-31): - `ds4_metal.m` -- `ds4.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 DS4 checkout. Rust owns the model, graph, session, sampling, and lifecycle. Objective-C remains only at the platform Metal boundary. diff --git a/native/metal/ds4_gpu.h b/native/metal/ds4_gpu.h index 2fe34b3..a88fc09 100644 --- a/native/metal/ds4_gpu.h +++ b/native/metal/ds4_gpu.h @@ -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, 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); +/* 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_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); @@ -668,6 +673,17 @@ int ds4_gpu_matmul_q8_0_pair_tensor( const ds4_gpu_tensor *x, 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. */ int ds4_gpu_matmul_q8_0_decode_rows_exact_tensor( ds4_gpu_tensor *out, @@ -1276,6 +1292,37 @@ int ds4_gpu_glm_store_indexer_k_tensor( float beta_slow, 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( ds4_gpu_tensor *key_cache, ds4_gpu_tensor *value_cache, @@ -1388,6 +1435,20 @@ int ds4_gpu_glm_indexer_scores_batch_tensor( float scale, 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( ds4_gpu_tensor *qk_low, const ds4_gpu_tensor *q, @@ -1698,6 +1759,21 @@ int ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor( float beta_fast, 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( ds4_gpu_tensor *lora_out, const ds4_gpu_tensor *q, @@ -2366,6 +2442,7 @@ int ds4_gpu_glm_routed_moe_one_tensor( const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, + float swiglu_clamp, uint32_t layer_index, const ds4_gpu_tensor *x, bool force_resident); @@ -2394,6 +2471,7 @@ int ds4_gpu_glm_routed_moe_batch_tensor( const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, + float swiglu_clamp, uint32_t layer_index, const ds4_gpu_tensor *x, uint32_t n_tokens, @@ -2424,6 +2502,7 @@ int ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor( const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, + float swiglu_clamp, uint32_t layer_index, const ds4_gpu_tensor *x, 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); 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( 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_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 * out and stay eager). Design ported from the Entrpi/ds4 batched-serving * fork's per-layer decode graph capture. The key identifies a captured diff --git a/native/metal/ds4_metal.m b/native/metal/ds4_metal.m index 88586b5..adae805 100644 --- a/native/metal/ds4_metal.m +++ b/native/metal/ds4_metal.m @@ -206,12 +206,14 @@ static id g_glm_k_b_project_pipeline; static id g_glm_store_compact_kv_pipeline; static id g_glm_qkv_norm_store_compact_kv_pipeline; static id g_glm_store_indexer_k_pipeline; +static id g_glm53_indexer_pool_update_pipeline; static id g_glm_build_kv_cache_pipeline; static id g_glm_build_kv_cache_decode_group4_pipeline; static id g_glm_build_kv_cache_flash_pipeline; static id g_glm_attention_full_pipeline; static id g_glm_fill_selected_range_pipeline; static id g_glm_fill_selected_range_batch_pipeline; +static id g_glm53_expand_pool_selection_pipeline; static id g_glm_indexer_rope_tail_pipeline; static id g_glm_indexer_score_one_pipeline; static id g_glm_indexer_score_one_direct_pipeline; @@ -4351,6 +4353,9 @@ static NSString *ds4_gpu_full_source(void) { NSArray *> *required_sources = @[ @[@"DS4_METAL_FLASH_ATTN_SOURCE", @"metal/flash_attn.metal"], @[@"DS4_METAL_DENSE_SOURCE", @"metal/dense.metal"], + @[@"DS4_METAL_GLM53_BF16_SOURCE", @"metal/glm53_bf16.metal"], + @[@"DS4_METAL_GLM53_VISION_SOURCE", @"metal/glm53_vision.metal"], + @[@"DS4_METAL_GLM53_KDA_SOURCE", @"metal/glm53_kda.metal"], @[@"DS4_METAL_MOE_SOURCE", @"metal/moe.metal"], @[@"DS4_METAL_DSV4_HC_SOURCE", @"metal/dsv4_hc.metal"], @[@"DS4_METAL_UNARY_SOURCE", @"metal/unary.metal"], @@ -6096,6 +6101,17 @@ typedef struct { float pad1; } ds4_gpu_glm_store_indexer_k_args; +typedef struct { + 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; +} ds4_gpu_glm53_indexer_pool_update_args; + typedef struct { uint32_t pos0; uint32_t n_tokens; @@ -6143,6 +6159,15 @@ typedef struct { uint32_t pad_row; } ds4_gpu_glm_fill_selected_range_batch_args; +typedef struct { + uint32_t n_tokens; + uint32_t pos0; + uint32_t selected_pools; + uint32_t index_topk; + uint32_t pool_size; + uint32_t output_width; +} ds4_gpu_glm53_expand_pool_selection_args; + typedef struct { uint32_t n_tokens; uint32_t n_head; @@ -6174,6 +6199,8 @@ typedef struct { uint32_t head_dim; uint32_t pos0; uint32_t cache_f16; + uint32_t row_group_size; + uint32_t pad0; uint64_t q_token_stride; uint64_t q_head_stride; uint64_t weights_token_stride; @@ -6280,6 +6307,7 @@ typedef struct { uint32_t n_tokens; uint32_t mid_token_stride; uint32_t down_type; + float swiglu_clamp; int32_t tp_rank; int32_t tp_world; int32_t tp_expert_base; @@ -8472,6 +8500,8 @@ int ds4_gpu_init(void) { ds4_gpu_get_pipeline("kernel_glm_qkv_norm_store_compact_kv"); g_glm_store_indexer_k_pipeline = ds4_gpu_get_pipeline("kernel_glm_store_indexer_k"); + g_glm53_indexer_pool_update_pipeline = + ds4_gpu_get_pipeline("kernel_glm53_indexer_pool_update"); g_glm_build_kv_cache_pipeline = ds4_gpu_get_pipeline("kernel_glm_build_kv_cache"); g_glm_build_kv_cache_decode_group4_pipeline = @@ -8484,6 +8514,8 @@ int ds4_gpu_init(void) { ds4_gpu_get_pipeline("kernel_glm_fill_selected_range"); g_glm_fill_selected_range_batch_pipeline = ds4_gpu_get_pipeline("kernel_glm_fill_selected_range_batch"); + g_glm53_expand_pool_selection_pipeline = + ds4_gpu_get_pipeline("kernel_glm53_expand_pool_selection"); g_glm_indexer_rope_tail_pipeline = ds4_gpu_get_pipeline("kernel_glm_indexer_rope_tail_f32"); g_glm_indexer_score_one_pipeline = @@ -8599,12 +8631,14 @@ int ds4_gpu_init(void) { !g_glm_store_compact_kv_pipeline || !g_glm_qkv_norm_store_compact_kv_pipeline || !g_glm_store_indexer_k_pipeline || + !g_glm53_indexer_pool_update_pipeline || !g_glm_build_kv_cache_pipeline || !g_glm_build_kv_cache_decode_group4_pipeline || !g_glm_build_kv_cache_flash_pipeline || !g_glm_attention_full_pipeline || !g_glm_fill_selected_range_pipeline || !g_glm_fill_selected_range_batch_pipeline || + !g_glm53_expand_pool_selection_pipeline || !g_glm_indexer_rope_tail_pipeline || !g_glm_indexer_score_one_pipeline || !g_glm_indexer_score_one_direct_pipeline || @@ -10342,12 +10376,14 @@ void ds4_gpu_cleanup(void) { g_glm_store_compact_kv_pipeline = nil; g_glm_qkv_norm_store_compact_kv_pipeline = nil; g_glm_store_indexer_k_pipeline = nil; + g_glm53_indexer_pool_update_pipeline = nil; g_glm_build_kv_cache_pipeline = nil; g_glm_build_kv_cache_decode_group4_pipeline = nil; g_glm_build_kv_cache_flash_pipeline = nil; g_glm_attention_full_pipeline = nil; g_glm_fill_selected_range_pipeline = nil; g_glm_fill_selected_range_batch_pipeline = nil; + g_glm53_expand_pool_selection_pipeline = nil; g_glm_indexer_rope_tail_pipeline = nil; g_glm_indexer_score_one_pipeline = nil; g_glm_indexer_score_one_direct_pipeline = nil; @@ -17282,41 +17318,41 @@ static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_address_table( if (!expert_buffers) return nil; for (uint32_t first = 0; first < n_total_expert; first += table_group_size) { - const uint32_t remaining = n_total_expert - first; - const uint32_t group_n = - remaining < table_group_size ? remaining : table_group_size; - if ((uint64_t)first > UINT64_MAX / expert_bytes || - (uint64_t)group_n > UINT64_MAX / expert_bytes) { - fprintf(stderr, "ds4: Metal Q4 expert address table group byte overflow\n"); - return nil; - } - const uint64_t rel = (uint64_t)first * expert_bytes; - const uint64_t group_bytes = (uint64_t)group_n * expert_bytes; - if (rel > UINT64_MAX - tensor_offset) { - fprintf(stderr, "ds4: Metal Q4 expert address table group offset overflow\n"); - return nil; - } - uint64_t inner = 0; - id group_buf = nil; - if (table_group_size > 1) { - group_buf = ds4_gpu_wrap_model_range(model_map, - model_size, - tensor_offset + rel, - group_bytes, - &inner); - } else { - group_buf = ds4_gpu_wrap_model_exact_range_owned(model_map, - model_size, - tensor_offset + rel, - expert_bytes, - &inner); - } - if (!group_buf) return nil; - const uint64_t base_address = (uint64_t)[group_buf gpuAddress] + inner; - for (uint32_t j = 0; j < group_n; j++) { - addresses[first + j] = base_address + (uint64_t)j * expert_bytes; - } - [expert_buffers addObject:group_buf]; + const uint32_t remaining = n_total_expert - first; + const uint32_t group_n = + remaining < table_group_size ? remaining : table_group_size; + if ((uint64_t)first > UINT64_MAX / expert_bytes || + (uint64_t)group_n > UINT64_MAX / expert_bytes) { + fprintf(stderr, "ds4: Metal Q4 expert address table group byte overflow\n"); + return nil; + } + const uint64_t rel = (uint64_t)first * expert_bytes; + const uint64_t group_bytes = (uint64_t)group_n * expert_bytes; + if (rel > UINT64_MAX - tensor_offset) { + fprintf(stderr, "ds4: Metal Q4 expert address table group offset overflow\n"); + return nil; + } + uint64_t inner = 0; + id group_buf = nil; + if (table_group_size > 1) { + group_buf = ds4_gpu_wrap_model_range(model_map, + model_size, + tensor_offset + rel, + group_bytes, + &inner); + } else { + group_buf = ds4_gpu_wrap_model_exact_range_owned(model_map, + model_size, + tensor_offset + rel, + expert_bytes, + &inner); + } + if (!group_buf) return nil; + const uint64_t base_address = (uint64_t)[group_buf gpuAddress] + inner; + for (uint32_t j = 0; j < group_n; j++) { + addresses[first + j] = base_address + (uint64_t)j * expert_bytes; + } + [expert_buffers addObject:group_buf]; } [address_buffer didModifyRange:NSMakeRange(0, @@ -24814,6 +24850,100 @@ int ds4_gpu_matmul_q8_0_kslice_tensor( } } +int ds4_gpu_matmul_q8_0_kslice_rows_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t full_in_dim, + uint64_t out_dim, + uint64_t k_off, + uint64_t k_cnt, + const ds4_gpu_tensor *x, + uint64_t n_rows) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!ds4_gpu_mpp_available() || + !out || !x || !model_map || + (full_in_dim % 64u) != 0 || (out_dim % 64u) != 0 || + (k_off % 32u) != 0 || (k_cnt % 64u) != 0 || + k_cnt == 0 || k_off + k_cnt > full_in_dim || + (n_rows % 32u) != 0 || n_rows == 0 || + full_in_dim > UINT32_MAX || out_dim > UINT32_MAX || + n_rows > UINT32_MAX) { + return 0; + } + + @autoreleasepool { + const uint64_t row_bytes = (full_in_dim / 32u) * 34u; + const uint64_t weight_bytes = out_dim * row_bytes; + const uint64_t x_bytes = n_rows * full_in_dim * sizeof(float); + const uint64_t out_bytes = n_rows * out_dim * sizeof(float); + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes || + weight_offset > model_size || + weight_bytes > model_size - weight_offset) { + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset); + if (!wbuf) return 0; + + uint64_t tile_rows = 32u; + if ((n_rows % 128u) == 0) tile_rows = 128u; + else if ((n_rows % 64u) == 0) tile_rows = 64u; + const char *fn = tile_rows == 128u + ? "kernel_mul_mm_q8_0_f32_nax_direct_rhs_n128" + : (tile_rows == 64u + ? "kernel_mul_mm_q8_0_f32_nax_direct_rhs_n64" + : "kernel_mul_mm_q8_0_f32_nax_direct_rhs"); + id pipeline = + ds4_gpu_get_mul_mm_pipeline(fn, false, false); + if (!pipeline) return 0; + + ds4_gpu_mul_mm_args args = + ds4_gpu_make_mm_args(k_cnt, out_dim, n_rows, row_bytes); + args.nb01 = row_bytes; + args.nb02 = row_bytes * out_dim; + args.nb03 = row_bytes * out_dim; + args.nb11 = full_in_dim * sizeof(float); + args.nb12 = n_rows * full_in_dim * sizeof(float); + args.nb13 = args.nb12; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf + offset:(NSUInteger)(inner_offset + (k_off / 32u) * 34u) + atIndex:1]; + [enc setBuffer:xbuf + offset:(NSUInteger)(ds4_gpu_tensor_offset(x) + + k_off * sizeof(float)) + atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:2u * 64u * 32u * sizeof(uint16_t) + atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_rows / tile_rows), + (NSUInteger)(out_dim / 64u), + 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, + owned, + "Q8_0 NAX kslice rows matmul"); + } +} + int ds4_gpu_matmul_quant_kslice_tensor( ds4_gpu_tensor *out, const void *model_map, @@ -26967,7 +27097,10 @@ static int ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec( uint32_t n_kv, uint32_t window, uint32_t n_head, - uint32_t head_dim) { + uint32_t head_dim, + bool raw_kv_f16, + bool has_sinks, + float scale) { if (!cbp || !*cbp) return 0; id cb = *cbp; if (head_dim != 512 || n_head == 0 || n_q == 0 || n_kv == 0) { @@ -26978,12 +27111,13 @@ static int ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec( id rawbuf = ds4_gpu_tensor_buffer(raw_kv); id headsbuf = ds4_gpu_tensor_buffer(heads); const uint64_t q_bytes = (uint64_t)n_q * n_head * head_dim * sizeof(float); - const uint64_t raw_bytes = (uint64_t)n_kv * head_dim * sizeof(float); - if (!qbuf || !rawbuf || !headsbuf || !sinks_buf || + const uint64_t raw_bytes = (uint64_t)n_kv * head_dim * + (raw_kv_f16 ? sizeof(uint16_t) : sizeof(float)); + if (!qbuf || !rawbuf || !headsbuf || (has_sinks && !sinks_buf) || ds4_gpu_tensor_bytes(q) < q_bytes || ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || ds4_gpu_tensor_bytes(heads) < q_bytes) { - fprintf(stderr, "ds4: Metal prefill raw DS4 non-vector FlashAttention received undersized buffers\n"); + fprintf(stderr, "ds4: Metal prefill shared-KV FlashAttention received undersized buffers\n"); return 0; } @@ -27088,18 +27222,19 @@ static int ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec( ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg); id attn_pipeline = ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512", - true, true, false, false, has_kvpad, bc_mask, + true, has_sinks, false, false, has_kvpad, bc_mask, (int32_t)head_dim, (int32_t)head_dim, (int32_t)nsg); if (!blk_pipeline || !attn_pipeline) return 0; - if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, - rawbuf, - ds4_gpu_tensor_offset(raw_kv), - g_flash_attn_kv_buffer, - 0, - n_kv * head_dim)) { + if (!ds4_gpu_encode_copy_to_f16_1d(cb, + rawbuf, + ds4_gpu_tensor_offset(raw_kv), + raw_kv_f16, + g_flash_attn_kv_buffer, + 0, + n_kv * head_dim)) { return 0; } DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw"); @@ -27188,7 +27323,7 @@ static int ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec( .ne1 = (int32_t)n_head, .ne2 = (int32_t)n_q, .ne3 = 1, - .scale = 1.0f / sqrtf((float)head_dim), + .scale = scale, .max_bias = 0.0f, .m0 = 0.0f, .m1 = 0.0f, @@ -27208,7 +27343,9 @@ static int ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec( [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; [enc setBuffer:mask_buffer offset:0 atIndex:4]; - [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; + [enc setBuffer:has_sinks ? sinks_buf : qbuf + offset:has_sinks ? sinks_offset : ds4_gpu_tensor_offset(q) + atIndex:5]; [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; [enc setBuffer:blk_buffer offset:0 atIndex:7]; [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8]; @@ -27254,7 +27391,10 @@ static int ds4_gpu_encode_flash_attention_prefill_raw_heads( n_kv, window, n_head, - head_dim); + head_dim, + false, + true, + 1.0f / sqrtf((float)head_dim)); } const uint32_t n_tokens = n_q; @@ -28363,6 +28503,57 @@ int ds4_gpu_attention_prefill_raw_heads_tensor( head_dim); } +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) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!lora_out || !qk_low || !kv_lora_cache || + n_q == 0 || n_kv == 0 || n_kv > cache_cap || + q_row0 >= n_kv || n_q > n_kv - q_row0 || + n_head == 0 || kv_lora_dim != 512u || qk_dim == 0) { + return 0; + } + + @autoreleasepool { + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec( + &cb, + lora_out, + nil, + 0, + qk_low, + kv_lora_cache, + q_row0, + n_q, + n_kv, + 0, + n_head, + kv_lora_dim, + cache_f16, + false, + 1.0f / sqrtf((float)qk_dim))) { + return 0; + } + if (!ds4_gpu_finish_command_buffer( + cb, owned, "GLM dense compact prefill attention")) { + return 0; + } + } + return 1; +} + int ds4_gpu_attention_decode_raw_batch_heads_tensor( ds4_gpu_tensor *heads, const void *model_map, @@ -32420,7 +32611,7 @@ int ds4_gpu_glm_store_compact_kv_tensor( if (!g_initialized && !ds4_gpu_init()) return 0; if (!kv_lora_cache || !k_rope_cache || !kv_norm || !kv_raw || n_tokens == 0 || cache_cap == 0 || - kv_raw_dim == 0 || kv_lora_dim == 0 || qk_rope == 0 || + kv_raw_dim == 0 || kv_lora_dim == 0 || kv_lora_dim > kv_raw_dim || qk_rope > kv_raw_dim - kv_lora_dim || pos0 > cache_cap || n_tokens > cache_cap - pos0) { @@ -32509,7 +32700,7 @@ int ds4_gpu_glm_qkv_norm_store_compact_kv_tensor( if (!g_initialized && !ds4_gpu_init()) return 0; if (!q_out || !q || !kv_lora_cache || !k_rope_cache || !kv_raw || !model_map || n_tokens == 0 || cache_cap == 0 || - q_n == 0 || kv_raw_dim == 0 || kv_lora_dim == 0 || qk_rope == 0 || + q_n == 0 || kv_raw_dim == 0 || kv_lora_dim == 0 || (q_n & 3u) != 0 || (kv_lora_dim & 3u) != 0 || kv_lora_dim > kv_raw_dim || qk_rope > kv_raw_dim - kv_lora_dim || @@ -32713,6 +32904,212 @@ int ds4_gpu_glm_store_indexer_k_tensor( return 1; } +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) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!pool_cache || !tail_k || !tail_gate || !raw_k || !gate || !model_map || + n_tokens == 0 || cache_cap == 0 || head_dim == 0 || pool_size != 4u || + head_dim > 1024u || !isfinite(eps) || eps <= 0.0f || + pos0 > cache_cap || n_tokens > cache_cap - pos0) { + return 0; + } + + @autoreleasepool { + id poolbuf = ds4_gpu_tensor_buffer(pool_cache); + id tailkbuf = ds4_gpu_tensor_buffer(tail_k); + id tailgatebuf = ds4_gpu_tensor_buffer(tail_gate); + id rawbuf = ds4_gpu_tensor_buffer(raw_k); + id gatebuf = ds4_gpu_tensor_buffer(gate); + const uint64_t pool_cap = + ((uint64_t)cache_cap + pool_size - 1u) / pool_size; + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t pool_bytes = pool_cap * head_dim * cache_elem_bytes; + const uint64_t tail_bytes = (uint64_t)pool_size * head_dim * sizeof(float); + const uint64_t input_bytes = (uint64_t)n_tokens * head_dim * sizeof(float); + const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); + const uint64_t ape_bytes = + (uint64_t)pool_size * head_dim * sizeof(uint16_t); + if (!poolbuf || !tailkbuf || !tailgatebuf || !rawbuf || !gatebuf || + ds4_gpu_tensor_bytes(pool_cache) < pool_bytes || + ds4_gpu_tensor_bytes(tail_k) < tail_bytes || + ds4_gpu_tensor_bytes(tail_gate) < tail_bytes || + ds4_gpu_tensor_bytes(raw_k) < input_bytes || + ds4_gpu_tensor_bytes(gate) < input_bytes) { + fprintf(stderr, "ds4: Metal GLM-5.3 indexer pool received undersized buffers\n"); + return 0; + } + if (norm_weight_offset > model_size || + norm_bytes > model_size - norm_weight_offset || + norm_bias_offset > model_size || + norm_bytes > model_size - norm_bias_offset || + ape_offset > model_size || ape_bytes > model_size - ape_offset) { + fprintf(stderr, "ds4: Metal GLM-5.3 indexer pool model range is invalid\n"); + return 0; + } + + uint64_t norm_weight_inner = 0; + uint64_t norm_bias_inner = 0; + uint64_t ape_inner = 0; + id norm_weightbuf = ds4_gpu_wrap_model_range( + model_map, model_size, norm_weight_offset, norm_bytes, &norm_weight_inner); + id norm_biasbuf = ds4_gpu_wrap_model_range( + model_map, model_size, norm_bias_offset, norm_bytes, &norm_bias_inner); + id apebuf = ds4_gpu_wrap_model_range( + model_map, model_size, ape_offset, ape_bytes, &ape_inner); + if (!norm_weightbuf || !norm_biasbuf || !apebuf) return 0; + + id pipeline = ds4_gpu_hot_pipeline( + g_glm53_indexer_pool_update_pipeline, + "kernel_glm53_indexer_pool_update"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBuffer:norm_weightbuf offset:(NSUInteger)norm_weight_inner atIndex:3]; + [enc setBuffer:norm_biasbuf offset:(NSUInteger)norm_bias_inner atIndex:4]; + [enc setBuffer:apebuf offset:(NSUInteger)ape_inner atIndex:5]; + [enc setBuffer:poolbuf offset:ds4_gpu_tensor_offset(pool_cache) atIndex:6]; + [enc setBuffer:tailkbuf offset:ds4_gpu_tensor_offset(tail_k) atIndex:7]; + [enc setBuffer:tailgatebuf offset:ds4_gpu_tensor_offset(tail_gate) atIndex:8]; + [enc setThreadgroupMemoryLength: + ((NSUInteger)pool_size * head_dim + 2u * pool_size) * sizeof(float) + atIndex:0]; + + uint32_t done = 0; + const uint32_t leading = pos0 % pool_size; + if (leading != 0u) { + uint32_t chunk = pool_size - leading; + if (chunk > n_tokens) chunk = n_tokens; + ds4_gpu_glm53_indexer_pool_update_args args = { + .pos0 = pos0, + .n_tokens = chunk, + .cache_cap = cache_cap, + .head_dim = head_dim, + .pool_size = pool_size, + .cache_f16 = cache_f16 ? 1u : 0u, + .eps = eps, + .pad0 = 0, + }; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_k) atIndex:1]; + [enc setBuffer:gatebuf offset:ds4_gpu_tensor_offset(gate) atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(head_dim, 1, 1)]; + done = chunk; + } + if (done < n_tokens) { + const uint32_t chunk = n_tokens - done; + ds4_gpu_glm53_indexer_pool_update_args args = { + .pos0 = pos0 + done, + .n_tokens = chunk, + .cache_cap = cache_cap, + .head_dim = head_dim, + .pool_size = pool_size, + .cache_f16 = cache_f16 ? 1u : 0u, + .eps = eps, + .pad0 = 0, + }; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:rawbuf + offset:ds4_gpu_tensor_offset(raw_k) + + (NSUInteger)((uint64_t)done * head_dim * sizeof(float)) + atIndex:1]; + [enc setBuffer:gatebuf + offset:ds4_gpu_tensor_offset(gate) + + (NSUInteger)((uint64_t)done * head_dim * sizeof(float)) + atIndex:2]; + const uint32_t groups = (chunk + pool_size - 1u) / pool_size; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(head_dim, 1, 1)]; + } + ds4_gpu_end_compute_encoder(cb, enc); + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 indexer pool update")) { + return 0; + } + } + return 1; +} + +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) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!raw_selected || !pool_selected || n_tokens == 0 || + selected_pools == 0 || index_topk == 0 || pool_size == 0 || + selected_pools > index_topk / pool_size || + output_width < index_topk || output_width < index_topk + pool_size - 1u) { + return 0; + } + @autoreleasepool { + id rawbuf = ds4_gpu_tensor_buffer(raw_selected); + id poolbuf = ds4_gpu_tensor_buffer(pool_selected); + const uint64_t raw_bytes = + (uint64_t)n_tokens * output_width * sizeof(uint32_t); + const uint64_t pool_bytes = + (uint64_t)n_tokens * selected_pools * sizeof(uint32_t); + if (!rawbuf || !poolbuf || + ds4_gpu_tensor_bytes(raw_selected) < raw_bytes || + ds4_gpu_tensor_bytes(pool_selected) < pool_bytes) { + fprintf(stderr, "ds4: Metal GLM-5.3 pool expansion received undersized buffers\n"); + return 0; + } + id pipeline = ds4_gpu_hot_pipeline( + g_glm53_expand_pool_selection_pipeline, + "kernel_glm53_expand_pool_selection"); + if (!pipeline) return 0; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + ds4_gpu_glm53_expand_pool_selection_args args = { + .n_tokens = n_tokens, + .pos0 = pos0, + .selected_pools = selected_pools, + .index_topk = index_topk, + .pool_size = pool_size, + .output_width = output_width, + }; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:poolbuf offset:ds4_gpu_tensor_offset(pool_selected) atIndex:1]; + [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_selected) atIndex:2]; + const NSUInteger total = (NSUInteger)n_tokens * output_width; + const NSUInteger nth = MIN((NSUInteger)256, pipeline.maxTotalThreadsPerThreadgroup); + [enc dispatchThreads:MTLSizeMake(total, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 pool selection expansion")) { + return 0; + } + } + return 1; +} + int ds4_gpu_glm_build_kv_cache_tensor( ds4_gpu_tensor *key_cache, ds4_gpu_tensor *value_cache, @@ -32741,12 +33138,12 @@ int ds4_gpu_glm_build_kv_cache_tensor( if (!key_cache || !value_cache || !kv_raw || !k_nope || !value || n_tokens == 0 || cache_cap == 0 || n_head == 0 || kv_raw_dim == 0 || kv_lora_dim == 0 || - qk_nope == 0 || qk_rope == 0 || value_dim == 0 || + qk_nope == 0 || value_dim == 0 || kv_lora_dim + qk_rope > kv_raw_dim || qk_dim < qk_nope || (qk_rope & 1u) != 0 || pos0 > cache_cap || n_tokens > cache_cap - pos0 || - !isfinite(freq_base) || freq_base <= 0.0f || - !isfinite(freq_scale) || freq_scale <= 0.0f || + (qk_rope != 0 && (!isfinite(freq_base) || freq_base <= 0.0f)) || + (qk_rope != 0 && (!isfinite(freq_scale) || freq_scale <= 0.0f)) || !isfinite(ext_factor) || !isfinite(attn_factor) || !isfinite(beta_fast) || !isfinite(beta_slow)) { return 0; @@ -32860,12 +33257,12 @@ int ds4_gpu_glm_build_kv_cache_flash_tensor( if (!key_cache || !value_cache || !kv_raw || !k_nope || !value || pos0 != 0 || n_tokens == 0 || cache_cap == 0 || n_head == 0 || kv_raw_dim == 0 || kv_lora_dim == 0 || - qk_nope == 0 || qk_rope == 0 || value_dim == 0 || + qk_nope == 0 || value_dim == 0 || kv_lora_dim + qk_rope > kv_raw_dim || qk_dim < qk_nope || (qk_rope & 1u) != 0 || n_tokens > cache_cap || - !isfinite(freq_base) || freq_base <= 0.0f || - !isfinite(freq_scale) || freq_scale <= 0.0f || + (qk_rope != 0 && (!isfinite(freq_base) || freq_base <= 0.0f)) || + (qk_rope != 0 && (!isfinite(freq_scale) || freq_scale <= 0.0f)) || !isfinite(ext_factor) || !isfinite(attn_factor) || !isfinite(beta_fast) || !isfinite(beta_slow)) { return 0; @@ -33721,7 +34118,7 @@ int ds4_gpu_glm_indexer_score_one_tensor( return 1; } -int ds4_gpu_glm_indexer_scores_batch_tensor( +static int ds4_gpu_glm_indexer_scores_batch_grouped_tensor( ds4_gpu_tensor *scores, const ds4_gpu_tensor *q, const ds4_gpu_tensor *weights, @@ -33729,14 +34126,17 @@ int ds4_gpu_glm_indexer_scores_batch_tensor( uint32_t n_rows, uint32_t n_tokens, uint32_t pos0, + uint32_t row_group_size, uint32_t n_head, uint32_t head_dim, float scale, bool cache_f16) { if (!g_initialized && !ds4_gpu_init()) return 0; if (!scores || !q || !weights || !indexer_key_cache || - n_rows == 0 || n_tokens == 0 || n_head == 0 || head_dim != 128 || - pos0 >= n_rows || n_tokens > n_rows - pos0 || + n_rows == 0 || n_tokens == 0 || row_group_size == 0 || + n_head == 0 || head_dim != 128 || + pos0 > UINT32_MAX - n_tokens || + (uint64_t)(pos0 + n_tokens) / row_group_size > n_rows || !isfinite(scale) || scale <= 0.0f) { return 0; } @@ -33785,6 +34185,8 @@ int ds4_gpu_glm_indexer_scores_batch_tensor( .head_dim = head_dim, .pos0 = pos0, .cache_f16 = cache_f16 ? 1u : 0u, + .row_group_size = row_group_size, + .pad0 = 0, .q_token_stride = (uint64_t)n_head * head_dim * sizeof(float), .q_head_stride = (uint64_t)head_dim * sizeof(float), .weights_token_stride = (uint64_t)n_head * sizeof(float), @@ -33830,6 +34232,42 @@ int ds4_gpu_glm_indexer_scores_batch_tensor( return 1; } +int ds4_gpu_glm_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 n_head, + uint32_t head_dim, + float scale, + bool cache_f16) { + return ds4_gpu_glm_indexer_scores_batch_grouped_tensor( + scores, q, weights, indexer_key_cache, n_rows, n_tokens, pos0, 1u, + n_head, head_dim, scale, 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) { + if (pool_size != 4u) return 0; + return ds4_gpu_glm_indexer_scores_batch_grouped_tensor( + scores, q, weights, indexer_key_cache, n_rows, n_tokens, pos0, + pool_size, n_head, head_dim, scale, cache_f16); +} + int ds4_gpu_glm_qk_lowrank_typed_tensor( ds4_gpu_tensor *qk_low, const ds4_gpu_tensor *q, @@ -33894,7 +34332,7 @@ int ds4_gpu_glm_qk_lowrank_typed_tensor( * shape for both Q8_0 k_b and the DenseQ4 GGUF's Q4_0 k_b (the Q8 * fast path above never engaged there). */ const int use_glm52_sg = - n_head == 64u && + (n_head == 32u || n_head == 64u) && kv_lora_dim == 512u && qk_nope == 192u && qk_dim == 256u && @@ -34291,10 +34729,10 @@ int ds4_gpu_glm_attention_indexed_decode_typed_tensor( !model_map || !selected || n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || n_head == 0 || kv_lora_dim == 0 || - qk_nope == 0 || qk_rope == 0 || (qk_rope & 1u) != 0 || + qk_nope == 0 || (qk_rope & 1u) != 0 || value_dim == 0 || qk_dim < qk_nope || - !isfinite(freq_base) || freq_base <= 0.0f || - !isfinite(freq_scale) || freq_scale <= 0.0f || + (qk_rope != 0 && (!isfinite(freq_base) || freq_base <= 0.0f)) || + (qk_rope != 0 && (!isfinite(freq_scale) || freq_scale <= 0.0f)) || !isfinite(ext_factor) || !isfinite(attn_factor) || !isfinite(beta_fast) || !isfinite(beta_slow)) { return 0; @@ -34495,8 +34933,8 @@ int ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor( block_rows == 0u || needed_blocks == 0u || n_blocks < needed_blocks || n_blocks > 64u || !cache_f16 || - !isfinite(freq_base) || freq_base <= 0.0f || - !isfinite(freq_scale) || freq_scale <= 0.0f || + (qk_rope != 0 && (!isfinite(freq_base) || freq_base <= 0.0f)) || + (qk_rope != 0 && (!isfinite(freq_scale) || freq_scale <= 0.0f)) || !isfinite(ext_factor) || !isfinite(attn_factor) || !isfinite(beta_fast) || !isfinite(beta_slow)) { return 0; @@ -34740,10 +35178,10 @@ int ds4_gpu_glm_attention_indexed_batch_typed_tensor( !model_map || !selected || n_tokens == 0 || n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || n_head == 0 || kv_lora_dim == 0 || - qk_nope == 0 || qk_rope == 0 || (qk_rope & 1u) != 0 || + qk_nope == 0 || (qk_rope & 1u) != 0 || value_dim == 0 || qk_dim < qk_nope || - !isfinite(freq_base) || freq_base <= 0.0f || - !isfinite(freq_scale) || freq_scale <= 0.0f || + (qk_rope != 0 && (!isfinite(freq_base) || freq_base <= 0.0f)) || + (qk_rope != 0 && (!isfinite(freq_scale) || freq_scale <= 0.0f)) || !isfinite(ext_factor) || !isfinite(attn_factor) || !isfinite(beta_fast) || !isfinite(beta_slow)) { return 0; @@ -35049,10 +35487,10 @@ static int ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor( if (!lora_out || !q || !qk_low || !kv_lora_cache || !k_rope_cache || !selected || n_tokens == 0 || n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || n_head == 0 || kv_lora_dim == 0 || - qk_nope == 0 || qk_rope == 0 || (qk_rope & 1u) != 0 || + qk_nope == 0 || (qk_rope & 1u) != 0 || qk_dim < qk_nope || - !isfinite(freq_base) || freq_base <= 0.0f || - !isfinite(freq_scale) || freq_scale <= 0.0f || + (qk_rope != 0 && (!isfinite(freq_base) || freq_base <= 0.0f)) || + (qk_rope != 0 && (!isfinite(freq_scale) || freq_scale <= 0.0f)) || !isfinite(ext_factor) || !isfinite(attn_factor) || !isfinite(beta_fast) || !isfinite(beta_slow)) { return 0; @@ -35090,7 +35528,8 @@ static int ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor( } const bool use_vec_lora = - cache_f16 && kv_lora_dim == 512u && qk_rope == 64u; + cache_f16 && kv_lora_dim == 512u && + (qk_rope == 0u || qk_rope == 64u); const bool full_head_groups = (n_head % 8u) == 0u; id pipeline = nil; if (use_vec_lora && selected_rows_valid && full_head_groups) { @@ -35201,10 +35640,10 @@ int ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor( n_tokens == 0 || n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || pos0 > n_selected || n_tokens > n_selected - pos0 || n_head == 0 || kv_lora_dim != 512u || - qk_nope == 0 || qk_rope != 64u || - qk_dim < qk_nope || !cache_f16 || - !isfinite(freq_base) || freq_base <= 0.0f || - !isfinite(freq_scale) || freq_scale <= 0.0f || + qk_nope == 0 || (qk_rope != 0u && qk_rope != 64u) || + qk_dim < qk_nope || + (qk_rope != 0u && (!isfinite(freq_base) || freq_base <= 0.0f)) || + (qk_rope != 0u && (!isfinite(freq_scale) || freq_scale <= 0.0f)) || !isfinite(ext_factor) || !isfinite(attn_factor) || !isfinite(beta_fast) || !isfinite(beta_slow)) { return 0; @@ -35216,7 +35655,8 @@ int ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor( id lowbuf = ds4_gpu_tensor_buffer(qk_low); id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); - const uint64_t cache_elem_bytes = sizeof(uint16_t); + const uint64_t cache_elem_bytes = + cache_f16 ? sizeof(uint16_t) : sizeof(float); const uint64_t lora_bytes = (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); const uint64_t q_bytes = @@ -35255,7 +35695,7 @@ int ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor( .n_tokens = n_tokens, .n_selected = n_selected, .cache_cap = cache_cap, - .cache_f16 = 1u, + .cache_f16 = cache_f16 ? 1u : 0u, .n_head = n_head, .kv_lora_dim = kv_lora_dim, .qk_nope = qk_nope, @@ -35406,7 +35846,7 @@ int ds4_gpu_glm_router_select_tensor( float expert_weight_scale) { if (!g_initialized && !ds4_gpu_init()) return 0; if (!selected || !weights || !probs || !logits || !model_map || - n_expert == 0 || n_expert > 256u || + n_expert == 0 || n_expert > 512u || n_expert_used == 0 || n_expert_used > n_expert) { return 0; } @@ -35471,9 +35911,11 @@ int ds4_gpu_glm_router_select_tensor( [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:3]; [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:4]; [enc setBuffer:probsbuf offset:ds4_gpu_tensor_offset(probs) atIndex:5]; - [enc setThreadgroupMemoryLength:256u * sizeof(float) + 256u * sizeof(int32_t) atIndex:0]; + const NSUInteger router_threads = n_expert > 256u ? 512u : 256u; + [enc setThreadgroupMemoryLength:router_threads * + (sizeof(float) + sizeof(int32_t)) atIndex:0]; [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + threadsPerThreadgroup:MTLSizeMake(router_threads, 1, 1)]; ds4_gpu_end_compute_encoder(cb, enc); if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM router select")) return 0; @@ -35497,7 +35939,7 @@ int ds4_gpu_glm_router_select_batch_tensor( if (!g_initialized && !ds4_gpu_init()) return 0; if (!selected || !weights || !probs || !logits || !model_map || n_tokens == 0 || - n_expert == 0 || n_expert > 256u || + n_expert == 0 || n_expert > 512u || n_expert_used == 0 || n_expert_used > n_expert) { return 0; } @@ -35555,9 +35997,11 @@ int ds4_gpu_glm_router_select_batch_tensor( [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:3]; [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:4]; [enc setBuffer:probsbuf offset:ds4_gpu_tensor_offset(probs) atIndex:5]; - [enc setThreadgroupMemoryLength:256u * sizeof(float) + 256u * sizeof(int32_t) atIndex:0]; + const NSUInteger router_threads = n_expert > 256u ? 512u : 256u; + [enc setThreadgroupMemoryLength:router_threads * + (sizeof(float) + sizeof(int32_t)) atIndex:0]; [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + threadsPerThreadgroup:MTLSizeMake(router_threads, 1, 1)]; ds4_gpu_end_compute_encoder(cb, enc); if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM batch router select")) return 0; @@ -35606,6 +36050,7 @@ int ds4_gpu_glm_routed_moe_one_tensor( const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, + float swiglu_clamp, uint32_t layer_index, const ds4_gpu_tensor *x, bool force_resident) { @@ -35616,6 +36061,15 @@ int ds4_gpu_glm_routed_moe_one_tensor( uint32_t first_expert = 0; uint32_t n_bind_expert = 0; ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert); + /* Streaming address tables point at individually cached experts by their + * global id. Keep the source offsets global; ownership is enforced by the + * kernels before they dereference an address-table slot. */ + const bool streaming_tp = + g_ssd_streaming_mode && g_tp_split_world == 2; + if (streaming_tp) { + first_expert = 0; + n_bind_expert = n_total_expert; + } const int32_t tp_expert_base_host = (int32_t)first_expert; gate_offset += (uint64_t)first_expert * gate_expert_bytes; up_offset += (uint64_t)first_expert * up_expert_bytes; @@ -36081,6 +36535,7 @@ int ds4_gpu_glm_routed_moe_one_tensor( .n_tokens = 1, .mid_token_stride = n_expert * expert_mid_dim, .down_type = down_type, + .swiglu_clamp = swiglu_clamp, .gate_expert_bytes = gate_expert_bytes, .gate_row_bytes = gate_row_bytes, .up_expert_bytes = up_expert_bytes, @@ -36360,6 +36815,7 @@ static int ds4_gpu_glm_routed_moe_batch_grouped_tensor( const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, + float swiglu_clamp, uint32_t layer_index, const ds4_gpu_tensor *x, uint32_t n_tokens) { @@ -36594,7 +37050,7 @@ static int ds4_gpu_glm_routed_moe_batch_grouped_tensor( ds4_gpu_tensor_offset(weights), expert_mid_dim, pair_rows, - 0.0f, + swiglu_clamp, mid_f16); } DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("activation_weight"); @@ -36661,6 +37117,7 @@ static int ds4_gpu_glm_routed_moe_batch_grouped_addr_tensor( const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, + float swiglu_clamp, uint32_t layer_index, const ds4_gpu_tensor *x, uint32_t n_tokens, @@ -36873,7 +37330,7 @@ static int ds4_gpu_glm_routed_moe_batch_grouped_addr_tensor( ds4_gpu_tensor_offset(weights), expert_mid_dim, pair_rows, - 0.0f, + swiglu_clamp, mid_f16); } DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("activation_weight"); @@ -36943,6 +37400,7 @@ static int ds4_gpu_glm_routed_moe_batch_tensor_impl( const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, + float swiglu_clamp, uint32_t layer_index, const ds4_gpu_tensor *x, uint32_t n_tokens, @@ -37030,6 +37488,7 @@ static int ds4_gpu_glm_routed_moe_batch_tensor_impl( weights, n_total_expert, n_expert, + swiglu_clamp, layer_index, x, n_tokens); @@ -37038,6 +37497,15 @@ static int ds4_gpu_glm_routed_moe_batch_tensor_impl( uint32_t first_expert = 0; uint32_t n_bind_expert = 0; ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert); + /* Streaming address tables are indexed by global expert id. As in the + * one-token path, keep model offsets global and let the kernels enforce + * rank ownership before dereferencing the selected address. */ + const bool streaming_tp = + g_ssd_streaming_mode && g_tp_split_world == 2; + if (streaming_tp) { + first_expert = 0; + n_bind_expert = n_total_expert; + } gate_offset += (uint64_t)first_expert * gate_expert_bytes; up_offset += (uint64_t)first_expert * up_expert_bytes; down_offset += (uint64_t)first_expert * down_expert_bytes; @@ -37243,6 +37711,7 @@ static int ds4_gpu_glm_routed_moe_batch_tensor_impl( weights, n_total_expert, n_expert, + swiglu_clamp, layer_index, x, n_tokens, @@ -37335,6 +37804,7 @@ static int ds4_gpu_glm_routed_moe_batch_tensor_impl( .n_tokens = n_tokens, .mid_token_stride = mid_token_stride, .down_type = down_type, + .swiglu_clamp = swiglu_clamp, .gate_expert_bytes = gate_expert_bytes, .gate_row_bytes = gate_row_bytes, .up_expert_bytes = up_expert_bytes, @@ -37460,6 +37930,7 @@ int ds4_gpu_glm_routed_moe_batch_tensor( const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, + float swiglu_clamp, uint32_t layer_index, const ds4_gpu_tensor *x, uint32_t n_tokens, @@ -37489,6 +37960,7 @@ int ds4_gpu_glm_routed_moe_batch_tensor( weights, n_total_expert, n_expert, + swiglu_clamp, layer_index, x, n_tokens, @@ -37521,6 +37993,7 @@ int ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor( const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, + float swiglu_clamp, uint32_t layer_index, const ds4_gpu_tensor *x, uint32_t n_tokens, @@ -37548,6 +38021,7 @@ int ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor( weights, n_total_expert, n_expert, + swiglu_clamp, layer_index, x, n_tokens, @@ -42792,8 +43266,14 @@ int ds4_gpu_hc_expand_split_tensor( const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out_hc || !block_out || !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0; + if (!g_initialized && !ds4_gpu_init()) { + fprintf(stderr, "ds4: Metal HC expand split could not initialize the backend\n"); + return 0; + } + if (!out_hc || !block_out || !residual_hc || !split || n_embd == 0 || n_hc == 0) { + fprintf(stderr, "ds4: Metal HC expand split received invalid arguments\n"); + return 0; + } @autoreleasepool { id blockbuf = ds4_gpu_tensor_buffer(block_out); @@ -42865,12 +43345,18 @@ int ds4_gpu_hc_expand_split_tensor( "kernel_dsv4_hc_expand4"); n_elem = (uint64_t)n_embd * n_tokens64; } - if (!expand_pipeline) return 0; + if (!expand_pipeline) { + fprintf(stderr, "ds4: Metal HC expand split pipeline is unavailable\n"); + return 0; + } const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem)); const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth; int owned = 0; id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; + if (!cb) { + fprintf(stderr, "ds4: Metal HC expand split could not obtain a command buffer\n"); + return 0; + } id enc = ds4_gpu_compute_encoder(cb); [enc setComputePipelineState:expand_pipeline]; @@ -43271,6 +43757,1002 @@ int ds4_gpu_matmul_q8_0_hc_expand_tensor( return 1; } + + +/* Kimi Delta Attention primitive shared with the GLM-5.3 graph. */ +static int glm53_gpu_mul_u64(uint64_t a, uint64_t b, uint64_t *out) { + if (!out || (a != 0 && b > UINT64_MAX / a)) return 0; + *out = a * b; + return 1; +} + +static int glm53_gpu_tensor_has( + const ds4_gpu_tensor *tensor, + uint64_t elements, + uint64_t element_size) { + uint64_t bytes = 0; + return tensor && + glm53_gpu_mul_u64(elements, element_size, &bytes) && + ds4_gpu_tensor_buffer(tensor) != nil && + ds4_gpu_tensor_bytes(tensor) >= bytes; +} + +static id glm53_gpu_weight_buffer( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t bytes, + uint64_t *inner_offset, + const char *label) { + if (!model_map || offset > model_size || bytes > model_size - offset) { + fprintf(stderr, "ds4: GLM-5.3 %s range is outside the mapped model\n", + label ? label : "weight"); + return nil; + } + return ds4_gpu_wrap_model_range(model_map, model_size, offset, bytes, + inner_offset); +} + +typedef struct { + uint32_t in_dim; + uint32_t out_dim; + uint32_t n_rows; +} glm53_gpu_bf16_matmul_args; + +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) { + if (!g_initialized && !ds4_gpu_init()) return 0; + uint64_t weights = 0, output = 0; + if (n_tokens == 0 || n_embd == 0 || n_vocab == 0 || + !glm53_gpu_mul_u64(n_embd, n_vocab, &weights) || + !glm53_gpu_mul_u64(n_embd, n_tokens, &output) || + !glm53_gpu_tensor_has(token_ids, n_tokens, sizeof(int32_t)) || + !glm53_gpu_tensor_has(out, output, sizeof(float))) { + fprintf(stderr, "ds4: GLM-5.3 BF16 embedding received invalid buffers\n"); + return 0; + } + @autoreleasepool { + uint64_t inner = 0; + id weightbuf = glm53_gpu_weight_buffer( + model_map, model_size, weight_offset, weights * sizeof(uint16_t), + &inner, "BF16 embedding"); + id pipeline = + ds4_gpu_get_pipeline("kernel_glm53_embedding_bf16"); + if (!weightbuf || !pipeline) return 0; + glm53_gpu_bf16_matmul_args args = { + .in_dim = n_embd, + .out_dim = n_vocab, + .n_rows = n_tokens, + }; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weightbuf offset:(NSUInteger)inner atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(token_ids) + offset:ds4_gpu_tensor_offset(token_ids) atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(out) + offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc dispatchThreads:MTLSizeMake(n_embd, n_tokens, 1) + threadsPerThreadgroup:MTLSizeMake(MIN((NSUInteger)n_embd, 256u), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 BF16 embedding"); + } +} + +static uint32_t glm53_gpu_bf16_mv_nsg(void) { + return ds4_gpu_device_name_contains("M3 Ultra") && + getenv("DS4_METAL_DISABLE_M3_ULTRA_GLM53_DECODE") == NULL && + getenv("DS4_METAL_DISABLE_M3_ULTRA_GLM53_BF16_NSG4") == NULL + ? 4u : 8u; +} + +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) { + if (!g_initialized && !ds4_gpu_init()) return 0; + uint64_t weights = 0, input = 0, output = 0; + if (in_dim == 0 || out_dim == 0 || n_rows == 0 || + !glm53_gpu_mul_u64(in_dim, out_dim, &weights) || + !glm53_gpu_mul_u64(in_dim, n_rows, &input) || + !glm53_gpu_mul_u64(out_dim, n_rows, &output) || + !glm53_gpu_tensor_has(x, input, sizeof(float)) || + !glm53_gpu_tensor_has(out, output, sizeof(float))) { + fprintf(stderr, "ds4: GLM-5.3 BF16 matmul received invalid buffers\n"); + return 0; + } + @autoreleasepool { + uint64_t inner = 0; + id weightbuf = glm53_gpu_weight_buffer( + model_map, model_size, weight_offset, weights * sizeof(uint16_t), + &inner, "BF16 matrix"); + if (!weightbuf) return 0; + const bool use_mv = n_rows <= 8u; + const bool bc_inp = (in_dim % 32u) != 0u; + const bool bc_out = (out_dim % 64u) != 0u || (n_rows % 32u) != 0u; + id pipeline = use_mv + ? ds4_gpu_get_pipeline("kernel_glm53_mul_mv_bf16_f32") + : ds4_gpu_get_mul_mm_pipeline( + "kernel_glm53_mul_mm_bf16_f32", bc_inp, bc_out); + if (!pipeline) return 0; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + if (use_mv) { + const uint32_t nsg = glm53_gpu_bf16_mv_nsg(); + glm53_gpu_bf16_matmul_args args = { + .in_dim = in_dim, + .out_dim = out_dim, + .n_rows = n_rows, + }; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weightbuf offset:(NSUInteger)inner atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(x) + offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(out) + offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake((out_dim + nsg - 1u) / nsg, + n_rows, 1) + threadsPerThreadgroup:MTLSizeMake(32u * nsg, 1, 1)]; + } else { + ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args( + in_dim, out_dim, n_rows, (uint64_t)in_dim * sizeof(uint16_t)); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weightbuf offset:(NSUInteger)inner atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(x) + offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(out) + offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((n_rows + 31u) / 32u, + (out_dim + 63u) / 64u, 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + } + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 BF16 matmul"); + } +} + +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) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!ds4_gpu_device_name_contains("M3 Ultra")) return 0; + uint64_t weights = 0; + if (in_dim == 0 || out_dim == 0 || + !glm53_gpu_mul_u64(in_dim, out_dim, &weights) || + !glm53_gpu_tensor_has(x, in_dim, sizeof(float)) || + !glm53_gpu_tensor_has(out_q, out_dim, sizeof(float)) || + !glm53_gpu_tensor_has(out_k, out_dim, sizeof(float)) || + !glm53_gpu_tensor_has(out_v, out_dim, sizeof(float))) { + return 0; + } + + @autoreleasepool { + const uint64_t weight_bytes = weights * sizeof(uint16_t); + uint64_t inner_q = 0, inner_k = 0, inner_v = 0; + id weight_q = glm53_gpu_weight_buffer( + model_map, model_size, weight_q_offset, weight_bytes, + &inner_q, "BF16 Q matrix"); + id weight_k = glm53_gpu_weight_buffer( + model_map, model_size, weight_k_offset, weight_bytes, + &inner_k, "BF16 K matrix"); + id weight_v = glm53_gpu_weight_buffer( + model_map, model_size, weight_v_offset, weight_bytes, + &inner_v, "BF16 V matrix"); + id pipeline = + ds4_gpu_get_pipeline("kernel_glm53_mul_mv_bf16_f32_qkv"); + if (!weight_q || !weight_k || !weight_v || !pipeline) return 0; + + const uint32_t nsg = glm53_gpu_bf16_mv_nsg(); + glm53_gpu_bf16_matmul_args args = { + .in_dim = in_dim, + .out_dim = out_dim, + .n_rows = 1u, + }; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weight_q offset:(NSUInteger)inner_q atIndex:1]; + [enc setBuffer:weight_k offset:(NSUInteger)inner_k atIndex:2]; + [enc setBuffer:weight_v offset:(NSUInteger)inner_v atIndex:3]; + [enc setBuffer:ds4_gpu_tensor_buffer(x) + offset:ds4_gpu_tensor_offset(x) atIndex:4]; + [enc setBuffer:ds4_gpu_tensor_buffer(out_q) + offset:ds4_gpu_tensor_offset(out_q) atIndex:5]; + [enc setBuffer:ds4_gpu_tensor_buffer(out_k) + offset:ds4_gpu_tensor_offset(out_k) atIndex:6]; + [enc setBuffer:ds4_gpu_tensor_buffer(out_v) + offset:ds4_gpu_tensor_offset(out_v) atIndex:7]; + [enc dispatchThreadgroups:MTLSizeMake((out_dim + nsg - 1u) / nsg, + 1u, 3u) + threadsPerThreadgroup:MTLSizeMake(32u * nsg, 1u, 1u)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, + "GLM-5.3 BF16 QKV matmul"); + } +} + +typedef struct { + uint32_t width; + uint32_t rows; + float eps; +} glm53_vision_rows_args; + +typedef struct { + uint32_t rows; + uint32_t grid_h; + uint32_t grid_w; + float eps; +} glm53_vision_qkv_args; + +typedef struct { + uint32_t rows; + float scale; +} glm53_vision_attention_args; + +typedef struct { + uint32_t dst_row; + uint32_t image_row; + uint32_t rows; + uint32_t total_rows; + uint32_t width; + uint32_t hc; +} glm53_vision_scatter_args; + +static int glm53_vision_dispatch_rows( + const char *kernel, + ds4_gpu_tensor *out, + const ds4_gpu_tensor *x, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t width, + uint32_t rows, + float eps, + const char *label) { + uint64_t inner = 0; + id weight = glm53_gpu_weight_buffer( + model_map, model_size, weight_offset, + (uint64_t)width * sizeof(uint16_t), &inner, label); + id pipeline = ds4_gpu_get_pipeline(kernel); + if (!weight || !pipeline) return 0; + glm53_vision_rows_args args = { width, rows, eps }; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(x) + offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setBuffer:weight offset:(NSUInteger)inner atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(out) + offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:8u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256u, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, label); +} + +static int glm53_vision_dispatch_bias( + const char *kernel, + ds4_gpu_tensor *x, + const ds4_gpu_tensor *residual, + const void *model_map, + uint64_t model_size, + uint64_t bias_offset, + uint32_t width, + uint32_t rows, + const char *label) { + uint64_t inner = 0; + id bias = glm53_gpu_weight_buffer( + model_map, model_size, bias_offset, + (uint64_t)width * sizeof(uint16_t), &inner, label); + id pipeline = ds4_gpu_get_pipeline(kernel); + if (!bias || !pipeline) return 0; + glm53_vision_rows_args args = { width, rows, 0.0f }; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(x) + offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setBuffer:bias offset:(NSUInteger)inner atIndex:2]; + if (residual) { + [enc setBuffer:ds4_gpu_tensor_buffer(residual) + offset:ds4_gpu_tensor_offset(residual) atIndex:3]; + } + [enc dispatchThreads:MTLSizeMake(width, rows, 1) + threadsPerThreadgroup:MTLSizeMake(MIN((NSUInteger)width, 256u), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, label); +} + +static int glm53_vision_dispatch_qkv( + ds4_gpu_tensor *q, + ds4_gpu_tensor *k, + ds4_gpu_tensor *v, + const ds4_gpu_tensor *qkv, + const void *model_map, + uint64_t model_size, + const ds4_glm53_vision_layer_weights *w, + uint32_t rows, + uint32_t grid_h, + uint32_t grid_w) { + uint64_t bias_inner = 0, q_inner = 0, k_inner = 0; + id bias = glm53_gpu_weight_buffer( + model_map, model_size, w->qkv_bias, 3072u * sizeof(uint16_t), + &bias_inner, "vision QKV bias"); + id qw = glm53_gpu_weight_buffer( + model_map, model_size, w->q_norm, 64u * sizeof(uint16_t), + &q_inner, "vision Q norm"); + id kw = glm53_gpu_weight_buffer( + model_map, model_size, w->k_norm, 64u * sizeof(uint16_t), + &k_inner, "vision K norm"); + id pipeline = + ds4_gpu_get_pipeline("kernel_glm53_vision_qkv_rope"); + if (!bias || !qw || !kw || !pipeline) return 0; + glm53_vision_qkv_args args = { rows, grid_h, grid_w, 1.0e-5f }; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(qkv) + offset:ds4_gpu_tensor_offset(qkv) atIndex:1]; + [enc setBuffer:bias offset:(NSUInteger)bias_inner atIndex:2]; + [enc setBuffer:qw offset:(NSUInteger)q_inner atIndex:3]; + [enc setBuffer:kw offset:(NSUInteger)k_inner atIndex:4]; + [enc setBuffer:ds4_gpu_tensor_buffer(q) + offset:ds4_gpu_tensor_offset(q) atIndex:5]; + [enc setBuffer:ds4_gpu_tensor_buffer(k) + offset:ds4_gpu_tensor_offset(k) atIndex:6]; + [enc setBuffer:ds4_gpu_tensor_buffer(v) + offset:ds4_gpu_tensor_offset(v) atIndex:7]; + [enc dispatchThreadgroups:MTLSizeMake(rows, 16u, 1) + threadsPerThreadgroup:MTLSizeMake(32u, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 vision QKV"); +} + +static int glm53_vision_dispatch_attention( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *k, + const ds4_gpu_tensor *v, + uint32_t rows) { + id pipeline = + ds4_gpu_get_pipeline("kernel_glm53_vision_attention"); + if (!pipeline) return 0; + glm53_vision_attention_args args = { rows, 0.125f }; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(q) + offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(k) + offset:ds4_gpu_tensor_offset(k) atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(v) + offset:ds4_gpu_tensor_offset(v) atIndex:3]; + [enc setBuffer:ds4_gpu_tensor_buffer(out) + offset:ds4_gpu_tensor_offset(out) atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(rows, 16u, 1) + threadsPerThreadgroup:MTLSizeMake(32u, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 vision attention"); +} + +static int glm53_vision_dispatch_swiglu_bias( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *gate, + const ds4_gpu_tensor *up, + const void *model_map, + uint64_t model_size, + uint64_t gate_bias_offset, + uint64_t up_bias_offset, + uint32_t width, + uint32_t rows) { + uint64_t gate_inner = 0, up_inner = 0; + id gate_bias = glm53_gpu_weight_buffer( + model_map, model_size, gate_bias_offset, + (uint64_t)width * sizeof(uint16_t), &gate_inner, + "vision gate bias"); + id up_bias = glm53_gpu_weight_buffer( + model_map, model_size, up_bias_offset, + (uint64_t)width * sizeof(uint16_t), &up_inner, + "vision up bias"); + id pipeline = + ds4_gpu_get_pipeline("kernel_glm53_vision_swiglu_bias"); + if (!gate_bias || !up_bias || !pipeline) return 0; + glm53_vision_rows_args args = { width, rows, 0.0f }; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(gate) + offset:ds4_gpu_tensor_offset(gate) atIndex:1]; + [enc setBuffer:gate_bias offset:(NSUInteger)gate_inner atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(up) + offset:ds4_gpu_tensor_offset(up) atIndex:3]; + [enc setBuffer:up_bias offset:(NSUInteger)up_inner atIndex:4]; + [enc setBuffer:ds4_gpu_tensor_buffer(out) + offset:ds4_gpu_tensor_offset(out) atIndex:5]; + [enc dispatchThreads:MTLSizeMake(width, rows, 1) + threadsPerThreadgroup:MTLSizeMake(MIN((NSUInteger)width, 256u), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 vision SwiGLU"); +} + +static int glm53_vision_dispatch_reorder( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *x, + uint32_t rows) { + id pipeline = + ds4_gpu_get_pipeline("kernel_glm53_vision_downsample_reorder"); + if (!pipeline) return 0; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBuffer:ds4_gpu_tensor_buffer(x) + offset:ds4_gpu_tensor_offset(x) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(out) + offset:ds4_gpu_tensor_offset(out) atIndex:1]; + [enc dispatchThreads:MTLSizeMake(4096u, rows, 1) + threadsPerThreadgroup:MTLSizeMake(256u, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 vision downsample reorder"); +} + +static int glm53_vision_dispatch_layernorm_gelu( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *x, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t bias_offset, + uint32_t rows) { + uint64_t weight_inner = 0, bias_inner = 0; + id weight = glm53_gpu_weight_buffer( + model_map, model_size, weight_offset, 4096u * sizeof(uint16_t), + &weight_inner, "vision merger norm"); + id bias = glm53_gpu_weight_buffer( + model_map, model_size, bias_offset, 4096u * sizeof(uint16_t), + &bias_inner, "vision merger norm bias"); + id pipeline = + ds4_gpu_get_pipeline("kernel_glm53_vision_layernorm_gelu"); + if (!weight || !bias || !pipeline) return 0; + glm53_vision_rows_args args = { 4096u, rows, 1.0e-5f }; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(x) + offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setBuffer:weight offset:(NSUInteger)weight_inner atIndex:2]; + [enc setBuffer:bias offset:(NSUInteger)bias_inner atIndex:3]; + [enc setBuffer:ds4_gpu_tensor_buffer(out) + offset:ds4_gpu_tensor_offset(out) atIndex:4]; + [enc setThreadgroupMemoryLength:8u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256u, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 vision merger norm"); +} + +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) { + if (!out || !patches || !model_map || !weights || grid_h == 0 || grid_w == 0 || + (grid_h & 1u) != 0 || (grid_w & 1u) != 0 || + grid_h > UINT32_MAX / grid_w || ds4_gpu_commands_active()) { + fprintf(stderr, "ds4: invalid GLM-5.3 vision encoder input\n"); + return 0; + } + if (!g_initialized && !ds4_gpu_init()) return 0; + const uint32_t rows = grid_h * grid_w; + const uint32_t merged_rows = rows / 4u; + const uint64_t row1024 = (uint64_t)rows * 1024u; + const uint64_t row3072 = (uint64_t)rows * 3072u; + const uint64_t row4096 = (uint64_t)rows * 4096u; + const uint64_t merged10240 = (uint64_t)merged_rows * 10240u; + if (row4096 > SIZE_MAX / sizeof(float) || + merged10240 > SIZE_MAX / sizeof(float)) return 0; + + ds4_gpu_tensor *patch = NULL, *a = NULL, *b = NULL, *qkv = NULL; + ds4_gpu_tensor *q = NULL, *k = NULL, *v = NULL, *attn = NULL; + ds4_gpu_tensor *gate = NULL, *up = NULL, *mid = NULL; + int ok = 0; +#define VISION_ALLOC(name_, count_) do { \ + name_ = ds4_gpu_tensor_alloc((count_) * sizeof(float)); \ + if (!(name_)) goto cleanup; \ + } while (0) + VISION_ALLOC(patch, (uint64_t)rows * 1176u); + VISION_ALLOC(a, row1024); + VISION_ALLOC(b, row1024); + VISION_ALLOC(qkv, row3072); + VISION_ALLOC(q, row1024); + VISION_ALLOC(k, row1024); + VISION_ALLOC(v, row1024); + VISION_ALLOC(attn, row1024); + VISION_ALLOC(gate, row4096); + VISION_ALLOC(up, row4096); + VISION_ALLOC(mid, row4096); +#undef VISION_ALLOC + if (!ds4_gpu_tensor_write(patch, 0, patches, + (uint64_t)rows * 1176u * sizeof(float))) goto cleanup; + if (!ds4_gpu_begin_commands()) goto cleanup; + ok = ds4_gpu_glm53_matmul_bf16(a, model_map, model_size, + weights->patch_weight, + 1176u, 1024u, patch, rows); + if (ok) ok = glm53_vision_dispatch_bias( + "kernel_glm53_vision_add_bias", a, NULL, model_map, model_size, + weights->patch_bias, 1024u, rows, "vision patch bias"); + + ds4_gpu_tensor *cur = a; + ds4_gpu_tensor *tmp = b; + for (uint32_t il = 0; ok && il < DS4_GLM53_VISION_LAYERS; il++) { + const ds4_glm53_vision_layer_weights *w = &weights->layer[il]; + ok = glm53_vision_dispatch_rows( + "kernel_glm53_vision_rms_bf16", tmp, cur, + model_map, model_size, w->norm1, 1024u, rows, 1.0e-5f, + "vision norm1"); + if (ok) ok = ds4_gpu_glm53_matmul_bf16( + qkv, model_map, model_size, w->qkv_weight, + 1024u, 3072u, tmp, rows); + if (ok) ok = glm53_vision_dispatch_qkv( + q, k, v, qkv, model_map, model_size, w, rows, grid_h, grid_w); + if (ok) ok = glm53_vision_dispatch_attention(attn, q, k, v, rows); + if (ok) ok = ds4_gpu_glm53_matmul_bf16( + tmp, model_map, model_size, w->attn_proj_weight, + 1024u, 1024u, attn, rows); + if (ok) ok = glm53_vision_dispatch_bias( + "kernel_glm53_vision_bias_residual", tmp, cur, + model_map, model_size, w->attn_proj_bias, + 1024u, rows, "vision attention bias"); + ds4_gpu_tensor *swap = cur; cur = tmp; tmp = swap; + + if (ok) ok = glm53_vision_dispatch_rows( + "kernel_glm53_vision_rms_bf16", tmp, cur, + model_map, model_size, w->norm2, 1024u, rows, 1.0e-5f, + "vision norm2"); + if (ok) ok = ds4_gpu_glm53_matmul_bf16( + gate, model_map, model_size, w->gate_weight, + 1024u, 4096u, tmp, rows); + if (ok) ok = ds4_gpu_glm53_matmul_bf16( + up, model_map, model_size, w->up_weight, + 1024u, 4096u, tmp, rows); + if (ok) ok = glm53_vision_dispatch_swiglu_bias( + mid, gate, up, model_map, model_size, + w->gate_bias, w->up_bias, 4096u, rows); + if (ok) ok = ds4_gpu_glm53_matmul_bf16( + tmp, model_map, model_size, w->down_weight, + 4096u, 1024u, mid, rows); + if (ok) ok = glm53_vision_dispatch_bias( + "kernel_glm53_vision_bias_residual", tmp, cur, + model_map, model_size, w->down_bias, + 1024u, rows, "vision down bias"); + swap = cur; cur = tmp; tmp = swap; + } + if (ok) ok = glm53_vision_dispatch_rows( + "kernel_glm53_vision_rms_bf16", tmp, cur, + model_map, model_size, weights->post_norm, + 1024u, rows, 1.0e-5f, "vision post norm"); + if (ok) ok = glm53_vision_dispatch_reorder(cur, tmp, merged_rows); + if (ok) ok = ds4_gpu_glm53_matmul_bf16( + tmp, model_map, model_size, weights->downsample_weight, + 4096u, 4096u, cur, merged_rows); + if (ok) ok = glm53_vision_dispatch_bias( + "kernel_glm53_vision_add_bias", tmp, NULL, + model_map, model_size, weights->downsample_bias, + 4096u, merged_rows, "vision downsample bias"); + if (ok) ok = ds4_gpu_glm53_matmul_bf16( + cur, model_map, model_size, weights->merger_proj, + 4096u, 4096u, tmp, merged_rows); + if (ok) ok = glm53_vision_dispatch_layernorm_gelu( + tmp, cur, model_map, model_size, + weights->merger_norm, weights->merger_norm_bias, merged_rows); + if (ok) ok = ds4_gpu_glm53_matmul_bf16( + gate, model_map, model_size, weights->merger_gate, + 4096u, 10240u, tmp, merged_rows); + if (ok) ok = ds4_gpu_glm53_matmul_bf16( + up, model_map, model_size, weights->merger_up, + 4096u, 10240u, tmp, merged_rows); + if (ok) ok = ds4_gpu_swiglu_tensor( + mid, gate, up, merged_rows * 10240u, 10.0f, 1.0f); + if (ok) ok = ds4_gpu_glm53_matmul_bf16( + cur, model_map, model_size, weights->merger_down, + 10240u, 4096u, mid, merged_rows); + if (ok) ok = ds4_gpu_end_commands(); + else (void)ds4_gpu_end_commands(); + if (ok) ok = ds4_gpu_tensor_read( + cur, 0, out, (uint64_t)merged_rows * 4096u * sizeof(float)); + +cleanup: + ds4_gpu_tensor_free(mid); + ds4_gpu_tensor_free(up); + ds4_gpu_tensor_free(gate); + ds4_gpu_tensor_free(attn); + ds4_gpu_tensor_free(v); + ds4_gpu_tensor_free(k); + ds4_gpu_tensor_free(q); + ds4_gpu_tensor_free(qkv); + ds4_gpu_tensor_free(b); + ds4_gpu_tensor_free(a); + ds4_gpu_tensor_free(patch); + return ok; +} + +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) { + if (!hc || !image || rows == 0 || n_embd == 0 || n_hc == 0 || + dst_row > total_rows || rows > total_rows - dst_row || + !ds4_gpu_commands_active()) return 0; + id pipeline = + ds4_gpu_get_pipeline("kernel_glm53_vision_scatter_hc"); + if (!pipeline) return 0; + glm53_vision_scatter_args args = { + dst_row, image_row, rows, total_rows, n_embd, n_hc + }; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(hc) + offset:ds4_gpu_tensor_offset(hc) atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(image) + offset:ds4_gpu_tensor_offset(image) atIndex:2]; + [enc dispatchThreads:MTLSizeMake(n_embd, rows * n_hc, 1) + threadsPerThreadgroup:MTLSizeMake(MIN((NSUInteger)n_embd, 256u), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 vision scatter"); +} + +typedef struct { + uint32_t n_heads; + uint32_t n_rows; + float lower_bound; + float norm_eps; +} glm53_gpu_kda_args; + +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) { + enum { GLM53_KDA_DIM = 128, GLM53_KDA_HISTORY = 3 }; + if (!g_initialized && !ds4_gpu_init()) return 0; + uint64_t projection = 0, activation_elements = 0; + uint64_t conv_elements = 0, state_elements = 0; + if (n_heads == 0 || n_rows == 0 || gate_lower_bound >= 0.0f || + !glm53_gpu_mul_u64(n_heads, GLM53_KDA_DIM, &projection) || + !glm53_gpu_mul_u64(projection, n_rows, &activation_elements) || + !glm53_gpu_mul_u64(activation_elements, + 3u * GLM53_KDA_HISTORY, &conv_elements) || + !glm53_gpu_mul_u64(activation_elements, GLM53_KDA_DIM, &state_elements) || + !glm53_gpu_tensor_has(q, activation_elements, sizeof(float)) || + !glm53_gpu_tensor_has(k, activation_elements, sizeof(float)) || + !glm53_gpu_tensor_has(v, activation_elements, sizeof(float)) || + !glm53_gpu_tensor_has(raw_gate, activation_elements, sizeof(float)) || + !glm53_gpu_tensor_has(output_gate, activation_elements, sizeof(float)) || + !glm53_gpu_tensor_has(raw_beta, + (uint64_t)n_rows * n_heads, sizeof(float)) || + !glm53_gpu_tensor_has(out, activation_elements, sizeof(float)) || + !glm53_gpu_tensor_has(conv_state, conv_elements, sizeof(float)) || + !glm53_gpu_tensor_has(recurrent_state, state_elements, sizeof(float))) { + fprintf(stderr, "ds4: GLM-5.3 KDA decode received invalid buffers\n"); + return 0; + } + + uint64_t conv_bytes = 0, dt_bytes = 0; + if (!glm53_gpu_mul_u64(projection, 4u * sizeof(float), &conv_bytes) || + !glm53_gpu_mul_u64(projection, sizeof(float), &dt_bytes)) { + return 0; + } + const uint64_t a_log_bytes = (uint64_t)n_heads * sizeof(float); + const uint64_t norm_bytes = GLM53_KDA_DIM * sizeof(float); + + @autoreleasepool { + uint64_t qw_inner = 0, kw_inner = 0, vw_inner = 0; + uint64_t a_inner = 0, dt_inner = 0, norm_inner = 0; + id qw = glm53_gpu_weight_buffer( + model_map, model_size, q_conv_offset, conv_bytes, + &qw_inner, "KDA Q convolution"); + id kw = glm53_gpu_weight_buffer( + model_map, model_size, k_conv_offset, conv_bytes, + &kw_inner, "KDA K convolution"); + id vw = glm53_gpu_weight_buffer( + model_map, model_size, v_conv_offset, conv_bytes, + &vw_inner, "KDA V convolution"); + id a_log = glm53_gpu_weight_buffer( + model_map, model_size, a_log_offset, a_log_bytes, + &a_inner, "KDA A_log"); + id dt_bias = glm53_gpu_weight_buffer( + model_map, model_size, dt_bias_offset, dt_bytes, + &dt_inner, "KDA dt bias"); + id output_norm = glm53_gpu_weight_buffer( + model_map, model_size, output_norm_offset, norm_bytes, + &norm_inner, "KDA output norm"); + id pipeline = + ds4_gpu_get_pipeline("kernel_glm53_kda_decode"); + if (!qw || !kw || !vw || !a_log || !dt_bias || !output_norm || + !pipeline) { + return 0; + } + + glm53_gpu_kda_args args = { + .n_heads = n_heads, + .n_rows = n_rows, + .lower_bound = gate_lower_bound, + .norm_eps = norm_eps, + }; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(q) + offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(k) + offset:ds4_gpu_tensor_offset(k) atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(v) + offset:ds4_gpu_tensor_offset(v) atIndex:3]; + [enc setBuffer:ds4_gpu_tensor_buffer(raw_gate) + offset:ds4_gpu_tensor_offset(raw_gate) atIndex:4]; + [enc setBuffer:ds4_gpu_tensor_buffer(raw_beta) + offset:ds4_gpu_tensor_offset(raw_beta) atIndex:5]; + [enc setBuffer:ds4_gpu_tensor_buffer(output_gate) + offset:ds4_gpu_tensor_offset(output_gate) atIndex:6]; + [enc setBuffer:qw offset:(NSUInteger)qw_inner atIndex:7]; + [enc setBuffer:kw offset:(NSUInteger)kw_inner atIndex:8]; + [enc setBuffer:vw offset:(NSUInteger)vw_inner atIndex:9]; + [enc setBuffer:a_log offset:(NSUInteger)a_inner atIndex:10]; + [enc setBuffer:dt_bias offset:(NSUInteger)dt_inner atIndex:11]; + [enc setBuffer:output_norm offset:(NSUInteger)norm_inner atIndex:12]; + [enc setBuffer:ds4_gpu_tensor_buffer(conv_state) + offset:ds4_gpu_tensor_offset(conv_state) atIndex:13]; + [enc setBuffer:ds4_gpu_tensor_buffer(recurrent_state) + offset:ds4_gpu_tensor_offset(recurrent_state) atIndex:14]; + [enc setBuffer:ds4_gpu_tensor_buffer(out) + offset:ds4_gpu_tensor_offset(out) atIndex:15]; + [enc setThreadgroupMemoryLength:656u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_rows, n_heads, 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer( + cb, owned, "GLM-5.3 fused KDA decode"); + } +} + +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) { + enum { GLM53_KDA_DIM = 128, GLM53_KDA_HISTORY = 3 }; + if (!g_initialized && !ds4_gpu_init()) return 0; + uint64_t projection = 0, activation_elements = 0; + uint64_t conv_elements = 0, state_elements = 0; + if (n_heads == 0 || n_tokens == 0 || gate_lower_bound >= 0.0f || + !glm53_gpu_mul_u64(n_heads, GLM53_KDA_DIM, &projection) || + !glm53_gpu_mul_u64(projection, n_tokens, &activation_elements) || + !glm53_gpu_mul_u64(projection, + 3u * GLM53_KDA_HISTORY, &conv_elements) || + !glm53_gpu_mul_u64(projection, GLM53_KDA_DIM, &state_elements) || + !glm53_gpu_tensor_has(q, activation_elements, sizeof(float)) || + !glm53_gpu_tensor_has(k, activation_elements, sizeof(float)) || + !glm53_gpu_tensor_has(v, activation_elements, sizeof(float)) || + !glm53_gpu_tensor_has(raw_gate, activation_elements, sizeof(float)) || + !glm53_gpu_tensor_has(output_gate, activation_elements, sizeof(float)) || + !glm53_gpu_tensor_has(raw_beta, + (uint64_t)n_tokens * n_heads, sizeof(float)) || + !glm53_gpu_tensor_has(out, activation_elements, sizeof(float)) || + !glm53_gpu_tensor_has(conv_state, conv_elements, sizeof(float)) || + !glm53_gpu_tensor_has(recurrent_state, state_elements, sizeof(float))) { + fprintf(stderr, "ds4: GLM-5.3 KDA prefill received invalid buffers\n"); + return 0; + } + + uint64_t conv_bytes = 0, dt_bytes = 0; + if (!glm53_gpu_mul_u64(projection, 4u * sizeof(float), &conv_bytes) || + !glm53_gpu_mul_u64(projection, sizeof(float), &dt_bytes)) { + return 0; + } + const uint64_t a_log_bytes = (uint64_t)n_heads * sizeof(float); + const uint64_t norm_bytes = GLM53_KDA_DIM * sizeof(float); + + @autoreleasepool { + uint64_t qw_inner = 0, kw_inner = 0, vw_inner = 0; + uint64_t a_inner = 0, dt_inner = 0, norm_inner = 0; + id qw = glm53_gpu_weight_buffer( + model_map, model_size, q_conv_offset, conv_bytes, + &qw_inner, "KDA Q convolution"); + id kw = glm53_gpu_weight_buffer( + model_map, model_size, k_conv_offset, conv_bytes, + &kw_inner, "KDA K convolution"); + id vw = glm53_gpu_weight_buffer( + model_map, model_size, v_conv_offset, conv_bytes, + &vw_inner, "KDA V convolution"); + id a_log = glm53_gpu_weight_buffer( + model_map, model_size, a_log_offset, a_log_bytes, + &a_inner, "KDA A_log"); + id dt_bias = glm53_gpu_weight_buffer( + model_map, model_size, dt_bias_offset, dt_bytes, + &dt_inner, "KDA dt bias"); + id output_norm = glm53_gpu_weight_buffer( + model_map, model_size, output_norm_offset, norm_bytes, + &norm_inner, "KDA output norm"); + id prep_pipeline = + ds4_gpu_get_pipeline("kernel_glm53_kda_prefill_prepare"); + id recurrence_pipeline = + ds4_gpu_get_pipeline("kernel_glm53_kda_prefill_recurrence"); + id output_pipeline = + ds4_gpu_get_pipeline("kernel_glm53_kda_prefill_output"); + if (!qw || !kw || !vw || !a_log || !dt_bias || !output_norm || + !prep_pipeline || !recurrence_pipeline || !output_pipeline) { + return 0; + } + + glm53_gpu_kda_args args = { + .n_heads = n_heads, + .n_rows = n_tokens, + .lower_bound = gate_lower_bound, + .norm_eps = norm_eps, + }; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + + [enc setComputePipelineState:prep_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(q) + offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(k) + offset:ds4_gpu_tensor_offset(k) atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(v) + offset:ds4_gpu_tensor_offset(v) atIndex:3]; + [enc setBuffer:ds4_gpu_tensor_buffer(raw_gate) + offset:ds4_gpu_tensor_offset(raw_gate) atIndex:4]; + [enc setBuffer:qw offset:(NSUInteger)qw_inner atIndex:5]; + [enc setBuffer:kw offset:(NSUInteger)kw_inner atIndex:6]; + [enc setBuffer:vw offset:(NSUInteger)vw_inner atIndex:7]; + [enc setBuffer:a_log offset:(NSUInteger)a_inner atIndex:8]; + [enc setBuffer:dt_bias offset:(NSUInteger)dt_inner atIndex:9]; + [enc setBuffer:ds4_gpu_tensor_buffer(conv_state) + offset:ds4_gpu_tensor_offset(conv_state) atIndex:10]; + [enc setThreadgroupMemoryLength:264u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_heads, 1, 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + + [enc setComputePipelineState:recurrence_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(q) + offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(k) + offset:ds4_gpu_tensor_offset(k) atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(v) + offset:ds4_gpu_tensor_offset(v) atIndex:3]; + [enc setBuffer:ds4_gpu_tensor_buffer(raw_gate) + offset:ds4_gpu_tensor_offset(raw_gate) atIndex:4]; + [enc setBuffer:ds4_gpu_tensor_buffer(raw_beta) + offset:ds4_gpu_tensor_offset(raw_beta) atIndex:5]; + [enc setBuffer:ds4_gpu_tensor_buffer(recurrent_state) + offset:ds4_gpu_tensor_offset(recurrent_state) atIndex:6]; + [enc setBuffer:ds4_gpu_tensor_buffer(out) + offset:ds4_gpu_tensor_offset(out) atIndex:7]; + [enc dispatchThreadgroups:MTLSizeMake(n_heads, 32, 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + + [enc setComputePipelineState:output_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(out) + offset:ds4_gpu_tensor_offset(out) atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(output_gate) + offset:ds4_gpu_tensor_offset(output_gate) atIndex:2]; + [enc setBuffer:output_norm offset:(NSUInteger)norm_inner atIndex:3]; + [enc setThreadgroupMemoryLength:4u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_tokens, n_heads, 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer( + cb, owned, "GLM-5.3 KDA layer-major prefill"); + } +} + void ds4_gpu_set_glm_mtp_verify_mode(bool enabled) { (void)enabled; } diff --git a/src/agent.rs b/src/agent.rs index 7b5dd72..9da4c55 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -3612,7 +3612,7 @@ pub(crate) fn parse_tool_calls( model: ModelChoice, text: &str, ) -> Result<(String, Vec), String> { - let (content, calls) = if model == ModelChoice::Glm52 { + let (content, calls) = if model.is_glm() { parse_glm_calls(text)? } else { crate::dsml::parse_tool_calls(text)? @@ -3643,7 +3643,7 @@ fn system_prompt_with_tools( ralph_child: bool, ) -> String { let schemas = tool_schemas(dev_brain, ralph_child); - let tools = if model == ModelChoice::Glm52 { + let tools = if model.is_glm() { 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\n{schemas}\n\n\nFor a function call, output exactly: function-namekeyvalue\nTool calls are not allowed inside . 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." ) diff --git a/src/app.rs b/src/app.rs index 1f56cbf..6e09930 100644 --- a/src/app.rs +++ b/src/app.rs @@ -35,6 +35,7 @@ use crate::settings::{ ReasoningMode, RuntimePreferences, SpeculativePreferences, SsdPreferences, SteeringPreferences, StreamingCacheBudget, }; +use base64::Engine as _; use iced::widget::{markdown, scrollable, text_editor}; use iced::{Size, Subscription, Task, keyboard, mouse, window}; use rfd::AsyncFileDialog; @@ -115,6 +116,7 @@ pub(crate) struct App { project_name_input: String, model_download: ModelDownload, pub(super) composer: text_editor::Content, + pub(super) pending_vision_image: Option, pub(super) queued_inputs: VecDeque, pub(super) conversation: Vec, /// Follow appended chat content until the user scrolls away from the tail. @@ -191,6 +193,7 @@ struct ChatSnapshot { selected_session: Option, permission_mode: PermissionMode, composer: text_editor::Content, + pending_vision_image: Option, queued_inputs: VecDeque, conversation: Vec, chat_follow_tail: bool, @@ -231,6 +234,11 @@ struct ChatSnapshot { skip_compaction_once: bool, } +pub(super) struct PendingVisionImage { + pub(super) name: String, + marker: String, +} + #[derive(Clone, Debug, Eq, PartialEq)] pub(super) struct ProjectChoice { id: i32, @@ -480,6 +488,9 @@ pub(crate) enum Message { StopModelDownload, DownloadProgressTick, ComposerAction(text_editor::Action), + ChooseVisionImage, + VisionImagePicked(Result)>, String>), + RemoveVisionImage, TranscriptAction(usize, text_editor::Action), ToggleReasoning(usize), OpenLink(markdown::Uri), @@ -641,6 +652,7 @@ impl App { project_name_input: String::new(), model_download: ModelDownload::Idle, composer: text_editor::Content::new(), + pending_vision_image: None, queued_inputs: VecDeque::new(), conversation: Vec::new(), chat_follow_tail: true, @@ -805,6 +817,7 @@ impl App { project_name_input: String::new(), model_download: ModelDownload::Idle, composer: text_editor::Content::new(), + pending_vision_image: None, queued_inputs: VecDeque::new(), conversation: Vec::new(), chat_follow_tail: true, @@ -882,6 +895,7 @@ impl App { selected_session: self.selected_session.take(), permission_mode: self.permission_mode, composer: std::mem::take(&mut self.composer), + pending_vision_image: self.pending_vision_image.take(), queued_inputs: std::mem::take(&mut self.queued_inputs), conversation: std::mem::take(&mut self.conversation), 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.permission_mode = snapshot.permission_mode; self.composer = snapshot.composer; + self.pending_vision_image = snapshot.pending_vision_image; self.queued_inputs = snapshot.queued_inputs; self.conversation = snapshot.conversation; self.chat_follow_tail = snapshot.chat_follow_tail; @@ -1290,6 +1305,55 @@ impl App { } Message::MetricsTick => self.sample_metrics(), 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) => { if !action.is_edit() && let Some(message) = self.conversation.get_mut(index) diff --git a/src/app/generation.rs b/src/app/generation.rs index e56290b..da2b6e1 100644 --- a/src/app/generation.rs +++ b/src/app/generation.rs @@ -517,7 +517,10 @@ impl App { if self.selected_project.is_none() { 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() { return; } @@ -659,13 +662,20 @@ impl App { #[cfg(target_os = "macos")] let assistant_reasoning = effective.turn.reasoning_mode != ReasoningMode::Direct; #[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 { format!( - "{prompt}\n\nA2UI client metadata:\n{}", + "{base_model_prompt}\n\nA2UI client metadata:\n{}", self.a2ui.client_metadata() ) } else { - prompt.clone() + base_model_prompt }; #[cfg(target_os = "macos")] let mut injected_system = self.workspace_instruction_messages(&[], opening_turn); @@ -750,8 +760,8 @@ impl App { }; let mut saved = match database.start_chat_turn( session_id, - &prompt, - (model_prompt != prompt).then_some(model_prompt.as_str()), + &visible_prompt, + (model_prompt != visible_prompt).then_some(model_prompt.as_str()), &injected_system, assistant_reasoning, ) { @@ -811,6 +821,7 @@ impl App { } assistant.reasoning_open = assistant_reasoning; self.composer = text_editor::Content::new(); + self.pending_vision_image = None; self.conversation.push(user); self.conversation.push(assistant); 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 /// usable came back. pub(super) fn session_title(reply: &str) -> Option { @@ -2485,7 +2508,7 @@ mod tests { ChatMessage, TOOL_PROTOCOL_CORRECTION, TurnSummary, chat_turn, compacted_context_start, correction_already_sent, extension_context_visibility, has_chat_after_last_compaction, 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::model::ModelChoice; @@ -2553,6 +2576,20 @@ mod tests { 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] fn legacy_generation_rows_collapse_into_their_user_turn() { let mut user = assistant(None, "question"); diff --git a/src/app/preferences.rs b/src/app/preferences.rs index 18ff261..87790e1 100644 --- a/src/app/preferences.rs +++ b/src/app/preferences.rs @@ -593,6 +593,9 @@ impl App { return; } 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; #[cfg(target_os = "macos")] { @@ -653,9 +656,11 @@ impl App { self.preference_draft.default_reasoning_mode = mode; self.preference_draft.load_generation(model, mode); self.preference_draft.load_acceleration(model); - if model == ModelChoice::Glm52 { + if model.is_glm() { self.preference_draft.power_percent.clear(); self.preference_draft.prefill_chunk.clear(); + } + if model == ModelChoice::Glm52 { self.preference_draft.directional_steering_file.clear(); self.preference_draft.directional_steering_ffn.clear(); self.preference_draft.directional_steering_attn.clear(); @@ -851,7 +856,7 @@ impl App { } Message::PreferenceGlmMtpChanged(value) => { 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 { self.preference_draft.glm_mtp_timing = false; } @@ -859,7 +864,7 @@ impl App { } Message::PreferenceGlmMtpTimingChanged(value) => { 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 { self.preference_draft.glm_mtp = true; } @@ -907,7 +912,7 @@ impl App { self.preference_error = None; } 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_error = None; diff --git a/src/app/view/chat.rs b/src/app/view/chat.rs index 5ccbecf..0956a95 100644 --- a/src/app/view/chat.rs +++ b/src/app/view/chat.rs @@ -250,7 +250,8 @@ impl App { .padding(8) .style(stop_button_style) .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) } else { 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 }; 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 { let mut queued = queued.replace('\n', " "); if queued.chars().count() > 120 { @@ -322,10 +337,32 @@ impl App { .padding([2, 6]) .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.push( 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( context_pie(context_fraction, 19), container( diff --git a/src/app/view/preferences.rs b/src/app/view/preferences.rs index ab56ea8..8095502 100644 --- a/src/app/view/preferences.rs +++ b/src/app/view/preferences.rs @@ -14,12 +14,16 @@ impl App { .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.", ); - let glm_mtp_toggle: Option Message> = - (self.preference_draft.acceleration_model == ModelChoice::Glm52) - .then_some(Message::PreferenceGlmMtpChanged); - let glm_mtp_timing_toggle: Option Message> = - (self.preference_draft.acceleration_model == ModelChoice::Glm52) - .then_some(Message::PreferenceGlmMtpTimingChanged); + let glm_mtp_toggle: Option Message> = self + .preference_draft + .acceleration_model + .supports_glm_mtp() + .then_some(Message::PreferenceGlmMtpChanged); + let glm_mtp_timing_toggle: Option Message> = self + .preference_draft + .acceleration_model + .supports_glm_mtp() + .then_some(Message::PreferenceGlmMtpTimingChanged); let dspark_strict_toggle: Option Message> = self .preference_draft .acceleration_model @@ -71,14 +75,16 @@ impl App { "0.8 (DeepSeek V4 Flash default)", &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); prefill = prefill.on_input(Message::PreferencePrefillChunkChanged); + } + if self.preference_draft.model != ModelChoice::Glm52 { steering_file = steering_file.on_input(Message::PreferenceSteeringFileChanged); steering_ffn = steering_ffn.on_input(Message::PreferenceSteeringFfnChanged); 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); } if self.preference_draft.acceleration_model.supports_dspark() { @@ -467,6 +473,12 @@ impl App { text_input("32768", &self.preference_draft.context_tokens) .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( "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.", @@ -550,8 +562,8 @@ impl App { .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.", ), - text(if self.preference_draft.model == ModelChoice::Glm52 { - "GLM 5.2 uses full GPU power and selects prefill chunks automatically." + text(if self.preference_draft.model.is_glm() { + "GLM uses full GPU power and selects prefill chunks automatically." } else { "Blank numeric values preserve DS4's automatic engine behavior." }) @@ -595,7 +607,7 @@ impl App { toggle(self.preference_draft.glm_mtp) .label("Enable integrated GLM MTP") .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( toggle(self.preference_draft.glm_mtp_timing) @@ -623,7 +635,7 @@ impl App { ), text(if self.preference_draft.acceleration_model.supports_dspark() { "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." } else { "No speculative-decoding support is available for this model." diff --git a/src/engine.rs b/src/engine.rs index 87dccb8..3bde425 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -12,8 +12,10 @@ use crate::model::{ModelChoice, validate_engine_artifacts}; #[cfg(target_os = "macos")] use crate::settings::TurnSettings; use crate::settings::{EngineSettings, ReasoningMode}; +#[cfg(target_os = "macos")] +use base64::Engine as _; 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")] use kvstore::{KvStore, StoreReason}; @@ -32,6 +34,19 @@ use std::time::Instant; use tokenizer::Tokenizer; 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")] unsafe extern "C" { fn mmap( @@ -47,14 +62,14 @@ unsafe extern "C" { 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")] pub(crate) use kvstore::sweep_unreachable as sweep_transient_cache; #[cfg(target_os = "macos")] 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 PLAIN: &[u32] = &[F16, F32]; const DSPARK_DENSE: &[u32] = &[F16, F32, Q8_0]; @@ -211,9 +226,53 @@ const GLM: Shape = Shape { 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 { main: Gguf, support: Option, + vision: Option, support_kind: Option, shape: Shape, tokenizer: Tokenizer, @@ -226,6 +285,7 @@ pub(crate) struct ModelSummary { pub(crate) tensor_count: usize, pub(crate) vocabulary_size: usize, pub(crate) support_loaded: bool, + pub(crate) vision_loaded: bool, } impl Model { @@ -249,6 +309,14 @@ impl Model { model.support = Some(support); 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) } @@ -266,6 +334,7 @@ impl Model { Ok(Self { main, support: None, + vision: None, support_kind: None, shape, tokenizer, @@ -275,14 +344,21 @@ impl Model { pub(crate) fn summary(&self) -> ModelSummary { ModelSummary { 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() + self .support .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(), support_loaded: self.support.is_some(), + vision_loaded: self.vision.is_some(), } } @@ -293,6 +369,9 @@ impl Model { if let Some(support) = &self.support { hash.update(support.checkpoint_identity()); } + if let Some(vision) = &self.vision { + hash.update(vision.checkpoint_identity()); + } hash.finalize().into() } @@ -732,7 +811,7 @@ impl Generator { } 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 { return Err("generation cancelled while updating the system prompt cache".into()); } @@ -1055,11 +1134,12 @@ impl Generator { &mut self, tokens: &[i32], reused: usize, + has_vision: bool, cancelled: &AtomicBool, progress: &mut impl FnMut(u32, u32, Option), ) -> Result { 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(); self.executor.prefill(suffix, |used| { progress(used, context, None); @@ -1079,6 +1159,87 @@ impl Generator { } } + fn render_multimodal_conversation( + &mut self, + system: &str, + messages: &[ChatTurn], + reasoning: ReasoningMode, + ) -> Result<(Vec, 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( &mut self, messages: &[ChatTurn], @@ -1087,29 +1248,41 @@ impl Generator { emit: &mut impl FnMut(bool, String), progress: &mut impl FnMut(u32, u32, Option), ) -> Result<(GenerationOutput, bool), String> { - 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( + let has_vision = messages + .iter() + .any(|message| message.content.contains(VISION_DATA_START)); + let (tokens, overlays) = if has_vision { + self.render_multimodal_conversation( &settings.system_prompt, messages, 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() { return Err("the rendered prompt is empty".into()); @@ -1122,6 +1295,7 @@ impl Generator { )); } let reused = self.executor.align_prompt(&tokens)?; + self.executor.set_vision_overlays(overlays)?; self.metrics.kv_prefix_reused(reused); progress(self.executor.position(), self.executor.context(), None); 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 prompt_tokens = tokens.len(); 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(); if completed != suffix.len() { return Ok(( diff --git a/src/engine/gguf.rs b/src/engine/gguf.rs index cb86ef0..3eadaeb 100644 --- a/src/engine/gguf.rs +++ b/src/engine/gguf.rs @@ -19,6 +19,7 @@ pub(super) const Q5_K: u32 = 13; pub(super) const Q6_K: u32 = 14; pub(super) const IQ2_XXS: u32 = 16; pub(super) const I32: u32 = 26; +pub(super) const BF16: u32 = 30; pub(super) const MXFP4: u32 = 39; #[derive(Clone, Debug)] @@ -488,6 +489,7 @@ impl<'a> Cursor<'a> { | "deepseek4.dspark.target_layer_ids" | "deepseek4.dspark_target_layer_ids" | "dspark.target_layer_ids" + | "glm5-next.layer_types" ); if !keep { for _ in 0..len { diff --git a/src/engine/metal.rs b/src/engine/metal.rs index cd8b0e3..d85ddef 100644 --- a/src/engine/metal.rs +++ b/src/engine/metal.rs @@ -3,12 +3,14 @@ mod glm; mod gpu; mod hotlist; mod profile; +mod vision; +pub(super) use vision::VisionEmbedding; use glm::GlmExecutor; use gpu::*; 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::{Model, ModelFamily, Rng, exact_delta_sample}; use crate::model::ModelChoice; @@ -42,7 +44,7 @@ fn environment_present(name: &CStr) -> bool { !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_DENSE_SOURCE", "dense.metal"), ("DS4_METAL_MOE_SOURCE", "moe.metal"), @@ -62,6 +64,9 @@ const SOURCES: [(&str, &str); 19] = [ ("DS4_METAL_NORM_SOURCE", "norm.metal"), ("DS4_METAL_BIN_SOURCE", "bin.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 @@ -1593,7 +1598,7 @@ impl Dspark { } } -struct Steering { +pub(super) struct Steering { directions: Buffer, attention_scale: f32, ffn_scale: f32, @@ -1615,6 +1620,7 @@ impl Steering { let expected = model .shape .layers + .saturating_sub(model.shape.nextn) .checked_mul(model.shape.embd as u32) .and_then(|values| values.checked_mul(4)) .ok_or("directional steering size overflow")? as usize; @@ -2168,7 +2174,9 @@ impl SsdPlan { let hotlist = match model.shape.model { ModelChoice::DeepSeekV4Flash0731 => hotlist::FLASH, 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 { if loaded == self.preload_experts { @@ -4435,6 +4443,7 @@ impl Executor { quality, ssd, speculative, + steering, expert_profile_path, ) .map(Box::new) @@ -4481,6 +4490,24 @@ impl Executor { } } + pub(super) fn encode_vision(&self, encoded: &[u8]) -> Result { + 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( &mut self, token: i32, @@ -7648,7 +7675,7 @@ fn compression_ratio(shape: super::Shape, layer: u32) -> u32 { } crate::model::ModelChoice::DeepSeekV4Flash0731 | crate::model::ModelChoice::DeepSeekV4Pro => 128, - crate::model::ModelChoice::Glm52 => 0, + crate::model::ModelChoice::Glm52 | crate::model::ModelChoice::Glm53Flash => 0, } } diff --git a/src/engine/metal/glm.rs b/src/engine/metal/glm.rs index fa3f8c5..598afed 100644 --- a/src/engine/metal/glm.rs +++ b/src/engine/metal/glm.rs @@ -1,11 +1,13 @@ use super::checkpoint::{read_buffer, read_u32, read_u64, write_buffer, write_u32, write_u64}; +use super::vision::{VisionEmbedding, VisionEncoder}; use super::*; -use crate::settings::EngineSsdSettings; +use crate::settings::{EngineSsdSettings, EngineSteeringSettings}; const CHECKPOINT_MAGIC: &[u8; 8] = b"DS4GLM01"; -const CHECKPOINT_VERSION: u32 = 1; +const CHECKPOINT_VERSION: u32 = 2; const CACHE_F16: bool = true; const DECODE_FLUSH_LAYERS: usize = 4; +const PREFILL_DRAIN_LAYERS: usize = 16; const AUTO_CACHE_BYTES: u64 = 12 * 1024 * 1024 * 1024; const STREAMING_TOKEN_PREFILL_MAX: u32 = 64; @@ -14,11 +16,17 @@ fn live_prefix_rewind_target(live: &[i32], incoming: &[i32]) -> Option { .then_some(incoming.len() - 1) } const STREAMING_FULL_ATTN_CONTEXT: u32 = 8192; +const RESIDENT_FULL_ATTN_CONTEXT: u32 = 4096; const LONG_CONTEXT_THRESHOLD: u32 = 65_536; const LONG_CONTEXT_FULL_ATTN_CONTEXT: u32 = 4096; -const INDEXED_PREFILL_CHUNK: u32 = 4096; +const INDEXED_PREFILL_CHUNK: u32 = 2048; const INDEXED_PREFILL_SCORE_BYTES: u64 = 256 * 1024 * 1024; +const GLM53_INDEXED_PREFILL_SCORE_BYTES: u64 = 64 * 1024 * 1024; const INDEXED_PREFILL_ATTN_SLICE: u32 = 2048; +const GLM53_KDA_HEADS: u64 = 64; +const GLM53_KDA_HEAD_DIM: u64 = 128; +const GLM53_KDA_CONV: u64 = 4; +const GLM53_INDEX_POOL: u64 = 4; #[derive(Clone, Copy)] struct SparseWeights { @@ -47,8 +55,18 @@ struct NextnWeights { shared_head_norm: Weight, } -struct GlmLayer { - attn_norm: Weight, +#[derive(Clone, Copy)] +struct HcWeights { + attn_fn: Weight, + attn_scale: Weight, + attn_base: Weight, + ffn_fn: Weight, + ffn_scale: Weight, + ffn_base: Weight, +} + +#[derive(Clone, Copy)] +struct DsaWeights { q_a: Weight, q_a_norm: Weight, q_b: Weight, @@ -62,6 +80,39 @@ struct GlmLayer { indexer_k_norm: Weight, indexer_k_bias: Weight, indexer_proj: Weight, + indexer_pool_ape: Option, + indexer_pool_gate: Option, +} + +#[derive(Clone, Copy)] +struct KdaWeights { + q: Weight, + k: Weight, + v: Weight, + q_conv: Weight, + k_conv: Weight, + v_conv: Weight, + f_a: Weight, + f_b: Weight, + dt_bias: Weight, + a_log: Weight, + beta: Weight, + g_a: Weight, + g_b: Weight, + o_norm: Weight, + output: Weight, +} + +#[derive(Clone, Copy)] +enum AttentionWeights { + Dsa(DsaWeights), + Kda(KdaWeights), +} + +struct GlmLayer { + attn_norm: Weight, + hc: Option, + attention: AttentionWeights, ffn_norm: Weight, dense: Option, sparse: Option, @@ -116,21 +167,65 @@ impl GlmWeights { }) }) .transpose()?; + let hc = (model.shape.model == ModelChoice::Glm53Flash && index < normal_layers) + .then(|| { + Ok::<_, String>(HcWeights { + attn_fn: required("hc_attn_fn.weight")?, + attn_scale: required("hc_attn_scale.weight")?, + attn_base: required("hc_attn_base.weight")?, + ffn_fn: required("hc_ffn_fn.weight")?, + ffn_scale: required("hc_ffn_scale.weight")?, + ffn_base: required("hc_ffn_base.weight")?, + }) + }) + .transpose()?; + let attention = if glm53_kda_layer(model.shape, index as usize) { + AttentionWeights::Kda(KdaWeights { + q: required("kda_q.weight")?, + k: required("kda_k.weight")?, + v: required("kda_v.weight")?, + q_conv: required("kda_q_conv.weight")?, + k_conv: required("kda_k_conv.weight")?, + v_conv: required("kda_v_conv.weight")?, + f_a: required("kda_f_a.weight")?, + f_b: required("kda_f_b.weight")?, + dt_bias: required("kda_dt_bias.weight")?, + a_log: required("kda_a_log.weight")?, + beta: required("kda_beta.weight")?, + g_a: required("kda_g_a.weight")?, + g_b: required("kda_g_b.weight")?, + o_norm: required("kda_o_norm.weight")?, + output: required("kda_output.weight")?, + }) + } else { + AttentionWeights::Dsa(DsaWeights { + q_a: required("attn_q_a.weight")?, + q_a_norm: required("attn_q_a_norm.weight")?, + q_b: required("attn_q_b.weight")?, + kv_a: required("attn_kv_a_mqa.weight")?, + kv_norm: required("attn_kv_a_norm.weight")?, + k_b: required("attn_k_b.weight")?, + v_b: required("attn_v_b.weight")?, + output: required("attn_output.weight")?, + indexer_k: required("indexer.attn_k.weight")?, + indexer_q: required("indexer.attn_q_b.weight")?, + indexer_k_norm: required("indexer.k_norm.weight")?, + indexer_k_bias: required("indexer.k_norm.bias")?, + indexer_proj: required("indexer.proj.weight")?, + indexer_pool_ape: Weight::optional( + main, + &format!("blk.{index}.indexer.pool_ape.weight"), + ), + indexer_pool_gate: Weight::optional( + main, + &format!("blk.{index}.indexer.pool_gate.weight"), + ), + }) + }; Ok(GlmLayer { attn_norm: required("attn_norm.weight")?, - q_a: required("attn_q_a.weight")?, - q_a_norm: required("attn_q_a_norm.weight")?, - q_b: required("attn_q_b.weight")?, - kv_a: required("attn_kv_a_mqa.weight")?, - kv_norm: required("attn_kv_a_norm.weight")?, - k_b: required("attn_k_b.weight")?, - v_b: required("attn_v_b.weight")?, - output: required("attn_output.weight")?, - indexer_k: required("indexer.attn_k.weight")?, - indexer_q: required("indexer.attn_q_b.weight")?, - indexer_k_norm: required("indexer.k_norm.weight")?, - indexer_k_bias: required("indexer.k_norm.bias")?, - indexer_proj: required("indexer.proj.weight")?, + hc, + attention, ffn_norm: required("ffn_norm.weight")?, dense, sparse, @@ -150,40 +245,122 @@ impl GlmWeights { } } +impl GlmLayer { + fn dsa(&self) -> Result { + match self.attention { + AttentionWeights::Dsa(weights) => Ok(weights), + AttentionWeights::Kda(_) => Err("GLM KDA layer reached the DSA graph".into()), + } + } +} + +struct KdaCache { + conv: Buffer, + recurrent: Buffer, +} + struct LayerCache { kv: Buffer, rope: Buffer, indexer: Option, + indexer_tail_k: Option, + indexer_tail_gate: Option, + kda: Option, } impl LayerCache { fn allocate(shape: super::super::Shape, layer: usize, context: u32) -> Result { - Ok(Self { - kv: Buffer::bytes(u64::from(context) * shape.kv_lora * 2)?, - rope: Buffer::bytes(u64::from(context) * shape.rot * 2)?, - indexer: full_indexer_layer(shape, layer) - .then(|| Buffer::bytes(u64::from(context) * shape.indexer_head_dim * 2)) + let kda = glm53_kda_layer(shape, layer); + let indexer = full_indexer_layer(shape, layer); + let kda = kda + .then(|| { + let projection = GLM53_KDA_HEADS * GLM53_KDA_HEAD_DIM; + let cache = KdaCache { + conv: Buffer::floats(3 * (GLM53_KDA_CONV - 1) * projection)?, + recurrent: Buffer::floats(projection * GLM53_KDA_HEAD_DIM)?, + }; + cache + .conv + .fill(0.0, 3 * (GLM53_KDA_CONV - 1) * projection)?; + cache.recurrent.fill(0.0, projection * GLM53_KDA_HEAD_DIM)?; + Ok::<_, String>(cache) + }) + .transpose()?; + let cache = Self { + kv: Buffer::bytes(if kda.is_some() { + 1 + } else { + u64::from(context) * shape.kv_lora * 2 + })?, + rope: Buffer::bytes(if kda.is_some() { + 1 + } else { + (u64::from(context) * shape.rot * 2).max(1) + })?, + indexer: indexer + .then(|| { + let rows = if shape.model == ModelChoice::Glm53Flash { + u64::from(context).div_ceil(GLM53_INDEX_POOL) + } else { + u64::from(context) + }; + Buffer::bytes(rows * shape.indexer_head_dim * 2) + }) .transpose()?, - }) + indexer_tail_k: (indexer && shape.model == ModelChoice::Glm53Flash) + .then(|| Buffer::floats(GLM53_INDEX_POOL * shape.indexer_head_dim)) + .transpose()?, + indexer_tail_gate: (indexer && shape.model == ModelChoice::Glm53Flash) + .then(|| Buffer::floats(GLM53_INDEX_POOL * shape.indexer_head_dim)) + .transpose()?, + kda, + }; + if let Some(tail) = &cache.indexer_tail_k { + tail.fill(0.0, GLM53_INDEX_POOL * shape.indexer_head_dim)?; + } + if let Some(tail) = &cache.indexer_tail_gate { + tail.fill(0.0, GLM53_INDEX_POOL * shape.indexer_head_dim)?; + } + Ok(cache) } } struct GlmScratch { + token: Buffer, current: Buffer, next: Buffer, + hc_current: Buffer, + hc_next: Buffer, + hc_after_attn: Buffer, + hc_flat: Buffer, + hc_mix: Buffer, + hc_split: Buffer, + hc_output: Buffer, + mtp_hc: Buffer, + hc_mean: Buffer, attn_norm: Buffer, q_rank: Buffer, q_rank_norm: Buffer, q: Buffer, kv_raw: Buffer, indexer_k: Buffer, + indexer_gate: Buffer, indexer_q: Buffer, indexer_weights: Buffer, indexer_scores: Buffer, indexer_selected: Buffer, + indexer_pool_selected: Buffer, qk_low: Buffer, heads: Buffer, attn_out: Buffer, + kda_q: Buffer, + kda_k: Buffer, + kda_v: Buffer, + kda_lowrank: Buffer, + kda_raw_gate: Buffer, + kda_raw_beta: Buffer, + kda_output_gate: Buffer, + kda_out: Buffer, after_attn: Buffer, ffn_norm: Buffer, ffn_gate: Buffer, @@ -204,6 +381,12 @@ struct GlmBatchScratch { tokens: Buffer, current: Buffer, next: Buffer, + hc_current: Buffer, + hc_next: Buffer, + hc_after_attn: Buffer, + hc_flat: Buffer, + hc_mix: Buffer, + hc_split: Buffer, attn_norm: Buffer, q_rank: Buffer, q_rank_norm: Buffer, @@ -211,14 +394,24 @@ struct GlmBatchScratch { kv_raw: Buffer, kv_norm: Buffer, indexer_k: Buffer, + indexer_gate: Buffer, indexer_q: Buffer, indexer_weights: Buffer, indexer_scores: Buffer, indexer_selected: Buffer, + indexer_pool_selected: Buffer, qk_low: Buffer, attn_lora: Buffer, heads: Buffer, attn_out: Buffer, + kda_q: Buffer, + kda_k: Buffer, + kda_v: Buffer, + kda_lowrank: Buffer, + kda_raw_gate: Buffer, + kda_raw_beta: Buffer, + kda_output_gate: Buffer, + kda_out: Buffer, after_attn: Buffer, ffn_norm: Buffer, ffn_gate: Buffer, @@ -239,6 +432,7 @@ struct GlmMtp { selected: Buffer, concat: Buffer, pending: Option, + parent: Option, min_pos: Option, cycles: u64, drafted: u64, @@ -248,6 +442,12 @@ struct GlmMtp { timing: bool, } +struct VisionOverlay { + token_start: u32, + tokens: u32, + embedding: Buffer, +} + #[derive(Clone, Copy)] struct GlmStreamingPlan { settings: EngineSsdSettings, @@ -261,22 +461,53 @@ impl GlmScratch { let shape = model.shape; let q = shape.heads * shape.key_mla; let heads = shape.heads * shape.value_mla; - Ok(Self { + let hc_dim = (shape.hc * shape.embd).max(1); + let hc_mix = (shape.hc * (shape.hc + 2)).max(1); + let kda_projection = GLM53_KDA_HEADS * GLM53_KDA_HEAD_DIM; + let scratch = Self { + token: Buffer::bytes(4)?, current: Buffer::floats(shape.embd)?, next: Buffer::floats(shape.embd)?, + hc_current: Buffer::floats(hc_dim)?, + hc_next: Buffer::floats(hc_dim)?, + hc_after_attn: Buffer::floats(hc_dim)?, + hc_flat: Buffer::floats(hc_dim)?, + hc_mix: Buffer::floats(hc_mix)?, + hc_split: Buffer::floats(hc_mix)?, + hc_output: Buffer::floats(shape.embd)?, + mtp_hc: Buffer::floats(2 * shape.embd)?, + hc_mean: Buffer::floats(shape.hc.max(1))?, attn_norm: Buffer::floats(shape.embd)?, q_rank: Buffer::floats(shape.lora_q)?, q_rank_norm: Buffer::floats(shape.lora_q)?, q: Buffer::floats(q)?, kv_raw: Buffer::floats(shape.head_dim)?, indexer_k: Buffer::floats(shape.indexer_head_dim)?, + indexer_gate: Buffer::floats(shape.indexer_head_dim)?, indexer_q: Buffer::floats(shape.indexer_heads * shape.indexer_head_dim)?, indexer_weights: Buffer::floats(shape.indexer_heads)?, indexer_scores: Buffer::floats(u64::from(context))?, - indexer_selected: Buffer::bytes(shape.indexer_top_k * 4)?, + indexer_selected: Buffer::bytes( + (shape.indexer_top_k + + if shape.model == ModelChoice::Glm53Flash { + 3 + } else { + 0 + }) + * 4, + )?, + indexer_pool_selected: Buffer::bytes(shape.indexer_top_k.max(4))?, qk_low: Buffer::floats(shape.heads * shape.kv_lora)?, heads: Buffer::floats(heads)?, attn_out: Buffer::floats(shape.embd)?, + kda_q: Buffer::floats(kda_projection)?, + kda_k: Buffer::floats(kda_projection)?, + kda_v: Buffer::floats(kda_projection)?, + kda_lowrank: Buffer::floats(GLM53_KDA_HEAD_DIM)?, + kda_raw_gate: Buffer::floats(kda_projection)?, + kda_raw_beta: Buffer::floats(GLM53_KDA_HEADS)?, + kda_output_gate: Buffer::floats(kda_projection)?, + kda_out: Buffer::floats(kda_projection)?, after_attn: Buffer::floats(shape.embd)?, ffn_norm: Buffer::floats(shape.embd)?, ffn_gate: Buffer::floats(shape.ff_dense.max(shape.experts_used * shape.ff_expert))?, @@ -291,7 +522,11 @@ impl GlmScratch { router_weights: Buffer::floats(shape.experts_used)?, output_norm: Buffer::floats(shape.embd)?, logits: Buffer::floats(shape.vocab)?, - }) + }; + if shape.hc != 0 { + scratch.hc_mean.fill(1.0 / shape.hc as f32, shape.hc)?; + } + Ok(scratch) } } @@ -304,13 +539,30 @@ impl GlmBatchScratch { let heads = shape.heads * shape.value_mla; let hidden = shape.ff_dense.max(shape.ff_expert); let routed_mid = shape.experts_used * shape.ff_expert; - let score_rows = (INDEXED_PREFILL_SCORE_BYTES / (u64::from(context) * 4)) - .max(1) - .min(rows); + let score_columns = if shape.model == ModelChoice::Glm53Flash { + u64::from(context).div_ceil(GLM53_INDEX_POOL) + } else { + u64::from(context) + }; + let score_bytes = if shape.model == ModelChoice::Glm53Flash { + GLM53_INDEXED_PREFILL_SCORE_BYTES + } else { + INDEXED_PREFILL_SCORE_BYTES + }; + let score_rows = (score_bytes / (score_columns * 4)).max(1).min(rows); + let hc_dim = (shape.hc * shape.embd).max(1); + let hc_mix = (shape.hc * (shape.hc + 2)).max(1); + let kda_projection = GLM53_KDA_HEADS * GLM53_KDA_HEAD_DIM; Ok(Self { tokens: Buffer::bytes(rows * 4)?, current: Buffer::floats(rows * shape.embd)?, next: Buffer::floats(rows * shape.embd)?, + hc_current: Buffer::floats(rows * hc_dim)?, + hc_next: Buffer::floats(rows * hc_dim)?, + hc_after_attn: Buffer::floats(rows * hc_dim)?, + hc_flat: Buffer::floats(rows * hc_dim)?, + hc_mix: Buffer::floats(rows * hc_mix)?, + hc_split: Buffer::floats(rows * hc_mix)?, attn_norm: Buffer::floats(rows * shape.embd)?, q_rank: Buffer::floats(rows * shape.lora_q)?, q_rank_norm: Buffer::floats(rows * shape.lora_q)?, @@ -318,14 +570,32 @@ impl GlmBatchScratch { kv_raw: Buffer::floats(rows * shape.head_dim)?, kv_norm: Buffer::floats(rows * shape.kv_lora)?, indexer_k: Buffer::floats(rows * shape.indexer_head_dim)?, + indexer_gate: Buffer::floats(rows * shape.indexer_head_dim)?, indexer_q: Buffer::floats(rows * shape.indexer_heads * shape.indexer_head_dim)?, indexer_weights: Buffer::floats(rows * shape.indexer_heads)?, - indexer_scores: Buffer::floats(score_rows * u64::from(context))?, - indexer_selected: Buffer::bytes(rows * shape.indexer_top_k * 4)?, + indexer_scores: Buffer::floats(score_rows * score_columns)?, + indexer_selected: Buffer::bytes( + rows * (shape.indexer_top_k + + if shape.model == ModelChoice::Glm53Flash { + 3 + } else { + 0 + }) + * 4, + )?, + indexer_pool_selected: Buffer::bytes(rows * shape.indexer_top_k.max(4))?, qk_low: Buffer::floats(rows * qk_low)?, attn_lora: Buffer::floats(rows * qk_low)?, heads: Buffer::floats(rows * heads)?, attn_out: Buffer::floats(rows * shape.embd)?, + kda_q: Buffer::floats(rows * kda_projection)?, + kda_k: Buffer::floats(rows * kda_projection)?, + kda_v: Buffer::floats(rows * kda_projection)?, + kda_lowrank: Buffer::floats(rows * GLM53_KDA_HEAD_DIM)?, + kda_raw_gate: Buffer::floats(rows * kda_projection)?, + kda_raw_beta: Buffer::floats(rows * GLM53_KDA_HEADS)?, + kda_output_gate: Buffer::floats(rows * kda_projection)?, + kda_out: Buffer::floats(rows * kda_projection)?, after_attn: Buffer::floats(rows * shape.embd)?, ffn_norm: Buffer::floats(rows * shape.embd)?, ffn_gate: Buffer::floats(rows * hidden)?, @@ -342,19 +612,26 @@ impl GlmBatchScratch { }) } - fn score_rows(&self, context: u32, rows: u32) -> u32 { - u32::try_from( - (INDEXED_PREFILL_SCORE_BYTES / (u64::from(context) * 4)) - .max(1) - .min(u64::from(rows)), - ) - .expect("GLM score row cap is bounded by the batch") + fn score_rows(&self, context: u32, rows: u32, pooled: bool) -> u32 { + let columns = if pooled { + u64::from(context).div_ceil(GLM53_INDEX_POOL) + } else { + u64::from(context) + }; + let bytes = if pooled { + GLM53_INDEXED_PREFILL_SCORE_BYTES + } else { + INDEXED_PREFILL_SCORE_BYTES + }; + u32::try_from((bytes / (columns * 4)).max(1).min(u64::from(rows))) + .expect("GLM score row cap is bounded by the batch") } } pub(in crate::engine) struct GlmExecutor { weights: GlmWeights, scratch: GlmScratch, + batch: Option, caches: Vec, logits: Vec, tokens: Vec, @@ -362,6 +639,7 @@ pub(in crate::engine) struct GlmExecutor { quality: bool, ssd: EngineSsdSettings, profile: Option, + steering: Option, mtp: Option, ssd_resident_bytes: u64, ssd_cache_bytes: u64, @@ -371,6 +649,8 @@ pub(in crate::engine) struct GlmExecutor { model_modified: (u64, u32), model_identity: [u8; 32], streaming_spans: Option>, + vision: Option, + vision_overlays: Vec, _context: Context, model: Model, } @@ -406,6 +686,11 @@ impl GlmExecutor { dspark_strict: false, dspark_exact_sampling: false, }, + EngineSteeringSettings { + file: None, + ffn_scale: 0.0, + attention_scale: 0.0, + }, None, ) } @@ -416,6 +701,7 @@ impl GlmExecutor { quality: bool, ssd: EngineSsdSettings, speculative: EngineSpeculativeSettings, + steering: EngineSteeringSettings, expert_profile_path: Option<&str>, ) -> Result { if context == 0 || u64::from(context) > model.shape.original_context { @@ -445,6 +731,9 @@ impl GlmExecutor { let model_spans = model_spans.map(|(spans, _)| spans); configure_streaming(&model, &weights, streaming.as_ref())?; let scratch = GlmScratch::allocate(&model, context)?; + let batch = (model.shape.model == ModelChoice::Glm53Flash) + .then(|| GlmBatchScratch::allocate(&model, context.min(INDEXED_PREFILL_CHUNK), context)) + .transpose()?; let caches = (0..weights.layers.len()) .map(|layer| LayerCache::allocate(model.shape, layer, context)) .collect::>()?; @@ -454,7 +743,21 @@ impl GlmExecutor { .and_then(|modified| modified.duration_since(UNIX_EPOCH).ok()) .map(|duration| (duration.as_secs(), duration.subsec_nanos())) .unwrap_or_default(); - let model_identity = model.checkpoint_identity(); + if model.shape.model != ModelChoice::Glm53Flash + && (steering.attention_scale != 0.0 || steering.ffn_scale != 0.0) + { + return Err("directional steering is supported only for GLM 5.3 Flash".into()); + } + let steering = Steering::load(&model, steering)?; + let model_identity = if let Some(steering) = &steering { + let mut hash = Sha256::new(); + hash.update(model.checkpoint_identity()); + hash.update(steering.identity); + hash.finalize().into() + } else { + model.checkpoint_identity() + }; + let vision = model.vision.as_ref().map(VisionEncoder::bind).transpose()?; let profile = ExpertProfile::new( expert_profile_path, model.shape.model, @@ -474,6 +777,7 @@ impl GlmExecutor { selected: Buffer::bytes(u64::from(context) * 4)?, concat: Buffer::floats(2 * model.shape.embd)?, pending: None, + parent: None, min_pos: None, cycles: 0, drafted: 0, @@ -503,6 +807,7 @@ impl GlmExecutor { Ok(Self { weights, scratch, + batch, caches, logits: vec![0.0; model.shape.vocab as usize], tokens: Vec::new(), @@ -510,6 +815,7 @@ impl GlmExecutor { quality, ssd: effective_ssd, profile, + steering, mtp, ssd_resident_bytes, ssd_cache_bytes, @@ -519,16 +825,55 @@ impl GlmExecutor { model_modified, model_identity, streaming_spans: model_spans, + vision, + vision_overlays: Vec::new(), _context: context_handle, model, }) } + pub(super) fn encode_vision(&self, encoded: &[u8]) -> Result { + let encoder = self + .vision + .as_ref() + .ok_or("GLM 5.3 vision encoder is not loaded")?; + let vision = self + .model + .vision + .as_ref() + .ok_or("GLM 5.3 vision GGUF is not loaded")?; + encoder.encode(vision, encoded) + } + + pub(super) fn set_vision_overlays( + &mut self, + overlays: Vec<(u32, VisionEmbedding)>, + ) -> Result<(), String> { + self.vision_overlays = overlays + .into_iter() + .map(|(token_start, embedding)| { + let buffer = Buffer::floats(u64::from(embedding.tokens) * self.model.shape.embd)?; + buffer.write_f32(&embedding.values)?; + debug_assert!(embedding.content_width <= embedding.width); + debug_assert!(embedding.content_height <= embedding.height); + Ok(VisionOverlay { + token_start, + tokens: embedding.tokens, + embedding: buffer, + }) + }) + .collect::>()?; + Ok(()) + } + pub(super) fn eval(&mut self, token: i32) -> Result<(), String> { let shape = self.model.shape; if token < 0 || token as u64 >= shape.vocab { return Err(format!("token {token} is outside the vocabulary")); } + if shape.model == ModelChoice::Glm53Flash { + return self.eval_glm53(token); + } let pos = u32::try_from(self.tokens.len()).map_err(|_| "GLM position overflow")?; if pos >= self.context { return Err(format!( @@ -634,6 +979,7 @@ impl GlmExecutor { selected_ready: &mut bool, selected_count: &mut u32, ) -> Result<(), String> { + let dsa = layer.dsa()?; let shape = self.model.shape; let map = self.model.main.map_ptr().cast(); let size = self.model.main.len(); @@ -651,7 +997,7 @@ impl GlmExecutor { )?; project_rows( &batch.q_rank, - layer.q_a, + dsa.q_a, shape.embd, shape.lora_q, &batch.attn_norm, @@ -662,7 +1008,7 @@ impl GlmExecutor { norm_rows( &batch.q_rank_norm, &batch.q_rank, - layer.q_a_norm, + dsa.q_a_norm, shape.lora_q as u32, rows, shape.rms_epsilon, @@ -671,7 +1017,7 @@ impl GlmExecutor { )?; project_rows( &batch.q, - layer.q_b, + dsa.q_b, shape.lora_q, q_dim, &batch.q_rank_norm, @@ -707,7 +1053,7 @@ impl GlmExecutor { .ok_or("GLM full-indexer layer is missing its key cache")?; project_rows( &batch.indexer_k, - layer.indexer_k, + dsa.indexer_k, shape.embd, shape.indexer_head_dim, &batch.current, @@ -722,8 +1068,8 @@ impl GlmExecutor { batch.indexer_k.raw(), map, size, - layer.indexer_k_norm.offset, - layer.indexer_k_bias.offset, + dsa.indexer_k_norm.offset, + dsa.indexer_k_bias.offset, pos, rows, self.context, @@ -746,7 +1092,7 @@ impl GlmExecutor { project_rows( &batch.kv_raw, - layer.kv_a, + dsa.kv_a, shape.embd, shape.head_dim, &batch.attn_norm, @@ -761,7 +1107,7 @@ impl GlmExecutor { batch.kv_raw.raw(), map, size, - layer.kv_norm.offset, + dsa.kv_norm.offset, rows, shape.head_dim as u32, shape.kv_lora as u32, @@ -821,8 +1167,8 @@ impl GlmExecutor { batch.q.raw(), map, size, - layer.k_b.offset, - layer.k_b.kind, + dsa.k_b.offset, + dsa.k_b.kind, rows, shape.heads as u32, shape.kv_lora as u32, @@ -843,7 +1189,7 @@ impl GlmExecutor { )?; project_rows( &batch.attn_out, - layer.output, + dsa.output, shape.heads * shape.value_mla, shape.embd, &batch.heads, @@ -877,12 +1223,13 @@ impl GlmExecutor { rows: u32, visible: u32, ) -> Result<(), String> { + let dsa = layer.dsa()?; let shape = self.model.shape; let map = self.model.main.map_ptr().cast(); let size = self.model.main.len(); project_rows( &batch.indexer_q, - layer.indexer_q, + dsa.indexer_q, shape.lora_q, shape.indexer_heads * shape.indexer_head_dim, &batch.q_rank_norm, @@ -912,7 +1259,7 @@ impl GlmExecutor { )?; f32_project_rows( &batch.indexer_weights, - layer.indexer_proj, + dsa.indexer_proj, shape.embd, shape.indexer_heads, &batch.current, @@ -924,7 +1271,7 @@ impl GlmExecutor { .indexer .as_ref() .ok_or("GLM full-indexer layer is missing its key cache")?; - let score_rows = batch.score_rows(self.context, rows); + let score_rows = batch.score_rows(self.context, rows, false); let q_row_bytes = shape.indexer_heads * shape.indexer_head_dim * 4; let weight_row_bytes = shape.indexer_heads * 4; let selected_row_bytes = shape.indexer_top_k * 4; @@ -989,6 +1336,21 @@ impl GlmExecutor { rows: u32, selected_count: u32, causal_range: bool, + ) -> Result<(), String> { + let dsa = layer.dsa()?; + self.encode_batch_attention_raw(batch, dsa, cache, pos, rows, selected_count, causal_range) + } + + #[allow(clippy::too_many_arguments)] + fn encode_batch_attention_raw( + &self, + batch: &GlmBatchScratch, + dsa: DsaWeights, + cache: &LayerCache, + pos: u32, + rows: u32, + selected_count: u32, + causal_range: bool, ) -> Result<(), String> { let shape = self.model.shape; let map = self.model.main.map_ptr().cast(); @@ -997,6 +1359,10 @@ impl GlmExecutor { let qk_low_dim = shape.heads * shape.kv_lora; let heads_dim = shape.heads * shape.value_mla; let selected_row_bytes = u64::from(selected_count) * 4; + call( + unsafe { ds4_gpu_flush_encoder() }, + "starting batched GLM attention", + )?; let mut start = 0; while start < rows { let slice = (rows - start).min(INDEXED_PREFILL_ATTN_SLICE); @@ -1015,7 +1381,26 @@ impl GlmExecutor { u64::from(start) * heads_dim * 4, u64::from(slice) * heads_dim * 4, )?; - if causal_range { + if causal_range && glm53_dense_compact_prefill(slice) { + call( + unsafe { + ds4_gpu_glm_attention_dense_compact_lora_causal_tensor( + lora.raw(), + qk_low.raw(), + cache.kv.raw(), + pos + start, + slice, + selected_count, + self.context, + CACHE_F16, + shape.heads as u32, + shape.kv_lora as u32, + (shape.key_mla - shape.rot) as u32, + ) + }, + "running dense compact GLM attention", + )?; + } else if causal_range { call( unsafe { ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor( @@ -1085,8 +1470,8 @@ impl GlmExecutor { lora.raw(), map, size, - layer.v_b.offset, - layer.v_b.kind, + dsa.v_b.offset, + dsa.v_b.kind, slice, shape.heads as u32, shape.kv_lora as u32, @@ -1097,7 +1482,10 @@ impl GlmExecutor { )?; start += slice; } - Ok(()) + call( + unsafe { ds4_gpu_flush_encoder() }, + "finishing batched GLM attention", + ) } fn encode_batch_ffn( @@ -1351,17 +1739,63 @@ impl GlmExecutor { if let Some(mtp) = &mut self.mtp { mtp.cycles += 1; } - let pending = self.mtp.as_ref().and_then(|mtp| mtp.pending); let started = Instant::now(); + if self.model.shape.model == ModelChoice::Glm53Flash { + let pos = self.position(); + let pending = self.mtp.as_ref().and_then(|mtp| mtp.pending); + let parent = self.mtp.as_ref().and_then(|mtp| mtp.parent); + if let Some(draft) = pending + && parent == Some(token) + && max_tokens > 1 + && pos + 1 < self.context + && !cancelled.load(std::sync::atomic::Ordering::Relaxed) + { + let first = self.scratch.mtp_hc.view(0, self.model.shape.embd * 4)?; + let mut selected = 0; + self.eval_glm53_capture(token, Some(&first), Some(&mut selected))?; + if let Some(mtp) = &mut self.mtp { + mtp.verifier_passes += 1; + } + if selected == draft { + let second = self + .scratch + .mtp_hc + .view(self.model.shape.embd * 4, self.model.shape.embd * 4)?; + self.eval_glm53_capture(draft, Some(&second), None)?; + let _ = self.mtp_step(draft, pos, Some(&first), false)?; + let next = argmax(&self.logits); + let pending = self.mtp_step(next, pos + 1, Some(&second), true)?; + if let Some(mtp) = &mut self.mtp { + mtp.pending = Some(pending); + mtp.parent = Some(next); + mtp.accepted += 1; + } + self.note_mtp_time(started); + return Ok(vec![token, draft]); + } + self.scratch.logits.read_f32(&mut self.logits)?; + let next = argmax(&self.logits); + let pending = self.mtp_step(next, pos, Some(&first), true)?; + if let Some(mtp) = &mut self.mtp { + mtp.pending = Some(pending); + mtp.parent = Some(next); + } + self.note_mtp_time(started); + return Ok(vec![token]); + } + } + + let pending = self.mtp.as_ref().and_then(|mtp| mtp.pending); self.eval(token)?; if let Some(mtp) = &mut self.mtp { mtp.verifier_passes += 1; } let pos = self.position() - 1; let next = argmax(&self.logits); - let mut draft = self.mtp_step(next, pos)?; + let mut draft = self.mtp_step(next, pos, None, true)?; let mut accepted = vec![token]; - if pending == Some(next) + if self.model.shape.model != ModelChoice::Glm53Flash + && pending == Some(next) && max_tokens > 1 && self.position() < self.context && !cancelled.load(std::sync::atomic::Ordering::Relaxed) @@ -1372,18 +1806,31 @@ impl GlmExecutor { mtp.accepted += 1; mtp.verifier_passes += 1; } - draft = self.mtp_step(argmax(&self.logits), pos + 1)?; + draft = self.mtp_step(argmax(&self.logits), pos + 1, None, true)?; } if let Some(mtp) = &mut self.mtp { mtp.pending = Some(draft); + mtp.parent = Some(argmax(&self.logits)); + } + self.note_mtp_time(started); + Ok(accepted) + } + + fn note_mtp_time(&mut self, started: Instant) { + if let Some(mtp) = &mut self.mtp { mtp.verifier_ns = mtp .verifier_ns .saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX)); } - Ok(accepted) } - fn mtp_step(&mut self, next_token: i32, pos: u32) -> Result { + fn mtp_step( + &mut self, + next_token: i32, + pos: u32, + target_hidden: Option<&Buffer>, + draft_requested: bool, + ) -> Result { let mut mtp = self.mtp.take().ok_or("GLM MTP is not configured")?; let result = (|| { let layer = self @@ -1400,10 +1847,12 @@ impl GlmExecutor { if selected_count > self.context { return Err("GLM MTP attention range exceeds the context".into()); } - let selected = (min_pos..=pos) - .map(|position| position as i32) - .collect::>(); - mtp.selected.write_i32(&selected)?; + if draft_requested { + let selected = (min_pos..=pos) + .map(|position| position as i32) + .collect::>(); + mtp.selected.write_i32(&selected)?; + } let shape = self.model.shape; let map = self.model.main.map_ptr().cast(); let size = self.model.main.len(); @@ -1411,21 +1860,40 @@ impl GlmExecutor { let hnorm = mtp.concat.view(shape.embd * 4, shape.embd * 4)?; let started = Instant::now(); let commands = Commands::begin()?; - call( - unsafe { - ds4_gpu_embed_token_quant_tensor( - self.scratch.next.raw(), - map, - size, - self.weights.embedding.offset, - self.weights.embedding.kind, - shape.vocab as u32, - next_token as u32, - shape.embd as u32, - ) - }, - "GLM MTP token embedding", - )?; + if shape.model == ModelChoice::Glm53Flash && self.weights.embedding.kind == BF16 { + self.scratch.token.write_i32(&[next_token])?; + call( + unsafe { + ds4_gpu_glm53_embedding_bf16( + self.scratch.next.raw(), + map, + size, + self.weights.embedding.offset, + self.scratch.token.raw(), + 1, + shape.embd as u32, + shape.vocab as u32, + ) + }, + "GLM 5.3 MTP token embedding", + )?; + } else { + call( + unsafe { + ds4_gpu_embed_token_quant_tensor( + self.scratch.next.raw(), + map, + size, + self.weights.embedding.offset, + self.weights.embedding.kind, + shape.vocab as u32, + next_token as u32, + shape.embd as u32, + ) + }, + "GLM MTP token embedding", + )?; + } norm( &enorm, &self.scratch.next, @@ -1437,33 +1905,62 @@ impl GlmExecutor { )?; norm( &hnorm, - &self.scratch.current, + if let Some(target_hidden) = target_hidden { + target_hidden + } else if shape.model == ModelChoice::Glm53Flash { + &self.scratch.hc_output + } else { + &self.scratch.current + }, nextn.hnorm, shape.embd as u32, shape.rms_epsilon, map, size, )?; - project( - &self.scratch.current, - nextn.eh_proj, - 2 * shape.embd, - shape.embd, - &mtp.concat, - map, - size, - self.ssd.enabled, - )?; - let mut count = selected_count; - self.encode_layer( - layer, - &mtp.cache, - (shape.layers - 1) as usize, - shape.layers - 1, - pos, - &mut count, - Some((&mtp.selected, selected_count)), - )?; + if shape.model == ModelChoice::Glm53Flash { + glm53_project( + &self.scratch.current, + nextn.eh_proj, + 2 * shape.embd, + shape.embd, + &mtp.concat, + map, + size, + )?; + if draft_requested { + self.encode_glm53_mtp_layer( + layer, + &mtp.cache, + pos, + &mtp.selected, + selected_count, + )?; + } else { + self.store_glm53_mtp_cache(layer, &mtp.cache, pos, false)?; + } + } else { + project( + &self.scratch.current, + nextn.eh_proj, + 2 * shape.embd, + shape.embd, + &mtp.concat, + map, + size, + self.ssd.enabled, + )?; + let mut count = selected_count; + self.encode_layer( + layer, + &mtp.cache, + (shape.layers - 1) as usize, + shape.layers - 1, + pos, + &mut count, + Some((&mtp.selected, selected_count)), + )?; + } if let Some(profile) = &mut self.profile { profile.record( (shape.layers - 1) as usize, @@ -1474,32 +1971,59 @@ impl GlmExecutor { false, )?; } - norm( - &self.scratch.output_norm, - &self.scratch.next, - nextn.shared_head_norm, - shape.embd as u32, - shape.rms_epsilon, - map, - size, - )?; - project( - &self.scratch.logits, - self.weights.output, - shape.embd, - shape.vocab, - &self.scratch.output_norm, - map, - size, - self.ssd.enabled, - )?; + if draft_requested { + norm( + &self.scratch.output_norm, + &self.scratch.next, + nextn.shared_head_norm, + shape.embd as u32, + shape.rms_epsilon, + map, + size, + )?; + if shape.model == ModelChoice::Glm53Flash { + glm53_project( + &self.scratch.logits, + self.weights.output, + shape.embd, + shape.vocab, + &self.scratch.output_norm, + map, + size, + )?; + } else { + project( + &self.scratch.logits, + self.weights.output, + shape.embd, + shape.vocab, + &self.scratch.output_norm, + map, + size, + self.ssd.enabled, + )?; + } + call( + unsafe { + ds4_gpu_argmax_tensor( + self.scratch.token.raw(), + self.scratch.logits.raw(), + shape.vocab as u32, + ) + }, + "selecting a GLM MTP draft", + )?; + } commands.finish()?; - let mut logits = vec![0.0; shape.vocab as usize]; - self.scratch.logits.read_f32(&mut logits)?; + if !draft_requested { + return Ok(-1); + } + let mut draft = [0]; + self.scratch.token.read_i32(&mut draft)?; if let Some(profile) = &self.profile { profile.write()?; } - let draft = argmax(&logits); + let draft = draft[0]; mtp.drafted += 1; if mtp.timing { eprintln!( @@ -1513,6 +2037,1739 @@ impl GlmExecutor { result } + fn encode_glm53_mtp_layer( + &self, + layer: &GlmLayer, + cache: &LayerCache, + pos: u32, + selected: &Buffer, + selected_count: u32, + ) -> Result<(), String> { + let dsa = self.store_glm53_mtp_cache(layer, cache, pos, true)?; + let shape = self.model.shape; + let map = self.model.main.map_ptr().cast(); + let size = self.model.main.len(); + glm53_decode_project( + &self.scratch.q, + dsa.q_b, + shape.lora_q, + shape.heads * shape.key_mla, + &self.scratch.q_rank_norm, + map, + size, + self.ssd.enabled, + )?; + call( + unsafe { + ds4_gpu_glm_qk_lowrank_typed_tensor( + self.scratch.qk_low.raw(), + self.scratch.q.raw(), + map, + size, + dsa.k_b.offset, + dsa.k_b.kind, + shape.heads as u32, + shape.kv_lora as u32, + shape.key_mla as u32, + shape.key_mla as u32, + ) + }, + "projecting the GLM 5.3 MTP low-rank query", + )?; + call( + unsafe { + ds4_gpu_glm_attention_indexed_decode_typed_tensor( + self.scratch.heads.raw(), + self.scratch.q.raw(), + self.scratch.qk_low.raw(), + cache.kv.raw(), + cache.rope.raw(), + map, + size, + dsa.v_b.offset, + dsa.v_b.kind, + selected.raw(), + selected_count, + self.context, + CACHE_F16, + shape.heads as u32, + shape.kv_lora as u32, + shape.key_mla as u32, + shape.rot as u32, + shape.value_mla as u32, + 0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + ) + }, + "running GLM 5.3 MTP attention", + )?; + glm53_decode_project( + &self.scratch.attn_out, + dsa.output, + shape.heads * shape.value_mla, + shape.embd, + &self.scratch.heads, + map, + size, + self.ssd.enabled, + )?; + call( + unsafe { + ds4_gpu_add_tensor( + self.scratch.after_attn.raw(), + self.scratch.current.raw(), + self.scratch.attn_out.raw(), + shape.embd as u32, + ) + }, + "adding the GLM 5.3 MTP attention residual", + )?; + norm( + &self.scratch.ffn_norm, + &self.scratch.after_attn, + layer.ffn_norm, + shape.embd as u32, + shape.rms_epsilon, + map, + size, + )?; + self.encode_glm53_ffn(layer, shape.layers - 1)?; + call( + unsafe { + ds4_gpu_add_tensor( + self.scratch.next.raw(), + self.scratch.after_attn.raw(), + self.scratch.next.raw(), + shape.embd as u32, + ) + }, + "adding the GLM 5.3 MTP FFN residual", + ) + } + + fn store_glm53_mtp_cache( + &self, + layer: &GlmLayer, + cache: &LayerCache, + pos: u32, + query_needed: bool, + ) -> Result { + let dsa = layer.dsa()?; + let shape = self.model.shape; + let map = self.model.main.map_ptr().cast(); + let size = self.model.main.len(); + norm( + &self.scratch.attn_norm, + &self.scratch.current, + layer.attn_norm, + shape.embd as u32, + shape.rms_epsilon, + map, + size, + )?; + if query_needed { + glm53_decode_project( + &self.scratch.q_rank, + dsa.q_a, + shape.embd, + shape.lora_q, + &self.scratch.attn_norm, + map, + size, + self.ssd.enabled, + )?; + } + glm53_decode_project( + &self.scratch.kv_raw, + dsa.kv_a, + shape.embd, + shape.head_dim, + &self.scratch.attn_norm, + map, + size, + self.ssd.enabled, + )?; + if query_needed { + call( + unsafe { + ds4_gpu_glm_qkv_norm_store_compact_kv_tensor( + self.scratch.q_rank_norm.raw(), + self.scratch.q_rank.raw(), + map, + size, + dsa.q_a_norm.offset, + shape.lora_q as u32, + cache.kv.raw(), + cache.rope.raw(), + self.scratch.kv_raw.raw(), + dsa.kv_norm.offset, + pos, + 1, + self.context, + shape.head_dim as u32, + shape.kv_lora as u32, + shape.rot as u32, + CACHE_F16, + shape.rms_epsilon, + ) + }, + "storing GLM 5.3 MTP compact KV", + )?; + } else { + call( + unsafe { + ds4_gpu_glm_kv_lora_rms_norm_tensor( + self.scratch.q_rank_norm.raw(), + self.scratch.kv_raw.raw(), + map, + size, + dsa.kv_norm.offset, + 1, + shape.head_dim as u32, + shape.kv_lora as u32, + shape.rms_epsilon, + ) + }, + "normalizing cache-only GLM 5.3 MTP KV", + )?; + call( + unsafe { + ds4_gpu_glm_store_compact_kv_tensor( + cache.kv.raw(), + cache.rope.raw(), + self.scratch.q_rank_norm.raw(), + self.scratch.kv_raw.raw(), + pos, + 1, + self.context, + shape.head_dim as u32, + shape.kv_lora as u32, + shape.rot as u32, + CACHE_F16, + ) + }, + "storing cache-only GLM 5.3 MTP KV", + )?; + } + Ok(dsa) + } + + fn encode_glm53_kda(&self, weights: KdaWeights, cache: &KdaCache) -> Result<(), String> { + let map = self.model.main.map_ptr().cast(); + let size = self.model.main.len(); + let projection = GLM53_KDA_HEADS * GLM53_KDA_HEAD_DIM; + let fused = weights.q.kind == BF16 + && weights.k.kind == BF16 + && weights.v.kind == BF16 + && unsafe { + ds4_gpu_glm53_matmul_bf16_qkv( + self.scratch.kda_q.raw(), + self.scratch.kda_k.raw(), + self.scratch.kda_v.raw(), + map, + size, + weights.q.offset, + weights.k.offset, + weights.v.offset, + self.model.shape.embd as u32, + projection as u32, + self.scratch.attn_norm.raw(), + ) + } != 0; + if !fused { + for (out, weight) in [ + (&self.scratch.kda_q, weights.q), + (&self.scratch.kda_k, weights.k), + (&self.scratch.kda_v, weights.v), + ] { + glm53_project( + out, + weight, + self.model.shape.embd, + projection, + &self.scratch.attn_norm, + map, + size, + )?; + } + } + glm53_project( + &self.scratch.kda_lowrank, + weights.f_a, + self.model.shape.embd, + GLM53_KDA_HEAD_DIM, + &self.scratch.attn_norm, + map, + size, + )?; + glm53_project( + &self.scratch.kda_raw_gate, + weights.f_b, + GLM53_KDA_HEAD_DIM, + projection, + &self.scratch.kda_lowrank, + map, + size, + )?; + glm53_project( + &self.scratch.kda_raw_beta, + weights.beta, + self.model.shape.embd, + GLM53_KDA_HEADS, + &self.scratch.attn_norm, + map, + size, + )?; + glm53_project( + &self.scratch.kda_lowrank, + weights.g_a, + self.model.shape.embd, + GLM53_KDA_HEAD_DIM, + &self.scratch.attn_norm, + map, + size, + )?; + glm53_project( + &self.scratch.kda_output_gate, + weights.g_b, + GLM53_KDA_HEAD_DIM, + projection, + &self.scratch.kda_lowrank, + map, + size, + )?; + call( + unsafe { + ds4_gpu_glm53_kda_decode( + self.scratch.kda_out.raw(), + cache.conv.raw(), + cache.recurrent.raw(), + self.scratch.kda_q.raw(), + self.scratch.kda_k.raw(), + self.scratch.kda_v.raw(), + self.scratch.kda_raw_gate.raw(), + self.scratch.kda_raw_beta.raw(), + self.scratch.kda_output_gate.raw(), + map, + size, + weights.q_conv.offset, + weights.k_conv.offset, + weights.v_conv.offset, + weights.a_log.offset, + weights.dt_bias.offset, + weights.o_norm.offset, + GLM53_KDA_HEADS as u32, + 1, + -5.0, + self.model.shape.rms_epsilon, + ) + }, + "running GLM 5.3 KDA", + )?; + glm53_project( + &self.scratch.attn_out, + weights.output, + projection, + self.model.shape.embd, + &self.scratch.kda_out, + map, + size, + ) + } + + fn encode_glm53_dsa( + &self, + weights: DsaWeights, + cache: &LayerCache, + pos: u32, + ) -> Result<(), String> { + let shape = self.model.shape; + let map = self.model.main.map_ptr().cast(); + let size = self.model.main.len(); + let q_dim = shape.heads * shape.key_mla; + glm53_decode_project( + &self.scratch.q_rank, + weights.q_a, + shape.embd, + shape.lora_q, + &self.scratch.attn_norm, + map, + size, + self.ssd.enabled, + )?; + glm53_decode_project( + &self.scratch.kv_raw, + weights.kv_a, + shape.embd, + shape.head_dim, + &self.scratch.attn_norm, + map, + size, + self.ssd.enabled, + )?; + call( + unsafe { + ds4_gpu_glm_qkv_norm_store_compact_kv_tensor( + self.scratch.q_rank_norm.raw(), + self.scratch.q_rank.raw(), + map, + size, + weights.q_a_norm.offset, + shape.lora_q as u32, + cache.kv.raw(), + cache.rope.raw(), + self.scratch.kv_raw.raw(), + weights.kv_norm.offset, + pos, + 1, + self.context, + shape.head_dim as u32, + shape.kv_lora as u32, + shape.rot as u32, + CACHE_F16, + shape.rms_epsilon, + ) + }, + "storing GLM 5.3 compact KV", + )?; + glm53_decode_project( + &self.scratch.q, + weights.q_b, + shape.lora_q, + q_dim, + &self.scratch.q_rank_norm, + map, + size, + self.ssd.enabled, + )?; + + let indexer_cache = cache + .indexer + .as_ref() + .ok_or("GLM 5.3 DSA layer is missing its indexer cache")?; + let tail_k = cache + .indexer_tail_k + .as_ref() + .ok_or("GLM 5.3 DSA layer is missing its indexer key tail")?; + let tail_gate = cache + .indexer_tail_gate + .as_ref() + .ok_or("GLM 5.3 DSA layer is missing its indexer gate tail")?; + let ape = weights + .indexer_pool_ape + .ok_or("GLM 5.3 DSA layer is missing pool APE weights")?; + let gate = weights + .indexer_pool_gate + .ok_or("GLM 5.3 DSA layer is missing pool gate weights")?; + glm53_project( + &self.scratch.indexer_k, + weights.indexer_k, + shape.embd, + shape.indexer_head_dim, + &self.scratch.attn_norm, + map, + size, + )?; + glm53_project( + &self.scratch.indexer_gate, + gate, + shape.embd, + shape.indexer_head_dim, + &self.scratch.attn_norm, + map, + size, + )?; + call( + unsafe { + ds4_gpu_glm53_indexer_pool_update_tensor( + indexer_cache.raw(), + tail_k.raw(), + tail_gate.raw(), + self.scratch.indexer_k.raw(), + self.scratch.indexer_gate.raw(), + map, + size, + weights.indexer_k_norm.offset, + weights.indexer_k_bias.offset, + ape.offset, + pos, + 1, + self.context, + shape.indexer_head_dim as u32, + GLM53_INDEX_POOL as u32, + 1.0e-6, + CACHE_F16, + ) + }, + "updating the GLM 5.3 pooled indexer", + )?; + let visible = pos + 1; + let dense_limit = glm53_dense_limit(self.context, self.ssd.enabled); + let selected_count = if visible <= dense_limit { + call( + unsafe { + ds4_gpu_glm_fill_selected_range_tensor( + self.scratch.indexer_selected.raw(), + visible, + ) + }, + "selecting the dense GLM 5.3 context", + )?; + visible + } else { + glm53_project( + &self.scratch.indexer_q, + weights.indexer_q, + shape.lora_q, + shape.indexer_heads * shape.indexer_head_dim, + &self.scratch.q_rank_norm, + map, + size, + )?; + glm53_project( + &self.scratch.indexer_weights, + weights.indexer_proj, + shape.embd, + shape.indexer_heads, + &self.scratch.attn_norm, + map, + size, + )?; + let pool_rows = visible / GLM53_INDEX_POOL as u32; + let selected_pools = shape.indexer_top_k as u32 / GLM53_INDEX_POOL as u32; + call( + unsafe { + ds4_gpu_glm_indexer_score_one_tensor( + self.scratch.indexer_scores.raw(), + self.scratch.indexer_q.raw(), + self.scratch.indexer_weights.raw(), + indexer_cache.raw(), + pool_rows, + shape.indexer_heads as u32, + shape.indexer_head_dim as u32, + ((shape.indexer_heads * shape.indexer_head_dim) as f32) + .sqrt() + .recip(), + CACHE_F16, + ) + }, + "scoring GLM 5.3 indexer pools", + )?; + call( + unsafe { + ds4_gpu_indexer_topk_tensor( + self.scratch.indexer_pool_selected.raw(), + self.scratch.indexer_scores.raw(), + pool_rows, + 1, + selected_pools, + ) + }, + "selecting GLM 5.3 indexer pools", + )?; + let width = shape.indexer_top_k as u32 + GLM53_INDEX_POOL as u32 - 1; + call( + unsafe { + ds4_gpu_glm53_expand_pool_selection_tensor( + self.scratch.indexer_selected.raw(), + self.scratch.indexer_pool_selected.raw(), + 1, + pos, + selected_pools, + shape.indexer_top_k as u32, + GLM53_INDEX_POOL as u32, + width, + ) + }, + "expanding GLM 5.3 indexer pools", + )?; + width + }; + call( + unsafe { + ds4_gpu_glm_qk_lowrank_typed_tensor( + self.scratch.qk_low.raw(), + self.scratch.q.raw(), + map, + size, + weights.k_b.offset, + weights.k_b.kind, + shape.heads as u32, + shape.kv_lora as u32, + shape.key_mla as u32, + shape.key_mla as u32, + ) + }, + "projecting the GLM 5.3 low-rank query", + )?; + call( + unsafe { + ds4_gpu_glm_attention_indexed_decode_typed_tensor( + self.scratch.heads.raw(), + self.scratch.q.raw(), + self.scratch.qk_low.raw(), + cache.kv.raw(), + cache.rope.raw(), + map, + size, + weights.v_b.offset, + weights.v_b.kind, + self.scratch.indexer_selected.raw(), + selected_count, + self.context, + CACHE_F16, + shape.heads as u32, + shape.kv_lora as u32, + shape.key_mla as u32, + shape.rot as u32, + shape.value_mla as u32, + 0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + ) + }, + "running GLM 5.3 indexed attention", + )?; + glm53_decode_project( + &self.scratch.attn_out, + weights.output, + shape.heads * shape.value_mla, + shape.embd, + &self.scratch.heads, + map, + size, + self.ssd.enabled, + ) + } + + fn encode_glm53_kda_batch( + &self, + batch: &GlmBatchScratch, + weights: KdaWeights, + cache: &KdaCache, + rows: u32, + ) -> Result<(), String> { + let shape = self.model.shape; + let map = self.model.main.map_ptr().cast(); + let size = self.model.main.len(); + let projection = GLM53_KDA_HEADS * GLM53_KDA_HEAD_DIM; + for (out, weight) in [ + (&batch.kda_q, weights.q), + (&batch.kda_k, weights.k), + (&batch.kda_v, weights.v), + ] { + glm53_project_rows( + out, + weight, + shape.embd, + projection, + &batch.attn_norm, + rows, + map, + size, + )?; + } + glm53_project_rows( + &batch.kda_lowrank, + weights.f_a, + shape.embd, + GLM53_KDA_HEAD_DIM, + &batch.attn_norm, + rows, + map, + size, + )?; + glm53_project_rows( + &batch.kda_raw_gate, + weights.f_b, + GLM53_KDA_HEAD_DIM, + projection, + &batch.kda_lowrank, + rows, + map, + size, + )?; + glm53_project_rows( + &batch.kda_raw_beta, + weights.beta, + shape.embd, + GLM53_KDA_HEADS, + &batch.attn_norm, + rows, + map, + size, + )?; + glm53_project_rows( + &batch.kda_lowrank, + weights.g_a, + shape.embd, + GLM53_KDA_HEAD_DIM, + &batch.attn_norm, + rows, + map, + size, + )?; + glm53_project_rows( + &batch.kda_output_gate, + weights.g_b, + GLM53_KDA_HEAD_DIM, + projection, + &batch.kda_lowrank, + rows, + map, + size, + )?; + call( + unsafe { + ds4_gpu_glm53_kda_prefill( + batch.kda_out.raw(), + cache.conv.raw(), + cache.recurrent.raw(), + batch.kda_q.raw(), + batch.kda_k.raw(), + batch.kda_v.raw(), + batch.kda_raw_gate.raw(), + batch.kda_raw_beta.raw(), + batch.kda_output_gate.raw(), + map, + size, + weights.q_conv.offset, + weights.k_conv.offset, + weights.v_conv.offset, + weights.a_log.offset, + weights.dt_bias.offset, + weights.o_norm.offset, + GLM53_KDA_HEADS as u32, + rows, + -5.0, + shape.rms_epsilon, + ) + }, + "running batched GLM 5.3 KDA", + )?; + glm53_project_rows( + &batch.attn_out, + weights.output, + projection, + shape.embd, + &batch.kda_out, + rows, + map, + size, + ) + } + + #[allow(clippy::too_many_arguments)] + fn encode_glm53_dsa_batch( + &self, + batch: &GlmBatchScratch, + weights: DsaWeights, + cache: &LayerCache, + pos: u32, + rows: u32, + dense_limit: u32, + ) -> Result<(), String> { + let shape = self.model.shape; + let map = self.model.main.map_ptr().cast(); + let size = self.model.main.len(); + let q_dim = shape.heads * shape.key_mla; + glm53_project_rows( + &batch.q_rank, + weights.q_a, + shape.embd, + shape.lora_q, + &batch.attn_norm, + rows, + map, + size, + )?; + glm53_project_rows( + &batch.kv_raw, + weights.kv_a, + shape.embd, + shape.head_dim, + &batch.attn_norm, + rows, + map, + size, + )?; + call( + unsafe { + ds4_gpu_rms_norm_weight_rows_tensor( + batch.q_rank_norm.raw(), + batch.q_rank.raw(), + map, + size, + weights.q_a_norm.offset, + shape.lora_q as u32, + rows, + shape.rms_epsilon, + ) + }, + "normalizing batched GLM 5.3 queries", + )?; + call( + unsafe { + ds4_gpu_glm_kv_lora_rms_norm_tensor( + batch.kv_norm.raw(), + batch.kv_raw.raw(), + map, + size, + weights.kv_norm.offset, + rows, + shape.head_dim as u32, + shape.kv_lora as u32, + shape.rms_epsilon, + ) + }, + "normalizing batched GLM 5.3 compact KV", + )?; + call( + unsafe { + ds4_gpu_glm_store_compact_kv_tensor( + cache.kv.raw(), + cache.rope.raw(), + batch.kv_norm.raw(), + batch.kv_raw.raw(), + pos, + rows, + self.context, + shape.head_dim as u32, + shape.kv_lora as u32, + shape.rot as u32, + CACHE_F16, + ) + }, + "storing batched GLM 5.3 compact KV", + )?; + glm53_project_rows( + &batch.q, + weights.q_b, + shape.lora_q, + q_dim, + &batch.q_rank_norm, + rows, + map, + size, + )?; + + let indexer_cache = cache + .indexer + .as_ref() + .ok_or("GLM 5.3 DSA layer is missing its indexer cache")?; + let tail_k = cache + .indexer_tail_k + .as_ref() + .ok_or("GLM 5.3 DSA layer is missing its indexer key tail")?; + let tail_gate = cache + .indexer_tail_gate + .as_ref() + .ok_or("GLM 5.3 DSA layer is missing its indexer gate tail")?; + let ape = weights + .indexer_pool_ape + .ok_or("GLM 5.3 DSA layer is missing pool APE weights")?; + let gate = weights + .indexer_pool_gate + .ok_or("GLM 5.3 DSA layer is missing pool gate weights")?; + glm53_project_rows( + &batch.indexer_k, + weights.indexer_k, + shape.embd, + shape.indexer_head_dim, + &batch.attn_norm, + rows, + map, + size, + )?; + glm53_project_rows( + &batch.indexer_gate, + gate, + shape.embd, + shape.indexer_head_dim, + &batch.attn_norm, + rows, + map, + size, + )?; + call( + unsafe { + ds4_gpu_glm53_indexer_pool_update_tensor( + indexer_cache.raw(), + tail_k.raw(), + tail_gate.raw(), + batch.indexer_k.raw(), + batch.indexer_gate.raw(), + map, + size, + weights.indexer_k_norm.offset, + weights.indexer_k_bias.offset, + ape.offset, + pos, + rows, + self.context, + shape.indexer_head_dim as u32, + GLM53_INDEX_POOL as u32, + 1.0e-6, + CACHE_F16, + ) + }, + "updating batched GLM 5.3 indexer pools", + )?; + + let visible = pos + rows; + let causal = pos < dense_limit; + let selected_count = if causal { + visible.min(dense_limit) + } else { + glm53_project_rows( + &batch.indexer_q, + weights.indexer_q, + shape.lora_q, + shape.indexer_heads * shape.indexer_head_dim, + &batch.q_rank_norm, + rows, + map, + size, + )?; + glm53_project_rows( + &batch.indexer_weights, + weights.indexer_proj, + shape.embd, + shape.indexer_heads, + &batch.attn_norm, + rows, + map, + size, + )?; + let pools = visible / GLM53_INDEX_POOL as u32; + let selected_pools = shape.indexer_top_k as u32 / GLM53_INDEX_POOL as u32; + let score_rows = batch.score_rows(self.context, rows, true); + let q_row = shape.indexer_heads * shape.indexer_head_dim * 4; + let weight_row = shape.indexer_heads * 4; + let pool_row = selected_pools as u64 * 4; + let mut start = 0; + while start < rows { + let slice = (rows - start).min(score_rows); + let q = batch + .indexer_q + .view(u64::from(start) * q_row, u64::from(slice) * q_row)?; + let projected = batch + .indexer_weights + .view(u64::from(start) * weight_row, u64::from(slice) * weight_row)?; + let selected = batch + .indexer_pool_selected + .view(u64::from(start) * pool_row, u64::from(slice) * pool_row)?; + call( + unsafe { + ds4_gpu_glm53_indexer_scores_batch_tensor( + batch.indexer_scores.raw(), + q.raw(), + projected.raw(), + indexer_cache.raw(), + pools, + slice, + pos + start, + GLM53_INDEX_POOL as u32, + shape.indexer_heads as u32, + shape.indexer_head_dim as u32, + ((shape.indexer_heads * shape.indexer_head_dim) as f32) + .sqrt() + .recip(), + CACHE_F16, + ) + }, + "scoring batched GLM 5.3 indexer pools", + )?; + call( + unsafe { + ds4_gpu_indexer_topk_tensor( + selected.raw(), + batch.indexer_scores.raw(), + pools, + slice, + selected_pools, + ) + }, + "selecting batched GLM 5.3 indexer pools", + )?; + start += slice; + } + let width = shape.indexer_top_k as u32 + GLM53_INDEX_POOL as u32 - 1; + call( + unsafe { + ds4_gpu_glm53_expand_pool_selection_tensor( + batch.indexer_selected.raw(), + batch.indexer_pool_selected.raw(), + rows, + pos, + selected_pools, + shape.indexer_top_k as u32, + GLM53_INDEX_POOL as u32, + width, + ) + }, + "expanding batched GLM 5.3 indexer pools", + )?; + width + }; + call( + unsafe { + ds4_gpu_glm_qk_lowrank_typed_batch_tensor( + batch.qk_low.raw(), + batch.q.raw(), + map, + size, + weights.k_b.offset, + weights.k_b.kind, + rows, + shape.heads as u32, + shape.kv_lora as u32, + shape.key_mla as u32, + shape.key_mla as u32, + ) + }, + "projecting batched GLM 5.3 low-rank queries", + )?; + self.encode_batch_attention_raw(batch, weights, cache, pos, rows, selected_count, causal)?; + glm53_project_rows( + &batch.attn_out, + weights.output, + shape.heads * shape.value_mla, + shape.embd, + &batch.heads, + rows, + map, + size, + ) + } + + fn eval_glm53(&mut self, token: i32) -> Result<(), String> { + self.eval_glm53_capture(token, None, None) + } + + fn eval_glm53_capture( + &mut self, + token: i32, + hidden: Option<&Buffer>, + selected: Option<&mut i32>, + ) -> Result<(), String> { + let pos = self.position(); + if pos >= self.context { + return Err(format!( + "the GLM Metal executor supports {} tokens", + self.context + )); + } + self.scratch.token.write_i32(&[token])?; + let commands = Commands::begin()?; + self.encode_glm53_token(token, pos, self.scratch.token.raw())?; + if let Some(hidden) = hidden { + hidden.copy_from( + 0, + &self.scratch.hc_output, + 0, + self.model.shape.embd * 4, + "saving a GLM 5.3 MTP hidden row", + )?; + } + if selected.is_some() { + call( + unsafe { + ds4_gpu_argmax_tensor( + self.scratch.token.raw(), + self.scratch.logits.raw(), + self.model.shape.vocab as u32, + ) + }, + "selecting the GLM 5.3 verifier token", + )?; + } + commands.finish()?; + if let Some(selected) = selected { + self.scratch + .token + .read_i32(std::slice::from_mut(selected))?; + } else { + self.scratch.logits.read_f32(&mut self.logits)?; + } + self.tokens.push(token); + Ok(()) + } + + fn encode_glm53_token( + &mut self, + token: i32, + pos: u32, + token_tensor: *const GpuTensor, + ) -> Result<(), String> { + let shape = self.model.shape; + let map = self.model.main.map_ptr().cast(); + let size = self.model.main.len(); + if self.weights.embedding.kind == BF16 { + call( + unsafe { + ds4_gpu_glm53_embedding_bf16( + self.scratch.current.raw(), + map, + size, + self.weights.embedding.offset, + token_tensor, + 1, + shape.embd as u32, + shape.vocab as u32, + ) + }, + "embedding a GLM 5.3 token", + )?; + } else { + call( + unsafe { + ds4_gpu_embed_token_quant_tensor( + self.scratch.current.raw(), + map, + size, + self.weights.embedding.offset, + self.weights.embedding.kind, + shape.vocab as u32, + token as u32, + shape.embd as u32, + ) + }, + "embedding a GLM 5.3 token", + )?; + } + call( + unsafe { + ds4_gpu_repeat_hc_tensor( + self.scratch.hc_current.raw(), + self.scratch.current.raw(), + shape.embd as u32, + shape.hc as u32, + ) + }, + "opening GLM 5.3 hyperconnections", + )?; + for (ordinal, (layer, cache)) in self.weights.layers.iter().zip(&self.caches).enumerate() { + let hc = layer + .hc + .ok_or("GLM 5.3 layer has no hyperconnection weights")?; + glm53_hc_pre( + &self.scratch.current, + &self.scratch.attn_norm, + &self.scratch.hc_flat, + &self.scratch.hc_mix, + &self.scratch.hc_split, + &self.scratch.hc_current, + hc.attn_fn, + hc.attn_scale, + hc.attn_base, + layer.attn_norm, + shape, + 1, + map, + size, + )?; + match layer.attention { + AttentionWeights::Kda(weights) => self.encode_glm53_kda( + weights, + cache.kda.as_ref().ok_or("GLM 5.3 KDA cache is missing")?, + )?, + AttentionWeights::Dsa(weights) => self.encode_glm53_dsa(weights, cache, pos)?, + } + if let Some(steering) = &self.steering { + steering.apply(&self.scratch.attn_out, ordinal as u32, 1, true)?; + } + let post = self.scratch.hc_split.view(shape.hc * 4, shape.hc * 4)?; + let combine = self + .scratch + .hc_split + .view(2 * shape.hc * 4, shape.hc * shape.hc * 4)?; + call( + unsafe { + ds4_gpu_hc_expand_tensor( + self.scratch.hc_after_attn.raw(), + self.scratch.attn_out.raw(), + self.scratch.hc_current.raw(), + post.raw(), + combine.raw(), + shape.embd as u32, + shape.hc as u32, + ) + }, + "expanding GLM 5.3 attention hyperconnections", + )?; + glm53_hc_pre( + &self.scratch.after_attn, + &self.scratch.ffn_norm, + &self.scratch.hc_flat, + &self.scratch.hc_mix, + &self.scratch.hc_split, + &self.scratch.hc_after_attn, + hc.ffn_fn, + hc.ffn_scale, + hc.ffn_base, + layer.ffn_norm, + shape, + 1, + map, + size, + )?; + self.encode_glm53_ffn(layer, ordinal as u32)?; + if let Some(steering) = &self.steering { + steering.apply(&self.scratch.next, ordinal as u32, 1, false)?; + } + let post = self.scratch.hc_split.view(shape.hc * 4, shape.hc * 4)?; + let combine = self + .scratch + .hc_split + .view(2 * shape.hc * 4, shape.hc * shape.hc * 4)?; + call( + unsafe { + ds4_gpu_hc_expand_tensor( + self.scratch.hc_next.raw(), + self.scratch.next.raw(), + self.scratch.hc_after_attn.raw(), + post.raw(), + combine.raw(), + shape.embd as u32, + shape.hc as u32, + ) + }, + "expanding GLM 5.3 FFN hyperconnections", + )?; + std::mem::swap(&mut self.scratch.hc_current, &mut self.scratch.hc_next); + if (ordinal + 1).is_multiple_of(DECODE_FLUSH_LAYERS) + && ordinal + 1 < self.weights.layers.len() + { + call( + unsafe { ds4_gpu_flush_commands() }, + "flushing GLM 5.3 decode", + )?; + } + } + call( + unsafe { + ds4_gpu_hc_weighted_sum_tensor( + self.scratch.hc_output.raw(), + self.scratch.hc_current.raw(), + self.scratch.hc_mean.raw(), + shape.embd as u32, + shape.hc as u32, + ) + }, + "collapsing GLM 5.3 output hyperconnections", + )?; + norm( + &self.scratch.output_norm, + &self.scratch.hc_output, + self.weights.output_norm, + shape.embd as u32, + shape.rms_epsilon, + map, + size, + )?; + glm53_project( + &self.scratch.logits, + self.weights.output, + shape.embd, + shape.vocab, + &self.scratch.output_norm, + map, + size, + ) + } + + fn encode_glm53_ffn(&self, layer: &GlmLayer, ordinal: u32) -> Result<(), String> { + let shape = self.model.shape; + let map = self.model.main.map_ptr().cast(); + let size = self.model.main.len(); + if let Some(dense) = layer.dense { + let fused = dense.gate.kind == Q8_0 + && dense.up.kind == Q8_0 + && unsafe { + ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor( + self.scratch.ffn_gate.raw(), + self.scratch.ffn_up.raw(), + self.scratch.ffn_mid.raw(), + map, + size, + dense.gate.offset, + dense.up.offset, + shape.embd, + shape.ff_dense, + self.scratch.ffn_norm.raw(), + shape.swiglu_clamp, + ) + } != 0; + if !fused { + glm53_decode_project( + &self.scratch.ffn_gate, + dense.gate, + shape.embd, + shape.ff_dense, + &self.scratch.ffn_norm, + map, + size, + self.ssd.enabled, + )?; + glm53_decode_project( + &self.scratch.ffn_up, + dense.up, + shape.embd, + shape.ff_dense, + &self.scratch.ffn_norm, + map, + size, + self.ssd.enabled, + )?; + call( + unsafe { + ds4_gpu_swiglu_tensor( + self.scratch.ffn_mid.raw(), + self.scratch.ffn_gate.raw(), + self.scratch.ffn_up.raw(), + shape.ff_dense as u32, + shape.swiglu_clamp, + 1.0, + ) + }, + "activating the dense GLM 5.3 FFN", + )?; + } + glm53_decode_project( + &self.scratch.next, + dense.down, + shape.ff_dense, + shape.embd, + &self.scratch.ffn_mid, + map, + size, + self.ssd.enabled, + )?; + return Ok(()); + } + let sparse = layer.sparse.ok_or("GLM 5.3 layer has no FFN weights")?; + f32_project( + &self.scratch.router_logits, + sparse.router, + shape.embd, + shape.experts, + &self.scratch.ffn_norm, + map, + size, + )?; + call( + unsafe { + ds4_gpu_glm_router_select_tensor( + self.scratch.router_selected.raw(), + self.scratch.router_weights.raw(), + self.scratch.router_probs.raw(), + map, + size, + sparse.bias.offset, + self.scratch.router_logits.raw(), + shape.experts as u32, + shape.experts_used as u32, + shape.expert_weight_scale, + ) + }, + "routing GLM 5.3 experts", + )?; + let force_resident = + self.ssd.enabled && ordinal.saturating_sub(shape.leading_dense) < self.ssd.full_layers; + let (gate_expert, gate_row) = expert_layout(sparse.gate, shape.experts); + let (up_expert, up_row) = expert_layout(sparse.up, shape.experts); + let (down_expert, down_row) = expert_layout(sparse.down, shape.experts); + let routed = unsafe { + if sparse.gate.kind == IQ2_XXS { + ds4_gpu_routed_moe_one_tensor( + self.scratch.ffn_out.raw(), + self.scratch.ffn_gate.raw(), + self.scratch.ffn_up.raw(), + self.scratch.ffn_mid.raw(), + self.scratch.routed_down.raw(), + map, + size, + sparse.gate.offset, + sparse.up.offset, + sparse.down.offset, + sparse.gate.kind, + sparse.down.kind, + gate_expert, + gate_row, + down_expert, + down_row, + shape.embd as u32, + shape.ff_expert as u32, + shape.embd as u32, + self.scratch.router_selected.raw(), + self.scratch.router_weights.raw(), + shape.experts as u32, + shape.experts_used as u32, + shape.swiglu_clamp, + self.scratch.ffn_norm.raw(), + std::ptr::null(), + ordinal, + force_resident, + ) + } else { + ds4_gpu_glm_routed_moe_one_tensor( + self.scratch.ffn_out.raw(), + self.scratch.ffn_mid.raw(), + map, + size, + sparse.gate.offset, + sparse.up.offset, + sparse.down.offset, + sparse.gate.kind, + sparse.up.kind, + sparse.down.kind, + gate_expert, + gate_row, + up_expert, + up_row, + down_expert, + down_row, + shape.embd as u32, + shape.ff_expert as u32, + shape.embd as u32, + self.scratch.router_selected.raw(), + self.scratch.router_weights.raw(), + shape.experts as u32, + shape.experts_used as u32, + ordinal, + self.scratch.ffn_norm.raw(), + force_resident, + ) + } + }; + call(routed, "running routed GLM 5.3 experts")?; + let fused = sparse.shared_gate.kind == Q8_0 + && sparse.shared_up.kind == Q8_0 + && unsafe { + ds4_gpu_shared_mid_swiglu_q8_0_tensor( + self.scratch.ffn_mid.raw(), + map, + size, + sparse.shared_gate.offset, + sparse.shared_up.offset, + shape.embd, + shape.ff_expert, + self.scratch.ffn_norm.raw(), + shape.swiglu_clamp, + ) + } != 0; + if !fused { + glm53_decode_project( + &self.scratch.ffn_gate, + sparse.shared_gate, + shape.embd, + shape.ff_expert, + &self.scratch.ffn_norm, + map, + size, + self.ssd.enabled, + )?; + glm53_decode_project( + &self.scratch.ffn_up, + sparse.shared_up, + shape.embd, + shape.ff_expert, + &self.scratch.ffn_norm, + map, + size, + self.ssd.enabled, + )?; + call( + unsafe { + ds4_gpu_swiglu_tensor( + self.scratch.ffn_mid.raw(), + self.scratch.ffn_gate.raw(), + self.scratch.ffn_up.raw(), + shape.ff_expert as u32, + shape.swiglu_clamp, + 1.0, + ) + }, + "activating the shared GLM 5.3 expert", + )?; + } + glm53_decode_project( + &self.scratch.ffn_sum, + sparse.shared_down, + shape.ff_expert, + shape.embd, + &self.scratch.ffn_mid, + map, + size, + self.ssd.enabled, + )?; + call( + unsafe { + ds4_gpu_add_tensor( + self.scratch.next.raw(), + self.scratch.ffn_out.raw(), + self.scratch.ffn_sum.raw(), + shape.embd as u32, + ) + }, + "combining GLM 5.3 experts", + ) + } + + fn encode_glm53_ffn_batch( + &self, + batch: &GlmBatchScratch, + layer: &GlmLayer, + ordinal: u32, + rows: u32, + ) -> Result<(), String> { + let shape = self.model.shape; + let map = self.model.main.map_ptr().cast(); + let size = self.model.main.len(); + if let Some(dense) = layer.dense { + if !glm53_fused_gate_up_rows( + &batch.ffn_gate, + &batch.ffn_up, + &batch.ffn_mid, + dense.gate, + dense.up, + shape.embd, + shape.ff_dense, + &batch.ffn_norm, + rows, + shape.swiglu_clamp, + map, + size, + ) { + glm53_project_rows( + &batch.ffn_gate, + dense.gate, + shape.embd, + shape.ff_dense, + &batch.ffn_norm, + rows, + map, + size, + )?; + glm53_project_rows( + &batch.ffn_up, + dense.up, + shape.embd, + shape.ff_dense, + &batch.ffn_norm, + rows, + map, + size, + )?; + call( + unsafe { + ds4_gpu_swiglu_tensor( + batch.ffn_mid.raw(), + batch.ffn_gate.raw(), + batch.ffn_up.raw(), + rows * shape.ff_dense as u32, + shape.swiglu_clamp, + 1.0, + ) + }, + "activating a batched dense GLM 5.3 FFN", + )?; + } + return glm53_project_rows( + &batch.next, + dense.down, + shape.ff_dense, + shape.embd, + &batch.ffn_mid, + rows, + map, + size, + ); + } + + let sparse = layer.sparse.ok_or("GLM 5.3 layer has no FFN weights")?; + f32_project_rows( + &batch.router_logits, + sparse.router, + shape.embd, + shape.experts, + &batch.ffn_norm, + rows, + map, + size, + )?; + call( + unsafe { + ds4_gpu_glm_router_select_batch_tensor( + batch.router_selected.raw(), + batch.router_weights.raw(), + batch.router_probs.raw(), + map, + size, + sparse.bias.offset, + batch.router_logits.raw(), + shape.experts as u32, + shape.experts_used as u32, + shape.expert_weight_scale, + rows, + ) + }, + "routing a batched GLM 5.3 FFN", + )?; + let (gate_expert, gate_row) = expert_layout(sparse.gate, shape.experts); + let (up_expert, up_row) = expert_layout(sparse.up, shape.experts); + let (down_expert, down_row) = expert_layout(sparse.down, shape.experts); + let force_resident = + self.ssd.enabled && ordinal.saturating_sub(shape.leading_dense) < self.ssd.full_layers; + let mut mid_is_f16 = false; + let routed = unsafe { + if sparse.gate.kind == IQ2_XXS { + ds4_gpu_routed_moe_batch_tensor( + batch.ffn_out.raw(), + batch.routed_gate.raw(), + batch.routed_up.raw(), + batch.ffn_mid.raw(), + batch.routed_down.raw(), + map, + size, + sparse.gate.offset, + sparse.up.offset, + sparse.down.offset, + sparse.gate.kind, + sparse.down.kind, + gate_expert, + gate_row, + down_expert, + down_row, + shape.embd as u32, + shape.ff_expert as u32, + shape.embd as u32, + batch.router_selected.raw(), + batch.router_weights.raw(), + shape.experts as u32, + shape.experts_used as u32, + shape.swiglu_clamp, + batch.ffn_norm.raw(), + ordinal, + rows, + &mut mid_is_f16, + force_resident, + ) + } else { + ds4_gpu_glm_routed_moe_batch_tensor( + batch.ffn_out.raw(), + batch.ffn_mid.raw(), + map, + size, + sparse.gate.offset, + sparse.up.offset, + sparse.down.offset, + sparse.gate.kind, + sparse.up.kind, + sparse.down.kind, + gate_expert, + gate_row, + up_expert, + up_row, + down_expert, + down_row, + shape.embd as u32, + shape.ff_expert as u32, + shape.embd as u32, + batch.router_selected.raw(), + batch.router_weights.raw(), + shape.experts as u32, + shape.experts_used as u32, + ordinal, + batch.ffn_norm.raw(), + rows, + (shape.experts_used * shape.ff_expert) as u32, + force_resident, + ) + } + }; + call(routed, "running batched routed GLM 5.3 experts")?; + if !glm53_fused_gate_up_rows( + &batch.ffn_gate, + &batch.ffn_up, + &batch.ffn_mid, + sparse.shared_gate, + sparse.shared_up, + shape.embd, + shape.ff_expert, + &batch.ffn_norm, + rows, + shape.swiglu_clamp, + map, + size, + ) { + glm53_project_rows( + &batch.ffn_gate, + sparse.shared_gate, + shape.embd, + shape.ff_expert, + &batch.ffn_norm, + rows, + map, + size, + )?; + glm53_project_rows( + &batch.ffn_up, + sparse.shared_up, + shape.embd, + shape.ff_expert, + &batch.ffn_norm, + rows, + map, + size, + )?; + call( + unsafe { + ds4_gpu_swiglu_tensor( + batch.ffn_mid.raw(), + batch.ffn_gate.raw(), + batch.ffn_up.raw(), + rows * shape.ff_expert as u32, + shape.swiglu_clamp, + 1.0, + ) + }, + "activating batched shared GLM 5.3 experts", + )?; + } + glm53_project_rows( + &batch.after_attn, + sparse.shared_down, + shape.ff_expert, + shape.embd, + &batch.ffn_mid, + rows, + map, + size, + )?; + call( + unsafe { + ds4_gpu_add_tensor( + batch.next.raw(), + batch.ffn_out.raw(), + batch.after_attn.raw(), + rows * shape.embd as u32, + ) + }, + "combining batched GLM 5.3 experts", + ) + } + #[allow(clippy::too_many_arguments)] fn encode_layer( &self, @@ -1524,6 +3781,7 @@ impl GlmExecutor { selected_count: &mut u32, attention_selected: Option<(&Buffer, u32)>, ) -> Result<(), String> { + let dsa = layer.dsa()?; let shape = self.model.shape; let map = self.model.main.map_ptr().cast(); let size = self.model.main.len(); @@ -1540,7 +3798,7 @@ impl GlmExecutor { )?; project( &self.scratch.q_rank, - layer.q_a, + dsa.q_a, shape.embd, shape.lora_q, &self.scratch.attn_norm, @@ -1550,7 +3808,7 @@ impl GlmExecutor { )?; project( &self.scratch.kv_raw, - layer.kv_a, + dsa.kv_a, shape.embd, shape.head_dim, &self.scratch.attn_norm, @@ -1565,12 +3823,12 @@ impl GlmExecutor { self.scratch.q_rank.raw(), map, size, - layer.q_a_norm.offset, + dsa.q_a_norm.offset, shape.lora_q as u32, cache.kv.raw(), cache.rope.raw(), self.scratch.kv_raw.raw(), - layer.kv_norm.offset, + dsa.kv_norm.offset, pos, 1, self.context, @@ -1585,7 +3843,7 @@ impl GlmExecutor { )?; project( &self.scratch.q, - layer.q_b, + dsa.q_b, shape.lora_q, q_dim, &self.scratch.q_rank_norm, @@ -1619,7 +3877,7 @@ impl GlmExecutor { } else if let Some(indexer_cache) = &cache.indexer { project( &self.scratch.indexer_k, - layer.indexer_k, + dsa.indexer_k, shape.embd, shape.indexer_head_dim, &self.scratch.current, @@ -1634,8 +3892,8 @@ impl GlmExecutor { self.scratch.indexer_k.raw(), map, size, - layer.indexer_k_norm.offset, - layer.indexer_k_bias.offset, + dsa.indexer_k_norm.offset, + dsa.indexer_k_bias.offset, pos, 1, self.context, @@ -1669,7 +3927,7 @@ impl GlmExecutor { } else { project( &self.scratch.indexer_q, - layer.indexer_q, + dsa.indexer_q, shape.lora_q, shape.indexer_heads * shape.indexer_head_dim, &self.scratch.q_rank_norm, @@ -1699,7 +3957,7 @@ impl GlmExecutor { )?; f32_project( &self.scratch.indexer_weights, - layer.indexer_proj, + dsa.indexer_proj, shape.embd, shape.indexer_heads, &self.scratch.current, @@ -1747,8 +4005,8 @@ impl GlmExecutor { self.scratch.q.raw(), map, size, - layer.k_b.offset, - layer.k_b.kind, + dsa.k_b.offset, + dsa.k_b.kind, shape.heads as u32, shape.kv_lora as u32, q_nope as u32, @@ -1767,8 +4025,8 @@ impl GlmExecutor { cache.rope.raw(), map, size, - layer.v_b.offset, - layer.v_b.kind, + dsa.v_b.offset, + dsa.v_b.kind, attention_selected.map_or_else( || self.scratch.indexer_selected.raw(), |value| value.0.raw(), @@ -1794,7 +4052,7 @@ impl GlmExecutor { )?; project( &self.scratch.attn_out, - layer.output, + dsa.output, shape.heads * shape.value_mla, shape.embd, &self.scratch.heads, @@ -2079,9 +4337,19 @@ impl GlmExecutor { } let pos = self.position(); - let rows = + let mut rows = indexed_prefill_rows(pos, remaining.len(), self.model.shape.indexer_top_k as u32); - let cancelled = self.eval_batch(&remaining[..rows], &mut progress)?; + if self.model.shape.model == ModelChoice::Glm53Flash { + let dense = glm53_dense_limit(self.context, self.ssd.enabled); + if pos < dense { + rows = rows.min((dense - pos) as usize); + } + } + let cancelled = self.eval_batch( + &remaining[..rows], + &mut progress, + completed + rows == tokens.len(), + )?; completed += rows; if cancelled { break; @@ -2114,6 +4382,7 @@ impl GlmExecutor { &mut self, tokens: &[i32], progress: &mut impl FnMut(u32) -> bool, + final_batch: bool, ) -> Result { let rows = u32::try_from(tokens.len()).map_err(|_| "GLM prefill batch is too large")?; if rows == 0 || rows > INDEXED_PREFILL_CHUNK { @@ -2125,6 +4394,9 @@ impl GlmExecutor { { return Err("GLM prefill contains a token outside the vocabulary".into()); } + if self.model.shape.model == ModelChoice::Glm53Flash { + return self.eval_glm53_batch(tokens, progress, final_batch); + } let pos = self.position(); let mut batch = GlmBatchScratch::allocate(&self.model, rows, self.context)?; @@ -2234,6 +4506,311 @@ impl GlmExecutor { Ok(cancelled) } + fn eval_glm53_batch( + &mut self, + tokens: &[i32], + progress: &mut impl FnMut(u32) -> bool, + read_logits: bool, + ) -> Result { + self.eval_glm53_batch_inner_entry(tokens, progress, read_logits) + } + + fn eval_glm53_batch_inner_entry( + &mut self, + tokens: &[i32], + progress: &mut impl FnMut(u32) -> bool, + read_logits: bool, + ) -> Result { + let mut batch = self + .batch + .take() + .ok_or("GLM 5.3 prefill workspace is missing")?; + let result = self.eval_glm53_batch_inner(tokens, progress, read_logits, &mut batch); + self.batch = Some(batch); + result + } + + fn eval_glm53_batch_inner( + &mut self, + tokens: &[i32], + progress: &mut impl FnMut(u32) -> bool, + read_logits: bool, + batch: &mut GlmBatchScratch, + ) -> Result { + let rows = u32::try_from(tokens.len()).map_err(|_| "GLM 5.3 prefill is too large")?; + let pos = self.position(); + let shape = self.model.shape; + let map = self.model.main.map_ptr().cast(); + let size = self.model.main.len(); + let dense_limit = glm53_dense_limit(self.context, self.ssd.enabled); + if pos < dense_limit && pos + rows > dense_limit { + return Err("GLM 5.3 prefill crossed the dense-attention boundary".into()); + } + batch.tokens.write_i32(tokens)?; + let mut commands = Commands::begin()?; + if self.weights.embedding.kind == BF16 { + call( + unsafe { + ds4_gpu_glm53_embedding_bf16( + batch.current.raw(), + map, + size, + self.weights.embedding.offset, + batch.tokens.raw(), + rows, + shape.embd as u32, + shape.vocab as u32, + ) + }, + "embedding a GLM 5.3 prefill batch", + )?; + } else { + call( + unsafe { + ds4_gpu_embed_tokens_quant_tensor( + batch.current.raw(), + batch.tokens.raw(), + map, + size, + self.weights.embedding.offset, + self.weights.embedding.kind, + shape.vocab as u32, + rows, + shape.embd as u32, + ) + }, + "embedding a GLM 5.3 prefill batch", + )?; + } + call( + unsafe { + ds4_gpu_repeat_hc_rows_tensor( + batch.hc_current.raw(), + batch.current.raw(), + rows, + shape.embd as u32, + shape.hc as u32, + ) + }, + "opening batched GLM 5.3 hyperconnections", + )?; + for overlay in &self.vision_overlays { + let start = overlay.token_start.max(pos); + let end = overlay + .token_start + .saturating_add(overlay.tokens) + .min(pos + rows); + if start < end { + call( + unsafe { + ds4_gpu_glm53_scatter_image_hc( + batch.hc_current.raw(), + overlay.embedding.raw(), + start - pos, + start - overlay.token_start, + end - start, + rows, + shape.embd as u32, + shape.hc as u32, + ) + }, + "overlaying GLM 5.3 image embeddings", + )?; + } + } + let mut cancelled = false; + for index in 0..self.weights.layers.len() { + if self.ssd.enabled + && self.weights.layers[index] + .sparse + .is_some_and(|sparse| sparse.gate.kind == IQ2_XXS) + && (index as u32).saturating_sub(shape.leading_dense) >= self.ssd.full_layers + { + install_glm_model_spans( + &self.model, + &glm_layer_model_spans(&self.model, index as u32)?, + "GLM 5.3 indexed-prefill layer", + )?; + } + let layer = &self.weights.layers[index]; + let hc = layer + .hc + .ok_or("GLM 5.3 layer has no hyperconnection weights")?; + glm53_hc_pre( + &batch.current, + &batch.attn_norm, + &batch.hc_flat, + &batch.hc_mix, + &batch.hc_split, + &batch.hc_current, + hc.attn_fn, + hc.attn_scale, + hc.attn_base, + layer.attn_norm, + shape, + rows, + map, + size, + )?; + match layer.attention { + AttentionWeights::Kda(weights) => self.encode_glm53_kda_batch( + batch, + weights, + self.caches[index] + .kda + .as_ref() + .ok_or("GLM 5.3 KDA cache is missing")?, + rows, + )?, + AttentionWeights::Dsa(weights) => self.encode_glm53_dsa_batch( + batch, + weights, + &self.caches[index], + pos, + rows, + dense_limit, + )?, + } + if let Some(steering) = &self.steering { + steering.apply(&batch.attn_out, index as u32, rows, true)?; + } + call( + unsafe { + ds4_gpu_hc_expand_split_tensor( + batch.hc_after_attn.raw(), + batch.attn_out.raw(), + batch.hc_current.raw(), + batch.hc_split.raw(), + shape.embd as u32, + shape.hc as u32, + ) + }, + "expanding batched GLM 5.3 attention hyperconnections", + )?; + glm53_hc_pre( + &batch.after_attn, + &batch.ffn_norm, + &batch.hc_flat, + &batch.hc_mix, + &batch.hc_split, + &batch.hc_after_attn, + hc.ffn_fn, + hc.ffn_scale, + hc.ffn_base, + layer.ffn_norm, + shape, + rows, + map, + size, + )?; + self.encode_glm53_ffn_batch(batch, layer, index as u32, rows)?; + if let Some(steering) = &self.steering { + steering.apply(&batch.next, index as u32, rows, false)?; + } + call( + unsafe { + ds4_gpu_hc_expand_split_tensor( + batch.hc_next.raw(), + batch.next.raw(), + batch.hc_after_attn.raw(), + batch.hc_split.raw(), + shape.embd as u32, + shape.hc as u32, + ) + }, + "expanding batched GLM 5.3 FFN hyperconnections", + )?; + if glm_prefill_flush_layers(rows, read_logits) { + if (index + 1).is_multiple_of(PREFILL_DRAIN_LAYERS) + && index + 1 < self.weights.layers.len() + { + commands.finish()?; + commands = Commands::begin()?; + } else { + commands.flush()?; + } + } + if layer.sparse.is_some() + && let Some(profile) = &mut self.profile + { + profile.record( + index, + pos, + &batch.router_selected, + &batch.router_weights, + rows, + false, + )?; + } + std::mem::swap(&mut batch.hc_current, &mut batch.hc_next); + let done = u32::try_from( + u64::from(rows) * (index as u64 + 1) / self.weights.layers.len() as u64, + ) + .expect("GLM 5.3 prefill progress is bounded by the batch"); + cancelled |= !progress(pos + done); + } + commands.finish()?; + + if let Some(spans) = &self.streaming_spans { + install_glm_model_spans(&self.model, spans, "GLM 5.3 decode/output")?; + } + if read_logits { + let last_hc = batch.hc_current.view( + u64::from(rows - 1) * shape.hc * shape.embd * 4, + shape.hc * shape.embd * 4, + )?; + let commands = Commands::begin()?; + call( + unsafe { + ds4_gpu_hc_weighted_sum_tensor( + self.scratch.hc_output.raw(), + last_hc.raw(), + self.scratch.hc_mean.raw(), + shape.embd as u32, + shape.hc as u32, + ) + }, + "collapsing batched GLM 5.3 output hyperconnections", + )?; + commands.finish()?; + let mut logits = vec![0.0; shape.vocab as usize]; + self.read_glm53_output(&self.scratch.hc_output, &mut logits)?; + self.logits = logits; + } + self.tokens.extend_from_slice(tokens); + if let Some(profile) = &self.profile { + profile.write()?; + } + Ok(cancelled) + } + + fn read_glm53_output(&self, hidden: &Buffer, logits: &mut [f32]) -> Result<(), String> { + let shape = self.model.shape; + let map = self.model.main.map_ptr().cast(); + let size = self.model.main.len(); + let commands = Commands::begin()?; + norm( + &self.scratch.output_norm, + hidden, + self.weights.output_norm, + shape.embd as u32, + shape.rms_epsilon, + map, + size, + )?; + glm53_project( + &self.scratch.logits, + self.weights.output, + shape.embd, + shape.vocab, + &self.scratch.output_norm, + map, + size, + )?; + commands.finish()?; + self.scratch.logits.read_f32(logits) + } + pub(super) fn logits(&self) -> &[f32] { &self.logits } @@ -2313,6 +4890,7 @@ impl GlmExecutor { self.context, )?; mtp.pending = None; + mtp.parent = None; mtp.min_pos = None; } Ok(()) @@ -2332,6 +4910,7 @@ impl GlmExecutor { selected: Buffer::bytes(u64::from(self.context) * 4)?, concat: Buffer::floats(2 * self.model.shape.embd)?, pending: None, + parent: None, min_pos: None, cycles: 0, drafted: 0, @@ -2376,6 +4955,7 @@ impl GlmExecutor { self.tokens.truncate(rewind); if let Some(mtp) = &mut self.mtp { mtp.pending = None; + mtp.parent = None; mtp.min_pos = None; } return Ok(rewind); @@ -2428,31 +5008,68 @@ impl GlmExecutor { let mut chunk = vec![0; CHECKPOINT_IO_CHUNK]; let rows = u64::from(self.position()); for cache in &self.caches { - write_buffer( - &mut file, - &cache.kv, - 0, - rows * shape.kv_lora * 2, - &mut chunk, - progress, - )?; - write_buffer( - &mut file, - &cache.rope, - 0, - rows * shape.rot * 2, - &mut chunk, - progress, - )?; + if let Some(kda) = &cache.kda { + write_buffer( + &mut file, + &kda.conv, + 0, + 3 * (GLM53_KDA_CONV - 1) * GLM53_KDA_HEADS * GLM53_KDA_HEAD_DIM * 4, + &mut chunk, + progress, + )?; + write_buffer( + &mut file, + &kda.recurrent, + 0, + GLM53_KDA_HEADS * GLM53_KDA_HEAD_DIM * GLM53_KDA_HEAD_DIM * 4, + &mut chunk, + progress, + )?; + } else { + write_buffer( + &mut file, + &cache.kv, + 0, + rows * shape.kv_lora * 2, + &mut chunk, + progress, + )?; + write_buffer( + &mut file, + &cache.rope, + 0, + rows * shape.rot * 2, + &mut chunk, + progress, + )?; + } if let Some(indexer) = &cache.indexer { + let index_rows = if shape.model == ModelChoice::Glm53Flash { + rows.div_ceil(GLM53_INDEX_POOL) + } else { + rows + }; write_buffer( &mut file, indexer, 0, - rows * shape.indexer_head_dim * 2, + index_rows * shape.indexer_head_dim * 2, &mut chunk, progress, )?; + for tail in [&cache.indexer_tail_k, &cache.indexer_tail_gate] + .into_iter() + .flatten() + { + write_buffer( + &mut file, + tail, + 0, + GLM53_INDEX_POOL * shape.indexer_head_dim * 4, + &mut chunk, + progress, + )?; + } } } file.sync_all().map_err(|e| e.to_string())?; @@ -2517,31 +5134,68 @@ impl GlmExecutor { let mut chunk = vec![0; CHECKPOINT_IO_CHUNK]; let rows = u64::from(token_count); for cache in &self.caches { - read_buffer( - &mut file, - &cache.kv, - 0, - rows * shape.kv_lora * 2, - &mut chunk, - progress, - )?; - read_buffer( - &mut file, - &cache.rope, - 0, - rows * shape.rot * 2, - &mut chunk, - progress, - )?; + if let Some(kda) = &cache.kda { + read_buffer( + &mut file, + &kda.conv, + 0, + 3 * (GLM53_KDA_CONV - 1) * GLM53_KDA_HEADS * GLM53_KDA_HEAD_DIM * 4, + &mut chunk, + progress, + )?; + read_buffer( + &mut file, + &kda.recurrent, + 0, + GLM53_KDA_HEADS * GLM53_KDA_HEAD_DIM * GLM53_KDA_HEAD_DIM * 4, + &mut chunk, + progress, + )?; + } else { + read_buffer( + &mut file, + &cache.kv, + 0, + rows * shape.kv_lora * 2, + &mut chunk, + progress, + )?; + read_buffer( + &mut file, + &cache.rope, + 0, + rows * shape.rot * 2, + &mut chunk, + progress, + )?; + } if let Some(indexer) = &cache.indexer { + let index_rows = if shape.model == ModelChoice::Glm53Flash { + rows.div_ceil(GLM53_INDEX_POOL) + } else { + rows + }; read_buffer( &mut file, indexer, 0, - rows * shape.indexer_head_dim * 2, + index_rows * shape.indexer_head_dim * 2, &mut chunk, progress, )?; + for tail in [&cache.indexer_tail_k, &cache.indexer_tail_gate] + .into_iter() + .flatten() + { + read_buffer( + &mut file, + tail, + 0, + GLM53_INDEX_POOL * shape.indexer_head_dim * 4, + &mut chunk, + progress, + )?; + } } } let mut trailing = [0]; @@ -2565,6 +5219,25 @@ fn streaming_token_prefill_eligible(context: u32, pos: u32, rows: u32, max: u32) max != 0 && rows != 0 && rows <= max && pos < full_cap && rows <= full_cap - pos } +fn glm53_dense_limit(context: u32, ssd_streaming: bool) -> u32 { + if context >= LONG_CONTEXT_THRESHOLD { + LONG_CONTEXT_FULL_ATTN_CONTEXT + } else if ssd_streaming { + STREAMING_FULL_ATTN_CONTEXT + } else { + RESIDENT_FULL_ATTN_CONTEXT + } + .min(context) +} + +fn glm53_dense_compact_prefill(rows: u32) -> bool { + rows >= 24 +} + +fn glm_prefill_flush_layers(rows: u32, logits_requested: bool) -> bool { + logits_requested && rows > 8 +} + fn indexed_prefill_rows(pos: u32, remaining: usize, top_k: u32) -> usize { let mut rows = remaining.min(INDEXED_PREFILL_CHUNK as usize); if pos < top_k { @@ -2574,10 +5247,19 @@ fn indexed_prefill_rows(pos: u32, remaining: usize, top_k: u32) -> usize { } fn full_indexer_layer(shape: super::super::Shape, layer: usize) -> bool { + if shape.model == ModelChoice::Glm53Flash { + return layer < (shape.layers - shape.nextn) as usize && !glm53_kda_layer(shape, layer); + } layer < (shape.layers - shape.nextn) as usize && (layer < shape.leading_dense as usize || (layer >= 6 && (layer - 6).is_multiple_of(4))) } +fn glm53_kda_layer(shape: super::super::Shape, layer: usize) -> bool { + shape.model == ModelChoice::Glm53Flash + && layer + (shape.nextn as usize) < shape.layers as usize + && layer % 4 != 3 +} + fn expert_layout(weight: Weight, experts: u64) -> (u64, u64) { let expert = weight.bytes / experts; (expert, expert / weight.dims[1]) @@ -2962,6 +5644,16 @@ fn install_glm_model_spans( ) } +fn glm53_kda_state_bytes(shape: super::super::Shape) -> u64 { + let layers = (0..(shape.layers - shape.nextn) as usize) + .filter(|layer| glm53_kda_layer(shape, *layer)) + .count() as u64; + layers + * (3 * (GLM53_KDA_CONV - 1) * GLM53_KDA_HEADS * GLM53_KDA_HEAD_DIM + + GLM53_KDA_HEADS * GLM53_KDA_HEAD_DIM * GLM53_KDA_HEAD_DIM) + * 4 +} + fn admission_bytes( model: &Model, weights: &GlmWeights, @@ -2970,17 +5662,34 @@ fn admission_bytes( ) -> Result { let shape = model.shape; let normal_layers = u64::from(shape.layers - shape.nextn); + let kda_layers = (0..normal_layers as usize) + .filter(|layer| glm53_kda_layer(shape, *layer)) + .count() as u64; + let dsa_layers = normal_layers - kda_layers; let indexer_layers = (0..normal_layers as usize) .filter(|layer| full_indexer_layer(shape, *layer)) .count() as u64; - let per_layer = u64::from(context) - .checked_mul((shape.kv_lora + shape.rot) * 2) + let per_dsa = u64::from(context) + .checked_mul(shape.kv_lora * 2 + shape.rot * 2) .ok_or("GLM compact-cache size overflow")?; - let kv = normal_layers - .checked_mul(per_layer) + let index_rows = if shape.model == ModelChoice::Glm53Flash { + u64::from(context).div_ceil(GLM53_INDEX_POOL) + } else { + u64::from(context) + }; + let index_tails = if shape.model == ModelChoice::Glm53Flash { + indexer_layers * 2 * GLM53_INDEX_POOL * shape.indexer_head_dim * 4 + } else { + 0 + }; + let kda_state = glm53_kda_state_bytes(shape); + let kv = dsa_layers + .checked_mul(per_dsa) .and_then(|bytes| { - bytes.checked_add(indexer_layers * u64::from(context) * shape.indexer_head_dim * 2) + bytes.checked_add(indexer_layers * index_rows * shape.indexer_head_dim * 2) }) + .and_then(|bytes| bytes.checked_add(index_tails)) + .and_then(|bytes| bytes.checked_add(kda_state)) .ok_or("GLM compact-cache size overflow")?; let resident = if let Some(plan) = streaming { let off_slab = weights @@ -3003,6 +5712,12 @@ fn admission_bytes( } else { model.main.len() - model.main.data_offset() }; + let resident = resident.saturating_add( + model + .vision + .as_ref() + .map_or(0, |vision| vision.len() - vision.data_offset()), + ); let cache = streaming.map_or(0, |plan| plan.planned_expert_bytes); let scratch = glm_batch_scratch_bytes(shape, context.min(INDEXED_PREFILL_CHUNK), context)? .max(512 * 1024 * 1024); @@ -3019,15 +5734,23 @@ fn glm_batch_scratch_bytes( context: u32, ) -> Result { let rows = u64::from(rows); - let score_rows = (INDEXED_PREFILL_SCORE_BYTES / (u64::from(context) * 4)) - .max(1) - .min(rows); + let score_columns = if shape.model == ModelChoice::Glm53Flash { + u64::from(context).div_ceil(GLM53_INDEX_POOL) + } else { + u64::from(context) + }; + let score_bytes = if shape.model == ModelChoice::Glm53Flash { + GLM53_INDEXED_PREFILL_SCORE_BYTES + } else { + INDEXED_PREFILL_SCORE_BYTES + }; + let score_rows = (score_bytes / (score_columns * 4)).max(1).min(rows); let q = shape.heads * shape.key_mla; let qk_low = shape.heads * shape.kv_lora; let heads = shape.heads * shape.value_mla; let hidden = shape.ff_dense.max(shape.ff_expert); let routed = shape.experts_used * shape.ff_expert; - let row_floats = 7 * shape.embd + let mut row_floats = 7 * shape.embd + 2 * shape.lora_q + q + shape.head_dim @@ -3043,12 +5766,37 @@ fn glm_batch_scratch_bytes( + shape.experts_used * shape.embd + 2 * shape.experts + shape.experts_used; + if shape.model == ModelChoice::Glm53Flash { + let projection = GLM53_KDA_HEADS * GLM53_KDA_HEAD_DIM; + row_floats += 4 * shape.hc * shape.embd + + 2 * shape.hc * (shape.hc + 2) + + 6 * projection + + 2 * GLM53_KDA_HEAD_DIM + + GLM53_KDA_HEADS + + shape.indexer_head_dim; + } rows.checked_mul(row_floats) .and_then(|floats| floats.checked_mul(4)) .and_then(|bytes| bytes.checked_add(rows * 4)) - .and_then(|bytes| bytes.checked_add(rows * shape.indexer_top_k * 4)) + .and_then(|bytes| { + let width = shape.indexer_top_k + + if shape.model == ModelChoice::Glm53Flash { + GLM53_INDEX_POOL - 1 + } else { + 0 + }; + bytes.checked_add(rows * width * 4) + }) + .and_then(|bytes| { + let pools = if shape.model == ModelChoice::Glm53Flash { + shape.indexer_top_k / GLM53_INDEX_POOL + } else { + 0 + }; + bytes.checked_add(rows * pools * 4) + }) .and_then(|bytes| bytes.checked_add(rows * shape.experts_used * 4)) - .and_then(|bytes| bytes.checked_add(score_rows * u64::from(context) * 4)) + .and_then(|bytes| bytes.checked_add(score_rows * score_columns * 4)) .ok_or_else(|| "GLM indexed-prefill scratch size overflow".into()) } @@ -3175,6 +5923,182 @@ fn project_rows( ) } +#[allow(clippy::too_many_arguments)] +fn glm53_project_rows( + out: &Buffer, + weight: Weight, + input: u64, + output: u64, + x: &Buffer, + rows: u32, + map: *const c_void, + size: u64, +) -> Result<(), String> { + if weight.kind == BF16 { + call( + unsafe { + ds4_gpu_glm53_matmul_bf16( + out.raw(), + map, + size, + weight.offset, + input as u32, + output as u32, + x.raw(), + rows, + ) + }, + "projecting GLM 5.3 BF16 activations", + ) + } else { + project_rows(out, weight, input, output, x, rows, map, size) + } +} + +#[allow(clippy::too_many_arguments)] +fn glm53_fused_gate_up_rows( + gate: &Buffer, + up: &Buffer, + mid: &Buffer, + gate_weight: Weight, + up_weight: Weight, + input: u64, + output: u64, + x: &Buffer, + rows: u32, + clamp: f32, + map: *const c_void, + size: u64, +) -> bool { + gate_weight.kind == Q8_0 + && up_weight.kind == Q8_0 + && unsafe { + if rows == 1 { + ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor( + gate.raw(), + up.raw(), + mid.raw(), + map, + size, + gate_weight.offset, + up_weight.offset, + input, + output, + x.raw(), + clamp, + ) + } else { + ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor( + gate.raw(), + up.raw(), + mid.raw(), + map, + size, + gate_weight.offset, + up_weight.offset, + input, + output, + x.raw(), + u64::from(rows), + clamp, + ) + } + } != 0 +} + +fn glm53_project( + out: &Buffer, + weight: Weight, + input: u64, + output: u64, + x: &Buffer, + map: *const c_void, + size: u64, +) -> Result<(), String> { + glm53_project_rows(out, weight, input, output, x, 1, map, size) +} + +#[allow(clippy::too_many_arguments)] +fn glm53_decode_project( + out: &Buffer, + weight: Weight, + input: u64, + output: u64, + x: &Buffer, + map: *const c_void, + size: u64, + streaming: bool, +) -> Result<(), String> { + if weight.kind == BF16 { + glm53_project(out, weight, input, output, x, map, size) + } else { + project(out, weight, input, output, x, map, size, streaming) + } +} + +#[allow(clippy::too_many_arguments)] +fn glm53_hc_pre( + collapsed: &Buffer, + normalized: &Buffer, + flat: &Buffer, + mix: &Buffer, + split: &Buffer, + residual: &Buffer, + function: Weight, + scale: Weight, + base: Weight, + norm_weight: Weight, + shape: super::super::Shape, + rows: u32, + map: *const c_void, + size: u64, +) -> Result<(), String> { + let hc_dim = shape.hc * shape.embd; + let hc_mix = shape.hc * (shape.hc + 2); + call( + unsafe { + ds4_gpu_rms_norm_plain_rows_tensor( + flat.raw(), + residual.raw(), + hc_dim as u32, + rows, + shape.rms_epsilon, + ) + }, + "normalizing GLM 5.3 hyperconnections", + )?; + glm53_project_rows(mix, function, hc_dim, hc_mix, flat, rows, map, size)?; + call( + unsafe { + ds4_gpu_hc_split_weighted_sum_tensor( + collapsed.raw(), + split.raw(), + mix.raw(), + residual.raw(), + map, + size, + scale.offset, + base.offset, + shape.embd as u32, + shape.hc as u32, + shape.hc_sinkhorn as u32, + shape.hc_epsilon, + ) + }, + "mixing GLM 5.3 hyperconnections", + )?; + norm_rows( + normalized, + collapsed, + norm_weight, + shape.embd as u32, + rows, + shape.rms_epsilon, + map, + size, + ) +} + #[allow(clippy::too_many_arguments)] fn f32_project( out: &Buffer, @@ -3233,12 +6157,13 @@ fn f32_project_rows( #[cfg(test)] mod tests { use super::{ - GlmExecutor, argmax, dynamic_expert_budget, full_indexer_layer, indexed_prefill_rows, + GlmExecutor, argmax, dynamic_expert_budget, full_indexer_layer, glm_prefill_flush_layers, + glm53_dense_compact_prefill, glm53_dense_limit, indexed_prefill_rows, live_prefix_rewind_target, streaming_token_prefill_eligible, }; use crate::engine::{GLM, Model, ReasoningMode}; use crate::model::ModelChoice; - use crate::settings::EngineSsdSettings; + use crate::settings::{EngineSsdSettings, EngineSteeringSettings}; fn installed_glm_path() -> std::path::PathBuf { crate::model::engine_artifacts(ModelChoice::Glm52, false, &crate::app::models_path()).model @@ -3262,8 +6187,417 @@ mod tests { assert!(!streaming_token_prefill_eligible(8192, 8180, 13, 64)); assert!(!streaming_token_prefill_eligible(65_536, 4096, 1, 64)); assert_eq!(indexed_prefill_rows(0, 5000, 2048), 2048); - assert_eq!(indexed_prefill_rows(2048, 5000, 2048), 4096); + assert_eq!(indexed_prefill_rows(2048, 5000, 2048), 2048); assert_eq!(indexed_prefill_rows(0, 17, 2048), 17); + assert_eq!(glm53_dense_limit(32_768, false), 4096); + assert_eq!(glm53_dense_limit(32_768, true), 8192); + assert_eq!(glm53_dense_limit(65_536, true), 4096); + assert!(!glm53_dense_compact_prefill(23)); + assert!(glm53_dense_compact_prefill(24)); + assert!(!glm_prefill_flush_layers(8, true)); + assert!(glm_prefill_flush_layers(9, true)); + assert!(!glm_prefill_flush_layers(2048, false)); + } + + #[test] + #[ignore = "requires the installed GLM 5.3 Flash Q2 checkpoint and Apple Metal"] + fn glm53_m5_text_performance_gate() { + use super::{INDEXED_PREFILL_CHUNK, glm_batch_scratch_bytes}; + use crate::engine::ChatTurn; + use crate::engine::metal::configure_sources; + use crate::engine::metal::gpu::ds4_gpu_print_memory_report; + use crate::settings::{EngineSpeculativeSettings, ReasoningMode}; + use sha2::{Digest, Sha256}; + use std::sync::atomic::AtomicBool; + use std::time::Instant; + + configure_sources().unwrap(); + let path = std::env::var_os("DS4SERVER_GLM53_MODEL") + .map(std::path::PathBuf::from) + .unwrap_or_else(|| { + crate::model::engine_artifacts( + ModelChoice::Glm53Flash, + false, + &crate::app::models_path(), + ) + .model + }); + let prompt_content = std::env::var_os("DS4SERVER_GLM53_PROMPT_FILE") + .map(|path| std::fs::read_to_string(path).unwrap()) + .unwrap_or_else(|| { + format!( + "Summarize the following numbered observations, preserving every number:\n{}", + (0..1_024) + .map(|index| format!( + "Observation {index}: parity requires identical work.\n" + )) + .collect::() + ) + }); + let context = std::env::var("DS4SERVER_GLM53_CONTEXT") + .ok() + .map_or(32_768, |value| value.parse().unwrap()); + let frontier = std::env::var("DS4SERVER_GLM53_FRONTIER") + .ok() + .map_or(8_192, |value| value.parse::().unwrap()); + let measured = std::env::var("DS4SERVER_GLM53_MEASURED") + .ok() + .map_or(64, |value| value.parse::().unwrap()); + let mtp = std::env::var_os("DS4SERVER_GLM53_MTP").is_some(); + let reference_prefill = std::env::var("DS4_REFERENCE_GLM53_PREFILL_TPS") + .ok() + .map(|value| value.parse::().unwrap()); + let reference_decode = std::env::var("DS4_REFERENCE_GLM53_DECODE_TPS") + .ok() + .map(|value| value.parse::().unwrap()); + let reference_hash = std::env::var("DS4_REFERENCE_GLM53_TOKEN_SHA256").ok(); + let runs = if reference_prefill.is_some() || reference_decode.is_some() { + 3 + } else { + 1 + }; + let mut results = Vec::with_capacity(runs); + for _ in 0..runs { + let model = Model::open_main(&path, ModelChoice::Glm53Flash).unwrap(); + let mut prompt = model.render_conversation( + "", + &[ChatTurn { + user: true, + tool: false, + system: false, + skip_previous_eos: false, + reasoning: None, + reasoning_complete: true, + content: prompt_content.clone(), + }], + ReasoningMode::Direct, + ); + assert!(prompt.len() >= frontier); + prompt.truncate(frontier); + assert!(prompt.len() < context as usize); + let eos = model.eos_token(); + if std::env::var_os("DS4SERVER_GLM53_MEMORY_REPORT").is_some() { + eprintln!( + "DS4SERVER_GLM53_BATCH_SCRATCH bytes={}", + glm_batch_scratch_bytes( + model.shape, + context.min(INDEXED_PREFILL_CHUNK), + context, + ) + .unwrap() + ); + } + let mut executor = GlmExecutor::open_profile( + model, + context, + false, + EngineSsdSettings { + enabled: false, + cold: false, + cache_experts: 0, + cache_bytes: 0, + full_layers: 0, + full_layers_set: false, + preload_experts: 0, + }, + EngineSpeculativeSettings { + glm_mtp: mtp, + glm_mtp_timing: std::env::var_os("DS4SERVER_GLM53_MTP_TIMING").is_some(), + dspark: false, + dspark_confidence_threshold: 0.9, + dspark_confidence_threshold_set: false, + dspark_strict: false, + dspark_exact_sampling: false, + }, + EngineSteeringSettings { + file: None, + ffn_scale: 0.0, + attention_scale: 0.0, + }, + None, + ) + .unwrap(); + if std::env::var_os("DS4SERVER_GLM53_MEMORY_REPORT").is_some() { + unsafe { ds4_gpu_print_memory_report(c"GLM 5.3 performance workspace".as_ptr()) }; + } + let prefill_started = Instant::now(); + assert_eq!(executor.prefill(&prompt, |_| true).unwrap(), prompt.len()); + let prefill_tps = prompt.len() as f64 / prefill_started.elapsed().as_secs_f64(); + let cancelled = AtomicBool::new(false); + let started = Instant::now(); + let mut generated = Vec::new(); + while generated.len() < measured as usize { + let token = argmax(executor.logits()); + assert_ne!(token, eos, "benchmark reached EOS before {measured} tokens"); + let cycle = executor + .eval_speculative_greedy(token, measured - generated.len() as u32, &cancelled) + .unwrap(); + generated.extend(cycle); + } + let decode_tps = measured as f64 / started.elapsed().as_secs_f64(); + let hash = Sha256::digest( + generated + .iter() + .flat_map(|token| token.to_le_bytes()) + .collect::>(), + ); + let hash = hash + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + eprintln!( + "DS4SERVER_GLM53_TEXT_PERF mtp={mtp} context={context} prompt={} prefill_tps={prefill_tps:.6} measured={measured} decode_tps={decode_tps:.6} token_sha256={hash}", + prompt.len(), + ); + if let Some(expected) = &reference_hash { + assert_eq!( + &hash, expected, + "DS4 and DS4Server generated different tokens" + ); + } + results.push((prefill_tps, decode_tps)); + if std::env::var_os("DS4SERVER_GLM53_MEMORY_REPORT").is_some() { + drop(executor); + unsafe { ds4_gpu_print_memory_report(c"after GLM 5.3 performance run".as_ptr()) }; + } + } + results.sort_by(|left, right| left.0.total_cmp(&right.0)); + let prefill = results[results.len() / 2].0; + results.sort_by(|left, right| left.1.total_cmp(&right.1)); + let decode = results[results.len() / 2].1; + if let Some(reference) = reference_prefill { + assert!( + prefill >= reference * 0.98, + "DS4Server prefill {prefill:.3} t/s is more than 2% below DS4 {reference:.3} t/s" + ); + } + if let Some(reference) = reference_decode { + assert!( + decode >= reference * 0.98, + "DS4Server decode {decode:.3} t/s is more than 2% below DS4 {reference:.3} t/s" + ); + } + } + + #[test] + #[ignore = "requires GLM 5.3 Q2, its vision sidecar, a fixed image, and Apple Metal"] + fn glm53_m5_vision_performance_and_oracle_gate() { + use crate::engine::metal::configure_sources; + use crate::engine::validation::validate_vision_artifact; + use crate::engine::{ChatTurn, VISION_TOKEN_END, VISION_TOKEN_START}; + use crate::settings::{EngineSpeculativeSettings, ReasoningMode}; + use sha2::{Digest, Sha256}; + use std::time::Instant; + + configure_sources().unwrap(); + let artifacts = crate::model::engine_artifacts( + ModelChoice::Glm53Flash, + false, + &crate::app::models_path(), + ); + let model_path = std::env::var_os("DS4SERVER_GLM53_MODEL") + .map(std::path::PathBuf::from) + .unwrap_or(artifacts.model); + let vision_path = std::env::var_os("DS4SERVER_GLM53_VISION") + .map(std::path::PathBuf::from) + .or(artifacts.vision) + .expect("set DS4SERVER_GLM53_VISION to the installed vision sidecar"); + let image_path = std::env::var_os("DS4SERVER_GLM53_IMAGE") + .map(std::path::PathBuf::from) + .expect("set DS4SERVER_GLM53_IMAGE to the fixed PNG/JPEG fixture"); + validate_vision_artifact(&vision_path).unwrap(); + let mut model = Model::open_main(&model_path, ModelChoice::Glm53Flash).unwrap(); + let eos = model.eos_token(); + model.vision = Some(crate::engine::gguf::Gguf::open(&vision_path).unwrap()); + let mut executor = GlmExecutor::open_profile( + model, + 4_096, + true, + EngineSsdSettings { + enabled: false, + cold: false, + cache_experts: 0, + cache_bytes: 0, + full_layers: 0, + full_layers_set: false, + preload_experts: 0, + }, + EngineSpeculativeSettings { + glm_mtp: false, + glm_mtp_timing: false, + dspark: false, + dspark_confidence_threshold: 0.9, + dspark_confidence_threshold_set: false, + dspark_strict: false, + dspark_exact_sampling: false, + }, + EngineSteeringSettings { + file: None, + ffn_scale: 0.0, + attention_scale: 0.0, + }, + None, + ) + .unwrap(); + let encoded = std::fs::read(image_path).unwrap(); + let mut runs = Vec::new(); + let mut embedding = None; + for _ in 0..3 { + let started = Instant::now(); + embedding = Some(executor.encode_vision(&encoded).unwrap()); + runs.push(started.elapsed().as_secs_f64()); + } + runs.sort_by(f64::total_cmp); + let vision_seconds = runs[1]; + let embedding = embedding.unwrap(); + assert!(embedding.values.iter().all(|value| value.is_finite())); + if let Some(reference_path) = std::env::var_os("DS4_REFERENCE_GLM53_VISION_EMBEDDING") { + let reference = std::fs::read(reference_path).unwrap(); + assert_eq!(reference.len(), embedding.values.len() * 4); + let mut dot = 0.0_f64; + let mut left = 0.0_f64; + let mut right = 0.0_f64; + let mut absolute = 0.0_f64; + let mut maximum = 0.0_f64; + for (expected, actual) in reference.chunks_exact(4).zip(&embedding.values) { + let expected = f32::from_le_bytes(expected.try_into().unwrap()) as f64; + let actual = f64::from(*actual); + dot += expected * actual; + left += expected * expected; + right += actual * actual; + let delta = (expected - actual).abs(); + absolute += delta; + maximum = maximum.max(delta); + } + let cosine = dot / (left.sqrt() * right.sqrt()); + let mean = absolute / embedding.values.len() as f64; + eprintln!( + "DS4SERVER_GLM53_VISION_ORACLE cosine={cosine:.9} mae={mean:.9} max={maximum:.9}" + ); + assert!(cosine >= 0.995 && mean <= 0.001 && maximum <= 0.06); + } + if let Ok(reference) = std::env::var("DS4_REFERENCE_GLM53_VISION_SECONDS") { + let reference = reference.parse::().unwrap(); + assert!( + vision_seconds <= reference * 1.02, + "DS4Server vision encode {vision_seconds:.3}s is more than 2% slower than DS4 {reference:.3}s" + ); + } + + let content = format!( + "{VISION_TOKEN_START}{}{VISION_TOKEN_END}\nDescribe the image briefly and state its dominant colors.", + embedding.tokens + ); + let prompt = executor.model().render_conversation( + "", + &[ChatTurn { + user: true, + tool: false, + system: false, + skip_previous_eos: false, + reasoning: None, + reasoning_complete: true, + content, + }], + ReasoningMode::Direct, + ); + let start = prompt + .iter() + .position(|token| *token == crate::engine::VISION_START_TOKEN) + .unwrap() as u32 + + 1; + let embedding_shape = ( + embedding.tokens, + embedding.width, + embedding.height, + embedding.content_width, + embedding.content_height, + ); + let embedding_values = embedding.values.clone(); + executor + .set_vision_overlays(vec![(start, embedding)]) + .unwrap(); + let prefill_started = Instant::now(); + assert_eq!(executor.prefill(&prompt, |_| true).unwrap(), prompt.len()); + let prefill_tps = prompt.len() as f64 / prefill_started.elapsed().as_secs_f64(); + let image_logits = executor.logits().to_vec(); + let make_embedding = |values| crate::engine::metal::vision::VisionEmbedding { + values, + tokens: embedding_shape.0, + width: embedding_shape.1, + height: embedding_shape.2, + content_width: embedding_shape.3, + content_height: embedding_shape.4, + }; + executor.reset().unwrap(); + executor + .set_vision_overlays(vec![( + start, + make_embedding(vec![0.0; embedding_values.len()]), + )]) + .unwrap(); + assert_eq!(executor.prefill(&prompt, |_| true).unwrap(), prompt.len()); + let zero_delta = image_logits + .iter() + .zip(executor.logits()) + .map(|(left, right)| (left - right).abs()) + .fold(0.0_f32, f32::max); + assert!( + zero_delta > 1.0e-4, + "visual embeddings did not affect logits" + ); + executor.reset().unwrap(); + executor + .set_vision_overlays(vec![(start, make_embedding(embedding_values))]) + .unwrap(); + assert_eq!(executor.prefill(&prompt, |_| true).unwrap(), prompt.len()); + let replay_delta = image_logits + .iter() + .zip(executor.logits()) + .map(|(left, right)| (left - right).abs()) + .fold(0.0_f32, f32::max); + assert!( + replay_delta <= 1.0e-6, + "restored image-conditioned logits changed by {replay_delta}" + ); + let started = Instant::now(); + let mut generated = Vec::new(); + for _ in 0..32 { + let token = argmax(executor.logits()); + assert_ne!(token, eos, "vision benchmark reached EOS before 32 tokens"); + generated.push(token); + executor.eval(token).unwrap(); + } + let decode_tps = 32.0 / started.elapsed().as_secs_f64(); + let hash = Sha256::digest( + generated + .iter() + .flat_map(|token| token.to_le_bytes()) + .collect::>(), + ); + let hash = hash + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + eprintln!( + "DS4SERVER_GLM53_VISION_PERF encode_seconds={vision_seconds:.6} image_tokens={} prompt={} prefill_tps={prefill_tps:.6} decode_tps={decode_tps:.6} token_sha256={hash}", + executor.vision_overlays[0].tokens, + prompt.len(), + ); + if let Ok(expected) = std::env::var("DS4_REFERENCE_GLM53_VISION_TOKEN_SHA256") { + assert_eq!( + hash, expected, + "DS4 and DS4Server vision output tokens differ" + ); + } + if let Ok(reference) = std::env::var("DS4_REFERENCE_GLM53_VISION_PREFILL_TPS") { + let reference = reference.parse::().unwrap(); + assert!(prefill_tps >= reference * 0.98); + } + if let Ok(reference) = std::env::var("DS4_REFERENCE_GLM53_VISION_DECODE_TPS") { + let reference = reference.parse::().unwrap(); + assert!(decode_tps >= reference * 0.98); + } } #[test] @@ -3438,6 +6772,11 @@ mod tests { dspark_strict: false, dspark_exact_sampling: false, }, + EngineSteeringSettings { + file: None, + ffn_scale: 0.0, + attention_scale: 0.0, + }, None, ) .unwrap(); diff --git a/src/engine/metal/gpu.rs b/src/engine/metal/gpu.rs index e48123e..6892558 100644 --- a/src/engine/metal/gpu.rs +++ b/src/engine/metal/gpu.rs @@ -5,6 +5,42 @@ pub(super) struct GpuTensor { _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)] pub(super) struct StreamExpertTable { pub(super) model_map: *const c_void, @@ -93,6 +129,12 @@ unsafe extern "C" { count: u32, ) -> 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_m5_apple_silicon() -> i32; #[cfg(test)] @@ -227,6 +269,65 @@ unsafe extern "C" { embd: u32, hc: u32, ) -> 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( out: *mut GpuTensor, map: *const c_void, @@ -417,6 +518,35 @@ unsafe extern "C" { beta_slow: f32, cache_f16: bool, ) -> 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( selected: *mut GpuTensor, count: u32, @@ -467,6 +597,66 @@ unsafe extern "C" { scale: f32, cache_f16: bool, ) -> 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( out: *mut GpuTensor, q: *const GpuTensor, @@ -554,6 +744,19 @@ unsafe extern "C" { beta_fast: f32, beta_slow: f32, ) -> 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( out: *mut GpuTensor, q: *const GpuTensor, @@ -777,6 +980,20 @@ unsafe extern "C" { eps: f32, norm_eps: f32, ) -> 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_tensor( out: *mut GpuTensor, @@ -1359,6 +1576,15 @@ unsafe extern "C" { embd: u32, hc: u32, ) -> 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( out: *mut GpuTensor, block: *const GpuTensor, @@ -1463,6 +1689,44 @@ unsafe extern "C" { x: *const GpuTensor, clamp: f32, ) -> 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( router_logits: *mut GpuTensor, gate: *mut GpuTensor, @@ -1598,6 +1862,21 @@ impl Context { 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) }; let model_file = File::open(model.main.path()).map_err(|error| { unsafe { ds4_gpu_cleanup() }; diff --git a/src/engine/metal/vision.rs b/src/engine/metal/vision.rs new file mode 100644 index 0000000..45525ef --- /dev/null +++ b/src/engine/metal/vision.rs @@ -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, + 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 { + 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 { + 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, + content_width: u32, + content_height: u32, + grid_width: u32, + grid_height: u32, +} + +fn preprocess(rgb: &[u8], width: u32, height: u32) -> Result { + 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))); + } +} diff --git a/src/engine/tokenizer.rs b/src/engine/tokenizer.rs index a17efa3..7ddf705 100644 --- a/src/engine/tokenizer.rs +++ b/src/engine/tokenizer.rs @@ -1,5 +1,8 @@ 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 std::collections::HashMap; @@ -160,6 +163,22 @@ impl Tokenizer { let mut span = 0; let mut position = 0; 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::() { + 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 .rendered_specials .iter() diff --git a/src/engine/validation.rs b/src/engine/validation.rs index 123522b..ad29d3a 100644 --- a/src/engine/validation.rs +++ b/src/engine/validation.rs @@ -9,7 +9,7 @@ pub(crate) fn validate_model_artifact( let model = Gguf::open(path)?; let shape = match expected { 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")); } }; @@ -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)] pub(super) enum SupportKind { DSpark, @@ -116,14 +189,11 @@ pub(super) fn validate_support(model: &Gguf, shape: &Shape) -> Result Result { - let family = if model.bytes("general.architecture").ok() == Some(b"glm-dsa") { - ModelFamily::Glm - } else { - ModelFamily::DeepSeek - }; - let shape = match family { - ModelFamily::Glm => GLM, - ModelFamily::DeepSeek => match (model.u32("deepseek4.block_count")?, expected) { + let architecture = model.bytes("general.architecture")?; + let shape = match architecture { + b"glm-dsa" => GLM, + b"glm5-next" => GLM53_FLASH, + _ => match (model.u32("deepseek4.block_count")?, expected) { (43, ModelChoice::DeepSeekV4Flash0731) => FLASH_0731, (43, _) => FLASH_0731, (61, _) => PRO, @@ -143,6 +213,9 @@ pub(super) fn validate_main(model: &Gguf, expected: ModelChoice) -> Result Result<(), String> { + if shape.model == ModelChoice::Glm53Flash { + return validate_glm53_metadata(model, shape); + } let prefix = if shape.family == ModelFamily::Glm { "glm-dsa" } else { @@ -266,13 +339,102 @@ fn validate_metadata(model: &Gguf, shape: &Shape) -> Result<(), String> { 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> { match shape.family { 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), } } +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> { let hc_dim = shape.embd * 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 } 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(); } } + + #[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(); + } + } } diff --git a/src/metrics.rs b/src/metrics.rs index d4c0d04..c2554c3 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -847,6 +847,7 @@ fn model_code(model: ModelChoice) -> u8 { ModelChoice::DeepSeekV4Pro => 2, ModelChoice::Glm52 => 3, ModelChoice::DeepSeekV4Flash0731 => 4, + ModelChoice::Glm53Flash => 5, } } @@ -856,6 +857,7 @@ fn model_name(value: u8) -> &'static str { 2 => "DeepSeek V4 Pro", 3 => "GLM 5.2", 4 => "DeepSeek V4 Flash 0731", + 5 => "GLM 5.3 Flash Q2", _ => "No model loaded", } } diff --git a/src/model.rs b/src/model.rs index bc57eb2..55173d2 100644 --- a/src/model.rs +++ b/src/model.rs @@ -9,20 +9,24 @@ use std::fmt; use std::fs; use std::path::{Path, PathBuf}; -pub(crate) const MODEL_CHOICES: [ModelChoice; 3] = [ +pub(crate) const MODEL_CHOICES: [ModelChoice; 4] = [ ModelChoice::DeepSeekV4Flash0731, ModelChoice::DeepSeekV4Pro, ModelChoice::Glm52, + ModelChoice::Glm53Flash, ]; -pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 4] = [ +pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 6] = [ ManagedArtifactId::DeepSeekV4Flash0731, ManagedArtifactId::DeepSeekV4Flash0731Dspark, ManagedArtifactId::DeepSeekV4Pro, ManagedArtifactId::Glm52, + ManagedArtifactId::Glm53Flash, + ManagedArtifactId::Glm53FlashVision, ]; const DEEPSEEK_REPOSITORY: &str = "antirez/deepseek-v4-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 { label: "DeepSeek V4 Flash 0731 model", @@ -31,6 +35,7 @@ const FLASH_0731: Artifact = Artifact { size: 86_720_111_488, sha256: "ca22ae2f838e14077c22bc1c1417b71b45b5e5a3687bd96c2ac6e17fdb6261c0", support: Some(false), + vision: false, }; const FLASH_0731_DSPARK: Artifact = Artifact { label: "DeepSeek V4 Flash 0731 DSpark support", @@ -39,6 +44,7 @@ const FLASH_0731_DSPARK: Artifact = Artifact { size: 5_989_114_272, sha256: "7e319924541db3f7a163ed7e11d7532a70d48228ab59d36cb81e1d4511885360", support: Some(true), + vision: false, }; const PRO: Artifact = Artifact { label: "DeepSeek V4 Pro 0813 model", @@ -47,6 +53,7 @@ const PRO: Artifact = Artifact { size: 464_627_334_560, sha256: "c4d997ab9894b6c78b759f7869fe1726b6314b6515f6ff82607df3797c5eb193", support: Some(false), + vision: false, }; const GLM: Artifact = Artifact { label: "GLM 5.2 model", @@ -55,6 +62,25 @@ const GLM: Artifact = Artifact { size: 211_075_856_448, sha256: "a49de64c5020432bdae23de36a423a9660a5621bc0db8d12b66bd8814b07fea0", 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)] @@ -66,6 +92,8 @@ pub(crate) enum ModelChoice { DeepSeekV4Pro, #[serde(rename = "glm-5.2")] Glm52, + #[serde(rename = "glm-5.3-flash")] + Glm53Flash, } impl ModelChoice { @@ -74,6 +102,7 @@ impl ModelChoice { Self::DeepSeekV4Flash0731 => "deepseek-v4-flash-0731", Self::DeepSeekV4Pro => "deepseek-v4-pro", Self::Glm52 => "glm-5.2", + Self::Glm53Flash => "glm-5.3-flash", } } @@ -85,26 +114,40 @@ impl ModelChoice { 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 { match self { Self::DeepSeekV4Flash0731 => &FLASH_0731, Self::DeepSeekV4Pro => &PRO, Self::Glm52 => &GLM, + Self::Glm53Flash => &GLM53_FLASH, } } fn dspark_artifact(self) -> Option<&'static Artifact> { match self { 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)] fn artifacts(self, dspark_enabled: bool) -> impl Iterator { [ Some(self.main_artifact()), dspark_enabled.then(|| self.dspark_artifact()).flatten(), + self.vision_artifact(), ] .into_iter() .flatten() @@ -115,6 +158,7 @@ impl ModelChoice { pub(crate) struct EngineArtifacts { pub(crate) model: PathBuf, pub(crate) support: Option, + pub(crate) vision: Option, } pub(crate) fn engine_artifacts( @@ -131,6 +175,11 @@ pub(crate) fn engine_artifacts( } else { 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", 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)] @@ -167,6 +223,8 @@ pub(crate) enum ManagedArtifactId { DeepSeekV4Flash0731Dspark, DeepSeekV4Pro, Glm52, + Glm53Flash, + Glm53FlashVision, } impl ManagedArtifactId { @@ -177,6 +235,7 @@ impl ManagedArtifactId { } Self::DeepSeekV4Pro => ModelChoice::DeepSeekV4Pro, Self::Glm52 => ModelChoice::Glm52, + Self::Glm53Flash | Self::Glm53FlashVision => ModelChoice::Glm53Flash, } } @@ -186,6 +245,8 @@ impl ManagedArtifactId { Self::DeepSeekV4Flash0731Dspark => &FLASH_0731_DSPARK, Self::DeepSeekV4Pro => &PRO, 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::DeepSeekV4Pro => "DeepSeek V4 Pro 0813", Self::Glm52 => "GLM 5.2", + Self::Glm53Flash => "GLM 5.3 Flash Q2", }) } } @@ -291,6 +353,7 @@ struct Artifact { size: u64, sha256: &'static str, support: Option, + vision: bool, } impl Artifact { diff --git a/src/model/transfer.rs b/src/model/transfer.rs index 86ece2d..dc81cec 100644 --- a/src/model/transfer.rs +++ b/src/model/transfer.rs @@ -205,7 +205,9 @@ fn verify( 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)?; } Ok(DownloadOutcome::Complete) @@ -330,8 +332,8 @@ mod tests { Some(ModelChoice::DeepSeekV4Flash0731) ); assert!(ModelChoice::from_id("unknown").is_none()); - assert_eq!(MODEL_CHOICES.len(), 3); - assert_eq!(MANAGED_ARTIFACTS.len(), 4); + assert_eq!(MODEL_CHOICES.len(), 4); + assert_eq!(MANAGED_ARTIFACTS.len(), 6); assert_eq!(ModelChoice::Glm52.main_artifact().size, 211_075_856_448); assert_eq!( ModelChoice::DeepSeekV4Flash0731.main_artifact().size, @@ -339,6 +341,7 @@ mod tests { ); assert_eq!(ModelChoice::DeepSeekV4Flash0731.artifacts(true).count(), 2); assert_eq!(ModelChoice::Glm52.artifacts(true).count(), 1); + assert_eq!(ModelChoice::Glm53Flash.artifacts(true).count(), 2); let id = SystemTime::now() .duration_since(UNIX_EPOCH) @@ -366,6 +369,7 @@ mod tests { size: 0, sha256: "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", support: None, + vision: false, }; let partial = empty.partial_path(ModelChoice::DeepSeekV4Flash0731, &models_path); fs::create_dir_all(partial.parent().unwrap()).unwrap(); @@ -393,6 +397,7 @@ mod tests { &EngineArtifacts { model: installed.clone(), support: None, + vision: None, }, ) .is_err() @@ -404,6 +409,7 @@ mod tests { &EngineArtifacts { model: installed, support: None, + vision: None, }, ) .is_err() @@ -428,6 +434,7 @@ mod tests { size: 3, sha256: "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad", support: None, + vision: false, }; let verified_bytes = AtomicU64::new(999); @@ -593,6 +600,7 @@ mod tests { size: 10, sha256: "unused", support: None, + vision: false, }; let partial = artifact.partial_path(ModelChoice::DeepSeekV4Flash0731, &directory); fs::create_dir_all(partial.parent().unwrap()).unwrap(); diff --git a/src/server.rs b/src/server.rs index 5ab98e3..e110667 100644 --- a/src/server.rs +++ b/src/server.rs @@ -554,6 +554,12 @@ fn model_alias(id: &str) -> Option { | "zai/glm-5.2" | "zai/glm-5.2-chat" | "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), } } @@ -900,6 +906,13 @@ mod tests { content_text(&json!([{"type": "text", "text": "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] diff --git a/src/server/request.rs b/src/server/request.rs index a2b2423..dfe783c 100644 --- a/src/server/request.rs +++ b/src/server/request.rs @@ -54,6 +54,23 @@ pub(super) fn anthropic_request(value: Value) -> Result 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( block .get("thinking") @@ -382,6 +399,24 @@ fn responses_content_text(value: &Value) -> Result { 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())), } } @@ -574,7 +609,12 @@ fn request_reasoning( || (explicit_thinking.is_none() && matches!( 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; diff --git a/src/server/tools.rs b/src/server/tools.rs index a5f4143..d73f1b3 100644 --- a/src/server/tools.rs +++ b/src/server/tools.rs @@ -365,6 +365,7 @@ pub(super) fn render_messages( let mut turns = Vec::::new(); for message in messages { + validate_inline_images(&message.content)?; let content = content_text(&message.content); match message.role.as_str() { "system" | "developer" => { @@ -426,6 +427,30 @@ pub(super) fn render_messages( 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( state: &State, messages: &[ApiMessage], @@ -779,8 +804,18 @@ pub(super) fn content_text(value: &Value) -> String { Value::Array(parts) => parts .iter() .filter_map(|part| match part { - Value::String(text) => Some(text.as_str()), - Value::Object(object) => object.get("text").and_then(Value::as_str), + Value::String(text) => Some(text.clone()), + 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, }) .collect(), diff --git a/src/settings.rs b/src/settings.rs index 10b5bc1..6a9696c 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -34,8 +34,8 @@ impl SpeculativePreferences { if self.glm_mtp_timing && !self.glm_mtp { return Err("GLM MTP timing requires GLM MTP.".into()); } - if model != ModelChoice::Glm52 && (self.glm_mtp || self.glm_mtp_timing) { - return Err("GLM MTP is available only for GLM 5.2.".into()); + if !model.supports_glm_mtp() && (self.glm_mtp || self.glm_mtp_timing) { + return Err("GLM MTP is available only for GLM models.".into()); } if self.dspark_enabled && !model.supports_dspark() { 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()); } - if self.full_layers.is_some_and(|layers| layers > 0) && model != ModelChoice::Glm52 { - return Err("Fully resident SSD layers are available only for GLM 5.2.".into()); + if self.full_layers.is_some_and(|layers| layers > 0) && !model.is_glm() { + return Err("Fully resident SSD layers are available only for GLM models.".into()); } if let Some(StreamingCacheBudget::Gib(gib)) = self.cache { validate_gib("SSD cache budget", gib)?; @@ -443,12 +443,12 @@ impl ExecutionPreferences { { 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) { - 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() { - return Err("GLM 5.2 selects its prefill chunk automatically.".into()); + return Err("GLM selects its prefill chunk automatically.".into()); } } Ok(()) @@ -569,7 +569,7 @@ impl GenerationPreferences { model: ModelChoice, kv_cache: KvCacheSettings, ) -> TurnSettings { - let glm = model == ModelChoice::Glm52; + let glm = model.is_glm(); TurnSettings { kv_cache, context_tokens: self.context_tokens,