diff --git a/CHANGELOG.md b/CHANGELOG.md index acf9673a39..655a664261 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/zig/kernels/hip/tiles/matrix_gemm.hpp b/zig/kernels/hip/tiles/matrix_gemm.hpp index f06ae9f42a..e896edd02c 100644 --- a/zig/kernels/hip/tiles/matrix_gemm.hpp +++ b/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(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 = [&](std::integral_constant, int st, int buf) { u32x4* dst = reinterpret_cast(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(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,7 +149,7 @@ __device__ __forceinline__ void wmma_gemm_tile(typename Dec::Args& a) { __syncthreads(); auto step = [&](std::integral_constant, 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 @@ -153,19 +157,23 @@ __device__ __forceinline__ void wmma_gemm_tile(typename Dec::Args& a) { 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{}, s); if (s + 1 < stages) step(std::integral_constant{}, 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