cuda: Nvfp4Experts.packGateUp packs gate and up in place, no scratch

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 };

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| {

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];