cuda qmmf: half tiles for decode calls whose 64-column grid leaves SMs idle

zig/kernels/cuda/qmmf.cu

@@ -64,7 +64,7 @@ __global__ void __launch_bounds__(WM * WN * 32) qmmf_kernel(
         float scale, void* __restrict__ out, float* __restrict__ part, int M, int N, int K, int SK, int npad, int ldx,
         int group) {
     using T = Tile<MODE, BM, BN, WM, WN, STAGES>;
-    static_assert(BN == 64, "a block reads one 64-column tile of words and scales");
+    static_assert(BN == 64 || BN == 32, "a block reads one 64-column tile of words and scales, or its half");
     extern __shared__ __align__(128) unsigned char buf[];
     const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5;
     const int wm = warp / WN, wn = warp % WN;
@@ -83,6 +83,17 @@ __global__ void __launch_bounds__(WM * WN * 32) qmmf_kernel(
         }
         unsigned char* pw = p + T::X;
         constexpr int TILE_BYTES = MODE == FP4 ? 64 * GS / 2 : MODE == BF16 ? 64 * GS * 2 : 64 * GS;
+        if constexpr (BN == 32) {
+            // half a packed 64-column tile: its n8 tiles 0-3 or 4-7, contiguous in every format's word order
+            const int half = (n0 % 64) / 32;
+            for (int c = tid; c < T::W / 16; c += T::THREADS)
+                cp16(pw + c * 16, w + (static_cast<size_t>(n0 / 64) * KG + g) * TILE_BYTES + half * (TILE_BYTES / 2) + c * 16);
+            if constexpr (T::S > 0) {
+                unsigned char* ps = pw + T::W;
+                for (int c = tid; c < T::S / 16; c += T::THREADS)
+                    cp16(ps + c * 16, bs + (static_cast<size_t>(n0 / 64) * KG + g) * (2 * T::S) + half * T::S + c * 16);
+            }
+        } else {
         for (int c = tid; c < T::W / 16; c += T::THREADS) {
             const int t = c / (TILE_BYTES / 16), off = c % (TILE_BYTES / 16);
             cp16(pw + c * 16, w + (static_cast<size_t>(n0 / 64 + t) * KG + g) * TILE_BYTES + off * 16);
@@ -92,6 +103,7 @@ __global__ void __launch_bounds__(WM * WN * 32) qmmf_kernel(
             for (int c = tid; c < T::S / 16; c += T::THREADS)
                 cp16(ps + c * 16, bs + (static_cast<size_t>(n0 / 64) * KG + g) * T::S + c * 16);
         }
+        }
     };
 
     float acc[T::MT][T::NT][4];
@@ -386,6 +398,17 @@ template __global__ void qmmf_kernel<FP8, 64, 64, 1, 4, 4, true, true, false>(co
 template __global__ void qmmf_kernel<FP8, 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<BF16, 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<BF16, 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);
+
+// Half tiles (32 columns a block) of the 16-row tile: the 64-column tile's sums, twice the blocks for a small grid.
+template __global__ void qmmf_kernel<MXFP8, 16, 32, 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<MXFP8, 16, 32, 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<MXFP8, 16, 32, 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<MXFP8, 16, 32, 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<BF16, 16, 32, 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<BF16, 16, 32, 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, 16, 32, 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<BF16, 16, 32, 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<BF16, 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<BF16, 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<BF16, 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);

zig/src/core/contract/lane_half_records.zig

@@ -0,0 +1,13 @@
+//! The lane's half tiles (16 rows by 32 columns, MXFP8 and BF16), in registry.cuda; nvcc 13.0.88 cubins.
+const CudaRecord = @import("registry.zig").CudaRecord;
+
+pub const cuda = [_]CudaRecord{
+    .{ .launch = .{ .name = "qmmf MXFP8 16-row half tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi16ELi32ELi1ELi4ELi4ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 16896 }, .usage = .{ .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 47, .static = 0, .bound = 128 }, .{ .regs = 47, .static = 0, .bound = 128 } } },
+    .{ .launch = .{ .name = "qmmf MXFP8 16-row half tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi16ELi32ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 16896, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 56, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 } } },
+    .{ .launch = .{ .name = "qmmf MXFP8 16-row half tile, fp32 out", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi16ELi32ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 16896 }, .usage = .{ .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 47, .static = 0, .bound = 128 }, .{ .regs = 47, .static = 0, .bound = 128 } } },
+    .{ .launch = .{ .name = "qmmf MXFP8 16-row half tile, fp32 out, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi16ELi32ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 16896, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 56, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 } } },
+    .{ .launch = .{ .name = "qmmf BF16 16-row half tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi4ELi16ELi32ELi1ELi4ELi4ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 24576 }, .usage = .{ .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 39, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 } } },
+    .{ .launch = .{ .name = "qmmf BF16 16-row half tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi4ELi16ELi32ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 24576, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 } } },
+    .{ .launch = .{ .name = "qmmf BF16 16-row half tile, fp32 out", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi4ELi16ELi32ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 24576 }, .usage = .{ .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 39, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 } } },
+    .{ .launch = .{ .name = "qmmf BF16 16-row half tile, fp32 out, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi4ELi16ELi32ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 24576, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 } } },
+};

zig/src/core/contract/registry.zig

@@ -221,7 +221,7 @@ pub const cuda = [_]CudaRecord{
     .{ .launch = .{ .name = "experts3_prefill down", .symbol = "_ZN19tf_experts3_prefill14prefill_kernelILi1ELi3ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_PvifS6_i", .image = "experts3_prefill", .threads = 256, .dynamic = 47104, .opt_in = true }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 } } },
     .{ .launch = .{ .name = "experts3 tf_experts3_widen", .symbol = "tf_experts3_widen", .image = "experts3", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 } } },
     .{ .launch = .{ .name = "experts3 tf_experts3_planes", .symbol = "tf_experts3_planes", .image = "experts3", .threads = 256 }, .usage = .{ .{ .regs = 26, .static = 0, .bound = 0 }, .{ .regs = 28, .static = 0, .bound = 0 }, .{ .regs = 28, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 } } },
-} ++ @import("registry_media.zig").cuda ++ @import("nemotron_prompt_records.zig").cuda ++ @import("glm_cuda_records.zig").cuda;
+} ++ @import("registry_media.zig").cuda ++ @import("nemotron_prompt_records.zig").cuda ++ @import("glm_cuda_records.zig").cuda ++ @import("lane_half_records.zig").cuda;
 /// The nvcc release `cuda`'s usage records come from: the compiler that builds the release fatbins.
 pub const cuda_release: limits.Nvcc = .{ .major = 13, .minor = 0 };
 /// One launch's usage under a later nvcc, by the record's name.

zig/src/cuda/qmmf.zig

@@ -79,6 +79,37 @@ fn symbol(
     return std.fmt.comptimePrint(name, .{ @backingInt(mode), bm, flags[0], flags[1], flags[2] });
 }
 
+/// The 16-row tile at 32 columns a block (MXFP8 and BF16): plain, then with a cluster.
+fn halfSymbol(comptime mode: Mode, comptime cluster: bool, comptime wide: bool) [:0]const u8 {
+    const name = "_ZN7tf_qmmf11qmmf_kernelILi{d}ELi16ELi32ELi1ELi4ELi4ELb{d}ELb{d}ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii";
+    return std.fmt.comptimePrint(name, .{ @backingInt(mode), @intFromBool(wide), @intFromBool(cluster) });
+}
+
+/// Dynamic shared bytes of the half tile: a stage holds half the weights (and scales) of the 64-column tile's.
+pub fn halfShared(mode: Mode) u32 {
+    const w: u32 = if (mode == .bf16) 32 * group * 2 else 32 * group;
+    const s: u32 = if (mode == .mxfp8) 32 * 2 else 0;
+    const stage = (16 * group * 2 + w + s + 127) / 128 * 128;
+    const partials = 1 * 1 * 4 * threads * 4;
+    return @max(4 * stage, partials);
+}
+
+/// MXFP8/BF16 decode calls take half tiles below this many 64-column blocks (TF_QMMF_HALF_BELOW; 0 never): same sums.
+pub const half_below_default: usize = 96;
+
+fn halfBelow() usize {
+    const v = std.c.getenv("TF_QMMF_HALF_BELOW") orelse return half_below_default;
+    return std.fmt.parseInt(usize, std.mem.span(v), 10) catch half_below_default;
+}
+
+pub fn halfModeIndex(mode: Mode) ?usize {
+    return switch (mode) {
+        .mxfp8 => 0,
+        .bf16 => 1,
+        else => null,
+    };
+}
+
 pub const reduce_symbols = [2][:0]const u8{
     "_ZN7tf_qmmf13reduce_kernelILb0EEEvPKfPvxif",
     "_ZN7tf_qmmf13reduce_kernelILb1EEEvPKfPvxif",
@@ -105,6 +136,8 @@ pub const Lane = struct {
     fns: [modes.len][7]Function, // per mode: lane16, lane16c, lane32, lane32c, lane64, lane64c, fused
     f32_fns: [modes.len][7]Function,
     reduce: [2]Function, // by Out
+    half: [2][2][2]Function, // by halfModeIndex, Out, then plain / cluster
+    half_below: usize = half_below_default,
     clusters: bool, // sm_90 on: slices meet in a cluster's shared memory
 
     pub fn load(d: *const Driver, capability_major: u32) !Lane {
@@ -116,13 +149,20 @@ pub const Lane = struct {
             .fns = undefined,
             .f32_fns = undefined,
             .reduce = undefined,
+            .half = undefined,
             .clusters = capability_major >= 9,
         };
         inline for (modes, 0..) |mode, mi| {
             l.fns[mi] = try loadTiles(mod, mode, false);
             l.f32_fns[mi] = try loadTiles(mod, mode, true);
         }
         for (&l.reduce, reduce_symbols) |*f, name| f.* = try mod.function(name);
+        inline for (.{ Mode.mxfp8, Mode.bf16 }, 0..) |mode, mi| inline for (.{ false, true }, 0..) |wide, oi| inline for (.{ false, true }, 0..) |cl, ci| {
+            const f = try mod.function(halfSymbol(mode, cl, wide));
+            try f.allowDynamicShared(halfShared(mode));
+            l.half[mi][oi][ci] = f;
+        };
+        l.half_below = halfBelow();
         return l;
     }
 
@@ -152,8 +192,11 @@ pub const Lane = struct {
         if (p.reduce and part == 0) return error.MissingLanePart;
         const set = (if (out == .f32) l.f32_fns 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;
+        const hm = halfModeIndex(w.mode);
+        const half = hm != null and p.bm == 16 and !p.fused and rows_t * ((w.n + bn - 1) / bn) * p.sk < l.half_below;
+        const f = if (half) l.half[hm.?][@backingInt(out)][@intFromBool(p.cluster)] else if (p.fused) set[6] else set[ti * 2 + @intFromBool(p.cluster)];
+        const cols: usize = if (half) 32 else bn;
         const band = @max(1, @min(rows_t, (12 << 20) / (@as(usize, p.bm) * w.k * 2)));
         var args: launch_.Args = .{};
         args.add(x);
@@ -170,9 +213,9 @@ pub const Lane = struct {
         args.add(@as(i32, @intCast(if (m == 1) w.k else ldx)));
         args.add(@as(i32, @intCast(band)));
         try launch_.launch(f, .{
-            .grid = .{ .x = @intCast(rows_t * ((w.n + bn - 1) / bn)), .y = 1, .z = if (p.fused) 1 else p.sk },
+            .grid = .{ .x = @intCast(rows_t * ((w.n + cols - 1) / cols)), .y = 1, .z = if (p.fused) 1 else p.sk },
             .block = .{ .x = threads, .y = 1, .z = 1 },
-            .shared = sharedBytes(w.mode, p.bm),
+            .shared = if (half) halfShared(w.mode) else sharedBytes(w.mode, p.bm),
             .cluster = if (p.cluster) .{ .x = 1, .y = 1, .z = p.sk } else null,
         }, s, &args);
         if (p.reduce) {

zig/src/families/nemotron/cuda_sm89_test.zig

@@ -69,6 +69,8 @@ const Fake = struct {
             f.* = .{ .d = d, .handle = @ptrFromInt(17 + mi * 16 + i) };
         };
         for (&l.reduce, 0..) |*f, i| f.* = .{ .d = d, .handle = @ptrFromInt(90 + i) };
+        l.half = @splat(@splat(@splat(.{ .d = d, .handle = @ptrFromInt(120) })));
+        l.half_below = 0; // the released launch parameters: the 64-column tile's
         return l;
     }
     fn kernels(d: *const cuda.Driver, cluster: bool) kern.Kernels {

zig/src/families/nemotron/cuda_subset_batch_test.zig

@@ -192,7 +192,7 @@ test "restored head: 15 and 16 token transitions reach mixed shrinking real chai
     e.k.ck = null;
     e.k.ggml = null;
     const f: cuda.Function = .{ .d = &d, .handle = @ptrFromInt(1) };
-    const lane: cuda.qmmf.Lane = .{ .mod = undefined, .fns = @splat(@splat(f)), .f32_fns = @splat(@splat(f)), .reduce = @splat(f), .clusters = false };
+    const lane: cuda.qmmf.Lane = .{ .mod = undefined, .fns = @splat(@splat(f)), .f32_fns = @splat(@splat(f)), .reduce = @splat(f), .half = @splat(@splat(@splat(f))), .half_below = 0, .clusters = false };
     e.k.modelopt = .{ .lane = lane, .gemm = undefined, .experts = undefined, .pack = undefined };
     e.k.draws.draw_fn = .{ .d = &d, .handle = @ptrFromInt(2) };
     e.k.draws.draw_ids_fn = .{ .d = &d, .handle = @ptrFromInt(3) };