From b117a8bcf214c4a3cdd57b2a96874b2bf31be0e7 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 20 Oct 2026 12:36:37 -0500 Subject: [PATCH 01/21] qwen27: serve qwen3_5_moe (Ornith-1.5-122B) on Metal: routed experts, 32-over-2 attention WIP: the text path matches an fp32-routing reference at every depth on the 4-bit checkpoint (KL 0.0023 at 48 layers; decode top-1 98-100%); host registration, DFlash2 and images follow. --- zig/build/qwen27.zig | 4 +- zig/kernels/metal/qwen27/attention_mpp.metal | 10 +- .../metal/qwen27/attention_prompt.metal | 4 +- zig/kernels/metal/qwen27/moe.metal | 140 +++++++++++++ zig/qwen27_core.zig | 2 + zig/src/core/lane_projection.zig | 2 + zig/src/families/qwen27/attention.zig | 96 +++++---- zig/src/families/qwen27/attention_plan.zig | 4 +- zig/src/families/qwen27/config.zig | 31 ++- zig/src/families/qwen27/decode_round.zig | 20 +- zig/src/families/qwen27/forward.zig | 9 + zig/src/families/qwen27/gpu_frame.zig | 19 +- zig/src/families/qwen27/gpu_weights.zig | 73 ++++++- zig/src/families/qwen27/model.zig | 15 +- zig/src/families/qwen27/moe_gpu.zig | 194 ++++++++++++++++++ zig/src/families/qwen27/widths.zig | 58 +++++- zig/tests/qwen27/run.zig | 18 ++ 17 files changed, 619 insertions(+), 80 deletions(-) create mode 100644 zig/kernels/metal/qwen27/moe.metal create mode 100644 zig/src/families/qwen27/moe_gpu.zig diff --git a/zig/build/qwen27.zig b/zig/build/qwen27.zig index 3a0b7ccab2..8e4dcd5d6c 100644 --- a/zig/build/qwen27.zig +++ b/zig/build/qwen27.zig @@ -16,7 +16,7 @@ pub const Shared = struct { }; /// The 27B's Metal sources, embedded by name (zig/kernels/metal/qwen27). -const parts = [_][]const u8{ "xsum", "norm_xs", "norm_input", "mlp_xs", "mlp_ane", "glue", "copy", "attention_mpp", "attention_io", "attention_prompt", "prompt_glue", "prompt_state" }; +const parts = [_][]const u8{ "xsum", "norm_xs", "norm_input", "mlp_xs", "mlp_ane", "glue", "copy", "moe", "attention_mpp", "attention_io", "attention_prompt", "prompt_glue", "prompt_state" }; /// The shared modules, imported into `engine` (whose `core` is itself) for the server and the release binary. pub fn engineModules(b: *std.Build, target: std.Build.ResolvedTarget, metal: *std.Build.Module, sources: *std.Build.Module, engine: *std.Build.Module) Shared { @@ -99,7 +99,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu const see = program(b, "tf-lanes-see-check", "Build synthetic shared-context checks, no GPU execution", "zig/tests/lanes_see.zig", target, false, &.{ .{ .name = "metal", .module = metal }, .{ .name = "shared_attention", .module = attention }, .{ .name = "shared_kv", .module = kv_layout } }); b.step("test-lanes-see", "Run small shared-context, mask, position and KV append device checks").dependOn(&b.addRunArtifact(see).step); - const qwen_core = b.createModule(.{ .root_source_file = b.path("zig/qwen27_core.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true, .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tokenizer", .module = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true }) } } }); + const qwen_core = b.createModule(.{ .root_source_file = b.path("zig/qwen27_core.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true, .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "kernel_sources", .module = sources }, .{ .name = "tokenizer", .module = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true }) } } }); const qwen27 = b.createModule(.{ .root_source_file = b.path("zig/src/families/qwen27/qwen27.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true, .imports = &.{ .{ .name = "core", .module = qwen_core }, .{ .name = "metal", .module = metal } } }); for ([_]struct { []const u8, *std.Build.Module }{ .{ "core", qwen_core }, .{ "core_lane", core_lane }, .{ "core_row", core_row }, .{ "bf16_topk", s.bf16_topk }, .{ "bf16_topk_gpu", s.bf16_topk_gpu }, .{ "tree_round", s.tree_round }, .{ "tree_round_gpu", s.tree_round_gpu }, .{ "tree_commit_gpu", s.tree_commit }, .{ "draft_ops", s.draft_ops }, .{ "kernel_sources", sources } }) |m| qwen_core.addImport(m[0], m[1]); for ([_]struct { []const u8, *std.Build.Module }{ .{ "core_lane", core_lane }, .{ "qwen_runtime_sources", s.runtime_sources }, .{ "core_delta", s.delta }, .{ "qwen_gdn_source", s.gdn_source } }) |m| qwen27.addImport(m[0], m[1]); diff --git a/zig/kernels/metal/qwen27/attention_mpp.metal b/zig/kernels/metal/qwen27/attention_mpp.metal index b81b68cba8..397505cbc3 100644 --- a/zig/kernels/metal/qwen27/attention_mpp.metal +++ b/zig/kernels/metal/qwen27/attention_mpp.metal @@ -3,7 +3,11 @@ #include using namespace metal; using namespace mpp::tensor_ops; -constant constexpr int G = 6, D = 256, TK = 64, MAXD = 128, MG = 8, MS = 12, NW = 4; // a row's key chunk is its node word 2 +#ifndef Q27_G +#define Q27_G 6 // query heads a KV head (the 27B's); the model's host prepends its own +#define Q27_KVH 4 +#endif +constant constexpr int G = Q27_G, KVH = Q27_KVH, D = 256, TK = 64, MAXD = 128, MG = 8, MS = 12, NW = 4; // a row's key chunk is its node word 2 struct Q27AttnReadPtrs { device const bfloat* keys; device const bfloat* values; }; // SG tiles a threadgroup, a simdgroup pair each: both run scores and softmax, one owns each half of the head dim. @@ -317,10 +321,10 @@ kernel void q27_attn_merge(device const float* POA [[buffer(0)]], device const f } const int h = hk * G + g; for (int i = 0; i < DP; i++) { - const int64_t index = ((int64_t)node * (G * 4) + h) * D + lane * DP + i; + const int64_t index = ((int64_t)node * (G * KVH) + h) * D + lane * DP + i; const bfloat v = static_cast(o[i] / l); if (gated) { - const float gate = float(QG[((int64_t)node * (G * 4) + h) * 2 * D + D + lane * DP + i]); + const float gate = float(QG[((int64_t)node * (G * KVH) + h) * 2 * D + D + lane * DP + i]); OUT[index] = static_cast(float(v) * q27_round(1.0f / (1.0f + exp(-gate)), 0)); } else OUT[index] = v; } diff --git a/zig/kernels/metal/qwen27/attention_prompt.metal b/zig/kernels/metal/qwen27/attention_prompt.metal index 8ba07e0501..e9b732d7ed 100644 --- a/zig/kernels/metal/qwen27/attention_prompt.metal +++ b/zig/kernels/metal/qwen27/attention_prompt.metal @@ -3,7 +3,7 @@ #include using namespace metal; using namespace mpp::tensor_ops; -constant constexpr int PG = 6, PD = 256, PH = 128, PCK = 512, PTK = 32; +constant constexpr int PG = Q27_G, PKVH = Q27_KVH, PD = 256, PH = 128, PCK = 512, PTK = 32; // after attention_mpp.metal's G struct Q27PromptArgs { int start, rows, chunks, rp, key_stride, value_stride; float scale2; }; typedef vec q27f8; typedef vec q27b8; @@ -166,5 +166,5 @@ kernel void q27_attn_prompt_merge(device const float* PO [[buffer(0)]], device c m = nm; } const int node = r / PG, head = hk * PG + r % PG; - for (int i = 0; i < DP; i++) OUT[((int64_t)node * (PG * 4) + head) * PD + lane * DP + i] = static_cast(o[i] / l); + for (int i = 0; i < DP; i++) OUT[((int64_t)node * (PG * PKVH) + head) * PD + lane * DP + i] = static_cast(o[i] / l); } diff --git a/zig/kernels/metal/qwen27/moe.metal b/zig/kernels/metal/qwen27/moe.metal new file mode 100644 index 0000000000..df69486477 --- /dev/null +++ b/zig/kernels/metal/qwen27/moe.metal @@ -0,0 +1,140 @@ +// The routed block: a row's softmax over its router logits, its top-k, the slots sorted for the tensor units' gathers, +// the experts' SwiGLU and the combine with the shared expert, all with the reference's bf16 boundaries. +#include +using namespace metal; + +struct RouteArgs { int experts, topk, ld, renorm; }; +struct SortArgs { int slots, experts; }; +struct GluArgs { int width; }; +struct CombineArgs { int hidden, topk, ld, experts; }; + +// A threadgroup a row: softmax (fp32, rounded to bf16 as the reference's) over LOGITS[row, 0..experts), the top-k by +// logit with the lower id on ties, in descending order; PICK and WEIGHT [rows, topk]. Column `experts` of the row is +// the shared expert's gate logit: its sigmoid, rounded to bf16, is left in column experts + 1 for the combine. +[[kernel]] void q27_route(device bfloat* LOGITS [[buffer(0)]], device uint* PICK [[buffer(1)]], + device float* WEIGHT [[buffer(2)]], constant RouteArgs& a [[buffer(3)]], + uint r [[threadgroup_position_in_grid]], uint t [[thread_position_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]], + uint nsg [[simdgroups_per_threadgroup]]) { + threadgroup float part[32]; + threadgroup float best_v[32]; + threadgroup int best_i[32]; + threadgroup int sel[16]; + threadgroup float w[16]; + device bfloat* z = LOGITS + size_t(r) * a.ld; + const bool mine = int(t) < a.experts; + const float x = mine ? float(z[t]) : -INFINITY; + // the row's max, then the sum of exp + float m = simd_max(x); + if (lane == 0) part[sg] = m; + threadgroup_barrier(mem_flags::mem_threadgroup); + m = part[0]; + for (uint g = 1; g < nsg; g++) m = max(m, part[g]); + threadgroup_barrier(mem_flags::mem_threadgroup); + const float ex = mine ? precise::exp(x - m) : 0.0f; + float s = simd_sum(ex); + if (lane == 0) part[sg] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + s = 0.0f; + for (uint g = 0; g < nsg; g++) s += part[g]; + const float p = float(bfloat(ex / s)); + float key = x; + for (int k = 0; k < a.topk; k++) { + float v = key; + int id = mine ? int(t) : 0x7fffffff; + for (ushort o = 16; o > 0; o >>= 1) { + const float ov = simd_shuffle_down(v, o); + const int oi = simd_shuffle_down(id, o); + if (ov > v || (ov == v && oi < id)) { + v = ov; + id = oi; + } + } + if (lane == 0) { + best_v[sg] = v; + best_i[sg] = id; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t == 0) { + float bv = best_v[0]; + int bi = best_i[0]; + for (uint g = 1; g < nsg; g++) + if (best_v[g] > bv || (best_v[g] == bv && best_i[g] < bi)) { + bv = best_v[g]; + bi = best_i[g]; + } + sel[k] = bi; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (int(t) == sel[k]) { + w[k] = p; + key = -INFINITY; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (t == 0) { + float sum = 0.0f; + for (int k = 0; k < a.topk; k++) sum += w[k]; + const float norm = float(bfloat(sum)); + for (int k = 0; k < a.topk; k++) { + PICK[size_t(r) * a.topk + k] = uint(sel[k]); + WEIGHT[size_t(r) * a.topk + k] = a.renorm ? float(bfloat(w[k] / norm)) : w[k]; + } + const float g = float(z[a.experts]); + z[a.experts + 1] = bfloat(1.0f / (1.0f + precise::exp(-g))); + } +} + +// A pass's routed slots sorted by expert (stable) for core/affine_mm's gathers: OFFS (each expert's first sorted slot, +// then the end), ORDER (sorted position to slot) and WS (the slot's weight); one thread. +[[kernel]] void q27_route_sort(const device uint* PICK [[buffer(0)]], const device float* WEIGHT [[buffer(1)]], + device int* OFFS [[buffer(2)]], device int* ORDER [[buffer(3)]], + device float* WS [[buffer(4)]], constant SortArgs& a [[buffer(5)]], + uint t [[thread_position_in_grid]]) { + if (t != 0) return; + for (int e = 0; e <= a.experts; e++) OFFS[e] = 0; + for (int i = 0; i < a.slots; i++) OFFS[PICK[i]]++; + int at = 0; + for (int e = 0; e < a.experts; e++) { + const int n = OFFS[e]; + OFFS[e] = at; + at += n; + } + for (int i = 0; i < a.slots; i++) { + const int at2 = OFFS[PICK[i]]++; + ORDER[at2] = i; + WS[at2] = WEIGHT[i]; + } + for (int e = a.experts; e > 0; e--) OFFS[e] = OFFS[e - 1]; + OFFS[0] = 0; +} + +// Sorted slot i's input row: X[ORDER[i] / topk] (a row for each of its picks). +[[kernel]] void q27_route_gather(const device bfloat* X [[buffer(0)]], const device int* ORDER [[buffer(1)]], + device bfloat* Y [[buffer(2)]], constant int2& wk [[buffer(3)]], + uint2 g [[thread_position_in_grid]]) { + Y[size_t(g.y) * wk.x + g.x] = X[size_t(ORDER[g.y] / wk.y) * wk.x + g.x]; +} + +// act = silu(gate) * up in fp32, rounded once, as the dense block's q27_mlp. +[[kernel]] void q27_expert_glu(const device bfloat* G [[buffer(0)]], const device bfloat* U [[buffer(1)]], + device bfloat* OUT [[buffer(2)]], constant GluArgs& a [[buffer(3)]], + uint2 g [[thread_position_in_grid]]) { + const size_t i = size_t(g.y) * a.width + g.x; + const float gate = float(G[i]); + OUT[i] = bfloat(gate / (1.0f + exp(-gate)) * float(U[i])); +} + +// OUT[r] (the shared expert's output) becomes bf16(bf16(sum_k bf16(Y[r*topk+k] * w_k)) + bf16(sigma_r * OUT[r])), +// the reference's roundings: each weighted expert row, their fp32 sum, the gated shared row, the sum of the two. +[[kernel]] void q27_expert_combine(const device bfloat* Y [[buffer(0)]], const device float* WEIGHT [[buffer(1)]], + const device bfloat* LOGITS [[buffer(2)]], device bfloat* OUT [[buffer(3)]], + constant CombineArgs& a [[buffer(4)]], uint2 g [[thread_position_in_grid]]) { + const size_t r = g.y, j = g.x; + float routed = 0.0f; + for (int k = 0; k < a.topk; k++) + routed += float(bfloat(float(Y[(r * a.topk + k) * a.hidden + j]) * WEIGHT[r * a.topk + k])); + const float sigma = float(LOGITS[r * a.ld + a.experts + 1]); + const float shared = float(bfloat(sigma * float(OUT[r * a.hidden + j]))); + OUT[r * a.hidden + j] = bfloat(float(bfloat(routed)) + shared); +} diff --git a/zig/qwen27_core.zig b/zig/qwen27_core.zig index ebd50da149..09dc047d96 100644 --- a/zig/qwen27_core.zig +++ b/zig/qwen27_core.zig @@ -16,6 +16,8 @@ pub const lane_projection = @import("core_lane"); pub const gpu_profile = @import("src/core/gpu_profile.zig"); pub const row_projection = @import("core_row"); +/// The tensor units' affine products by expert (the routed block's prompt path). +pub const affine_mm = @import("core_lane").affine_mm; pub const quant = @import("src/core/quant/quant.zig"); pub const draft_ops = @import("draft_ops"); diff --git a/zig/src/core/lane_projection.zig b/zig/src/core/lane_projection.zig index 001b684892..d515d45cf6 100644 --- a/zig/src/core/lane_projection.zig +++ b/zig/src/core/lane_projection.zig @@ -3,6 +3,8 @@ const std = @import("std"); const mtl = @import("metal"); const ks = @import("kernel_sources"); const frags = @import("frags.zig"); +/// The tensor units' affine products by expert, beside the lane layouts (one module for their shared fragments). +pub const affine_mm = @import("affine_mm.zig"); pub const weight_layout = "u32[ceil(N/32)][K/group][32][group*bits/32], low-bit-first codes, columns past N zero"; pub const reg_layout = "u32[N/32][K/64][2][32 lanes][4] 4-bit codes, regBlock order"; diff --git a/zig/src/families/qwen27/attention.zig b/zig/src/families/qwen27/attention.zig index a29fbdde5e..e6a8197ea7 100644 --- a/zig/src/families/qwen27/attention.zig +++ b/zig/src/families/qwen27/attention.zig @@ -10,7 +10,7 @@ pub const Cache = struct { keys: Ref, values: Ref, capacity: u32, key_stride: u3 const CachePtrs = extern struct { keys: u64, values: u64 }; const PackArgs = extern struct { rows: u32, packed_rows: u32, heads: u32 = 24, kv_heads: u32 = 4, dim: u32 = 256, group: u32 = 6 }; const PromptArgs = extern struct { start: i32, rows: i32, chunks: i32, rp: i32, key_stride: i32, value_stride: i32, scale2: f32 }; -pub const Limits = struct { rows: u32 = 128, streams: u32 = 8, keys: u32 = 32768, shared_prefix: bool = true }; +pub const Limits = struct { rows: u32 = 128, streams: u32 = 8, keys: u32 = 32768, shared_prefix: bool = true, heads: u32 = 24, kv_heads: u32 = 4 }; pub const Field = enum { tile_stream, q_rows, nodes, paths, positions, qa, qb, poa, pma, pla, pob, pmb, plb, keep_k, keep_v }; pub const BindingField = enum { meta, caches, keep_map }; // Each layer owns immutable descriptors until its command buffer is terminal; arithmetic workspace is shared serially. @@ -51,11 +51,12 @@ pub const Binding = struct { const ks = if (cache.key_stride == 0) cache.capacity * 256 else cache.key_stride; const vs = if (cache.value_stride == 0) cache.capacity * 256 else cache.value_stride; if (ks < @as(u64, cache.capacity) * 256 or vs < @as(u64, cache.capacity) * 256 or ks > std.math.maxInt(i32) or vs > std.math.maxInt(i32)) return error.AttentionCacheStride; - try spanAligned(cache.keys, cacheBytes(cache, true), 16); - try spanAligned(cache.values, cacheBytes(cache, false), 16); - try disjoint(cache.keys, cacheBytes(cache, true), cache.values, cacheBytes(cache, false)); + const kvh = b.limits.kv_heads; + try spanAligned(cache.keys, cacheBytes(cache, true, kvh), 16); + try spanAligned(cache.values, cacheBytes(cache, false, kvh), 16); + try disjoint(cache.keys, cacheBytes(cache, true, kvh), cache.values, cacheBytes(cache, false, kvh)); for (caches[0..i]) |prior| { - for ([_]bool{ true, false }) |key| for ([_]bool{ true, false }) |prior_key| try disjoint(if (key) cache.keys else cache.values, cacheBytes(cache, key), if (prior_key) prior.keys else prior.values, cacheBytes(prior, prior_key)); + for ([_]bool{ true, false }) |key| for ([_]bool{ true, false }) |prior_key| try disjoint(if (key) cache.keys else cache.values, cacheBytes(cache, key, kvh), if (prior_key) prior.keys else prior.values, cacheBytes(prior, prior_key, kvh)); } table[i] = .{ .keys = cache.keys.buffer.gpuAddress() + cache.keys.offset, .values = cache.values.buffer.gpuAddress() + cache.values.offset }; meta[at + 7] = @intCast(ks); @@ -111,13 +112,14 @@ pub const Scratch = struct { uploaded: ?[32]u8 = null, pub fn init(device: mtl.Device, limits: Limits) !Scratch { try validLimits(limits); - const tiles = (limits.rows * 6 + 15) / 16 + limits.streams; + const tiles = (limits.rows * (limits.heads / limits.kv_heads) + 15) / 16 + limits.streams; const chunks = @max((limits.keys + 511) / 512, (planning.fine_until + 128 + planning.fine_chunk - 1) / planning.fine_chunk); // fine chunks below fine_until const rp: usize = tiles * 16; - const a: usize = if (limits.shared_prefix) 4 * @as(usize, chunks) * rp else 0; + const kv: usize = limits.kv_heads; + const a: usize = if (limits.shared_prefix) kv * @as(usize, chunks) * rp else 0; const tail_chunks = if (limits.shared_prefix) 3 else chunks; // a tail of up to 191 keys over 128-key chunks - const b: usize = 4 * @as(usize, tail_chunks) * @as(usize, limits.rows) * 16; - const sizes = [_]usize{ tiles * 4, rp * 4, limits.rows * planning.node_words * 4, limits.rows * 128 * 4, limits.rows * 4, 4 * rp * 256 * 2, 4 * @as(usize, limits.rows) * 16 * 256 * 2, a * 256 * 4, a * 4, a * 4, b * 256 * 4, b * 4, b * 4, limits.rows * 4 * 256 * 2, limits.rows * 4 * 256 * 2 }; + const b: usize = kv * @as(usize, tail_chunks) * @as(usize, limits.rows) * 16; + const sizes = [_]usize{ tiles * 4, rp * 4, limits.rows * planning.node_words * 4, limits.rows * 128 * 4, limits.rows * 4, kv * rp * 256 * 2, kv * @as(usize, limits.rows) * 16 * 256 * 2, a * 256 * 4, a * 4, a * 4, b * 256 * 4, b * 4, b * 4, limits.rows * 4 * 256 * 2, limits.rows * 4 * 256 * 2 }; comptime std.debug.assert(sizes.len == std.enums.values(Field).len); var buffers: [sizes.len]mtl.Buffer = undefined; var made: usize = 0; @@ -173,16 +175,17 @@ fn fingerprint(plan: Plan) [32]u8 { hash.final(&digest); return digest; } -fn cacheBytes(cache: Cache, key: bool) usize { +/// A cache's bytes: `kv_heads` planes of `stride` elements, the last holding `capacity` rows. +fn cacheBytes(cache: Cache, key: bool, kv_heads: u32) usize { const requested = if (key) cache.key_stride else cache.value_stride; const stride: usize = if (requested == 0) @as(usize, cache.capacity) * 256 else requested; - return (3 * stride + @as(usize, cache.capacity) * 256) * 2; + return ((kv_heads - 1) * stride + @as(usize, cache.capacity) * 256) * 2; } fn bind(e: mtl.ComputeEncoder, refs: anytype) void { inline for (refs, 0..) |r, i| e.setBuffer(r.buffer, r.offset, i); } /// h_end and w_end above 0: an M-RoPE round, the h and w position planes `plane` rows past each other in `positions`. -pub const Preprocess = struct { qg: Ref, key: Ref, q_gain: Ref, k_gain: Ref, norm_q: Ref, norm_k: Ref, query: Ref, rotated_key: Ref, positions: Ref, rows: u32, eps: f32, theta: f64, h_end: u32 = 0, w_end: u32 = 0, plane: u32 = 0 }; +pub const Preprocess = struct { heads: u32 = 24, kv_heads: u32 = 4, qg: Ref, key: Ref, q_gain: Ref, k_gain: Ref, norm_q: Ref, norm_k: Ref, query: Ref, rotated_key: Ref, positions: Ref, rows: u32, eps: f32, theta: f64, h_end: u32 = 0, w_end: u32 = 0, plane: u32 = 0 }; pub const Attention = struct { profiler: ?*profile.Trace = null, prefix: [16]mtl.Pipeline, @@ -219,14 +222,16 @@ pub const Attention = struct { const e = timing.encoder; if (p.rows == 0 or p.rows > @import("gpu_frame.zig").max_rows or !std.math.isFinite(p.eps) or p.eps <= 0 or !std.math.isFinite(p.theta) or p.theta <= 0) return error.TargetAttentionShape; - try span(p.qg, @as(usize, p.rows) * 12288 * 2); - try span(p.key, @as(usize, p.rows) * 1024 * 2); + const qw: usize = @as(usize, p.heads) * 256; + const kw: usize = @as(usize, p.kv_heads) * 256; + try span(p.qg, @as(usize, p.rows) * 2 * qw * 2); + try span(p.key, @as(usize, p.rows) * kw * 2); try span(p.q_gain, 256 * 2); try span(p.k_gain, 256 * 2); - try span(p.norm_q, @as(usize, p.rows) * 6144 * 2); - try span(p.norm_k, @as(usize, p.rows) * 1024 * 2); - try span(p.query, @as(usize, p.rows) * 6144 * 2); - try span(p.rotated_key, @as(usize, p.rows) * 1024 * 2); + try span(p.norm_q, @as(usize, p.rows) * qw * 2); + try span(p.norm_k, @as(usize, p.rows) * kw * 2); + try span(p.query, @as(usize, p.rows) * qw * 2); + try span(p.rotated_key, @as(usize, p.rows) * kw * 2); if ((p.h_end > 0 or p.w_end > 0) and p.plane < p.rows) return error.TargetAttentionShape; const plane_rows: usize = if (p.h_end > 0 or p.w_end > 0) 2 * @as(usize, p.plane) + p.rows else p.rows; try spanAligned(p.positions, plane_rows * 4, 4); @@ -238,16 +243,16 @@ pub const Attention = struct { } e.setPipeline(a.io[7]); bind(e, .{ p.qg, p.q_gain }); - e.setValue(glue.Head{ .width = 256, .rows = p.rows * 24, .stride = 512, .activation = .bf16, .gain = .bf16, .eps = p.eps }, 2); + e.setValue(glue.Head{ .width = 256, .rows = p.rows * p.heads, .stride = 512, .activation = .bf16, .gain = .bf16, .eps = p.eps }, 2); e.setBuffer(p.norm_q.buffer, p.norm_q.offset, 3); - e.dispatchGroups(mtl.Size.of(p.rows * 24, 1, 1), mtl.Size.of(64, 1, 1)); + e.dispatchGroups(mtl.Size.of(p.rows * p.heads, 1, 1), mtl.Size.of(64, 1, 1)); e.setPipeline(a.io[7]); bind(e, .{ p.key, p.k_gain }); - e.setValue(glue.Head{ .width = 256, .rows = p.rows * 4, .stride = 256, .activation = .bf16, .gain = .bf16, .eps = p.eps }, 2); + e.setValue(glue.Head{ .width = 256, .rows = p.rows * p.kv_heads, .stride = 256, .activation = .bf16, .gain = .bf16, .eps = p.eps }, 2); e.setBuffer(p.norm_k.buffer, p.norm_k.offset, 3); - e.dispatchGroups(mtl.Size.of(p.rows * 4, 1, 1), mtl.Size.of(64, 1, 1)); + e.dispatchGroups(mtl.Size.of(p.rows * p.kv_heads, 1, 1), mtl.Size.of(64, 1, 1)); e.barrier(); - for ([_]struct { x: Ref, y: Ref, heads: u32 }{ .{ .x = p.norm_q, .y = p.query, .heads = 24 }, .{ .x = p.norm_k, .y = p.rotated_key, .heads = 4 } }) |r| { + for ([_]struct { x: Ref, y: Ref, heads: u32 }{ .{ .x = p.norm_q, .y = p.query, .heads = p.heads }, .{ .x = p.norm_k, .y = p.rotated_key, .heads = p.kv_heads } }) |r| { e.setPipeline(a.io[8]); bind(e, .{ r.x, p.positions }); e.setValue(glue.Rope{ .width = 256, .heads = r.heads, .rows = p.rows, .rotary = 64, .activation = .bf16, .theta = @floatCast(p.theta), .h_end = p.h_end, .w_end = p.w_end, .plane = p.plane }, 2); @@ -265,18 +270,21 @@ pub const Attention = struct { try scratch.check(plan); try binding.check(plan, caches); - try spanAligned(q, @as(usize, plan.rows) * 6144 * 2, 16); - try span(k, @as(usize, plan.rows) * 1024 * 2); - try span(v, @as(usize, plan.rows) * 1024 * 2); - try span(out, @as(usize, plan.rows) * 6144 * 2); + const g = plan.geometry; + const qw: usize = @as(usize, g.heads) * 256; + const kw: usize = @as(usize, g.kv_heads) * 256; + try spanAligned(q, @as(usize, plan.rows) * qw * 2, 16); + try span(k, @as(usize, plan.rows) * kw * 2); + try span(v, @as(usize, plan.rows) * kw * 2); + try span(out, @as(usize, plan.rows) * qw * 2); if (raw_gate) |gate| { - try span(gate, @as(usize, plan.rows) * 12288 * 2); - try disjoint(out, @as(usize, plan.rows) * 6144 * 2, gate, @as(usize, plan.rows) * 12288 * 2); + try span(gate, @as(usize, plan.rows) * 2 * qw * 2); + try disjoint(out, @as(usize, plan.rows) * qw * 2, gate, @as(usize, plan.rows) * 2 * qw * 2); } const operands = [_]struct { ref: Ref, bytes: usize }{ .{ .ref = q, .bytes = @as(usize, plan.rows) * 6144 * 2 }, .{ .ref = k, .bytes = @as(usize, plan.rows) * 1024 * 2 }, .{ .ref = v, .bytes = @as(usize, plan.rows) * 1024 * 2 }, .{ .ref = out, .bytes = @as(usize, plan.rows) * 6144 * 2 } }; for (caches) |cache| for (operands) |operand| { - try disjoint(cache.keys, cacheBytes(cache, true), operand.ref, operand.bytes); - try disjoint(cache.values, cacheBytes(cache, false), operand.ref, operand.bytes); + try disjoint(cache.keys, cacheBytes(cache, true, g.kv_heads), operand.ref, operand.bytes); + try disjoint(cache.values, cacheBytes(cache, false, g.kv_heads), operand.ref, operand.bytes); }; for (scratch.buffers) |buffer| e.useResource(buffer, mtl.ResourceUsage.read | mtl.ResourceUsage.write); for (binding.buffers) |buffer| e.useResource(buffer, mtl.ResourceUsage.read | mtl.ResourceUsage.write); @@ -284,24 +292,24 @@ pub const Attention = struct { e.useResource(cache.keys.buffer, mtl.ResourceUsage.read | mtl.ResourceUsage.write); e.useResource(cache.values.buffer, mtl.ResourceUsage.read | mtl.ResourceUsage.write); } - const args = PackArgs{ .rows = plan.rows, .packed_rows = plan.tiles * 16 }; + const args = PackArgs{ .rows = plan.rows, .packed_rows = plan.tiles * 16, .heads = g.heads, .kv_heads = g.kv_heads, .group = g.group() }; e.setPipeline(a.io[2]); bind(e, .{ k, v, binding.get(.caches), binding.get(.meta), scratch.get(.nodes) }); e.setValue(args, 5); - e.dispatchThreads(mtl.Size.of(256, 4, plan.rows), mtl.Size.of(32, 4, 1)); + e.dispatchThreads(mtl.Size.of(256, g.kv_heads, plan.rows), mtl.Size.of(32, 4, 1)); e.setPipeline(a.io[0]); bind(e, .{ q, scratch.get(.q_rows), scratch.get(.qa) }); e.setValue(args, 3); - e.dispatchThreads(mtl.Size.of(256, args.packed_rows, 4), mtl.Size.of(32, 4, 1)); + e.dispatchThreads(mtl.Size.of(256, args.packed_rows, g.kv_heads), mtl.Size.of(32, 4, 1)); if (prompt and a.prompt != null and plan.geometry.shared_prefix and plan.streams == 1) { e.barrier(); try a.promptRows(e, scratch, binding, plan, out); if (raw_gate) |gate| { e.setPipeline(a.io[9]); bind(e, .{ out, gate }); - e.setValue(glue.Element{ .width = 256, .heads = 24, .rows = plan.rows, .activation = .bf16 }, 2); + e.setValue(glue.Element{ .width = 256, .heads = g.heads, .rows = plan.rows, .activation = .bf16 }, 2); e.setBuffer(out.buffer, out.offset, 3); - e.dispatchThreads(mtl.Size.of(256, 24, plan.rows), mtl.Size.of(32, 1, 1)); + e.dispatchThreads(mtl.Size.of(256, g.heads, plan.rows), mtl.Size.of(32, 1, 1)); e.barrier(); } } else try a.treeRows(e, scratch, binding, plan, q, out, raw_gate, args); @@ -321,13 +329,13 @@ pub const Attention = struct { e.setBuffer(scratch.get(.poa).buffer, 0, 3); e.setBuffer(scratch.get(.pma).buffer, 0, 4); e.setBuffer(scratch.get(.pla).buffer, 0, 5); - e.dispatchGroups(mtl.Size.of(4, chunks, (plan.tiles + pairs - 1) / pairs), mtl.Size.of(32 * 2 * pairs, 1, 1)); + e.dispatchGroups(mtl.Size.of(plan.geometry.kv_heads, chunks, (plan.tiles + pairs - 1) / pairs), mtl.Size.of(32 * 2 * pairs, 1, 1)); e.barrier(); e.setPipeline(a.prompt.?[1]); bind(e, .{ scratch.get(.poa), scratch.get(.pma), scratch.get(.pla) }); e.setValue(args, 3); e.setBuffer(out.buffer, out.offset, 4); - e.dispatchGroups(mtl.Size.of(4, plan.rows * 6, 1), mtl.Size.of(32, 1, 1)); + e.dispatchGroups(mtl.Size.of(plan.geometry.kv_heads, plan.rows * plan.geometry.group(), 1), mtl.Size.of(32, 1, 1)); e.barrier(); } /// Trees and decoded rows: cached keys by chunk (prefix), the round's own by path (tail), merged and gated. @@ -336,7 +344,7 @@ pub const Attention = struct { e.setPipeline(a.io[1]); bind(e, .{ q, scratch.get(.qb) }); e.setValue(args, 2); - e.dispatchThreads(mtl.Size.of(256, 16, 4 * plan.rows), mtl.Size.of(32, 4, 1)); + e.dispatchThreads(mtl.Size.of(256, 16, plan.geometry.kv_heads * plan.rows), mtl.Size.of(32, 4, 1)); } e.barrier(); const scale: f32 = 0.0625; @@ -351,7 +359,7 @@ pub const Attention = struct { e.setBuffer(scratch.get(.pma).buffer, 0, 6); e.setBuffer(scratch.get(.pla).buffer, 0, 7); e.setBuffer(scratch.get(.nodes).buffer, 0, 8); - e.dispatchGroups(mtl.Size.of(4, plan.prefix_chunks, (plan.tiles + sg - 1) / sg), mtl.Size.of(64 * sg, 1, 1)); + e.dispatchGroups(mtl.Size.of(plan.geometry.kv_heads, plan.prefix_chunks, (plan.tiles + sg - 1) / sg), mtl.Size.of(64 * sg, 1, 1)); e.barrier(); } if (!plan.fused) { @@ -363,13 +371,13 @@ pub const Attention = struct { const r = scratch.get(field); e.setBuffer(r.buffer, 0, i); } - e.dispatchGroups(mtl.Size.of(4, plan.tail_chunks, plan.rows), mtl.Size.of(128, 1, 1)); + e.dispatchGroups(mtl.Size.of(plan.geometry.kv_heads, plan.tail_chunks, plan.rows), mtl.Size.of(128, 1, 1)); e.barrier(); } e.setPipeline(a.io[4]); bind(e, .{ scratch.get(.poa), scratch.get(.pma), scratch.get(.pla), scratch.get(.pob), scratch.get(.pmb), scratch.get(.plb), binding.get(.meta), scratch.get(.nodes), out, raw_gate orelse out }); e.setValue(@as(u32, @intFromBool(raw_gate != null)), 10); - e.dispatchGroups(mtl.Size.of(4, plan.rows * 6, 1), mtl.Size.of(32, 1, 1)); + e.dispatchGroups(mtl.Size.of(plan.geometry.kv_heads, plan.rows * plan.geometry.group(), 1), mtl.Size.of(32, 1, 1)); e.barrier(); } pub fn keep(a: *const Attention, e: mtl.ComputeEncoder, scratch: *const Scratch, binding: *const Binding, kept: u32, caches: []const Cache) !void { @@ -385,12 +393,12 @@ pub const Attention = struct { e.setPipeline(a.io[5]); bind(e, .{ binding.get(.caches), binding.get(.meta), binding.get(.keep_map), scratch.get(.keep_k), scratch.get(.keep_v) }); e.setValue(kept, 5); - e.dispatchThreads(mtl.Size.of(256, 4, kept), mtl.Size.of(32, 4, 1)); + e.dispatchThreads(mtl.Size.of(256, scratch.limits.kv_heads, kept), mtl.Size.of(32, 4, 1)); e.barrier(); e.setPipeline(a.io[6]); bind(e, .{ scratch.get(.keep_k), scratch.get(.keep_v), binding.get(.caches), binding.get(.meta), binding.get(.keep_map) }); e.setValue(kept, 5); - e.dispatchThreads(mtl.Size.of(256, 4, kept), mtl.Size.of(32, 4, 1)); + e.dispatchThreads(mtl.Size.of(256, scratch.limits.kv_heads, kept), mtl.Size.of(32, 4, 1)); e.barrier(); } }; diff --git a/zig/src/families/qwen27/attention_plan.zig b/zig/src/families/qwen27/attention_plan.zig index b6bb100d43..2751a41c56 100644 --- a/zig/src/families/qwen27/attention_plan.zig +++ b/zig/src/families/qwen27/attention_plan.zig @@ -29,7 +29,9 @@ pub const Geometry = struct { /// Chains of rows take the prefix kernel to their last key; a test turns this off to compare with the tail kernel. fuse_chain: bool = true, pub fn check(g: Geometry) !void { - if (g.heads != 24 or g.kv_heads != 4 or g.dim != 256 or g.rotary != 64) return error.TargetAttentionShape; + // the kernels' 16-row tiles hold one KV group of up to 16 query heads: the 27B's 24 over 4 and Ornith's 32 over 2 + const grouped = (g.heads == 24 and g.kv_heads == 4) or (g.heads == 32 and g.kv_heads == 2); + if (!grouped or g.dim != 256 or g.rotary != 64) return error.TargetAttentionShape; } pub fn promptRows(g: Geometry, requested: u32) u32 { return @min(requested, @as(u32, if (g.shared_prefix and g.fuse_chain) chain_rows else max_rows)); diff --git a/zig/src/families/qwen27/config.zig b/zig/src/families/qwen27/config.zig index e4122d631e..1c275b78c1 100644 --- a/zig/src/families/qwen27/config.zig +++ b/zig/src/families/qwen27/config.zig @@ -1,4 +1,4 @@ -//! Dense Qwen text shapes, layer schedule, rotary metadata and ordered checkpoint EOS ids. +//! Qwen3.5-class text shapes (dense, or routed experts), layer schedule, rotary metadata and ordered checkpoint EOS ids. const std = @import("std"); pub const Kind = enum { linear, attention }; @@ -27,6 +27,17 @@ pub const Config = struct { eos_count: usize = 0, max_position: usize = 262144, mrope_interleaved: bool = true, + /// Routed experts (qwen3_5_moe): their count and the picks a row makes; 0: a dense MLP of `intermediate`. + experts: usize = 0, + top_k: usize = 0, + /// A routed expert's width; `intermediate` is then the shared expert's (the same, as the CUDA family requires). + moe_intermediate: usize = 0, + /// Whether a row's top-k softmax weights are renormalized to sum to one (`norm_topk_prob`). + renorm: bool = false, + + pub fn moe(c: Config) bool { + return c.experts != 0; + } pub fn kind(c: Config, layer: usize) Kind { return c.kinds[layer]; @@ -119,8 +130,13 @@ pub fn parse(gpa: std.mem.Allocator, json: []const u8, generation: ?[]const u8) const root = try object(parsed.value); const t = if (root.get("text_config")) |v| try object(v) else root; const model = t.get("model_type") orelse root.get("model_type") orelse return error.BadConfig; - if (model != .string or (!std.mem.eql(u8, model.string, "qwen3_5") and !std.mem.eql(u8, model.string, "qwen3_5_text"))) return error.NotQwenDense; - if (t.get("num_experts")) |v| if (try integer(v, true) != 0) return error.NotQwenDense; + const kinds_of = [_][]const u8{ "qwen3_5", "qwen3_5_text", "qwen3_5_moe", "qwen3_5_moe_text" }; + if (model != .string) return error.NotQwenDense; + const routed = for (kinds_of, 0..) |k, i| { + if (std.mem.eql(u8, model.string, k)) break i >= 2; + } else return error.NotQwenDense; + const experts = if (t.get("num_experts")) |v| try integer(v, true) else 0; + if (routed != (experts != 0)) return error.NotQwenDense; if (try flag(root, "tie_word_embeddings") or try flag(t, "tie_word_embeddings")) return error.TiedEmbeddingHead; if (t.get("attn_output_gate")) |v| if (v != .bool or !v.bool) return error.UnsupportedAttentionGate; inline for (.{ .{ "hidden_act", "silu" }, .{ "output_gate_type", "swish" }, .{ "linear_attn_state_dtype", "float32" } }) |entry| { @@ -137,7 +153,7 @@ pub fn parse(gpa: std.mem.Allocator, json: []const u8, generation: ?[]const u8) if (factor <= 0 or factor > 1) return error.BadConfig; var c = Config{ .hidden = hidden, - .intermediate = try required(t, "intermediate_size"), + .intermediate = if (routed) try required(t, "shared_expert_intermediate_size") else try required(t, "intermediate_size"), .layers = try required(t, "num_hidden_layers"), .vocab = try required(t, "vocab_size"), .heads = heads, @@ -152,6 +168,13 @@ pub fn parse(gpa: std.mem.Allocator, json: []const u8, generation: ?[]const u8) .rope_dims = @intFromFloat(@as(f64, @floatFromInt(head_dim)) * factor), .rope_theta = if (rope) |r| try optional(r, "rope_theta", try optional(t, "rope_theta", 10000000)) else try optional(t, "rope_theta", 10000000), }; + if (routed) { + c.experts = experts; + c.top_k = try required(t, "num_experts_per_tok"); + c.moe_intermediate = try required(t, "moe_intermediate_size"); + c.renorm = try flag(t, "norm_topk_prob"); + if (c.experts > 256 or c.top_k >= c.experts or c.top_k > 15 or c.moe_intermediate != c.intermediate) return error.UnsupportedRouting; + } if (t.get("max_position_embeddings")) |v| c.max_position = try integer(v, false); if (rope) |r| if (r.get("mrope_interleaved")) |v| { if (v != .bool or !v.bool) return error.UnsupportedMrope; diff --git a/zig/src/families/qwen27/decode_round.zig b/zig/src/families/qwen27/decode_round.zig index 4bfd54920a..4b5e921987 100644 --- a/zig/src/families/qwen27/decode_round.zig +++ b/zig/src/families/qwen27/decode_round.zig @@ -84,13 +84,18 @@ pub const Runner = struct { const units = m.device.tensorUnits(); // before M5 the tail attention reads its tensor-op results in the simdgroup-matrix layout (simd_attention.zig) if (!units) try @import("simd_attention.zig").check(m.device, m.queue); - const text = if (units) try a.dupe(u8, source.glue ++ source.attention_mpp ++ source.attention_io ++ source.attention_prompt) else try @import("simd_attention.zig").rewrite(a, source.glue ++ source.attention_mpp ++ source.attention_io); + // the kernels' KV group: query heads a KV head, prepended as the defines attention_mpp.metal reads + const defines = try std.fmt.allocPrint(a, "#define Q27_G {d}\n#define Q27_KVH {d}\n", .{ m.config.heads / m.config.kv_heads, m.config.kv_heads }); + defer a.free(defines); + const body = if (units) try std.mem.concat(a, u8, &.{ source.glue, source.attention_mpp, source.attention_io, source.attention_prompt }) else try @import("simd_attention.zig").rewrite(a, source.glue ++ source.attention_mpp ++ source.attention_io); + defer a.free(body); + const text = try std.mem.concat(a, u8, &.{ defines, body }); defer a.free(text); const library = try mtl.Library.fromSource(m.device, text, mtl.CompileOptions.mlx()); defer library.deinit(); var attention = try attn.Attention.init(m.device, library, units); errdefer attention.deinit(); - const limits = attn.Limits{ .rows = m.frame.capacity, .streams = slots, .keys = capacity, .shared_prefix = m.device.tensorUnits() }; + const limits = attn.Limits{ .rows = m.frame.capacity, .streams = slots, .keys = capacity, .shared_prefix = m.device.tensorUnits(), .heads = @intCast(m.config.heads), .kv_heads = @intCast(m.config.kv_heads) }; var scratch = try attn.Scratch.init(m.device, limits); errdefer scratch.deinit(); const bindings = try a.alloc(attn.Binding, na); @@ -354,14 +359,14 @@ pub const Runner = struct { } pub fn promptRows(r: *const Runner, requested: u32) u32 { - return (attn.planning.Geometry{ .shared_prefix = r.model.device.tensorUnits() }).promptRows(@min(requested, r.model.frame.capacity)); + return (attn.planning.Geometry{ .shared_prefix = r.model.device.tensorUnits(), .heads = @intCast(r.model.config.heads), .kv_heads = @intCast(r.model.config.kv_heads) }).promptRows(@min(requested, r.model.frame.capacity)); } fn attentionPlan(r: *Runner, round: *const Round) !attn.Plan { var windows: [attn.planning.max_streams]attn.planning.Window = undefined; if (round.windows.len > windows.len) return error.AttentionStreams; for (round.windows, 0..) |window, i| windows[i] = .{ .parents = window.parents, .start = window.start, .capacity = r.capacity }; - return attn.Plan.init(r.allocator, windows[0..round.windows.len], .{ .shared_prefix = r.model.device.tensorUnits() }); + return attn.Plan.init(r.allocator, windows[0..round.windows.len], .{ .shared_prefix = r.model.device.tensorUnits(), .heads = @intCast(r.model.config.heads), .kv_heads = @intCast(r.model.config.kv_heads) }); } fn layerCaches(r: *Runner, layer: usize, round: *const Round, views: *[attn.planning.max_streams]attn.Cache) []const attn.Cache { @@ -385,10 +390,11 @@ pub const Runner = struct { try r.model.kernels.quant(e, weights.q, f.get(.input), forward.sums(f, .sums, rows), f.get(.dims), f.get(.q_gate), rows); try r.model.kernels.quant(e, weights.kv, f.get(.input), forward.sums(f, .sums, rows), f.get(.dims), f.get(.mixed), rows); e.barrier(); - try r.model.glue.unstack(e, f.get(.mixed), f.get(.key), .{ .rows = rows, .width = 1024, .stride = 2048, .offset = 0 }); - try r.model.glue.unstack(e, f.get(.mixed), f.get(.value), .{ .rows = rows, .width = 1024, .stride = 2048, .offset = 1024 }); + const kv: u32 = @intCast(r.model.config.kvDim()); + try r.model.glue.unstack(e, f.get(.mixed), f.get(.key), .{ .rows = rows, .width = kv, .stride = 2 * kv, .offset = 0 }); + try r.model.glue.unstack(e, f.get(.mixed), f.get(.value), .{ .rows = rows, .width = kv, .stride = 2 * kv, .offset = kv }); e.barrier(); - try r.attention.preprocess(e, .{ .qg = f.get(.q_gate), .key = f.get(.key), .q_gain = wts.ref(weights.q_norm), .k_gain = wts.ref(weights.k_norm), .norm_q = f.get(.norm_query), .norm_k = f.get(.norm_key), .query = f.get(.query), .rotated_key = f.get(.rotated_key), .positions = f.get(.positions), .rows = rows, .eps = r.model.config.eps, .theta = r.model.config.rope_theta, .h_end = r.current_placed.h_end, .w_end = r.current_placed.w_end, .plane = f.capacity }); + try r.attention.preprocess(e, .{ .heads = @intCast(r.model.config.heads), .kv_heads = @intCast(r.model.config.kv_heads), .qg = f.get(.q_gate), .key = f.get(.key), .q_gain = wts.ref(weights.q_norm), .k_gain = wts.ref(weights.k_norm), .norm_q = f.get(.norm_query), .norm_k = f.get(.norm_key), .query = f.get(.query), .rotated_key = f.get(.rotated_key), .positions = f.get(.positions), .rows = rows, .eps = r.model.config.eps, .theta = r.model.config.rope_theta, .h_end = r.current_placed.h_end, .w_end = r.current_placed.w_end, .plane = f.capacity }); try r.attention.encode(e, &r.scratch, &r.bindings[layer], plan, f.get(.query), f.get(.rotated_key), f.get(.value), f.get(.mixer_out), caches, f.get(.q_gate), r.model.kernels.prompt); } diff --git a/zig/src/families/qwen27/forward.zig b/zig/src/families/qwen27/forward.zig index 7cf8e564bc..ab2f647505 100644 --- a/zig/src/families/qwen27/forward.zig +++ b/zig/src/families/qwen27/forward.zig @@ -103,6 +103,10 @@ pub const Stages = struct { try m.kernels.addAne(lane.enc, .{ .buffer = share.partial }, .{ .buffer = share.outputs[slice].buffer }, f.get(.output), h, width, share.columns, @min(ane_mlp.width, s.rows - base), base); } lane.enc.barrier(); + if (layer.moe) |*mo| { + try m.moe.?.block(lane.enc, m.config, mo, f, s.rows, m.kernels.prompt); + lane.enc.barrier(); + } s.pending = f.get(.output); return; } else { @@ -113,6 +117,11 @@ pub const Stages = struct { lane.enc.barrier(); try m.kernels.quant(lane.enc, layer.down, f.get(.activated), sums(f, .act_sums, s.rows), f.get(.dims), f.get(.output), s.rows); lane.enc.barrier(); + if (layer.moe) |*mo| { + // the shared expert's rows in f.output take the routed experts' weighted rows and the shared gate + if (std.c.getenv("TF_NO_ROUTED") == null) try m.moe.?.block(lane.enc, m.config, mo, f, s.rows, m.kernels.prompt); + lane.enc.barrier(); + } s.pending = f.get(.output); } pub fn finish(s: *Stages, lane: *@import("core").segments.Lane) !void { diff --git a/zig/src/families/qwen27/gpu_frame.zig b/zig/src/families/qwen27/gpu_frame.zig index 6c068ca305..9171196e92 100644 --- a/zig/src/families/qwen27/gpu_frame.zig +++ b/zig/src/families/qwen27/gpu_frame.zig @@ -4,10 +4,14 @@ const mtl = @import("metal"); const Config = @import("config.zig").Config; const Ref = @import("projection.zig").Ref; -pub const Field = enum { ids, dims, positions, hidden, residual, input, sums, mixed, mixer_out, output, z, a, b, query, key, value, norm_query, norm_key, q_gate, gate, up, activated, act_sums, logits, drawn, eps, parents, windows, row_slots, segments, keeps, kept_rows, rotated_key, takes, slot_steps }; +pub const Field = enum { ids, dims, positions, hidden, residual, input, sums, mixed, mixer_out, output, z, a, b, query, key, value, norm_query, norm_key, q_gate, gate, up, activated, act_sums, logits, drawn, eps, parents, windows, row_slots, segments, keeps, kept_rows, rotated_key, takes, slot_steps, router, picks, route_weights, offsets, order, sorted_weights, sorted_rows, expert_gate, expert_up, expert_act, expert_out, expert_sums }; /// A prompt pass's rows at most (a chain); rounds with trees and every logits row stay within `logit_rows`. pub const max_rows = 2048; + +fn routerRows(config: Config) usize { + return if (config.moe()) @import("widths.zig").routerRows(config.experts) else 0; +} pub const logit_rows = 128; pub const Frame = struct { @@ -18,6 +22,7 @@ pub const Frame = struct { pub fn init(device: mtl.Device, config: Config, capacity: u32) !Frame { if (capacity == 0 or capacity > max_rows) return error.BadFrameRows; const mp = 16 * ((capacity + 15) / 16); + const slots: usize = if (config.moe()) capacity * config.top_k else 0; const widths = [_]usize{ capacity * 4, 8 * 4, 3 * capacity * 4, capacity * config.hidden * 2, capacity * config.hidden * 2, capacity * config.hidden * 2, @@ -31,6 +36,13 @@ pub const Frame = struct { capacity * config.conv_kernel * 4, capacity * 4, capacity * 16, capacity * 16, capacity * 4, capacity * config.kvDim() * 2, capacity * 4, capacity * @sizeOf(@import("core").tree_slots.Step), + // the routed block: a row's logits, its picks and weights, the sorted slots, the experts' rows + capacity * routerRows(config) * 2, slots * 4, + slots * 4, (config.experts + 1) * 4, + slots * 4, slots * 4, + slots * config.hidden * 2, slots * config.moe_intermediate * 2, + slots * config.moe_intermediate * 2, slots * config.moe_intermediate * 2, + slots * config.hidden * 2, slots * @max(config.hidden, config.moe_intermediate) / 64 * 4, }; comptime std.debug.assert(widths.len == std.enums.values(Field).len); var buffers: [widths.len]mtl.Buffer = undefined; @@ -44,6 +56,11 @@ pub const Frame = struct { return .{ .buffers = buffers, .capacity = capacity, .config = config }; } + /// Routed slots a pass of `rows` makes: each row's top-k picks. + pub fn slotsOf(frame: *const Frame, rows: u32) u32 { + return rows * @as(u32, @intCast(frame.config.top_k)); + } + pub fn deinit(frame: *Frame) void { for (frame.buffers) |b| b.deinit(); } diff --git a/zig/src/families/qwen27/gpu_weights.zig b/zig/src/families/qwen27/gpu_weights.zig index bb0fba68e3..7cb521c1ee 100644 --- a/zig/src/families/qwen27/gpu_weights.zig +++ b/zig/src/families/qwen27/gpu_weights.zig @@ -5,6 +5,7 @@ const cp = @import("core").checkpoint_metal; const Config = @import("config.zig").Config; const projection = @import("projection.zig"); const lane = @import("core").lane_projection; +const row = @import("core").row_projection; const widths = @import("widths.zig"); const Width = widths.Width; const Ref = projection.Ref; @@ -13,7 +14,10 @@ pub const Tensor = cp.Tensor; pub const Gdn = struct { qkv: Linear, zba: Linear, out: Linear, conv: Tensor, a_log: Tensor, dt: Tensor, norm: Tensor }; pub const Attention = struct { q: Linear, kv: Linear, out: Linear, q_norm: Tensor, k_norm: Tensor }; -pub const Layer = struct { norm: Tensor, post_norm: Tensor, gu: Linear, down: Linear, mixer: union(@import("config.zig").Kind) { linear: Gdn, attention: Attention } }; +/// The routed block: the router's rows (the experts, then the shared expert's gate, then zero rows) and the stacked experts. +pub const Moe = struct { router: row.Weights, gate: row.Experts, up: row.Experts, down: row.Experts }; +/// With experts, `gu` and `down` are the shared expert's. +pub const Layer = struct { norm: Tensor, post_norm: Tensor, gu: Linear, down: Linear, moe: ?Moe = null, mixer: union(@import("config.zig").Kind) { linear: Gdn, attention: Attention } }; pub const Weights = struct { allocator: std.mem.Allocator, @@ -44,7 +48,7 @@ fn uniform(p: Linear) bool { } /// Modules the graph reads only through repacked copies: their tensors stay in the files until load stages them. -const repacked_modules = [_][]const u8{ ".mlp.gate_proj", ".mlp.up_proj", ".mlp.down_proj", ".linear_attn.in_proj_qkv", ".linear_attn.in_proj_z", ".linear_attn.in_proj_b", ".linear_attn.in_proj_a", ".linear_attn.out_proj", ".self_attn.q_proj", ".self_attn.k_proj", ".self_attn.v_proj", ".self_attn.o_proj" }; +const repacked_modules = [_][]const u8{ ".mlp.gate_proj", ".mlp.up_proj", ".mlp.down_proj", ".mlp.shared_expert.gate_proj", ".mlp.shared_expert.up_proj", ".mlp.shared_expert.down_proj", ".linear_attn.in_proj_qkv", ".linear_attn.in_proj_z", ".linear_attn.in_proj_b", ".linear_attn.in_proj_a", ".linear_attn.out_proj", ".self_attn.q_proj", ".self_attn.k_proj", ".self_attn.v_proj", ".self_attn.o_proj" }; /// Whether `name` is a weight, scales or biases tensor of a repacked module or the LM head. pub fn repacked(name: []const u8, _: usize) bool { @@ -89,6 +93,54 @@ const Builder = struct { return .{ .weight = tensors[0], .scales = tensors[1], .biases = tensors[2], .bits = w.bits, .group = w.group }; } + /// `experts` [n, k] matrices stacked in `module` at width `w`, borrowed where the files packed them. + fn experts(b: Builder, module: []const u8, count: usize, n: usize, k: usize, w: Width, repeat: usize) !row.Experts { + const t = try b.quantAt(module, &.{count}, n, k, w); + const out: row.Experts = .{ .one = weightsOf(t, n, k), .experts = count, .repeat = repeat }; + try out.validate(); + return out; + } + + /// A quantized module's three tensors of `lead` ++ [n, packed k] shape. + fn quantAt(b: Builder, module: []const u8, lead: []const usize, n: usize, k: usize, w: Width) !projection.Triple { + if (k * w.bits % 32 != 0 or k % w.group != 0) return error.BadNativeTensor; + const suffixes = [_][]const u8{ ".weight", ".scales", ".biases" }; + var tensors: [3]Tensor = undefined; + var shape: [3]usize = undefined; + @memcpy(shape[0..lead.len], lead); + shape[lead.len] = n; + for (suffixes, &tensors, 0..) |suffix, *t, i| { + const name = try std.mem.concat(b.allocator, u8, &.{ module, suffix }); + defer b.allocator.free(name); + shape[lead.len + 1] = if (i == 0) k * w.bits / 32 else k / w.group; + t.* = try b.tensor(name, if (i == 0) .u32 else .bf16, shape[0 .. lead.len + 2]); + } + return .{ .weight = tensors[0], .scales = tensors[1], .biases = tensors[2], .bits = w.bits, .group = w.group }; + } + + /// The router's [experts, k] rows and the shared expert's gate row in one owned matrix of widths.routerRows rows. + fn router(b: Builder, prefix: []const u8, c: Config, w: Width) !row.Weights { + const k = c.hidden; + const r = try b.quantAt(try std.mem.concat(b.allocator, u8, &.{ prefix, "mlp.gate" }), &.{}, c.experts, k, w); + const g = try b.quantAt(try std.mem.concat(b.allocator, u8, &.{ prefix, "mlp.shared_expert_gate" }), &.{}, 1, k, w); + const n = widths.routerRows(c.experts); + const words = k * w.bits / 32 * 4; + const groups = k / w.group * 2; + const parts = [3][2]usize{ .{ words, 0 }, .{ groups, words * n }, .{ groups, words * n + groups * n } }; + const buffer = try b.device.buffer(parts[2][1] + groups * n, mtl.ResourceOptions.shared); + errdefer buffer.deinit(); + const bytes = buffer.slice(u8, parts[2][1] + groups * n); + @memset(bytes, 0); + for ([_]Tensor{ r.weight, r.scales, r.biases }, [_]Tensor{ g.weight, g.scales, g.biases }, parts) |rt, gt, part| { + @memcpy(bytes[part[1]..][0 .. part[0] * c.experts], rt.host(u8)[0 .. part[0] * c.experts]); + @memcpy(bytes[part[1] + part[0] * c.experts ..][0..part[0]], gt.host(u8)[0..part[0]]); + } + try b.owned.append(b.allocator, buffer); + const out: row.Weights = .{ .w = buffer, .w_off = parts[0][1], .scales = buffer, .s_off = parts[1][1], .biases = buffer, .b_off = parts[2][1], .n = n, .k = k, .group = w.group, .bits = w.bits }; + try out.validate(); + return out; + } + /// `parts` (stacked in order, `n` rows of `k` each, at their widths) staged from the files and repacked. fn linear(b: Builder, parts: []const Part, k: usize) !Linear { var arena = std.heap.ArenaAllocator.init(b.allocator); @@ -168,9 +220,15 @@ const Builder = struct { .k_norm = try b.tensor(try Name.full(a, prefix, "self_attn.k_norm.weight"), .bf16, &.{c.head_dim}), } }, }; - const gu = try b.linear(&.{ .{ .module = try Name.full(a, prefix, "mlp.gate_proj"), .n = c.intermediate, .w = w[@backingInt(M.gate)] }, .{ .module = try Name.full(a, prefix, "mlp.up_proj"), .n = c.intermediate, .w = w[@backingInt(M.up)] } }, c.hidden); - const down = try b.linear(&.{.{ .module = try Name.full(a, prefix, "mlp.down_proj"), .n = c.hidden, .w = w[@backingInt(M.down)] }}, c.intermediate); - return .{ .norm = try b.tensor(try Name.full(a, prefix, "input_layernorm.weight"), .bf16, &.{c.hidden}), .post_norm = try b.tensor(try Name.full(a, prefix, "post_attention_layernorm.weight"), .bf16, &.{c.hidden}), .gu = gu, .down = down, .mixer = mixer }; + const moe: ?Moe = if (!c.moe()) null else .{ + .router = try b.router(prefix, c, w[@backingInt(M.router)]), + .gate = try b.experts(try Name.full(a, prefix, widths.leaf(c, .egate)), c.experts, c.moe_intermediate, c.hidden, w[@backingInt(M.egate)], c.top_k), + .up = try b.experts(try Name.full(a, prefix, widths.leaf(c, .eup)), c.experts, c.moe_intermediate, c.hidden, w[@backingInt(M.eup)], c.top_k), + .down = try b.experts(try Name.full(a, prefix, widths.leaf(c, .edown)), c.experts, c.hidden, c.moe_intermediate, w[@backingInt(M.edown)], 1), + }; + const gu = try b.linear(&.{ .{ .module = try Name.full(a, prefix, widths.leaf(c, .gate)), .n = c.intermediate, .w = w[@backingInt(M.gate)] }, .{ .module = try Name.full(a, prefix, widths.leaf(c, .up)), .n = c.intermediate, .w = w[@backingInt(M.up)] } }, c.hidden); + const down = try b.linear(&.{.{ .module = try Name.full(a, prefix, widths.leaf(c, .down)), .n = c.hidden, .w = w[@backingInt(M.down)] }}, c.intermediate); + return .{ .norm = try b.tensor(try Name.full(a, prefix, "input_layernorm.weight"), .bf16, &.{c.hidden}), .post_norm = try b.tensor(try Name.full(a, prefix, "post_attention_layernorm.weight"), .bf16, &.{c.hidden}), .gu = gu, .down = down, .moe = moe, .mixer = mixer }; } }; @@ -229,6 +287,11 @@ fn regJobs(a: std.mem.Allocator, jobs: *std.ArrayList(lane.Order), p: Linear) st if (p.runs.len == 0 and p.bits == 4 and p.group == 64 and p.n % 32 == 0) try jobs.append(a, .{ .w = .{ .buf = p.words.buffer, .off = p.words.offset }, .n = p.n, .k = p.k }); } +/// A borrowed quantized tensor as the row kernels' matrix. +fn weightsOf(t: projection.Triple, n: usize, k: usize) row.Weights { + return .{ .w = t.weight.buffer, .w_off = t.weight.offset, .scales = t.scales.buffer, .s_off = t.scales.offset, .biases = t.biases.buffer, .b_off = t.biases.offset, .n = n, .k = k, .group = t.group, .bits = t.bits }; +} + pub fn ref(t: Tensor) Ref { return .{ .buffer = t.buffer, .offset = t.offset }; } diff --git a/zig/src/families/qwen27/model.zig b/zig/src/families/qwen27/model.zig index 68de84bf0b..a864547bc4 100644 --- a/zig/src/families/qwen27/model.zig +++ b/zig/src/families/qwen27/model.zig @@ -9,6 +9,7 @@ const widths = @import("widths.zig"); const wts = @import("gpu_weights.zig"); const quant = @import("quant_gpu.zig"); const glue = @import("glue_gpu.zig"); +const moe_gpu = @import("moe_gpu.zig"); const Frame = @import("gpu_frame.zig").Frame; const ane_mlp = @import("ane_mlp.zig"); @@ -23,6 +24,8 @@ pub const Model = struct { weights: wts.Weights, kernels: quant.Kernels, glue: glue.Kernels, + /// The routed block's kernels, with experts. + moe: ?moe_gpu.Kernels = null, frame: Frame, /// The Neural Engine's share of prompt-chunk MLPs, when it loaded. ane: ?*ane_mlp.Share = null, @@ -35,8 +38,10 @@ pub const Model = struct { defer a.free(path); const text = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(16 << 20)); defer a.free(text); - const c = try cfg.parse(a, text, null); + var c = try cfg.parse(a, text, null); try supported(c); + // TF_LAYERS=N runs the first N layers only (a port's layer-by-layer check against a reference) + if (std.c.getenv("TF_LAYERS")) |v| c.layers = @min(c.layers, std.fmt.parseInt(usize, std.mem.span(v), 10) catch c.layers); var formats = try affine.Formats.init(a, text); defer formats.deinit(); var ws = try widths.Widths.resolve(a, &formats, c); @@ -68,6 +73,8 @@ pub const Model = struct { errdefer m.kernels.deinit(); m.glue = try glue.Kernels.init(m.device); errdefer m.glue.deinit(); + m.moe = if (c.moe()) try moe_gpu.Kernels.init(a, m.device) else null; + errdefer if (m.moe) |*k| k.deinit(); m.frame = try Frame.init(m.device, c, @max(rows, promptRows())); errdefer m.frame.deinit(); m.ane = null; @@ -88,6 +95,7 @@ pub const Model = struct { pub fn deinit(m: *Model) void { if (m.ane) |share| share.deinit(); m.frame.deinit(); + if (m.moe) |*k| k.deinit(); m.glue.deinit(); m.kernels.deinit(); m.weights.deinit(); @@ -108,7 +116,10 @@ pub fn promptRows() u32 { /// The qualified shape and layer schedule; per-module widths are admitted by widths.admit once the chip is known. pub fn supported(c: cfg.Config) !void { - if (c.hidden != 5120 or c.intermediate != 17408 or c.layers != 64 or c.vocab != 248320 or c.heads != 24 or c.kv_heads != 4 or c.head_dim != 256 or c.k_heads != 16 or c.v_heads != 48 or c.dk != 128 or c.dv != 128 or c.conv_kernel != 4 or c.rope_dims != 64) return error.UnqualifiedQwenShape; + const dense = c.hidden == 5120 and c.intermediate == 17408 and c.layers == 64 and c.heads == 24 and c.kv_heads == 4 and c.k_heads == 16 and c.v_heads == 48 and !c.moe(); + // Ornith-1.5-122B: 256 experts of 1,024 and a shared expert of 1,024 over 3,072, 32 heads over 2 of 256 + const routed = c.hidden == 3072 and c.intermediate == 1024 and c.layers == 48 and c.heads == 32 and c.kv_heads == 2 and c.k_heads == 16 and c.v_heads == 64 and c.experts == 256 and c.top_k == 8 and c.moe_intermediate == 1024; + if (!(dense or routed) or c.vocab != 248320 or c.head_dim != 256 or c.dk != 128 or c.dv != 128 or c.conv_kernel != 4 or c.rope_dims != 64) return error.UnqualifiedQwenShape; for (c.kinds[0..c.layers], 0..) |kind, i| if (kind != (if ((i + 1) % 4 == 0) cfg.Kind.attention else cfg.Kind.linear)) return error.UnqualifiedLayerSchedule; } diff --git a/zig/src/families/qwen27/moe_gpu.zig b/zig/src/families/qwen27/moe_gpu.zig new file mode 100644 index 0000000000..d89e820811 --- /dev/null +++ b/zig/src/families/qwen27/moe_gpu.zig @@ -0,0 +1,194 @@ +//! The routed block on the GPU: the router's row kernel, the route, the experts by id (row kernels) or sorted by +//! expert on the tensor units (core.affine_mm), their SwiGLU, and the combine over the shared expert's output. +const std = @import("std"); +const mtl = @import("metal"); +const core = @import("core"); +const row = core.row_projection; +const am = core.affine_mm; +const source = @import("qwen_runtime_sources"); +const profile = core.gpu_profile; +const Config = @import("config.zig").Config; +const wts = @import("gpu_weights.zig"); +const Frame = @import("gpu_frame.zig").Frame; +const Ref = @import("projection.zig").Ref; + +const RouteArgs = extern struct { experts: i32, topk: i32, ld: i32, renorm: i32 }; +const SortArgs = extern struct { slots: i32, experts: i32 }; +const GluArgs = extern struct { width: i32 }; +const CombineArgs = extern struct { hidden: i32, topk: i32, ld: i32, experts: i32 }; +const R = struct { buf: mtl.Buffer, off: usize = 0 }; +const Matrix = struct { w: R, s: R, b: R, n: u32, k: u32 }; +const Key = struct { bits: u8, group: u16 }; + +pub const Kernels = struct { + allocator: std.mem.Allocator, + device: mtl.Device, + route: mtl.Pipeline, + sort: mtl.Pipeline, + gather: mtl.Pipeline, + glu: mtl.Pipeline, + combine: mtl.Pipeline, + rows: ?row.Pipelines = null, + /// core.affine_mm's pipelines by the experts' width, built when a prompt first takes the tensor units. + nax: std.AutoHashMapUnmanaged(Key, am.Pipes) = .empty, + profiler: ?*profile.Trace = null, + + pub fn init(a: std.mem.Allocator, device: mtl.Device) !Kernels { + const lib = try mtl.Library.fromSource(device, source.moe, mtl.CompileOptions.mlx()); + defer lib.deinit(); + var pipes: [5]mtl.Pipeline = undefined; + var made: usize = 0; + errdefer for (pipes[0..made]) |p| p.deinit(); + for ([_][:0]const u8{ "q27_route", "q27_route_sort", "q27_route_gather", "q27_expert_glu", "q27_expert_combine" }, &pipes) |name, *p| { + p.* = try mtl.Pipeline.init(device, lib, name, false); + made += 1; + } + return .{ .allocator = a, .device = device, .route = pipes[0], .sort = pipes[1], .gather = pipes[2], .glu = pipes[3], .combine = pipes[4] }; + } + + pub fn deinit(k: *Kernels) void { + inline for (.{ k.route, k.sort, k.gather, k.glu, k.combine }) |p| p.deinit(); + if (k.rows) |*p| p.deinit(); + var it = k.nax.valueIterator(); + while (it.next()) |pipes| for (pipes) |p| p.deinit(); + k.nax.deinit(k.allocator); + } + + fn rowPipelines(k: *Kernels) !*const row.Pipelines { + if (k.rows == null) k.rows = try row.Pipelines.load(k.device); + return &k.rows.?; + } + + fn naxPipes(k: *Kernels, key: Key) !am.Pipes { + if (k.nax.get(key)) |p| return p; + const src = try am.source(k.device, k.allocator, .{ .bits = key.bits, .group = key.group }); + defer k.allocator.free(src); + const lib = try mtl.Library.fromSource(k.device, src, mtl.CompileOptions.mlx()); + defer lib.deinit(); + var pipes: am.Pipes = undefined; + var made: usize = 0; + errdefer for (pipes[0..made]) |p| p.deinit(); + for (am.names, &pipes) |name, *p| { + p.* = try mtl.Pipeline.init(k.device, lib, name, false); + made += 1; + } + try k.nax.put(k.allocator, key, pipes); + return pipes; + } + + /// Whether a pass of `rows` takes the sorted tensor-unit path: a prompt chunk of whole tiles, experts of one width. + pub fn sorted(k: *Kernels, mo: *const wts.Moe, rows: u32, prompt: bool) bool { + if (!prompt or rows < 16 or !k.device.tensorUnits()) return false; + const g = mo.gate.one; + return g.bits == mo.up.one.bits and g.group == mo.up.one.group and g.bits == mo.down.one.bits and g.group == mo.down.one.group; + } + + /// f.output, the shared expert's rows of f.input's normed rows, becomes the block's: the routed experts' weighted + /// rows plus the gated shared one. + pub fn block(k: *Kernels, e: mtl.ComputeEncoder, c: Config, mo: *const wts.Moe, f: *const Frame, rows: u32, prompt: bool) !void { + var span = profile.begin(k.profiler, e, .activation); + errdefer span.cancel(); + const enc = span.encoder; + const slots = f.slotsOf(rows); + const experts: u32 = @intCast(c.experts); + const topk: u32 = @intCast(c.top_k); + const hidden: u32 = @intCast(c.hidden); + const inter: u32 = @intCast(c.moe_intermediate); + const ld: u32 = @intCast(mo.router.n); + // the router: bf16 logits [rows, ld], the shared gate's in column `experts` + const pipes = try k.rowPipelines(); + const call = try row.call(pipes, mo.router, rows, false); + enc.setPipeline(call.pipeline); + enc.setBuffer(f.get(.input).buffer, 0, 0); + enc.setBuffer(mo.router.w, mo.router.w_off, 1); + enc.setBuffer(mo.router.scales, mo.router.s_off, 2); + enc.setBuffer(mo.router.biases, mo.router.b_off, 3); + enc.setValue(call.dims, 4); + enc.setBuffer(f.get(.router).buffer, 0, 5); + enc.dispatchGroups(.{ .width = call.groups }, .{ .width = call.threads }); + enc.barrier(); + enc.setPipeline(k.route); + enc.setBuffer(f.get(.router).buffer, 0, 0); + enc.setBuffer(f.get(.picks).buffer, 0, 1); + enc.setBuffer(f.get(.route_weights).buffer, 0, 2); + enc.setValue(RouteArgs{ .experts = @intCast(experts), .topk = @intCast(topk), .ld = @intCast(ld), .renorm = @intFromBool(c.renorm) }, 3); + enc.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(@max(32, (experts + 31) / 32 * 32), 1, 1)); + enc.barrier(); + if (k.sorted(mo, rows, prompt)) { + const p = try k.naxPipes(.{ .bits = mo.gate.one.bits, .group = @intCast(mo.gate.one.group) }); + enc.setPipeline(k.sort); + enc.setBuffer(f.get(.picks).buffer, 0, 0); + enc.setBuffer(f.get(.route_weights).buffer, 0, 1); + enc.setBuffer(f.get(.offsets).buffer, 0, 2); + enc.setBuffer(f.get(.order).buffer, 0, 3); + enc.setBuffer(f.get(.sorted_weights).buffer, 0, 4); + enc.setValue(SortArgs{ .slots = @intCast(slots), .experts = @intCast(experts) }, 5); + enc.dispatchThreads(mtl.Size.of(1, 1, 1), mtl.Size.of(1, 1, 1)); + enc.barrier(); + enc.setPipeline(k.gather); + enc.setBuffer(f.get(.input).buffer, 0, 0); + enc.setBuffer(f.get(.order).buffer, 0, 1); + enc.setBuffer(f.get(.sorted_rows).buffer, 0, 2); + enc.setValue([2]i32{ @intCast(hidden), @intCast(topk) }, 3); + enc.dispatchThreads(mtl.Size.of(hidden, slots, 1), mtl.Size.of(256, 1, 1)); + enc.barrier(); + const sums = R{ .buf = f.get(.expert_sums).buffer }; + const offsets = R{ .buf = f.get(.offsets).buffer }; + const x = R{ .buf = f.get(.sorted_rows).buffer }; + am.rowSums(enc, p, @intCast(mo.gate.one.group), x, sums, slots, hidden); + enc.barrier(); + am.gather(enc, p, x, sums, matrix(mo.gate.one), offsets, R{ .buf = f.get(.expert_gate).buffer }, slots, experts); + am.gather(enc, p, x, sums, matrix(mo.up.one), offsets, R{ .buf = f.get(.expert_up).buffer }, slots, experts); + enc.barrier(); + k.activation(enc, f, slots, inter); + enc.barrier(); + const act = R{ .buf = f.get(.expert_act).buffer }; + am.rowSums(enc, p, @intCast(mo.down.one.group), act, sums, slots, inter); + enc.barrier(); + am.gatherTo(enc, p, act, sums, matrix(mo.down.one), offsets, R{ .buf = f.get(.expert_out).buffer }, R{ .buf = f.get(.order).buffer }, slots, experts); + } else { + try k.indexed(enc, pipes, mo.gate, slots, f.get(.input), f.get(.picks), f.get(.expert_gate)); + try k.indexed(enc, pipes, mo.up, slots, f.get(.input), f.get(.picks), f.get(.expert_up)); + enc.barrier(); + k.activation(enc, f, slots, inter); + enc.barrier(); + try k.indexed(enc, pipes, mo.down, slots, f.get(.expert_act), f.get(.picks), f.get(.expert_out)); + } + enc.barrier(); + enc.setPipeline(k.combine); + enc.setBuffer(f.get(.expert_out).buffer, 0, 0); + enc.setBuffer(f.get(.route_weights).buffer, 0, 1); + enc.setBuffer(f.get(.router).buffer, 0, 2); + enc.setBuffer(f.get(.output).buffer, 0, 3); + enc.setValue(CombineArgs{ .hidden = @intCast(hidden), .topk = @intCast(topk), .ld = @intCast(ld), .experts = @intCast(experts) }, 4); + enc.dispatchThreads(mtl.Size.of(hidden, rows, 1), mtl.Size.of(256, 1, 1)); + try span.end(); + } + + fn activation(k: *Kernels, e: mtl.ComputeEncoder, f: *const Frame, slots: u32, inter: u32) void { + e.setPipeline(k.glu); + e.setBuffer(f.get(.expert_gate).buffer, 0, 0); + e.setBuffer(f.get(.expert_up).buffer, 0, 1); + e.setBuffer(f.get(.expert_act).buffer, 0, 2); + e.setValue(GluArgs{ .width = @intCast(inter) }, 3); + e.dispatchThreads(mtl.Size.of(inter, slots, 1), mtl.Size.of(256, 1, 1)); + } + + /// Slot i takes expert ids[i] on input row i / repeat. + fn indexed(_: *Kernels, e: mtl.ComputeEncoder, pipes: *const row.Pipelines, x: row.Experts, slots: u32, in: Ref, ids: Ref, out: Ref) !void { + const call = try row.callIndexed(pipes, x, slots, false); + e.setPipeline(call.pipeline); + e.setBuffer(in.buffer, in.offset, 0); + e.setBuffer(x.one.w, x.one.w_off, 1); + e.setBuffer(x.one.scales, x.one.s_off, 2); + e.setBuffer(x.one.biases, x.one.b_off, 3); + e.setValue(call.dims, 4); + e.setBuffer(out.buffer, out.offset, 5); + e.setBuffer(ids.buffer, ids.offset, 6); + e.dispatchThreads(mtl.Size.of(call.groups * call.threads, call.slots, 1), mtl.Size.of(call.threads, 1, 1)); + } +}; + +fn matrix(w: row.Weights) Matrix { + return .{ .w = .{ .buf = w.w, .off = w.w_off }, .s = .{ .buf = w.scales, .off = w.s_off }, .b = .{ .buf = w.biases, .off = w.b_off }, .n = @intCast(w.n), .k = @intCast(w.k) }; +} diff --git a/zig/src/families/qwen27/widths.zig b/zig/src/families/qwen27/widths.zig index e1f9b7a64f..d9c84124d1 100644 --- a/zig/src/families/qwen27/widths.zig +++ b/zig/src/families/qwen27/widths.zig @@ -17,14 +17,22 @@ pub fn isUniform(w: Width) bool { } /// A layer's projections; a layer stores entries for both mixers and uses those of its kind. -pub const Member = enum { gate, up, down, qkv, z, b, a, out, q, k, v, o }; -const members = @backingInt(Member.o) + 1; +pub const Member = enum { gate, up, down, qkv, z, b, a, out, q, k, v, o, router, shared_gate, egate, eup, edown }; +const members = @backingInt(Member.edown) + 1; +/// The routed block's members: the router, the shared expert's gate and the stacked experts. +pub const routed = [_]Member{ .router, .shared_gate, .egate, .eup, .edown }; -pub fn leaf(m: Member) []const u8 { +/// A member's module under its layer; with experts, `gate`, `up` and `down` are the shared expert's. +pub fn leaf(c: cfg.Config, m: Member) []const u8 { return switch (m) { - .gate => "mlp.gate_proj", - .up => "mlp.up_proj", - .down => "mlp.down_proj", + .gate => if (c.moe()) "mlp.shared_expert.gate_proj" else "mlp.gate_proj", + .up => if (c.moe()) "mlp.shared_expert.up_proj" else "mlp.up_proj", + .down => if (c.moe()) "mlp.shared_expert.down_proj" else "mlp.down_proj", + .router => "mlp.gate", + .shared_gate => "mlp.shared_expert_gate", + .egate => "mlp.switch_mlp.gate_proj", + .eup => "mlp.switch_mlp.up_proj", + .edown => "mlp.switch_mlp.down_proj", .qkv => "linear_attn.in_proj_qkv", .z => "linear_attn.in_proj_z", .b => "linear_attn.in_proj_b", @@ -37,11 +45,15 @@ pub fn leaf(m: Member) []const u8 { }; } -/// Output rows and input width of a member. +/// Output rows and input width of a member (one expert's, for the stacked experts). pub fn shape(c: cfg.Config, m: Member) [2]usize { return switch (m) { .gate, .up => .{ c.intermediate, c.hidden }, .down => .{ c.hidden, c.intermediate }, + .router => .{ c.experts, c.hidden }, + .shared_gate => .{ 1, c.hidden }, + .egate, .eup => .{ c.moe_intermediate, c.hidden }, + .edown => .{ c.hidden, c.moe_intermediate }, .qkv => .{ c.gdnQkvDim(), c.hidden }, .z => .{ c.gdnValueDim(), c.hidden }, .b, .a => .{ c.v_heads, c.hidden }, @@ -85,7 +97,10 @@ pub const Widths = struct { for (layers, 0..) |*l, i| { l.* = @splat(uniform); for (stacks(c.kind(i))) |stack| for (stack) |m| { - l[@backingInt(m)] = try one(f, try std.fmt.bufPrint(&name, "model.layers.{d}.{s}", .{ i, leaf(m) })); + l[@backingInt(m)] = try one(f, try std.fmt.bufPrint(&name, "model.layers.{d}.{s}", .{ i, leaf(c, m) })); + }; + if (c.moe()) for (routed) |m| { + l[@backingInt(m)] = try one(f, try std.fmt.bufPrint(&name, "model.layers.{d}.{s}", .{ i, leaf(c, m) })); }; } return .{ .allocator = a, .layers = layers, .embed = try one(f, "model.embed_tokens"), .head = try one(f, "lm_head") }; @@ -151,11 +166,18 @@ pub fn admit(w: *const Widths, c: cfg.Config, chip: Chip, why: *Why) !void { var n: usize = 0; for (stack[r[0]..][0..r[1]]) |m| n += shape(c, m)[0]; check(chip, n, shape(c, stack[r[0]])[1], ws[r[0]]) catch |err| { - why.* = .{ .layer = i, .module = leaf(stack[r[0]]), .width = ws[r[0]], .reason = err }; + why.* = .{ .layer = i, .module = leaf(c, stack[r[0]]), .width = ws[r[0]], .reason = err }; return error.UnqualifiedQwenFormat; }; } }; + if (c.moe()) for (0..c.layers) |i| for (routed) |m| { + const x = w.of(i, m); + checkRouted(c, chip, m, x) catch |err| { + why.* = .{ .layer = i, .module = leaf(c, m), .width = x, .reason = err }; + return error.UnqualifiedQwenFormat; + }; + }; glue.checkEmbed(.{ .width = @intCast(c.hidden), .rows = 1, .vocab = @intCast(c.vocab), .bits = w.embed.bits, .group = w.embed.group, .metadata = .bf16 }) catch |err| { why.* = .{ .module = "embed_tokens", .width = w.embed, .reason = err }; return error.UnqualifiedQwenFormat; @@ -177,6 +199,24 @@ fn check(chip: Chip, n: usize, k: usize, w: Width) !void { return (row.Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = n, .k = k, .group = w.group }).validate(); } +/// The routed block's kernels: the router and the shared gate share one row matrix; the experts take the row kernels +/// by id and, on tensor units, core.affine_mm's gathers (whose format admits 2-6 and 8 bits at groups 32, 64 and 128). +fn checkRouted(c: cfg.Config, chip: Chip, m: Member, w: Width) !void { + const s = shape(c, m); + if (std.mem.indexOfScalar(u8, chip.row_bits, w.bits) == null) return error.UnsupportedRowBits; + const n = if (m == .router or m == .shared_gate) routerRows(c.experts) else s[0]; + try (row.Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = n, .k = s[1], .group = w.group, .bits = w.bits }).validate(); + if (m == .egate or m == .eup or m == .edown) { + if (chip.tensor_units and (w.group != 32 and w.group != 64 and w.group != 128)) return error.UnsupportedLaneGroup; + if (chip.tensor_units and (s[0] % 64 != 0 or s[1] % 64 != 0)) return error.UnsupportedExpertShape; + } +} + +/// The router matrix's rows: the experts, the shared expert's gate, then zero rows to a multiple of 8. +pub fn routerRows(experts: usize) usize { + return (experts + 1 + 7) / 8 * 8; +} + const all_rows = [_]u8{ 2, 3, 4, 5, 6, 8 }; fn testConfig(a: std.mem.Allocator, quantization: []const u8) ![]u8 { diff --git a/zig/tests/qwen27/run.zig b/zig/tests/qwen27/run.zig index 5d31a8efd7..faa32fbe23 100644 --- a/zig/tests/qwen27/run.zig +++ b/zig/tests/qwen27/run.zig @@ -48,6 +48,8 @@ fn emit(a: std.mem.Allocator, io: std.Io, dir: []const u8) !void { } const Difference = struct { position: usize, expected: ?u32, actual: u32, top2: q.session.TopTwo }; const Generated = struct { tokens: []u32, first_difference: ?Difference, top2: ?[]q.session.TopTwo }; +/// TF_LOGITS_DUMP names a file that takes every step's bf16 logits row (a check's teacher-forced distributions). +var logits_dump: std.ArrayList(u8) = .empty; fn generated(a: std.mem.Allocator, s: q.session.Session, count: usize, force: bool, expected: ?[]const u32, trace: ?*profile.Trace, profile_steps: usize, teacher: bool) !Generated { var out: std.ArrayList(u32) = .empty; var difference: ?Difference = null; @@ -55,6 +57,7 @@ fn generated(a: std.mem.Allocator, s: q.session.Session, count: usize, force: bo for (0..count) |i| { const token = try s.greedy(); if (margins) |values| values[i] = try q.session.topTwo(s.runner.model.frame.get(.logits).buffer.slice(u16, s.runner.model.config.vocab)); + if (std.c.getenv("TF_LOGITS_DUMP") != null) try logits_dump.appendSlice(a, std.mem.sliceAsBytes(s.runner.model.frame.get(.logits).buffer.slice(u16, s.runner.model.config.vocab))); if (expected) |wanted| if (difference == null and (i >= wanted.len or token != wanted[i])) { difference = .{ .position = i, .expected = if (i < wanted.len) wanted[i] else null, .actual = token, .top2 = try q.session.topTwo(s.runner.model.frame.get(.logits).buffer.slice(u16, s.runner.model.config.vocab)) }; }; @@ -144,6 +147,21 @@ pub fn main(init: std.process.Init) !void { if (profile_steps > 0) runner.profile(&trace); const decode_start = std.Io.Clock.awake.now(io); const generation = try generated(a, s, count, force, expected, if (profile_steps > 0) &trace else null, profile_steps, teacher); + if (std.c.getenv("TF_LOGITS_DUMP")) |v| try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = std.mem.span(v), .data = logits_dump.items }); + if (std.c.getenv("TF_ROUTE_DUMP")) |v| { + // the frame's last pass's picks and weights (its last layer): a port's routing check + const f = &s.runner.model.frame; + const n = f.slotsOf(@intCast(ids.len)); + const picks = std.mem.sliceAsBytes(f.get(.picks).buffer.slice(u32, n)); + const weights = std.mem.sliceAsBytes(f.get(.route_weights).buffer.slice(f32, n)); + const h = s.runner.model.config.hidden; + const out_rows = std.mem.sliceAsBytes(f.get(.output).buffer.slice(u16, ids.len * h)); + const expert_rows = std.mem.sliceAsBytes(f.get(.expert_out).buffer.slice(u16, n * h)); + const hid = std.mem.sliceAsBytes(f.get(.hidden).buffer.slice(u16, ids.len * h)); + const res = std.mem.sliceAsBytes(f.get(.residual).buffer.slice(u16, ids.len * h)); + const both = try std.mem.concat(a, u8, &.{ picks, weights, out_rows, expert_rows, hid, res }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = std.mem.span(v), .data = both }); + } const out = generation.tokens; const decode_ns = decode_start.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds(); const decoded = if (compare > 0) try s.fingerprint() else digest; From 68fa10522dd0ba8edf54cad6d69196bf9284eba4 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 20 Oct 2026 12:56:15 -0500 Subject: [PATCH 02/21] qwen27: register qwen3_5_moe on Metal; DFlash2 drafters of up to eight taps and a block of 16 (Ornith's) --- zig/src/cli/hub.zig | 1 + zig/src/families/qwen27/dflash/config.zig | 39 +++++++++++++------ zig/src/families/qwen27/dflash/generation.zig | 1 + zig/src/families/qwen27/dflash/session.zig | 4 +- zig/src/families/qwen27/glue_gpu.zig | 2 +- zig/src/families/qwen27/taps.zig | 21 +++++----- zig/src/native/metal.zig | 10 +++-- zig/src/native/qwen27_draft.zig | 8 +++- zig/src/native/qwen27_served.zig | 2 +- 9 files changed, 57 insertions(+), 31 deletions(-) diff --git a/zig/src/cli/hub.zig b/zig/src/cli/hub.zig index 4a90bace27..6f17d01012 100644 --- a/zig/src/cli/hub.zig +++ b/zig/src/cli/hub.zig @@ -14,6 +14,7 @@ const titles = [_]struct { model_type: []const u8, title: []const u8 }{ .{ .model_type = "qwen4_exp", .title = "Qwen3.8 Flash Next" }, .{ .model_type = "glm5_next", .title = "GLM-5.3-Flash" }, .{ .model_type = "qwen3_5", .title = "Qwen3.5-2B or Qwen3.8-27B" }, + .{ .model_type = "qwen3_5_moe", .title = "Qwen3.5-MoE (Ornith-1.5-122B)" }, .{ .model_type = "gemma4", .title = "Gemma 4 31B" }, .{ .model_type = "gemma4_unified", .title = "Gemma 4 12B" }, .{ .model_type = "step5_vl", .title = "Step-5-Preview" }, diff --git a/zig/src/families/qwen27/dflash/config.zig b/zig/src/families/qwen27/dflash/config.zig index 92a5dea74f..59621fb860 100644 --- a/zig/src/families/qwen27/dflash/config.zig +++ b/zig/src/families/qwen27/dflash/config.zig @@ -1,7 +1,12 @@ -//! DFlash2 is a separate five-layer GQA contract; the target's GDN and draft conventions are not interchangeable. +//! DFlash2 is a separate GQA contract (five layers tapping the 27B at five blocks, six tapping Ornith at eight); the +//! target's GDN and draft conventions are not interchangeable. const std = @import("std"); pub const repository = "z-lab/Qwen3.8-27B-DFlash2"; -pub const tap_ids = [5]u32{ 5, 19, 33, 47, 61 }; +pub const max_taps = 8; +/// The 27B drafter's taps (five); `tap_count` says how many of a config's `taps` are in use. +pub const tap_ids = [max_taps]u32{ 5, 19, 33, 47, 61, 0, 0, 0 }; +/// incoai/Ornith-1.5-122B-DFlash2: six layers over Ornith-1.5-122B's 48, a block of 16. +pub const ornith_tap_ids = [max_taps]u32{ 1, 7, 14, 20, 26, 32, 39, 45 }; pub const Config = struct { hidden: u32 = 5120, intermediate: u32 = 17408, @@ -20,7 +25,8 @@ pub const Config = struct { topk: u32 = 16, eps: f32 = 1e-6, theta: f64 = 10000000, - taps: [5]u32 = tap_ids, + taps: [max_taps]u32 = tap_ids, + tap_count: u32 = 5, pub fn qWidth(c: Config) u32 { return c.heads * c.head_dim; } @@ -28,7 +34,10 @@ pub const Config = struct { return c.kv_heads * c.head_dim; } pub fn tapWidth(c: Config) u32 { - return c.hidden * 5; + return c.hidden * c.tap_count; + } + pub fn tapIds(c: *const Config) []const u32 { + return c.taps[0..c.tap_count]; } pub fn dynamicWidth(c: Config) u32 { return 4 * (c.hidden / c.group); @@ -45,20 +54,24 @@ pub const Config = struct { pub fn check(c: Config) !void { const dimensions = [_]u32{ c.hidden, c.intermediate, c.layers, c.vocab, c.heads, c.kv_heads, c.head_dim, c.target_layers, c.group, c.rank, c.topk }; for (dimensions) |n| if (n == 0) return error.BadDraftConfig; - if (c.hidden % 64 != 0 or c.intermediate % 64 != 0 or c.hidden % c.group != 0 or c.heads % c.kv_heads != 0 or c.head_dim % 2 != 0 or c.block != 8 or c.conv != 2 or c.layers > 8 or c.mask >= c.vocab or c.topk > c.vocab or c.window < c.block) return error.BadDraftConfig; + if (c.hidden % 64 != 0 or c.intermediate % 64 != 0 or c.hidden % c.group != 0 or c.heads % c.kv_heads != 0 or c.head_dim % 2 != 0 or (c.block != 8 and c.block != 16) or c.conv != 2 or c.layers > 8 or c.mask >= c.vocab or c.topk > c.vocab or c.window < c.block) return error.BadDraftConfig; if (!std.math.isFinite(c.eps) or c.eps <= 0 or !std.math.isFinite(c.theta) or c.theta <= 0) return error.BadDraftConfig; - if (c.target_layers == 64 and !std.mem.eql(u32, &c.taps, &tap_ids)) return error.UnsupportedDraftTaps; - for (c.taps, 0..) |tap, i| for (c.taps[0..i]) |previous| if (tap == previous) return error.UnsupportedDraftTaps; - for (c.taps) |i| if (i >= c.target_layers) return error.BadDraftConfig; + if (c.tap_count == 0 or c.tap_count > max_taps) return error.BadDraftConfig; + if (c.target_layers == 64 and !std.mem.eql(u32, c.tapIds(), tap_ids[0..5])) return error.UnsupportedDraftTaps; + for (c.tapIds(), 0..) |tap, i| for (c.taps[0..i]) |previous| if (tap == previous) return error.UnsupportedDraftTaps; + for (c.tapIds()) |i| if (i >= c.target_layers) return error.BadDraftConfig; const q = try std.math.mul(u32, c.heads, c.head_dim); _ = try std.math.mul(u32, c.kv_heads, c.head_dim); - _ = try std.math.mul(u32, c.hidden, 5); + _ = try std.math.mul(u32, c.hidden, c.tap_count); _ = try std.math.mul(u32, c.hidden / c.group, 4); if (q % 64 != 0) return error.BadDraftConfig; } + /// The 27B's stock drafter, or Ornith-1.5-122B's: the two contracts the Metal runtime has run. pub fn requireStock(c: Config) !void { try c.check(); - if (c.hidden != 5120 or c.intermediate != 17408 or c.layers != 5 or c.vocab != 248320 or c.heads != 32 or c.kv_heads != 8 or c.head_dim != 128 or c.target_layers != 64 or c.mask != 248070 or c.window != 2048 or c.group != 16 or c.rank != 256 or c.topk != 16 or c.eps != 1e-6 or c.theta != 10000000) return error.NotStockDFlash2; + const stock = c.hidden == 5120 and c.intermediate == 17408 and c.layers == 5 and c.heads == 32 and c.kv_heads == 8 and c.target_layers == 64 and c.mask == 248070 and c.block == 8 and c.tap_count == 5; + const ornith = c.hidden == 3072 and c.intermediate == 9216 and c.layers == 6 and c.heads == 32 and c.kv_heads == 8 and c.target_layers == 48 and c.mask == 248077 and c.block == 16 and c.tap_count == 8 and std.mem.eql(u32, c.tapIds(), ornith_tap_ids[0..8]); + if (c.vocab != 248320 or c.head_dim != 128 or c.window != 2048 or c.group != 16 or c.conv != 2 or c.rank != 256 or c.topk != 16 or !(stock or ornith)) return error.UnsupportedDraftShape; } }; fn object(v: std.json.Value) !std.json.ObjectMap { @@ -102,7 +115,9 @@ pub fn parse(gpa: std.mem.Allocator, raw: []const u8) !Config { try text(rope, "rope_type", "default"); var c = Config{ .hidden = try integer(o, "hidden_size"), .intermediate = try integer(o, "intermediate_size"), .layers = try integer(o, "num_hidden_layers"), .vocab = try integer(o, "vocab_size"), .heads = try integer(o, "num_attention_heads"), .kv_heads = try integer(o, "num_key_value_heads"), .head_dim = try integer(o, "head_dim"), .target_layers = try integer(o, "num_target_layers"), .block = try integer(d, "block_size"), .mask = try integer(d, "mask_token_id"), .window = try integer(o, "sliding_window"), .group = try integer(d, "conv_group_size"), .conv = try integer(d, "conv_kernel_size"), .rank = try integer(d, "selector_rank"), .topk = try integer(d, "selector_top_k"), .eps = @floatCast(try number(o, "rms_norm_eps")), .theta = try number(rope, "rope_theta") }; const taps = d.get("target_layer_ids") orelse return error.BadDraftConfig; - if (taps != .array or taps.array.items.len != 5) return error.BadDraftConfig; + if (taps != .array or taps.array.items.len == 0 or taps.array.items.len > max_taps) return error.BadDraftConfig; + c.tap_count = @intCast(taps.array.items.len); + c.taps = @splat(0); for (taps.array.items, 0..) |v, i| { if (v != .integer or v.integer < 0 or v.integer > std.math.maxInt(u32)) return error.BadDraftConfig; c.taps[i] = @intCast(v.integer); @@ -128,7 +143,7 @@ test "tap convention is fixed after decoder blocks and cannot be shifted" { } test "block slots and finite dimensions are admitted before operators" { var c = Config{}; - c.block = 16; + c.block = 12; try std.testing.expectError(error.BadDraftConfig, c.check()); c = .{}; c.hidden = 0; diff --git a/zig/src/families/qwen27/dflash/generation.zig b/zig/src/families/qwen27/dflash/generation.zig index 489bda71ce..e2fad6160a 100644 --- a/zig/src/families/qwen27/dflash/generation.zig +++ b/zig/src/families/qwen27/dflash/generation.zig @@ -25,6 +25,7 @@ pub const Generation = struct { if (runner.model != draft.backend.target or runner.model.config.hidden != draft.graph.config.hidden or runner.model.config.layers != draft.graph.config.target_layers) return error.TargetBinding; runner.capture_taps = true; for (draft.graph.config.taps, &runner.taps.ids) |id, *to| to.* = id; + runner.taps.count = draft.graph.config.tap_count; const options = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; const kept = try runner.model.device.buffer(128 * 4, options); errdefer kept.deinit(); diff --git a/zig/src/families/qwen27/dflash/session.zig b/zig/src/families/qwen27/dflash/session.zig index 76c1853488..e07d05f663 100644 --- a/zig/src/families/qwen27/dflash/session.zig +++ b/zig/src/families/qwen27/dflash/session.zig @@ -5,7 +5,7 @@ const op = @import("operators.zig"); const weights = @import("weights.zig"); const execution = @import("execution.zig"); const selector = @import("selector.zig"); -pub const Batch = struct { values: op.View, positions: []const u64, layers: [5]u32, target_end: u64 }; +pub const Batch = struct { values: op.View, positions: []const u64, layers: [cfg.max_taps]u32, target_end: u64 }; pub const Session = struct { config: cfg.Config, context: op.Context, @@ -16,7 +16,7 @@ pub const Session = struct { return .{ .config = c, .context = .{ .cache = cache, .begin = first, .end = first } }; } pub fn admit(s: Session, batch: Batch) !op.Context { - if (!std.mem.eql(u32, &batch.layers, &s.config.taps) or batch.positions.len == 0) return error.BadCommittedTaps; + if (!std.mem.eql(u32, batch.layers[0..s.config.tap_count], s.config.tapIds()) or batch.positions.len == 0) return error.BadCommittedTaps; try batch.values.check(@intCast(batch.positions.len), s.config.tapWidth(), .bf16); const end = try std.math.add(u64, s.context.end, batch.positions.len); if (end != batch.target_end) return error.BadCommittedTaps; diff --git a/zig/src/families/qwen27/glue_gpu.zig b/zig/src/families/qwen27/glue_gpu.zig index d746877919..d1ee208a3a 100644 --- a/zig/src/families/qwen27/glue_gpu.zig +++ b/zig/src/families/qwen27/glue_gpu.zig @@ -104,7 +104,7 @@ pub const Kernels = struct { } pub fn tapRows(k: Kernels, e: mtl.ComputeEncoder, x: Ref, kept: Ref, y: Ref, p: Gather) !void { - if (p.rows == 0 or p.rows > p.capacity or p.planes != 5 or p.width == 0) return error.BadTapGather; + if (p.rows == 0 or p.rows > p.capacity or p.planes == 0 or p.planes > 8 or p.width == 0) return error.BadTapGather; if (x.offset > x.buffer.length() or @as(usize, p.capacity) * p.width * p.planes * 2 > x.buffer.length() - x.offset or kept.offset > kept.buffer.length() or @as(usize, p.rows) * 4 > kept.buffer.length() - kept.offset or y.offset > y.buffer.length() or @as(usize, p.rows) * p.width * p.planes * 2 > y.buffer.length() - y.offset) return error.BadTapGather; e.setPipeline(k.gather); bind(e, .{ x, kept }); diff --git a/zig/src/families/qwen27/taps.zig b/zig/src/families/qwen27/taps.zig index 23e04b5751..51237d2bd2 100644 --- a/zig/src/families/qwen27/taps.zig +++ b/zig/src/families/qwen27/taps.zig @@ -1,22 +1,25 @@ -//! Tap planes store native post-block residuals until accepted rows and their successor ids are known. +//! Tap planes (up to eight) store native post-block residuals until accepted rows and their successor ids are known. const std = @import("std"); const mtl = @import("metal"); const Ref = @import("projection.zig").Ref; const Glue = @import("glue_gpu.zig").Kernels; +/// The 27B drafter's tap layers, the default before a drafter sets its own. pub const layers = [_]usize{ 5, 19, 33, 47, 61 }; +pub const max_planes = 8; pub const Taps = struct { buffer: mtl.Buffer, capacity: u32, width: u32, - seen: u8 = 0, - ids: [5]usize = layers, + seen: u16 = 0, + ids: [max_planes]usize = layers ++ [_]usize{ 0, 0, 0 }, + count: usize = layers.len, rows: u32 = 0, pub fn init(device: mtl.Device, capacity: u32, width: u32) !Taps { if (capacity == 0 or capacity > @import("gpu_frame.zig").max_rows or width == 0) return error.BadTapShape; - const bytes = try std.math.mul(usize, try std.math.mul(usize, capacity, width), 2 * layers.len); + const bytes = try std.math.mul(usize, try std.math.mul(usize, capacity, width), 2 * max_planes); return .{ .buffer = try device.buffer(bytes, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked), .capacity = capacity, .width = width }; } @@ -31,9 +34,9 @@ pub const Taps = struct { pub fn record(t: *Taps, glue: Glue, e: mtl.ComputeEncoder, layer: usize, hidden: Ref, rows: u32) !void { if (rows == 0 or rows > t.capacity) return error.BadTapShape; - const index = std.mem.indexOfScalar(usize, &t.ids, layer) orelse return; + const index = std.mem.indexOfScalar(usize, t.ids[0..t.count], layer) orelse return; if (t.seen != 0 and rows != t.rows) return error.BadTapShape; - const bit = @as(u8, 1) << @intCast(index); + const bit = @as(u16, 1) << @intCast(index); if (t.seen & bit != 0) return error.DuplicateTap; try glue.unstack(e, hidden, try t.plane(index), .{ .rows = rows, .width = t.width, .stride = t.width, .offset = 0 }); e.barrier(); @@ -42,12 +45,12 @@ pub const Taps = struct { } pub fn plane(t: Taps, index: usize) !Ref { - if (index >= layers.len) return error.BadTapIndex; + if (index >= t.count) return error.BadTapIndex; return .{ .buffer = t.buffer, .offset = index * @as(usize, t.capacity) * t.width * 2 }; } pub fn complete(t: Taps) !void { - if (t.seen != (1 << layers.len) - 1) return error.IncompleteTaps; + if (t.count == 0 or t.count > max_planes or t.seen != (@as(u16, 1) << @intCast(t.count)) - 1) return error.IncompleteTaps; } pub fn accepted(t: Taps, glue: Glue, e: mtl.ComputeEncoder, path: []const u32, kept: Ref, out: Ref) !void { @@ -55,7 +58,7 @@ pub const Taps = struct { if (path.len == 0) return; if (path.len > t.rows) return error.BadTapGather; for (path) |row| if (row >= t.rows) return error.BadTapGather; - try glue.tapRows(e, .{ .buffer = t.buffer }, kept, out, .{ .rows = @intCast(path.len), .width = t.width, .capacity = t.capacity, .planes = layers.len }); + try glue.tapRows(e, .{ .buffer = t.buffer }, kept, out, .{ .rows = @intCast(path.len), .width = t.width, .capacity = t.capacity, .planes = @intCast(t.count) }); e.barrier(); } }; diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index c778274196..e152da2284 100644 --- a/zig/src/native/metal.zig +++ b/zig/src/native/metal.zig @@ -38,6 +38,7 @@ pub const families: []const api.Family = &.{ .{ .model_type = "qwen4_exp", .formats = affine_widths }, .{ .model_type = "glm5_next", .formats = &.{"mlx-q4g64"} }, .{ .model_type = "qwen3_5", .formats = affine_widths }, + .{ .model_type = "qwen3_5_moe", .formats = affine_widths }, .{ .model_type = "lfm2_moe", .formats = &.{ "mlx-q4g64", "mlx-q6g64", "mlx-q8g64" } }, .{ .model_type = "glm_moe_dsa", .formats = &.{"exlr8"} }, .{ .model_type = "gemma4", .formats = gemma4_widths }, @@ -153,13 +154,14 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c defer report.deinit(); if (std.mem.eql(u8, o.model_type, "fish_omni")) return @import("drama_open.zig").openMetal(a, gpa, io, o, problem); if (@import("save_admission.zig").refused(io, o.dir, problem)) return null; - const dense = std.mem.eql(u8, o.model_type, "qwen3_5") and qwen27Engine(a, io, o.dir); + // the 27B engine serves the dense 27B class and the routed Qwen3.5-MoE class (Ornith-1.5-122B) + const dense = (std.mem.eql(u8, o.model_type, "qwen3_5") and qwen27Engine(a, io, o.dir)) or std.mem.eql(u8, o.model_type, "qwen3_5_moe"); if (try refusal(a, o, dense)) |why| { problem.* = why; return null; } if (o.drafter != null and !dense and !isGemma4(o.model_type)) { - problem.* = "--drafter is supported for Qwen3.8-27B-class checkpoints (qwen3_5 with its own output head) and Gemma 4 dense"; + problem.* = "--drafter is supported for Qwen3.8-27B-class checkpoints (qwen3_5 with its own output head), Qwen3.5-MoE (qwen3_5_moe) and Gemma 4 dense"; return null; } if (dense) return openQwen27(a, gpa, io, o, problem); @@ -340,7 +342,7 @@ fn qwen27Engine(a: Allocator, io: std.Io, dir: []const u8) bool { fn openQwen27(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { const context = o.context orelse 32768; if (context <= 0 or context > 262144) { - problem.* = try std.fmt.allocPrint(a, "--context {d} is outside Qwen3.8-27B's 1 to 262,144 tokens", .{context}); + problem.* = try std.fmt.allocPrint(a, "--context {d} is outside this Qwen checkpoint's 1 to 262,144 tokens", .{context}); return null; } if (o.drafts and o.drafter != null and o.drafter_bits != 0 and o.drafter_bits != 4) { @@ -351,7 +353,7 @@ fn openQwen27(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[ const drafter = if (o.drafts) o.drafter else null; const learn: ?@import("qwen27_disk.zig").Learn = if (o.learn_conversations) if (o.learn) |root| .{ .root = root, .dir = o.dir, .drafter = drafter, .cap = gibBytes(o.learn_gib), .floor = gibBytes(o.learn_min_free_gib), .daily = gibBytes(o.learn_daily_gib) } else null else null; const h = qwen27.open(gpa, io, o.dir, .{ .context = @intCast(context), .streams = o.lanes, .streams_fixed = o.lanes_fixed, .drafter = drafter, .drafter_bits = o.drafter_bits, .cache_gib = o.prompt_cache_gib, .cache_over_cap = o.prompt_cache_over_cap, .learn = learn, .vision = o.vision }, a, &why) catch |err| { - problem.* = if (why.len > 0) why else try std.fmt.allocPrint(a, "the native Qwen3.8-27B engine cannot load {s} ({s})", .{ o.dir, @errorName(err) }); + problem.* = if (why.len > 0) why else try std.fmt.allocPrint(a, "the native Qwen engine (27B class) cannot load {s} ({s})", .{ o.dir, @errorName(err) }); return null; }; return .{ .engine = h.engine(), .close = qwen27.close, .ctx = h }; diff --git a/zig/src/native/qwen27_draft.zig b/zig/src/native/qwen27_draft.zig index 2b62315e8c..fb075431a1 100644 --- a/zig/src/native/qwen27_draft.zig +++ b/zig/src/native/qwen27_draft.zig @@ -10,7 +10,8 @@ pub const Draft = struct { owns_model: bool, generation: df.generation.Generation, previous_capture: bool, - previous_taps: [5]usize, + previous_taps: [q.taps.max_planes]usize, + previous_count: usize, pub fn open(gpa: std.mem.Allocator, io: std.Io, target: *q.model.Model, runner: *q.decode_round.Runner, dir: []const u8, mode: df.operators.Mode) !*Draft { if (runner.model != target) return error.TargetBinding; @@ -28,12 +29,14 @@ pub const Draft = struct { errdefer gpa.destroy(d); const previous_capture = runner.capture_taps; const previous_taps = runner.taps.ids; + const previous_count = runner.taps.count; const generation = df.generation.Generation.init(gpa, runner, model) catch |err| { runner.capture_taps = previous_capture; runner.taps.ids = previous_taps; + runner.taps.count = previous_count; return err; }; - d.* = .{ .gpa = gpa, .runner = runner, .model = model, .owns_model = owns_model, .generation = generation, .previous_capture = previous_capture, .previous_taps = previous_taps }; + d.* = .{ .gpa = gpa, .runner = runner, .model = model, .owns_model = owns_model, .generation = generation, .previous_capture = previous_capture, .previous_taps = previous_taps, .previous_count = previous_count }; return d; } pub fn reset(d: *Draft) !void { @@ -63,6 +66,7 @@ pub const Draft = struct { d.generation.deinit(); d.runner.capture_taps = d.previous_capture; d.runner.taps.ids = d.previous_taps; + d.runner.taps.count = d.previous_count; if (d.owns_model) d.model.deinit(); d.gpa.destroy(d); } diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index 2e33bc3168..4ea2f2d480 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -177,7 +177,7 @@ fn slotCost(model: *const q.model.Model, context: u32, drafting: bool) usize { var attention: usize = 0; for (0..c.layers) |i| attention += @intFromBool(c.kind(i) == .attention); const kv = attention * 2 * @as(usize, context) * c.kvDim() * 2; - const ring_bytes: usize = if (drafting) 2048 * 5 * c.hidden * 2 else 0; + const ring_bytes: usize = if (drafting) 2048 * 8 * c.hidden * 2 else 0; // a drafter's window of up to eight taps return kv + ring_bytes + (256 << 20); // DeltaNet states, conv tails and scratch, generously } From 15595537a8fe5204f96b20c80729fdf72ad49658 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 20 Oct 2026 13:00:35 -0500 Subject: [PATCH 03/21] qwen27: drop the port's debug switches; fmt --- zig/src/families/qwen27/forward.zig | 2 +- zig/src/families/qwen27/gpu_frame.zig | 11 +++++------ zig/tests/qwen27/run.zig | 15 +-------------- 3 files changed, 7 insertions(+), 21 deletions(-) diff --git a/zig/src/families/qwen27/forward.zig b/zig/src/families/qwen27/forward.zig index ab2f647505..c1e8831607 100644 --- a/zig/src/families/qwen27/forward.zig +++ b/zig/src/families/qwen27/forward.zig @@ -119,7 +119,7 @@ pub const Stages = struct { lane.enc.barrier(); if (layer.moe) |*mo| { // the shared expert's rows in f.output take the routed experts' weighted rows and the shared gate - if (std.c.getenv("TF_NO_ROUTED") == null) try m.moe.?.block(lane.enc, m.config, mo, f, s.rows, m.kernels.prompt); + try m.moe.?.block(lane.enc, m.config, mo, f, s.rows, m.kernels.prompt); lane.enc.barrier(); } s.pending = f.get(.output); diff --git a/zig/src/families/qwen27/gpu_frame.zig b/zig/src/families/qwen27/gpu_frame.zig index 9171196e92..d21456d6dc 100644 --- a/zig/src/families/qwen27/gpu_frame.zig +++ b/zig/src/families/qwen27/gpu_frame.zig @@ -36,12 +36,11 @@ pub const Frame = struct { capacity * config.conv_kernel * 4, capacity * 4, capacity * 16, capacity * 16, capacity * 4, capacity * config.kvDim() * 2, capacity * 4, capacity * @sizeOf(@import("core").tree_slots.Step), - // the routed block: a row's logits, its picks and weights, the sorted slots, the experts' rows - capacity * routerRows(config) * 2, slots * 4, - slots * 4, (config.experts + 1) * 4, - slots * 4, slots * 4, - slots * config.hidden * 2, slots * config.moe_intermediate * 2, - slots * config.moe_intermediate * 2, slots * config.moe_intermediate * 2, + // the routed block: a row's logits, its picks and weights, the sorted slots, the experts' rows + capacity * routerRows(config) * 2, + slots * 4, slots * 4, (config.experts + 1) * 4, + slots * 4, slots * 4, slots * config.hidden * 2, + slots * config.moe_intermediate * 2, slots * config.moe_intermediate * 2, slots * config.moe_intermediate * 2, slots * config.hidden * 2, slots * @max(config.hidden, config.moe_intermediate) / 64 * 4, }; comptime std.debug.assert(widths.len == std.enums.values(Field).len); diff --git a/zig/tests/qwen27/run.zig b/zig/tests/qwen27/run.zig index faa32fbe23..b29413f2d6 100644 --- a/zig/tests/qwen27/run.zig +++ b/zig/tests/qwen27/run.zig @@ -18,6 +18,7 @@ const help = \\--compare-chunk reruns the whole prompt and decode and checks prompt state/KV/logits bytes. \\--seams N (with --compare-chunk): state/KV/logits bytes every N prompt rows under both chunks and after decode. \\CPU only: --help or --emit-sources EXISTING_DIR, no Metal device or model load. + \\TF_LAYERS=N runs the first N layers only; TF_LOGITS_DUMP=FILE writes every step's bf16 logits row (a port's checks). ; fn tokens(a: std.mem.Allocator, io: std.Io, file: []const u8) ![]const u32 { @@ -148,20 +149,6 @@ pub fn main(init: std.process.Init) !void { const decode_start = std.Io.Clock.awake.now(io); const generation = try generated(a, s, count, force, expected, if (profile_steps > 0) &trace else null, profile_steps, teacher); if (std.c.getenv("TF_LOGITS_DUMP")) |v| try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = std.mem.span(v), .data = logits_dump.items }); - if (std.c.getenv("TF_ROUTE_DUMP")) |v| { - // the frame's last pass's picks and weights (its last layer): a port's routing check - const f = &s.runner.model.frame; - const n = f.slotsOf(@intCast(ids.len)); - const picks = std.mem.sliceAsBytes(f.get(.picks).buffer.slice(u32, n)); - const weights = std.mem.sliceAsBytes(f.get(.route_weights).buffer.slice(f32, n)); - const h = s.runner.model.config.hidden; - const out_rows = std.mem.sliceAsBytes(f.get(.output).buffer.slice(u16, ids.len * h)); - const expert_rows = std.mem.sliceAsBytes(f.get(.expert_out).buffer.slice(u16, n * h)); - const hid = std.mem.sliceAsBytes(f.get(.hidden).buffer.slice(u16, ids.len * h)); - const res = std.mem.sliceAsBytes(f.get(.residual).buffer.slice(u16, ids.len * h)); - const both = try std.mem.concat(a, u8, &.{ picks, weights, out_rows, expert_rows, hid, res }); - try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = std.mem.span(v), .data = both }); - } const out = generation.tokens; const decode_ns = decode_start.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds(); const decoded = if (compare > 0) try s.fingerprint() else digest; From 5ef0dbce572073c988cf8787ef44d3b2b3036e98 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 20 Oct 2026 13:08:02 -0500 Subject: [PATCH 04/21] qwen27: prompt rows route on the tensor units; the profile times the routed block (route, experts) and TF_PROFILE_PROMPT --- zig/src/core/gpu_profile.zig | 2 +- zig/src/families/qwen27/decode_round.zig | 1 + zig/src/families/qwen27/gpu_frame.zig | 3 +- zig/src/families/qwen27/moe_gpu.zig | 123 +++++++++++++---------- zig/tests/qwen27/run.zig | 7 +- 5 files changed, 78 insertions(+), 58 deletions(-) diff --git a/zig/src/core/gpu_profile.zig b/zig/src/core/gpu_profile.zig index 531bb24fd3..aac21fe3bf 100644 --- a/zig/src/core/gpu_profile.zig +++ b/zig/src/core/gpu_profile.zig @@ -1,7 +1,7 @@ //! Diagnostic class timings split command buffers, wait for each span and keep ordinary encoders when disabled. const std = @import("std"); const mtl = @import("metal"); -pub const Class = enum { dense, norm, xsum, embedding, copy, activation, gdn_prepare, gdn_scan, gdn_gate, state_publish, attention_prepare, attention }; +pub const Class = enum { dense, norm, xsum, embedding, copy, activation, gdn_prepare, gdn_scan, gdn_gate, state_publish, attention_prepare, attention, route, experts }; pub const Shape = struct { n: usize = 0, k: usize = 0, sk: usize = 0, bits: u8 = 0, calls: u64 = 0, gpu_seconds: f64 = 0 }; pub const Trace = struct { queue: mtl.Queue, diff --git a/zig/src/families/qwen27/decode_round.zig b/zig/src/families/qwen27/decode_round.zig index 4b5e921987..76abf3afea 100644 --- a/zig/src/families/qwen27/decode_round.zig +++ b/zig/src/families/qwen27/decode_round.zig @@ -173,6 +173,7 @@ pub const Runner = struct { r.model.glue.profiler = trace; r.gdn.profiler = trace; r.attention.profiler = trace; + if (r.model.moe) |*k| k.profiler = trace; } pub fn verify(r: *Runner, round: *const Round) !void { diff --git a/zig/src/families/qwen27/gpu_frame.zig b/zig/src/families/qwen27/gpu_frame.zig index d21456d6dc..0fb86d8684 100644 --- a/zig/src/families/qwen27/gpu_frame.zig +++ b/zig/src/families/qwen27/gpu_frame.zig @@ -9,8 +9,9 @@ pub const Field = enum { ids, dims, positions, hidden, residual, input, sums, mi /// A prompt pass's rows at most (a chain); rounds with trees and every logits row stay within `logit_rows`. pub const max_rows = 2048; +/// The router's output pitch: its rows rounded up to the tensor units' 64 (core.affine_mm.dense). fn routerRows(config: Config) usize { - return if (config.moe()) @import("widths.zig").routerRows(config.experts) else 0; + return if (config.moe()) std.mem.alignForward(usize, @import("widths.zig").routerRows(config.experts), 64) else 0; } pub const logit_rows = 128; diff --git a/zig/src/families/qwen27/moe_gpu.zig b/zig/src/families/qwen27/moe_gpu.zig index d89e820811..4ad5c72e53 100644 --- a/zig/src/families/qwen27/moe_gpu.zig +++ b/zig/src/families/qwen27/moe_gpu.zig @@ -86,7 +86,7 @@ pub const Kernels = struct { /// f.output, the shared expert's rows of f.input's normed rows, becomes the block's: the routed experts' weighted /// rows plus the gated shared one. pub fn block(k: *Kernels, e: mtl.ComputeEncoder, c: Config, mo: *const wts.Moe, f: *const Frame, rows: u32, prompt: bool) !void { - var span = profile.begin(k.profiler, e, .activation); + var span = profile.begin(k.profiler, e, .route); errdefer span.cancel(); const enc = span.encoder; const slots = f.slotsOf(rows); @@ -94,18 +94,29 @@ pub const Kernels = struct { const topk: u32 = @intCast(c.top_k); const hidden: u32 = @intCast(c.hidden); const inter: u32 = @intCast(c.moe_intermediate); - const ld: u32 = @intCast(mo.router.n); - // the router: bf16 logits [rows, ld], the shared gate's in column `experts` + const units = k.sorted(mo, rows, prompt); + // the router: bf16 logits [rows, ld], the shared gate's in column `experts`; a prompt chunk's rows on the + // tensor units (y's pitch is n rounded up to 64), decode rows through the row kernel (pitch n) + const ld: u32 = if (units) std.mem.alignForward(u32, @intCast(mo.router.n), 64) else @intCast(mo.router.n); const pipes = try k.rowPipelines(); - const call = try row.call(pipes, mo.router, rows, false); - enc.setPipeline(call.pipeline); - enc.setBuffer(f.get(.input).buffer, 0, 0); - enc.setBuffer(mo.router.w, mo.router.w_off, 1); - enc.setBuffer(mo.router.scales, mo.router.s_off, 2); - enc.setBuffer(mo.router.biases, mo.router.b_off, 3); - enc.setValue(call.dims, 4); - enc.setBuffer(f.get(.router).buffer, 0, 5); - enc.dispatchGroups(.{ .width = call.groups }, .{ .width = call.threads }); + if (units) { + const p = try k.naxPipes(.{ .bits = mo.router.bits, .group = @intCast(mo.router.group) }); + const x = R{ .buf = f.get(.input).buffer }; + const sums = R{ .buf = f.get(.expert_sums).buffer }; + am.rowSums(enc, p, @intCast(mo.router.group), x, sums, rows, hidden); + enc.barrier(); + am.dense(enc, p, x, sums, matrix(mo.router), R{ .buf = f.get(.router).buffer }, rows); + } else { + const call = try row.call(pipes, mo.router, rows, false); + enc.setPipeline(call.pipeline); + enc.setBuffer(f.get(.input).buffer, 0, 0); + enc.setBuffer(mo.router.w, mo.router.w_off, 1); + enc.setBuffer(mo.router.scales, mo.router.s_off, 2); + enc.setBuffer(mo.router.biases, mo.router.b_off, 3); + enc.setValue(call.dims, 4); + enc.setBuffer(f.get(.router).buffer, 0, 5); + enc.dispatchGroups(.{ .width = call.groups }, .{ .width = call.threads }); + } enc.barrier(); enc.setPipeline(k.route); enc.setBuffer(f.get(.router).buffer, 0, 0); @@ -114,55 +125,59 @@ pub const Kernels = struct { enc.setValue(RouteArgs{ .experts = @intCast(experts), .topk = @intCast(topk), .ld = @intCast(ld), .renorm = @intFromBool(c.renorm) }, 3); enc.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(@max(32, (experts + 31) / 32 * 32), 1, 1)); enc.barrier(); - if (k.sorted(mo, rows, prompt)) { + try span.end(); + var products = profile.begin(k.profiler, e, .experts); + errdefer products.cancel(); + const enc2 = products.encoder; + if (units) { const p = try k.naxPipes(.{ .bits = mo.gate.one.bits, .group = @intCast(mo.gate.one.group) }); - enc.setPipeline(k.sort); - enc.setBuffer(f.get(.picks).buffer, 0, 0); - enc.setBuffer(f.get(.route_weights).buffer, 0, 1); - enc.setBuffer(f.get(.offsets).buffer, 0, 2); - enc.setBuffer(f.get(.order).buffer, 0, 3); - enc.setBuffer(f.get(.sorted_weights).buffer, 0, 4); - enc.setValue(SortArgs{ .slots = @intCast(slots), .experts = @intCast(experts) }, 5); - enc.dispatchThreads(mtl.Size.of(1, 1, 1), mtl.Size.of(1, 1, 1)); - enc.barrier(); - enc.setPipeline(k.gather); - enc.setBuffer(f.get(.input).buffer, 0, 0); - enc.setBuffer(f.get(.order).buffer, 0, 1); - enc.setBuffer(f.get(.sorted_rows).buffer, 0, 2); - enc.setValue([2]i32{ @intCast(hidden), @intCast(topk) }, 3); - enc.dispatchThreads(mtl.Size.of(hidden, slots, 1), mtl.Size.of(256, 1, 1)); - enc.barrier(); + enc2.setPipeline(k.sort); + enc2.setBuffer(f.get(.picks).buffer, 0, 0); + enc2.setBuffer(f.get(.route_weights).buffer, 0, 1); + enc2.setBuffer(f.get(.offsets).buffer, 0, 2); + enc2.setBuffer(f.get(.order).buffer, 0, 3); + enc2.setBuffer(f.get(.sorted_weights).buffer, 0, 4); + enc2.setValue(SortArgs{ .slots = @intCast(slots), .experts = @intCast(experts) }, 5); + enc2.dispatchThreads(mtl.Size.of(1, 1, 1), mtl.Size.of(1, 1, 1)); + enc2.barrier(); + enc2.setPipeline(k.gather); + enc2.setBuffer(f.get(.input).buffer, 0, 0); + enc2.setBuffer(f.get(.order).buffer, 0, 1); + enc2.setBuffer(f.get(.sorted_rows).buffer, 0, 2); + enc2.setValue([2]i32{ @intCast(hidden), @intCast(topk) }, 3); + enc2.dispatchThreads(mtl.Size.of(hidden, slots, 1), mtl.Size.of(256, 1, 1)); + enc2.barrier(); const sums = R{ .buf = f.get(.expert_sums).buffer }; const offsets = R{ .buf = f.get(.offsets).buffer }; const x = R{ .buf = f.get(.sorted_rows).buffer }; - am.rowSums(enc, p, @intCast(mo.gate.one.group), x, sums, slots, hidden); - enc.barrier(); - am.gather(enc, p, x, sums, matrix(mo.gate.one), offsets, R{ .buf = f.get(.expert_gate).buffer }, slots, experts); - am.gather(enc, p, x, sums, matrix(mo.up.one), offsets, R{ .buf = f.get(.expert_up).buffer }, slots, experts); - enc.barrier(); - k.activation(enc, f, slots, inter); - enc.barrier(); + am.rowSums(enc2, p, @intCast(mo.gate.one.group), x, sums, slots, hidden); + enc2.barrier(); + am.gather(enc2, p, x, sums, matrix(mo.gate.one), offsets, R{ .buf = f.get(.expert_gate).buffer }, slots, experts); + am.gather(enc2, p, x, sums, matrix(mo.up.one), offsets, R{ .buf = f.get(.expert_up).buffer }, slots, experts); + enc2.barrier(); + k.activation(enc2, f, slots, inter); + enc2.barrier(); const act = R{ .buf = f.get(.expert_act).buffer }; - am.rowSums(enc, p, @intCast(mo.down.one.group), act, sums, slots, inter); - enc.barrier(); - am.gatherTo(enc, p, act, sums, matrix(mo.down.one), offsets, R{ .buf = f.get(.expert_out).buffer }, R{ .buf = f.get(.order).buffer }, slots, experts); + am.rowSums(enc2, p, @intCast(mo.down.one.group), act, sums, slots, inter); + enc2.barrier(); + am.gatherTo(enc2, p, act, sums, matrix(mo.down.one), offsets, R{ .buf = f.get(.expert_out).buffer }, R{ .buf = f.get(.order).buffer }, slots, experts); } else { - try k.indexed(enc, pipes, mo.gate, slots, f.get(.input), f.get(.picks), f.get(.expert_gate)); - try k.indexed(enc, pipes, mo.up, slots, f.get(.input), f.get(.picks), f.get(.expert_up)); - enc.barrier(); - k.activation(enc, f, slots, inter); - enc.barrier(); - try k.indexed(enc, pipes, mo.down, slots, f.get(.expert_act), f.get(.picks), f.get(.expert_out)); + try k.indexed(enc2, pipes, mo.gate, slots, f.get(.input), f.get(.picks), f.get(.expert_gate)); + try k.indexed(enc2, pipes, mo.up, slots, f.get(.input), f.get(.picks), f.get(.expert_up)); + enc2.barrier(); + k.activation(enc2, f, slots, inter); + enc2.barrier(); + try k.indexed(enc2, pipes, mo.down, slots, f.get(.expert_act), f.get(.picks), f.get(.expert_out)); } - enc.barrier(); - enc.setPipeline(k.combine); - enc.setBuffer(f.get(.expert_out).buffer, 0, 0); - enc.setBuffer(f.get(.route_weights).buffer, 0, 1); - enc.setBuffer(f.get(.router).buffer, 0, 2); - enc.setBuffer(f.get(.output).buffer, 0, 3); - enc.setValue(CombineArgs{ .hidden = @intCast(hidden), .topk = @intCast(topk), .ld = @intCast(ld), .experts = @intCast(experts) }, 4); - enc.dispatchThreads(mtl.Size.of(hidden, rows, 1), mtl.Size.of(256, 1, 1)); - try span.end(); + enc2.barrier(); + enc2.setPipeline(k.combine); + enc2.setBuffer(f.get(.expert_out).buffer, 0, 0); + enc2.setBuffer(f.get(.route_weights).buffer, 0, 1); + enc2.setBuffer(f.get(.router).buffer, 0, 2); + enc2.setBuffer(f.get(.output).buffer, 0, 3); + enc2.setValue(CombineArgs{ .hidden = @intCast(hidden), .topk = @intCast(topk), .ld = @intCast(ld), .experts = @intCast(experts) }, 4); + enc2.dispatchThreads(mtl.Size.of(hidden, rows, 1), mtl.Size.of(256, 1, 1)); + try products.end(); } fn activation(k: *Kernels, e: mtl.ComputeEncoder, f: *const Frame, slots: u32, inter: u32) void { diff --git a/zig/tests/qwen27/run.zig b/zig/tests/qwen27/run.zig index b29413f2d6..8ad7848036 100644 --- a/zig/tests/qwen27/run.zig +++ b/zig/tests/qwen27/run.zig @@ -140,12 +140,15 @@ pub fn main(init: std.process.Init) !void { try s.prefill(ids[0..@min(ids.len, 128)], 128); try runner.reset(0); } + var trace = profile.Trace{ .queue = m.queue }; + if (profile_steps > 0) runner.profile(&trace); + // TF_PROFILE_PROMPT=1 (with --profile-steps) counts the prompt pass's kernels into the profile too + trace.enabled = profile_steps > 0 and std.c.getenv("TF_PROFILE_PROMPT") != null; const start = std.Io.Clock.awake.now(io); try s.prefill(ids, chunk); const prompt_ns = start.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds(); + trace.enabled = false; const digest = try s.fingerprint(); - var trace = profile.Trace{ .queue = m.queue }; - if (profile_steps > 0) runner.profile(&trace); const decode_start = std.Io.Clock.awake.now(io); const generation = try generated(a, s, count, force, expected, if (profile_steps > 0) &trace else null, profile_steps, teacher); if (std.c.getenv("TF_LOGITS_DUMP")) |v| try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = std.mem.span(v), .data = logits_dump.items }); From 64698105faf91d348c4e08c306a3483bf9796f1b Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 20 Oct 2026 13:10:37 -0500 Subject: [PATCH 05/21] qwen27: routed checkpoints prefill in 1,024-row chunks --- zig/src/families/qwen27/model.zig | 10 ++++++---- zig/src/native/qwen27_served.zig | 2 +- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/zig/src/families/qwen27/model.zig b/zig/src/families/qwen27/model.zig index a864547bc4..674fb396ac 100644 --- a/zig/src/families/qwen27/model.zig +++ b/zig/src/families/qwen27/model.zig @@ -75,7 +75,7 @@ pub const Model = struct { errdefer m.glue.deinit(); m.moe = if (c.moe()) try moe_gpu.Kernels.init(a, m.device) else null; errdefer if (m.moe) |*k| k.deinit(); - m.frame = try Frame.init(m.device, c, @max(rows, promptRows())); + m.frame = try Frame.init(m.device, c, @max(rows, promptRows(c))); errdefer m.frame.deinit(); m.ane = null; const columns = try ane_policy.columns(if (platform.c.getenv("TF_ANE_COLUMNS")) |v| std.mem.span(v) else null); @@ -106,9 +106,11 @@ pub const Model = struct { } }; -/// A prompt pass's rows: 512 (TF_PROMPT_ROWS takes another multiple of 16 up to the frame's 2048). -pub fn promptRows() u32 { - const default = 512; +/// A prompt pass's rows: 512, or 1,024 with routed experts (a chunk reads every expert's weights once, so wider +/// chunks prefill faster: 2,230, 2,580 and 2,760 tok/s at 512, 1,024 and 2,048 on Ornith-1.5-122B); TF_PROMPT_ROWS +/// takes another multiple of 16 up to the frame's 2048. +pub fn promptRows(c: cfg.Config) u32 { + const default: u32 = if (c.moe()) 1024 else 512; const v = std.c.getenv("TF_PROMPT_ROWS") orelse return default; const n = std.fmt.parseInt(u32, std.mem.span(v), 10) catch return default; return if (n >= 16 and n <= @import("gpu_frame.zig").max_rows and n % 16 == 0) n else default; diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index 4ea2f2d480..ef73190ee1 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -311,7 +311,7 @@ pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.df errdefer if (h.draft) |d| d.deinit(); h.back = try lb.Metal.init(gpa, &h.runner, h.draft); errdefer h.back.deinit(); - h.back.chunk = h.runner.promptRows(q.model.promptRows()); + h.back.chunk = h.runner.promptRows(q.model.promptRows(h.runner.model.config)); h.back.grow = .{ .ctx = h, .room = Host.roomFor, .added = Host.added }; h.back.vision = o.images; if (h.draft != null) try h.back.measure(io); From 1cf5b96487ea6d21ca7731888cf999487f91937a Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 20 Oct 2026 13:16:17 -0500 Subject: [PATCH 06/21] qwen27: one-line comments for the lean check --- zig/kernels/metal/qwen27/moe.metal | 15 +++++---------- zig/src/families/qwen27/attention_plan.zig | 2 +- zig/src/families/qwen27/config.zig | 2 +- zig/src/families/qwen27/dflash/config.zig | 3 +-- zig/src/families/qwen27/gpu_weights.zig | 2 +- zig/src/families/qwen27/model.zig | 4 +--- zig/src/families/qwen27/moe_gpu.zig | 9 +++------ zig/src/families/qwen27/widths.zig | 3 +-- 8 files changed, 14 insertions(+), 26 deletions(-) diff --git a/zig/kernels/metal/qwen27/moe.metal b/zig/kernels/metal/qwen27/moe.metal index df69486477..f03f4d918a 100644 --- a/zig/kernels/metal/qwen27/moe.metal +++ b/zig/kernels/metal/qwen27/moe.metal @@ -1,16 +1,13 @@ -// The routed block: a row's softmax over its router logits, its top-k, the slots sorted for the tensor units' gathers, -// the experts' SwiGLU and the combine with the shared expert, all with the reference's bf16 boundaries. +// The routed block: route, the sort for the tensor units, the experts' SwiGLU and the combine (bf16 as the reference). #include using namespace metal; -struct RouteArgs { int experts, topk, ld, renorm; }; +struct RouteArgs { int experts, topk, ld, renorm; }; // PICK, WEIGHT [rows, topk]; column experts + 1 takes the shared gate's bf16 sigmoid struct SortArgs { int slots, experts; }; struct GluArgs { int width; }; struct CombineArgs { int hidden, topk, ld, experts; }; -// A threadgroup a row: softmax (fp32, rounded to bf16 as the reference's) over LOGITS[row, 0..experts), the top-k by -// logit with the lower id on ties, in descending order; PICK and WEIGHT [rows, topk]. Column `experts` of the row is -// the shared expert's gate logit: its sigmoid, rounded to bf16, is left in column experts + 1 for the combine. +// A row a threadgroup: fp32 softmax (to bf16) over LOGITS[row, 0..experts), top-k by logit (lower id on ties). [[kernel]] void q27_route(device bfloat* LOGITS [[buffer(0)]], device uint* PICK [[buffer(1)]], device float* WEIGHT [[buffer(2)]], constant RouteArgs& a [[buffer(3)]], uint r [[threadgroup_position_in_grid]], uint t [[thread_position_in_threadgroup]], @@ -85,8 +82,7 @@ struct CombineArgs { int hidden, topk, ld, experts; }; } } -// A pass's routed slots sorted by expert (stable) for core/affine_mm's gathers: OFFS (each expert's first sorted slot, -// then the end), ORDER (sorted position to slot) and WS (the slot's weight); one thread. +// Slots sorted by expert (stable) for core/affine_mm: OFFS (first slots, then the end), ORDER, WS; one thread. [[kernel]] void q27_route_sort(const device uint* PICK [[buffer(0)]], const device float* WEIGHT [[buffer(1)]], device int* OFFS [[buffer(2)]], device int* ORDER [[buffer(3)]], device float* WS [[buffer(4)]], constant SortArgs& a [[buffer(5)]], @@ -125,8 +121,7 @@ struct CombineArgs { int hidden, topk, ld, experts; }; OUT[i] = bfloat(gate / (1.0f + exp(-gate)) * float(U[i])); } -// OUT[r] (the shared expert's output) becomes bf16(bf16(sum_k bf16(Y[r*topk+k] * w_k)) + bf16(sigma_r * OUT[r])), -// the reference's roundings: each weighted expert row, their fp32 sum, the gated shared row, the sum of the two. +// OUT[r] (the shared expert's row) becomes bf16(bf16(sum_k bf16(Y[r*topk+k] * w_k)) + bf16(sigma_r * OUT[r])). [[kernel]] void q27_expert_combine(const device bfloat* Y [[buffer(0)]], const device float* WEIGHT [[buffer(1)]], const device bfloat* LOGITS [[buffer(2)]], device bfloat* OUT [[buffer(3)]], constant CombineArgs& a [[buffer(4)]], uint2 g [[thread_position_in_grid]]) { diff --git a/zig/src/families/qwen27/attention_plan.zig b/zig/src/families/qwen27/attention_plan.zig index 2751a41c56..a89ad5d262 100644 --- a/zig/src/families/qwen27/attention_plan.zig +++ b/zig/src/families/qwen27/attention_plan.zig @@ -29,7 +29,7 @@ pub const Geometry = struct { /// Chains of rows take the prefix kernel to their last key; a test turns this off to compare with the tail kernel. fuse_chain: bool = true, pub fn check(g: Geometry) !void { - // the kernels' 16-row tiles hold one KV group of up to 16 query heads: the 27B's 24 over 4 and Ornith's 32 over 2 + // a 16-row tile holds one KV group of up to 16 query heads: the 27B's 24 over 4, Ornith's 32 over 2 const grouped = (g.heads == 24 and g.kv_heads == 4) or (g.heads == 32 and g.kv_heads == 2); if (!grouped or g.dim != 256 or g.rotary != 64) return error.TargetAttentionShape; } diff --git a/zig/src/families/qwen27/config.zig b/zig/src/families/qwen27/config.zig index 1c275b78c1..91cfdf88f8 100644 --- a/zig/src/families/qwen27/config.zig +++ b/zig/src/families/qwen27/config.zig @@ -1,4 +1,4 @@ -//! Qwen3.5-class text shapes (dense, or routed experts), layer schedule, rotary metadata and ordered checkpoint EOS ids. +//! Qwen3.5-class text shapes (dense or routed), layer schedule, rotary metadata and ordered checkpoint EOS ids. const std = @import("std"); pub const Kind = enum { linear, attention }; diff --git a/zig/src/families/qwen27/dflash/config.zig b/zig/src/families/qwen27/dflash/config.zig index 59621fb860..5bd662f747 100644 --- a/zig/src/families/qwen27/dflash/config.zig +++ b/zig/src/families/qwen27/dflash/config.zig @@ -1,5 +1,4 @@ -//! DFlash2 is a separate GQA contract (five layers tapping the 27B at five blocks, six tapping Ornith at eight); the -//! target's GDN and draft conventions are not interchangeable. +//! DFlash2 is a separate GQA contract (five layers over the 27B's taps, six over Ornith's eight), not the target's. const std = @import("std"); pub const repository = "z-lab/Qwen3.8-27B-DFlash2"; pub const max_taps = 8; diff --git a/zig/src/families/qwen27/gpu_weights.zig b/zig/src/families/qwen27/gpu_weights.zig index 7cb521c1ee..479b3285be 100644 --- a/zig/src/families/qwen27/gpu_weights.zig +++ b/zig/src/families/qwen27/gpu_weights.zig @@ -14,7 +14,7 @@ pub const Tensor = cp.Tensor; pub const Gdn = struct { qkv: Linear, zba: Linear, out: Linear, conv: Tensor, a_log: Tensor, dt: Tensor, norm: Tensor }; pub const Attention = struct { q: Linear, kv: Linear, out: Linear, q_norm: Tensor, k_norm: Tensor }; -/// The routed block: the router's rows (the experts, then the shared expert's gate, then zero rows) and the stacked experts. +/// The routed block: the router's rows (the experts, the shared gate, then zero rows) and the stacked experts. pub const Moe = struct { router: row.Weights, gate: row.Experts, up: row.Experts, down: row.Experts }; /// With experts, `gu` and `down` are the shared expert's. pub const Layer = struct { norm: Tensor, post_norm: Tensor, gu: Linear, down: Linear, moe: ?Moe = null, mixer: union(@import("config.zig").Kind) { linear: Gdn, attention: Attention } }; diff --git a/zig/src/families/qwen27/model.zig b/zig/src/families/qwen27/model.zig index 674fb396ac..e5a0c912f8 100644 --- a/zig/src/families/qwen27/model.zig +++ b/zig/src/families/qwen27/model.zig @@ -106,9 +106,7 @@ pub const Model = struct { } }; -/// A prompt pass's rows: 512, or 1,024 with routed experts (a chunk reads every expert's weights once, so wider -/// chunks prefill faster: 2,230, 2,580 and 2,760 tok/s at 512, 1,024 and 2,048 on Ornith-1.5-122B); TF_PROMPT_ROWS -/// takes another multiple of 16 up to the frame's 2048. +/// A prompt pass's rows: 512, or 1,024 with experts (a chunk reads every expert once); TF_PROMPT_ROWS: 16 to 2048. pub fn promptRows(c: cfg.Config) u32 { const default: u32 = if (c.moe()) 1024 else 512; const v = std.c.getenv("TF_PROMPT_ROWS") orelse return default; diff --git a/zig/src/families/qwen27/moe_gpu.zig b/zig/src/families/qwen27/moe_gpu.zig index 4ad5c72e53..999c952acc 100644 --- a/zig/src/families/qwen27/moe_gpu.zig +++ b/zig/src/families/qwen27/moe_gpu.zig @@ -1,5 +1,4 @@ -//! The routed block on the GPU: the router's row kernel, the route, the experts by id (row kernels) or sorted by -//! expert on the tensor units (core.affine_mm), their SwiGLU, and the combine over the shared expert's output. +//! The routed block on the GPU: router, route, experts by id (row kernels) or sorted (core.affine_mm), SwiGLU, combine. const std = @import("std"); const mtl = @import("metal"); const core = @import("core"); @@ -83,8 +82,7 @@ pub const Kernels = struct { return g.bits == mo.up.one.bits and g.group == mo.up.one.group and g.bits == mo.down.one.bits and g.group == mo.down.one.group; } - /// f.output, the shared expert's rows of f.input's normed rows, becomes the block's: the routed experts' weighted - /// rows plus the gated shared one. + /// f.output (the shared expert's rows of f.input) becomes the block's: weighted experts plus the gated shared row. pub fn block(k: *Kernels, e: mtl.ComputeEncoder, c: Config, mo: *const wts.Moe, f: *const Frame, rows: u32, prompt: bool) !void { var span = profile.begin(k.profiler, e, .route); errdefer span.cancel(); @@ -95,8 +93,7 @@ pub const Kernels = struct { const hidden: u32 = @intCast(c.hidden); const inter: u32 = @intCast(c.moe_intermediate); const units = k.sorted(mo, rows, prompt); - // the router: bf16 logits [rows, ld], the shared gate's in column `experts`; a prompt chunk's rows on the - // tensor units (y's pitch is n rounded up to 64), decode rows through the row kernel (pitch n) + // the router's bf16 logits [rows, ld] (column `experts`: the shared gate), on tensor units or the row kernel const ld: u32 = if (units) std.mem.alignForward(u32, @intCast(mo.router.n), 64) else @intCast(mo.router.n); const pipes = try k.rowPipelines(); if (units) { diff --git a/zig/src/families/qwen27/widths.zig b/zig/src/families/qwen27/widths.zig index d9c84124d1..526e923d17 100644 --- a/zig/src/families/qwen27/widths.zig +++ b/zig/src/families/qwen27/widths.zig @@ -199,8 +199,7 @@ fn check(chip: Chip, n: usize, k: usize, w: Width) !void { return (row.Weights{ .w = undefined, .scales = undefined, .biases = undefined, .n = n, .k = k, .group = w.group }).validate(); } -/// The routed block's kernels: the router and the shared gate share one row matrix; the experts take the row kernels -/// by id and, on tensor units, core.affine_mm's gathers (whose format admits 2-6 and 8 bits at groups 32, 64 and 128). +/// The routed block's kernels: the router's row matrix; the experts' row kernels and (tensor units) affine_mm gathers. fn checkRouted(c: cfg.Config, chip: Chip, m: Member, w: Width) !void { const s = shape(c, m); if (std.mem.indexOfScalar(u8, chip.row_bits, w.bits) == null) return error.UnsupportedRowBits; From 5eeb4121079c9c8962e8fd0983284fcb18151562 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 20 Oct 2026 13:17:09 -0500 Subject: [PATCH 07/21] qwen27: a config test for the routed shape --- zig/src/families/qwen27/config.zig | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/zig/src/families/qwen27/config.zig b/zig/src/families/qwen27/config.zig index 91cfdf88f8..213a1db312 100644 --- a/zig/src/families/qwen27/config.zig +++ b/zig/src/families/qwen27/config.zig @@ -261,6 +261,30 @@ test "rope parameters preserve precision and section metadata" { try std.testing.expectEqualSlices(usize, &.{ 2, 1, 1 }, &c.mrope_section); } +test "a routed config reads its experts, the shared expert's width as the layer's and the routing flag" { + const a = std.testing.allocator; + const json = + \\{"model_type":"qwen3_5_moe","text_config":{"model_type":"qwen3_5_moe_text","hidden_size":3072,"intermediate_size":1024, + \\"moe_intermediate_size":1024,"shared_expert_intermediate_size":1024,"num_experts":256,"num_experts_per_tok":8, + \\"num_hidden_layers":48,"num_attention_heads":32,"num_key_value_heads":2,"head_dim":256,"vocab_size":248320, + \\"linear_num_key_heads":16,"linear_num_value_heads":64,"linear_key_head_dim":128,"linear_value_head_dim":128, + \\"linear_conv_kernel_dim":4,"partial_rotary_factor":0.25,"eos_token_id":248044}} + ; + const c = try parse(a, json, null); + try std.testing.expect(c.moe()); + try std.testing.expectEqual(@as(usize, 256), c.experts); + try std.testing.expectEqual(@as(usize, 8), c.top_k); + try std.testing.expectEqual(@as(usize, 1024), c.moe_intermediate); + try std.testing.expectEqual(@as(usize, 1024), c.intermediate); + try std.testing.expect(!c.renorm); + try std.testing.expectEqual(Kind.attention, c.kind(3)); + try std.testing.expectEqual(Kind.linear, c.kind(2)); + // a routed model type without experts, and a dense one with them, are refused + const dense_typed = try std.mem.replaceOwned(u8, a, json, "\"num_experts\":256", "\"num_experts\":0"); + defer a.free(dense_typed); + try std.testing.expectError(error.NotQwenDense, parse(a, dense_typed, null)); +} + test "dense config refuses MoE and tied embedding heads" { const a = std.testing.allocator; const moe = try edit(a, "num_experts", .{ .integer = 8 }); From 4acc699aa38f66d31b1af778871046a6389897bf Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 20 Oct 2026 13:29:24 -0500 Subject: [PATCH 08/21] qwen27: routed models let plain rounds compete with drafted depths --- zig/src/native/qwen27_lanes.zig | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/zig/src/native/qwen27_lanes.zig b/zig/src/native/qwen27_lanes.zig index 819d88d62c..c156b3e685 100644 --- a/zig/src/native/qwen27_lanes.zig +++ b/zig/src/native/qwen27_lanes.zig @@ -128,7 +128,9 @@ pub const Metal = struct { const shared = b.runner.slots > 1; // a lone stream drafts DFlash2 trees, as in 1.0.5: copies land less often there (TF_QWEN27_COPIES=1 opts in) const lone_copies = if (platform.c.getenv("TF_QWEN27_COPIES")) |v| std.mem.eql(u8, std.mem.span(v), "1") else false; - return .{ .exact_width = if (drafting or shared) max_rows else 1, .mtp = drafting, .speculate = drafting, .speculate_early = false, .drafts = b.nodes, .draft_probabilities = true, .window_costs = b.costs[0..b.timed], .hidden_rows = shared, .streams_exact = true, .batch_rows = if (shared) batch_rows else max_rows, .max_streams = b.runner.slots, .lone_copies = lone_copies, .draft_streams = drafting and b.draft_slots > 1 }; + // routed experts make a drafted window read most expert weights: plain rounds compete with drafted depths + const routed = b.runner.model.config.moe(); + return .{ .exact_width = if (drafting or shared) max_rows else 1, .mtp = drafting, .speculate = drafting, .speculate_early = false, .drafts = b.nodes, .draft_probabilities = true, .plain_guard = routed, .window_costs = b.costs[0..b.timed], .hidden_rows = shared, .streams_exact = true, .batch_rows = if (shared) batch_rows else max_rows, .max_streams = b.runner.slots, .lone_copies = lone_copies, .draft_streams = drafting and b.draft_slots > 1 }; } /// The slot holding `s`. From 32dad79b171d270167db62bdbc1d6ce62a9fe6bc Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 20 Oct 2026 13:34:59 -0500 Subject: [PATCH 09/21] qwen27: the router matrix is padded to whole 64-row tiles for the tensor units' product --- zig/src/families/qwen27/gpu_frame.zig | 4 ++-- zig/src/families/qwen27/moe_gpu.zig | 2 +- zig/src/families/qwen27/widths.zig | 5 +++-- zig/tests/qwen27/run.zig | 3 ++- 4 files changed, 8 insertions(+), 6 deletions(-) diff --git a/zig/src/families/qwen27/gpu_frame.zig b/zig/src/families/qwen27/gpu_frame.zig index 0fb86d8684..5b25e82833 100644 --- a/zig/src/families/qwen27/gpu_frame.zig +++ b/zig/src/families/qwen27/gpu_frame.zig @@ -9,9 +9,9 @@ pub const Field = enum { ids, dims, positions, hidden, residual, input, sums, mi /// A prompt pass's rows at most (a chain); rounds with trees and every logits row stay within `logit_rows`. pub const max_rows = 2048; -/// The router's output pitch: its rows rounded up to the tensor units' 64 (core.affine_mm.dense). +/// The router's output pitch: its rows (a multiple of 64, the tensor units' dense product's pitch). fn routerRows(config: Config) usize { - return if (config.moe()) std.mem.alignForward(usize, @import("widths.zig").routerRows(config.experts), 64) else 0; + return if (config.moe()) @import("widths.zig").routerRows(config.experts) else 0; } pub const logit_rows = 128; diff --git a/zig/src/families/qwen27/moe_gpu.zig b/zig/src/families/qwen27/moe_gpu.zig index 999c952acc..20a16b4f81 100644 --- a/zig/src/families/qwen27/moe_gpu.zig +++ b/zig/src/families/qwen27/moe_gpu.zig @@ -94,7 +94,7 @@ pub const Kernels = struct { const inter: u32 = @intCast(c.moe_intermediate); const units = k.sorted(mo, rows, prompt); // the router's bf16 logits [rows, ld] (column `experts`: the shared gate), on tensor units or the row kernel - const ld: u32 = if (units) std.mem.alignForward(u32, @intCast(mo.router.n), 64) else @intCast(mo.router.n); + const ld: u32 = @intCast(mo.router.n); // a multiple of 64: both paths write that pitch const pipes = try k.rowPipelines(); if (units) { const p = try k.naxPipes(.{ .bits = mo.router.bits, .group = @intCast(mo.router.group) }); diff --git a/zig/src/families/qwen27/widths.zig b/zig/src/families/qwen27/widths.zig index 526e923d17..3d8ae74563 100644 --- a/zig/src/families/qwen27/widths.zig +++ b/zig/src/families/qwen27/widths.zig @@ -211,9 +211,10 @@ fn checkRouted(c: cfg.Config, chip: Chip, m: Member, w: Width) !void { } } -/// The router matrix's rows: the experts, the shared expert's gate, then zero rows to a multiple of 8. +/// The router matrix's rows: the experts, the shared expert's gate, then zero rows to a multiple of 64 (the tensor +/// units' dense product reads whole 64-row tiles of the weights). pub fn routerRows(experts: usize) usize { - return (experts + 1 + 7) / 8 * 8; + return (experts + 1 + 63) / 64 * 64; } const all_rows = [_]u8{ 2, 3, 4, 5, 6, 8 }; diff --git a/zig/tests/qwen27/run.zig b/zig/tests/qwen27/run.zig index 8ad7848036..8f0265b032 100644 --- a/zig/tests/qwen27/run.zig +++ b/zig/tests/qwen27/run.zig @@ -18,7 +18,8 @@ const help = \\--compare-chunk reruns the whole prompt and decode and checks prompt state/KV/logits bytes. \\--seams N (with --compare-chunk): state/KV/logits bytes every N prompt rows under both chunks and after decode. \\CPU only: --help or --emit-sources EXISTING_DIR, no Metal device or model load. - \\TF_LAYERS=N runs the first N layers only; TF_LOGITS_DUMP=FILE writes every step's bf16 logits row (a port's checks). + \\TF_LAYERS=N runs the first N layers only; TF_LOGITS_DUMP=FILE writes every step's bf16 logits row (a port's checks); + \\TF_PROFILE_PROMPT=1 with --profile-steps counts the prompt pass's kernels into the profile. ; fn tokens(a: std.mem.Allocator, io: std.Io, file: []const u8) ![]const u32 { From e077860c0e94bf328c0ad4e5990cc1cc0281420b Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 20 Oct 2026 13:59:23 -0500 Subject: [PATCH 10/21] qwen27: tap commits gather every tap plane (the drafter's context held five of Ornith's eight); selector options by env --- zig/src/families/qwen27/dflash/generation.zig | 2 +- zig/src/families/qwen27/dflash/runtime_model.zig | 15 ++++++++++++--- zig/src/families/qwen27/moe_gpu.zig | 1 + zig/src/native/qwen27_lanes.zig | 2 +- 4 files changed, 15 insertions(+), 5 deletions(-) diff --git a/zig/src/families/qwen27/dflash/generation.zig b/zig/src/families/qwen27/dflash/generation.zig index e2fad6160a..ceedc84d91 100644 --- a/zig/src/families/qwen27/dflash/generation.zig +++ b/zig/src/families/qwen27/dflash/generation.zig @@ -137,7 +137,7 @@ pub const Generation = struct { const width: u32 = @intCast(g.runner.model.config.vocab); const shape = try g.matcher.encode(e, .{ .buf = logits.buffer, .off = logits.offset }, ids, parents, width, budget, eos); const result = @import("core").tree_commit_gpu.Ref{ .buf = g.matcher.buffers[4] }; - try g.commit.taps(e, result, .{ .buf = g.runner.taps.buffer }, .{ .buf = g.taps }, .{ .rows = shape.rows, .width = g.runner.taps.width, .capacity = g.runner.taps.capacity, .planes = 5 }); + try g.commit.taps(e, result, .{ .buf = g.runner.taps.buffer }, .{ .buf = g.taps }, .{ .rows = shape.rows, .width = g.runner.taps.width, .capacity = g.runner.taps.capacity, .planes = @intCast(g.runner.taps.count) }); try g.commit.head(e, result, .{ .buf = logits.buffer, .off = logits.offset }, .{ .rows = shape.rows, .width = width, .stride = width }); e.end(); ended = true; diff --git a/zig/src/families/qwen27/dflash/runtime_model.zig b/zig/src/families/qwen27/dflash/runtime_model.zig index e22377fafb..996eda41e9 100644 --- a/zig/src/families/qwen27/dflash/runtime_model.zig +++ b/zig/src/families/qwen27/dflash/runtime_model.zig @@ -110,6 +110,15 @@ pub const Model = struct { try m.session.absorb(m.backend.ops(), &m.graph, m.preparation, .{ .values = view, .positions = positions[0..rows], .layers = m.graph.config.taps, .target_end = first + rows }); } } + /// The tree search's options for `nodes` drafts: TF_DRAFT_CHILDREN, TF_DRAFT_TAU and TF_DRAFT_EDGE override the defaults. + pub fn selectorOptions(m: *const Model, nodes: u32) selector.Options { + _ = m; + var o = selector.Options{ .nodes = nodes }; + if (std.c.getenv("TF_DRAFT_CHILDREN")) |v| o.children = std.fmt.parseInt(u32, std.mem.span(v), 10) catch o.children; + if (std.c.getenv("TF_DRAFT_TAU")) |v| o.tau = std.fmt.parseFloat(f64, std.mem.span(v)) catch o.tau; + if (std.c.getenv("TF_DRAFT_EDGE")) |v| o.edge = std.fmt.parseFloat(f64, std.mem.span(v)) catch o.edge; + return o; + } pub fn setTreeBlock(m: *Model, block: u32) !void { if (block == 0 or block > 16 or m.backend.failed or m.backend.frame.active or m.backend.transaction != null) return error.BadRuntimeBlock; m.tree_block = block; @@ -181,7 +190,7 @@ pub const Model = struct { var made: usize = 0; errdefer for (out[0..made]) |*t| t.deinit(); for (entries, streams[0..n], lattices[0..n], out) |en, s2, l, *t| { - t.* = try selector.propose(m.allocator, &m.graph, l, en.pending, s2.context.end + 1, .{ .nodes = en.nodes }); + t.* = try selector.propose(m.allocator, &m.graph, l, en.pending, s2.context.end + 1, m.selectorOptions(en.nodes)); made += 1; } } @@ -210,11 +219,11 @@ pub const Model = struct { for (positions[0..rows], 0..) |*position, i| position.* = m.session.context.end + i; const stride = @as(usize, m.graph.config.tapWidth()) * 2; const view = try m.backend.frame.borrow(.{ .buf = m.pending, .off = @as(usize, slot) * stride }, rows, m.graph.config.tapWidth()); - return m.session.absorbTreeBlock(m.allocator, m.backend.ops(), &m.graph, m.preparation, .{ .values = view, .positions = positions[0..rows], .layers = m.graph.config.taps, .target_end = m.session.context.end + rows }, pending, block, .{ .nodes = nodes }); + return m.session.absorbTreeBlock(m.allocator, m.backend.ops(), &m.graph, m.preparation, .{ .values = view, .positions = positions[0..rows], .layers = m.graph.config.taps, .target_end = m.session.context.end + rows }, pending, block, m.selectorOptions(nodes)); } try m.flush(); const first = try m.session.firstDraw(); - return m.session.treeBlock(m.allocator, m.backend.ops(), &m.graph, m.preparation, pending, first, block, .{ .nodes = nodes }); + return m.session.treeBlock(m.allocator, m.backend.ops(), &m.graph, m.preparation, pending, first, block, m.selectorOptions(nodes)); } }; /// A stream's own drafter state, parked while another stream drafts: its taps ring, rolling K/V and session. diff --git a/zig/src/families/qwen27/moe_gpu.zig b/zig/src/families/qwen27/moe_gpu.zig index 20a16b4f81..153b5f9afe 100644 --- a/zig/src/families/qwen27/moe_gpu.zig +++ b/zig/src/families/qwen27/moe_gpu.zig @@ -76,6 +76,7 @@ pub const Kernels = struct { } /// Whether a pass of `rows` takes the sorted tensor-unit path: a prompt chunk of whole tiles, experts of one width. + /// Decode windows keep the row kernels, whose sums match a lone row's bit for bit (drafted replies stay exact). pub fn sorted(k: *Kernels, mo: *const wts.Moe, rows: u32, prompt: bool) bool { if (!prompt or rows < 16 or !k.device.tensorUnits()) return false; const g = mo.gate.one; diff --git a/zig/src/native/qwen27_lanes.zig b/zig/src/native/qwen27_lanes.zig index c156b3e685..1202a2eced 100644 --- a/zig/src/native/qwen27_lanes.zig +++ b/zig/src/native/qwen27_lanes.zig @@ -386,7 +386,7 @@ pub const Metal = struct { var ended = false; errdefer if (!ended) e.end(); try r.keepDevice(&b.round.?, e, g.commit, ref); - try g.commit.taps(e, ref, .{ .buf = r.taps.buffer }, .{ .buf = g.taps }, .{ .rows = rows, .width = r.taps.width, .capacity = r.taps.capacity, .planes = 5 }); + try g.commit.taps(e, ref, .{ .buf = r.taps.buffer }, .{ .buf = g.taps }, .{ .rows = rows, .width = r.taps.width, .capacity = r.taps.capacity, .planes = @intCast(r.taps.count) }); try g.commit.head(e, ref, .{ .buf = logits.buffer, .off = logits.offset }, .{ .rows = rows, .width = vocab, .stride = vocab }); e.end(); ended = true; From 2f995f3570af65c8c92bd21365c0abc656b25542 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 20 Oct 2026 14:01:06 -0500 Subject: [PATCH 11/21] qwen27: one-line comments for the lean check --- zig/src/families/qwen27/dflash/runtime_model.zig | 2 +- zig/src/families/qwen27/moe_gpu.zig | 3 +-- zig/src/families/qwen27/widths.zig | 3 +-- 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/zig/src/families/qwen27/dflash/runtime_model.zig b/zig/src/families/qwen27/dflash/runtime_model.zig index 996eda41e9..09df7d1135 100644 --- a/zig/src/families/qwen27/dflash/runtime_model.zig +++ b/zig/src/families/qwen27/dflash/runtime_model.zig @@ -110,7 +110,7 @@ pub const Model = struct { try m.session.absorb(m.backend.ops(), &m.graph, m.preparation, .{ .values = view, .positions = positions[0..rows], .layers = m.graph.config.taps, .target_end = first + rows }); } } - /// The tree search's options for `nodes` drafts: TF_DRAFT_CHILDREN, TF_DRAFT_TAU and TF_DRAFT_EDGE override the defaults. + /// The tree search's options for `nodes` drafts; TF_DRAFT_CHILDREN, TF_DRAFT_TAU, TF_DRAFT_EDGE override them. pub fn selectorOptions(m: *const Model, nodes: u32) selector.Options { _ = m; var o = selector.Options{ .nodes = nodes }; diff --git a/zig/src/families/qwen27/moe_gpu.zig b/zig/src/families/qwen27/moe_gpu.zig index 153b5f9afe..03633e2cb6 100644 --- a/zig/src/families/qwen27/moe_gpu.zig +++ b/zig/src/families/qwen27/moe_gpu.zig @@ -75,8 +75,7 @@ pub const Kernels = struct { return pipes; } - /// Whether a pass of `rows` takes the sorted tensor-unit path: a prompt chunk of whole tiles, experts of one width. - /// Decode windows keep the row kernels, whose sums match a lone row's bit for bit (drafted replies stay exact). + /// The sorted tensor-unit path: prompt chunks of whole tiles only; decode windows keep the row kernels (exact). pub fn sorted(k: *Kernels, mo: *const wts.Moe, rows: u32, prompt: bool) bool { if (!prompt or rows < 16 or !k.device.tensorUnits()) return false; const g = mo.gate.one; diff --git a/zig/src/families/qwen27/widths.zig b/zig/src/families/qwen27/widths.zig index 3d8ae74563..938fa6dc5a 100644 --- a/zig/src/families/qwen27/widths.zig +++ b/zig/src/families/qwen27/widths.zig @@ -211,8 +211,7 @@ fn checkRouted(c: cfg.Config, chip: Chip, m: Member, w: Width) !void { } } -/// The router matrix's rows: the experts, the shared expert's gate, then zero rows to a multiple of 64 (the tensor -/// units' dense product reads whole 64-row tiles of the weights). +/// The router matrix's rows: the experts, the shared gate, then zero rows to a multiple of 64 (whole weight tiles). pub fn routerRows(experts: usize) usize { return (experts + 1 + 63) / 64 * 64; } From 0d37dcb24cce23da957035a16c749dd441cd672f Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 20 Oct 2026 14:13:32 -0500 Subject: [PATCH 12/21] qwen27: decode rows run the shared expert as one more slot of the experts' kernel (core: indexed_extra row kernel) --- zig/kernels/metal/core/row_projection.metal | 49 +++++++++++++++++++++ zig/kernels/metal/qwen27/moe.metal | 38 +++++++++------- zig/src/core/row_projection.zig | 13 ++++++ zig/src/families/qwen27/forward.zig | 7 +++ zig/src/families/qwen27/gpu_frame.zig | 6 +-- zig/src/families/qwen27/gpu_weights.zig | 24 +++++++--- zig/src/families/qwen27/moe_gpu.zig | 47 +++++++++++--------- 7 files changed, 136 insertions(+), 48 deletions(-) diff --git a/zig/kernels/metal/core/row_projection.metal b/zig/kernels/metal/core/row_projection.metal index 0be006be24..457cebf246 100644 --- a/zig/kernels/metal/core/row_projection.metal +++ b/zig/kernels/metal/core/row_projection.metal @@ -190,53 +190,102 @@ template if (lane == 0) for (int j = 0; j < RPS; j++) out[j] = tf_rp_out(acc[0][j]); } +// `tf_row_projection_indexed` with an extra matrix at id `experts` (buffers 7-9): a shared expert rides with the routed ones. +template +[[kernel]] void tf_row_projection_indexed_extra(const device bfloat* X [[buffer(0)]], + const device uint32_t* W [[buffer(1)]], + const device bfloat* S [[buffer(2)]], + const device bfloat* B [[buffer(3)]], + constant int* dims [[buffer(4)]], + device bfloat* OUT [[buffer(5)]], + const device uint* ids [[buffer(6)]], + const device uint32_t* W2 [[buffer(7)]], + const device bfloat* S2 [[buffer(8)]], + const device bfloat* B2 [[buffer(9)]], + uint simdgroup_index_in_threadgroup [[simdgroup_index_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]], + uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]]) { + const int K = dims[0], N = dims[1], GS = dims[2], experts = dims[4], repeat = dims[5]; + const uint lane = thread_index_in_simdgroup; + const int row0 = (int(threadgroup_position_in_grid.x) * SG + int(simdgroup_index_in_threadgroup)) * RPS; + if (row0 >= N) return; + const uint slot = threadgroup_position_in_grid.y; + const uint expert = ids[slot]; + device bfloat* out = OUT + size_t(slot) * N + row0; + if (expert > uint(experts)) { + if (lane < uint(RPS)) out[lane] = bfloat(0.0f); + return; + } + // id `experts` is the extra matrix (W2, S2, B2: one more [N, K] of the same width), the shared expert's + const bool extra = expert == uint(experts); + const size_t at = (extra ? 0 : size_t(expert) * N) + size_t(row0); + const device uint8_t* wr = (const device uint8_t*)(extra ? W2 : W) + at * (K * BITS / 8); + const device bfloat* sr = (extra ? S2 : S) + at * (K / GS); + const device bfloat* br = (extra ? B2 : B) + at * (K / GS); + float acc[1][RPS]; + tf_rp_rows(wr, sr, br, X + size_t(slot / repeat) * K, K, GS, lane, acc); + if (lane == 0) + for (int j = 0; j < RPS; j++) out[j] = tf_rp_out(acc[0][j]); +} template [[host_name("tf_row_projection_q2_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 2, false>) tf_row_projection<4, 2, 2, false, 2, false>; template [[host_name("tf_row_projection_indexed_q2_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 2, false>) tf_row_projection_indexed<4, 2, false, 2, false>; +template [[host_name("tf_row_projection_indexed_extra_q2_f32")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 2, false>) tf_row_projection_indexed_extra<4, 2, false, 2, false>; template [[host_name("tf_row_projection_relu2_q2_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 2, false>) tf_row_projection<4, 2, 2, true, 2, false>; template [[host_name("tf_row_projection_indexed_relu2_q2_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 2, false>) tf_row_projection_indexed<4, 2, true, 2, false>; template [[host_name("tf_row_projection_q2_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 2, true>) tf_row_projection<4, 2, 2, false, 2, true>; template [[host_name("tf_row_projection_indexed_q2_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 2, true>) tf_row_projection_indexed<4, 2, false, 2, true>; +template [[host_name("tf_row_projection_indexed_extra_q2_bf16")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 2, true>) tf_row_projection_indexed_extra<4, 2, false, 2, true>; template [[host_name("tf_row_projection_relu2_q2_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 2, true>) tf_row_projection<4, 2, 2, true, 2, true>; template [[host_name("tf_row_projection_indexed_relu2_q2_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 2, true>) tf_row_projection_indexed<4, 2, true, 2, true>; template [[host_name("tf_row_projection_q4_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 4, false>) tf_row_projection<4, 2, 2, false, 4, false>; template [[host_name("tf_row_projection_indexed_q4_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 4, false>) tf_row_projection_indexed<4, 2, false, 4, false>; +template [[host_name("tf_row_projection_indexed_extra_q4_f32")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 4, false>) tf_row_projection_indexed_extra<4, 2, false, 4, false>; template [[host_name("tf_row_projection_relu2_q4_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 4, false>) tf_row_projection<4, 2, 2, true, 4, false>; template [[host_name("tf_row_projection_indexed_relu2_q4_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 4, false>) tf_row_projection_indexed<4, 2, true, 4, false>; template [[host_name("tf_row_projection_q4_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 4, true>) tf_row_projection<4, 2, 2, false, 4, true>; template [[host_name("tf_row_projection_indexed_q4_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 4, true>) tf_row_projection_indexed<4, 2, false, 4, true>; +template [[host_name("tf_row_projection_indexed_extra_q4_bf16")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 4, true>) tf_row_projection_indexed_extra<4, 2, false, 4, true>; template [[host_name("tf_row_projection_relu2_q4_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 4, true>) tf_row_projection<4, 2, 2, true, 4, true>; template [[host_name("tf_row_projection_indexed_relu2_q4_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 4, true>) tf_row_projection_indexed<4, 2, true, 4, true>; template [[host_name("tf_row_projection_q6_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 6, false>) tf_row_projection<4, 2, 2, false, 6, false>; template [[host_name("tf_row_projection_indexed_q6_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 6, false>) tf_row_projection_indexed<4, 2, false, 6, false>; +template [[host_name("tf_row_projection_indexed_extra_q6_f32")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 6, false>) tf_row_projection_indexed_extra<4, 2, false, 6, false>; template [[host_name("tf_row_projection_relu2_q6_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 6, false>) tf_row_projection<4, 2, 2, true, 6, false>; template [[host_name("tf_row_projection_indexed_relu2_q6_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 6, false>) tf_row_projection_indexed<4, 2, true, 6, false>; template [[host_name("tf_row_projection_q6_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 6, true>) tf_row_projection<4, 2, 2, false, 6, true>; template [[host_name("tf_row_projection_indexed_q6_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 6, true>) tf_row_projection_indexed<4, 2, false, 6, true>; +template [[host_name("tf_row_projection_indexed_extra_q6_bf16")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 6, true>) tf_row_projection_indexed_extra<4, 2, false, 6, true>; template [[host_name("tf_row_projection_relu2_q6_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 6, true>) tf_row_projection<4, 2, 2, true, 6, true>; template [[host_name("tf_row_projection_indexed_relu2_q6_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 6, true>) tf_row_projection_indexed<4, 2, true, 6, true>; template [[host_name("tf_row_projection_q8_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 8, false>) tf_row_projection<4, 2, 2, false, 8, false>; template [[host_name("tf_row_projection_indexed_q8_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 8, false>) tf_row_projection_indexed<4, 2, false, 8, false>; +template [[host_name("tf_row_projection_indexed_extra_q8_f32")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 8, false>) tf_row_projection_indexed_extra<4, 2, false, 8, false>; template [[host_name("tf_row_projection_relu2_q8_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 8, false>) tf_row_projection<4, 2, 2, true, 8, false>; template [[host_name("tf_row_projection_indexed_relu2_q8_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 8, false>) tf_row_projection_indexed<4, 2, true, 8, false>; template [[host_name("tf_row_projection_q8_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 8, true>) tf_row_projection<4, 2, 2, false, 8, true>; template [[host_name("tf_row_projection_indexed_q8_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 8, true>) tf_row_projection_indexed<4, 2, false, 8, true>; +template [[host_name("tf_row_projection_indexed_extra_q8_bf16")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 8, true>) tf_row_projection_indexed_extra<4, 2, false, 8, true>; template [[host_name("tf_row_projection_relu2_q8_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 8, true>) tf_row_projection<4, 2, 2, true, 8, true>; template [[host_name("tf_row_projection_indexed_relu2_q8_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 8, true>) tf_row_projection_indexed<4, 2, true, 8, true>; template [[host_name("tf_row_projection_q3_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 3, false>) tf_row_projection<4, 2, 2, false, 3, false>; template [[host_name("tf_row_projection_indexed_q3_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 3, false>) tf_row_projection_indexed<4, 2, false, 3, false>; +template [[host_name("tf_row_projection_indexed_extra_q3_f32")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 3, false>) tf_row_projection_indexed_extra<4, 2, false, 3, false>; template [[host_name("tf_row_projection_relu2_q3_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 3, false>) tf_row_projection<4, 2, 2, true, 3, false>; template [[host_name("tf_row_projection_indexed_relu2_q3_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 3, false>) tf_row_projection_indexed<4, 2, true, 3, false>; template [[host_name("tf_row_projection_q3_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 3, true>) tf_row_projection<4, 2, 2, false, 3, true>; template [[host_name("tf_row_projection_indexed_q3_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 3, true>) tf_row_projection_indexed<4, 2, false, 3, true>; +template [[host_name("tf_row_projection_indexed_extra_q3_bf16")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 3, true>) tf_row_projection_indexed_extra<4, 2, false, 3, true>; template [[host_name("tf_row_projection_relu2_q3_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 3, true>) tf_row_projection<4, 2, 2, true, 3, true>; template [[host_name("tf_row_projection_indexed_relu2_q3_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 3, true>) tf_row_projection_indexed<4, 2, true, 3, true>; template [[host_name("tf_row_projection_q5_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 5, false>) tf_row_projection<4, 2, 2, false, 5, false>; template [[host_name("tf_row_projection_indexed_q5_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 5, false>) tf_row_projection_indexed<4, 2, false, 5, false>; +template [[host_name("tf_row_projection_indexed_extra_q5_f32")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 5, false>) tf_row_projection_indexed_extra<4, 2, false, 5, false>; template [[host_name("tf_row_projection_relu2_q5_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 5, false>) tf_row_projection<4, 2, 2, true, 5, false>; template [[host_name("tf_row_projection_indexed_relu2_q5_f32")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 5, false>) tf_row_projection_indexed<4, 2, true, 5, false>; template [[host_name("tf_row_projection_q5_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 5, true>) tf_row_projection<4, 2, 2, false, 5, true>; template [[host_name("tf_row_projection_indexed_q5_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, false, 5, true>) tf_row_projection_indexed<4, 2, false, 5, true>; +template [[host_name("tf_row_projection_indexed_extra_q5_bf16")]] [[kernel]] decltype(tf_row_projection_indexed_extra<4, 2, false, 5, true>) tf_row_projection_indexed_extra<4, 2, false, 5, true>; template [[host_name("tf_row_projection_relu2_q5_bf16")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, true, 5, true>) tf_row_projection<4, 2, 2, true, 5, true>; template [[host_name("tf_row_projection_indexed_relu2_q5_bf16")]] [[kernel]] decltype(tf_row_projection_indexed<4, 2, true, 5, true>) tf_row_projection_indexed<4, 2, true, 5, true>; template [[host_name("tf_row_projection_strided_q2_f32")]] [[kernel]] decltype(tf_row_projection<4, 2, 2, false, 2, false, true>) tf_row_projection<4, 2, 2, false, 2, false, true>; diff --git a/zig/kernels/metal/qwen27/moe.metal b/zig/kernels/metal/qwen27/moe.metal index f03f4d918a..68dbd99504 100644 --- a/zig/kernels/metal/qwen27/moe.metal +++ b/zig/kernels/metal/qwen27/moe.metal @@ -2,13 +2,13 @@ #include using namespace metal; -struct RouteArgs { int experts, topk, ld, renorm; }; // PICK, WEIGHT [rows, topk]; column experts + 1 takes the shared gate's bf16 sigmoid +struct RouteArgs { int experts, topk, ld, renorm; }; // PICK, WEIGHT [rows, topk + 1]: the picks, then the shared expert (id experts, its sigmoid) struct SortArgs { int slots, experts; }; struct GluArgs { int width; }; -struct CombineArgs { int hidden, topk, ld, experts; }; +struct CombineArgs { int hidden, topk, shared_in_out, pad; }; // A row a threadgroup: fp32 softmax (to bf16) over LOGITS[row, 0..experts), top-k by logit (lower id on ties). -[[kernel]] void q27_route(device bfloat* LOGITS [[buffer(0)]], device uint* PICK [[buffer(1)]], +[[kernel]] void q27_route(const device bfloat* LOGITS [[buffer(0)]], device uint* PICK [[buffer(1)]], device float* WEIGHT [[buffer(2)]], constant RouteArgs& a [[buffer(3)]], uint r [[threadgroup_position_in_grid]], uint t [[thread_position_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]], @@ -18,7 +18,7 @@ struct CombineArgs { int hidden, topk, ld, experts; }; threadgroup int best_i[32]; threadgroup int sel[16]; threadgroup float w[16]; - device bfloat* z = LOGITS + size_t(r) * a.ld; + const device bfloat* z = LOGITS + size_t(r) * a.ld; const bool mine = int(t) < a.experts; const float x = mine ? float(z[t]) : -INFINITY; // the row's max, then the sum of exp @@ -73,23 +73,25 @@ struct CombineArgs { int hidden, topk, ld, experts; }; float sum = 0.0f; for (int k = 0; k < a.topk; k++) sum += w[k]; const float norm = float(bfloat(sum)); + const int stride = a.topk + 1; for (int k = 0; k < a.topk; k++) { - PICK[size_t(r) * a.topk + k] = uint(sel[k]); - WEIGHT[size_t(r) * a.topk + k] = a.renorm ? float(bfloat(w[k] / norm)) : w[k]; + PICK[size_t(r) * stride + k] = uint(sel[k]); + WEIGHT[size_t(r) * stride + k] = a.renorm ? float(bfloat(w[k] / norm)) : w[k]; } const float g = float(z[a.experts]); - z[a.experts + 1] = bfloat(1.0f / (1.0f + precise::exp(-g))); + PICK[size_t(r) * stride + a.topk] = uint(a.experts); + WEIGHT[size_t(r) * stride + a.topk] = float(bfloat(1.0f / (1.0f + precise::exp(-g)))); } } -// Slots sorted by expert (stable) for core/affine_mm: OFFS (first slots, then the end), ORDER, WS; one thread. +// Routed slots (id < experts) sorted by expert (stable) for core/affine_mm: OFFS (first slots, then the end), ORDER, WS. [[kernel]] void q27_route_sort(const device uint* PICK [[buffer(0)]], const device float* WEIGHT [[buffer(1)]], device int* OFFS [[buffer(2)]], device int* ORDER [[buffer(3)]], device float* WS [[buffer(4)]], constant SortArgs& a [[buffer(5)]], uint t [[thread_position_in_grid]]) { if (t != 0) return; for (int e = 0; e <= a.experts; e++) OFFS[e] = 0; - for (int i = 0; i < a.slots; i++) OFFS[PICK[i]]++; + for (int i = 0; i < a.slots; i++) if (PICK[i] < uint(a.experts)) OFFS[PICK[i]]++; int at = 0; for (int e = 0; e < a.experts; e++) { const int n = OFFS[e]; @@ -97,6 +99,7 @@ struct CombineArgs { int hidden, topk, ld, experts; }; at += n; } for (int i = 0; i < a.slots; i++) { + if (PICK[i] >= uint(a.experts)) continue; const int at2 = OFFS[PICK[i]]++; ORDER[at2] = i; WS[at2] = WEIGHT[i]; @@ -121,15 +124,16 @@ struct CombineArgs { int hidden, topk, ld, experts; }; OUT[i] = bfloat(gate / (1.0f + exp(-gate)) * float(U[i])); } -// OUT[r] (the shared expert's row) becomes bf16(bf16(sum_k bf16(Y[r*topk+k] * w_k)) + bf16(sigma_r * OUT[r])). +// OUT[r] becomes bf16(bf16(sum_k bf16(Y[slot k] * w_k)) + bf16(sigma * shared)): the shared row is OUT[r] itself +// (shared_in_out, a prompt chunk's lane product) or slot topk of Y (a decode row's extra slot); sigma is its weight. [[kernel]] void q27_expert_combine(const device bfloat* Y [[buffer(0)]], const device float* WEIGHT [[buffer(1)]], - const device bfloat* LOGITS [[buffer(2)]], device bfloat* OUT [[buffer(3)]], - constant CombineArgs& a [[buffer(4)]], uint2 g [[thread_position_in_grid]]) { - const size_t r = g.y, j = g.x; + device bfloat* OUT [[buffer(2)]], constant CombineArgs& a [[buffer(3)]], + uint2 g [[thread_position_in_grid]]) { + const size_t r = g.y, j = g.x, stride = a.topk + 1; float routed = 0.0f; for (int k = 0; k < a.topk; k++) - routed += float(bfloat(float(Y[(r * a.topk + k) * a.hidden + j]) * WEIGHT[r * a.topk + k])); - const float sigma = float(LOGITS[r * a.ld + a.experts + 1]); - const float shared = float(bfloat(sigma * float(OUT[r * a.hidden + j]))); - OUT[r * a.hidden + j] = bfloat(float(bfloat(routed)) + shared); + routed += float(bfloat(float(Y[(r * stride + k) * a.hidden + j]) * WEIGHT[r * stride + k])); + const float sigma = WEIGHT[r * stride + a.topk]; + const float row = a.shared_in_out ? float(OUT[r * a.hidden + j]) : float(Y[(r * stride + a.topk) * a.hidden + j]); + OUT[r * a.hidden + j] = bfloat(float(bfloat(routed)) + float(bfloat(sigma * row))); } diff --git a/zig/src/core/row_projection.zig b/zig/src/core/row_projection.zig index ca484fa6bb..6d9def3b9f 100644 --- a/zig/src/core/row_projection.zig +++ b/zig/src/core/row_projection.zig @@ -65,6 +65,7 @@ pub const Pipelines = struct { relu2: [variants]mtl.Pipeline, indexed: [variants]mtl.Pipeline, indexed_relu2: [variants]mtl.Pipeline, + indexed_extra: [variants]mtl.Pipeline, strided: [variants]mtl.Pipeline, pub fn load(device: mtl.Device) !Pipelines { @@ -79,6 +80,7 @@ pub const Pipelines = struct { p.relu2[i] = try mtl.Pipeline.init(device, lib, std.fmt.comptimePrint("tf_row_projection_relu2_q{d}_{s}", .{ bits, sum }), false); p.indexed[i] = try mtl.Pipeline.init(device, lib, std.fmt.comptimePrint("tf_row_projection_indexed_q{d}_{s}", .{ bits, sum }), false); p.indexed_relu2[i] = try mtl.Pipeline.init(device, lib, std.fmt.comptimePrint("tf_row_projection_indexed_relu2_q{d}_{s}", .{ bits, sum }), false); + p.indexed_extra[i] = try mtl.Pipeline.init(device, lib, std.fmt.comptimePrint("tf_row_projection_indexed_extra_q{d}_{s}", .{ bits, sum }), false); p.strided[i] = try mtl.Pipeline.init(device, lib, std.fmt.comptimePrint("tf_row_projection_strided_q{d}_{s}", .{ bits, sum }), false); }; return p; @@ -89,6 +91,7 @@ pub const Pipelines = struct { for (p.relu2) |x| x.deinit(); for (p.indexed) |x| x.deinit(); for (p.indexed_relu2) |x| x.deinit(); + for (p.indexed_extra) |x| x.deinit(); for (p.strided) |x| x.deinit(); } }; @@ -120,6 +123,16 @@ pub fn callIndexed(p: *const Pipelines, x: Experts, slots: usize, relu2: bool) ! return .{ .pipeline = if (relu2) p.indexed_relu2[i] else p.indexed[i], .dims = .{ @intCast(x.one.k), @intCast(x.one.n), @intCast(x.one.group), @intCast(slots), @intCast(x.experts), @intCast(x.repeat) }, .groups = x.one.n / (columns * groups), .slots = slots }; } +/// `callIndexed` with one more matrix `extra` of the experts' shape at id `experts` (buffers 7, 8, 9: its W, S, B). +pub fn callIndexedExtra(p: *const Pipelines, x: Experts, extra: Weights, slots: usize) !Call { + try x.validate(); + try extra.validate(); + const one = x.one; + if (extra.n != one.n or extra.k != one.k or extra.group != one.group or extra.bits != one.bits or extra.sum != one.sum) return error.UnsupportedExpertShape; + if (slots == 0 or slots % x.repeat != 0 or slots > std.math.maxInt(u16)) return error.NoRows; + return .{ .pipeline = p.indexed_extra[one.variant()], .dims = .{ @intCast(one.k), @intCast(one.n), @intCast(one.group), @intCast(slots), @intCast(x.experts), @intCast(x.repeat) }, .groups = one.n / (columns * groups), .slots = slots }; +} + test { _ = simd; _ = simd_bits; diff --git a/zig/src/families/qwen27/forward.zig b/zig/src/families/qwen27/forward.zig index c1e8831607..ecee016e74 100644 --- a/zig/src/families/qwen27/forward.zig +++ b/zig/src/families/qwen27/forward.zig @@ -79,6 +79,13 @@ pub const Stages = struct { try m.kernels.norm(e, s.hidden, f.get(.output), wts.ref(layer.post_norm), f.get(.eps), f.get(.dims), s.residual, f.get(.input), f.get(.sums), h, s.rows); e.barrier(); std.mem.swap(Ref, &s.hidden, &s.residual); + if (layer.moe) |*mo| if (!m.moe.?.sorted(mo, s.rows, m.kernels.prompt)) { + // decode rows: the routed block runs the shared expert as one more slot of the experts' kernel + try m.moe.?.block(lane.enc, m.config, mo, f, s.rows, m.kernels.prompt); + lane.enc.barrier(); + s.pending = f.get(.output); + return; + }; var up = f.get(.mixed); up.offset += @as(usize, m.config.intermediate) * 2; const width: u32 = @intCast(m.config.intermediate); diff --git a/zig/src/families/qwen27/gpu_frame.zig b/zig/src/families/qwen27/gpu_frame.zig index 5b25e82833..913d073259 100644 --- a/zig/src/families/qwen27/gpu_frame.zig +++ b/zig/src/families/qwen27/gpu_frame.zig @@ -23,7 +23,7 @@ pub const Frame = struct { pub fn init(device: mtl.Device, config: Config, capacity: u32) !Frame { if (capacity == 0 or capacity > max_rows) return error.BadFrameRows; const mp = 16 * ((capacity + 15) / 16); - const slots: usize = if (config.moe()) capacity * config.top_k else 0; + const slots: usize = if (config.moe()) capacity * (config.top_k + 1) else 0; // a row's picks and its shared expert const widths = [_]usize{ capacity * 4, 8 * 4, 3 * capacity * 4, capacity * config.hidden * 2, capacity * config.hidden * 2, capacity * config.hidden * 2, @@ -56,9 +56,9 @@ pub const Frame = struct { return .{ .buffers = buffers, .capacity = capacity, .config = config }; } - /// Routed slots a pass of `rows` makes: each row's top-k picks. + /// Slots a pass of `rows` makes: each row's top-k picks and its shared expert. pub fn slotsOf(frame: *const Frame, rows: u32) u32 { - return rows * @as(u32, @intCast(frame.config.top_k)); + return rows * @as(u32, @intCast(frame.config.top_k + 1)); } pub fn deinit(frame: *Frame) void { diff --git a/zig/src/families/qwen27/gpu_weights.zig b/zig/src/families/qwen27/gpu_weights.zig index 479b3285be..594fa7a167 100644 --- a/zig/src/families/qwen27/gpu_weights.zig +++ b/zig/src/families/qwen27/gpu_weights.zig @@ -14,8 +14,9 @@ pub const Tensor = cp.Tensor; pub const Gdn = struct { qkv: Linear, zba: Linear, out: Linear, conv: Tensor, a_log: Tensor, dt: Tensor, norm: Tensor }; pub const Attention = struct { q: Linear, kv: Linear, out: Linear, q_norm: Tensor, k_norm: Tensor }; -/// The routed block: the router's rows (the experts, the shared gate, then zero rows) and the stacked experts. -pub const Moe = struct { router: row.Weights, gate: row.Experts, up: row.Experts, down: row.Experts }; +/// The routed block: the router's rows (the experts, the shared gate, then zero rows), the stacked experts, and the +/// shared expert's own matrices (decode rows run it as one more slot of the experts' kernel). +pub const Moe = struct { router: row.Weights, gate: row.Experts, up: row.Experts, down: row.Experts, shared_gate: row.Weights, shared_up: row.Weights, shared_down: row.Weights }; /// With experts, `gu` and `down` are the shared expert's. pub const Layer = struct { norm: Tensor, post_norm: Tensor, gu: Linear, down: Linear, moe: ?Moe = null, mixer: union(@import("config.zig").Kind) { linear: Gdn, attention: Attention } }; @@ -48,7 +49,7 @@ fn uniform(p: Linear) bool { } /// Modules the graph reads only through repacked copies: their tensors stay in the files until load stages them. -const repacked_modules = [_][]const u8{ ".mlp.gate_proj", ".mlp.up_proj", ".mlp.down_proj", ".mlp.shared_expert.gate_proj", ".mlp.shared_expert.up_proj", ".mlp.shared_expert.down_proj", ".linear_attn.in_proj_qkv", ".linear_attn.in_proj_z", ".linear_attn.in_proj_b", ".linear_attn.in_proj_a", ".linear_attn.out_proj", ".self_attn.q_proj", ".self_attn.k_proj", ".self_attn.v_proj", ".self_attn.o_proj" }; +const repacked_modules = [_][]const u8{ ".mlp.gate_proj", ".mlp.up_proj", ".mlp.down_proj", ".linear_attn.in_proj_qkv", ".linear_attn.in_proj_z", ".linear_attn.in_proj_b", ".linear_attn.in_proj_a", ".linear_attn.out_proj", ".self_attn.q_proj", ".self_attn.k_proj", ".self_attn.v_proj", ".self_attn.o_proj" }; /// Whether `name` is a weight, scales or biases tensor of a repacked module or the LM head. pub fn repacked(name: []const u8, _: usize) bool { @@ -220,14 +221,23 @@ const Builder = struct { .k_norm = try b.tensor(try Name.full(a, prefix, "self_attn.k_norm.weight"), .bf16, &.{c.head_dim}), } }, }; + const shared: ?[3]projection.Triple = if (!c.moe()) null else .{ + try b.quant(try Name.full(a, prefix, widths.leaf(c, .gate)), c.intermediate, c.hidden, w[@backingInt(M.gate)]), + try b.quant(try Name.full(a, prefix, widths.leaf(c, .up)), c.intermediate, c.hidden, w[@backingInt(M.up)]), + try b.quant(try Name.full(a, prefix, widths.leaf(c, .down)), c.hidden, c.intermediate, w[@backingInt(M.down)]), + }; const moe: ?Moe = if (!c.moe()) null else .{ + .shared_gate = weightsOf(shared.?[0], c.intermediate, c.hidden), + .shared_up = weightsOf(shared.?[1], c.intermediate, c.hidden), + .shared_down = weightsOf(shared.?[2], c.hidden, c.intermediate), .router = try b.router(prefix, c, w[@backingInt(M.router)]), - .gate = try b.experts(try Name.full(a, prefix, widths.leaf(c, .egate)), c.experts, c.moe_intermediate, c.hidden, w[@backingInt(M.egate)], c.top_k), - .up = try b.experts(try Name.full(a, prefix, widths.leaf(c, .eup)), c.experts, c.moe_intermediate, c.hidden, w[@backingInt(M.eup)], c.top_k), + .gate = try b.experts(try Name.full(a, prefix, widths.leaf(c, .egate)), c.experts, c.moe_intermediate, c.hidden, w[@backingInt(M.egate)], c.top_k + 1), + .up = try b.experts(try Name.full(a, prefix, widths.leaf(c, .eup)), c.experts, c.moe_intermediate, c.hidden, w[@backingInt(M.eup)], c.top_k + 1), .down = try b.experts(try Name.full(a, prefix, widths.leaf(c, .edown)), c.experts, c.hidden, c.moe_intermediate, w[@backingInt(M.edown)], 1), }; - const gu = try b.linear(&.{ .{ .module = try Name.full(a, prefix, widths.leaf(c, .gate)), .n = c.intermediate, .w = w[@backingInt(M.gate)] }, .{ .module = try Name.full(a, prefix, widths.leaf(c, .up)), .n = c.intermediate, .w = w[@backingInt(M.up)] } }, c.hidden); - const down = try b.linear(&.{.{ .module = try Name.full(a, prefix, widths.leaf(c, .down)), .n = c.hidden, .w = w[@backingInt(M.down)] }}, c.intermediate); + // with experts the shared expert's lane layouts come from its resident tensors, else the dense MLP's are staged + const gu = if (shared) |t| try b.stack(&.{ t[0], t[1] }) else try b.linear(&.{ .{ .module = try Name.full(a, prefix, widths.leaf(c, .gate)), .n = c.intermediate, .w = w[@backingInt(M.gate)] }, .{ .module = try Name.full(a, prefix, widths.leaf(c, .up)), .n = c.intermediate, .w = w[@backingInt(M.up)] } }, c.hidden); + const down = if (shared) |t| try b.stack(&.{t[2]}) else try b.linear(&.{.{ .module = try Name.full(a, prefix, widths.leaf(c, .down)), .n = c.hidden, .w = w[@backingInt(M.down)] }}, c.intermediate); return .{ .norm = try b.tensor(try Name.full(a, prefix, "input_layernorm.weight"), .bf16, &.{c.hidden}), .post_norm = try b.tensor(try Name.full(a, prefix, "post_attention_layernorm.weight"), .bf16, &.{c.hidden}), .gu = gu, .down = down, .moe = moe, .mixer = mixer }; } }; diff --git a/zig/src/families/qwen27/moe_gpu.zig b/zig/src/families/qwen27/moe_gpu.zig index 03633e2cb6..b4d9bb0f06 100644 --- a/zig/src/families/qwen27/moe_gpu.zig +++ b/zig/src/families/qwen27/moe_gpu.zig @@ -14,7 +14,7 @@ const Ref = @import("projection.zig").Ref; const RouteArgs = extern struct { experts: i32, topk: i32, ld: i32, renorm: i32 }; const SortArgs = extern struct { slots: i32, experts: i32 }; const GluArgs = extern struct { width: i32 }; -const CombineArgs = extern struct { hidden: i32, topk: i32, ld: i32, experts: i32 }; +const CombineArgs = extern struct { hidden: i32, topk: i32, shared_in_out: i32, pad: i32 = 0 }; const R = struct { buf: mtl.Buffer, off: usize = 0 }; const Matrix = struct { w: R, s: R, b: R, n: u32, k: u32 }; const Key = struct { bits: u8, group: u16 }; @@ -82,16 +82,18 @@ pub const Kernels = struct { return g.bits == mo.up.one.bits and g.group == mo.up.one.group and g.bits == mo.down.one.bits and g.group == mo.down.one.group; } - /// f.output (the shared expert's rows of f.input) becomes the block's: weighted experts plus the gated shared row. + /// f.output becomes the block's rows: the weighted experts plus the gated shared expert, which a prompt chunk's lane + /// product left in f.output and a decode row runs as one more slot of the experts' kernel. pub fn block(k: *Kernels, e: mtl.ComputeEncoder, c: Config, mo: *const wts.Moe, f: *const Frame, rows: u32, prompt: bool) !void { var span = profile.begin(k.profiler, e, .route); errdefer span.cancel(); const enc = span.encoder; - const slots = f.slotsOf(rows); + const slots = f.slotsOf(rows); // each row's picks and its shared expert const experts: u32 = @intCast(c.experts); const topk: u32 = @intCast(c.top_k); const hidden: u32 = @intCast(c.hidden); const inter: u32 = @intCast(c.moe_intermediate); + const routed = rows * topk; // the sorted path's slots: the picks alone const units = k.sorted(mo, rows, prompt); // the router's bf16 logits [rows, ld] (column `experts`: the shared gate), on tensor units or the row kernel const ld: u32 = @intCast(mo.router.n); // a multiple of 64: both paths write that pitch @@ -134,45 +136,45 @@ pub const Kernels = struct { enc2.setBuffer(f.get(.offsets).buffer, 0, 2); enc2.setBuffer(f.get(.order).buffer, 0, 3); enc2.setBuffer(f.get(.sorted_weights).buffer, 0, 4); - enc2.setValue(SortArgs{ .slots = @intCast(slots), .experts = @intCast(experts) }, 5); + enc2.setValue(SortArgs{ .slots = @intCast(slots), .experts = @intCast(experts) }, 5); // the shared slots skipped enc2.dispatchThreads(mtl.Size.of(1, 1, 1), mtl.Size.of(1, 1, 1)); enc2.barrier(); enc2.setPipeline(k.gather); enc2.setBuffer(f.get(.input).buffer, 0, 0); enc2.setBuffer(f.get(.order).buffer, 0, 1); enc2.setBuffer(f.get(.sorted_rows).buffer, 0, 2); - enc2.setValue([2]i32{ @intCast(hidden), @intCast(topk) }, 3); - enc2.dispatchThreads(mtl.Size.of(hidden, slots, 1), mtl.Size.of(256, 1, 1)); + enc2.setValue([2]i32{ @intCast(hidden), @intCast(topk + 1) }, 3); + enc2.dispatchThreads(mtl.Size.of(hidden, routed, 1), mtl.Size.of(256, 1, 1)); enc2.barrier(); const sums = R{ .buf = f.get(.expert_sums).buffer }; const offsets = R{ .buf = f.get(.offsets).buffer }; const x = R{ .buf = f.get(.sorted_rows).buffer }; - am.rowSums(enc2, p, @intCast(mo.gate.one.group), x, sums, slots, hidden); + am.rowSums(enc2, p, @intCast(mo.gate.one.group), x, sums, routed, hidden); enc2.barrier(); - am.gather(enc2, p, x, sums, matrix(mo.gate.one), offsets, R{ .buf = f.get(.expert_gate).buffer }, slots, experts); - am.gather(enc2, p, x, sums, matrix(mo.up.one), offsets, R{ .buf = f.get(.expert_up).buffer }, slots, experts); + am.gather(enc2, p, x, sums, matrix(mo.gate.one), offsets, R{ .buf = f.get(.expert_gate).buffer }, routed, experts); + am.gather(enc2, p, x, sums, matrix(mo.up.one), offsets, R{ .buf = f.get(.expert_up).buffer }, routed, experts); enc2.barrier(); - k.activation(enc2, f, slots, inter); + k.activation(enc2, f, routed, inter); enc2.barrier(); const act = R{ .buf = f.get(.expert_act).buffer }; - am.rowSums(enc2, p, @intCast(mo.down.one.group), act, sums, slots, inter); + am.rowSums(enc2, p, @intCast(mo.down.one.group), act, sums, routed, inter); enc2.barrier(); - am.gatherTo(enc2, p, act, sums, matrix(mo.down.one), offsets, R{ .buf = f.get(.expert_out).buffer }, R{ .buf = f.get(.order).buffer }, slots, experts); + am.gatherTo(enc2, p, act, sums, matrix(mo.down.one), offsets, R{ .buf = f.get(.expert_out).buffer }, R{ .buf = f.get(.order).buffer }, routed, experts); } else { - try k.indexed(enc2, pipes, mo.gate, slots, f.get(.input), f.get(.picks), f.get(.expert_gate)); - try k.indexed(enc2, pipes, mo.up, slots, f.get(.input), f.get(.picks), f.get(.expert_up)); + // every slot through the experts' kernel, the shared expert as the extra matrix of each product + try k.indexed(enc2, pipes, mo.gate, mo.shared_gate, slots, f.get(.input), f.get(.picks), f.get(.expert_gate)); + try k.indexed(enc2, pipes, mo.up, mo.shared_up, slots, f.get(.input), f.get(.picks), f.get(.expert_up)); enc2.barrier(); k.activation(enc2, f, slots, inter); enc2.barrier(); - try k.indexed(enc2, pipes, mo.down, slots, f.get(.expert_act), f.get(.picks), f.get(.expert_out)); + try k.indexed(enc2, pipes, mo.down, mo.shared_down, slots, f.get(.expert_act), f.get(.picks), f.get(.expert_out)); } enc2.barrier(); enc2.setPipeline(k.combine); enc2.setBuffer(f.get(.expert_out).buffer, 0, 0); enc2.setBuffer(f.get(.route_weights).buffer, 0, 1); - enc2.setBuffer(f.get(.router).buffer, 0, 2); - enc2.setBuffer(f.get(.output).buffer, 0, 3); - enc2.setValue(CombineArgs{ .hidden = @intCast(hidden), .topk = @intCast(topk), .ld = @intCast(ld), .experts = @intCast(experts) }, 4); + enc2.setBuffer(f.get(.output).buffer, 0, 2); + enc2.setValue(CombineArgs{ .hidden = @intCast(hidden), .topk = @intCast(topk), .shared_in_out = @intFromBool(units) }, 3); enc2.dispatchThreads(mtl.Size.of(hidden, rows, 1), mtl.Size.of(256, 1, 1)); try products.end(); } @@ -186,9 +188,9 @@ pub const Kernels = struct { e.dispatchThreads(mtl.Size.of(inter, slots, 1), mtl.Size.of(256, 1, 1)); } - /// Slot i takes expert ids[i] on input row i / repeat. - fn indexed(_: *Kernels, e: mtl.ComputeEncoder, pipes: *const row.Pipelines, x: row.Experts, slots: u32, in: Ref, ids: Ref, out: Ref) !void { - const call = try row.callIndexed(pipes, x, slots, false); + /// Slot i takes expert ids[i] on input row i / repeat; id `experts` is the shared expert's `extra` matrix. + fn indexed(_: *Kernels, e: mtl.ComputeEncoder, pipes: *const row.Pipelines, x: row.Experts, extra: row.Weights, slots: u32, in: Ref, ids: Ref, out: Ref) !void { + const call = try row.callIndexedExtra(pipes, x, extra, slots); e.setPipeline(call.pipeline); e.setBuffer(in.buffer, in.offset, 0); e.setBuffer(x.one.w, x.one.w_off, 1); @@ -197,6 +199,9 @@ pub const Kernels = struct { e.setValue(call.dims, 4); e.setBuffer(out.buffer, out.offset, 5); e.setBuffer(ids.buffer, ids.offset, 6); + e.setBuffer(extra.w, extra.w_off, 7); + e.setBuffer(extra.scales, extra.s_off, 8); + e.setBuffer(extra.biases, extra.b_off, 9); e.dispatchThreads(mtl.Size.of(call.groups * call.threads, call.slots, 1), mtl.Size.of(call.threads, 1, 1)); } }; From f11bef4ec463efd35333ba703f4c2559d5563c8d Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 20 Oct 2026 14:15:44 -0500 Subject: [PATCH 13/21] qwen27: one-line comments for the lean check --- zig/kernels/metal/core/row_projection.metal | 2 +- zig/kernels/metal/qwen27/moe.metal | 5 ++--- zig/src/families/qwen27/gpu_weights.zig | 3 +-- zig/src/families/qwen27/moe_gpu.zig | 3 +-- 4 files changed, 5 insertions(+), 8 deletions(-) diff --git a/zig/kernels/metal/core/row_projection.metal b/zig/kernels/metal/core/row_projection.metal index 457cebf246..842f96434f 100644 --- a/zig/kernels/metal/core/row_projection.metal +++ b/zig/kernels/metal/core/row_projection.metal @@ -190,7 +190,7 @@ template if (lane == 0) for (int j = 0; j < RPS; j++) out[j] = tf_rp_out(acc[0][j]); } -// `tf_row_projection_indexed` with an extra matrix at id `experts` (buffers 7-9): a shared expert rides with the routed ones. +// tf_row_projection_indexed with an extra matrix at id `experts` (buffers 7-9): a shared expert among the routed ones. template [[kernel]] void tf_row_projection_indexed_extra(const device bfloat* X [[buffer(0)]], const device uint32_t* W [[buffer(1)]], diff --git a/zig/kernels/metal/qwen27/moe.metal b/zig/kernels/metal/qwen27/moe.metal index 68dbd99504..21c4493fc5 100644 --- a/zig/kernels/metal/qwen27/moe.metal +++ b/zig/kernels/metal/qwen27/moe.metal @@ -84,7 +84,7 @@ struct CombineArgs { int hidden, topk, shared_in_out, pad; }; } } -// Routed slots (id < experts) sorted by expert (stable) for core/affine_mm: OFFS (first slots, then the end), ORDER, WS. +// Routed slots (id < experts) sorted by expert (stable) for core/affine_mm: OFFS (first slots, the end), ORDER, WS. [[kernel]] void q27_route_sort(const device uint* PICK [[buffer(0)]], const device float* WEIGHT [[buffer(1)]], device int* OFFS [[buffer(2)]], device int* ORDER [[buffer(3)]], device float* WS [[buffer(4)]], constant SortArgs& a [[buffer(5)]], @@ -124,8 +124,7 @@ struct CombineArgs { int hidden, topk, shared_in_out, pad; }; OUT[i] = bfloat(gate / (1.0f + exp(-gate)) * float(U[i])); } -// OUT[r] becomes bf16(bf16(sum_k bf16(Y[slot k] * w_k)) + bf16(sigma * shared)): the shared row is OUT[r] itself -// (shared_in_out, a prompt chunk's lane product) or slot topk of Y (a decode row's extra slot); sigma is its weight. +// OUT[r] = bf16(bf16(sum_k bf16(Y[k] * w_k)) + bf16(sigma * shared)); shared: OUT[r] (shared_in_out) or Y's slot topk. [[kernel]] void q27_expert_combine(const device bfloat* Y [[buffer(0)]], const device float* WEIGHT [[buffer(1)]], device bfloat* OUT [[buffer(2)]], constant CombineArgs& a [[buffer(3)]], uint2 g [[thread_position_in_grid]]) { diff --git a/zig/src/families/qwen27/gpu_weights.zig b/zig/src/families/qwen27/gpu_weights.zig index 594fa7a167..743f627691 100644 --- a/zig/src/families/qwen27/gpu_weights.zig +++ b/zig/src/families/qwen27/gpu_weights.zig @@ -14,8 +14,7 @@ pub const Tensor = cp.Tensor; pub const Gdn = struct { qkv: Linear, zba: Linear, out: Linear, conv: Tensor, a_log: Tensor, dt: Tensor, norm: Tensor }; pub const Attention = struct { q: Linear, kv: Linear, out: Linear, q_norm: Tensor, k_norm: Tensor }; -/// The routed block: the router's rows (the experts, the shared gate, then zero rows), the stacked experts, and the -/// shared expert's own matrices (decode rows run it as one more slot of the experts' kernel). +/// The routed block: the router (experts, shared gate, zero rows), the stacked experts, the shared expert's matrices. pub const Moe = struct { router: row.Weights, gate: row.Experts, up: row.Experts, down: row.Experts, shared_gate: row.Weights, shared_up: row.Weights, shared_down: row.Weights }; /// With experts, `gu` and `down` are the shared expert's. pub const Layer = struct { norm: Tensor, post_norm: Tensor, gu: Linear, down: Linear, moe: ?Moe = null, mixer: union(@import("config.zig").Kind) { linear: Gdn, attention: Attention } }; diff --git a/zig/src/families/qwen27/moe_gpu.zig b/zig/src/families/qwen27/moe_gpu.zig index b4d9bb0f06..f400871d94 100644 --- a/zig/src/families/qwen27/moe_gpu.zig +++ b/zig/src/families/qwen27/moe_gpu.zig @@ -82,8 +82,7 @@ pub const Kernels = struct { return g.bits == mo.up.one.bits and g.group == mo.up.one.group and g.bits == mo.down.one.bits and g.group == mo.down.one.group; } - /// f.output becomes the block's rows: the weighted experts plus the gated shared expert, which a prompt chunk's lane - /// product left in f.output and a decode row runs as one more slot of the experts' kernel. + /// f.output becomes the weighted experts plus the gated shared expert (f.output for a prompt chunk, else a slot). pub fn block(k: *Kernels, e: mtl.ComputeEncoder, c: Config, mo: *const wts.Moe, f: *const Frame, rows: u32, prompt: bool) !void { var span = profile.begin(k.profiler, e, .route); errdefer span.cancel(); From 30e5dcb4bfbbb03bf24a0d15332cf739a72a0a3b Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 20 Oct 2026 14:51:05 -0500 Subject: [PATCH 14/21] qwen27: an absent norm_topk_prob means renormalized top-k weights, as the reference config and the CUDA family read it --- zig/src/families/qwen27/config.zig | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/zig/src/families/qwen27/config.zig b/zig/src/families/qwen27/config.zig index 213a1db312..b0aa5a7248 100644 --- a/zig/src/families/qwen27/config.zig +++ b/zig/src/families/qwen27/config.zig @@ -172,7 +172,8 @@ pub fn parse(gpa: std.mem.Allocator, json: []const u8, generation: ?[]const u8) c.experts = experts; c.top_k = try required(t, "num_experts_per_tok"); c.moe_intermediate = try required(t, "moe_intermediate_size"); - c.renorm = try flag(t, "norm_topk_prob"); + // absent means true, as Qwen3.5-MoE's reference config and the CUDA family read it + c.renorm = if (t.get("norm_topk_prob")) |v| (if (v == .bool) v.bool else return error.BadConfig) else true; if (c.experts > 256 or c.top_k >= c.experts or c.top_k > 15 or c.moe_intermediate != c.intermediate) return error.UnsupportedRouting; } if (t.get("max_position_embeddings")) |v| c.max_position = try integer(v, false); @@ -276,7 +277,7 @@ test "a routed config reads its experts, the shared expert's width as the layer' try std.testing.expectEqual(@as(usize, 8), c.top_k); try std.testing.expectEqual(@as(usize, 1024), c.moe_intermediate); try std.testing.expectEqual(@as(usize, 1024), c.intermediate); - try std.testing.expect(!c.renorm); + try std.testing.expect(c.renorm); // absent: renormalized try std.testing.expectEqual(Kind.attention, c.kind(3)); try std.testing.expectEqual(Kind.linear, c.kind(2)); // a routed model type without experts, and a dense one with them, are refused From 146a658847ea004102d05ea988fd9583e4bb70e3 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Tue, 20 Oct 2026 18:58:57 -1000 Subject: [PATCH 15/21] qwen27: score and decide passes queued onto the lane host's idle engine thread Ornith serves /v1/decisions and /v1/systemone while ordinary text submits keep running: info.decides stays off and the hooks appear only when a decider is set. A lane host's runner is owned by its engine thread, so HTTP threads cannot run the Metal passes themselves. engine_api gains a DecisionTask/Decider pair; the host queues a Work per score/decide call, the round and generic driver loops run the oldest once no stream is admitted, filling or active, and shutdown or a stale device fails every waiter instead of sleeping past the thread's end. qwen27's passes run one prompt each on the existing LM head path: score reads the labels' logits at the row after its prompt plus the vocabulary logsumexp, and decide keeps the trunk's slot state once, restoring it per branch so every branch's last-row label logits stand on the same trunk. Pictures and requests past the context window are refused. --- zig/src/core/engine_api.zig | 14 +++ zig/src/core/lane_host.zig | 23 +++- zig/src/core/lane_host_decide.zig | 80 ++++++++++++ zig/src/core/lane_host_decide_test.zig | 162 +++++++++++++++++++++++++ zig/src/core/lane_host_tests.zig | 1 + zig/src/core/lane_rounds.zig | 3 + zig/src/native/qwen27_decide.zig | 104 ++++++++++++++++ zig/src/native/qwen27_served.zig | 8 ++ 8 files changed, 394 insertions(+), 1 deletion(-) create mode 100644 zig/src/core/lane_host_decide.zig create mode 100644 zig/src/core/lane_host_decide_test.zig create mode 100644 zig/src/native/qwen27_decide.zig diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 7a8bfd4edd..09d37295ec 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -302,6 +302,20 @@ pub const Decision = struct { arena: ?std.mem.Allocator = null, // the request's, for what the engine hands back }; +/// One decision pass a host queues for its engine thread: a score's single row, or a decide's packed set. +pub const DecisionTask = union(enum) { + /// `labels`' logits at the row after `prompt`; `out` takes the vocabulary logsumexp when `mass`. + score: struct { prompt: []const u32, labels: []const u32, logits: []f64, mass: bool, out: f64 = 0 }, + /// The packed decision: each branch's label logits at its last row after the shared trunk. + decide: *Decision, +}; + +/// A lane-hosted backend's decision passes, run on the engine thread once no stream is live. +pub const Decider = struct { + ctx: *anyopaque, + run: *const fn (ctx: *anyopaque, task: *DecisionTask) error{ Failed, Refused }!void, +}; + /// An opened engine; ``close`` stops its thread and frees its backend. pub const Opened = struct { engine: Engine, diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 7f24ff457e..63f7b00748 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -5,6 +5,7 @@ const api = @import("engine_api.zig"); const pc = @import("prompt_cache.zig"); const status = @import("lane_host_status.zig"); const learn = @import("lane_host_learn.zig"); +const decide = @import("lane_host_decide.zig"); const lone = @import("lane_host_lone.zig"); const fill = @import("lane_host_fill.zig"); const admit = @import("lane_host_admit.zig"); @@ -63,6 +64,9 @@ pub const LaneHost = struct { memory: ?api.MemorySource = null, // the backend's memory counts; null: Engine.memory reports none explain: ?api.Explain = null, // the backend's words for a request it refuses; null: the error's name learner: ?api.Learner = null, // the family's Sliding Weights learner; null: learn requests are refused + decider: ?api.Decider = null, // the backend's decision passes, run while the engine is idle + decisions: std.ArrayList(*decide.Work) = .empty, // score/decide calls waiting for an idle engine, in order + decided: std.Io.Condition = .init, // a queued decision finished; its requester wakes lessons: std.ArrayList(Lesson) = .empty, // learn requests waiting for an idle engine lesson_open: bool = false, // the learner holds a begun job (engine thread only) admissions: u64 = 0, // jobs admitted so far: each job's `seq` @@ -102,6 +106,7 @@ pub const LaneHost = struct { h.lessons.deinit(h.gpa); h.queued.deinit(h.gpa); h.admitted.deinit(h.gpa); + h.decisions.deinit(h.gpa); h.filling.deinit(h.gpa); h.cancels.deinit(h.gpa); h.decoded.deinit(h.gpa); @@ -113,8 +118,22 @@ pub const LaneHost = struct { /// A drain's deadline: the round in flight completes on every rank, so the ranks stop at the same boundary. pub const halt = lane_halt.halt; + /// The host's engine; a backend that set `decider` serves score/decide, else the hooks stay absent. pub fn engine(h: *LaneHost) Engine { - return .{ .ctx = h, .vtable = &.{ + return .{ .ctx = h, .vtable = if (h.decider != null) &.{ + .info = infoFn, + .submit = submitFn, + .cancel = cancelFn, + .status = status.read, + .memory = memoryFn, + .keepalive = keepaliveFn, + .score = decide.score, + .decide = decide.decide, + .learn = learn.submit, + .recall = recalls.attach, + .imprint = rows_host.ask, + .drain = rows_host.drain, + } else &.{ .info = infoFn, .submit = submitFn, .cancel = cancelFn, @@ -507,6 +526,7 @@ pub const LaneHost = struct { ends.takeCancels(h); rows_host.answer(h); learn.refuse(h); + decide.step(h); // a queued decision takes an idle engine before new work; a stale one fails it while (h.admitOne()) {} h.fillOne(); h.core.sweep() catch |err| { @@ -519,6 +539,7 @@ pub const LaneHost = struct { if (h.closing) { const left = h.queued.items.len + h.admitted.items.len; h.unlock(); + decide.failAll(h); // their requesters must not sleep past the engine thread's end if (h.lesson_open) { h.learner.?.abort(h.learner.?.ctx); h.lesson_open = false; diff --git a/zig/src/core/lane_host_decide.zig b/zig/src/core/lane_host_decide.zig new file mode 100644 index 0000000000..8e31a73251 --- /dev/null +++ b/zig/src/core/lane_host_decide.zig @@ -0,0 +1,80 @@ +//! A lane host's decision passes: HTTP threads queue them; the engine thread runs them while no stream lives. +const std = @import("std"); +const api = @import("engine_api.zig"); +const LaneHost = @import("lane_host.zig").LaneHost; + +/// One queued decision pass: its task in, the failure and completion its requester waits on. +pub const Work = struct { + task: api.DecisionTask, + failed: ?error{ Failed, Refused } = null, + finished: bool = false, // written under the host's mutex; `decided` is broadcast with it +}; + +/// Engine.score: the prompt's label logits queued for the engine thread (a score has no Refused). +pub fn score(ctx: *anyopaque, prompt: []const u32, labels: []const u32, logits: []f64, mass: bool) error{Failed}!f64 { + var work = Work{ .task = .{ .score = .{ .prompt = prompt, .labels = labels, .logits = logits, .mass = mass } } }; + wait(self(ctx), &work) catch return error.Failed; + return work.task.score.out; +} + +/// Engine.decide: the packed decision queued for the engine thread; its refusal comes back as given. +pub fn decide(ctx: *anyopaque, d: *api.Decision) error{ Failed, Refused }!void { + var work = Work{ .task = .{ .decide = d } }; + return wait(self(ctx), &work); +} + +/// Queue `work` and sleep until the engine thread's pass lands; it stays valid on this stack frame. +pub fn wait(h: *LaneHost, work: *Work) error{ Failed, Refused }!void { + h.mutex.lockUncancelable(h.io); + if (h.closing or h.stale or h.decider == null) { + h.mutex.unlock(h.io); + return error.Failed; + } + h.decisions.append(h.gpa, work) catch { + h.mutex.unlock(h.io); + return error.Failed; + }; + h.mutex.unlock(h.io); + h.wake.signal(h.io); + h.mutex.lockUncancelable(h.io); + defer h.mutex.unlock(h.io); + while (!work.finished) h.decided.waitUncancelable(h.io, &h.mutex); + if (work.failed) |e| return e; +} + +/// The oldest queued decision's pass: an idle engine runs it, a stale one fails the waiters instead. +pub fn step(h: *LaneHost) void { + h.lock(); + if (h.stale) { + h.unlock(); + return failAll(h); + } + const idle = !h.closing and h.decider != null and h.admitted.items.len == 0 and h.filling.items.len == 0 and h.core.activeCount() == 0; + const work: ?*Work = if (idle and h.decisions.items.len > 0) h.decisions.orderedRemove(0) else null; + h.unlock(); + const w = work orelse return; + const decider = h.decider.?; + decider.run(decider.ctx, &w.task) catch |e| { + w.failed = e; + }; + h.lock(); + w.finished = true; + h.unlock(); + h.decided.broadcast(h.io); +} + +/// Every queued decision failed (the host is closing or the device went stale): requesters wake. +pub fn failAll(h: *LaneHost) void { + h.lock(); + while (h.decisions.items.len > 0) { + const work = h.decisions.orderedRemove(0); + work.failed = error.Failed; + work.finished = true; + } + h.unlock(); + h.decided.broadcast(h.io); +} + +fn self(ctx: *anyopaque) *LaneHost { + return @ptrCast(@alignCast(ctx)); +} diff --git a/zig/src/core/lane_host_decide_test.zig b/zig/src/core/lane_host_decide_test.zig new file mode 100644 index 0000000000..70c8930144 --- /dev/null +++ b/zig/src/core/lane_host_decide_test.zig @@ -0,0 +1,162 @@ +//! The lane host's decision hooks: queued passes on an idle engine, in the decider's own words. +const std = @import("std"); +const lanes = @import("lanes"); +const api = @import("engine_api.zig"); +const LaneHost = @import("lane_host.zig").LaneHost; +const Id = api.Id; +const Event = api.Event; +const Reason = api.Reason; +const Request = api.Request; + +/// A decider that echoes its task: score gives each label its id as a logit; decide adds the branch's length. +const Record = struct { + core: *lanes.Engine, + calls: usize = 0, + idle: bool = true, // activeCount at every pass: decisions only ever run on an idle engine + fail: ?error{ Failed, Refused } = null, + + fn run(ctx: *anyopaque, task: *api.DecisionTask) error{ Failed, Refused }!void { + const r: *Record = @ptrCast(@alignCast(ctx)); + defer r.calls += 1; + r.idle = r.idle and r.core.activeCount() == 0; + if (r.fail) |e| return e; + switch (task.*) { + .score => |*s| { + for (s.labels, s.logits) |label, *x| x.* = @floatFromInt(label); + s.out = if (s.mass) 42.5 else 0; + }, + .decide => |d| { + var at: usize = 0; + for (d.branches, d.counts) |b, n| { + for (d.labels[at..][0..n], d.logits[at..][0..n]) |label, *x| x.* = @floatFromInt(label + b.len); + d.read += b.len; + at += n; + } + d.read += d.trunk.len; + }, + } + } +}; + +fn config(gpa: std.mem.Allocator) !lanes.Config { + return lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7); +} + +test "a lane host without a decider keeps score/decide absent; one with it serves both" { + const gpa = std.testing.allocator; + var cfg = try config(gpa); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.backend(), clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 1 }); + try std.testing.expect(host.engine().vtable.score == null); + try std.testing.expect(host.engine().vtable.decide == null); + var record = Record{ .core = &core }; + host.decider = .{ .ctx = &record, .run = Record.run }; + const e = host.engine(); + try std.testing.expect(e.vtable.score != null); + try std.testing.expect(e.vtable.decide != null); + try host.start(); + defer host.stop(); + const prompt = [_]u32{ 2, 7, 1, 8 }; + const labels = [_]u32{ 5, 9 }; + var logits = [_]f64{ 0, 0 }; + try std.testing.expectEqual(@as(f64, 42.5), try e.score(&prompt, &labels, &logits, true)); + try std.testing.expectEqualSlices(f64, &.{ 5, 9 }, &logits); + const branches = [_][]const u32{ &.{ 4, 4 }, &.{6} }; + const counts = [_]usize{ 1, 1 }; + const wanted = [_]u32{ 30, 40 }; + var out = [_]f64{ 0, 0 }; + var d = api.Decision{ .trunk = &prompt, .branches = &branches, .counts = &counts, .labels = &wanted, .logits = &out }; + try e.decide(&d); + try std.testing.expectEqualSlices(f64, &.{ 32, 41 }, &out); + try std.testing.expectEqual(@as(usize, prompt.len + 3), d.read); + try std.testing.expectEqual(@as(usize, 2), record.calls); + try std.testing.expect(record.idle); +} + +test "a decider's refusal and failure reach the requester" { + const gpa = std.testing.allocator; + var cfg = try config(gpa); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.backend(), clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 1 }); + var record = Record{ .core = &core }; + host.decider = .{ .ctx = &record, .run = Record.run }; + try host.start(); + defer host.stop(); + const e = host.engine(); + const prompt = [_]u32{ 2, 7, 1, 8 }; + const labels = [_]u32{5}; + var logits = [_]f64{0}; + const branches = [_][]const u32{&.{4}}; + const counts = [_]usize{1}; + var d = api.Decision{ .trunk = &prompt, .branches = &branches, .counts = &counts, .labels = &labels, .logits = &logits }; + record.fail = error.Refused; + try std.testing.expectError(error.Refused, e.decide(&d)); + try std.testing.expectError(error.Failed, e.score(&prompt, &labels, &logits, false)); // a score has no Refused + record.fail = error.Failed; + try std.testing.expectError(error.Failed, e.decide(&d)); + try std.testing.expectError(error.Failed, e.score(&prompt, &labels, &logits, false)); +} + +test "a queued decision runs once the live stream ends, and the text reply still lands" { + const gpa = std.testing.allocator; + var cfg = try config(gpa); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.backend(), clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 1 }); + var record = Record{ .core = &core }; + host.decider = .{ .ctx = &record, .run = Record.run }; + try host.start(); + defer host.stop(); + const Box = struct { + mutex: std.Io.Mutex = .init, + done: ?Reason = null, + tokens: usize = 0, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens += t.len, + .finished => |f| b.done = f.reason, + else => {}, + } + } + }; + const prompt = [_]u32{ 2, 7, 1, 8 }; + var box: Box = .{}; + const request: Request = .{ .prompt = &prompt, .max_tokens = 64 }; + const e = host.engine(); + try e.submit(1, &request, .{ .ctx = &box, .event = Box.event }); + const labels = [_]u32{5}; + var logits = [_]f64{0}; + // whatever round the reply is on, the pass waits for it to leave the engine + try std.testing.expectEqual(@as(f64, 0), try e.score(&prompt, &labels, &logits, false)); + try std.testing.expectEqualSlices(f64, &.{5}, &logits); + while (true) { + box.mutex.lockUncancelable(std.testing.io); + const done = box.done; + box.mutex.unlock(std.testing.io); + if (done) |r| { + try std.testing.expectEqual(Reason.length, r); + break; + } + std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + } + try std.testing.expectEqual(@as(usize, 64), box.tokens); + try std.testing.expect(record.idle); + try std.testing.expectEqual(@as(usize, 1), record.calls); +} diff --git a/zig/src/core/lane_host_tests.zig b/zig/src/core/lane_host_tests.zig index d121279474..4c1f68c876 100644 --- a/zig/src/core/lane_host_tests.zig +++ b/zig/src/core/lane_host_tests.zig @@ -2,6 +2,7 @@ test { _ = @import("lane_host_test.zig"); + _ = @import("lane_host_decide_test.zig"); _ = @import("lane_rounds_test.zig"); _ = @import("lane_halt_test.zig"); _ = @import("lane_host_end_test.zig"); diff --git a/zig/src/core/lane_rounds.zig b/zig/src/core/lane_rounds.zig index 74a476d540..161eb40c7f 100644 --- a/zig/src/core/lane_rounds.zig +++ b/zig/src/core/lane_rounds.zig @@ -5,6 +5,7 @@ const api = @import("engine_api.zig"); const host = @import("lane_host.zig"); const lane_stream = @import("lane_stream.zig"); const learn = @import("lane_host_learn.zig"); +const decide = @import("lane_host_decide.zig"); const status = @import("lane_host_status.zig"); const ends = @import("lane_host_end.zig"); const rows_host = @import("lane_host_rows.zig"); @@ -265,10 +266,12 @@ pub fn run(h: *LaneHost, rd: Rounds) void { ends.takeCancels(h); rows_host.answer(h); learn.refuse(h); + decide.step(h); // a queued decision takes an idle engine before new work; a stale one fails it h.lock(); if (h.closing) { const left = h.queued.items.len + h.admitted.items.len; h.unlock(); + decide.failAll(h); // their requesters must not sleep past the engine thread's end if (h.lesson_open) { h.learner.?.abort(h.learner.?.ctx); h.lesson_open = false; diff --git a/zig/src/native/qwen27_decide.zig b/zig/src/native/qwen27_decide.zig new file mode 100644 index 0000000000..0753616cab --- /dev/null +++ b/zig/src/native/qwen27_decide.zig @@ -0,0 +1,104 @@ +//! The 27B lane backend's decision passes on an idle engine: the trunk's state kept once, each branch read over it. +const std = @import("std"); +const api = @import("engine_api"); +const tf = @import("tensorfold"); +const q = tf.qwen27; +const lb = @import("qwen27_lanes.zig"); + +/// A queued decision's GPU work (LaneHost runs it alone): score reads the row after its prompt; decide, each branch's last row. +pub fn run(b: *lb.Metal, window: u32, task: *api.DecisionTask) error{ Failed, Refused }!void { + switch (task.*) { + .score => |*s| s.out = try score(b, s.prompt, s.labels, s.logits, s.mass), + .decide => |d| try decide(b, window, d), + } +} + +/// `labels`' logits at the row after `prompt`: a prefill's `.last` head leaves them in logits row 0. +fn score(b: *lb.Metal, prompt: []const u32, labels: []const u32, logits: []f64, mass: bool) error{Failed}!f64 { + const r = b.runner; + if (prompt.len == 0 or prompt.len > r.capacity or logits.len < labels.len) return error.Failed; + try reset(b); + const session = q.session.Session{ .runner = r }; + session.prefill(prompt, b.chunk) catch |e| return failed("a score's prompt pass", e); + return read(row0(r), labels, logits, mass); +} + +/// Each branch's label logits at its last row after the shared trunk: the trunk's state is kept once +/// and restored a branch at a time; every branch's rows stand on the same trunk state. +fn decide(b: *lb.Metal, window: u32, d: *api.Decision) error{ Failed, Refused }!void { + const r = b.runner; + if (d.pictures.len > 0) return refuse(d, "decisions with pictures are not served by this engine yet"); + if (d.trunk.len == 0 or d.branches.len == 0 or d.counts.len != d.branches.len) return error.Failed; + var wanted: usize = 0; + for (d.branches, d.counts) |br, n| { + if (br.len == 0) return error.Failed; + if (d.trunk.len + br.len > window) + return refuse(d, std.fmt.allocPrint(d.arena orelse return error.Failed, "a decision sequence exceeds the {d}-token context", .{window}) catch return error.Failed); + wanted = std.math.add(usize, wanted, n) catch return error.Failed; + } + if (d.labels.len < wanted or d.logits.len < wanted) return error.Failed; + try reset(b); + const session = q.session.Session{ .runner = r }; + session.prefill(d.trunk, b.chunk) catch |e| return failed("the trunk's prompt pass", e); + const st = q.snapshot.saveSlot(b.gpa, r, 0, null, @intCast(d.trunk.len)) catch |e| return failed("keeping the trunk's state", e); + defer q.snapshot.drop(b.gpa, st); + d.read = d.trunk.len; + var at: usize = 0; + for (d.branches, d.counts) |br, n| { + q.snapshot.restoreSlot(b.gpa, r, 0, null, st) catch |e| return failed("restoring the trunk's state", e); + session.prefill(br, b.chunk) catch |e| return failed("a branch's prompt pass", e); + _ = try read(row0(r), d.labels[at..][0..n], d.logits[at..][0..n], false); + d.read += br.len; + at += n; + } + reset(b) catch {}; // slot 0 stands empty for the next stream, whatever its prompt was +} + +/// Slot 0 empty for a pass: a pending keep ends first, and the drafter's reset covers its session too. +fn reset(b: *lb.Metal) error{Failed}!void { + b.settle() catch {}; + b.dropRound(); + b.runner.rope[0] = .{}; + b.runner.images = null; + if (b.draft) |d| d.reset() catch |e| return failed("resetting slot 0", e); + return b.runner.reset(0) catch |e| failed("resetting slot 0", e); +} + +/// The last pass's logits row 0: bf16, `vocab` wide; the `.last` head leaves its one row there. +fn row0(r: *const q.decode_round.Runner) []const u16 { + const logits = r.model.frame.get(.logits); + const words: [*]const u16 = @ptrCast(@alignCast(logits.buffer.contents() + logits.offset)); + return words[0..r.model.config.vocab]; +} + +/// The row's `labels` logits out, then the vocabulary logsumexp when `mass` (else 0). +fn read(row: []const u16, labels: []const u32, logits: []f64, mass: bool) error{Failed}!f64 { + for (labels, logits) |label, *x| { + if (label >= row.len) return error.Failed; + x.* = try value(row[@intCast(label)]); + } + if (!mass) return 0; + var peak: f64 = -std.math.inf(f64); + for (row) |bits| peak = @max(peak, try value(bits)); + var total: f64 = 0; + for (row) |bits| total += @exp(try value(bits) - peak); + return peak + @log(total); +} + +/// A bf16 logit as f64; a nonfinite one fails the pass, as a draw's argmax check does. +fn value(bits: u16) error{Failed}!f64 { + const v: f64 = @as(f32, @bitCast(@as(u32, bits) << 16)); + return if (std.math.isFinite(v)) v else error.Failed; +} + +/// A pass's error logged once, the decision hooks' Failed. +fn failed(what: []const u8, err: anyerror) error{Failed} { + std.log.err("qwen27 decision pass failed ({s}): {s}", .{ what, @errorName(err) }); + return error.Failed; +} + +/// The decision refused, in `problem`'s words for the requester. +fn refuse(d: *api.Decision, problem: []const u8) error{Refused} { + d.problem = problem; + return error.Refused; +} diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index ef73190ee1..b47f59dcb7 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -12,6 +12,7 @@ const cache_fit = @import("cache_fit.zig"); const Resident = @import("resident.zig").Resident; const lb = @import("qwen27_lanes.zig"); pub const lanes_backend = lb; +const decide = @import("qwen27_decide.zig"); const disk = @import("qwen27_disk.zig"); const images = @import("qwen27_images.zig"); const Allocator = std.mem.Allocator; @@ -82,6 +83,12 @@ pub const Host = struct { const h = of(ptr); return images.prepare(h.vision.?, h.clock.io, a, bytes, max_tokens); } + + /// LaneHost's decision hook: Ornith's label reads over a shared trunk, on an idle engine. + fn decideRun(ptr: *anyopaque, task: *api.DecisionTask) error{ Failed, Refused }!void { + const h = of(ptr); + return decide.run(h.back, h.host.info_.context_window, task); + } }; /// The reply's state at `stands` (its end or one short), born with its request's prompt: that request's marks stay. @@ -324,6 +331,7 @@ pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.df h.host = api.LaneHost.init(gpa, io, &h.core, .{ .name = "qwen27", .structures = @import("qwen27_masks.zig").supported, .lanes = slots, .context_window = capacity, .prefill_step = @intCast(h.back.chunk), .vision = offer }); h.host.memory = .{ .ctx = h, .read = readMemory }; h.host.explain = .{ .text = words }; + h.host.decider = .{ .ctx = h, .run = Host.decideRun }; // the engine still answers text: info_.decides stays off h.host.decoded_rows = true; // a reply's rows prefill with the bits decoding gave them, with the cache on or off // a lone greedy request takes the serial loop on tensor units, where it is faster (TF_QWEN27_LANES overrides) const units = model.device.tensorUnits(); From 227b961d223b8f71ab5ce8e06b57e8232eca03f3 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Wed, 21 Oct 2026 16:50:17 -1000 Subject: [PATCH 16/21] qwen27: one-line comments for the lean check --- zig/src/native/qwen27_decide.zig | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/zig/src/native/qwen27_decide.zig b/zig/src/native/qwen27_decide.zig index 0753616cab..8a1ca2a31c 100644 --- a/zig/src/native/qwen27_decide.zig +++ b/zig/src/native/qwen27_decide.zig @@ -5,7 +5,7 @@ const tf = @import("tensorfold"); const q = tf.qwen27; const lb = @import("qwen27_lanes.zig"); -/// A queued decision's GPU work (LaneHost runs it alone): score reads the row after its prompt; decide, each branch's last row. +/// A queued decision's GPU work on the idle host: score the row after its prompt; decide, each branch's last row. pub fn run(b: *lb.Metal, window: u32, task: *api.DecisionTask) error{ Failed, Refused }!void { switch (task.*) { .score => |*s| s.out = try score(b, s.prompt, s.labels, s.logits, s.mass), @@ -23,8 +23,7 @@ fn score(b: *lb.Metal, prompt: []const u32, labels: []const u32, logits: []f64, return read(row0(r), labels, logits, mass); } -/// Each branch's label logits at its last row after the shared trunk: the trunk's state is kept once -/// and restored a branch at a time; every branch's rows stand on the same trunk state. +/// Each branch's label logits at its last row after the shared trunk; the trunk state kept, restored per branch. fn decide(b: *lb.Metal, window: u32, d: *api.Decision) error{ Failed, Refused }!void { const r = b.runner; if (d.pictures.len > 0) return refuse(d, "decisions with pictures are not served by this engine yet"); From ee2ed5089b3d267c1852c9b011336ed2673cd267 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:34:01 -1000 Subject: [PATCH 17/21] qwen27: a JointHead side-head scores options on post-norm hidden rows at their final-token positions --- zig/src/core/engine_api.zig | 2 + zig/src/native/metal.zig | 3 +- zig/src/native/qwen27_decide.zig | 67 +++++++- zig/src/native/qwen27_joint.zig | 272 +++++++++++++++++++++++++++++++ zig/src/native/qwen27_served.zig | 18 +- zig/src/server/cli.zig | 6 +- 6 files changed, 356 insertions(+), 12 deletions(-) create mode 100644 zig/src/native/qwen27_joint.zig diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 09d37295ec..9f480db1f2 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -286,6 +286,8 @@ pub const Open = struct { policy: []const u8 = "", /// --vision: load the checkpoint's image tower (null: text only; an engine without one refuses). vision: ?VisionOpen = null, + /// --joint-head: a decision side-head's folder; when set, score/decide labels are option-final token positions. + joint_head: ?[]const u8 = null, }; /// A packed decision: a trunk (pictures expand at `at`), a branch a question, and each branch's labels (`counts`). diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index e152da2284..552fb54604 100644 --- a/zig/src/native/metal.zig +++ b/zig/src/native/metal.zig @@ -352,7 +352,7 @@ fn openQwen27(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[ var why: []const u8 = ""; const drafter = if (o.drafts) o.drafter else null; const learn: ?@import("qwen27_disk.zig").Learn = if (o.learn_conversations) if (o.learn) |root| .{ .root = root, .dir = o.dir, .drafter = drafter, .cap = gibBytes(o.learn_gib), .floor = gibBytes(o.learn_min_free_gib), .daily = gibBytes(o.learn_daily_gib) } else null else null; - const h = qwen27.open(gpa, io, o.dir, .{ .context = @intCast(context), .streams = o.lanes, .streams_fixed = o.lanes_fixed, .drafter = drafter, .drafter_bits = o.drafter_bits, .cache_gib = o.prompt_cache_gib, .cache_over_cap = o.prompt_cache_over_cap, .learn = learn, .vision = o.vision }, a, &why) catch |err| { + const h = qwen27.open(gpa, io, o.dir, .{ .context = @intCast(context), .streams = o.lanes, .streams_fixed = o.lanes_fixed, .drafter = drafter, .drafter_bits = o.drafter_bits, .cache_gib = o.prompt_cache_gib, .cache_over_cap = o.prompt_cache_over_cap, .learn = learn, .vision = o.vision, .joint_head = o.joint_head }, a, &why) catch |err| { problem.* = if (why.len > 0) why else try std.fmt.allocPrint(a, "the native Qwen engine (27B class) cannot load {s} ({s})", .{ o.dir, @errorName(err) }); return null; }; @@ -414,6 +414,7 @@ test { _ = glm; _ = glm_dsa; _ = qwen27; + _ = @import("qwen27_joint.zig"); _ = @import("cache_fit.zig"); _ = @import("resident.zig"); _ = @import("metal_identity.zig"); diff --git a/zig/src/native/qwen27_decide.zig b/zig/src/native/qwen27_decide.zig index 8a1ca2a31c..e023b2dc87 100644 --- a/zig/src/native/qwen27_decide.zig +++ b/zig/src/native/qwen27_decide.zig @@ -4,27 +4,36 @@ const api = @import("engine_api"); const tf = @import("tensorfold"); const q = tf.qwen27; const lb = @import("qwen27_lanes.zig"); +const joint = @import("qwen27_joint.zig"); -/// A queued decision's GPU work on the idle host: score the row after its prompt; decide, each branch's last row. -pub fn run(b: *lb.Metal, window: u32, task: *api.DecisionTask) error{ Failed, Refused }!void { +/// A queued decision's GPU work (LaneHost runs it alone): score reads the row after its prompt; decide, each branch's last row. +/// `head` loaded: `labels` are instead option-final token positions, each read of the pass's post-final-norm hidden rows. +pub fn run(b: *lb.Metal, head: ?*joint.Head, window: u32, task: *api.DecisionTask) error{ Failed, Refused }!void { switch (task.*) { - .score => |*s| s.out = try score(b, s.prompt, s.labels, s.logits, s.mass), - .decide => |d| try decide(b, window, d), + .score => |*s| s.out = try score(b, head, s.prompt, s.labels, s.logits, s.mass), + .decide => |d| try decide(b, head, window, d), } } /// `labels`' logits at the row after `prompt`: a prefill's `.last` head leaves them in logits row 0. -fn score(b: *lb.Metal, prompt: []const u32, labels: []const u32, logits: []f64, mass: bool) error{Failed}!f64 { +/// A JointHead instead: `labels` are each option's final-token position, scored on that row's hidden state. +fn score(b: *lb.Metal, head: ?*joint.Head, prompt: []const u32, labels: []const u32, logits: []f64, mass: bool) error{Failed}!f64 { const r = b.runner; if (prompt.len == 0 or prompt.len > r.capacity or logits.len < labels.len) return error.Failed; try reset(b); const session = q.session.Session{ .runner = r }; + if (head) |h| { + try forward(h, b, session, 0, prompt, labels, logits); + return massOf(logits[0..labels.len], mass); + } session.prefill(prompt, b.chunk) catch |e| return failed("a score's prompt pass", e); return read(row0(r), labels, logits, mass); } -/// Each branch's label logits at its last row after the shared trunk; the trunk state kept, restored per branch. -fn decide(b: *lb.Metal, window: u32, d: *api.Decision) error{ Failed, Refused }!void { +/// Each branch's label logits at its last row after the shared trunk: the trunk's state is kept once +/// and restored a branch at a time; every branch's rows stand on the same trunk state. +/// A JointHead instead: `labels` are option-final positions in the whole sequence (trunk's length on a branch's rows). +fn decide(b: *lb.Metal, head: ?*joint.Head, window: u32, d: *api.Decision) error{ Failed, Refused }!void { const r = b.runner; if (d.pictures.len > 0) return refuse(d, "decisions with pictures are not served by this engine yet"); if (d.trunk.len == 0 or d.branches.len == 0 or d.counts.len != d.branches.len) return error.Failed; @@ -45,14 +54,38 @@ fn decide(b: *lb.Metal, window: u32, d: *api.Decision) error{ Failed, Refused }! var at: usize = 0; for (d.branches, d.counts) |br, n| { q.snapshot.restoreSlot(b.gpa, r, 0, null, st) catch |e| return failed("restoring the trunk's state", e); - session.prefill(br, b.chunk) catch |e| return failed("a branch's prompt pass", e); - _ = try read(row0(r), d.labels[at..][0..n], d.logits[at..][0..n], false); + if (head) |h| { + try forward(h, b, session, d.trunk.len, br, d.labels[at..][0..n], d.logits[at..][0..n]); + } else { + session.prefill(br, b.chunk) catch |e| return failed("a branch's prompt pass", e); + _ = try read(row0(r), d.labels[at..][0..n], d.logits[at..][0..n], false); + } d.read += br.len; at += n; } reset(b) catch {}; // slot 0 stands empty for the next stream, whatever its prompt was } +/// `ids` prefilled in chunks (each chunk's rows alone stand in `input`): a `positions` inside the chunk +/// is that option's final-token position, base the rows' offset in the whole sequence (0 for a prompt, +/// the trunk's length on a branch); the JointHead's score lands in `logits`. +fn forward(h: *joint.Head, b: *lb.Metal, session: q.session.Session, base: usize, ids: []const u32, positions: []const u32, logits: []f64) error{Failed}!void { + const r = b.runner; + for (positions) |p| if (p < base or p - base >= ids.len) return error.Failed; + const dim: usize = @intCast(r.model.config.hidden); + var first: usize = 0; + while (first < ids.len) { + const end = @min(ids.len, first + @as(usize, r.promptRows(@intCast(b.chunk)))); + session.promptChunk(ids[first..end], false) catch |e| return failed("a decision's prompt pass", e); + const words = hiddenRows(r); + for (positions, logits) |p, *x| { + if (p - base < first or p - base >= end) continue; + x.* = h.score(words[(p - base - first) * dim ..][0..dim]) catch |e| return failed("a JointHead option", e); + } + first = end; + } +} + /// Slot 0 empty for a pass: a pending keep ends first, and the drafter's reset covers its session too. fn reset(b: *lb.Metal) error{Failed}!void { b.settle() catch {}; @@ -70,6 +103,12 @@ fn row0(r: *const q.decode_round.Runner) []const u16 { return words[0..r.model.config.vocab]; } +/// The last pass's post-final-norm hidden rows: the norm the head's input reads leaves every row there. +fn hiddenRows(r: *const q.decode_round.Runner) [*]const u16 { + const input = r.model.frame.get(.input); + return @ptrCast(@alignCast(input.buffer.contents() + input.offset)); +} + /// The row's `labels` logits out, then the vocabulary logsumexp when `mass` (else 0). fn read(row: []const u16, labels: []const u32, logits: []f64, mass: bool) error{Failed}!f64 { for (labels, logits) |label, *x| { @@ -84,6 +123,16 @@ fn read(row: []const u16, labels: []const u32, logits: []f64, mass: bool) error{ return peak + @log(total); } +/// A JointHead pass has no vocabulary: the option scores' own logsumexp stands as the label mass (it reads 1). +fn massOf(logits: []const f64, mass: bool) f64 { + if (!mass or logits.len == 0) return 0; + var peak: f64 = -std.math.inf(f64); + for (logits) |x| peak = @max(peak, x); + var total: f64 = 0; + for (logits) |x| total += @exp(x - peak); + return peak + @log(total); +} + /// A bf16 logit as f64; a nonfinite one fails the pass, as a draw's argmax check does. fn value(bits: u16) error{Failed}!f64 { const v: f64 = @as(f32, @bitCast(@as(u32, bits) << 16)); diff --git a/zig/src/native/qwen27_joint.zig b/zig/src/native/qwen27_joint.zig new file mode 100644 index 0000000000..64a5fa5455 --- /dev/null +++ b/zig/src/native/qwen27_joint.zig @@ -0,0 +1,272 @@ +//! A JointHead beside the 27B's trunk: post-final-norm hidden rows at option-final token positions in, one score an option out. +const std = @import("std"); +const tf = @import("tensorfold"); +const st = tf.safetensors; +const Allocator = std.mem.Allocator; + +/// `joint_head_config.json`: the head's shape and its weight file's format. +pub const Config = struct { + dim: usize, // the trunk's hidden width + width: usize, // the head's + layers: usize, // Linear→GELU residual blocks +}; + +/// The head: `proj` (dim→width), `layers` of Linear→GELU + residual, `out` (width→1); bf16 words copied out of the file. +pub const Head = struct { + gpa: Allocator, + io: std.Io, + file: st.File, + cfg: Config, + data: []u16, // every weight, one owned row: proj, each block, out + proj_w: []const u16, // [width, dim] + proj_b: []const u16, // [width] + block_w: [][]const u16, // a [width, width] a block + block_b: [][]const u16, // a [width] a block + out_w: []const u16, // [width] + out_b: u16, + x: []f32, // a row after proj and each residual + y: []f32, // a block's Linear→GELU half + h: []f32, // a hidden row as f32 + + /// `path`'s head: `joint_head_config.json` plus its one safetensors (a directory), or the file beside the config. + pub fn open(gpa: Allocator, io: std.Io, path: []const u8) !*Head { + var dir_path = path; + var name: ?[]const u8 = null; + if (std.mem.endsWith(u8, path, ".safetensors")) { + name = std.fs.path.basename(path); + dir_path = std.fs.path.dirname(path) orelse "."; + } + var dir = try std.Io.Dir.cwd().openDir(io, dir_path, .{ .iterate = true }); + defer dir.close(io); + const text = try dir.readFileAlloc(io, "joint_head_config.json", gpa, .limited(1 << 20)); + defer gpa.free(text); + const cfg = try parseConfig(gpa, text); + const file_name = name orelse try oneShard(gpa, io, dir_path); + defer if (name == null) gpa.free(file_name); + const file_path = try std.fs.path.join(gpa, &.{ dir_path, file_name }); + defer gpa.free(file_path); + const h = try gpa.create(Head); + errdefer gpa.destroy(h); + h.* = .{ .gpa = gpa, .io = io, .file = try st.File.open(gpa, io, file_path), .cfg = cfg, .data = &.{}, .proj_w = &.{}, .proj_b = &.{}, .block_w = &.{}, .block_b = &.{}, .out_w = &.{}, .out_b = 0, .x = &.{}, .y = &.{}, .h = &.{} }; + errdefer h.file.close(io); + h.h = try gpa.alloc(f32, cfg.dim); + errdefer gpa.free(h.h); + h.x = try gpa.alloc(f32, cfg.width); + errdefer gpa.free(h.x); + h.y = try gpa.alloc(f32, cfg.width); + errdefer gpa.free(h.y); + h.data = try gpa.alloc(u16, cfg.width * cfg.dim + cfg.width + cfg.layers * cfg.width * (cfg.width + 1) + cfg.width + 1); + errdefer gpa.free(h.data); + h.block_w = try gpa.alloc([]const u16, cfg.layers); + errdefer gpa.free(h.block_w); + h.block_b = try gpa.alloc([]const u16, cfg.layers); + errdefer gpa.free(h.block_b); + var off: usize = 0; + h.proj_w = try h.part(&off, "proj.weight", &.{ cfg.width, cfg.dim }); + h.proj_b = try h.part(&off, "proj.bias", &.{cfg.width}); + for (0..cfg.layers) |i| { + var buf: [64]u8 = undefined; + h.block_w[i] = try h.part(&off, try std.fmt.bufPrint(&buf, "blocks.{d}.0.weight", .{i}), &.{ cfg.width, cfg.width }); + h.block_b[i] = try h.part(&off, try std.fmt.bufPrint(&buf, "blocks.{d}.0.bias", .{i}), &.{cfg.width}); + } + h.out_w = try h.part(&off, "out.weight", &.{ 1, cfg.width }); + h.out_b = (try h.part(&off, "out.bias", &.{1}))[0]; + return h; + } + + pub fn close(h: *Head) void { + h.gpa.free(h.block_b); + h.gpa.free(h.block_w); + h.gpa.free(h.data); + h.gpa.free(h.y); + h.gpa.free(h.x); + h.gpa.free(h.h); + h.file.close(h.io); + h.gpa.destroy(h); + } + + /// One option's score: `hidden` a post-final-norm row of the trunk (`dim` bf16 words). + pub fn score(h: *Head, hidden: []const u16) error{Failed}!f64 { + if (hidden.len != h.cfg.dim) return error.Failed; + for (h.h, hidden) |*d, bits| { + const v: f32 = @bitCast(@as(u32, bits) << 16); + if (!std.math.isFinite(v)) return error.Failed; + d.* = v; + } + matvec(h.proj_w, h.cfg.dim, h.h, h.proj_b, h.x); + for (h.block_w, h.block_b) |w, bias| { + matvec(w, h.cfg.width, h.x, bias, h.y); + for (h.x, h.y) |*a, b| a.* += gelu(b); + } + var s: f64 = f32of(h.out_b); + for (h.x, h.out_w) |a, bits| s += @as(f64, f32of(bits)) * a; + return if (std.math.isFinite(s)) s else error.Failed; + } + + /// `data`'s next `shape` elements taken as the bf16 tensor `name`, checked. + fn part(h: *Head, off: *usize, name: []const u8, shape: []const usize) ![]const u16 { + const t = h.file.get(name) orelse return error.MissingTensor; + if (!t.is(.bf16, shape)) return error.UnexpectedTensor; + const dest = h.data[off.*..][0..t.numel()]; + @memcpy(std.mem.sliceAsBytes(dest), t.bytes); + off.* += t.numel(); + return dest; + } +}; + +fn f32of(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +/// y[j] = bias[j] + row j of `w` ([out, in] bf16) times `x`, f64 accumulation. +fn matvec(w: []const u16, in: usize, x: []const f32, bias: []const u16, y: []f32) void { + for (y, bias, 0..) |*o, b, j| { + var s: f64 = f32of(b); + for (w[j * in ..][0..in], x) |wb, xv| s += @as(f64, f32of(wb)) * xv; + o.* = @floatCast(s); + } +} + +/// nn.GELU: x·Φ(x); Φ's tail by Abramowitz–Stegun 26.2.17 (|ε| ≤ 7.5e-8), 1/√(2π) = 0.39894228040143268. +fn gelu(x: f32) f32 { + const t = 1.0 / (1.0 + 0.2316419 * @abs(x)); + const poly = t * (0.319381530 + t * (-0.356563782 + t * (1.781477937 + t * (-1.821255978 + t * 1.330274429)))); + const tail = poly * @exp(-x * x / 2) * 0.39894228040143268; + return x * (if (x < 0) tail else 1.0 - tail); +} + +/// The config's dim, width and layers; format must be bf16-safetensors-v1. +fn parseConfig(a: Allocator, text: []const u8) !Config { + const parsed = try std.json.parseFromSlice(std.json.Value, a, text, .{}); + defer parsed.deinit(); + if (parsed.value != .object) return error.BadJointHeadConfig; + const o = parsed.value.object; + const format = o.get("format") orelse return error.BadJointHeadConfig; + if (format != .string or !std.mem.eql(u8, format.string, "bf16-safetensors-v1")) return error.BadJointHeadConfig; + return .{ .dim = try field(o, "dim"), .width = try field(o, "width"), .layers = try field(o, "layers") }; +} + +fn field(o: std.json.ObjectMap, name: []const u8) !usize { + const v = o.get(name) orelse return error.BadJointHeadConfig; + if (v != .integer or v.integer <= 0) return error.BadJointHeadConfig; + return @intCast(v.integer); +} + +/// The directory's one safetensors file's name; zero or several are refused. +fn oneShard(gpa: Allocator, io: std.Io, dir_path: []const u8) ![]const u8 { + var dir = try std.Io.Dir.cwd().openDir(io, dir_path, .{ .iterate = true }); + defer dir.close(io); + var found: ?[]const u8 = null; + var it = dir.iterate(); + while (try it.next(io)) |e| { + if (e.kind != .file or !std.mem.endsWith(u8, e.name, ".safetensors")) continue; + if (found != null) return error.JointHeadShardCount; + found = try gpa.dupe(u8, e.name); + } + return found orelse error.JointHeadShardCount; +} + +fn bf16(x: f32) u16 { + const bits: u32 = @bitCast(x); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +/// A head fixture in `tmp` (dim 4, width 4, 2 blocks; every weight exactly bf16): its directory's path. +fn fixture(gpa: Allocator, tmp: *const std.testing.TmpDir) ![]u8 { + const header = + \\{"proj.weight":{"dtype":"BF16","shape":[4,4],"data_offsets":[0,32]}, + \\ "proj.bias":{"dtype":"BF16","shape":[4],"data_offsets":[32,40]}, + \\ "blocks.0.0.weight":{"dtype":"BF16","shape":[4,4],"data_offsets":[40,72]}, + \\ "blocks.0.0.bias":{"dtype":"BF16","shape":[4],"data_offsets":[72,80]}, + \\ "blocks.1.0.weight":{"dtype":"BF16","shape":[4,4],"data_offsets":[80,112]}, + \\ "blocks.1.0.bias":{"dtype":"BF16","shape":[4],"data_offsets":[112,120]}, + \\ "out.weight":{"dtype":"BF16","shape":[1,4],"data_offsets":[120,128]}, + \\ "out.bias":{"dtype":"BF16","shape":[1],"data_offsets":[128,130]}} + ; + const weights = [_][]const f32{ + &.{ 0.5, -0.25, 0.125, 0 }, &.{ 0, 0.5, -0.25, 0.125 }, &.{ 0.125, 0, 0.5, -0.25 }, &.{ -0.25, 0.125, 0, 0.5 }, // proj + &.{ 0.125, -0.25, 0.375, -0.5 }, // proj bias + &.{ 0.25, 0.125, 0, 0 }, &.{ 0, 0.25, 0.125, 0 }, &.{ 0, 0, 0.25, 0.125 }, &.{ 0.125, 0, 0, 0.25 }, // block 0 + &.{ 0.125, -0.125, 0.125, -0.125 }, // block 0 bias + &.{ 0.5, 0.25, 0, 0 }, &.{ 0, 0.5, 0.25, 0 }, &.{ 0, 0, 0.5, 0.25 }, &.{ 0.25, 0, 0, 0.5 }, // block 1 + &.{ -0.25, 0.25, -0.25, 0.25 }, // block 1 bias + &.{ 0.5, -0.5, 0.25, -0.25 }, // out + &.{0.125}, // out bias + }; + var data: [8 + header.len + 130]u8 = undefined; + std.mem.writeInt(u64, data[0..8], header.len, .little); + @memcpy(data[8..][0..header.len], header); + var at: usize = 8 + header.len; + for (weights) |w| for (w) |x| { + std.mem.writeInt(u16, data[at..][0..2], bf16(x), .little); + at += 2; + }; + try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "head.safetensors", .data = &data }); + try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "joint_head_config.json", .data = "{\"dim\":4,\"width\":4,\"layers\":2,\"format\":\"bf16-safetensors-v1\"}" }); + const root = try tmp.dir.realPathFileAlloc(std.testing.io, ".", gpa); + defer gpa.free(root); + return std.fs.path.join(gpa, &.{root}); +} + +fn row(gpa: Allocator, xs: []const f32) ![]u16 { + const out = try gpa.alloc(u16, xs.len); + for (xs, out) |x, *b| b.* = bf16(x); + return out; +} + +test "a JointHead loads its config and tensors and scores a hidden row" { + const a = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try fixture(a, &tmp); + defer a.free(dir); + const h = try Head.open(a, std.testing.io, dir); + defer h.close(); + try std.testing.expectEqual(@as(usize, 4), h.cfg.dim); + try std.testing.expectEqual(@as(usize, 4), h.cfg.width); + try std.testing.expectEqual(@as(usize, 2), h.cfg.layers); + const hidden = try row(a, &.{ 0.5, -0.25, 0.125, -0.5 }); + defer a.free(hidden); + try std.testing.expectApproxEqAbs(1.035910918342501, try h.score(hidden), 1e-4); + const second = try row(a, &.{ 0.25, 0.5, -0.125, 0.75 }); + defer a.free(second); + try std.testing.expectApproxEqAbs(0.09842275576323394, try h.score(second), 1e-4); + const third = try row(a, &.{ -1.0, 0.0, 1.0, 0.5 }); + defer a.free(third); + try std.testing.expectApproxEqAbs(0.3339105568875796, try h.score(third), 1e-4); + const file = try std.fs.path.join(a, &.{ dir, "head.safetensors" }); + defer a.free(file); + const other = try Head.open(a, std.testing.io, file); + defer other.close(); + try std.testing.expectApproxEqAbs(1.035910918342501, try other.score(hidden), 1e-4); // the file path itself also opens +} + +test "a JointHead refuses a wrong-width row and a nonfinite one" { + const a = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try fixture(a, &tmp); + defer a.free(dir); + const h = try Head.open(a, std.testing.io, dir); + defer h.close(); + const hidden = try row(a, &.{ 0.5, -0.25, 0.125 }); + defer a.free(hidden); + try std.testing.expectError(error.Failed, h.score(hidden)); + const bad = try row(a, &.{ 0.5, -0.25, 0.125, -0.5 }); + defer a.free(bad); + bad[2] = 0x7FC0; // a NaN hidden word ends the pass + try std.testing.expectError(error.Failed, h.score(bad)); +} + +test "a JointHead refuses a wrong-format config and a missing tensor" { + const a = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try fixture(a, &tmp); + defer a.free(dir); + try tmp.dir.writeFile(io, .{ .sub_path = "joint_head_config.json", .data = "{\"dim\":4,\"width\":4,\"layers\":2,\"format\":\"fp32\"}" }); + try std.testing.expectError(error.BadJointHeadConfig, Head.open(a, io, dir)); + try tmp.dir.writeFile(io, .{ .sub_path = "joint_head_config.json", .data = "{\"dim\":8,\"width\":4,\"layers\":2,\"format\":\"bf16-safetensors-v1\"}" }); + try std.testing.expectError(error.UnexpectedTensor, Head.open(a, io, dir)); // proj.weight is [4, 4], not [4, 8] +} diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index b47f59dcb7..f8d90cc7c2 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -13,6 +13,7 @@ const Resident = @import("resident.zig").Resident; const lb = @import("qwen27_lanes.zig"); pub const lanes_backend = lb; const decide = @import("qwen27_decide.zig"); +const joint = @import("qwen27_joint.zig"); const disk = @import("qwen27_disk.zig"); const images = @import("qwen27_images.zig"); const Allocator = std.mem.Allocator; @@ -28,6 +29,7 @@ pub const Options = struct { learn: ?disk.Learn = null, // --learn-conversations: kept states written behind to disk and read back vision: ?api.VisionOpen = null, // --vision: open loads the image tower, and requests may hold images images: ?*images.Vision = null, // the loaded tower attach serves (and the host then owns) + joint_head: ?[]const u8 = null, // --joint-head: a decision side-head; score/decide read its option-final rows }; pub const Host = struct { @@ -46,6 +48,7 @@ pub const Host = struct { warm: mtl.keepalive.Target, disk_stage: ?mtl.Buffer = null, // copies out of private kept states for the disk (qwen27_disk.zig) vision: ?*images.Vision = null, // --vision's tower and image rows + joint: ?*joint.Head = null, // --joint-head's loaded side-head pub fn engine(h: *Host) api.Engine { return h.host.engine(); @@ -87,7 +90,7 @@ pub const Host = struct { /// LaneHost's decision hook: Ornith's label reads over a shared trunk, on an idle engine. fn decideRun(ptr: *anyopaque, task: *api.DecisionTask) error{ Failed, Refused }!void { const h = of(ptr); - return decide.run(h.back, h.host.info_.context_window, task); + return decide.run(h.back, h.joint, h.host.info_.context_window, task); } }; @@ -316,6 +319,7 @@ pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.df return error.DrafterLoad; }; errdefer if (h.draft) |d| d.deinit(); + h.joint = null; h.back = try lb.Metal.init(gpa, &h.runner, h.draft); errdefer h.back.deinit(); h.back.chunk = h.runner.promptRows(q.model.promptRows(h.runner.model.config)); @@ -346,6 +350,17 @@ pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.df if (o.learn) |l| try disk.attach(h, io, a, l); errdefer if (h.cache) |*store| if (store.disk) |r| r.deinit(); h.host.keepalive_target = .{ .ctx = &h.warm, .tick = mtl.keepalive.Target.tick }; + if (o.joint_head) |p| { + h.joint = joint.Head.open(gpa, io, p) catch |err| { + why.* = try std.fmt.allocPrint(a, "the native Qwen3.8-27B joint head cannot load {s} ({s})", .{ p, @errorName(err) }); + return error.JointHeadLoad; + }; + if (h.joint.?.cfg.dim != model.config.hidden) { + why.* = try std.fmt.allocPrint(a, "--joint-head's dim {d} does not match this checkpoint's hidden width {d}", .{ h.joint.?.cfg.dim, model.config.hidden }); + return error.JointHeadDim; + } + } + errdefer if (h.joint) |j| j.close(); try h.host.start(); std.log.info("Qwen3.8-27B: {d} streams of {d} tokens{s}{s}", .{ slots, capacity, if (h.draft != null) ", DFlash2 drafts for one of them" else "", if (o.images != null) ", images read" else "" }); return h; @@ -401,6 +416,7 @@ pub fn close(ptr: *anyopaque) void { h.core.deinit(); h.cfg.deinit(h.gpa); h.back.deinit(); + if (h.joint) |j| j.close(); if (h.draft) |d| d.deinit(); h.runner.deinit(); if (h.vision) |v| v.deinit(); diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 8afebaebe9..a0c34697f2 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -125,6 +125,8 @@ pub const flags = [_]Flag{ .{ .name = "--recall-dir", .native = true }, .{ .name = "--recall-gib", .native = true }, .{ .name = "--recall-read-only", .kind = .store_true, .native = true }, + // the native server's own: a decision side-head (joint_head_config.json + bf16 safetensors) scoring option-final rows + .{ .name = "--joint-head", .native = true }, // the native server's own: Cortheon's reasoning program and server-run web research for every chat request .{ .name = "--cortheon", .kind = .store_true, .native = true }, }; @@ -178,6 +180,7 @@ pub const Args = struct { compact_memory: ?[]const u8 = null, drafter: ?[]const u8 = null, drafter_bits: u8 = 4, + joint_head: ?[]const u8 = null, // --joint-head: a decision side-head; score/decide labels become option-final token positions record: ?[]const u8 = null, // drafter training rows written under this folder kv_cache: []const u8 = "bf16", // fp8: e4m3 attention caches with a scale a row, about twice the context slide: bool = false, // Sliding Weights learns what /v1/slide/learn is told into the served weights, live @@ -319,7 +322,7 @@ fn apply(a: Allocator, out: *Args, name: []const u8, value: ?[]const u8, u: *Usa out.learn_daily_gib = try gib(u, a, name, v); } else if (is(name, "--prompt-seed")) { out.prompt_seed = true; - } else if (is(name, "--max-tokens")) out.max_tokens = try int(u, a, name, v) else if (is(name, "--temperature")) out.temperature = try float(u, a, name, v) else if (is(name, "--top-p")) out.top_p = try float(u, a, name, v) else if (is(name, "--top-k")) out.top_k = try int(u, a, name, v) else if (is(name, "--min-p")) out.min_p = try float(u, a, name, v) else if (is(name, "--thinking")) out.thinking = true else if (is(name, "--no-thinking")) out.thinking = false else if (is(name, "--reasoning-effort")) out.reasoning_effort = v else if (is(name, "--thinking-budget")) out.thinking_budget = try int(u, a, name, v) else if (is(name, "--loop-guard")) out.loop_guard = true else if (is(name, "--drafter")) out.drafter = v else if (is(name, "--drafter-bits")) out.drafter_bits = @intCast(try int(u, a, name, v)) else if (is(name, "--record")) out.record = v else if (is(name, "--kv-cache")) out.kv_cache = v else if (is(name, "--no-drafts")) out.no_drafts = true else if (is(name, "--keep-warm")) out.keep_warm = try int(u, a, name, v) else if (is(name, "--compact-at")) { + } else if (is(name, "--max-tokens")) out.max_tokens = try int(u, a, name, v) else if (is(name, "--temperature")) out.temperature = try float(u, a, name, v) else if (is(name, "--top-p")) out.top_p = try float(u, a, name, v) else if (is(name, "--top-k")) out.top_k = try int(u, a, name, v) else if (is(name, "--min-p")) out.min_p = try float(u, a, name, v) else if (is(name, "--thinking")) out.thinking = true else if (is(name, "--no-thinking")) out.thinking = false else if (is(name, "--reasoning-effort")) out.reasoning_effort = v else if (is(name, "--thinking-budget")) out.thinking_budget = try int(u, a, name, v) else if (is(name, "--loop-guard")) out.loop_guard = true else if (is(name, "--drafter")) out.drafter = v else if (is(name, "--drafter-bits")) out.drafter_bits = @intCast(try int(u, a, name, v)) else if (is(name, "--joint-head")) out.joint_head = v else if (is(name, "--record")) out.record = v else if (is(name, "--kv-cache")) out.kv_cache = v else if (is(name, "--no-drafts")) out.no_drafts = true else if (is(name, "--keep-warm")) out.keep_warm = try int(u, a, name, v) else if (is(name, "--compact-at")) { if (std.mem.eql(u8, v, "auto")) out.compact_auto = true else { const f = try float(u, a, name, v); if (!(f > 0 and f <= 1)) return fail(u, a, "argument --compact-at: expected auto or a fraction in (0, 1]: '{s}'", .{v}); @@ -438,6 +441,7 @@ pub fn request(a: Allocator, dir: []const u8, model_type: []const u8, args: Args .drafts = !args.no_drafts, .drafter = args.drafter, .drafter_bits = args.drafter_bits, + .joint_head = args.joint_head, .record = args.record, .kv8 = std.mem.eql(u8, args.kv_cache, "fp8"), .speed_up = args.speed_up, From 9a2f2eff75220f0cfb006ce8d7eec5c3552d444c Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:56:42 -1000 Subject: [PATCH 18/21] server: a side head's options are read at their lines' final-token positions, derived from the render --- zig/src/core/engine_api.zig | 2 ++ zig/src/native/qwen27_served.zig | 1 + zig/src/server/decisions.zig | 49 +++++++++++++++++++++++++++----- zig/src/server/model_text.zig | 20 +++++++++++++ zig/src/server/systemone.zig | 49 +++++++++++++++++++++++++------- 5 files changed, 103 insertions(+), 18 deletions(-) diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 9f480db1f2..e7e1f6d8b0 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -200,6 +200,8 @@ pub const Info = struct { prompt_cache_plan: ?PromptCachePlan = null, /// The engine scores decisions (`score`) and writes no text: the server refuses text requests. decides: bool = false, + /// `score`/`decide` labels are option-final token positions, not token ids: a side head reads there. + decision_positions: bool = false, /// Decisions may carry pictures (`Decision.pictures`). pictures: bool = false, /// Image input (--vision); null: text only. diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index f8d90cc7c2..d26657afc3 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -361,6 +361,7 @@ pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.df } } errdefer if (h.joint) |j| j.close(); + if (h.joint != null) h.host.info_.decision_positions = true; // score/decide labels are option-final rows try h.host.start(); std.log.info("Qwen3.8-27B: {d} streams of {d} tokens{s}{s}", .{ slots, capacity, if (h.draft != null) ", DFlash2 drafts for one of them" else "", if (o.images != null) ", images read" else "" }); return h; diff --git a/zig/src/server/decisions.zig b/zig/src/server/decisions.zig index 215de8416a..ad48f84ff8 100644 --- a/zig/src/server/decisions.zig +++ b/zig/src/server/decisions.zig @@ -157,10 +157,11 @@ const Wording = struct { kind: Kind, names: []const []const u8, labels: []const []const u8, + lines: []const []const u8, // every rendered line, the option lines between the question and the closing content: []const u8, }; -fn wording(cx: *Cx, input: []const u8, question: Value) errors.Refused!Wording { +fn wording(cx: *Cx, input: []const u8, question: Value, positions: bool) errors.Refused!Wording { const kind_value = question.get("type") orelse .null; const kind: Kind = if (kind_value == .string and std.mem.eql(u8, kind_value.string, "choice")) .choice @@ -189,7 +190,10 @@ fn wording(cx: *Cx, input: []const u8, question: Value) errors.Refused!Wording { try rejectUnknown(cx, question.object, &score_fields); const listed = try scoreList(cx, question); try lines.append(cx.a, try questionLine(cx, question)); - for (listed.labels, listed.details) |label, detail| try lines.append(cx.a, try std.fmt.allocPrint(cx.a, "{s}: {s}", .{ label, detail })); + for (listed.labels, listed.details) |label, detail| { + const line = try std.fmt.allocPrint(cx.a, "{s}: {s}", .{ label, detail }); + try lines.append(cx.a, line); + } break :blk .{ listed.names, listed.labels }; }, .yes_no => blk: { @@ -199,8 +203,15 @@ fn wording(cx: *Cx, input: []const u8, question: Value) errors.Refused!Wording { const lead = try renderText(cx, question.get("question")); if (reply_text.pyStrip(lead).len == 0) return cx.refuse("a question must not be blank"); try lines.append(cx.a, try std.fmt.allocPrint(cx.a, "Is the following true? {s}", .{lead})); - if (yes.len > 0) try lines.append(cx.a, try std.fmt.allocPrint(cx.a, "yes: {s}", .{yes})); - if (no.len > 0) try lines.append(cx.a, try std.fmt.allocPrint(cx.a, "no: {s}", .{no})); + // a side head needs each option rendered: the bare answer words stand in when no text is given + if (yes.len > 0 or positions) { + const line = try std.fmt.allocPrint(cx.a, "yes: {s}", .{if (yes.len > 0) yes else "yes"}); + try lines.append(cx.a, line); + } + if (no.len > 0 or positions) { + const line = try std.fmt.allocPrint(cx.a, "no: {s}", .{if (no.len > 0) no else "no"}); + try lines.append(cx.a, line); + } const names = [_][]const u8{ "yes", "no" }; break :blk .{ try cx.a.dupe([]const u8, &names), try cx.a.dupe([]const u8, &names) }; }, @@ -211,7 +222,7 @@ fn wording(cx: *Cx, input: []const u8, question: Value) errors.Refused!Wording { .yes_no => "Answer with yes or no only.", }; try lines.append(cx.a, closing); - return .{ .kind = kind, .names = names, .labels = labels, .content = try contentOf(cx.a, input, lines.items) }; + return .{ .kind = kind, .names = names, .labels = labels, .lines = lines.items, .content = try contentOf(cx.a, input, lines.items) }; } const Listed = struct { @@ -291,7 +302,8 @@ fn contentOf(a: Allocator, input: []const u8, lines: []const []const u8) Allocat } fn prepareOne(srv: *Server, cx: *Cx, input: []const u8, question: Value) errors.Refused!Prepared { - const word = try wording(cx, input, question); + const positions = srv.engine.info().decision_positions; + const word = try wording(cx, input, question, positions); const prompt_text = try renderPrompt(srv, cx, word.content); const prompt_ids = try encodePrompt(srv, cx, prompt_text); try thinkOpen(cx, prompt_text); @@ -301,7 +313,10 @@ fn prepareOne(srv: *Server, cx: *Cx, input: []const u8, question: Value) errors. .kind = word.kind, .names = word.names, .prompt_ids = prompt_ids, - .label_ids = try labelIds(srv, cx, prompt_text, prompt_ids, word.labels), + .label_ids = if (positions) + try optionPositions(srv, cx, prompt_text, prompt_ids, input.len, word.lines, word.content) + else + try labelIds(srv, cx, prompt_text, prompt_ids, word.labels), }; } @@ -354,6 +369,26 @@ fn labelRefuse(cx: *Cx, label: []const u8) errors.Refused { return cx.fail(.request, "the answer label {s} is not one distinct token after the chat prompt for this tokenizer, so this model is not supported", .{pyrepr.repr(cx.a, .{ .string = label }) catch label}); } +/// A side head's `labels`: each option's final-token position in `prompt`, its line's last token; +/// the lines sit one a row between `input` and the closing inside `content` (`base` it in `prompt`). +fn optionPositions(srv: *Server, cx: *Cx, prompt: []const u8, prompt_ids: []const u32, input_len: usize, lines: []const []const u8, content: []const u8) errors.Refused![]u32 { + const a = cx.a; + const base = std.mem.indexOf(u8, prompt, content) orelse return positionRefuse(cx, lines[0]); + const ats = try a.alloc(usize, lines.len - 2); + var at = base + input_len + 2 + lines[0].len + 1; // past the input, the question line and their rows + for (lines[1 .. lines.len - 1], ats) |line, *end| { + end.* = at + line.len; + at = end.* + 1; + } + const out = try a.alloc(u32, ats.len); + if (!(srv.text.indexAt(a, prompt, prompt_ids, ats, out) catch return error.OutOfMemory)) return positionRefuse(cx, lines[0]); + return out; +} + +fn positionRefuse(cx: *Cx, line: []const u8) errors.Refused { + return cx.fail(.request, "the option line {s} is not found at a token boundary in the chat prompt for this tokenizer, so this model's decision head cannot read it", .{pyrepr.repr(cx.a, .{ .string = line }) catch line}); +} + fn scoreAll(srv: *Server, cx: *Cx, prepared: []const Prepared) errors.Refused![]Scored { const scored = try cx.a.alloc(Scored, prepared.len); for (prepared, scored) |item, *row| { diff --git a/zig/src/server/model_text.zig b/zig/src/server/model_text.zig index fd04340b14..f8da2c8f42 100644 --- a/zig/src/server/model_text.zig +++ b/zig/src/server/model_text.zig @@ -73,6 +73,26 @@ pub const Text = struct { return f(t.ctx, a, id); } + /// `ids` are `text`'s encoding: for each of `ats` (a span's end offset, ascending) the index of + /// the token covering its last byte, or false when a token's own text and `text` disagree. + pub fn indexAt(t: Text, a: Allocator, text: []const u8, ids: []const u32, ats: []const usize, out: []u32) Allocator.Error!bool { + var i: usize = 0; + var end: usize = 0; + for (ats, out) |at, *o| { + if (at == 0 or at > text.len) return false; + while (end < at) { + if (i == ids.len) return false; + const piece = try t.decode(a, ids[i .. i + 1]); + defer a.free(piece); + if (piece.len == 0 or !std.mem.startsWith(u8, text[end..], piece)) return false; + end += piece.len; + i += 1; + } + o.* = @intCast(i - 1); + } + return true; + } + pub fn eosIds(t: Text) []const u32 { return t.vtable.eos_ids(t.ctx); } diff --git a/zig/src/server/systemone.zig b/zig/src/server/systemone.zig index 57b03f33d5..22e4c0073a 100644 --- a/zig/src/server/systemone.zig +++ b/zig/src/server/systemone.zig @@ -76,12 +76,15 @@ pub fn reply(srv: *Server, cx: *Cx, body: Value) errors.Refused!Value { const trunk = if (seen.urls.len > 0) try std.mem.concat(cx.a, u32, &.{ head, try f.encode(try prefix(cx.a, seen.state, false)) }) else try f.encode(try prefix(cx.a, seen.state, true)); const codes = try cx.a.alloc([]const Code, questions.len); const branches = try cx.a.alloc([]const u32, questions.len); - for (questions, codes, branches) |q, *c, *b| { + const groups = try cx.a.alloc([]const []const u32, questions.len); + const positions = srv.engine.info().decision_positions; // a side head reads option lines' last rows + for (questions, codes, branches, groups) |q, *c, *b, *g| { c.* = if (q.kind == .choice) try f.codes(q.keys) else &.{}; - b.* = try f.encode(try suffix(cx.a, q, c.*)); + var ends: std.ArrayList(usize) = .empty; + const rendered = try suffix(cx.a, q, c.*, positions, &ends); + b.* = try f.encode(rendered); + g.* = if (positions) try f.positions(rendered, b.*, ends.items, trunk.len) else try f.readout(q, c.*); } - const groups = try cx.a.alloc([]const []const u32, questions.len); - for (questions, codes, groups) |q, c, *g| g.* = try f.readout(q, c); const window = srv.engine.info().context_window; var read: usize = trunk.len; for (branches) |b| { @@ -270,6 +273,21 @@ const Format = struct { } } + /// A side head's readout: one group an option, the absolute index of its line's last token (`base` the trunk's length). + fn positions(f: *Format, text_: []const u8, ids: []const u32, ends: []const usize, base: usize) errors.Refused![]const []const u32 { + const a = f.cx.a; + const ats = try a.alloc(u32, ends.len); + if (!(f.srv.text.indexAt(a, text_, ids, ends, ats) catch return error.OutOfMemory)) + return f.cx.refuse("an option's line is not found at a token boundary in this decision's text, so this model's decision head cannot read it"); + const groups = try a.alloc([]const u32, ends.len); + for (ats, groups) |p, *g| { + const one = try a.alloc(u32, 1); + one[0] = @intCast(base + p); + g.* = one; + } + return groups; + } + fn forms(f: *Format, texts: []const []const u8) errors.Refused![]const u32 { var seen: std.ArrayList(u32) = .empty; for (texts) |t| if (try f.single(t)) |id| if (std.mem.indexOfScalar(u32, seen.items, id) == null) try seen.append(f.cx.a, id); @@ -334,14 +352,16 @@ fn indented(w: *std.Io.Writer, v: Value, level: usize) std.Io.Writer.Error!void } } -fn suffix(a: Allocator, q: Question, codes: []const Code) Allocator.Error![]const u8 { +fn suffix(a: Allocator, q: Question, codes: []const Code, positions: bool, ends: *std.ArrayList(usize)) Allocator.Error![]const u8 { var out: std.Io.Writer.Allocating = .init(a); - questionText(a, &out.writer, q, codes) catch return error.OutOfMemory; + questionText(a, &out, q, codes, positions, ends) catch return error.OutOfMemory; out.writer.writeAll(im_end ++ "\n" ++ im_start ++ "assistant\n") catch return error.OutOfMemory; return out.toOwnedSlice(); } -fn questionText(a: Allocator, w: *std.Io.Writer, q: Question, codes: []const Code) !void { +/// The question's text; `ends` takes each option line's end offset for a side head's positions. +fn questionText(a: Allocator, out: *std.Io.Writer.Allocating, q: Question, codes: []const Code, positions: bool, ends: *std.ArrayList(usize)) !void { + const w = &out.writer; try w.writeAll(q.instructions); switch (q.kind) { .choice => { @@ -351,21 +371,28 @@ fn questionText(a: Allocator, w: *std.Io.Writer, q: Question, codes: []const Cod try w.print("{s} ", .{c.text}); // a description as written (JSON when it is not text); without one, the name if (desc != null and desc.?.truthy()) try w.writeAll(try text(a, desc.?)) else for (key) |ch| try w.writeByte(if (ch == '_') ' ' else ch); + try ends.append(a, out.written().len); } try w.writeAll("\n\nReply with the option code only."); }, .noul => { - // the template shows the descriptions only when one side has one, then both + // the template shows the descriptions only when one side has one, then both; + // a side head needs each option rendered: the bare answer words stand in when no text is given const any = q.texts.len == 2 and (q.texts[0] != null and q.texts[0].? != .null or q.texts[1] != null and q.texts[1].? != .null); - if (any) for ([_][]const u8{ "Yes", "No" }, q.texts) |name, t| { - const shown: Value = if (t != null and t.? != .null) t.? else .{ .string = "None" }; + if (any or positions) for ([_][]const u8{ "Yes", "No" }, 0..) |name, i| { + const t: ?Value = if (i < q.texts.len) q.texts[i] else null; + const shown: Value = if (t != null and t.? != .null) t.? else .{ .string = if (any) "None" else if (i == 0) "yes" else "no" }; try w.print("\n{s}: {s}", .{ name, try text(a, shown) }); + try ends.append(a, out.written().len); }; try w.writeAll("\n\nReply with yes or no only."); }, .score => { try w.writeAll("\n\n"); - for (q.keys, 0..) |name, i| try w.print("{s}{d} {s}", .{ if (i > 0) "\n" else "", i, name }); + for (q.keys, 0..) |name, i| { + try w.print("{s}{d} {s}", .{ if (i > 0) "\n" else "", i, name }); + try ends.append(a, out.written().len); + } try w.print("\n\nReply with a single digit 0-{d} only.", .{q.keys.len - 1}); }, } From 55b51c0ba7dd8e4fab1195669cb070225e2f11b9 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Wed, 21 Oct 2026 16:51:16 -1000 Subject: [PATCH 19/21] qwen27: one-line comments for the lean check --- zig/src/native/qwen27_decide.zig | 14 ++++---------- zig/src/native/qwen27_joint.zig | 4 ++-- zig/src/server/cli.zig | 2 +- zig/src/server/decisions.zig | 3 +-- zig/src/server/model_text.zig | 3 +-- zig/src/server/systemone.zig | 5 ++--- 6 files changed, 11 insertions(+), 20 deletions(-) diff --git a/zig/src/native/qwen27_decide.zig b/zig/src/native/qwen27_decide.zig index e023b2dc87..7f3cafcdc3 100644 --- a/zig/src/native/qwen27_decide.zig +++ b/zig/src/native/qwen27_decide.zig @@ -6,8 +6,7 @@ const q = tf.qwen27; const lb = @import("qwen27_lanes.zig"); const joint = @import("qwen27_joint.zig"); -/// A queued decision's GPU work (LaneHost runs it alone): score reads the row after its prompt; decide, each branch's last row. -/// `head` loaded: `labels` are instead option-final token positions, each read of the pass's post-final-norm hidden rows. +/// A queued decision's GPU work on the idle host: score the row after its prompt; decide, each branch's last row. pub fn run(b: *lb.Metal, head: ?*joint.Head, window: u32, task: *api.DecisionTask) error{ Failed, Refused }!void { switch (task.*) { .score => |*s| s.out = try score(b, head, s.prompt, s.labels, s.logits, s.mass), @@ -15,8 +14,7 @@ pub fn run(b: *lb.Metal, head: ?*joint.Head, window: u32, task: *api.DecisionTas } } -/// `labels`' logits at the row after `prompt`: a prefill's `.last` head leaves them in logits row 0. -/// A JointHead instead: `labels` are each option's final-token position, scored on that row's hidden state. +/// `labels`' logits at the row after `prompt` (`.last` head, logits row 0); a JointHead scores option-final rows. fn score(b: *lb.Metal, head: ?*joint.Head, prompt: []const u32, labels: []const u32, logits: []f64, mass: bool) error{Failed}!f64 { const r = b.runner; if (prompt.len == 0 or prompt.len > r.capacity or logits.len < labels.len) return error.Failed; @@ -30,9 +28,7 @@ fn score(b: *lb.Metal, head: ?*joint.Head, prompt: []const u32, labels: []const return read(row0(r), labels, logits, mass); } -/// Each branch's label logits at its last row after the shared trunk: the trunk's state is kept once -/// and restored a branch at a time; every branch's rows stand on the same trunk state. -/// A JointHead instead: `labels` are option-final positions in the whole sequence (trunk's length on a branch's rows). +/// Each branch's label logits at its last row after the shared trunk; a `head` reads option-final positions. fn decide(b: *lb.Metal, head: ?*joint.Head, window: u32, d: *api.Decision) error{ Failed, Refused }!void { const r = b.runner; if (d.pictures.len > 0) return refuse(d, "decisions with pictures are not served by this engine yet"); @@ -66,9 +62,7 @@ fn decide(b: *lb.Metal, head: ?*joint.Head, window: u32, d: *api.Decision) error reset(b) catch {}; // slot 0 stands empty for the next stream, whatever its prompt was } -/// `ids` prefilled in chunks (each chunk's rows alone stand in `input`): a `positions` inside the chunk -/// is that option's final-token position, base the rows' offset in the whole sequence (0 for a prompt, -/// the trunk's length on a branch); the JointHead's score lands in `logits`. +/// `ids` prefilled in chunks: each `positions` entry is an option's final-token position offset by `base`. fn forward(h: *joint.Head, b: *lb.Metal, session: q.session.Session, base: usize, ids: []const u32, positions: []const u32, logits: []f64) error{Failed}!void { const r = b.runner; for (positions) |p| if (p < base or p - base >= ids.len) return error.Failed; diff --git a/zig/src/native/qwen27_joint.zig b/zig/src/native/qwen27_joint.zig index 64a5fa5455..5314ed2732 100644 --- a/zig/src/native/qwen27_joint.zig +++ b/zig/src/native/qwen27_joint.zig @@ -1,4 +1,4 @@ -//! A JointHead beside the 27B's trunk: post-final-norm hidden rows at option-final token positions in, one score an option out. +//! A JointHead beside the 27B's trunk: option-final hidden rows in, one score an option out. const std = @import("std"); const tf = @import("tensorfold"); const st = tf.safetensors; @@ -11,7 +11,7 @@ pub const Config = struct { layers: usize, // Linear→GELU residual blocks }; -/// The head: `proj` (dim→width), `layers` of Linear→GELU + residual, `out` (width→1); bf16 words copied out of the file. +/// The head: `proj` (dim→width), Linear→GELU residual layers, `out` (width→1); bf16 words read from the file. pub const Head = struct { gpa: Allocator, io: std.Io, diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index a0c34697f2..86649ce179 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -125,7 +125,7 @@ pub const flags = [_]Flag{ .{ .name = "--recall-dir", .native = true }, .{ .name = "--recall-gib", .native = true }, .{ .name = "--recall-read-only", .kind = .store_true, .native = true }, - // the native server's own: a decision side-head (joint_head_config.json + bf16 safetensors) scoring option-final rows + // the native server's own: a side head (joint_head_config.json + bf16 safetensors) on option-final rows .{ .name = "--joint-head", .native = true }, // the native server's own: Cortheon's reasoning program and server-run web research for every chat request .{ .name = "--cortheon", .kind = .store_true, .native = true }, diff --git a/zig/src/server/decisions.zig b/zig/src/server/decisions.zig index ad48f84ff8..a431055e99 100644 --- a/zig/src/server/decisions.zig +++ b/zig/src/server/decisions.zig @@ -369,8 +369,7 @@ fn labelRefuse(cx: *Cx, label: []const u8) errors.Refused { return cx.fail(.request, "the answer label {s} is not one distinct token after the chat prompt for this tokenizer, so this model is not supported", .{pyrepr.repr(cx.a, .{ .string = label }) catch label}); } -/// A side head's `labels`: each option's final-token position in `prompt`, its line's last token; -/// the lines sit one a row between `input` and the closing inside `content` (`base` it in `prompt`). +/// A side head's `labels`: each option's final-token position in `prompt` (its line's last token, `base`d). fn optionPositions(srv: *Server, cx: *Cx, prompt: []const u8, prompt_ids: []const u32, input_len: usize, lines: []const []const u8, content: []const u8) errors.Refused![]u32 { const a = cx.a; const base = std.mem.indexOf(u8, prompt, content) orelse return positionRefuse(cx, lines[0]); diff --git a/zig/src/server/model_text.zig b/zig/src/server/model_text.zig index f8da2c8f42..76c2954d1d 100644 --- a/zig/src/server/model_text.zig +++ b/zig/src/server/model_text.zig @@ -73,8 +73,7 @@ pub const Text = struct { return f(t.ctx, a, id); } - /// `ids` are `text`'s encoding: for each of `ats` (a span's end offset, ascending) the index of - /// the token covering its last byte, or false when a token's own text and `text` disagree. + /// `ids` are `text`'s encoding: for each of `ats` (ascending ends) the covering token's index, or false. pub fn indexAt(t: Text, a: Allocator, text: []const u8, ids: []const u32, ats: []const usize, out: []u32) Allocator.Error!bool { var i: usize = 0; var end: usize = 0; diff --git a/zig/src/server/systemone.zig b/zig/src/server/systemone.zig index 22e4c0073a..175cd94a80 100644 --- a/zig/src/server/systemone.zig +++ b/zig/src/server/systemone.zig @@ -273,7 +273,7 @@ const Format = struct { } } - /// A side head's readout: one group an option, the absolute index of its line's last token (`base` the trunk's length). + /// A side head's readout: one group an option, its line's last token's index (`base` the trunk's length). fn positions(f: *Format, text_: []const u8, ids: []const u32, ends: []const usize, base: usize) errors.Refused![]const []const u32 { const a = f.cx.a; const ats = try a.alloc(u32, ends.len); @@ -376,8 +376,7 @@ fn questionText(a: Allocator, out: *std.Io.Writer.Allocating, q: Question, codes try w.writeAll("\n\nReply with the option code only."); }, .noul => { - // the template shows the descriptions only when one side has one, then both; - // a side head needs each option rendered: the bare answer words stand in when no text is given + // the template shows descriptions when both sides carry one; a side head renders text or bare words const any = q.texts.len == 2 and (q.texts[0] != null and q.texts[0].? != .null or q.texts[1] != null and q.texts[1].? != .null); if (any or positions) for ([_][]const u8{ "Yes", "No" }, 0..) |name, i| { const t: ?Value = if (i < q.texts.len) q.texts[i] else null; From df7f0af13300a9cee06003832ad29ca52c111a28 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Tue, 20 Oct 2026 23:58:07 -1000 Subject: [PATCH 20/21] lanes: a queued decision takes a free slot beside the live streams, as a new stream's prompt pass would MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A score/decide call no longer waits for an idle engine: the decider's `ready` hook says when a slot is free, and the queued pass runs beside the live rounds — its prompt prefilled in the backend's usual chunks, its label rows read, its slot unbound without a drawn token. A queued request a lane is free for still goes first; a slot that will not fit yet leaves the pass queued (error.Busy) for the next gap. An idle engine runs a decision exactly as before, and a decider without `ready` keeps the idle-only rule. qwen27 binds a synthetic stream (id "decision", max_new 0, drafts off) to a free slot for the pass's length — the same bind an admission makes: buffers grown where needed, held drafts dropped, runner slot reset, drafter context cleared — and `unbind` frees it, promoting a live stream into slot 0 when the pass held the drafter's. The pass never overlaps a planned round, so draft-slot ownership and acceptance are untouched. --- zig/src/core/engine_api.zig | 8 +- zig/src/core/lane_host.zig | 6 +- zig/src/core/lane_host_decide.zig | 36 ++++++++- zig/src/core/lane_host_decide_test.zig | 60 ++++++++++++++- zig/src/core/lane_host_lone.zig | 3 +- zig/src/core/lane_rounds.zig | 2 +- zig/src/native/qwen27_decide.zig | 102 ++++++++++++++++++------- zig/src/native/qwen27_lanes.zig | 6 +- zig/src/native/qwen27_served.zig | 11 ++- zig/src/native/qwen27_slots.zig | 4 +- 10 files changed, 190 insertions(+), 48 deletions(-) diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index e7e1f6d8b0..93876ba566 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -314,10 +314,14 @@ pub const DecisionTask = union(enum) { decide: *Decision, }; -/// A lane-hosted backend's decision passes, run on the engine thread once no stream is live. +/// A lane-hosted backend's decision passes: an idle engine runs them; one whose ``ready`` says a +/// slot is free runs them beside the live streams, as a new stream's prompt pass would. pub const Decider = struct { ctx: *anyopaque, - run: *const fn (ctx: *anyopaque, task: *DecisionTask) error{ Failed, Refused }!void, + /// A pass on the engine thread; ``error.Busy`` leaves it queued for the next gap (its slot went). + run: *const fn (ctx: *anyopaque, task: *DecisionTask) error{ Failed, Refused, Busy }!void, + /// Engine thread, the host locked: may a pass take a slot beside the live streams (null: idle only)? + ready: ?*const fn (ctx: *anyopaque) bool = null, }; /// An opened engine; ``close`` stops its thread and frees its backend. diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 63f7b00748..5d10a967dc 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -64,8 +64,8 @@ pub const LaneHost = struct { memory: ?api.MemorySource = null, // the backend's memory counts; null: Engine.memory reports none explain: ?api.Explain = null, // the backend's words for a request it refuses; null: the error's name learner: ?api.Learner = null, // the family's Sliding Weights learner; null: learn requests are refused - decider: ?api.Decider = null, // the backend's decision passes, run while the engine is idle - decisions: std.ArrayList(*decide.Work) = .empty, // score/decide calls waiting for an idle engine, in order + decider: ?api.Decider = null, // the backend's decision passes: an idle engine's, or beside the live streams when its `ready` says a slot is free + decisions: std.ArrayList(*decide.Work) = .empty, // score/decide calls waiting for a pass, in order decided: std.Io.Condition = .init, // a queued decision finished; its requester wakes lessons: std.ArrayList(Lesson) = .empty, // learn requests waiting for an idle engine lesson_open: bool = false, // the learner holds a begun job (engine thread only) @@ -526,7 +526,7 @@ pub const LaneHost = struct { ends.takeCancels(h); rows_host.answer(h); learn.refuse(h); - decide.step(h); // a queued decision takes an idle engine before new work; a stale one fails it + decide.step(h); // a queued decision's pass: an idle engine's, or a free slot beside the live streams while (h.admitOne()) {} h.fillOne(); h.core.sweep() catch |err| { diff --git a/zig/src/core/lane_host_decide.zig b/zig/src/core/lane_host_decide.zig index 8e31a73251..4f631fb7da 100644 --- a/zig/src/core/lane_host_decide.zig +++ b/zig/src/core/lane_host_decide.zig @@ -1,6 +1,7 @@ //! A lane host's decision passes: HTTP threads queue them; the engine thread runs them while no stream lives. const std = @import("std"); const api = @import("engine_api.zig"); +const admit = @import("lane_host_admit.zig"); const LaneHost = @import("lane_host.zig").LaneHost; /// One queued decision pass: its task in, the failure and completion its requester waits on. @@ -42,27 +43,54 @@ pub fn wait(h: *LaneHost, work: *Work) error{ Failed, Refused }!void { if (work.failed) |e| return e; } -/// The oldest queued decision's pass: an idle engine runs it, a stale one fails the waiters instead. +/// The oldest queued decision's pass: an idle engine runs it; a decider whose ``ready`` says a +/// slot is free runs it beside the live streams. A stale host fails the waiters instead. pub fn step(h: *LaneHost) void { h.lock(); if (h.stale) { h.unlock(); return failAll(h); } - const idle = !h.closing and h.decider != null and h.admitted.items.len == 0 and h.filling.items.len == 0 and h.core.activeCount() == 0; - const work: ?*Work = if (idle and h.decisions.items.len > 0) h.decisions.orderedRemove(0) else null; + const work = head(h); h.unlock(); const w = work orelse return; const decider = h.decider.?; decider.run(decider.ctx, &w.task) catch |e| { - w.failed = e; + switch (e) { + error.Busy => return, // the slot the pass wanted went: it keeps the queue's head for the next gap + else => |x| w.failed = @errorCast(x), + } }; h.lock(); + _ = h.decisions.orderedRemove(0); // it is still the head: only this thread takes works out w.finished = true; h.unlock(); h.decided.broadcast(h.io); } +/// The queue's head when a pass may run now (the caller holds the lock): always on an idle engine; +/// beside the live streams when the decider's ``ready`` says a slot is free and no queued request +/// the lanes would serve waits for that turn (a decision takes the same turn a new stream would). +fn head(h: *LaneHost) ?*Work { + if (h.closing or h.decisions.items.len == 0) return null; + const decider = h.decider orelse return null; + const idle = h.admitted.items.len == 0 and h.filling.items.len == 0 and h.core.activeCount() == 0; + if (idle) return h.decisions.items[0]; + const ready = decider.ready orelse return null; + if (h.filling.items.len > 0) return null; + const lane_free = h.admitted.items.len < @min(h.info_.lanes, h.memory_lanes orelse h.info_.lanes); + if (lane_free and h.queued.items.len > 0 and admit.pick(h) != null) return null; // a request's turn + return if (ready(decider.ctx)) h.decisions.items[0] else null; +} + +/// A queued decision a free slot could start now (the caller holds the lock): the lone driver hands +/// its stream to the lane core for it rather than hold the engine through the whole reply. +pub fn waiting(h: *LaneHost) bool { + const decider = h.decider orelse return false; + const ready = decider.ready orelse return false; + return h.decisions.items.len > 0 and ready(decider.ctx); +} + /// Every queued decision failed (the host is closing or the device went stale): requesters wake. pub fn failAll(h: *LaneHost) void { h.lock(); diff --git a/zig/src/core/lane_host_decide_test.zig b/zig/src/core/lane_host_decide_test.zig index 70c8930144..485b50fec8 100644 --- a/zig/src/core/lane_host_decide_test.zig +++ b/zig/src/core/lane_host_decide_test.zig @@ -12,9 +12,15 @@ const Request = api.Request; const Record = struct { core: *lanes.Engine, calls: usize = 0, - idle: bool = true, // activeCount at every pass: decisions only ever run on an idle engine + idle: bool = true, // activeCount at every pass: false once a pass ran beside a live stream + live_ok: bool = false, // what `ready` answers when the host asks beside live work (false: idle only) fail: ?error{ Failed, Refused } = null, + fn ready(ctx: *anyopaque) bool { + const r: *Record = @ptrCast(@alignCast(ctx)); + return r.live_ok; + } + fn run(ctx: *anyopaque, task: *api.DecisionTask) error{ Failed, Refused }!void { const r: *Record = @ptrCast(@alignCast(ctx)); defer r.calls += 1; @@ -160,3 +166,55 @@ test "a queued decision runs once the live stream ends, and the text reply still try std.testing.expect(record.idle); try std.testing.expectEqual(@as(usize, 1), record.calls); } + +test "a decider that can take a slot beside the rounds answers a queued decision while a reply is live" { + const gpa = std.testing.allocator; + var cfg = try config(gpa); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.backend(), clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 1 }); + var record = Record{ .core = &core, .live_ok = true }; + host.decider = .{ .ctx = &record, .run = Record.run, .ready = Record.ready }; + try host.start(); + defer host.stop(); + const Box = struct { + mutex: std.Io.Mutex = .init, + done: ?Reason = null, + tokens: usize = 0, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens += t.len, + .finished => |f| b.done = f.reason, + else => {}, + } + } + fn read(b: *@This()) struct { done: ?Reason, tokens: usize } { + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + return .{ .done = b.done, .tokens = b.tokens }; + } + }; + const prompt = [_]u32{ 2, 7, 1, 8 }; + var box: Box = .{}; + const request: Request = .{ .prompt = &prompt, .max_tokens = 1_000_000 }; + const e = host.engine(); + try e.submit(1, &request, .{ .ctx = &box, .event = Box.event }); + // the reply is live (its first tokens landed) before the decision is queued: it must not wait for the end + while (box.read().tokens == 0) std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + const labels = [_]u32{5}; + var logits = [_]f64{0}; + try std.testing.expectEqual(@as(f64, 0), try e.score(&prompt, &labels, &logits, false)); + try std.testing.expectEqualSlices(f64, &.{5}, &logits); + try std.testing.expect(!record.idle); // the pass ran beside the reply, not after it + try std.testing.expectEqual(@as(usize, 1), record.calls); + e.cancel(1); + while (box.read().done == null) std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + try std.testing.expectEqual(Reason.cancelled, box.read().done.?); +} diff --git a/zig/src/core/lane_host_lone.zig b/zig/src/core/lane_host_lone.zig index e7cf1f17cf..b54dcefe69 100644 --- a/zig/src/core/lane_host_lone.zig +++ b/zig/src/core/lane_host_lone.zig @@ -1,6 +1,7 @@ //! A lane host's lone driver: the backend's own loop runs a lone stream until a foreground request or a cancel waits. const LaneHost = @import("lane_host.zig").LaneHost; const status = @import("lane_host_status.zig"); +const decide = @import("lane_host_decide.zig"); const Job = LaneHost.Job; /// An idle backend driver takes a lone drafted request, sampled, with logprobs, a budget or a grammar when supported. @@ -55,7 +56,7 @@ pub fn run(h: *LaneHost, job: *Job, began: i96, resumed: bool) bool { const host: *LaneHost = @ptrCast(@alignCast(ctx)); host.lock(); defer host.unlock(); - return host.foregroundWaiting() or host.cancels.items.len > 0 or host.closing; + return host.foregroundWaiting() or host.cancels.items.len > 0 or host.closing or decide.waiting(host); } }; const paused = lone.run(lone.ctx, &job.stream, .{ .ctx = h, .committed = Hooks.committed, .yield = Hooks.yield, .resumed = resumed }); diff --git a/zig/src/core/lane_rounds.zig b/zig/src/core/lane_rounds.zig index 161eb40c7f..46f158c82c 100644 --- a/zig/src/core/lane_rounds.zig +++ b/zig/src/core/lane_rounds.zig @@ -266,7 +266,7 @@ pub fn run(h: *LaneHost, rd: Rounds) void { ends.takeCancels(h); rows_host.answer(h); learn.refuse(h); - decide.step(h); // a queued decision takes an idle engine before new work; a stale one fails it + decide.step(h); // a queued decision's pass: an idle engine's, or a free slot beside the live streams h.lock(); if (h.closing) { const left = h.queued.items.len + h.admitted.items.len; diff --git a/zig/src/native/qwen27_decide.zig b/zig/src/native/qwen27_decide.zig index 7f3cafcdc3..1222429300 100644 --- a/zig/src/native/qwen27_decide.zig +++ b/zig/src/native/qwen27_decide.zig @@ -1,25 +1,75 @@ -//! The 27B lane backend's decision passes on an idle engine: the trunk's state kept once, each branch read over it. +//! The 27B lane backend's decision passes: a synthetic stream's admission on a free slot, beside the +//! live rounds as a new stream's prompt pass would, or alone on the idle engine as before. const std = @import("std"); const api = @import("engine_api"); const tf = @import("tensorfold"); const q = tf.qwen27; +const lanes = tf.lanes; const lb = @import("qwen27_lanes.zig"); +const slots = @import("qwen27_slots.zig"); const joint = @import("qwen27_joint.zig"); -/// A queued decision's GPU work on the idle host: score the row after its prompt; decide, each branch's last row. -pub fn run(b: *lb.Metal, head: ?*joint.Head, window: u32, task: *api.DecisionTask) error{ Failed, Refused }!void { +/// The LaneHost's busy gate: a decision pass may take a slot beside the live streams when one is free. +pub fn ready(b: *const lb.Metal) bool { + return b.freeSlot() != null; +} + +/// A queued decision's GPU work (LaneHost runs it alone): score reads the row after its prompt; decide, each branch's last row. +/// `head` loaded: `labels` are instead option-final token positions, each read of the pass's post-final-norm hidden rows. +/// The pass binds a free slot as a synthetic stream — its rows join the rounds as any admitted prompt's do, +/// it never draws a token — or stays queued (error.Busy) while the batch gives no slot. +pub fn run(b: *lb.Metal, head: ?*joint.Head, window: u32, task: *api.DecisionTask) error{ Failed, Refused, Busy }!void { switch (task.*) { - .score => |*s| s.out = try score(b, head, s.prompt, s.labels, s.logits, s.mass), - .decide => |d| try decide(b, head, window, d), + .score => |*s| { + var stream = try synthetic(b, s.prompt); + defer stream.deinit(b.gpa); + const slot = try bind(b, &stream); + defer b.unbind(&stream); + s.out = try score(b, head, slot, s.prompt, s.labels, s.logits, s.mass); + }, + .decide => |d| { + try check(window, d); + var stream = try synthetic(b, d.trunk); + defer stream.deinit(b.gpa); + const slot = try bind(b, &stream); + defer b.unbind(&stream); + try decide(b, head, slot, d); + }, } } -/// `labels`' logits at the row after `prompt` (`.last` head, logits row 0); a JointHead scores option-final rows. -fn score(b: *lb.Metal, head: ?*joint.Head, prompt: []const u32, labels: []const u32, logits: []f64, mass: bool) error{Failed}!f64 { +/// The pass's synthetic stream: a slot's owner for the pass's rows; drafts off, it never generates. +fn synthetic(b: *lb.Metal, ids: []const u32) error{Failed}!lanes.Stream { + return lanes.Stream.init(b.gpa, .{ .id = "decision", .prompt = ids, .max_new = 0, .drafts = false }) catch return error.Failed; +} + +/// `s` bound to a free slot the way a new stream's admission binds one: its buffers grown where needed, +/// the slot's held drafts dropped and its runner state reset. `b.unbind(s)` frees it (the caller's defer). +/// error.Busy: no slot is free, or one's buffers do not fit yet — the pass stays queued for the next gap. +fn bind(b: *lb.Metal, s: *lanes.Stream) error{ Failed, Busy }!u32 { + b.settle() catch {}; // a pending verify's rows are kept before the pass touches the shared frame + b.dropRound(); + const slot = b.freeSlot() orelse return error.Busy; + b.ready(slot, s) catch |e| return switch (e) { + error.LaneMemory => error.Busy, + else => failed("a decision's slot", e), + }; + b.dropTreeOf(slot); + b.owners[slot] = s; + b.runner.rope[slot] = .{}; + b.runner.images = null; + slots.recover(b, s) catch |e| return failed("recovering the device", e); + if (b.drafter(slot)) |d| d.reset() catch |e| return failed("resetting a decision's slot", e) else b.runner.reset(slot) catch |e| return failed("resetting a decision's slot", e); + slots.clearContext(b, slot) catch |e| return failed("clearing a decision's slot", e); + return slot; +} + +/// `labels`' logits at the row after `prompt` on `slot`: a prefill's `.last` head leaves them in logits row 0. +/// A JointHead instead: `labels` are each option's final-token position, scored on that row's hidden state. +fn score(b: *lb.Metal, head: ?*joint.Head, slot: u32, prompt: []const u32, labels: []const u32, logits: []f64, mass: bool) error{Failed}!f64 { const r = b.runner; if (prompt.len == 0 or prompt.len > r.capacity or logits.len < labels.len) return error.Failed; - try reset(b); - const session = q.session.Session{ .runner = r }; + const session = q.session.Session{ .runner = r, .slot = slot }; if (head) |h| { try forward(h, b, session, 0, prompt, labels, logits); return massOf(logits[0..labels.len], mass); @@ -28,9 +78,8 @@ fn score(b: *lb.Metal, head: ?*joint.Head, prompt: []const u32, labels: []const return read(row0(r), labels, logits, mass); } -/// Each branch's label logits at its last row after the shared trunk; a `head` reads option-final positions. -fn decide(b: *lb.Metal, head: ?*joint.Head, window: u32, d: *api.Decision) error{ Failed, Refused }!void { - const r = b.runner; +/// What a packed decision needs before it takes a slot (its refusal words reach the requester). +fn check(window: u32, d: *api.Decision) error{ Failed, Refused }!void { if (d.pictures.len > 0) return refuse(d, "decisions with pictures are not served by this engine yet"); if (d.trunk.len == 0 or d.branches.len == 0 or d.counts.len != d.branches.len) return error.Failed; var wanted: usize = 0; @@ -41,15 +90,21 @@ fn decide(b: *lb.Metal, head: ?*joint.Head, window: u32, d: *api.Decision) error wanted = std.math.add(usize, wanted, n) catch return error.Failed; } if (d.labels.len < wanted or d.logits.len < wanted) return error.Failed; - try reset(b); - const session = q.session.Session{ .runner = r }; +} + +/// Each branch's label logits at its last row after the shared trunk on `slot`: the trunk's state is kept once +/// and restored a branch at a time; every branch's rows stand on the same trunk state. +/// A JointHead instead: `labels` are option-final positions in the whole sequence (trunk's length on a branch's rows). +fn decide(b: *lb.Metal, head: ?*joint.Head, slot: u32, d: *api.Decision) error{Failed}!void { + const r = b.runner; + const session = q.session.Session{ .runner = r, .slot = slot }; session.prefill(d.trunk, b.chunk) catch |e| return failed("the trunk's prompt pass", e); - const st = q.snapshot.saveSlot(b.gpa, r, 0, null, @intCast(d.trunk.len)) catch |e| return failed("keeping the trunk's state", e); + const st = q.snapshot.saveSlot(b.gpa, r, slot, null, @intCast(d.trunk.len)) catch |e| return failed("keeping the trunk's state", e); defer q.snapshot.drop(b.gpa, st); d.read = d.trunk.len; var at: usize = 0; for (d.branches, d.counts) |br, n| { - q.snapshot.restoreSlot(b.gpa, r, 0, null, st) catch |e| return failed("restoring the trunk's state", e); + q.snapshot.restoreSlot(b.gpa, r, slot, null, st) catch |e| return failed("restoring the trunk's state", e); if (head) |h| { try forward(h, b, session, d.trunk.len, br, d.labels[at..][0..n], d.logits[at..][0..n]); } else { @@ -59,10 +114,11 @@ fn decide(b: *lb.Metal, head: ?*joint.Head, window: u32, d: *api.Decision) error d.read += br.len; at += n; } - reset(b) catch {}; // slot 0 stands empty for the next stream, whatever its prompt was } -/// `ids` prefilled in chunks: each `positions` entry is an option's final-token position offset by `base`. +/// `ids` prefilled in chunks (each chunk's rows alone stand in `input`): a `positions` inside the chunk +/// is that option's final-token position, base the rows' offset in the whole sequence (0 for a prompt, +/// the trunk's length on a branch); the JointHead's score lands in `logits`. fn forward(h: *joint.Head, b: *lb.Metal, session: q.session.Session, base: usize, ids: []const u32, positions: []const u32, logits: []f64) error{Failed}!void { const r = b.runner; for (positions) |p| if (p < base or p - base >= ids.len) return error.Failed; @@ -80,16 +136,6 @@ fn forward(h: *joint.Head, b: *lb.Metal, session: q.session.Session, base: usize } } -/// Slot 0 empty for a pass: a pending keep ends first, and the drafter's reset covers its session too. -fn reset(b: *lb.Metal) error{Failed}!void { - b.settle() catch {}; - b.dropRound(); - b.runner.rope[0] = .{}; - b.runner.images = null; - if (b.draft) |d| d.reset() catch |e| return failed("resetting slot 0", e); - return b.runner.reset(0) catch |e| failed("resetting slot 0", e); -} - /// The last pass's logits row 0: bf16, `vocab` wide; the `.last` head leaves its one row there. fn row0(r: *const q.decode_round.Runner) []const u16 { const logits = r.model.frame.get(.logits); diff --git a/zig/src/native/qwen27_lanes.zig b/zig/src/native/qwen27_lanes.zig index 1202a2eced..08d139d0b0 100644 --- a/zig/src/native/qwen27_lanes.zig +++ b/zig/src/native/qwen27_lanes.zig @@ -112,10 +112,10 @@ pub const Metal = struct { pub const lone = slots.lone; const prefillSlot = slots.prefill; - const freeSlot = slots.free; + pub const freeSlot = slots.free; const ringChunk = slots.ringChunk; - const unbind = slots.unbind; - const ready = slots.ready; + pub const unbind = slots.unbind; + pub const ready = slots.ready; const refill = slots.refill; pub fn backend(b: *Metal) be.Backend { diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index d26657afc3..b028a174e7 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -87,11 +87,16 @@ pub const Host = struct { return images.prepare(h.vision.?, h.clock.io, a, bytes, max_tokens); } - /// LaneHost's decision hook: Ornith's label reads over a shared trunk, on an idle engine. - fn decideRun(ptr: *anyopaque, task: *api.DecisionTask) error{ Failed, Refused }!void { + /// LaneHost's decision hook: Ornith's label reads over a shared trunk, beside live rounds on a free slot. + fn decideRun(ptr: *anyopaque, task: *api.DecisionTask) error{ Failed, Refused, Busy }!void { const h = of(ptr); return decide.run(h.back, h.joint, h.host.info_.context_window, task); } + + /// The host's busy gate: a queued decision may take a slot beside the live streams when one is free. + fn decideReady(ptr: *anyopaque) bool { + return decide.ready(of(ptr).back); + } }; /// The reply's state at `stands` (its end or one short), born with its request's prompt: that request's marks stay. @@ -335,7 +340,7 @@ pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.df h.host = api.LaneHost.init(gpa, io, &h.core, .{ .name = "qwen27", .structures = @import("qwen27_masks.zig").supported, .lanes = slots, .context_window = capacity, .prefill_step = @intCast(h.back.chunk), .vision = offer }); h.host.memory = .{ .ctx = h, .read = readMemory }; h.host.explain = .{ .text = words }; - h.host.decider = .{ .ctx = h, .run = Host.decideRun }; // the engine still answers text: info_.decides stays off + h.host.decider = .{ .ctx = h, .run = Host.decideRun, .ready = Host.decideReady }; // the engine still answers text: info_.decides stays off h.host.decoded_rows = true; // a reply's rows prefill with the bits decoding gave them, with the cache on or off // a lone greedy request takes the serial loop on tensor units, where it is faster (TF_QWEN27_LANES overrides) const units = model.device.tensorUnits(); diff --git a/zig/src/native/qwen27_slots.zig b/zig/src/native/qwen27_slots.zig index 4b70deb42a..a6a0919e2b 100644 --- a/zig/src/native/qwen27_slots.zig +++ b/zig/src/native/qwen27_slots.zig @@ -77,7 +77,7 @@ pub fn prefill(b: *Metal, s: *lanes.Stream, slot: u32) !void { } /// After a GPU failure whose streams have all left, the runner and drafter take rounds again; `s` takes its slot now. -fn recover(b: *Metal, s: *const lanes.Stream) !void { +pub fn recover(b: *Metal, s: *const lanes.Stream) !void { const drafter = if (b.draft) |d| d.model.backend.failed else false; if (!b.runner.failed and !drafter) return; for (b.owners[0..b.runner.slots]) |o| if (o != null and o != s) return error.RoundGpuFailure; @@ -90,7 +90,7 @@ fn recover(b: *Metal, s: *const lanes.Stream) !void { } /// Empty `slot`'s own drafter state past slot 0 (its taps ring and session), as a fresh prompt starts it. -fn clearContext(b: *Metal, slot: u32) !void { +pub fn clearContext(b: *Metal, slot: u32) !void { if (slot == 0) return; const d = b.draft orelse return; if (b.contexts[slot]) |x| try x.reset(d.model.graph.config); From 7f82411cfe6624db8d6805eaa60bfa8c070f05d4 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Wed, 21 Oct 2026 16:51:41 -1000 Subject: [PATCH 21/21] lanes+qwen27: one-line comments for the lean check --- zig/src/core/engine_api.zig | 3 +-- zig/src/core/lane_host_decide.zig | 10 +++------- zig/src/native/qwen27_decide.zig | 23 ++++++----------------- 3 files changed, 10 insertions(+), 26 deletions(-) diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 93876ba566..a5215cc8b3 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -314,8 +314,7 @@ pub const DecisionTask = union(enum) { decide: *Decision, }; -/// A lane-hosted backend's decision passes: an idle engine runs them; one whose ``ready`` says a -/// slot is free runs them beside the live streams, as a new stream's prompt pass would. +/// A lane-hosted backend's decision passes: on an idle engine, or beside the streams when ``ready`` frees a slot. pub const Decider = struct { ctx: *anyopaque, /// A pass on the engine thread; ``error.Busy`` leaves it queued for the next gap (its slot went). diff --git a/zig/src/core/lane_host_decide.zig b/zig/src/core/lane_host_decide.zig index 4f631fb7da..5b282d6c45 100644 --- a/zig/src/core/lane_host_decide.zig +++ b/zig/src/core/lane_host_decide.zig @@ -43,8 +43,7 @@ pub fn wait(h: *LaneHost, work: *Work) error{ Failed, Refused }!void { if (work.failed) |e| return e; } -/// The oldest queued decision's pass: an idle engine runs it; a decider whose ``ready`` says a -/// slot is free runs it beside the live streams. A stale host fails the waiters instead. +/// The oldest queued decision's pass (idle engine, or beside the streams on a free slot); a stale host fails waiters. pub fn step(h: *LaneHost) void { h.lock(); if (h.stale) { @@ -68,9 +67,7 @@ pub fn step(h: *LaneHost) void { h.decided.broadcast(h.io); } -/// The queue's head when a pass may run now (the caller holds the lock): always on an idle engine; -/// beside the live streams when the decider's ``ready`` says a slot is free and no queued request -/// the lanes would serve waits for that turn (a decision takes the same turn a new stream would). +/// The queue's head when a pass may run now: an idle engine, or a free slot after queued requests' turns. fn head(h: *LaneHost) ?*Work { if (h.closing or h.decisions.items.len == 0) return null; const decider = h.decider orelse return null; @@ -83,8 +80,7 @@ fn head(h: *LaneHost) ?*Work { return if (ready(decider.ctx)) h.decisions.items[0] else null; } -/// A queued decision a free slot could start now (the caller holds the lock): the lone driver hands -/// its stream to the lane core for it rather than hold the engine through the whole reply. +/// A queued decision a free slot could start now: the lone driver hands its stream to the lane core. pub fn waiting(h: *LaneHost) bool { const decider = h.decider orelse return false; const ready = decider.ready orelse return false; diff --git a/zig/src/native/qwen27_decide.zig b/zig/src/native/qwen27_decide.zig index 1222429300..fd281e70e0 100644 --- a/zig/src/native/qwen27_decide.zig +++ b/zig/src/native/qwen27_decide.zig @@ -1,5 +1,4 @@ -//! The 27B lane backend's decision passes: a synthetic stream's admission on a free slot, beside the -//! live rounds as a new stream's prompt pass would, or alone on the idle engine as before. +//! Decision passes on the 27B lane backend: a free-slot synthetic stream beside live rounds, or alone when idle. const std = @import("std"); const api = @import("engine_api"); const tf = @import("tensorfold"); @@ -14,10 +13,7 @@ pub fn ready(b: *const lb.Metal) bool { return b.freeSlot() != null; } -/// A queued decision's GPU work (LaneHost runs it alone): score reads the row after its prompt; decide, each branch's last row. -/// `head` loaded: `labels` are instead option-final token positions, each read of the pass's post-final-norm hidden rows. -/// The pass binds a free slot as a synthetic stream — its rows join the rounds as any admitted prompt's do, -/// it never draws a token — or stays queued (error.Busy) while the batch gives no slot. +/// A queued decision's GPU work: a synthetic stream on a free slot beside the rounds, or `error.Busy` to wait. pub fn run(b: *lb.Metal, head: ?*joint.Head, window: u32, task: *api.DecisionTask) error{ Failed, Refused, Busy }!void { switch (task.*) { .score => |*s| { @@ -43,9 +39,7 @@ fn synthetic(b: *lb.Metal, ids: []const u32) error{Failed}!lanes.Stream { return lanes.Stream.init(b.gpa, .{ .id = "decision", .prompt = ids, .max_new = 0, .drafts = false }) catch return error.Failed; } -/// `s` bound to a free slot the way a new stream's admission binds one: its buffers grown where needed, -/// the slot's held drafts dropped and its runner state reset. `b.unbind(s)` frees it (the caller's defer). -/// error.Busy: no slot is free, or one's buffers do not fit yet — the pass stays queued for the next gap. +/// `s` bound to a free slot as a new stream's admission binds it; `b.unbind(s)` frees it (`error.Busy` = wait). fn bind(b: *lb.Metal, s: *lanes.Stream) error{ Failed, Busy }!u32 { b.settle() catch {}; // a pending verify's rows are kept before the pass touches the shared frame b.dropRound(); @@ -64,8 +58,7 @@ fn bind(b: *lb.Metal, s: *lanes.Stream) error{ Failed, Busy }!u32 { return slot; } -/// `labels`' logits at the row after `prompt` on `slot`: a prefill's `.last` head leaves them in logits row 0. -/// A JointHead instead: `labels` are each option's final-token position, scored on that row's hidden state. +/// `labels`' logits at the row after `prompt` on `slot` (`.last` → logits row 0); a JointHead scores option-final rows. fn score(b: *lb.Metal, head: ?*joint.Head, slot: u32, prompt: []const u32, labels: []const u32, logits: []f64, mass: bool) error{Failed}!f64 { const r = b.runner; if (prompt.len == 0 or prompt.len > r.capacity or logits.len < labels.len) return error.Failed; @@ -92,9 +85,7 @@ fn check(window: u32, d: *api.Decision) error{ Failed, Refused }!void { if (d.labels.len < wanted or d.logits.len < wanted) return error.Failed; } -/// Each branch's label logits at its last row after the shared trunk on `slot`: the trunk's state is kept once -/// and restored a branch at a time; every branch's rows stand on the same trunk state. -/// A JointHead instead: `labels` are option-final positions in the whole sequence (trunk's length on a branch's rows). +/// Each branch's label logits at its last row after the shared trunk on `slot`; a `head` reads option-final positions. fn decide(b: *lb.Metal, head: ?*joint.Head, slot: u32, d: *api.Decision) error{Failed}!void { const r = b.runner; const session = q.session.Session{ .runner = r, .slot = slot }; @@ -116,9 +107,7 @@ fn decide(b: *lb.Metal, head: ?*joint.Head, slot: u32, d: *api.Decision) error{F } } -/// `ids` prefilled in chunks (each chunk's rows alone stand in `input`): a `positions` inside the chunk -/// is that option's final-token position, base the rows' offset in the whole sequence (0 for a prompt, -/// the trunk's length on a branch); the JointHead's score lands in `logits`. +/// `ids` prefilled in chunks: each `positions` entry is an option's final-token position offset by `base`. fn forward(h: *joint.Head, b: *lb.Metal, session: q.session.Session, base: usize, ids: []const u32, positions: []const u32, logits: []f64) error{Failed}!void { const r = b.runner; for (positions) |p| if (p < base or p - base >= ids.len) return error.Failed;