Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions zig/kernels/cuda/prompt16.cu
Original file line number Diff line number Diff line change
Expand Up @@ -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<FP4, 128, 128, 2, 2, 2, false>(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<FP4, 128, 128, 2, 2, 2, true>(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<FP8, 128, 128, 2, 2, 2, false>(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, int, int, int, int, int, int);
}
8 changes: 8 additions & 0 deletions zig/kernels/cuda/qmmf.cu
Original file line number Diff line number Diff line change
Expand Up @@ -308,6 +308,13 @@ template __global__ void qmmf_kernel<FP4, 32, 64, 1, 4, 4, false, true, false>(c
template __global__ void qmmf_kernel<FP4, 64, 64, 1, 4, 4, false, false, false>(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<FP4, 64, 64, 1, 4, 4, false, true, false>(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<FP4, 64, 64, 1, 4, 4, false, false, true>(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<FP4, 16, 64, 1, 4, 4, true, false, false>(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<FP4, 16, 64, 1, 4, 4, true, true, false>(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<FP4, 32, 64, 1, 4, 4, true, false, false>(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<FP4, 32, 64, 1, 4, 4, true, true, false>(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<FP4, 64, 64, 1, 4, 4, true, false, false>(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<FP4, 64, 64, 1, 4, 4, true, true, false>(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<FP4, 64, 64, 1, 4, 4, true, false, true>(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<FP8, 16, 64, 1, 4, 4, false, false, false>(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<FP8, 16, 64, 1, 4, 4, false, true, false>(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<FP8, 32, 64, 1, 4, 4, false, false, false>(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int);
Expand All @@ -323,4 +330,5 @@ template __global__ void qmmf_kernel<BF16, 64, 64, 1, 4, 4, false, false, false>
template __global__ void qmmf_kernel<BF16, 64, 64, 1, 4, 4, false, true, false>(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<BF16, 64, 64, 1, 4, 4, false, false, true>(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<false>(const float* __restrict__, void* __restrict__, long long, int, float);
template __global__ void reduce_kernel<true>(const float* __restrict__, void* __restrict__, long long, int, float);
}
19 changes: 14 additions & 5 deletions zig/src/cuda/nvfp4.zig
Original file line number Diff line number Diff line change
Expand Up @@ -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<MODE, 128, 128, 2, 2, 2>::SMEM.
pub fn sharedBytes(mode: qmmf.Mode) u32 {
Expand All @@ -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 {
Expand All @@ -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;
Expand Down Expand Up @@ -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" {
Expand Down
33 changes: 31 additions & 2 deletions zig/src/cuda/qlinear.zig
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
};

Expand Down Expand Up @@ -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));
}
55 changes: 35 additions & 20 deletions zig/src/cuda/qmmf.zig
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand All @@ -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;
Expand Down Expand Up @@ -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);
}
}
};
Expand Down Expand Up @@ -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));
}
21 changes: 21 additions & 0 deletions zig/tests/cuda/quant_ref_tests.zig
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
Loading