diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 25315ca223..653698c856 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -61,6 +61,7 @@ pub const kernels = [_]Kernel{ .{ .name = "qwen35_keep", .flags = &.{"-O3"} }, // ours: a verified window's kept path applied to its stream's state .{ .name = "qwen_vision", .flags = &.{ "-O3", "--fmad=false" } }, // ours: the Qwen3.5 vision tower's glue and attention .{ .name = "qwen35_kv8", .flags = &.{"-O3"} }, // ours: --kv-cache fp8's e4m3 attention rows for the Qwen3.5 hybrids + .{ .name = "qwen35_bf16_rows", .flags = &.{"-O3"} }, // ours: bf16 token rows and small bf16 dense projections (qwen3_5/cuda_bf16_rows.zig) .{ .name = "nemotron_norms", .flags = glue }, // ours, each with a host reference (glue_ref.zig): Nemotron's glue .{ .name = "nemotron_route", .flags = glue }, .{ .name = "nemotron_mamba", .flags = glue }, diff --git a/zig/kernels/cuda/qwen35_bf16_rows.cu b/zig/kernels/cuda/qwen35_bf16_rows.cu new file mode 100644 index 0000000000..19e21f9277 --- /dev/null +++ b/zig/kernels/cuda/qwen35_bf16_rows.cu @@ -0,0 +1,49 @@ +// Qwen3.5 bf16-weight pieces: token rows of a bf16 table and small dense projections (GDN in_proj_a / in_proj_b). + +#include +#include + +// P: device counts as qwen35.cu's; rows at or past P[2] are padding and left untouched; may be null. +__device__ __forceinline__ bool q35b_past_rows(const int* P, int m) { return P != nullptr && m >= P[2]; } + +// out[r] = table[ids[r]], D wide, a bit copy; grid: rows; 16-byte vectors when D % 8 == 0 and aligned. +extern "C" __global__ void __launch_bounds__(256) tf_q35_embed_bf16(const int* __restrict__ ids, const __nv_bfloat16* __restrict__ table, + __nv_bfloat16* __restrict__ out, int D, const int* __restrict__ P) { + if (q35b_past_rows(P, blockIdx.x)) return; + const size_t id = (size_t)ids[blockIdx.x]; + const __nv_bfloat16* src = table + id * (size_t)D; + __nv_bfloat16* dst = out + (size_t)blockIdx.x * D; + if ((D & 7) == 0 && ((reinterpret_cast(src) | reinterpret_cast(dst)) & 15) == 0) { + for (int i = threadIdx.x; i < D / 8; i += blockDim.x) reinterpret_cast(dst)[i] = reinterpret_cast(src)[i]; + } else { + for (int i = threadIdx.x; i < D; i += blockDim.x) dst[i] = src[i]; + } +} + +// y[r][j] = bf16(sum_k x[r][k] * W[j][k]) in a fixed order whatever the row count; grid (n, rows), k % 8 == 0, aligned. +extern "C" __global__ void __launch_bounds__(256) tf_q35_dense_bf16(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ W, + __nv_bfloat16* __restrict__ Y, int n, int k) { + __shared__ float red[8]; + const int j = blockIdx.x; + const size_t r = blockIdx.y; + const uint4* x = reinterpret_cast(X + r * (size_t)k); + const uint4* w = reinterpret_cast(W + (size_t)j * k); + float acc = 0.0f; + for (int c = threadIdx.x; c < k / 8; c += 256) { + const uint4 xv = x[c], wv = w[c]; + const __nv_bfloat16* xb = reinterpret_cast(&xv); + const __nv_bfloat16* wb = reinterpret_cast(&wv); +#pragma unroll + for (int e = 0; e < 8; e++) acc = __fmaf_rn(__bfloat162float(xb[e]), __bfloat162float(wb[e]), acc); + } +#pragma unroll + for (int o = 16; o > 0; o >>= 1) acc += __shfl_xor_sync(0xffffffffu, acc, o); + if ((threadIdx.x & 31) == 0) red[threadIdx.x >> 5] = acc; + __syncthreads(); + if (threadIdx.x == 0) { + float s = red[0]; +#pragma unroll + for (int i = 1; i < 8; i++) s += red[i]; + Y[r * (size_t)n + j] = __float2bfloat16_rn(s); + } +} diff --git a/zig/src/core/contract/cuda_lane_pending.zig b/zig/src/core/contract/cuda_lane_pending.zig index ebabbe6e5b..22ac3209a1 100644 --- a/zig/src/core/contract/cuda_lane_pending.zig +++ b/zig/src/core/contract/cuda_lane_pending.zig @@ -26,4 +26,6 @@ pub const records = [_]CudaRecord{ .{ .launch = .{ .name = "qwen35_draft tf_q35d_mask_rows", .symbol = "tf_q35d_mask_rows", .image = "qwen35_draft", .threads = 256 }, .usage = none }, .{ .launch = .{ .name = "qwen35_draft tf_q35d_attn_partial", .symbol = "tf_q35d_attn_partial", .image = "qwen35_draft", .threads = 32 }, .usage = none }, .{ .launch = .{ .name = "qwen35_draft tf_q35d_attn_merge", .symbol = "tf_q35d_attn_merge", .image = "qwen35_draft", .threads = 32 }, .usage = none }, + .{ .launch = .{ .name = "qwen35_bf16_rows tf_q35_embed_bf16", .symbol = "tf_q35_embed_bf16", .image = "qwen35_bf16_rows", .threads = 256 }, .usage = none }, + .{ .launch = .{ .name = "qwen35_bf16_rows tf_q35_dense_bf16", .symbol = "tf_q35_dense_bf16", .image = "qwen35_bf16_rows", .threads = 256 }, .usage = none }, }; diff --git a/zig/src/cuda/exl3.zig b/zig/src/cuda/exl3.zig index c5bb710e5a..54176e9f64 100644 --- a/zig/src/cuda/exl3.zig +++ b/zig/src/cuda/exl3.zig @@ -216,6 +216,13 @@ pub const Dense = struct { return slot.*.?.ptr; } + /// Grow every scratch buffer to what 128 rows of `w` need, so a call inside a graph capture allocates nothing. + pub fn reserve(de: *Dense, w: Weight) !void { + if (plan(w.k, w.n, 192, 8).sk > 1) _ = try de.scratch(&de.z, Linear.zBytes(w, 128)); + _ = try de.scratch(&de.counters, Linear.counterBytes(w)); + _ = try de.scratch(&de.xh, 128 * w.k * 2); + } + /// Decode rows: one call (the kernel takes 1..128), any count above in 128-row blocks. pub fn decode(de: *Dense, s: Stream, w: Weight, x: u64, x_dtype: DType, rows: usize, y: u64, y_dtype: DType) !void { if (de.implementation == .tiles) return de.tiles(s, w, x, x_dtype, rows, y, y_dtype, false); diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 80c5f50d6f..c670328d9b 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -45,6 +45,7 @@ pub const qwen35_draft: []const u8 = if (available) &Blob("fatbin_qwen35_draft") pub const qwen35_keep: []const u8 = if (available) &Blob("fatbin_qwen35_keep").bytes else &.{}; pub const qwen_vision: []const u8 = if (available) &Blob("fatbin_qwen_vision").bytes else &.{}; pub const qwen35_kv8: []const u8 = if (available) &Blob("fatbin_qwen35_kv8").bytes else &.{}; +pub const qwen35_bf16_rows: []const u8 = if (available) &Blob("fatbin_qwen35_bf16_rows").bytes else &.{}; pub const qmmf: []const u8 = if (available) &Blob("fatbin_qmmf").bytes else &.{}; pub const prompt16: []const u8 = if (available) &Blob("fatbin_prompt16").bytes else &.{}; pub const tree_attention: []const u8 = if (available) &Blob("fatbin_tree_attention").bytes else &.{}; diff --git a/zig/src/families/qwen3_5/cuda_bf16_rows.zig b/zig/src/families/qwen3_5/cuda_bf16_rows.zig new file mode 100644 index 0000000000..1f69e9e72d --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_bf16_rows.zig @@ -0,0 +1,48 @@ +//! qwen35_bf16_rows.cu: bf16 token rows and small bf16 dense projections (GDN in_proj_a and in_proj_b). +const std = @import("std"); +const cuda = @import("cuda"); + +const threads = 256; + +pub const Kernels = struct { + mod: cuda.Module, + embed_fn: cuda.Function, + dense_fn: cuda.Function, + + pub fn resolve(mod: cuda.Module) !Kernels { + return .{ .mod = mod, .embed_fn = try mod.function("tf_q35_embed_bf16"), .dense_fn = try mod.function("tf_q35_dense_bf16") }; + } + + /// Load the embedded module (owned: call unload) and resolve its functions. + pub fn load(d: *const cuda.Driver) !Kernels { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var mod = try cuda.Module.load(d, cuda.kernels.qwen35_bf16_rows); + errdefer mod.unload(); + return resolve(mod); + } + + pub fn unload(k: *Kernels) void { + k.mod.unload(); + } + + /// out[r] = table[ids[r]] (ids i32 [rows], table bf16 [vocab][hidden]); params as tf_q35_embed's, or 0. + pub fn embedRows(k: *const Kernels, s: cuda.Stream, ids: u64, table: u64, out: u64, rows: usize, hidden: usize, params: u64) !void { + if (rows == 0) return; + if (hidden == 0 or hidden > std.math.maxInt(i32) or rows > std.math.maxInt(u32)) return error.RowShape; + var a: cuda.launch.Args = .{}; + for ([_]u64{ ids, table, out }) |v| a.add(v); + a.add(@as(i32, @intCast(hidden))); + a.add(params); + try cuda.launch.launch(k.embed_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = threads } }, s, &a); + } + + /// out[rows][n] = x[rows][k] @ w[n][k]^T in bf16 with fixed-order fp32 sums (rows <= 128, k % 8 == 0, aligned). + pub fn denseRows(kn: *const Kernels, s: cuda.Stream, x: u64, w: u64, out: u64, rows: usize, n: usize, k: usize) !void { + if (rows == 0 or n == 0) return; + if (rows > 128 or k == 0 or k % 8 != 0 or k > std.math.maxInt(i32) or n > std.math.maxInt(i32) or x % 16 != 0 or w % 16 != 0) return error.RowShape; + var a: cuda.launch.Args = .{}; + for ([_]u64{ x, w, out }) |v| a.add(v); + for ([_]usize{ n, k }) |v| a.add(@as(i32, @intCast(v))); + try cuda.launch.launch(kn.dense_fn, .{ .grid = .{ .x = @intCast(n), .y = @intCast(rows), .z = 1 }, .block = .{ .x = threads } }, s, &a); + } +}; diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index 431c4ae085..f5b752e294 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -10,6 +10,7 @@ const routed = @import("cuda_routed.zig"); const hooks = @import("cuda_learn_hooks.zig"); const pocket = @import("cuda_pocket.zig"); const profile = @import("cuda_startup_profile.zig"); +const x3w = @import("cuda_exl3_weights.zig"); pub const Geometry = geometry.Geometry; pub const Cache = caches.Cache; @@ -101,7 +102,8 @@ pub const Engine = struct { embed_w: u64, embed_s: u64, embed_b: u64, - embed_bits: u8, // the embedding's rows' width: 4 or 8 + embed_bits: u8, // the embedding's rows' width: 4 or 8 (16: bf16 rows, an EXL3 checkpoint's) + kit: ?*x3w.Kit = null, // an EXL3 checkpoint's kernels (g.exl3) head: Proj, final_norm: u64, layers: []Layer, @@ -178,6 +180,7 @@ pub const Engine = struct { e.learning = null; e.mlp = null; e.taught = 0; + e.kit = null; e.d = d; e.g = g; e.geo = .{ .hid = @intCast(g.hidden), .nk = @intCast(g.k_heads), .nv = @intCast(g.v_heads), .qh = @intCast(g.q_heads), .kvh = @intCast(g.kv_heads), .nl = @intCast(g.deltas()) }; @@ -190,7 +193,11 @@ pub const Engine = struct { e.features = 0; e.stream = try cuda.Stream.init(d, true); errdefer e.stream.deinit(); - e.arena = try DeviceBuffer.alloc(d, weights.total(g) + try e.scratch(.measure)); + var ck = try core.Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + // an EXL3 checkpoint's weights are sized by a dry run of its loader (trellis widths come from the tensors) + const wbytes = if (g.exl3) try x3w.measure(e, &ck) else weights.total(g); + e.arena = try DeviceBuffer.alloc(d, wbytes + try e.scratch(.measure)); errdefer e.arena.free(); try e.arena.fill8(0, null); // touched once now: the driver commits memory on first touch, slowly const kk = cuda.kernels; @@ -217,9 +224,11 @@ pub const Engine = struct { e.fp8 = .{ .rope = try e.mods[9].function("tf_q35_rope8"), .attn_partial = try e.mods[9].function("tf_q35_attn_partial8"), .compact = try e.mods[9].function("tf_q35k_compact8") }; e.layers = try gpa.alloc(Layer, g.layers); errdefer gpa.free(e.layers); - var ck = try core.Checkpoint.openModel(gpa, io, dir); - defer ck.close(); - { + if (g.exl3) { + e.kit = try x3w.Kit.init(gpa, d, g); + errdefer e.kit.?.deinit(gpa); + try x3w.load(e, &ck); + } else { // each projection as stored, packed from here; freed before the server counts what streams can use var stage = try DeviceBuffer.alloc(d, weights.stageBytes(g)); defer stage.free(); @@ -254,6 +263,7 @@ pub const Engine = struct { pub fn deinit(e: *Engine) void { e.pockets.deinit(); + if (e.kit) |k| k.deinit(e.gpa); var it = e.graphs.valueIterator(); while (it.next()) |x| x.deinit(); e.graphs.deinit(e.gpa); @@ -356,14 +366,22 @@ pub const Engine = struct { fn project(e: *Engine, l: Proj, x: u64, xs: u64, out: u64, rows: usize, prompt: bool, params: u64, sums: bool) !void { const n_out = l.outputs(); const k = l.inputs(); + // EXL3 trellis rows and bf16 matrices need no group sums: rows go through their own kernels at any count + switch (l) { + .x3 => |q| return if (prompt and rows > 16) e.kit.?.dense.prompt(e.stream, q, x, .bf16, rows, out, .bf16) else e.kit.?.dense.decode(e.stream, q, x, .bf16, rows, out, .bf16), + .bf => |q| return e.kit.?.b16.denseRows(e.stream, x, q.w, out, rows, q.n, q.k), + else => {}, + } if (prompt and rows > 16) return switch (l) { .a4 => |q| e.affine.prompt(e.stream, x, q, out, rows), .an => |q| e.wide.prompt(e.stream, x, q, out, rows), + .x3, .bf => unreachable, }; if (sums) try e.run(e.fns.sums, .{ rows, (k / 64 + 63) / 64, 1 }, .{ 64, 1, 1 }, .{ x, xs, @as(c_int, @intCast(k)), params }); const tile: usize = switch (l) { .a4 => e.affine.verifyRows(), .an => 16, + .x3, .bf => unreachable, }; var at: usize = 0; while (at < rows) : (at += tile) { @@ -374,12 +392,14 @@ pub const Engine = struct { switch (l) { .a4 => |q| try e.affine.verify(e.stream, xa, sa, q, oa, n), .an => |q| try e.wide.decode(e.stream, xa, sa, q, oa, n), + .x3, .bf => unreachable, } } } /// Token rows of the embedding into `out` (rows, hidden) bf16, at the width its words are served. pub fn embed(e: *Engine, ids: u64, out: u64, rows: usize, params: u64) !void { + if (e.kit) |kit| return kit.b16.embedRows(e.stream, ids, e.embed_w, out, rows, e.g.hidden, params); const f = if (e.embed_bits == 8) e.fns.embed8 else e.fns.embed; try e.run(f, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ ids, e.embed_w, e.embed_s, e.embed_b, out, @as(c_int, @intCast(e.g.hidden)), params }); } diff --git a/zig/src/families/qwen3_5/cuda_exl3_admission_test.zig b/zig/src/families/qwen3_5/cuda_exl3_admission_test.zig new file mode 100644 index 0000000000..c5dd01d4bd --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_exl3_admission_test.zig @@ -0,0 +1,34 @@ +//! An EXL3 Qwen3.5-MoE config is admitted by its quant_method alone; anything else it cannot serve is refused. +const std = @import("std"); +const parse = @import("cuda_geometry.zig").parse; +const a = std.testing.allocator; +const text = + \\"text_config":{"model_type":"qwen3_5_moe_text","attention_bias":false,"attn_output_gate":true,"hidden_act":"silu", + \\ "head_dim":256,"linear_key_head_dim":128,"linear_value_head_dim":128,"linear_conv_kernel_dim":4, + \\ "hidden_size":2048,"moe_intermediate_size":512,"shared_expert_intermediate_size":512,"num_experts":256, + \\ "num_experts_per_tok":8,"num_hidden_layers":4,"vocab_size":248320,"linear_num_key_heads":16, + \\ "linear_num_value_heads":32,"num_attention_heads":16,"num_key_value_heads":2,"rms_norm_eps":1e-06, + \\ "max_position_embeddings":262144, + \\ "layer_types":["linear_attention","linear_attention","linear_attention","full_attention"], + \\ "rope_parameters":{"rope_type":"default","rope_theta":10000000,"partial_rotary_factor":0.25}} +; + +fn config(quant: []const u8) ![]u8 { + return std.fmt.allocPrint(a, "{{\"model_type\":\"qwen3_5_moe\",{s}\"tie_word_embeddings\":false,{s}}}", .{ quant, text }); +} + +test "quant_method exl3 is admitted without an MLX quantization table" { + const c = try config("\"quantization_config\":{\"quant_method\":\"exl3\",\"bits\":4.0},"); + defer a.free(c); + const g = try parse(a, c); + try std.testing.expect(g.exl3); + try std.testing.expectEqual(@as(usize, 256), g.experts); +} + +test "a config with neither an MLX table nor quant_method exl3 is refused" { + for ([_][]const u8{ "", "\"quantization_config\":{\"quant_method\":\"gptq\"}," }) |q| { + const c = try config(q); + defer a.free(c); + try std.testing.expectError(error.UnsupportedQwenQuantization, parse(a, c)); + } +} diff --git a/zig/src/families/qwen3_5/cuda_exl3_weights.zig b/zig/src/families/qwen3_5/cuda_exl3_weights.zig new file mode 100644 index 0000000000..5367a93461 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_exl3_weights.zig @@ -0,0 +1,311 @@ +//! ExLlamaV3 (EXL3, mul1 codebook) Qwen3.5-MoE checkpoints on the CUDA engine, read under their HF tensor names. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const geometry = @import("cuda_geometry.zig"); +const engine = @import("cuda_engine.zig"); +const b16 = @import("cuda_bf16_rows.zig"); + +const Engine = engine.Engine; +const Proj = engine.Proj; +const aligned = engine.aligned; +const ex = cuda.exl3_experts; + +/// The mul1 codebook marker's value. +const mul1_magic: u32 = 0x83DCD12D; + +/// The EXL3 kernels beside the engine's own: dense rows, the grouped experts and the bf16 rows. +pub const Kit = struct { + dense: cuda.exl3.Dense, + experts: ex.Experts, + grouped: cuda.exl3_grouped.Grouped, + b16: b16.Kernels, + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, g: geometry.Geometry) !*Kit { + const k = try gpa.create(Kit); + errdefer gpa.destroy(k); + k.dense = try cuda.exl3.Dense.load(d); + errdefer k.dense.unload(); + k.experts = try ex.Experts.load(d); + errdefer k.experts.unload(); + k.b16 = try b16.Kernels.load(d); + errdefer k.b16.unload(); + k.grouped = .{ .d = d, .kernels = &k.experts }; + // the grouped experts' scratch grows on first use: allocate it before any graph capture, for the widest call + const pairs = engine.max_rows * g.slots(); + const items = cuda.grouped.maxItems(pairs, g.experts + 1, 16); + const up = try cuda.exl3_grouped.Grouped.bytes(g.hidden, g.intermediate, items, try ex.plan(g.hidden, g.intermediate, true), 2); + const down = try cuda.exl3_grouped.Grouped.bytes(g.intermediate, g.hidden, items, try ex.plan(g.intermediate, g.hidden, false), 1); + k.grouped.storage = try cuda.DeviceBuffer.alloc(d, @max(up, down)); + return k; + } + + pub fn deinit(k: *Kit, gpa: std.mem.Allocator) void { + k.grouped.deinit(); + k.b16.unload(); + k.experts.unload(); + k.dense.unload(); + gpa.destroy(k); + } +}; + +/// Where tensors go: the engine's arena, or (measure) just counted with the same alignment. +const Sink = struct { + e: *Engine, + measure: bool, + at: usize = 0, + + fn alloc(s: *Sink, bytes: usize) !u64 { + if (!s.measure) return s.e.alloc(bytes); + s.at = aligned(s.at) + bytes; + return 0; + } + + fn put(s: *Sink, bytes: []const u8) !u64 { + const p = try s.alloc(bytes.len); + if (!s.measure) try s.e.arena.upload(p - s.e.arena.ptr, bytes); + return p; + } +}; + +fn name(buf: []u8, comptime fmt: []const u8, args: anytype) ![]const u8 { + return std.fmt.bufPrint(buf, fmt, args); +} + +fn bf16Of(x: f32) u16 { + const u: u32 = @bitCast(x); + return @intCast((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} + +/// f16 values (as stored) to bf16, round to nearest even. +fn f16Rows(gpa: std.mem.Allocator, bytes: []const u8) ![]u16 { + const out = try gpa.alloc(u16, bytes.len / 2); + for (out, 0..) |*o, i| { + const h: f16 = @bitCast(std.mem.readInt(u16, bytes[2 * i ..][0..2], .little)); + o.* = bf16Of(@floatCast(h)); + } + return out; +} + +/// Zero-centred RMSNorm weights (HF: y = x * (1 + w)) as the engine's norm kernels take them (y = x * w): bf16(1 + w). +fn shifted(gpa: std.mem.Allocator, bytes: []const u8) ![]u16 { + const out = try gpa.alloc(u16, bytes.len / 2); + for (out, 0..) |*o, i| { + const w: f32 = @bitCast(@as(u32, std.mem.readInt(u16, bytes[2 * i ..][0..2], .little)) << 16); + o.* = bf16Of(1.0 + w); + } + return out; +} + +fn plain(s: *Sink, ck: *core.Checkpoint, comptime fmt: []const u8, args: anytype, dtype: core.checkpoint.DType, shape: []const usize) !u64 { + var buf: [200]u8 = undefined; + return s.put((try ck.expect(try name(&buf, fmt, args), dtype, shape)).bytes); +} + +/// A norm weight, shifted by one (see `shifted`). +fn norm1(s: *Sink, ck: *core.Checkpoint, comptime fmt: []const u8, args: anytype, n: usize) !u64 { + var buf: [200]u8 = undefined; + const t = try ck.expect(try name(&buf, fmt, args), .bf16, &.{n}); + const w = try shifted(s.e.gpa, t.bytes); + defer s.e.gpa.free(w); + return s.put(std.mem.sliceAsBytes(w)); +} + +/// A vector the kernels read as f32, stored as bf16 (A_log). +fn f32vec(s: *Sink, ck: *core.Checkpoint, comptime fmt: []const u8, args: anytype, n: usize) !u64 { + var buf: [200]u8 = undefined; + const t = try ck.get(try name(&buf, fmt, args)); + if (!t.is(.bf16, &.{n})) return error.UnexpectedTensor; + const words = try s.e.gpa.alloc(u32, n); + defer s.e.gpa.free(words); + for (words, 0..) |*w, i| w.* = @as(u32, std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little)) << 16; + return s.put(std.mem.sliceAsBytes(words)); +} + +/// One EXL3 group's tensors, checked: trellis [k/16][n/16][16 * bits] i16, the mul1 marker, fp16 suh [k] and svh [n]. +const Group = struct { trellis: core.checkpoint.Tensor, suh: core.checkpoint.Tensor, svh: core.checkpoint.Tensor, halves: u32 }; + +fn group(ck: *core.Checkpoint, base: []const u8, n: usize, k: usize) !Group { + var buf: [240]u8 = undefined; + const t = try ck.get(try name(&buf, "{s}.trellis", .{base})); + if (t.dtype != .i16 or t.rank != 3 or t.dim(0) != k / 16 or t.dim(1) != n / 16 or t.dim(2) % 8 != 0) { + std.log.err("{s}: unexpected trellis {any}", .{ base, t.shape[0..t.rank] }); + return error.UnexpectedTensor; + } + const marker = try ck.get(try name(&buf, "{s}.mul1", .{base})); + if (marker.bytes.len != 4 or std.mem.readInt(u32, marker.bytes[0..4], .little) != mul1_magic) return error.UnsupportedQwenQuantization; + return .{ + .trellis = t, + .suh = try ck.expect(try name(&buf, "{s}.suh", .{base}), .f16, &.{k}), + .svh = try ck.expect(try name(&buf, "{s}.svh", .{base}), .f16, &.{n}), + .halves = @intCast(t.dim(2) / 8), + }; +} + +/// An n-by-k EXL3 projection, its trellis kept as stored (Dense reads the stored layout). +fn dense(s: *Sink, kit: ?*Kit, ck: *core.Checkpoint, prefix: []const u8, mod: []const u8, n: usize, k: usize) !Proj { + var buf: [200]u8 = undefined; + const gr = try group(ck, try name(&buf, "{s}{s}", .{ prefix, mod }), n, k); + const st = cuda.exl3.strides(k, n, gr.halves, false); + const w: cuda.exl3.Weight = .{ + .trellis = try s.put(gr.trellis.bytes), + .suh = try s.put(gr.suh.bytes), + .svh = try s.put(gr.svh.bytes), + .k = @intCast(k), + .n = @intCast(n), + .halves = gr.halves, + .codebook = .mul1, + .stride_k = st.k, + .stride_nb = st.nb, + }; + if (kit) |x| try x.dense.reserve(w); + return .{ .x3 = w }; +} + +/// A small f16 matrix [n][k] (in_proj_a / in_proj_b) as bf16 for denseRows. +fn smallBf16(s: *Sink, ck: *core.Checkpoint, prefix: []const u8, mod: []const u8, n: usize, k: usize) !Proj { + var buf: [200]u8 = undefined; + const t = try ck.expect(try name(&buf, "{s}{s}.weight", .{ prefix, mod }), .f16, &.{ n, k }); + const rows = try f16Rows(s.e.gpa, t.bytes); + defer s.e.gpa.free(rows); + return .{ .bf = .{ .w = try s.put(std.mem.sliceAsBytes(rows)), .n = n, .k = k } }; +} + +/// The router as bf16 [experts + 1][hidden]: the gate's rows, then shared_expert_gate's (both stored f16). +fn router(s: *Sink, ck: *core.Checkpoint, prefix: []const u8) !u64 { + const g = s.e.g; + const rows = try s.e.gpa.alloc(u16, (g.experts + 1) * g.hidden); + defer s.e.gpa.free(rows); + var buf: [200]u8 = undefined; + const a = try ck.expect(try name(&buf, "{s}mlp.gate.weight", .{prefix}), .f16, &.{ g.experts, g.hidden }); + const b = try ck.expect(try name(&buf, "{s}mlp.shared_expert_gate.weight", .{prefix}), .f16, &.{ 1, g.hidden }); + for (rows[0 .. g.experts * g.hidden], 0..) |*o, i| o.* = bf16Of(@floatCast(@as(f16, @bitCast(std.mem.readInt(u16, a.bytes[2 * i ..][0..2], .little))))); + for (rows[g.experts * g.hidden ..], 0..) |*o, i| o.* = bf16Of(@floatCast(@as(f16, @bitCast(std.mem.readInt(u16, b.bytes[2 * i ..][0..2], .little))))); + return s.put(std.mem.sliceAsBytes(rows)); +} + +/// A layer's experts for the grouped EXL3 kernels: trellis pointers and widths per expert, the shared one last. +fn experts(s: *Sink, ck: *core.Checkpoint, p: []const u8) !cuda.experts.Layer { + const g = s.e.g; + const gpa = s.e.gpa; + const count = g.experts + 1; + const mods = [3][]const u8{ "gate_proj", "up_proj", "down_proj" }; + var ptrs: [3]u64 = undefined; + var k2: [3]u64 = undefined; + var suh: [3]u64 = undefined; + var svh: [3]u64 = undefined; + var buf: [200]u8 = undefined; + for (mods, 0..) |mod, m| { + const k = if (m == 2) g.intermediate else g.hidden; + const n = if (m == 2) g.hidden else g.intermediate; + const tp = try gpa.alloc(u64, count); + defer gpa.free(tp); + const widths = try gpa.alloc(i32, count); + defer gpa.free(widths); + const su = try gpa.alloc(u8, count * k * 2); + defer gpa.free(su); + const sv = try gpa.alloc(u8, count * n * 2); + defer gpa.free(sv); + for (0..count) |e| { + const base = if (e < g.experts) try name(&buf, "{s}mlp.experts.{d}.{s}", .{ p, e, mod }) else try name(&buf, "{s}mlp.shared_expert.{s}", .{ p, mod }); + const gr = try group(ck, base, n, k); + tp[e] = try s.put(gr.trellis.bytes); + widths[e] = @intCast(gr.halves); + @memcpy(su[e * k * 2 ..][0 .. k * 2], gr.suh.bytes); + @memcpy(sv[e * n * 2 ..][0 .. n * 2], gr.svh.bytes); + } + ptrs[m] = try s.put(std.mem.sliceAsBytes(tp)); + k2[m] = try s.put(std.mem.sliceAsBytes(widths)); + suh[m] = try s.put(su); + svh[m] = try s.put(sv); + } + const lay: ex.Layer = .{ + .gate = ptrs[0], + .up = ptrs[1], + .down = ptrs[2], + .gate_k2 = k2[0], + .up_k2 = k2[1], + .down_k2 = k2[2], + .suh_g = suh[0], + .suh_u = suh[1], + .svh_g = svh[0], + .svh_u = svh[1], + .suh_d = suh[2], + .svh_d = svh[2], + .experts = count, + .dims = g.hidden, + .width = g.intermediate, + .codebook = .mul1, + .limit = 0, // Qwen's SwiGLU is unclamped + .act_mode = 1, // fp32 silu(gate) * up + }; + return .{ .format = .exl3, .exl3 = lay, .up = 0, .down = 0, .width = g.intermediate, .dims = g.hidden, .experts = count, .swiglu = true }; +} + +/// The whole checkpoint through `s`; the engine's fields are set only in the real pass. +fn all(s: *Sink, ck: *core.Checkpoint) !void { + const e = s.e; + const g = e.g; + const kit = e.kit; + const m = "model.language_model."; + // embeddings stay bf16 rows (embedRows); the head is an EXL3 projection + const embed = try plain(s, ck, "{s}embed_tokens.weight", .{m}, .bf16, &.{ g.vocab, g.hidden }); + const head = try dense(s, kit, ck, "", "lm_head", g.vocab, g.hidden); + const final = try norm1(s, ck, "{s}norm.weight", .{m}, g.hidden); + if (!s.measure) { + e.embed_w = embed; + e.embed_s = 0; + e.embed_b = 0; + e.embed_bits = 16; + e.head = head; + e.final_norm = final; + } + const conv = g.convWidth(); + var scratch_layer: engine.Layer = undefined; + for (0..g.layers) |i| { + var pbuf: [96]u8 = undefined; + const p = try std.fmt.bufPrint(&pbuf, "{s}layers.{d}.", .{ m, i }); + const l = if (s.measure) &scratch_layer else &e.layers[i]; + l.input_norm = try norm1(s, ck, "{s}input_layernorm.weight", .{p}, g.hidden); + l.post_norm = try norm1(s, ck, "{s}post_attention_layernorm.weight", .{p}, g.hidden); + if (g.linear[i]) { + l.mixer = .{ + .delta = .{ + .qkv = try dense(s, kit, ck, p, "linear_attn.in_proj_qkv", conv, g.hidden), + .z = try dense(s, kit, ck, p, "linear_attn.in_proj_z", g.valueWidth(), g.hidden), + .a = try smallBf16(s, ck, p, "linear_attn.in_proj_a", g.v_heads, g.hidden), + .b = try smallBf16(s, ck, p, "linear_attn.in_proj_b", g.v_heads, g.hidden), + .out = try dense(s, kit, ck, p, "linear_attn.out_proj", g.hidden, g.valueWidth()), + // HF's conv1d is [channels][1][taps]: the same bytes as the engine's [channels][taps] + .conv = try plain(s, ck, "{s}linear_attn.conv1d.weight", .{p}, .bf16, &.{ conv, 1, geometry.taps }), + .norm = try plain(s, ck, "{s}linear_attn.norm.weight", .{p}, .bf16, &.{geometry.dk}), // gated norm: unshifted + .a_log = try f32vec(s, ck, "{s}linear_attn.A_log", .{p}, g.v_heads), + .dt_bias = try plain(s, ck, "{s}linear_attn.dt_bias", .{p}, .bf16, &.{g.v_heads}), + }, + }; + } else { + l.mixer = .{ .attention = .{ + .q = try dense(s, kit, ck, p, "self_attn.q_proj", 2 * g.attnWidth(), g.hidden), + .k = try dense(s, kit, ck, p, "self_attn.k_proj", g.kv_heads * geometry.head_dim, g.hidden), + .v = try dense(s, kit, ck, p, "self_attn.v_proj", g.kv_heads * geometry.head_dim, g.hidden), + .o = try dense(s, kit, ck, p, "self_attn.o_proj", g.hidden, g.attnWidth()), + .q_norm = try norm1(s, ck, "{s}self_attn.q_norm.weight", .{p}, geometry.head_dim), + .k_norm = try norm1(s, ck, "{s}self_attn.k_norm.weight", .{p}, geometry.head_dim), + } }; + } + l.mlp = .{ .moe = .{ .router = try router(s, ck, p), .experts = try experts(s, ck, p) } }; + } +} + +/// The arena bytes the weights take (the scratch is added by the caller). +pub fn measure(e: *Engine, ck: *core.Checkpoint) !usize { + var s: Sink = .{ .e = e, .measure = true }; + try all(&s, ck); + return aligned(s.at) + 256; +} + +pub fn load(e: *Engine, ck: *core.Checkpoint) !void { + var s: Sink = .{ .e = e, .measure = false }; + try all(&s, ck); + try e.stream.synchronize(); +} diff --git a/zig/src/families/qwen3_5/cuda_geometry.zig b/zig/src/families/qwen3_5/cuda_geometry.zig index 740b71adff..e300ab3e61 100644 --- a/zig/src/families/qwen3_5/cuda_geometry.zig +++ b/zig/src/families/qwen3_5/cuda_geometry.zig @@ -68,6 +68,8 @@ pub const Geometry = struct { top_k: usize = 0, /// The routed experts' code width: 4, or 3 where every layer's switch_mlp overrides say so. expert_bits: u8 = 4, + /// An EXL3 checkpoint (quant_method "exl3", HF tensor names): trellis projections, bf16 or f16 the rest. + exl3: bool = false, pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Geometry { const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); @@ -177,6 +179,14 @@ fn expertBits(q: std.json.ObjectMap, layers: usize) !struct { bits: u8, count: u return .{ .bits = @intCast(bits orelse 4), .count = count }; } +/// Whether config.json's quantization_config says quant_method "exl3" (Mia-AiLab's HF-layout checkpoints). +fn isExl3(root: std.json.ObjectMap) bool { + const q = root.get("quantization_config") orelse return false; + if (q != .object) return false; + const m = q.object.get("quant_method") orelse return false; + return m == .string and std.mem.eql(u8, m.string, "exl3"); +} + fn object(v: std.json.Value) !std.json.ObjectMap { return if (v == .object) v.object else error.BadQwenConfig; } @@ -229,7 +239,10 @@ pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) !Geometry { const rope = try object(t.get("rope_parameters") orelse return error.UnsupportedQwenRotary); try is(rope, if (rope.get("rope_type") != null) "rope_type" else "type", []const u8, "default"); // "type": the older key if (try number(rope, "rope_theta") != 10000000 or try number(rope, "partial_rotary_factor") != 0.25) return error.UnsupportedQwenRotary; - const q = try object(root.get("quantization") orelse return error.UnsupportedQwenQuantization); + // EXL3: no MLX table; per-tensor widths come from the trellis shapes at load (cuda_exl3_weights.zig). + const exl3 = isExl3(root); + var q: std.json.ObjectMap = undefined; + if (!exl3) q = try object(root.get("quantization") orelse return error.UnsupportedQwenQuantization); var g: Geometry = .{ .hidden = try integer(t, "hidden_size"), .intermediate = try integer(t, if (moe) "moe_intermediate_size" else "intermediate_size"), @@ -247,6 +260,7 @@ pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) !Geometry { }, .context = try integer(t, "max_position_embeddings"), .linear = @splat(false), + .exl3 = exl3, }; if (t.get("tensorfold_learned_neurons")) |pocket| { if (pocket != .integer or pocket.integer < 0 or pocket.integer > 128 or @mod(pocket.integer, 64) != 0) @@ -262,9 +276,11 @@ pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) !Geometry { // the shared expert runs in the routed experts' table, so it needs their width if (try integer(t, "shared_expert_intermediate_size") != g.intermediate) return error.UnsupportedQwenGeometry; if (g.experts > 256 or g.top_k > 31 or g.top_k >= g.experts) return error.UnsupportedQwenGeometry; - const w = try expertBits(q, g.layers); - if (w.count != 0 and w.count != 3 * g.layers) return error.UnsupportedQwenQuantization; - g.expert_bits = w.bits; + if (!exl3) { + const w = try expertBits(q, g.layers); + if (w.count != 0 and w.count != 3 * g.layers) return error.UnsupportedQwenQuantization; + g.expert_bits = w.bits; + } } if (g.v_heads % g.k_heads != 0 or g.q_heads % g.kv_heads != 0 or g.q_heads > max_query_heads) return error.UnsupportedQwenGeometry; const types = t.get("layer_types") orelse return error.UnsupportedQwenGeometry; @@ -273,6 +289,10 @@ pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) !Geometry { if (kind != .string) return error.UnsupportedQwenGeometry; if (std.mem.eql(u8, kind.string, "linear_attention")) g.linear[i] = true else if (!std.mem.eql(u8, kind.string, "full_attention")) return error.UnsupportedQwenGeometry; } + if (exl3) { + if (!moe or g.pocket != 0) return error.UnsupportedQwenQuantization; + return g; + } g.bits = try widths(gpa, root, q, g); if (g.pocket > 0) { if (g.moe()) return error.UnsupportedLearningFormat; @@ -420,4 +440,5 @@ test "fp8 caches take e4m3 rows and a scale each: 520 bytes a position and head test { _ = @import("cuda_quant_admission_test.zig"); + _ = @import("cuda_exl3_admission_test.zig"); } diff --git a/zig/src/families/qwen3_5/cuda_layers.zig b/zig/src/families/qwen3_5/cuda_layers.zig index 42b779704a..0e7d24b8dd 100644 --- a/zig/src/families/qwen3_5/cuda_layers.zig +++ b/zig/src/families/qwen3_5/cuda_layers.zig @@ -8,6 +8,10 @@ const AffineN = cuda.qlinear.AffineN; pub const Proj = union(enum) { a4: Affine4, an: AffineN, + /// EXL3 trellis projection (cuda.exl3.Dense, bf16 in and out); only in an EXL3 checkpoint's engine. + x3: cuda.exl3.Weight, + /// A small unquantized bf16 matrix [n][k] (the GDN in_proj_a / in_proj_b), cuda_bf16.zig's denseRows. + bf: Bf16, pub fn outputs(p: Proj) usize { return switch (p) { @@ -22,6 +26,8 @@ pub const Proj = union(enum) { } }; +pub const Bf16 = struct { w: u64, n: usize, k: usize }; + pub const Delta = struct { qkv: Proj, z: Proj, a: Proj, b: Proj, out: Proj, conv: u64, norm: u64, a_log: u64, dt_bias: u64 }; pub const Attention = struct { q: Proj, k: Proj, v: Proj, o: Proj, q_norm: u64, k_norm: u64 }; /// A dense SwiGLU MLP, or routed experts (the shared one last in their tables) behind a router of bf16 rows. diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index d783449cb0..7cf0db5cab 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -106,6 +106,7 @@ fn reserve(p: *anyopaque, streams: u32) anyerror!void { pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []const u8, bits: u8) !struct { drafter: lanes.drafter.Drafter, ctx: *anyopaque, close: *const fn (*anyopaque) void, reserve: *const fn (*anyopaque, u32) anyerror!void, prepare: struct { ctx: *anyopaque, run: *const fn (*anyopaque) anyerror!lanes.drafter.Facts, bytes: usize } } { if (bits != 4) return error.DrafterBitsUnsupported; const l: *Lanes = @ptrCast(@alignCast(p)); + if (l.e.g.exl3) return error.Exl3DrafterUnsupported; if (l.caches.count() != 0 or l.pool.items.len != 0) return error.TapsAfterStreams; const d = try Draft.init(gpa, io, l.e, dir); errdefer d.deinit(); @@ -120,6 +121,7 @@ pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []con pub fn openMtp(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []const u8, bits: u8) !struct { drafter: lanes.drafter.Drafter, ctx: *anyopaque, close: *const fn (*anyopaque) void, reserve: *const fn (*anyopaque, u32) anyerror!void } { if (bits != 4) return error.DrafterBitsUnsupported; const l: *Lanes = @ptrCast(@alignCast(p)); + if (l.e.g.exl3) return error.Exl3DrafterUnsupported; if (l.caches.count() != 0 or l.pool.items.len != 0) return error.TapsAfterStreams; const d = try Mtp.init(gpa, io, l.e, dir, l.capacity); errdefer d.deinit(); diff --git a/zig/src/families/qwen3_5/cuda_routed.zig b/zig/src/families/qwen3_5/cuda_routed.zig index 89788cfe26..70abdcff02 100644 --- a/zig/src/families/qwen3_5/cuda_routed.zig +++ b/zig/src/families/qwen3_5/cuda_routed.zig @@ -44,7 +44,7 @@ pub fn rows(e: *E.Engine, router: u64, layer: cuda.experts.Layer, x: u64, out: u const tile: usize = if (wide) 64 else 16; const items = cuda.grouped.maxItems(pairs, count, tile); const plan: cuda.grouped.Plan = .{ .members = s.members, .items = s.items, .counts = s.counts, .rank = s.rank, .hist = s.hist }; - const grouped: cuda.experts.Grouped = .{ .affine4 = &e.experts, .affine3 = &e.experts3 }; + const grouped: cuda.experts.Grouped = .{ .affine4 = &e.experts, .affine3 = &e.experts3, .exl3 = if (e.kit) |k| &k.grouped else null }; try e.moe.route(e.stream, x, g.hidden, total, router, s.logits, g.hidden, count); try e.moe.pick(e.stream, s.logits, s.picks, s.wts, total, g.experts, g.top_k); try e.plan.route(e.stream, s.picks, pairs, count, tile, plan);