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`: checkpoints do not learn on 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. 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/kernels/hip/attention/prefill.hip b/zig/kernels/hip/attention/prefill.hip index 87e39f57d4..1f8a3dc977 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. @@ -51,7 +52,46 @@ __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 +__device__ __forceinline__ word4 load16(const Cache* base, long long stride, int row, int rows, int col) { + return *reinterpret_cast(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 +__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 +__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 __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, @@ -59,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; 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(k0) * k_ss; - const Cache* vtile = PAGED ? vbase + page_row(table[k0 / kPageTokens], count, kv, 0, d) : vbase + static_cast(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(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(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 = [&](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(ktile, kstride, key, rows, 8 * (c0 + c)); + } + }; + // a value chunk: key pair `pair`, dims 8g .. 8g + 7 of both keys + const auto load_v = [&](const Cache* vtile, int rows, int c0) __attribute__((always_inline)) { + const int pair = tid / kKChunks, g = tid % kKChunks; + va = load16(vtile, vstride, 2 * pair, rows, 8 * (c0 + g)); + vb = load16(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()(ktile_at(0, page), span, 0); + const auto tile = [&](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(ktile + static_cast(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()(ktile, rows, c0 + kKChunks); + else load_v.template operator()(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,35 +241,30 @@ __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(p[0]), static_cast(p[1])}; - sH[(2 * sk + 1) * kPLd + row] = pair_t{static_cast(p[2]), static_cast(p[3])}; - if (sk == 0) sA[row] = alpha; + ph[i][0] = pair_t{static_cast(p[0]), static_cast(p[1])}; + ph[i][1] = pair_t{static_cast(p[2]), static_cast(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(vtile + static_cast(2 * pair) * vstride + 8 * (c0 + g)) - : uint4{0, 0, 0, 0}; - const uint4 b = key + 1 < span - ? *reinterpret_cast(vtile + static_cast(2 * pair + 1) * vstride + 8 * (c0 + g)) - : uint4{0, 0, 0, 0}; - __syncthreads(); - const elem* ha = reinterpret_cast(&a); - const elem* hb = reinterpret_cast(&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()(vtile, rows, (at + 1) * kKChunks); + else if (more) load_k.template operator()(ktile_at(k0 + kKeys, next), rows - kKeys, 0); float pv[4][4]; #pragma unroll for (int i = 0; i < 4; ++i) @@ -211,38 +272,40 @@ __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]); } } - } - 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()(k0); + for (; k0 < max_vis; k0 += kKeys) tile.template operator()(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(row) * d + 64 * at + 4 * pc + m] = o[i][4 * at + m] * inv; + orows[static_cast(row) * d + 64 * at + 4 * sk + m] = o[i][4 * at + m] * inv; } } } 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/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); +} 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); + } +};