perf(hip): keep the matrix tile's inner loop free of selects and divisions

CHANGELOG.md

@@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel
 
 ## Unreleased
 
+- HIP: the matrix-core prompt tile of gfx11 GPUs no longer spills to scratch inside its K loop: a group's index and bounds are a shift and a mask instead of a run-time division, and a group's first product starts from zero on a uniform branch instead of a select per accumulator, so the tile fits in 248 registers. Every sum is the same in the same order, so replies are token-equal to before. On a Ryzen AI Max+ 395 (gfx1151) with Qwen3.8-27B MLX 4-bit, the seven prompt products at 2,048 rows take 63.7 ms instead of 105.8 ms, and cold prefill is 32% faster at 2k tokens (8.60 to 5.83 s) and 24.5% faster at 64k; decode, Metal and CUDA are unchanged (#1663). Thanks to @AbrahamCete.
+
 - The paged prompt store (`prompt_radix`) can now serve a family whose cached rows read tokens past their page, as GLM's MTP rows do: with a lookahead, a page is matched and shared only while a prompt repeats the tokens past it, and a planned family cuts and resumes only on its grid, a whole number of pages. A store given rules it cannot keep (a zero page, or a planned grid that is no multiple of the page) now fails with `error.BadPromptRadixRules` in every build instead of a debug-only assert. No family turns these rules on yet, so HIP's families and every other path keep, resume and evict exactly as before, which a new golden test pins (#1659, part of #1510). Thanks to @Bizuayeu.
 
 - MOSS-Image: `moss_image.noise` makes the pipeline's seeded starting latents (`torch.randn` on a seeded CPU generator) in the family's host code, byte for byte what x86 torch makes, on any CPU. ARM torch, on a Mac or a GB10, gives slightly different bits: for seed 42 at 1024x1024, 200,873 of 524,288 values differ, by at most 2.2e-6. `zig build test` checks 69 cases, 24 seeds and 9 sizes from 16x16 to 2048x2048, against torch 2.13.0's recorded latents, and `moss_image_case.py noise` records them. `tf-moss-image` still reads its latents from the case file, and nothing else changes (#1668, fixes #1530). Thanks to @BobClawblaw.

zig/kernels/hip/tiles/matrix_gemm.hpp

@@ -31,6 +31,10 @@ __device__ __forceinline__ void wmma_gemm_tile(typename Dec::Args& a) {
     const int groups = a.k / a.group;
     const int stages = a.k / kGemmK;
     const int per = a.group / kGemmK;
+    const int per_log2 = __builtin_ctz(per);  // groups are 32, 64 or 128 codes: one, two or four stages
+    auto group_of = [&](int st) { return st >> per_log2; };
+    auto group_start = [&](int st) { return (st & (per - 1)) == 0; };
+    auto group_end = [&](int st) { return (st & (per - 1)) == per - 1; };
     const int rows_left = a.m - m0;
     const int line = tid & (kGemmN - 1);
     const int side = __builtin_amdgcn_readfirstlane(tid >> 7);
@@ -62,8 +66,8 @@ __device__ __forceinline__ void wmma_gemm_tile(typename Dec::Args& a) {
             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);
+        term[0] = Dec::term(a, base + group_of(st));
+        term[1] = Dec::term(a, base + group_of(st1));
     };
     auto stash = [&]<int J>(std::integral_constant<int, J>, int st, int buf) {
         u32x4* dst = reinterpret_cast<u32x4*>(ws[buf] + line * kLdw + side * 8);
@@ -85,13 +89,13 @@ __device__ __forceinline__ void wmma_gemm_tile(typename Dec::Args& a) {
 #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]);
+        if (side == 0 && group_end(st)) {
+            sc[group_of(st) % 3][line] = Dec::scale_value(a, term[J]);
+            bi[group_of(st) % 3][line] = Dec::bias_value(a, term[J]);
         }
     };
     auto sumx = [&](int st, int buf) {
-        if (st % per == 0) run = 0.f;
+        if (group_start(st)) run = 0.f;
         const u32x4* src = reinterpret_cast<const u32x4*>(xs[buf] + line * kLdw);
 #pragma unroll
         for (int v = 0; v < 4; ++v) {
@@ -104,7 +108,7 @@ __device__ __forceinline__ void wmma_gemm_tile(typename Dec::Args& a) {
                 run += T::hi(p);
             }
         }
-        if (st % per == per - 1) sx[(st / per) & 1][line] = run;
+        if (group_end(st)) sx[group_of(st) & 1][line] = run;
     };
 
     // Wave (wm, wn) owns rows 32 wm.., columns 64 wn.. (2 x 4 tiles); a lane: column lane % 16, rows 2 i + lane / 16.
@@ -145,27 +149,31 @@ __device__ __forceinline__ void wmma_gemm_tile(typename Dec::Args& a) {
     __syncthreads();
     auto step = [&]<int P>(std::integral_constant<int, P>, int s) {
         constexpr int cur = P;
-        if (live0 && s > 0 && (s - 1) % per == per - 1) apply_bias(((s - 1) / per) % 3, ((s - 1) / per) & 1);
+        if (live0 && s > 0 && group_end(s - 1)) apply_bias(group_of(s - 1) % 3, group_of(s - 1) & 1);
         if constexpr (P == 1) fetch(s + 1);
         if (live0) {
 #pragma unroll
             for (int t = 0; t < 2; ++t) {
                 bf16x16 bf[4];
 #pragma unroll
                 for (int nt = 0; nt < 4; ++nt) bf[nt] = frag16(ws[cur] + (wn * 64 + nt * 16 + rl) * kLdw + t * 8);
-                const bool first = t == 0 && s % per == 0;
+                // A group's first product starts from zero on a uniform branch, not a select per accumulator.
+                const bool first = t == 0 && group_start(s);
 #pragma unroll
                 for (int mt = 0; mt < 2; ++mt) {
                     if (mt == 1 && !live1) continue;
                     const bf16x16 af = frag16(xs[cur] + (wm * 32 + mt * 16 + rl) * kLdw + t * 8);
+                    if (first) {
 #pragma unroll
-                    for (int nt = 0; nt < 4; ++nt) {
-                        d[mt][nt] = T::mma(af, bf[nt], first ? zero : d[mt][nt]);
+                        for (int nt = 0; nt < 4; ++nt) d[mt][nt] = T::mma(af, bf[nt], zero);
+                    } else {
+#pragma unroll
+                        for (int nt = 0; nt < 4; ++nt) d[mt][nt] = T::mma(af, bf[nt], d[mt][nt]);
                     }
                 }
             }
-            if (s % per == per - 1) {
-                const int g = (s / per) % 3;
+            if (group_end(s)) {
+                const int g = group_of(s) % 3;
 #pragma unroll
                 for (int mt = 0; mt < 2; ++mt) {
                     if (mt == 1 && !live1) continue;
@@ -186,7 +194,7 @@ __device__ __forceinline__ void wmma_gemm_tile(typename Dec::Args& a) {
         step(std::integral_constant<int, 0>{}, s);
         if (s + 1 < stages) step(std::integral_constant<int, 1>{}, s + 1);
     }
-    if (live0 && (stages - 1) % per == per - 1) apply_bias(((stages - 1) / per) % 3, ((stages - 1) / per) & 1);
+    if (live0 && group_end(stages - 1)) apply_bias(group_of(stages - 1) % 3, group_of(stages - 1) & 1);
 #pragma unroll
     for (int mt = 0; mt < 2; ++mt) {
 #pragma unroll