From 44917d6694a39886751007d98c938fd5b079bb45 Mon Sep 17 00:00:00 2001
From: Abraham <119352736+AbrahamCete@users.noreply.github.com>
Date: Sat, 31 Oct 2026 16:59:54 +0100
Subject: [PATCH 1/7] test(hip): pin the paged prompt tile's output bits on 27B
 shapes

The paged tile was only compared with the flat tile of the same template, so a change to both left the test green. The new test digests its output on 24 heads over 4 kv heads at 256 dims, with partial pages, a short last row block, diagonal tiles and later chunks, and pins the digest on gfx1151.
---
 zig/src/hip/kernel_tests.zig               |   1 +
 zig/src/hip/kernel_tests/attention_pin.zig | 122 +++++++++++++++++++++
 2 files changed, 123 insertions(+)
 create mode 100644 zig/src/hip/kernel_tests/attention_pin.zig

diff --git a/zig/src/hip/kernel_tests.zig b/zig/src/hip/kernel_tests.zig
index 6c8dfcb91d..934db3f04a 100644
--- a/zig/src/hip/kernel_tests.zig
+++ b/zig/src/hip/kernel_tests.zig
@@ -23,6 +23,7 @@ test {
     _ = @import("kernel_tests/decode.zig");
     _ = @import("kernel_tests/chain.zig");
     _ = @import("kernel_tests/attention.zig");
+    _ = @import("kernel_tests/attention_pin.zig");
     _ = @import("kernel_tests/ops.zig");
     _ = @import("kernel_tests/argmax.zig");
     _ = @import("kernel_tests/tp.zig");
diff --git a/zig/src/hip/kernel_tests/attention_pin.zig b/zig/src/hip/kernel_tests/attention_pin.zig
new file mode 100644
index 0000000000..51ba5126ce
--- /dev/null
+++ b/zig/src/hip/kernel_tests/attention_pin.zig
@@ -0,0 +1,122 @@
+//! The paged prompt tile's exact output bits on fixed integer-derived inputs, pinned by a digest per architecture.
+const std = @import("std");
+const rig = @import("rig.zig");
+const Rig = rig.Rig;
+const gpa = rig.gpa;
+const at = rig.at;
+const kernels = @import("../kernels.zig");
+
+const page = 64;
+const heads = 24;
+const kv_heads = 4;
+const d = 256;
+
+/// The digest of every case's output on an architecture whose bits were recorded.
+const pins = [_]struct { arch: []const u8, digest: u64 }{
+    .{ .arch = "gfx1151", .digest = 0x30a7911a6c016c2b },
+};
+
+/// A prompt chunk: `qlen` queries from position `q_pos0`.
+const Case = struct { qlen: usize, q_pos0: usize };
+
+/// Partial last pages, a short last row block, diagonal tiles and, from a later chunk, whole tiles before them.
+const cases = [_]Case{
+    .{ .qlen = 200, .q_pos0 = 0 },
+    .{ .qlen = 150, .q_pos0 = 320 },
+    .{ .qlen = 96, .q_pos0 = 100 },
+};
+
+fn mix(x: u32) u32 {
+    var h = x;
+    h ^= h >> 16;
+    h *%= 0x7feb352d;
+    h ^= h >> 15;
+    h *%= 0x846ca68b;
+    h ^= h >> 16;
+    return h;
+}
+
+/// A bf16 of any sign and mantissa with magnitude in [2^-4, 2^2).
+fn cacheBits(i: usize, salt: u32) u16 {
+    const h = mix(@as(u32, @truncate(i)) *% 0x9E3779B1 +% salt);
+    const sign: u16 = @intCast(h >> 31);
+    const exp: u16 = @intCast(123 + (h >> 8) % 6);
+    return sign << 15 | exp << 7 | @as(u16, @intCast(h & 0x7f));
+}
+
+/// An fp32 query element with a full mantissa, so its rounding to the cache type matters.
+fn queryValue(i: usize) f32 {
+    const h = mix(@as(u32, @truncate(i)) *% 0x85EBCA77 +% 0x51);
+    const sign: u32 = h >> 31;
+    const exp: u32 = 123 + (h >> 23) % 6;
+    const mant: u32 = mix(h) & 0x7fffff;
+    return @bitCast(sign << 31 | exp << 23 | mant);
+}
+
+/// FNV-1a over `bytes`, continuing from `h`.
+fn fnv(h: u64, bytes: []const u8) u64 {
+    var x = h;
+    for (bytes) |b| {
+        x ^= b;
+        x *%= 0x100000001b3;
+    }
+    return x;
+}
+
+/// One case's output through the paged tile, folded into `h`; slots past the span hold NaN so a read of them shows.
+fn run(t: *Rig, c: Case, h: u64) !u64 {
+    const span = c.q_pos0 + c.qlen;
+    const count = (span + page - 1) / page;
+    const pool = kv_heads * count * page * d;
+    const k = try gpa.alloc(u16, pool);
+    defer gpa.free(k);
+    const v = try gpa.alloc(u16, pool);
+    defer gpa.free(v);
+    const table = try gpa.alloc(u32, count);
+    defer gpa.free(table);
+    // a permutation, so a walk that ignored the table would read the wrong keys
+    for (table, 0..) |*p, i| p.* = @intCast((i * 5 + 3) % count);
+    @memset(k, 0x7FC0);
+    @memset(v, 0x7FC0);
+    for (0..kv_heads) |kh| for (0..span) |pos| {
+        const to = ((kh * count + table[pos / page]) * page + pos % page) * d;
+        for (0..d) |j| {
+            const i = (kh * span + pos) * d + j;
+            k[to + j] = cacheBits(i, 0x4B);
+            v[to + j] = cacheBits(i, 0x56);
+        }
+    };
+    const hq = try gpa.alloc(f32, heads * c.qlen * d);
+    defer gpa.free(hq);
+    for (hq, 0..) |*x, i| x.* = queryValue(i + c.q_pos0 * heads * d);
+    var q = try t.upload(hq);
+    defer q.free();
+    var pool_k = try t.upload(k);
+    defer pool_k.free();
+    var pool_v = try t.upload(v);
+    defer pool_v.free();
+    var dev_table = try t.upload(table);
+    defer dev_table.free();
+    var out = try t.alloc(heads * c.qlen * d * 4);
+    defer out.free();
+    const scale: f32 = 1 / @sqrt(@as(f32, d));
+    try t.on.pagedCausal(at(q), at(pool_k), at(pool_v), at(dev_table), at(out), c.qlen, span, heads, kv_heads, d, scale, c.q_pos0, count, 1, t.stream.handle);
+    try t.stream.synchronize();
+    const got = try rig.download(f32, out);
+    defer gpa.free(got);
+    for (got) |x| try std.testing.expect(std.math.isFinite(x));
+    return fnv(h, std.mem.sliceAsBytes(got));
+}
+
+test "the paged prompt tile writes its pinned bits for 24 heads over 4 kv heads at 256 dims, partial pages included" {
+    const t = try Rig.open(1 << 20);
+    defer t.close();
+    if (!t.bf16) return error.SkipZigTest;
+    const want = for (pins) |p| {
+        if (std.mem.eql(u8, p.arch, kernels.arch)) break p.digest;
+    } else return error.SkipZigTest;
+    var h: u64 = 0xcbf29ce484222325;
+    for (cases) |c| h = try run(t, c, h);
+    if (h != want) std.debug.print("paged tile digest on {s}: 0x{x:0>16}\n", .{ kernels.arch, h });
+    try std.testing.expectEqual(want, h);
+}

From de0b7649a91635d2a5c633833064cb47e4df4287 Mon Sep 17 00:00:00 2001
From: Abraham <119352736+AbrahamCete@users.noreply.github.com>
Date: Sat, 31 Oct 2026 17:27:33 +0100
Subject: [PATCH 2/7] perf(hip): skip the masks and range guards on whole
 prompt tiles

A key tile that ends inside the span and before the first row of the 64-row block is seen in full by every stored row. Such tiles now take a path without the per-score mask and without the guarded key and value loads. The tiles are visited in the same order and every stored value is the same, so the output bits do not change.
---
 zig/kernels/hip/attention/prefill.hip | 27 +++++++++++++++++----------
 1 file changed, 17 insertions(+), 10 deletions(-)

diff --git a/zig/kernels/hip/attention/prefill.hip b/zig/kernels/hip/attention/prefill.hip
index 87e39f57d4..35f46fb206 100644
--- a/zig/kernels/hip/attention/prefill.hip
+++ b/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.
@@ -107,7 +108,10 @@ __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) {
+    // 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;
+    const auto tile = [&]<bool FULL>(int k0) __attribute__((always_inline)) {
         // 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;
@@ -117,19 +121,19 @@ __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k
 #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];
+            word4 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};
+                staged[t] = (FULL || k0 + key < span)
+                                ? *reinterpret_cast<const word4*>(ktile + static_cast<long long>(key) * kstride + 8 * (c0 + c))
+                                : word4{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))] = __builtin_bit_cast(uint4, staged[t]);
             }
             __syncthreads();
 #pragma unroll 2
@@ -162,7 +166,7 @@ __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
@@ -190,9 +194,9 @@ __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k
             // a value chunk: key pair `pair`, dims 8g .. 8g + 7 interleaved as (key, key + 1) fp16 pairs
             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))
+            const uint4 a = (FULL || 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
+            const uint4 b = (FULL || 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();
@@ -227,7 +231,10 @@ __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k
                     acc_o = fmaf(sA[pr + i], acc_o, pv[i][m]);
                 }
         }
-    }
+    };
+    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);
     if (sk == 0)
 #pragma unroll
         for (int i = 0; i < 4; ++i) sL[sr + 16 * i] = l_run[i];

From 709184e2990e12cf6c165ab43479c7af9324002f Mon Sep 17 00:00:00 2001
From: Abraham <119352736+AbrahamCete@users.noreply.github.com>
Date: Sat, 31 Oct 2026 17:37:34 +0100
Subject: [PATCH 3/7] perf(hip): reduce the prompt tile's softmax rows with DPP
 butterflies

The 16-lane max and sum of each softmax row now pair lanes with DPP row_xmask moves instead of ds_bpermute shuffles. The pairing over lane distances 8, 4, 2 and 1 and each lane's own-then-partner operand order are unchanged, so the output bits are the same.
---
 zig/kernels/hip/attention/prefill.hip | 27 +++++++++++++++++++++++----
 1 file changed, 23 insertions(+), 4 deletions(-)

diff --git a/zig/kernels/hip/attention/prefill.hip b/zig/kernels/hip/attention/prefill.hip
index 35f46fb206..035d6ca3ee 100644
--- a/zig/kernels/hip/attention/prefill.hip
+++ b/zig/kernels/hip/attention/prefill.hip
@@ -52,6 +52,27 @@ __device__ __forceinline__ void pairs4(const void* at, Pair out[4]) {
     out[3] = __builtin_bit_cast(Pair, w.w);
 }
 
+// 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);
+}
+
 // 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.
 template <typename Cache, bool PAGED>
 __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k, const Cache* __restrict__ v,
@@ -169,8 +190,7 @@ __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k
                 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;
@@ -179,8 +199,7 @@ __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])};

From 751f80f87e374973f734fe37c9579bdebc45009e Mon Sep 17 00:00:00 2001
From: Abraham <119352736+AbrahamCete@users.noreply.github.com>
Date: Sat, 31 Oct 2026 17:47:31 +0100
Subject: [PATCH 4/7] perf(hip): prefetch the prompt tile's next key and value
 chunks

Each staged chunk's global loads now issue right after the previous chunk reaches LDS, before that chunk's products; the first value chunk and the next tile's first key chunk issue the same way, and the next tile's page id loads one chunk ahead. A row past the span reads the span's last row and is zeroed when staged, instead of a guarded load. The arithmetic and its order are unchanged, so the output bits are the same.
---
 zig/kernels/hip/attention/prefill.hip | 80 +++++++++++++++++++--------
 1 file changed, 57 insertions(+), 23 deletions(-)

diff --git a/zig/kernels/hip/attention/prefill.hip b/zig/kernels/hip/attention/prefill.hip
index 035d6ca3ee..2479f4b4a5 100644
--- a/zig/kernels/hip/attention/prefill.hip
+++ b/zig/kernels/hip/attention/prefill.hip
@@ -52,6 +52,17 @@ __device__ __forceinline__ void pairs4(const void* at, Pair out[4]) {
     out[3] = __builtin_bit_cast(Pair, w.w);
 }
 
+// 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) {
@@ -132,31 +143,57 @@ __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k
     // 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)) {
-        // 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;
+        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) {
-            word4 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] = (FULL || k0 + key < span)
-                                ? *reinterpret_cast<const word4*>(ktile + static_cast<long long>(key) * kstride + 8 * (c0 + c))
-                                : word4{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))] = __builtin_bit_cast(uint4, 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];
@@ -209,24 +246,20 @@ __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k
 #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 = (FULL || 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 = (FULL || 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)
@@ -250,6 +283,7 @@ __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k
                     acc_o = fmaf(sA[pr + i], acc_o, pv[i][m]);
                 }
         }
+        page = next;
     };
     int k0 = 0;
     for (; k0 + kKeys <= full_end; k0 += kKeys) tile.template operator()<true>(k0);

From 139d0477334ca54c8ac81b22796c7a417364d1a9 Mon Sep 17 00:00:00 2001
From: Abraham <119352736+AbrahamCete@users.noreply.github.com>
Date: Sat, 31 Oct 2026 18:10:35 +0100
Subject: [PATCH 5/7] perf(hip): keep the prompt tile's probabilities in
 registers

Each lane keeps its rows' bf16 probability pairs and rescale factors in registers, and the PV products take a key pair's probabilities from the lane that holds them with v_permlane16 instead of an LDS round trip. The PV work covers rows sr + 16i at dims 4sk .. 4sk + 3, the rows each lane already holds, and the epilogue divides by the lane's own running sum. Without the probability, rescale and sum arrays a block uses 40,960 bytes of LDS instead of 50,176, so three blocks fit a workgroup processor and gfx1151 runs 6 waves per SIMD instead of 4. Every sum keeps its operands and order, so the output bits are the same.
---
 zig/kernels/hip/attention/prefill.hip | 49 ++++++++++++++-------------
 1 file changed, 26 insertions(+), 23 deletions(-)

diff --git a/zig/kernels/hip/attention/prefill.hip b/zig/kernels/hip/attention/prefill.hip
index 2479f4b4a5..1f8a3dc977 100644
--- a/zig/kernels/hip/attention/prefill.hip
+++ b/zig/kernels/hip/attention/prefill.hip
@@ -84,7 +84,14 @@ __device__ __forceinline__ float row_sum16(float x) {
     return x + row_xor<1>(x);
 }
 
-// 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.
+// 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,
@@ -92,14 +99,11 @@ __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k
                         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;
@@ -128,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;
@@ -215,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;
@@ -239,9 +244,9 @@ __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k
             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) {
@@ -267,20 +272,22 @@ __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k
                 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]);
                 }
         }
         page = next;
@@ -288,21 +295,17 @@ __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k
     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);
-    if (sk == 0)
-#pragma unroll
-        for (int i = 0; i < 4; ++i) sL[sr + 16 * i] = l_run[i];
-    __syncthreads();
 #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;
         }
     }
 }

From 5c3ba8935a23d7da7d63a85a89fef1d1c600de2b Mon Sep 17 00:00:00 2001
From: Abraham <119352736+AbrahamCete@users.noreply.github.com>
Date: Sat, 31 Oct 2026 21:38:43 +0100
Subject: [PATCH 6/7] test(hip): register the prompt attention tiles with the
 kernel contract

The four symbols of prefill.hip's 64-row tile now have contract specs and launch records. tf_fa_paged_bf16 and tf_fa_paged_f16 run 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 gate 4's equal keys prove fp32 sums. tf_fa_wide_bf16 and tf_fa_wide_f16, which no model calls, must equal the paged tile's bits on the same keys. Gate 1 is a known finding on both, and v1.2.0 shows the same one: a row whose every score is NaN writes 0, and an infinite value at a key whose fp32 weight is 0 gives NaN, as P x V takes every key of its tile. The records hold what ROCm 7.13's hipcc compiles for gfx1030, gfx1100, gfx1151 and gfx1201.
---
 docs/kernel-contract.md                      |   8 +
 zig/src/core/contract/hip_records.zig        |   2 +-
 zig/src/core/contract/hip_records_prompt.zig |  38 ++
 zig/tests/hip/contract_attn.zig              | 144 ++++++++
 zig/tests/hip/contract_main.zig              |   5 +
 zig/tests/hip/contract_prefill.zig           | 359 +++++++++++++++++++
 6 files changed, 555 insertions(+), 1 deletion(-)
 create mode 100644 zig/src/core/contract/hip_records_prompt.zig
 create mode 100644 zig/tests/hip/contract_attn.zig
 create mode 100644 zig/tests/hip/contract_prefill.zig

diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md
index b690ec8f9e..205d44fa11 100644
--- a/docs/kernel-contract.md
+++ b/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.
diff --git a/zig/src/core/contract/hip_records.zig b/zig/src/core/contract/hip_records.zig
index af2514abd2..57b3f8de51 100644
--- a/zig/src/core/contract/hip_records.zig
+++ b/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;
 
diff --git a/zig/src/core/contract/hip_records_prompt.zig b/zig/src/core/contract/hip_records_prompt.zig
new file mode 100644
index 0000000000..dbf91f4871
--- /dev/null
+++ b/zig/src/core/contract/hip_records_prompt.zig
@@ -0,0 +1,38 @@
+//! Gate 5's records of the 64-row prompt attention tiles (prefill.hip's tf_fa_wide_* and tf_fa_paged_*).
+const limits = @import("limits.zig");
+const Record = @import("hip_records.zig").Record;
+
+const U = limits.HipUsage;
+
+/// The 64-row prompt attention tile (prefill.hip), flat or over pages, at 256 threads and no dynamic LDS.
+fn promptTile(comptime symbol: []const u8, usage: [limits.rdna.len]U) Record {
+    const name = "prompt attention " ++ symbol;
+    return .{ .launch = .{ .name = name, .symbol = symbol, .image = "prefill", .threads = 256 }, .usage = usage };
+}
+
+pub const records = [_]Record{
+    promptTile("tf_fa_wide_f16", .{
+        .{ .vgprs = 231, .sgprs = 45, .lds = 40960, .private = 0, .bound = 256, .wave = 32 },
+        .{ .vgprs = 231, .sgprs = 45, .lds = 40960, .private = 0, .bound = 256, .wave = 32 },
+        .{ .vgprs = 231, .sgprs = 45, .lds = 40960, .private = 0, .bound = 256, .wave = 32 },
+        .{ .vgprs = 231, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 },
+    }),
+    promptTile("tf_fa_wide_bf16", .{
+        .{ .vgprs = 247, .sgprs = 45, .lds = 40960, .private = 0, .bound = 256, .wave = 32 },
+        .{ .vgprs = 231, .sgprs = 45, .lds = 40960, .private = 0, .bound = 256, .wave = 32 },
+        .{ .vgprs = 231, .sgprs = 45, .lds = 40960, .private = 0, .bound = 256, .wave = 32 },
+        .{ .vgprs = 231, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 },
+    }),
+    promptTile("tf_fa_paged_f16", .{
+        .{ .vgprs = 232, .sgprs = 51, .lds = 40960, .private = 0, .bound = 256, .wave = 32 },
+        .{ .vgprs = 232, .sgprs = 51, .lds = 40960, .private = 0, .bound = 256, .wave = 32 },
+        .{ .vgprs = 232, .sgprs = 51, .lds = 40960, .private = 0, .bound = 256, .wave = 32 },
+        .{ .vgprs = 232, .sgprs = 51, .lds = 40960, .private = 0, .bound = 256, .wave = 32 },
+    }),
+    promptTile("tf_fa_paged_bf16", .{
+        .{ .vgprs = 250, .sgprs = 50, .lds = 40960, .private = 0, .bound = 256, .wave = 32 },
+        .{ .vgprs = 232, .sgprs = 51, .lds = 40960, .private = 0, .bound = 256, .wave = 32 },
+        .{ .vgprs = 232, .sgprs = 51, .lds = 40960, .private = 0, .bound = 256, .wave = 32 },
+        .{ .vgprs = 232, .sgprs = 51, .lds = 40960, .private = 0, .bound = 256, .wave = 32 },
+    }),
+};
diff --git a/zig/tests/hip/contract_attn.zig b/zig/tests/hip/contract_attn.zig
new file mode 100644
index 0000000000..7b34949869
--- /dev/null
+++ b/zig/tests/hip/contract_attn.zig
@@ -0,0 +1,144 @@
+//! Attention kernels under the contract on HIP: their fp32 outputs against core/attention_ref.zig, head by head.
+const std = @import("std");
+const core = @import("core");
+
+const contract = core.contract;
+const c = contract.spec;
+const g = contract.gates;
+const rows = contract.rows;
+const ref = core.attention_ref;
+
+/// Why gates 1 and 3 run attention_ref's bound: P enters P x V in the cache type; the spec's bound takes fp32 terms.
+pub const why = "P is rounded to the cache type before the P x V dot2: gates 1 and 3 hold each output to " ++
+    "attention_ref's float64 bound";
+
+/// One output head's float64 inputs, slices of the test's: its query, its visible keys' K and V rows ([keys][d]).
+pub const Head = struct { q: []const f64 = &.{}, k: []const f64 = &.{}, v: []const f64 = &.{}, n: usize = 0 };
+
+/// The float64 side of an attention test: `t.load(i, *Head)` points output head i at its inputs (fp32 outputs).
+pub const Ref = struct {
+    a: std.mem.Allocator,
+    d: usize,
+    head: Head = .{},
+    keys: []usize,
+    p: []f64,
+    want: []f64,
+    tol: []f64,
+    msg: [240]u8 = undefined,
+
+    pub fn init(a: std.mem.Allocator, d: usize, most: usize) !Ref {
+        const keys = try a.alloc(usize, most);
+        errdefer a.free(keys);
+        for (keys, 0..) |*k, i| k.* = i;
+        const p = try a.alloc(f64, most);
+        errdefer a.free(p);
+        const want = try a.alloc(f64, d);
+        errdefer a.free(want);
+        return .{ .a = a, .d = d, .keys = keys, .p = p, .want = want, .tol = try a.alloc(f64, d) };
+    }
+
+    pub fn deinit(r: *Ref) void {
+        for ([_][]f64{ r.p, r.want, r.tol }) |x| r.a.free(x);
+        r.a.free(r.keys);
+    }
+
+    /// Gates 1 and 3: every output's class equals float64's, and a finite one is within attention_ref's bound.
+    pub fn check(r: *Ref, t: anytype, heads: usize, scale: f64, out: []const u8, at: g.Gate) !?[]const u8 {
+        var planted: usize = 0;
+        var differ: usize = 0;
+        var first: struct { i: usize = 0, e: usize = 0, got: c.Class = .finite, want: c.Class = .finite } = .{};
+        for (0..heads) |i| {
+            try t.load(i, &r.head);
+            ref.head(r.head.q, r.head.k, r.head.v, r.keys[0..r.head.n], r.d, scale, r.want, r.tol);
+            for (r.want, r.tol, 0..) |w, tl, e| {
+                const got = c.read(.f32, out, i * r.d + e);
+                if (c.classOf(w) != .finite) planted += 1;
+                if (c.classOf(got) != c.classOf(w)) {
+                    if (differ == 0) first = .{ .i = i, .e = e, .got = c.classOf(got), .want = c.classOf(w) };
+                    differ += 1;
+                    continue;
+                }
+                if (at != .nonfinite and std.math.isFinite(w) and !(@abs(got - w) <= tl)) {
+                    const say = "head {d} element {d} is {e:.6} where float64 gives {e:.6}: {d:.2} of " ++
+                        "attention_ref's bound";
+                    return try std.fmt.bufPrint(&r.msg, say, .{ i, e, got, w, @abs(got - w) / tl });
+                }
+            }
+        }
+        if (differ > 0) {
+            const say = "{d} of {d} outputs differ in class; first at head {d} element {d}: {t} where float64 " ++
+                "gives {t}";
+            return try std.fmt.bufPrint(&r.msg, say, .{ differ, heads * r.d, first.i, first.e, first.got, first.want });
+        }
+        if (at == .nonfinite and (planted == 0 or planted == heads * r.d)) {
+            const say = "the planted values make {d} of {d} reference outputs non-finite: no proof";
+            return try std.fmt.bufPrint(&r.msg, say, .{ planted, heads * r.d });
+        }
+        return null;
+    }
+
+    fn weights(r: *Ref, scale: f64) f64 {
+        const h = r.head;
+        var m: f64 = -std.math.inf(f64);
+        for (r.p[0..h.n], 0..) |*p, j| {
+            var s: f64 = 0;
+            for (0..r.d) |e| s += h.q[e] * h.k[j * r.d + e];
+            p.* = s * scale;
+            m = @max(m, p.*);
+        }
+        var l: f64 = 0;
+        for (r.p[0..h.n]) |*p| {
+            p.* = @exp(p.* - m);
+            l += p.*;
+        }
+        return l;
+    }
+
+    /// Float64 outputs and each one's term magnitude, sum p_j |v_j| / sum p_j (gate 4 reads them).
+    pub fn reference(r: *Ref, t: anytype, heads: usize, scale: f64, want: []f64, mag: []f64) !void {
+        for (0..heads) |i| {
+            try t.load(i, &r.head);
+            const l = r.weights(scale);
+            for (0..r.d) |e| {
+                var s: f64 = 0;
+                var abs: f64 = 0;
+                for (r.p[0..r.head.n], 0..) |p, j| {
+                    s += p * r.head.v[j * r.d + e];
+                    abs += p * @abs(r.head.v[j * r.d + e]);
+                }
+                want[i * r.d + e] = s / l;
+                mag[i * r.d + e] = abs / l;
+            }
+        }
+    }
+
+    /// Gate 4's sums in an accumulator: its keys score alike, so every p is 1 and the sums are V's terms and l's ones.
+    pub fn simulate(r: *Ref, t: anytype, heads: usize, scale: f64, acc: c.Accum, order: g.Order, out: []f64) !void {
+        const terms = try r.a.alloc(f64, r.p.len);
+        defer r.a.free(terms);
+        const ones = try r.a.alloc(f64, r.p.len);
+        defer r.a.free(ones);
+        @memset(ones, 1);
+        for (0..heads) |i| {
+            try t.load(i, &r.head);
+            _ = r.weights(scale);
+            for (r.p[0..r.head.n]) |p| if (p != 1) return error.KeysScoreApart;
+            const l = rows.sumIn(acc, order, ones[0..r.head.n]);
+            for (0..r.d) |e| {
+                for (terms[0..r.head.n], 0..) |*x, j| x.* = r.head.v[j * r.d + e];
+                out[i * r.d + e] = rows.sumIn(acc, order, terms[0..r.head.n]) / l;
+            }
+        }
+    }
+};
+
+/// Gate 4's values: odd integers, every other one times 64, the second half cancelling the first (before rounding).
+pub fn mirrored(rng: std.Random, out: []f64) void {
+    const h = out.len / 2;
+    for (0..h) |i| {
+        const v: f64 = @floatFromInt((1 + 2 * rng.uintLessThan(u32, 128)) * @as(u32, if (i % 2 == 1) 64 else 1));
+        out[i] = v;
+        out[h + i] = (if (i < 8) @as(f64, 1) else 0) - v;
+    }
+    if (out.len % 2 == 1) out[out.len - 1] = 0;
+}
diff --git a/zig/tests/hip/contract_main.zig b/zig/tests/hip/contract_main.zig
index 3f091c3ffd..44e0dca46d 100644
--- a/zig/tests/hip/contract_main.zig
+++ b/zig/tests/hip/contract_main.zig
@@ -11,6 +11,7 @@ const mamba_k = @import("contract_mamba.zig");
 const scan_k = @import("contract_scan.zig");
 const moe_k = @import("contract_nemo_moe.zig");
 const walk_k = @import("contract_walk.zig");
+const prefill_k = @import("contract_prefill.zig");
 const add_rms_k = @import("contract_add_rms.zig");
 const window_k = @import("contract_window.zig");
 const relu2_k = @import("contract_relu2.zig");
@@ -112,6 +113,10 @@ fn run(a: std.mem.Allocator, io: std.Io) !usize {
     }
     defer for (&attends) |*w| w.deinit();
     for (&attends) |*w| try kernels.append(a, w.kernel());
+    var prefills: [2]prefill_k.Prefill = undefined;
+    for (&prefills, [_]bool{ true, false }) |*p, paged| p.* = .init(a, &r, s, &l, paged);
+    defer for (&prefills) |*p| p.deinit();
+    for (&prefills) |*p| try kernels.append(a, p.kernel());
     var add_rms: [3]add_rms_k.AddRms = undefined;
     for (&add_rms, std.enums.values(add_rms_k.Mode)) |*m, w| m.* = .init(a, &r, s, &l, w);
     defer for (&add_rms) |*m| m.deinit();
diff --git a/zig/tests/hip/contract_prefill.zig b/zig/tests/hip/contract_prefill.zig
new file mode 100644
index 0000000000..f1758173de
--- /dev/null
+++ b/zig/tests/hip/contract_prefill.zig
@@ -0,0 +1,359 @@
+//! The 64-row prompt attention tile (prefill.hip) under the contract on HIP: over pages, and flat bit-equal to it.
+const std = @import("std");
+const hip = @import("hip");
+const core = @import("core");
+const rig_ = @import("contract_rig.zig");
+const kit = @import("contract_kit.zig");
+const attn = @import("contract_attn.zig");
+
+const contract = core.contract;
+const c = contract.spec;
+const g = contract.gates;
+
+const page = 64;
+
+/// A caller's query heads (a shape's n) over its kv heads at its head dim.
+const Geometry = struct { heads: usize, kv: usize, d: usize };
+
+/// The 27B's and the 35B-A3B's full-attention layers, and Nemotron's.
+const geometries = [_]Geometry{
+    .{ .heads = 24, .kv = 4, .d = 256 },
+    .{ .heads = 16, .kv = 2, .d = 256 },
+    .{ .heads = 32, .kv = 2, .d = 128 },
+};
+
+/// Prompt chunks (rows from position k): a first chunk with partial pages, later chunks over whole tiles before them.
+const served = [_]c.Shape{
+    .{ .rows = 200, .n = 24, .k = 0 },
+    .{ .rows = 150, .n = 24, .k = 320 },
+    .{ .rows = 96, .n = 24, .k = 100 },
+    .{ .rows = 64, .n = 16, .k = 2048 },
+    .{ .rows = 128, .n = 32, .k = 1024 },
+};
+
+const names = [2][2][]const u8{
+    .{
+        "hip prompt tile tf_fa_wide_f16 (flat fp16 cache, no served caller)",
+        "hip prompt tile tf_fa_wide_bf16 (flat bf16 cache, no served caller)",
+    },
+    .{ "hip prompt tile tf_fa_paged_f16 (fp16 cache pages)", "hip prompt tile tf_fa_paged_bf16 (bf16 cache pages)" },
+};
+
+/// What gate 1 finds on the tile: NaN scores leave its max, and its P x V takes every key of a tile.
+const finding = "a row whose every score is NaN writes 0 where float64 gives NaN (its running max is fmaxf, " ++
+    "which drops NaN, 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 tile, or under fp32's least) gives NaN where float64 gives inf or a finite " ++
+    "value, as 0 x inf enters P x V";
+
+/// A row's P x V chain: a dot2 a key pair and a fold a tile, under its keys plus a tile.
+fn keyChain(s: c.Shape) usize {
+    return s.k + s.rows + page;
+}
+
+/// Rows of a prompt chunk (rows), query heads (n) and the chunk's first position (k); the flat tile equals the paged.
+pub fn spec(k: kit.Kind, paged: bool) c.Spec {
+    const bf16 = k == .bf16;
+    return .{
+        .name = names[@intFromBool(paged)][@intFromBool(bf16)],
+        .accum = .f32,
+        .depth = keyChain,
+        .ops = 1,
+        .out = .f32,
+        .replaces = if (paged) null else if (bf16) "tf_fa_paged_bf16" else "tf_fa_paged_f16",
+        .tiles = .{ .{ .size = 64 }, .{}, .{ .size = page, .min = 0 } },
+        .served = &served,
+        .sums = .{ .why = attn.why },
+        .known = .{ .gate = .nonfinite, .finding = finding },
+    };
+}
+
+/// shape.rows prompt rows of shape.n query heads at positions shape.k on, each over every key up to its own.
+pub const Prefill = struct {
+    rig: rig_.Rig,
+    l: *const hip.Launcher,
+    k: kit.Kind,
+    /// tf_fa_paged_* over a permuted page table; else tf_fa_wide_* over flat keys, the paged tile its baseline.
+    paged: bool,
+    arena: std.heap.ArenaAllocator,
+    r: ?attn.Ref = null,
+    shape: c.Shape = undefined,
+    geo: Geometry = undefined,
+    q: []f64 = &.{}, // [heads][rows][d]
+    kc: []f64 = &.{}, // [kv][span][d]
+    vc: []f64 = &.{},
+    q32: u64 = 0,
+    flat: [2]u64 = @splat(0), // k and v as [kv][span][d]
+    pool: [2]u64 = @splat(0), // k and v as decode/pages.hpp lays a pool out
+    table: u64 = 0,
+    pages: usize = 0,
+    out: u64 = 0,
+    base: u64 = 0, // the paged tile's output on the flat tile's inputs
+
+    pub fn init(
+        a: std.mem.Allocator,
+        rt: *const hip.Runtime,
+        s: hip.Stream,
+        l: *const hip.Launcher,
+        paged: bool,
+    ) Prefill {
+        return .{ .rig = .{ .a = a, .r = rt, .s = s }, .l = l, .k = kit.act(), .paged = paged, .arena = .init(a) };
+    }
+
+    pub fn deinit(self: *Prefill) void {
+        if (self.r) |*r| r.deinit();
+        self.arena.deinit();
+        self.rig.deinit();
+    }
+
+    pub fn kernel(self: *Prefill) g.Kernel {
+        return .{ .spec = spec(self.k, self.paged), .ctx = self, .vt = &vt };
+    }
+
+    const vt = g.hooks(Prefill);
+
+    fn span(self: *const Prefill) usize {
+        return self.shape.k + self.shape.rows;
+    }
+
+    /// Element e of kv head kh's key `key` in kc and vc.
+    fn at(self: *const Prefill, kh: usize, key: usize, e: usize) usize {
+        return (kh * self.span() + key) * self.geo.d + e;
+    }
+
+    /// Row r's position.
+    fn pos(self: *const Prefill, r: usize) usize {
+        return self.shape.k + r;
+    }
+
+    pub fn prepare(self: *Prefill, shape: c.Shape, fill: g.Fill, seed: u64) !bool {
+        const geo = for (geometries) |x| {
+            if (x.heads == shape.n) break x;
+        } else return false;
+        if (shape.rows == 0 or (!self.paged and !self.l.wide)) return false;
+        self.rig.reset();
+        _ = self.arena.reset(.retain_capacity);
+        self.shape = shape;
+        self.geo = geo;
+        const a = self.arena.allocator();
+        var prng = std.Random.DefaultPrng.init(seed);
+        const rng = prng.random();
+        const k = self.k;
+        const d = geo.d;
+        const keys = self.span();
+        self.q = try a.alloc(f64, geo.heads * shape.rows * d);
+        self.kc = try a.alloc(f64, geo.kv * keys * d);
+        self.vc = try a.alloc(f64, self.kc.len);
+        // q exact in the cache type, so the tile's rounding of it is exact
+        for ([_][]f64{ self.q, self.kc, self.vc }) |x| for (x) |*v| {
+            v.* = kit.round(k, rng.floatNorm(f64));
+        };
+        if (self.r) |*r| r.deinit();
+        self.r = null;
+        self.r = try attn.Ref.init(self.rig.a, d, keys);
+        switch (fill) {
+            .random => {},
+            .nonfinite => self.plant(),
+            .stress => try self.stress(rng),
+        }
+        self.q32 = try self.rig.input(try kit.f32Bytes(a, self.q));
+        // the pages under a permuted table, slots past the span NaN so a read of one shows
+        self.pages = (keys + page - 1) / page;
+        const n = self.pages;
+        const table = try a.alloc(u32, n);
+        for (table, 0..) |*p, i| p.* = @intCast((i * 5 + 3) % n);
+        if (n % 5 == 0) for (table, 0..) |*p, i| {
+            p.* = @intCast(n - 1 - i);
+        };
+        for ([_][]const f64{ self.kc, self.vc }, &self.pool) |src, *dev| {
+            const bytes = try a.alloc(u8, geo.kv * n * page * d * 2);
+            for (0..bytes.len / 2) |i| kit.put(k, bytes, i, std.math.nan(f64));
+            for (0..geo.kv) |h| for (0..keys) |p| for (0..d) |e| {
+                kit.put(k, bytes, ((h * n + table[p / page]) * page + p % page) * d + e, src[self.at(h, p, e)]);
+            };
+            dev.* = try self.rig.input(bytes);
+        }
+        self.table = try self.rig.input(std.mem.sliceAsBytes(table));
+        if (!self.paged) for ([_][]const f64{ self.kc, self.vc }, &self.flat) |src, *dev| {
+            const bytes = try a.alloc(u8, src.len * 2);
+            for (src, 0..) |v, i| kit.put(k, bytes, i, v);
+            dev.* = try self.rig.input(bytes);
+        };
+        const len = self.count() * 4;
+        self.out = try self.rig.output(len);
+        if (!self.paged) self.base = try self.rig.scratch(len);
+        return true;
+    }
+
+    /// NaN query; +inf and -inf values, both in one column; NaN key; +inf at a row's first key and at tiny weights.
+    fn plant(self: *Prefill) void {
+        const n = self.shape.rows;
+        const hk = self.geo.kv;
+        const inf = std.math.inf(f64);
+        self.q[3] = std.math.nan(f64); // row 0, head 0
+        self.vc[self.at(0, self.pos(0), 7)] = inf;
+        if (n > 1) self.vc[self.at(1 % hk, self.pos(1), 9)] = -inf;
+        if (n > 2) {
+            self.vc[self.at(2 % hk, self.pos(2), 11)] = inf;
+            self.vc[self.at(2 % hk, self.pos(2) - 1, 11)] = -inf;
+        }
+        if (n > 3) self.kc[self.at(3 % hk, self.pos(3), 0)] = std.math.nan(f64);
+        if (n > 4) self.vc[self.at(4 % hk, 0, 13)] = inf;
+        // +inf at float64 weight 0 (row 5's own key: NaN), and under fp32's least (inf)
+        if (n > 6) {
+            if (hk > 2) self.reach(5, 1);
+            self.under(n - 1, 2 % hk);
+        }
+    }
+
+    /// Query row r of kv head kh's first head.
+    fn query(self: *Prefill, r: usize, kh: usize) []f64 {
+        const head = kh * (self.geo.heads / self.geo.kv);
+        return self.q[(head * self.shape.rows + r) * self.geo.d ..][0..self.geo.d];
+    }
+
+    fn keyRow(self: *Prefill, kh: usize, key: usize) []f64 {
+        return self.kc[self.at(kh, key, 0)..][0..self.geo.d];
+    }
+
+    /// Row r, key head kh's first head: query all 8, key p - 1 all 8, key p all -8 with +inf at element 17.
+    fn reach(self: *Prefill, r: usize, kh: usize) void {
+        const p = self.pos(r);
+        @memset(self.query(r, kh), 8);
+        @memset(self.keyRow(kh, p - 1), 8);
+        @memset(self.keyRow(kh, p), -8);
+        self.vc[self.at(kh, p, 17)] = std.math.inf(f64);
+    }
+
+    /// Row r, key head kh's first head: query all 2, key p - 1 all 8, its other keys all -8 with +inf at element 21.
+    fn under(self: *Prefill, r: usize, kh: usize) void {
+        const p = self.pos(r);
+        @memset(self.query(r, kh), 2);
+        for (0..p + 1) |key| {
+            @memset(self.keyRow(kh, key), @as(f64, if (key == p - 1) 8 else -8));
+            if (key != p - 1) self.vc[self.at(kh, key, 21)] = std.math.inf(f64);
+        }
+    }
+
+    /// Gate 4: every key alike (each p is 1); values mirrored over the keys every row sees, 0 elsewhere.
+    fn stress(self: *Prefill, rng: std.Random) !void {
+        const a = self.arena.allocator();
+        const d = self.geo.d;
+        for (0..self.geo.kv) |kh| {
+            const key = self.keyRow(kh, 0);
+            for (1..self.span()) |p| @memcpy(self.keyRow(kh, p), key);
+        }
+        @memset(self.vc, 0);
+        const col = try a.alloc(f64, self.pos(0) + 1);
+        for (0..self.geo.kv) |kh| for (0..d) |e| {
+            attn.mirrored(rng, col);
+            for (col, 0..) |y, key| self.vc[self.at(kh, key, e)] = kit.round(self.k, y);
+        };
+    }
+
+    /// pagedCausal's cache kinds are 0 fp16 and 1 bf16, not kit.Kind's.
+    fn cacheKind(self: *const Prefill) c_int {
+        return if (self.k == .bf16) 1 else 0;
+    }
+
+    fn scale(self: *const Prefill) f32 {
+        return 1 / @sqrt(@as(f32, @floatFromInt(self.geo.d)));
+    }
+
+    fn launchPaged(self: *Prefill, out: u64) !void {
+        const s = self.shape;
+        const geo = self.geo;
+        try self.l.pagedCausal(
+            self.q32,
+            self.pool[0],
+            self.pool[1],
+            self.table,
+            out,
+            s.rows,
+            self.span(),
+            geo.heads,
+            geo.kv,
+            geo.d,
+            self.scale(),
+            s.k,
+            self.pages,
+            self.cacheKind(),
+            self.rig.s.handle,
+        );
+    }
+
+    pub fn run(self: *Prefill) !void {
+        if (self.paged) {
+            try self.launchPaged(self.out);
+            return self.rig.s.synchronize();
+        }
+        const s = self.shape;
+        const geo = self.geo;
+        const ss: c_longlong = @intCast(geo.d);
+        const sh: c_longlong = @intCast(self.span() * geo.d);
+        // tf_causal takes the 64-row tile for a 16-bit cache of whole 64-wide heads
+        try self.l.tf_causal(
+            @ptrFromInt(self.q32),
+            @ptrFromInt(self.flat[0]),
+            @ptrFromInt(self.flat[1]),
+            @ptrFromInt(self.out),
+            1,
+            @intCast(s.rows),
+            @intCast(self.span()),
+            @intCast(geo.heads),
+            @intCast(geo.kv),
+            @intCast(geo.d),
+            self.scale(),
+            @intCast(s.k),
+            0,
+            sh,
+            ss,
+            0,
+            sh,
+            ss,
+            self.cacheKind(),
+            null,
+            null,
+            null,
+            self.rig.s.handle,
+            null,
+        );
+        try self.rig.s.synchronize();
+    }
+
+    /// The paged tile on the same keys, laid out in pages.
+    pub fn replaced(self: *Prefill) ![]const u8 {
+        try self.launchPaged(self.base);
+        return self.rig.bytes(self.base, self.count() * 4);
+    }
+
+    pub fn count(self: *Prefill) usize {
+        return self.geo.heads * self.shape.rows * self.geo.d;
+    }
+
+    /// Output head i (head i / rows, row i % rows): its query and keys 0 to its position, at their kv head.
+    pub fn load(self: *Prefill, i: usize, hd: *attn.Head) !void {
+        const d = self.geo.d;
+        const r = i % self.shape.rows;
+        const kh = (i / self.shape.rows) / (self.geo.heads / self.geo.kv);
+        hd.n = self.pos(r) + 1;
+        hd.q = self.q[i * d ..][0..d];
+        hd.k = self.kc[self.at(kh, 0, 0)..][0 .. hd.n * d];
+        hd.v = self.vc[self.at(kh, 0, 0)..][0 .. hd.n * d];
+    }
+
+    fn heads(self: *Prefill) usize {
+        return self.shape.rows * self.geo.heads;
+    }
+
+    pub fn check(self: *Prefill, gate: g.Gate) !?[]const u8 {
+        return self.r.?.check(self, self.heads(), self.scale(), try self.rig.outputBytes(), gate);
+    }
+
+    pub fn reference(self: *Prefill, want: []f64, mag: []f64) !void {
+        return self.r.?.reference(self, self.heads(), self.scale(), want, mag);
+    }
+
+    pub fn simulate(self: *Prefill, acc: c.Accum, order: g.Order, out: []f64) !void {
+        return self.r.?.simulate(self, self.heads(), self.scale(), acc, order, out);
+    }
+};

From 7217ce7ce14b360d912b36fa4992f7979705ad28 Mon Sep 17 00:00:00 2001
From: ashhart <71890694+ashhart@users.noreply.github.com>
Date: Sun, 1 Nov 2026 03:29:46 +0000
Subject: [PATCH 7/7] changelog: #1788 under Unreleased

---
 CHANGELOG.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 87b9e1029c..bbde9d2318 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 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.