From b117a8bcf214c4a3cdd57b2a96874b2bf31be0e7 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Tue, 20 Oct 2026 12:36:37 -0500 Subject: [PATCH 01/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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");