diff --git a/zig/src/cuda/experts.zig b/zig/src/cuda/experts.zig index c940202bea..c5f48c462e 100644 --- a/zig/src/cuda/experts.zig +++ b/zig/src/cuda/experts.zig @@ -374,11 +374,22 @@ pub const Nvfp4Experts = struct { /// One expert's words [n, k/2] (low nibble first) and e4m3 scales [n, k/16] as blocks [n/32][k/32][144] (_pack). pub fn packOne(out: []u32, codes: []const u8, scales: []const u8, n: usize, k: usize) void { + packPart(out, 1, 0, codes, scales, n, k); + } + + /// Gate and up interleaved as make() stacks them, [n/32][k/32][2][144], packed in place (a loader packs thousands). + pub fn packGateUp(out: []u32, gate: [2][]const u8, up: [2][]const u8, n: usize, k: usize) void { + packPart(out, 2, 0, gate[0], gate[1], n, k); + packPart(out, 2, 1, up[0], up[1], n, k); + } + + /// packOne into matrix `part` of `parts` stacked per (column block, k32 group). + fn packPart(out: []u32, parts: usize, part: usize, codes: []const u8, scales: []const u8, n: usize, k: usize) void { const kg = k / 32; - std.debug.assert(n % cols == 0 and k % 32 == 0 and out.len == n / cols * kg * words); + std.debug.assert(n % cols == 0 and k % 32 == 0 and out.len == n / cols * kg * parts * words and part < parts); std.debug.assert(codes.len == n * k / 2 and scales.len == n * k / 16); for (0..n / cols) |cb| for (0..kg) |g| { - const block = out[(cb * kg + g) * words ..][0..words]; + const block = out[((cb * kg + g) * parts + part) * words ..][0..words]; for (0..8) |gq| for (0..4) |t| for (0..4) |j| { const row = cb * cols + j * 8 + gq; var w: u32 = 0; @@ -397,22 +408,6 @@ pub const Nvfp4Experts = struct { }; }; } - - /// Gate and up interleaved per (column block, k32 group) as make() stacks them: [n/32][k/32][2][144]. - pub fn packGateUp(gpa: std.mem.Allocator, out: []u32, gate: [2][]const u8, up: [2][]const u8, n: usize, k: usize) !void { - const half = n / cols * (k / 32) * words; - std.debug.assert(out.len == 2 * half); - const g = try gpa.alloc(u32, half); - defer gpa.free(g); - const u = try gpa.alloc(u32, half); - defer gpa.free(u); - packOne(g, gate[0], gate[1], n, k); - packOne(u, up[0], up[1], n, k); - for (0..half / words) |b| { - @memcpy(out[(2 * b) * words ..][0..words], g[b * words ..][0..words]); - @memcpy(out[(2 * b + 1) * words ..][0..words], u[b * words ..][0..words]); - } - } }; test "a packed nvfp4 block puts each nibble and scale where nvfp4/experts.py's _pack does" { @@ -435,6 +430,25 @@ test "a packed nvfp4 block puts each nibble and scale where nvfp4/experts.py's _ try std.testing.expectEqual(scales[44 * 4 + 3], sc[((2 * 2 + 1) * 4 + 1) * 2]); } +test "packGateUp stacks packOne's gate and up blocks per (column block, k32 group)" { + const n = 64; + const k = 64; + const words = Nvfp4Experts.words; + var codes: [2][n * k / 2]u8 = undefined; + var scales: [2][n * k / 16]u8 = undefined; + for (0..2) |m| { + for (&codes[m], 0..) |*c, i| c.* = @truncate(i * 29 + 3 + m * 101); + for (&scales[m], 0..) |*c, i| c.* = @truncate(i * 7 + 1 + m * 13); + } + var one: [2][n / 32 * (k / 32) * words]u32 = undefined; + for (0..2) |m| Nvfp4Experts.packOne(&one[m], &codes[m], &scales[m], n, k); + var both: [2 * one[0].len]u32 = undefined; + Nvfp4Experts.packGateUp(&both, .{ &codes[0], &scales[0] }, .{ &codes[1], &scales[1] }, n, k); + for (0..one[0].len / words) |b| for (0..2) |m| { + try std.testing.expectEqualSlices(u32, one[m][b * words ..][0..words], both[(2 * b + m) * words ..][0..words]); + }; +} + test "nvfp4 rows of either kind are refused when the format is not loaded" { const g: Grouped = .{}; const l: Layer = .{ .format = .nvfp4, .up = 0, .down = 0, .width = 64, .dims = 64, .experts = 2, .relu2 = true }; diff --git a/zig/tests/cuda/nvfp4_experts_tests.zig b/zig/tests/cuda/nvfp4_experts_tests.zig index 7c8b7bdb2e..2031017836 100644 --- a/zig/tests/cuda/nvfp4_experts_tests.zig +++ b/zig/tests/cuda/nvfp4_experts_tests.zig @@ -39,7 +39,7 @@ pub fn experts(gpu: Gpu, dir: []const u8) !void { for (0..e) |x| { const gate = [2][]const u8{ raw[0][0][x * wb ..][0..wb], raw[0][1][x * sbytes ..][0..sbytes] }; const up = [2][]const u8{ raw[1][0][x * wb ..][0..wb], raw[1][1][x * sbytes ..][0..sbytes] }; - try fx4.packGateUp(a, pu[x * 2 * blk ..][0 .. 2 * blk], gate, up, ni, d); + fx4.packGateUp(pu[x * 2 * blk ..][0 .. 2 * blk], gate, up, ni, d); fx4.packOne(pd[x * blk ..][0..blk], raw[2][0][x * wb ..][0..wb], raw[2][1][x * sbytes ..][0..sbytes], d, ni); } inline for (.{ .{ "packed_up", pu }, .{ "packed_down", pd } }) |c| { diff --git a/zig/tests/cuda/quant_ref_experts.zig b/zig/tests/cuda/quant_ref_experts.zig index 4b9dad472d..2b3971df02 100644 --- a/zig/tests/cuda/quant_ref_experts.zig +++ b/zig/tests/cuda/quant_ref_experts.zig @@ -195,7 +195,7 @@ pub fn run(r: *Rig) !void { const sb = NI * D / 16; const us = try r.a.alloc(f32, 2 * E); for (0..E) |e| { - try W.packGateUp(r.a, pu[e * 2 * blk ..][0 .. 2 * blk], .{ gm.codes[e * wb ..][0..wb], gm.scales[e * sb ..][0..sb] }, .{ um.codes[e * wb ..][0..wb], um.scales[e * sb ..][0..sb] }, NI, D); + W.packGateUp(pu[e * 2 * blk ..][0 .. 2 * blk], .{ gm.codes[e * wb ..][0..wb], gm.scales[e * sb ..][0..sb] }, .{ um.codes[e * wb ..][0..wb], um.scales[e * sb ..][0..sb] }, NI, D); W.packOne(pd[e * blk ..][0..blk], dm.codes[e * wb ..][0..wb], dm.scales[e * sb ..][0..sb], D, NI); us[2 * e] = gm.global[e]; us[2 * e + 1] = um.global[e];