From bbfd4ac7ef4dbf0313484533df61dce810b1c51e Mon Sep 17 00:00:00 2001 From: MiaAI-Lab <83042094+MiaAI-Lab@users.noreply.github.com> Date: Wed, 21 Oct 2026 12:15:44 +0300 Subject: [PATCH 1/3] qwen3_5 cuda: bf16 token-row and small bf16 dense kernels Adds qwen35_bf16_rows.cu and its launcher (cuda_bf16_rows.zig). tf_q35_embed_bf16 copies the rows of a bf16 embedding table, with the engine's padding-row convention (rows at or past params[2] are left untouched). tf_q35_dense_bf16 is a bf16-in, bf16-out dense projection for small matrices such as the GDN in_proj_a and in_proj_b: one block reduces one output in a fixed order (fp32 fused multiply-adds per thread, a fixed shuffle tree, then the eight warp sums in order), so a row's bits do not depend on how many rows share the launch. Nothing calls them yet; the EXL3 loader in a later commit of this series does. Both launches are listed as pending contract records beside the other Qwen3.5 launches. --- zig/build/cuda_kernels.zig | 1 + zig/kernels/cuda/qwen35_bf16_rows.cu | 49 +++++++++++++++++++++ zig/src/core/contract/cuda_lane_pending.zig | 2 + zig/src/cuda/kernels.zig | 1 + zig/src/families/qwen3_5/cuda_bf16_rows.zig | 48 ++++++++++++++++++++ 5 files changed, 101 insertions(+) create mode 100644 zig/kernels/cuda/qwen35_bf16_rows.cu create mode 100644 zig/src/families/qwen3_5/cuda_bf16_rows.zig 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/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); + } +}; From 24bb1a5f5a79efd9d568b556a6a85114a013c3a1 Mon Sep 17 00:00:00 2001 From: MiaAI-Lab <83042094+MiaAI-Lab@users.noreply.github.com> Date: Wed, 21 Oct 2026 12:15:44 +0300 Subject: [PATCH 2/3] exl3 dense: reserve the scratch a weight needs before graph capture Dense.reserve grows the split-K sums, the counters and the fp16 input staging to what 128 rows of a given weight need. The buffers otherwise grow on first use, which allocates inside a CUDA graph capture when a weight is first called there. The EXL3 loader in the next commit calls it for every projection it loads. --- zig/src/cuda/exl3.zig | 7 +++++++ 1 file changed, 7 insertions(+) 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); From f8dc42c54162019cb65bab2e0f86ff118212a9f2 Mon Sep 17 00:00:00 2001 From: MiaAI-Lab <83042094+MiaAI-Lab@users.noreply.github.com> Date: Wed, 21 Oct 2026 12:15:44 +0300 Subject: [PATCH 3/3] qwen3_5 cuda: serve ExLlamaV3 (EXL3) Qwen3.5-MoE checkpoints Reads a Hugging Face-layout Qwen3.5-MoE checkpoint whose config.json says quantization_config.quant_method "exl3" (ExLlamaV3's trellis format, mul1 codebook, with each tensor's width taken from its trellis shape). Projections and every routed and shared expert stay in their stored layout and run on the existing cuda.exl3 Dense and exl3_experts grouped kernels; embeddings stay bf16, the f16 router and GDN in_proj_a/in_proj_b become bf16 at load, and the zero-centred RMSNorm weights are shifted by one into the engine's form. - cuda_geometry.zig accepts quant_method "exl3" in place of the MLX `quantization` table. Only the routed (MoE) model without a learning pocket is admitted; per-tensor widths come from the trellis shapes at load. - cuda_exl3_weights.zig loads the checkpoint twice over one code path: a dry run sizes the arena, then the real pass uploads. The kit holds the EXL3 dense and grouped-expert kernels and the bf16 rows; their scratch is reserved before any graph capture. - cuda_engine.zig routes `.x3` and `.bf` projections to those kernels at any row count and embeds from the bf16 table. - An EXL3 checkpoint refuses --drafter for now (Exl3DrafterUnsupported): its MTP layer and DFlash2 drafter are not read by this change. The format and the kernel design are turboderp's ExLlamaV3 (https://github.com/turboderp-org/exllamav3); this repository already carries its Dense and grouped-expert kernels in cuda/exl3.zig. The checkpoint this was exercised on is Mia-AiLab/Ornith-1.5-122B-EXL3-EAQv1-4-bpw. Prompt and decode speed use the existing generic EXL3 kernels and are not tuned for this model. --- zig/src/families/qwen3_5/cuda_engine.zig | 30 +- .../qwen3_5/cuda_exl3_admission_test.zig | 34 ++ .../families/qwen3_5/cuda_exl3_weights.zig | 311 ++++++++++++++++++ zig/src/families/qwen3_5/cuda_geometry.zig | 29 +- zig/src/families/qwen3_5/cuda_layers.zig | 6 + zig/src/families/qwen3_5/cuda_native.zig | 2 + zig/src/families/qwen3_5/cuda_routed.zig | 2 +- 7 files changed, 404 insertions(+), 10 deletions(-) create mode 100644 zig/src/families/qwen3_5/cuda_exl3_admission_test.zig create mode 100644 zig/src/families/qwen3_5/cuda_exl3_weights.zig 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);