diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index b70dfe431f..cc6ce7e31f 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -27,6 +27,7 @@ const kernels = [_]Kernel{ .{ .name = "fp8_experts", .flags = &.{"-O3"} }, // fp8/experts.cu's device code, tensorfold_fp8_experts_v6 .{ .name = "exlr8", .flags = &.{ "-O3", "--fmad=false" } }, // ours: EXLR8 lane projections, the Metal kernels' order .{ .name = "exlr8_dense", .flags = &.{ "-O3", "--fmad=false" } }, // ours: EXLR8 dense matrices by row, the Metal order + .{ .name = "exl3", .flags = &.{ "-O3", "--fmad=false" } }, // ours: EXL3 trellis projections behind qlinear .{ .name = "nvfp4_experts", .flags = &.{"-O3"} }, // nvfp4/experts.cu's device code, the NVFP4 expert kernels .{ .name = "lane_gemv", .flags = &.{"-O3"} }, // ours: qmm_group's arithmetic, a column tile's K slices in one CTA .{ .name = "sample", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // ours: the Metal engine's keyed draws @@ -130,6 +131,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu runner.addImport("core", mods.core); runner.addImport("lanes", mods.lanes); runner.addImport("nemotron", mods.nemotron); + runner.addImport("core", mods.core); const exlr8 = exlr8_build.coreAlias(b, target, optimize, mods.core); runner.addImport("exlr8", exlr8); const exlr8_tests = exlr8_build.testModules(b, target, optimize, exlr8); diff --git a/zig/kernels/cuda/exl3.cu b/zig/kernels/cuda/exl3.cu new file mode 100644 index 0000000000..fdd523a323 --- /dev/null +++ b/zig/kernels/cuda/exl3.cu @@ -0,0 +1,268 @@ +// EXL3 projections on CUDA: trellis tiles as stored ([K/16, N/16, 16 bits] int16), core/exl3_format.zig's codebooks. +#include "exl3.cuh" + +struct TfExl3Rows { + unsigned rows, x_stride, y_stride, y_f32; +}; +struct TfExl3Matrix { + unsigned k, n, bits, tiles_n, split; +}; + +// A block's tail from zs (rows x 128 sums): a split's parts summed in order by the last part, transform, scale. +__device__ __forceinline__ void tf_exl3_tail(unsigned rows, const __half* svh, float out_scale, const TfExl3Rows a, + const TfExl3Matrix m, uint8_t* Y, float* ws, unsigned* tickets, unsigned ob, + unsigned sg, unsigned lane, unsigned tid, float* zs) { + __shared__ unsigned last; + if (m.split > 1) { + float* part = ws + ((uint64_t)ob * m.split + blockIdx.y) * (16 * 128); + for (unsigned i = tid; i < rows * 128; i += 256) part[i] = zs[i]; + __threadfence(); + __syncthreads(); + if (tid == 0) last = atomicAdd(tickets + ob, 1u) == m.split - 1; + __syncthreads(); + if (!last) return; + __threadfence(); + const float* parts = ws + (uint64_t)ob * m.split * (16 * 128); + for (unsigned i = tid; i < rows * 128; i += 256) { + float v = __ldcg(parts + i); + for (unsigned p = 1; p < m.split; ++p) v = __fadd_rn(v, __ldcg(parts + p * (16 * 128) + i)); + zs[i] = v; + } + if (tid == 0) tickets[ob] = 0; + __syncthreads(); + } + for (unsigned i = sg; i < rows; i += 8) { + float v[4]; + for (int q = 0; q < 4; ++q) v[q] = zs[i * 128 + lane + 32 * q]; + tf_exl3_fwht(v, lane); + for (int q = 0; q < 4; ++q) { + const unsigned c = ob * 128 + lane + 32 * q; + // svh times a power of two is exact; one rounding remains + const float scale = __fmul_rn(__half2float(svh[c]), out_scale); + tf_exl3_store(Y, (uint64_t)i * a.y_stride + c, __fmul_rn(v[q], scale), a.y_f32); + } + } +} + +// Qualification: every state's weight under codebook c, for the host to compare with exl3_format.codebookValue. +extern "C" __global__ void tf_exl3_codebook(unsigned c, __half* out) { + const unsigned i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= 65536) return; + out[i] = c == 2 ? tf_exl3_value<2>(i) : c == 1 ? tf_exl3_value<1>(i) : tf_exl3_value<0>(i); +} + +template +__device__ __forceinline__ void tf_exl3_unpack_tile(const uint32_t* tile, unsigned c, unsigned lane, __half (&w)[8]) { + if (c == 2) tf_exl3_decode(tile, lane, w); + else if (c == 1) tf_exl3_decode(tile, lane, w); + else tf_exl3_decode(tile, lane, w); +} + +// Qualification: W_q [K, N] fp16 in the rotated domain, a warp per tile, as exl3_format.unpack places it. +extern "C" __global__ void __launch_bounds__(32) tf_exl3_unpack(const uint8_t* T, TfExl3Matrix m, unsigned c, __half* out) { + const unsigned lane = threadIdx.x, tn = blockIdx.x, tk = blockIdx.y; + const uint32_t* tile = reinterpret_cast(T + ((uint64_t)tk * m.tiles_n + tn) * 32 * m.bits); + __half w[8]; + switch (m.bits) { + case 1: tf_exl3_unpack_tile<1>(tile, c, lane, w); break; + case 2: tf_exl3_unpack_tile<2>(tile, c, lane, w); break; + case 3: tf_exl3_unpack_tile<3>(tile, c, lane, w); break; + case 4: tf_exl3_unpack_tile<4>(tile, c, lane, w); break; + case 5: tf_exl3_unpack_tile<5>(tile, c, lane, w); break; + case 6: tf_exl3_unpack_tile<6>(tile, c, lane, w); break; + case 7: tf_exl3_unpack_tile<7>(tile, c, lane, w); break; + default: tf_exl3_unpack_tile<8>(tile, c, lane, w); break; + } + for (unsigned j = 0; j < 8; ++j) { + const unsigned row = 16 * tk + 2 * (lane % 4) + (j & 1) + 8 * ((j >> 1) & 1), col = 16 * tn + lane / 4 + 8 * (j >> 2); + out[(uint64_t)row * m.n + col] = w[j]; + } +} + +// Prompt rows (docs/recipes/exl3.md, Prompts): (x * suh) H times 2^-4 rounded once to fp16; 2^-3 goes on the output. +extern "C" __global__ void __launch_bounds__(32) tf_exl3_rotate_h(const __nv_bfloat16* X, unsigned x_stride, const __half* suh, + __half* R, unsigned r_stride) { + const unsigned lane = threadIdx.x, row = blockIdx.y, base = blockIdx.x * 128; + float v[4]; + for (int q = 0; q < 4; ++q) { + const unsigned c = base + lane + 32 * q; + v[q] = __fmul_rn(__bfloat162float(X[(uint64_t)row * x_stride + c]), __half2float(suh[c])); + } + tf_exl3_fwht(v, lane); + for (int q = 0; q < 4; ++q) R[(uint64_t)row * r_stride + base + lane + 32 * q] = __float2half_rn(__fmul_rn(v[q], 0.0625f)); +} + +// 64 prompt rows x 128 columns a block: each tile decodes once into mma.m16n8k16 B fragments (their value order). +template +__device__ __forceinline__ void tf_exl3_prompt_job(const __half* Xh, unsigned rows, const uint8_t* T, const __half* svh, + const TfExl3Matrix m, uint8_t* Y, unsigned y_f32, float* buf) { + constexpr unsigned TB = 32 * K; + const unsigned tid = threadIdx.x, sg = tid / 32, lane = tid % 32, g8 = lane / 4, t4 = lane % 4, ob = blockIdx.x; + const unsigned row0 = blockIdx.y * 64, kb = m.k / 128; + __half* xs = reinterpret_cast<__half*>(buf); + const uint64_t row_bytes = (uint64_t)m.tiles_n * TB; + const uint8_t* col = T + (uint64_t)(8 * ob + sg) * TB; + float acc[4][2][4]; +#pragma unroll + for (int i = 0; i < 4; ++i) + for (int n = 0; n < 2; ++n) acc[i][n][0] = acc[i][n][1] = acc[i][n][2] = acc[i][n][3] = 0.0f; + uint4 next = tf_exl3_words(reinterpret_cast(col), lane); + for (unsigned ib = 0; ib < kb; ++ib) { + // the block's 64 rows of this 128-input block, fp16, zeros past the last row + for (unsigned i = tid; i < 64 * 16; i += 256) { + const unsigned r = i / 16, c = (i % 16) * 8; + const uint4 v = row0 + r < rows ? *reinterpret_cast(Xh + (uint64_t)(row0 + r) * m.k + ib * 128 + c) : make_uint4(0, 0, 0, 0); + *reinterpret_cast(xs + r * TF_EXL3_XS + c) = v; + } + __syncthreads(); + for (unsigned r = 0; r < 8; ++r) { + const uint4 now = next; + next = tf_exl3_words(reinterpret_cast(col + (uint64_t)(r < 7 ? 8 * ib + r + 1 : (ib + 1 < kb ? 8 * ib + 8 : 8 * ib + 7)) * row_bytes), lane); + __half w[8]; + tf_exl3_unpack(now, lane, w); + const uint32_t b[2][2] = {{tf_exl3_pair(w[0], w[1]), tf_exl3_pair(w[2], w[3])}, {tf_exl3_pair(w[4], w[5]), tf_exl3_pair(w[6], w[7])}}; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const __half* x = xs + (16 * i + g8) * TF_EXL3_XS + 16 * r + 2 * t4; + const uint32_t a[4] = {*reinterpret_cast(x), *reinterpret_cast(x + 8 * TF_EXL3_XS), + *reinterpret_cast(x + 8), *reinterpret_cast(x + 8 * TF_EXL3_XS + 8)}; + tf_exl3_mma(acc[i][0], a, b[0][0], b[0][1]); + tf_exl3_mma(acc[i][1], a, b[1][0], b[1][1]); + } + } + __syncthreads(); + } + // the sums by row into shared memory (over the rows' halves), then a warp a row: transform, svh, 2^-3 +#pragma unroll + for (int i = 0; i < 4; ++i) + for (int n = 0; n < 2; ++n) + for (int h = 0; h < 2; ++h) { + const unsigned c = 16 * sg + 8 * n + 2 * t4; + buf[(16 * i + g8 + 8 * h) * 128 + c] = acc[i][n][2 * h]; + buf[(16 * i + g8 + 8 * h) * 128 + c + 1] = acc[i][n][2 * h + 1]; + } + __syncthreads(); + for (unsigned i = sg; i < 64 && row0 + i < rows; i += 8) { + float v[4]; + for (int q = 0; q < 4; ++q) v[q] = buf[i * 128 + lane + 32 * q]; + tf_exl3_fwht(v, lane); + for (int q = 0; q < 4; ++q) { + const unsigned c = ob * 128 + lane + 32 * q; + tf_exl3_store(Y, (uint64_t)(row0 + i) * m.n + c, __fmul_rn(v[q], __fmul_rn(__half2float(svh[c]), 0.125f)), y_f32); + } + } +} + +// One prompt kernel a codebook; blocks (128-column block, 64-row block). +#define TF_EXL3_PROMPT(C) \ + extern "C" __global__ void __launch_bounds__(256) tf_exl3_prompt_##C(const __half* Xh, unsigned rows, \ + const uint8_t* T, const __half* svh, \ + const TfExl3Matrix m, uint8_t* Y, \ + unsigned y_f32) { \ + __shared__ __align__(16) float buf[64 * 128]; \ + switch (m.bits) { \ + case 1: tf_exl3_prompt_job<1, C>(Xh, rows, T, svh, m, Y, y_f32, buf); break; \ + case 2: tf_exl3_prompt_job<2, C>(Xh, rows, T, svh, m, Y, y_f32, buf); break; \ + case 3: tf_exl3_prompt_job<3, C>(Xh, rows, T, svh, m, Y, y_f32, buf); break; \ + case 4: tf_exl3_prompt_job<4, C>(Xh, rows, T, svh, m, Y, y_f32, buf); break; \ + case 5: tf_exl3_prompt_job<5, C>(Xh, rows, T, svh, m, Y, y_f32, buf); break; \ + case 6: tf_exl3_prompt_job<6, C>(Xh, rows, T, svh, m, Y, y_f32, buf); break; \ + case 7: tf_exl3_prompt_job<7, C>(Xh, rows, T, svh, m, Y, y_f32, buf); break; \ + case 8: tf_exl3_prompt_job<8, C>(Xh, rows, T, svh, m, Y, y_f32, buf); break; \ + default: break; \ + } \ + } +TF_EXL3_PROMPT(0) +TF_EXL3_PROMPT(1) +TF_EXL3_PROMPT(2) + + +// Decode rows for the tensor units: tf_exl3_rotate_h's values, each 128-block in fragment order (t4, r, 4 halves). +extern "C" __global__ void __launch_bounds__(32) tf_exl3_rotate_f(const __nv_bfloat16* X, unsigned x_stride, const __half* suh, + __half* R, unsigned r_stride) { + const unsigned lane = threadIdx.x, row = blockIdx.y, base = blockIdx.x * 128; + float v[4]; + for (int q = 0; q < 4; ++q) { + const unsigned c = base + lane + 32 * q; + v[q] = __fmul_rn(__bfloat162float(X[(uint64_t)row * x_stride + c]), __half2float(suh[c])); + } + tf_exl3_fwht(v, lane); + for (int q = 0; q < 4; ++q) { + const unsigned j = lane + 32 * q, r = j / 16, c = j % 16; + const unsigned at = ((c % 8) / 2) * 32 + r * 4 + (c % 2) + 2 * (c / 8); + R[(uint64_t)row * r_stride + base + at] = __float2half_rn(__fmul_rn(v[q], 0.0625f)); + } +} + +// Decode rows on the tensor units, 1-16 a launch: tf_exl3_rotate_f's rows as A fragments, the same at every count. +template +__device__ __forceinline__ void tf_exl3_tc_job(const __half* Xh, const uint8_t* T, const __half* svh, const TfExl3Rows a, + const TfExl3Matrix m, uint8_t* Y, float* ws, unsigned* tickets, float* zs) { + constexpr unsigned TB = 32 * K; + const unsigned tid = threadIdx.x, sg = tid / 32, lane = tid % 32, g8 = lane / 4, t4 = lane % 4, ob = blockIdx.x; + const unsigned rows = a.rows < 16 ? a.rows : 16, kb = m.k / 128, per = (kb + m.split - 1) / m.split; + const unsigned b0 = blockIdx.y * per, b1 = b0 + per < kb ? b0 + per : kb; + const uint64_t row_bytes = (uint64_t)m.tiles_n * TB; + const uint8_t* col = T + (uint64_t)(8 * ob + sg) * TB; + const bool lo = g8 < rows, hi = g8 + 8 < rows; + // a lane's A words of a 128-block, rows g8 and g8 + 8 (r 0-7: columns 2t4, 2t4 + 8), loaded a block ahead + const uint4* x0 = reinterpret_cast(Xh + (uint64_t)(lo ? g8 : 0) * m.k + 32 * t4); + const uint4* x1 = reinterpret_cast(Xh + (uint64_t)(hi ? g8 + 8 : 0) * m.k + 32 * t4); + float acc[2][4] = {{0.0f, 0.0f, 0.0f, 0.0f}, {0.0f, 0.0f, 0.0f, 0.0f}}; + uint4 next = tf_exl3_words(reinterpret_cast(col + (uint64_t)(8 * b0) * row_bytes), lane); + uint4 an[2][4]; + for (int q = 0; q < 4; ++q) { + an[0][q] = lo ? __ldg(x0 + b0 * 16 + q) : make_uint4(0, 0, 0, 0); + an[1][q] = hi ? __ldg(x1 + b0 * 16 + q) : make_uint4(0, 0, 0, 0); + } + for (unsigned ib = b0; ib < b1; ++ib) { + uint32_t au[2][16]; + for (int q = 0; q < 4; ++q) { + au[0][4 * q] = an[0][q].x, au[0][4 * q + 1] = an[0][q].y, au[0][4 * q + 2] = an[0][q].z, au[0][4 * q + 3] = an[0][q].w; + au[1][4 * q] = an[1][q].x, au[1][4 * q + 1] = an[1][q].y, au[1][4 * q + 2] = an[1][q].z, au[1][4 * q + 3] = an[1][q].w; + } + if (ib + 1 < b1) + for (int q = 0; q < 4; ++q) { + an[0][q] = lo ? __ldg(x0 + (ib + 1) * 16 + q) : make_uint4(0, 0, 0, 0); + an[1][q] = hi ? __ldg(x1 + (ib + 1) * 16 + q) : make_uint4(0, 0, 0, 0); + } +#pragma unroll + for (unsigned r = 0; r < 8; ++r) { + const uint4 now = next; + next = tf_exl3_words(reinterpret_cast(col + (uint64_t)(r < 7 ? 8 * ib + r + 1 : (ib + 1 < b1 ? 8 * ib + 8 : 8 * ib + 7)) * row_bytes), lane); + const uint32_t af[4] = {au[0][2 * r], au[1][2 * r], au[0][2 * r + 1], au[1][2 * r + 1]}; + __half w[8]; + tf_exl3_unpack(now, lane, w); + tf_exl3_mma(acc[0], af, tf_exl3_pair(w[0], w[1]), tf_exl3_pair(w[2], w[3])); + tf_exl3_mma(acc[1], af, tf_exl3_pair(w[4], w[5]), tf_exl3_pair(w[6], w[7])); + } + } +#pragma unroll + for (int n = 0; n < 2; ++n) + for (int h = 0; h < 2; ++h) { + const unsigned c = 16 * sg + 8 * n + 2 * t4, row = g8 + 8 * h; + if (row < rows) { + zs[row * 128 + c] = acc[n][2 * h]; + zs[row * 128 + c + 1] = acc[n][2 * h + 1]; + } + } + __syncthreads(); + tf_exl3_tail(rows, svh, 0.125f, a, m, Y, ws, tickets, ob, sg, lane, tid, zs); +} + +#define TF_EXL3_TC_CASE(B, C) \ + case B: tf_exl3_tc_job(Xh, T, svh, a, m, Y, ws, tickets, zs); break; +#define TF_EXL3_TC(C) \ + extern "C" __global__ void __launch_bounds__(256) tf_exl3_tc_##C(const __half* Xh, const uint8_t* T, const __half* svh, \ + const TfExl3Rows a, const TfExl3Matrix m, uint8_t* Y, \ + float* ws, unsigned* tickets) { \ + __shared__ __align__(16) float zs[16 * 128]; \ + switch (m.bits) { \ + TF_EXL3_TC_CASE(1, C) TF_EXL3_TC_CASE(2, C) TF_EXL3_TC_CASE(3, C) TF_EXL3_TC_CASE(4, C) \ + TF_EXL3_TC_CASE(5, C) TF_EXL3_TC_CASE(6, C) TF_EXL3_TC_CASE(7, C) TF_EXL3_TC_CASE(8, C) \ + default: break; \ + } \ + } +TF_EXL3_TC(0) +TF_EXL3_TC(1) +TF_EXL3_TC(2) diff --git a/zig/kernels/cuda/exl3.cuh b/zig/kernels/cuda/exl3.cuh new file mode 100644 index 0000000000..b3dd3a6fd7 --- /dev/null +++ b/zig/kernels/cuda/exl3.cuh @@ -0,0 +1,102 @@ +// EXL3 on CUDA, shared by exl3.cu and exl3_experts.cu: codebooks, trellis states, the transform, a tile row's sums. +#pragma once +#include +#include +#include + +#define TF_FULL 0xFFFFFFFFu + +// The unnormalized Sylvester transform of 128 values held as lane + 32 q, in one fixed order of fp32 adds. +__device__ __forceinline__ void tf_exl3_fwht(float (&v)[4], unsigned lane) { + for (unsigned d = 1; d < 32; d <<= 1) { + for (int q = 0; q < 4; ++q) { + const float other = __shfl_xor_sync(TF_FULL, v[q], d); + v[q] = (lane & d) ? __fsub_rn(other, v[q]) : __fadd_rn(v[q], other); + } + } + const float a0 = __fadd_rn(v[0], v[1]), a1 = __fsub_rn(v[0], v[1]); + const float a2 = __fadd_rn(v[2], v[3]), a3 = __fsub_rn(v[2], v[3]); + v[0] = __fadd_rn(a0, a2); + v[1] = __fadd_rn(a1, a3); + v[2] = __fsub_rn(a0, a2); + v[3] = __fsub_rn(a1, a3); +} + +// A state's fp16 weight, one rounding: mul1's sum is an integer times 2^-18 under 2^20, exact in fp32 before it. +template +__device__ __forceinline__ __half tf_exl3_value(unsigned state) { + if (C == 2) { + const float h = (float)__dp4a(state * 0x83DCD12Du, 0x01010101u, 1024u); + return __float2half_rn(__fmaf_rn(h, __half2float(__ushort_as_half(0x1EEE)), __half2float(__ushort_as_half(0xC931)))); + } + const unsigned x = C == 1 ? state * 0xCBAC1FEDu : state * 89226354u + 64248484u; + const unsigned m = (x & 0x8FFF8FFFu) ^ 0x3B603B60u; + // the two halves' exponents differ by at most a few binades, so the fp32 sum is exact + const float s = __fadd_rn(__half2float(__ushort_as_half((unsigned short)(m & 0xFFFFu))), __half2float(__ushort_as_half((unsigned short)(m >> 16)))); + return __float2half_rn(s); +} + +// Lane l's stream words for its eight values 8 l .. 8 l + 7: the four u32 words its states span, MSB first. +template +__device__ __forceinline__ uint4 tf_exl3_words(const uint32_t* tile, unsigned lane) { + constexpr unsigned ring = 256 * K, words = 8 * K; + const unsigned i0 = ((8 * lane * K + K + ring - 16) % ring) / 32; + return make_uint4(tile[i0], tile[(i0 + 1) % words], tile[(i0 + 2) % words], tile[(i0 + 3) % words]); +} + +// The eight states out of those words: the 16 stream bits that end at each value. +template +__device__ __forceinline__ void tf_exl3_states(uint4 t, unsigned lane, unsigned (&st)[8]) { + const unsigned off = (8 * lane * K + K + 256 * K - 16) % 32; + // one variable shift lines the window up (a, b, c: its first 96 bits); every state after it is a constant shift + const unsigned a = __funnelshift_l(t.y, t.x, off), b = __funnelshift_l(t.z, t.y, off), c = __funnelshift_l(t.w, t.z, off); +#pragma unroll + for (int n = 0; n < 8; ++n) { + const unsigned p = n * K; + st[n] = __funnelshift_l(p < 32 ? b : c, p < 32 ? a : b, p % 32) >> 16; + } +} + +template +__device__ __forceinline__ void tf_exl3_unpack(uint4 t, unsigned lane, __half (&w)[8]) { + unsigned st[8]; + tf_exl3_states(t, lane, st); + if constexpr (C == 2) { + // mul1 two states at a time: the sums (<= 2044) are exact in fp16, and one fused fp16 multiply-add rounds once + const __half2 scale = __half2half2(__ushort_as_half(0x1EEE)), bias = __half2half2(__ushort_as_half(0xC931)); +#pragma unroll + for (int n = 0; n < 8; n += 2) { + const __half2 h = __halves2half2(__uint2half_rn(__dp4a(st[n] * 0x83DCD12Du, 0x01010101u, 1024u)), + __uint2half_rn(__dp4a(st[n + 1] * 0x83DCD12Du, 0x01010101u, 1024u))); + const __half2 v = __hfma2(h, scale, bias); + w[n] = __low2half(v); + w[n + 1] = __high2half(v); + } + } else { +#pragma unroll + for (int n = 0; n < 8; ++n) w[n] = tf_exl3_value(st[n]); + } +} + +template +__device__ __forceinline__ void tf_exl3_decode(const uint32_t* tile, unsigned lane, __half (&w)[8]) { + tf_exl3_unpack(tf_exl3_words(tile, lane), lane, w); +} + +__device__ __forceinline__ void tf_exl3_store(uint8_t* y, uint64_t at, float v, unsigned f32) { + if (f32) reinterpret_cast(y)[at] = v; + else reinterpret_cast<__nv_bfloat16*>(y)[at] = __float2bfloat16_rn(v); +} + +#define TF_EXL3_XS 136 // shared prompt rows 136 halves apart: a fragment's eight rows hit eight bank groups + +// Prompt rows on the tensor units: fp16 A and B fragments, fp32 sums. +__device__ __forceinline__ void tf_exl3_mma(float (&d)[4], const uint32_t (&a)[4], uint32_t b0, uint32_t b1) { + asm("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} + +__device__ __forceinline__ uint32_t tf_exl3_pair(__half lo, __half hi) { + return (uint32_t)__half_as_ushort(lo) | ((uint32_t)__half_as_ushort(hi) << 16); +} diff --git a/zig/kernels/cuda/prompt16.cu b/zig/kernels/cuda/prompt16.cu index c2c6c2f3ab..01749ba86d 100644 --- a/zig/kernels/cuda/prompt16.cu +++ b/zig/kernels/cuda/prompt16.cu @@ -192,5 +192,6 @@ __global__ void __launch_bounds__(WM * WN * 32) prompt_kernel( // The instance the Zig launcher loads by name (PROMPT_TILE 4 in linear.py). template __global__ void prompt_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, int, int, int, int, int, int); +template __global__ void prompt_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, int, int, int, int, int, int); template __global__ void prompt_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, int, int, int, int, int, int); } diff --git a/zig/kernels/cuda/qmmf.cu b/zig/kernels/cuda/qmmf.cu index 4703273410..57a7c1e582 100644 --- a/zig/kernels/cuda/qmmf.cu +++ b/zig/kernels/cuda/qmmf.cu @@ -308,6 +308,13 @@ template __global__ void qmmf_kernel(c template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); @@ -323,4 +330,5 @@ template __global__ void qmmf_kernel template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void reduce_kernel(const float* __restrict__, void* __restrict__, long long, int, float); +template __global__ void reduce_kernel(const float* __restrict__, void* __restrict__, long long, int, float); } diff --git a/zig/src/cuda/exl3_linear.zig b/zig/src/cuda/exl3_linear.zig new file mode 100644 index 0000000000..c894f33096 --- /dev/null +++ b/zig/src/cuda/exl3_linear.zig @@ -0,0 +1,211 @@ +//! exl3.cu's EXL3 projections: bf16 rows rotated into fp32, then 1-16 rows a launch over the trellis tiles as stored. +const std = @import("std"); +const driver = @import("driver.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const memory = @import("memory.zig"); +const stream_ = @import("stream.zig"); + +const Driver = driver.Driver; +const Function = module.Function; +const Stream = stream_.Stream; + +/// core/exl3_format.zig's Codebook, in the order exl3.cu numbers them. +pub const Codebook = enum(u32) { inst3 = 0, mcg = 1, mul1 = 2 }; + +/// One matrix on the device as stored: trellis int16 [k/16, n/16, 16 bits], suh fp16 [k], svh fp16 [n]. +pub const Matrix = struct { + trellis: u64, + suh: u64, + svh: u64, + k: usize, + n: usize, + bits: u32, + codebook: Codebook, + + /// Whole widths 1-8 over whole 128-blocks; exl3.cu takes no half widths and no partial blocks. + pub fn check(m: Matrix) !void { + if (m.k == 0 or m.n == 0 or m.k % 128 != 0 or m.n % 128 != 0 or m.bits < 1 or m.bits > 8) return error.InvalidExl3Matrix; + } + + pub fn trellisBytes(k: usize, n: usize, bits: u32) usize { + return k / 16 * (n / 16) * 32 * @as(usize, bits); + } + + fn args(m: Matrix) MatrixArgs { + return .{ .k = @intCast(m.k), .n = @intCast(m.n), .bits = m.bits, .tiles_n = @intCast(m.n / 16), .split = @intCast(split(m.k, m.n)) }; + } +}; + +/// exl3.cu's TfExl3Rows and TfExl3Matrix, passed by value. +pub const RowsArgs = extern struct { rows: u32, x_stride: u32, y_stride: u32, y_f32: u32 }; +pub const MatrixArgs = extern struct { k: u32, n: u32, bits: u32, tiles_n: u32, split: u32 }; + +/// Decode rows a launch; longer decode runs go window by window. +pub const window = 16; +/// Prompt rows a chunk: rotated to fp16 together, then 64 rows a block on the tensor units. +pub const prompt_chunk = 256; +/// Blocks a launch aims for: a matrix with fewer 128-column blocks splits its input blocks into parts. +const target_blocks = 96; + +/// Parts of the input blocks, from the shape alone: enough for ~target_blocks blocks, at least 4 input blocks a part. +pub fn split(k: usize, n: usize) usize { + const cols = n / 128; + const kb = k / 128; + if (cols == 0 or cols >= target_blocks) return 1; + const want = @min((target_blocks + cols - 1) / cols, kb / 4); + if (want <= 1) return 1; + const per = (kb + want - 1) / want; + return (kb + per - 1) / per; +} + +/// The kernels and their scratch (fp16 rotated rows up to `k_max` inputs): one stream's projections at a time. +pub const Kernels = struct { + decode_fns: [3]Function, // 1-16 rows on the tensor units, by codebook + prompt_fns: [3]Function, // 64 rows a block on the tensor units, by codebook + rotate_f_fn: Function, // decode rows, each 128-block in fragment order + rotate_h_fn: Function, // prompt rows, plain order + codebook_fn: Function, + unpack_fn: Function, + rotated: memory.DeviceBuffer, // a decode window's or prompt chunk's rotated rows (fp16, prompt_chunk x k_max) + parts: memory.DeviceBuffer, // a split matrix's part totals: 16 rows x 128 columns per (column block, part) + tickets: memory.DeviceBuffer, // per column block, parts finished; the last part resets its count to zero + k_max: usize, + + pub fn resolve(d: *const Driver, mod: module.Module, k_max: usize) !Kernels { + var k: Kernels = .{ .decode_fns = undefined, .prompt_fns = undefined, .rotate_f_fn = try mod.function("tf_exl3_rotate_f"), .rotate_h_fn = try mod.function("tf_exl3_rotate_h"), .codebook_fn = try mod.function("tf_exl3_codebook"), .unpack_fn = try mod.function("tf_exl3_unpack"), .rotated = undefined, .parts = undefined, .tickets = undefined, .k_max = k_max }; + for (&k.decode_fns, &k.prompt_fns, 0..) |*f, *p, c| { + var name: [32]u8 = undefined; + f.* = try mod.function(try std.fmt.bufPrintSentinel(&name, "tf_exl3_tc_{d}", .{c}, 0)); + p.* = try mod.function(try std.fmt.bufPrintSentinel(&name, "tf_exl3_prompt_{d}", .{c}, 0)); + } + k.rotated = try memory.DeviceBuffer.alloc(d, prompt_chunk * k_max * 2); + errdefer k.rotated.free(); + // split * cols stays under target_blocks + cols, so 2 * target_blocks parts always fit + k.parts = try memory.DeviceBuffer.alloc(d, 2 * target_blocks * window * 128 * 4); + errdefer k.parts.free(); + k.tickets = try memory.DeviceBuffer.alloc(d, target_blocks * 4); + errdefer k.tickets.free(); + try d.check(d.api.cuMemsetD32_v2(k.tickets.ptr, 0, target_blocks), "cuMemsetD32"); + return k; + } + + pub fn deinit(k: *Kernels) void { + k.rotated.free(); + k.parts.free(); + k.tickets.free(); + } + + /// Decode rows: bf16 rows `ldx` apart -> out (rows, n), 16 a window on the tensor units, the same at every count. + pub fn rows(k: *const Kernels, s: Stream, x: u64, ldx: usize, m: Matrix, out: u64, count: usize, f32_out: bool) !void { + try k.fits(m, ldx); + const elem: usize = if (f32_out) 4 else 2; + var at: usize = 0; + while (at < count) : (at += window) { + const r = @min(window, count - at); + try k.rotate(k.rotate_f_fn, s, x + at * ldx * 2, ldx, m, r); + var a: launch_.Args = .{}; + a.add(k.rotated.ptr); + a.add(m.trellis); + a.add(m.svh); + a.add(RowsArgs{ .rows = @intCast(r), .x_stride = @intCast(m.k), .y_stride = @intCast(m.n), .y_f32 = @intFromBool(f32_out) }); + a.add(m.args()); + a.add(out + at * m.n * elem); + a.add(k.parts.ptr); + a.add(k.tickets.ptr); + try launch_.launch(k.decode_fns[@backingInt(m.codebook)], .{ .grid = .{ .x = @intCast(m.n / 128), .y = @intCast(split(m.k, m.n)), .z = 1 }, .block = .{ .x = 256 } }, s, &a); + } + } + + /// Prompt rows (docs/recipes/exl3.md, Prompts): fp16 rows on the tensor units, fp32 sums, bf16 or fp32 out. + pub fn prompt(k: *const Kernels, s: Stream, x: u64, ldx: usize, m: Matrix, out: u64, count: usize, f32_out: bool) !void { + const elem: usize = if (f32_out) 4 else 2; + try k.fits(m, ldx); + var at: usize = 0; + while (at < count) : (at += prompt_chunk) { + const r = @min(prompt_chunk, count - at); + try k.rotate(k.rotate_h_fn, s, x + at * ldx * 2, ldx, m, r); + var a: launch_.Args = .{}; + a.add(k.rotated.ptr); + a.add(@as(u32, @intCast(r))); + a.add(m.trellis); + a.add(m.svh); + a.add(m.args()); + a.add(out + at * m.n * elem); + a.add(@as(u32, @intFromBool(f32_out))); + try launch_.launch(k.prompt_fns[@backingInt(m.codebook)], .{ .grid = .{ .x = @intCast(m.n / 128), .y = @intCast((r + 63) / 64), .z = 1 }, .block = .{ .x = 256 } }, s, &a); + } + } + + fn fits(k: *const Kernels, m: Matrix, ldx: usize) !void { + try m.check(); + if (m.k > k.k_max) return error.Exl3InputsOverScratch; + if (ldx < m.k) return error.RowStride; + } + + /// (x * suh) H times 2^-4 as fp16 rows into the scratch, one warp per (row, 128-block), in `f`'s order. + fn rotate(k: *const Kernels, f: Function, s: Stream, x: u64, ldx: usize, m: Matrix, count: usize) !void { + var a: launch_.Args = .{}; + a.add(x); + a.add(@as(u32, @intCast(ldx))); + a.add(m.suh); + a.add(k.rotated.ptr); + a.add(@as(u32, @intCast(m.k))); + try launch_.launch(f, .{ .grid = .{ .x = @intCast(m.k / 128), .y = @intCast(count), .z = 1 }, .block = .{ .x = 32 } }, s, &a); + } + + /// Every state's fp16 weight under `c` (65536 halves) into `out`. + pub fn codebook(k: *const Kernels, s: Stream, c: Codebook, out: u64) !void { + var a: launch_.Args = .{}; + a.add(@backingInt(c)); + a.add(out); + try launch_.launch(k.codebook_fn, .{ .grid = .{ .x = 256, .y = 1, .z = 1 }, .block = .{ .x = 256 } }, s, &a); + } + + /// W_q [k, n] fp16 (the rotated domain) into `out`, a warp per tile; suh and svh are not read. + pub fn unpack(k: *const Kernels, s: Stream, m: Matrix, out: u64) !void { + if (m.k % 16 != 0 or m.n % 16 != 0 or m.bits < 1 or m.bits > 8) return error.InvalidExl3Matrix; + var a: launch_.Args = .{}; + a.add(m.trellis); + a.add(m.args()); + a.add(@backingInt(m.codebook)); + a.add(out); + try launch_.launch(k.unpack_fn, .{ .grid = .{ .x = @intCast(m.n / 16), .y = @intCast(m.k / 16), .z = 1 }, .block = .{ .x = 32 } }, s, &a); + } +}; + +comptime { + std.debug.assert(@sizeOf(RowsArgs) == 16 and @sizeOf(MatrixArgs) == 20); +} + +test "a matrix exl3.cu cannot take is refused before any launch" { + const base: Matrix = .{ .trellis = 0, .suh = 0, .svh = 0, .k = 4096, .n = 2048, .bits = 4, .codebook = .mul1 }; + try base.check(); + var m = base; + m.n = 2000; + try std.testing.expectError(error.InvalidExl3Matrix, m.check()); + m = base; + m.bits = 9; + try std.testing.expectError(error.InvalidExl3Matrix, m.check()); +} + +test "the trellis size is 16 bits a 16-value word row, 32 bytes a bit" { + // GLM-5.3-Flash's routed gate_proj at K=4: trellis [256, 128, 64] int16 + try std.testing.expectEqual(@as(usize, 256 * 128 * 64 * 2), Matrix.trellisBytes(4096, 2048, 4)); +} + +test "the split comes from the shape alone and leaves no part empty" { + try std.testing.expectEqual(@as(usize, 6), split(4096, 2048)); // a routed gate_proj: 16 column blocks, 32 input blocks + try std.testing.expectEqual(@as(usize, 3), split(16384, 4096)); // o_proj + try std.testing.expectEqual(@as(usize, 1), split(1536, 16384)); // q_b_proj: 128 column blocks already + try std.testing.expectEqual(@as(usize, 1), split(512, 384)); + var k: usize = 128; + while (k <= 16384) : (k += 128) { + var n: usize = 128; + while (n <= 16384) : (n += 128) { + const p = split(k, n); + const per = (k / 128 + p - 1) / p; + try std.testing.expect(p * (n / 128) <= 2 * target_blocks and (p - 1) * per < k / 128); + } + } +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 0fa8bf6e28..b17d2725b6 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -27,6 +27,7 @@ pub const affine4_pack: []const u8 = if (available) &Blob("fatbin_affine4_pack") pub const fp8_experts: []const u8 = if (available) &Blob("fatbin_fp8_experts").bytes else &.{}; pub const exlr8: []const u8 = if (available) &Blob("fatbin_exlr8").bytes else &.{}; pub const exlr8_dense: []const u8 = if (available) &Blob("fatbin_exlr8_dense").bytes else &.{}; +pub const exl3: []const u8 = if (available) &Blob("fatbin_exl3").bytes else &.{}; pub const nvfp4_experts: []const u8 = if (available) &Blob("fatbin_nvfp4_experts").bytes else &.{}; pub const sample: []const u8 = if (available) &Blob("fatbin_sample").bytes else &.{}; pub const train: []const u8 = if (available) &Blob("fatbin_train").bytes else &.{}; diff --git a/zig/src/cuda/nvfp4.zig b/zig/src/cuda/nvfp4.zig index d9502895d3..e0d3c838f4 100644 --- a/zig/src/cuda/nvfp4.zig +++ b/zig/src/cuda/nvfp4.zig @@ -64,12 +64,13 @@ pub const Prompt = struct { mod: Module, fp4: Function, fp8: Function, + fp4_f32: Function, // FP4 writing the fp32 sums unrounded const bm = 128; const bn = 128; const threads = 128; - const name = "_ZN11tf_prompt1613prompt_kernelILi{d}ELi128ELi128ELi2ELi2ELi2ELb0EEEvPK13__nv_bfloat16PKhS5_fPviiiiii"; - pub const symbols = [2][:0]const u8{ std.fmt.comptimePrint(name, .{0}), std.fmt.comptimePrint(name, .{1}) }; + const name = "_ZN11tf_prompt1613prompt_kernelILi{d}ELi128ELi128ELi2ELi2ELi2ELb{d}EEEvPK13__nv_bfloat16PKhS5_fPviiiiii"; + pub const symbols = [3][:0]const u8{ std.fmt.comptimePrint(name, .{ 0, 0 }), std.fmt.comptimePrint(name, .{ 1, 0 }), std.fmt.comptimePrint(name, .{ 0, 1 }) }; /// prompt16.cu Tile::SMEM. pub fn sharedBytes(mode: qmmf.Mode) u32 { @@ -86,7 +87,9 @@ pub const Prompt = struct { try fp4.allowDynamicShared(sharedBytes(.fp4)); const fp8 = try mod.function(symbols[1]); try fp8.allowDynamicShared(sharedBytes(.fp8)); - return .{ .mod = mod, .fp4 = fp4, .fp8 = fp8 }; + const fp4_f32 = try mod.function(symbols[2]); + try fp4_f32.allowDynamicShared(sharedBytes(.fp4)); + return .{ .mod = mod, .fp4 = fp4, .fp8 = fp8, .fp4_f32 = fp4_f32 }; } pub fn unload(p: *Prompt) void { @@ -95,9 +98,14 @@ pub const Prompt = struct { /// y [m, n] bf16 = prompt rows x [m, k] bf16 (rows `ldx` apart) times an FP4 or FP8 weight (linear._prompt). pub fn matmul(p: Prompt, s: Stream, x: u64, ldx: usize, m: usize, w: qmmf.Weight, y: u64) !void { + return p.matmulOut(s, x, ldx, m, w, y, .bf16); + } + + /// matmul() writing y as `out` says: bf16, or fp32 (FP4 only). + pub fn matmulOut(p: Prompt, s: Stream, x: u64, ldx: usize, m: usize, w: qmmf.Weight, y: u64, out: qmmf.Out) !void { const f = switch (w.mode) { - .fp4 => p.fp4, - .fp8 => p.fp8, + .fp4 => if (out == .f32) p.fp4_f32 else p.fp4, + .fp8 => if (out == .f32) return error.EpilogueNotBuilt else p.fp8, else => return error.FormatMismatch, }; const rows_t = (m + bm - 1) / bm; @@ -126,6 +134,7 @@ test "the prompt tile's shared bytes and instance names" { try std.testing.expectEqual(@as(u32, 41984), Prompt.sharedBytes(.fp4)); try std.testing.expectEqual(@as(u32, 49152), Prompt.sharedBytes(.fp8)); // 2 x (16384 + 8192): no scale tile try std.testing.expectEqualStrings("_ZN11tf_prompt1613prompt_kernelILi0ELi128ELi128ELi2ELi2ELi2ELb0EEEvPK13__nv_bfloat16PKhS5_fPviiiiii", Prompt.symbols[0]); + try std.testing.expectEqualStrings("_ZN11tf_prompt1613prompt_kernelILi0ELi128ELi128ELi2ELi2ELi2ELb1EEEvPK13__nv_bfloat16PKhS5_fPviiiiii", Prompt.symbols[2]); } test "word tiling reads each nibble where qmm.pack puts it" { diff --git a/zig/src/cuda/qlinear.zig b/zig/src/cuda/qlinear.zig index 6d6a56f5a3..5dfb6ce871 100644 --- a/zig/src/cuda/qlinear.zig +++ b/zig/src/cuda/qlinear.zig @@ -7,6 +7,7 @@ const memory = @import("memory.zig"); const stream_ = @import("stream.zig"); const qmmf = @import("qmmf.zig"); const nvfp4 = @import("nvfp4.zig"); +const exl3 = @import("exl3_linear.zig"); const Driver = driver.Driver; const Module = module.Module; @@ -36,11 +37,13 @@ pub const Weight = union(enum) { nvfp4: qmmf.Weight, fp8: qmmf.Weight, // e4m3 with one fp32 scale (ModelOpt FP8) bf16: qmmf.Weight, // unquantized: the checkpoint's bf16 values + exl3: exl3.Matrix, // trellis tiles as stored, with suh and svh pub fn outputs(w: Weight) usize { return switch (w) { .affine4 => |q| q.n, .fp8g, .nvfp4, .fp8, .bf16 => |q| q.n, + .exl3 => |q| q.n, }; } @@ -48,6 +51,7 @@ pub const Weight = union(enum) { return switch (w) { .affine4 => |q| q.k, .fp8g, .nvfp4, .fp8, .bf16 => |q| q.k, + .exl3 => |q| q.k, }; } }; @@ -60,6 +64,7 @@ pub const Linear = struct { affine4: ?*const Affine4Kernels = null, lane: ?*const qmmf.Lane = null, // decode rows of every lane format; fp8g and bf16 prompt rows too nvfp4_prompt: ?*const nvfp4.Prompt = null, // the prompt GEMM: nvfp4 and fp8 prompt rows (linear.py's prefill) + exl3: ?*const exl3.Kernels = null, // decode rows 16 a window; prompt rows on the tensor units /// Decode rows: out (rows, n) bf16, each row's bits the same at every row count the format takes. pub fn decode(l: Linear, s: Stream, w: Weight, in: Rows, out: u64, rows: usize) !void { @@ -73,6 +78,7 @@ pub const Linear = struct { .nvfp4 => |q| return l.onLane(s, q, .fp4, in, out, rows), .fp8 => |q| return l.onLane(s, q, .fp8, in, out, rows), .bf16 => |q| return l.onLane(s, q, .bf16, in, out, rows), + .exl3 => |q| return l.onExl3(s, q, in, out, rows), } } @@ -88,19 +94,53 @@ pub const Linear = struct { .nvfp4 => |q| return l.onPrompt(s, q, .fp4, in, out, rows), .fp8 => |q| return l.onPrompt(s, q, .fp8, in, out, rows), .bf16 => |q| return l.onLane(s, q, .bf16, in, out, rows), + .exl3 => |q| return (l.exl3 orelse return error.FormatNotLoaded).prompt(s, in.x, if (in.ldx == 0) q.k else in.ldx, q, out, rows, false), } } + /// Decode rows as decode(), out (rows, n) fp32 unrounded: a tensor-parallel rank's partial rows (nvfp4, exl3). + pub fn decodeF32(l: Linear, s: Stream, w: Weight, in: Rows, out: u64, rows: usize) !void { + const q = switch (w) { + .nvfp4 => |q| q, + .exl3 => |q| return (l.exl3 orelse return error.FormatNotLoaded).rows(s, in.x, if (in.ldx == 0) q.k else in.ldx, q, out, rows, true), + else => return error.EpilogueNotBuilt, + }; + return l.onLaneAs(s, q, .fp4, in, out, .f32, rows); + } + + /// Prompt rows as prompt(), out (rows, n) fp32 unrounded (nvfp4, exl3). + pub fn promptF32(l: Linear, s: Stream, w: Weight, in: Rows, out: u64, rows: usize) !void { + const q = switch (w) { + .nvfp4 => |q| q, + .exl3 => |q| return (l.exl3 orelse return error.FormatNotLoaded).prompt(s, in.x, if (in.ldx == 0) q.k else in.ldx, q, out, rows, true), + else => return error.EpilogueNotBuilt, + }; + return l.onPromptAs(s, q, .fp4, in, out, .f32, rows); + } + + fn onExl3(l: Linear, s: Stream, q: exl3.Matrix, in: Rows, out: u64, rows: usize) !void { + const e = l.exl3 orelse return error.FormatNotLoaded; + return e.rows(s, in.x, if (in.ldx == 0) q.k else in.ldx, q, out, rows, false); + } + fn onPrompt(l: Linear, s: Stream, q: qmmf.Weight, mode: qmmf.Mode, in: Rows, out: u64, rows: usize) !void { + return l.onPromptAs(s, q, mode, in, out, .bf16, rows); + } + + fn onPromptAs(l: Linear, s: Stream, q: qmmf.Weight, mode: qmmf.Mode, in: Rows, out: u64, kind: qmmf.Out, rows: usize) !void { const p = l.nvfp4_prompt orelse return error.FormatNotLoaded; if (q.mode != mode) return error.FormatMismatch; - return p.matmul(s, in.x, if (in.ldx == 0) q.k else in.ldx, rows, q, out); + return p.matmulOut(s, in.x, if (in.ldx == 0) q.k else in.ldx, rows, q, out, kind); } fn onLane(l: Linear, s: Stream, q: qmmf.Weight, mode: qmmf.Mode, in: Rows, out: u64, rows: usize) !void { + return l.onLaneAs(s, q, mode, in, out, .bf16, rows); + } + + fn onLaneAs(l: Linear, s: Stream, q: qmmf.Weight, mode: qmmf.Mode, in: Rows, out: u64, kind: qmmf.Out, rows: usize) !void { const lane = l.lane orelse return error.FormatNotLoaded; if (q.mode != mode) return error.FormatMismatch; - return lane.matmul(s, in.x, if (in.ldx == 0) q.k else in.ldx, rows, q, out, in.part); + return lane.matmulOut(s, in.x, if (in.ldx == 0) q.k else in.ldx, rows, q, out, kind, in.part); } }; @@ -315,6 +355,9 @@ test "a format that is not loaded is refused before any launch" { const s: Stream = undefined; // never reached: the refusal comes first try std.testing.expectError(error.FormatNotLoaded, l.decode(s, w, .{ .x = 0 }, 0, 1)); try std.testing.expectError(error.FormatNotLoaded, l.prompt(s, w, .{ .x = 0 }, 0, 1)); + const e: Weight = .{ .exl3 = .{ .trellis = 0, .suh = 0, .svh = 0, .k = 128, .n = 128, .bits = 4, .codebook = .mul1 } }; + try std.testing.expectError(error.FormatNotLoaded, l.decode(s, e, .{ .x = 0 }, 0, 1)); + try std.testing.expectError(error.FormatNotLoaded, l.prompt(s, e, .{ .x = 0 }, 0, 1)); } test "the affine-4 tiles' sizes follow qmm_fast.tile" { @@ -331,4 +374,7 @@ test "a lane weight whose mode is not its format's is refused before any launch" const s: Stream = undefined; try std.testing.expectError(error.FormatMismatch, l.decode(s, .{ .fp8g = fp4 }, .{ .x = 0 }, 0, 1)); try std.testing.expectError(error.FormatNotLoaded, l.prompt(s, .{ .nvfp4 = fp4 }, .{ .x = 0 }, 0, 1)); + const f8: qmmf.Weight = .{ .mode = .fp8g, .codes = 0, .scales = 0, .n = 64, .k = 64, .npad = 128 }; + try std.testing.expectError(error.EpilogueNotBuilt, l.decodeF32(s, .{ .fp8g = f8 }, .{ .x = 0 }, 0, 1)); + try std.testing.expectError(error.EpilogueNotBuilt, l.promptF32(s, .{ .bf16 = f8 }, .{ .x = 0 }, 0, 1)); } diff --git a/zig/src/cuda/qmmf.zig b/zig/src/cuda/qmmf.zig index 054e2c34bb..38e3e43a33 100644 --- a/zig/src/cuda/qmmf.zig +++ b/zig/src/cuda/qmmf.zig @@ -66,35 +66,43 @@ pub fn plan(m: usize, n: usize, k: usize, clusters: bool) Plan { } /// qmmf.cu's instance names, per mode: row tiles 16, 32, 64 (each plain, then with a cluster), then the fused kernel. -fn symbol(comptime mode: Mode, comptime bm: u32, comptime cluster: bool, comptime fused: bool) [:0]const u8 { - const name = "_ZN7tf_qmmf11qmmf_kernelILi{d}ELi{d}ELi64ELi1ELi4ELi4ELb0ELb{d}ELb{d}EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii"; - return std.fmt.comptimePrint(name, .{ @backingInt(mode), bm, @intFromBool(cluster), @intFromBool(fused) }); +fn symbol(comptime mode: Mode, comptime bm: u32, comptime cluster: bool, comptime fused: bool, comptime wide: bool) [:0]const u8 { + const name = "_ZN7tf_qmmf11qmmf_kernelILi{d}ELi{d}ELi64ELi1ELi4ELi4ELb{d}ELb{d}ELb{d}EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii"; + return std.fmt.comptimePrint(name, .{ @backingInt(mode), bm, @intFromBool(wide), @intFromBool(cluster), @intFromBool(fused) }); } -pub const reduce_symbol = "_ZN7tf_qmmf13reduce_kernelILb0EEEvPKfPvxif"; +pub const reduce_symbols = [2][:0]const u8{ "_ZN7tf_qmmf13reduce_kernelILb0EEEvPKfPvxif", "_ZN7tf_qmmf13reduce_kernelILb1EEEvPKfPvxif" }; const row_tiles = [_]u32{ 16, 32, 64 }; +/// What a call writes: bf16 rows, or the fp32 sums unrounded (FP4 only, a tensor-parallel rank's partial rows). +pub const Out = enum { bf16, f32 }; + +fn loadTiles(mod: Module, comptime mode: Mode, comptime wide: bool) ![7]Function { + var set: [7]Function = undefined; + inline for (row_tiles, 0..) |bm, ti| inline for (.{ false, true }, 0..) |cl, ci| { + const f = try mod.function(symbol(mode, bm, cl, false, wide)); + try f.allowDynamicShared(sharedBytes(mode, bm)); + set[ti * 2 + ci] = f; + }; + set[6] = try mod.function(symbol(mode, 64, false, true, wide)); + try set[6].allowDynamicShared(sharedBytes(mode, 64)); + return set; +} + pub const Lane = struct { mod: Module, fns: [modes.len][7]Function, // per mode: lane16, lane16c, lane32, lane32c, lane64, lane64c, fused - reduce: Function, + fp4_f32: [7]Function, // FP4's tiles writing fp32 + reduce: [2]Function, // by Out clusters: bool, // sm_90 on: slices meet in a cluster's shared memory pub fn load(d: *const Driver, capability_major: u32) !Lane { if (!kernels.available) return error.BuiltWithoutKernels; var mod = try Module.load(d, kernels.qmmf); errdefer mod.unload(); - var l: Lane = .{ .mod = mod, .fns = undefined, .reduce = try mod.function(reduce_symbol), .clusters = capability_major >= 9 }; - inline for (modes, 0..) |mode, mi| { - inline for (row_tiles, 0..) |bm, ti| inline for (.{ false, true }, 0..) |cl, ci| { - const f = try mod.function(symbol(mode, bm, cl, false)); - try f.allowDynamicShared(sharedBytes(mode, bm)); - l.fns[mi][ti * 2 + ci] = f; - }; - const f = try mod.function(symbol(mode, 64, false, true)); - try f.allowDynamicShared(sharedBytes(mode, 64)); - l.fns[mi][6] = f; - } + var l: Lane = .{ .mod = mod, .fns = undefined, .fp4_f32 = try loadTiles(mod, .fp4, true), .reduce = undefined, .clusters = capability_major >= 9 }; + inline for (modes, 0..) |mode, mi| l.fns[mi] = try loadTiles(mod, mode, false); + for (&l.reduce, reduce_symbols) |*f, name| f.* = try mod.function(name); return l; } @@ -110,8 +118,14 @@ pub const Lane = struct { /// y [m, n] bf16 = x [m, k] bf16 (rows `ldx` apart) times the weight, fp32 sums; `part` holds partBytes. pub fn matmul(l: Lane, s: Stream, x: u64, ldx: usize, m: usize, w: Weight, y: u64, part: u64) !void { + return l.matmulOut(s, x, ldx, m, w, y, .bf16, part); + } + + /// matmul() writing y as `out` says: bf16, or fp32 (FP4 only). + pub fn matmulOut(l: Lane, s: Stream, x: u64, ldx: usize, m: usize, w: Weight, y: u64, out: Out, part: u64) !void { + if (out == .f32 and w.mode != .fp4) return error.EpilogueNotBuilt; const p = plan(m, w.n, w.k, l.clusters); - const set = l.fns[modeIndex(w.mode)]; + const set = if (out == .f32) l.fp4_f32 else l.fns[modeIndex(w.mode)]; const ti: usize = if (p.bm == 16) 0 else if (p.bm == 32) 1 else 2; const f = if (p.fused) set[6] else set[ti * 2 + @intFromBool(p.cluster)]; const rows_t = (m + p.bm - 1) / p.bm; @@ -144,7 +158,7 @@ pub const Lane = struct { r.add(@as(i64, @intCast(total))); r.add(@as(i32, @intCast(p.sk))); r.add(w.scale); - try launch_.launch(l.reduce, .{ .grid = .{ .x = @intCast((total + 255) / 256), .y = 1, .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, s, &r); + try launch_.launch(l.reduce[@backingInt(out)], .{ .grid = .{ .x = @intCast((total + 255) / 256), .y = 1, .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, s, &r); } } }; @@ -209,6 +223,7 @@ test "tile shared bytes, slices and row tiles as the Python host picks them" { } test "instance names are qmmf.cu's" { - try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi3ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp8g, 16, true, false)); - try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi0ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp4, 64, false, true)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi3ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp8g, 16, true, false, false)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi0ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp4, 64, false, true, false)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi0ELi32ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp4, 32, true, false, true)); } diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 3053c09a6b..85f0ffc066 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -27,6 +27,7 @@ pub const kernels = @import("kernels.zig"); pub const qmmf = @import("qmmf.zig"); pub const fp8 = @import("fp8.zig"); pub const exlr8 = @import("exlr8.zig"); +pub const exl3_linear = @import("exl3_linear.zig"); pub const nvfp4 = @import("nvfp4.zig"); pub const segments = @import("segments.zig"); pub const grouped = @import("grouped.zig"); @@ -43,6 +44,7 @@ test { _ = @import("qlinear.zig"); _ = @import("norms.zig"); _ = @import("embed_rows.zig"); + _ = @import("exl3_linear.zig"); _ = @import("experts.zig"); _ = @import("sample.zig"); _ = launch; diff --git a/zig/tests/cuda/exl3_tests.zig b/zig/tests/cuda/exl3_tests.zig new file mode 100644 index 0000000000..7243552bcb --- /dev/null +++ b/zig/tests/cuda/exl3_tests.zig @@ -0,0 +1,557 @@ +//! EXL3 projections on CUDA: codebooks and tiles bit for bit against core/exl3_format.zig, FP64 bounds, row windows. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const Gpu = check.Gpu; +const fmt = core.exl3_format; +const E = cuda.exl3_linear; +const st = core.safetensors; + +const k_max = 16384; // GLM-5.3-Flash's widest input (o_proj: 64 heads of 256) +const u = std.math.pow(f64, 2, -24); + +/// Device buffers for one test, freed together. +pub const Rig = struct { + gpu: Gpu, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + + pub fn put(r: *Rig, items: anytype) !u64 { + const b = try cuda.DeviceBuffer.fromHost(r.gpu.d, std.mem.sliceAsBytes(items)); + try r.bufs.append(r.gpu.gpa, b); + return b.ptr; + } + pub fn zeros(r: *Rig, len: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(r.gpu.d, @max(len, 16)); + try b.fill8(0, null); + try r.gpu.d.check(r.gpu.d.api.cuStreamSynchronize(null), "cuStreamSynchronize"); + try r.bufs.append(r.gpu.gpa, b); + return b.ptr; + } + pub fn get(r: *Rig, ptr: u64, out: []u8) !void { + const b: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = ptr, .len = out.len }; + try b.download(0, out); + } + pub fn deinit(r: *Rig) void { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(r.gpu.gpa); + } +}; + +fn refBook(c: E.Codebook) fmt.Codebook { + return switch (c) { + .inst3 => .inst3, + .mcg => .mcg, + .mul1 => .mul1, + }; +} + +pub fn bf16(v: f32) u16 { + return @intCast(@as(u32, @bitCast(v)) >> 16); +} + +fn ofBf16(b: u16) f64 { + return @as(f32, @bitCast(@as(u32, b) << 16)); +} + +fn gamma(n: usize) f64 { + return gammaU(n, u); +} + +fn gammaU(n: usize, unit: f64) f64 { + const nu = @as(f64, @floatFromInt(n)) * unit; + return nu / (1 - nu); +} + +/// `tf-cuda-test exl3 [DIR]`: codebooks, tiles, FP64 bounds and row windows; DIR adds a checkpoint's real matrices. +pub fn run(gpu: Gpu, dir: ?[]const u8) !void { + var mod = try cuda.Module.load(gpu.d, cuda.kernels.exl3); + defer mod.unload(); + var k = try E.Kernels.resolve(gpu.d, mod, k_max); + defer k.deinit(); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + try codebooks(gpu, &k, s); + try tiles(gpu, &k, s); + var b: Bound = .{}; + const cases = [_]Case{ + .{ .k = 256, .n = 256, .bits = 1, .book = .mul1 }, .{ .k = 256, .n = 384, .bits = 2, .book = .mcg }, + .{ .k = 384, .n = 256, .bits = 3, .book = .inst3 }, .{ .k = 256, .n = 256, .bits = 4, .book = .mul1 }, + .{ .k = 512, .n = 256, .bits = 5, .book = .mul1 }, .{ .k = 256, .n = 512, .bits = 6, .book = .mul1 }, + .{ .k = 256, .n = 256, .bits = 7, .book = .mcg }, .{ .k = 256, .n = 256, .bits = 8, .book = .mul1 }, + .{ .k = 1024, .n = 256, .bits = 4, .book = .inst3 }, .{ .k = 2048, .n = 128, .bits = 4, .book = .mul1 }, + .{ .k = 1024, .n = 256, .bits = 6, .book = .mcg }, + }; + for (cases, 0..) |c, i| try synthetic(gpu, &k, s, c, 0xE3 + i, &b); + check.pass("{d} synthetic matrices, widths 1-8, three codebooks: {d} FP64 values, max error {e}, max {d:.4} of the bound, worst column {e} of the largest output", .{ cases.len, b.values, b.max_error, b.max_fraction, b.worst_column }); + try windows(gpu, &k, s); + try fp32Out(gpu, &k, s); + if (dir) |d| try real(gpu, &k, s, d); +} + +fn codebooks(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream) !void { + var rig: Rig = .{ .gpu = gpu }; + defer rig.deinit(); + const out = try rig.zeros(65536 * 2); + const got = try gpu.gpa.alloc(u16, 65536); + defer gpu.gpa.free(got); + for ([_]E.Codebook{ .inst3, .mcg, .mul1 }) |c| { + try k.codebook(s, c, out); + try s.synchronize(); + try rig.get(out, std.mem.sliceAsBytes(got)); + for (got, 0..) |g, i| try check.expect(g == @as(u16, @bitCast(fmt.codebookValue(refBook(c), @intCast(i)))), "exl3 {t} state {d}", .{ c, i }); + } + check.pass("3 codebooks x 65536 states: the GPU's fp16 weights equal exl3_format.codebookValue bit for bit", .{}); +} + +fn tiles(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream) !void { + var prng = std.Random.DefaultPrng.init(0x7E11); + const r = prng.random(); + const tk = 2; + const tn = 3; + var count: usize = 0; + for (1..9) |bits| for ([_]E.Codebook{ .inst3, .mcg, .mul1 }) |c| { + var rig: Rig = .{ .gpu = gpu }; + defer rig.deinit(); + const words = try gpu.gpa.alloc(i16, tk * tn * 16 * bits); + defer gpu.gpa.free(words); + for (words) |*w| w.* = r.int(i16); + const m: E.Matrix = .{ .trellis = try rig.put(words), .suh = 0, .svh = 0, .k = tk * 16, .n = tn * 16, .bits = @intCast(bits), .codebook = c }; + const out = try rig.zeros(m.k * m.n * 2); + try k.unpack(s, m, out); + try s.synchronize(); + var got: [tk * tn * 256]f16 = undefined; + var want: [tk * tn * 256]f16 = undefined; + try rig.get(out, std.mem.sliceAsBytes(&got)); + fmt.unpack(.{ .halves = @intCast(2 * bits) }, refBook(c), words, tk, tn, &want); + try check.sameBytes("exl3 unpack", std.mem.sliceAsBytes(&got), std.mem.sliceAsBytes(&want)); + count += 1; + }; + check.pass("{d} width/codebook pairs: random tiles unpack to exl3_format.unpack's W_q bit for bit", .{count}); +} + +pub const Case = struct { k: usize, n: usize, bits: u32, book: E.Codebook }; + +/// Worst observed error, as a value and as a fraction of the fp32 bound, over every value checked. +pub const Bound = struct { + values: usize = 0, + max_error: f64 = 0, + max_fraction: f64 = 0, + worst_column: f64 = 0, +}; + +/// One matrix's host copy (a column block of a real one, or a whole synthetic one) with its rows of input. +pub const Host = struct { + k: usize, + n: usize, + bits: u32, + book: E.Codebook, + trellis: []const i16, + suh: []const f16, + svh: []const f16, +}; + +/// FP64 y = x W from exl3_format.dequantize, and per value the bound of exl3.cu's order (rotate, chain, transform). +pub fn reference(a: std.mem.Allocator, h: Host, x: []const u16, rows: usize, y: []f64, bound: []f64, prompt: bool) !void { + const w = try a.alloc(f64, h.k * h.n); + defer a.free(w); + const scratch = try a.alloc(f64, 2 * h.k * h.n); + defer a.free(scratch); + fmt.dequantize(.{ .halves = @intCast(2 * h.bits) }, refBook(h.book), h.trellis, h.k / 16, h.n / 16, h.suh, h.svh, w, scratch); + const wq = try a.alloc(f16, h.k * h.n); + defer a.free(wq); + fmt.unpack(.{ .halves = @intCast(2 * h.bits) }, refBook(h.book), h.trellis, h.k / 16, h.n / 16, wq); + const rot = try a.alloc(f64, 2 * h.k); + defer a.free(rot); + const z = try a.alloc(f64, 2 * h.n); + defer a.free(z); + for (0..rows) |i| { + const xr = x[i * h.k ..][0..h.k]; + for (0..h.n) |j| { + var acc: f64 = 0; + for (0..h.k) |l| acc += ofBf16(xr[l]) * w[l * h.n + j]; + y[i * h.n + j] = acc; + } + // r = (x suh) H unnormalized, |r| and its error from one product and seven adds + const r = rot[0..h.k]; + const re = rot[h.k..]; + for (0..h.k) |l| r[l] = ofBf16(xr[l]) * @as(f64, @floatCast(h.suh[l])); + var blk: usize = 0; + while (blk < h.k) : (blk += 128) { + var sum: f64 = 0; + for (r[blk..][0..128]) |v| sum += @abs(v); + fwht(r[blk..][0..128]); + for (re[blk..][0..128], r[blk..][0..128]) |*e, v| { + e.* = gamma(8) * sum; + // prompt rows: r / 16 rounded to fp16 (relative 2^-11, 2^-25 absolute among subnormals) + if (prompt) e.* += std.math.pow(f64, 2, -11) * (@abs(v) + e.*) + std.math.pow(f64, 2, -21); + } + } + const za = z[0..h.n]; + const ze = z[h.n..]; + for (0..h.n) |j| { + var mag: f64 = 0; + var carried: f64 = 0; + for (0..h.k) |l| { + const q = @abs(@as(f64, @floatCast(wq[l * h.n + j]))); + mag += (@abs(r[l]) + re[l]) * q; + carried += re[l] * q; + } + ze[j] = (if (prompt) gammaU(h.k + 3, std.math.pow(f64, 2, -22)) else gamma(h.k + 3)) * mag + carried; + za[j] = mag; + } + var ob: usize = 0; + while (ob < h.n) : (ob += 128) { + var zsum: f64 = 0; + var esum: f64 = 0; + for (ob..ob + 128) |j| { + zsum += za[j] + ze[j]; + esum += ze[j]; + } + for (ob..ob + 128) |j| { + const e = std.math.pow(f64, 2, -7) * @abs(@as(f64, @floatCast(h.svh[j]))) * (gamma(7) * zsum + esum); + bound[i * h.n + j] = e * (1 + u) + (if (prompt) std.math.pow(f64, 2, -8) else u) * (@abs(y[i * h.n + j]) + e); + } + } + } +} + +fn fwht(v: []f64) void { + var d: usize = 1; + while (d < v.len) : (d <<= 1) { + for (0..v.len) |i| { + if (i & d != 0) continue; + const a = v[i]; + const b = v[i + d]; + v[i] = a + b; + v[i + d] = a - b; + } + } +} + +/// Every value within its bound; the largest error by column against that column's largest output. +pub fn compare(what: []const u8, got: []const f32, y: []const f64, bound: []const f64, rows: usize, n: usize, b: *Bound) !void { + var scale: f64 = 0; + for (y[0 .. rows * n]) |v| scale = @max(scale, @abs(v)); + for (0..n) |j| { + var col_err: f64 = 0; + for (0..rows) |i| { + const at = i * n + j; + const err = @abs(@as(f64, got[at]) - y[at]); + try check.expect(err <= bound[at], "{s}: row {d} column {d} error {e} over its bound {e}", .{ what, i, j, err, bound[at] }); + b.max_error = @max(b.max_error, err); + if (bound[at] > 0) b.max_fraction = @max(b.max_fraction, err / bound[at]); + col_err = @max(col_err, err); + } + if (scale > 0) b.worst_column = @max(b.worst_column, col_err / scale); + } + b.values += rows * n; +} + +pub fn randomMatrix(a: std.mem.Allocator, r: std.Random, c: Case) !Host { + const trellis = try a.alloc(i16, E.Matrix.trellisBytes(c.k, c.n, c.bits) / 2); + for (trellis) |*w| w.* = r.int(i16); + const suh = try a.alloc(f16, c.k); + const svh = try a.alloc(f16, c.n); + for (suh) |*v| v.* = @floatCast((0.25 + 1.75 * r.float(f32)) * @as(f32, if (r.boolean()) -1 else 1)); + for (svh) |*v| v.* = @floatCast((0.002 + 0.02 * r.float(f32)) * @as(f32, if (r.boolean()) -1 else 1)); + return .{ .k = c.k, .n = c.n, .bits = c.bits, .book = c.book, .trellis = trellis, .suh = suh, .svh = svh }; +} + +pub fn freeHost(a: std.mem.Allocator, h: Host) void { + a.free(h.trellis); + a.free(h.suh); + a.free(h.svh); +} + +pub fn randomRows(a: std.mem.Allocator, r: std.Random, rows: usize, k: usize) ![]u16 { + const x = try a.alloc(u16, rows * k); + for (x) |*v| v.* = bf16(r.floatNorm(f32) * 2); + return x; +} + +fn upload(rig: *Rig, h: Host) !E.Matrix { + return .{ .trellis = try rig.put(h.trellis), .suh = try rig.put(h.suh), .svh = try rig.put(h.svh), .k = h.k, .n = h.n, .bits = h.bits, .codebook = h.book }; +} + +fn synthetic(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream, c: Case, seed: u64, b: *Bound) !void { + const a = gpu.gpa; + var prng = std.Random.DefaultPrng.init(seed); + const h = try randomMatrix(a, prng.random(), c); + defer freeHost(a, h); + const rows = 16; + const x = try randomRows(a, prng.random(), rows, c.k); + defer a.free(x); + var rig: Rig = .{ .gpu = gpu }; + defer rig.deinit(); + const m = try upload(&rig, h); + const out = try rig.zeros(rows * c.n * 4); + try k.rows(s, try rig.put(x), c.k, m, out, rows, true); + try s.synchronize(); + const got = try a.alloc(f32, rows * c.n); + defer a.free(got); + try rig.get(out, std.mem.sliceAsBytes(got)); + const y = try a.alloc(f64, rows * c.n); + defer a.free(y); + const bound = try a.alloc(f64, rows * c.n); + defer a.free(bound); + try reference(a, h, x, rows, y, bound, true); + try compare("exl3 synthetic", got, y, bound, rows, c.n, b); + // the same rows as a prompt: the tensor units' arithmetic, bf16 out + try k.prompt(s, try rig.put(x), c.k, m, out, rows, false); + try s.synchronize(); + try getBf16(&rig, out, got); + try reference(a, h, x, rows, y, bound, true); + try compare("exl3 synthetic prompt", got, y, bound, rows, c.n, b); +} + +fn getBf16(rig: *Rig, ptr: u64, out: []f32) !void { + const raw = try rig.gpu.gpa.alloc(u16, out.len); + defer rig.gpu.gpa.free(raw); + try rig.get(ptr, std.mem.sliceAsBytes(raw)); + for (out, raw) |*o, v| o.* = @floatCast(ofBf16(v)); +} + +/// Each row alone, then every window 1-16 and a 40-row prompt through qlinear: each row's bf16 bytes as when alone. +fn windows(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream) !void { + const a = gpu.gpa; + var prng = std.Random.DefaultPrng.init(0x3D16); + const h = try randomMatrix(a, prng.random(), .{ .k = 2048, .n = 256, .bits = 4, .book = .mul1 }); + defer freeHost(a, h); + const rows = 40; + const prompt_rows = 300; // past a 256-row chunk and a 64-row block + const x = try randomRows(a, prng.random(), prompt_rows, h.k); + defer a.free(x); + var rig: Rig = .{ .gpu = gpu }; + defer rig.deinit(); + const w: cuda.qlinear.Weight = .{ .exl3 = try upload(&rig, h) }; + const l: cuda.qlinear.Linear = .{ .exl3 = k }; + const xd = try rig.put(x); + const out = try rig.zeros(prompt_rows * h.n * 2); + const row_bytes = h.n * 2; + const alone = try a.alloc(u8, rows * row_bytes); + defer a.free(alone); + const got = try a.alloc(u8, rows * row_bytes); + defer a.free(got); + for (0..rows) |i| { + try l.decode(s, w, .{ .x = xd + i * h.k * 2 }, out, 1); + try s.synchronize(); + try rig.get(out, alone[i * row_bytes ..][0..row_bytes]); + } + for (1..E.window + 1) |m| { + try l.decode(s, w, .{ .x = xd }, out, m); + try s.synchronize(); + try rig.get(out, got[0 .. m * row_bytes]); + try check.sameBytes("exl3 decode window", got[0 .. m * row_bytes], alone[0 .. m * row_bytes]); + } + const p_alone = try a.alloc(u8, prompt_rows * row_bytes); + defer a.free(p_alone); + const p_got = try a.alloc(u8, prompt_rows * row_bytes); + defer a.free(p_got); + for (0..prompt_rows) |i| { + try l.prompt(s, w, .{ .x = xd + i * h.k * 2 }, out, 1); + try s.synchronize(); + try rig.get(out, p_alone[i * row_bytes ..][0..row_bytes]); + } + for ([_]usize{ 2, 63, 64, 65, 255, 256, 257, prompt_rows }) |m| { + try l.prompt(s, w, .{ .x = xd }, out, m); + try s.synchronize(); + try rig.get(out, p_got[0 .. m * row_bytes]); + try check.sameBytes("exl3 prompt rows", p_got[0 .. m * row_bytes], p_alone[0 .. m * row_bytes]); + } + check.pass("decode windows of 1-16 rows ({d} input parts) and prompts of 2-{d} rows through qlinear: every row's bf16 bytes equal the row alone on its path", .{ E.split(h.k, h.n), prompt_rows }); +} + +/// decodeF32 and promptF32 through qlinear: within FP64 bounds, the same bits at every row count, the bf16 out their rounding. +fn fp32Out(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream) !void { + const a = gpu.gpa; + var prng = std.Random.DefaultPrng.init(0xF32); + const h = try randomMatrix(a, prng.random(), .{ .k = 2048, .n = 256, .bits = 4, .book = .mcg }); + defer freeHost(a, h); + const rows = 300; + const x = try randomRows(a, prng.random(), rows, h.k); + defer a.free(x); + var rig: Rig = .{ .gpu = gpu }; + defer rig.deinit(); + const w: cuda.qlinear.Weight = .{ .exl3 = try upload(&rig, h) }; + const l: cuda.qlinear.Linear = .{ .exl3 = k }; + const in: cuda.qlinear.Rows = .{ .x = try rig.put(x) }; + const o32 = try rig.zeros(rows * h.n * 4); + const o16 = try rig.zeros(rows * h.n * 2); + const wide = try a.alloc(f32, rows * h.n); + defer a.free(wide); + const got = try a.alloc(f32, rows * h.n); + defer a.free(got); + const b16 = try a.alloc(u16, rows * h.n); + defer a.free(b16); + const y = try a.alloc(f64, rows * h.n); + defer a.free(y); + const bound = try a.alloc(f64, rows * h.n); + defer a.free(bound); + var b: Bound = .{}; + for ([_]bool{ false, true }) |prompt| { + const what = if (prompt) "exl3 promptF32" else "exl3 decodeF32"; + const top: usize = if (prompt) rows else E.window; + const counts: []const usize = if (prompt) &.{ 1, 63, 64, 65, 257 } else &.{ 1, 2, 3, 7, 8, 15 }; + if (prompt) try l.promptF32(s, w, in, o32, top) else try l.decodeF32(s, w, in, o32, top); + if (prompt) try l.prompt(s, w, in, o16, top) else try l.decode(s, w, in, o16, top); + try s.synchronize(); + try rig.get(o32, std.mem.sliceAsBytes(wide[0 .. top * h.n])); + try rig.get(o16, std.mem.sliceAsBytes(b16[0 .. top * h.n])); + try reference(a, h, x, top, y, bound, true); + try compare(what, wide, y, bound, top, h.n, &b); + var differ: usize = 0; + for (wide[0 .. top * h.n], b16[0 .. top * h.n]) |v, o| differ += @intFromBool(roundBf16(v) != o); + try check.expect(differ == 0, "{s}: the bf16 out is the fp32 out rounded ({d} of {d} differ)", .{ what, differ, top * h.n }); + for (counts) |m| { + if (prompt) try l.promptF32(s, w, in, o32, m) else try l.decodeF32(s, w, in, o32, m); + try s.synchronize(); + try rig.get(o32, std.mem.sliceAsBytes(got[0 .. m * h.n])); + try check.sameBytes(what, std.mem.sliceAsBytes(got[0 .. m * h.n]), std.mem.sliceAsBytes(wide[0 .. m * h.n])); + } + } + check.pass("decodeF32 (1-{d} rows) and promptF32 (1-{d} rows) through qlinear: {d} FP64 values within bounds (max {d:.4} of the bound), the same bits at every row count, the bf16 out their rounding", .{ E.window, rows, b.values, b.max_fraction }); +} + +fn roundBf16(v: f32) u16 { + const w: u32 = @bitCast(v); + return @intCast((w + 0x7FFF + ((w >> 16) & 1)) >> 16); +} + +/// The checkpoint's real matrices: FP64 on three column blocks, then GB/s at 1 and 16 rows. +fn real(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream, dir: []const u8) !void { + const a = gpu.gpa; + const index_path = try std.fmt.allocPrint(a, "{s}/model.safetensors.index.json", .{dir}); + defer a.free(index_path); + const index_bytes = try std.Io.Dir.cwd().readFileAlloc(gpu.io, index_path, a, .limited(64 << 20)); + defer a.free(index_bytes); + const index = try std.json.parseFromSlice(std.json.Value, a, index_bytes, .{}); + defer index.deinit(); + const map = index.value.object.get("weight_map").?.object; + const P = "model.language_model.layers."; + const names = [_][]const u8{ P ++ "3.mlp.experts.0.gate_proj", P ++ "3.mlp.experts.0.down_proj", P ++ "0.mlp.down_proj", P ++ "3.mlp.shared_experts.up_proj", P ++ "3.self_attn.q_b_proj", P ++ "3.self_attn.o_proj", "lm_head" }; + var b: Bound = .{}; + var checked: usize = 0; + for (names) |name| { + const tname = try std.fmt.allocPrint(a, "{s}.trellis", .{name}); + defer a.free(tname); + const wname = try std.fmt.allocPrint(a, "{s}.weight", .{name}); + defer a.free(wname); + const shard = map.get(tname) orelse { + // a checkpoint may quantize only some matrices (experts-only EXL3): a plain .weight is skipped, a missing tensor fails + if (map.get(wname) == null) return error.TensorNotInIndex; + std.debug.print("SKIP exl3 {s}: not EXL3 in this checkpoint (.weight)\n", .{name}); + continue; + }; + const path = try std.fmt.allocPrint(a, "{s}/{s}", .{ dir, shard.string }); + defer a.free(path); + var f = try st.File.open(a, gpu.io, path); + defer f.close(gpu.io); + try realOne(gpu, k, s, &f, name, &b); + checked += 1; + } + try check.expect(checked > 0, "exl3: none of the {d} selected matrices is EXL3 in this checkpoint", .{names.len}); + check.pass("{d} real matrices: {d} FP64 values, max error {e}, max {d:.4} of the bound, worst column {e} of the largest output", .{ checked, b.values, b.max_error, b.max_fraction, b.worst_column }); +} + +fn tensor(f: *const st.File, name: []const u8, part: []const u8) ?st.Tensor { + var buf: [256]u8 = undefined; + return f.get(std.fmt.bufPrint(&buf, "{s}.{s}", .{ name, part }) catch return null); +} + +fn realOne(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream, f: *const st.File, name: []const u8, b: *Bound) !void { + const a = gpu.gpa; + const t = tensor(f, name, "trellis") orelse return error.MissingTrellis; + const suh_t = tensor(f, name, "suh") orelse return error.MissingSuh; + const svh_t = tensor(f, name, "svh") orelse return error.MissingSvh; + const book: E.Codebook = if (tensor(f, name, "mul1") != null) .mul1 else if (tensor(f, name, "mcg") != null) .mcg else .inst3; + if (t.dtype != .i16 or t.rank != 3 or t.dim(2) % 16 != 0) return error.UnsupportedTrellis; + const kk = t.dim(0) * 16; + const n = t.dim(1) * 16; + const bits: u32 = @intCast(t.dim(2) / 16); + const words: []align(1) const i16 = std.mem.bytesAsSlice(i16, t.bytes); + const suh: []align(1) const f16 = std.mem.bytesAsSlice(f16, suh_t.bytes); + const svh: []align(1) const f16 = std.mem.bytesAsSlice(f16, svh_t.bytes); + var rig: Rig = .{ .gpu = gpu }; + defer rig.deinit(); + const m: E.Matrix = .{ .trellis = try rig.put(t.bytes), .suh = try rig.put(suh_t.bytes), .svh = try rig.put(svh_t.bytes), .k = kk, .n = n, .bits = bits, .codebook = book }; + var prng = std.Random.DefaultPrng.init(0xF1A5); + const rows = 4; + const prompt_rows = 512; + const x = try randomRows(a, prng.random(), prompt_rows, kk); + defer a.free(x); + const xd = try rig.put(x); + const out = try rig.zeros(prompt_rows * n * 4); + try k.rows(s, xd, kk, m, out, rows, true); + try s.synchronize(); + const got = try a.alloc(f32, rows * n); + defer a.free(got); + try rig.get(out, std.mem.sliceAsBytes(got)); + // three 128-column blocks (first, middle, last), each a trellis of its own 8 tile columns + const tile_words = 16 * @as(usize, bits); + const sub = try a.alloc(i16, kk / 16 * 8 * tile_words); + defer a.free(sub); + const suh_c = try a.alloc(f16, kk); + defer a.free(suh_c); + for (suh_c, suh) |*d, v| d.* = v; + const svh_c = try a.alloc(f16, 128); + defer a.free(svh_c); + const y = try a.alloc(f64, rows * 128); + defer a.free(y); + const bound = try a.alloc(f64, rows * 128); + defer a.free(bound); + const g = try a.alloc(f32, rows * 128); + defer a.free(g); + const prompt_got = try a.alloc(f32, rows * n); + defer a.free(prompt_got); + + try k.prompt(s, xd, kk, m, out, rows, false); + try s.synchronize(); + try getBf16(&rig, out, prompt_got); + const blocks = n / 128; + for ([_]usize{ 0, blocks / 2, blocks - 1 }) |ob| { + for (0..kk / 16) |tk| for (0..8) |c| for (0..tile_words) |wi| { + sub[(tk * 8 + c) * tile_words + wi] = words[(tk * (n / 16) + 8 * ob + c) * tile_words + wi]; + }; + for (svh_c, 0..) |*v, j| v.* = svh[ob * 128 + j]; + const sub_h: Host = .{ .k = kk, .n = 128, .bits = bits, .book = book, .trellis = sub, .suh = suh_c, .svh = svh_c }; + try reference(a, sub_h, x, rows, y, bound, true); + for (0..rows) |i| @memcpy(g[i * 128 ..][0..128], got[i * n + ob * 128 ..][0..128]); + try compare(name, g, y, bound, rows, 128, b); + for (0..rows) |i| @memcpy(g[i * 128 ..][0..128], prompt_got[i * n + ob * 128 ..][0..128]); + try compare(name, g, y, bound, rows, 128, b); + } + const one = try time(gpu, k, s, xd, m, out, 1); + const sixteen = try time(gpu, k, s, xd, m, out, E.window); + const many = try timePrompt(gpu, k, s, xd, m, out, prompt_rows); + const gb = @as(f64, @floatFromInt(t.bytes.len)) / 1e9; + check.pass("{s}: K={d} N={d} {d} bits {t}, 3 column blocks within FP64 bounds (decode and prompt); 1 row {d:.1} us ({d:.0} GB/s), 16 rows {d:.1} us, a {d}-row prompt {d:.0} rows/s", .{ name, kk, n, bits, book, one * 1e6, gb / one, sixteen * 1e6, prompt_rows, prompt_rows / many }); +} + +/// Seconds a call, the median of 7 runs of 20 after a warm-up, bf16 out. +fn time(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream, x: u64, m: E.Matrix, out: u64, rows: usize) !f64 { + try k.rows(s, x, m.k, m, out, rows, false); + try s.synchronize(); + var runs: [7]f64 = undefined; + for (&runs) |*r| { + const t0 = check.now(gpu.io); + for (0..20) |_| try k.rows(s, x, m.k, m, out, rows, false); + try s.synchronize(); + r.* = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 20e9; + } + return check.median(&runs); +} + +/// Seconds a prompt call, the median of 7 runs of 5 after a warm-up. +fn timePrompt(gpu: Gpu, k: *const E.Kernels, s: cuda.Stream, x: u64, m: E.Matrix, out: u64, rows: usize) !f64 { + try k.prompt(s, x, m.k, m, out, rows, false); + try s.synchronize(); + var runs: [7]f64 = undefined; + for (&runs) |*r| { + const t0 = check.now(gpu.io); + for (0..5) |_| try k.prompt(s, x, m.k, m, out, rows, false); + try s.synchronize(); + r.* = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 5e9; + } + return check.median(&runs); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 7b8a089947..f6b99a8ca3 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -27,6 +27,7 @@ const exlr8_tests = @import("exlr8_tests.zig"); const exlr8_dense_tests = @import("exlr8_dense_tests.zig"); const norms_tests = @import("norms_tests.zig"); const norm_rows_tests = @import("norm_rows_tests.zig"); +const exl3_tests = @import("exl3_tests.zig"); const usage = \\usage: tf-cuda-test @@ -64,6 +65,7 @@ const usage = \\ exlr8 EXLR8 projections: every codebook state, FP64, row windows, the Metal frozen digests \\ exlr8-bench DIR [MANIFEST] [LAYER] [SETS] a real EXLR8 layer: FP64, row windows, then GB/s by stage \\ exlr8-dense [DIR] EXLR8 dense formats: decode bit-equal, FP64, row windows; DIR adds the real matrices + \\ exl3 [DIR] EXL3 projections: codebooks and tiles bit-equal, FP64, row windows; DIR adds real matrices \\ train-pack MODEL EXAMPLES [COUNT] training passes holding several sequences against one a pass: differences, time \\ ; @@ -121,6 +123,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8")) return exlr8_tests.run(gpu); + if (std.mem.eql(u8, cmd, "exl3")) return exl3_tests.run(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exlr8-dense")) return exlr8_dense_tests.run(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exlr8-bench")) return exlr8_tests.bench(gpu, try arg(rest, 0), if (rest.len > 1) rest[1] else "home", if (rest.len > 2) try std.fmt.parseInt(u16, rest[2], 10) else 40, if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 4); if (std.mem.eql(u8, cmd, "sample")) return sample_tests.draws(gpu); diff --git a/zig/tests/cuda/quant_ref_tests.zig b/zig/tests/cuda/quant_ref_tests.zig index a4552e898b..f33487c538 100644 --- a/zig/tests/cuda/quant_ref_tests.zig +++ b/zig/tests/cuda/quant_ref_tests.zig @@ -168,11 +168,32 @@ fn lanes(r: *Rig, lin: qlinear.Linear, n: usize, k: usize) !void { const path = if (prompt) "prompt" else "decode"; try r.close(try std.fmt.bufPrint(&name, "{s} {s} {d} rows x [{d}, {d}]", .{ f[0], path, wide, n, k }), outs[rows.len - 1], false, ref.want, ref.tol); for (rows[0 .. rows.len - 1], outs[0 .. rows.len - 1]) |m, o| try sameRows(r, try std.fmt.bufPrint(&name, "{s} {s}: {d} rows equal the {d}-row call's", .{ f[0], path, m, wide }), outs[rows.len - 1], o, m * n * 2); + if (comptime std.mem.eql(u8, f[0], "nvfp4")) try fp32Out(r, lin, view, dx, part, n, k, prompt, ref.want, ref.tol, outs[rows.len - 1]); } } check.pass("qlinear fp8g and nvfp4, [{d}, {d}]: decode and prompt within float64 tolerance, every row count's bits equal", .{ n, k }); } +/// nvfp4's decodeF32 and promptF32: in float64 tolerance, the same bits at any row count, the bf16 out its rounding. +fn fp32Out(r: *Rig, lin: qlinear.Linear, view: qlinear.Weight, dx: u64, part: u64, n: usize, k: usize, prompt: bool, want: []const f64, tol: []const f64, bf16_out: u64) !void { + const wide = rows[rows.len - 1]; + const path = if (prompt) "prompt" else "decode"; + var name: [96]u8 = undefined; + const outs = try r.a.alloc(u64, rows.len); + for (rows, outs) |m, *o| { + o.* = try r.zeros(m * n * 4); + const in: qlinear.Rows = .{ .x = dx, .ldx = k, .part = part }; + if (prompt) try lin.promptF32(r.s, view, in, o.*, m) else try lin.decodeF32(r.s, view, in, o.*, m); + } + try r.close(try std.fmt.bufPrint(&name, "nvfp4 {s} fp32 {d} rows x [{d}, {d}]", .{ path, wide, n, k }), outs[rows.len - 1], true, want, tol); + for (rows[0 .. rows.len - 1], outs[0 .. rows.len - 1]) |m, o| try sameRows(r, try std.fmt.bufPrint(&name, "nvfp4 {s} fp32: {d} rows equal the {d}-row call's", .{ path, m, wide }), outs[rows.len - 1], o, m * n * 4); + const f = std.mem.bytesAsSlice(f32, try r.back(outs[rows.len - 1], wide * n * 4)); + const b = std.mem.bytesAsSlice(u16, try r.back(bf16_out, wide * n * 2)); + var differ: usize = 0; + for (f, b) |x, y| differ += @intFromBool(bf16(x) != y); + try check.expect(differ == 0, "nvfp4 {s}: the bf16 out is the fp32 out rounded ({d} of {d} differ)", .{ path, differ, f.len }); +} + pub fn run(gpu: Gpu) !void { var arena: std.heap.ArenaAllocator = .init(gpu.gpa); defer arena.deinit();