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