perf(hip): prompt attention tile keeps its probabilities in registers and fits 6 waves, same bits

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 prompt attention tile keeps its probabilities in registers and runs 6 waves a SIMD on gfx1151, with the same output bits; cold first token is 1% to 5% faster from 2K to 64K tokens on Qwen3.8-27B, and the tile is registered with the kernel contract (#1788). Thanks to @AbrahamCete.
+
 - `--slide` on a checkpoint whose family doesn't learn on that backend now stops before loading with "`--slide`: <model type> checkpoints do not learn on <backend> yet", on CUDA, Metal, HIP and XPU, instead of serving and learning nothing. Metal learners register by family, as CUDA's do, and read their widths from the family's Living Weights description, so a new family's learner plugs in without editing the Metal host.
 
 - DeepSeek-V4.1-Flash-1017's release checkpoint is read as published (FP8 blocks, MXFP4 experts and the FP8 Engram table), with a manifest-checked Engram pack; the CUDA engine refuses it by name, as it runs EXL3 (#1705, part of #1676). Thanks to @gilby.

docs/kernel-contract.md

@@ -117,6 +117,10 @@ MLX's formula keeps the infinity, because its codes enter the MMAs as q + 128 an
 `tf_glm_scores` scores a NaN dot as 0 where float64 keeps the NaN: its relu is `fmaxf(x, 0)`, as `tl.maximum` is in the
 Python family's Triton glue it was ported from. Nemotron's Metal `tf_xup_w3` gives 0 for a NaN dot where float64 keeps
 the NaN: its relu2 is `metal::max(bfloat(acc), 0.0f)`; the other up pipelines give its bytes on the same inputs.
+HIP's prompt attention tile (`tf_fa_paged_*`, `tf_fa_wide_*`) writes 0 for a row whose every score is NaN where
+float64 gives NaN, as its running max is `fmaxf` and a row at max -inf takes p = 0; and an infinite value at a key
+whose fp32 weight is 0 (past the row's causal edge in its 64-key tile, or under fp32's least) gives NaN where float64
+gives inf or a finite value, as its P x V takes every key of the tile.
 
 HIP’s MLX stream tile keeps its finite dot2 sums unchanged. A non-finite result with an actual non-finite
 operand is recomputed by full quantization groups using the factored float64 formula. The plain tile is
@@ -184,6 +188,10 @@ each row then L2-normed, against float64's ranges at their bf16 points and expf'
 CUDA also checks `dense_rows`' scaled epilogue (Dory's logits, fp32(bf16(x . w)) times a per-column scale), bit-equal to its
 unscaled tile's outputs rounded and scaled on the host.
 HIP: the RMS norm rows, a wave a row (fp32, fp16 and bf16 out) and a block a row (fp16 and bf16 out).
+HIP also checks the 64-row prompt attention tile (`prefill.hip`): `tf_fa_paged_bf16` and `tf_fa_paged_f16` over a
+permuted page table at the 27B's, the 35B-A3B's and Nemotron's head counts, each output within attention_ref's float64
+bound, and `tf_fa_wide_bf16` and `tf_fa_wide_f16`, which no model calls, bit-equal to them on flat keys. A build runs
+its chip's cache type (fp16 on gfx1030, bf16 on the rest); gate 5 holds all four.
 
 RADIO's fused bias, activation and residual (`tf_vit_bias_act`) replaces `tf_vit_bias` then
 `tf_vit_pointwise`, byte for byte.

zig/kernels/hip/attention/prefill.hip

@@ -12,6 +12,7 @@ constexpr int kRows = 64;     // query rows a block
 constexpr int kKeys = 64;     // keys a tile
 constexpr int kChunk = 64;    // head dimensions a staged key or value chunk
 constexpr int kThreads = 256;
+using word4 = unsigned __attribute__((ext_vector_type(4)));   // a 16-byte load held as one value
 static_assert(kKeys == kPageTokens, "a key tile is one page");
 
 // A cache element type's pair for dot2, the rounding into it and the product.
@@ -51,22 +52,58 @@ __device__ __forceinline__ void pairs4(const void* at, Pair out[4]) {
     out[3] = __builtin_bit_cast(Pair, w.w);
 }
 
-// q rounds to the cache type once and stays in LDS; with PAGED, k and v are a layer's pools and a 64-key tile a page.
+// A 16-byte load of row `row` at element `col`; unless FULL, a row from `rows` on reads row `rows` - 1 instead.
+template <bool FULL, typename Cache>
+__device__ __forceinline__ word4 load16(const Cache* base, long long stride, int row, int rows, int col) {
+    return *reinterpret_cast<const word4*>(base + (FULL || row < rows ? row : rows - 1) * stride + col);
+}
+
+// A staged piece, zero for a row past the span.
+__device__ __forceinline__ uint4 kept(word4 w, bool in) {
+    return __builtin_bit_cast(uint4, in ? w : word4{0, 0, 0, 0});
+}
+
+// x from lane ^ M of its 16-lane row (M < 16): a DPP row_xmask move, all lanes active.
+template <int M>
+__device__ __forceinline__ float row_xor(float x) {
+    return __builtin_bit_cast(float, __builtin_amdgcn_mov_dpp(__builtin_bit_cast(int, x), 0x160 + M, 0xf, 0xf, false));
+}
+
+// A 16-lane row's max and sum by xor butterflies over lane distances 8, 4, 2, 1, each lane taking own then partner.
+__device__ __forceinline__ float row_max16(float x) {
+    x = fmaxf(x, row_xor<8>(x));
+    x = fmaxf(x, row_xor<4>(x));
+    x = fmaxf(x, row_xor<2>(x));
+    return fmaxf(x, row_xor<1>(x));
+}
+
+__device__ __forceinline__ float row_sum16(float x) {
+    x += row_xor<8>(x);
+    x += row_xor<4>(x);
+    x += row_xor<2>(x);
+    return x + row_xor<1>(x);
+}
+
+// One lane's x on every lane of its 16-lane row, `sel` naming that lane in each nibble: v_permlane16.
+template <typename P>
+__device__ __forceinline__ P row_lane(P x, unsigned sel) {
+    const int w = __builtin_bit_cast(int, x);
+    return __builtin_bit_cast(P, __builtin_amdgcn_permlane16(w, w, sel, sel, false, false));
+}
+
+// q rounds to the cache type once into LDS, p stays in registers; with PAGED, k and v are pools and a key tile a page.
 template <typename Cache, bool PAGED>
 __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k, const Cache* __restrict__ v,
                         float* __restrict__ out, int qlen, int span, int heads, int kv_heads, int d, float scale,
                         int q_pos0, long long k_sb, long long k_sh, long long k_ss, long long v_sb, long long v_sh,
                         long long v_ss, const unsigned* table, int count) {
     constexpr int kQChunks = 256 / 8;      // 16-byte chunks a q row
     constexpr int kKChunks = kChunk / 8;   // a staged key row's
-    constexpr int kPLd = kRows + 4;        // a key pair's probabilities, row-indexed
     __shared__ uint4 sQ[kRows * kQChunks];
     __shared__ uint4 sKV[kKeys * kKChunks];   // a key chunk (key, 8 dims) or a value chunk (key pair, 4 dims)
     using T = Dot<Cache>;
     using elem = typename T::elem;
     using pair_t = typename T::pair;
-    __shared__ alignas(16) pair_t sH[kKeys / 2 * kPLd];
-    __shared__ float sA[kRows], sL[kRows];
 
     const int q0 = blockIdx.x * kRows;
     const int head = blockIdx.y;
@@ -95,9 +132,8 @@ __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k
         sQ[row * kQChunks + (c ^ (row & 7))] = __builtin_bit_cast(uint4, h);
     }
 
-    // scores: rows sr + 16i against keys 4sk .. 4sk + 3; PV: rows pr .. pr + 3 at dims 4pc .. 4pc + 3 of each chunk
+    // rows sr + 16i: scores against keys 4sk .. 4sk + 3, PV at dims 4sk .. 4sk + 3 of each chunk
     const int sr = tid / 16, sk = tid % 16;
-    const int pr = (tid / 16) * 4, pc = tid % 16;
     float m_run[4], l_run[4];
 #pragma unroll
     for (int i = 0; i < 4; ++i) m_run[i] = -INFINITY, l_run[i] = 0.f;
@@ -107,31 +143,60 @@ __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k
 #pragma unroll
         for (int m = 0; m < 16; ++m) o[i][m] = 0.f;
 
-    for (int k0 = 0; k0 < max_vis; k0 += kKeys) {
-        // the tile's first key row in the keys' and the values' storage
-        const Cache* ktile = PAGED ? kbase + page_row(table[k0 / kPageTokens], count, kv, 0, d) : kbase + static_cast<long long>(k0) * k_ss;
-        const Cache* vtile = PAGED ? vbase + page_row(table[k0 / kPageTokens], count, kv, 0, d) : vbase + static_cast<long long>(k0) * v_ss;
+    // a tile ending by full_end is in range and visible to every stored row, so FULL drops its guards and masks
+    int full_end = q_pos0 + q0 + 1;
+    if (full_end > span) full_end = span;
+    // the first key row of the tile at k0 (on page `pg` when paged) in the keys' and the values' storage
+    const auto ktile_at = [&](int k0, unsigned pg) __attribute__((always_inline)) {
+        return PAGED ? kbase + page_row(pg, count, kv, 0, d) : kbase + static_cast<long long>(k0) * k_ss;
+    };
+    const auto vtile_at = [&](int k0, unsigned pg) __attribute__((always_inline)) {
+        return PAGED ? vbase + page_row(pg, count, kv, 0, d) : vbase + static_cast<long long>(k0) * v_ss;
+    };
+    // each staged chunk's loads issue before the previous chunk's products, which hide their latency
+    word4 kst[2], va, vb;
+    const auto load_k = [&]<bool FULL>(const Cache* ktile, int rows, int c0) __attribute__((always_inline)) {
+#pragma unroll
+        for (int t = 0; t < 2; ++t) {
+            const int e = tid + t * kThreads, key = e / kKChunks, c = e % kKChunks;
+            kst[t] = load16<FULL>(ktile, kstride, key, rows, 8 * (c0 + c));
+        }
+    };
+    // a value chunk: key pair `pair`, dims 8g .. 8g + 7 of both keys
+    const auto load_v = [&]<bool FULL>(const Cache* vtile, int rows, int c0) __attribute__((always_inline)) {
+        const int pair = tid / kKChunks, g = tid % kKChunks;
+        va = load16<FULL>(vtile, vstride, 2 * pair, rows, 8 * (c0 + g));
+        vb = load16<FULL>(vtile, vstride, 2 * pair + 1, rows, 8 * (c0 + g));
+    };
+    unsigned page = PAGED && max_vis > 0 ? table[0] : 0;
+    if (max_vis > 0) load_k.template operator()<false>(ktile_at(0, page), span, 0);
+    const auto tile = [&]<bool FULL>(int k0) __attribute__((always_inline)) {
+        const Cache* ktile = ktile_at(k0, page);
+        const Cache* vtile = vtile_at(k0, page);
+        const int rows = span - k0;
+        const bool more = k0 + kKeys < max_vis;
+        unsigned next = page;
         float acc[4][4];
 #pragma unroll
         for (int i = 0; i < 4; ++i)
 #pragma unroll
             for (int j = 0; j < 4; ++j) acc[i][j] = 0.f;
         for (int c0 = 0; c0 < chunks; c0 += kKChunks) {
-            uint4 staged[2];
-#pragma unroll
-            for (int t = 0; t < 2; ++t) {
-                const int e = tid + t * kThreads, key = e / kKChunks, c = e % kKChunks;
-                staged[t] = k0 + key < span
-                                ? *reinterpret_cast<const uint4*>(ktile + static_cast<long long>(key) * kstride + 8 * (c0 + c))
-                                : uint4{0, 0, 0, 0};
-            }
             __syncthreads();
 #pragma unroll
             for (int t = 0; t < 2; ++t) {
                 const int e = tid + t * kThreads, key = e / kKChunks, c = e % kKChunks;
-                sKV[key * kKChunks + (c ^ ((key >> 2) & 7))] = staged[t];
+                sKV[key * kKChunks + (c ^ ((key >> 2) & 7))] = kept(kst[t], FULL || key < rows);
             }
             __syncthreads();
+            if (PAGED && c0 == 0 && more) {
+                // a vector index keeps the page id out of scalar registers, so its load is awaited only where used
+                int at_page = (k0 + kKeys) / kPageTokens;
+                asm volatile("" : "+v"(at_page));
+                next = table[at_page];
+            }
+            if (c0 + kKChunks < chunks) load_k.template operator()<FULL>(ktile, rows, c0 + kKChunks);
+            else load_v.template operator()<FULL>(vtile, rows, 0);
 #pragma unroll 2
             for (int c = 0; c < kKChunks; ++c) {
                 pair_t qv[4][4], kf[4][4];
@@ -153,7 +218,9 @@ __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k
                         for (int m = 0; m < 4; ++m) acc[i][j] = T::dot(qv[i][m], kf[j][m], acc[i][j]);
             }
         }
-        // the online softmax in registers: a row's 64 keys sit on its 16 lanes
+        // the online softmax in registers: a row's 64 keys sit on its 16 lanes, key pair pp's p on lane pp / 2
+        pair_t ph[4][2];
+        float alpha_r[4];
 #pragma unroll
         for (int i = 0; i < 4; ++i) {
             const int row = sr + 16 * i;
@@ -162,11 +229,10 @@ __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k
             float tile_max = -INFINITY;
 #pragma unroll
             for (int j = 0; j < 4; ++j) {
-                acc[i][j] = (row < q_hi && k0 + 4 * sk + j < vis) ? acc[i][j] * scale : -INFINITY;
+                acc[i][j] = (FULL || (row < q_hi && k0 + 4 * sk + j < vis)) ? acc[i][j] * scale : -INFINITY;
                 tile_max = fmaxf(tile_max, acc[i][j]);
             }
-#pragma unroll
-            for (int lane = 8; lane > 0; lane >>= 1) tile_max = fmaxf(tile_max, __shfl_xor(tile_max, lane, 16));
+            tile_max = row_max16(tile_max);
             const float next_m = fmaxf(m_run[i], tile_max);
             const float alpha = next_m == -INFINITY ? 1.f : expf(m_run[i] - next_m);
             float p[4], sum = 0.f;
@@ -175,74 +241,71 @@ __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k
                 p[j] = next_m == -INFINITY ? 0.f : expf(acc[i][j] - next_m);
                 sum += p[j];
             }
-#pragma unroll
-            for (int lane = 8; lane > 0; lane >>= 1) sum += __shfl_xor(sum, lane, 16);
+            sum = row_sum16(sum);
             l_run[i] = alpha * l_run[i] + sum;
             m_run[i] = next_m;
-            sH[(2 * sk) * kPLd + row] = pair_t{static_cast<elem>(p[0]), static_cast<elem>(p[1])};
-            sH[(2 * sk + 1) * kPLd + row] = pair_t{static_cast<elem>(p[2]), static_cast<elem>(p[3])};
-            if (sk == 0) sA[row] = alpha;
+            ph[i][0] = pair_t{static_cast<elem>(p[0]), static_cast<elem>(p[1])};
+            ph[i][1] = pair_t{static_cast<elem>(p[2]), static_cast<elem>(p[3])};
+            alpha_r[i] = alpha;
         }
 #pragma unroll
         for (int at = 0; at < 256 / kChunk; ++at) {
             if (at * kKChunks >= chunks) break;
-            const int c0 = at * kKChunks;
-            // a value chunk: key pair `pair`, dims 8g .. 8g + 7 interleaved as (key, key + 1) fp16 pairs
+            // the chunk's key pairs as (key, key + 1) pairs at each of its dims
             const int pair = tid / kKChunks, g = tid % kKChunks;
-            const int key = k0 + 2 * pair;
-            const uint4 a = key < span ? *reinterpret_cast<const uint4*>(vtile + static_cast<long long>(2 * pair) * vstride + 8 * (c0 + g))
-                                       : uint4{0, 0, 0, 0};
-            const uint4 b = key + 1 < span
-                                ? *reinterpret_cast<const uint4*>(vtile + static_cast<long long>(2 * pair + 1) * vstride + 8 * (c0 + g))
-                                : uint4{0, 0, 0, 0};
-            __syncthreads();
-            const elem* ha = reinterpret_cast<const elem*>(&a);
-            const elem* hb = reinterpret_cast<const elem*>(&b);
+            using elem8 = elem __attribute__((ext_vector_type(8)));
+            const elem8 ha = __builtin_bit_cast(elem8, kept(va, FULL || 2 * pair < rows));
+            const elem8 hb = __builtin_bit_cast(elem8, kept(vb, FULL || 2 * pair + 1 < rows));
             pair_t lo[4], hi[4];
 #pragma unroll
             for (int m = 0; m < 4; ++m) lo[m] = pair_t{ha[m], hb[m]}, hi[m] = pair_t{ha[4 + m], hb[4 + m]};
+            __syncthreads();
             sKV[pair * 16 + 2 * g] = __builtin_bit_cast(uint4, lo);
             sKV[pair * 16 + 2 * g + 1] = __builtin_bit_cast(uint4, hi);
             __syncthreads();
+            if ((at + 1) * kKChunks < chunks) load_v.template operator()<FULL>(vtile, rows, (at + 1) * kKChunks);
+            else if (more) load_k.template operator()<false>(ktile_at(k0 + kKeys, next), rows - kKeys, 0);
             float pv[4][4];
 #pragma unroll
             for (int i = 0; i < 4; ++i)
 #pragma unroll
                 for (int m = 0; m < 4; ++m) pv[i][m] = 0.f;
 #pragma unroll 4
             for (int pp = 0; pp < kKeys / 2; ++pp) {
-                pair_t pr4[4], vv[4];
-                pairs4(sH + pp * kPLd + pr, pr4);
-                pairs4(&sKV[pp * 16 + pc], vv);
+                pair_t vv[4];
+                pairs4(&sKV[pp * 16 + sk], vv);
+                const unsigned lane = (pp / 2) * 0x11111111u;
 #pragma unroll
-                for (int i = 0; i < 4; ++i)
+                for (int i = 0; i < 4; ++i) {
+                    const pair_t pp4 = row_lane(ph[i][pp % 2], lane);
 #pragma unroll
-                    for (int m = 0; m < 4; ++m) pv[i][m] = T::dot(pr4[i], vv[m], pv[i][m]);
+                    for (int m = 0; m < 4; ++m) pv[i][m] = T::dot(pp4, vv[m], pv[i][m]);
+                }
             }
 #pragma unroll
             for (int i = 0; i < 4; ++i)
 #pragma unroll
                 for (int m = 0; m < 4; ++m) {
                     float& acc_o = o[i][4 * at + m];
-                    acc_o = fmaf(sA[pr + i], acc_o, pv[i][m]);
+                    acc_o = fmaf(alpha_r[i], acc_o, pv[i][m]);
                 }
         }
-    }
-    if (sk == 0)
-#pragma unroll
-        for (int i = 0; i < 4; ++i) sL[sr + 16 * i] = l_run[i];
-    __syncthreads();
+        page = next;
+    };
+    int k0 = 0;
+    for (; k0 + kKeys <= full_end; k0 += kKeys) tile.template operator()<true>(k0);
+    for (; k0 < max_vis; k0 += kKeys) tile.template operator()<false>(k0);
 #pragma unroll
     for (int i = 0; i < 4; ++i) {
-        const int row = pr + i;
+        const int row = sr + 16 * i;
         if (row >= q_hi) continue;
-        const float inv = sL[row] == 0.f ? 0.f : 1.f / sL[row];
+        const float inv = l_run[i] == 0.f ? 0.f : 1.f / l_run[i];
 #pragma unroll
         for (int at = 0; at < 256 / kChunk; ++at) {
             if (at * kKChunks >= chunks) break;
 #pragma unroll
             for (int m = 0; m < 4; ++m)
-                orows[static_cast<long long>(row) * d + 64 * at + 4 * pc + m] = o[i][4 * at + m] * inv;
+                orows[static_cast<long long>(row) * d + 64 * at + 4 * sk + m] = o[i][4 * at + m] * inv;
         }
     }
 }

zig/src/core/contract/hip_records.zig

@@ -479,7 +479,7 @@ const own_records = [_]Record{
             .{ .vgprs = 36, .sgprs = 32, .lds = 50688, .private = 0, .bound = 1024, .wave = 32 },
         },
     },
-} ++ @import("hip_records_pair.zig").records;
+} ++ @import("hip_records_pair.zig").records ++ @import("hip_records_prompt.zig").records;
 
 pub const records = own_records ++ @import("hip_records_stream.zig").records;