qwen3_5 cuda: serve ExLlamaV3 (EXL3) Qwen3.5-MoE checkpoints

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 },

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 <cuda_bf16.h>
+#include <stdint.h>
+
+// 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<uintptr_t>(src) | reinterpret_cast<uintptr_t>(dst)) & 15) == 0) {
+        for (int i = threadIdx.x; i < D / 8; i += blockDim.x) reinterpret_cast<uint4*>(dst)[i] = reinterpret_cast<const uint4*>(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<const uint4*>(X + r * (size_t)k);
+    const uint4* w = reinterpret_cast<const uint4*>(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<const __nv_bfloat16*>(&xv);
+        const __nv_bfloat16* wb = reinterpret_cast<const __nv_bfloat16*>(&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);
+    }
+}

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 },
 };

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);

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 &.{};

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);
+    }
+};

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 });
     }

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));
+    }
+}