From 8a8a86e3245c112700774abc6b32b00748a511c7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:52:53 +0200 Subject: [PATCH 1/3] cuda qmmf: half tiles for decode calls whose 64-column grid leaves SMs idle A 16-row MXFP8 or BF16 lane call whose 64-column grid holds fewer than 96 blocks (TF_QMMF_HALF_BELOW, 0 never) takes a 32-column tile: half a packed tile's n8 tiles, contiguous in both formats' word and scale order, so each column's mma chain, K slices and cluster sum are the 64-column tile's and its bits the same, with twice the blocks. Flash Next's hyper-connection down projection ([324, 10240] BF16, 6 tiles x 8 slices) reads its weights at 183 GB/s instead of 158 (36.3 vs 42.0 us cold), the router and the shared expert's MXFP8 projections 4-8 % faster. The existing instances are untouched (their loads stay behind `if constexpr (BN == 64)`). tf-cuda-test lane-half: MXFP8 and BF16, fp32 and bf16 out, bytes equal to the 64-column tile at Flash Next's decode shapes and the column edges, FP64-bounded, timed on weights past the L2. Served (FP4-Local, no drafts): 23.95 / 23.83 tok/s one stream (main 23.58 / 23.66), replies equal. --- zig/kernels/cuda/qmmf.cu | 25 +++- zig/src/cuda/qmmf.zig | 49 +++++++- .../nemotron/cuda_subset_batch_test.zig | 2 +- zig/tests/cuda/lane_half_tests.zig | 109 ++++++++++++++++++ zig/tests/cuda/main.zig | 3 + 5 files changed, 183 insertions(+), 5 deletions(-) create mode 100644 zig/tests/cuda/lane_half_tests.zig diff --git a/zig/kernels/cuda/qmmf.cu b/zig/kernels/cuda/qmmf.cu index 6ef58c2406..46291eda63 100644 --- a/zig/kernels/cuda/qmmf.cu +++ b/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; - 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(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(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(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(n0 / 64) * KG + g) * T::S + c * 16); } + } }; float acc[T::MT][T::NT][4]; @@ -386,6 +398,17 @@ template __global__ void qmmf_kernel(co template __global__ void qmmf_kernel(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(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(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(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(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(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(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(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(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(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(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(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(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(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); diff --git a/zig/src/cuda/qmmf.zig b/zig/src/cuda/qmmf.zig index 166b057750..0b4774e55d 100644 --- a/zig/src/cuda/qmmf.zig +++ b/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,6 +149,7 @@ pub const Lane = struct { .fns = undefined, .f32_fns = undefined, .reduce = undefined, + .half = undefined, .clusters = capability_major >= 9, }; inline for (modes, 0..) |mode, mi| { @@ -123,6 +157,12 @@ pub const Lane = struct { 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) { diff --git a/zig/src/families/nemotron/cuda_subset_batch_test.zig b/zig/src/families/nemotron/cuda_subset_batch_test.zig index 5e99addd94..ca131f480d 100644 --- a/zig/src/families/nemotron/cuda_subset_batch_test.zig +++ b/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) }; diff --git a/zig/tests/cuda/lane_half_tests.zig b/zig/tests/cuda/lane_half_tests.zig new file mode 100644 index 0000000000..34500e5cfa --- /dev/null +++ b/zig/tests/cuda/lane_half_tests.zig @@ -0,0 +1,109 @@ +//! The lane's half tiles (16x32, MXFP8 and BF16): the 64-column tile's bytes, FP64-bounded, timed on DRAM weights. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const rig = @import("flashnext_rig.zig"); + +const Mode = cuda.qmmf.Mode; +const Shape = struct { mode: Mode, m: usize, n: usize, k: usize }; + +/// Flash Next decode projections whose 64-column grids leave GB10's SMs idle, the 64-column edge, wider row counts. +const shapes = [_]Shape{ + .{ .mode = .bf16, .m = 1, .n = 324, .k = 10240 }, .{ .mode = .bf16, .m = 1, .n = 320, .k = 10240 }, + .{ .mode = .bf16, .m = 1, .n = 512, .k = 2560 }, .{ .mode = .mxfp8, .m = 1, .n = 640, .k = 2560 }, + .{ .mode = .mxfp8, .m = 1, .n = 2560, .k = 640 }, .{ .mode = .bf16, .m = 7, .n = 324, .k = 10240 }, + .{ .mode = .mxfp8, .m = 16, .n = 640, .k = 2560 }, .{ .mode = .bf16, .m = 3, .n = 96, .k = 512 }, + .{ .mode = .mxfp8, .m = 5, .n = 32, .k = 128 }, .{ .mode = .bf16, .m = 1, .n = 10240, .k = 320 }, + .{ .mode = .mxfp8, .m = 1, .n = 2560, .k = 6144 }, .{ .mode = .mxfp8, .m = 1, .n = 6144, .k = 2560 }, +}; + +const Built = struct { w: cuda.qmmf.Weight, deq: []f64 }; + +/// A random weight in `mode`, packed for the lane, and its exact float64 values [n][k]. +fn build(r: *rig.Rig, a: std.mem.Allocator, rng: std.Random, sh: Shape) !Built { + const deq = try a.alloc(f64, sh.n * sh.k); + if (sh.mode == .bf16) { + const w = r.bfs(sh.n * sh.k, 0.05); + for (deq, w) |*o, v| o.* = rig.bf(v); + const tiled = try a.alloc(u16, cuda.nvfp4.padded(sh.n) * sh.k); + cuda.qmmf.packBf16(tiled, w, sh.n, sh.k); + return .{ .w = cuda.modelopt.bf16Weight(try r.dev(std.mem.sliceAsBytes(tiled)), sh.n, sh.k), .deq = deq }; + } + const codes = try a.alloc(u8, sh.n * sh.k); + for (codes) |*c| c.* = rng.int(u8) & 0xF7; // no e4m3 NaN (0x7F, 0xFF) + const scales = try a.alloc(u8, sh.n * sh.k / 32); + for (scales) |*s| s.* = 120 + rng.uintLessThan(u8, 8); + for (deq, 0..) |*o, i| o.* = @import("core").quant.fp8.e4m3(codes[i]) * std.math.pow(f64, 2, @as(f64, @floatFromInt(@as(i32, scales[i / 32]) - 127))); + const pc = try a.alloc(u8, cuda.mxfp8.padded(sh.n) * sh.k); + cuda.mxfp8.packCodes(pc, codes, sh.n, sh.k); + const ps = try a.alloc(u8, cuda.mxfp8.padded(sh.n) * (sh.k / 32)); + cuda.mxfp8.packScales(ps, scales, sh.n, sh.k); + return .{ .w = cuda.mxfp8.weight(try r.dev(pc), try r.dev(ps), sh.n, sh.k), .deq = deq }; +} + +pub fn run(gpu: check.Gpu) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var prng = std.Random.DefaultPrng.init(0x4a1f); + var r = try rig.Rig.init(gpu, a, prng.random()); + defer r.deinit(); + var lane = try cuda.qmmf.Lane.load(gpu.d, @intCast(try gpu.ctx.capability() / 10)); + defer lane.unload(); + const keep = lane.half_below; + var halves: usize = 0; + for (shapes) |sh| { + const built = try build(&r, a, prng.random(), sh); + const weight = built.w; + const xs = r.bfs(sh.m * sh.k, 1.0); + const x = try r.dev(std.mem.sliceAsBytes(xs)); + const part = try r.out(@max(16, lane.partBytes(sh.m, weight))); + const p = cuda.qmmf.plan(sh.m, sh.n, sh.k, lane.clusters); + const takes = p.bm == 16 and !p.fused and ((sh.n + 63) / 64) * p.sk < keep; + if (takes) halves += 1; + var times: [2]f64 = .{ 0, 0 }; + var f32_rows: []f32 = &.{}; + inline for (.{ cuda.qmmf.Out.f32, cuda.qmmf.Out.bf16 }) |out| { + const bytes = sh.m * sh.n * (if (out == .f32) @as(usize, 4) else 2); + const ys = [2]u64{ try r.out(bytes), try r.out(bytes) }; + for (0..2) |on| { + lane.half_below = if (on == 1) keep else 0; + try lane.matmulOut(r.s, x, sh.k, sh.m, weight, ys[on], out, part); + try r.s.synchronize(); + } + const got = try r.back(u8, ys[1], bytes); + const want = try r.back(u8, ys[0], bytes); + try check.sameBytes("half tiles equal the 64-column tile", got, want); + if (out == .f32) f32_rows = std.mem.bytesAsSlice(f32, @as([]align(4) u8, @alignCast(got))); + } + // FP64: every fp32 output within the lane's sum bound + for (0..sh.m) |row| for (0..sh.n) |j| { + var want: f64 = 0; + var mag: f64 = 0; + for (0..sh.k) |i| { + const t = rig.bf(xs[row * sh.k + i]) * built.deq[j * sh.k + i]; + want += t; + mag += @abs(t); + } + try r.within("half-tile lane fp32 vs FP64", f32_rows[row * sh.n + j], want, rig.sumTol(mag, sh.k) + 1e-6); + }; + // time on cold weights: copies past the L2, cycled + const wbytes: usize = if (sh.mode == .bf16) sh.n * sh.k * 2 else sh.n * sh.k + sh.n * sh.k / 32; + const copies = @min(64, @max(2, (128 << 20) / wbytes)); + var cold: [64]cuda.qmmf.Weight = undefined; + for (0..copies) |c| cold[c] = (try build(&r, a, prng.random(), sh)).w; + const y = try r.out(sh.m * sh.n * 4); + for (0..2) |on| { + lane.half_below = if (on == 1) keep else 0; + const reps = 300; + const t0 = check.now(gpu.io); + for (0..reps) |i| try lane.matmulOut(r.s, x, sh.k, sh.m, cold[i % copies], y, .f32, part); + try r.s.synchronize(); + times[on] = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e3 / reps; + } + std.debug.print("lane-half {t} m={d} n={d} k={d} (sk {d}{s}): 64-col {d:.1} us ({d:.0} GB/s), half {d:.1} us ({d:.0} GB/s), bytes equal\n", .{ sh.mode, sh.m, sh.n, sh.k, p.sk, if (takes) ", takes half tiles" else ", stays 64-col", times[0], @as(f64, @floatFromInt(wbytes)) / 1e3 / times[0], times[1], @as(f64, @floatFromInt(wbytes)) / 1e3 / times[1] }); + } + lane.half_below = keep; + try check.expect(halves >= 5, "the decode shapes take half tiles ({d})", .{halves}); + r.report("lane half tiles: MXFP8 and BF16, f32 and bf16 out, bit-equal to the 64-column tile, FP64-bounded"); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 837f1bb070..3ee7413ea6 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -74,6 +74,7 @@ const flashnext_rank_tests = @import("flashnext_rank_tests.zig"); const flashnext_draft_affine_tests = @import("flashnext_draft_affine_tests.zig"); const flashnext_shared_tests = @import("flashnext_shared_tests.zig"); const flashnext_nvfp4_tp_tests = @import("flashnext_nvfp4_tp_tests.zig"); +const lane_half_tests = @import("lane_half_tests.zig"); const flashnext_dense_hc_tests = @import("flashnext_dense_hc_tests.zig"); const flashnext_hc_tests = @import("flashnext_hc_tests.zig"); const flashnext_attn_tests = @import("flashnext_attn_tests.zig"); @@ -103,6 +104,7 @@ const usage = \\ flashnext-draft-affine proposal-head affine projection against FP64 and row-count invariance \\ flashnext-shared separate MXFP8 shared expert and reserved-slot scatter against FP64 \\ flashnext-nvfp4-tp rank-local expert cuts and FP32 partials against FP64 (one GPU) + \\ lane-half lane half tiles (MXFP8, BF16): bytes equal to the 64-column tile, FP64, cold-weight time \\ flashnext-dense-hc dense connection weights and FP32 normalization against FP64 \\ flashnext-hc Flash Next's own hyper-connection and n-gram glue against float64 (synthetic rows) \\ flashnext-attn Flash Next's own attention glue, pooling and block scores against float64 (synthetic) @@ -249,6 +251,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "flashnext-draft-affine")) return flashnext_draft_affine_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-shared")) return flashnext_shared_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-nvfp4-tp")) return flashnext_nvfp4_tp_tests.run(gpu); + if (std.mem.eql(u8, cmd, "lane-half")) return lane_half_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-dense-hc")) return flashnext_dense_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-hc")) return flashnext_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-attn")) return flashnext_attn_tests.run(gpu); From af619a83e1f31fda85eb098ad82d9452e85eba65 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 21 Oct 2026 07:19:43 +0200 Subject: [PATCH 2/3] contract: the lane's half tiles replace the 64-column 16-row tile (MXFP8 and BF16), a 1.5x-scale control, gate 5 records --- zig/src/core/contract/lane_half_records.zig | 13 ++ zig/src/core/contract/registry.zig | 2 +- zig/tests/cuda/contract_lane_half.zig | 160 ++++++++++++++++++++ zig/tests/cuda/contract_tests.zig | 6 + 4 files changed, 180 insertions(+), 1 deletion(-) create mode 100644 zig/src/core/contract/lane_half_records.zig create mode 100644 zig/tests/cuda/contract_lane_half.zig diff --git a/zig/src/core/contract/lane_half_records.zig b/zig/src/core/contract/lane_half_records.zig new file mode 100644 index 0000000000..24a2bb7024 --- /dev/null +++ b/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 } } }, +}; diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 2f307c6753..ef8a5a9128 100644 --- a/zig/src/core/contract/registry.zig +++ b/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. diff --git a/zig/tests/cuda/contract_lane_half.zig b/zig/tests/cuda/contract_lane_half.zig new file mode 100644 index 0000000000..ff575821f6 --- /dev/null +++ b/zig/tests/cuda/contract_lane_half.zig @@ -0,0 +1,160 @@ +//! The lane's half tiles (16x32, MXFP8 and BF16) under the contract: the 64-column tile's bits; control: 1.5x scale. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const Dense = contract.dense.Dense; +const Rig = @import("contract_rig.zig").Rig; + +pub const Kind = enum { mxfp8, bf16, control }; + +fn depth(shape: c.Shape) usize { + return shape.k / 16 + 24; +} + +/// Flash Next's decode calls whose 64-column grids take half tiles, the column edges, and 16 rows. +const served = [_]c.Shape{ + .{ .rows = 1, .n = 640, .k = 2560 }, .{ .rows = 1, .n = 2560, .k = 640 }, .{ .rows = 16, .n = 640, .k = 2560 }, + .{ .rows = 1, .n = 324, .k = 10240 }, .{ .rows = 7, .n = 512, .k = 2560 }, .{ .rows = 3, .n = 96, .k = 512 }, +}; + +pub fn spec(kind: Kind) c.Spec { + return .{ + .name = switch (kind) { + .mxfp8 => "cuda MXFP8 lane half tiles (32 columns, 16 rows)", + .bf16 => "cuda BF16 lane half tiles (32 columns, 16 rows)", + .control => "control: MXFP8 half tiles at a 1.5x tensor scale posing as the 64-column tile", + }, + .accum = .f32, + .depth = depth, + .out = .bf16, + .replaces = "the 64-column 16-row tile", + .tiles = .{ .{ .size = 16 }, .{ .size = 32, .step = 32, .min = 32 }, .{ .size = 512, .step = 64, .min = 64 } }, + .served = &served, + }; +} + +pub const Half = struct { + rig: Rig, + lane: *cuda.qmmf.Lane, + kind: Kind, + d: ?Dense = null, + weight: cuda.qmmf.Weight = undefined, + x: u64 = 0, + out: u64 = 0, + base: u64 = 0, + part: u64 = 0, + + pub fn deinit(self: *Half) void { + if (self.d) |d| d.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Half) g.Kernel { + return .{ .spec = spec(self.kind), .ctx = self, .vt = &vt, .control = if (self.kind == .control) .shapes else null }; + } + + const vt = g.hooks(Half); + + fn mode(self: *const Half) cuda.qmmf.Mode { + return if (self.kind == .bf16) .bf16 else .mxfp8; + } + + pub fn prepare(self: *Half, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (self.d) |d| d.deinit(); + self.d = null; + self.rig.reset(); + self.part = 0; + if (shape.rows == 0 or shape.rows > 16 or shape.n == 0 or shape.k == 0 or shape.k % 64 != 0) return false; + const a = self.rig.gpu.gpa; + const d = try Dense.init(a, shape, 1); + self.d = d; + d.fillX(fill, seed, 16); + var prng = std.Random.DefaultPrng.init(seed ^ 0x4a1f); + const random = prng.random(); + const n = shape.n; + const k = shape.k; + if (self.mode() == .bf16) { + const w = try a.alloc(u16, n * k); + defer a.free(w); + for (w, 0..) |*v, i| { + const f: f32 = if (fill == .stress) @floatFromInt(@as(i32, 1) + @as(i32, @intCast(random.uintLessThan(u32, 8)))) else random.float(f32) * 0.1 - 0.05; + v.* = c.toBf16(f); + _ = i; + } + if (fill == .stress) for (0..n) |j| for (k / 2..k) |i| { + w[j * k + i] = w[j * k + i - k / 2] | 0x8000; + }; + for (d.w, w) |*o, v| o.* = c.fromBf16(v); + const tiled = try a.alloc(u16, cuda.nvfp4.padded(n) * k); + defer a.free(tiled); + cuda.qmmf.packBf16(tiled, w, n, k); + self.weight = cuda.modelopt.bf16Weight(try self.rig.input(std.mem.sliceAsBytes(tiled)), n, k); + } else { + const codes = try a.alloc(u8, n * k); + defer a.free(codes); + const scales = try a.alloc(u8, n * k / 32); + defer a.free(scales); + for (scales) |*s| s.* = if (fill == .stress) 127 else 123 + random.uintLessThan(u8, 8); + for (0..n) |j| for (0..k) |i| { + codes[j * k + i] = if (fill == .stress) 0x38 + random.uintLessThan(u8, 24) else random.uintLessThan(u8, 0x7f); + if (fill != .stress and random.boolean()) codes[j * k + i] |= 0x80; + }; + if (fill == .stress) for (0..n) |j| for (k / 2..k) |i| { + codes[j * k + i] = codes[j * k + i - k / 2] | 0x80; + }; + for (0..n) |j| for (0..k) |i| { + const scale = std.math.ldexp(@as(f64, 1), @as(i32, scales[j * (k / 32) + i / 32]) - 127); + d.w[j * k + i] = core.quant.fp8.e4m3(codes[j * k + i]) * scale; + }; + const npad = cuda.mxfp8.padded(n); + const tiled_codes = try a.alloc(u8, npad * k); + defer a.free(tiled_codes); + const packed_scales = try a.alloc(u8, npad * k / 32); + defer a.free(packed_scales); + cuda.mxfp8.packCodes(tiled_codes, codes, n, k); + cuda.mxfp8.packScales(packed_scales, scales, n, k); + self.weight = cuda.mxfp8.weight(try self.rig.input(tiled_codes), try self.rig.input(packed_scales), n, k); + } + self.x = try self.rig.input(std.mem.sliceAsBytes(d.x)); + const bytes = self.lane.partBytes(shape.rows, self.weight); + if (bytes > 0) self.part = try self.rig.scratch(bytes); + self.out = try self.rig.output(shape.rows * n * 2); + self.base = try self.rig.scratch(shape.rows * n * 2); + return true; + } + + fn launch(self: *Half, half: bool, out: u64) !void { + const d = self.d.?; + const keep = self.lane.half_below; + defer self.lane.half_below = keep; + self.lane.half_below = if (half) std.math.maxInt(usize) else 0; + var w = self.weight; + if (half and self.kind == .control) w.scale = 1.5; + try self.lane.matmul(self.rig.s, self.x, d.shape.k, d.shape.rows, w, out, self.part); + try self.rig.s.synchronize(); + } + + pub fn run(self: *Half) !void { + try self.launch(true, self.out); + } + + pub fn replaced(self: *Half) ![]const u8 { + try self.launch(false, self.base); + return self.rig.bytes(self.base, self.count() * 2); + } + + pub fn count(self: *Half) usize { + return self.d.?.outputs(); + } + + pub fn reference(self: *Half, want: []f64, magnitude: []f64) !void { + try self.d.?.reference(want, magnitude); + } + + pub fn simulate(self: *Half, accumulator: c.Accum, order: g.Order, out: []f64) !void { + try self.d.?.simulate(accumulator, order, out); + } +}; diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index 86aa6eb966..6f4e0e4e4e 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -9,6 +9,7 @@ const rows_k = @import("contract_rows.zig"); const router_k = @import("contract_router.zig"); const mm_k = @import("contract_mm.zig"); const mx_k = @import("contract_mxfp8.zig"); +const half_k = @import("contract_lane_half.zig"); const wide_k = @import("contract_projection_f32.zig"); const rank_k = @import("contract_rank.zig"); const writeback_k = @import("contract_flashnext_writeback.zig"); @@ -150,6 +151,11 @@ pub fn run(gpu: check.Gpu) !void { }; defer for (&mx) |*m| m.deinit(); for (&mx) |*m| try kernels.append(a, m.kernel()); + var half_lane = lane; + var halves: [3]half_k.Half = undefined; + for (std.enums.values(half_k.Kind), &halves) |kind, *h| h.* = .{ .rig = .{ .gpu = gpu, .s = s }, .lane = &half_lane, .kind = kind }; + defer for (&halves) |*h| h.deinit(); + for (&halves) |*h| try kernels.append(a, h.kernel()); var wide: [wide_k.modes.len * 2]wide_k.Product = undefined; for (wide_k.modes, 0..) |mode, i| { wide[i * 2] = .{ .rig = .{ .gpu = gpu, .s = s }, .lane = &lane, .mode = mode, .clusters = true }; From b6d7cd466468022f2008490a67235f50673dde19 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 21 Oct 2026 17:08:18 +0200 Subject: [PATCH 3/3] tests: the host lanes built by hand name the half tiles, at threshold 0 for the released 64-column launches projection_f32_host_tests and the nemotron sm_89 host tests build a qmmf.Lane field by field; they now set its half tiles and half_below = 0, so their launch parameters stay the 64-column tile's they check. --- zig/src/families/nemotron/cuda_sm89_test.zig | 2 ++ zig/tests/cuda/projection_f32_host_tests.zig | 2 ++ 2 files changed, 4 insertions(+) diff --git a/zig/src/families/nemotron/cuda_sm89_test.zig b/zig/src/families/nemotron/cuda_sm89_test.zig index e3e544ff9b..1a4e1160ef 100644 --- a/zig/src/families/nemotron/cuda_sm89_test.zig +++ b/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 { diff --git a/zig/tests/cuda/projection_f32_host_tests.zig b/zig/tests/cuda/projection_f32_host_tests.zig index 5d38e2e29c..41bcb98b9d 100644 --- a/zig/tests/cuda/projection_f32_host_tests.zig +++ b/zig/tests/cuda/projection_f32_host_tests.zig @@ -60,6 +60,8 @@ test "every F32 and BF16 fixture dispatch rejects a no-op after a valid stale pr .fns = @splat(@splat(f)), .f32_fns = @splat(@splat(f)), .reduce = @splat(f), + .half = @splat(@splat(@splat(f))), + .half_below = 0, .clusters = false, }; const weight: cuda.qmmf.Weight = .{ .mode = .fp4, .codes = 0, .scales = 0, .n = 70, .k = 1024, .npad = 128 };