diff --git a/docs/recipes/exl3.md b/docs/recipes/exl3.md index 2b36354f58..1756bd4908 100644 --- a/docs/recipes/exl3.md +++ b/docs/recipes/exl3.md @@ -172,4 +172,30 @@ reproduces the numpy oracle on every codebook and width and on real checkpoint b (the oracle vectors ship beside the test). Its real checkpoint bytes are one 128x128 block from `UnstableLlama/Qwen3.5-2B-exl3-8.00bpw` (mcg codebook, 8.00 bpw), in `zig/tests/fixtures/exl3_oracle.bin`. -The scalar decoder and its oracle tests do not register an EXL3 model or provide a native GPU projection. +The scalar decoder and its oracle tests do not register an EXL3 model on their own; the native CUDA kernels below +are checked against them. + +## Native CUDA (Zig) + +The Zig engine carries the same arithmetic as hand-written CUDA, ported from the Python engine's kernels so a +row's bits match it: + +- `zig/kernels/cuda/exl3_decode.cuh`, `exl3_linear.cu`, `zig/src/cuda/exl3.zig`: the input rotation and the linear + for 1 to 128 rows, every codebook and width (the `mul1` half steps included), with the K split planned from the + layer's shape alone. `exl3.Dense` runs prompt rows in 128-row blocks of the same kernel, so a prompt row's bits + equal its decode-row bits wherever its chunk starts; `qlinear.Weight` has an `.exl3` arm for both. +- `zig/kernels/cuda/exl3_experts.cu`, `zig/src/cuda/exl3_experts.zig`: grouped routed experts (member lists, + each expert's rotation, gate and up, GLM's SwiGLU roundings, down and the weighted combine), widths mixed within + a layer; an expert's member count never changes a row's bits. +- `zig/src/core/exl3_load.zig`: every trellis group of a checkpoint from its safetensors headers; + `zig/src/cuda/exl3_upload.zig` puts a group, or a layer's experts, on the device. + +`tf-cuda-test exl3 DIR` checks them: every width and codebook against float64 references built on +`exl3_format.zig`, row and chunk invariance, and, with `DIR` from +`zig/tests/cuda/oracle/exl3_oracle.py` (run with the Python engine on a CUDA box), real checkpoint layers bit for +bit against the Python engine, two whole layers of RED-SNOW-5.3-Flash (a KDA and a DSA layer, all 288 experts under +real routing) included. `tf-cuda-test exl3-tp2 RANK HOST:PORT DIR` runs the GLM split rules on two ranks over NCCL +(output-column cuts, input-row cuts with fp32 partials summed rank 0 first, and the experts' intermediate cut) and +checks every rank's bytes and the sums the same way. + +No family serves EXL3 natively yet: a family's engine takes these pieces through `qlinear` and `exl3_experts`. diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index b70dfe431f..dafa9544b7 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -41,6 +41,8 @@ const kernels = [_]Kernel{ .{ .name = "qmmf", .flags = &.{"-O3"} }, // nvfp4/qmmf.cu's device code (FP8G, FP4), tensorfold_nvfp4_v3 .{ .name = "prompt16", .flags = &.{"-O3"} }, // nvfp4/prompt.cu's FP4 and FP8 tile 4, tensorfold_nvfp4_prompt_v1 .{ .name = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts + .{ .name = "exl3_linear", .flags = &.{ "-O3", "--expt-relaxed-constexpr" } }, // exl3/linear.cu's device code, tensorfold_exl3_linear_v5 + .{ .name = "exl3_experts", .flags = &.{ "-O3", "--expt-relaxed-constexpr" } }, // exl3/experts.cu's device code, tensorfold_exl3_experts .{ .name = "torch_argmax", .src = "torch_ops/argmax", .flags = torch_ops }, .{ .name = "torch_topk", .src = "torch_ops/topk", .flags = torch_ops }, .{ .name = "torch_pointwise", .src = "torch_ops/pointwise", .flags = torch_ops }, diff --git a/zig/kernels/cuda/exl3_decode.cuh b/zig/kernels/cuda/exl3_decode.cuh new file mode 100644 index 0000000000..8b84470827 --- /dev/null +++ b/zig/kernels/cuda/exl3_decode.cuh @@ -0,0 +1,141 @@ +// EXL3 tiles: lane L decodes values 8L..8L+7 straight into its two mma B fragments; port of tensorfold's decode.cuh. +#pragma once + +#include +#include + +namespace tf_exl3 { + +enum Codebook : int { CB_3INST = 0, CB_MCG = 1, CB_MUL1 = 2 }; + +// 32-bit words of a tile: 256 values of K2 / 2 bits. +template +__host__ __device__ constexpr int tile_words() { + return 4 * K2; +} + +// E(p): one past the last stream bit of value p's 16-bit window (exl3_format.zig's streamEnds). +template +__host__ __device__ constexpr int stream_end(int p) { + return (K2 & 1) ? ((p + 1) * K2 - ((p + 1) & 1)) / 2 : (p + 1) * (K2 / 2); +} + +// Words a lane reads for its 8 values: the most any lane's windows span (2, or 3 for 3.5 bits and 5 to 8 bits). +template +__host__ __device__ constexpr int lane_words() { + int most = 0; + for (int l = 0; l < 32; ++l) { + const int first = (stream_end(8 * l) - 16 + 128 * K2) % 32; + const int need = first + stream_end(8 * l + 7) - stream_end(8 * l) + 16; + most = need > most ? need : most; + } + return (most + 31) / 32; +} + +// The index of the first word lane `lane` reads, and the bit offset of its first window in that word. +template +__device__ __forceinline__ void lane_start(int lane, int& word, int& offset) { + const int first = 4 * lane * K2 + K2 / 2 - 16 + 128 * K2; // E(8 lane) - 16, made non-negative + word = (first >> 5) % tile_words(); + offset = first & 31; +} + +// The same from global memory through the read-only cache. +template +__device__ __forceinline__ void ldg_lane_words(const uint32_t* tile, int lane, uint32_t (&w)[lane_words()]) { + int word, offset; + lane_start(lane, word, offset); +#pragma unroll + for (int i = 0; i < lane_words(); ++i) w[i] = __ldg(tile + (word + i) % tile_words()); +} + +// Bits [d, d + 16) of the 96-bit big-endian stream hi:mid:lo (d a compile-time constant after unrolling). +__device__ __forceinline__ uint32_t window16(uint32_t hi, uint32_t mid, uint32_t lo, int d) { + if (d <= 16) return (hi >> (16 - d)) & 0xffffu; + if (d < 32) return __funnelshift_l(mid, hi, d) >> 16; + if (d <= 48) return (mid >> (48 - d)) & 0xffffu; + return __funnelshift_l(lo, mid, d - 32) >> 16; +} + +// The 16-bit states of this lane's 8 values (8 lane + j, j = 0..7) from its words. +template +__device__ __forceinline__ void lane_states(const uint32_t (&w)[lane_words()], int lane, uint32_t (&s)[8]) { + int word, offset; + lane_start(lane, word, offset); + const uint32_t c = lane_words() > 2 ? w[lane_words() > 2 ? 2 : 0] : 0u; + const uint32_t hi = __funnelshift_l(w[1], w[0], offset); // the stream from the lane's first window on + const uint32_t mid = __funnelshift_l(c, w[1], offset); + const uint32_t lo = c << offset; +#pragma unroll + for (int j = 0; j < 8; ++j) s[j] = window16(hi, mid, lo, stream_end(j) - stream_end(0)); +} + +// Two codebook values from two 16-bit states, as a half2 in a uint32 (the first state in the low half). +template +__device__ __forceinline__ uint32_t decode2(uint32_t s0, uint32_t s1) { + if constexpr (CB == CB_MUL1) { + const uint32_t x0 = s0 * 0x83DCD12Du, x1 = s1 * 0x83DCD12Du; + const uint32_t h0 = __dp4a(x0, 0x01010101u, 0x6400u); // fp16 bits of 1024 + the byte sum + const uint32_t h1 = __dp4a(x1, 0x01010101u, 0x6400u); + const uint32_t hh = (h0 & 0xffffu) | (h1 << 16); + const half2 r = __hfma2(*reinterpret_cast(&hh), __half2half2(__ushort_as_half(0x1eee)), + __half2half2(__ushort_as_half(0xc931))); + return *reinterpret_cast(&r); + } else { + uint32_t x0, x1; + if constexpr (CB == CB_MCG) { + x0 = s0 * 0xCBAC1FEDu; + x1 = s1 * 0xCBAC1FEDu; + } else { + x0 = s0 * 89226354u + 64248484u; + x1 = s1 * 89226354u + 64248484u; + } + x0 = (x0 & 0x8FFF8FFFu) ^ 0x3B603B60u; + x1 = (x1 & 0x8FFF8FFFu) ^ 0x3B603B60u; + const uint32_t lo = __byte_perm(x0, x1, 0x5410); + const uint32_t hi = __byte_perm(x0, x1, 0x7632); + const half2 r = __hadd2(*reinterpret_cast(&lo), *reinterpret_cast(&hi)); + return *reinterpret_cast(&r); + } +} + +// This lane's 8 values of a tile as the B fragments of the tile's two n8 halves (columns 0-7 in b0, 8-15 in b1). +template +__device__ __forceinline__ void decode_lane(const uint32_t (&w)[lane_words()], int lane, uint32_t (&b0)[2], + uint32_t (&b1)[2]) { + uint32_t s[8]; + lane_states(w, lane, s); + b0[0] = decode2(s[0], s[1]); + b0[1] = decode2(s[2], s[3]); + b1[0] = decode2(s[4], s[5]); + b1[1] = decode2(s[6], s[7]); +} + +// mma.m16n8k16, fp16 inputs, fp32 accumulators: d += a @ b. +__device__ __forceinline__ void mma16816(float (&d)[4], const uint32_t (&a)[4], const uint32_t (&b)[2]) { + asm volatile("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"(b[0]), "r"(b[1])); +} + +// Unscaled Walsh-Hadamard transform of 128 fp32 values, 4 a lane, in a fixed butterfly order (multiply by 1/sqrt(128)). +__device__ __forceinline__ void fwht128(float (&v)[4], int lane) { + const float a = v[0] + v[1], b = v[0] - v[1], c = v[2] + v[3], d = v[2] - v[3]; + v[0] = a + c; + v[1] = b + d; + v[2] = a - c; + v[3] = b - d; +#pragma unroll + for (int m = 1; m < 32; m <<= 1) { +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float o = __shfl_xor_sync(0xffffffffu, v[j], m); + v[j] = (lane & m) ? o - v[j] : v[j] + o; + } + } +} + +constexpr float HAD_SCALE = 0.08838834764831845f; // 1 / sqrt(128) + +} // namespace tf_exl3 diff --git a/zig/kernels/cuda/exl3_experts.cu b/zig/kernels/cuda/exl3_experts.cu new file mode 100644 index 0000000000..8147a68604 --- /dev/null +++ b/zig/kernels/cuda/exl3_experts.cu @@ -0,0 +1,368 @@ +// EXL3 routed experts: grouping, grouped GEMMs, epilogues; port of tensorfold's experts.cu (ExLlamaV3, Turboderp). + +#include +#include +#include + +#include "exl3_decode.cuh" + +namespace tf_exl3 { + +// One projection's grouped arguments; the fields must stay in launch.zig add order. +struct GroupedArgs { + const half* x0; // fp16 [P, K], the rotated gate (or down) inputs + const half* x1; // fp16 [P, K], the rotated up inputs + const long long* tp0; // int64 [E], one trellis pointer per expert + const long long* tp1; + const int* k2_0; // int32 [E], half-bits a value + const int* k2_1; + const int* uids; // int32 [nexp_max], distinct routed experts in id order + const int* ucount; // int32 [1] + const int* members; // int32 [nexp_max, maxm], row * 32 + slot, -1 after the last member + float* z; // fp32 partials, [(mat * sk + split) * P + p] * N + col + int k, n, p, sk, maxm, slots, nexp_max, mats, ns; +}; + +// One warp's k tiles [kt0, kt0 + nkt) of an expert matrix into acc: per-lane loads and the shared decode (same bits). +template +__device__ __forceinline__ void warp_tiles(const uint32_t* T, int ntiles, int kt0, int nkt, int nt0, const half* x0, + const half* x1, bool ok0, bool ok1, int lane, float (&acc)[NT][2][4]) { + constexpr int TW = tile_words(), LW = lane_words(); + const size_t kstride = (size_t)ntiles * TW; + for (int i = 0; i < nkt; ++i) { + const uint32_t* t0 = T + (size_t)(kt0 + i) * kstride + (size_t)nt0 * TW; + uint32_t w[NT][LW]; +#pragma unroll + for (int j = 0; j < NT; ++j) ldg_lane_words(t0 + j * TW, lane, w[j]); + const int k = (kt0 + i) * 16; + const uint32_t a[4] = {ok0 ? *reinterpret_cast(x0 + k) : 0u, + ok1 ? *reinterpret_cast(x1 + k) : 0u, + ok0 ? *reinterpret_cast(x0 + k + 8) : 0u, + ok1 ? *reinterpret_cast(x1 + k + 8) : 0u}; +#pragma unroll + for (int j = 0; j < NT; ++j) { + uint32_t b0[2], b1[2]; + decode_lane(w[j], lane, b0, b1); + mma16816(acc[j][0], a, b0); + mma16816(acc[j][1], a, b1); + } + } +} + +// Program (expert u, n block, split, member tile): 16 members times W_q over the split's K range, summed in warp order. +template +__device__ __forceinline__ void grouped_body(const GroupedArgs a, int u, int nblock, int zz) { + if (u >= a.ucount[0]) return; + const int MT = (a.maxm + 15) / 16; + const int mtile = zz % MT; + const int split = (zz / MT) % a.sk; + const int mat = zz / MT / a.sk; + const half* X = mat ? a.x1 : a.x0; + const int e = a.uids[u]; + const uint32_t* T = reinterpret_cast(mat ? a.tp1[e] : a.tp0[e]); + const int k2 = mat ? a.k2_1[e] : a.k2_0[e]; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int g = lane >> 2, t = lane & 3; + const int NTILES = a.ns > 0 ? a.ns : (a.n >> 4); // row stride of the trellis, in tiles + + __shared__ int rows_sh[16]; + if (threadIdx.x < 16) { + const int m = mtile * 16 + threadIdx.x; + const int code = m < a.maxm ? a.members[u * a.maxm + m] : -1; + rows_sh[threadIdx.x] = code >= 0 ? (code >> 5) * a.slots + (code & 31) : -1; + } + __syncthreads(); + if (rows_sh[0] < 0) return; // members come first, so this tile is empty + const int r0 = rows_sh[g], r1 = rows_sh[g + 8]; + const half* x0 = X + (size_t)(r0 < 0 ? 0 : r0) * a.k + 2 * t; + const half* x1 = X + (size_t)(r1 < 0 ? 0 : r1) * a.k + 2 * t; + const int per_warp = (a.k >> 4) / a.sk / W; + const int kt0 = split * (per_warp * W) + warp * per_warp; + const int nt0 = nblock * NT; + + float acc[NT][2][4]; +#pragma unroll + for (int i = 0; i < NT; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[i][h][c] = 0.f; + + __shared__ __align__(16) float red[W][16][NT * 16]; + switch (k2) { +#define TF_EXL3X_CASE(K2_) \ + case K2_: \ + warp_tiles(T, NTILES, kt0, per_warp, nt0, x0, x1, r0 >= 0, r1 >= 0, lane, acc); \ + break; + TF_EXL3X_CASE(2) + TF_EXL3X_CASE(3) + TF_EXL3X_CASE(4) + TF_EXL3X_CASE(5) + TF_EXL3X_CASE(6) + TF_EXL3X_CASE(7) + TF_EXL3X_CASE(8) + TF_EXL3X_CASE(9) + TF_EXL3X_CASE(10) + TF_EXL3X_CASE(11) + TF_EXL3X_CASE(12) + TF_EXL3X_CASE(13) + TF_EXL3X_CASE(14) + TF_EXL3X_CASE(15) + TF_EXL3X_CASE(16) +#undef TF_EXL3X_CASE + default: + __trap(); + } + +#pragma unroll + for (int i = 0; i < NT; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int col = i * 16 + h * 8 + 2 * t; + red[warp][g][col] = acc[i][h][0]; + red[warp][g][col + 1] = acc[i][h][1]; + red[warp][g + 8][col] = acc[i][h][2]; + red[warp][g + 8][col + 1] = acc[i][h][3]; + } + __syncthreads(); + for (int idx = threadIdx.x; idx < 16 * NT * 16; idx += W * 32) { + const int row = idx / (NT * 16), col = idx % (NT * 16); + const int r = rows_sh[row]; + if (r < 0) continue; + float s = red[0][row][col]; +#pragma unroll + for (int w = 1; w < W; ++w) s += red[w][row][col]; + a.z[(((size_t)mat * a.sk + split) * a.p + r) * a.n + nt0 * 16 + col] = s; + } +} + +} // namespace tf_exl3 + +// --------------------------------------------------------------------------------------------------------------- + +extern "C" __global__ void tf_exl3_grouped_0(tf_exl3::GroupedArgs a) { // 3inst, any width + tf_exl3::grouped_body<0, 8, 4>(a, blockIdx.x, blockIdx.y, blockIdx.z); +} + +extern "C" __global__ void tf_exl3_grouped_1(tf_exl3::GroupedArgs a) { // mcg, any width + tf_exl3::grouped_body<1, 8, 4>(a, blockIdx.x, blockIdx.y, blockIdx.z); +} + +extern "C" __global__ void tf_exl3_grouped_2(tf_exl3::GroupedArgs a) { // mul1, any width (RED-SNOW's experts) + tf_exl3::grouped_body<2, 8, 4>(a, blockIdx.x, blockIdx.y, blockIdx.z); +} + +namespace tf_exl3 { + +// Grouping in one block: distinct experts (< E) in id order, members row * 32 + slot in row order, -1 after the last. +constexpr int GROUP_THREADS = 1024; +constexpr int GROUP_PER_THREAD = 4; + +struct GroupArgs { + const int* pick; // int32 [rows * slots] + int* uids; // int32 [nexp_max] + int* ucount; // int32 [1] + int* members; // int32 [nexp_max, maxm] + int rows, slots, experts, maxm; +}; + +extern "C" __global__ void __launch_bounds__(GROUP_THREADS) tf_exl3_group(const GroupArgs a) { + extern __shared__ int sh_pick[]; + __shared__ int warp_tot[GROUP_THREADS / 32]; + const int n = a.rows * a.slots; + for (int i = threadIdx.x; i < n; i += GROUP_THREADS) sh_pick[i] = a.pick[i]; + __syncthreads(); + int cnt[GROUP_PER_THREAD]; + int used = 0; +#pragma unroll + for (int q = 0; q < GROUP_PER_THREAD; ++q) { + const int e = threadIdx.x * GROUP_PER_THREAD + q; + int c = 0; + if (e < a.experts) + for (int i = 0; i < n; ++i) c += sh_pick[i] == e; + cnt[q] = c; + used += c > 0; + } + // exclusive scan of `used` over threads + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + int inc = used; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) { + int v = __shfl_up_sync(0xffffffffu, inc, o); + if (lane >= o) inc += v; + } + if (lane == 31) warp_tot[warp] = inc; + __syncthreads(); + if (warp == 0) { + int v = warp_tot[lane]; + int s = v; + for (int o = 1; o < 32; o <<= 1) { + int x = __shfl_up_sync(0xffffffffu, s, o); + if (lane >= o) s += x; + } + warp_tot[lane] = s - v; // exclusive per warp + if (lane == 31) a.ucount[0] = s; + } + __syncthreads(); + int place = warp_tot[warp] + inc - used; +#pragma unroll + for (int q = 0; q < GROUP_PER_THREAD; ++q) { + if (cnt[q] == 0) continue; + const int e = threadIdx.x * GROUP_PER_THREAD + q; + a.uids[place] = e; + int j = 0; + for (int i = 0; i < n && j < a.maxm; ++i) + if (sh_pick[i] == e) a.members[place * a.maxm + j++] = (i / a.slots) * 32 + (i % a.slots); + for (; j < a.maxm; ++j) a.members[place * a.maxm + j] = -1; + ++place; + } +} + +// Program (slot, 128-block of K, matrix): Xh = fp16((x * suh) @ H) for gate and up of every routed slot (pick < E). +struct RotInArgs { + const void* x; // bf16 or fp16 [rows, x_stride] + int x_stride; + int x_dtype; // 0: fp16, 1: bf16 (exl3.zig DType) + const int* pick; // int32 [rows * slots] + const half* suh0; // fp16 [E, K], gate + const half* suh1; // fp16 [E, K], up + half* out0; // fp16 [P, K] + half* out1; + int k, slots, experts; +}; + +extern "C" __global__ void tf_exl3_experts_rot_in(const RotInArgs a) { + const int p = blockIdx.x, blk = blockIdx.y, mat = blockIdx.z; + const int e = a.pick[p]; + if (e < 0 || e >= a.experts) return; + const int lane = threadIdx.x; + const int row = p / a.slots; + const half* suh = (mat ? a.suh1 : a.suh0) + (size_t)e * a.k + blk * 128 + 4 * lane; + const uint16_t* xr = reinterpret_cast(a.x) + (size_t)row * a.x_stride + blk * 128 + 4 * lane; + float v[4]; + if (a.x_dtype == 1) { + const __nv_bfloat16* xv = reinterpret_cast(xr); +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = __bfloat162float(xv[j]) * __half2float(suh[j]); + } else { + const half* xv = reinterpret_cast(xr); +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = __half2float(xv[j]) * __half2float(suh[j]); + } + fwht128(v, lane); + half* o = (mat ? a.out1 : a.out0) + (size_t)p * a.k + blk * 128 + 4 * lane; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); +} + +__device__ __forceinline__ float bf16r(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +// Program (slot, 128-block of the width): splits summed in order, rotated, * svh, SwiGLU (0: bf16, 1: fp32), then Xd. +struct GateUpArgs { + const float* z; // [(mat * sk + split) * P + p] * N + col + const int* pick; // int32 [P] + const half* svh_g; // fp16 [E, N] + const half* svh_u; // fp16 [E, N] + const half* suh_d; // fp16 [E, N] + half* xd; // fp16 [P, N] + int p, n, sk, experts; + float limit; + int act_mode; +}; + +extern "C" __global__ void tf_exl3_gateup_epilogue(const GateUpArgs a) { + const int p = blockIdx.x, blk = blockIdx.y; + const int e = a.pick[p]; + if (e < 0 || e >= a.experts) return; + const int lane = threadIdx.x; + const int n = blk * 128 + 4 * lane; + float gv[4], uv[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float sg = 0.f, su = 0.f; + for (int s = 0; s < a.sk; ++s) { + sg += a.z[((size_t)(0 * a.sk + s) * a.p + p) * a.n + n + j]; + su += a.z[((size_t)(1 * a.sk + s) * a.p + p) * a.n + n + j]; + } + gv[j] = sg; + uv[j] = su; + } + fwht128(gv, lane); + fwht128(uv, lane); + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float act; + if (a.act_mode == 0) { + float gg = fminf(bf16r(gv[j] * HAD_SCALE * __half2float(a.svh_g[(size_t)e * a.n + n + j])), a.limit); + float uu = fminf(fmaxf(bf16r(uv[j] * HAD_SCALE * __half2float(a.svh_u[(size_t)e * a.n + n + j])), + -a.limit), a.limit); + act = bf16r(bf16r(gg / (1.f + expf(-gg))) * uu); + } else { + float gg = fminf(gv[j] * HAD_SCALE * __half2float(a.svh_g[(size_t)e * a.n + n + j]), a.limit); + float uu = fminf(fmaxf(uv[j] * HAD_SCALE * __half2float(a.svh_u[(size_t)e * a.n + n + j]), -a.limit), + a.limit); + act = gg / (1.f + expf(-gg)) * uu; + } + v[j] = act * __half2float(a.suh_d[(size_t)e * a.n + n + j]); + } + fwht128(v, lane); + half* o = a.xd + (size_t)p * a.n + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); +} + +// Program (member row, 128-block of the width): down epilogue and combine in one launch, the same bits as the two. +struct DownArgs { + const float* z; // [split * P + p] * D + col + const int* pick; // int32 [P] + const half* svh_d; // fp16 [E, D] + float* y; // fp32 [P, D] + const float* wts; // fp32 [rows, slots] + float* out; // fp32 [rows, D] + int p, d, sk, experts, slots; +}; + +extern "C" __global__ void tf_exl3_down_combine(const DownArgs a) { + __shared__ float4 part[32][32]; // [slot][lane]: the slot's 4 outputs of the lane + const int r = blockIdx.x, blk = blockIdx.y; + const int k = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int n = blk * 128 + 4 * lane; + const int p = r * a.slots + k; + const int e = a.pick[p]; + float o[4]; + if (e >= 0 && e < a.experts) { + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float s = 0.f; + for (int q = 0; q < a.sk; ++q) s += a.z[((size_t)q * a.p + p) * a.d + n + j]; + v[j] = s; + } + fwht128(v, lane); +#pragma unroll + for (int j = 0; j < 4; ++j) { + o[j] = v[j] * HAD_SCALE * __half2float(a.svh_d[(size_t)e * a.d + n + j]); + a.y[(size_t)p * a.d + n + j] = o[j]; + } + } else { +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = a.y[(size_t)p * a.d + n + j]; + } + part[k][lane] = make_float4(o[0], o[1], o[2], o[3]); + __syncthreads(); + if (k != 0) return; + float acc[4] = {0.f, 0.f, 0.f, 0.f}; + for (int q = 0; q < a.slots; ++q) { + const float w = a.wts[r * a.slots + q]; + const float4 u = part[q][lane]; + acc[0] = fmaf(w, u.x, acc[0]); + acc[1] = fmaf(w, u.y, acc[1]); + acc[2] = fmaf(w, u.z, acc[2]); + acc[3] = fmaf(w, u.w, acc[3]); + } +#pragma unroll + for (int j = 0; j < 4; ++j) a.out[(size_t)r * a.d + n + j] = acc[j]; +} + +} // namespace tf_exl3 diff --git a/zig/kernels/cuda/exl3_linear.cu b/zig/kernels/cuda/exl3_linear.cu new file mode 100644 index 0000000000..6b805e4390 --- /dev/null +++ b/zig/kernels/cuda/exl3_linear.cu @@ -0,0 +1,274 @@ +// EXL3 decode-row linear: any codebook, width, 1-128 rows, a row's bits only its own; port of tensorfold's linear.cu. +#include +#include +#include + +#include "exl3_decode.cuh" + +namespace tf_exl3 { + +enum DType : int { F16 = 0, BF16 = 1, F32 = 2 }; + +__device__ __forceinline__ int imin(int a, int b) { return a < b ? a : b; } + +struct LinearArgs { + const void* xh; // fp16 [m, k], the hadamard-rotated inputs + const uint32_t* t; // trellis words + long long stride_k; // words between k tiles + long long stride_nb; // words between 128-column blocks + const half* svh; // fp16 [n], output hadamard scales + const half* bias; // fp16 [n] or null + void* y; // output + float* z; // fp32 split partials [sk, m, n] or null + int* counters; // per (pass, block) arrival counts, zeroed before the launch + int y_dtype; + int n, sk, m, k; +}; + +struct RotInArgs { + const void* x; // f16 / bf16 / f32 [m, k] + const half* suh; // fp16 [k], input hadamard scales + half* xh; // fp16 [m, k] out + int x_dtype; + int m, k; +}; + +__device__ __forceinline__ void load4(const void* p, int dtype, size_t i, float (&v)[4]) { + if (dtype == F32) { + const float4 u = *reinterpret_cast(static_cast(p) + i); + v[0] = u.x; v[1] = u.y; v[2] = u.z; v[3] = u.w; + } else if (dtype == BF16) { + const uint2 u = *reinterpret_cast(static_cast(p) + i); + const float2 a = __bfloat1622float2(*reinterpret_cast(&u.x)); + const float2 b = __bfloat1622float2(*reinterpret_cast(&u.y)); + v[0] = a.x; v[1] = a.y; v[2] = b.x; v[3] = b.y; + } else { + const uint2 u = *reinterpret_cast(static_cast(p) + i); + const float2 a = __half22float2(*reinterpret_cast(&u.x)); + const float2 b = __half22float2(*reinterpret_cast(&u.y)); + v[0] = a.x; v[1] = a.y; v[2] = b.x; v[3] = b.y; + } +} + +__device__ __forceinline__ void store4(void* p, int dtype, size_t i, const float (&v)[4]) { + if (dtype == F32) { + *reinterpret_cast(static_cast(p) + i) = make_float4(v[0], v[1], v[2], v[3]); + } else if (dtype == BF16) { + __nv_bfloat162 a = __floats2bfloat162_rn(v[0], v[1]), b = __floats2bfloat162_rn(v[2], v[3]); + uint2 u; + u.x = *reinterpret_cast(&a); + u.y = *reinterpret_cast(&b); + *reinterpret_cast(static_cast<__nv_bfloat16*>(p) + i) = u; + } else { + half2 a = __floats2half2_rn(v[0], v[1]), b = __floats2half2_rn(v[2], v[3]); + uint2 u; + u.x = *reinterpret_cast(&a); + u.y = *reinterpret_cast(&b); + *reinterpret_cast(static_cast(p) + i) = u; + } +} + +// The finished outputs of one row's 128 columns from their fp32 sums (4 a lane): H / sqrt(128), * svh, + bias. +__device__ __forceinline__ void finish(float (&v)[4], int lane, const half* svh, const half* bias, int col) { + fwht128(v, lane); +#pragma unroll + for (int j = 0; j < 4; ++j) { + v[j] = v[j] * HAD_SCALE * __half2float(__ldg(svh + col + j)); + // Speculative loads must have a valid address even when the split-K result has no bias. + const half bv = __ldg((bias ? bias : svh) + col + j); + if (bias) v[j] += __half2float(bv); + } +} + +// A k step's words for one lane: the step is 8 tiles back to back, each lane takes the same slice of every tile. +template +__device__ __forceinline__ void load_step(const uint32_t* step, int lane, uint32_t (&raw)[8 * lane_words()]) { + constexpr int TW = tile_words(), LW = lane_words(); + int word, offset; + lane_start(lane, word, offset); +#pragma unroll + for (int j = 0; j < 8; ++j) +#pragma unroll + for (int q = 0; q < LW; ++q) raw[j * LW + q] = __ldg(step + j * TW + (word + q) % TW); +} + +// Tile j's lane words from a loaded step. +template +__device__ __forceinline__ void step_lane_words(const uint32_t (&raw)[8 * lane_words()], int j, + uint32_t (&w)[lane_words()]) { +#pragma unroll + for (int q = 0; q < lane_words(); ++q) w[q] = raw[j * lane_words() + q]; +} + +// xh = fp16(((x * suh) @ H) / sqrt(128)): one warp per 128 columns, one block of 4 warps per 512. +extern "C" __global__ void __launch_bounds__(128) tf_exl3_rot_in(RotInArgs a) { + const int blk = blockIdx.x * 4 + (threadIdx.x >> 5), row = blockIdx.y, lane = threadIdx.x & 31; + if (blk * 128 >= a.k) return; + const int k = blk * 128 + 4 * lane; + float v[4], s[4]; + load4(a.x, a.x_dtype, (size_t)row * a.k + k, v); + load4(a.suh, F16, k, s); +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] *= s[j]; + fwht128(v, lane); +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] *= HAD_SCALE; + store4(a.xh, F16, (size_t)row * a.k + k, v); +} + + +// One program: (N / 128) * SK blocks, split-major like a (N / 128, SK) grid, WK warps each. +template +__device__ __forceinline__ void linear_body(LinearArgs a) { + constexpr int TW = tile_words(), LW = lane_words(); + constexpr bool PF = K2 <= 12; // the next k step's words load while this one decodes + constexpr int SR = 8 * LW; + extern __shared__ __align__(16) float red[]; // WK * RH * 128 floats + __shared__ int last; + const int M = a.m, K = a.k, N = a.n, SK = a.sk, NB = N >> 7; + const int RH = imin(M, 8); // rows of red a warp + const int nb = (int)blockIdx.x % NB, split = (int)blockIdx.x / NB; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int g = lane >> 2, t = lane & 3; + const int per_warp = (K >> 4) / SK / WK; + const int kt0 = split * (per_warp * WK) + warp * per_warp; + const uint32_t* tiles = a.t + nb * a.stride_nb; + const int col0 = nb * 128; + + for (int m0 = 0, pass = 0; m0 < M; m0 += 16, ++pass) { + const int R = imin(16, M - m0); + + float acc[8][2][4]; +#pragma unroll + for (int i = 0; i < 8; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[i][h][c] = 0.f; + + // the walk's rows for the mma, clamped so rows past the pass read inside the buffer (their outputs are dropped) + const int r0 = m0 + (g < R ? g : R - 1), r1 = m0 + (g + 8 < R ? g + 8 : R - 1); + const half* x0 = static_cast(a.xh) + (size_t)r0 * K; + const half* x1 = static_cast(a.xh) + (size_t)r1 * K; + const uint32_t* tile = tiles + (size_t)kt0 * a.stride_k; + uint32_t cur[PF ? SR : 1], nxt[PF ? SR : 1]; + if constexpr (PF) load_step(tile, lane, cur); // the weights do not wait for the previous kernel + + // x's fragments of k step kt: rows g (and g + 8): [0] [2] k 2t.., 2t + 8..; [1] [3] the same of row g + 8 + auto load_x = [&](int kt, uint32_t (&u)[4]) { + u[0] = __ldg(reinterpret_cast(x0 + kt * 16 + 2 * t)); + u[2] = __ldg(reinterpret_cast(x0 + kt * 16 + 2 * t + 8)); + u[1] = __ldg(reinterpret_cast(x1 + kt * 16 + 2 * t)); + u[3] = __ldg(reinterpret_cast(x1 + kt * 16 + 2 * t + 8)); + }; +#pragma unroll 1 + for (int i = 0; i < per_warp; ++i) { + uint32_t u[4]; + load_x(kt0 + i, u); + if constexpr (PF) { + if (i + 1 < per_warp) load_step(tile + (size_t)(i + 1) * a.stride_k, lane, nxt); + } +#pragma unroll + for (int j = 0; j < 8; ++j) { + uint32_t w[LW]; + if constexpr (PF) step_lane_words(cur, j, w); + else ldg_lane_words(tile + (size_t)i * a.stride_k + j * TW, lane, w); + uint32_t b0[2], b1[2]; + decode_lane(w, lane, b0, b1); + mma16816(acc[j][0], u, b0); + mma16816(acc[j][1], u, b1); + } + if constexpr (PF) { +#pragma unroll + for (int q = 0; q < SR; ++q) cur[q] = nxt[q]; + } + } + + // the warps' sums, added in warp order, rows 0-7 of the pass and then rows 8-15 + for (int rlo = 0; rlo < R; rlo += 8) { + const int rn = imin(R - rlo, 8); + __syncthreads(); // red is reused by every half and pass + if (g < RH) { +#pragma unroll + for (int i = 0; i < 8; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int col = i * 16 + h * 8 + 2 * t; + *reinterpret_cast(red + (warp * RH + g) * 128 + col) = + rlo ? make_float2(acc[i][h][2], acc[i][h][3]) + : make_float2(acc[i][h][0], acc[i][h][1]); + } + } + __syncthreads(); + + if (SK == 1) { + for (int r = warp; r < rn; r += WK) { + float v[4]; + const float4 u = *reinterpret_cast(red + r * 128 + 4 * lane); + v[0] = u.x; v[1] = u.y; v[2] = u.z; v[3] = u.w; +#pragma unroll + for (int w = 1; w < WK; ++w) { + const float4 q = *reinterpret_cast(red + (w * RH + r) * 128 + 4 * lane); + v[0] += q.x; v[1] += q.y; v[2] += q.z; v[3] += q.w; + } + finish(v, lane, a.svh, a.bias, col0 + 4 * lane); + store4(a.y, a.y_dtype, (size_t)(m0 + rlo + r) * N + col0 + 4 * lane, v); + } + } else { + for (int idx = threadIdx.x; idx < rn * 32; idx += WK * 32) { + const int r = idx >> 5, c = 4 * (idx & 31); + float4 s = *reinterpret_cast(red + r * 128 + c); +#pragma unroll + for (int w = 1; w < WK; ++w) { + const float4 q = *reinterpret_cast(red + (w * RH + r) * 128 + c); + s.x += q.x; s.y += q.y; s.z += q.z; s.w += q.w; + } + *reinterpret_cast(a.z + ((size_t)split * M + m0 + rlo + r) * N + col0 + c) = s; + } + } + } + if (SK > 1) { + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) last = atomicAdd(a.counters + pass * NB + nb, 1) == SK - 1; + __syncthreads(); + if (last) { + __threadfence(); + for (int r = warp; r < R; r += WK) { + const size_t at = ((size_t)m0 + r) * N + col0 + 4 * lane; + float4 s = __ldcg(reinterpret_cast(a.z + at)); + for (int q = 1; q < SK; ++q) { + const float4 u = __ldcg(reinterpret_cast(a.z + (size_t)q * M * N + at)); + s.x += u.x; s.y += u.y; s.z += u.z; s.w += u.w; + } + float v[4] = {s.x, s.y, s.z, s.w}; + finish(v, lane, a.svh, a.bias, col0 + 4 * lane); + store4(a.y, a.y_dtype, (size_t)(m0 + r) * N + col0 + 4 * lane, v); + } + if (threadIdx.x == 0) a.counters[pass * NB + nb] = 0; // every program of the block has arrived + } + } + __syncthreads(); // red is reused in the next pass + } +} + +// One extern "C" kernel per instantiation: the fatbin carries __global__ functions, the host resolves them by name. +#define TF_EXL3_LINEAR(K2, CB, WK) \ + extern "C" __global__ void __launch_bounds__((WK) * 32) \ + tf_exl3_linear_##K2##_##CB##_##WK(tf_exl3::LinearArgs a) { \ + tf_exl3::linear_body(a); \ + } + +// Every width the format serves (linear.cu TF_EXL3_ALL plus 4.5-6.5): whole bits, all codebooks; 7.5 is not a width. +#define TF_EXL3_WIDTHS(K2, CB) TF_EXL3_LINEAR(K2, CB, 4) TF_EXL3_LINEAR(K2, CB, 8) +#define TF_EXL3_CODEBOOK(CB) \ + TF_EXL3_WIDTHS(2, CB) TF_EXL3_WIDTHS(4, CB) TF_EXL3_WIDTHS(6, CB) TF_EXL3_WIDTHS(8, CB) \ + TF_EXL3_WIDTHS(10, CB) TF_EXL3_WIDTHS(12, CB) TF_EXL3_WIDTHS(14, CB) TF_EXL3_WIDTHS(16, CB) +TF_EXL3_CODEBOOK(2) // mul1 +TF_EXL3_WIDTHS(3, 2) TF_EXL3_WIDTHS(5, 2) TF_EXL3_WIDTHS(7, 2) TF_EXL3_WIDTHS(9, 2) +TF_EXL3_WIDTHS(11, 2) TF_EXL3_WIDTHS(13, 2) +TF_EXL3_CODEBOOK(1) // mcg +TF_EXL3_CODEBOOK(0) // inst3 +#undef TF_EXL3_CODEBOOK +#undef TF_EXL3_WIDTHS +} // namespace tf_exl3 diff --git a/zig/src/core/exl3_load.zig b/zig/src/core/exl3_load.zig new file mode 100644 index 0000000000..1bf23fa2d8 --- /dev/null +++ b/zig/src/core/exl3_load.zig @@ -0,0 +1,227 @@ +//! EXL3 checkpoint groups: every `{prefix}.trellis` and its parts, admitted from the safetensors headers (format.py). +const std = @import("std"); +const checkpoint = @import("checkpoint.zig"); +const st = @import("safetensors.zig"); + +pub const Codebook = enum { inst3, mcg, mul1 }; + +/// format.py's MARKERS: a scalar int32 tensor per codebook, whose value must be the magic. +pub const Marker = struct { name: []const u8, codebook: Codebook, value: u32 }; +pub const markers = [_]Marker{ + .{ .name = "mcg", .codebook = .mcg, .value = 0xCBAC1FED }, + .{ .name = "mul1", .codebook = .mul1, .value = 0x83DCD12D }, +}; + +/// One admitted group: a trellis int16 [K/16, N/16, 16 * bits], its scales, and the optional bias. +pub const Group = struct { + prefix: []const u8, + halves: u16, // 2 * bits; odd halves are the mul1 half steps + codebook: Codebook, + k: usize, + n: usize, + suh: bool, // fp16 scales; false: the packed sign words su and sv + bias: bool, +}; + +/// Counts of the admitted groups, by halves and for the routed experts' down projections. +pub const Summary = struct { + groups: usize = 0, + inst3: usize = 0, + mcg: usize = 0, + mul1: usize = 0, + suh_svh: usize = 0, + su_sv: usize = 0, + bias: usize = 0, + bits: [17]usize = @splat(0), // by halves, 2 to 16 + experts_down: [17]usize = @splat(0), // *.experts.*.down_proj by halves + + fn bump(s: *Summary, g: Group) void { + s.groups += 1; + switch (g.codebook) { + .inst3 => s.inst3 += 1, + .mcg => s.mcg += 1, + .mul1 => s.mul1 += 1, + } + if (g.suh) s.suh_svh += 1 else s.su_sv += 1; + if (g.bias) s.bias += 1; + s.bits[g.halves] += 1; + if (std.mem.endsWith(u8, g.prefix, ".down_proj") and std.mem.indexOf(u8, g.prefix, ".experts.") != null) s.experts_down[g.halves] += 1; + } +}; + +/// One group's tensors: the trellis words, its scales, and the optional bias. +pub const Parts = struct { + trellis: st.Tensor, + suh: st.Tensor, + svh: st.Tensor, + bias: ?st.Tensor = null, + suh_packed: bool = false, // su/sv sign words instead of fp16 scales +}; + +pub const Scan = struct { + gpa: std.mem.Allocator, + ck: checkpoint.Checkpoint, + groups: []Group, + summary: Summary = .{}, + bad: usize = 0, // trellis tensors refused: another format's tensors, a broken or half-written group + + pub fn deinit(s: *Scan) void { + for (s.groups) |g| s.gpa.free(g.prefix); + s.gpa.free(s.groups); + s.ck.close(); + } + + /// The group a tensor prefix names, if any. + pub fn find(s: *const Scan, prefix: []const u8) ?usize { + for (s.groups, 0..) |g, i| if (std.mem.eql(u8, g.prefix, prefix)) return i; + return null; + } + + /// One group's tensors, their layouts checked on the way out. + pub fn parts(s: *Scan, a: std.mem.Allocator, gi: usize) !Parts { + const g = s.groups[gi]; + const trellis = try s.ck.get(try std.mem.concat(a, u8, &.{ g.prefix, ".trellis" })); + if (!trellis.is(.i16, &.{ g.k / 16, g.n / 16, @as(usize, g.halves) * 8 })) return error.BadTrellis; + const su_name = if (g.suh) "suh" else "su"; + const sv_name = if (g.suh) "svh" else "sv"; + const suh = try s.ck.get(try std.mem.concat(a, u8, &.{ g.prefix, ".", su_name })); + const svh = try s.ck.get(try std.mem.concat(a, u8, &.{ g.prefix, ".", sv_name })); + const in_shape: []const usize = if (g.suh) &.{g.k} else &.{g.k / 16}; + const out_shape: []const usize = if (g.suh) &.{g.n} else &.{g.n / 16}; + if (!suh.is(.f16, in_shape) or !svh.is(.f16, out_shape)) return error.BadScales; + var bias: ?st.Tensor = null; + if (g.bias) { + bias = try s.ck.get(try std.mem.concat(a, u8, &.{ g.prefix, ".bias" })); + if (!bias.?.is(.f16, &.{g.n})) return error.BadBias; + } + return .{ .trellis = trellis, .suh = suh, .svh = svh, .bias = bias, .suh_packed = !g.suh }; + } +}; + +/// Every trellis tensor of `dir`'s shards, admitted or counted as refused. +pub fn scan(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Scan { + var ck = try checkpoint.Checkpoint.openModel(gpa, io, dir); + errdefer ck.close(); + var groups: std.ArrayList(Group) = .empty; + errdefer groups.deinit(gpa); + var summary: Summary = .{}; + var bad: usize = 0; + var arena_state = std.heap.ArenaAllocator.init(gpa); + defer arena_state.deinit(); + const a = arena_state.allocator(); + for (ck.files.items) |*f| for (f.names.keys(), f.names.values()) |name, e| { + if (!std.mem.endsWith(u8, name, ".trellis")) continue; + const prefix = name[0 .. name.len - ".trellis".len]; + if (admit(a, &ck, prefix, e)) |g| { + try groups.append(gpa, .{ .prefix = try gpa.dupe(u8, g.prefix), .halves = g.halves, .codebook = g.codebook, .k = g.k, .n = g.n, .suh = g.suh, .bias = g.bias }); + summary.bump(g); + } else |_| { + bad += 1; + } + }; + return .{ .gpa = gpa, .ck = ck, .groups = try groups.toOwnedSlice(gpa), .summary = summary, .bad = bad }; +} + +/// parse_group's admission: the trellis layout, the codebook marker's value, and the group's parts. +fn admit(a: std.mem.Allocator, ck: *checkpoint.Checkpoint, prefix: []const u8, e: st.Entry) !Group { + if (e.dtype != .i16 or e.rank != 3) return error.Refuse; + const halves = e.shape[2] / 8; // the trellis's last dim is 16 * bits + if (e.shape[2] % 8 != 0 or halves < 2 or halves > 16) return error.Refuse; + const k = e.shape[0] * 16; + const n = e.shape[1] * 16; + if (k % 128 != 0 or n % 128 != 0) return error.Refuse; + var codebook: Codebook = .inst3; + var marker: ?Marker = null; + for (markers) |m| if (ck.has(try std.mem.concat(a, u8, &.{ prefix, ".", m.name }))) { + if (marker != null) return error.Refuse; // both codebooks' markers: not a loadable group + marker = m; + codebook = m.codebook; + }; + if (halves % 2 == 1 and codebook != .mul1) return error.Refuse; // odd halves are the mul1 half steps + if (marker) |m| { + const t = try ck.get(try std.mem.concat(a, u8, &.{ prefix, ".", m.name })); + if (t.bytes.len != 4 or std.mem.readInt(u32, t.bytes[0..4], .little) != m.value) return error.Refuse; + } + const su = try scale(a, ck, prefix, "suh", k, "su", k / 16); + const sv = try scale(a, ck, prefix, "svh", n, "sv", n / 16); + if (su == null or sv == null or su.? != sv.?) return error.Refuse; // one fp16, one packed: a half-written group + const bias = try scale(a, ck, prefix, "bias", n, "", 0); + return .{ .prefix = prefix, .halves = @intCast(halves), .codebook = codebook, .k = k, .n = n, .suh = su.? == .f16, .bias = bias != null }; +} + +const Scales = enum { f16, words }; // fp16 per row, or the packed sign words + +/// One part's admission: the fp16 name, or the packed-words fallback; a missing input scale (or fallback) is refused. +fn scale(a: std.mem.Allocator, ck: *checkpoint.Checkpoint, prefix: []const u8, name: []const u8, len: usize, packed_name: []const u8, packed_len: usize) !?Scales { + const full = try std.mem.concat(a, u8, &.{ prefix, ".", name }); + if (!ck.has(full)) { + if (packed_name.len == 0) return null; // an optional part + const packed_full = try std.mem.concat(a, u8, &.{ prefix, ".", packed_name }); + if (!ck.has(packed_full)) return error.Refuse; + const p = try ck.get(packed_full); + if (p.dtype != .i16 or p.numel() != packed_len) return error.Refuse; + return .words; + } + const t = try ck.get(full); + if (t.dtype != .f16 or t.numel() != len) return error.Refuse; + return .f16; +} + +test "a hand-built safetensors file admits its one group and refuses the scaleless one" { + const io = std.testing.io; + const gpa = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + // one full group (trellis [8, 8, 32] i16, fp16 scales, the mul1 marker) and one without scales + const trellis: [4096]u8 = @splat(0); + const scales: [256]u8 = @splat(0); + const marker = [_]u8{ 0x2D, 0xD1, 0xDC, 0x83 }; // 0x83DCD12D little-endian + const Entry = st.Entry; + const entries = [_]Entry{ + .{ .dtype = .i16, .rank = 3, .shape = .{ 8, 8, 32, 1, 1 }, .begin = 0, .end = 4096 }, + .{ .dtype = .f16, .rank = 1, .shape = .{ 128, 1, 1, 1, 1 }, .begin = 4096, .end = 4352 }, + .{ .dtype = .f16, .rank = 1, .shape = .{ 128, 1, 1, 1, 1 }, .begin = 4352, .end = 4608 }, + .{ .dtype = .i32, .rank = 1, .shape = .{ 1, 1, 1, 1, 1 }, .begin = 4608, .end = 4612 }, + .{ .dtype = .i16, .rank = 3, .shape = .{ 8, 8, 32, 1, 1 }, .begin = 4612, .end = 4612 + 4096 }, + }; + const names = [_][]const u8{ "a.trellis", "a.suh", "a.svh", "a.mul1", "b.trellis" }; + const dts = [_][]const u8{ "I16", "F16", "F16", "I32", "I16" }; + const data = try std.mem.concat(gpa, u8, &.{ &trellis, &scales, &scales, &marker, &trellis }); + defer gpa.free(data); + var header: std.ArrayListUnmanaged(u8) = .empty; + defer header.deinit(gpa); + try header.appendSlice(gpa, "{"); + for (names, entries, dts, 0..) |name, e, dt, i| { + if (i != 0) try header.appendSlice(gpa, ","); + try header.print(gpa, "\"{s}\":{{\"dtype\":\"{s}\",\"shape\":[", .{ name, dt }); + for (e.shape[0..e.rank], 0..) |d, j| { + if (j != 0) try header.appendSlice(gpa, ","); + try header.print(gpa, "{d}", .{d}); + } + try header.print(gpa, "],\"data_offsets\":[{d},{d}]}}", .{ e.begin, e.end }); + } + try header.appendSlice(gpa, "}"); + var file: std.ArrayListUnmanaged(u8) = .empty; + defer file.deinit(gpa); + try file.appendNTimes(gpa, 0, 8); + std.mem.writeInt(u64, file.items[0..8], header.items.len, .little); + try file.appendSlice(gpa, header.items); + try file.appendSlice(gpa, data); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = file.items }); + const dir = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(dir); + + var s = try scan(gpa, io, dir); + defer s.deinit(); + try std.testing.expectEqual(@as(usize, 1), s.bad); + try std.testing.expectEqual(@as(usize, 1), s.summary.mul1); + try std.testing.expectEqual(@as(usize, 1), s.summary.suh_svh); + try std.testing.expectEqual(@as(usize, 1), s.summary.bits[4]); // a 32-word trellis is 4 halves + const gi = s.find("a") orelse return error.TestUnexpectedResult; + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const p = try s.parts(arena.allocator(), gi); + try std.testing.expectEqual(@as(usize, 128), p.trellis.dim(0) * 16); + try std.testing.expect(!p.suh_packed); + try std.testing.expectEqual(@as(usize, 128), p.svh.numel()); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 736248a93d..a901666ffd 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -4,6 +4,7 @@ pub const safetensors = @import("safetensors.zig"); pub const exl3_format = @import("exl3_format.zig"); pub const exlr8 = @import("exlr8.zig"); pub const checkpoint = @import("checkpoint.zig"); +pub const exl3_load = @import("exl3_load.zig"); pub const direct_io = @import("direct_io.zig"); pub const Checkpoint = checkpoint.Checkpoint; pub const draft_depth = @import("draft_depth.zig"); @@ -31,6 +32,7 @@ test { _ = act_ref; _ = exl3_format; _ = exlr8; + _ = exl3_load; _ = @import("exl3_rect_test.zig"); _ = checkpoint; _ = direct_io; diff --git a/zig/src/cuda/exl3.zig b/zig/src/cuda/exl3.zig new file mode 100644 index 0000000000..2a42bce18a --- /dev/null +++ b/zig/src/cuda/exl3.zig @@ -0,0 +1,274 @@ +//! The EXL3 decode-row linear (exl3/linear.cu's device code): the plan, a layer's tensors, and the two kernels. +const std = @import("std"); +const driver = @import("driver.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); +const kernels = @import("kernels.zig"); +const memory = @import("memory.zig"); +const Driver = driver.Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +/// x's and y's element types (linear.cu DType). +pub const DType = enum(i32) { f16 = 0, bf16 = 1, f32 = 2 }; + +/// The three codebooks (exl3_format.zig's, numbered as the kernel's). +pub const Codebook = enum(u3) { inst3 = 0, mcg = 1, mul1 = 2 }; + +/// One compiled kernel: 2 * bits halves, a codebook, and the warps one block runs. +pub const Variant = struct { halves: u32, codebook: Codebook, warps: u32 }; + +/// Every instantiation of exl3_linear.cu: every width (whole bits all codebooks, the mul1 half steps), 4 and 8 warps. +pub const variants = blk: { + var vs: [60]Variant = undefined; + var i: usize = 0; + for (0..2) |w| { + const warps: u32 = if (w == 0) 4 else 8; + for ([_]u32{ 2, 1, 0 }) |cbi| { + const cb: Codebook = @fromBackingInt(@intCast(cbi)); + var halves: u32 = 2; + while (halves <= 16) : (halves += 2) { + vs[i] = .{ .halves = halves, .codebook = cb, .warps = warps }; + i += 1; + } + if (cb != .mul1) continue; // the half steps exist only with mul1 + halves = 3; + while (halves <= 13) : (halves += 2) { // the mul1 half steps, 1.5 to 6.5 bits (format.py's BITS) + vs[i] = .{ .halves = halves, .codebook = cb, .warps = warps }; + i += 1; + } + } + } + break :blk vs; +}; + +/// The kernel's extern "C" name for an instantiation (linear.cu's TF_EXL3_LINEAR macro). +pub fn symbol(comptime halves: u32, comptime codebook: Codebook, comptime warps: u32) [:0]const u8 { + return std.fmt.comptimePrint("tf_exl3_linear_{d}_{d}_{d}", .{ halves, @backingInt(codebook), warps }); +} + +pub const rot_in_symbol = "tf_exl3_rot_in"; + +pub fn variantIndex(halves: u32, codebook: Codebook, warps: u32) ?usize { + for (variants, 0..) |v, i| { + if (v.halves == halves and v.codebook == codebook and v.warps == warps) return i; + } + return null; +} + +/// What one shape launches, from the shape alone so a row's bits never depend on the row count (linear.py plan). +pub fn plan(k: usize, n: usize, blocks: usize, min_tiles: usize) struct { sk: u32, wk: u32 } { + const kt = k / 16; + const nb = n / 128; + var sk: usize = 1; + var wk: usize = if (nb >= 64) 8 else 4; + if (nb < 64) { + while (nb * sk < blocks and sk < 64) sk *= 2; + } + while ((wk > 4 or sk > 1) and (kt % (sk * wk) != 0 or kt / (sk * wk) < min_tiles)) { + if (wk > 4) wk = 4 else if (sk > 1) sk /= 2 else break; + } + return .{ .sk = @intCast(sk), .wk = @intCast(wk) }; +} + +/// Words between k tiles and 128-column blocks of one layout: strips [N/128, K/16, 8, W] or stored [K/16, N/16, W]. +pub fn strides(k: usize, n: usize, halves: usize, strips: bool) struct { k: usize, nb: usize } { + const words = halves * 4; + return if (strips) + .{ .k = 8 * words, .nb = (k / 16) * 8 * words } + else + .{ .k = (n / 16) * words, .nb = 8 * words }; +} + +/// One layer's tensors on the device: the trellis words, the Hadamard scales, the optional bias, shape and layout. +pub const Weight = struct { + trellis: u64, // u32 words, `stride_k` and `stride_nb` apart + suh: u64, // fp16 [k] + svh: u64, // fp16 [n] + bias: u64 = 0, // fp16 [n], 0 when the layer has none + k: u32, + n: u32, + halves: u32, // 2 * bits, odd for the mul1 half steps + codebook: Codebook, + stride_k: u64, + stride_nb: u64, +}; + +/// The module and its kernels. +pub const Linear = struct { + mod: Module, + fns: [variants.len]Function, + rot_in: Function, + + pub fn load(d: *const Driver) !Linear { + if (!kernels.available) return error.BuiltWithoutKernels; + var l: Linear = .{ .mod = try Module.load(d, kernels.exl3_linear), .fns = undefined, .rot_in = undefined }; + errdefer l.mod.unload(); + l.rot_in = try l.mod.function(rot_in_symbol); + inline for (0..variants.len) |i| { + l.fns[i] = try l.mod.function(symbol(variants[i].halves, variants[i].codebook, variants[i].warps)); + } + return l; + } + + pub fn unload(l: *Linear) void { + l.mod.unload(); + } + + /// Shared bytes one call needs: WK * min(m, 8) * 128 floats of warp sums, at most 32 KB. + pub fn sharedBytes(w: Weight, m: usize) u32 { + const p = plan(w.k, w.n, 192, 8); + return @intCast(p.wk * @min(m, 8) * 128 * 4); + } + + /// Split partials one call needs (fp32 [sk, m, n]) when the plan splits K, else 0. + pub fn zBytes(w: Weight, m: usize) usize { + const p = plan(w.k, w.n, 192, 8); + return if (p.sk > 1) @as(usize, p.sk) * m * w.n * 4 else 0; + } + + /// Counter bytes one call needs: per pass and block, zeroed (fill8(0), never zeros()). + pub fn counterBytes(w: Weight) usize { + return 8 * (w.n / 128) * 4; + } + + /// y [m, n] = x [m, k] through the weight; `xh` is fp16 [m, k] scratch, `z` the split partials, `counters` zeroed. + pub fn matmul(l: *const Linear, s: Stream, w: Weight, x: u64, x_dtype: DType, m: usize, y: u64, y_dtype: DType, xh: u64, z: u64, counters: u64) !void { + // Shapes the loader admits (linear.py's k % 128 and n % 128, 1..128 rows) are returned errors here too. + if (m < 1 or m > 128) return error.Exl3Shape; + if (w.k % 128 != 0 or w.n % 128 != 0) return error.Exl3Shape; + const p = plan(w.k, w.n, 192, 8); + if ((w.k / 16) % (p.sk * p.wk) != 0) return error.Exl3Shape; + const idx = variantIndex(w.halves, w.codebook, p.wk) orelse return error.UnsupportedExl3Width; + if (p.sk > 1 and z == 0) return error.Exl3Scratch; + + var r: launch_.Args = .{}; + r.add(x); + r.add(w.suh); + r.add(xh); + r.add(@as(i32, @backingInt(x_dtype))); + r.add(@as(i32, @intCast(m))); + r.add(@as(i32, @intCast(w.k))); + try launch_.launch(l.rot_in, .{ + .grid = .{ .x = @intCast((w.k / 128 + 3) / 4), .y = @intCast(m), .z = 1 }, + .block = .{ .x = 128, .y = 1, .z = 1 }, + .shared = 0, + }, s, &r); + + var args: launch_.Args = .{}; + args.add(xh); + args.add(w.trellis); + args.add(@as(i64, @intCast(w.stride_k))); + args.add(@as(i64, @intCast(w.stride_nb))); + args.add(w.svh); + args.add(w.bias); + args.add(y); + args.add(if (p.sk > 1) z else @as(u64, 0)); + args.add(counters); + args.add(@as(i32, @backingInt(y_dtype))); + args.add(@as(i32, @intCast(w.n))); + args.add(@as(i32, @intCast(p.sk))); + args.add(@as(i32, @intCast(m))); + args.add(@as(i32, @intCast(w.k))); + try launch_.launch(l.fns[idx], .{ + .grid = .{ .x = @intCast(w.n / 128 * p.sk), .y = 1, .z = 1 }, + .block = .{ .x = @intCast(p.wk * 32), .y = 1, .z = 1 }, + .shared = sharedBytes(w, m), + }, s, &args); + } +}; + +/// The loaded kernels plus the scratch one stream's dense calls share; prompt rows run in 128-row blocks (forward.py). +pub const Dense = struct { + d: *const Driver, + lin: Linear, + xh: ?memory.DeviceBuffer = null, // the rotated inputs, 128 rows of the widest layer at most + z: ?memory.DeviceBuffer = null, // the split partials, 128 rows at most + counters: ?memory.DeviceBuffer = null, // per pass and block, zeroed before every call + + pub fn load(d: *const Driver) !Dense { + return .{ .d = d, .lin = try Linear.load(d) }; + } + + pub fn unload(de: *Dense) void { + de.lin.unload(); + const slots = [_]*?memory.DeviceBuffer{ &de.xh, &de.z, &de.counters }; + for (slots) |slot| { + if (slot.* != null) slot.*.?.free(); + slot.* = null; + } + } + + /// A scratch buffer's address, grown when a layer needs more than it holds. + fn scratch(de: *Dense, slot: *?memory.DeviceBuffer, bytes: usize) !u64 { + if (slot.* != null and slot.*.?.len >= bytes) return slot.*.?.ptr; + if (slot.* != null) slot.*.?.free(); + slot.* = try memory.DeviceBuffer.alloc(de.d, bytes); + return slot.*.?.ptr; + } + + /// Decode rows: one call (the kernel takes 1..128), any count above in 128-row blocks. + pub fn decode(de: *Dense, s: Stream, w: Weight, x: u64, x_dtype: DType, rows: usize, y: u64, y_dtype: DType) !void { + return de.blocks(s, w, x, x_dtype, rows, y, y_dtype); + } + + /// Prompt rows: 128-row blocks of `matmul`, so a row's bits never depend on where its chunk starts. + pub fn prompt(de: *Dense, s: Stream, w: Weight, x: u64, x_dtype: DType, rows: usize, y: u64, y_dtype: DType) !void { + return de.blocks(s, w, x, x_dtype, rows, y, y_dtype); + } + fn blocks(de: *Dense, s: Stream, w: Weight, x: u64, x_dtype: DType, rows: usize, y: u64, y_dtype: DType) !void { + const p = plan(w.k, w.n, 192, 8); + const z: u64 = if (p.sk > 1) try de.scratch(&de.z, Linear.zBytes(w, 128)) else 0; + const counters = try de.scratch(&de.counters, Linear.counterBytes(w)); + const xe: u64 = if (x_dtype == .f32) 4 else 2; + const ye: u64 = if (y_dtype == .f32) 4 else 2; + var r: usize = 0; + while (r < rows) : (r += 128) { + const m = @min(rows - r, 128); + const xh = try de.scratch(&de.xh, m * w.k * 2); + try de.counters.?.fill8(0, s.handle); + try de.lin.matmul(s, w, x + @as(u64, r) * w.k * xe, x_dtype, m, y + @as(u64, r) * w.n * ye, y_dtype, xh, z, counters); + } + } +}; + +test "the plan matches linear.py's for known shapes" { + try std.testing.expectEqual(@as(@TypeOf(plan(0, 0, 192, 8)), .{ .sk = 8, .wk = 4 }), plan(4096, 1024, 192, 8)); + try std.testing.expectEqual(@as(@TypeOf(plan(0, 0, 192, 8)), .{ .sk = 2, .wk = 4 }), plan(1024, 4096, 192, 8)); + try std.testing.expectEqual(@as(@TypeOf(plan(0, 0, 192, 8)), .{ .sk = 4, .wk = 4 }), plan(2048, 512, 192, 8)); + try std.testing.expectEqual(@as(@TypeOf(plan(0, 0, 192, 8)), .{ .sk = 1, .wk = 4 }), plan(512, 1024, 192, 8)); + try std.testing.expectEqual(@as(@TypeOf(plan(0, 0, 192, 8)), .{ .sk = 1, .wk = 4 }), plan(512, 512, 192, 8)); + try std.testing.expectEqual(@as(@TypeOf(plan(0, 0, 192, 8)), .{ .sk = 4, .wk = 4 }), plan(2048, 4096, 192, 8)); +} + +test "the stored and strips strides of one layout" { + try std.testing.expectEqual(@as(@TypeOf(strides(0, 0, 0, false)), .{ .k = 4096, .nb = 256 }), strides(4096, 2048, 8, false)); + try std.testing.expectEqual(@as(@TypeOf(strides(0, 0, 0, true)), .{ .k = 256, .nb = 65536 }), strides(4096, 2048, 8, true)); +} + +test "every variant's symbol is its instantiation's name" { + try std.testing.expectEqualStrings("tf_exl3_linear_8_2_4", symbol(8, .mul1, 4)); + try std.testing.expectEqualStrings("tf_exl3_linear_7_2_8", symbol(7, .mul1, 8)); + try std.testing.expectEqualStrings("tf_exl3_linear_8_0_4", symbol(8, .inst3, 4)); + try std.testing.expectEqual(@as(?usize, 5), variantIndex(12, .mul1, 4)); + try std.testing.expectEqual(@as(?usize, 2), variantIndex(6, .mul1, 4)); + try std.testing.expectEqual(@as(?usize, null), variantIndex(3, .mcg, 4)); // the half steps are mul1's alone +} + +test "matmul refuses shapes the loader would refuse, before any launch" { + const l: Linear = .{ .mod = undefined, .fns = undefined, .rot_in = undefined }; + const s: Stream = undefined; // never reached: the refusal comes first + const st = strides(4096, 1024, 8, false); + const w: Weight = .{ .trellis = 0, .suh = 0, .svh = 0, .bias = 0, .k = 4096, .n = 1024, .halves = 8, .codebook = .mul1, .stride_k = st.k, .stride_nb = st.nb }; + try std.testing.expectError(error.Exl3Shape, l.matmul(s, w, 0, .bf16, 0, 0, .bf16, 0, 0, 0)); + try std.testing.expectError(error.Exl3Shape, l.matmul(s, w, 0, .bf16, 129, 0, .bf16, 0, 0, 0)); + var kbad = w; + kbad.k = 144; // k % 16 == 0 but not a multiple of 128 + try std.testing.expectError(error.Exl3Shape, l.matmul(s, kbad, 0, .bf16, 8, 0, .bf16, 0, 0, 0)); + var nbad = w; + nbad.n = 100; + try std.testing.expectError(error.Exl3Shape, l.matmul(s, nbad, 0, .bf16, 8, 0, .bf16, 0, 0, 0)); + try std.testing.expectError(error.Exl3Scratch, l.matmul(s, w, 0, .bf16, 8, 0, .bf16, 0, 0, 0)); // plan(4096, 1024) splits K +} diff --git a/zig/src/cuda/exl3_experts.zig b/zig/src/cuda/exl3_experts.zig new file mode 100644 index 0000000000..e665dd3a02 --- /dev/null +++ b/zig/src/cuda/exl3_experts.zig @@ -0,0 +1,299 @@ +//! The EXL3 routed experts (exl3/experts.cu's device code): the group, the grouped GEMMs, the epilogues. +const std = @import("std"); +const driver = @import("driver.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); +const kernels = @import("kernels.zig"); +const exl3 = @import("exl3.zig"); +const Driver = driver.Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +pub const DType = exl3.DType; +pub const Codebook = exl3.Codebook; + +/// One layer's routed experts (experts.py's Exl3RoutedExperts): per-expert trellis pointers, stacked scales, the clamp. +pub const Layer = struct { + gate: u64, // int64 [E], one trellis pointer per expert + up: u64, + down: u64, + gate_k2: u64, // int32 [E], half-bits a value (2 * bits) + up_k2: u64, + down_k2: u64, + suh_g: u64, // fp16 [E, D] + suh_u: u64, + svh_g: u64, // fp16 [E, I] + svh_u: u64, + suh_d: u64, // fp16 [E, I] + svh_d: u64, // fp16 [E, D] + experts: usize, // E + dims: usize, // D (model width) + width: usize, // I (expert width) + codebook: Codebook, + gu_stride: usize = 0, // gate/up trellis row stride in tiles (0: contiguous per projection) + limit: f32 = 0, // the SwiGLU clamp (0: none) + act_mode: i32 = 0, // 0: GLM's bf16 roundings (the route RED-SNOW runs), 1: fp32 +}; + +/// What one projection launches (experts.py default_config): the 8-tile 4-warp K splits, shape-fixed; n % 128 required. +pub fn plan(k: usize, n: usize, gateup: bool) !u32 { + const cands = [_]usize{ if (gateup) 4 else 1, 2, 1 }; + for (cands) |sk| { + if (k % (16 * sk * 4) == 0 and n % 128 == 0) return @intCast(sk); + } + return error.TileNotBuilt; +} + +/// One call's buffers (exl3/experts.py Scratch), laid out in scratch()'s order, each 256-byte aligned. +pub const Scratch = struct { + xg: u64, // fp16 [P, D], the rotated gate inputs + xu: u64, // fp16 [P, D], the rotated up inputs + xd: u64, // fp16 [P, I], the rotated gate and up outputs (the down inputs) + z: u64, // fp32, gate and up write 2 * sk_gu * P * I partials, down sk_d * P * D + y: u64, // fp32 [P, D], the per-slot outputs + ids: u64, // int32 [min(P, E)], distinct routed experts in id order + count: u64, // int32 [1] + members: u64, // int32 [min(P, E), rows], row * 32 + slot, -1 after the last member + rows: usize, // the row capacity the buffers hold + slots: usize, // slots a row, at most 32 +}; + +/// The byte size of the scratch buffers for a layer and row capacity, or error.TileNotBuilt when plan() fails. +pub fn scratchBytes(l: Layer, rows: usize, slots: usize) !usize { + const g_sk = try plan(l.dims, l.width, true); + const d_sk = try plan(l.width, l.dims, false); + const p = rows * slots; + const maxu = @min(p, l.experts); + return align256(p * l.dims * 2) + align256(p * l.dims * 2) + align256(p * l.width * 2) + + align256(@max(2 * g_sk * l.width, d_sk * l.dims) * p * 4) + align256(p * l.dims * 4) + + align256(maxu * 4) + 256 + align256(maxu * rows * 4); +} + +fn align256(n: usize) usize { + return (n + 255) / 256 * 256; +} + +/// The pointers into one allocation of scratchBytes(), or error.TileNotBuilt when plan() fails. +pub fn scratch(base: u64, l: Layer, rows: usize, slots: usize) !Scratch { + const p = rows * slots; + const g_sk = try plan(l.dims, l.width, true); + const d_sk = try plan(l.width, l.dims, false); + const maxu = @min(p, l.experts); + var at = base; + const take = struct { + fn f(a: *u64, bytes: usize) u64 { + const r = a.*; + a.* = r + align256(bytes); + return r; + } + }.f; + const xg = take(&at, p * l.dims * 2); + const xu = take(&at, p * l.dims * 2); + const xd = take(&at, p * l.width * 2); + const z = take(&at, @max(2 * g_sk * l.width, d_sk * l.dims) * p * 4); + const y = take(&at, p * l.dims * 4); + const ids = take(&at, maxu * 4); + const count = take(&at, 4); + const members = take(&at, maxu * rows * 4); + return .{ .xg = xg, .xu = xu, .xd = xd, .z = z, .y = y, .ids = ids, .count = count, .members = members, .rows = rows, .slots = slots }; +} + +/// The loaded kernels. +pub const Experts = struct { + mod: Module, + group: Function, + rot_in: Function, + grouped: [3]Function, // one per codebook + gateup: Function, + down_combine: Function, + + pub fn load(d: *const Driver) !Experts { + if (!kernels.available) return error.BuiltWithoutKernels; + var e: Experts = .{ + .mod = try Module.load(d, kernels.exl3_experts), + .group = undefined, + .rot_in = undefined, + .grouped = undefined, + .gateup = undefined, + .down_combine = undefined, + }; + errdefer e.mod.unload(); + e.group = try e.mod.function("tf_exl3_group"); + try e.group.allowDynamicShared(64 * 1024); + e.rot_in = try e.mod.function("tf_exl3_experts_rot_in"); + e.grouped = .{ + try e.mod.function("tf_exl3_grouped_0"), + try e.mod.function("tf_exl3_grouped_1"), + try e.mod.function("tf_exl3_grouped_2"), + }; + e.gateup = try e.mod.function("tf_exl3_gateup_epilogue"); + e.down_combine = try e.mod.function("tf_exl3_down_combine"); + return e; + } + + pub fn unload(e: *Experts) void { + e.mod.unload(); + } + + /// The routed experts of `rows` rows: group, rotate in, grouped gate/up, SwiGLU, grouped down, wts-weighted sum. + pub fn run(e: *const Experts, s: Stream, l: Layer, x: u64, x_stride: usize, x_dtype: DType, pick: u64, wts: u64, rows: usize, out: u64, sc: Scratch) !void { + const d = l.dims; + const w = l.width; + const slots = sc.slots; + std.debug.assert(rows >= 1 and rows <= sc.rows); + std.debug.assert(slots >= 1 and slots <= 32); + std.debug.assert(l.experts >= 1 and l.experts <= 4096); + const g_sk = try plan(d, w, true); + const d_sk = try plan(w, d, false); + const p = rows * slots; + const maxu = @min(p, l.experts); + const mt = (rows + 15) / 16; + const i32d = struct { + fn f(v: usize) i32 { + return @intCast(v); + } + }.f; + + var a: launch_.Args = .{}; + a.add(pick); + a.add(sc.ids); + a.add(sc.count); + a.add(sc.members); + a.add(i32d(rows)); + a.add(i32d(slots)); + a.add(i32d(l.experts)); + a.add(i32d(rows)); + try launch_.launch(e.group, .{ .grid = .{ .x = 1, .y = 1, .z = 1 }, .block = .{ .x = 1024, .y = 1, .z = 1 }, .shared = @intCast(rows * slots * 4) }, s, &a); + + var r: launch_.Args = .{}; + r.add(x); + r.add(i32d(x_stride)); + r.add(@as(i32, @backingInt(x_dtype))); + r.add(pick); + r.add(l.suh_g); + r.add(l.suh_u); + r.add(sc.xg); + r.add(sc.xu); + r.add(i32d(d)); + r.add(i32d(slots)); + r.add(i32d(l.experts)); + try launch_.launch(e.rot_in, .{ .grid = .{ .x = @intCast(p), .y = @intCast(d / 128), .z = 2 }, .block = .{ .x = 32, .y = 1, .z = 1 } }, s, &r); + + var g: launch_.Args = .{}; + g.add(sc.xg); + g.add(sc.xu); + g.add(l.gate); + g.add(l.up); + g.add(l.gate_k2); + g.add(l.up_k2); + g.add(sc.ids); + g.add(sc.count); + g.add(sc.members); + g.add(sc.z); + g.add(i32d(d)); + g.add(i32d(w)); + g.add(i32d(p)); + g.add(i32d(g_sk)); + g.add(i32d(rows)); + g.add(i32d(slots)); + g.add(i32d(maxu)); + g.add(@as(i32, 2)); + g.add(i32d(l.gu_stride)); + try launch_.launch(e.grouped[@backingInt(l.codebook)], .{ .grid = .{ .x = @intCast(maxu), .y = @intCast(w / 128), .z = @intCast(2 * g_sk * mt) }, .block = .{ .x = 128, .y = 1, .z = 1 } }, s, &g); + + var ep: launch_.Args = .{}; + ep.add(sc.z); + ep.add(pick); + ep.add(l.svh_g); + ep.add(l.svh_u); + ep.add(l.suh_d); + ep.add(sc.xd); + ep.add(i32d(p)); + ep.add(i32d(w)); + ep.add(i32d(g_sk)); + ep.add(i32d(l.experts)); + ep.add(l.limit); + ep.add(l.act_mode); + try launch_.launch(e.gateup, .{ .grid = .{ .x = @intCast(p), .y = @intCast(w / 128), .z = 1 }, .block = .{ .x = 32, .y = 1, .z = 1 } }, s, &ep); + + var dn: launch_.Args = .{}; + dn.add(sc.xd); + dn.add(sc.xd); + dn.add(l.down); + dn.add(l.down); + dn.add(l.down_k2); + dn.add(l.down_k2); + dn.add(sc.ids); + dn.add(sc.count); + dn.add(sc.members); + dn.add(sc.z); + dn.add(i32d(w)); + dn.add(i32d(d)); + dn.add(i32d(p)); + dn.add(i32d(d_sk)); + dn.add(i32d(rows)); + dn.add(i32d(slots)); + dn.add(i32d(maxu)); + dn.add(@as(i32, 1)); + dn.add(@as(i32, 0)); + try launch_.launch(e.grouped[@backingInt(l.codebook)], .{ .grid = .{ .x = @intCast(maxu), .y = @intCast(d / 128), .z = @intCast(d_sk * mt) }, .block = .{ .x = 128, .y = 1, .z = 1 } }, s, &dn); + + var cb: launch_.Args = .{}; + cb.add(sc.z); + cb.add(pick); + cb.add(l.svh_d); + cb.add(sc.y); + cb.add(wts); + cb.add(out); + cb.add(i32d(p)); + cb.add(i32d(d)); + cb.add(i32d(d_sk)); + cb.add(i32d(l.experts)); + cb.add(i32d(slots)); + try launch_.launch(e.down_combine, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(d / 128), .z = 1 }, .block = .{ .x = @intCast(32 * slots), .y = 1, .z = 1 } }, s, &cb); + } +}; + +test "the plan matches experts.py's for known shapes" { + try std.testing.expectEqual(@as(u32, 4), try plan(4096, 2048, true)); // GLM_GATEUP + try std.testing.expectEqual(@as(u32, 1), try plan(2048, 4096, false)); // GLM_DOWN + try std.testing.expectEqual(@as(u32, 4), try plan(512, 256, true)); + try std.testing.expectEqual(@as(u32, 1), try plan(256, 512, false)); + try std.testing.expectError(error.TileNotBuilt, plan(300, 256, true)); +} + +test "the scratch layout is 256-aligned and covered by scratchBytes" { + const l: Layer = .{ + .gate = 0, + .up = 0, + .down = 0, + .gate_k2 = 0, + .up_k2 = 0, + .down_k2 = 0, + .suh_g = 0, + .suh_u = 0, + .svh_g = 0, + .svh_u = 0, + .suh_d = 0, + .svh_d = 0, + .experts = 9, + .dims = 512, + .width = 256, + .codebook = .mul1, + }; + const bytes = try scratchBytes(l, 64, 8); + var buf: [4096]u8 align(256) = undefined; + const sc = try scratch(@intFromPtr(&buf), l, 64, 8); + try std.testing.expectEqual(@as(u64, @intFromPtr(&buf)), sc.xg); + try std.testing.expectEqual(@as(u64, 0), sc.xg % 256); + try std.testing.expectEqual(@as(u64, 0), sc.xu % 256); + try std.testing.expectEqual(@as(u64, 0), sc.xd % 256); + try std.testing.expectEqual(@as(u64, 0), sc.z % 256); + try std.testing.expectEqual(@as(u64, 0), sc.y % 256); + try std.testing.expectEqual(@as(u64, 0), sc.ids % 256); + try std.testing.expectEqual(@as(u64, 0), sc.count % 256); + try std.testing.expectEqual(@as(u64, 0), sc.members % 256); + try std.testing.expectEqual(@as(usize, @intFromPtr(&buf) + bytes), sc.members + @as(usize, 9) * 64 * 4); +} diff --git a/zig/src/cuda/exl3_upload.zig b/zig/src/cuda/exl3_upload.zig new file mode 100644 index 0000000000..ef0e434a09 --- /dev/null +++ b/zig/src/cuda/exl3_upload.zig @@ -0,0 +1,105 @@ +//! Uploading core exl3_load groups to the device as the exl3 kernels' weights (the engine's load path). +const std = @import("std"); +const driver = @import("driver.zig"); +const memory = @import("memory.zig"); +const exl3 = @import("exl3.zig"); +const exl3e = @import("exl3_experts.zig"); +const Driver = driver.Driver; + +/// The tensors one upload holds: appended to `bufs` and freed together. +pub const Uploader = struct { + d: *const Driver, + gpa: std.mem.Allocator, + bufs: std.ArrayList(memory.DeviceBuffer) = .empty, + bytes: usize = 0, // the tensors' total device bytes + + pub fn deinit(u: *Uploader) void { + for (u.bufs.items) |*b| b.free(); + u.bufs.deinit(u.gpa); + } + + pub fn dev(u: *Uploader, bytes: []const u8) !u64 { + const b = try memory.DeviceBuffer.fromHost(u.d, bytes); + try u.bufs.append(u.gpa, b); + u.bytes += bytes.len; + return b.ptr; + } + + /// One dense group: `parts` (exl3_load.Parts) and `g` (exl3_load.Group) as the Weight; the scales must be fp16. + pub fn dense(u: *Uploader, parts: anytype, g: anytype) !exl3.Weight { + if (parts.suh_packed) return error.PackedScales; + const st = exl3.strides(g.k, g.n, g.halves, false); + return .{ + .trellis = try u.dev(parts.trellis.bytes), + .suh = try u.dev(parts.suh.bytes), + .svh = try u.dev(parts.svh.bytes), + .bias = if (parts.bias) |b| try u.dev(b.bytes) else 0, + .k = @intCast(g.k), + .n = @intCast(g.n), + .halves = g.halves, + .codebook = switch (g.codebook) { + .inst3 => .inst3, + .mcg => .mcg, + .mul1 => .mul1, + }, + .stride_k = st.k, + .stride_nb = st.nb, + }; + } +}; + +/// One expert projection's tensors (fp16 scales) and half-bit width: the exl3_load Parts' bytes, the group's halves. +pub const Proj = struct { trellis: []const u8, suh: []const u8, svh: []const u8, k2: u32 }; + +/// One layer's routed experts: trellises in place (pointers stacked), scales stacked [E, n] fp16, widths int32 [E]. +pub fn layer(u: *Uploader, a: std.mem.Allocator, dims: usize, width: usize, cb: exl3.Codebook, limit: f32, act_mode: i32, gate: []const Proj, up: []const Proj, down: []const Proj) !exl3e.Layer { + const e = gate.len; + if (up.len != e or down.len != e or e == 0) return error.BadExperts; + const gt = try pointers(u, a, gate); + const ut = try pointers(u, a, up); + const dt = try pointers(u, a, down); + return .{ + .gate = try u.dev(std.mem.sliceAsBytes(gt)), + .up = try u.dev(std.mem.sliceAsBytes(ut)), + .down = try u.dev(std.mem.sliceAsBytes(dt)), + .gate_k2 = try k2s(u, a, gate), + .up_k2 = try k2s(u, a, up), + .down_k2 = try k2s(u, a, down), + .suh_g = try stacked(u, a, gate, dims, "suh"), + .suh_u = try stacked(u, a, up, dims, "suh"), + .svh_g = try stacked(u, a, gate, width, "svh"), + .svh_u = try stacked(u, a, up, width, "svh"), + .suh_d = try stacked(u, a, down, width, "suh"), + .svh_d = try stacked(u, a, down, dims, "svh"), + .experts = e, + .dims = dims, + .width = width, + .codebook = cb, + .gu_stride = 0, + .limit = limit, + .act_mode = act_mode, + }; +} + +fn pointers(u: *Uploader, a: std.mem.Allocator, projs: []const Proj) ![]u64 { + const out = try a.alloc(u64, projs.len); + for (projs, 0..) |p, i| out[i] = try u.dev(p.trellis); + return out; +} + +fn k2s(u: *Uploader, a: std.mem.Allocator, projs: []const Proj) !u64 { + const host = try a.alloc(i32, projs.len); + for (projs, 0..) |p, i| host[i] = @intCast(p.k2); + return u.dev(std.mem.sliceAsBytes(host)); +} + +/// Every expert's fp16 scale of `name` ("suh" or "svh"), `n` values each, stacked [E, n]. +fn stacked(u: *Uploader, a: std.mem.Allocator, projs: []const Proj, n: usize, comptime name: []const u8) !u64 { + const host = try a.alloc(u8, projs.len * n * 2); + for (projs, 0..) |p, i| { + const s = @field(p, name); + if (s.len != n * 2) return error.BadScales; + @memcpy(host[i * n * 2 ..][0..s.len], s); + } + return u.dev(host); +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 0fa8bf6e28..10fd9e55ab 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -33,6 +33,8 @@ pub const train: []const u8 = if (available) &Blob("fatbin_train").bytes else &. pub const qmmf: []const u8 = if (available) &Blob("fatbin_qmmf").bytes else &.{}; pub const prompt16: []const u8 = if (available) &Blob("fatbin_prompt16").bytes else &.{}; pub const modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack").bytes else &.{}; +pub const exl3_linear: []const u8 = if (available) &Blob("fatbin_exl3_linear").bytes else &.{}; +pub const exl3_experts: []const u8 = if (available) &Blob("fatbin_exl3_experts").bytes else &.{}; pub const torch_argmax: []const u8 = if (available) &Blob("fatbin_torch_argmax").bytes else &.{}; pub const torch_topk: []const u8 = if (available) &Blob("fatbin_torch_topk").bytes else &.{}; pub const torch_pointwise: []const u8 = if (available) &Blob("fatbin_torch_pointwise").bytes else &.{}; diff --git a/zig/src/cuda/qlinear.zig b/zig/src/cuda/qlinear.zig index 6d6a56f5a3..0ce5060684 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.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.Weight, // the trellis: decode and prompt rows through the same row-invariant kernel 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: ?*exl3.Dense = null, // the EXL3 kernels and the scratch the dense calls share /// 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,11 @@ 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| { + const e = l.exl3 orelse return error.FormatNotLoaded; + if (in.ldx != 0 and in.ldx != q.k) return error.RowStride; + return e.decode(s, q, in.x, .bf16, rows, out, .bf16); + }, } } @@ -88,6 +98,11 @@ 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| { + const e = l.exl3 orelse return error.FormatNotLoaded; + if (in.ldx != 0 and in.ldx != q.k) return error.RowStride; + return e.prompt(s, q, in.x, .bf16, rows, out, .bf16); + }, } } @@ -309,6 +324,16 @@ test "a weight view reports its shape in every format" { try std.testing.expectEqual(@as(usize, 2560), f.inputs()); } +test "the exl3 weight view reports its shape and is refused when not loaded" { + const w: Weight = .{ .exl3 = .{ .trellis = 0, .suh = 0, .svh = 0, .k = 4096, .n = 1024, .halves = 8, .codebook = .mul1, .stride_k = 4096, .stride_nb = 256 } }; + try std.testing.expectEqual(@as(usize, 1024), w.outputs()); + try std.testing.expectEqual(@as(usize, 4096), w.inputs()); + const l: Linear = .{}; + 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)); +} + test "a format that is not loaded is refused before any launch" { const l: Linear = .{}; const w: Weight = .{ .affine4 = .{ .w = 0, .s = 0, .b = 0, .n = 64, .k = 64, .npad = 64 } }; diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 3053c09a6b..fea0c8bb94 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -36,6 +36,9 @@ pub const embed_rows = @import("embed_rows.zig"); pub const experts = @import("experts.zig"); pub const sample = @import("sample.zig"); pub const modelopt = @import("modelopt.zig"); +pub const exl3 = @import("exl3.zig"); +pub const exl3_experts = @import("exl3_experts.zig"); +pub const exl3_upload = @import("exl3_upload.zig"); test { _ = @import("memory.zig"); @@ -53,5 +56,8 @@ test { _ = @import("fp8.zig"); _ = carveout; _ = @import("nvfp4.zig"); + _ = exl3; + _ = exl3_experts; + _ = exl3_upload; _ = @import("modelopt.zig"); } diff --git a/zig/tests/cuda/exl3_dense_tests.zig b/zig/tests/cuda/exl3_dense_tests.zig new file mode 100644 index 0000000000..877d9ed791 --- /dev/null +++ b/zig/tests/cuda/exl3_dense_tests.zig @@ -0,0 +1,329 @@ +//! EXL3 dense prompt path: qlinear's dispatch, float64 references, chunk invariance, the engine's prompt bytes. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const exl3fmt = @import("core").exl3_format; +const exl3_load = @import("core").exl3_load; +const exl3 = cuda.exl3; +const qlinear = cuda.qlinear; +const exl3_tests = @import("exl3_tests.zig"); +const Run = exl3_tests.Run; +const Rig = quant_ref_tests.RigT; +const bf = quant_ref_tests.helpers.bfOf; + +pub fn run(r: *Rig, state: *Run, dir: []const u8) !void { + var dense = try exl3.Dense.load(r.gpu.d); + defer dense.unload(); + const ql: qlinear.Linear = .{ .exl3 = &dense }; + { + var ok = true; + qlinearCheck(r, &ql, &dense.lin) catch |e| { + state.fail(e, "exl3.qlinear"); + ok = false; + }; + if (ok) check.pass("exl3.qlinear", .{}); + } + { + var ok = true; + promptF64(r, &dense.lin) catch |e| { + state.fail(e, "exl3.prompt.linear.f64"); + ok = false; + }; + if (ok) check.pass("exl3.prompt.linear.f64", .{}); + } + { + var ok = true; + chunks(r, &dense) catch |e| { + state.fail(e, "exl3.prompt.linear.chunks"); + ok = false; + }; + if (ok) check.pass("exl3.prompt.linear.chunks", .{}); + } + if (pythonPrompt(r, &dense, dir) catch |e| blk: { + state.fail(e, "exl3.prompt.python"); + break :blk false; + }) check.pass("exl3.prompt.python", .{}); + if (loadRedsnow(r, &dense, dir) catch |e| blk: { + state.fail(e, "exl3.load.redsnow"); + break :blk false; + }) check.pass("exl3.load.redsnow", .{}); +} + +/// Through qlinear.Linear's decode and prompt dispatch, 1, 17 and 300 rows byte-equal the direct block loop. +fn qlinearCheck(r: *Rig, ql: *const qlinear.Linear, lin: *const exl3.Linear) !void { + const k = 1024; + const n = 512; + const trellis, const suh, const svh = try tensors(r, k, n, .{ .halves = 8 }); + const w = try weight(r, trellis, suh, svh, k, n, 8, .mul1); + const x = try r.bfs(300 * k, 1.0); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const counters = try exl3_tests.zeroed(r, exl3.Linear.counterBytes(w)); + const xh = try exl3_tests.zeroed(r, 128 * k * 2); + const z: u64 = if (exl3.plan(k, n, 192, 8).sk > 1) try exl3_tests.zeroed(r, exl3.Linear.zBytes(w, 128)) else 0; + for ([_]usize{ 1, 17, 300 }) |m| { + const want = try r.zeros(m * n * 2); + var off: usize = 0; + while (off < m) : (off += 128) { + try lin.matmul(r.s, w, dx + off * k * 2, .bf16, @min(m - off, 128), want + off * n * 2, .bf16, xh, z, counters); + } + const wantraw = try r.back(want, m * n * 2); + const Kind = enum { decode, prompt }; + for ([_]Kind{ .decode, .prompt }) |kind| { + const got = try r.zeros(m * n * 2); + switch (kind) { + .decode => try ql.decode(r.s, .{ .exl3 = w }, .{ .x = dx }, got, m), + .prompt => try ql.prompt(r.s, .{ .exl3 = w }, .{ .x = dx }, got, m), + } + const raw = try r.back(got, m * n * 2); + const id = try std.fmt.allocPrint(r.a, "exl3.qlinear.{s} m={d}", .{ @tagName(kind), m }); + try check.sameBytes(id, raw, wantraw); + } + } +} + +/// 300 and 1000 prompt rows against the float64 reference (exl3_tests' reference with m a parameter). +fn promptF64(r: *Rig, lin: *const exl3.Linear) !void { + const cfgs = [_]struct { halves: u32, codebook: exl3fmt.Codebook, k: usize, n: usize }{ + .{ .halves = 8, .codebook = .mul1, .k = 512, .n = 512 }, + .{ .halves = 10, .codebook = .mul1, .k = 1024, .n = 512 }, + }; + for (cfgs) |c| { + const width: exl3fmt.Width = .{ .halves = @intCast(c.halves) }; + const trellis, const suh, const svh = try tensors(r, c.k, c.n, width); + const w = try weight(r, trellis, suh, svh, c.k, c.n, c.halves, cbOf(@tagName(c.codebook))); + const counters = try exl3_tests.zeroed(r, exl3.Linear.counterBytes(w)); + const xh = try exl3_tests.zeroed(r, 128 * c.k * 2); + const z: u64 = if (exl3.plan(c.k, c.n, 192, 8).sk > 1) try exl3_tests.zeroed(r, exl3.Linear.zBytes(w, 128)) else 0; + for ([_]usize{ 300, 1000 }) |m| { + const x = try r.bfs(m * c.k, 1.0); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const ref = try promptReference(r, width, c.codebook, trellis, c.k, c.n, x, suh, svh, m); + const id = try std.fmt.allocPrint(r.a, "exl3.prompt.linear.f64 k={d} n={d} m={d}", .{ c.k, c.n, m }); + const dy = try r.zeros(m * c.n * 2); + var off: usize = 0; + while (off < m) : (off += 128) { + try lin.matmul(r.s, w, dx + off * c.k * 2, .bf16, @min(m - off, 128), dy + off * c.n * 2, .bf16, xh, z, counters); + } + try exl3_tests.closeOut(r, id, dy, m * c.n, ref.want, ref.tol); + } + } +} + +/// A row's bytes are identical whether the prompt runs as one call, 128-row calls, a split at 37, or 2048 chunks. +fn chunks(r: *Rig, de: *exl3.Dense) !void { + const k = 512; + const n = 512; + const rows = 2185; // 17 whole 128-blocks and a 9-row tail + const trellis, const suh, const svh = try tensors(r, k, n, .{ .halves = 8 }); + const w = try weight(r, trellis, suh, svh, k, n, 8, .mul1); + const x = try r.bfs(rows * k, 1.0); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const y0 = try r.zeros(rows * n * 2); + try de.prompt(r.s, w, dx, .bf16, rows, y0, .bf16); + const want = try r.back(y0, rows * n * 2); + var b128: [17]usize = undefined; + for (&b128, 1..) |*v, i| v.* = i * 128; + for ([_][]const usize{ &.{}, &b128, &.{37}, &.{2048} }) |bounds| { + const y = try r.zeros(rows * n * 2); + var prev: usize = 0; + for (bounds) |b| { + try de.prompt(r.s, w, dx + prev * k * 2, .bf16, b - prev, y + prev * n * 2, .bf16); + prev = b; + } + try de.prompt(r.s, w, dx + prev * k * 2, .bf16, rows - prev, y + prev * n * 2, .bf16); + const got = try r.back(y, rows * n * 2); + const id = try std.fmt.allocPrint(r.a, "exl3.prompt.linear.chunks {d} splits", .{bounds.len}); + try check.sameBytes(id, got, want); + } +} + +/// The Python engine's recorded prompt bytes for one real dense projection, bit for bit. +fn pythonPrompt(r: *Rig, de: *exl3.Dense, dir: []const u8) !bool { + if (std.mem.eql(u8, dir, "none")) { + std.debug.print("SKIP exl3.prompt.python\n", .{}); + return false; + } + const d = (try denseEntry(r, dir)) orelse { + std.debug.print("SKIP exl3.prompt.python: no dense entry\n", .{}); + return false; + }; + const k: usize = @intCast(d.get("k").?.integer); + const n: usize = @intCast(d.get("n").?.integer); + const bits: u32 = bitsOf(d.get("bits").?); + const m: usize = @intCast(d.get("rows").?.integer); + const bias = d.get("bias").?; + const st = exl3.strides(k, n, 2 * bits, false); + const w: exl3.Weight = .{ + .trellis = try r.dev(try exl3_tests.readFile(r, dir, d.get("trellis").?.string)), + .suh = try r.dev(try exl3_tests.readFile(r, dir, d.get("suh").?.string)), + .svh = try r.dev(try exl3_tests.readFile(r, dir, d.get("svh").?.string)), + .bias = if (std.meta.activeTag(bias) == .null) 0 else try r.dev(try exl3_tests.readFile(r, dir, bias.string)), + .k = @intCast(k), + .n = @intCast(n), + .halves = 2 * bits, + .codebook = cbOf(d.get("codebook").?.string), + .stride_k = st.k, + .stride_nb = st.nb, + }; + const x = try exl3_tests.readFile(r, dir, d.get("x").?.string); + const want = try exl3_tests.readFile(r, dir, d.get("y").?.string); + const dy = try r.zeros(m * n * 2); + try de.prompt(r.s, w, try r.dev(x), .bf16, m, dy, .bf16); + const got = try r.back(dy, want.len); + try check.sameBytes("exl3.prompt.python", got, want); + return true; +} + +/// The checkpoint loader over real RED-SNOW: every EXL3 group admitted, one layer's prompt vs the fixture bytes. +fn loadRedsnow(r: *Rig, de: *exl3.Dense, dir: []const u8) !bool { + if (std.mem.eql(u8, dir, "none")) { + std.debug.print("SKIP exl3.load.redsnow\n", .{}); + return false; + } + const o = try manifest(r, dir); + const model = o.get("model") orelse return error.TestFailed; // the fixture names its model folder + const d = (try denseEntry(r, dir)) orelse { + std.debug.print("SKIP exl3.load.redsnow: no dense entry\n", .{}); + return false; + }; + var s = try exl3_load.scan(r.gpu.gpa, r.gpu.io, model.string); + defer s.deinit(); + const sum = s.summary; + // the inspect tool's counts: every group suh/svh mul1, the whole-bits histogram, the experts' downs + try check.expect(sum.groups == 37408, "exl3.load.redsnow: {d} groups, expected 37408", .{sum.groups}); + try check.expect(sum.mul1 == 37408, "exl3.load.redsnow: {d} mul1, expected 37408", .{sum.mul1}); + try check.expect(sum.suh_svh == 37408, "exl3.load.redsnow: {d} suh/svh, expected 37408", .{sum.suh_svh}); + const hist = [_]usize{ 15425, 11801, 10054, 43, 85 }; // whole bits 4 to 8 + const downs = [_]usize{ 10283, 1448, 653, 0, 0 }; + for (hist, downs, 4..) |want, down, bits| { + try check.expect(sum.bits[2 * bits] == want, "exl3.load.redsnow: {d} groups at {d} bits, expected {d}", .{ sum.bits[2 * bits], bits, want }); + try check.expect(sum.experts_down[2 * bits] == down, "exl3.load.redsnow: {d} expert down projections at {d} bits, expected {d}", .{ sum.experts_down[2 * bits], bits, down }); + } + try check.expect(s.bad == 0, "exl3.load.redsnow: {d} tensors refused", .{s.bad}); + const gi = s.find(d.get("name").?.string) orelse return error.TestFailed; // the fixture's layer + const g = s.groups[gi]; + const p = try s.parts(r.a, gi); + const st = exl3.strides(g.k, g.n, g.halves, false); + const w: exl3.Weight = .{ + .trellis = try r.dev(p.trellis.bytes), + .suh = try r.dev(p.suh.bytes), + .svh = try r.dev(p.svh.bytes), + .bias = if (p.bias) |b| try r.dev(b.bytes) else 0, + .k = @intCast(g.k), + .n = @intCast(g.n), + .halves = g.halves, + .codebook = .mul1, + .stride_k = st.k, + .stride_nb = st.nb, + }; + try check.expect(g.halves == 2 * bitsOf(d.get("bits").?) and g.k == @as(usize, @intCast(d.get("k").?.integer)) and + g.n == @as(usize, @intCast(d.get("n").?.integer)), "exl3.load.redsnow: the loaded layer is not the fixture's", .{}); + const m: usize = @intCast(d.get("rows").?.integer); + const x = try exl3_tests.readFile(r, dir, d.get("x").?.string); + const want = try exl3_tests.readFile(r, dir, d.get("y").?.string); + const dy = try r.zeros(m * g.n * 2); + try de.prompt(r.s, w, try r.dev(x), .bf16, m, dy, .bf16); + const got = try r.back(dy, want.len); + try check.sameBytes("exl3.load.redsnow", got, want); + return true; +} + +/// The manifest's "dense" entry, null when an older fixture dir has none. +fn denseEntry(r: *Rig, dir: []const u8) !?std.json.ObjectMap { + const o = try manifest(r, dir); + return if (o.get("dense")) |v| v.object else null; +} + +fn manifest(r: *Rig, dir: []const u8) !std.json.ObjectMap { + const text = try std.Io.Dir.cwd().readFileAlloc(r.gpu.io, try std.fs.path.join(r.a, &.{ dir, "manifest.json" }), r.a, .limited(1 << 24)); + const parsed = try std.json.parseFromSlice(std.json.Value, r.a, text, .{}); + return parsed.value.object; +} + +fn bitsOf(v: std.json.Value) u32 { + return switch (v) { + .integer => |i| @intCast(i), + .float => |f| @intFromFloat(f), + else => unreachable, + }; +} + +fn cbOf(name: []const u8) exl3.Codebook { + if (std.mem.eql(u8, name, "mcg")) return .mcg; + if (std.mem.eql(u8, name, "inst3")) return .inst3; + return .mul1; +} + +/// Random trellis, scales at one shape (the same layout as exl3_tests' f64 cfgs). +fn tensors(r: *Rig, k: usize, n: usize, width: exl3fmt.Width) !struct { []i16, []u16, []u16 } { + const trellis = try r.a.alloc(i16, (k / 16) * (n / 16) * width.tileWords()); + for (trellis) |*v| v.* = r.rng.int(i16); + const suh = try r.a.alloc(u16, k); + const svh = try r.a.alloc(u16, n); + for (suh) |*v| v.* = @bitCast(@as(f16, @floatCast(r.rng.float(f32) * 2 - 1))); + for (svh) |*v| v.* = @bitCast(@as(f16, @floatCast(r.rng.float(f32) * 2 - 1))); + return .{ trellis, suh, svh }; +} + +fn weight(r: *Rig, trellis: []const i16, suh: []const u16, svh: []const u16, k: usize, n: usize, halves: u32, cb: exl3.Codebook) !exl3.Weight { + const st = exl3.strides(k, n, halves, false); + return .{ + .trellis = try r.dev(std.mem.sliceAsBytes(trellis)), + .suh = try r.dev(std.mem.sliceAsBytes(suh)), + .svh = try r.dev(std.mem.sliceAsBytes(svh)), + .bias = 0, + .k = @intCast(k), + .n = @intCast(n), + .halves = halves, + .codebook = cb, + .stride_k = st.k, + .stride_nb = st.nb, + }; +} + +/// The Python engine's f64 reference, with the row count a parameter. +fn promptReference(r: *Rig, width: exl3fmt.Width, cb: exl3fmt.Codebook, trellis: []const i16, k: usize, n: usize, x: []const u16, suh: []const u16, svh: []const u16, m: usize) !struct { want: []f64, tol: []f64 } { + const wq = try r.a.alloc(f16, k * n); + exl3fmt.unpack(width, cb, trellis, k / 16, n / 16, wq); + const xs = try r.a.alloc(f64, m * k); + for (0..m) |row| { + for (0..k) |i| xs[row * k + i] = bf(x[row * k + i]) * @as(f64, @floatCast(@as(f16, @bitCast(suh[i])))); + } + exl3fmt.rotate(xs, m, k, 1); + const xh = try r.a.alloc(f16, m * k); + for (xh, xs) |*o, v| o.* = @floatCast(v); + const ypre = try r.a.alloc(f64, m * n); // [n][m], rotated along n + @memset(ypre, 0); + const s1 = try r.a.alloc(f64, m * n); // each output's sum of |xh * wq| over k + @memset(s1, 0); + var xv: [1000]f64 = undefined; + var xa: [1000]f64 = undefined; + for (0..k) |ki| { + for (0..m) |row| { + xv[row] = @floatCast(xh[row * k + ki]); + xa[row] = @abs(xv[row]); + } + for (0..n) |col| { + const wv: f64 = @floatCast(wq[ki * n + col]); + const wa = @abs(wv); + const base = col * m; + for (0..m) |row| { + ypre[base + row] += xv[row] * wv; + s1[base + row] += xa[row] * wa; + } + } + } + exl3fmt.rotate(ypre, n, m, 0); + const want = try r.a.alloc(f64, m * n); + const tol = try r.a.alloc(f64, m * n); + for (0..n) |col| { + const scale: f64 = @floatCast(@as(f16, @bitCast(svh[col]))); + for (0..m) |row| { + want[row * n + col] = ypre[col * m + row] * scale; + // bf16 output ulp (2^-7 of the value) plus fp32 accumulation and fp16 xh cases, 2^-11 of the |t| sum + tol[row * n + col] = @abs(want[row * n + col]) / 128 + s1[row * n + col] / 2048 + 1e-30; + } + } + return .{ .want = want, .tol = tol }; +} diff --git a/zig/tests/cuda/exl3_experts_f64_tests.zig b/zig/tests/cuda/exl3_experts_f64_tests.zig new file mode 100644 index 0000000000..6e0c162cbb --- /dev/null +++ b/zig/tests/cuda/exl3_experts_f64_tests.zig @@ -0,0 +1,463 @@ +//! The EXL3 routed experts' float64 reference: the mirrored chain and each output's worst-case error bound. +const std = @import("std"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const cuda = @import("cuda"); +const exl3fmt = @import("core").exl3_format; +const exl3e = cuda.exl3_experts; +const exl3t = @import("exl3_tests.zig"); +const Rig = quant_ref_tests.RigT; +const bf = quant_ref_tests.helpers.bfOf; +const Run = exl3t.Run; + +pub const E = 9; // experts of a synthetic layer +pub const D = 512; // model width +const I = 256; // expert width +pub const SLOTS = 8; +const LIMIT: f32 = 10.0; // the SwiGLU clamp +const CAP = 300; // rows of x, pick and wts the builders make (the f64 checks use the first 128) +const rows = [_]usize{ 1, 3, 8, 16, 17, 64, 128 }; // the f64 checks' row counts (the reference builds 128 rows) + +const f64_ids = [_][]const u8{ "exl3.experts.f64.mul1.4", "exl3.experts.f64.mul1.5", "exl3.experts.f64.mul1.6" }; + +/// One synthetic layer: the device tensors, the host trellises the reference unpacks, and the seeded routing. +pub const Built = struct { + layer: exl3e.Layer, + bits: []const u32, + gate: [E][]i16, + up: [E][]i16, + down: [E][]i16, + suh_g: []u16, // fp16 [E, D] + suh_u: []u16, + svh_g: []u16, // fp16 [E, I] + svh_u: []u16, + suh_d: []u16, // fp16 [E, I] + svh_d: []u16, // fp16 [E, D] + dx: u64, + dpick: u64, + dwts: u64, + x: []u16, // bf16 [CAP, D] + pick: []i32, // [CAP, SLOTS] + wts: []f32, // [CAP, SLOTS] +}; + +/// The f64 id checks: every width alone, then the mixed layer folded into the .6 id (the contract's ids). +pub fn checks(r: *Rig, exp: *const exl3e.Experts, state: *Run) void { + for (f64_ids, 0..) |id, i| { + var ok = true; + f64Check(r, exp, &allBits(@intCast(i + 4))) catch |e| { + state.fail(e, id); + ok = false; + }; + if (ok) check.pass("{s}", .{id}); + } + { + var ok = true; + f64Check(r, exp, &.{ 4, 5, 6, 4, 5, 6, 4, 5, 6 }) catch |e| { + state.fail(e, f64_ids[2]); + ok = false; + }; + if (ok) check.pass("{s}", .{f64_ids[2]}); + } +} + +fn allBits(bits: u32) [E]u32 { + return .{ bits, bits, bits, bits, bits, bits, bits, bits, bits }; +} + +fn widthOf(bits: u32) exl3fmt.Width { + return .{ .halves = @intCast(2 * bits) }; +} + +fn f16v(v: u16) f64 { + return @floatCast(@as(f16, @bitCast(v))); +} + +/// bf16 round-to-nearest-even of an f64 (the kernel rounds its fp32 value; the double round is within the tolerance). +fn bf16r(v: f64) f64 { + const u: u32 = @bitCast(@as(f32, @floatCast(v))); + return bf(@intCast((u + 0x7fff + ((u >> 16) & 1)) >> 16)); +} + +pub fn buildLayer(r: *Rig, bits: []const u32) !Built { + var b: Built = .{ + .layer = undefined, + .bits = bits, + .gate = undefined, + .up = undefined, + .down = undefined, + .dx = 0, + .dpick = 0, + .dwts = 0, + .suh_g = undefined, + .suh_u = undefined, + .svh_g = undefined, + .svh_u = undefined, + .suh_d = undefined, + .svh_d = undefined, + .x = undefined, + .pick = undefined, + .wts = undefined, + }; + var gate_p: [E]u64 = undefined; + var up_p: [E]u64 = undefined; + var down_p: [E]u64 = undefined; + var gk2: [E]i32 = undefined; + var uk2: [E]i32 = undefined; + var dk2: [E]i32 = undefined; + const suh_g = try r.a.alloc(u16, E * D); + const suh_u = try r.a.alloc(u16, E * D); + const svh_g = try r.a.alloc(u16, E * I); + const svh_u = try r.a.alloc(u16, E * I); + const suh_d = try r.a.alloc(u16, E * I); + const svh_d = try r.a.alloc(u16, E * D); + for (0..E) |e| { + b.gate[e] = try randI16(r, (D / 16) * (I / 16) * widthOf(bits[e]).tileWords()); + b.up[e] = try randI16(r, (D / 16) * (I / 16) * widthOf(bits[e]).tileWords()); + b.down[e] = try randI16(r, (I / 16) * (D / 16) * widthOf(bits[e]).tileWords()); + gate_p[e] = try r.dev(std.mem.sliceAsBytes(b.gate[e])); + up_p[e] = try r.dev(std.mem.sliceAsBytes(b.up[e])); + down_p[e] = try r.dev(std.mem.sliceAsBytes(b.down[e])); + gk2[e] = @intCast(2 * bits[e]); + uk2[e] = gk2[e]; + dk2[e] = gk2[e]; + randScale(r, suh_g[e * D ..][0..D]); + randScale(r, suh_u[e * D ..][0..D]); + randScale(r, svh_g[e * I ..][0..I]); + randScale(r, svh_u[e * I ..][0..I]); + randScale(r, suh_d[e * I ..][0..I]); + randScale(r, svh_d[e * D ..][0..D]); + } + b.suh_g = suh_g; + b.suh_u = suh_u; + b.svh_g = svh_g; + b.svh_u = svh_u; + b.suh_d = suh_d; + b.svh_d = svh_d; + b.layer = .{ + .gate = try r.dev(std.mem.sliceAsBytes(gate_p[0..])), + .up = try r.dev(std.mem.sliceAsBytes(up_p[0..])), + .down = try r.dev(std.mem.sliceAsBytes(down_p[0..])), + .gate_k2 = try r.dev(std.mem.sliceAsBytes(gk2[0..])), + .up_k2 = try r.dev(std.mem.sliceAsBytes(uk2[0..])), + .down_k2 = try r.dev(std.mem.sliceAsBytes(dk2[0..])), + .suh_g = try r.dev(std.mem.sliceAsBytes(suh_g)), + .suh_u = try r.dev(std.mem.sliceAsBytes(suh_u)), + .svh_g = try r.dev(std.mem.sliceAsBytes(svh_g)), + .svh_u = try r.dev(std.mem.sliceAsBytes(svh_u)), + .suh_d = try r.dev(std.mem.sliceAsBytes(suh_d)), + .svh_d = try r.dev(std.mem.sliceAsBytes(svh_d)), + .experts = E, + .dims = D, + .width = I, + .codebook = .mul1, + .limit = LIMIT, + .act_mode = 0, + }; + b.x = try r.bfs(CAP * D, 0.5); // half-range x, with the small scales, keeps the SwiGLU inputs moderate + b.pick = try picks(r, CAP); + b.wts = try wtsOf(r, CAP); + b.dx = try r.dev(std.mem.sliceAsBytes(b.x)); + b.dpick = try r.dev(std.mem.sliceAsBytes(b.pick)); + b.dwts = try r.dev(std.mem.sliceAsBytes(b.wts)); + return b; +} + +fn randI16(r: *Rig, n: usize) ![]i16 { + const out = try r.a.alloc(i16, n); + for (out) |*v| v.* = r.rng.int(i16); + return out; +} + +/// Small scales (with the half-range x): the |t||w| sums and the SwiGLU inputs stay moderate for the bound. +fn randScale(r: *Rig, out: []u16) void { + for (out) |*v| v.* = @bitCast(@as(f16, @floatCast(r.rng.float(f32) * 0.25 - 0.125))); +} + +/// Distinct experts a row: a partial Fisher-Yates over the layer's experts, seeded. +fn picks(r: *Rig, cap: usize) ![]i32 { + const out = try r.a.alloc(i32, cap * SLOTS); + for (0..cap) |row| { + var idx: [E]u32 = undefined; + for (0..E, 0..) |_, i| idx[i] = @intCast(i); + var n: usize = E; + for (0..SLOTS) |k| { + const j = r.rng.uintLessThan(usize, n); + out[row * SLOTS + k] = @intCast(idx[j]); + idx[j] = idx[n - 1]; + n -= 1; + } + } + return out; +} + +fn wtsOf(r: *Rig, cap: usize) ![]f32 { + const out = try r.a.alloc(f32, cap * SLOTS); + for (out) |*v| v.* = 0.1 + 0.9 * r.rng.float(f32); + return out; +} + +/// One routed call over `rows` rows starting at row 0 of the built routing; the scratch and out are fresh canaries. +pub fn runRows(r: *Rig, exp: *const exl3e.Experts, b: *const Built, rows_n: usize) !u64 { + return (try runSlice(r, exp, b.layer, b.dx, b.dpick, b.dwts, rows_n)).out; +} + +pub fn runSlice(r: *Rig, exp: *const exl3e.Experts, layer: exl3e.Layer, x: u64, pick: u64, wts: u64, rows_n: usize) !struct { out: u64, y: u64 } { + const sb = try r.zeros(try exl3e.scratchBytes(layer, rows_n, SLOTS)); + const sc = try exl3e.scratch(sb, layer, rows_n, SLOTS); + const out = try r.zeros(rows_n * D * 4); + try exp.run(r.s, layer, x, D, .bf16, pick, wts, rows_n, out, sc); + return .{ .out = out, .y = sc.y }; +} + +fn f64Check(r: *Rig, exp: *const exl3e.Experts, bits: []const u32) !void { + const b = try buildLayer(r, bits); + const want = try reference(r, &b, 128); + for (rows) |m| try closeOut(r, b.bits, m, (try runRows(r, exp, &b, m)), want.want, want.tol); +} + +/// fp32 outputs within the f64 reference's tolerance, with the exact FAIL line. +fn closeOut(r: *Rig, bits: []const u32, m: usize, ptr: u64, want: []const f64, tol: []const f64) !void { + const raw = try r.back(ptr, m * D * 4); + var most: f64 = 0; + var at: usize = 0; + for (want[0 .. m * D], tol[0 .. m * D], 0..) |wv, t, i| { + const g: f64 = @floatCast(std.mem.bytesToValue(f32, raw[4 * i ..][0..4])); + const q = @abs(g - wv) / t; + if (q > most) { + most = q; + at = i; + } + } + std.debug.print("RESULT exl3.experts.f64.mul1.{d} m={d}: largest difference {d:.3} of its tolerance (element {d})\n", .{ bits[0], m, most, at }); + if (most > 1) { + std.debug.print("FAIL exl3.experts.f64.mul1.{d}: largest difference {d:.3} of its tolerance\n", .{ bits[0], most }); + return error.TestFailed; + } +} + +/// A 300-row chunk against the f64 reference (the mixed-width layer). +pub fn prompt(r: *Rig, exp: *const exl3e.Experts) !void { + var bits: [E]u32 = .{ 4, 5, 6, 4, 5, 6, 4, 5, 6 }; + const b = try buildLayer(r, &bits); + const want = try reference(r, &b, 300); + const res = try runRows(r, exp, &b, 300); + const raw = try r.back(res, 300 * D * 4); + var most: f64 = 0; + var at: usize = 0; + for (want.want, want.tol, 0..) |wv, t, i| { + const g: f64 = @floatCast(std.mem.bytesToValue(f32, raw[4 * i ..][0..4])); + const q = @abs(g - wv) / t; + if (q > most) { + most = q; + at = i; + } + } + std.debug.print("RESULT exl3.prompt.experts.f64: largest difference {d:.3} of its tolerance (element {d})\n", .{ most, at }); + if (most > 1) return error.TestFailed; +} + +/// The f64 reference: the kernel's chain mirrored in f64, each output's error its per-element worst case. +fn reference(r: *Rig, b: *const Built, rows_n: usize) !struct { want: []f64, tol: []f64 } { + const p = rows_n * SLOTS; + // xg: fp16((x * suh_g) @ H) per slot; |xg| and its error <= V, the block's sum of |x * suh| over sqrt(128) + const xs = try r.a.alloc(f64, p * D); + for (0..p) |s0| { + const ex: usize = @intCast(b.pick[s0]); + for (0..D) |i| xs[s0 * D + i] = bf(b.x[(s0 / SLOTS) * D + i]) * f16v(b.suh_g[ex * D + i]); + } + exl3fmt.rotate(xs, p, D, 1); + const xg = try r.a.alloc(f16, p * D); + for (xg, xs) |*o, v| o.* = @floatCast(v); + const xgerr = try r.a.alloc(f64, p * D); + for (xgerr, xs) |*o, v| o.* = @abs(v); + absRotateCols(xgerr, p, D); + for (xgerr) |*v| v.* *= 1.0 / 2048.0 + 8.0 / 16777216.0; // the fp16 rounding (2^-11) plus the fp32 rotation's (8 * 2^-24), both of V + + const wq = try r.a.alloc(f16, @max(D, I) * @max(D, I)); + const zg = try r.a.alloc(f64, I * p); // [n][p] + const zu = try r.a.alloc(f64, I * p); + const s1g = try r.a.alloc(f64, I * p); + const s1u = try r.a.alloc(f64, I * p); + const eg = try r.a.alloc(f64, I * p); // the GEMM's propagated input error, before the output rotation + const eu = try r.a.alloc(f64, I * p); + @memset(zg, 0); + @memset(zu, 0); + @memset(s1g, 0); + @memset(s1u, 0); + @memset(eg, 0); + @memset(eu, 0); + for (0..E) |ex| { + exl3fmt.unpack(widthOf(b.bits[ex]), .mul1, b.gate[ex], D / 16, I / 16, wq[0 .. D * I]); + for (0..p) |s0| { + if (b.pick[s0] != @as(i32, @intCast(ex))) continue; + gemmK(xg[s0 * D ..][0..D], wq, zg[s0..], s1g[s0..], eg[s0..], xgerr[s0 * D ..][0..D], D, I, p); + } + exl3fmt.unpack(widthOf(b.bits[ex]), .mul1, b.up[ex], D / 16, I / 16, wq[0 .. D * I]); + for (0..p) |s0| { + if (b.pick[s0] != @as(i32, @intCast(ex))) continue; + gemmK(xg[s0 * D ..][0..D], wq, zu[s0..], s1u[s0..], eu[s0..], xgerr[s0 * D ..][0..D], D, I, p); + } + } + // the fp16 rounding (2^-11) and the fp32 accumulation (2^-17 of 128 roundings), both on the |t||w| sum + for (eg, s1g) |*e, s| e.* += s * (1.0 / 2048.0 + 1.0 / 131072.0); + for (eu, s1u) |*e, s| e.* += s * (1.0 / 2048.0 + 1.0 / 131072.0); + exl3fmt.rotate(zg, I, p, 0); + exl3fmt.rotate(zu, I, p, 0); + // the output Hadamard: worst case is the block's sum of |errors| over sqrt(128), plus the fp32 rotation's own + absRotate(eg, I, p); + absRotate(eu, I, p); + const zgerr = try r.a.alloc(f64, I * p); + for (zgerr, zg) |*o, v| o.* = @abs(v); + absRotate(zgerr, I, p); + for (zgerr) |*v| v.* *= 8.0 / 16777216.0; + const zuerr = try r.a.alloc(f64, I * p); + for (zuerr, zu) |*o, v| o.* = @abs(v); + absRotate(zuerr, I, p); + for (zuerr) |*v| v.* *= 8.0 / 16777216.0; + for (eg, zgerr) |*e, v| e.* += v; + for (eu, zuerr) |*e, v| e.* += v; + + // the SwiGLU (bf16 roundings, f64 exp): g's error is the rotated GEMM error through svh plus a bf16 double rounding + const v = try r.a.alloc(f64, I * p); + const verr = try r.a.alloc(f64, I * p); + for (0..I) |n| { + for (0..p) |s0| { + const ex: usize = @intCast(b.pick[s0]); + const vg = zg[n * p + s0] * f16v(b.svh_g[ex * I + n]); + const gg = @min(bf16r(vg), @as(f64, LIMIT)); + const vu = zu[n * p + s0] * f16v(b.svh_u[ex * I + n]); + const uu = @min(@max(bf16r(vu), -@as(f64, LIMIT)), @as(f64, LIMIT)); + const inner = bf16r(gg / (1.0 + std.math.exp(-gg))); + const act = bf16r(inner * uu); + v[n * p + s0] = act * f16v(b.suh_d[ex * I + n]); + const egk = eg[n * p + s0] * @abs(f16v(b.svh_g[ex * I + n])) + (@abs(vg) + @abs(gg)) / 256; + const euk = eu[n * p + s0] * @abs(f16v(b.svh_u[ex * I + n])) + (@abs(vu) + @abs(uu)) / 256; + const einner = egk + (@abs(gg) + 1) / 524288 + @abs(inner) / 256; + verr[n * p + s0] = (@abs(uu) * einner + @abs(inner) * euk + @abs(act) / 256) * @abs(f16v(b.suh_d[ex * I + n])) + + @abs(v[n * p + s0]) / 16777216; // act's: |silu'| (<= 1) g's, |silu(g)| (<= |g|) u's, expf and roundings + } + } + // the down input: the fp32 rotation's roundings (of the block's sum of |v|) and the fp16 store (2^-11) + exl3fmt.rotate(v, I, p, 0); + absRotate(verr, I, p); + const vrot = try r.a.alloc(f64, I * p); + for (vrot, v) |*o, val| o.* = @abs(val); + absRotate(vrot, I, p); + for (vrot) |*o| o.* *= 8.0 / 16777216.0; + const xd = try r.a.alloc(f16, p * I); + const xderr = try r.a.alloc(f64, p * I); + for (0..p) |s0| { + for (0..I) |ki| { + xd[s0 * I + ki] = @floatCast(v[ki * p + s0]); + xderr[s0 * I + ki] = verr[ki * p + s0] + vrot[ki * p + s0] + @abs(xd[s0 * I + ki]) / 2048; + } + } + + // down: the GEMM, its |t| and |w| sums, and the propagated input error + const zd = try r.a.alloc(f64, D * p); // [d][p] + const s1d = try r.a.alloc(f64, D * p); + const de = try r.a.alloc(f64, D * p); + @memset(zd, 0); + @memset(s1d, 0); + @memset(de, 0); + for (0..E) |ex| { + exl3fmt.unpack(widthOf(b.bits[ex]), .mul1, b.down[ex], I / 16, D / 16, wq[0 .. I * D]); + for (0..p) |s0| { + if (b.pick[s0] != @as(i32, @intCast(ex))) continue; + for (0..I) |ki| { + const t: f64 = @floatCast(xd[s0 * I + ki]); + const ta = @abs(t); + const row = wq[ki * D ..][0..D]; + const xek = xderr[s0 * I + ki]; + for (0..D) |d| { + const wv: f64 = @floatCast(row[d]); + zd[d * p + s0] += t * wv; + s1d[d * p + s0] += ta * @abs(wv); + de[d * p + s0] += @abs(wv) * xek; + } + } + } + } + // the fp16 xd rounding (2^-11 of the |t||w| sum) and the fp32 accumulation (2^-17), then the output Hadamard's mix + for (de, s1d) |*e, s| e.* += s * (1.0 / 2048.0 + 1.0 / 131072.0); + exl3fmt.rotate(zd, D, p, 0); + absRotate(de, D, p); + const zderr = try r.a.alloc(f64, D * p); + for (zderr, zd) |*o, val| o.* = @abs(val); + absRotate(zderr, D, p); + for (zderr) |*o| o.* *= 8.0 / 16777216.0; + for (de, zderr) |*e, val| e.* += val; + + // the per-slot outputs and the weighted combine: the fp32 fma chain (8 roundings of the running sum) and the store + const want = try r.a.alloc(f64, rows_n * D); + const tol = try r.a.alloc(f64, rows_n * D); + for (0..rows_n) |row| { + for (0..D) |d| { + var acc: f32 = 0; + var err: f64 = 0; + var swy: f64 = 0; + for (0..SLOTS) |k| { + const s0 = row * SLOTS + k; + const ex: usize = @intCast(b.pick[s0]); + const y: f32 = @floatCast(zd[d * p + s0] * f16v(b.svh_d[ex * D + d])); + acc = @mulAdd(f32, b.wts[row * SLOTS + k], y, acc); + const yerr = de[d * p + s0] * @abs(f16v(b.svh_d[ex * D + d])) + @abs(y) / 16777216; + err += @abs(b.wts[row * SLOTS + k]) * yerr; + swy += @abs(b.wts[row * SLOTS + k] * y); + } + want[row * D + d] = acc; + tol[row * D + d] = err + swy / 2097152 + @abs(want[row * D + d]) / 16777216 + 1e-30; + } + } + return .{ .want = want, .tol = tol }; +} + +/// One output column's GEMM over k: the sums into [n][p], their |t| sums, the propagated input error. +fn gemmK(xrow: []const f16, wq: []const f16, z: []f64, s1: []f64, e: []f64, xerr: []const f64, k: usize, n: usize, p: usize) void { + for (0..k) |ki| { + const t: f64 = @floatCast(xrow[ki]); + const ta = @abs(t); + const wrow = wq[ki * n ..][0..n]; + const xek = xerr[ki]; + for (0..n) |col| { + const wv: f64 = @floatCast(wrow[col]); + z[col * p] += t * wv; + s1[col * p] += ta * @abs(wv); + e[col * p] += @abs(wv) * xek; + } + } +} + +/// Each 128-block along the row axis replaced by the sum of its |entries| over sqrt(128) (the worst case). +fn absRotate(x: []f64, nrows: usize, cols: usize) void { + const inv = 1.0 / @sqrt(@as(f64, 128.0)); + var r: usize = 0; + while (r < nrows) : (r += 128) { + for (0..cols) |c| { + var s: f64 = 0; + for (0..128) |n| s += @abs(x[(r + n) * cols + c]); + s *= inv; + for (0..128) |n| x[(r + n) * cols + c] = s; + } + } +} + +test "the ids are the contract's" { + try std.testing.expectEqualStrings("exl3.experts.f64.mul1.4", f64_ids[0]); + try std.testing.expectEqualStrings("exl3.experts.f64.mul1.6", f64_ids[2]); +} + +/// The column-axis twin: each 128-run along the columns replaced by the sum of its |entries| over sqrt(128). +fn absRotateCols(x: []f64, nrows: usize, cols: usize) void { + const inv = 1.0 / @sqrt(@as(f64, 128.0)); + for (0..nrows) |row| { + var c: usize = 0; + while (c < cols) : (c += 128) { + var s: f64 = 0; + for (0..128) |n| s += @abs(x[row * cols + c + n]); + s *= inv; + for (0..128) |n| x[row * cols + c + n] = s; + } + } +} diff --git a/zig/tests/cuda/exl3_experts_tests.zig b/zig/tests/cuda/exl3_experts_tests.zig new file mode 100644 index 0000000000..20710d303f --- /dev/null +++ b/zig/tests/cuda/exl3_experts_tests.zig @@ -0,0 +1,197 @@ +//! The EXL3 routed experts: the f64 checks, row and chunk invariance, and the Python engine's recorded bytes. +const std = @import("std"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const cuda = @import("cuda"); +const exl3 = cuda.exl3; +const exl3e = cuda.exl3_experts; +const exl3t = @import("exl3_tests.zig"); +const ref = @import("exl3_experts_f64_tests.zig"); +const Rig = quant_ref_tests.RigT; +const Run = exl3t.Run; + +pub fn run(r: *Rig, state: *Run, dir: []const u8) !void { + var exp = try exl3e.Experts.load(r.gpu.d); + defer exp.unload(); + + // every expert one width: 4, 5 and 6 bits (the mixed run folds into .6) + ref.checks(r, &exp, state); + { + var ok = true; + rowsInvariant(r, &exp) catch |e| { + state.fail(e, "exl3.experts.rows"); + ok = false; + }; + if (ok) check.pass("exl3.experts.rows", .{}); + } + { + var ok = true; + ref.prompt(r, &exp) catch |e| { + state.fail(e, "exl3.prompt.experts.f64"); + ok = false; + }; + if (ok) check.pass("exl3.prompt.experts.f64", .{}); + } + { + var ok = true; + chunks(r, &exp) catch |e| { + state.fail(e, "exl3.prompt.experts.chunks"); + ok = false; + }; + if (ok) check.pass("exl3.prompt.experts.chunks", .{}); + } + { + var ok = true; + python(r, &exp, dir, state) catch |e| { + state.fail(e, "exl3.experts.python"); + ok = false; + }; + if (ok) check.pass("exl3.experts.python", .{}); + } +} + +/// A row's output bytes are its own: solo, with 3 rows, in 64-row batches, shared or not. +fn rowsInvariant(r: *Rig, exp: *const exl3e.Experts) !void { + var bits: [ref.E]u32 = .{ 6, 4, 5, 6, 4, 5, 6, 4, 5 }; + const b = try ref.buildLayer(r, &bits); + const solo = try r.back(try ref.runRows(r, exp, &b, 1), ref.D * 4); + // 3 rows: two more rows on the same picks as row 0 + const p3 = try r.a.alloc(i32, 3 * ref.SLOTS); + @memcpy(p3[0 .. 2 * ref.SLOTS], b.pick[0 .. 2 * ref.SLOTS]); + @memcpy(p3[2 * ref.SLOTS ..], b.pick[0..ref.SLOTS]); + const w3 = try r.a.alloc(f32, 3 * ref.SLOTS); + @memcpy(w3, b.wts[0 .. 3 * ref.SLOTS]); + const got3 = try r.back((try ref.runSlice(r, exp, b.layer, b.dx, try r.dev(std.mem.sliceAsBytes(p3)), try r.dev(std.mem.sliceAsBytes(w3)), 3)).out, ref.D * 4); + try check.sameBytes("exl3.experts.rows 3 rows", got3, solo); + // 64 rows: the others share row 0's experts (the same picks), or run their own + const p64 = try r.a.alloc(i32, 64 * ref.SLOTS); + const w64 = try r.a.alloc(f32, 64 * ref.SLOTS); + @memcpy(p64[0..ref.SLOTS], b.pick[0..ref.SLOTS]); + @memcpy(w64[0..ref.SLOTS], b.wts[0..ref.SLOTS]); + for (1..64) |row| { + @memcpy(p64[row * ref.SLOTS ..][0..ref.SLOTS], b.pick[0..ref.SLOTS]); + @memcpy(w64[row * ref.SLOTS ..][0..ref.SLOTS], b.wts[0..ref.SLOTS]); + } + const shared = try r.back((try ref.runSlice(r, exp, b.layer, b.dx, try r.dev(std.mem.sliceAsBytes(p64)), try r.dev(std.mem.sliceAsBytes(w64)), 64)).out, ref.D * 4); + try check.sameBytes("exl3.experts.rows 64 rows shared", shared[0 .. ref.D * 4], solo); + const own = try r.back((try ref.runSlice(r, exp, b.layer, b.dx, b.dpick, b.dwts, 64)).out, ref.D * 4); + try check.sameBytes("exl3.experts.rows 64 rows own", own[0 .. ref.D * 4], solo); + // and a 64-row batch whose other rows interleave with row 0's experts in both orders + const p64b = try r.a.alloc(i32, 64 * ref.SLOTS); + const w64b = try r.a.alloc(f32, 64 * ref.SLOTS); + @memcpy(p64b[0..ref.SLOTS], b.pick[0..ref.SLOTS]); + @memcpy(w64b[0..ref.SLOTS], b.wts[0..ref.SLOTS]); + for (1..64) |row| { + const src = if (row % 2 == 0) @as(usize, 0) else row % ref.SLOTS; + @memcpy(p64b[row * ref.SLOTS ..][0..ref.SLOTS], b.pick[src * ref.SLOTS ..][0..ref.SLOTS]); + @memcpy(w64b[row * ref.SLOTS ..][0..ref.SLOTS], b.wts[src * ref.SLOTS ..][0..ref.SLOTS]); + } + const mixed = try r.back((try ref.runSlice(r, exp, b.layer, b.dx, try r.dev(std.mem.sliceAsBytes(p64b)), try r.dev(std.mem.sliceAsBytes(w64b)), 64)).out, ref.D * 4); + try check.sameBytes("exl3.experts.rows 64 rows mixed", mixed[0 .. ref.D * 4], solo); +} + +/// The same 300 rows as one call, split at 128 and at 37: a row's bytes never move. +fn chunks(r: *Rig, exp: *const exl3e.Experts) !void { + var bits: [ref.E]u32 = .{ 4, 5, 6, 4, 5, 6, 4, 5, 6 }; + const b = try ref.buildLayer(r, &bits); + const whole = try r.back(try ref.runRows(r, exp, &b, 300), 300 * ref.D * 4); + var at: usize = 0; + while (at < 300) : (at += 128) { + const n: usize = @min(300 - at, 128); + const res = try ref.runSlice(r, exp, b.layer, b.dx + @as(u64, at) * ref.D * 2, b.dpick + @as(u64, at) * ref.SLOTS * 4, b.dwts + @as(u64, at) * ref.SLOTS * 4, n); + const got = try r.back(res.out, n * ref.D * 4); + try check.sameBytes("exl3.prompt.experts.chunks 128", got, whole[at * ref.D * 4 ..][0 .. n * ref.D * 4]); + } + at = 0; + while (at < 300) : (at += 37) { + const n: usize = @min(300 - at, 37); + const res = try ref.runSlice(r, exp, b.layer, b.dx + @as(u64, at) * ref.D * 2, b.dpick + @as(u64, at) * ref.SLOTS * 4, b.dwts + @as(u64, at) * ref.SLOTS * 4, n); + const got = try r.back(res.out, n * ref.D * 4); + try check.sameBytes("exl3.prompt.experts.chunks 37", got, whole[at * ref.D * 4 ..][0 .. n * ref.D * 4]); + } +} + +/// The Python engine's recorded y and out, bit for bit (exl3_oracle_experts.py's manifest and raw files). +fn python(r: *Rig, exp: *const exl3e.Experts, dir: []const u8, state: *Run) !void { + if (std.mem.eql(u8, dir, "none")) { + std.debug.print("SKIP exl3.experts.python\n", .{}); + return; + } + const manifest_path = try std.fs.path.join(r.a, &.{ dir, "manifest.json" }); + const text = std.Io.Dir.cwd().readFileAlloc(r.gpu.io, manifest_path, r.a, .limited(1 << 24)) catch |e| { + std.debug.print("FAIL exl3.experts.python: {t}\n", .{e}); + state.failed = true; + return; + }; + var parsed = try std.json.parseFromSlice(std.json.Value, r.a, text, .{}); + const o = parsed.value.object.get("experts").?.object; + const e_count: usize = @intCast(o.get("count").?.integer); + const dims: usize = @intCast(o.get("dims").?.integer); + const width: usize = @intCast(o.get("width").?.integer); + const slots: usize = @intCast(o.get("slots").?.integer); + const cb: exl3.Codebook = if (std.mem.eql(u8, o.get("codebook").?.string, "mcg")) + .mcg + else if (std.mem.eql(u8, o.get("codebook").?.string, "inst3")) + .inst3 + else + .mul1; + const trellis = o.get("trellis").?.object; + var gate_p = try r.a.alloc(u64, e_count); + var up_p = try r.a.alloc(u64, e_count); + var down_p = try r.a.alloc(u64, e_count); + var gk2 = try r.a.alloc(i32, e_count); + var uk2 = try r.a.alloc(i32, e_count); + var dk2 = try r.a.alloc(i32, e_count); + for (o.get("gate_k2").?.array.items, 0..) |v, i| gk2[i] = @intCast(v.integer); + for (o.get("up_k2").?.array.items, 0..) |v, i| uk2[i] = @intCast(v.integer); + for (o.get("down_k2").?.array.items, 0..) |v, i| dk2[i] = @intCast(v.integer); + for (0..e_count) |i| { + gate_p[i] = try r.dev(try readFile(r, dir, trellis.get("gate").?.array.items[i].string)); + up_p[i] = try r.dev(try readFile(r, dir, trellis.get("up").?.array.items[i].string)); + down_p[i] = try r.dev(try readFile(r, dir, trellis.get("down").?.array.items[i].string)); + } + const layer: exl3e.Layer = .{ + .gate = try r.dev(std.mem.sliceAsBytes(gate_p)), + .up = try r.dev(std.mem.sliceAsBytes(up_p)), + .down = try r.dev(std.mem.sliceAsBytes(down_p)), + .gate_k2 = try r.dev(std.mem.sliceAsBytes(gk2)), + .up_k2 = try r.dev(std.mem.sliceAsBytes(uk2)), + .down_k2 = try r.dev(std.mem.sliceAsBytes(dk2)), + .suh_g = try r.dev(try readFile(r, dir, o.get("suh_g").?.string)), + .suh_u = try r.dev(try readFile(r, dir, o.get("suh_u").?.string)), + .svh_g = try r.dev(try readFile(r, dir, o.get("svh_g").?.string)), + .svh_u = try r.dev(try readFile(r, dir, o.get("svh_u").?.string)), + .suh_d = try r.dev(try readFile(r, dir, o.get("suh_d").?.string)), + .svh_d = try r.dev(try readFile(r, dir, o.get("svh_d").?.string)), + .experts = e_count, + .dims = dims, + .width = width, + .codebook = cb, + .limit = @floatCast(o.get("limit").?.float), + .act_mode = @intCast(o.get("act_mode").?.integer), + }; + const xmap = o.get("x").?.object; + const pickmap = o.get("pick").?.object; + const wtsmap = o.get("wts").?.object; + const ymap = o.get("y").?.object; + const outmap = o.get("out").?.object; + for (o.get("rows").?.array.items) |rv| { + const m: usize = @intCast(rv.integer); + const key = try std.fmt.allocPrint(r.a, "{d}", .{m}); + const dx = try r.dev(try readFile(r, dir, xmap.get(key).?.string)); + const dpick = try r.dev(try readFile(r, dir, pickmap.get(key).?.string)); + const dwts = try r.dev(try readFile(r, dir, wtsmap.get(key).?.string)); + const sb = try r.zeros(try exl3e.scratchBytes(layer, m, slots)); + const sc = try exl3e.scratch(sb, layer, m, slots); + const out = try r.zeros(m * dims * 4); + try exp.run(r.s, layer, dx, dims, .bf16, dpick, dwts, m, out, sc); + const got_y = try r.back(sc.y, m * slots * dims * 4); + try check.sameBytes("exl3.experts.python y", got_y, try readFile(r, dir, ymap.get(key).?.string)); + const got_out = try r.back(out, m * dims * 4); + try check.sameBytes("exl3.experts.python out", got_out, try readFile(r, dir, outmap.get(key).?.string)); + } +} + +fn readFile(r: *Rig, dir: []const u8, name: []const u8) ![]u8 { + return std.Io.Dir.cwd().readFileAlloc(r.gpu.io, try std.fs.path.join(r.a, &.{ dir, name }), r.a, .limited(1 << 32)); +} diff --git a/zig/tests/cuda/exl3_layer_tests.zig b/zig/tests/cuda/exl3_layer_tests.zig new file mode 100644 index 0000000000..4282f83ecd --- /dev/null +++ b/zig/tests/cuda/exl3_layer_tests.zig @@ -0,0 +1,156 @@ +//! EXL3 whole layers: a KDA and a DSA MoE layer of real RED-SNOW through core/exl3_load vs the engine's fixture bytes. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const exl3_load = @import("core").exl3_load; +const exl3 = cuda.exl3; +const exl3e = cuda.exl3_experts; +const exl3_upload = cuda.exl3_upload; +const exl3_tests = @import("exl3_tests.zig"); +const Run = exl3_tests.Run; +const Rig = quant_ref_tests.RigT; + +pub fn run(r: *Rig, state: *Run, dir: []const u8) !void { + if (std.mem.eql(u8, dir, "none")) { + std.debug.print("SKIP exl3.layer.python\n", .{}); + return; + } + const o = manifest(r, dir) catch |e| { + std.debug.print("FAIL exl3.layer.python: {t}\n", .{e}); + state.failed = true; + return; + }; + const layers = o.get("layers") orelse { + std.debug.print("SKIP exl3.layer.python: no layers entry\n", .{}); + return; + }; + const model = o.get("model").?.string; + var scan = exl3_load.scan(r.gpu.gpa, r.gpu.io, model) catch |e| { + std.debug.print("FAIL exl3.layer.python: {t}\n", .{e}); + state.failed = true; + return; + }; + defer scan.deinit(); + var idx = std.StringHashMapUnmanaged(usize).empty; + for (scan.groups, 0..) |g, i| try idx.put(r.a, g.prefix, i); + var dense = try exl3.Dense.load(r.gpu.d); + defer dense.unload(); + var exp = try exl3e.Experts.load(r.gpu.d); + defer exp.unload(); + var ok = true; + layerChecks(r, &scan, &idx, &dense, &exp, dir, layers.array.items) catch |e| { + state.fail(e, "exl3.layer.python"); + ok = false; + }; + if (ok) check.pass("exl3.layer.python", .{}); +} + +fn layerChecks(r: *Rig, scan: *exl3_load.Scan, idx: *std.StringHashMapUnmanaged(usize), dense: *exl3.Dense, exp: *const exl3e.Experts, dir: []const u8, layers: []std.json.Value) !void { + const t0 = check.now(r.gpu.io); + var total: usize = 0; + for (layers) |item| { + const o = item.object; + const id = o.get("id").?.string; + var up: exl3_upload.Uploader = .{ .d = r.gpu.d, .gpa = r.gpu.gpa }; + defer up.deinit(); + for (o.get("groups").?.array.items) |gitem| try denseGroup(r, scan, idx, dense, dir, &up, gitem.object); + try expertLayer(r, scan, idx, exp, dir, &up, o.get("experts").?.object); + std.debug.print("RESULT exl3.layer.python.{s}: {d:.1} MiB of weights on the device\n", .{ id, @as(f64, @floatFromInt(up.bytes)) / (1 << 20) }); + total += up.bytes; + } + const dt = @as(f64, @floatFromInt(check.now(r.gpu.io) - t0)) / 1e9; + std.debug.print("RESULT exl3.layer.python: {d} layers, {d:.1} MiB uploaded, {d:.2} s\n", .{ layers.len, @as(f64, @floatFromInt(total)) / (1 << 20), dt }); +} + +/// One non-routed group: checkpoint tensors through the uploader, fixture inputs and bytes through the prompt path. +fn denseGroup(r: *Rig, scan: *exl3_load.Scan, idx: *std.StringHashMapUnmanaged(usize), dense: *exl3.Dense, dir: []const u8, up: *exl3_upload.Uploader, o: std.json.ObjectMap) !void { + const name = o.get("name").?.string; + const gi = idx.get(name) orelse return error.TestFailed; + const g = scan.groups[gi]; + try check.expect(g.halves == 2 * bitsOf(o.get("bits").?) and g.k == @as(usize, @intCast(o.get("k").?.integer)) and + g.n == @as(usize, @intCast(o.get("n").?.integer)), "exl3.layer.python {s}: the checkpoint group is not the fixture's", .{name}); + const parts = try scan.parts(r.a, gi); + const w = try up.dense(parts, g); + const xmap = o.get("x").?.object; + const ymap = o.get("y").?.object; + for (o.get("rows").?.array.items) |rv| { + const m: usize = @intCast(rv.integer); + const key = try std.fmt.allocPrint(r.a, "{d}", .{m}); + const dx = try r.dev(try readFile(r, dir, xmap.get(key).?.string)); + const want = try readFile(r, dir, ymap.get(key).?.string); + const dy = try r.zeros(m * g.n * 2); + try dense.prompt(r.s, w, dx, .bf16, m, dy, .bf16); + const got = try r.back(dy, want.len); + try check.sameBytes(try std.fmt.allocPrint(r.a, "exl3.layer.python {s} {d} rows", .{ name, m }), got, want); + } +} + +/// One layer's routed experts: all E experts' projections uploaded, the fixture's routing and inputs per row count. +fn expertLayer(r: *Rig, scan: *exl3_load.Scan, idx: *std.StringHashMapUnmanaged(usize), exp: *const exl3e.Experts, dir: []const u8, up: *exl3_upload.Uploader, o: std.json.ObjectMap) !void { + const base = o.get("name").?.string; + const e_count: usize = @intCast(o.get("count").?.integer); + const dims: usize = @intCast(o.get("dims").?.integer); + const width: usize = @intCast(o.get("width").?.integer); + const slots: usize = @intCast(o.get("slots").?.integer); + var gate = try r.a.alloc(exl3_upload.Proj, e_count); + var upj = try r.a.alloc(exl3_upload.Proj, e_count); + var down = try r.a.alloc(exl3_upload.Proj, e_count); + const shorts = [_][]const u8{ "gate", "up", "down" }; + const parts_names = [_][]const u8{ "gate_proj", "up_proj", "down_proj" }; + for (shorts, parts_names) |short, part| { + const k2map = o.get(try std.fmt.allocPrint(r.a, "{s}_k2", .{short})).?.array; + const list: *[]exl3_upload.Proj = switch (short[0]) { + 'g' => &gate, + 'u' => &upj, + else => &down, + }; + for (0..e_count) |e| { + const prefix = try std.fmt.allocPrint(r.a, "{s}.{d}.{s}", .{ base, e, part }); + const gi = idx.get(prefix) orelse return error.TestFailed; + const g = scan.groups[gi]; + try check.expect(g.halves == bitsOf(k2map.items[e]), "exl3.layer.python {s}: expert {d} is {d} halves, the fixture says {d}", .{ prefix, e, g.halves, bitsOf(k2map.items[e]) }); + const parts = try scan.parts(r.a, gi); + list.*[e] = .{ .trellis = parts.trellis.bytes, .suh = parts.suh.bytes, .svh = parts.svh.bytes, .k2 = g.halves }; + } + } + const l = try exl3_upload.layer(up, r.a, dims, width, .mul1, @floatCast(o.get("limit").?.float), @intCast(o.get("act_mode").?.integer), gate, upj, down); + const xmap = o.get("x").?.object; + const pickmap = o.get("pick").?.object; + const wtsmap = o.get("wts").?.object; + const ymap = o.get("y").?.object; + const outmap = o.get("out").?.object; + for (o.get("rows").?.array.items) |rv| { + const m: usize = @intCast(rv.integer); + const key = try std.fmt.allocPrint(r.a, "{d}", .{m}); + const dx = try r.dev(try readFile(r, dir, xmap.get(key).?.string)); + const dpick = try r.dev(try readFile(r, dir, pickmap.get(key).?.string)); + const dwts = try r.dev(try readFile(r, dir, wtsmap.get(key).?.string)); + const sb = try r.zeros(try exl3e.scratchBytes(l, m, slots)); + const sc = try exl3e.scratch(sb, l, m, slots); + const out = try r.zeros(m * dims * 4); + try exp.run(r.s, l, dx, dims, .bf16, dpick, dwts, m, out, sc); + const got_out = try r.back(out, m * dims * 4); + try check.sameBytes(try std.fmt.allocPrint(r.a, "exl3.layer.python {s} experts {d} rows out", .{ base, m }), got_out, try readFile(r, dir, outmap.get(key).?.string)); + const got_y = try r.back(sc.y, m * slots * dims * 4); + try check.sameBytes(try std.fmt.allocPrint(r.a, "exl3.layer.python {s} experts {d} rows y", .{ base, m }), got_y, try readFile(r, dir, ymap.get(key).?.string)); + } +} + +fn manifest(r: *Rig, dir: []const u8) !std.json.ObjectMap { + const text = try std.Io.Dir.cwd().readFileAlloc(r.gpu.io, try std.fs.path.join(r.a, &.{ dir, "manifest.json" }), r.a, .limited(1 << 24)); + const parsed = try std.json.parseFromSlice(std.json.Value, r.a, text, .{}); + return parsed.value.object; +} + +fn bitsOf(v: std.json.Value) u32 { + return switch (v) { + .integer => |i| @intCast(i), + .float => |f| @intFromFloat(f), + else => unreachable, + }; +} + +fn readFile(r: *Rig, dir: []const u8, name: []const u8) ![]u8 { + return std.Io.Dir.cwd().readFileAlloc(r.gpu.io, try std.fs.path.join(r.a, &.{ dir, name }), r.a, .limited(1 << 30)); +} diff --git a/zig/tests/cuda/exl3_tests.zig b/zig/tests/cuda/exl3_tests.zig new file mode 100644 index 0000000000..2f21c5e16e --- /dev/null +++ b/zig/tests/cuda/exl3_tests.zig @@ -0,0 +1,400 @@ +//! EXL3: the decode-row linear against float64 references, row invariance, the engine's fixture bytes, bandwidth. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const exl3fmt = @import("core").exl3_format; +const exl3 = cuda.exl3; +const Gpu = check.Gpu; +const Rig = quant_ref_tests.RigT; +const bf = quant_ref_tests.helpers.bfOf; + +const rows = [_]usize{ 1, 3, 8, 16, 17, 64, 128 }; // every row count in the f64 checks (128 covers the rest) + +const Cfg = struct { id: []const u8, halves: u32, codebook: exl3fmt.Codebook, k: usize, n: usize }; + +const f64_cfgs = [_]Cfg{ + .{ .id = "exl3.linear.f64.mul1.4", .halves = 8, .codebook = .mul1, .k = 4096, .n = 1024 }, + .{ .id = "exl3.linear.f64.mul1.4", .halves = 2, .codebook = .mul1, .k = 512, .n = 512 }, // 1 bit + .{ .id = "exl3.linear.f64.mul1.4", .halves = 4, .codebook = .mul1, .k = 512, .n = 512 }, // 2 bits + .{ .id = "exl3.linear.f64.mul1.5", .halves = 10, .codebook = .mul1, .k = 1024, .n = 4096 }, + .{ .id = "exl3.linear.f64.mul1.5", .halves = 6, .codebook = .mul1, .k = 512, .n = 512 }, // 3 bits + .{ .id = "exl3.linear.f64.mul1.6", .halves = 12, .codebook = .mul1, .k = 2048, .n = 8192 }, // 8 warps a block + .{ .id = "exl3.linear.f64.mul1.7", .halves = 14, .codebook = .mul1, .k = 512, .n = 1024 }, + .{ .id = "exl3.linear.f64.mul1.8", .halves = 16, .codebook = .mul1, .k = 2048, .n = 4096 }, + .{ .id = "exl3.linear.f64.mcg", .halves = 8, .codebook = .mcg, .k = 1024, .n = 512 }, + .{ .id = "exl3.linear.f64.mcg", .halves = 2, .codebook = .mcg, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.mcg", .halves = 4, .codebook = .mcg, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.mcg", .halves = 6, .codebook = .mcg, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.mcg", .halves = 10, .codebook = .mcg, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.mcg", .halves = 12, .codebook = .mcg, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.mcg", .halves = 14, .codebook = .mcg, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.mcg", .halves = 16, .codebook = .mcg, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.inst3", .halves = 8, .codebook = .inst3, .k = 1024, .n = 1024 }, + .{ .id = "exl3.linear.f64.inst3", .halves = 2, .codebook = .inst3, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.inst3", .halves = 4, .codebook = .inst3, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.inst3", .halves = 6, .codebook = .inst3, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.inst3", .halves = 10, .codebook = .inst3, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.inst3", .halves = 12, .codebook = .inst3, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.inst3", .halves = 14, .codebook = .inst3, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.inst3", .halves = 16, .codebook = .inst3, .k = 512, .n = 512 }, + // the mul1 half steps, 1.5 to 6.5 bits, under one check + .{ .id = "exl3.linear.f64.halfstep", .halves = 3, .codebook = .mul1, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.halfstep", .halves = 5, .codebook = .mul1, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.halfstep", .halves = 7, .codebook = .mul1, .k = 512, .n = 512 }, + .{ .id = "exl3.linear.f64.halfstep", .halves = 9, .codebook = .mul1, .k = 1024, .n = 1024 }, + .{ .id = "exl3.linear.f64.halfstep", .halves = 11, .codebook = .mul1, .k = 2048, .n = 1024 }, + .{ .id = "exl3.linear.f64.halfstep", .halves = 13, .codebook = .mul1, .k = 512, .n = 512 }, +}; + +/// One check's outcome: a failure prints its FAIL line (TestFailed ones already printed theirs) and fails the run. +pub const Run = struct { + failed: bool = false, + + pub fn fail(self: *Run, e: anyerror, id: []const u8) void { + if (e != error.TestFailed) std.debug.print("FAIL {s}: {t}\n", .{ id, e }); + self.failed = true; + } +}; + +pub fn run(gpu: Gpu, dir: []const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(11); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + var lin = try exl3.Linear.load(gpu.d); + defer lin.unload(); + + var state = Run{}; + var last: ?[]const u8 = null; + var ok = true; + for (f64_cfgs) |c| { + if (last) |prev| { + if (!std.mem.eql(u8, prev, c.id)) { + if (ok) check.pass("{s}", .{prev}); + ok = true; + } + } + f64Linear(&r, &lin, c) catch |e| { + state.fail(e, c.id); + ok = false; + }; + last = c.id; + } + if (last) |prev| { + if (ok) check.pass("{s}", .{prev}); + } + { + var rows_ok = true; + rowsInvariant(&r, &lin) catch |e| { + state.fail(e, "exl3.linear.rows"); + rows_ok = false; + }; + if (rows_ok) check.pass("exl3.linear.rows", .{}); + } + pythonChecks(&r, &lin, dir, &state) catch |e| { + state.fail(e, "exl3.linear.python"); + }; + try @import("exl3_dense_tests.zig").run(&r, &state, dir); + try @import("exl3_experts_tests.zig").run(&r, &state, dir); + try @import("exl3_layer_tests.zig").run(&r, &state, dir); + if (state.failed) return error.TestFailed; + try bench(&r, &lin); +} + +/// A zeroed device buffer (counters must read zero, unlike zeros()'s 0xff canaries). +pub fn zeroed(r: *Rig, bytes: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(r.gpu.d, @max(bytes, 16)); + try r.bufs.append(r.gpu.gpa, b); + try b.fill8(0, r.s.handle); + return b.ptr; +} + +/// One codebook and width at one shape: random trellis, scales and 128 rows of x, then every row count of `rows`. +fn f64Linear(r: *Rig, lin: *const exl3.Linear, c: Cfg) !void { + const width = exl3fmt.Width{ .halves = @intCast(c.halves) }; + const cb: exl3.Codebook = switch (c.codebook) { + .inst3 => .inst3, + .mcg => .mcg, + .mul1 => .mul1, + }; + const k = c.k; + const n = c.n; + const trellis = try r.a.alloc(i16, (k / 16) * (n / 16) * width.tileWords()); + for (trellis) |*v| v.* = r.rng.int(i16); + const suh = try r.a.alloc(u16, k); + const svh = try r.a.alloc(u16, n); + for (suh) |*v| v.* = @bitCast(@as(f16, @floatCast(r.rng.float(f32) * 2 - 1))); + for (svh) |*v| v.* = @bitCast(@as(f16, @floatCast(r.rng.float(f32) * 2 - 1))); + const x = try r.bfs(128 * k, 1.0); + const want = try reference(r, width, c.codebook, trellis, k, n, x, suh, svh); + const st = exl3.strides(k, n, c.halves, false); + const w: exl3.Weight = .{ + .trellis = try r.dev(std.mem.sliceAsBytes(trellis)), + .suh = try r.dev(std.mem.sliceAsBytes(suh)), + .svh = try r.dev(std.mem.sliceAsBytes(svh)), + .bias = 0, + .k = @intCast(k), + .n = @intCast(n), + .halves = c.halves, + .codebook = cb, + .stride_k = st.k, + .stride_nb = st.nb, + }; + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const counters = try zeroed(r, exl3.Linear.counterBytes(w)); + for (rows) |m| { + const dy = try r.zeros(m * n * 2); + const xh = try r.zeros(m * k * 2); + const z: u64 = if (exl3.plan(k, n, 192, 8).sk > 1) try r.zeros(exl3.Linear.zBytes(w, m)) else 0; + try lin.matmul(r.s, w, dx, .bf16, m, dy, .bf16, xh, z, counters); + try closeOut(r, c.id, dy, m * n, want.want, want.tol); + } +} + +/// The Python engine's f64 reference: Wq fp16 from the trellis, fp16 xh, the y sums, and each element's |t| sum. +fn reference(r: *Rig, width: exl3fmt.Width, cb: exl3fmt.Codebook, trellis: []const i16, k: usize, n: usize, x: []const u16, suh: []const u16, svh: []const u16) !struct { want: []f64, tol: []f64 } { + const m = 128; + const wq = try r.a.alloc(f16, k * n); + exl3fmt.unpack(width, cb, trellis, k / 16, n / 16, wq); + const xs = try r.a.alloc(f64, m * k); + for (0..m) |row| { + for (0..k) |i| xs[row * k + i] = bf(x[row * k + i]) * @as(f64, @floatCast(@as(f16, @bitCast(suh[i])))); + } + exl3fmt.rotate(xs, m, k, 1); + const xh = try r.a.alloc(f16, m * k); + for (xh, xs) |*o, v| o.* = @floatCast(v); + const ypre = try r.a.alloc(f64, m * n); // [n][m], rotated along n + @memset(ypre, 0); + const s1 = try r.a.alloc(f64, m * n); // each output's sum of |xh * wq| over k + @memset(s1, 0); + var xv: [128]f64 = undefined; + var xa: [128]f64 = undefined; + for (0..k) |ki| { + for (0..m) |row| { + xv[row] = @floatCast(xh[row * k + ki]); + xa[row] = @abs(xv[row]); + } + for (0..n) |col| { + const wv: f64 = @floatCast(wq[ki * n + col]); + const wa = @abs(wv); + const base = col * m; + for (0..m) |row| { + ypre[base + row] += xv[row] * wv; + s1[base + row] += xa[row] * wa; + } + } + } + exl3fmt.rotate(ypre, n, m, 0); + const want = try r.a.alloc(f64, m * n); + const tol = try r.a.alloc(f64, m * n); + for (0..n) |col| { + const scale: f64 = @floatCast(@as(f16, @bitCast(svh[col]))); + for (0..m) |row| { + want[row * n + col] = ypre[col * m + row] * scale; + // bf16 output ulp (2^-7 of the value) plus fp32 accumulation and fp16 xh cases, 2^-11 of the |t| sum + tol[row * n + col] = @abs(want[row * n + col]) / 128 + s1[row * n + col] / 2048 + 1e-30; + } + } + return .{ .want = want, .tol = tol }; +} + +/// bf16 outputs within the f64 reference's tolerance, with the exact FAIL line. +pub fn closeOut(r: *Rig, id: []const u8, ptr: u64, len: usize, want: []const f64, tol: []const f64) !void { + const raw = try r.back(ptr, len * 2); + var most: f64 = 0; + var at: usize = 0; + for (want[0..len], tol[0..len], 0..) |wv, t, i| { + const g = bf(std.mem.bytesToValue(u16, raw[2 * i ..][0..2])); + const q = @abs(g - wv) / t; + if (q > most) { + most = q; + at = i; + } + } + std.debug.print("RESULT {s}: largest difference {d:.3} of its tolerance (element {d})\n", .{ id, most, at }); + if (most > 1) { + std.debug.print("FAIL {s}: largest difference {d:.3} of its tolerance\n", .{ id, most }); + return error.TestFailed; + } +} + +/// The first m rows of a 128-row call byte-equal every narrower call, at a split shape and a single-split one. +fn rowsInvariant(r: *Rig, lin: *const exl3.Linear) !void { + const shapes = [_][2]usize{ .{ 4096, 1024 }, .{ 512, 4096 } }; + for (shapes) |sh| { + const k = sh[0]; + const n = sh[1]; + const width: exl3fmt.Width = .{ .halves = 8 }; + const trellis = try r.a.alloc(i16, (k / 16) * (n / 16) * width.tileWords()); + for (trellis) |*v| v.* = r.rng.int(i16); + const suh = try r.a.alloc(u16, k); + const svh = try r.a.alloc(u16, n); + for (suh) |*v| v.* = @bitCast(@as(f16, @floatCast(r.rng.float(f32) * 2 - 1))); + for (svh) |*v| v.* = @bitCast(@as(f16, @floatCast(r.rng.float(f32) * 2 - 1))); + const st = exl3.strides(k, n, 8, false); + const w: exl3.Weight = .{ + .trellis = try r.dev(std.mem.sliceAsBytes(trellis)), + .suh = try r.dev(std.mem.sliceAsBytes(suh)), + .svh = try r.dev(std.mem.sliceAsBytes(svh)), + .bias = 0, + .k = @intCast(k), + .n = @intCast(n), + .halves = 8, + .codebook = .mul1, + .stride_k = st.k, + .stride_nb = st.nb, + }; + const x = try r.bfs(128 * k, 1.0); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const counters = try zeroed(r, exl3.Linear.counterBytes(w)); + const y128 = try r.zeros(128 * n * 2); + { + const xh = try r.zeros(128 * k * 2); + const z: u64 = if (exl3.plan(k, n, 192, 8).sk > 1) try r.zeros(exl3.Linear.zBytes(w, 128)) else 0; + try lin.matmul(r.s, w, dx, .bf16, 128, y128, .bf16, xh, z, counters); + } + const got128 = try r.back(y128, 128 * n * 2); + for ([_]usize{ 1, 2, 3, 16, 17, 64 }) |m| { + const ym = try r.zeros(m * n * 2); + const xh = try r.zeros(m * k * 2); + const z: u64 = if (exl3.plan(k, n, 192, 8).sk > 1) try r.zeros(exl3.Linear.zBytes(w, m)) else 0; + try lin.matmul(r.s, w, dx, .bf16, m, ym, .bf16, xh, z, counters); + const got = try r.back(ym, m * n * 2); + const id = try std.fmt.allocPrint(r.a, "exl3.linear.rows k={d} n={d} m={d}", .{ k, n, m }); + try check.sameBytes(id, got, got128[0 .. m * n * 2]); + } + } +} + +/// The Python engine's recorded outputs, bit for bit (exl3_oracle.py's manifest and raw files). +fn pythonChecks(r: *Rig, lin: *const exl3.Linear, dir: []const u8, state: *Run) !void { + if (std.mem.eql(u8, dir, "none")) { + for (4..9) |bits| std.debug.print("SKIP exl3.linear.python.{d}\n", .{bits}); + return; + } + const manifest_path = try std.fs.path.join(r.a, &.{ dir, "manifest.json" }); + const text = std.Io.Dir.cwd().readFileAlloc(r.gpu.io, manifest_path, r.a, .limited(1 << 24)) catch |e| { + for (4..9) |bits| { + std.debug.print("FAIL exl3.linear.python.{d}: {t}\n", .{ bits, e }); + state.failed = true; + } + return; + }; + var parsed = try std.json.parseFromSlice(std.json.Value, r.a, text, .{}); + var seen = [_]bool{ false, false, false, false, false, false, false, false, false }; + for (parsed.value.object.get("linears").?.array.items) |item| { + const o = item.object; + const bits: u32 = bitsOf(o.get("bits").?); + seen[bits] = true; + const id = try std.fmt.allocPrint(r.a, "exl3.linear.python.{d}", .{bits}); + var ok = true; + pythonLinear(r, lin, dir, o, id) catch |e| { + state.fail(e, id); + ok = false; + }; + if (ok) check.pass("{s}", .{id}); + } + for (4..9) |bits| { + if (!seen[bits]) std.debug.print("SKIP exl3.linear.python.{d}\n", .{bits}); + } +} + +fn bitsOf(v: std.json.Value) u32 { + return switch (v) { + .integer => |i| @intCast(i), + .float => |f| @intFromFloat(f), + else => unreachable, + }; +} + +fn pythonLinear(r: *Rig, lin: *const exl3.Linear, dir: []const u8, o: std.json.ObjectMap, id: []const u8) !void { + const k: usize = @intCast(o.get("k").?.integer); + const n: usize = @intCast(o.get("n").?.integer); + const bits: u32 = bitsOf(o.get("bits").?); + const cb: exl3.Codebook = if (std.mem.eql(u8, o.get("codebook").?.string, "mcg")) + .mcg + else if (std.mem.eql(u8, o.get("codebook").?.string, "inst3")) + .inst3 + else + .mul1; + const bias = o.get("bias").?; + const st = exl3.strides(k, n, 2 * bits, false); + const w: exl3.Weight = .{ + .trellis = try r.dev(try readFile(r, dir, o.get("trellis").?.string)), + .suh = try r.dev(try readFile(r, dir, o.get("suh").?.string)), + .svh = try r.dev(try readFile(r, dir, o.get("svh").?.string)), + .bias = if (std.meta.activeTag(bias) == .null) 0 else try r.dev(try readFile(r, dir, bias.string)), + .k = @intCast(k), + .n = @intCast(n), + .halves = 2 * bits, + .codebook = cb, + .stride_k = st.k, + .stride_nb = st.nb, + }; + const xmap = o.get("x").?.object; + const ymap = o.get("y").?.object; + const counters = try zeroed(r, exl3.Linear.counterBytes(w)); + for (o.get("rows").?.array.items) |rv| { + const m: usize = @intCast(rv.integer); + const key = try std.fmt.allocPrint(r.a, "{d}", .{m}); + const x = try readFile(r, dir, xmap.get(key).?.string); + const want = try readFile(r, dir, ymap.get(key).?.string); + const dy = try r.zeros(m * n * 2); + const xh = try r.zeros(m * k * 2); + const z: u64 = if (exl3.plan(k, n, 192, 8).sk > 1) try r.zeros(exl3.Linear.zBytes(w, m)) else 0; + try lin.matmul(r.s, w, try r.dev(x), .bf16, m, dy, .bf16, xh, z, counters); + const got = try r.back(dy, want.len); + try check.sameBytes(id, got, want); + } +} + +pub fn readFile(r: *Rig, dir: []const u8, name: []const u8) ![]u8 { + return std.Io.Dir.cwd().readFileAlloc(r.gpu.io, try std.fs.path.join(r.a, &.{ dir, name }), r.a, .limited(1 << 30)); +} + +/// Trellis bytes a second: one row of a 4-bit 4096x32768 layer (64 MiB, past L2), decode-dominated. +fn bench(r: *Rig, lin: *const exl3.Linear) !void { + const k = 4096; + const n = 32768; + const bytes = k * n / 2; + const st = exl3.strides(k, n, 8, false); + const suh = try r.a.alloc(u16, k); + const svh = try r.a.alloc(u16, n); + @memset(suh, 0x3C00); // 1.0 + @memset(svh, 0x3C00); + const x = try r.bfs(k, 1.0); + const w: exl3.Weight = .{ + .trellis = try r.zeros(bytes), + .suh = try r.dev(std.mem.sliceAsBytes(suh)), + .svh = try r.dev(std.mem.sliceAsBytes(svh)), + .bias = 0, + .k = k, + .n = n, + .halves = 8, + .codebook = .mul1, + .stride_k = st.k, + .stride_nb = st.nb, + }; + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const dy = try r.zeros(n * 2); + const xh = try r.zeros(k * 2); + const z = try r.zeros(exl3.Linear.zBytes(w, 1)); + const counters = try zeroed(r, exl3.Linear.counterBytes(w)); + for (0..3) |_| try lin.matmul(r.s, w, dx, .bf16, 1, dy, .bf16, xh, z, counters); + try r.s.synchronize(); + const reps = 200; + const t0 = check.now(r.gpu.io); + for (0..reps) |_| try lin.matmul(r.s, w, dx, .bf16, 1, dy, .bf16, xh, z, counters); + try r.s.synchronize(); + const t1 = check.now(r.gpu.io); + const sec = @as(f64, @floatFromInt(t1 - t0)) / 1e9; + const gb = @as(f64, @floatFromInt(bytes * reps)) / sec / 1e9; + std.debug.print("RESULT exl3.bench.decode {d:.1}\n", .{gb}); +} diff --git a/zig/tests/cuda/exl3_tp2_tests.zig b/zig/tests/cuda/exl3_tp2_tests.zig new file mode 100644 index 0000000000..2eb2091dc6 --- /dev/null +++ b/zig/tests/cuda/exl3_tp2_tests.zig @@ -0,0 +1,430 @@ +//! EXL3 tensor parallel: each rank's shard of real RED-SNOW projections and experts vs the engine's, all-gathered. +const std = @import("std"); +const posix = std.posix; +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const exl3 = cuda.exl3; +const exl3e = cuda.exl3_experts; +const exl3t = @import("exl3_tests.zig"); +const Gpu = check.Gpu; +const Rig = quant_ref_tests.RigT; + +/// NCCL's 128-byte unique id across the cable, rank 0 serving it on host:port (the benchmark starts rank 1 first). +pub fn run(gpu: Gpu, rank_arg: []const u8, addr: []const u8, dir: []const u8) !void { + const rank = try std.fmt.parseInt(u8, rank_arg, 10); + if (rank > 1) { + std.debug.print("usage: tf-cuda-test exl3-tp2 \n", .{}); + return error.BadArguments; + } + var lib = try cuda.nccl.Library.open(); + defer lib.close(); + var id: cuda.nccl.UniqueId = undefined; + const id_bytes = std.mem.asBytes(&id); + if (rank == 0) { + try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); + try serve(addr, id_bytes); + } else { + try receive(gpu, addr, id_bytes); + } + var comm: cuda.nccl.Comm = null; + try lib.check(lib.api.ncclCommInitRank(&comm, 2, id, @intCast(rank)), "ncclCommInitRank"); + defer _ = lib.api.ncclCommDestroy(comm); + + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(11); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + var lin = try exl3.Linear.load(gpu.d); + defer lin.unload(); + var exp = try exl3e.Experts.load(gpu.d); + defer exp.unload(); + + var failed = false; + { + var ok = true; + commSanity(&r, &lib, comm) catch |e| { + fail0(rank, "exl3.tp2.comm", e); + ok = false; + }; + if (!ok) failed = true; + } + if (std.mem.eql(u8, dir, "none")) { + if (rank == 0) { + std.debug.print("SKIP exl3.tp2.linear\n", .{}); + std.debug.print("SKIP exl3.tp2.experts\n", .{}); + } + } else { + linearChecks(&r, &lib, comm, &lin, dir, rank, &failed); + expertsCheck(&r, &lib, comm, &exp, dir, rank, &failed); + } + if (failed) return error.TestFailed; + std.debug.print("RESULT exl3.tp2 rank {d}: every collective and comparison agreed\n", .{rank}); +} + +/// One f32 all-gather: the communicator moves bits before anything else is trusted. +fn commSanity(r: *Rig, lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm) !void { + const one = [_]u8{ 0, 0, 0x80, 0x3f }; + const send = try r.dev(&one); + const recv = try r.zeros(8); + try lib.check(lib.api.ncclAllGather(send, recv, 1, .f32, comm, r.s.handle), "ncclAllGather"); + try r.s.synchronize(); + const got = try r.back(recv, 8); + try same(0, "exl3.tp2.comm all-gather", got, &.{ 0, 0, 0x80, 0x3f, 0, 0, 0x80, 0x3f }); +} + +fn fail0(rank: u8, id: []const u8, e: anyerror) void { + if (rank == 0) std.debug.print("FAIL {s}: {t}\n", .{ id, e }); +} + +/// Equal bytes, or the first difference; the FAIL line is rank 0's (the contract's one printer). +fn same(rank: u8, what: []const u8, got: []const u8, want: []const u8) !void { + if (std.mem.eql(u8, got, want)) return; + if (rank == 0) { + if (got.len != want.len) { + std.debug.print("FAIL {s}: {d} bytes, expected {d}\n", .{ what, got.len, want.len }); + } else { + std.debug.print("FAIL {s}: first difference at byte {d}\n", .{ what, std.mem.indexOfDiff(u8, got, want).? }); + } + } + return error.TestFailed; +} + +fn bitsOf(v: std.json.Value) u32 { + return switch (v) { + .integer => |i| @intCast(i), + .float => |f| @intFromFloat(f), + else => unreachable, + }; +} + +fn codebook(name: []const u8) exl3.Codebook { + if (std.mem.eql(u8, name, "mcg")) return .mcg; + if (std.mem.eql(u8, name, "inst3")) return .inst3; + return .mul1; +} + +fn tp2Section(r: *Rig, dir: []const u8) !std.json.ObjectMap { + const text = try std.Io.Dir.cwd().readFileAlloc(r.gpu.io, try std.fs.path.join(r.a, &.{ dir, "manifest.json" }), r.a, .limited(1 << 24)); + var parsed = try std.json.parseFromSlice(std.json.Value, r.a, text, .{}); + return parsed.value.object.get("tp2").?.object; +} + +/// The two fixture linears, the out-split ('dim1') and the in-split ('row') one: each reports its own line. +fn linearChecks(r: *Rig, lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, lin: *const exl3.Linear, dir: []const u8, rank: u8, failed: *bool) void { + const tp2 = tp2Section(r, dir) catch |e| { + fail0(rank, "exl3.tp2.linear", e); + failed.* = true; + return; + }; + var all = true; + for (tp2.get("linears").?.array.items) |item| { + const o = item.object; + const id = std.fmt.allocPrint(r.a, "exl3.tp2.linear.{s}", .{o.get("id").?.string}) catch unreachable; + var ok = true; + linearCheck(r, lib, comm, lin, dir, rank, o, id) catch |e| { + if (e != error.TestFailed) fail0(rank, id, e); + ok = false; + }; + if (ok and rank == 0) check.pass("{s}", .{id}); + if (!ok) all = false; + } + if (all and rank == 0) check.pass("exl3.tp2.linear", .{}); + if (!all) failed.* = true; +} + +/// One projection: the rank's trellis/scale shard of the EXL3_RULES cut, its output, and the all-gather. +fn linearCheck(r: *Rig, lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, lin: *const exl3.Linear, dir: []const u8, rank: u8, o: std.json.ObjectMap, id: []const u8) !void { + const dim1 = std.mem.eql(u8, o.get("split").?.string, "dim1"); + const k: usize = @intCast(o.get("k").?.integer); + const n: usize = @intCast(o.get("n").?.integer); + const bits: u32 = bitsOf(o.get("bits").?); + const halves: usize = 2 * bits; + const kt = k / 16; + const nt = n / 16; + const tw = halves * 16; // bytes a trellis tile (8 * halves i16 values) + const full = try exl3t.readFile(r, dir, o.get("trellis").?.string); + const suh_full = try exl3t.readFile(r, dir, o.get("suh").?.string); + const svh_full = try exl3t.readFile(r, dir, o.get("svh").?.string); + var k_use = k; + var n_use = n; + // dim1: trellis tile columns and svh rows in half, suh whole; row: tile rows, suh rows and x's columns + const trellis = if (dim1) try cutDim1(r, full, kt, nt, tw, rank) else full[rank * (kt / 2) * nt * tw ..][0 .. (kt / 2) * nt * tw]; + const suh = if (dim1) suh_full else suh_full[rank * (k / 2) * 2 ..][0 .. (k / 2) * 2]; + const svh = if (dim1) svh_full[rank * (n / 2) * 2 ..][0 .. (n / 2) * 2] else svh_full; + if (dim1) n_use = n / 2 else k_use = k / 2; + const biasv = o.get("bias").?; + var bias: u64 = 0; + if (std.meta.activeTag(biasv) == .string) { + const bfull = try exl3t.readFile(r, dir, biasv.string); + bias = try r.dev(if (dim1) bfull[rank * (n / 2) * 2 ..][0 .. (n / 2) * 2] else bfull); + } + const st = exl3.strides(k_use, n_use, halves, false); + const w: exl3.Weight = .{ + .trellis = try r.dev(trellis), + .suh = try r.dev(suh), + .svh = try r.dev(svh), + .bias = bias, + .k = @intCast(k_use), + .n = @intCast(n_use), + .halves = @intCast(halves), + .codebook = codebook(o.get("codebook").?.string), + .stride_k = st.k, + .stride_nb = st.nb, + }; + const counters = try exl3t.zeroed(r, exl3.Linear.counterBytes(w)); + const ydt: exl3.DType = if (dim1) .bf16 else .f32; + const esz: usize = if (dim1) 2 else 4; + const dt: cuda.nccl.DataType = if (dim1) .bf16 else .f32; + const xmap = o.get("x").?.object; + + for (o.get("rows").?.array.items) |rv| { + const m: usize = @intCast(rv.integer); + const key = try std.fmt.allocPrint(r.a, "{d}", .{m}); + const what = try std.fmt.allocPrint(r.a, "{s} m={d}", .{ id, m }); + const xf = try exl3t.readFile(r, dir, xmap.get(key).?.string); + const x_use = if (dim1) xf else try cutDim1(r, xf, m, k, 2, rank); + const count = m * n_use; // elements this rank all-gathers + const send = try r.zeros(count * esz); + const xh = try r.zeros(m * k_use * 2); + const z: u64 = if (exl3.plan(k_use, n_use, 192, 8).sk > 1) try r.zeros(exl3.Linear.zBytes(w, m)) else 0; + try lin.matmul(r.s, w, try r.dev(x_use), .bf16, m, send, ydt, xh, z, counters); + const recv = try r.zeros(2 * count * esz); + try lib.check(lib.api.ncclAllGather(send, recv, count, dt, comm, r.s.handle), "ncclAllGather"); + try r.s.synchronize(); + const got = try r.back(recv, 2 * count * esz); + if (dim1) { + // the two ranks' bf16 halves concatenated are the full output + const it = o.get("shards").?.object.get(key).?.array.items; + const want = try std.mem.concat(r.a, u8, &.{ + try exl3t.readFile(r, dir, it[0].string), + try exl3t.readFile(r, dir, it[1].string), + }); + try same(rank, what, got, want); + } else { + // fp32 partials bit for bit, then bf16(sum rank 0 first) as forward.py's hc_post rounds + const it = o.get("partials").?.object.get(key).?.array.items; + try same(rank, what, got[0 .. m * n * 4], try exl3t.readFile(r, dir, it[0].string)); + try same(rank, what, got[m * n * 4 ..], try exl3t.readFile(r, dir, it[1].string)); + const want = try r.a.alloc(u8, m * n * 2); + for (0..m * n) |i| { + const g0 = std.mem.bytesToValue(f32, got[4 * i ..][0..4]); + const g1 = std.mem.bytesToValue(f32, got[4 * (m * n + i) ..][0..4]); + std.mem.writeInt(u16, want[2 * i ..][0..2], bf16r(g0 + g1), .little); + } + try same(rank, what, want, try exl3t.readFile(r, dir, o.get("y").?.object.get(key).?.string)); + } + } +} + +/// The routed experts of one MoE layer: each rank's gate/up/down shards, its fp32 partials, the all-gathered sum. +fn expertsCheck(r: *Rig, lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, exp: *const exl3e.Experts, dir: []const u8, rank: u8, failed: *bool) void { + const id = "exl3.tp2.experts"; + const o = blk: { + const tp2 = tp2Section(r, dir) catch |e| { + fail0(rank, id, e); + failed.* = true; + return; + }; + break :blk tp2.get("experts").?.object; + }; + var ok = true; + expertsRun(r, lib, comm, exp, dir, rank, o, id) catch |e| { + if (e != error.TestFailed) fail0(rank, id, e); + ok = false; + }; + if (ok and rank == 0) check.pass("{s}", .{id}); + if (!ok) failed.* = true; +} + +fn expertsRun(r: *Rig, lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, exp: *const exl3e.Experts, dir: []const u8, rank: u8, o: std.json.ObjectMap, id: []const u8) !void { + const e_count: usize = @intCast(o.get("count").?.integer); + const dims: usize = @intCast(o.get("dims").?.integer); + const width: usize = @intCast(o.get("width").?.integer); // the full layer's; each rank holds width / 2 + const slots: usize = @intCast(o.get("slots").?.integer); + const trellis = o.get("trellis").?.object; + var gate_p = try r.a.alloc(u64, e_count); + var up_p = try r.a.alloc(u64, e_count); + var down_p = try r.a.alloc(u64, e_count); + var gk2 = try r.a.alloc(i32, e_count); + var uk2 = try r.a.alloc(i32, e_count); + var dk2 = try r.a.alloc(i32, e_count); + for (o.get("gate_k2").?.array.items, 0..) |v, i| gk2[i] = @intCast(v.integer); + for (o.get("up_k2").?.array.items, 0..) |v, i| uk2[i] = @intCast(v.integer); + for (o.get("down_k2").?.array.items, 0..) |v, i| dk2[i] = @intCast(v.integer); + // gate/up: trellis tile columns and svh rows in half; down: tile rows and suh rows in half + for (0..e_count) |i| { + const gw = @as(usize, @intCast(gk2[i])) * 16; // bytes a gate/up tile + const dw = @as(usize, @intCast(dk2[i])) * 16; + gate_p[i] = try r.dev(try cutDim1(r, try exl3t.readFile(r, dir, trellis.get("gate").?.array.items[i].string), dims / 16, width / 16, gw, rank)); + up_p[i] = try r.dev(try cutDim1(r, try exl3t.readFile(r, dir, trellis.get("up").?.array.items[i].string), dims / 16, width / 16, gw, rank)); + const dfull = try exl3t.readFile(r, dir, trellis.get("down").?.array.items[i].string); + down_p[i] = try r.dev(dfull[rank * (width / 32) * (dims / 16) * dw ..][0 .. (width / 32) * (dims / 16) * dw]); + } + const layer: exl3e.Layer = .{ + .gate = try r.dev(std.mem.sliceAsBytes(gate_p)), + .up = try r.dev(std.mem.sliceAsBytes(up_p)), + .down = try r.dev(std.mem.sliceAsBytes(down_p)), + .gate_k2 = try r.dev(std.mem.sliceAsBytes(gk2)), + .up_k2 = try r.dev(std.mem.sliceAsBytes(uk2)), + .down_k2 = try r.dev(std.mem.sliceAsBytes(dk2)), + .suh_g = try r.dev(try exl3t.readFile(r, dir, o.get("suh_g").?.string)), + .suh_u = try r.dev(try exl3t.readFile(r, dir, o.get("suh_u").?.string)), + .svh_g = try r.dev(try cutDim1(r, try exl3t.readFile(r, dir, o.get("svh_g").?.string), e_count, width, 2, rank)), + .svh_u = try r.dev(try cutDim1(r, try exl3t.readFile(r, dir, o.get("svh_u").?.string), e_count, width, 2, rank)), + .suh_d = try r.dev(try cutDim1(r, try exl3t.readFile(r, dir, o.get("suh_d").?.string), e_count, width, 2, rank)), + .svh_d = try r.dev(try exl3t.readFile(r, dir, o.get("svh_d").?.string)), + .experts = e_count, + .dims = dims, + .width = width / 2, + .codebook = codebook(o.get("codebook").?.string), + .limit = @floatCast(o.get("limit").?.float), + .act_mode = @intCast(o.get("act_mode").?.integer), + }; + const xmap = o.get("x").?.object; + const pickmap = o.get("pick").?.object; + const wtsmap = o.get("wts").?.object; + for (o.get("rows").?.array.items) |rv| { + const m: usize = @intCast(rv.integer); + const key = try std.fmt.allocPrint(r.a, "{d}", .{m}); + const what = try std.fmt.allocPrint(r.a, "{s} m={d}", .{ id, m }); + const dx = try r.dev(try exl3t.readFile(r, dir, xmap.get(key).?.string)); + const dpick = try r.dev(try exl3t.readFile(r, dir, pickmap.get(key).?.string)); + const dwts = try r.dev(try exl3t.readFile(r, dir, wtsmap.get(key).?.string)); + const sb = try r.zeros(try exl3e.scratchBytes(layer, m, slots)); + const sc = try exl3e.scratch(sb, layer, m, slots); + const out = try r.zeros(m * dims * 4); + try exp.run(r.s, layer, dx, dims, .bf16, dpick, dwts, m, out, sc); + const recv = try r.zeros(2 * m * dims * 4); + try lib.check(lib.api.ncclAllGather(out, recv, m * dims, .f32, comm, r.s.handle), "ncclAllGather"); + try r.s.synchronize(); + const got = try r.back(recv, 2 * m * dims * 4); + const it = o.get("partials").?.object.get(key).?.array.items; + try same(rank, what, got[0 .. m * dims * 4], try exl3t.readFile(r, dir, it[0].string)); + try same(rank, what, got[m * dims * 4 ..], try exl3t.readFile(r, dir, it[1].string)); + const want = try r.a.alloc(u8, m * dims * 2); + for (0..m * dims) |i| { + const g0 = std.mem.bytesToValue(f32, got[4 * i ..][0..4]); + const g1 = std.mem.bytesToValue(f32, got[4 * (m * dims + i) ..][0..4]); + std.mem.writeInt(u16, want[2 * i ..][0..2], bf16r(g0 + g1), .little); + } + try same(rank, what, want, try exl3t.readFile(r, dir, o.get("y").?.object.get(key).?.string)); + } +} + +/// The rank's half of an even 2-D row-major tensor's second axis, repacked contiguous (split.py's dim1 cut). +fn cutDim1(r: *Rig, src: []const u8, rows: usize, cols: usize, item: usize, rank: u8) ![]u8 { + const per = cols / 2 * item; + const out = try r.a.alloc(u8, rows * per); + for (0..rows) |i| { + @memcpy(out[i * per ..][0..per], src[i * cols * item + @as(usize, rank) * per ..][0..per]); + } + return out; +} + +/// bf16 round-to-nearest-even of an fp32 sum (hc_post's `acc.to(torch.bfloat16)`). +fn bf16r(v: f32) u16 { + const u: u32 = @bitCast(v); + return @intCast((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +/// Rank 0 binds host:port and sends its unique id to the one peer (70 s for rank 1's retries). +fn serve(addr_text: []const u8, bytes: []const u8) !void { + const a = try parseAddr(addr_text); + var storage: std.Io.Threaded.PosixAddress = undefined; + const len = try posixLen(a.host, a.port, &storage); + const rc = posix.system.socket(posix.AF.INET, posix.SOCK.STREAM, 0); + if (posix.errno(rc) != .SUCCESS) return error.SocketFailed; + const fd: posix.socket_t = @intCast(rc); + defer _ = posix.system.close(fd); + const one: c_int = 1; + try posix.setsockopt(fd, posix.SOL.SOCKET, posix.SO.REUSEADDR, std.mem.asBytes(&one)); + switch (posix.errno(posix.system.bind(fd, &storage.any, len))) { + .SUCCESS => {}, + .ADDRINUSE => return error.AddressInUse, + else => return error.BindFailed, + } + if (posix.errno(posix.system.listen(fd, 1)) != .SUCCESS) return error.ListenFailed; + var waited: u64 = 0; + while (true) { + var fds = [_]posix.pollfd{.{ .fd = fd, .events = posix.POLL.IN, .revents = 0 }}; + const ready = posix.poll(&fds, 500) catch return error.AcceptFailed; + if (ready > 0) break; + waited += 500; + if (waited > 70_000) return error.PeerTimeout; + } + var peer: std.Io.Threaded.PosixAddress = undefined; + var peer_len: posix.socklen_t = @sizeOf(std.Io.Threaded.PosixAddress); + const rc2 = posix.system.accept(fd, &peer.any, &peer_len); + if (posix.errno(rc2) != .SUCCESS) return error.AcceptFailed; + const cfd: posix.socket_t = @intCast(rc2); + defer _ = posix.system.close(cfd); + var off: usize = 0; + while (off < bytes.len) { + const n = posix.system.write(cfd, bytes.ptr + off, bytes.len - off); + if (posix.errno(n) != .SUCCESS) return error.WriteFailed; + off += @intCast(n); + } +} + +/// Rank 1 connects to host:port, retrying for 60 s (rank 0 may not have bound yet), and reads the id. +fn receive(gpu: Gpu, addr_text: []const u8, bytes: []u8) !void { + const a = try parseAddr(addr_text); + const ip = try std.Io.net.IpAddress.parseIp4(a.host, a.port); + var fd: posix.socket_t = -1; + defer { + if (fd >= 0) _ = posix.system.close(fd); + } + var waited: u64 = 0; + while (fd < 0) { + var storage: std.Io.Threaded.PosixAddress = undefined; + const len = std.Io.Threaded.addressToPosix(&ip, &storage); + const rc = posix.system.socket(posix.AF.INET, posix.SOCK.STREAM, 0); + if (posix.errno(rc) != .SUCCESS) return error.SocketFailed; + const c: posix.socket_t = @intCast(rc); + if (posix.errno(posix.system.connect(c, &storage.any, len)) == .SUCCESS) { + fd = c; + } else { + _ = posix.system.close(c); + waited += 200; + if (waited > 60_000) return error.PeerTimeout; + std.Io.sleep(gpu.io, .fromMilliseconds(200), .awake) catch {}; + } + } + var off: usize = 0; + var read_waited: u64 = 0; + while (off < bytes.len) { + var fds = [_]posix.pollfd{.{ .fd = fd, .events = posix.POLL.IN, .revents = 0 }}; + const ready = posix.poll(&fds, 500) catch return error.ReadFailed; + if (ready == 0) { + read_waited += 500; + if (read_waited > 60_000) return error.PeerTimeout; + continue; + } + const n = posix.system.read(fd, bytes.ptr + off, bytes.len - off); + switch (posix.errno(n)) { + .SUCCESS => { + if (n == 0) return error.PeerClosed; // the peer died before its id arrived + off += @intCast(n); + }, + .INTR => {}, + else => return error.ReadFailed, + } + } +} + +fn posixLen(host: []const u8, port: u16, storage: *std.Io.Threaded.PosixAddress) !posix.socklen_t { + const ip = try std.Io.net.IpAddress.parseIp4(host, port); + return std.Io.Threaded.addressToPosix(&ip, storage); +} + +fn parseAddr(text: []const u8) !struct { host: []const u8, port: u16 } { + const at = std.mem.lastIndexOfScalar(u8, text, ':') orelse return error.BadAddress; + return .{ .host = text[0..at], .port = try std.fmt.parseInt(u16, text[at + 1 ..], 10) }; +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 7b8a089947..421611fb82 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -22,6 +22,8 @@ const quant_ref_tests = @import("quant_ref_tests.zig"); const carveout_tests = @import("carveout_tests.zig"); const nvfp4_tests = @import("nvfp4_tests.zig"); const modelopt_tests = @import("modelopt_tests.zig"); +const exl3_tests = @import("exl3_tests.zig"); +const exl3_tp2_tests = @import("exl3_tp2_tests.zig"); const train_pack = @import("train_pack.zig"); const exlr8_tests = @import("exlr8_tests.zig"); const exlr8_dense_tests = @import("exlr8_dense_tests.zig"); @@ -64,6 +66,8 @@ 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 decode-row linear: float64 checks, row invariance, the Python engine's fixture (dir none skips it) + \\ exl3-tp2 [dir] EXL3 tensor parallel: rank shards vs the Python engine's, all-gathered (dir none skips them) \\ train-pack MODEL EXAMPLES [COUNT] training passes holding several sequences against one a pass: differences, time \\ ; @@ -123,6 +127,8 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "exlr8")) return exlr8_tests.run(gpu); 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, "exl3")) return exl3_tests.run(gpu, if (rest.len > 0) rest[0] else "none"); + if (std.mem.eql(u8, cmd, "exl3-tp2")) return exl3_tp2_tests.run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else "none"); if (std.mem.eql(u8, cmd, "sample")) return sample_tests.draws(gpu); if (std.mem.eql(u8, cmd, "norms")) return norms_tests.run(gpu); if (std.mem.eql(u8, cmd, "norm-rows")) return norm_rows_tests.run(gpu); diff --git a/zig/tests/cuda/oracle/exl3_oracle.py b/zig/tests/cuda/oracle/exl3_oracle.py new file mode 100644 index 0000000000..17374dd307 --- /dev/null +++ b/zig/tests/cuda/oracle/exl3_oracle.py @@ -0,0 +1,124 @@ +"""Fixtures for `tf-cuda-test exl3`: the Python engine's EXL3 linear outputs on real RED-SNOW tensors.""" +from __future__ import annotations + +import argparse +import hashlib +import json +import sys +from pathlib import Path + +import numpy as np +import torch + +# the other fixtures: siblings in this folder, which leads sys.path when the script runs +import exl3_oracle_dense +import exl3_oracle_experts +import exl3_oracle_layer +import exl3_oracle_tp2 + +ROWS = (1, 3, 16) +SEED = 1234 +SIZE_LIMIT = 8192 # prefer layers with N, K at or under this + + +def rank(t, bits: float) -> int: + """Which layers to prefer for a width: a gate_proj for 4..6 bits, a shared-expert or attention one for 7..8.""" + if bits <= 6 and "gate_proj" in t.prefix: + return 0 + if bits >= 7 and ("shared_experts" in t.prefix or "self_attn" in t.prefix): + return 0 + return 1 + + +def pick(model_dir: Path) -> list: + """One EXL3 linear per whole bits value 4..8, first by sorted name among the preferred ones.""" + from tensorfold.cuda.exl3 import format as fmt + + cp = fmt.scan(model_dir) + chosen = [] + for bits in (4.0, 5.0, 6.0, 7.0, 8.0): + cands = [t for t in cp.groups.values() + if t.bits == bits and t.in_scales == "suh" and t.out_scales == "svh"] + if not cands: + print(f"no EXL3 linear at {bits:g} bits in {model_dir}") + continue + modest = [t for t in cands if t.k <= SIZE_LIMIT and t.n <= SIZE_LIMIT] + t = min(modest or cands, key=lambda t: (rank(t, bits), t.prefix)) + marker = cp.markers.get(f"{t.prefix}.{t.codebook}") + if t.codebook in fmt.MARKERS and marker is not None and marker != fmt.MARKERS[t.codebook]: + raise ValueError(f"{t.prefix}: {t.codebook} marker is 0x{marker:08X}, want 0x{fmt.MARKERS[t.codebook]:08X}") + chosen.append((t, marker)) + return chosen + + +def save(out: Path, name: str, t: torch.Tensor) -> None: + (out / f"{name}.bin").write_bytes(t.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes()) + + +def one(out: Path, model_dir: Path, t, marker: int | None) -> dict: + from tensorfold.cuda.exl3 import format as fmt + from tensorfold.cuda.exl3.linear import Exl3Linear + + lid = f"python.{t.bits:g}" + # "stored" keeps the checkpoint's trellis word order; the forward path is identical either way + lin = Exl3Linear.load(model_dir, t.prefix, layout="stored") + g = torch.Generator().manual_seed(SEED) + x = torch.randn((max(ROWS), lin.k), generator=g).to(torch.bfloat16) + + save(out, f"{lid}.trellis", lin.words) + save(out, f"{lid}.suh", lin.suh) + save(out, f"{lid}.svh", lin.svh) + if lin.bias is not None: + save(out, f"{lid}.bias", lin.bias) + + ys, xfiles, yfiles, sha = [], {}, {}, {} + for m in ROWS: + y = lin(x[:m].cuda()) + torch.cuda.synchronize() + save(out, f"{lid}.x{m}", x[:m]) + save(out, f"{lid}.y{m}", y) + xfiles[str(m)] = f"{lid}.x{m}.bin" + yfiles[str(m)] = f"{lid}.y{m}.bin" + sha[str(m)] = hashlib.sha256(y.cpu().view(torch.uint8).numpy().tobytes()).hexdigest() + ys.append(y) + + rel, ref = 0.0, fmt.forward(x.float().numpy().astype(np.float64), lin.words.view(torch.int16), lin.suh, + lin.svh, t.bits, t.codebook, lin.bias) + for m, y in zip(ROWS, ys): + want = ref[:m] + rel = max(rel, np.linalg.norm(y.cpu().float().numpy().astype(np.float64) - want) + / max(np.linalg.norm(want), 1e-30)) + + entry = {"id": lid, "name": t.prefix, "bits": int(t.bits), "codebook": t.codebook, + "k": lin.k, "n": lin.n, "rows": list(ROWS), "trellis": f"{lid}.trellis.bin", + "suh": f"{lid}.suh.bin", "svh": f"{lid}.svh.bin", + "bias": None if lin.bias is None else f"{lid}.bias.bin", + "x": xfiles, "y": yfiles, "x_dtype": "bf16", "y_dtype": "bf16", "sha256": sha} + if marker is not None: + entry["marker"] = f"0x{marker:08X}" + print(f"{lid}: {t.prefix} EXL3 {t.bits:g}-bit {t.codebook} K={lin.k} N={lin.n} " + f"rel={rel:.2e} sha256={sha['16'][:16]}…") + return entry + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__.partition("\n")[0]) + ap.add_argument("--model", required=True, help="the EXL3 checkpoint folder") + ap.add_argument("--out", required=True, help="the fixture folder to (over)write") + a = ap.parse_args() + out, model_dir = Path(a.out), Path(a.model) + out.mkdir(parents=True, exist_ok=True) + linears = [one(out, model_dir, t, marker) for t, marker in pick(model_dir)] + dense = exl3_oracle_dense.write(out, model_dir) + experts = exl3_oracle_experts.write(out, model_dir) + layers = exl3_oracle_layer.write(out, model_dir) + tp2 = exl3_oracle_tp2.write(out, model_dir) + manifest = {"version": 1, "linears": linears, "dense": dense, "experts": experts, "layers": layers, "tp2": tp2, + "model": str(model_dir)} + (out / "manifest.json").write_text(json.dumps(manifest, indent=1) + "\n") + print(f"{len(linears)} linears -> {out / 'manifest.json'}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/zig/tests/cuda/oracle/exl3_oracle_dense.py b/zig/tests/cuda/oracle/exl3_oracle_dense.py new file mode 100644 index 0000000000..afb5b620e8 --- /dev/null +++ b/zig/tests/cuda/oracle/exl3_oracle_dense.py @@ -0,0 +1,67 @@ +"""The dense prompt fixture: one real projection's 300-row prompt chunk through forward.py's exl3_mm path.""" +from __future__ import annotations + +import hashlib +from pathlib import Path + +import numpy as np +import torch + +ROWS = 300 +SEED = 1235 +SIZE_LIMIT = 8192 # prefer layers with N, K at or under this +PREFER = ("self_attn.qkv_proj", "shared_experts.gate_proj", "self_attn.q_b_proj") + + +def save(out: Path, name: str, t: torch.Tensor) -> None: + (out / f"{name}.bin").write_bytes(t.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes()) + + +def write(out: Path, model_dir: Path) -> dict: + from tensorfold.cuda.exl3 import format as fmt + from tensorfold.cuda.exl3.linear import Exl3Linear + + cp = fmt.scan(model_dir) + cands = [t for t in cp.groups.values() + if t.bits in (4.0, 5.0, 6.0, 7.0, 8.0) and t.in_scales == "suh" and t.out_scales == "svh"] + + def pref(t): + rank = next((i for i, pat in enumerate(PREFER) if pat in t.prefix), len(PREFER)) + return (rank, t.k * t.n) + + modest = [t for t in cands if t.k <= SIZE_LIMIT and t.n <= SIZE_LIMIT] + t = min(modest or cands, key=pref) + marker = cp.markers.get(f"{t.prefix}.{t.codebook}") + if t.codebook in fmt.MARKERS and marker is not None and marker != fmt.MARKERS[t.codebook]: + raise ValueError(f"{t.prefix}: {t.codebook} marker is 0x{marker:08X}, want 0x{fmt.MARKERS[t.codebook]:08X}") + lin = Exl3Linear.load(model_dir, t.prefix, layout="stored") + g = torch.Generator().manual_seed(SEED) + x = torch.randn((ROWS, lin.k), generator=g).to(torch.bfloat16) + + # forward.py's exl3_mm: the decode-row linear over 128-row chunks, a fresh output tile each chunk + y = torch.empty((ROWS, lin.n), dtype=torch.bfloat16) + o = torch.empty((min(ROWS, 128), lin.n), dtype=torch.bfloat16, device="cuda") + for r in range(0, ROWS, 128): + e = min(ROWS - r, 128) + lin(x[r:r + e].cuda(), out=o[:e]) + torch.cuda.synchronize() + y[r:r + e] = o[:e].cpu() + + lid = "dense" + save(out, f"{lid}.trellis", lin.words) + save(out, f"{lid}.suh", lin.suh) + save(out, f"{lid}.svh", lin.svh) + if lin.bias is not None: + save(out, f"{lid}.bias", lin.bias) + save(out, f"{lid}.x", x) + save(out, f"{lid}.y", y) + sha = hashlib.sha256(y.cpu().view(torch.uint8).numpy().tobytes()).hexdigest() + entry = {"id": lid, "name": t.prefix, "bits": int(t.bits), "codebook": t.codebook, + "k": lin.k, "n": lin.n, "rows": ROWS, "trellis": f"{lid}.trellis.bin", + "suh": f"{lid}.suh.bin", "svh": f"{lid}.svh.bin", + "bias": None if lin.bias is None else f"{lid}.bias.bin", + "x": f"{lid}.x.bin", "y": f"{lid}.y.bin", "x_dtype": "bf16", "y_dtype": "bf16", "sha256": sha} + if marker is not None: + entry["marker"] = f"0x{marker:08X}" + print(f"{lid}: {t.prefix} EXL3 {t.bits:g}-bit {t.codebook} K={lin.k} N={lin.n} sha256={sha[:16]}…") + return entry diff --git a/zig/tests/cuda/oracle/exl3_oracle_experts.py b/zig/tests/cuda/oracle/exl3_oracle_experts.py new file mode 100644 index 0000000000..54acb4a9a9 --- /dev/null +++ b/zig/tests/cuda/oracle/exl3_oracle_experts.py @@ -0,0 +1,115 @@ +"""The routed-experts fixture: nine real experts of one MoE layer through the Python engine's generic.routed.""" +from __future__ import annotations + +import hashlib +import json +import re +from pathlib import Path + +import torch + +ROWS = (1, 3, 16) +SEED = 1236 +SLOTS = 8 +LIMIT = 10.0 +POOL = 3 # experts per width + + +def save(out: Path, name: str, t: torch.Tensor) -> None: + (out / f"{name}.bin").write_bytes(t.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes()) + + +def load_group(model_dir: Path, prefix: str) -> tuple: + """(trellis, suh, svh) of one expert projection, on the GPU (the checkpoint's stored trellis order).""" + from safetensors import safe_open + + root = Path(model_dir) + index = root / "model.safetensors.index.json" + rel = None + if index.exists(): + weight_map = json.loads(index.read_text())["weight_map"] + rel = weight_map.get(f"{prefix}.trellis") + if rel is None: + for part in sorted(root.glob("*.safetensors")): + with safe_open(str(part), framework="pt") as g: + if f"{prefix}.trellis" in set(g.keys()): + rel = part.name + break + if rel is None: + raise ValueError(f"{prefix}.trellis is in no file of {root}") + file = root / (rel or "model.safetensors") + with safe_open(str(file), framework="pt", device="cuda") as f: + return (f.get_tensor(f"{prefix}.trellis"), f.get_tensor(f"{prefix}.suh"), f.get_tensor(f"{prefix}.svh")) + + +def pick_layer(cp) -> str: + """The first layer (sorted) whose routed experts cover 4, 5 and 6 bits.""" + layers: dict[str, dict[float, list[int]]] = {} + for prefix, t in cp.groups.items(): + m = re.match(r"^(.*)\.mlp\.experts\.(\d+)\.gate_proj$", prefix) + if m is not None and t.in_scales == "suh" and t.out_scales == "svh": + layers.setdefault(m.group(1), {}).setdefault(t.bits, []).append(int(m.group(2))) + for base in sorted(layers): + if all(len(layers[base].get(b, ())) >= POOL for b in (4.0, 5.0, 6.0)): + return base + raise ValueError("no MoE layer has 4, 5 and 6-bit routed experts") + +def write(out: Path, model_dir: Path) -> dict: + from tensorfold.cuda.exl3 import experts as generic + from tensorfold.cuda.exl3 import format as fmt + + cp = fmt.scan(model_dir) + base = pick_layer(cp) + + by_bits: dict[float, list[int]] = {} + for prefix, t in cp.groups.items(): + m = re.match(rf"^{re.escape(base)}\.mlp\.experts\.(\d+)\.gate_proj$", prefix) + if m is not None: + by_bits.setdefault(t.bits, []).append(int(m.group(1))) + pool = [e for b in (4.0, 5.0, 6.0) for e in sorted(by_bits[b])[:POOL]] + + gate, up, down = [], [], [] + trellis = {"gate": [], "up": [], "down": []} + for slot, e in enumerate(pool): + for part, short, mats, names in (("gate_proj", "gate", gate, trellis["gate"]), + ("up_proj", "up", up, trellis["up"]), + ("down_proj", "down", down, trellis["down"])): + mats.append(load_group(model_dir, f"{base}.mlp.experts.{e}.{part}")) + names.append(f"experts.{short}{slot}.bin") + ex = generic.prepare(gate, up, down, "mul1", device="cuda") + + suh_files = {} + for key, t in (("suh_g", ex.suh_g), ("suh_u", ex.suh_u), ("svh_g", ex.svh_g), + ("svh_u", ex.svh_u), ("suh_d", ex.suh_d), ("svh_d", ex.svh_d)): + save(out, f"experts.{key}", t) + suh_files[key] = f"experts.{key}.bin" + for slot, mats in enumerate(zip(gate, up, down)): + for name, t in zip(("gate", "up", "down"), mats): + save(out, f"experts.{name}{slot}", t[0]) + + g = torch.Generator().manual_seed(SEED) + x = torch.randn((max(ROWS), ex.dims), generator=g).to(torch.bfloat16) + pick = torch.stack([torch.randperm(ex.count, generator=g)[:SLOTS] for _ in range(max(ROWS))]).int() + wts = (0.1 + 0.9 * torch.rand((max(ROWS), SLOTS), generator=g)).float() + + scratch = generic.Scratch(ex, max(ROWS), SLOTS) + xmap, pickmap, wtsmap, ymap, outmap, sha = {}, {}, {}, {}, {}, {} + for m in ROWS: + y = generic.routed(x[:m].cuda(), pick[:m].cuda(), wts[:m].cuda(), ex, scratch, None, m, + limit=LIMIT, act_mode=generic.ACT_BF16) + torch.cuda.synchronize() + saved = {"x": x[:m], "pick": pick[:m], "wts": wts[:m], "y": scratch.y[:m * SLOTS], "out": y} + maps = {"x": xmap, "pick": pickmap, "wts": wtsmap, "y": ymap, "out": outmap} + for key, t in saved.items(): + save(out, f"experts.{key}{m}", t) + maps[key][str(m)] = f"experts.{key}{m}.bin" + sha[str(m)] = hashlib.sha256(y.cpu().view(torch.uint8).numpy().tobytes()).hexdigest() + + entry = {"id": "experts", "name": base, "count": ex.count, "dims": ex.dims, "width": ex.width, + "slots": SLOTS, "codebook": "mul1", "limit": LIMIT, "act_mode": generic.ACT_BF16, + "rows": list(ROWS), "gate_k2": [generic.k2_of(m[0][0]) for m in gate], + "up_k2": [generic.k2_of(m[0]) for m in up], "down_k2": [generic.k2_of(m[0]) for m in down], + "trellis": trellis, **suh_files, "x": xmap, "pick": pickmap, "wts": wtsmap, + "y": ymap, "out": outmap, "sha256": sha} + print(f"experts: {base} {len(pool)} experts (4/5/6-bit pool), rows {ROWS}, sha256={sha['16'][:16]}…") + return entry diff --git a/zig/tests/cuda/oracle/exl3_oracle_layer.py b/zig/tests/cuda/oracle/exl3_oracle_layer.py new file mode 100644 index 0000000000..469f12e0ba --- /dev/null +++ b/zig/tests/cuda/oracle/exl3_oracle_layer.py @@ -0,0 +1,170 @@ +"""The whole-layer fixtures: every EXL3 group of one KDA and one DSA MoE layer through the Python engine.""" +from __future__ import annotations + +import hashlib +import json +import zlib +from pathlib import Path + +import torch + +ROWS = (1, 16, 300) # the dense groups' row counts +EXPERT_ROWS = (1, 16) +SEED = 1239 # + layer index + + +def save(out: Path, name: str, t: torch.Tensor) -> None: + (out / f"{name}.bin").write_bytes(t.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes()) + + +def sha_of(y: torch.Tensor) -> str: + return hashlib.sha256(y.cpu().contiguous().view(torch.uint8).numpy().tobytes()).hexdigest() + + +def load_tensor(model_dir: Path, name: str, device: str = "cpu") -> torch.Tensor: + from safetensors import safe_open + + root = Path(model_dir) + index = root / "model.safetensors.index.json" + if index.exists(): + rel = json.loads(index.read_text())["weight_map"][name] + else: + rel = "model.safetensors" + with safe_open(str(root / rel), framework="pt", device=device) as f: + return f.get_tensor(name) + + +def mm_rows(x: torch.Tensor, lin, m: int) -> torch.Tensor: + """lin on m rows: one call to 128, a prompt chunk beyond that as 128-row blocks (forward.py's exl3_mm).""" + if m <= 128: + y = lin(x[:m].cuda()) + torch.cuda.synchronize() + return y + y = torch.empty((m, lin.n), dtype=torch.bfloat16) + o = torch.empty((min(m, 128), lin.n), dtype=torch.bfloat16, device="cuda") + for r in range(0, m, 128): + e = min(m - r, 128) + lin(x[r:r + e].cuda(), out=o[:e]) + torch.cuda.synchronize() + y[r:r + e] = o[:e].cpu() + return y + + +def lid_of(base: str) -> str: + return base.rsplit(".", 1)[1] + + +def pick_layers(model_dir: Path) -> list[tuple[str, int]]: + """(kind, layer index) of the first KDA and the first DSA MoE layer.""" + cfg = json.loads((Path(model_dir) / "config.json").read_text())["text_config"] + keys = {"kda": "linear_attention", "dsa": "deepseek_sparse_attention"} + out = [] + for kind, key in keys.items(): + for i, (attn, mlp) in enumerate(zip(cfg["layer_types"], cfg["mlp_layer_types"])): + if attn == key and mlp == "sparse": + out.append((kind, i)) + break + else: + raise ValueError(f"no {kind} MoE layer in {model_dir}") + return out + + +def one_dense(out: Path, model_dir: Path, cp, prefix: str, base: str, hidden: torch.Tensor) -> dict: + """One non-routed group's inputs and the engine's outputs, bit for bit.""" + from tensorfold.cuda.exl3 import linear + + t = cp.groups[prefix] + lin = linear.Exl3Linear.load(model_dir, prefix, layout="stored") + gid = prefix[len(base) + 1:] + if lin.k == hidden.shape[1]: + x = hidden + else: + g = torch.Generator().manual_seed(SEED + zlib.crc32(gid.encode())) + x = torch.randn((max(ROWS), lin.k), generator=g).to(torch.bfloat16) + xmap, ymap, sha = {}, {}, {} + for m in ROWS: + y = mm_rows(x, lin, m) + save(out, f"{lid_of(base)}.{gid}.x{m}", x[:m]) + save(out, f"{lid_of(base)}.{gid}.y{m}", y) + xmap[str(m)] = f"{lid_of(base)}.{gid}.x{m}.bin" + ymap[str(m)] = f"{lid_of(base)}.{gid}.y{m}.bin" + sha[str(m)] = sha_of(y) + entry = {"id": gid, "name": prefix, "bits": int(t.bits), "codebook": t.codebook, + "k": lin.k, "n": lin.n, "rows": list(ROWS), "x": xmap, "y": ymap, "sha256": sha} + print(f" {gid}: EXL3 {t.bits:g}-bit {t.codebook} K={lin.k} N={lin.n}") + return entry + + +def expert_part(out: Path, model_dir: Path, cp, base: str, hidden: torch.Tensor, cfg: dict) -> dict: + """All experts as one grouped call with the engine's routing, bit for bit.""" + from tensorfold.cuda.exl3 import experts as generic + from tensorfold.families.glm5_next.cuda import glue + + codebook = next(t.codebook for p, t in cp.groups.items() if p.startswith(f"{base}.mlp.experts.0.")) + gate, up, down = [], [], [] + k2 = {"gate": [], "up": [], "down": []} + for e in range(cfg["n_routed_experts"]): + for part, mats in (("gate_proj", gate), ("up_proj", up), ("down_proj", down)): + mats.append(tuple(load_tensor(model_dir, f"{base}.mlp.experts.{e}.{part}.{x}", "cuda") + for x in ("trellis", "suh", "svh"))) + for mats, key in ((gate, "gate"), (up, "up"), (down, "down")): + k2[key] = [generic.k2_of(m[0]) for m in mats] + ex = generic.prepare(gate, up, down, codebook, device="cuda") + nbytes = sum(m[0].element_size() * m[0].numel() for mats in (gate, up, down) for m in mats) \ + + sum(t.element_size() * t.numel() for t in (ex.suh_g, ex.suh_u, ex.svh_g, ex.svh_u, ex.suh_d, ex.svh_d)) + print(f" experts: {ex.count} experts, {ex.dims}x{ex.width}, {nbytes / (1 << 30):.2f} GiB") + + # the routing, as the GLM engine's moe_block routes: glue.router on the bf16 gate weight, glue.select on the bias + gate_w = load_tensor(model_dir, f"{base}.mlp.gate.weight", "cuda").to(torch.bfloat16).contiguous() + bias = load_tensor(model_dir, f"{base}.mlp.gate.e_score_correction_bias", "cuda").float().contiguous() + logits = torch.empty((hidden.shape[0], ex.count), dtype=torch.float32, device="cuda") + glue.router(hidden.cuda(), gate_w, logits) + topk = cfg["num_experts_per_tok"] + pick = torch.empty((hidden.shape[0], topk + 1), dtype=torch.int32, device="cuda") + wts = torch.empty((hidden.shape[0], topk + 1), dtype=torch.float32, device="cuda") + glue.select(logits, bias, pick, wts, topk, ex.count, cfg["routed_scaling_factor"], + cfg.get("norm_topk_prob", True)) + pick, wts = pick[:, :topk].contiguous(), wts[:, :topk].contiguous() + + scratch = generic.Scratch(ex, max(EXPERT_ROWS), topk) + xmap, pmap, wmap, ymap, omap, sha = {}, {}, {}, {}, {}, {} + for m in EXPERT_ROWS: + y = generic.routed(hidden[:m].cuda(), pick[:m].cuda(), wts[:m].cuda(), ex, scratch, None, m, + limit=cfg.get("swiglu_limit", 10.0), act_mode=generic.ACT_BF16) + torch.cuda.synchronize() + for key, t in (("x", hidden[:m]), ("pick", pick[:m]), ("wts", wts[:m]), + ("y", scratch.y[:m * topk]), ("out", y)): + save(out, f"{lid_of(base)}.experts.{key}{m}", t) + xmap[str(m)] = f"{lid_of(base)}.experts.x{m}.bin" + pmap[str(m)] = f"{lid_of(base)}.experts.pick{m}.bin" + wmap[str(m)] = f"{lid_of(base)}.experts.wts{m}.bin" + ymap[str(m)] = f"{lid_of(base)}.experts.y{m}.bin" + omap[str(m)] = f"{lid_of(base)}.experts.out{m}.bin" + sha[str(m)] = sha_of(y) + return {"name": f"{base}.mlp.experts", "count": ex.count, "dims": ex.dims, "width": ex.width, + "slots": topk, "codebook": codebook, "limit": cfg.get("swiglu_limit", 10.0), + "act_mode": generic.ACT_BF16, "rows": list(EXPERT_ROWS), "gate_k2": k2["gate"], + "up_k2": k2["up"], "down_k2": k2["down"], "x": xmap, "pick": pmap, "wts": wmap, + "y": ymap, "out": omap, "sha256": sha} + + +def one_layer(out: Path, model_dir: Path, cp, cfg: dict, li: int, kind: str) -> dict: + base = f"model.language_model.layers.{li}" + print(f"{kind} layer {li}:") + g = torch.Generator().manual_seed(SEED + li) + hidden = torch.randn((max(ROWS), cfg["hidden_size"]), generator=g).to(torch.bfloat16) + groups = [one_dense(out, model_dir, cp, prefix, base, hidden) + for prefix in sorted(cp.groups) + if prefix.startswith(base + ".") and ".experts." not in prefix] + return {"id": kind, "layer": li, "name": base, "dims": cfg["hidden_size"], + "groups": groups, "experts": expert_part(out, model_dir, cp, base, hidden, cfg)} + + +def write(out: Path, model_dir: Path) -> list: + from tensorfold.cuda.exl3 import format as fmt + + cp = fmt.scan(model_dir) + cfg = json.loads((Path(model_dir) / "config.json").read_text())["text_config"] + layers = [one_layer(out, model_dir, cp, cfg, li, kind) for kind, li in pick_layers(model_dir)] + print(f"{len(layers)} layers -> {out}") + return layers diff --git a/zig/tests/cuda/oracle/exl3_oracle_tp2.py b/zig/tests/cuda/oracle/exl3_oracle_tp2.py new file mode 100644 index 0000000000..961df5f5b9 --- /dev/null +++ b/zig/tests/cuda/oracle/exl3_oracle_tp2.py @@ -0,0 +1,232 @@ +"""The tensor-parallel fixture: two real projections and one MoE layer's routed experts, split as the ranks are.""" +from __future__ import annotations + +import hashlib +import re +from pathlib import Path + +import torch + +import exl3_oracle_experts + +ROWS = (1, 3, 16) +SEED = 1237 +EXPERTS_SEED = 1238 +POOL = 3 # experts per width, as the routed fixture's pool +SLOTS = 8 +LIMIT = 10.0 + + +def save(out: Path, name: str, t: torch.Tensor) -> None: + (out / f"{name}.bin").write_bytes(t.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes()) + + +def cut_dim1(raw: torch.Tensor, rank: int) -> torch.Tensor: + """A tensor's second axis in half (the rank's tile columns / scale rows), contiguous.""" + half = raw.shape[1] // 2 + return raw[:, rank * half:(rank + 1) * half].contiguous() + + +def cut_row(raw: torch.Tensor, rank: int) -> torch.Tensor: + """A tensor's leading axis in half (the rank's tile rows), contiguous.""" + half = raw.shape[0] // 2 + return raw[rank * half:(rank + 1) * half].contiguous() + + +def load_group(model_dir: Path, prefix: str) -> tuple: + """(trellis int16 as stored, suh, svh) of one linear, from the checkpoint (on the GPU).""" + from safetensors import safe_open + + root = Path(model_dir) + index = root / "model.safetensors.index.json" + rel = None + if index.exists(): + weight_map = __import__("json").loads(index.read_text())["weight_map"] + rel = weight_map.get(f"{prefix}.trellis") + if rel is None: + for part in sorted(root.glob("*.safetensors")): + with safe_open(str(part), framework="pt") as g: + if f"{prefix}.trellis" in set(g.keys()): + rel = part.name + break + if rel is None: + raise ValueError(f"{prefix}.trellis is in no file of {root}") + with safe_open(str(root / (rel or "model.safetensors")), framework="pt", device="cuda") as f: + return (f.get_tensor(f"{prefix}.trellis"), f.get_tensor(f"{prefix}.suh"), f.get_tensor(f"{prefix}.svh")) + + +def pick_out_linear(cp) -> object: + """The out-split layer: a shared_experts.gate_proj (or any gate/up one) whose cut stays on 128 boundaries.""" + cands = [t for t in cp.groups.values() + if re.search(r"\.(mlp\.shared_experts\.)?(gate|up)_proj$", t.prefix) + and t.in_scales == "suh" and t.out_scales == "svh" + and t.k % 256 == 0 and t.n % 256 == 0 and t.k <= 8192 and t.n <= 8192] + if not cands: + raise ValueError("no EXL3 gate/up projection splits on 128-boundaries at world 2") + shared = [t for t in cands if "shared_experts" in t.prefix] + return min(shared or cands, key=lambda t: (t.k * t.n, t.prefix)) + + +def pick_in_linear(cp, out) -> object: + """The in-split layer: the same layer's shared_experts.down_proj, else a down_proj, else an o_proj.""" + base = out.prefix.rsplit(".mlp.", 1)[0] if ".mlp." in out.prefix else out.prefix.rsplit(".self_attn.", 1)[0] + prefs = (f"{base}.mlp.shared_experts.down_proj", f"{base}.mlp.down_proj") + for name in prefs: + t = cp.groups.get(name) + if t is not None and t.in_scales == "suh" and t.out_scales == "svh" and t.k % 256 == 0 and t.n % 256 == 0: + return t + cands = [t for t in cp.groups.values() + if (t.prefix.endswith(".mlp.down_proj") or t.prefix.endswith(".self_attn.o_proj")) + and t.in_scales == "suh" and t.out_scales == "svh" and t.k % 256 == 0 and t.n % 256 == 0 + and t.k <= 8192 and t.n <= 8192] + if not cands: + raise ValueError("no EXL3 down/o projection splits on 128-boundaries at world 2") + return min(cands, key=lambda t: (t.k * t.n, t.prefix)) + + +def linear_entry(out: Path, model_dir: Path, t, split: str, lid: str) -> dict: + """One projection's full tensors, x rows, and per rank the engine's output on its shard.""" + from tensorfold.cuda.exl3.linear import Exl3Linear + + trellis, suh, svh = load_group(model_dir, t.prefix) + save(out, f"{lid}.trellis", trellis) + save(out, f"{lid}.suh", suh) + save(out, f"{lid}.svh", svh) + g = torch.Generator().manual_seed(SEED) + x = torch.randn((max(ROWS), t.k), generator=g).to(torch.bfloat16) + xmap, sha = {}, {} + for m in ROWS: + save(out, f"{lid}.x{m}", x[:m]) + xmap[str(m)] = f"{lid}.x{m}.bin" + + if split == "dim1": # out-split: the rank's bf16 half of the outputs + shards = {} + for rank in (0, 1): + lin = Exl3Linear.from_tensors(cut_dim1(trellis, rank), suh, cut_row(svh, rank), t.codebook, + None, layout="stored") + for m in ROWS: + y = lin(x[:m].cuda()) + torch.cuda.synchronize() + save(out, f"{lid}.shard{m}.r{rank}", y) + shards.setdefault(str(m), []).append(f"{lid}.shard{m}.r{rank}.bin") + entry = {"shards": shards} + else: # in-split: the rank's fp32 partial over the full outputs + partials, ymap = {}, {} + plist = {0: [], 1: []} + for rank in (0, 1): + lin = Exl3Linear.from_tensors(cut_row(trellis, rank), cut_row(suh, rank), svh, t.codebook, + None, layout="stored") + xr = x[:, rank * (t.k // 2):(rank + 1) * (t.k // 2)].contiguous() + for m in ROWS: + part = lin(xr[:m].cuda(), out_dtype=torch.float32) + torch.cuda.synchronize() + save(out, f"{lid}.part{m}.r{rank}", part) + partials.setdefault(str(m), []).append(f"{lid}.part{m}.r{rank}.bin") + plist[rank].append(part) + for m, p0, p1 in zip(ROWS, plist[0], plist[1]): + y = (p0 + p1).to(torch.bfloat16) # rank 0 first, fp32, then bf16 + save(out, f"{lid}.y{m}", y) + ymap[str(m)] = f"{lid}.y{m}.bin" + sha[str(m)] = hashlib.sha256(y.cpu().view(torch.uint8).numpy().tobytes()).hexdigest() + entry = {"partials": partials, "y": ymap, "sha256": sha} + + print(f"{lid}: {t.prefix} EXL3 {t.bits:g}-bit {t.codebook} K={t.k} N={t.n} split={split} " + f"sha256={sha.get('16', '')[:16]}…") + return {"id": lid.rsplit(".", 1)[1], "name": t.prefix, "split": split, "bits": int(t.bits), + "codebook": t.codebook, "k": t.k, "n": t.n, "rows": list(ROWS), + "trellis": f"{lid}.trellis.bin", "suh": f"{lid}.suh.bin", "svh": f"{lid}.svh.bin", + "bias": None, "x": xmap, **entry} + + +def experts_entry(out: Path, model_dir: Path, cp) -> dict: + """One MoE layer's routed experts, full tensors and each rank's fp32 partials (its gate/up/down shards).""" + from tensorfold.cuda.exl3 import experts as generic + + base = exl3_oracle_experts.pick_layer(cp) + by_bits: dict[float, list[int]] = {} + for prefix, t in cp.groups.items(): + m = re.match(rf"^{re.escape(base)}\.mlp\.experts\.(\d+)\.gate_proj$", prefix) + if m is not None: + by_bits.setdefault(t.bits, []).append(int(m.group(1))) + pool = [e for b in (4.0, 5.0, 6.0) for e in sorted(by_bits[b])[:POOL]] + + gate, up, down = [], [], [] + trellis = {"gate": [], "up": [], "down": []} + for slot, e in enumerate(pool): + for part, short, mats, names in (("gate_proj", "gate", gate, trellis["gate"]), + ("up_proj", "up", up, trellis["up"]), + ("down_proj", "down", down, trellis["down"])): + mats.append(load_group(model_dir, f"{base}.mlp.experts.{e}.{part}")) + names.append(f"tp2.experts.{short}{slot}.bin") + save(out, "tp2.experts.suh_g", torch.stack([m[1] for m in gate])) + save(out, "tp2.experts.suh_u", torch.stack([m[1] for m in up])) + save(out, "tp2.experts.svh_g", torch.stack([m[2] for m in gate])) + save(out, "tp2.experts.svh_u", torch.stack([m[2] for m in up])) + save(out, "tp2.experts.suh_d", torch.stack([m[1] for m in down])) + save(out, "tp2.experts.svh_d", torch.stack([m[2] for m in down])) + for slot, mats in enumerate(zip(gate, up, down)): + for name, t in zip(("gate", "up", "down"), mats): + save(out, f"tp2.experts.{name}{slot}", t[0]) + + g = torch.Generator().manual_seed(EXPERTS_SEED) + dims, width = gate[0][0].shape[0] * 16, gate[0][0].shape[1] * 16 + x = torch.randn((max(ROWS), dims), generator=g).to(torch.bfloat16) + pick = torch.stack([torch.randperm(len(pool), generator=g)[:SLOTS] for _ in range(max(ROWS))]).int() + wts = (0.1 + 0.9 * torch.rand((max(ROWS), SLOTS), generator=g)).float() + xmap, pickmap, wtsmap = {}, {}, {} + for m in ROWS: + save(out, f"tp2.experts.x{m}", x[:m]) + save(out, f"tp2.experts.pick{m}", pick[:m]) + save(out, f"tp2.experts.wts{m}", wts[:m]) + xmap[str(m)] = f"tp2.experts.x{m}.bin" + pickmap[str(m)] = f"tp2.experts.pick{m}.bin" + wtsmap[str(m)] = f"tp2.experts.wts{m}.bin" + + partials, ymap, sha = {}, {}, {} + plist = {0: [], 1: []} + for rank in (0, 1): + def shard(mats, part): + """The rank's (trellis, suh, svh) of one expert projection: out-split for gate/up, in for down.""" + cut = cut_dim1 if part in ("gate", "up") else (lambda t, r: cut_row(t, r)) + scale_at = 2 if part in ("gate", "up") else 1 + return [(cut(t, rank), suh if scale_at == 2 else cut_row(suh, rank), + cut_row(svh, rank) if scale_at == 2 else svh) for t, suh, svh in mats] + + ex = generic.prepare(shard(gate, "gate"), shard(up, "up"), shard(down, "down"), "mul1") + scratch = generic.Scratch(ex, max(ROWS), SLOTS) + for m in ROWS: + part = generic.routed(x[:m].cuda(), pick[:m].cuda(), wts[:m].cuda(), ex, scratch, None, m, + limit=LIMIT, act_mode=generic.ACT_BF16) + torch.cuda.synchronize() + save(out, f"tp2.experts.part{m}.r{rank}", part) + partials.setdefault(str(m), []).append(f"tp2.experts.part{m}.r{rank}.bin") + plist[rank].append(part) + for m, p0, p1 in zip(ROWS, plist[0], plist[1]): + y = (p0 + p1).to(torch.bfloat16) # rank 0 first, fp32, then bf16 + save(out, f"tp2.experts.y{m}", y) + ymap[str(m)] = f"tp2.experts.y{m}.bin" + sha[str(m)] = hashlib.sha256(y.cpu().view(torch.uint8).numpy().tobytes()).hexdigest() + + print(f"tp2 experts: {base} {len(pool)} experts (4/5/6-bit pool), width {width}, " + f"sha256={sha['16'][:16]}…") + return {"name": base, "count": len(pool), "dims": dims, "width": width, "slots": SLOTS, + "codebook": "mul1", "limit": LIMIT, "act_mode": generic.ACT_BF16, "rows": list(ROWS), + "gate_k2": [generic.k2_of(m[0]) for m in gate], + "up_k2": [generic.k2_of(m[0]) for m in up], + "down_k2": [generic.k2_of(m[0]) for m in down], + "trellis": trellis, + "suh_g": "tp2.experts.suh_g.bin", "suh_u": "tp2.experts.suh_u.bin", + "svh_g": "tp2.experts.svh_g.bin", "svh_u": "tp2.experts.svh_u.bin", + "suh_d": "tp2.experts.suh_d.bin", "svh_d": "tp2.experts.svh_d.bin", + "x": xmap, "pick": pickmap, "wts": wtsmap, "partials": partials, "y": ymap, "sha256": sha} + + +def write(out: Path, model_dir: Path) -> dict: + from tensorfold.cuda.exl3 import format as fmt + + cp = fmt.scan(model_dir) + t_out = pick_out_linear(cp) + t_in = pick_in_linear(cp, t_out) + return {"linears": [linear_entry(out, model_dir, t_out, "dim1", "tp2.linear.out"), + linear_entry(out, model_dir, t_in, "row", "tp2.linear.in")], + "experts": experts_entry(out, model_dir, cp)}