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..e7e1f6d8b0 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -200,6 +200,8 @@ pub const Info = struct { prompt_cache_plan: ?PromptCachePlan = null, /// The engine scores decisions (`score`) and writes no text: the server refuses text requests. decides: bool = false, + /// `score`/`decide` labels are option-final token positions, not token ids: a side head reads there. + decision_positions: bool = false, /// Decisions may carry pictures (`Decision.pictures`). pictures: bool = false, /// Image input (--vision); null: text only. @@ -286,6 +288,8 @@ pub const Open = struct { policy: []const u8 = "", /// --vision: load the checkpoint's image tower (null: text only; an engine without one refuses). vision: ?VisionOpen = null, + /// --joint-head: a decision side-head's folder; when set, score/decide labels are option-final token positions. + joint_head: ?[]const u8 = null, }; /// A packed decision: a trunk (pictures expand at `at`), a branch a question, and each branch's labels (`counts`). @@ -302,6 +306,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..552fb54604 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) { @@ -350,8 +352,8 @@ fn openQwen27(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[ var why: []const u8 = ""; const drafter = if (o.drafts) o.drafter else null; const learn: ?@import("qwen27_disk.zig").Learn = if (o.learn_conversations) if (o.learn) |root| .{ .root = root, .dir = o.dir, .drafter = drafter, .cap = gibBytes(o.learn_gib), .floor = gibBytes(o.learn_min_free_gib), .daily = gibBytes(o.learn_daily_gib) } else null else null; - const h = qwen27.open(gpa, io, o.dir, .{ .context = @intCast(context), .streams = o.lanes, .streams_fixed = o.lanes_fixed, .drafter = drafter, .drafter_bits = o.drafter_bits, .cache_gib = o.prompt_cache_gib, .cache_over_cap = o.prompt_cache_over_cap, .learn = learn, .vision = o.vision }, a, &why) catch |err| { - 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) }); + const h = qwen27.open(gpa, io, o.dir, .{ .context = @intCast(context), .streams = o.lanes, .streams_fixed = o.lanes_fixed, .drafter = drafter, .drafter_bits = o.drafter_bits, .cache_gib = o.prompt_cache_gib, .cache_over_cap = o.prompt_cache_over_cap, .learn = learn, .vision = o.vision, .joint_head = o.joint_head }, a, &why) catch |err| { + problem.* = if (why.len > 0) why else try std.fmt.allocPrint(a, "the native Qwen engine (27B class) cannot load {s} ({s})", .{ o.dir, @errorName(err) }); return null; }; return .{ .engine = h.engine(), .close = qwen27.close, .ctx = h }; @@ -412,6 +414,7 @@ test { _ = glm; _ = glm_dsa; _ = qwen27; + _ = @import("qwen27_joint.zig"); _ = @import("cache_fit.zig"); _ = @import("resident.zig"); _ = @import("metal_identity.zig"); diff --git a/zig/src/native/qwen27_decide.zig b/zig/src/native/qwen27_decide.zig new file mode 100644 index 0000000000..7f3cafcdc3 --- /dev/null +++ b/zig/src/native/qwen27_decide.zig @@ -0,0 +1,146 @@ +//! 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"); +const joint = @import("qwen27_joint.zig"); + +/// A queued decision's GPU work on the idle host: score the row after its prompt; decide, each branch's last row. +pub fn run(b: *lb.Metal, head: ?*joint.Head, window: u32, task: *api.DecisionTask) error{ Failed, Refused }!void { + switch (task.*) { + .score => |*s| s.out = try score(b, head, s.prompt, s.labels, s.logits, s.mass), + .decide => |d| try decide(b, head, window, d), + } +} + +/// `labels`' logits at the row after `prompt` (`.last` head, logits row 0); a JointHead scores option-final rows. +fn score(b: *lb.Metal, head: ?*joint.Head, prompt: []const u32, labels: []const u32, logits: []f64, mass: bool) error{Failed}!f64 { + const r = b.runner; + if (prompt.len == 0 or prompt.len > r.capacity or logits.len < labels.len) return error.Failed; + try reset(b); + const session = q.session.Session{ .runner = r }; + if (head) |h| { + try forward(h, b, session, 0, prompt, labels, logits); + return massOf(logits[0..labels.len], mass); + } + session.prefill(prompt, b.chunk) catch |e| return failed("a score's prompt pass", e); + return read(row0(r), labels, logits, mass); +} + +/// Each branch's label logits at its last row after the shared trunk; a `head` reads option-final positions. +fn decide(b: *lb.Metal, head: ?*joint.Head, window: u32, d: *api.Decision) error{ Failed, Refused }!void { + const r = b.runner; + if (d.pictures.len > 0) return refuse(d, "decisions with pictures are not served by this engine yet"); + 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); + if (head) |h| { + try forward(h, b, session, d.trunk.len, br, d.labels[at..][0..n], d.logits[at..][0..n]); + } else { + session.prefill(br, b.chunk) catch |e| return failed("a branch's prompt pass", e); + _ = try read(row0(r), d.labels[at..][0..n], d.logits[at..][0..n], false); + } + d.read += br.len; + at += n; + } + reset(b) catch {}; // slot 0 stands empty for the next stream, whatever its prompt was +} + +/// `ids` prefilled in chunks: each `positions` entry is an option's final-token position offset by `base`. +fn forward(h: *joint.Head, b: *lb.Metal, session: q.session.Session, base: usize, ids: []const u32, positions: []const u32, logits: []f64) error{Failed}!void { + const r = b.runner; + for (positions) |p| if (p < base or p - base >= ids.len) return error.Failed; + const dim: usize = @intCast(r.model.config.hidden); + var first: usize = 0; + while (first < ids.len) { + const end = @min(ids.len, first + @as(usize, r.promptRows(@intCast(b.chunk)))); + session.promptChunk(ids[first..end], false) catch |e| return failed("a decision's prompt pass", e); + const words = hiddenRows(r); + for (positions, logits) |p, *x| { + if (p - base < first or p - base >= end) continue; + x.* = h.score(words[(p - base - first) * dim ..][0..dim]) catch |e| return failed("a JointHead option", e); + } + first = end; + } +} + +/// Slot 0 empty for a pass: a pending keep ends first, and the drafter's reset covers its session too. +fn reset(b: *lb.Metal) error{Failed}!void { + b.settle() catch {}; + 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 last pass's post-final-norm hidden rows: the norm the head's input reads leaves every row there. +fn hiddenRows(r: *const q.decode_round.Runner) [*]const u16 { + const input = r.model.frame.get(.input); + return @ptrCast(@alignCast(input.buffer.contents() + input.offset)); +} + +/// The row's `labels` logits out, then the vocabulary logsumexp when `mass` (else 0). +fn read(row: []const u16, labels: []const u32, logits: []f64, mass: bool) error{Failed}!f64 { + for (labels, logits) |label, *x| { + 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 JointHead pass has no vocabulary: the option scores' own logsumexp stands as the label mass (it reads 1). +fn massOf(logits: []const f64, mass: bool) f64 { + if (!mass or logits.len == 0) return 0; + var peak: f64 = -std.math.inf(f64); + for (logits) |x| peak = @max(peak, x); + var total: f64 = 0; + for (logits) |x| total += @exp(x - peak); + return peak + @log(total); +} + +/// A bf16 logit as f64; a nonfinite one fails the pass, as a draw's argmax check does. +fn value(bits: u16) error{Failed}!f64 { + const v: f64 = @as(f32, @bitCast(@as(u32, bits) << 16)); + 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_joint.zig b/zig/src/native/qwen27_joint.zig new file mode 100644 index 0000000000..5314ed2732 --- /dev/null +++ b/zig/src/native/qwen27_joint.zig @@ -0,0 +1,272 @@ +//! A JointHead beside the 27B's trunk: option-final hidden rows in, one score an option out. +const std = @import("std"); +const tf = @import("tensorfold"); +const st = tf.safetensors; +const Allocator = std.mem.Allocator; + +/// `joint_head_config.json`: the head's shape and its weight file's format. +pub const Config = struct { + dim: usize, // the trunk's hidden width + width: usize, // the head's + layers: usize, // Linear→GELU residual blocks +}; + +/// The head: `proj` (dim→width), Linear→GELU residual layers, `out` (width→1); bf16 words read from the file. +pub const Head = struct { + gpa: Allocator, + io: std.Io, + file: st.File, + cfg: Config, + data: []u16, // every weight, one owned row: proj, each block, out + proj_w: []const u16, // [width, dim] + proj_b: []const u16, // [width] + block_w: [][]const u16, // a [width, width] a block + block_b: [][]const u16, // a [width] a block + out_w: []const u16, // [width] + out_b: u16, + x: []f32, // a row after proj and each residual + y: []f32, // a block's Linear→GELU half + h: []f32, // a hidden row as f32 + + /// `path`'s head: `joint_head_config.json` plus its one safetensors (a directory), or the file beside the config. + pub fn open(gpa: Allocator, io: std.Io, path: []const u8) !*Head { + var dir_path = path; + var name: ?[]const u8 = null; + if (std.mem.endsWith(u8, path, ".safetensors")) { + name = std.fs.path.basename(path); + dir_path = std.fs.path.dirname(path) orelse "."; + } + var dir = try std.Io.Dir.cwd().openDir(io, dir_path, .{ .iterate = true }); + defer dir.close(io); + const text = try dir.readFileAlloc(io, "joint_head_config.json", gpa, .limited(1 << 20)); + defer gpa.free(text); + const cfg = try parseConfig(gpa, text); + const file_name = name orelse try oneShard(gpa, io, dir_path); + defer if (name == null) gpa.free(file_name); + const file_path = try std.fs.path.join(gpa, &.{ dir_path, file_name }); + defer gpa.free(file_path); + const h = try gpa.create(Head); + errdefer gpa.destroy(h); + h.* = .{ .gpa = gpa, .io = io, .file = try st.File.open(gpa, io, file_path), .cfg = cfg, .data = &.{}, .proj_w = &.{}, .proj_b = &.{}, .block_w = &.{}, .block_b = &.{}, .out_w = &.{}, .out_b = 0, .x = &.{}, .y = &.{}, .h = &.{} }; + errdefer h.file.close(io); + h.h = try gpa.alloc(f32, cfg.dim); + errdefer gpa.free(h.h); + h.x = try gpa.alloc(f32, cfg.width); + errdefer gpa.free(h.x); + h.y = try gpa.alloc(f32, cfg.width); + errdefer gpa.free(h.y); + h.data = try gpa.alloc(u16, cfg.width * cfg.dim + cfg.width + cfg.layers * cfg.width * (cfg.width + 1) + cfg.width + 1); + errdefer gpa.free(h.data); + h.block_w = try gpa.alloc([]const u16, cfg.layers); + errdefer gpa.free(h.block_w); + h.block_b = try gpa.alloc([]const u16, cfg.layers); + errdefer gpa.free(h.block_b); + var off: usize = 0; + h.proj_w = try h.part(&off, "proj.weight", &.{ cfg.width, cfg.dim }); + h.proj_b = try h.part(&off, "proj.bias", &.{cfg.width}); + for (0..cfg.layers) |i| { + var buf: [64]u8 = undefined; + h.block_w[i] = try h.part(&off, try std.fmt.bufPrint(&buf, "blocks.{d}.0.weight", .{i}), &.{ cfg.width, cfg.width }); + h.block_b[i] = try h.part(&off, try std.fmt.bufPrint(&buf, "blocks.{d}.0.bias", .{i}), &.{cfg.width}); + } + h.out_w = try h.part(&off, "out.weight", &.{ 1, cfg.width }); + h.out_b = (try h.part(&off, "out.bias", &.{1}))[0]; + return h; + } + + pub fn close(h: *Head) void { + h.gpa.free(h.block_b); + h.gpa.free(h.block_w); + h.gpa.free(h.data); + h.gpa.free(h.y); + h.gpa.free(h.x); + h.gpa.free(h.h); + h.file.close(h.io); + h.gpa.destroy(h); + } + + /// One option's score: `hidden` a post-final-norm row of the trunk (`dim` bf16 words). + pub fn score(h: *Head, hidden: []const u16) error{Failed}!f64 { + if (hidden.len != h.cfg.dim) return error.Failed; + for (h.h, hidden) |*d, bits| { + const v: f32 = @bitCast(@as(u32, bits) << 16); + if (!std.math.isFinite(v)) return error.Failed; + d.* = v; + } + matvec(h.proj_w, h.cfg.dim, h.h, h.proj_b, h.x); + for (h.block_w, h.block_b) |w, bias| { + matvec(w, h.cfg.width, h.x, bias, h.y); + for (h.x, h.y) |*a, b| a.* += gelu(b); + } + var s: f64 = f32of(h.out_b); + for (h.x, h.out_w) |a, bits| s += @as(f64, f32of(bits)) * a; + return if (std.math.isFinite(s)) s else error.Failed; + } + + /// `data`'s next `shape` elements taken as the bf16 tensor `name`, checked. + fn part(h: *Head, off: *usize, name: []const u8, shape: []const usize) ![]const u16 { + const t = h.file.get(name) orelse return error.MissingTensor; + if (!t.is(.bf16, shape)) return error.UnexpectedTensor; + const dest = h.data[off.*..][0..t.numel()]; + @memcpy(std.mem.sliceAsBytes(dest), t.bytes); + off.* += t.numel(); + return dest; + } +}; + +fn f32of(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +/// y[j] = bias[j] + row j of `w` ([out, in] bf16) times `x`, f64 accumulation. +fn matvec(w: []const u16, in: usize, x: []const f32, bias: []const u16, y: []f32) void { + for (y, bias, 0..) |*o, b, j| { + var s: f64 = f32of(b); + for (w[j * in ..][0..in], x) |wb, xv| s += @as(f64, f32of(wb)) * xv; + o.* = @floatCast(s); + } +} + +/// nn.GELU: x·Φ(x); Φ's tail by Abramowitz–Stegun 26.2.17 (|ε| ≤ 7.5e-8), 1/√(2π) = 0.39894228040143268. +fn gelu(x: f32) f32 { + const t = 1.0 / (1.0 + 0.2316419 * @abs(x)); + const poly = t * (0.319381530 + t * (-0.356563782 + t * (1.781477937 + t * (-1.821255978 + t * 1.330274429)))); + const tail = poly * @exp(-x * x / 2) * 0.39894228040143268; + return x * (if (x < 0) tail else 1.0 - tail); +} + +/// The config's dim, width and layers; format must be bf16-safetensors-v1. +fn parseConfig(a: Allocator, text: []const u8) !Config { + const parsed = try std.json.parseFromSlice(std.json.Value, a, text, .{}); + defer parsed.deinit(); + if (parsed.value != .object) return error.BadJointHeadConfig; + const o = parsed.value.object; + const format = o.get("format") orelse return error.BadJointHeadConfig; + if (format != .string or !std.mem.eql(u8, format.string, "bf16-safetensors-v1")) return error.BadJointHeadConfig; + return .{ .dim = try field(o, "dim"), .width = try field(o, "width"), .layers = try field(o, "layers") }; +} + +fn field(o: std.json.ObjectMap, name: []const u8) !usize { + const v = o.get(name) orelse return error.BadJointHeadConfig; + if (v != .integer or v.integer <= 0) return error.BadJointHeadConfig; + return @intCast(v.integer); +} + +/// The directory's one safetensors file's name; zero or several are refused. +fn oneShard(gpa: Allocator, io: std.Io, dir_path: []const u8) ![]const u8 { + var dir = try std.Io.Dir.cwd().openDir(io, dir_path, .{ .iterate = true }); + defer dir.close(io); + var found: ?[]const u8 = null; + var it = dir.iterate(); + while (try it.next(io)) |e| { + if (e.kind != .file or !std.mem.endsWith(u8, e.name, ".safetensors")) continue; + if (found != null) return error.JointHeadShardCount; + found = try gpa.dupe(u8, e.name); + } + return found orelse error.JointHeadShardCount; +} + +fn bf16(x: f32) u16 { + const bits: u32 = @bitCast(x); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +/// A head fixture in `tmp` (dim 4, width 4, 2 blocks; every weight exactly bf16): its directory's path. +fn fixture(gpa: Allocator, tmp: *const std.testing.TmpDir) ![]u8 { + const header = + \\{"proj.weight":{"dtype":"BF16","shape":[4,4],"data_offsets":[0,32]}, + \\ "proj.bias":{"dtype":"BF16","shape":[4],"data_offsets":[32,40]}, + \\ "blocks.0.0.weight":{"dtype":"BF16","shape":[4,4],"data_offsets":[40,72]}, + \\ "blocks.0.0.bias":{"dtype":"BF16","shape":[4],"data_offsets":[72,80]}, + \\ "blocks.1.0.weight":{"dtype":"BF16","shape":[4,4],"data_offsets":[80,112]}, + \\ "blocks.1.0.bias":{"dtype":"BF16","shape":[4],"data_offsets":[112,120]}, + \\ "out.weight":{"dtype":"BF16","shape":[1,4],"data_offsets":[120,128]}, + \\ "out.bias":{"dtype":"BF16","shape":[1],"data_offsets":[128,130]}} + ; + const weights = [_][]const f32{ + &.{ 0.5, -0.25, 0.125, 0 }, &.{ 0, 0.5, -0.25, 0.125 }, &.{ 0.125, 0, 0.5, -0.25 }, &.{ -0.25, 0.125, 0, 0.5 }, // proj + &.{ 0.125, -0.25, 0.375, -0.5 }, // proj bias + &.{ 0.25, 0.125, 0, 0 }, &.{ 0, 0.25, 0.125, 0 }, &.{ 0, 0, 0.25, 0.125 }, &.{ 0.125, 0, 0, 0.25 }, // block 0 + &.{ 0.125, -0.125, 0.125, -0.125 }, // block 0 bias + &.{ 0.5, 0.25, 0, 0 }, &.{ 0, 0.5, 0.25, 0 }, &.{ 0, 0, 0.5, 0.25 }, &.{ 0.25, 0, 0, 0.5 }, // block 1 + &.{ -0.25, 0.25, -0.25, 0.25 }, // block 1 bias + &.{ 0.5, -0.5, 0.25, -0.25 }, // out + &.{0.125}, // out bias + }; + var data: [8 + header.len + 130]u8 = undefined; + std.mem.writeInt(u64, data[0..8], header.len, .little); + @memcpy(data[8..][0..header.len], header); + var at: usize = 8 + header.len; + for (weights) |w| for (w) |x| { + std.mem.writeInt(u16, data[at..][0..2], bf16(x), .little); + at += 2; + }; + try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "head.safetensors", .data = &data }); + try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "joint_head_config.json", .data = "{\"dim\":4,\"width\":4,\"layers\":2,\"format\":\"bf16-safetensors-v1\"}" }); + const root = try tmp.dir.realPathFileAlloc(std.testing.io, ".", gpa); + defer gpa.free(root); + return std.fs.path.join(gpa, &.{root}); +} + +fn row(gpa: Allocator, xs: []const f32) ![]u16 { + const out = try gpa.alloc(u16, xs.len); + for (xs, out) |x, *b| b.* = bf16(x); + return out; +} + +test "a JointHead loads its config and tensors and scores a hidden row" { + const a = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try fixture(a, &tmp); + defer a.free(dir); + const h = try Head.open(a, std.testing.io, dir); + defer h.close(); + try std.testing.expectEqual(@as(usize, 4), h.cfg.dim); + try std.testing.expectEqual(@as(usize, 4), h.cfg.width); + try std.testing.expectEqual(@as(usize, 2), h.cfg.layers); + const hidden = try row(a, &.{ 0.5, -0.25, 0.125, -0.5 }); + defer a.free(hidden); + try std.testing.expectApproxEqAbs(1.035910918342501, try h.score(hidden), 1e-4); + const second = try row(a, &.{ 0.25, 0.5, -0.125, 0.75 }); + defer a.free(second); + try std.testing.expectApproxEqAbs(0.09842275576323394, try h.score(second), 1e-4); + const third = try row(a, &.{ -1.0, 0.0, 1.0, 0.5 }); + defer a.free(third); + try std.testing.expectApproxEqAbs(0.3339105568875796, try h.score(third), 1e-4); + const file = try std.fs.path.join(a, &.{ dir, "head.safetensors" }); + defer a.free(file); + const other = try Head.open(a, std.testing.io, file); + defer other.close(); + try std.testing.expectApproxEqAbs(1.035910918342501, try other.score(hidden), 1e-4); // the file path itself also opens +} + +test "a JointHead refuses a wrong-width row and a nonfinite one" { + const a = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try fixture(a, &tmp); + defer a.free(dir); + const h = try Head.open(a, std.testing.io, dir); + defer h.close(); + const hidden = try row(a, &.{ 0.5, -0.25, 0.125 }); + defer a.free(hidden); + try std.testing.expectError(error.Failed, h.score(hidden)); + const bad = try row(a, &.{ 0.5, -0.25, 0.125, -0.5 }); + defer a.free(bad); + bad[2] = 0x7FC0; // a NaN hidden word ends the pass + try std.testing.expectError(error.Failed, h.score(bad)); +} + +test "a JointHead refuses a wrong-format config and a missing tensor" { + const a = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try fixture(a, &tmp); + defer a.free(dir); + try tmp.dir.writeFile(io, .{ .sub_path = "joint_head_config.json", .data = "{\"dim\":4,\"width\":4,\"layers\":2,\"format\":\"fp32\"}" }); + try std.testing.expectError(error.BadJointHeadConfig, Head.open(a, io, dir)); + try tmp.dir.writeFile(io, .{ .sub_path = "joint_head_config.json", .data = "{\"dim\":8,\"width\":4,\"layers\":2,\"format\":\"bf16-safetensors-v1\"}" }); + try std.testing.expectError(error.UnexpectedTensor, Head.open(a, io, dir)); // proj.weight is [4, 4], not [4, 8] +} diff --git a/zig/src/native/qwen27_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..d26657afc3 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -12,6 +12,8 @@ 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 joint = @import("qwen27_joint.zig"); const disk = @import("qwen27_disk.zig"); const images = @import("qwen27_images.zig"); const Allocator = std.mem.Allocator; @@ -27,6 +29,7 @@ pub const Options = struct { learn: ?disk.Learn = null, // --learn-conversations: kept states written behind to disk and read back vision: ?api.VisionOpen = null, // --vision: open loads the image tower, and requests may hold images images: ?*images.Vision = null, // the loaded tower attach serves (and the host then owns) + joint_head: ?[]const u8 = null, // --joint-head: a decision side-head; score/decide read its option-final rows }; pub const Host = struct { @@ -45,6 +48,7 @@ pub const Host = struct { warm: mtl.keepalive.Target, disk_stage: ?mtl.Buffer = null, // copies out of private kept states for the disk (qwen27_disk.zig) vision: ?*images.Vision = null, // --vision's tower and image rows + joint: ?*joint.Head = null, // --joint-head's loaded side-head pub fn engine(h: *Host) api.Engine { return h.host.engine(); @@ -82,6 +86,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.joint, 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 +187,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 } @@ -309,9 +319,10 @@ pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.df return error.DrafterLoad; }; errdefer if (h.draft) |d| d.deinit(); + h.joint = null; h.back = try lb.Metal.init(gpa, &h.runner, h.draft); errdefer h.back.deinit(); - h.back.chunk = h.runner.promptRows(q.model.promptRows()); + h.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 +335,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(); @@ -338,6 +350,18 @@ pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.df if (o.learn) |l| try disk.attach(h, io, a, l); errdefer if (h.cache) |*store| if (store.disk) |r| r.deinit(); h.host.keepalive_target = .{ .ctx = &h.warm, .tick = mtl.keepalive.Target.tick }; + if (o.joint_head) |p| { + h.joint = joint.Head.open(gpa, io, p) catch |err| { + why.* = try std.fmt.allocPrint(a, "the native Qwen3.8-27B joint head cannot load {s} ({s})", .{ p, @errorName(err) }); + return error.JointHeadLoad; + }; + if (h.joint.?.cfg.dim != model.config.hidden) { + why.* = try std.fmt.allocPrint(a, "--joint-head's dim {d} does not match this checkpoint's hidden width {d}", .{ h.joint.?.cfg.dim, model.config.hidden }); + return error.JointHeadDim; + } + } + errdefer if (h.joint) |j| j.close(); + if (h.joint != null) h.host.info_.decision_positions = true; // score/decide labels are option-final rows try h.host.start(); std.log.info("Qwen3.8-27B: {d} streams of {d} tokens{s}{s}", .{ slots, capacity, if (h.draft != null) ", DFlash2 drafts for one of them" else "", if (o.images != null) ", images read" else "" }); return h; @@ -393,6 +417,7 @@ pub fn close(ptr: *anyopaque) void { h.core.deinit(); h.cfg.deinit(h.gpa); h.back.deinit(); + if (h.joint) |j| j.close(); if (h.draft) |d| d.deinit(); h.runner.deinit(); if (h.vision) |v| v.deinit(); diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 8afebaebe9..86649ce179 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -125,6 +125,8 @@ pub const flags = [_]Flag{ .{ .name = "--recall-dir", .native = true }, .{ .name = "--recall-gib", .native = true }, .{ .name = "--recall-read-only", .kind = .store_true, .native = true }, + // the native server's own: a side head (joint_head_config.json + bf16 safetensors) on option-final rows + .{ .name = "--joint-head", .native = true }, // the native server's own: Cortheon's reasoning program and server-run web research for every chat request .{ .name = "--cortheon", .kind = .store_true, .native = true }, }; @@ -178,6 +180,7 @@ pub const Args = struct { compact_memory: ?[]const u8 = null, drafter: ?[]const u8 = null, drafter_bits: u8 = 4, + joint_head: ?[]const u8 = null, // --joint-head: a decision side-head; score/decide labels become option-final token positions record: ?[]const u8 = null, // drafter training rows written under this folder kv_cache: []const u8 = "bf16", // fp8: e4m3 attention caches with a scale a row, about twice the context slide: bool = false, // Sliding Weights learns what /v1/slide/learn is told into the served weights, live @@ -319,7 +322,7 @@ fn apply(a: Allocator, out: *Args, name: []const u8, value: ?[]const u8, u: *Usa out.learn_daily_gib = try gib(u, a, name, v); } else if (is(name, "--prompt-seed")) { out.prompt_seed = true; - } else if (is(name, "--max-tokens")) out.max_tokens = try int(u, a, name, v) else if (is(name, "--temperature")) out.temperature = try float(u, a, name, v) else if (is(name, "--top-p")) out.top_p = try float(u, a, name, v) else if (is(name, "--top-k")) out.top_k = try int(u, a, name, v) else if (is(name, "--min-p")) out.min_p = try float(u, a, name, v) else if (is(name, "--thinking")) out.thinking = true else if (is(name, "--no-thinking")) out.thinking = false else if (is(name, "--reasoning-effort")) out.reasoning_effort = v else if (is(name, "--thinking-budget")) out.thinking_budget = try int(u, a, name, v) else if (is(name, "--loop-guard")) out.loop_guard = true else if (is(name, "--drafter")) out.drafter = v else if (is(name, "--drafter-bits")) out.drafter_bits = @intCast(try int(u, a, name, v)) else if (is(name, "--record")) out.record = v else if (is(name, "--kv-cache")) out.kv_cache = v else if (is(name, "--no-drafts")) out.no_drafts = true else if (is(name, "--keep-warm")) out.keep_warm = try int(u, a, name, v) else if (is(name, "--compact-at")) { + } else if (is(name, "--max-tokens")) out.max_tokens = try int(u, a, name, v) else if (is(name, "--temperature")) out.temperature = try float(u, a, name, v) else if (is(name, "--top-p")) out.top_p = try float(u, a, name, v) else if (is(name, "--top-k")) out.top_k = try int(u, a, name, v) else if (is(name, "--min-p")) out.min_p = try float(u, a, name, v) else if (is(name, "--thinking")) out.thinking = true else if (is(name, "--no-thinking")) out.thinking = false else if (is(name, "--reasoning-effort")) out.reasoning_effort = v else if (is(name, "--thinking-budget")) out.thinking_budget = try int(u, a, name, v) else if (is(name, "--loop-guard")) out.loop_guard = true else if (is(name, "--drafter")) out.drafter = v else if (is(name, "--drafter-bits")) out.drafter_bits = @intCast(try int(u, a, name, v)) else if (is(name, "--joint-head")) out.joint_head = v else if (is(name, "--record")) out.record = v else if (is(name, "--kv-cache")) out.kv_cache = v else if (is(name, "--no-drafts")) out.no_drafts = true else if (is(name, "--keep-warm")) out.keep_warm = try int(u, a, name, v) else if (is(name, "--compact-at")) { if (std.mem.eql(u8, v, "auto")) out.compact_auto = true else { const f = try float(u, a, name, v); if (!(f > 0 and f <= 1)) return fail(u, a, "argument --compact-at: expected auto or a fraction in (0, 1]: '{s}'", .{v}); @@ -438,6 +441,7 @@ pub fn request(a: Allocator, dir: []const u8, model_type: []const u8, args: Args .drafts = !args.no_drafts, .drafter = args.drafter, .drafter_bits = args.drafter_bits, + .joint_head = args.joint_head, .record = args.record, .kv8 = std.mem.eql(u8, args.kv_cache, "fp8"), .speed_up = args.speed_up, diff --git a/zig/src/server/decisions.zig b/zig/src/server/decisions.zig index 215de8416a..a431055e99 100644 --- a/zig/src/server/decisions.zig +++ b/zig/src/server/decisions.zig @@ -157,10 +157,11 @@ const Wording = struct { kind: Kind, names: []const []const u8, labels: []const []const u8, + lines: []const []const u8, // every rendered line, the option lines between the question and the closing content: []const u8, }; -fn wording(cx: *Cx, input: []const u8, question: Value) errors.Refused!Wording { +fn wording(cx: *Cx, input: []const u8, question: Value, positions: bool) errors.Refused!Wording { const kind_value = question.get("type") orelse .null; const kind: Kind = if (kind_value == .string and std.mem.eql(u8, kind_value.string, "choice")) .choice @@ -189,7 +190,10 @@ fn wording(cx: *Cx, input: []const u8, question: Value) errors.Refused!Wording { try rejectUnknown(cx, question.object, &score_fields); const listed = try scoreList(cx, question); try lines.append(cx.a, try questionLine(cx, question)); - for (listed.labels, listed.details) |label, detail| try lines.append(cx.a, try std.fmt.allocPrint(cx.a, "{s}: {s}", .{ label, detail })); + for (listed.labels, listed.details) |label, detail| { + const line = try std.fmt.allocPrint(cx.a, "{s}: {s}", .{ label, detail }); + try lines.append(cx.a, line); + } break :blk .{ listed.names, listed.labels }; }, .yes_no => blk: { @@ -199,8 +203,15 @@ fn wording(cx: *Cx, input: []const u8, question: Value) errors.Refused!Wording { const lead = try renderText(cx, question.get("question")); if (reply_text.pyStrip(lead).len == 0) return cx.refuse("a question must not be blank"); try lines.append(cx.a, try std.fmt.allocPrint(cx.a, "Is the following true? {s}", .{lead})); - if (yes.len > 0) try lines.append(cx.a, try std.fmt.allocPrint(cx.a, "yes: {s}", .{yes})); - if (no.len > 0) try lines.append(cx.a, try std.fmt.allocPrint(cx.a, "no: {s}", .{no})); + // a side head needs each option rendered: the bare answer words stand in when no text is given + if (yes.len > 0 or positions) { + const line = try std.fmt.allocPrint(cx.a, "yes: {s}", .{if (yes.len > 0) yes else "yes"}); + try lines.append(cx.a, line); + } + if (no.len > 0 or positions) { + const line = try std.fmt.allocPrint(cx.a, "no: {s}", .{if (no.len > 0) no else "no"}); + try lines.append(cx.a, line); + } const names = [_][]const u8{ "yes", "no" }; break :blk .{ try cx.a.dupe([]const u8, &names), try cx.a.dupe([]const u8, &names) }; }, @@ -211,7 +222,7 @@ fn wording(cx: *Cx, input: []const u8, question: Value) errors.Refused!Wording { .yes_no => "Answer with yes or no only.", }; try lines.append(cx.a, closing); - return .{ .kind = kind, .names = names, .labels = labels, .content = try contentOf(cx.a, input, lines.items) }; + return .{ .kind = kind, .names = names, .labels = labels, .lines = lines.items, .content = try contentOf(cx.a, input, lines.items) }; } const Listed = struct { @@ -291,7 +302,8 @@ fn contentOf(a: Allocator, input: []const u8, lines: []const []const u8) Allocat } fn prepareOne(srv: *Server, cx: *Cx, input: []const u8, question: Value) errors.Refused!Prepared { - const word = try wording(cx, input, question); + const positions = srv.engine.info().decision_positions; + const word = try wording(cx, input, question, positions); const prompt_text = try renderPrompt(srv, cx, word.content); const prompt_ids = try encodePrompt(srv, cx, prompt_text); try thinkOpen(cx, prompt_text); @@ -301,7 +313,10 @@ fn prepareOne(srv: *Server, cx: *Cx, input: []const u8, question: Value) errors. .kind = word.kind, .names = word.names, .prompt_ids = prompt_ids, - .label_ids = try labelIds(srv, cx, prompt_text, prompt_ids, word.labels), + .label_ids = if (positions) + try optionPositions(srv, cx, prompt_text, prompt_ids, input.len, word.lines, word.content) + else + try labelIds(srv, cx, prompt_text, prompt_ids, word.labels), }; } @@ -354,6 +369,25 @@ fn labelRefuse(cx: *Cx, label: []const u8) errors.Refused { return cx.fail(.request, "the answer label {s} is not one distinct token after the chat prompt for this tokenizer, so this model is not supported", .{pyrepr.repr(cx.a, .{ .string = label }) catch label}); } +/// A side head's `labels`: each option's final-token position in `prompt` (its line's last token, `base`d). +fn optionPositions(srv: *Server, cx: *Cx, prompt: []const u8, prompt_ids: []const u32, input_len: usize, lines: []const []const u8, content: []const u8) errors.Refused![]u32 { + const a = cx.a; + const base = std.mem.indexOf(u8, prompt, content) orelse return positionRefuse(cx, lines[0]); + const ats = try a.alloc(usize, lines.len - 2); + var at = base + input_len + 2 + lines[0].len + 1; // past the input, the question line and their rows + for (lines[1 .. lines.len - 1], ats) |line, *end| { + end.* = at + line.len; + at = end.* + 1; + } + const out = try a.alloc(u32, ats.len); + if (!(srv.text.indexAt(a, prompt, prompt_ids, ats, out) catch return error.OutOfMemory)) return positionRefuse(cx, lines[0]); + return out; +} + +fn positionRefuse(cx: *Cx, line: []const u8) errors.Refused { + return cx.fail(.request, "the option line {s} is not found at a token boundary in the chat prompt for this tokenizer, so this model's decision head cannot read it", .{pyrepr.repr(cx.a, .{ .string = line }) catch line}); +} + fn scoreAll(srv: *Server, cx: *Cx, prepared: []const Prepared) errors.Refused![]Scored { const scored = try cx.a.alloc(Scored, prepared.len); for (prepared, scored) |item, *row| { diff --git a/zig/src/server/model_text.zig b/zig/src/server/model_text.zig index fd04340b14..76c2954d1d 100644 --- a/zig/src/server/model_text.zig +++ b/zig/src/server/model_text.zig @@ -73,6 +73,25 @@ pub const Text = struct { return f(t.ctx, a, id); } + /// `ids` are `text`'s encoding: for each of `ats` (ascending ends) the covering token's index, or false. + pub fn indexAt(t: Text, a: Allocator, text: []const u8, ids: []const u32, ats: []const usize, out: []u32) Allocator.Error!bool { + var i: usize = 0; + var end: usize = 0; + for (ats, out) |at, *o| { + if (at == 0 or at > text.len) return false; + while (end < at) { + if (i == ids.len) return false; + const piece = try t.decode(a, ids[i .. i + 1]); + defer a.free(piece); + if (piece.len == 0 or !std.mem.startsWith(u8, text[end..], piece)) return false; + end += piece.len; + i += 1; + } + o.* = @intCast(i - 1); + } + return true; + } + pub fn eosIds(t: Text) []const u32 { return t.vtable.eos_ids(t.ctx); } diff --git a/zig/src/server/systemone.zig b/zig/src/server/systemone.zig index 57b03f33d5..175cd94a80 100644 --- a/zig/src/server/systemone.zig +++ b/zig/src/server/systemone.zig @@ -76,12 +76,15 @@ pub fn reply(srv: *Server, cx: *Cx, body: Value) errors.Refused!Value { const trunk = if (seen.urls.len > 0) try std.mem.concat(cx.a, u32, &.{ head, try f.encode(try prefix(cx.a, seen.state, false)) }) else try f.encode(try prefix(cx.a, seen.state, true)); const codes = try cx.a.alloc([]const Code, questions.len); const branches = try cx.a.alloc([]const u32, questions.len); - for (questions, codes, branches) |q, *c, *b| { + const groups = try cx.a.alloc([]const []const u32, questions.len); + const positions = srv.engine.info().decision_positions; // a side head reads option lines' last rows + for (questions, codes, branches, groups) |q, *c, *b, *g| { c.* = if (q.kind == .choice) try f.codes(q.keys) else &.{}; - b.* = try f.encode(try suffix(cx.a, q, c.*)); + var ends: std.ArrayList(usize) = .empty; + const rendered = try suffix(cx.a, q, c.*, positions, &ends); + b.* = try f.encode(rendered); + g.* = if (positions) try f.positions(rendered, b.*, ends.items, trunk.len) else try f.readout(q, c.*); } - const groups = try cx.a.alloc([]const []const u32, questions.len); - for (questions, codes, groups) |q, c, *g| g.* = try f.readout(q, c); const window = srv.engine.info().context_window; var read: usize = trunk.len; for (branches) |b| { @@ -270,6 +273,21 @@ const Format = struct { } } + /// A side head's readout: one group an option, its line's last token's index (`base` the trunk's length). + fn positions(f: *Format, text_: []const u8, ids: []const u32, ends: []const usize, base: usize) errors.Refused![]const []const u32 { + const a = f.cx.a; + const ats = try a.alloc(u32, ends.len); + if (!(f.srv.text.indexAt(a, text_, ids, ends, ats) catch return error.OutOfMemory)) + return f.cx.refuse("an option's line is not found at a token boundary in this decision's text, so this model's decision head cannot read it"); + const groups = try a.alloc([]const u32, ends.len); + for (ats, groups) |p, *g| { + const one = try a.alloc(u32, 1); + one[0] = @intCast(base + p); + g.* = one; + } + return groups; + } + fn forms(f: *Format, texts: []const []const u8) errors.Refused![]const u32 { var seen: std.ArrayList(u32) = .empty; for (texts) |t| if (try f.single(t)) |id| if (std.mem.indexOfScalar(u32, seen.items, id) == null) try seen.append(f.cx.a, id); @@ -334,14 +352,16 @@ fn indented(w: *std.Io.Writer, v: Value, level: usize) std.Io.Writer.Error!void } } -fn suffix(a: Allocator, q: Question, codes: []const Code) Allocator.Error![]const u8 { +fn suffix(a: Allocator, q: Question, codes: []const Code, positions: bool, ends: *std.ArrayList(usize)) Allocator.Error![]const u8 { var out: std.Io.Writer.Allocating = .init(a); - questionText(a, &out.writer, q, codes) catch return error.OutOfMemory; + questionText(a, &out, q, codes, positions, ends) catch return error.OutOfMemory; out.writer.writeAll(im_end ++ "\n" ++ im_start ++ "assistant\n") catch return error.OutOfMemory; return out.toOwnedSlice(); } -fn questionText(a: Allocator, w: *std.Io.Writer, q: Question, codes: []const Code) !void { +/// The question's text; `ends` takes each option line's end offset for a side head's positions. +fn questionText(a: Allocator, out: *std.Io.Writer.Allocating, q: Question, codes: []const Code, positions: bool, ends: *std.ArrayList(usize)) !void { + const w = &out.writer; try w.writeAll(q.instructions); switch (q.kind) { .choice => { @@ -351,21 +371,27 @@ fn questionText(a: Allocator, w: *std.Io.Writer, q: Question, codes: []const Cod try w.print("{s} ", .{c.text}); // a description as written (JSON when it is not text); without one, the name if (desc != null and desc.?.truthy()) try w.writeAll(try text(a, desc.?)) else for (key) |ch| try w.writeByte(if (ch == '_') ' ' else ch); + try ends.append(a, out.written().len); } try w.writeAll("\n\nReply with the option code only."); }, .noul => { - // the template shows the descriptions only when one side has one, then both + // the template shows descriptions when both sides carry one; a side head renders text or bare words const any = q.texts.len == 2 and (q.texts[0] != null and q.texts[0].? != .null or q.texts[1] != null and q.texts[1].? != .null); - if (any) for ([_][]const u8{ "Yes", "No" }, q.texts) |name, t| { - const shown: Value = if (t != null and t.? != .null) t.? else .{ .string = "None" }; + if (any or positions) for ([_][]const u8{ "Yes", "No" }, 0..) |name, i| { + const t: ?Value = if (i < q.texts.len) q.texts[i] else null; + const shown: Value = if (t != null and t.? != .null) t.? else .{ .string = if (any) "None" else if (i == 0) "yes" else "no" }; try w.print("\n{s}: {s}", .{ name, try text(a, shown) }); + try ends.append(a, out.written().len); }; try w.writeAll("\n\nReply with yes or no only."); }, .score => { try w.writeAll("\n\n"); - for (q.keys, 0..) |name, i| try w.print("{s}{d} {s}", .{ if (i > 0) "\n" else "", i, name }); + for (q.keys, 0..) |name, i| { + try w.print("{s}{d} {s}", .{ if (i > 0) "\n" else "", i, name }); + try ends.append(a, out.written().len); + } try w.print("\n\nReply with a single digit 0-{d} only.", .{q.keys.len - 1}); }, } 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;