From a055fb6c8dcdcc3f3121d11326a406e72f1fbfa2 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Sun, 11 Oct 2026 10:25:50 +0900 Subject: [PATCH] cuda: Nvfp4Experts.packGateUp packs gate and up in place, no scratch packGateUp packed each matrix into a scratch half and copied the halves into their interleaved places, two allocations and a copy per expert. A loader packs every routed expert of every layer this way (GLM at TP=2: 45 x 256 experts, 2.4 MB a half), and the fresh pages and the copies put about 20 s on the load. packOne's loop now takes the matrix's place among the stacked parts, so packGateUp writes gate and up straight where they go, and it needs no allocator. The words are the same: a host test checks packGateUp against packOne's blocks interleaved. --- zig/src/cuda/experts.zig | 50 ++++++++++++++++---------- zig/tests/cuda/nvfp4_experts_tests.zig | 2 +- zig/tests/cuda/quant_ref_experts.zig | 2 +- 3 files changed, 34 insertions(+), 20 deletions(-) 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];