From 9a14da6a34b3b4aa66a3629e58c45196bf9c70f9 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Tue, 13 Oct 2026 06:37:55 +0900 Subject: [PATCH 1/4] cuda: the NVFP4 lane matmul and prompt GEMM write fp32 too (qlinear decodeF32 and promptF32) A tensor-parallel family sums its ranks' partial rows before rounding, so it needs a projection's fp32 sums as the kernel holds them. qmmf.cu and prompt16.cu already take F32 as a template argument; this instantiates FP4's fp32 tiles (and reduce_kernel for the split path), loads them beside the bf16 ones, and adds Lane.matmulOut, Prompt.matmulOut and qlinear's decodeF32 and promptF32. Other formats refuse fp32 out with EpilogueNotBuilt, as the experts do. matmul, decode and prompt keep their signatures and bf16 bits. --- zig/kernels/cuda/prompt16.cu | 1 + zig/kernels/cuda/qmmf.cu | 8 +++++ zig/src/cuda/nvfp4.zig | 19 ++++++++--- zig/src/cuda/qlinear.zig | 33 ++++++++++++++++-- zig/src/cuda/qmmf.zig | 55 +++++++++++++++++++----------- zig/tests/cuda/quant_ref_tests.zig | 21 ++++++++++++ 6 files changed, 110 insertions(+), 27 deletions(-) diff --git a/zig/kernels/cuda/prompt16.cu b/zig/kernels/cuda/prompt16.cu index c2c6c2f3ab..01749ba86d 100644 --- a/zig/kernels/cuda/prompt16.cu +++ b/zig/kernels/cuda/prompt16.cu @@ -192,5 +192,6 @@ __global__ void __launch_bounds__(WM * WN * 32) prompt_kernel( // The instance the Zig launcher loads by name (PROMPT_TILE 4 in linear.py). template __global__ void prompt_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, int, int, int, int, int, int); +template __global__ void prompt_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, int, int, int, int, int, int); template __global__ void prompt_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, int, int, int, int, int, int); } diff --git a/zig/kernels/cuda/qmmf.cu b/zig/kernels/cuda/qmmf.cu index 4703273410..57a7c1e582 100644 --- a/zig/kernels/cuda/qmmf.cu +++ b/zig/kernels/cuda/qmmf.cu @@ -308,6 +308,13 @@ template __global__ void qmmf_kernel(c template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); @@ -323,4 +330,5 @@ template __global__ void qmmf_kernel template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void reduce_kernel(const float* __restrict__, void* __restrict__, long long, int, float); +template __global__ void reduce_kernel(const float* __restrict__, void* __restrict__, long long, int, float); } diff --git a/zig/src/cuda/nvfp4.zig b/zig/src/cuda/nvfp4.zig index d9502895d3..e0d3c838f4 100644 --- a/zig/src/cuda/nvfp4.zig +++ b/zig/src/cuda/nvfp4.zig @@ -64,12 +64,13 @@ pub const Prompt = struct { mod: Module, fp4: Function, fp8: Function, + fp4_f32: Function, // FP4 writing the fp32 sums unrounded const bm = 128; const bn = 128; const threads = 128; - const name = "_ZN11tf_prompt1613prompt_kernelILi{d}ELi128ELi128ELi2ELi2ELi2ELb0EEEvPK13__nv_bfloat16PKhS5_fPviiiiii"; - pub const symbols = [2][:0]const u8{ std.fmt.comptimePrint(name, .{0}), std.fmt.comptimePrint(name, .{1}) }; + const name = "_ZN11tf_prompt1613prompt_kernelILi{d}ELi128ELi128ELi2ELi2ELi2ELb{d}EEEvPK13__nv_bfloat16PKhS5_fPviiiiii"; + pub const symbols = [3][:0]const u8{ std.fmt.comptimePrint(name, .{ 0, 0 }), std.fmt.comptimePrint(name, .{ 1, 0 }), std.fmt.comptimePrint(name, .{ 0, 1 }) }; /// prompt16.cu Tile::SMEM. pub fn sharedBytes(mode: qmmf.Mode) u32 { @@ -86,7 +87,9 @@ pub const Prompt = struct { try fp4.allowDynamicShared(sharedBytes(.fp4)); const fp8 = try mod.function(symbols[1]); try fp8.allowDynamicShared(sharedBytes(.fp8)); - return .{ .mod = mod, .fp4 = fp4, .fp8 = fp8 }; + const fp4_f32 = try mod.function(symbols[2]); + try fp4_f32.allowDynamicShared(sharedBytes(.fp4)); + return .{ .mod = mod, .fp4 = fp4, .fp8 = fp8, .fp4_f32 = fp4_f32 }; } pub fn unload(p: *Prompt) void { @@ -95,9 +98,14 @@ pub const Prompt = struct { /// y [m, n] bf16 = prompt rows x [m, k] bf16 (rows `ldx` apart) times an FP4 or FP8 weight (linear._prompt). pub fn matmul(p: Prompt, s: Stream, x: u64, ldx: usize, m: usize, w: qmmf.Weight, y: u64) !void { + return p.matmulOut(s, x, ldx, m, w, y, .bf16); + } + + /// matmul() writing y as `out` says: bf16, or fp32 (FP4 only). + pub fn matmulOut(p: Prompt, s: Stream, x: u64, ldx: usize, m: usize, w: qmmf.Weight, y: u64, out: qmmf.Out) !void { const f = switch (w.mode) { - .fp4 => p.fp4, - .fp8 => p.fp8, + .fp4 => if (out == .f32) p.fp4_f32 else p.fp4, + .fp8 => if (out == .f32) return error.EpilogueNotBuilt else p.fp8, else => return error.FormatMismatch, }; const rows_t = (m + bm - 1) / bm; @@ -126,6 +134,7 @@ test "the prompt tile's shared bytes and instance names" { try std.testing.expectEqual(@as(u32, 41984), Prompt.sharedBytes(.fp4)); try std.testing.expectEqual(@as(u32, 49152), Prompt.sharedBytes(.fp8)); // 2 x (16384 + 8192): no scale tile try std.testing.expectEqualStrings("_ZN11tf_prompt1613prompt_kernelILi0ELi128ELi128ELi2ELi2ELi2ELb0EEEvPK13__nv_bfloat16PKhS5_fPviiiiii", Prompt.symbols[0]); + try std.testing.expectEqualStrings("_ZN11tf_prompt1613prompt_kernelILi0ELi128ELi128ELi2ELi2ELi2ELb1EEEvPK13__nv_bfloat16PKhS5_fPviiiiii", Prompt.symbols[2]); } test "word tiling reads each nibble where qmm.pack puts it" { diff --git a/zig/src/cuda/qlinear.zig b/zig/src/cuda/qlinear.zig index 6d6a56f5a3..ff70a0d88a 100644 --- a/zig/src/cuda/qlinear.zig +++ b/zig/src/cuda/qlinear.zig @@ -91,16 +91,42 @@ pub const Linear = struct { } } + /// Decode rows as decode(), out (rows, n) fp32 unrounded: a tensor-parallel rank's partial rows (nvfp4 only). + pub fn decodeF32(l: Linear, s: Stream, w: Weight, in: Rows, out: u64, rows: usize) !void { + const q = switch (w) { + .nvfp4 => |q| q, + else => return error.EpilogueNotBuilt, + }; + return l.onLaneAs(s, q, .fp4, in, out, .f32, rows); + } + + /// Prompt rows as prompt(), out (rows, n) fp32 unrounded (nvfp4 only). + pub fn promptF32(l: Linear, s: Stream, w: Weight, in: Rows, out: u64, rows: usize) !void { + const q = switch (w) { + .nvfp4 => |q| q, + else => return error.EpilogueNotBuilt, + }; + return l.onPromptAs(s, q, .fp4, in, out, .f32, rows); + } + fn onPrompt(l: Linear, s: Stream, q: qmmf.Weight, mode: qmmf.Mode, in: Rows, out: u64, rows: usize) !void { + return l.onPromptAs(s, q, mode, in, out, .bf16, rows); + } + + fn onPromptAs(l: Linear, s: Stream, q: qmmf.Weight, mode: qmmf.Mode, in: Rows, out: u64, kind: qmmf.Out, rows: usize) !void { const p = l.nvfp4_prompt orelse return error.FormatNotLoaded; if (q.mode != mode) return error.FormatMismatch; - return p.matmul(s, in.x, if (in.ldx == 0) q.k else in.ldx, rows, q, out); + return p.matmulOut(s, in.x, if (in.ldx == 0) q.k else in.ldx, rows, q, out, kind); } fn onLane(l: Linear, s: Stream, q: qmmf.Weight, mode: qmmf.Mode, in: Rows, out: u64, rows: usize) !void { + return l.onLaneAs(s, q, mode, in, out, .bf16, rows); + } + + fn onLaneAs(l: Linear, s: Stream, q: qmmf.Weight, mode: qmmf.Mode, in: Rows, out: u64, kind: qmmf.Out, rows: usize) !void { const lane = l.lane orelse return error.FormatNotLoaded; if (q.mode != mode) return error.FormatMismatch; - return lane.matmul(s, in.x, if (in.ldx == 0) q.k else in.ldx, rows, q, out, in.part); + return lane.matmulOut(s, in.x, if (in.ldx == 0) q.k else in.ldx, rows, q, out, kind, in.part); } }; @@ -331,4 +357,7 @@ test "a lane weight whose mode is not its format's is refused before any launch" const s: Stream = undefined; try std.testing.expectError(error.FormatMismatch, l.decode(s, .{ .fp8g = fp4 }, .{ .x = 0 }, 0, 1)); try std.testing.expectError(error.FormatNotLoaded, l.prompt(s, .{ .nvfp4 = fp4 }, .{ .x = 0 }, 0, 1)); + const f8: qmmf.Weight = .{ .mode = .fp8g, .codes = 0, .scales = 0, .n = 64, .k = 64, .npad = 128 }; + try std.testing.expectError(error.EpilogueNotBuilt, l.decodeF32(s, .{ .fp8g = f8 }, .{ .x = 0 }, 0, 1)); + try std.testing.expectError(error.EpilogueNotBuilt, l.promptF32(s, .{ .bf16 = f8 }, .{ .x = 0 }, 0, 1)); } diff --git a/zig/src/cuda/qmmf.zig b/zig/src/cuda/qmmf.zig index 054e2c34bb..38e3e43a33 100644 --- a/zig/src/cuda/qmmf.zig +++ b/zig/src/cuda/qmmf.zig @@ -66,35 +66,43 @@ pub fn plan(m: usize, n: usize, k: usize, clusters: bool) Plan { } /// qmmf.cu's instance names, per mode: row tiles 16, 32, 64 (each plain, then with a cluster), then the fused kernel. -fn symbol(comptime mode: Mode, comptime bm: u32, comptime cluster: bool, comptime fused: bool) [:0]const u8 { - const name = "_ZN7tf_qmmf11qmmf_kernelILi{d}ELi{d}ELi64ELi1ELi4ELi4ELb0ELb{d}ELb{d}EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii"; - return std.fmt.comptimePrint(name, .{ @backingInt(mode), bm, @intFromBool(cluster), @intFromBool(fused) }); +fn symbol(comptime mode: Mode, comptime bm: u32, comptime cluster: bool, comptime fused: bool, comptime wide: bool) [:0]const u8 { + const name = "_ZN7tf_qmmf11qmmf_kernelILi{d}ELi{d}ELi64ELi1ELi4ELi4ELb{d}ELb{d}ELb{d}EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii"; + return std.fmt.comptimePrint(name, .{ @backingInt(mode), bm, @intFromBool(wide), @intFromBool(cluster), @intFromBool(fused) }); } -pub const reduce_symbol = "_ZN7tf_qmmf13reduce_kernelILb0EEEvPKfPvxif"; +pub const reduce_symbols = [2][:0]const u8{ "_ZN7tf_qmmf13reduce_kernelILb0EEEvPKfPvxif", "_ZN7tf_qmmf13reduce_kernelILb1EEEvPKfPvxif" }; const row_tiles = [_]u32{ 16, 32, 64 }; +/// What a call writes: bf16 rows, or the fp32 sums unrounded (FP4 only, a tensor-parallel rank's partial rows). +pub const Out = enum { bf16, f32 }; + +fn loadTiles(mod: Module, comptime mode: Mode, comptime wide: bool) ![7]Function { + var set: [7]Function = undefined; + inline for (row_tiles, 0..) |bm, ti| inline for (.{ false, true }, 0..) |cl, ci| { + const f = try mod.function(symbol(mode, bm, cl, false, wide)); + try f.allowDynamicShared(sharedBytes(mode, bm)); + set[ti * 2 + ci] = f; + }; + set[6] = try mod.function(symbol(mode, 64, false, true, wide)); + try set[6].allowDynamicShared(sharedBytes(mode, 64)); + return set; +} + pub const Lane = struct { mod: Module, fns: [modes.len][7]Function, // per mode: lane16, lane16c, lane32, lane32c, lane64, lane64c, fused - reduce: Function, + fp4_f32: [7]Function, // FP4's tiles writing fp32 + reduce: [2]Function, // by Out clusters: bool, // sm_90 on: slices meet in a cluster's shared memory pub fn load(d: *const Driver, capability_major: u32) !Lane { if (!kernels.available) return error.BuiltWithoutKernels; var mod = try Module.load(d, kernels.qmmf); errdefer mod.unload(); - var l: Lane = .{ .mod = mod, .fns = undefined, .reduce = try mod.function(reduce_symbol), .clusters = capability_major >= 9 }; - inline for (modes, 0..) |mode, mi| { - inline for (row_tiles, 0..) |bm, ti| inline for (.{ false, true }, 0..) |cl, ci| { - const f = try mod.function(symbol(mode, bm, cl, false)); - try f.allowDynamicShared(sharedBytes(mode, bm)); - l.fns[mi][ti * 2 + ci] = f; - }; - const f = try mod.function(symbol(mode, 64, false, true)); - try f.allowDynamicShared(sharedBytes(mode, 64)); - l.fns[mi][6] = f; - } + var l: Lane = .{ .mod = mod, .fns = undefined, .fp4_f32 = try loadTiles(mod, .fp4, true), .reduce = undefined, .clusters = capability_major >= 9 }; + inline for (modes, 0..) |mode, mi| l.fns[mi] = try loadTiles(mod, mode, false); + for (&l.reduce, reduce_symbols) |*f, name| f.* = try mod.function(name); return l; } @@ -110,8 +118,14 @@ pub const Lane = struct { /// y [m, n] bf16 = x [m, k] bf16 (rows `ldx` apart) times the weight, fp32 sums; `part` holds partBytes. pub fn matmul(l: Lane, s: Stream, x: u64, ldx: usize, m: usize, w: Weight, y: u64, part: u64) !void { + return l.matmulOut(s, x, ldx, m, w, y, .bf16, part); + } + + /// matmul() writing y as `out` says: bf16, or fp32 (FP4 only). + pub fn matmulOut(l: Lane, s: Stream, x: u64, ldx: usize, m: usize, w: Weight, y: u64, out: Out, part: u64) !void { + if (out == .f32 and w.mode != .fp4) return error.EpilogueNotBuilt; const p = plan(m, w.n, w.k, l.clusters); - const set = l.fns[modeIndex(w.mode)]; + const set = if (out == .f32) l.fp4_f32 else l.fns[modeIndex(w.mode)]; const ti: usize = if (p.bm == 16) 0 else if (p.bm == 32) 1 else 2; const f = if (p.fused) set[6] else set[ti * 2 + @intFromBool(p.cluster)]; const rows_t = (m + p.bm - 1) / p.bm; @@ -144,7 +158,7 @@ pub const Lane = struct { r.add(@as(i64, @intCast(total))); r.add(@as(i32, @intCast(p.sk))); r.add(w.scale); - try launch_.launch(l.reduce, .{ .grid = .{ .x = @intCast((total + 255) / 256), .y = 1, .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, s, &r); + try launch_.launch(l.reduce[@backingInt(out)], .{ .grid = .{ .x = @intCast((total + 255) / 256), .y = 1, .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, s, &r); } } }; @@ -209,6 +223,7 @@ test "tile shared bytes, slices and row tiles as the Python host picks them" { } test "instance names are qmmf.cu's" { - try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi3ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp8g, 16, true, false)); - try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi0ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp4, 64, false, true)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi3ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp8g, 16, true, false, false)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi0ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp4, 64, false, true, false)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi0ELi32ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp4, 32, true, false, true)); } diff --git a/zig/tests/cuda/quant_ref_tests.zig b/zig/tests/cuda/quant_ref_tests.zig index a4552e898b..f33487c538 100644 --- a/zig/tests/cuda/quant_ref_tests.zig +++ b/zig/tests/cuda/quant_ref_tests.zig @@ -168,11 +168,32 @@ fn lanes(r: *Rig, lin: qlinear.Linear, n: usize, k: usize) !void { const path = if (prompt) "prompt" else "decode"; try r.close(try std.fmt.bufPrint(&name, "{s} {s} {d} rows x [{d}, {d}]", .{ f[0], path, wide, n, k }), outs[rows.len - 1], false, ref.want, ref.tol); for (rows[0 .. rows.len - 1], outs[0 .. rows.len - 1]) |m, o| try sameRows(r, try std.fmt.bufPrint(&name, "{s} {s}: {d} rows equal the {d}-row call's", .{ f[0], path, m, wide }), outs[rows.len - 1], o, m * n * 2); + if (comptime std.mem.eql(u8, f[0], "nvfp4")) try fp32Out(r, lin, view, dx, part, n, k, prompt, ref.want, ref.tol, outs[rows.len - 1]); } } check.pass("qlinear fp8g and nvfp4, [{d}, {d}]: decode and prompt within float64 tolerance, every row count's bits equal", .{ n, k }); } +/// nvfp4's decodeF32 and promptF32: in float64 tolerance, the same bits at any row count, the bf16 out its rounding. +fn fp32Out(r: *Rig, lin: qlinear.Linear, view: qlinear.Weight, dx: u64, part: u64, n: usize, k: usize, prompt: bool, want: []const f64, tol: []const f64, bf16_out: u64) !void { + const wide = rows[rows.len - 1]; + const path = if (prompt) "prompt" else "decode"; + var name: [96]u8 = undefined; + const outs = try r.a.alloc(u64, rows.len); + for (rows, outs) |m, *o| { + o.* = try r.zeros(m * n * 4); + const in: qlinear.Rows = .{ .x = dx, .ldx = k, .part = part }; + if (prompt) try lin.promptF32(r.s, view, in, o.*, m) else try lin.decodeF32(r.s, view, in, o.*, m); + } + try r.close(try std.fmt.bufPrint(&name, "nvfp4 {s} fp32 {d} rows x [{d}, {d}]", .{ path, wide, n, k }), outs[rows.len - 1], true, want, tol); + for (rows[0 .. rows.len - 1], outs[0 .. rows.len - 1]) |m, o| try sameRows(r, try std.fmt.bufPrint(&name, "nvfp4 {s} fp32: {d} rows equal the {d}-row call's", .{ path, m, wide }), outs[rows.len - 1], o, m * n * 4); + const f = std.mem.bytesAsSlice(f32, try r.back(outs[rows.len - 1], wide * n * 4)); + const b = std.mem.bytesAsSlice(u16, try r.back(bf16_out, wide * n * 2)); + var differ: usize = 0; + for (f, b) |x, y| differ += @intFromBool(bf16(x) != y); + try check.expect(differ == 0, "nvfp4 {s}: the bf16 out is the fp32 out rounded ({d} of {d} differ)", .{ path, differ, f.len }); +} + pub fn run(gpu: Gpu) !void { var arena: std.heap.ArenaAllocator = .init(gpu.gpa); defer arena.deinit(); From 241ac1bbc7656e053c443defa6451615ef35232f Mon Sep 17 00:00:00 2001 From: BHCC2025 <241330250+BHCC2025@users.noreply.github.com> Date: Mon, 12 Oct 2026 12:47:11 -0500 Subject: [PATCH 2/4] cuda: EXL3 projections behind qlinear kernels/cuda/exl3.cu decodes EXL3 trellis tiles as checkpoints store them ([K/16, N/16, 16 bits] int16), with core/exl3_format.zig's codebooks (3inst, mcg, mul1) and whole widths 1-8. Rows are rotated by suh and the 128-point Hadamard transform, scaled by 2^-4 and rounded once to fp16; each tile decodes once straight into the B fragments of two mma.m16n8k16 (the tile's value order is theirs), fp32 sums, then the output transform, svh and 2^-3 (docs/recipes/exl3.md, Prompts). Decode rows go 1-16 a launch, read as A fragments from rows written in fragment order; matrices with fewer than 96 column blocks split their input blocks into a shape-fixed number of parts that the last part sums in order. Prompt rows go 64 a block from shared memory. Every row is independent, so a row's bits are the same at every row count on its path. qlinear takes it as Weight.exl3, refused with FormatNotLoaded when the module is absent. tf-cuda-test exl3 [DIR]: every codebook state and random tiles at every width bit-equal to exl3_format; FP64 references from exl3_format.dequantize with a bound for the kernels' rounding order, worst column and aggregate error; decode windows of 1-16 rows and prompts of 2-300 rows byte-equal to each row alone; DIR adds real matrices of a checkpoint and times them. --- zig/build/cuda.zig | 2 + zig/kernels/cuda/exl3.cu | 267 ++++++++++++++++++ zig/kernels/cuda/exl3.cuh | 102 +++++++ zig/src/cuda/exl3_linear.zig | 209 +++++++++++++++ zig/src/cuda/kernels.zig | 1 + zig/src/cuda/qlinear.zig | 15 ++ zig/src/cuda/root.zig | 2 + zig/tests/cuda/exl3_tests.zig | 490 ++++++++++++++++++++++++++++++++++ zig/tests/cuda/main.zig | 3 + 9 files changed, 1091 insertions(+) create mode 100644 zig/kernels/cuda/exl3.cu create mode 100644 zig/kernels/cuda/exl3.cuh create mode 100644 zig/src/cuda/exl3_linear.zig create mode 100644 zig/tests/cuda/exl3_tests.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index b70dfe431f..cc6ce7e31f 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -27,6 +27,7 @@ const kernels = [_]Kernel{ .{ .name = "fp8_experts", .flags = &.{"-O3"} }, // fp8/experts.cu's device code, tensorfold_fp8_experts_v6 .{ .name = "exlr8", .flags = &.{ "-O3", "--fmad=false" } }, // ours: EXLR8 lane projections, the Metal kernels' order .{ .name = "exlr8_dense", .flags = &.{ "-O3", "--fmad=false" } }, // ours: EXLR8 dense matrices by row, the Metal order + .{ .name = "exl3", .flags = &.{ "-O3", "--fmad=false" } }, // ours: EXL3 trellis projections behind qlinear .{ .name = "nvfp4_experts", .flags = &.{"-O3"} }, // nvfp4/experts.cu's device code, the NVFP4 expert kernels .{ .name = "lane_gemv", .flags = &.{"-O3"} }, // ours: qmm_group's arithmetic, a column tile's K slices in one CTA .{ .name = "sample", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // ours: the Metal engine's keyed draws @@ -130,6 +131,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu runner.addImport("core", mods.core); runner.addImport("lanes", mods.lanes); runner.addImport("nemotron", mods.nemotron); + runner.addImport("core", mods.core); const exlr8 = exlr8_build.coreAlias(b, target, optimize, mods.core); runner.addImport("exlr8", exlr8); const exlr8_tests = exlr8_build.testModules(b, target, optimize, exlr8); diff --git a/zig/kernels/cuda/exl3.cu b/zig/kernels/cuda/exl3.cu new file mode 100644 index 0000000000..a27f1b7486 --- /dev/null +++ b/zig/kernels/cuda/exl3.cu @@ -0,0 +1,267 @@ +// EXL3 projections on CUDA: trellis tiles as stored ([K/16, N/16, 16 bits] int16), core/exl3_format.zig's codebooks. +#include "exl3.cuh" + +struct TfExl3Rows { + unsigned rows, x_stride, y_stride, y_f32; +}; +struct TfExl3Matrix { + unsigned k, n, bits, tiles_n, split; +}; + +// A block's tail from zs (rows x 128 sums): a split's parts summed in order by the last part, transform, scale. +__device__ __forceinline__ void tf_exl3_tail(unsigned rows, const __half* svh, float out_scale, const TfExl3Rows a, + const TfExl3Matrix m, uint8_t* Y, float* ws, unsigned* tickets, unsigned ob, + unsigned sg, unsigned lane, unsigned tid, float* zs) { + __shared__ unsigned last; + if (m.split > 1) { + float* part = ws + ((uint64_t)ob * m.split + blockIdx.y) * (16 * 128); + for (unsigned i = tid; i < rows * 128; i += 256) part[i] = zs[i]; + __threadfence(); + __syncthreads(); + if (tid == 0) last = atomicAdd(tickets + ob, 1u) == m.split - 1; + __syncthreads(); + if (!last) return; + __threadfence(); + const float* parts = ws + (uint64_t)ob * m.split * (16 * 128); + for (unsigned i = tid; i < rows * 128; i += 256) { + float v = __ldcg(parts + i); + for (unsigned p = 1; p < m.split; ++p) v = __fadd_rn(v, __ldcg(parts + p * (16 * 128) + i)); + zs[i] = v; + } + if (tid == 0) tickets[ob] = 0; + __syncthreads(); + } + for (unsigned i = sg; i < rows; i += 8) { + float v[4]; + for (int q = 0; q < 4; ++q) v[q] = zs[i * 128 + lane + 32 * q]; + tf_exl3_fwht(v, lane); + for (int q = 0; q < 4; ++q) { + const unsigned c = ob * 128 + lane + 32 * q; + // svh times a power of two is exact; one rounding remains + const float scale = __fmul_rn(__half2float(svh[c]), out_scale); + tf_exl3_store(Y, (uint64_t)i * a.y_stride + c, __fmul_rn(v[q], scale), a.y_f32); + } + } +} + +// Qualification: every state's weight under codebook c, for the host to compare with exl3_format.codebookValue. +extern "C" __global__ void tf_exl3_codebook(unsigned c, __half* out) { + const unsigned i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= 65536) return; + out[i] = c == 2 ? tf_exl3_value<2>(i) : c == 1 ? tf_exl3_value<1>(i) : tf_exl3_value<0>(i); +} + +template +__device__ __forceinline__ void tf_exl3_unpack_tile(const uint32_t* tile, unsigned c, unsigned lane, __half (&w)[8]) { + if (c == 2) tf_exl3_decode(tile, lane, w); + else if (c == 1) tf_exl3_decode(tile, lane, w); + else tf_exl3_decode(tile, lane, w); +} + +// Qualification: W_q [K, N] fp16 in the rotated domain, a warp per tile, as exl3_format.unpack places it. +extern "C" __global__ void __launch_bounds__(32) tf_exl3_unpack(const uint8_t* T, TfExl3Matrix m, unsigned c, __half* out) { + const unsigned lane = threadIdx.x, tn = blockIdx.x, tk = blockIdx.y; + const uint32_t* tile = reinterpret_cast(T + ((uint64_t)tk * m.tiles_n + tn) * 32 * m.bits); + __half w[8]; + switch (m.bits) { + case 1: tf_exl3_unpack_tile<1>(tile, c, lane, w); break; + case 2: tf_exl3_unpack_tile<2>(tile, c, lane, w); break; + case 3: tf_exl3_unpack_tile<3>(tile, c, lane, w); break; + case 4: tf_exl3_unpack_tile<4>(tile, c, lane, w); break; + case 5: tf_exl3_unpack_tile<5>(tile, c, lane, w); break; + case 6: tf_exl3_unpack_tile<6>(tile, c, lane, w); break; + case 7: tf_exl3_unpack_tile<7>(tile, c, lane, w); break; + default: tf_exl3_unpack_tile<8>(tile, c, lane, w); break; + } + for (unsigned j = 0; j < 8; ++j) { + const unsigned row = 16 * tk + 2 * (lane % 4) + (j & 1) + 8 * ((j >> 1) & 1), col = 16 * tn + lane / 4 + 8 * (j >> 2); + out[(uint64_t)row * m.n + col] = w[j]; + } +} + +// Prompt rows (docs/recipes/exl3.md, Prompts): (x * suh) H times 2^-4 rounded once to fp16; 2^-3 goes on the output. +extern "C" __global__ void __launch_bounds__(32) tf_exl3_rotate_h(const __nv_bfloat16* X, unsigned x_stride, const __half* suh, + __half* R, unsigned r_stride) { + const unsigned lane = threadIdx.x, row = blockIdx.y, base = blockIdx.x * 128; + float v[4]; + for (int q = 0; q < 4; ++q) { + const unsigned c = base + lane + 32 * q; + v[q] = __fmul_rn(__bfloat162float(X[(uint64_t)row * x_stride + c]), __half2float(suh[c])); + } + tf_exl3_fwht(v, lane); + for (int q = 0; q < 4; ++q) R[(uint64_t)row * r_stride + base + lane + 32 * q] = __float2half_rn(__fmul_rn(v[q], 0.0625f)); +} + +// 64 prompt rows x 128 columns a block: each tile decodes once into mma.m16n8k16 B fragments (their value order). +template +__device__ __forceinline__ void tf_exl3_prompt_job(const __half* Xh, unsigned rows, const uint8_t* T, const __half* svh, + const TfExl3Matrix m, __nv_bfloat16* Y, float* buf) { + constexpr unsigned TB = 32 * K; + const unsigned tid = threadIdx.x, sg = tid / 32, lane = tid % 32, g8 = lane / 4, t4 = lane % 4, ob = blockIdx.x; + const unsigned row0 = blockIdx.y * 64, kb = m.k / 128; + __half* xs = reinterpret_cast<__half*>(buf); + const uint64_t row_bytes = (uint64_t)m.tiles_n * TB; + const uint8_t* col = T + (uint64_t)(8 * ob + sg) * TB; + float acc[4][2][4]; +#pragma unroll + for (int i = 0; i < 4; ++i) + for (int n = 0; n < 2; ++n) acc[i][n][0] = acc[i][n][1] = acc[i][n][2] = acc[i][n][3] = 0.0f; + uint4 next = tf_exl3_words(reinterpret_cast(col), lane); + for (unsigned ib = 0; ib < kb; ++ib) { + // the block's 64 rows of this 128-input block, fp16, zeros past the last row + for (unsigned i = tid; i < 64 * 16; i += 256) { + const unsigned r = i / 16, c = (i % 16) * 8; + const uint4 v = row0 + r < rows ? *reinterpret_cast(Xh + (uint64_t)(row0 + r) * m.k + ib * 128 + c) : make_uint4(0, 0, 0, 0); + *reinterpret_cast(xs + r * TF_EXL3_XS + c) = v; + } + __syncthreads(); + for (unsigned r = 0; r < 8; ++r) { + const uint4 now = next; + next = tf_exl3_words(reinterpret_cast(col + (uint64_t)(r < 7 ? 8 * ib + r + 1 : (ib + 1 < kb ? 8 * ib + 8 : 8 * ib + 7)) * row_bytes), lane); + __half w[8]; + tf_exl3_unpack(now, lane, w); + const uint32_t b[2][2] = {{tf_exl3_pair(w[0], w[1]), tf_exl3_pair(w[2], w[3])}, {tf_exl3_pair(w[4], w[5]), tf_exl3_pair(w[6], w[7])}}; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const __half* x = xs + (16 * i + g8) * TF_EXL3_XS + 16 * r + 2 * t4; + const uint32_t a[4] = {*reinterpret_cast(x), *reinterpret_cast(x + 8 * TF_EXL3_XS), + *reinterpret_cast(x + 8), *reinterpret_cast(x + 8 * TF_EXL3_XS + 8)}; + tf_exl3_mma(acc[i][0], a, b[0][0], b[0][1]); + tf_exl3_mma(acc[i][1], a, b[1][0], b[1][1]); + } + } + __syncthreads(); + } + // the sums by row into shared memory (over the rows' halves), then a warp a row: transform, svh, 2^-3 +#pragma unroll + for (int i = 0; i < 4; ++i) + for (int n = 0; n < 2; ++n) + for (int h = 0; h < 2; ++h) { + const unsigned c = 16 * sg + 8 * n + 2 * t4; + buf[(16 * i + g8 + 8 * h) * 128 + c] = acc[i][n][2 * h]; + buf[(16 * i + g8 + 8 * h) * 128 + c + 1] = acc[i][n][2 * h + 1]; + } + __syncthreads(); + for (unsigned i = sg; i < 64 && row0 + i < rows; i += 8) { + float v[4]; + for (int q = 0; q < 4; ++q) v[q] = buf[i * 128 + lane + 32 * q]; + tf_exl3_fwht(v, lane); + for (int q = 0; q < 4; ++q) { + const unsigned c = ob * 128 + lane + 32 * q; + Y[(uint64_t)(row0 + i) * m.n + c] = __float2bfloat16_rn(__fmul_rn(v[q], __fmul_rn(__half2float(svh[c]), 0.125f))); + } + } +} + +// One prompt kernel a codebook; blocks (128-column block, 64-row block). +#define TF_EXL3_PROMPT(C) \ + extern "C" __global__ void __launch_bounds__(256) tf_exl3_prompt_##C(const __half* Xh, unsigned rows, \ + const uint8_t* T, const __half* svh, \ + const TfExl3Matrix m, __nv_bfloat16* Y) { \ + __shared__ __align__(16) float buf[64 * 128]; \ + switch (m.bits) { \ + case 1: tf_exl3_prompt_job<1, C>(Xh, rows, T, svh, m, Y, buf); break; \ + case 2: tf_exl3_prompt_job<2, C>(Xh, rows, T, svh, m, Y, buf); break; \ + case 3: tf_exl3_prompt_job<3, C>(Xh, rows, T, svh, m, Y, buf); break; \ + case 4: tf_exl3_prompt_job<4, C>(Xh, rows, T, svh, m, Y, buf); break; \ + case 5: tf_exl3_prompt_job<5, C>(Xh, rows, T, svh, m, Y, buf); break; \ + case 6: tf_exl3_prompt_job<6, C>(Xh, rows, T, svh, m, Y, buf); break; \ + case 7: tf_exl3_prompt_job<7, C>(Xh, rows, T, svh, m, Y, buf); break; \ + case 8: tf_exl3_prompt_job<8, C>(Xh, rows, T, svh, m, Y, buf); break; \ + default: break; \ + } \ + } +TF_EXL3_PROMPT(0) +TF_EXL3_PROMPT(1) +TF_EXL3_PROMPT(2) + + +// Decode rows for the tensor units: tf_exl3_rotate_h's values, each 128-block in fragment order (t4, r, 4 halves). +extern "C" __global__ void __launch_bounds__(32) tf_exl3_rotate_f(const __nv_bfloat16* X, unsigned x_stride, const __half* suh, + __half* R, unsigned r_stride) { + const unsigned lane = threadIdx.x, row = blockIdx.y, base = blockIdx.x * 128; + float v[4]; + for (int q = 0; q < 4; ++q) { + const unsigned c = base + lane + 32 * q; + v[q] = __fmul_rn(__bfloat162float(X[(uint64_t)row * x_stride + c]), __half2float(suh[c])); + } + tf_exl3_fwht(v, lane); + for (int q = 0; q < 4; ++q) { + const unsigned j = lane + 32 * q, r = j / 16, c = j % 16; + const unsigned at = ((c % 8) / 2) * 32 + r * 4 + (c % 2) + 2 * (c / 8); + R[(uint64_t)row * r_stride + base + at] = __float2half_rn(__fmul_rn(v[q], 0.0625f)); + } +} + +// Decode rows on the tensor units, 1-16 a launch: tf_exl3_rotate_f's rows as A fragments, the same at every count. +template +__device__ __forceinline__ void tf_exl3_tc_job(const __half* Xh, const uint8_t* T, const __half* svh, const TfExl3Rows a, + const TfExl3Matrix m, uint8_t* Y, float* ws, unsigned* tickets, float* zs) { + constexpr unsigned TB = 32 * K; + const unsigned tid = threadIdx.x, sg = tid / 32, lane = tid % 32, g8 = lane / 4, t4 = lane % 4, ob = blockIdx.x; + const unsigned rows = a.rows < 16 ? a.rows : 16, kb = m.k / 128, per = (kb + m.split - 1) / m.split; + const unsigned b0 = blockIdx.y * per, b1 = b0 + per < kb ? b0 + per : kb; + const uint64_t row_bytes = (uint64_t)m.tiles_n * TB; + const uint8_t* col = T + (uint64_t)(8 * ob + sg) * TB; + const bool lo = g8 < rows, hi = g8 + 8 < rows; + // a lane's A words of a 128-block, rows g8 and g8 + 8 (r 0-7: columns 2t4, 2t4 + 8), loaded a block ahead + const uint4* x0 = reinterpret_cast(Xh + (uint64_t)(lo ? g8 : 0) * m.k + 32 * t4); + const uint4* x1 = reinterpret_cast(Xh + (uint64_t)(hi ? g8 + 8 : 0) * m.k + 32 * t4); + float acc[2][4] = {{0.0f, 0.0f, 0.0f, 0.0f}, {0.0f, 0.0f, 0.0f, 0.0f}}; + uint4 next = tf_exl3_words(reinterpret_cast(col + (uint64_t)(8 * b0) * row_bytes), lane); + uint4 an[2][4]; + for (int q = 0; q < 4; ++q) { + an[0][q] = lo ? __ldg(x0 + b0 * 16 + q) : make_uint4(0, 0, 0, 0); + an[1][q] = hi ? __ldg(x1 + b0 * 16 + q) : make_uint4(0, 0, 0, 0); + } + for (unsigned ib = b0; ib < b1; ++ib) { + uint32_t au[2][16]; + for (int q = 0; q < 4; ++q) { + au[0][4 * q] = an[0][q].x, au[0][4 * q + 1] = an[0][q].y, au[0][4 * q + 2] = an[0][q].z, au[0][4 * q + 3] = an[0][q].w; + au[1][4 * q] = an[1][q].x, au[1][4 * q + 1] = an[1][q].y, au[1][4 * q + 2] = an[1][q].z, au[1][4 * q + 3] = an[1][q].w; + } + if (ib + 1 < b1) + for (int q = 0; q < 4; ++q) { + an[0][q] = lo ? __ldg(x0 + (ib + 1) * 16 + q) : make_uint4(0, 0, 0, 0); + an[1][q] = hi ? __ldg(x1 + (ib + 1) * 16 + q) : make_uint4(0, 0, 0, 0); + } +#pragma unroll + for (unsigned r = 0; r < 8; ++r) { + const uint4 now = next; + next = tf_exl3_words(reinterpret_cast(col + (uint64_t)(r < 7 ? 8 * ib + r + 1 : (ib + 1 < b1 ? 8 * ib + 8 : 8 * ib + 7)) * row_bytes), lane); + const uint32_t af[4] = {au[0][2 * r], au[1][2 * r], au[0][2 * r + 1], au[1][2 * r + 1]}; + __half w[8]; + tf_exl3_unpack(now, lane, w); + tf_exl3_mma(acc[0], af, tf_exl3_pair(w[0], w[1]), tf_exl3_pair(w[2], w[3])); + tf_exl3_mma(acc[1], af, tf_exl3_pair(w[4], w[5]), tf_exl3_pair(w[6], w[7])); + } + } +#pragma unroll + for (int n = 0; n < 2; ++n) + for (int h = 0; h < 2; ++h) { + const unsigned c = 16 * sg + 8 * n + 2 * t4, row = g8 + 8 * h; + if (row < rows) { + zs[row * 128 + c] = acc[n][2 * h]; + zs[row * 128 + c + 1] = acc[n][2 * h + 1]; + } + } + __syncthreads(); + tf_exl3_tail(rows, svh, 0.125f, a, m, Y, ws, tickets, ob, sg, lane, tid, zs); +} + +#define TF_EXL3_TC_CASE(B, C) \ + case B: tf_exl3_tc_job(Xh, T, svh, a, m, Y, ws, tickets, zs); break; +#define TF_EXL3_TC(C) \ + extern "C" __global__ void __launch_bounds__(256) tf_exl3_tc_##C(const __half* Xh, const uint8_t* T, const __half* svh, \ + const TfExl3Rows a, const TfExl3Matrix m, uint8_t* Y, \ + float* ws, unsigned* tickets) { \ + __shared__ __align__(16) float zs[16 * 128]; \ + switch (m.bits) { \ + TF_EXL3_TC_CASE(1, C) TF_EXL3_TC_CASE(2, C) TF_EXL3_TC_CASE(3, C) TF_EXL3_TC_CASE(4, C) \ + TF_EXL3_TC_CASE(5, C) TF_EXL3_TC_CASE(6, C) TF_EXL3_TC_CASE(7, C) TF_EXL3_TC_CASE(8, C) \ + default: break; \ + } \ + } +TF_EXL3_TC(0) +TF_EXL3_TC(1) +TF_EXL3_TC(2) diff --git a/zig/kernels/cuda/exl3.cuh b/zig/kernels/cuda/exl3.cuh new file mode 100644 index 0000000000..b3dd3a6fd7 --- /dev/null +++ b/zig/kernels/cuda/exl3.cuh @@ -0,0 +1,102 @@ +// EXL3 on CUDA, shared by exl3.cu and exl3_experts.cu: codebooks, trellis states, the transform, a tile row's sums. +#pragma once +#include +#include +#include + +#define TF_FULL 0xFFFFFFFFu + +// The unnormalized Sylvester transform of 128 values held as lane + 32 q, in one fixed order of fp32 adds. +__device__ __forceinline__ void tf_exl3_fwht(float (&v)[4], unsigned lane) { + for (unsigned d = 1; d < 32; d <<= 1) { + for (int q = 0; q < 4; ++q) { + const float other = __shfl_xor_sync(TF_FULL, v[q], d); + v[q] = (lane & d) ? __fsub_rn(other, v[q]) : __fadd_rn(v[q], other); + } + } + const float a0 = __fadd_rn(v[0], v[1]), a1 = __fsub_rn(v[0], v[1]); + const float a2 = __fadd_rn(v[2], v[3]), a3 = __fsub_rn(v[2], v[3]); + v[0] = __fadd_rn(a0, a2); + v[1] = __fadd_rn(a1, a3); + v[2] = __fsub_rn(a0, a2); + v[3] = __fsub_rn(a1, a3); +} + +// A state's fp16 weight, one rounding: mul1's sum is an integer times 2^-18 under 2^20, exact in fp32 before it. +template +__device__ __forceinline__ __half tf_exl3_value(unsigned state) { + if (C == 2) { + const float h = (float)__dp4a(state * 0x83DCD12Du, 0x01010101u, 1024u); + return __float2half_rn(__fmaf_rn(h, __half2float(__ushort_as_half(0x1EEE)), __half2float(__ushort_as_half(0xC931)))); + } + const unsigned x = C == 1 ? state * 0xCBAC1FEDu : state * 89226354u + 64248484u; + const unsigned m = (x & 0x8FFF8FFFu) ^ 0x3B603B60u; + // the two halves' exponents differ by at most a few binades, so the fp32 sum is exact + const float s = __fadd_rn(__half2float(__ushort_as_half((unsigned short)(m & 0xFFFFu))), __half2float(__ushort_as_half((unsigned short)(m >> 16)))); + return __float2half_rn(s); +} + +// Lane l's stream words for its eight values 8 l .. 8 l + 7: the four u32 words its states span, MSB first. +template +__device__ __forceinline__ uint4 tf_exl3_words(const uint32_t* tile, unsigned lane) { + constexpr unsigned ring = 256 * K, words = 8 * K; + const unsigned i0 = ((8 * lane * K + K + ring - 16) % ring) / 32; + return make_uint4(tile[i0], tile[(i0 + 1) % words], tile[(i0 + 2) % words], tile[(i0 + 3) % words]); +} + +// The eight states out of those words: the 16 stream bits that end at each value. +template +__device__ __forceinline__ void tf_exl3_states(uint4 t, unsigned lane, unsigned (&st)[8]) { + const unsigned off = (8 * lane * K + K + 256 * K - 16) % 32; + // one variable shift lines the window up (a, b, c: its first 96 bits); every state after it is a constant shift + const unsigned a = __funnelshift_l(t.y, t.x, off), b = __funnelshift_l(t.z, t.y, off), c = __funnelshift_l(t.w, t.z, off); +#pragma unroll + for (int n = 0; n < 8; ++n) { + const unsigned p = n * K; + st[n] = __funnelshift_l(p < 32 ? b : c, p < 32 ? a : b, p % 32) >> 16; + } +} + +template +__device__ __forceinline__ void tf_exl3_unpack(uint4 t, unsigned lane, __half (&w)[8]) { + unsigned st[8]; + tf_exl3_states(t, lane, st); + if constexpr (C == 2) { + // mul1 two states at a time: the sums (<= 2044) are exact in fp16, and one fused fp16 multiply-add rounds once + const __half2 scale = __half2half2(__ushort_as_half(0x1EEE)), bias = __half2half2(__ushort_as_half(0xC931)); +#pragma unroll + for (int n = 0; n < 8; n += 2) { + const __half2 h = __halves2half2(__uint2half_rn(__dp4a(st[n] * 0x83DCD12Du, 0x01010101u, 1024u)), + __uint2half_rn(__dp4a(st[n + 1] * 0x83DCD12Du, 0x01010101u, 1024u))); + const __half2 v = __hfma2(h, scale, bias); + w[n] = __low2half(v); + w[n + 1] = __high2half(v); + } + } else { +#pragma unroll + for (int n = 0; n < 8; ++n) w[n] = tf_exl3_value(st[n]); + } +} + +template +__device__ __forceinline__ void tf_exl3_decode(const uint32_t* tile, unsigned lane, __half (&w)[8]) { + tf_exl3_unpack(tf_exl3_words(tile, lane), lane, w); +} + +__device__ __forceinline__ void tf_exl3_store(uint8_t* y, uint64_t at, float v, unsigned f32) { + if (f32) reinterpret_cast(y)[at] = v; + else reinterpret_cast<__nv_bfloat16*>(y)[at] = __float2bfloat16_rn(v); +} + +#define TF_EXL3_XS 136 // shared prompt rows 136 halves apart: a fragment's eight rows hit eight bank groups + +// Prompt rows on the tensor units: fp16 A and B fragments, fp32 sums. +__device__ __forceinline__ void tf_exl3_mma(float (&d)[4], const uint32_t (&a)[4], uint32_t b0, uint32_t b1) { + asm("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} + +__device__ __forceinline__ uint32_t tf_exl3_pair(__half lo, __half hi) { + return (uint32_t)__half_as_ushort(lo) | ((uint32_t)__half_as_ushort(hi) << 16); +} diff --git a/zig/src/cuda/exl3_linear.zig b/zig/src/cuda/exl3_linear.zig new file mode 100644 index 0000000000..b254b9fb79 --- /dev/null +++ b/zig/src/cuda/exl3_linear.zig @@ -0,0 +1,209 @@ +//! exl3.cu's EXL3 projections: bf16 rows rotated into fp32, then 1-16 rows a launch over the trellis tiles as stored. +const std = @import("std"); +const driver = @import("driver.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const memory = @import("memory.zig"); +const stream_ = @import("stream.zig"); + +const Driver = driver.Driver; +const Function = module.Function; +const Stream = stream_.Stream; + +/// core/exl3_format.zig's Codebook, in the order exl3.cu numbers them. +pub const Codebook = enum(u32) { inst3 = 0, mcg = 1, mul1 = 2 }; + +/// One matrix on the device as stored: trellis int16 [k/16, n/16, 16 bits], suh fp16 [k], svh fp16 [n]. +pub const Matrix = struct { + trellis: u64, + suh: u64, + svh: u64, + k: usize, + n: usize, + bits: u32, + codebook: Codebook, + + /// Whole widths 1-8 over whole 128-blocks; exl3.cu takes no half widths and no partial blocks. + pub fn check(m: Matrix) !void { + if (m.k == 0 or m.n == 0 or m.k % 128 != 0 or m.n % 128 != 0 or m.bits < 1 or m.bits > 8) return error.InvalidExl3Matrix; + } + + pub fn trellisBytes(k: usize, n: usize, bits: u32) usize { + return k / 16 * (n / 16) * 32 * @as(usize, bits); + } + + fn args(m: Matrix) MatrixArgs { + return .{ .k = @intCast(m.k), .n = @intCast(m.n), .bits = m.bits, .tiles_n = @intCast(m.n / 16), .split = @intCast(split(m.k, m.n)) }; + } +}; + +/// exl3.cu's TfExl3Rows and TfExl3Matrix, passed by value. +pub const RowsArgs = extern struct { rows: u32, x_stride: u32, y_stride: u32, y_f32: u32 }; +pub const MatrixArgs = extern struct { k: u32, n: u32, bits: u32, tiles_n: u32, split: u32 }; + +/// Decode rows a launch; longer decode runs go window by window. +pub const window = 16; +/// Prompt rows a chunk: rotated to fp16 together, then 64 rows a block on the tensor units. +pub const prompt_chunk = 256; +/// Blocks a launch aims for: a matrix with fewer 128-column blocks splits its input blocks into parts. +const target_blocks = 96; + +/// Parts of the input blocks, from the shape alone: enough for ~target_blocks blocks, at least 4 input blocks a part. +pub fn split(k: usize, n: usize) usize { + const cols = n / 128; + const kb = k / 128; + if (cols == 0 or cols >= target_blocks) return 1; + const want = @min((target_blocks + cols - 1) / cols, kb / 4); + if (want <= 1) return 1; + const per = (kb + want - 1) / want; + return (kb + per - 1) / per; +} + +/// The kernels and their scratch (fp16 rotated rows up to `k_max` inputs): one stream's projections at a time. +pub const Kernels = struct { + decode_fns: [3]Function, // 1-16 rows on the tensor units, by codebook + prompt_fns: [3]Function, // 64 rows a block on the tensor units, by codebook + rotate_f_fn: Function, // decode rows, each 128-block in fragment order + rotate_h_fn: Function, // prompt rows, plain order + codebook_fn: Function, + unpack_fn: Function, + rotated: memory.DeviceBuffer, // a decode window's or prompt chunk's rotated rows (fp16, prompt_chunk x k_max) + parts: memory.DeviceBuffer, // a split matrix's part totals: 16 rows x 128 columns per (column block, part) + tickets: memory.DeviceBuffer, // per column block, parts finished; the last part resets its count to zero + k_max: usize, + + pub fn resolve(d: *const Driver, mod: module.Module, k_max: usize) !Kernels { + var k: Kernels = .{ .decode_fns = undefined, .prompt_fns = undefined, .rotate_f_fn = try mod.function("tf_exl3_rotate_f"), .rotate_h_fn = try mod.function("tf_exl3_rotate_h"), .codebook_fn = try mod.function("tf_exl3_codebook"), .unpack_fn = try mod.function("tf_exl3_unpack"), .rotated = undefined, .parts = undefined, .tickets = undefined, .k_max = k_max }; + for (&k.decode_fns, &k.prompt_fns, 0..) |*f, *p, c| { + var name: [32]u8 = undefined; + f.* = try mod.function(try std.fmt.bufPrintSentinel(&name, "tf_exl3_tc_{d}", .{c}, 0)); + p.* = try mod.function(try std.fmt.bufPrintSentinel(&name, "tf_exl3_prompt_{d}", .{c}, 0)); + } + k.rotated = try memory.DeviceBuffer.alloc(d, prompt_chunk * k_max * 2); + errdefer k.rotated.free(); + // split * cols stays under target_blocks + cols, so 2 * target_blocks parts always fit + k.parts = try memory.DeviceBuffer.alloc(d, 2 * target_blocks * window * 128 * 4); + errdefer k.parts.free(); + k.tickets = try memory.DeviceBuffer.alloc(d, target_blocks * 4); + errdefer k.tickets.free(); + try d.check(d.api.cuMemsetD32_v2(k.tickets.ptr, 0, target_blocks), "cuMemsetD32"); + return k; + } + + pub fn deinit(k: *Kernels) void { + k.rotated.free(); + k.parts.free(); + k.tickets.free(); + } + + /// Decode rows: bf16 rows `ldx` apart -> out (rows, n), 16 a window on the tensor units, the same at every count. + pub fn rows(k: *const Kernels, s: Stream, x: u64, ldx: usize, m: Matrix, out: u64, count: usize, f32_out: bool) !void { + try k.fits(m, ldx); + const elem: usize = if (f32_out) 4 else 2; + var at: usize = 0; + while (at < count) : (at += window) { + const r = @min(window, count - at); + try k.rotate(k.rotate_f_fn, s, x + at * ldx * 2, ldx, m, r); + var a: launch_.Args = .{}; + a.add(k.rotated.ptr); + a.add(m.trellis); + a.add(m.svh); + a.add(RowsArgs{ .rows = @intCast(r), .x_stride = @intCast(m.k), .y_stride = @intCast(m.n), .y_f32 = @intFromBool(f32_out) }); + a.add(m.args()); + a.add(out + at * m.n * elem); + a.add(k.parts.ptr); + a.add(k.tickets.ptr); + try launch_.launch(k.decode_fns[@backingInt(m.codebook)], .{ .grid = .{ .x = @intCast(m.n / 128), .y = @intCast(split(m.k, m.n)), .z = 1 }, .block = .{ .x = 256 } }, s, &a); + } + } + + /// Prompt rows (docs/recipes/exl3.md, Prompts): fp16 rows on the tensor units, fp32 sums, the same at every count. + pub fn prompt(k: *const Kernels, s: Stream, x: u64, ldx: usize, m: Matrix, out: u64, count: usize) !void { + try k.fits(m, ldx); + var at: usize = 0; + while (at < count) : (at += prompt_chunk) { + const r = @min(prompt_chunk, count - at); + try k.rotate(k.rotate_h_fn, s, x + at * ldx * 2, ldx, m, r); + var a: launch_.Args = .{}; + a.add(k.rotated.ptr); + a.add(@as(u32, @intCast(r))); + a.add(m.trellis); + a.add(m.svh); + a.add(m.args()); + a.add(out + at * m.n * 2); + try launch_.launch(k.prompt_fns[@backingInt(m.codebook)], .{ .grid = .{ .x = @intCast(m.n / 128), .y = @intCast((r + 63) / 64), .z = 1 }, .block = .{ .x = 256 } }, s, &a); + } + } + + fn fits(k: *const Kernels, m: Matrix, ldx: usize) !void { + try m.check(); + if (m.k > k.k_max) return error.Exl3InputsOverScratch; + if (ldx < m.k) return error.RowStride; + } + + /// (x * suh) H times 2^-4 as fp16 rows into the scratch, one warp per (row, 128-block), in `f`'s order. + fn rotate(k: *const Kernels, f: Function, s: Stream, x: u64, ldx: usize, m: Matrix, count: usize) !void { + var a: launch_.Args = .{}; + a.add(x); + a.add(@as(u32, @intCast(ldx))); + a.add(m.suh); + a.add(k.rotated.ptr); + a.add(@as(u32, @intCast(m.k))); + try launch_.launch(f, .{ .grid = .{ .x = @intCast(m.k / 128), .y = @intCast(count), .z = 1 }, .block = .{ .x = 32 } }, s, &a); + } + + /// Every state's fp16 weight under `c` (65536 halves) into `out`. + pub fn codebook(k: *const Kernels, s: Stream, c: Codebook, out: u64) !void { + var a: launch_.Args = .{}; + a.add(@backingInt(c)); + a.add(out); + try launch_.launch(k.codebook_fn, .{ .grid = .{ .x = 256, .y = 1, .z = 1 }, .block = .{ .x = 256 } }, s, &a); + } + + /// W_q [k, n] fp16 (the rotated domain) into `out`, a warp per tile; suh and svh are not read. + pub fn unpack(k: *const Kernels, s: Stream, m: Matrix, out: u64) !void { + if (m.k % 16 != 0 or m.n % 16 != 0 or m.bits < 1 or m.bits > 8) return error.InvalidExl3Matrix; + var a: launch_.Args = .{}; + a.add(m.trellis); + a.add(m.args()); + a.add(@backingInt(m.codebook)); + a.add(out); + try launch_.launch(k.unpack_fn, .{ .grid = .{ .x = @intCast(m.n / 16), .y = @intCast(m.k / 16), .z = 1 }, .block = .{ .x = 32 } }, s, &a); + } +}; + +comptime { + std.debug.assert(@sizeOf(RowsArgs) == 16 and @sizeOf(MatrixArgs) == 20); +} + +test "a matrix exl3.cu cannot take is refused before any launch" { + const base: Matrix = .{ .trellis = 0, .suh = 0, .svh = 0, .k = 4096, .n = 2048, .bits = 4, .codebook = .mul1 }; + try base.check(); + var m = base; + m.n = 2000; + try std.testing.expectError(error.InvalidExl3Matrix, m.check()); + m = base; + m.bits = 9; + try std.testing.expectError(error.InvalidExl3Matrix, m.check()); +} + +test "the trellis size is 16 bits a 16-value word row, 32 bytes a bit" { + // GLM-5.3-Flash's routed gate_proj at K=4: trellis [256, 128, 64] int16 + try std.testing.expectEqual(@as(usize, 256 * 128 * 64 * 2), Matrix.trellisBytes(4096, 2048, 4)); +} + +test "the split comes from the shape alone and leaves no part empty" { + try std.testing.expectEqual(@as(usize, 6), split(4096, 2048)); // a routed gate_proj: 16 column blocks, 32 input blocks + try std.testing.expectEqual(@as(usize, 3), split(16384, 4096)); // o_proj + try std.testing.expectEqual(@as(usize, 1), split(1536, 16384)); // q_b_proj: 128 column blocks already + try std.testing.expectEqual(@as(usize, 1), split(512, 384)); + var k: usize = 128; + while (k <= 16384) : (k += 128) { + var n: usize = 128; + while (n <= 16384) : (n += 128) { + const p = split(k, n); + const per = (k / 128 + p - 1) / p; + try std.testing.expect(p * (n / 128) <= 2 * target_blocks and (p - 1) * per < k / 128); + } + } +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 0fa8bf6e28..b17d2725b6 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -27,6 +27,7 @@ pub const affine4_pack: []const u8 = if (available) &Blob("fatbin_affine4_pack") pub const fp8_experts: []const u8 = if (available) &Blob("fatbin_fp8_experts").bytes else &.{}; pub const exlr8: []const u8 = if (available) &Blob("fatbin_exlr8").bytes else &.{}; pub const exlr8_dense: []const u8 = if (available) &Blob("fatbin_exlr8_dense").bytes else &.{}; +pub const exl3: []const u8 = if (available) &Blob("fatbin_exl3").bytes else &.{}; pub const nvfp4_experts: []const u8 = if (available) &Blob("fatbin_nvfp4_experts").bytes else &.{}; pub const sample: []const u8 = if (available) &Blob("fatbin_sample").bytes else &.{}; pub const train: []const u8 = if (available) &Blob("fatbin_train").bytes else &.{}; diff --git a/zig/src/cuda/qlinear.zig b/zig/src/cuda/qlinear.zig index ff70a0d88a..7c6f7f3e65 100644 --- a/zig/src/cuda/qlinear.zig +++ b/zig/src/cuda/qlinear.zig @@ -7,6 +7,7 @@ const memory = @import("memory.zig"); const stream_ = @import("stream.zig"); const qmmf = @import("qmmf.zig"); const nvfp4 = @import("nvfp4.zig"); +const exl3 = @import("exl3_linear.zig"); const Driver = driver.Driver; const Module = module.Module; @@ -36,11 +37,13 @@ pub const Weight = union(enum) { nvfp4: qmmf.Weight, fp8: qmmf.Weight, // e4m3 with one fp32 scale (ModelOpt FP8) bf16: qmmf.Weight, // unquantized: the checkpoint's bf16 values + exl3: exl3.Matrix, // trellis tiles as stored, with suh and svh pub fn outputs(w: Weight) usize { return switch (w) { .affine4 => |q| q.n, .fp8g, .nvfp4, .fp8, .bf16 => |q| q.n, + .exl3 => |q| q.n, }; } @@ -48,6 +51,7 @@ pub const Weight = union(enum) { return switch (w) { .affine4 => |q| q.k, .fp8g, .nvfp4, .fp8, .bf16 => |q| q.k, + .exl3 => |q| q.k, }; } }; @@ -60,6 +64,7 @@ pub const Linear = struct { affine4: ?*const Affine4Kernels = null, lane: ?*const qmmf.Lane = null, // decode rows of every lane format; fp8g and bf16 prompt rows too nvfp4_prompt: ?*const nvfp4.Prompt = null, // the prompt GEMM: nvfp4 and fp8 prompt rows (linear.py's prefill) + exl3: ?*const exl3.Kernels = null, // decode rows 16 a window; prompt rows on the tensor units /// Decode rows: out (rows, n) bf16, each row's bits the same at every row count the format takes. pub fn decode(l: Linear, s: Stream, w: Weight, in: Rows, out: u64, rows: usize) !void { @@ -73,6 +78,7 @@ pub const Linear = struct { .nvfp4 => |q| return l.onLane(s, q, .fp4, in, out, rows), .fp8 => |q| return l.onLane(s, q, .fp8, in, out, rows), .bf16 => |q| return l.onLane(s, q, .bf16, in, out, rows), + .exl3 => |q| return l.onExl3(s, q, in, out, rows), } } @@ -88,6 +94,7 @@ pub const Linear = struct { .nvfp4 => |q| return l.onPrompt(s, q, .fp4, in, out, rows), .fp8 => |q| return l.onPrompt(s, q, .fp8, in, out, rows), .bf16 => |q| return l.onLane(s, q, .bf16, in, out, rows), + .exl3 => |q| return (l.exl3 orelse return error.FormatNotLoaded).prompt(s, in.x, if (in.ldx == 0) q.k else in.ldx, q, out, rows), } } @@ -109,6 +116,11 @@ pub const Linear = struct { return l.onPromptAs(s, q, .fp4, in, out, .f32, rows); } + fn onExl3(l: Linear, s: Stream, q: exl3.Matrix, in: Rows, out: u64, rows: usize) !void { + const e = l.exl3 orelse return error.FormatNotLoaded; + return e.rows(s, in.x, if (in.ldx == 0) q.k else in.ldx, q, out, rows, false); + } + fn onPrompt(l: Linear, s: Stream, q: qmmf.Weight, mode: qmmf.Mode, in: Rows, out: u64, rows: usize) !void { return l.onPromptAs(s, q, mode, in, out, .bf16, rows); } @@ -341,6 +353,9 @@ test "a format that is not loaded is refused before any launch" { const s: Stream = undefined; // never reached: the refusal comes first try std.testing.expectError(error.FormatNotLoaded, l.decode(s, w, .{ .x = 0 }, 0, 1)); try std.testing.expectError(error.FormatNotLoaded, l.prompt(s, w, .{ .x = 0 }, 0, 1)); + const e: Weight = .{ .exl3 = .{ .trellis = 0, .suh = 0, .svh = 0, .k = 128, .n = 128, .bits = 4, .codebook = .mul1 } }; + try std.testing.expectError(error.FormatNotLoaded, l.decode(s, e, .{ .x = 0 }, 0, 1)); + try std.testing.expectError(error.FormatNotLoaded, l.prompt(s, e, .{ .x = 0 }, 0, 1)); } test "the affine-4 tiles' sizes follow qmm_fast.tile" { diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 3053c09a6b..85f0ffc066 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -27,6 +27,7 @@ pub const kernels = @import("kernels.zig"); pub const qmmf = @import("qmmf.zig"); pub const fp8 = @import("fp8.zig"); pub const exlr8 = @import("exlr8.zig"); +pub const exl3_linear = @import("exl3_linear.zig"); pub const nvfp4 = @import("nvfp4.zig"); pub const segments = @import("segments.zig"); pub const grouped = @import("grouped.zig"); @@ -43,6 +44,7 @@ test { _ = @import("qlinear.zig"); _ = @import("norms.zig"); _ = @import("embed_rows.zig"); + _ = @import("exl3_linear.zig"); _ = @import("experts.zig"); _ = @import("sample.zig"); _ = launch; diff --git a/zig/tests/cuda/exl3_tests.zig b/zig/tests/cuda/exl3_tests.zig new file mode 100644 index 0000000000..139468e77f --- /dev/null +++ b/zig/tests/cuda/exl3_tests.zig @@ -0,0 +1,490 @@ +//! EXL3 projections on CUDA: codebooks and tiles bit for bit against core/exl3_format.zig, FP64 bounds, row windows. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const Gpu = check.Gpu; +const fmt = core.exl3_format; +const E = cuda.exl3_linear; +const st = core.safetensors; + +const k_max = 16384; // GLM-5.3-Flash's widest input (o_proj: 64 heads of 256) +const u = std.math.pow(f64, 2, -24); + +/// Device buffers for one test, freed together. +pub const Rig = struct { + gpu: Gpu, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + + pub fn put(r: *Rig, items: anytype) !u64 { + const b = try cuda.DeviceBuffer.fromHost(r.gpu.d, std.mem.sliceAsBytes(items)); + try r.bufs.append(r.gpu.gpa, b); + return b.ptr; + } + pub fn zeros(r: *Rig, len: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(r.gpu.d, @max(len, 16)); + try b.fill8(0, null); + try r.gpu.d.check(r.gpu.d.api.cuStreamSynchronize(null), "cuStreamSynchronize"); + try r.bufs.append(r.gpu.gpa, b); + return b.ptr; + } + pub fn get(r: *Rig, ptr: u64, out: []u8) !void { + const b: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = ptr, .len = out.len }; + try b.download(0, out); + } + pub fn deinit(r: *Rig) void { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(r.gpu.gpa); + } +}; + +fn refBook(c: E.Codebook) fmt.Codebook { + return switch (c) { + .inst3 => .inst3, + .mcg => .mcg, + .mul1 => .mul1, + }; +} + +pub fn bf16(v: f32) u16 { + return @intCast(@as(u32, @bitCast(v)) >> 16); +} + +fn ofBf16(b: u16) f64 { + return @as(f32, @bitCast(@as(u32, b) << 16)); +} + +fn gamma(n: usize) f64 { + return gammaU(n, u); +} + +fn gammaU(n: usize, unit: f64) f64 { + const nu = @as(f64, @floatFromInt(n)) * unit; + return nu / (1 - nu); +} + +/// `tf-cuda-test exl3 [DIR]`: codebooks, tiles, FP64 bounds and row windows; DIR adds a checkpoint's real matrices. +pub fn run(gpu: Gpu, dir: ?[]const u8) !void { + var mod = try cuda.Module.load(gpu.d, cuda.kernels.exl3); + defer mod.unload(); + var k = try E.Kernels.resolve(gpu.d, mod, k_max); + defer k.deinit(); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + try codebooks(gpu, &k, s); + try tiles(gpu, &k, s); + var b: Bound = .{}; + const cases = [_]Case{ + .{ .k = 256, .n = 256, .bits = 1, .book = .mul1 }, .{ .k = 256, .n = 384, .bits = 2, .book = .mcg }, + .{ .k = 384, .n = 256, .bits = 3, .book = .inst3 }, .{ .k = 256, .n = 256, .bits = 4, .book = .mul1 }, + .{ .k = 512, .n = 256, .bits = 5, .book = .mul1 }, .{ .k = 256, .n = 512, .bits = 6, .book = .mul1 }, + .{ .k = 256, .n = 256, .bits = 7, .book = .mcg }, .{ .k = 256, .n = 256, .bits = 8, .book = .mul1 }, + .{ .k = 1024, .n = 256, .bits = 4, .book = .inst3 }, .{ .k = 2048, .n = 128, .bits = 4, .book = .mul1 }, + .{ .k = 1024, .n = 256, .bits = 6, .book = .mcg }, + }; + for (cases, 0..) |c, i| try synthetic(gpu, &k, s, c, 0xE3 + i, &b); + check.pass("{d} synthetic matrices, widths 1-8, three codebooks: {d} FP64 values, max error {e}, max {d:.4} of the bound, worst column {e} of the largest output", .{ cases.len, b.values, b.max_error, b.max_fraction, b.worst_column }); + try windows(gpu, &k, s); + if (dir) |d| try real(gpu, &k, s, d); +} + +fn codebooks(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream) !void { + var rig: Rig = .{ .gpu = gpu }; + defer rig.deinit(); + const out = try rig.zeros(65536 * 2); + const got = try gpu.gpa.alloc(u16, 65536); + defer gpu.gpa.free(got); + for ([_]E.Codebook{ .inst3, .mcg, .mul1 }) |c| { + try k.codebook(s, c, out); + try s.synchronize(); + try rig.get(out, std.mem.sliceAsBytes(got)); + for (got, 0..) |g, i| try check.expect(g == @as(u16, @bitCast(fmt.codebookValue(refBook(c), @intCast(i)))), "exl3 {t} state {d}", .{ c, i }); + } + check.pass("3 codebooks x 65536 states: the GPU's fp16 weights equal exl3_format.codebookValue bit for bit", .{}); +} + +fn tiles(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream) !void { + var prng = std.Random.DefaultPrng.init(0x7E11); + const r = prng.random(); + const tk = 2; + const tn = 3; + var count: usize = 0; + for (1..9) |bits| for ([_]E.Codebook{ .inst3, .mcg, .mul1 }) |c| { + var rig: Rig = .{ .gpu = gpu }; + defer rig.deinit(); + const words = try gpu.gpa.alloc(i16, tk * tn * 16 * bits); + defer gpu.gpa.free(words); + for (words) |*w| w.* = r.int(i16); + const m: E.Matrix = .{ .trellis = try rig.put(words), .suh = 0, .svh = 0, .k = tk * 16, .n = tn * 16, .bits = @intCast(bits), .codebook = c }; + const out = try rig.zeros(m.k * m.n * 2); + try k.unpack(s, m, out); + try s.synchronize(); + var got: [tk * tn * 256]f16 = undefined; + var want: [tk * tn * 256]f16 = undefined; + try rig.get(out, std.mem.sliceAsBytes(&got)); + fmt.unpack(.{ .halves = @intCast(2 * bits) }, refBook(c), words, tk, tn, &want); + try check.sameBytes("exl3 unpack", std.mem.sliceAsBytes(&got), std.mem.sliceAsBytes(&want)); + count += 1; + }; + check.pass("{d} width/codebook pairs: random tiles unpack to exl3_format.unpack's W_q bit for bit", .{count}); +} + +pub const Case = struct { k: usize, n: usize, bits: u32, book: E.Codebook }; + +/// Worst observed error, as a value and as a fraction of the fp32 bound, over every value checked. +pub const Bound = struct { + values: usize = 0, + max_error: f64 = 0, + max_fraction: f64 = 0, + worst_column: f64 = 0, +}; + +/// One matrix's host copy (a column block of a real one, or a whole synthetic one) with its rows of input. +pub const Host = struct { + k: usize, + n: usize, + bits: u32, + book: E.Codebook, + trellis: []const i16, + suh: []const f16, + svh: []const f16, +}; + +/// FP64 y = x W from exl3_format.dequantize, and per value the bound of exl3.cu's order (rotate, chain, transform). +pub fn reference(a: std.mem.Allocator, h: Host, x: []const u16, rows: usize, y: []f64, bound: []f64, prompt: bool) !void { + const w = try a.alloc(f64, h.k * h.n); + defer a.free(w); + const scratch = try a.alloc(f64, 2 * h.k * h.n); + defer a.free(scratch); + fmt.dequantize(.{ .halves = @intCast(2 * h.bits) }, refBook(h.book), h.trellis, h.k / 16, h.n / 16, h.suh, h.svh, w, scratch); + const wq = try a.alloc(f16, h.k * h.n); + defer a.free(wq); + fmt.unpack(.{ .halves = @intCast(2 * h.bits) }, refBook(h.book), h.trellis, h.k / 16, h.n / 16, wq); + const rot = try a.alloc(f64, 2 * h.k); + defer a.free(rot); + const z = try a.alloc(f64, 2 * h.n); + defer a.free(z); + for (0..rows) |i| { + const xr = x[i * h.k ..][0..h.k]; + for (0..h.n) |j| { + var acc: f64 = 0; + for (0..h.k) |l| acc += ofBf16(xr[l]) * w[l * h.n + j]; + y[i * h.n + j] = acc; + } + // r = (x suh) H unnormalized, |r| and its error from one product and seven adds + const r = rot[0..h.k]; + const re = rot[h.k..]; + for (0..h.k) |l| r[l] = ofBf16(xr[l]) * @as(f64, @floatCast(h.suh[l])); + var blk: usize = 0; + while (blk < h.k) : (blk += 128) { + var sum: f64 = 0; + for (r[blk..][0..128]) |v| sum += @abs(v); + fwht(r[blk..][0..128]); + for (re[blk..][0..128], r[blk..][0..128]) |*e, v| { + e.* = gamma(8) * sum; + // prompt rows: r / 16 rounded to fp16 (relative 2^-11, 2^-25 absolute among subnormals) + if (prompt) e.* += std.math.pow(f64, 2, -11) * (@abs(v) + e.*) + std.math.pow(f64, 2, -21); + } + } + const za = z[0..h.n]; + const ze = z[h.n..]; + for (0..h.n) |j| { + var mag: f64 = 0; + var carried: f64 = 0; + for (0..h.k) |l| { + const q = @abs(@as(f64, @floatCast(wq[l * h.n + j]))); + mag += (@abs(r[l]) + re[l]) * q; + carried += re[l] * q; + } + ze[j] = (if (prompt) gammaU(h.k + 3, std.math.pow(f64, 2, -22)) else gamma(h.k + 3)) * mag + carried; + za[j] = mag; + } + var ob: usize = 0; + while (ob < h.n) : (ob += 128) { + var zsum: f64 = 0; + var esum: f64 = 0; + for (ob..ob + 128) |j| { + zsum += za[j] + ze[j]; + esum += ze[j]; + } + for (ob..ob + 128) |j| { + const e = std.math.pow(f64, 2, -7) * @abs(@as(f64, @floatCast(h.svh[j]))) * (gamma(7) * zsum + esum); + bound[i * h.n + j] = e * (1 + u) + (if (prompt) std.math.pow(f64, 2, -8) else u) * (@abs(y[i * h.n + j]) + e); + } + } + } +} + +fn fwht(v: []f64) void { + var d: usize = 1; + while (d < v.len) : (d <<= 1) { + for (0..v.len) |i| { + if (i & d != 0) continue; + const a = v[i]; + const b = v[i + d]; + v[i] = a + b; + v[i + d] = a - b; + } + } +} + +/// Every value within its bound; the largest error by column against that column's largest output. +pub fn compare(what: []const u8, got: []const f32, y: []const f64, bound: []const f64, rows: usize, n: usize, b: *Bound) !void { + var scale: f64 = 0; + for (y[0 .. rows * n]) |v| scale = @max(scale, @abs(v)); + for (0..n) |j| { + var col_err: f64 = 0; + for (0..rows) |i| { + const at = i * n + j; + const err = @abs(@as(f64, got[at]) - y[at]); + try check.expect(err <= bound[at], "{s}: row {d} column {d} error {e} over its bound {e}", .{ what, i, j, err, bound[at] }); + b.max_error = @max(b.max_error, err); + if (bound[at] > 0) b.max_fraction = @max(b.max_fraction, err / bound[at]); + col_err = @max(col_err, err); + } + if (scale > 0) b.worst_column = @max(b.worst_column, col_err / scale); + } + b.values += rows * n; +} + +pub fn randomMatrix(a: std.mem.Allocator, r: std.Random, c: Case) !Host { + const trellis = try a.alloc(i16, E.Matrix.trellisBytes(c.k, c.n, c.bits) / 2); + for (trellis) |*w| w.* = r.int(i16); + const suh = try a.alloc(f16, c.k); + const svh = try a.alloc(f16, c.n); + for (suh) |*v| v.* = @floatCast((0.25 + 1.75 * r.float(f32)) * @as(f32, if (r.boolean()) -1 else 1)); + for (svh) |*v| v.* = @floatCast((0.002 + 0.02 * r.float(f32)) * @as(f32, if (r.boolean()) -1 else 1)); + return .{ .k = c.k, .n = c.n, .bits = c.bits, .book = c.book, .trellis = trellis, .suh = suh, .svh = svh }; +} + +pub fn freeHost(a: std.mem.Allocator, h: Host) void { + a.free(h.trellis); + a.free(h.suh); + a.free(h.svh); +} + +pub fn randomRows(a: std.mem.Allocator, r: std.Random, rows: usize, k: usize) ![]u16 { + const x = try a.alloc(u16, rows * k); + for (x) |*v| v.* = bf16(r.floatNorm(f32) * 2); + return x; +} + +fn upload(rig: *Rig, h: Host) !E.Matrix { + return .{ .trellis = try rig.put(h.trellis), .suh = try rig.put(h.suh), .svh = try rig.put(h.svh), .k = h.k, .n = h.n, .bits = h.bits, .codebook = h.book }; +} + +fn synthetic(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream, c: Case, seed: u64, b: *Bound) !void { + const a = gpu.gpa; + var prng = std.Random.DefaultPrng.init(seed); + const h = try randomMatrix(a, prng.random(), c); + defer freeHost(a, h); + const rows = 16; + const x = try randomRows(a, prng.random(), rows, c.k); + defer a.free(x); + var rig: Rig = .{ .gpu = gpu }; + defer rig.deinit(); + const m = try upload(&rig, h); + const out = try rig.zeros(rows * c.n * 4); + try k.rows(s, try rig.put(x), c.k, m, out, rows, true); + try s.synchronize(); + const got = try a.alloc(f32, rows * c.n); + defer a.free(got); + try rig.get(out, std.mem.sliceAsBytes(got)); + const y = try a.alloc(f64, rows * c.n); + defer a.free(y); + const bound = try a.alloc(f64, rows * c.n); + defer a.free(bound); + try reference(a, h, x, rows, y, bound, true); + try compare("exl3 synthetic", got, y, bound, rows, c.n, b); + // the same rows as a prompt: the tensor units' arithmetic, bf16 out + try k.prompt(s, try rig.put(x), c.k, m, out, rows); + try s.synchronize(); + try getBf16(&rig, out, got); + try reference(a, h, x, rows, y, bound, true); + try compare("exl3 synthetic prompt", got, y, bound, rows, c.n, b); +} + +fn getBf16(rig: *Rig, ptr: u64, out: []f32) !void { + const raw = try rig.gpu.gpa.alloc(u16, out.len); + defer rig.gpu.gpa.free(raw); + try rig.get(ptr, std.mem.sliceAsBytes(raw)); + for (out, raw) |*o, v| o.* = @floatCast(ofBf16(v)); +} + +/// Each row alone, then every window 1-16 and a 40-row prompt through qlinear: each row's bf16 bytes as when alone. +fn windows(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream) !void { + const a = gpu.gpa; + var prng = std.Random.DefaultPrng.init(0x3D16); + const h = try randomMatrix(a, prng.random(), .{ .k = 2048, .n = 256, .bits = 4, .book = .mul1 }); + defer freeHost(a, h); + const rows = 40; + const prompt_rows = 300; // past a 256-row chunk and a 64-row block + const x = try randomRows(a, prng.random(), prompt_rows, h.k); + defer a.free(x); + var rig: Rig = .{ .gpu = gpu }; + defer rig.deinit(); + const w: cuda.qlinear.Weight = .{ .exl3 = try upload(&rig, h) }; + const l: cuda.qlinear.Linear = .{ .exl3 = k }; + const xd = try rig.put(x); + const out = try rig.zeros(prompt_rows * h.n * 2); + const row_bytes = h.n * 2; + const alone = try a.alloc(u8, rows * row_bytes); + defer a.free(alone); + const got = try a.alloc(u8, rows * row_bytes); + defer a.free(got); + for (0..rows) |i| { + try l.decode(s, w, .{ .x = xd + i * h.k * 2 }, out, 1); + try s.synchronize(); + try rig.get(out, alone[i * row_bytes ..][0..row_bytes]); + } + for (1..E.window + 1) |m| { + try l.decode(s, w, .{ .x = xd }, out, m); + try s.synchronize(); + try rig.get(out, got[0 .. m * row_bytes]); + try check.sameBytes("exl3 decode window", got[0 .. m * row_bytes], alone[0 .. m * row_bytes]); + } + const p_alone = try a.alloc(u8, prompt_rows * row_bytes); + defer a.free(p_alone); + const p_got = try a.alloc(u8, prompt_rows * row_bytes); + defer a.free(p_got); + for (0..prompt_rows) |i| { + try l.prompt(s, w, .{ .x = xd + i * h.k * 2 }, out, 1); + try s.synchronize(); + try rig.get(out, p_alone[i * row_bytes ..][0..row_bytes]); + } + for ([_]usize{ 2, 63, 64, 65, 255, 256, 257, prompt_rows }) |m| { + try l.prompt(s, w, .{ .x = xd }, out, m); + try s.synchronize(); + try rig.get(out, p_got[0 .. m * row_bytes]); + try check.sameBytes("exl3 prompt rows", p_got[0 .. m * row_bytes], p_alone[0 .. m * row_bytes]); + } + check.pass("decode windows of 1-16 rows ({d} input parts) and prompts of 2-{d} rows through qlinear: every row's bf16 bytes equal the row alone on its path", .{ E.split(h.k, h.n), prompt_rows }); +} + +/// The checkpoint's real matrices: FP64 on three column blocks, then GB/s at 1 and 16 rows. +fn real(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream, dir: []const u8) !void { + const a = gpu.gpa; + const index_path = try std.fmt.allocPrint(a, "{s}/model.safetensors.index.json", .{dir}); + defer a.free(index_path); + const index_bytes = try std.Io.Dir.cwd().readFileAlloc(gpu.io, index_path, a, .limited(64 << 20)); + defer a.free(index_bytes); + const index = try std.json.parseFromSlice(std.json.Value, a, index_bytes, .{}); + defer index.deinit(); + const map = index.value.object.get("weight_map").?.object; + const P = "model.language_model.layers."; + const names = [_][]const u8{ P ++ "3.mlp.experts.0.gate_proj", P ++ "3.mlp.experts.0.down_proj", P ++ "0.mlp.down_proj", P ++ "3.mlp.shared_experts.up_proj", P ++ "3.self_attn.q_b_proj", P ++ "3.self_attn.o_proj", "lm_head" }; + var b: Bound = .{}; + for (names) |name| { + const tname = try std.fmt.allocPrint(a, "{s}.trellis", .{name}); + defer a.free(tname); + const shard = map.get(tname) orelse return error.TensorNotInIndex; + const path = try std.fmt.allocPrint(a, "{s}/{s}", .{ dir, shard.string }); + defer a.free(path); + var f = try st.File.open(a, gpu.io, path); + defer f.close(gpu.io); + try realOne(gpu, k, s, &f, name, &b); + } + check.pass("{d} real matrices: {d} FP64 values, max error {e}, max {d:.4} of the bound, worst column {e} of the largest output", .{ names.len, b.values, b.max_error, b.max_fraction, b.worst_column }); +} + +fn tensor(f: *const st.File, name: []const u8, part: []const u8) ?st.Tensor { + var buf: [256]u8 = undefined; + return f.get(std.fmt.bufPrint(&buf, "{s}.{s}", .{ name, part }) catch return null); +} + +fn realOne(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream, f: *const st.File, name: []const u8, b: *Bound) !void { + const a = gpu.gpa; + const t = tensor(f, name, "trellis") orelse return error.MissingTrellis; + const suh_t = tensor(f, name, "suh") orelse return error.MissingSuh; + const svh_t = tensor(f, name, "svh") orelse return error.MissingSvh; + const book: E.Codebook = if (tensor(f, name, "mul1") != null) .mul1 else if (tensor(f, name, "mcg") != null) .mcg else .inst3; + if (t.dtype != .i16 or t.rank != 3 or t.dim(2) % 16 != 0) return error.UnsupportedTrellis; + const kk = t.dim(0) * 16; + const n = t.dim(1) * 16; + const bits: u32 = @intCast(t.dim(2) / 16); + const words: []align(1) const i16 = std.mem.bytesAsSlice(i16, t.bytes); + const suh: []align(1) const f16 = std.mem.bytesAsSlice(f16, suh_t.bytes); + const svh: []align(1) const f16 = std.mem.bytesAsSlice(f16, svh_t.bytes); + var rig: Rig = .{ .gpu = gpu }; + defer rig.deinit(); + const m: E.Matrix = .{ .trellis = try rig.put(t.bytes), .suh = try rig.put(suh_t.bytes), .svh = try rig.put(svh_t.bytes), .k = kk, .n = n, .bits = bits, .codebook = book }; + var prng = std.Random.DefaultPrng.init(0xF1A5); + const rows = 4; + const prompt_rows = 512; + const x = try randomRows(a, prng.random(), prompt_rows, kk); + defer a.free(x); + const xd = try rig.put(x); + const out = try rig.zeros(prompt_rows * n * 4); + try k.rows(s, xd, kk, m, out, rows, true); + try s.synchronize(); + const got = try a.alloc(f32, rows * n); + defer a.free(got); + try rig.get(out, std.mem.sliceAsBytes(got)); + // three 128-column blocks (first, middle, last), each a trellis of its own 8 tile columns + const tile_words = 16 * @as(usize, bits); + const sub = try a.alloc(i16, kk / 16 * 8 * tile_words); + defer a.free(sub); + const suh_c = try a.alloc(f16, kk); + defer a.free(suh_c); + for (suh_c, suh) |*d, v| d.* = v; + const svh_c = try a.alloc(f16, 128); + defer a.free(svh_c); + const y = try a.alloc(f64, rows * 128); + defer a.free(y); + const bound = try a.alloc(f64, rows * 128); + defer a.free(bound); + const g = try a.alloc(f32, rows * 128); + defer a.free(g); + const prompt_got = try a.alloc(f32, rows * n); + defer a.free(prompt_got); + + try k.prompt(s, xd, kk, m, out, rows); + try s.synchronize(); + try getBf16(&rig, out, prompt_got); + const blocks = n / 128; + for ([_]usize{ 0, blocks / 2, blocks - 1 }) |ob| { + for (0..kk / 16) |tk| for (0..8) |c| for (0..tile_words) |wi| { + sub[(tk * 8 + c) * tile_words + wi] = words[(tk * (n / 16) + 8 * ob + c) * tile_words + wi]; + }; + for (svh_c, 0..) |*v, j| v.* = svh[ob * 128 + j]; + const sub_h: Host = .{ .k = kk, .n = 128, .bits = bits, .book = book, .trellis = sub, .suh = suh_c, .svh = svh_c }; + try reference(a, sub_h, x, rows, y, bound, true); + for (0..rows) |i| @memcpy(g[i * 128 ..][0..128], got[i * n + ob * 128 ..][0..128]); + try compare(name, g, y, bound, rows, 128, b); + for (0..rows) |i| @memcpy(g[i * 128 ..][0..128], prompt_got[i * n + ob * 128 ..][0..128]); + try compare(name, g, y, bound, rows, 128, b); + } + const one = try time(gpu, k, s, xd, m, out, 1); + const sixteen = try time(gpu, k, s, xd, m, out, E.window); + const many = try timePrompt(gpu, k, s, xd, m, out, prompt_rows); + const gb = @as(f64, @floatFromInt(t.bytes.len)) / 1e9; + check.pass("{s}: K={d} N={d} {d} bits {t}, 3 column blocks within FP64 bounds (decode and prompt); 1 row {d:.1} us ({d:.0} GB/s), 16 rows {d:.1} us, a {d}-row prompt {d:.0} rows/s", .{ name, kk, n, bits, book, one * 1e6, gb / one, sixteen * 1e6, prompt_rows, prompt_rows / many }); +} + +/// Seconds a call, the median of 7 runs of 20 after a warm-up, bf16 out. +fn time(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream, x: u64, m: E.Matrix, out: u64, rows: usize) !f64 { + try k.rows(s, x, m.k, m, out, rows, false); + try s.synchronize(); + var runs: [7]f64 = undefined; + for (&runs) |*r| { + const t0 = check.now(gpu.io); + for (0..20) |_| try k.rows(s, x, m.k, m, out, rows, false); + try s.synchronize(); + r.* = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 20e9; + } + return check.median(&runs); +} + +/// Seconds a prompt call, the median of 7 runs of 5 after a warm-up. +fn timePrompt(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream, x: u64, m: E.Matrix, out: u64, rows: usize) !f64 { + try k.prompt(s, x, m.k, m, out, rows); + try s.synchronize(); + var runs: [7]f64 = undefined; + for (&runs) |*r| { + const t0 = check.now(gpu.io); + for (0..5) |_| try k.prompt(s, x, m.k, m, out, rows); + try s.synchronize(); + r.* = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 5e9; + } + return check.median(&runs); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 7b8a089947..f6b99a8ca3 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -27,6 +27,7 @@ const exlr8_tests = @import("exlr8_tests.zig"); const exlr8_dense_tests = @import("exlr8_dense_tests.zig"); const norms_tests = @import("norms_tests.zig"); const norm_rows_tests = @import("norm_rows_tests.zig"); +const exl3_tests = @import("exl3_tests.zig"); const usage = \\usage: tf-cuda-test @@ -64,6 +65,7 @@ const usage = \\ exlr8 EXLR8 projections: every codebook state, FP64, row windows, the Metal frozen digests \\ exlr8-bench DIR [MANIFEST] [LAYER] [SETS] a real EXLR8 layer: FP64, row windows, then GB/s by stage \\ exlr8-dense [DIR] EXLR8 dense formats: decode bit-equal, FP64, row windows; DIR adds the real matrices + \\ exl3 [DIR] EXL3 projections: codebooks and tiles bit-equal, FP64, row windows; DIR adds real matrices \\ train-pack MODEL EXAMPLES [COUNT] training passes holding several sequences against one a pass: differences, time \\ ; @@ -121,6 +123,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8")) return exlr8_tests.run(gpu); + if (std.mem.eql(u8, cmd, "exl3")) return exl3_tests.run(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exlr8-dense")) return exlr8_dense_tests.run(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exlr8-bench")) return exlr8_tests.bench(gpu, try arg(rest, 0), if (rest.len > 1) rest[1] else "home", if (rest.len > 2) try std.fmt.parseInt(u16, rest[2], 10) else 40, if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 4); if (std.mem.eql(u8, cmd, "sample")) return sample_tests.draws(gpu); From a35cc23141c9c819ec09b8415e918029c11ccf1a Mon Sep 17 00:00:00 2001 From: BHCC2025 <241330250+BHCC2025@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:35:35 -0500 Subject: [PATCH 3/4] tests: exl3 skips a selected matrix a checkpoint stores as plain .weight, and fails when none is EXL3 --- zig/tests/cuda/exl3_tests.zig | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/zig/tests/cuda/exl3_tests.zig b/zig/tests/cuda/exl3_tests.zig index 139468e77f..54702a292e 100644 --- a/zig/tests/cuda/exl3_tests.zig +++ b/zig/tests/cuda/exl3_tests.zig @@ -374,17 +374,27 @@ fn real(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream, dir: []const u8) !void { const P = "model.language_model.layers."; const names = [_][]const u8{ P ++ "3.mlp.experts.0.gate_proj", P ++ "3.mlp.experts.0.down_proj", P ++ "0.mlp.down_proj", P ++ "3.mlp.shared_experts.up_proj", P ++ "3.self_attn.q_b_proj", P ++ "3.self_attn.o_proj", "lm_head" }; var b: Bound = .{}; + var checked: usize = 0; for (names) |name| { const tname = try std.fmt.allocPrint(a, "{s}.trellis", .{name}); defer a.free(tname); - const shard = map.get(tname) orelse return error.TensorNotInIndex; + const wname = try std.fmt.allocPrint(a, "{s}.weight", .{name}); + defer a.free(wname); + const shard = map.get(tname) orelse { + // a checkpoint may quantize only some matrices (experts-only EXL3): a plain .weight is skipped, a missing tensor fails + if (map.get(wname) == null) return error.TensorNotInIndex; + std.debug.print("SKIP exl3 {s}: not EXL3 in this checkpoint (.weight)\n", .{name}); + continue; + }; const path = try std.fmt.allocPrint(a, "{s}/{s}", .{ dir, shard.string }); defer a.free(path); var f = try st.File.open(a, gpu.io, path); defer f.close(gpu.io); try realOne(gpu, k, s, &f, name, &b); + checked += 1; } - check.pass("{d} real matrices: {d} FP64 values, max error {e}, max {d:.4} of the bound, worst column {e} of the largest output", .{ names.len, b.values, b.max_error, b.max_fraction, b.worst_column }); + try check.expect(checked > 0, "exl3: none of the {d} selected matrices is EXL3 in this checkpoint", .{names.len}); + check.pass("{d} real matrices: {d} FP64 values, max error {e}, max {d:.4} of the bound, worst column {e} of the largest output", .{ checked, b.values, b.max_error, b.max_fraction, b.worst_column }); } fn tensor(f: *const st.File, name: []const u8, part: []const u8) ?st.Tensor { From ff00ed307ddefe854c5178fe3b2833aac919bac6 Mon Sep 17 00:00:00 2001 From: BHCC2025 <241330250+BHCC2025@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:49:05 -0500 Subject: [PATCH 4/4] cuda: EXL3 projections write fp32 too (qlinear decodeF32 and promptF32) A tensor-parallel family sums its ranks' partial rows before rounding, so an EXL3 projection must hand out its fp32 sums as #757 does for NVFP4. The decode kernel already takes an fp32 output flag; the prompt kernel gains the same flag (tf_exl3_store), and qlinear's decodeF32 and promptF32 route .exl3 to them. decode and prompt keep their bf16 bits: the bf16 out is the fp32 out rounded. The new exl3 test checks both entries within FP64 bounds, the same bits at every row count, and the bf16 out as their rounding; it fails with EpilogueNotBuilt before this change. --- zig/kernels/cuda/exl3.cu | 23 +++++++------ zig/src/cuda/exl3_linear.zig | 8 +++-- zig/src/cuda/qlinear.zig | 8 +++-- zig/tests/cuda/exl3_tests.zig | 65 ++++++++++++++++++++++++++++++++--- 4 files changed, 83 insertions(+), 21 deletions(-) diff --git a/zig/kernels/cuda/exl3.cu b/zig/kernels/cuda/exl3.cu index a27f1b7486..fdd523a323 100644 --- a/zig/kernels/cuda/exl3.cu +++ b/zig/kernels/cuda/exl3.cu @@ -95,7 +95,7 @@ extern "C" __global__ void __launch_bounds__(32) tf_exl3_rotate_h(const __nv_bfl // 64 prompt rows x 128 columns a block: each tile decodes once into mma.m16n8k16 B fragments (their value order). template __device__ __forceinline__ void tf_exl3_prompt_job(const __half* Xh, unsigned rows, const uint8_t* T, const __half* svh, - const TfExl3Matrix m, __nv_bfloat16* Y, float* buf) { + const TfExl3Matrix m, uint8_t* Y, unsigned y_f32, float* buf) { constexpr unsigned TB = 32 * K; const unsigned tid = threadIdx.x, sg = tid / 32, lane = tid % 32, g8 = lane / 4, t4 = lane % 4, ob = blockIdx.x; const unsigned row0 = blockIdx.y * 64, kb = m.k / 128; @@ -148,7 +148,7 @@ __device__ __forceinline__ void tf_exl3_prompt_job(const __half* Xh, unsigned ro tf_exl3_fwht(v, lane); for (int q = 0; q < 4; ++q) { const unsigned c = ob * 128 + lane + 32 * q; - Y[(uint64_t)(row0 + i) * m.n + c] = __float2bfloat16_rn(__fmul_rn(v[q], __fmul_rn(__half2float(svh[c]), 0.125f))); + tf_exl3_store(Y, (uint64_t)(row0 + i) * m.n + c, __fmul_rn(v[q], __fmul_rn(__half2float(svh[c]), 0.125f)), y_f32); } } } @@ -157,17 +157,18 @@ __device__ __forceinline__ void tf_exl3_prompt_job(const __half* Xh, unsigned ro #define TF_EXL3_PROMPT(C) \ extern "C" __global__ void __launch_bounds__(256) tf_exl3_prompt_##C(const __half* Xh, unsigned rows, \ const uint8_t* T, const __half* svh, \ - const TfExl3Matrix m, __nv_bfloat16* Y) { \ + const TfExl3Matrix m, uint8_t* Y, \ + unsigned y_f32) { \ __shared__ __align__(16) float buf[64 * 128]; \ switch (m.bits) { \ - case 1: tf_exl3_prompt_job<1, C>(Xh, rows, T, svh, m, Y, buf); break; \ - case 2: tf_exl3_prompt_job<2, C>(Xh, rows, T, svh, m, Y, buf); break; \ - case 3: tf_exl3_prompt_job<3, C>(Xh, rows, T, svh, m, Y, buf); break; \ - case 4: tf_exl3_prompt_job<4, C>(Xh, rows, T, svh, m, Y, buf); break; \ - case 5: tf_exl3_prompt_job<5, C>(Xh, rows, T, svh, m, Y, buf); break; \ - case 6: tf_exl3_prompt_job<6, C>(Xh, rows, T, svh, m, Y, buf); break; \ - case 7: tf_exl3_prompt_job<7, C>(Xh, rows, T, svh, m, Y, buf); break; \ - case 8: tf_exl3_prompt_job<8, C>(Xh, rows, T, svh, m, Y, buf); break; \ + case 1: tf_exl3_prompt_job<1, C>(Xh, rows, T, svh, m, Y, y_f32, buf); break; \ + case 2: tf_exl3_prompt_job<2, C>(Xh, rows, T, svh, m, Y, y_f32, buf); break; \ + case 3: tf_exl3_prompt_job<3, C>(Xh, rows, T, svh, m, Y, y_f32, buf); break; \ + case 4: tf_exl3_prompt_job<4, C>(Xh, rows, T, svh, m, Y, y_f32, buf); break; \ + case 5: tf_exl3_prompt_job<5, C>(Xh, rows, T, svh, m, Y, y_f32, buf); break; \ + case 6: tf_exl3_prompt_job<6, C>(Xh, rows, T, svh, m, Y, y_f32, buf); break; \ + case 7: tf_exl3_prompt_job<7, C>(Xh, rows, T, svh, m, Y, y_f32, buf); break; \ + case 8: tf_exl3_prompt_job<8, C>(Xh, rows, T, svh, m, Y, y_f32, buf); break; \ default: break; \ } \ } diff --git a/zig/src/cuda/exl3_linear.zig b/zig/src/cuda/exl3_linear.zig index b254b9fb79..c894f33096 100644 --- a/zig/src/cuda/exl3_linear.zig +++ b/zig/src/cuda/exl3_linear.zig @@ -117,8 +117,9 @@ pub const Kernels = struct { } } - /// Prompt rows (docs/recipes/exl3.md, Prompts): fp16 rows on the tensor units, fp32 sums, the same at every count. - pub fn prompt(k: *const Kernels, s: Stream, x: u64, ldx: usize, m: Matrix, out: u64, count: usize) !void { + /// Prompt rows (docs/recipes/exl3.md, Prompts): fp16 rows on the tensor units, fp32 sums, bf16 or fp32 out. + pub fn prompt(k: *const Kernels, s: Stream, x: u64, ldx: usize, m: Matrix, out: u64, count: usize, f32_out: bool) !void { + const elem: usize = if (f32_out) 4 else 2; try k.fits(m, ldx); var at: usize = 0; while (at < count) : (at += prompt_chunk) { @@ -130,7 +131,8 @@ pub const Kernels = struct { a.add(m.trellis); a.add(m.svh); a.add(m.args()); - a.add(out + at * m.n * 2); + a.add(out + at * m.n * elem); + a.add(@as(u32, @intFromBool(f32_out))); try launch_.launch(k.prompt_fns[@backingInt(m.codebook)], .{ .grid = .{ .x = @intCast(m.n / 128), .y = @intCast((r + 63) / 64), .z = 1 }, .block = .{ .x = 256 } }, s, &a); } } diff --git a/zig/src/cuda/qlinear.zig b/zig/src/cuda/qlinear.zig index 7c6f7f3e65..5dfb6ce871 100644 --- a/zig/src/cuda/qlinear.zig +++ b/zig/src/cuda/qlinear.zig @@ -94,23 +94,25 @@ pub const Linear = struct { .nvfp4 => |q| return l.onPrompt(s, q, .fp4, in, out, rows), .fp8 => |q| return l.onPrompt(s, q, .fp8, in, out, rows), .bf16 => |q| return l.onLane(s, q, .bf16, in, out, rows), - .exl3 => |q| return (l.exl3 orelse return error.FormatNotLoaded).prompt(s, in.x, if (in.ldx == 0) q.k else in.ldx, q, out, rows), + .exl3 => |q| return (l.exl3 orelse return error.FormatNotLoaded).prompt(s, in.x, if (in.ldx == 0) q.k else in.ldx, q, out, rows, false), } } - /// Decode rows as decode(), out (rows, n) fp32 unrounded: a tensor-parallel rank's partial rows (nvfp4 only). + /// Decode rows as decode(), out (rows, n) fp32 unrounded: a tensor-parallel rank's partial rows (nvfp4, exl3). pub fn decodeF32(l: Linear, s: Stream, w: Weight, in: Rows, out: u64, rows: usize) !void { const q = switch (w) { .nvfp4 => |q| q, + .exl3 => |q| return (l.exl3 orelse return error.FormatNotLoaded).rows(s, in.x, if (in.ldx == 0) q.k else in.ldx, q, out, rows, true), else => return error.EpilogueNotBuilt, }; return l.onLaneAs(s, q, .fp4, in, out, .f32, rows); } - /// Prompt rows as prompt(), out (rows, n) fp32 unrounded (nvfp4 only). + /// Prompt rows as prompt(), out (rows, n) fp32 unrounded (nvfp4, exl3). pub fn promptF32(l: Linear, s: Stream, w: Weight, in: Rows, out: u64, rows: usize) !void { const q = switch (w) { .nvfp4 => |q| q, + .exl3 => |q| return (l.exl3 orelse return error.FormatNotLoaded).prompt(s, in.x, if (in.ldx == 0) q.k else in.ldx, q, out, rows, true), else => return error.EpilogueNotBuilt, }; return l.onPromptAs(s, q, .fp4, in, out, .f32, rows); diff --git a/zig/tests/cuda/exl3_tests.zig b/zig/tests/cuda/exl3_tests.zig index 54702a292e..7243552bcb 100644 --- a/zig/tests/cuda/exl3_tests.zig +++ b/zig/tests/cuda/exl3_tests.zig @@ -85,6 +85,7 @@ pub fn run(gpu: Gpu, dir: ?[]const u8) !void { for (cases, 0..) |c, i| try synthetic(gpu, &k, s, c, 0xE3 + i, &b); check.pass("{d} synthetic matrices, widths 1-8, three codebooks: {d} FP64 values, max error {e}, max {d:.4} of the bound, worst column {e} of the largest output", .{ cases.len, b.values, b.max_error, b.max_fraction, b.worst_column }); try windows(gpu, &k, s); + try fp32Out(gpu, &k, s); if (dir) |d| try real(gpu, &k, s, d); } @@ -297,7 +298,7 @@ fn synthetic(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream, c: Case, seed: u64, try reference(a, h, x, rows, y, bound, true); try compare("exl3 synthetic", got, y, bound, rows, c.n, b); // the same rows as a prompt: the tensor units' arithmetic, bf16 out - try k.prompt(s, try rig.put(x), c.k, m, out, rows); + try k.prompt(s, try rig.put(x), c.k, m, out, rows, false); try s.synchronize(); try getBf16(&rig, out, got); try reference(a, h, x, rows, y, bound, true); @@ -361,6 +362,62 @@ fn windows(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream) !void { check.pass("decode windows of 1-16 rows ({d} input parts) and prompts of 2-{d} rows through qlinear: every row's bf16 bytes equal the row alone on its path", .{ E.split(h.k, h.n), prompt_rows }); } +/// decodeF32 and promptF32 through qlinear: within FP64 bounds, the same bits at every row count, the bf16 out their rounding. +fn fp32Out(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream) !void { + const a = gpu.gpa; + var prng = std.Random.DefaultPrng.init(0xF32); + const h = try randomMatrix(a, prng.random(), .{ .k = 2048, .n = 256, .bits = 4, .book = .mcg }); + defer freeHost(a, h); + const rows = 300; + const x = try randomRows(a, prng.random(), rows, h.k); + defer a.free(x); + var rig: Rig = .{ .gpu = gpu }; + defer rig.deinit(); + const w: cuda.qlinear.Weight = .{ .exl3 = try upload(&rig, h) }; + const l: cuda.qlinear.Linear = .{ .exl3 = k }; + const in: cuda.qlinear.Rows = .{ .x = try rig.put(x) }; + const o32 = try rig.zeros(rows * h.n * 4); + const o16 = try rig.zeros(rows * h.n * 2); + const wide = try a.alloc(f32, rows * h.n); + defer a.free(wide); + const got = try a.alloc(f32, rows * h.n); + defer a.free(got); + const b16 = try a.alloc(u16, rows * h.n); + defer a.free(b16); + const y = try a.alloc(f64, rows * h.n); + defer a.free(y); + const bound = try a.alloc(f64, rows * h.n); + defer a.free(bound); + var b: Bound = .{}; + for ([_]bool{ false, true }) |prompt| { + const what = if (prompt) "exl3 promptF32" else "exl3 decodeF32"; + const top: usize = if (prompt) rows else E.window; + const counts: []const usize = if (prompt) &.{ 1, 63, 64, 65, 257 } else &.{ 1, 2, 3, 7, 8, 15 }; + if (prompt) try l.promptF32(s, w, in, o32, top) else try l.decodeF32(s, w, in, o32, top); + if (prompt) try l.prompt(s, w, in, o16, top) else try l.decode(s, w, in, o16, top); + try s.synchronize(); + try rig.get(o32, std.mem.sliceAsBytes(wide[0 .. top * h.n])); + try rig.get(o16, std.mem.sliceAsBytes(b16[0 .. top * h.n])); + try reference(a, h, x, top, y, bound, true); + try compare(what, wide, y, bound, top, h.n, &b); + var differ: usize = 0; + for (wide[0 .. top * h.n], b16[0 .. top * h.n]) |v, o| differ += @intFromBool(roundBf16(v) != o); + try check.expect(differ == 0, "{s}: the bf16 out is the fp32 out rounded ({d} of {d} differ)", .{ what, differ, top * h.n }); + for (counts) |m| { + if (prompt) try l.promptF32(s, w, in, o32, m) else try l.decodeF32(s, w, in, o32, m); + try s.synchronize(); + try rig.get(o32, std.mem.sliceAsBytes(got[0 .. m * h.n])); + try check.sameBytes(what, std.mem.sliceAsBytes(got[0 .. m * h.n]), std.mem.sliceAsBytes(wide[0 .. m * h.n])); + } + } + check.pass("decodeF32 (1-{d} rows) and promptF32 (1-{d} rows) through qlinear: {d} FP64 values within bounds (max {d:.4} of the bound), the same bits at every row count, the bf16 out their rounding", .{ E.window, rows, b.values, b.max_fraction }); +} + +fn roundBf16(v: f32) u16 { + const w: u32 = @bitCast(v); + return @intCast((w + 0x7FFF + ((w >> 16) & 1)) >> 16); +} + /// The checkpoint's real matrices: FP64 on three column blocks, then GB/s at 1 and 16 rows. fn real(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream, dir: []const u8) !void { const a = gpu.gpa; @@ -448,7 +505,7 @@ fn realOne(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream, f: *const st.File, nam const prompt_got = try a.alloc(f32, rows * n); defer a.free(prompt_got); - try k.prompt(s, xd, kk, m, out, rows); + try k.prompt(s, xd, kk, m, out, rows, false); try s.synchronize(); try getBf16(&rig, out, prompt_got); const blocks = n / 128; @@ -487,12 +544,12 @@ fn time(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream, x: u64, m: E.Matrix, out: /// Seconds a prompt call, the median of 7 runs of 5 after a warm-up. fn timePrompt(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream, x: u64, m: E.Matrix, out: u64, rows: usize) !f64 { - try k.prompt(s, x, m.k, m, out, rows); + try k.prompt(s, x, m.k, m, out, rows, false); try s.synchronize(); var runs: [7]f64 = undefined; for (&runs) |*r| { const t0 = check.now(gpu.io); - for (0..5) |_| try k.prompt(s, x, m.k, m, out, rows); + for (0..5) |_| try k.prompt(s, x, m.k, m, out, rows, false); try s.synchronize(); r.* = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 5e9; }