From e614eec6941b2ab6a5e92a49d8b86f7435fffd03 Mon Sep 17 00:00:00 2001 From: Abraham <119352736+AbrahamCete@users.noreply.github.com> Date: Tue, 27 Oct 2026 22:21:43 +0100 Subject: [PATCH 1/6] perf(hip): interleave a stream tile's row blocks so wide rounds reuse cached weights A lane round of more than four rows ran the stream tile as ceil(rows / 4) row blocks on the grid's y, so each later row block read every weight from DRAM again. The row blocks of one column block now run next to each other on x and the later ones find its weights in cache. Each row's arithmetic is unchanged, so a row still writes the same bytes alone and in a round. On gfx1151 (Radeon 8060S), Qwen3.8-27B MLX 4-bit, alternated A/B against main: verify of 8 rows 154 -> 124 ms, 16 rows 283 -> 208 ms; 1 to 4 rows and plain decode unchanged. Qwen3.6-35B-A3B MLX 4-bit MTP community-v1 unchanged. --- zig/kernels/hip/tiles/stream.hpp | 19 ++++++++++++------- zig/src/hip/launches/affine.zig | 4 +++- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/zig/kernels/hip/tiles/stream.hpp b/zig/kernels/hip/tiles/stream.hpp index e17d21ed26..1b92f6bdc4 100644 --- a/zig/kernels/hip/tiles/stream.hpp +++ b/zig/kernels/hip/tiles/stream.hpp @@ -53,13 +53,14 @@ __device__ inline void stream_dots(const typename T::pair (&xp)[RG][16], const u // Epi::kPair: a lane's CB columns are CB / 2 gate columns and the same up columns; out is silu(gate) * up. template -__device__ inline void stream_body(const typename Dec::Args& a, int lpc_log2, int gshift, int bx, int pair_cols, float limit) { +__device__ inline void stream_body(const typename Dec::Args& a, int lpc_log2, int gshift, int bx, int by, int pair_cols, + float limit) { using pair = typename T::pair; constexpr bool PAIR = Epi::kPair; constexpr int RG = R < 4 ? R : 4; // rows a pass keeps in registers constexpr int NG = R / RG; constexpr uint32_t kOne = T::ones; - const int row0 = blockIdx.y * R; // a block's rows: more than 8 rows are 8-row blocks on the grid's y + const int row0 = by * R; // a block's rows: row block `by` of the grid's y, or of x when interleaved const int lpc = 1 << lpc_log2; const int lane = threadIdx.x & 31; const int j = lane & (lpc - 1); @@ -213,7 +214,11 @@ __device__ inline void stream_body(const typename Dec::Args& a, int lpc_log2, in // Item z of the plan (or plain product, or a group's side) at CB columns a lane, R rows; chunk group is q >> gshift. template __device__ inline void stream_tile(typename Dec::Args& a, const Sides& sides, int lpc_log2, int gshift) { - int bx = blockIdx.x; + // a one-row grid over more than R rows interleaves: a column block's row blocks are neighbours on x + const int nby = (a.m + R - 1) / R; + const bool interleaved = gridDim.y == 1 && nby > 1; + int bx = interleaved ? static_cast(blockIdx.x) / nby : static_cast(blockIdx.x); + const int by = interleaved ? static_cast(blockIdx.x) % nby : static_cast(blockIdx.y); if (sides.count > 0) { int s = 0; while (s + 1 < sides.count && bx >= sides.first[s + 1]) ++s; @@ -226,15 +231,15 @@ __device__ inline void stream_tile(typename Dec::Args& a, const Sides& sides, in a.out = static_cast(sides.out[s]); } } - if (!Dec::take_item(a, blockIdx.z) || static_cast(blockIdx.y) * R >= a.m) return; + if (!Dec::take_item(a, blockIdx.z) || by * R >= a.m) return; if constexpr (Dec::kWideBytes == 4) { - stream_body(a, lpc_log2, gshift, bx, sides.pair_cols, sides.limit); + stream_body(a, lpc_log2, gshift, bx, by, sides.pair_cols, sides.limit); } else { // every chunk is wide-aligned when the words and the row stride are if (Dec::wide_aligned(a)) { - stream_body(a, lpc_log2, gshift, bx, sides.pair_cols, sides.limit); + stream_body(a, lpc_log2, gshift, bx, by, sides.pair_cols, sides.limit); } else { - stream_body(a, lpc_log2, gshift, bx, sides.pair_cols, sides.limit); + stream_body(a, lpc_log2, gshift, bx, by, sides.pair_cols, sides.limit); } } } diff --git a/zig/src/hip/launches/affine.zig b/zig/src/hip/launches/affine.zig index 4378596dc7..49ae0b32d4 100644 --- a/zig/src/hip/launches/affine.zig +++ b/zig/src/hip/launches/affine.zig @@ -377,7 +377,9 @@ pub const Kernels = struct { })); const table = if (sides.pair_cols > 0) &k.pairs else if (relu2) &k.relu2 else &k.stream; const f = table[b][ri][pick].?; - try go(r, f, .{ .x = blocks, .y = cdiv(a.m, stream_rows[ri]), .z = @intCast(items) }, .{ .x = 32 * stream_waves }, s, &args); + // a column block's row blocks run next to each other on x, so the later ones read its weights from cache + const ny = cdiv(a.m, stream_rows[ri]); + try go(r, f, .{ .x = blocks * ny, .y = 1, .z = @intCast(items) }, .{ .x = 32 * stream_waves }, s, &args); } /// One row on the row tile: a lane owns a group, a wave reads a column's row in one run. From de8773740954b7db9c72adf46eafc666a42521e7 Mon Sep 17 00:00:00 2001 From: Abraham <119352736+AbrahamCete@users.noreply.github.com> Date: Wed, 28 Oct 2026 00:42:00 +0100 Subject: [PATCH 2/6] test(hip): pin the stream tile's interleaved row blocks to the reference layout The interleave was an implicit contract between the stream launch and stream_tile. stream_layout.zig is now the one place that lays a round's row blocks on the grid (StreamSides moves there with it), Kernels.stream_interleave switches back to the reference layout on y, and a GPU test runs plain, relu2, swiglu pair, routed and grouped products at rounds of 5 to 32 rows in both layouts and requires equal bytes. A planted swap of the column and row block indices fails it. Naming the layout in a StreamSides field read by stream_tile was tried and dropped: on gfx1151 it doubled the stream kernels' VGPRs (4x4: 140 -> 256) and cost about 10% on wide verify rounds. --- zig/kernels/hip/tiles/stream.hpp | 2 +- zig/src/hip/kernel_tests.zig | 1 + zig/src/hip/kernel_tests/affine/layout.zig | 138 +++++++++++++++++++++ zig/src/hip/launches/affine.zig | 23 ++-- zig/src/hip/launches/stream_layout.zig | 32 +++++ 5 files changed, 180 insertions(+), 16 deletions(-) create mode 100644 zig/src/hip/kernel_tests/affine/layout.zig create mode 100644 zig/src/hip/launches/stream_layout.zig diff --git a/zig/kernels/hip/tiles/stream.hpp b/zig/kernels/hip/tiles/stream.hpp index 1b92f6bdc4..750c2efa95 100644 --- a/zig/kernels/hip/tiles/stream.hpp +++ b/zig/kernels/hip/tiles/stream.hpp @@ -214,7 +214,7 @@ __device__ inline void stream_body(const typename Dec::Args& a, int lpc_log2, in // Item z of the plan (or plain product, or a group's side) at CB columns a lane, R rows; chunk group is q >> gshift. template __device__ inline void stream_tile(typename Dec::Args& a, const Sides& sides, int lpc_log2, int gshift) { - // a one-row grid over more than R rows interleaves: a column block's row blocks are neighbours on x + // a one-row grid over more than R rows interleaves (launches/stream_layout.zig grid): row blocks neighbour on x const int nby = (a.m + R - 1) / R; const bool interleaved = gridDim.y == 1 && nby > 1; int bx = interleaved ? static_cast(blockIdx.x) / nby : static_cast(blockIdx.x); diff --git a/zig/src/hip/kernel_tests.zig b/zig/src/hip/kernel_tests.zig index 78dc1da06f..523d0100da 100644 --- a/zig/src/hip/kernel_tests.zig +++ b/zig/src/hip/kernel_tests.zig @@ -28,6 +28,7 @@ test { _ = @import("kernel_tests/tp.zig"); _ = @import("kernel_tests/affine/affine.zig"); _ = @import("kernel_tests/affine/rows.zig"); + _ = @import("kernel_tests/affine/layout.zig"); _ = @import("kernel_tests/graph_split.zig"); _ = @import("kernel_tests/flashnext.zig"); _ = @import("kernel_tests/flashnext_attn.zig"); diff --git a/zig/src/hip/kernel_tests/affine/layout.zig b/zig/src/hip/kernel_tests/affine/layout.zig new file mode 100644 index 0000000000..ac78ac5f73 --- /dev/null +++ b/zig/src/hip/kernel_tests/affine/layout.zig @@ -0,0 +1,138 @@ +//! A stream product writes the same bytes with a round's row blocks interleaved on x as with them laid on y. + +const std = @import("std"); +const base = @import("../rig.zig"); +const affine = @import("../../launches/affine.zig"); +const Buffer = base.Buffer; +const gpa = base.gpa; +const data = @import("data.zig"); +const Rig = @import("product.zig").Rig; +const Kernels = affine.Kernels; + +/// Round widths past one row block: a ragged last block, whole blocks and the widest rounds. +const rounds = [_]usize{ 5, 8, 9, 16, 31, 32 }; +const most_rows = 32; +const n = 1024; +const k_dim = 2048; +const group = 64; + +const Kind = enum { plain, relu2, pair, routed, group }; + +/// The tables of `experts` stacked products of `rows` columns each, on the device. +const Tables = struct { + words: Buffer, + scale: Buffer, + bias: Buffer, + + fn init(t: *Rig, rng: *data.Rng, rows: usize, bits: c_int) !Tables { + const words_row = k_dim * @as(usize, @intCast(bits)) / 32; + const groups = k_dim / group; + const hw = try data.fill(gpa, u32, rows * words_row, rng, data.makeWord); + defer gpa.free(hw); + const hs = try data.fill(gpa, u16, rows * groups, rng, data.makeScale); + defer gpa.free(hs); + const hb = try data.fill(gpa, u16, rows * groups, rng, data.makeBias); + defer gpa.free(hb); + var words = try Buffer.fromHost(t.base, hw); + errdefer words.free(); + var scale = try Buffer.fromHost(t.base, hs); + errdefer scale.free(); + return .{ .words = words, .scale = scale, .bias = try Buffer.fromHost(t.base, hb) }; + } + + fn deinit(self: *Tables) void { + self.words.free(); + self.scale.free(); + self.bias.free(); + } +}; + +/// One launch of `kind` through `k`: its output goes to `arg.out` (fp32) or `arg.out16` (activation type). +fn launch(t: *Rig, k: *const Kernels, kind: Kind, arg: affine.Arg, sides: []const affine.Side, items: c_int) !void { + const r = &t.base.r; + const s = t.base.stream.handle; + switch (kind) { + .plain => try k.streamGo(r, arg, .{}, arg.n, 1, s), + .relu2 => try k.relu2Go(r, arg, 1, s), + .pair => try k.streamGo(r, arg, .{ .pair_cols = @divExact(arg.n, 2) }, @divExact(arg.n, 2), 1, s), + .routed => try k.routed(r, arg, items, s), + .group => try k.groupGo(r, arg, sides, false, s), + } + try t.base.stream.synchronize(); +} + +/// `kind` at every round width, interleaved against the reference layout: the 32-bit words that differ. +fn compare(t: *Rig, rng: *data.Rng, kind: Kind, bits: c_int) !usize { + var reference = t.kernels.*; + reference.stream_interleave = false; + const experts: usize = if (kind == .routed) 2 else 1; + var tables = try Tables.init(t, rng, experts * n, bits); + defer tables.deinit(); + const hx = try data.fill(gpa, u16, most_rows * k_dim, rng, if (t.fp16) data.makeX16 else data.makeXB); + defer gpa.free(hx); + var x = try Buffer.fromHost(t.base, hx); + defer x.free(); + const bytes = most_rows * n * 4; + var outs: [2]Buffer = .{ try Buffer.alloc(t.base, bytes), undefined }; + defer outs[0].free(); + outs[1] = try Buffer.alloc(t.base, bytes); + defer outs[1].free(); + var members: [most_rows]c_int = undefined; + for (&members, 0..) |*mb, i| mb.* = @intCast(i); + var mem = try Buffer.fromHost(t.base, @as([]const c_int, &members)); + defer mem.free(); + const words_bytes = n * k_dim * @as(usize, @intCast(bits)) / 32 * 4; + const table_bytes = n * (k_dim / group) * 2; + + var diff: usize = 0; + for (rounds) |m| { + // two experts split the round's rows, the first holding the most + const first: c_int = @intCast((m + 1) / 2); + const plan = [_]c_int{ 0, 0, first, 1, first, @as(c_int, @intCast(m)) - first }; + var items = try Buffer.fromHost(t.base, @as([]const c_int, &plan)); + defer items.free(); + for (outs, [_]*const Kernels{ t.kernels, &reference }) |out, kern| { + try out.fill8(0xA5, null); + var arg: affine.Arg = .{ + .x = x.ptr, + .words = tables.words.ptr, + .scale = .{ .p = tables.scale.ptr, .kind = 1 }, + .bias = .{ .p = tables.bias.ptr, .kind = 1 }, + .out = if (kind == .relu2 or kind == .pair) 0 else out.ptr, + .out16 = if (kind == .relu2 or kind == .pair) out.ptr else 0, + .m = @intCast(m), + .n = n, + .k = k_dim, + .bits = bits, + .group = group, + .fp16 = @intFromBool(t.fp16), + }; + if (kind == .routed) { + arg.m = first; + arg.route = .{ .items = items.ptr, .members = mem.ptr, .x_div = 1 }; + } + // three sides of unequal widths over one x, each writing its own stretch of the output + const sides = [_]affine.Side{ + .{ .words = tables.words.ptr, .scale = tables.scale.ptr, .bias = tables.bias.ptr, .n = 512, .out = out.ptr }, + .{ .words = tables.words.ptr + words_bytes / 2, .scale = tables.scale.ptr + table_bytes / 2, .bias = tables.bias.ptr + table_bytes / 2, .n = 256, .out = out.ptr + most_rows * 512 * 4 }, + .{ .words = tables.words.ptr + words_bytes * 3 / 4, .scale = tables.scale.ptr + table_bytes * 3 / 4, .bias = tables.bias.ptr + table_bytes * 3 / 4, .n = 192, .out = out.ptr + most_rows * 768 * 4 }, + }; + try launch(t, kern, kind, arg, &sides, 2); + } + const d = try data.sameOnDevice(outs[0], outs[1], bytes); + if (d != 0) std.debug.print(" stream layout {t} {d}-bit, {d} rows: {d} words differ\n", .{ kind, bits, m, d }); + diff += d; + } + return diff; +} + +test "a stream product writes the bytes of the reference layout with its row blocks interleaved on x" { + const b = try base.Rig.open(1 << 20); + defer b.close(); + var rig = Rig.open(b); + var rng: data.Rng = .{ .state = 0x2545F4914F6CDD1D }; + var diff: usize = 0; + for ([_]Kind{ .plain, .relu2, .pair, .routed, .group }) |kind| diff += try compare(&rig, &rng, kind, 4); + diff += try compare(&rig, &rng, .plain, 3); + try base.expect(diff == 0, "stream layout: {d} words differ from the reference layout", .{diff}); +} diff --git a/zig/src/hip/launches/affine.zig b/zig/src/hip/launches/affine.zig index 49ae0b32d4..28fdcc1d18 100644 --- a/zig/src/hip/launches/affine.zig +++ b/zig/src/hip/launches/affine.zig @@ -14,6 +14,7 @@ const registry = @import("core").registry; const quant = @import("core").quant; const tuning = @import("../tuning/tuning.zig"); const mlx_entries = @import("mlx_entries.zig"); +const stream_layout = @import("stream_layout.zig"); const Error = runtime.Error; @@ -42,19 +43,7 @@ pub const Arg = extern struct { out16: u64 = 0, }; -/// quant/mlx_tiles.hpp's StreamSides: up to four products that share x, K, width and group in one launch. -pub const StreamSides = extern struct { - words: [4]u64 = @splat(0), - scale: [4]u64 = @splat(0), - bias: [4]u64 = @splat(0), - out: [4]u64 = @splat(0), - n: [4]c_int = @splat(0), - first: [5]c_int = @splat(0), - count: c_int = 0, - out_half: c_int = 0, - pair_cols: c_int = 0, - limit: f32 = 0, -}; +pub const StreamSides = stream_layout.StreamSides; /// One side of a group: its matrix and where its (rows, n) output goes. pub const Side = struct { words: u64, scale: u64, bias: u64, n: c_int, out: u64 }; @@ -134,6 +123,8 @@ pub const Kernels = struct { relu2_pair: [bit_widths.len][stream_rows.len][stream_cbs.len]?Function, /// Off keeps the decode tiles before the stream tile (the policy's reference switch). stream_on: bool, + /// Off lays a round's stream row blocks on the grid's y: the reference layout the interleaved one must equal. + stream_interleave: bool, block: [bit_widths.len]Function, gemm: [bit_widths.len]Function, wmma_gemm: [bit_widths.len]Function, @@ -156,6 +147,7 @@ pub const Kernels = struct { k.mode = choices.matrix; k.tile = if (choices.gemm) .gemm else .block; k.stream_on = choices.stream; + k.stream_interleave = true; k.reg = Registry.init(tableFor(caps), &mlx_entries.all); k.reg.verify(k.env(), .mlx) catch return error.Invalid; // the GEMM tile's rows of x a lane keeps (tiles/gemm.hpp GemmShape): 8 for BF16, 16 for FP16 @@ -378,8 +370,8 @@ pub const Kernels = struct { const table = if (sides.pair_cols > 0) &k.pairs else if (relu2) &k.relu2 else &k.stream; const f = table[b][ri][pick].?; // a column block's row blocks run next to each other on x, so the later ones read its weights from cache - const ny = cdiv(a.m, stream_rows[ri]); - try go(r, f, .{ .x = blocks * ny, .y = 1, .z = @intCast(items) }, .{ .x = 32 * stream_waves }, s, &args); + const grid = stream_layout.grid(k.stream_interleave, blocks, cdiv(a.m, stream_rows[ri]), items); + try go(r, f, grid, .{ .x = 32 * stream_waves }, s, &args); } /// One row on the row tile: a lane owns a group, a wave reads a column's row in one run. @@ -598,4 +590,5 @@ test "the lane groups of a row tile" { test { _ = mlx_entries; + _ = stream_layout; } diff --git a/zig/src/hip/launches/stream_layout.zig b/zig/src/hip/launches/stream_layout.zig new file mode 100644 index 0000000000..ad5876b7c7 --- /dev/null +++ b/zig/src/hip/launches/stream_layout.zig @@ -0,0 +1,32 @@ +//! The stream tile's launch: quant/mlx_tiles.hpp's StreamSides and where a round's row blocks lie on the grid. +const std = @import("std"); +const hl = @import("../launch.zig"); + +/// quant/mlx_tiles.hpp's StreamSides: up to four products that share x, K, width and group in one launch. +pub const StreamSides = extern struct { + words: [4]u64 = @splat(0), + scale: [4]u64 = @splat(0), + bias: [4]u64 = @splat(0), + out: [4]u64 = @splat(0), + n: [4]c_int = @splat(0), + first: [5]c_int = @splat(0), + count: c_int = 0, + out_half: c_int = 0, + pair_cols: c_int = 0, + limit: f32 = 0, +}; + +/// The stream tile's grid; stream_tile reads a one-row grid over more than one row block as interleaved on x. +pub fn grid(interleave: bool, blocks: c_uint, ny: c_uint, items: c_int) hl.Dim3 { + if (interleave and ny > 1) return .{ .x = blocks * ny, .y = 1, .z = @intCast(items) }; + return .{ .x = blocks, .y = ny, .z = @intCast(items) }; +} + +test "a round's row blocks interleave on x only when there are several and the launch asks" { + const on = grid(true, 30, 4, 2); + try std.testing.expect(on.x == 120 and on.y == 1 and on.z == 2); + const off = grid(false, 30, 4, 2); + try std.testing.expect(off.x == 30 and off.y == 4 and off.z == 2); + const one = grid(true, 30, 1, 1); + try std.testing.expect(one.x == 30 and one.y == 1 and one.z == 1); +} From ecff498ac0a4ab5dfdeb60672ceeef81aa5b0e5d Mon Sep 17 00:00:00 2001 From: Abraham <119352736+AbrahamCete@users.noreply.github.com> Date: Wed, 28 Oct 2026 06:41:28 +0100 Subject: [PATCH 3/6] test(hip): route the layout test's rounds top-2 so 5-8 row rounds interleave Each row goes to both experts (x_div 2), so every expert holds the whole round and the routed case reaches the interleaved grid from 5 rows; each round now asserts it spans more than one row block. --- zig/src/hip/kernel_tests/affine/layout.zig | 28 ++++++++++++---------- 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/zig/src/hip/kernel_tests/affine/layout.zig b/zig/src/hip/kernel_tests/affine/layout.zig index ac78ac5f73..9a6c71bd29 100644 --- a/zig/src/hip/kernel_tests/affine/layout.zig +++ b/zig/src/hip/kernel_tests/affine/layout.zig @@ -72,25 +72,32 @@ fn compare(t: *Rig, rng: *data.Rng, kind: Kind, bits: c_int) !usize { defer gpa.free(hx); var x = try Buffer.fromHost(t.base, hx); defer x.free(); - const bytes = most_rows * n * 4; + // a routed round sends each row to both experts, so its output holds two pairs a row + const bytes = @as(usize, if (kind == .routed) 2 else 1) * most_rows * n * 4; var outs: [2]Buffer = .{ try Buffer.alloc(t.base, bytes), undefined }; defer outs[0].free(); outs[1] = try Buffer.alloc(t.base, bytes); defer outs[1].free(); - var members: [most_rows]c_int = undefined; - for (&members, 0..) |*mb, i| mb.* = @intCast(i); - var mem = try Buffer.fromHost(t.base, @as([]const c_int, &members)); - defer mem.free(); const words_bytes = n * k_dim * @as(usize, @intCast(bits)) / 32 * 4; const table_bytes = n * (k_dim / group) * 2; var diff: usize = 0; for (rounds) |m| { - // two experts split the round's rows, the first holding the most - const first: c_int = @intCast((m + 1) / 2); - const plan = [_]c_int{ 0, 0, first, 1, first, @as(c_int, @intCast(m)) - first }; + // the launch reaches the interleaved grid only past one row block + const block_rows = affine.stream_rows[affine.stream_rows.len - 1]; + try base.expect(m > block_rows, "stream layout: {d} rows fit one row block", .{m}); + // top-2 routing: every row goes to both experts, pair 2r to the first and 2r + 1 to the second + const rows: c_int = @intCast(m); + const plan = [_]c_int{ 0, 0, rows, 1, rows, rows }; + var members: [2 * most_rows]c_int = undefined; + for (0..m) |i| { + members[i] = @intCast(2 * i); + members[m + i] = @intCast(2 * i + 1); + } var items = try Buffer.fromHost(t.base, @as([]const c_int, &plan)); defer items.free(); + var mem = try Buffer.fromHost(t.base, @as([]const c_int, members[0 .. 2 * m])); + defer mem.free(); for (outs, [_]*const Kernels{ t.kernels, &reference }) |out, kern| { try out.fill8(0xA5, null); var arg: affine.Arg = .{ @@ -107,10 +114,7 @@ fn compare(t: *Rig, rng: *data.Rng, kind: Kind, bits: c_int) !usize { .group = group, .fp16 = @intFromBool(t.fp16), }; - if (kind == .routed) { - arg.m = first; - arg.route = .{ .items = items.ptr, .members = mem.ptr, .x_div = 1 }; - } + if (kind == .routed) arg.route = .{ .items = items.ptr, .members = mem.ptr, .x_div = 2 }; // three sides of unequal widths over one x, each writing its own stretch of the output const sides = [_]affine.Side{ .{ .words = tables.words.ptr, .scale = tables.scale.ptr, .bias = tables.bias.ptr, .n = 512, .out = out.ptr }, From a3b937f3d2107d3d62b13a8d8bf784dce140bde7 Mon Sep 17 00:00:00 2001 From: Abraham <119352736+AbrahamCete@users.noreply.github.com> Date: Wed, 28 Oct 2026 10:50:09 +0100 Subject: [PATCH 4/6] perf(hip): stream tile makes a chunk's code pairs once and takes 4-bit BF16 pairs from bytes Each column's 16 code pairs are made once a chunk, then every row of the block dots them in turn, in the order a row alone takes, so the bytes stay. On the BF16 build a 4-bit word becomes two byte-per-code words and v_cvt_f32_ubyteN makes each code's float in one op; v_perm_b32 packs the top halves. The 4x4 tile's inner loop drops from 1149 to 852 instructions on gfx1151. gfx1151, A/B alternated twice of the previous commit's build against the same stream.hpp change on a local experiment build (1, 2 and 4 rows a block): Qwen3.8-27B verify 9-16 rows 204 -> 179 ms, decode 13.2 -> 13.6 tok/s; Qwen3.6-35B-A3B decode 71.4 -> 76.1 tok/s, MTP 95.9 -> 101.9 tok/s. --- zig/kernels/hip/tiles/stream.hpp | 105 ++++++++++++++++++------------- 1 file changed, 62 insertions(+), 43 deletions(-) diff --git a/zig/kernels/hip/tiles/stream.hpp b/zig/kernels/hip/tiles/stream.hpp index 750c2efa95..5277e61227 100644 --- a/zig/kernels/hip/tiles/stream.hpp +++ b/zig/kernels/hip/tiles/stream.hpp @@ -39,16 +39,43 @@ __device__ inline void stream_x_pairs(const uint32_t (&xn)[16], typename T::pair ...); } -// One column's chunk against RG rows: each code pair is made once and dotted with every row's x. -template -__device__ inline void stream_dots(const typename T::pair (&xp)[RG][16], const uint32_t (&w)[Dec::kWords], float (&d)[RG], - std::integer_sequence) { - (([&] { - const typename T::pair q = Dec::template stream_pair(w); +// One column's 16 code pairs of a chunk, made once for every row of the block. +template +__device__ inline void stream_codes(const uint32_t (&w)[Dec::kWords], typename T::pair (&qp)[16], + std::integer_sequence) { + if constexpr (Dec::kWords == 4 && Dec::kGapPairs && !std::is_same_v) { + // the byte path below hard-codes the decoder's pairing: pair 4 wd + k is codes 8 wd + k and 8 wd + k + 4 + static_assert(Dec::gap() == 4); + static_assert(((Dec::template lo() == 8 * (I >> 2) + (I & 3)) && ...)); + // 4-bit BF16 pairs from bytes: a byte's float is exact and its top half is the BF16, as stream_pair makes it + // v_cvt_f32_ubyteN reads byte N itself; written out, LLVM would shift and extract first + auto f = [](uint32_t v, int b) { + uint32_t r; + switch (b) { + case 0: asm("v_cvt_f32_ubyte0 %0, %1" : "=v"(r) : "v"(v)); break; + case 1: asm("v_cvt_f32_ubyte1 %0, %1" : "=v"(r) : "v"(v)); break; + case 2: asm("v_cvt_f32_ubyte2 %0, %1" : "=v"(r) : "v"(v)); break; + default: asm("v_cvt_f32_ubyte3 %0, %1" : "=v"(r) : "v"(v)); break; + } + return r; + }; + // the top halves of two floats as one pair, in one v_perm_b32 (stream_pick's selector) + auto pack = [](uint32_t lo, uint32_t hi) { + return __builtin_bit_cast(typename T::pair, stream_pick(lo, hi, 1)); + }; #pragma unroll - for (int r = 0; r < RG; ++r) d[r] = T::dot(xp[r][I], q, d[r]); - }()), - ...); + for (int wd = 0; wd < 4; ++wd) { + const uint32_t e = w[wd] & 0x0F0F0F0Fu; // nibbles 0, 2, 4, 6 as bytes 0..3 + const uint32_t o = (w[wd] >> 4) & 0x0F0F0F0Fu; // nibbles 1, 3, 5, 7 + // pair k of a word is nibbles k and k + 4 + qp[4 * wd + 0] = pack(f(e, 0), f(e, 2)); + qp[4 * wd + 1] = pack(f(o, 0), f(o, 2)); + qp[4 * wd + 2] = pack(f(e, 1), f(e, 3)); + qp[4 * wd + 3] = pack(f(o, 1), f(o, 3)); + } + } else { + ((qp[I] = Dec::template stream_pair(w)), ...); + } } // Epi::kPair: a lane's CB columns are CB / 2 gate columns and the same up columns; out is silu(gate) * up. @@ -57,8 +84,6 @@ __device__ inline void stream_body(const typename Dec::Args& a, int lpc_log2, in float limit) { using pair = typename T::pair; constexpr bool PAIR = Epi::kPair; - constexpr int RG = R < 4 ? R : 4; // rows a pass keeps in registers - constexpr int NG = R / RG; constexpr uint32_t kOne = T::ones; const int row0 = by * R; // a block's rows: row block `by` of the grid's y, or of x when interleaved const int lpc = 1 << lpc_log2; @@ -116,43 +141,37 @@ __device__ inline void stream_body(const typename Dec::Args& a, int lpc_log2, in sc[c] = live ? terms.widen(sr[c]) : 0.f; bi[c] = live ? terms.widen(br[c]) : 0.f; } + // every column's code pairs are made once, then each row in turn dots them in the order a row alone takes + pair qp[CB][16]; +#pragma unroll + for (int c = 0; c < CB; ++c) stream_codes(w[c], qp[c], std::make_integer_sequence{}); +#pragma unroll + for (int r = 0; r < R; ++r) { + if (row0 + r >= a.m) break; + pair xp[16]; + uint32_t xn[16]; + const u32x4* src = reinterpret_cast(xr[r] + static_cast(q) * 32); #pragma unroll - for (int g = 0; g < NG; ++g) { - if (row0 + g * RG >= a.m) break; - pair xp[RG][16]; - float sx[RG]; -#pragma unroll - for (int r = 0; r < RG; ++r) { - uint32_t xn[16]; - const u32x4* src = - reinterpret_cast(xr[g * RG + r] + static_cast(q) * 32); -#pragma unroll - for (int v = 0; v < 4; ++v) { - const u32x4 t = src[v]; - xn[4 * v] = t.x; - xn[4 * v + 1] = t.y; - xn[4 * v + 2] = t.z; - xn[4 * v + 3] = t.w; - } - float s = 0.f; -#pragma unroll - for (int i = 0; i < 16; ++i) { - s = T::dot(__builtin_bit_cast(pair, xn[i]), __builtin_bit_cast(pair, kOne), s); - } - sx[r] = s; - stream_x_pairs(xn, xp[r], std::make_integer_sequence{}); + for (int v = 0; v < 4; ++v) { + const u32x4 t = src[v]; + xn[4 * v] = t.x; + xn[4 * v + 1] = t.y; + xn[4 * v + 2] = t.z; + xn[4 * v + 3] = t.w; } + float sx = 0.f; #pragma unroll - for (int c = 0; c < CB; ++c) { - float d[RG]; + for (int i = 0; i < 16; ++i) { + sx = T::dot(__builtin_bit_cast(pair, xn[i]), __builtin_bit_cast(pair, kOne), sx); + } + stream_x_pairs(xn, xp, std::make_integer_sequence{}); #pragma unroll - for (int r = 0; r < RG; ++r) d[r] = 0.f; - stream_dots(xp, w[c], d, std::make_integer_sequence{}); + for (int c = 0; c < CB; ++c) { + float d = 0.f; #pragma unroll - for (int r = 0; r < RG; ++r) { - acc[g * RG + r][c] = Dec::fold_scale(acc[g * RG + r][c], d[r], sc[c]); - acc[g * RG + r][c] = Dec::fold_bias(acc[g * RG + r][c], sx[r], bi[c]); - } + for (int i = 0; i < 16; ++i) d = T::dot(xp[i], qp[c][i], d); + acc[r][c] = Dec::fold_scale(acc[r][c], d, sc[c]); + acc[r][c] = Dec::fold_bias(acc[r][c], sx, bi[c]); } } }; From 984b5ed64420c7c48a11cdd7df8ebd1149097660 Mon Sep 17 00:00:00 2001 From: Abraham <119352736+AbrahamCete@users.noreply.github.com> Date: Wed, 28 Oct 2026 11:21:32 +0100 Subject: [PATCH 5/6] test(hip): pin the 4-bit BF16 stream tile's bytes to the code pairs stream_pair made Plain, gate|up pair and top-2 routed launches of 1 to 16 rows hash to the digest the stream tile wrote before the previous commit, when its code pairs came from stream_pair. --- zig/src/hip/kernel_tests.zig | 1 + zig/src/hip/kernel_tests/affine/golden.zig | 110 +++++++++++++++++++++ 2 files changed, 111 insertions(+) create mode 100644 zig/src/hip/kernel_tests/affine/golden.zig diff --git a/zig/src/hip/kernel_tests.zig b/zig/src/hip/kernel_tests.zig index 523d0100da..c521496736 100644 --- a/zig/src/hip/kernel_tests.zig +++ b/zig/src/hip/kernel_tests.zig @@ -29,6 +29,7 @@ test { _ = @import("kernel_tests/affine/affine.zig"); _ = @import("kernel_tests/affine/rows.zig"); _ = @import("kernel_tests/affine/layout.zig"); + _ = @import("kernel_tests/affine/golden.zig"); _ = @import("kernel_tests/graph_split.zig"); _ = @import("kernel_tests/flashnext.zig"); _ = @import("kernel_tests/flashnext_attn.zig"); diff --git a/zig/src/hip/kernel_tests/affine/golden.zig b/zig/src/hip/kernel_tests/affine/golden.zig new file mode 100644 index 0000000000..a62703e8e1 --- /dev/null +++ b/zig/src/hip/kernel_tests/affine/golden.zig @@ -0,0 +1,110 @@ +//! The 4-bit BF16 stream tile writes the bytes its code pairs made one at a time wrote: a pinned digest of them. + +const std = @import("std"); +const base = @import("../rig.zig"); +const affine = @import("../../launches/affine.zig"); +const Buffer = base.Buffer; +const gpa = base.gpa; +const data = @import("data.zig"); +const prod = @import("product.zig"); +const Rig = prod.Rig; + +/// One row (the 1-row tile), each row block's ragged and whole sizes, and interleaved rounds. +const rounds = [_]usize{ 1, 2, 3, 4, 5, 8, 16 }; +const most_rows = 16; +const n = 512; +const k_dim = 2048; +const group = 64; +const bits = 4; + +/// The Wyhash of the bytes the tile wrote when its code pairs came from stream_pair; the byte path must reproduce it. +const want: u64 = 0x24742f9ea12d63cd; + +const Kind = enum { plain, pair, routed }; + +fn digestOf(t: *Rig, rng: *data.Rng, kind: Kind) !u64 { + const experts: usize = if (kind == .routed) 2 else 1; + const words_row = k_dim * bits / 32; + const groups = k_dim / group; + const hw = try data.fill(gpa, u32, experts * n * words_row, rng, data.makeWord); + defer gpa.free(hw); + const hs = try data.fill(gpa, u16, experts * n * groups, rng, data.makeScale); + defer gpa.free(hs); + const hb = try data.fill(gpa, u16, experts * n * groups, rng, data.makeBias); + defer gpa.free(hb); + const hx = try data.fill(gpa, u16, most_rows * k_dim, rng, data.makeXB); + defer gpa.free(hx); + var words = try Buffer.fromHost(t.base, hw); + defer words.free(); + var scale = try Buffer.fromHost(t.base, hs); + defer scale.free(); + var bias = try Buffer.fromHost(t.base, hb); + defer bias.free(); + var x = try Buffer.fromHost(t.base, hx); + defer x.free(); + const bytes = 2 * most_rows * n * 4; + var out = try Buffer.alloc(t.base, bytes); + defer out.free(); + + const r = &t.base.r; + const s = t.base.stream.handle; + var h: u64 = 0; + for (rounds) |m| { + // top-2 routing: every row goes to both experts, pair 2r to the first and 2r + 1 to the second + const rows: c_int = @intCast(m); + const plan = [_]c_int{ 0, 0, rows, 1, rows, rows }; + var members: [2 * most_rows]c_int = undefined; + for (0..m) |i| { + members[i] = @intCast(2 * i); + members[m + i] = @intCast(2 * i + 1); + } + var items = try Buffer.fromHost(t.base, @as([]const c_int, &plan)); + defer items.free(); + var mem = try Buffer.fromHost(t.base, @as([]const c_int, members[0 .. 2 * m])); + defer mem.free(); + try out.fill8(0xA5, null); + var arg: affine.Arg = .{ + .x = x.ptr, + .words = words.ptr, + .scale = .{ .p = scale.ptr, .kind = 1 }, + .bias = .{ .p = bias.ptr, .kind = 1 }, + .out = if (kind == .pair) 0 else out.ptr, + .out16 = if (kind == .pair) out.ptr else 0, + .m = rows, + .n = n, + .k = k_dim, + .bits = bits, + .group = group, + .fp16 = 0, + }; + switch (kind) { + .plain => try t.kernels.streamGo(r, arg, .{}, n, 1, s), + .pair => try t.kernels.streamGo(r, arg, .{ .pair_cols = n / 2 }, n / 2, 1, s), + .routed => { + arg.route = .{ .items = items.ptr, .members = mem.ptr, .x_div = 2 }; + try t.kernels.routed(r, arg, 2, s); + }, + } + try t.base.stream.synchronize(); + const written: usize = switch (kind) { + .plain => m * n * 4, + .pair => m * (n / 2) * 2, + .routed => 2 * m * n * 4, + }; + h = std.hash.Wyhash.hash(h, std.mem.asBytes(&try prod.digestOf(out, written))); + } + return h; +} + +test "the 4-bit BF16 stream tile writes the pinned bytes of its plain, pair and routed launches" { + const b = try base.Rig.open(1 << 20); + defer b.close(); + var rig = Rig.open(b); + if (rig.fp16) return error.SkipZigTest; // the FP16 build keeps stream_pair's magic-number pairs + var rng: data.Rng = .{ .state = 0x9E3779B97F4A7C15 }; + var h: u64 = 0; + for ([_]Kind{ .plain, .pair, .routed }) |kind| { + h = std.hash.Wyhash.hash(h, std.mem.asBytes(&try digestOf(&rig, &rng, kind))); + } + try base.expect(h == want, "stream golden: digest 0x{x:0>16}, pinned 0x{x:0>16}", .{ h, want }); +} From 675c33b2056ca14c2cbd44003344ed91c8f0f811 Mon Sep 17 00:00:00 2001 From: Abraham <119352736+AbrahamCete@users.noreply.github.com> Date: Fri, 30 Oct 2026 20:10:56 +0100 Subject: [PATCH 6/6] test(hip): give the layout test a routed item shorter than the launch's rows The second expert takes only the first half of the rows, so its row blocks past them return early in both the interleaved and the reference layout. --- zig/src/hip/kernel_tests/affine/layout.zig | 26 +++++++++++----------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/zig/src/hip/kernel_tests/affine/layout.zig b/zig/src/hip/kernel_tests/affine/layout.zig index 9a6c71bd29..0e3d26534b 100644 --- a/zig/src/hip/kernel_tests/affine/layout.zig +++ b/zig/src/hip/kernel_tests/affine/layout.zig @@ -16,7 +16,7 @@ const n = 1024; const k_dim = 2048; const group = 64; -const Kind = enum { plain, relu2, pair, routed, group }; +const Kind = enum { plain, relu2, pair, routed, uneven, group }; /// The tables of `experts` stacked products of `rows` columns each, on the device. const Tables = struct { @@ -55,7 +55,7 @@ fn launch(t: *Rig, k: *const Kernels, kind: Kind, arg: affine.Arg, sides: []cons .plain => try k.streamGo(r, arg, .{}, arg.n, 1, s), .relu2 => try k.relu2Go(r, arg, 1, s), .pair => try k.streamGo(r, arg, .{ .pair_cols = @divExact(arg.n, 2) }, @divExact(arg.n, 2), 1, s), - .routed => try k.routed(r, arg, items, s), + .routed, .uneven => try k.routed(r, arg, items, s), .group => try k.groupGo(r, arg, sides, false, s), } try t.base.stream.synchronize(); @@ -65,7 +65,8 @@ fn launch(t: *Rig, k: *const Kernels, kind: Kind, arg: affine.Arg, sides: []cons fn compare(t: *Rig, rng: *data.Rng, kind: Kind, bits: c_int) !usize { var reference = t.kernels.*; reference.stream_interleave = false; - const experts: usize = if (kind == .routed) 2 else 1; + const two = kind == .routed or kind == .uneven; + const experts: usize = if (two) 2 else 1; var tables = try Tables.init(t, rng, experts * n, bits); defer tables.deinit(); const hx = try data.fill(gpa, u16, most_rows * k_dim, rng, if (t.fp16) data.makeX16 else data.makeXB); @@ -73,7 +74,7 @@ fn compare(t: *Rig, rng: *data.Rng, kind: Kind, bits: c_int) !usize { var x = try Buffer.fromHost(t.base, hx); defer x.free(); // a routed round sends each row to both experts, so its output holds two pairs a row - const bytes = @as(usize, if (kind == .routed) 2 else 1) * most_rows * n * 4; + const bytes = @as(usize, if (two) 2 else 1) * most_rows * n * 4; var outs: [2]Buffer = .{ try Buffer.alloc(t.base, bytes), undefined }; defer outs[0].free(); outs[1] = try Buffer.alloc(t.base, bytes); @@ -86,17 +87,16 @@ fn compare(t: *Rig, rng: *data.Rng, kind: Kind, bits: c_int) !usize { // the launch reaches the interleaved grid only past one row block const block_rows = affine.stream_rows[affine.stream_rows.len - 1]; try base.expect(m > block_rows, "stream layout: {d} rows fit one row block", .{m}); - // top-2 routing: every row goes to both experts, pair 2r to the first and 2r + 1 to the second + // pair 2r of every row goes to the first expert, 2r + 1 to the second (the first half of the rows when uneven) const rows: c_int = @intCast(m); - const plan = [_]c_int{ 0, 0, rows, 1, rows, rows }; + const second: usize = if (kind == .uneven) m / 2 else m; + const plan = [_]c_int{ 0, 0, rows, 1, rows, @intCast(second) }; var members: [2 * most_rows]c_int = undefined; - for (0..m) |i| { - members[i] = @intCast(2 * i); - members[m + i] = @intCast(2 * i + 1); - } + for (0..m) |i| members[i] = @intCast(2 * i); + for (0..second) |i| members[m + i] = @intCast(2 * i + 1); var items = try Buffer.fromHost(t.base, @as([]const c_int, &plan)); defer items.free(); - var mem = try Buffer.fromHost(t.base, @as([]const c_int, members[0 .. 2 * m])); + var mem = try Buffer.fromHost(t.base, @as([]const c_int, members[0 .. m + second])); defer mem.free(); for (outs, [_]*const Kernels{ t.kernels, &reference }) |out, kern| { try out.fill8(0xA5, null); @@ -114,7 +114,7 @@ fn compare(t: *Rig, rng: *data.Rng, kind: Kind, bits: c_int) !usize { .group = group, .fp16 = @intFromBool(t.fp16), }; - if (kind == .routed) arg.route = .{ .items = items.ptr, .members = mem.ptr, .x_div = 2 }; + if (two) arg.route = .{ .items = items.ptr, .members = mem.ptr, .x_div = 2 }; // three sides of unequal widths over one x, each writing its own stretch of the output const sides = [_]affine.Side{ .{ .words = tables.words.ptr, .scale = tables.scale.ptr, .bias = tables.bias.ptr, .n = 512, .out = out.ptr }, @@ -136,7 +136,7 @@ test "a stream product writes the bytes of the reference layout with its row blo var rig = Rig.open(b); var rng: data.Rng = .{ .state = 0x2545F4914F6CDD1D }; var diff: usize = 0; - for ([_]Kind{ .plain, .relu2, .pair, .routed, .group }) |kind| diff += try compare(&rig, &rng, kind, 4); + for ([_]Kind{ .plain, .relu2, .pair, .routed, .uneven, .group }) |kind| diff += try compare(&rig, &rng, kind, 4); diff += try compare(&rig, &rng, .plain, 3); try base.expect(diff == 0, "stream layout: {d} words differ from the reference layout", .{diff}); }