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/core/row_projection.metal b/zig/kernels/metal/core/row_projection.metal index 0be006be24..842f96434f 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 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)]], + 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/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..21c4493fc5 --- /dev/null +++ b/zig/kernels/metal/qwen27/moe.metal @@ -0,0 +1,138 @@ +// 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; }; // 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, 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(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]], + 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]; + 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 + 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)); + const int stride = a.topk + 1; + for (int k = 0; k < a.topk; 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]); + 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)))); + } +} + +// 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)]], + 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++) 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]; + OFFS[e] = at; + 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]; + } + 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] = 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]]) { + 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 * 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/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/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/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/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/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_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/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/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/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..a89ad5d262 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; + // 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; } 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..b0aa5a7248 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), 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,14 @@ 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"); + // 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); if (rope) |r| if (r.get("mrope_interleaved")) |v| { if (v != .bool or !v.bool) return error.UnsupportedMrope; @@ -238,6 +262,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); // 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 + 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 }); diff --git a/zig/src/families/qwen27/decode_round.zig b/zig/src/families/qwen27/decode_round.zig index 4bfd54920a..76abf3afea 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); @@ -168,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 { @@ -354,14 +360,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 +391,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/dflash/config.zig b/zig/src/families/qwen27/dflash/config.zig index 92a5dea74f..5bd662f747 100644 --- a/zig/src/families/qwen27/dflash/config.zig +++ b/zig/src/families/qwen27/dflash/config.zig @@ -1,7 +1,11 @@ -//! 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 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 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 +24,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 +33,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 +53,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 +114,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 +142,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..ceedc84d91 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(); @@ -136,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..09df7d1135 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, TF_DRAFT_EDGE override them. + 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/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/forward.zig b/zig/src/families/qwen27/forward.zig index 7cf8e564bc..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); @@ -103,6 +110,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 +124,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 + 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/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/gpu_frame.zig b/zig/src/families/qwen27/gpu_frame.zig index 6c068ca305..913d073259 100644 --- a/zig/src/families/qwen27/gpu_frame.zig +++ b/zig/src/families/qwen27/gpu_frame.zig @@ -4,10 +4,15 @@ 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; + +/// 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()) @import("widths.zig").routerRows(config.experts) else 0; +} pub const logit_rows = 128; pub const Frame = struct { @@ -18,6 +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 + 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, @@ -31,6 +37,12 @@ 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 }; } + /// 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 + 1)); + } + 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..743f627691 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 (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 } }; pub const Weights = struct { allocator: std.mem.Allocator, @@ -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,24 @@ 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 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 + 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), + }; + // 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 }; } }; @@ -229,6 +296,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..e5a0c912f8 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,7 +73,9 @@ pub const Model = struct { errdefer m.kernels.deinit(); m.glue = try glue.Kernels.init(m.device); errdefer m.glue.deinit(); - m.frame = try Frame.init(m.device, c, @max(rows, promptRows())); + 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(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); @@ -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(); @@ -98,9 +106,9 @@ 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 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; 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; @@ -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..f400871d94 --- /dev/null +++ b/zig/src/families/qwen27/moe_gpu.zig @@ -0,0 +1,210 @@ +//! 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"); +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, 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 }; + +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; + } + + /// 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; + 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 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(); + const enc = span.encoder; + 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 + const pipes = try k.rowPipelines(); + 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); + 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(); + 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) }); + 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); // 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 + 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, routed, hidden); + enc2.barrier(); + 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, routed, inter); + enc2.barrier(); + const act = R{ .buf = f.get(.expert_act).buffer }; + 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 }, routed, experts); + } else { + // 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, 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(.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(); + } + + 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; 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); + 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.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)); + } +}; + +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/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/families/qwen27/widths.zig b/zig/src/families/qwen27/widths.zig index e1f9b7a64f..938fa6dc5a 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,23 @@ 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'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; + 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 gate, then zero rows to a multiple of 64 (whole weight tiles). +pub fn routerRows(experts: usize) usize { + return (experts + 1 + 63) / 64 * 64; +} + const all_rows = [_]u8{ 2, 3, 4, 5, 6, 8 }; fn testConfig(a: std.mem.Allocator, quantization: []const u8) ![]u8 { 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_decide.zig b/zig/src/native/qwen27_decide.zig new file mode 100644 index 0000000000..8a1ca2a31c --- /dev/null +++ b/zig/src/native/qwen27_decide.zig @@ -0,0 +1,103 @@ +//! 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 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), + .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 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"); + 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_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_lanes.zig b/zig/src/native/qwen27_lanes.zig index 819d88d62c..1202a2eced 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`. @@ -384,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; diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index 2e33bc3168..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. @@ -177,7 +184,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 } @@ -311,7 +318,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); @@ -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(); diff --git a/zig/tests/qwen27/run.zig b/zig/tests/qwen27/run.zig index 5d31a8efd7..8f0265b032 100644 --- a/zig/tests/qwen27/run.zig +++ b/zig/tests/qwen27/run.zig @@ -18,6 +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_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 { @@ -48,6 +50,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 +59,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)) }; }; @@ -136,14 +141,18 @@ 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 }); 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;