hip: a 64-row routed prompt tile for many-expert plans on gfx1151

CHANGELOG.md

@@ -5,6 +5,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel
 ## Unreleased
 
 - TensorFold Starlight moves to its own repository, [ashhart/tf-starlight](https://github.com/ashhart/tf-starlight), with its history: it pins the TensorFold release it is tested with and checks the engine contract against that release in CI (#1972, fixes #1970).
+- HIP on the gfx1151 iGPU routes a many-expert prompt through a 64-row tile, with up to 12 expert pairs an item and prompt chunks capped at 1,024 rows; other GPUs are unchanged (#1658). Thanks to @gilby.
 
 - DeepSeek-V4.1 on Metal, part one: its kernels, the decoder over every layer, Engram, prompt chunks and DSpark drafts, each checked against float64 references (#1710-#1719). Serving it on Metal follows with #1720-#1722. Thanks to @gilby.
 

zig/kernels/hip/quant/mlx_tiles.hpp

@@ -11,6 +11,7 @@
 #include "tiles/gemm.hpp"
 #include "tiles/gemm_kp.hpp"
 #include "tiles/matrix_gemm.hpp"
+#include "tiles/matrix_item.hpp"
 #include "tiles/stream.hpp"
 
 namespace tf {
@@ -37,6 +38,13 @@ __global__ void __launch_bounds__(256) affine_wmma_gemm(Affine a) {
     wmma_gemm_tile<MlxDecoder<BITS>, BF16Act, WmmaBF16, GemmOut<WmmaBF16>>(a);
 }
 
+// ---- the same on RM x 16 rows a block, for a routed item of few rows ----
+
+template <int BITS, int RM>
+__global__ void __launch_bounds__(256) affine_wmma_item(Affine a) {
+    wmma_item_tile<MlxDecoder<BITS>, BF16Act, WmmaBF16, GemmOut<WmmaBF16>, RM>(a);
+}
+
 // ---- the decode stream tile ----
 
 // Up to four products sharing x, K, width and group in one launch; side s owns blocks [first[s], first[s + 1]).

zig/kernels/hip/tiles/dot2_tiles.hip

@@ -366,7 +366,8 @@ using Dot = DotF16;
     TF_STREAM(BITS, 1, 2) TF_STREAM(BITS, 1, 4) TF_STREAM(BITS, 1, 8)                        \
     TF_STREAM(BITS, 2, 2) TF_STREAM(BITS, 2, 4) TF_STREAM(BITS, 2, 8)                        \
     TF_STREAM(BITS, 4, 2) TF_STREAM(BITS, 4, 4)                                               \
-    TF_WMMA(BITS)
+    TF_WMMA(BITS)                                                                             \
+    template __global__ void affine_wmma_item<BITS, 4>(Affine);
 
 TF_TILES(2)
 TF_TILES(3)

zig/kernels/hip/tiles/matrix_item.hpp

@@ -0,0 +1,184 @@
+#pragma once
+
+// gfx11 WMMA prefill GEMM for a routed item of few rows: RM x 16 rows a block, all 8 waves split the 128 columns.
+
+#include "common/wmma.hpp"
+#include "tiles/gemm.hpp"
+#include "tiles/plan.hpp"
+
+namespace tf {
+namespace rocm {
+
+// Same staging, stage order and per-element sums as wmma_gemm_tile, so each output equals that tile's byte for byte.
+template <class Dec, class Act, class T, class Epi, int RM>
+__device__ __forceinline__ void wmma_item_tile(typename Dec::Args& a) {
+#if TF_DEVICE_WMMA_GFX11
+    if (!Dec::take_item(a, blockIdx.z)) return;
+    constexpr int kRows = RM * 16;
+    const int bx = blockIdx.x;
+    const int m0 = blockIdx.y * kRows;
+    if (m0 >= a.m) return;
+    constexpr int kLdw = kGemmLd / 2;
+    __shared__ __attribute__((aligned(16))) uint32_t xs[2][kRows * kLdw];
+    __shared__ __attribute__((aligned(16))) uint32_t ws[2][kGemmN * kLdw];
+    __shared__ float sx[2][kRows];
+    __shared__ float sc[3][kGemmN];
+    __shared__ float bi[3][kGemmN];
+
+    const int tid = threadIdx.x;
+    const int n0 = bx * kGemmN;
+    const int words_row = Dec::template row_words<int>(a);
+    const int groups = a.k / a.group;
+    const int stages = a.k / kGemmK;
+    const int per = a.group / kGemmK;
+    const int rows_left = a.m - m0;
+    const int line = tid & (kGemmN - 1);
+    const int side = __builtin_amdgcn_readfirstlane(tid >> 7);
+    const int wave = __builtin_amdgcn_readfirstlane(tid >> 5);
+    const int lane = tid & 31;
+
+    // Weights as in wmma_gemm_tile; x is staged by the first 2 * kRows threads, two stages a fetch.
+    const int col_s = n0 + line < a.n ? n0 + line : a.n - 1;
+    const uint32_t* wsrc = Dec::words(a) + static_cast<long long>(col_s) * words_row;
+    const int xrow = tid >> 1;
+    const int xhalf = tid & 1;
+    const bool xlive = xrow < kRows && (xrow & ~15) < rows_left;
+    const bool sums = side == 1 && line < kRows;
+    const int row_s = xrow < kRows && m0 + xrow < a.m ? m0 + xrow : 0;
+    const typename Act::elem* xrow_ptr = Act::row(a, row_s);
+    const u32x4* xsrc = reinterpret_cast<const u32x4*>(xrow_ptr);
+    uint32_t wreg[2][Dec::kWords];
+    u32x4 xreg[4];
+    typename Dec::Term term[2];
+    float run = 0.f;
+    auto fetch = [&](int st0) {
+        const int st1 = st0 + 1 < stages ? st0 + 1 : stages - 1;
+        const int st = st0 < stages ? st0 : stages - 1;
+        Dec::template load<true>(wsrc + st * Dec::kWords, wreg[0]);
+        Dec::template load<true>(wsrc + st1 * Dec::kWords, wreg[1]);
+        const int sx_ = xhalf ? st1 : st;
+        if (xlive) {
+#pragma unroll
+            for (int v = 0; v < 4; ++v) xreg[v] = xsrc[sx_ * 4 + v];
+        }
+        const long long base = static_cast<long long>(col_s) * groups;
+        term[0] = Dec::term(a, base + st / per);
+        term[1] = Dec::term(a, base + st1 / per);
+    };
+    auto stash = [&]<int J>(std::integral_constant<int, J>, int st, int buf) {
+        u32x4* dst = reinterpret_cast<u32x4*>(ws[buf] + line * kLdw + side * 8);
+        auto decode = [&]<int S>(std::integral_constant<int, S>) {
+#pragma unroll
+            for (int v = 0; v < 2; ++v) {
+                uint32_t d[4];
+#pragma unroll
+                for (int h = 0; h < 4; ++h) {
+                    const int t = S * 16 + v * 8 + h * 2;
+                    d[h] = __builtin_bit_cast(uint32_t, Dec::template pair<T>(wreg[J], t));
+                }
+                dst[v] = u32x4{d[0], d[1], d[2], d[3]};
+            }
+        };
+        if (side == 0) decode(std::integral_constant<int, 0>{}); else decode(std::integral_constant<int, 1>{});
+        if (xhalf == J && xlive) {
+            u32x4* xdst = reinterpret_cast<u32x4*>(xs[buf] + xrow * kLdw);
+#pragma unroll
+            for (int v = 0; v < 4; ++v) xdst[v] = xreg[v];
+        }
+        if (side == 0 && st % per == per - 1) {
+            sc[(st / per) % 3][line] = Dec::scale_value(a, term[J]);
+            bi[(st / per) % 3][line] = Dec::bias_value(a, term[J]);
+        }
+    };
+    auto sumx = [&](int st, int buf) {
+        if (st % per == 0) run = 0.f;
+        const u32x4* src = reinterpret_cast<const u32x4*>(xs[buf] + line * kLdw);
+#pragma unroll
+        for (int v = 0; v < 4; ++v) {
+            const u32x4 u = src[v];
+            const uint32_t q[4] = {u.x, u.y, u.z, u.w};
+#pragma unroll
+            for (int h = 0; h < 4; ++h) {
+                const typename T::pair p = __builtin_bit_cast(typename T::pair, q[h]);
+                run += T::lo(p);
+                run += T::hi(p);
+            }
+        }
+        if (st % per == per - 1) sx[(st / per) & 1][line] = run;
+    };
+
+    // Wave w owns columns 16 w..; a lane: column lane % 16, rows 2 i + lane / 16 of each of the RM row tiles.
+    const int rl = lane & 15;
+    const int rh = lane >> 4;
+    const int cw = wave * 16;
+    f32x8 acc[RM];
+    f32x8 d[RM];
+#pragma unroll
+    for (int mt = 0; mt < RM; ++mt) {
+        acc[mt] = f32x8{0, 0, 0, 0, 0, 0, 0, 0};
+        d[mt] = acc[mt];
+    }
+    const f32x8 zero = {0, 0, 0, 0, 0, 0, 0, 0};
+    auto apply_bias = [&](int slot, int g) {
+        const float bias = bi[slot][cw + rl];
+#pragma unroll
+        for (int mt = 0; mt < RM; ++mt) {
+            if (mt * 16 >= rows_left) continue;
+#pragma unroll
+            for (int i = 0; i < 8; ++i) acc[mt][i] = Dec::fold_bias(acc[mt][i], sx[g][mt * 16 + 2 * i + rh], bias);
+        }
+    };
+    fetch(0);
+    stash(std::integral_constant<int, 0>{}, 0, 0);
+    __syncthreads();
+    auto step = [&]<int P>(std::integral_constant<int, P>, int s) {
+        constexpr int cur = P;
+        if (s > 0 && (s - 1) % per == per - 1) apply_bias(((s - 1) / per) % 3, ((s - 1) / per) & 1);
+        if constexpr (P == 1) fetch(s + 1);
+#pragma unroll
+        for (int t = 0; t < 2; ++t) {
+            const bf16x16 bf = frag16(ws[cur] + (cw + rl) * kLdw + t * 8);
+            const bool first = t == 0 && s % per == 0;
+#pragma unroll
+            for (int mt = 0; mt < RM; ++mt) {
+                if (mt * 16 >= rows_left) continue;
+                const bf16x16 af = frag16(xs[cur] + (mt * 16 + rl) * kLdw + t * 8);
+                d[mt] = T::mma(af, bf, first ? zero : d[mt]);
+            }
+        }
+        if (s % per == per - 1) {
+            const float scale = sc[(s / per) % 3][cw + rl];
+#pragma unroll
+            for (int mt = 0; mt < RM; ++mt) {
+                if (mt * 16 >= rows_left) continue;
+#pragma unroll
+                for (int i = 0; i < 8; ++i) acc[mt][i] = Dec::fold_scale(acc[mt][i], d[mt][i], scale);
+            }
+        }
+        if (sums) sumx(s, cur);
+        stash(std::integral_constant<int, P ^ 1>{}, s + 1, cur ^ 1);
+        __syncthreads();
+    };
+    for (int s = 0; s < stages; s += 2) {
+        step(std::integral_constant<int, 0>{}, s);
+        if (s + 1 < stages) step(std::integral_constant<int, 1>{}, s + 1);
+    }
+    if ((stages - 1) % per == per - 1) apply_bias(((stages - 1) / per) % 3, ((stages - 1) / per) & 1);
+    const int col = n0 + cw + rl;
+    if (col >= a.n) return;
+#pragma unroll
+    for (int mt = 0; mt < RM; ++mt) {
+#pragma unroll
+        for (int i = 0; i < 8; ++i) {
+            const int row = m0 + mt * 16 + 2 * i + rh;
+            if (row < a.m) Epi::store(a, row, col, acc[mt][i]);
+        }
+    }
+#else
+    // Never launched off gfx11; the lanes' own id keeps the stub a kernel the limits records can name.
+    if (threadIdx.x < 100u) __builtin_trap();
+#endif
+}
+
+}  // namespace rocm
+}  // namespace tf

zig/src/core/contract/hip_records.zig

@@ -498,7 +498,8 @@ const own_records = [_]Record{
 
 pub const records = own_records ++ @import("hip_records_stream.zig").records ++
     @import("hip_records_router.zig").records ++
-    @import("hip_records_glm_ocr.zig").records;
+    @import("hip_records_glm_ocr.zig").records ++
+    @import("hip_records_item.zig").records;
 
 /// Gate 5's planted HIP controls: a launch past its __launch_bounds__, and 72 KiB of LDS where a workgroup has 64.
 pub const controls = [_]Record{

zig/src/core/contract/hip_records_item.zig

@@ -0,0 +1,26 @@
+//! Gate 5's records of the routed prompt's few-row matrix-core tiles (affine_wmma_item), 4-bit.
+const std = @import("std");
+const limits = @import("limits.zig");
+const Record = @import("hip_records.zig").Record;
+
+/// The 4-bit tile of `rm` 16-row tiles a block; gfx11 launches it, the other chips only carry its symbol.
+fn itemTile(comptime rm: u32) limits.HipLaunch {
+    return .{
+        .name = std.fmt.comptimePrint("wmma item tile 4-bit {d} rows", .{16 * rm}),
+        .symbol = std.fmt.comptimePrint("_ZN2tf4rocm16affine_wmma_itemILi4ELi{d}EEEvNS0_6AffineE", .{rm}),
+        .image = "affine_tiles",
+        .threads = 256,
+    };
+}
+
+pub const records = [_]Record{
+    .{
+        .launch = itemTile(4),
+        .usage = .{
+            .{ .vgprs = 1, .sgprs = 1, .lds = 0, .private = 0, .bound = 256, .wave = 32 },
+            .{ .vgprs = 176, .sgprs = 55, .lds = 34304, .private = 0, .bound = 256, .wave = 32 },
+            .{ .vgprs = 176, .sgprs = 55, .lds = 34304, .private = 0, .bound = 256, .wave = 32 },
+            .{ .vgprs = 1, .sgprs = 1, .lds = 0, .private = 0, .bound = 256, .wave = 32 },
+        },
+    },
+};

zig/src/hip/kernel_tests/affine/rows.zig

@@ -4,6 +4,7 @@ const std = @import("std");
 const base = @import("../rig.zig");
 const affine = @import("../../launches/affine.zig");
 const registry = @import("core").registry;
+const mlx_entries = @import("../../launches/mlx_entries.zig");
 const runtime = @import("../../runtime.zig");
 const Buffer = base.Buffer;
 const gpa = base.gpa;
@@ -282,6 +283,8 @@ fn tierProduct(t: *Rig, m: usize, n: usize, k: usize, bits: c_int, group: c_int,
     // a routed plan's pairs are its members, one a row
     const plan_pairs: usize = if (routed) m else 0;
     try prod.pickable(k_ptr, arg, item_count, plan_pairs, .prefill, k_ptr.env(), &set, gpa);
+    // the few-row matrix tiles hold a plan's rows only under a pairs bound, so the sweep may not reach them
+    if (routed) for (&mlx_entries.all) |*e| if (std.mem.indexOf(u8, e.id, ".prefill.item") != null and std.mem.indexOfScalar(*const Entry, set.items, e) == null) try set.append(gpa, e);
     var ids: [registry.max_entries]*const Entry = undefined;
     var n_ids: usize = 0;
     for (set.items) |e| if (prod.takesArg(k_ptr, e, arg, item_count, plan_pairs, .prefill)) {

zig/src/hip/launches/affine.zig

@@ -6,7 +6,6 @@ const runtime = @import("../runtime.zig");
 const hl = @import("../launch.zig");
 const Module = @import("../module.zig").Module;
 const Function = @import("../module.zig").Function;
-const Stream = @import("../stream.zig").Stream;
 const util = @import("util.zig");
 
 const Caps = @import("../caps.zig").Caps;
@@ -138,6 +137,7 @@ pub const Kernels = struct {
     block: [bit_widths.len]Function,
     gemm: [bit_widths.len]Function,
     wmma_gemm: [bit_widths.len]Function,
+    item: @import("matrix_item.zig").Functions,
     kp: [bit_widths.len][kp_tiles.len]Function,
     span: [2]Function,
     fold: Function,
@@ -176,6 +176,7 @@ pub const Kernels = struct {
         inline for (bit_widths, 0..) |bits, b| {
             k.block[b] = try m.function(std.fmt.comptimePrint("_ZN2tf4rocm17affine_dot2_blockINS0_{s}ELi{d}EEEvNS0_6AffineE", .{ dot, bits }));
             if (k.matrix) k.wmma_gemm[b] = try m.function(std.fmt.comptimePrint("_ZN2tf4rocm16affine_wmma_gemmILi{d}EEEvNS0_6AffineE", .{bits}));
+            if (k.matrix) k.item[b] = try @import("matrix_item.zig").load(m, bits);
             k.gemm[b] = try m.function(std.fmt.comptimePrint("_ZN2tf4rocm17affine_gemm_blockINS0_{s}ELi{d}E{s}EEvNS0_6AffineE", .{ dot, bits, shape }));
             inline for (kp_tiles, 0..) |t, i| {
                 k.kp[b][i] = try m.function(std.fmt.comptimePrint("_ZN2tf4rocm14affine_gemm_kpINS0_{s}ELi{d}ELi{d}ELi{d}ELi{d}ELi{d}ELb{d}EEEvNS0_6AffineE", .{ dot, bits, t.cb, t.r, t.waves, t.rb, @intFromBool(t.loop) }));
@@ -290,7 +291,8 @@ pub const Kernels = struct {
     pub fn tableFor(caps: Caps) *const tuning.Table {
         return switch (caps.generation) {
             .rdna2 => &tuning.gfx1030,
-            .rdna3, .rdna3_5 => &tuning.gfx1100,
+            .rdna3 => &tuning.gfx1100,
+            .rdna3_5 => &tuning.gfx1151,
             .rdna4 => &tuning.gfx1201,
         };
     }

zig/src/hip/launches/matrix_item.zig

@@ -0,0 +1,36 @@
+//! The routed prompt's matrix-core tile of few rows an item (tiles/matrix_item.hpp): its functions and its launch.
+
+const std = @import("std");
+const hl = @import("../launch.zig");
+const Module = @import("../module.zig").Module;
+const Function = @import("../module.zig").Function;
+const affine = @import("affine.zig");
+const util = @import("util.zig");
+
+/// The 16-row tiles a block holds, by tier: 64 rows.
+pub const row_tiles = [_]usize{4};
+pub const Functions = [affine.bit_widths.len][row_tiles.len]Function;
+
+/// The tier's functions at a width, as the module names them.
+pub fn load(m: Module, comptime bits: c_int) util.Error![row_tiles.len]Function {
+    var out: [row_tiles.len]Function = undefined;
+    inline for (row_tiles, 0..) |rm, i| {
+        out[i] = try m.function(std.fmt.comptimePrint("_ZN2tf4rocm16affine_wmma_itemILi{d}ELi{d}EEEvNS0_6AffineE", .{ bits, rm }));
+    }
+    return out;
+}
+
+/// The rows of a tier's block.
+pub fn rows(tier: usize) usize {
+    return 16 * row_tiles[tier];
+}
+
+/// Tier `tier` over every item of a plan: 128 columns by the tier's rows a block, all waves across the columns.
+pub fn go(k: *const affine.Kernels, c: affine.Call, tier: usize) util.Error!void {
+    const b = std.mem.indexOfScalar(c_int, &affine.bit_widths, c.arg.bits) orelse return error.Invalid;
+    if (c.s == util.counting) return;
+    var args: hl.Args = .{};
+    try args.add(c.arg);
+    const grid = util.dim(util.cdiv(c.arg.n, 128), util.cdiv(c.arg.m, rows(tier)), c.items);
+    try hl.launch(k.item[b][tier], .{ .grid = grid, .block = .{ .x = 256 } }, .{ .r = c.r, .handle = c.s }, &args);
+}