diff --git a/docs/roce-gather.md b/docs/roce-gather.md new file mode 100644 index 0000000000..e0ff5ae82c --- /dev/null +++ b/docs/roce-gather.md @@ -0,0 +1,24 @@ +# RoCE all-gather for two ranks + +One-shot all-gather between two DGX Sparks over RoCE: the GPU side of b12x's "RoCEnante" exchange (https://github.com/local-inference-lab/b12x, b12x/comm/roce, Apache License 2.0), written for two ranks. The host side (zig/src/cuda/roce.zig) owns a pinned region both GPUs and both NICs reach: + +``` + recv[slot] the peer's shard, written by the peer's NIC (2 slots of slot_bytes) + flag[slot][hca] the sequence the peer's NIC writes after its stripe on that HCA (128 bytes apart) + send[slot] this rank's shard, staged here for its NIC (2 slots) + ctrl {seq, nbytes of slot 0, nbytes of slot 1, error seq, missing hca}: the proxy's doorbell +``` + +One launch: (1) every block stages its share of the shard into send[seq & 1]; (2) the last block to stage publishes the slot's size and rings ctrl.seq after a system fence, and the proxy thread RDMA-writes the shard and then a flag per HCA stripe on the same queue pair (so a flag never lands before its stripe); (3) each block waits for every stripe's flag; (4) every block writes both shards in rank order (the layout NCCL's all-gather gives); (5) the last block to finish advances the device epoch, so a CUDA graph replays the exchange with new sequences. A wait past the spin limit poisons the exchange: later launches return at once and the host reports the error. + +## Host side + +Small all-gathers between two ranks as one-shot RDMA writes over RoCE: b12x's "RoCEnante" exchange (https://github.com/local-inference-lab/b12x, b12x/comm/roce, Apache License 2.0), written for two ranks: the kernel (kernels/cuda/roce_gather.cu) stages this rank's shard in pinned host memory and rings a doorbell; a proxy thread here RDMA-writes it to the peer, striped over the HCAs, each stripe followed by its sequence flag on the same queue pair; the kernel waits for the peer's flags and copies both shards out. A decode step's gathers are a few rows of partials. On two GB10s with one ConnectX-7 cable (NCCL 2.30.7, `tf-cuda-test roce-gather`, 2,000 calls back to back) this path took 10.4 us a call at 4 KiB a rank against NCCL's 12.3, 11.6 against 13.5 at 8 KiB, 17.4 against 20.5 at 32 KiB, and 17.4 against 51.5 at 64 KiB. libibverbs is opened at run time (as the driver and NCCL are): without it, or without an active HCA, the caller keeps NCCL. + +## Setup and use + +`tp_group.Group.openRoce` brings the path up on a two-rank group: each rank opens its HCAs, and what each brings (its region, keys, queue pairs and the largest shard it would send) crosses the group's own NCCL communicator, as does each rank's verdict after connecting. Both ranks get the path or neither does (`TF_ROCE=0` on either rank keeps NCCL). From then on `Group.allGather` sends a gather of up to the smaller rank's cap (512 KiB by default) this way and leaves larger ones to NCCL; the size alone decides, so both ranks always take the same path. A rank whose buffers are not 16-byte aligned gets `error.Misaligned` instead of falling back alone. Liveness is `tp_watch`'s, as for every group. + +A failed exchange writes nothing into the gathered rows. Call `Group.check` after any synchronize whose result reads a gather: it returns the error the GPU or the proxy recorded. + +`tf-cuda-test roce-gather ` runs on both ranks (NCCL_COMM_ID set): every size up to the cap and past it against NCCL's all-gather, byte for byte, then each path's time a call. diff --git a/tools/glm_vision_ref.py b/tools/glm_vision_ref.py new file mode 100644 index 0000000000..25c1c4234a --- /dev/null +++ b/tools/glm_vision_ref.py @@ -0,0 +1,213 @@ +#!/usr/bin/env python3 +"""GLM-5.3-Flash's vision tower as mlx-vlm's glm5_next (vision.py, processing.py) defines it, in torch on the CPU +(fp32 math on the checkpoint's bf16 weights): the reference the native tower is checked against. + +usage: glm_vision_ref.py [--bf16] ... + --bf16: the same tower in bfloat16 (tf-cuda-test glm-vision bounds the native tower's error by this run's) + an image: its patches (the GLM processor's canvas, padding, normalization and merge order) -> features + a folder of frame_*.png: a video's sampled frames (already sampled) -> one [t, h, w] grid +writes /case/{patches.f32, grid.json, features.f32, block0.f32} +""" +import json, math, struct, sys +from pathlib import Path + +import numpy as np +import torch + +DT = torch.float32 +from PIL import Image + +MEAN = np.array([0.48145466, 0.4578275, 0.40821073], np.float32) +STD = np.array([0.26862954, 0.26130258, 0.27577711], np.float32) +PATCH, MERGE, TEMPORAL = 14, 2, 2 + + +def smart_resize(num_frames, height, width, temporal_factor, factor, min_pixels, max_pixels): + ppt = temporal_factor * factor ** 2 + min_pixels *= ppt + max_pixels *= ppt + + def align(v, f): + return math.ceil(v / f) * f + + aligned_frames = max(temporal_factor, round(num_frames / temporal_factor) * temporal_factor) + + def fit(): + low, high, best = 1, height, (factor, factor) + while low <= high: + ch = (low + high) // 2 + cw = max(1, math.floor(width * ch / height)) + cand = (align(ch, factor), align(cw, factor)) + if aligned_frames * cand[0] * cand[1] <= max_pixels: + best, low = cand, ch + 1 + else: + high = ch - 1 + return best + + ah, aw = align(height, factor), align(width, factor) + pixels = aligned_frames * ah * aw + if pixels < min_pixels: + scale = math.sqrt(min_pixels / (num_frames * height * width)) + ah, aw = align(max(1, math.ceil(height * scale)), factor), align(max(1, math.ceil(width * scale)), factor) + pixels = aligned_frames * ah * aw + if pixels > max_pixels: + ah, aw = fit() + return ah, aw + + +def geometry(num_frames, height, width, min_tokens, max_tokens): + factor = PATCH * MERGE + ppt = TEMPORAL * factor ** 2 + th, tw = smart_resize(num_frames, height, width, TEMPORAL, factor, min_tokens, max_tokens) + scale = min(th / height, tw / width) + if num_frames * height * width >= min_tokens * ppt: + scale = min(1.0, scale) + return th, tw, max(1, min(th, math.floor(height * scale))), max(1, min(tw, math.floor(width * scale))) + + +def frames_to_patches(frames, min_tokens, max_tokens): + """frames: list of RGB PIL images (1 for an image) -> (patches [n, 1176] fp32, grid [t, h, w]).""" + n = len(frames) + w0, h0 = frames[0].size + th, tw, ch, cw = geometry(TEMPORAL if n == 1 else n, h0, w0, min_tokens, max_tokens) + arr = [] + for f in frames: + f = f if f.size == (cw, ch) else f.resize((cw, ch), Image.BICUBIC) + a = np.asarray(f, np.uint8).transpose(2, 0, 1) + a = np.pad(a, ((0, 0), (0, th - ch), (0, tw - cw))).astype(np.float32) / np.float32(255.0) + arr.append((a - MEAN[:, None, None]) / STD[:, None, None]) + gh, gw = th // PATCH, tw // PATCH + if n == 1: # an image: its one frame twice (the temporal patch) + a = arr[0] + p = a.reshape(3, gh // MERGE, MERGE, PATCH, gw // MERGE, MERGE, PATCH).transpose(1, 4, 2, 5, 0, 3, 6) + p = np.broadcast_to(p[:, :, :, :, :, None], (*p.shape[:5], TEMPORAL, *p.shape[5:])) + return np.ascontiguousarray(p.reshape(gh * gw, 3 * TEMPORAL * PATCH * PATCH)), [1, gh, gw] + f = np.stack(arr) + if n % TEMPORAL: + f = np.concatenate([f, f[-1:]]) + gt = f.shape[0] // TEMPORAL + p = f.reshape(gt, TEMPORAL, 3, gh // MERGE, MERGE, PATCH, gw // MERGE, MERGE, PATCH).transpose(0, 3, 6, 4, 7, 2, 1, 5, 8) + return np.ascontiguousarray(p.reshape(gt * gh * gw, 3 * TEMPORAL * PATCH * PATCH)), [gt, gh, gw] + + +def load_visual(model_dir): + idx = json.loads((model_dir / "model.safetensors.index.json").read_text())["weight_map"] + files = sorted({f for k, f in idx.items() if k.startswith("model.visual.")}) + out = {} + for name in files: + raw = (model_dir / name).read_bytes() + n = struct.unpack(" {feats.shape[0]} tokens, |features| rms {feats.float().pow(2).mean().sqrt():.4f}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 684d6c3cc9..3530629e25 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -52,12 +52,16 @@ const kernels = [_]Kernel{ .{ .name = "glm_glue", .flags = &.{"-O3"} }, // ours: the Python family's Triton glue by hand, its PTX op for op .{ .name = "glm_ops", .flags = &.{"-O3"} }, // ours: GLM's integer work the Python family did in torch (the DSA token lists) .{ .name = "glm_quant", .flags = torch_ops }, // ours: GLM's draft-only weight quantizers, torch at load in the Python family + .{ .name = "glm_exl3", .flags = &.{"-O3"} }, // ours: GLM's EXL3 routed experts (Mia-AiLab's checkpoint), decode and prompt + .{ .name = "glm_vision", .flags = &.{"-O3"} }, // ours: GLM's image tower steps around cuBLASLt (norms, rotary, softmax, merger) + .{ .name = "glm_mm", .flags = &.{"-O3"} }, // ours: GLM's bf16 linears on tensor cores, weights tiled at load, row-stable .{ .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 }, .{ .name = "torch_indexing", .src = "torch_ops/indexing", .flags = torch_ops }, .{ .name = "torch_movement", .src = "torch_ops/movement", .flags = torch_ops }, .{ .name = "torch_nemotron_constants", .src = "torch_ops/nemotron_constants", .flags = torch_ops }, + .{ .name = "roce_gather", .src = "roce_gather", .flags = &.{"-O3"} }, // b12x's one-shot RoCE all-gather, two ranks }; /// torch.utils.cpp_extension's own nvcc flags (torch 2.13): C++20 and which half/bf16 operators the headers define. diff --git a/zig/kernels/cuda/glm_exl3.cu b/zig/kernels/cuda/glm_exl3.cu new file mode 100644 index 0000000000..4846e60371 --- /dev/null +++ b/zig/kernels/cuda/glm_exl3.cu @@ -0,0 +1,363 @@ +// GLM-5.3-Flash's EXL3 routed experts (4-bit mcg trellis, ExLlamaV3's format: MIT, (c) 2025 Turboderp): decode, prompt. + +#include +#include +#include + +namespace { + +// Lets a dependent launch start on its weights early (sm_90 and newer; older targets run in plain stream order). +__device__ __forceinline__ void launch_dependents() { +#if __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +#endif +} + +constexpr float HAD_SCALE = 0.08838834764831845f; // 1 / sqrt(128) + +// Two values of the "mcg" codebook from two 16-bit states, as a half2 (first state in .x). +__device__ __forceinline__ uint32_t mcg2(uint32_t s0, uint32_t s1) { + uint32_t x0 = s0 * 0xCBAC1FEDu; + uint32_t x1 = s1 * 0xCBAC1FEDu; + x0 = (x0 & 0x8FFF8FFFu) ^ 0x3B603B60u; + x1 = (x1 & 0x8FFF8FFFu) ^ 0x3B603B60u; + uint32_t lo = __byte_perm(x0, x1, 0x5410); + uint32_t hi = __byte_perm(x0, x1, 0x7632); + half2 r = __hadd2(*reinterpret_cast(&lo), *reinterpret_cast(&hi)); + return *reinterpret_cast(&r); +} + +// This lane's eight values of a 4-bit tile (word = tile[lane]) as the B fragments of its two n8 halves. +__device__ __forceinline__ void decode_tile(uint32_t w, int lane, uint32_t (&b0)[2], uint32_t (&b1)[2]) { + uint32_t p = __shfl_sync(0xffffffffu, w, (lane + 31) & 31); + uint32_t s = __funnelshift_r(w, p, 20); + b0[0] = mcg2((s >> 8) & 0xffffu, (s >> 4) & 0xffffu); + b0[1] = mcg2(s & 0xffffu, w >> 16); + b1[0] = mcg2((w >> 12) & 0xffffu, (w >> 8) & 0xffffu); + b1[1] = mcg2((w >> 4) & 0xffffu, w & 0xffffu); +} + +__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])); +} + +__device__ __forceinline__ uint32_t load_pair(const half* x, bool ok) { + return ok ? *reinterpret_cast(x) : 0u; +} + +// Grid (item, n block, mat * SK + split) -> Z[mat][split][pair][n]: an expert and <= 16 MT of its pairs; warp order. +template +__device__ __forceinline__ void grouped( + const half* __restrict__ X0, const half* __restrict__ X1, const uint32_t* __restrict__ T0, + const uint32_t* __restrict__ T1, const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, float* __restrict__ Z, int K, int N, int P, int SK) { + const int item = blockIdx.x; + if (item >= counts[0]) return; + const int e = items[3 * item], first = items[3 * item + 1], cnt = items[3 * item + 2]; + const int split = blockIdx.z % SK; + const int mat = blockIdx.z / SK; + const half* X = mat ? X1 : X0; + const uint32_t* T = mat ? T1 : T0; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int g = lane >> 2, t = lane & 3; + const int KT = K >> 4, NTILES = N >> 4; + constexpr int ROWS = 16 * MT; + + __shared__ int rows_sh[ROWS]; + for (int i = threadIdx.x; i < ROWS; i += W * 32) rows_sh[i] = i < cnt ? members[first + i] : -1; + __syncthreads(); + int r0[MT], r1[MT]; + const half* x0[MT]; + const half* x1[MT]; +#pragma unroll + for (int m = 0; m < MT; ++m) { + r0[m] = rows_sh[m * 16 + g]; + r1[m] = rows_sh[m * 16 + g + 8]; + x0[m] = X + (size_t)(r0[m] < 0 ? 0 : r0[m]) * K + 2 * t; + x1[m] = X + (size_t)(r1[m] < 0 ? 0 : r1[m]) * K + 2 * t; + } + + const int per_split = KT / SK, per_warp = per_split / W; + const int kt0 = split * per_split + warp * per_warp; + const int nt0 = blockIdx.y * NT; + const uint32_t* tile = T + (((size_t)e * KT + kt0) * NTILES + nt0) * 32 + lane; + + float acc[MT][NT][2][4]; +#pragma unroll + for (int m = 0; m < MT; ++m) +#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[m][i][h][c] = 0.f; + + for (int kt = kt0; kt < kt0 + per_warp; ++kt) { + uint32_t words[NT]; +#pragma unroll + for (int i = 0; i < NT; ++i) words[i] = __ldg(tile + i * 32); + const int k = kt * 16; + uint32_t a[MT][4]; +#pragma unroll + for (int m = 0; m < MT; ++m) { + a[m][0] = load_pair(x0[m] + k, r0[m] >= 0); + a[m][1] = load_pair(x1[m] + k, r1[m] >= 0); + a[m][2] = load_pair(x0[m] + k + 8, r0[m] >= 0); + a[m][3] = load_pair(x1[m] + k + 8, r1[m] >= 0); + } +#pragma unroll + for (int i = 0; i < NT; ++i) { + uint32_t b0[2], b1[2]; + decode_tile(words[i], lane, b0, b1); +#pragma unroll + for (int m = 0; m < MT; ++m) { + mma16816(acc[m][i][0], a[m], b0); + mma16816(acc[m][i][1], a[m], b1); + } + } + tile += (size_t)NTILES * 32; + } + + // warps' partial sums through shared memory, added in warp order + __shared__ float red[W][ROWS][NT * 16]; +#pragma unroll + for (int m = 0; m < MT; ++m) +#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][m * 16 + g][col] = acc[m][i][h][0]; + red[warp][m * 16 + g][col + 1] = acc[m][i][h][1]; + red[warp][m * 16 + g + 8][col] = acc[m][i][h][2]; + red[warp][m * 16 + g + 8][col + 1] = acc[m][i][h][3]; + } + __syncthreads(); + for (int idx = threadIdx.x; idx < ROWS * 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]; + Z[(((size_t)mat * SK + split) * P + r) * N + nt0 * 16 + col] = s; + } +} + +// Fast Walsh-Hadamard transform of 128 values, 4 a lane, natural order: strides 1, 2 in registers, 4..64 across lanes. +__device__ __forceinline__ void fwht128(float (&v)[4], int lane) { + 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) { + float o = __shfl_xor_sync(0xffffffffu, v[j], m); + v[j] = (lane & m) ? o - v[j] : v[j] + o; + } + } +} + +__device__ __forceinline__ float bf16r(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +// The SwiGLU limits as torch.clamp takes them: a NaN stays NaN (fminf and fmaxf would drop it for the bound). +__device__ __forceinline__ float clamp_hi(float x, float hi) { return x > hi ? hi : x; } +__device__ __forceinline__ float clamp_both(float x, float l) { return x < -l ? -l : (x > l ? l : x); } + + +// -- helpers the transposed prompt kernels share (glm_exl3_transposed.cuh) -- + +// decode_tile via PRMT (MiaAI-Lab GLM-5.3-Flash-EXL3-2x-DGX-Sparks-TensorFold patch 0020, Apache-2.0): same bits. +__device__ __forceinline__ void decode_tile_p(uint32_t w, int lane, uint32_t (&b0)[2], uint32_t (&b1)[2]) { + const uint32_t p = __shfl_sync(0xffffffffu, w, (lane + 31) & 31); + const uint32_t s = __funnelshift_r(w, p, 20); + const uint32_t t = w >> 4; + b0[0] = mcg2(__byte_perm(s, 0u, 0x4421), (s >> 4) & 0xffffu); // (s >> 8) & 0xffff, (s >> 4) & 0xffff + b0[1] = mcg2(s & 0xffffu, w >> 16); + b1[0] = mcg2(__byte_perm(t, 0u, 0x4421), __byte_perm(w, 0u, 0x4421)); // (w >> 12) & 0xffff, (w >> 8) & 0xffff + b1[1] = mcg2(__byte_perm(t, 0u, 0x4410), w & 0xffffu); // (w >> 4) & 0xffff, w & 0xffff +} + +__device__ __forceinline__ void cp_async16(void* dst, const void* src, bool valid) { + const unsigned d = (unsigned)__cvta_generic_to_shared(dst); + const int n = valid ? 16 : 0; // src-size 0: the 16 bytes come as zeros (a padding row) + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(d), "l"(src), "r"(n)); +} +__device__ __forceinline__ void cp_async_commit() { asm volatile("cp.async.commit_group;\n" ::); } +template +__device__ __forceinline__ void cp_async_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); } + +__device__ __forceinline__ void ldmatrix_x4(uint32_t (&a)[4], const half* p) { + const unsigned s = (unsigned)__cvta_generic_to_shared(p); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) : "r"(s)); +} + +} // namespace + +// Program (row, K block): Xh[row][block] = fp16((x[row] * suh) @ H) when a layer's experts share suh; rot_in's bits. +extern "C" __global__ void glm_exl3_rot_rows(const __nv_bfloat16* __restrict__ x, int x_stride, const half* __restrict__ suh, + half* __restrict__ out, int K) { + launch_dependents(); + const int row = blockIdx.x, blk = blockIdx.y; + const int lane = threadIdx.x; + const half* sh = suh + blk * 128 + 4 * lane; + const __nv_bfloat16* xr = x + (size_t)row * x_stride + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = __bfloat162float(xr[j]) * __half2float(sh[j]); + fwht128(v, lane); + half* o = out + (size_t)row * K + blk * 128 + 4 * lane; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); +} + +namespace { +} // namespace + +// decode windows: items of 16 pairs, 8 n-tiles a block, 4 warps +extern "C" __global__ void __launch_bounds__(128) glm_exl3_grouped(const half* __restrict__ X0, const half* __restrict__ X1, + const uint32_t* __restrict__ T0, const uint32_t* __restrict__ T1, const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, float* __restrict__ Z, int K, int N, int P, int SK) { + launch_dependents(); // a dependent matmul may start on its weights + grouped<1, 8, 4>(X0, X1, T0, T1, items, counts, members, Z, K, N, P, SK); +} + +// Program (pair, K block, matrix): Xh[mat][pair][block] = fp16((x[row] * suh[mat][e]) @ H), pair = row * 8 + slot. +extern "C" __global__ void glm_exl3_rot_in(const __nv_bfloat16* __restrict__ x, int x_stride, const int* __restrict__ pick, + const half* __restrict__ suh0, const half* __restrict__ suh1, half* __restrict__ out0, half* __restrict__ out1, int K) { + launch_dependents(); // a dependent matmul may start on its weights + const int p = blockIdx.x, blk = blockIdx.y, mat = blockIdx.z; + const int row = p >> 3; + const int lane = threadIdx.x; + const int e = pick[p]; + const half* suh = (mat ? suh1 : suh0) + (size_t)e * K + blk * 128 + 4 * lane; + const __nv_bfloat16* xr = x + (size_t)row * x_stride + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = __bfloat162float(xr[j]) * __half2float(suh[j]); + fwht128(v, lane); + half* o = (mat ? out1 : out0) + (size_t)p * K + blk * 128 + 4 * lane; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); +} + +// Program (pair, width block): gate, up over splits in order, rotated, svh, limited SwiGLU, Xd = fp16(act suh_d H). +extern "C" __global__ void glm_exl3_gateup(const float* __restrict__ Z, const int* __restrict__ pick, const half* __restrict__ svh_g, + const half* __restrict__ svh_u, const half* __restrict__ suh_d, half* __restrict__ xd, int P, int N, int SK, float limit) { + launch_dependents(); // a dependent matmul may start on its weights + const int p = blockIdx.x, blk = blockIdx.y; + const int lane = threadIdx.x; + const int e = pick[p]; + 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 < SK; ++s) { + sg += Z[((size_t)(0 * SK + s) * P + p) * N + n + j]; + su += Z[((size_t)(1 * SK + s) * P + p) * 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 gg = clamp_hi(bf16r(gv[j] * HAD_SCALE * __half2float(svh_g[(size_t)e * N + n + j])), limit); + float uu = clamp_both(bf16r(uv[j] * HAD_SCALE * __half2float(svh_u[(size_t)e * N + n + j])), limit); + float act = bf16r(bf16r(gg / (1.f + expf(-gg))) * uu); + v[j] = act * __half2float(suh_d[(size_t)e * N + n + j]); + } + fwht128(v, lane); + half* o = xd + (size_t)p * N + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); +} + +// Program (pair, model-width block): down over splits in order, rotated, svh -> y[row][slot] fp32; this rank's share. +extern "C" __global__ void glm_exl3_down(const float* __restrict__ Z, const int* __restrict__ pick, const half* __restrict__ svh_d, + float* __restrict__ y, int P, int D, int SK, int ys) { + launch_dependents(); // a dependent matmul may start on its weights + const int p = blockIdx.x, blk = blockIdx.y; + const int lane = threadIdx.x; + const int e = pick[p]; + const int n = blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float s = 0.f; + for (int k = 0; k < SK; ++k) s += Z[((size_t)k * P + p) * D + n + j]; + v[j] = s; + } + fwht128(v, lane); + float* o = y + ((size_t)(p >> 3) * ys + (p & 7)) * D + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = v[j] * HAD_SCALE * __half2float(svh_d[(size_t)e * D + n + j]); +} + +// -- prompt chunks, the transposed design: Y^T = W^T X^T, one fp32 mma chain an output over ascending k tiles ------- + +// After MiaAI-Lab's prompt kernels (GLM-5.3-Flash-EXL3-2x-DGX-Sparks-TensorFold patches 0004, 0009, Apache-2.0). +#include "glm_exl3_transposed.cuh" + +// Prompt items of 64 pairs: block (128 expert columns, item), 8 warps: Xd = gate/up epilogue; xshift 3 or 0 as above. +extern "C" __global__ void __launch_bounds__(256, 2) glm_exl3_prompt_gu(const half* __restrict__ Xg, const half* __restrict__ Xu, + const uint32_t* __restrict__ Tg, const uint32_t* __restrict__ Tu, const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, const half* __restrict__ svh_g, const half* __restrict__ svh_u, const half* __restrict__ suh_d, + half* __restrict__ xd, int K, int N, int xshift, float limit) { + launch_dependents(); + __shared__ __align__(16) unsigned char smem_raw[tsmem_bytes()]; // one buffer for either instantiation + __shared__ int rows_sh[TPM]; + if (xshift) tgateup(Xg, Xu, Tg, Tu, items, counts, members, svh_g, svh_u, suh_d, xd, K, N, xshift, limit, smem_raw, rows_sh); + else tgateup<2 * TPM>(Xg, Xu, Tg, Tu, items, counts, members, svh_g, svh_u, suh_d, xd, K, N, 0, limit, smem_raw, rows_sh); +} + +// Block (256 model-width columns, item): y[row][slot][:] = the down epilogue of the item's pairs (this rank's share). +extern "C" __global__ void __launch_bounds__(256, 2) glm_exl3_prompt_dn(const half* __restrict__ Xd, const uint32_t* __restrict__ Td, + const int* __restrict__ items, const int* __restrict__ counts, const int* __restrict__ members, const half* __restrict__ svh_d, + float* __restrict__ y, int K, int D, int ys) { + launch_dependents(); + __shared__ __align__(16) unsigned char smem_raw[tsmem_bytes()]; + __shared__ int rows_sh[TPM]; + const int cb = blockIdx.x, item = blockIdx.y; + if (item >= counts[0]) return; + const int e = items[3 * item], first = items[3 * item + 1], cnt = items[3 * item + 2]; + if (threadIdx.x < TPM) rows_sh[threadIdx.x] = (int)threadIdx.x < cnt ? members[first + threadIdx.x] : -1; + __syncthreads(); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + float acc[TMT][TNB][4]; + tmain<1, TPM>(Xd, Xd, Td, Td, rows_sh, cnt, e, K, D >> 4, cb, 0, reinterpret_cast(smem_raw), acc); + float* Ed = reinterpret_cast(smem_raw); + for (int hf = 0; hf < 2 && hf * TPH < cnt; ++hf) { + if (hf) __syncthreads(); + tspill(Ed, DN_ES, warp * TMT * 16, acc, hf); + __syncthreads(); + for (int i = warp; i < TPH && hf * TPH + i < cnt; i += TNW) { + const int p = rows_sh[hf * TPH + i]; +#pragma unroll + for (int blk = 0; blk < 2; ++blk) { + const int n = cb * 256 + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = Ed[(size_t)i * DN_ES + blk * 128 + 4 * lane + j]; + fwht128(v, lane); + float* o = y + ((size_t)(p >> 3) * ys + (p & 7)) * D + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = v[j] * HAD_SCALE * __half2float(svh_d[(size_t)e * D + n + j]); + } + } + } +} + +// Program (row): the routed picks of a row of `slots` (the shared expert last) as a row of `slots - 1`, for the plan. +extern "C" __global__ void glm_exl3_picks(const int* __restrict__ pick, int* __restrict__ out, int rows, int slots) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= rows * (slots - 1)) return; + out[i] = pick[(i / (slots - 1)) * slots + i % (slots - 1)]; +} diff --git a/zig/kernels/cuda/glm_exl3_transposed.cuh b/zig/kernels/cuda/glm_exl3_transposed.cuh new file mode 100644 index 0000000000..3e7919e84a --- /dev/null +++ b/zig/kernels/cuda/glm_exl3_transposed.cuh @@ -0,0 +1,168 @@ +// Transposed prompt helpers after MiaAI-Lab (GLM-5.3-Flash-EXL3-2x-DGX-Sparks-TensorFold 0004, 0009, Apache-2.0). +#pragma once +// Included by exl3.cu after its shared helpers (mcg decode, cp.async, ldmatrix), in the same translation unit. + +namespace { + +constexpr int TPM = 64, TNW = 8, TMT = 16 / TNW, TNB = TPM / 8, TT = TNW * 32, TPH = TPM / 2; +constexpr int TBK = 32, TKT = TBK / 16, TSTAGES = 3, TXS = TBK + 8; // a stage: 32 k values a row; +16 B: ldmatrix rows on distinct banks +constexpr int TWORDS = 16 * TKT * 32; // a stage's trellis words (16 tiles x 2 k tiles x 32) +constexpr int GU_ES = 128 + 4, DN_ES = 256 + 4; // the epilogue sums' row strides (floats) +template +constexpr int tstage_halves() { return XR * TXS + 2 * TWORDS; } +constexpr int tsmem_bytes() { + const int gu = TSTAGES * tstage_halves<2 * TPM>() * 2; // gate and up rows apart (no shared rotation) + const int ep = 2 * TPH * GU_ES * 4, dn = TPH * DN_ES * 4; + return gu > ep ? (gu > dn ? gu : dn) : (ep > dn ? ep : dn); +} + +// Which (matrix, weight tile) warp-tile q (0..15) of column block cb is. +template +__device__ __forceinline__ void ttile(int q, int cb, int& mat, int& nt) { + mat = MATS == 2 ? q >> 3 : 0; + nt = MATS == 2 ? cb * 8 + (q & 7) : cb * 16 + q; +} + +// Main loop: acc[i][nb] = warp tile i x members nb * 8 .. + 7 over K; XR rows a stage (TPM if X0 == X1, else 2 TPM). +template +__device__ __forceinline__ void tmain(const half* __restrict__ X0, const half* __restrict__ X1, const uint32_t* __restrict__ T0, + const uint32_t* __restrict__ T1, const int* rows_sh, int cnt, int e, int K, int NTILES, int cb, int xshift, + half* stage, float (&acc)[TMT][TNB][4]) { + constexpr int SH = tstage_halves(); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + int mat, nt_unused; + ttile(warp * TMT, cb, mat, nt_unused); + const int KT = K >> 4, chunks = K / TBK; + const int nblocks = (cnt + 15) >> 4; // 16-member blocks in use +#pragma unroll + for (int i = 0; i < TMT; ++i) +#pragma unroll + for (int j = 0; j < TNB; ++j) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[i][j][c] = 0.f; + auto load = [&](int chunk) { + half* dst = stage + (size_t)(chunk % TSTAGES) * SH; + const int k0 = chunk * TBK; + for (int idx = threadIdx.x; idx < XR * (TBK / 8); idx += TT) { + const int m = idx / (TPM * (TBK / 8)), rem = idx % (TPM * (TBK / 8)), i = rem / (TBK / 8), j = rem % (TBK / 8); + if (i >= nblocks * 16) continue; // a 16-member block without a member: never read + const int r = rows_sh[i]; + const half* src = (m ? X1 : X0) + (size_t)(r < 0 ? 0 : (r >> xshift)) * K + k0 + j * 8; + cp_async16(dst + ((size_t)m * TPM + i) * TXS + j * 8, src, r >= 0); + } + uint32_t* wdst = reinterpret_cast(dst + (size_t)XR * TXS); + for (int idx = threadIdx.x; idx < 16 * TKT * 8; idx += TT) { // 16-byte pieces of the 128-byte tile words + const int q = idx / (TKT * 8), kk = (idx / 8) % TKT, part = idx % 8; + int wm, nt; + ttile(q, cb, wm, nt); + const uint32_t* src = (wm ? T1 : T0) + (((size_t)e * KT + chunk * TKT + kk) * NTILES + nt) * 32 + part * 4; + cp_async16(wdst + ((size_t)q * TKT + kk) * 32 + part * 4, src, true); + } + }; +#pragma unroll + for (int c = 0; c < TSTAGES - 1; ++c) { + if (c < chunks) load(c); + cp_async_commit(); + } + for (int c = 0; c < chunks; ++c) { + cp_async_wait(); + __syncthreads(); + if (c + TSTAGES - 1 < chunks) load(c + TSTAGES - 1); + cp_async_commit(); + const half* base = stage + (size_t)(c % TSTAGES) * SH; + const half* xs = base + (size_t)(XR == TPM ? 0 : mat) * TPM * TXS; + const uint32_t* ws = reinterpret_cast(base + (size_t)XR * TXS); +#pragma unroll + for (int kk = 0; kk < TKT; ++kk) { + uint32_t a[TMT][4]; +#pragma unroll + for (int i = 0; i < TMT; ++i) { + uint32_t b0[2], b1[2]; + decode_tile_p(ws[((size_t)(warp * TMT + i) * TKT + kk) * 32 + lane], lane, b0, b1); + a[i][0] = b0[0]; a[i][1] = b1[0]; a[i][2] = b0[1]; a[i][3] = b1[1]; + } +#pragma unroll + for (int jb = 0; jb < TPM / 16; ++jb) { + if (jb < nblocks) { + uint32_t r[4]; + const int row = jb * 16 + (lane & 7) + ((lane >> 4) & 1) * 8; + ldmatrix_x4(r, xs + (size_t)row * TXS + kk * 16 + ((lane >> 3) & 1) * 8); + const uint32_t bl[2] = {r[0], r[1]}, bh[2] = {r[2], r[3]}; +#pragma unroll + for (int i = 0; i < TMT; ++i) { + mma16816(acc[i][2 * jb], a[i], bl); + mma16816(acc[i][2 * jb + 1], a[i], bh); + } + } + } + } + } + cp_async_wait<0>(); + __syncthreads(); // the stages become the epilogue's sums +} + +// acc (members TPH * hf .. + TPH - 1) -> E[member - TPH * hf][col] fp32, row stride ES; col0: first tile's column. +__device__ __forceinline__ void tspill(float* E, int ES, int col0, const float (&acc)[TMT][TNB][4], int hf) { + const int lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; +#pragma unroll + for (int i = 0; i < TMT; ++i) +#pragma unroll + for (int q = 0; q < TPM / 16; ++q) { + const int nb = hf * (TPM / 16) + q; + const int tok = q * 8 + 2 * t, col = col0 + i * 16 + g; + E[(size_t)tok * ES + col] = acc[i][nb][0]; + E[(size_t)(tok + 1) * ES + col] = acc[i][nb][1]; + E[(size_t)tok * ES + col + 8] = acc[i][nb][2]; + E[(size_t)(tok + 1) * ES + col + 8] = acc[i][nb][3]; + } +} + +template +__device__ __forceinline__ void tgateup(const half* __restrict__ Xg, const half* __restrict__ Xu, const uint32_t* __restrict__ Tg, + const uint32_t* __restrict__ Tu, const int* __restrict__ items, const int* __restrict__ counts, const int* __restrict__ members, + const half* __restrict__ svh_g, const half* __restrict__ svh_u, const half* __restrict__ suh_d, half* __restrict__ xd, + int K, int N, int xshift, float limit, unsigned char* smem_raw, int* rows_sh) { + const int cb = blockIdx.x, item = blockIdx.y; + if (item >= counts[0]) return; + const int e = items[3 * item], first = items[3 * item + 1], cnt = items[3 * item + 2]; + if (threadIdx.x < TPM) rows_sh[threadIdx.x] = (int)threadIdx.x < cnt ? members[first + threadIdx.x] : -1; + __syncthreads(); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + float acc[TMT][TNB][4]; + tmain<2, XR>(Xg, Xu, Tg, Tu, rows_sh, cnt, e, K, N >> 4, cb, xshift, reinterpret_cast(smem_raw), acc); + float* Eg = reinterpret_cast(smem_raw); + int mat, nt; + ttile<2>(warp * TMT, cb, mat, nt); + const int col0 = (nt - cb * 8) * 16; + const int n = cb * 128 + 4 * lane; + for (int hf = 0; hf < 2 && hf * TPH < cnt; ++hf) { + if (hf) __syncthreads(); // the first half's rows are read + tspill(Eg + (size_t)mat * TPH * GU_ES, GU_ES, col0, acc, hf); + __syncthreads(); + for (int i = warp; i < TPH && hf * TPH + i < cnt; i += TNW) { + const int p = rows_sh[hf * TPH + i]; + float gv[4], uv[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + gv[j] = Eg[(size_t)i * GU_ES + 4 * lane + j]; + uv[j] = Eg[(size_t)(TPH + i) * GU_ES + 4 * lane + j]; + } + fwht128(gv, lane); + fwht128(uv, lane); + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float gg = clamp_hi(bf16r(gv[j] * HAD_SCALE * __half2float(svh_g[(size_t)e * N + n + j])), limit); + float uu = clamp_both(bf16r(uv[j] * HAD_SCALE * __half2float(svh_u[(size_t)e * N + n + j])), limit); + float act = bf16r(bf16r(gg / (1.f + expf(-gg))) * uu); + v[j] = act * __half2float(suh_d[(size_t)e * N + n + j]); + } + fwht128(v, lane); + half* o = xd + (size_t)p * N + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); + } + } +} + +} // namespace diff --git a/zig/kernels/cuda/glm_mm.cu b/zig/kernels/cuda/glm_mm.cu new file mode 100644 index 0000000000..2fc8238929 --- /dev/null +++ b/zig/kernels/cuda/glm_mm.cu @@ -0,0 +1,295 @@ +// GLM dense linears, ours: one tensor-core kernel for any row count, 64-groups then K slices in order; row-stable bits. + +#include +#include +#include + +#include "glm_qmm_frag.cuh" + +namespace { + +using namespace qmm_frag; + +// 16 codes from a packed chunk (ops.cu's layout), stored + 2^(BITS-1) +template +__device__ __forceinline__ void codes16(const uint32_t* __restrict__ p, int (&q)[16]) { + if constexpr (BITS == 4) { +#pragma unroll + for (int i = 0; i < 16; ++i) q[i] = (int)((p[i / 8] >> (4 * (i % 8))) & 15u) - 8; + } else if constexpr (BITS == 8) { +#pragma unroll + for (int i = 0; i < 16; ++i) q[i] = (int)((p[i / 4] >> (8 * (i % 4))) & 255u) - 128; + } else { +#pragma unroll + for (int i = 0; i < 15; ++i) q[i] = (int)((p[i / 5] >> (6 * (i % 5))) & 63u) - 32; + q[15] = (int)((p[0] >> 30) | ((p[1] >> 30) << 2) | ((p[2] >> 30) << 4)) - 32; + } +} + +template +__device__ __forceinline__ void pack16(const int (&q)[16], uint32_t* __restrict__ out) { + constexpr int half = 1 << (BITS - 1); + uint32_t w[BITS / 2]; +#pragma unroll + for (int i = 0; i < BITS / 2; ++i) w[i] = 0u; + if constexpr (BITS == 4) { +#pragma unroll + for (int i = 0; i < 16; ++i) w[i / 8] |= (uint32_t)(q[i] + half) << (4 * (i % 8)); + } else if constexpr (BITS == 8) { +#pragma unroll + for (int i = 0; i < 16; ++i) w[i / 4] |= (uint32_t)(q[i] + half) << (8 * (i % 4)); + } else { +#pragma unroll + for (int i = 0; i < 15; ++i) w[i / 5] |= (uint32_t)(q[i] + half) << (6 * (i % 5)); + const uint32_t last = (uint32_t)(q[15] + half); +#pragma unroll + for (int j = 0; j < 3; ++j) w[j] |= ((last >> (2 * j)) & 3u) << 30; + } +#pragma unroll + for (int i = 0; i < BITS / 2; ++i) out[i] = w[i]; +} + +// input of value v (0..15) of lane ``lane``'s chunk within its group of 64 +__device__ __forceinline__ int frag_k(int lane, int v) { + const int t = lane & 3, kt = v >> 2, j = v & 3; + return kt * 16 + 2 * t + (j & 1) + (j >> 1) * 8; +} + +// the B-fragment bf16 pairs of a chunk: pairs 2 kt and 2 kt + 1 are step kt's b0 and b1 +template +__device__ __forceinline__ void bpairs(const uint32_t* __restrict__ p, uint32_t (&b)[8]) { + if constexpr (BITS == 16) { +#pragma unroll + for (int i = 0; i < 8; ++i) b[i] = p[i]; + } else { + int q[16]; + codes16(p, q); +#pragma unroll + for (int i = 0; i < 8; ++i) { + const __nv_bfloat162 h = __floats2bfloat162_rn((float)q[2 * i], (float)q[2 * i + 1]); // exact integers + b[i] = *reinterpret_cast(&h); + } + } +} + +template +struct Tile { + static constexpr int WN = 4, BN = 64, NT = 2, MT = BM / WM / 16, THREADS = WM * WN * 32; + static constexpr int WPC = BITS / 2; + static constexpr int ROW = 128, CHUNKS = 8; // a row's 64 inputs: 128 bytes, 8 chunks of 16 + static constexpr int X = BM * ROW; + static constexpr int W = 8 * 32 * WPC * 4; // 8 column tiles' chunks + static constexpr int S = 8 * 8 * 2; // their scales (fp16) + static constexpr int STAGE = X + W + S; + static constexpr int SMEM = STAGES * STAGE; +}; + +// ROWS_FIRST: row blocks fastest, so the blocks reading one weight tile run together and share it in L2 (prompts) +template +__device__ __forceinline__ void mm_kernel(const __nv_bfloat16* __restrict__ x, long long ldx, + const uint32_t* __restrict__ w, const __half* __restrict__ s, void* __restrict__ out, float* __restrict__ part, + int f32, int M, int N, int K, int SK) { + using T = Tile; + extern __shared__ __align__(128) unsigned char buf[]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int wm = warp / T::WN, wn = warp % T::WN; + const int KG = K >> 6, per = KG / SK, N8 = N >> 3; + const int n0 = (ROWS_FIRST ? blockIdx.y : blockIdx.x) * T::BN, m0 = (ROWS_FIRST ? blockIdx.x : blockIdx.y) * BM; + const int slice = blockIdx.z, g0 = slice * per; + const int t0 = n0 >> 3; // the block's first column tile + auto stage = [&](int st) { return buf + st * T::STAGE; }; + auto loadX = [&](int st, int g) { + unsigned char* p = stage(st); + for (int c = tid; c < BM * T::CHUNKS; c += T::THREADS) { + const int r = c / T::CHUNKS, ch = c % T::CHUNKS; + const int row = min(m0 + r, M - 1); + cp16z(p + r * T::ROW + swz(r, ch) * 16, x + (long long)row * ldx + g * 64 + ch * 8, m0 + r < M); + } + }; + auto loadW = [&](int st, int g) { + unsigned char* pw = stage(st) + T::X; + constexpr int PER_TILE = 32 * T::WPC * 4 / 16; // 16-byte pieces of one column tile's chunks + for (int c = tid; c < 8 * PER_TILE; c += T::THREADS) { + const int tt = c / PER_TILE, off = c % PER_TILE; + if (t0 + tt >= N8) continue; + const long long src = ((long long)(t0 + tt) * KG + g) * 32 * T::WPC * 4 + off * 16; + cp16(pw + c * 16, reinterpret_cast(w) + src); + } + if constexpr (BITS != 16) { + unsigned char* ps = pw + T::W; + for (int tt = tid; tt < 8; tt += T::THREADS) { + if (t0 + tt >= N8) continue; + cp16(ps + tt * 16, s + ((long long)(t0 + tt) * KG + g) * 8); + } + } + }; + float acc[T::MT][T::NT][4]; +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) acc[i][j][e] = 0.0f; + // dependent launch: first stages' weights stream in early; inputs and outputs only after griddepcontrol.wait + grid_launch(); +#pragma unroll + for (int st = 0; st < STAGES - 1; ++st) + if (st < per) loadW(st, g0 + st); + grid_wait(); +#pragma unroll + for (int st = 0; st < STAGES - 1; ++st) { + if (st < per) loadX(st, g0 + st); + commit(); + } + for (int it = 0; it < per; ++it) { + wait(); + __syncthreads(); + const int next = it + STAGES - 1; + if (next < per) { + loadX(next % STAGES, g0 + next); + loadW(next % STAGES, g0 + next); + } + commit(); + const unsigned char* p = stage(it % STAGES); + const uint32_t* pw = reinterpret_cast(p + T::X); + const __half* ps = reinterpret_cast(p + T::X + T::W); + uint32_t b[T::NT][8]; +#pragma unroll + for (int j = 0; j < T::NT; ++j) bpairs(pw + ((wn * T::NT + j) * 32 + lane) * T::WPC, b[j]); + float d[T::MT][T::NT][4]; +#pragma unroll + for (int kt = 0; kt < 4; ++kt) { + uint32_t a[T::MT][4]; +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + const int r = wm * (BM / WM) + i * 16 + (lane & 7) + ((lane >> 3) & 1) * 8; + const int ch = kt * 2 + (lane >> 4); + ldmatrix4(a[i], p + r * T::ROW + swz(r, ch) * 16); + } +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + if (kt == 0) mma0(d[i][j], a[i], b[j][0], b[j][1]); + else mma(d[i][j], a[i], b[j][2 * kt], b[j][2 * kt + 1]); + } + } +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + float s0 = 1.0f, s1 = 1.0f; + if constexpr (BITS != 16) { + const __half2 h = *reinterpret_cast(ps + (wn * T::NT + j) * 8 + (lane & 3) * 2); + s0 = __low2float(h); + s1 = __high2float(h); + } +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int e = 0; e < 4; ++e) acc[i][j][e] = __fmaf_rn(d[i][j][e], (e & 1) ? s1 : s0, acc[i][j][e]); + } + } + wait<0>(); +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const int col = n0 + (wn * T::NT + j) * 8 + (lane & 3) * 2; + if (col >= N) continue; +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int row = m0 + wm * (BM / WM) + i * 16 + (lane >> 2) + h * 8; + if (row >= M) continue; + const float v0 = acc[i][j][2 * h], v1 = acc[i][j][2 * h + 1]; + if (SK > 1) { + float* dst = part + ((long long)slice * M + row) * N + col; + dst[0] = v0; + dst[1] = v1; + } else if (f32) { + float* dst = reinterpret_cast(out) + (long long)row * N + col; + dst[0] = v0; + dst[1] = v1; + } else { + *reinterpret_cast<__nv_bfloat162*>(reinterpret_cast<__nv_bfloat16*>(out) + (long long)row * N + col) = + __floats2bfloat162_rn(v0, v1); + } + } + } +} + +} // namespace + +// The K slices' partials added in slice order, then rounded once. +extern "C" __global__ void glm_mm_reduce(const float* __restrict__ part, void* __restrict__ out, long long total, int SK, int f32) { + qmm_frag::grid_launch(); + const long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total) return; + float acc = part[i]; + for (int sl = 1; sl < SK; ++sl) acc = acc + part[sl * total + i]; + if (f32) reinterpret_cast(out)[i] = acc; + else reinterpret_cast<__nv_bfloat16*>(out)[i] = __float2bfloat16_rn(acc); +} + +// one instance a (bits, row bucket): (16, 1 warp row), (32, 1), (64, 2), (128, 2); three stages +#define GLM_MM_ALL(B) \ + extern "C" __global__ void __launch_bounds__(128) glm_mm##B##_16(const __nv_bfloat16* __restrict__ x, long long ldx, \ + const uint32_t* __restrict__ w, const __half* __restrict__ s, void* __restrict__ out, float* __restrict__ part, \ + int f32, int M, int N, int K, int SK) { mm_kernel(x, ldx, w, s, out, part, f32, M, N, K, SK); } \ + extern "C" __global__ void __launch_bounds__(128) glm_mm##B##_32(const __nv_bfloat16* __restrict__ x, long long ldx, \ + const uint32_t* __restrict__ w, const __half* __restrict__ s, void* __restrict__ out, float* __restrict__ part, \ + int f32, int M, int N, int K, int SK) { mm_kernel(x, ldx, w, s, out, part, f32, M, N, K, SK); } \ + extern "C" __global__ void __launch_bounds__(256) glm_mm##B##_64(const __nv_bfloat16* __restrict__ x, long long ldx, \ + const uint32_t* __restrict__ w, const __half* __restrict__ s, void* __restrict__ out, float* __restrict__ part, \ + int f32, int M, int N, int K, int SK) { mm_kernel(x, ldx, w, s, out, part, f32, M, N, K, SK); } \ + extern "C" __global__ void __launch_bounds__(256) glm_mm##B##_128(const __nv_bfloat16* __restrict__ x, long long ldx, \ + const uint32_t* __restrict__ w, const __half* __restrict__ s, void* __restrict__ out, float* __restrict__ part, \ + int f32, int M, int N, int K, int SK) { mm_kernel(x, ldx, w, s, out, part, f32, M, N, K, SK); } +// prompts: 128-row blocks, row blocks fastest (one slice: a prompt's rows fill the GPU) +extern "C" __global__ void __launch_bounds__(256) glm_mm16_rows(const __nv_bfloat16* __restrict__ x, long long ldx, + const uint32_t* __restrict__ w, const __half* __restrict__ s, void* __restrict__ out, float* __restrict__ part, + int f32, int M, int N, int K, int SK) { mm_kernel<16, 128, 2, 3, true>(x, ldx, w, s, out, part, f32, M, N, K, SK); } +GLM_MM_ALL(4) +GLM_MM_ALL(6) +GLM_MM_ALL(8) +GLM_MM_ALL(16) + +// -- tiling at load: ops.cu's packed rows and fp16 scales [N][K / 64] -> the tiled chunks and scales --------------- +template +__device__ __forceinline__ void tile_packed(const uint32_t* __restrict__ src, const __half* __restrict__ ss, + uint32_t* __restrict__ dst, __half* __restrict__ ds, int N, int K) { + constexpr int WPC = BITS / 2; + const int KG = K >> 6, N8 = N >> 3; + const long long c = (long long)blockIdx.x * blockDim.x + threadIdx.x; // a destination chunk: (tile, group, lane) + if (c >= (long long)N8 * KG * 32) return; + const int lane = (int)(c & 31); + const long long tg = c >> 5; + const int g = (int)(tg % KG), tt = (int)(tg / KG); + const int n = tt * 8 + (lane >> 2); + const uint32_t* row = src + (long long)n * (K >> 4) * WPC; + int q[16]; +#pragma unroll + for (int v = 0; v < 16; ++v) { + const int k = g * 64 + frag_k(lane, v); + int sq[16]; + codes16(row + (long long)(k >> 4) * WPC, sq); + q[v] = sq[k & 15]; + } + pack16(q, dst + c * WPC); + if (lane < 8) ds[(tg * 8) + lane] = ss[(long long)(tt * 8 + lane) * KG + g]; +} + +extern "C" __global__ void glm_mm_tile4(const uint32_t* src, const __half* ss, uint32_t* dst, __half* ds, int N, int K) { tile_packed<4>(src, ss, dst, ds, N, K); } +extern "C" __global__ void glm_mm_tile6(const uint32_t* src, const __half* ss, uint32_t* dst, __half* ds, int N, int K) { tile_packed<6>(src, ss, dst, ds, N, K); } +extern "C" __global__ void glm_mm_tile8(const uint32_t* src, const __half* ss, uint32_t* dst, __half* ds, int N, int K) { tile_packed<8>(src, ss, dst, ds, N, K); } + +// bf16 rows [N][K] -> the tiled chunks (16 bf16 a chunk, in pair order) +extern "C" __global__ void glm_mm_tile16(const __nv_bfloat16* __restrict__ src, __nv_bfloat16* __restrict__ dst, int N, int K) { + const int KG = K >> 6, N8 = N >> 3; + const long long c = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (c >= (long long)N8 * KG * 32) return; + const int lane = (int)(c & 31); + const long long tg = c >> 5; + const int g = (int)(tg % KG), tt = (int)(tg / KG); + const int n = tt * 8 + (lane >> 2); +#pragma unroll + for (int v = 0; v < 16; ++v) dst[c * 16 + v] = src[(long long)n * K + g * 64 + frag_k(lane, v)]; +} diff --git a/zig/kernels/cuda/glm_qmm_frag.cuh b/zig/kernels/cuda/glm_qmm_frag.cuh new file mode 100644 index 0000000000..77a424cffb --- /dev/null +++ b/zig/kernels/cuda/glm_qmm_frag.cuh @@ -0,0 +1,117 @@ +// Device code of cuda/kernels/qmm_frag.cuh (flashnext/gen_copies.py): cp.async, ldmatrix, bf16 mma, nibble decoding. +#pragma once + +#include +#include +#include + +namespace qmm_frag { + +// Lane matmul tile shapes: BM rows by BN columns a block, WM x WN warps, each warp (BM / WM) x (BN / WN). +template +struct LaneTile { + static constexpr int THREADS = WM * WN * 32; + static constexpr int MT = BM / WM / 16; // m16 tiles a warp + static constexpr int NT = BN / WN / 8; // n8 tiles a warp + static constexpr int ROW = GS * 2; // bytes of one input row a group + static constexpr int CHUNKS = ROW / 16; + static constexpr int X = BM * ROW; // stage bytes: inputs, + static constexpr int W = BN * GS / 2; // weights, + static constexpr int S = BN * 2; // scales, biases (bf16), + static constexpr int XS = BM * 4; // and input sums (fp32) + static constexpr int STAGE = X + W + 2 * S + XS; + static constexpr int PARTIALS = MT * NT * 4 * THREADS * 4; // a K slice's partial, parked for the cluster sum + static constexpr int SMEM = STAGES * STAGE > PARTIALS ? STAGES * STAGE : PARTIALS; +}; + +__device__ __forceinline__ uint32_t smem(const void* p) { + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ void cp16(void* dst, const void* src) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(smem(dst)), "l"(src)); +} + +// 16 bytes, or zeros without reading when ``read`` is false (a row past M: computed, never stored) +__device__ __forceinline__ void cp16z(void* dst, const void* src, bool read) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(smem(dst)), "l"(src), "r"(read ? 16 : 0)); +} + +__device__ __forceinline__ void cp8(void* dst, const void* src) { + asm volatile("cp.async.ca.shared.global [%0], [%1], 8;\n" ::"r"(smem(dst)), "l"(src)); +} + +__device__ __forceinline__ void cp4(void* dst, const void* src) { + asm volatile("cp.async.ca.shared.global [%0], [%1], 4;\n" ::"r"(smem(dst)), "l"(src)); +} + +__device__ __forceinline__ void commit() { asm volatile("cp.async.commit_group;\n" ::); } + +template +__device__ __forceinline__ void wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); } + +__device__ __forceinline__ void ldmatrix4(uint32_t (&r)[4], const void* p) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(smem(p))); +} + +__device__ __forceinline__ void ldmatrix2(uint32_t (&r)[2], const void* p) { + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" : "=r"(r[0]), "=r"(r[1]) : "r"(smem(p))); +} + +__device__ __forceinline__ void mma(float (&d)[4], const uint32_t (&a)[4], uint32_t b0, uint32_t b1) { + asm( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} + +// The same mma from a zero accumulator (a group's first k16 step): no registers to clear. +__device__ __forceinline__ void mma0(float (&d)[4], const uint32_t (&a)[4], uint32_t b0, uint32_t b1) { + const float z = 0.0f; + asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%10, %10, %10, %10};\n" + : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1), "f"(z)); +} + +// Nibbles at bits [s, s + 4) and [16 + s, 20 + s) as a bf16 pair: exponent bits make 128 + q, then minus 128. +__device__ __forceinline__ uint32_t pair(uint32_t w, int s) { + const uint32_t t = ((w >> s) & 0x000F000Fu) | 0x43004300u; + uint32_t r; +#if __CUDA_ARCH__ >= 900 + asm("sub.rn.bf16x2 %0, %1, %2;\n" : "=r"(r) : "r"(t), "r"(0x43004300u)); +#else + // no sub.bf16x2 before sm_90: t * 1 - 128 rounds the exact q once too (an exact value, the same bits) + asm("fma.rn.bf16x2 %0, %1, %2, %3;\n" : "=r"(r) : "r"(t), "r"(0x3F803F80u), "r"(0xC300C300u)); +#endif + return r; +} + +// Programmatic dependent launch (sm_90+, else no-ops): wait for the previous kernel, release the next one's blocks. +__device__ __forceinline__ void grid_wait() { +#if __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.wait;\n" ::: "memory"); +#endif +} + +__device__ __forceinline__ void grid_launch() { +#if __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.launch_dependents;\n" ::: "memory"); +#endif +} + +// Block b's (first row, first column), row tiles fastest in bands of ``group`` so blocks in flight share L2. +__device__ __forceinline__ int2 tile_of(int b, int M, int N, int BM, int BN, int group) { + const int rows_t = (M + BM - 1) / BM, cols_t = (N + BN - 1) / BN, band = group * cols_t; + const int first = b / band * group, in_band = b % band, height = min(group, rows_t - first); + return make_int2((first + in_band % height) * BM, in_band / height * BN); +} + +// A 16-byte chunk c of row r sits at chunk c ^ (r % CHUNKS) so ldmatrix's eight rows hit eight bank groups. +template +__device__ __forceinline__ int swz(int r, int c) { return c ^ (r % CHUNKS); } + +} // namespace qmm_frag diff --git a/zig/kernels/cuda/glm_vision.cu b/zig/kernels/cuda/glm_vision.cu new file mode 100644 index 0000000000..d144113051 --- /dev/null +++ b/zig/kernels/cuda/glm_vision.cu @@ -0,0 +1,181 @@ +// GLM vision tower's steps around cuBLASLt, as mlx-vlm glm5_next/vision.py: fp32 residual, bf16 outputs rounded once. + +#include +#include + +namespace { +__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } +__device__ __forceinline__ __nv_bfloat16 tobf(float x) { return __float2bfloat16_rn(x); } + +// a block's sum (blockDim.x a multiple of 32, at most 1024); every thread gets it +__device__ float block_sum(float v) { + __shared__ float part[32]; + __shared__ float total; + for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o); + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + __syncthreads(); + if (lane == 0) part[warp] = v; + __syncthreads(); + if (warp == 0) { + float s = lane < (int)(blockDim.x >> 5) ? part[lane] : 0.f; + for (int o = 16; o > 0; o >>= 1) s += __shfl_xor_sync(0xffffffffu, s, o); + if (lane == 0) total = s; + } + __syncthreads(); + return total; +} + +__device__ __forceinline__ float limited_swiglu(float g, float u, float limit) { + g = g > limit ? limit : g; // as torch.clamp: a NaN stays NaN (fminf and fmaxf would drop it for the bound) + u = u < -limit ? -limit : (u > limit ? limit : u); + return g / (1.f + expf(-g)) * u; +} +} // namespace + +// dst[i] = bf16(src[i]) (the processor's fp32 patches) +extern "C" __global__ void glm_vis_to_bf16(const float* __restrict__ src, __nv_bfloat16* __restrict__ dst, long long n) { + for (long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) + dst[i] = tobf(src[i]); +} + +// out[r, c] = bf16(acc[r, c] + bias[c]) (bias may be null); with x32: x32[r, c] += that bf16 value (a residual add) +extern "C" __global__ void glm_vis_bias(const float* __restrict__ acc, const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ out, float* __restrict__ x32, long long n, int cols) { + for (long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) { + const float v = acc[i] + (bias ? bf(bias[i % cols]) : 0.f); + const __nv_bfloat16 h = tobf(v); + if (out) out[i] = h; + if (x32) x32[i] += bf(h); + } +} + +// x32[i] = acc[i] + bias[i % cols] (the patch embedding: the residual stream's first value, in fp32) +extern "C" __global__ void glm_vis_embed32(const float* __restrict__ acc, const __nv_bfloat16* __restrict__ bias, + float* __restrict__ x32, long long n, int cols) { + for (long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) + x32[i] = bf(tobf(acc[i] + bf(bias[i % cols]))); +} + +// RMSNorm of fp32 rows: y = bf16(x rsqrt(mean(x^2) + eps) w), a block a row +extern "C" __global__ void glm_vis_rmsnorm(const float* __restrict__ x, const __nv_bfloat16* __restrict__ w, + __nv_bfloat16* __restrict__ y, int cols, float eps) { + const float* row = x + (long long)blockIdx.x * cols; + float s = 0.f; + for (int c = threadIdx.x; c < cols; c += blockDim.x) s += row[c] * row[c]; + const float inv = rsqrtf(block_sum(s) / cols + eps); + for (int c = threadIdx.x; c < cols; c += blockDim.x) y[(long long)blockIdx.x * cols + c] = tobf(row[c] * inv * bf(w[c])); +} + +// attention inputs: q, k head-RMSNormed and rotated (fp32, 2D angles) into qh, kh; v into vt; a warp a (row, head). +extern "C" __global__ void glm_vis_qkv(const __nv_bfloat16* __restrict__ qkv, const __nv_bfloat16* __restrict__ qn, + const __nv_bfloat16* __restrict__ kn, const int* __restrict__ pos, + const float* __restrict__ inv_freq, __nv_bfloat16* __restrict__ qh, + __nv_bfloat16* __restrict__ kh, __nv_bfloat16* __restrict__ vt, int N, int seg, + int H, float eps) { + const long long item = blockIdx.x * (long long)(blockDim.x >> 5) + (threadIdx.x >> 5); + if (item >= (long long)N * H) return; + const int lane = threadIdx.x & 31; + const int n = (int)(item / H), h = (int)(item % H); + const int g = n / seg, j = n % seg; + const long long blk = (long long)g * H + h; + const __nv_bfloat16* base = qkv + (long long)n * 3 * H * 64 + h * 64; + const int p = lane < 16 ? pos[2 * n] : pos[2 * n + 1]; + const float ang = (float)p * inv_freq[lane & 15]; + const float c = cosf(ang), s = sinf(ang); + for (int which = 0; which < 2; ++which) { + const __nv_bfloat16* src = base + which * H * 64; + const __nv_bfloat16* wn = which ? kn : qn; + float a = bf(src[lane]), b = bf(src[lane + 32]); + float ss = a * a + b * b; + for (int o = 16; o > 0; o >>= 1) ss += __shfl_xor_sync(0xffffffffu, ss, o); + const float inv = rsqrtf(ss / 64.f + eps); + a = bf(tobf(a * inv * bf(wn[lane]))); + b = bf(tobf(b * inv * bf(wn[lane + 32]))); + // rotate_half: dims l < 32 take -x[l + 32], dims l + 32 take x[l] + const float ra = a * c - b * s, rb = b * c + a * s; + __nv_bfloat16* dst = (which ? kh : qh) + (blk * seg + j) * 64; + dst[lane] = tobf(ra); + dst[lane + 32] = tobf(rb); + } + const __nv_bfloat16* v = base + 2 * H * 64; + vt[(blk * 64 + lane) * seg + j] = v[lane]; + vt[(blk * 64 + lane + 32) * seg + j] = v[lane + 32]; +} + +// A scores row's softmax: pr = bf16(exp(scale (s - max))) and its sum (of the bf16 values), a block a row +extern "C" __global__ void glm_vis_softmax(const float* __restrict__ sc, __nv_bfloat16* __restrict__ pr, + float* __restrict__ sums, int n, float scale) { + const float* row = sc + (long long)blockIdx.x * n; + __shared__ float mx_s[32]; + __shared__ float mx; + float m = -INFINITY; + for (int c = threadIdx.x; c < n; c += blockDim.x) m = fmaxf(m, row[c]); + for (int o = 16; o > 0; o >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffffu, m, o)); + if ((threadIdx.x & 31) == 0) mx_s[threadIdx.x >> 5] = m; + __syncthreads(); + if (threadIdx.x < 32) { + float v = threadIdx.x < (blockDim.x >> 5) ? mx_s[threadIdx.x] : -INFINITY; + for (int o = 16; o > 0; o >>= 1) v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, o)); + if (threadIdx.x == 0) mx = v; + } + __syncthreads(); + float s = 0.f; + for (int c = threadIdx.x; c < n; c += blockDim.x) { + const __nv_bfloat16 e = tobf(expf((row[c] - mx) * scale)); + pr[(long long)blockIdx.x * n + c] = e; + s += bf(e); + } + s = block_sum(s); + if (threadIdx.x == 0) sums[blockIdx.x] = s; +} + +// att[r, h * 64 + d] = bf16(oh[r, d] / sums[r]) (a head's values product, normalized, into its columns) +extern "C" __global__ void glm_vis_head_out(const float* __restrict__ oh, const float* __restrict__ sums, + __nv_bfloat16* __restrict__ att, int rows, int H, int h) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= rows * 64) return; + const int r = i / 64, d = i % 64; + att[(long long)r * H * 64 + h * 64 + d] = tobf(oh[i] / sums[r]); +} + +// out = bf16(limited_swiglu(bf16(g + gb), bf16(u + ub))) from the two fp32 products (biases may be null) +extern "C" __global__ void glm_vis_swiglu(const float* __restrict__ g, const float* __restrict__ u, + const __nv_bfloat16* __restrict__ gb, const __nv_bfloat16* __restrict__ ub, + __nv_bfloat16* __restrict__ out, long long n, int cols, float limit) { + for (long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) { + const int c = (int)(i % cols); + const float gv = bf(tobf(g[i] + (gb ? bf(gb[c]) : 0.f))); + const float uv = bf(tobf(u[i] + (ub ? bf(ub[c]) : 0.f))); + out[i] = tobf(limited_swiglu(gv, uv, limit)); + } +} + +// The merger's norm: y = bf16(gelu_erf(LayerNorm(bf16(acc)) w + b)), a block a row of ``cols`` +extern "C" __global__ void glm_vis_ln_gelu(const float* __restrict__ acc, const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ b, __nv_bfloat16* __restrict__ y, int cols, + float eps) { + const float* row = acc + (long long)blockIdx.x * cols; + float s = 0.f; + for (int c = threadIdx.x; c < cols; c += blockDim.x) s += bf(tobf(row[c])); + const float mean = block_sum(s) / cols; + float v = 0.f; + for (int c = threadIdx.x; c < cols; c += blockDim.x) { + const float d = bf(tobf(row[c])) - mean; + v += d * d; + } + const float inv = rsqrtf(block_sum(v) / cols + eps); + for (int c = threadIdx.x; c < cols; c += blockDim.x) { + const float x = bf(tobf((bf(tobf(row[c])) - mean) * inv * bf(w[c]) + bf(b[c]))); + y[(long long)blockIdx.x * cols + c] = tobf(0.5f * x * (1.f + erff(x * 0.70710678118654752f))); + } +} + +// Placeholder rows take the tower's rows in every stream: row r takes feature row map[r] (map[r] < 0: kept). +extern "C" __global__ void glm_vis_inject(const __nv_bfloat16* __restrict__ feats, __nv_bfloat16* __restrict__ out, + const int* __restrict__ map, int rows, int d, int copies) { + const int r = blockIdx.x; + if (r >= rows || map[r] < 0) return; + const __nv_bfloat16* src = feats + (long long)map[r] * d; + for (int c = threadIdx.x; c < d; c += blockDim.x) + for (int k = 0; k < copies; ++k) out[((long long)r * copies + k) * d + c] = src[c]; +} diff --git a/zig/kernels/cuda/roce_gather.cu b/zig/kernels/cuda/roce_gather.cu new file mode 100644 index 0000000000..832afbcf9b --- /dev/null +++ b/zig/kernels/cuda/roce_gather.cu @@ -0,0 +1,100 @@ +// Two-rank one-shot RoCE all-gather after b12x (github.com/local-inference-lab/b12x, Apache-2.0): docs/roce-gather.md. + +#include + +namespace { + +constexpr int C_SEQ = 0, C_NBYTES = 1, C_ERROR = 3, C_HCA = 4; // ctrl words (C_NBYTES + slot) +constexpr int FLAG_STRIDE = 128; +// counters (device memory, u32): epoch, poison, then stage and tail arrivals for each grid size class +constexpr int K_EPOCH = 0, K_POISON = 1, K_CLASSES = 4; + +__device__ __forceinline__ uint32_t ld_relaxed_gpu(const uint32_t* p) { + uint32_t v; + asm volatile("ld.relaxed.gpu.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} +__device__ __forceinline__ uint32_t ld_relaxed_sys(const uint32_t* p) { + uint32_t v; + asm volatile("ld.relaxed.sys.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} +__device__ __forceinline__ uint32_t ld_acquire_sys(const uint32_t* p) { + uint32_t v; + asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} +__device__ __forceinline__ uint4 ld_relaxed_sys4(const uint4* p) { + uint4 v; + asm volatile("ld.relaxed.sys.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p) : "memory"); + return v; +} +__device__ __forceinline__ void st_relaxed_sys(uint32_t* p, uint32_t v) { + asm volatile("st.relaxed.sys.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} +__device__ __forceinline__ void st_release_gpu(uint32_t* p, uint32_t v) { + asm volatile("st.release.gpu.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} +__device__ __forceinline__ uint32_t arrive(uint32_t* p) { + uint32_t old; + asm volatile("atom.relaxed.gpu.global.add.u32 %0, [%1], 1;" : "=r"(old) : "l"(p) : "memory"); + return old; +} +__device__ __forceinline__ void fence_sys() { asm volatile("fence.sc.sys;" ::: "memory"); } +__device__ __forceinline__ void fence_gpu() { asm volatile("fence.sc.gpu;" ::: "memory"); } + +} // namespace + +// in: this rank's shard (packs of 16 bytes); out: [rank 0's shard | rank 1's shard]. klass: log2(gridDim.x). +extern "C" __global__ void __launch_bounds__(512) roce_gather( + const uint4* __restrict__ in, uint4* __restrict__ out, int packs, int rank, char* recv, char* flags, char* send, + uint32_t* ctrl, long long slot_bytes, uint32_t* counters, int klass, uint32_t spin_limit, int n_hca) { + const int tid = threadIdx.x; + uint32_t* stage = counters + 2 + klass; + uint32_t* tail = counters + 2 + K_CLASSES + klass; + if (ld_relaxed_gpu(counters + K_POISON) != 0u) return; + const uint32_t seq = ld_relaxed_gpu(counters + K_EPOCH) + 1u; + const uint32_t slot = seq & 1u; + const int at = blockIdx.x * blockDim.x + tid, stride = gridDim.x * blockDim.x; + + uint4* mine = reinterpret_cast(send + (long long)slot * slot_bytes); // 1. stage + for (int i = at; i < packs; i += stride) mine[i] = in[i]; + __syncthreads(); + if (tid == 0) { // 2. the doorbell, by the last block to stage + fence_sys(); + if ((arrive(stage) + 1u) % gridDim.x == 0u) { + st_relaxed_sys(ctrl + C_NBYTES + slot, (uint32_t)packs * 16u); + fence_sys(); + st_relaxed_sys(ctrl + C_SEQ, seq); + } + } + if (tid < n_hca) { // 3. the peer's stripes + const uint32_t* flag = reinterpret_cast(flags + ((long long)slot * n_hca + tid) * FLAG_STRIDE); + uint32_t polls = 0; + while (ld_acquire_sys(flag) != seq) { + if (++polls >= spin_limit) { + st_relaxed_sys(ctrl + C_HCA, (uint32_t)tid); + st_relaxed_sys(ctrl + C_ERROR, seq); + st_release_gpu(counters + K_POISON, seq); + break; + } + } + } + __syncthreads(); + if (ld_relaxed_gpu(counters + K_POISON) == 0u) { // 4. both shards in rank order + const uint4* theirs = reinterpret_cast(recv + (long long)slot * slot_bytes); + uint4* own_out = out + (long long)rank * packs; + uint4* peer_out = out + (long long)(1 - rank) * packs; + for (int i = at; i < packs; i += stride) { + own_out[i] = in[i]; + peer_out[i] = ld_relaxed_sys4(theirs + i); + } + } + fence_gpu(); // 5. the epoch, by the last block to finish + __syncthreads(); + if (tid == 0 && (arrive(tail) + 1u) % gridDim.x == 0u) { + fence_gpu(); + if (ld_relaxed_sys(ctrl + C_ERROR) == 0u) st_release_gpu(counters + K_EPOCH, seq); + } +} diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index b79337b5c5..8ea6c97123 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -45,6 +45,14 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "vit tf_vit_pointwise", .symbol = "tf_vit_pointwise", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "vit tf_vit_bias", .symbol = "tf_vit_bias", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "vit tf_vit_unshuffle", .symbol = "tf_vit_unshuffle", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 29, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "glm_exl3 rot_rows", .symbol = "glm_exl3_rot_rows", .image = "glm_exl3", .threads = 32 }, .usage = .{ .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 20, .static = 0, .bound = 0 }, .{ .regs = 20, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "glm_exl3 rot_in", .symbol = "glm_exl3_rot_in", .image = "glm_exl3", .threads = 32 }, .usage = .{ .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 20, .static = 0, .bound = 0 }, .{ .regs = 20, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "glm_exl3 grouped", .symbol = "glm_exl3_grouped", .image = "glm_exl3", .threads = 128 }, .usage = .{ .{ .regs = 120, .static = 32832, .bound = 128 }, .{ .regs = 130, .static = 32832, .bound = 128 }, .{ .regs = 130, .static = 32832, .bound = 128 }, .{ .regs = 109, .static = 32832, .bound = 128 }, .{ .regs = 109, .static = 32832, .bound = 128 } } }, + .{ .launch = .{ .name = "glm_exl3 gateup", .symbol = "glm_exl3_gateup", .image = "glm_exl3", .threads = 32 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "glm_exl3 down", .symbol = "glm_exl3_down", .image = "glm_exl3", .threads = 32 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "glm_exl3 prompt_gu", .symbol = "glm_exl3_prompt_gu", .image = "glm_exl3", .threads = 256 }, .usage = .{ .{ .regs = 127, .static = 43264, .bound = 256 }, .{ .regs = 128, .static = 43264, .bound = 256 }, .{ .regs = 128, .static = 43264, .bound = 256 }, .{ .regs = 126, .static = 43264, .bound = 256 }, .{ .regs = 126, .static = 43264, .bound = 256 } } }, + .{ .launch = .{ .name = "glm_exl3 prompt_dn", .symbol = "glm_exl3_prompt_dn", .image = "glm_exl3", .threads = 256 }, .usage = .{ .{ .regs = 126, .static = 43264, .bound = 256 }, .{ .regs = 126, .static = 43264, .bound = 256 }, .{ .regs = 126, .static = 43264, .bound = 256 }, .{ .regs = 127, .static = 43264, .bound = 256 }, .{ .regs = 127, .static = 43264, .bound = 256 } } }, + .{ .launch = .{ .name = "glm_exl3 picks", .symbol = "glm_exl3_picks", .image = "glm_exl3", .threads = 256 }, .usage = .{ .{ .regs = 13, .static = 0, .bound = 0 }, .{ .regs = 13, .static = 0, .bound = 0 }, .{ .regs = 13, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 } } }, }; /// Gate 5's planted CUDA controls: a launch past its __launch_bounds__, one needing more registers than a block has. diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 4a8f7d3493..c2f27f4ec1 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -274,17 +274,31 @@ pub const LaneHost = struct { h.mutex.unlock(h.io); } - /// Prefills the next queued request into a free lane; false when none waits or no lane is free. + /// The next queued request into a free lane (or several, when the backend batches prompts); false when none can go. fn admitOne(h: *LaneHost) bool { + if (admit.batches(h)) return admit.many(h); + const job = switch (h.takeJob()) { + .none => return false, + .failed => return true, + .job => |j| j, + }; + const began = h.setup(job) orelse return true; + return h.launch(job, began); + } + + pub const Taken = union(enum) { none, failed, job: *Job }; + + /// The next queued request a lane is free for, moved to the admitted list (`failed`: it was finished instead). + pub fn takeJob(h: *LaneHost) Taken { h.lock(); // a stale device takes no reply: `learn.refuse` ends the queued ones instead if (h.stale or h.queued.items.len == 0 or h.admitted.items.len >= @min(h.info_.lanes, h.memory_lanes orelse h.info_.lanes)) { h.unlock(); - return false; + return .none; } const at = admit.pick(h) orelse { // a shared cache pool's head waits for room h.unlock(); - return false; + return .none; }; const job = h.queued.orderedRemove(at); job.seq = h.admissions; @@ -292,11 +306,16 @@ pub const LaneHost = struct { h.admitted.append(h.gpa, job) catch { h.unlock(); h.finish(job, .failed, "out of memory"); - return true; + return .failed; }; h.unlock(); + return .{ .job = job }; + } + + /// A taken job's stream, set up to start; null when the job ended instead (dropped or cancelled). + pub fn setup(h: *LaneHost, job: *Job) ?i96 { const r = job.request; - if (r.background and r.max_tokens == 0 and h.warm_quiet > 0 and !h.quiet(job)) return h.cancel(job); // other work came first: it never ran + if (r.background and r.max_tokens == 0 and h.warm_quiet > 0 and !h.quiet(job)) return h.ended(h.cancel(job)); // other work came first var reuse: lanes.stream.Reuse = .{}; // the entry stays alive until the backend restores it: nothing keeps between here and this stream's own pass if (h.cache) |store| if (r.prompt_logprobs == null and !r.untaught) if (store.lookupRewind(h.gpa, r.cacheKey(), r.bornLen(), r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| { @@ -305,7 +324,7 @@ pub const LaneHost = struct { job.marks = l.marks; reuse = .{ .saved = if (l.entry) |e| e.saved else null, .at = if (l.entry) |e| e.at else 0, .marks = l.marks, .hook = .{ .ptr = job, .at = Job.kept } }; } else |_| {}; - job.proposer = lanes.SuffixLookup.init(h.gpa, .{ .min_match = h.min_match }) catch return h.drop(job, "the drafter could not start"); + job.proposer = lanes.SuffixLookup.init(h.gpa, .{ .min_match = h.min_match }) catch return h.ended(h.drop(job, "the drafter could not start")); const recall = recalls.reader(h, job); job.stream = lanes.Stream.init(h.gpa, .{ .id = "request", @@ -332,7 +351,7 @@ pub const LaneHost = struct { .grammar = if (r.structure) |st| h.constraint(st) catch |e| { job.proposer.deinit(); recalls.release(job); - return h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"); + return h.ended(h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)")); } else null, .decode_spans = h.spans(r), .decide = r.decide, @@ -342,26 +361,46 @@ pub const LaneHost = struct { }) catch { job.proposer.deinit(); recalls.release(job); - return h.drop(job, "out of memory"); + return h.ended(h.drop(job, "out of memory")); }; job.started = true; job.fill_on.store(true, .release); // status readers count the fill from here until the pass reports const began = std.Io.Clock.awake.now(h.io).toNanoseconds(); job.began = began; + return began; + } + + fn ended(_: *LaneHost, _: bool) ?i96 { + return null; + } + + /// Whether a set-up job starts outside a shared prompt pass: a warm pass, or the lone driver's. + pub fn startsAlone(h: *LaneHost, job: *Job) bool { + return (job.request.max_tokens == 0 and h.info_.warm_turns) or lone.fits(h, job); + } + + /// A set-up job into the engine: a warm pass, the lone driver, the fill queue, or its own prompt pass. + pub fn launch(h: *LaneHost, job: *Job, began: i96) bool { + const r = job.request; if (r.max_tokens == 0 and h.info_.warm_turns) return h.warmPass(job, began); if (lone.fits(h, job)) return lone.run(h, job, began, false); if (h.core.fills()) { h.filling.append(h.gpa, job) catch return h.drop(job, "out of memory"); return true; } - h.core.addStream(&job.stream) catch |e| return switch (e) { + h.core.addStream(&job.stream) catch |e| return h.failedPass(job, e); + h.prefilled(job, began); + _ = h.deliver(job); + return true; + } + + /// A prompt pass that failed: cancelled, no lane memory yet (back to the queue), or the backend's words. + pub fn failedPass(h: *LaneHost, job: *Job, e: anyerror) bool { + return switch (e) { error.Cancelled => h.cancel(job), error.LaneMemory => h.requeue(job), else => h.drop(job, h.words(e)), }; - h.prefilled(job, began); - _ = h.deliver(job); - return true; } /// The next filling job's next prompt chunk: the other streams' rounds run between chunks, not after the prompt. diff --git a/zig/src/core/lane_host_admit.zig b/zig/src/core/lane_host_admit.zig index 016ab4398a..fde33f7010 100644 --- a/zig/src/core/lane_host_admit.zig +++ b/zig/src/core/lane_host_admit.zig @@ -2,6 +2,7 @@ const api = @import("engine_api.zig"); const LaneHost = @import("lane_host.zig").LaneHost; const Job = LaneHost.Job; +const lanes = @import("lanes"); /// The queued job to admit (its index), or null to hold; the caller holds the lock and a lane is free. pub fn pick(h: *LaneHost) ?usize { @@ -21,3 +22,53 @@ pub fn pick(h: *LaneHost) ?usize { fn fits(a: api.Admit, job: *const Job) bool { return a.fits(a.ctx, @intCast(job.request.prompt.len), job.request.max_tokens); } + +/// Whether to admit several at once: prompts batch, none fill by chunks, and two or more wait with two lanes free. +pub fn batches(h: *LaneHost) bool { + if (!h.core.batchesPrompts() or h.core.fills()) return false; + h.lock(); + defer h.unlock(); + const lanes_now = @min(h.info_.lanes, h.memory_lanes orelse h.info_.lanes); + return h.queued.items.len >= 2 and lanes_now -| h.admitted.items.len >= 2; +} + +/// Every queued request a lane is free for, prompts in one pass (lone and warm ones as before); false: none taken. +pub fn many(h: *LaneHost) bool { + var jobs: [lanes.Engine.max_together]*Job = undefined; + var began: [lanes.Engine.max_together]i96 = undefined; + var n: usize = 0; + var took = false; + while (n < jobs.len) { + const job = switch (h.takeJob()) { + .none => break, + .failed => { + took = true; + continue; + }, + .job => |j| j, + }; + took = true; + const at = h.setup(job) orelse continue; + if (h.startsAlone(job)) { + _ = h.launch(job, at); + continue; + } + jobs[n] = job; + began[n] = at; + n += 1; + } + if (n == 1) return h.launch(jobs[0], began[0]); + var streams: [lanes.Engine.max_together]*lanes.Stream = undefined; + for (jobs[0..n], 0..) |job, i| streams[i] = &job.stream; + var errs: [lanes.Engine.max_together]?anyerror = undefined; + h.core.addStreams(streams[0..n], errs[0..n]); + for (jobs[0..n], began[0..n], errs[0..n]) |job, at, err| { + if (err) |e| { + _ = h.failedPass(job, e); + continue; + } + h.prefilled(job, at); + _ = h.deliver(job); + } + return took; +} diff --git a/zig/src/core/lane_host_many_test.zig b/zig/src/core/lane_host_many_test.zig new file mode 100644 index 0000000000..25b6993b0c --- /dev/null +++ b/zig/src/core/lane_host_many_test.zig @@ -0,0 +1,103 @@ +//! Several queued requests in one prompt pass (a backend with prefill_many): the replies a pass alone would give. +const std = @import("std"); +const lanes = @import("lanes"); +const api = @import("engine_api.zig"); +const LaneHost = @import("lane_host.zig").LaneHost; +const Reason = api.Reason; +const Id = api.Id; +const Event = api.Event; +const Request = api.Request; +const gpa = std.testing.allocator; + +const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: std.ArrayList(u32) = .empty, + done: ?Reason = null, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *Box = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens.appendSlice(gpa, t) catch {}, + .finished => |f| b.done = f.reason, + else => {}, + } + } + fn wait(b: *Box) Reason { + while (true) { + b.mutex.lockUncancelable(std.testing.io); + const d = b.done; + b.mutex.unlock(std.testing.io); + if (d) |r| return r; + std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + } + } +}; + +/// Each reply the fake target's own continuation of its prompt. +fn expectOwn(prompt: []const u32, tokens: []const u32) !void { + var history: std.ArrayList(u32) = .empty; + defer history.deinit(gpa); + try history.appendSlice(gpa, prompt); + for (tokens) |t| { + try std.testing.expectEqual(lanes.fake.next(history.items, null, history.items.len), t); + try history.append(gpa, t); + } +} + +/// `n` requests queued before the host starts (`sampled`: one the fake refuses), replies checked, batched passes. +fn serve(n: usize, memory_lanes: ?usize, sampled: ?usize) !struct { passes: usize, streams: usize } { + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa, .memory_lanes = memory_lanes, .refuse_sampled = sampled != null }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.batched(), clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 4 }); + var prompts: [6][5]u32 = undefined; + var requests: [6]Request = undefined; + var boxes: [6]Box = @splat(.{}); + defer for (&boxes) |*b| b.tokens.deinit(gpa); + const e = host.engine(); + for (0..n) |i| { + prompts[i] = .{ 2, 7, 1, 8, @intCast(10 + i) }; + requests[i] = .{ .prompt = &prompts[i], .max_tokens = 12 }; + if (sampled == i) requests[i].sampling = .{ .seed = 3, .temperature = 0.7, .top_k = 5 }; + try e.submit(@intCast(i + 1), &requests[i], .{ .ctx = &boxes[i], .event = Box.event }); + } + try host.start(); + defer host.stop(); + for (0..n) |i| { + if (sampled == i) { + try std.testing.expectEqual(Reason.failed, boxes[i].wait()); + continue; + } + try std.testing.expectEqual(Reason.length, boxes[i].wait()); + try std.testing.expectEqual(@as(usize, 12), boxes[i].tokens.items.len); + try expectOwn(&prompts[i], boxes[i].tokens.items); + } + return .{ .passes = target.many_passes, .streams = target.many_streams }; +} + +test "queued requests share one prompt pass, each reply the target's own continuation" { + const r = try serve(4, null, null); + try std.testing.expectEqual(@as(usize, 1), r.passes); + try std.testing.expectEqual(@as(usize, 4), r.streams); +} + +test "more requests than lanes: the free lanes' prompts in one pass, the rest after" { + const r = try serve(6, null, null); + try std.testing.expect(r.passes >= 1); + try std.testing.expect(r.streams >= 4 and r.streams <= 6); +} + +test "a backend without room for them all: they go one at a time, a later pass only as many as fit, same replies" { + const r = try serve(4, 2, null); + try std.testing.expect(r.streams <= 2 * r.passes); // the fake refuses a pass past its room: none took more +} + +test "one request the backend refuses fails alone: the pass is retried a stream at a time, the others' replies kept" { + const r = try serve(4, null, 2); + try std.testing.expectEqual(@as(usize, 0), r.passes); +} diff --git a/zig/src/core/lane_host_test.zig b/zig/src/core/lane_host_test.zig index 85a34c0e15..56f8e353cd 100644 --- a/zig/src/core/lane_host_test.zig +++ b/zig/src/core/lane_host_test.zig @@ -570,4 +570,5 @@ test { _ = @import("lane_host_borrow_test.zig"); _ = @import("lane_host_warm_test.zig"); _ = @import("lane_host_rows_test.zig"); + _ = @import("lane_host_many_test.zig"); } diff --git a/zig/src/core/lanes/backend.zig b/zig/src/core/lanes/backend.zig index 8378c44900..31cf0f4786 100644 --- a/zig/src/core/lanes/backend.zig +++ b/zig/src/core/lanes/backend.zig @@ -68,6 +68,8 @@ pub const Backend = struct { prefill: *const fn (ptr: *anyopaque, s: *Stream) anyerror!void, /// One chunk of the prompt pass, false while chunks remain; null: the pass runs whole. prefill_step: ?*const fn (ptr: *anyopaque, s: *Stream) anyerror!bool = null, + /// Several prompts in one pass, each as after `prefill`; false or an error leaves none prefilled. + prefill_many: ?*const fn (ptr: *anyopaque, ss: []const *Stream) anyerror!bool = null, /// Draw the first token from the prompt's last row at `position`; a handle to it. first: *const fn (ptr: *anyopaque, s: *Stream, position: u64) anyerror!u64, /// Feed one token and queue the draw of the next at `position` (a one-token round); a handle to it. diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index f11d8dbec5..7e47830941 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -79,6 +79,13 @@ pub const Engine = struct { } pub const addStream = open.add; + pub const addStreams = open.addStreams; + pub const max_together = open.max_together; + + /// Whether the backend prefills several prompts in one pass (addStreams batches them). + pub fn batchesPrompts(e: *const Engine) bool { + return e.backend.vtable.prefill_many != null; + } pub const opened = open.opened; pub const fills = open.fills; pub const fillStream = open.fillStream; diff --git a/zig/src/core/lanes/fake.zig b/zig/src/core/lanes/fake.zig index 7827f286ac..234fe8d63e 100644 --- a/zig/src/core/lanes/fake.zig +++ b/zig/src/core/lanes/fake.zig @@ -53,6 +53,8 @@ pub const Fake = struct { answer_cycles: bool = false, prefill_chunks: usize = 0, prefill_count: usize = 0, + many_passes: usize = 0, // prefill_many calls that took their streams (batched()) + many_streams: usize = 0, // streams those calls prefilled prefill_hook: ?*const fn (ctx: *anyopaque, s: *Stream, chunk: usize) void = null, prefill_hook_ctx: ?*anyopaque = null, refuse_sampled: bool = false, // prefill refuses a sampled stream with error.SamplingRefused @@ -79,6 +81,26 @@ pub const Fake = struct { return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .first = first, .first_masked = firstMasked, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .features = features, .release = release, .first_row = firstRow, .alternatives = alternatives, .held = heldDrafts, .reports_fill = true } }; } + /// The same target taking several prompts in one pass (`prefill_many`). + pub fn batched(x: *Fake) be.Backend { + return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .prefill_many = prefillMany, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .features = features, .release = release, .first_row = firstRow, .reports_fill = true } }; + } + + fn prefillMany(ptr: *anyopaque, ss: []const *Stream) anyerror!bool { + const x = self(ptr); + if (x.memory_lanes) |n| if (x.lanes.count() + ss.len > n) return false; // not all at once: one at a time instead + for (ss, 0..) |s, i| prefill(ptr, s) catch |err| { // none left prefilled, as the contract says + for (ss[0..i]) |done| if (x.lanes.fetchRemove(done)) |kv| { + var l = kv.value; + freeLane(x.gpa, &l); + }; + return err; + }; + x.many_passes += 1; + x.many_streams += ss.len; + return true; + } + /// The same target filling its prompt a chunk a call (`prefill_chunks`), as a backend with `prefill_step`. pub fn stepped(x: *Fake) be.Backend { return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .prefill_step = prefillStep, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .features = features, .release = release, .first_row = firstRow, .reports_fill = true } }; diff --git a/zig/src/core/lanes/open.zig b/zig/src/core/lanes/open.zig index 7cddf9a064..e0b0ba45e8 100644 --- a/zig/src/core/lanes/open.zig +++ b/zig/src/core/lanes/open.zig @@ -23,6 +23,29 @@ pub fn add(e: *Engine, s: *Stream) !void { try opened(e, s); } +/// Streams a backend prefills in one pass at most. +pub const max_together = 16; + +/// `add` for several streams, prompts in one pass when the backend takes them; `errs[i]`: stream i's error, or null. +pub fn addStreams(e: *Engine, ss: []const *Stream, errs: []?anyerror) void { + for (errs) |*x| x.* = null; + const many = e.backend.vtable.prefill_many orelse return addEach(e, ss, errs); + if (ss.len < 2 or ss.len > max_together) return addEach(e, ss, errs); + _ = e.arena.reset(.retain_capacity); + for (ss) |s| trail.event(e, &.{ f("ev", str("add")), f("stream", str(s.id)) }) catch {}; + const together = many(e.backend.ptr, ss) catch false; // one stream's refusal is its own: each goes alone + if (!together) return addEach(e, ss, errs); + for (ss, errs) |s, *x| opened(e, s) catch |err| { // a stream cancelled in the pass gets no first token + x.* = err; + }; +} + +fn addEach(e: *Engine, ss: []const *Stream, errs: []?anyerror) void { + for (ss, errs) |s, *x| add(e, s) catch |err| { + x.* = err; + }; +} + /// After the prompt pass: the first token drawn and committed, the first drafts asked; rounds from the next. pub fn opened(e: *Engine, s: *Stream) !void { if (s.isCancelled()) { // cancelled in its last chunk: no first token diff --git a/zig/src/core/lanes/recorder.zig b/zig/src/core/lanes/recorder.zig index e62eda60fd..841424ddb7 100644 --- a/zig/src/core/lanes/recorder.zig +++ b/zig/src/core/lanes/recorder.zig @@ -229,7 +229,7 @@ pub const Recorded = struct { pub fn backend(x: *Recorded) be.Backend { const v = x.inner.vtable; - x.table = .{ .prefill = prefill, .prefill_step = if (v.prefill_step != null) prefillStep else null, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .unspeculate = if (v.unspeculate != null) unspeculate else null, .probabilities = if (v.probabilities != null) probabilities else null, .tree = if (v.tree != null) tree else null, .alternatives = if (v.alternatives != null) alternatives else null, .prepare_features = if (v.prepare_features != null) prepareFeatures else null, .features = features, .first_row = if (v.first_row != null) firstRow else null, .first_masked = if (v.first_masked != null) firstMasked else null, .held = if (v.held != null) held else null, .prompt_rows = v.prompt_rows, .release = release }; + x.table = .{ .prefill = prefill, .prefill_step = if (v.prefill_step != null) prefillStep else null, .prefill_many = if (v.prefill_many != null) prefillMany else null, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .unspeculate = if (v.unspeculate != null) unspeculate else null, .probabilities = if (v.probabilities != null) probabilities else null, .tree = if (v.tree != null) tree else null, .alternatives = if (v.alternatives != null) alternatives else null, .prepare_features = if (v.prepare_features != null) prepareFeatures else null, .features = features, .first_row = if (v.first_row != null) firstRow else null, .first_masked = if (v.first_masked != null) firstMasked else null, .held = if (v.held != null) held else null, .prompt_rows = v.prompt_rows, .release = release }; return .{ .ptr = x, .vtable = &x.table }; } @@ -247,6 +247,12 @@ pub const Recorded = struct { if (done) x.recorder.prompt(x.inner, s); return done; } + fn prefillMany(p: *anyopaque, ss: []const *Stream) anyerror!bool { + const x = of(p); + if (!try x.inner.vtable.prefill_many.?(x.inner.ptr, ss)) return false; + for (ss) |s| x.recorder.prompt(x.inner, s); // each prompt's rows, as after its own prefill + return true; + } fn first(p: *anyopaque, s: *Stream, position: u64) anyerror!u64 { return of(p).inner.first(s, position); } diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 18542b8cb4..096ea30aa5 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -39,6 +39,9 @@ pub const modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack pub const glm_glue: []const u8 = if (available) &Blob("fatbin_glm_glue").bytes else &.{}; pub const glm_ops: []const u8 = if (available) &Blob("fatbin_glm_ops").bytes else &.{}; pub const glm_quant: []const u8 = if (available) &Blob("fatbin_glm_quant").bytes else &.{}; +pub const glm_exl3: []const u8 = if (available) &Blob("fatbin_glm_exl3").bytes else &.{}; +pub const glm_vision: []const u8 = if (available) &Blob("fatbin_glm_vision").bytes else &.{}; +pub const glm_mm: []const u8 = if (available) &Blob("fatbin_glm_mm").bytes else &.{}; pub const glm_kda: []const u8 = if (available) &Blob("fatbin_glm_kda").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 &.{}; @@ -56,6 +59,7 @@ pub const ck_lane: []const u8 = if (available) &Blob("fatbin_ck_lane").bytes els pub const ck_experts: []const u8 = if (available) &Blob("fatbin_ck_experts").bytes else &.{}; pub const norms: []const u8 = if (available) &Blob("fatbin_norms").bytes else &.{}; pub const embed_rows: []const u8 = if (available) &Blob("fatbin_embed_rows").bytes else &.{}; +pub const roce_gather: []const u8 = if (available) &Blob("fatbin_roce_gather").bytes else &.{}; /// Symbols in the gdn image as cuobjdump lists them for the built fatbin (named namespace tf_gdn). pub const gdn_symbols = struct { diff --git a/zig/src/cuda/nccl.zig b/zig/src/cuda/nccl.zig index b0d527ce75..b8d8f39222 100644 --- a/zig/src/cuda/nccl.zig +++ b/zig/src/cuda/nccl.zig @@ -21,6 +21,8 @@ pub const Api = struct { ncclCommInitRank: *const fn (*Comm, c_int, UniqueId, c_int) callconv(.c) R, ncclCommInitAll: *const fn ([*]Comm, c_int, ?[*]const c_int) callconv(.c) R, ncclCommDestroy: *const fn (Comm) callconv(.c) R, + ncclCommAbort: *const fn (Comm) callconv(.c) R, + ncclCommGetAsyncError: *const fn (Comm, *R) callconv(.c) R, ncclGetErrorString: *const fn (R) callconv(.c) ?[*:0]const u8, ncclAllReduce: *const fn (D, D, usize, DataType, RedOp, Comm, abi.Stream) callconv(.c) R, ncclAllGather: *const fn (D, D, usize, DataType, Comm, abi.Stream) callconv(.c) R, diff --git a/zig/src/cuda/roce.zig b/zig/src/cuda/roce.zig new file mode 100644 index 0000000000..25295d9bcc --- /dev/null +++ b/zig/src/cuda/roce.zig @@ -0,0 +1,596 @@ +//! Two-rank all-gathers as one-shot RoCE writes after b12x's RoCEnante (local-inference-lab/b12x, Apache-2.0). + +const std = @import("std"); +const abi = @import("abi.zig"); +const Driver = @import("driver.zig").Driver; +const memory = @import("memory.zig"); +const Module = @import("module.zig").Module; +const Function = @import("module.zig").Function; +const launch = @import("launch.zig"); +const tp_group = @import("tp_group.zig"); +const Stream = @import("stream.zig").Stream; +const kernels = @import("kernels.zig"); + +const max_hcas = 2; +const slots = 2; +const flag_stride = 128; +const send_depth = 256; +const port_num: u8 = 1; +const threads = 512; +const max_blocks = 8; // the largest grid; smaller shards launch fewer (a power of two) +const classes = 4; // grid size classes 1, 2, 4, 8 (each its own arrival counters) +const idle_spins: u64 = 20_000_000; // doorbell polls before the proxy naps between them (a model's gaps are short) +const spin_limit: u32 = 100_000_000; // the kernel's flag polls before it gives up and poisons the exchange + +// -- libibverbs (the ABI of rdma-core's verbs.h; offsets checked against its headers) --------------------------------- +const Device = opaque {}; +const Context = opaque {}; +const Pd = opaque {}; +const Cq = extern struct { context: *Context }; +const Mr = extern struct { context: *Context, pd: *Pd, addr: ?*anyopaque, length: usize, handle: u32, lkey: u32, rkey: u32 }; +const Qp = extern struct { context: *Context, qp_context: ?*anyopaque, pd: *Pd, send_cq: *Cq, recv_cq: *Cq, srq: ?*anyopaque, handle: u32, qp_num: u32 }; +const QpCap = extern struct { max_send_wr: u32 = 0, max_recv_wr: u32 = 0, max_send_sge: u32 = 0, max_recv_sge: u32 = 0, max_inline_data: u32 = 0 }; +const QpInitAttr = extern struct { qp_context: ?*anyopaque = null, send_cq: ?*Cq = null, recv_cq: ?*Cq = null, srq: ?*anyopaque = null, cap: QpCap = .{}, qp_type: c_int = 0, sq_sig_all: c_int = 0 }; +const GlobalRoute = extern struct { dgid: [16]u8 = @splat(0), flow_label: u32 = 0, sgid_index: u8 = 0, hop_limit: u8 = 0, traffic_class: u8 = 0 }; +const AhAttr = extern struct { grh: GlobalRoute = .{}, dlid: u16 = 0, sl: u8 = 0, src_path_bits: u8 = 0, static_rate: u8 = 0, is_global: u8 = 0, port_num: u8 = 0 }; +const QpAttr = extern struct { + qp_state: u32 = 0, + cur_qp_state: u32 = 0, + path_mtu: u32 = 0, + path_mig_state: u32 = 0, + qkey: u32 = 0, + rq_psn: u32 = 0, + sq_psn: u32 = 0, + dest_qp_num: u32 = 0, + qp_access_flags: u32 = 0, + cap: QpCap = .{}, + ah_attr: AhAttr = .{}, + alt_ah_attr: AhAttr = .{}, + pkey_index: u16 = 0, + alt_pkey_index: u16 = 0, + en_sqd_async_notify: u8 = 0, + sq_draining: u8 = 0, + max_rd_atomic: u8 = 0, + max_dest_rd_atomic: u8 = 0, + min_rnr_timer: u8 = 0, + port_num: u8 = 0, + timeout: u8 = 0, + retry_cnt: u8 = 0, + rnr_retry: u8 = 0, + alt_port_num: u8 = 0, + alt_timeout: u8 = 0, + rate_limit: u32 = 0, + _tail: u32 = 0, +}; +const Sge = extern struct { addr: u64, length: u32, lkey: u32 }; +const SendWr = extern struct { + wr_id: u64 = 0, + next: ?*SendWr = null, + sg_list: ?*Sge = null, + num_sge: c_int = 0, + opcode: u32 = 0, + send_flags: u32 = 0, + imm_data: u32 = 0, + remote_addr: u64 = 0, + rkey: u32 = 0, + _tail: [128 - 52]u8 = @splat(0), +}; +const Wc = extern struct { wr_id: u64, status: u32, opcode: u32, vendor_err: u32, _rest: [48 - 20]u8 }; +const PortAttr = extern struct { bytes: [56]u8 align(8) }; + +comptime { + std.debug.assert(@sizeOf(QpInitAttr) == 64 and @offsetOf(QpInitAttr, "cap") == 32 and @offsetOf(QpInitAttr, "qp_type") == 52); + std.debug.assert(@sizeOf(AhAttr) == 32 and @offsetOf(AhAttr, "dlid") == 24 and @offsetOf(AhAttr, "port_num") == 30); + std.debug.assert(@sizeOf(QpAttr) == 144 and @offsetOf(QpAttr, "ah_attr") == 56 and @offsetOf(QpAttr, "pkey_index") == 120); + std.debug.assert(@offsetOf(QpAttr, "rnr_retry") == 132 and @offsetOf(QpAttr, "qp_access_flags") == 32); + std.debug.assert(@sizeOf(SendWr) == 128 and @offsetOf(SendWr, "remote_addr") == 40 and @offsetOf(SendWr, "rkey") == 48); + std.debug.assert(@sizeOf(Wc) == 48 and @offsetOf(Mr, "lkey") == 36 and @offsetOf(Qp, "qp_num") == 52); +} + +// struct ibv_context's ops: the two calls verbs.h inlines +const ops_poll_cq = 96; +const ops_post_send = 208; +const PollCq = *const fn (*Cq, c_int, [*]Wc) callconv(.c) c_int; +const PostSend = *const fn (*Qp, *SendWr, *?*SendWr) callconv(.c) c_int; + +const wr_rdma_write = 0; +const send_signaled = 2; +const send_inline = 8; +const qpt_rc = 2; +const qps_init = 1; +const qps_rtr = 2; +const qps_rts = 3; +const port_active = 4; +const access_local_write = 1; +const access_remote_write = 2; +const qp_state = 1; +const qp_access_flags = 8; +const qp_pkey_index = 16; +const qp_port = 32; +const qp_av = 128; +const qp_path_mtu = 256; +const qp_timeout = 512; +const qp_retry_cnt = 1024; +const qp_rnr_retry = 2048; +const qp_rq_psn = 4096; +const qp_max_qp_rd_atomic = 8192; +const qp_min_rnr_timer = 32768; +const qp_sq_psn = 65536; +const qp_max_dest_rd_atomic = 131072; +const qp_dest_qpn = 1048576; + +const Api = struct { + ibv_get_device_list: *const fn (?*c_int) callconv(.c) ?[*]?*Device, + ibv_free_device_list: *const fn ([*]?*Device) callconv(.c) void, + ibv_get_device_name: *const fn (*Device) callconv(.c) ?[*:0]const u8, + ibv_open_device: *const fn (*Device) callconv(.c) ?*Context, + ibv_close_device: *const fn (*Context) callconv(.c) c_int, + ibv_query_port: *const fn (*Context, u8, *PortAttr) callconv(.c) c_int, + ibv_query_gid: *const fn (*Context, u8, c_int, *[16]u8) callconv(.c) c_int, + ibv_alloc_pd: *const fn (*Context) callconv(.c) ?*Pd, + ibv_dealloc_pd: *const fn (*Pd) callconv(.c) c_int, + ibv_reg_mr: *const fn (*Pd, ?*anyopaque, usize, c_int) callconv(.c) ?*Mr, + ibv_dereg_mr: *const fn (*Mr) callconv(.c) c_int, + ibv_create_cq: *const fn (*Context, c_int, ?*anyopaque, ?*anyopaque, c_int) callconv(.c) ?*Cq, + ibv_destroy_cq: *const fn (*Cq) callconv(.c) c_int, + ibv_create_qp: *const fn (*Pd, *QpInitAttr) callconv(.c) ?*Qp, + ibv_destroy_qp: *const fn (*Qp) callconv(.c) c_int, + ibv_modify_qp: *const fn (*Qp, *QpAttr, c_int) callconv(.c) c_int, + ibv_wc_status_str: *const fn (u32) callconv(.c) ?[*:0]const u8, +}; + +fn opsCall(comptime T: type, ctx: *Context, offset: usize) T { + const slot: *const T = @ptrFromInt(@intFromPtr(ctx) + offset); + return slot.*; +} + +pub const Error = error{ VerbsUnavailable, NoHca, HcaFailed, PeerRefused, ProxyFailed, ExchangeTimedOut, Misaligned }; + +/// TF_ROCE=0 keeps every gather on NCCL; unset or any other value asks for this path (both ranks must ask). +pub fn wanted() bool { + const v = std.c.getenv("TF_ROCE") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// The peer's `mine`: both ranks' values cross the group's communicator, rank order. +fn swap(comptime T: type, group: *const tp_group.Group, d: *const Driver, s: Stream, mine: T) !T { + var both: [2]T = undefined; + try group.gather(T, d, s, mine, &both); + return both[1 - group.rank]; +} + +// -- the region's layout (the same on both ranks) --------------------------------------------------------------------- +const Layout = struct { + slot_bytes: usize, + recv: usize, + flags: usize, + send: usize, + ctrl: usize, + total: usize, + + fn of(max_bytes: usize) Layout { + const s = std.mem.alignForward(usize, max_bytes, 4096); + const flags = slots * s; + const send = flags + slots * max_hcas * flag_stride; + const ctrl = send + slots * s; + return .{ .slot_bytes = s, .recv = 0, .flags = flags, .send = send, .ctrl = ctrl, .total = ctrl + flag_stride }; + } +}; +// ctrl words +const c_seq = 0; +const c_nbytes = 1; // + slot +const c_error = 3; +const c_hca = 4; + +/// What a rank tells its peer to connect: its region, and per HCA its key, address and queue pair. +const Blob = extern struct { + ok: u32, + n_hca: u32, + slot_bytes: u64, + max_bytes: u64, // the largest shard this rank would send this way (both ranks use the smaller) + addr: u64, + rkey: [max_hcas]u32, + qpn: [max_hcas]u32, + mtu: [max_hcas]u32, + lid: [max_hcas]u16, + gid: [max_hcas][16]u8, +}; + +const Hca = struct { + ctx: ?*Context = null, + pd: ?*Pd = null, + mr: ?*Mr = null, + cq: ?*Cq = null, + qp: ?*Qp = null, + gid: [16]u8 = @splat(0), + gid_index: c_int = 0, + lid: u16 = 0, + mtu: u32 = 0, + outstanding: u32 = 0, +}; + +pub const Gather = struct { + gpa: std.mem.Allocator, + d: *const Driver, + lib: std.DynLib, + api: Api, + rank: u8, + hcas: [max_hcas]Hca = @splat(.{}), + n_hca: u32 = 0, + layout: Layout, + region: memory.HostBuffer, + region_dev: abi.DevicePtr, + counters: memory.DeviceBuffer, + module: Module, + f: Function, + max_bytes: usize, + peer_addr: u64 = 0, + peer_rkey: [max_hcas]u32 = @splat(0), + last_seq: u32 = 0, + running: std.atomic.Value(bool) = .init(false), + failed: std.atomic.Value(bool) = .init(false), + thread: ?std.Thread = null, + err: [256]u8 = @splat(0), + gathers: u64 = 0, + + /// Both ranks' exchange or null on both; setup crosses ``group``'s communicator on ``s``. ``max_bytes``: a shard. + pub fn init(gpa: std.mem.Allocator, io: std.Io, d: *const Driver, group: *const tp_group.Group, s: Stream, want: bool, max_bytes: usize) !?*Gather { + if (group.world != 2) return error.PeerRefused; + const g = try gpa.create(Gather); + var mine: Blob = std.mem.zeroes(Blob); + const ok = want and blk: { + g.setup(gpa, io, d, @intCast(group.rank), max_bytes) catch |e| { + std.log.warn("RoCE all-gather: {s} ({s}); NCCL carries every gather", .{ @errorName(e), std.mem.sliceTo(&g.err, 0) }); + break :blk false; + }; + break :blk true; + }; + if (ok) g.blob(&mine); + const theirs = swap(Blob, group, d, s, mine) catch |e| { + if (ok) g.teardown(); + gpa.destroy(g); + return e; + }; + if (!ok or theirs.ok == 0 or theirs.n_hca != mine.n_hca or theirs.slot_bytes != mine.slot_bytes) { + if (ok) g.teardown(); + gpa.destroy(g); + return null; + } + g.max_bytes = agreedCap(mine.max_bytes, theirs.max_bytes); // one cap for fits on both ranks + const connected = blk: { + g.connect(&theirs) catch |e| { + std.log.warn("RoCE all-gather: {s} ({s}); NCCL carries every gather", .{ @errorName(e), std.mem.sliceTo(&g.err, 0) }); + break :blk false; + }; + break :blk true; + }; + const peer_connected = swap(u32, group, d, s, @intFromBool(connected)) catch |e| { + g.teardown(); + gpa.destroy(g); + return e; + }; + if (!connected or peer_connected == 0) { + g.teardown(); + gpa.destroy(g); + return null; + } + std.log.info("RoCE all-gather: {d} HCA(s), shards up to {d} KiB by RDMA writes", .{ g.n_hca, g.max_bytes >> 10 }); + return g; + } + + fn fail(g: *Gather, comptime fmt: []const u8, args: anytype) void { + _ = std.fmt.bufPrint(g.err[0 .. g.err.len - 1], fmt, args) catch {}; + } + + fn setup(g: *Gather, gpa: std.mem.Allocator, io: std.Io, d: *const Driver, rank: u8, max_bytes: usize) !void { + g.* = .{ .gpa = gpa, .d = d, .lib = undefined, .api = undefined, .rank = rank, .layout = .of(max_bytes), .region = undefined, .region_dev = 0, .counters = undefined, .module = undefined, .f = undefined, .max_bytes = max_bytes }; + g.lib = std.DynLib.open("libibverbs.so.1") catch return error.VerbsUnavailable; + errdefer g.lib.close(); + const info = @typeInfo(Api).@"struct"; + inline for (info.field_names, info.field_types) |name, T| { + @field(g.api, name) = g.lib.lookup(T, name) orelse return error.VerbsUnavailable; + } + var names: [max_hcas][64]u8 = undefined; + var lens: [max_hcas]usize = undefined; + g.n_hca = try hcaNames(io, gpa, &names, &lens); + g.region = try memory.HostBuffer.allocMapped(d, g.layout.total); + errdefer g.region.free(); + @memset(g.region.bytes, 0); + g.region_dev = try g.region.device(); + g.counters = try memory.DeviceBuffer.alloc(d, 64 * 4); + errdefer g.counters.free(); + const zeros: [64 * 4]u8 = @splat(0); + try g.counters.upload(0, &zeros); + g.module = try Module.load(d, kernels.roce_gather); + errdefer g.module.unload(); + g.f = try g.module.function("roce_gather"); + errdefer g.closeHcas(); + for (0..g.n_hca) |h| try g.openHca(io, gpa, h, names[h][0..lens[h]]); + } + + fn closeHcas(g: *Gather) void { + for (&g.hcas) |*x| { + if (x.qp) |q| _ = g.api.ibv_destroy_qp(q); + if (x.cq) |q| _ = g.api.ibv_destroy_cq(q); + if (x.mr) |m| _ = g.api.ibv_dereg_mr(m); + if (x.pd) |p| _ = g.api.ibv_dealloc_pd(p); + if (x.ctx) |c| _ = g.api.ibv_close_device(c); + x.* = .{}; + } + } + + fn openHca(g: *Gather, io: std.Io, gpa: std.mem.Allocator, h: usize, name: []const u8) !void { + const a = &g.api; + var num: c_int = 0; + const list = a.ibv_get_device_list(&num) orelse return error.NoHca; + defer a.ibv_free_device_list(list); + var dev: ?*Device = null; + for (0..@intCast(num)) |i| { + const n = a.ibv_get_device_name(list[i].?) orelse continue; + if (std.mem.eql(u8, std.mem.span(n), name)) dev = list[i]; + } + const hca = &g.hcas[h]; + hca.ctx = a.ibv_open_device(dev orelse { + g.fail("RDMA device {s} not found", .{name}); + return error.NoHca; + }) orelse return error.HcaFailed; + var port: PortAttr = std.mem.zeroes(PortAttr); + if (a.ibv_query_port(hca.ctx.?, port_num, &port) != 0) return error.HcaFailed; + if (std.mem.readInt(u32, port.bytes[0..4], .little) != port_active) { + g.fail("{s} port 1 is not active", .{name}); + return error.NoHca; + } + hca.mtu = std.mem.readInt(u32, port.bytes[8..12], .little); + hca.lid = std.mem.readInt(u16, port.bytes[34..36], .little); + hca.gid_index = try gidIndex(io, gpa, name); + if (a.ibv_query_gid(hca.ctx.?, port_num, hca.gid_index, &hca.gid) != 0) return error.HcaFailed; + hca.pd = a.ibv_alloc_pd(hca.ctx.?) orelse return error.HcaFailed; + hca.mr = a.ibv_reg_mr(hca.pd.?, g.region.bytes.ptr, g.region.bytes.len, access_local_write | access_remote_write) orelse { + g.fail("{s}: registering the pinned region failed", .{name}); + return error.HcaFailed; + }; + hca.cq = a.ibv_create_cq(hca.ctx.?, send_depth, null, null, 0) orelse return error.HcaFailed; + var attr: QpInitAttr = .{ .send_cq = hca.cq, .recv_cq = hca.cq, .qp_type = qpt_rc, .cap = .{ .max_send_wr = send_depth, .max_recv_wr = 1, .max_send_sge = 1, .max_recv_sge = 1, .max_inline_data = 16 } }; + hca.qp = a.ibv_create_qp(hca.pd.?, &attr) orelse return error.HcaFailed; + var init_attr: QpAttr = .{ .qp_state = qps_init, .port_num = port_num, .qp_access_flags = access_remote_write }; + if (a.ibv_modify_qp(hca.qp.?, &init_attr, qp_state | qp_pkey_index | qp_port | qp_access_flags) != 0) return error.HcaFailed; + } + + fn blob(g: *Gather, b: *Blob) void { + b.ok = 1; + b.n_hca = g.n_hca; + b.slot_bytes = g.layout.slot_bytes; + b.max_bytes = g.max_bytes; + b.addr = @intFromPtr(g.region.bytes.ptr); + for (0..g.n_hca) |h| { + const x = &g.hcas[h]; + b.rkey[h] = x.mr.?.rkey; + b.qpn[h] = x.qp.?.qp_num; + b.mtu[h] = x.mtu; + b.lid[h] = x.lid; + b.gid[h] = x.gid; + } + } + + fn connect(g: *Gather, peer: *const Blob) !void { + g.peer_addr = peer.addr; + for (0..g.n_hca) |h| { + const x = &g.hcas[h]; + g.peer_rkey[h] = peer.rkey[h]; + var rtr: QpAttr = .{ .qp_state = qps_rtr, .path_mtu = @min(peer.mtu[h], x.mtu), .dest_qp_num = peer.qpn[h], .max_dest_rd_atomic = 1, .min_rnr_timer = 12, .ah_attr = .{ .is_global = 1, .dlid = peer.lid[h], .port_num = port_num, .grh = .{ .dgid = peer.gid[h], .sgid_index = @intCast(x.gid_index), .hop_limit = 64 } } }; + if (g.api.ibv_modify_qp(x.qp.?, &rtr, qp_state | qp_av | qp_path_mtu | qp_dest_qpn | qp_rq_psn | qp_max_dest_rd_atomic | qp_min_rnr_timer) != 0) { + g.fail("HCA {d}: the queue pair refused RTR", .{h}); + return error.PeerRefused; + } + var rts: QpAttr = .{ .qp_state = qps_rts, .timeout = 14, .retry_cnt = 7, .rnr_retry = 7, .max_rd_atomic = 1 }; + if (g.api.ibv_modify_qp(x.qp.?, &rts, qp_state | qp_timeout | qp_retry_cnt | qp_rnr_retry | qp_sq_psn | qp_max_qp_rd_atomic) != 0) { + g.fail("HCA {d}: the queue pair refused RTS", .{h}); + return error.PeerRefused; + } + } + g.running.store(true, .release); + g.thread = try std.Thread.spawn(.{}, proxy, .{g}); + } + + fn ctrl(g: *Gather) [*]u32 { + return @ptrCast(@alignCast(g.region.bytes.ptr + g.layout.ctrl)); + } + + fn drain(g: *Gather, h: usize) !void { + const x = &g.hcas[h]; + var wc: [32]Wc = undefined; + const n = opsCall(PollCq, x.ctx.?, ops_poll_cq)(x.cq.?, wc.len, &wc); + if (n < 0) return error.ProxyFailed; + for (wc[0..@intCast(n)]) |w| { + if (w.status != 0) { + const why = g.api.ibv_wc_status_str(w.status) orelse "?"; + g.fail("RDMA write on HCA {d} failed: {s} (vendor 0x{x}, seq {d})", .{ h, std.mem.span(why), w.vendor_err, g.last_seq }); + return error.ProxyFailed; + } + x.outstanding -= 1; + } + } + + /// One sequence's shard to the peer: per HCA a stripe, then its flag (inline, signalled) on the same queue pair. + fn post(g: *Gather, seq: u32, nbytes: u32) !void { + if (nbytes == 0 or nbytes % 16 != 0 or nbytes > g.layout.slot_bytes) { + g.fail("bad shard size {d} at seq {d}", .{ nbytes, seq }); + return error.ProxyFailed; + } + const slot: u64 = seq & 1; + const base = @intFromPtr(g.region.bytes.ptr); + const packs = nbytes / 16; + var seq_copy: u32 = seq; + var offset: u64 = 0; + for (0..g.n_hca) |h| { + const x = &g.hcas[h]; + while (x.outstanding >= send_depth / 4) { + try g.drain(h); + if (!g.running.load(.monotonic)) return error.ProxyFailed; + } + var n = packs / g.n_hca; + if (h < packs % g.n_hca) n += 1; + const bytes: u32 = n * 16; + var flag_sge: Sge = .{ .addr = @intFromPtr(&seq_copy), .length = 4, .lkey = 0 }; + var flag: SendWr = .{ .wr_id = h, .sg_list = &flag_sge, .num_sge = 1, .opcode = wr_rdma_write, .send_flags = send_signaled | send_inline, .remote_addr = g.peer_addr + g.layout.flags + (slot * g.n_hca + h) * flag_stride, .rkey = g.peer_rkey[h] }; + var data_sge: Sge = .{ .addr = base + g.layout.send + slot * g.layout.slot_bytes + offset, .length = bytes, .lkey = x.mr.?.lkey }; + var data: SendWr = .{ .wr_id = h, .next = &flag, .sg_list = &data_sge, .num_sge = 1, .opcode = wr_rdma_write, .remote_addr = g.peer_addr + g.layout.recv + slot * g.layout.slot_bytes + offset, .rkey = g.peer_rkey[h] }; + var bad: ?*SendWr = null; + const first: *SendWr = if (bytes != 0) &data else &flag; + const rc = opsCall(PostSend, x.ctx.?, ops_post_send)(x.qp.?, first, &bad); + if (rc != 0) { + g.fail("ibv_post_send: {d}", .{rc}); + return error.ProxyFailed; + } + x.outstanding += 1; + offset += bytes; + } + for (0..g.n_hca) |h| try g.drain(h); + } + + /// The proxy: spins on the doorbell and posts each sequence in order (a missed one from its slot's size). + fn proxy(g: *Gather) void { + const c = g.ctrl(); + var idle: u64 = 0; + const nap: std.c.timespec = .{ .sec = 0, .nsec = 20_000 }; + while (g.running.load(.monotonic)) { + const seq = @atomicLoad(u32, &c[c_seq], .acquire); + if (seq == g.last_seq) { + idle += 1; + if (idle % 64 == 0) for (0..g.n_hca) |h| g.drain(h) catch { + g.failed.store(true, .release); + return; + }; + if (idle >= idle_spins) _ = std.c.nanosleep(&nap, null); + continue; + } + idle = 0; + // a rank's next gather can ring before this thread saw the last one; peers stay within a slot's distance + var pending = seq -% g.last_seq; + if (pending > slots) { + g.fail("the doorbell skipped {d} sequences (last {d}, now {d})", .{ pending, g.last_seq, seq }); + g.failed.store(true, .release); + return; + } + var s = g.last_seq +% 1; + while (pending > 0) : ({ + s +%= 1; + pending -= 1; + }) { + const nbytes = @atomicLoad(u32, &c[c_nbytes + (s & 1)], .acquire); + g.post(s, nbytes) catch { + g.failed.store(true, .release); + return; + }; + g.last_seq = s; + } + } + } + + /// The cap both ranks test in ``fits``: the smaller of the two, which each rank's slot holds. + fn agreedCap(mine: u64, theirs: u64) usize { + return @intCast(@min(mine, theirs)); + } + + /// Whether a gather of ``bytes`` a rank goes this way: the size alone decides, so both ranks pick the same path. + pub fn fits(g: *const Gather, bytes: usize) bool { + return bytes > 0 and bytes <= g.max_bytes and bytes % 16 == 0; + } + + /// Both ranks' ``bytes`` from ``send`` into ``recv`` in rank order (as ncclAllGather); misaligned is an error. + pub fn gather(g: *Gather, stream: abi.Stream, send: abi.DevicePtr, recv: abi.DevicePtr, bytes: usize) !void { + if (send % 16 != 0 or recv % 16 != 0) return error.Misaligned; + const packs = bytes / 16; + const need = @max(1, std.math.divCeil(usize, packs, 2 * threads) catch unreachable); + const grid: u32 = @intCast(@min(std.math.ceilPowerOfTwo(usize, need) catch max_blocks, max_blocks)); + const klass: i32 = @intCast(std.math.log2_int(u32, grid)); + const base = g.region_dev; + var a: launch.Args = .{}; + a.add(send); + a.add(recv); + a.add(@as(i32, @intCast(packs))); + a.add(@as(i32, g.rank)); + a.add(base + g.layout.recv); + a.add(base + g.layout.flags); + a.add(base + g.layout.send); + a.add(base + g.layout.ctrl); + a.add(@as(i64, @intCast(g.layout.slot_bytes))); + a.add(g.counters.ptr); + a.add(klass); + a.add(spin_limit); + a.add(@as(i32, @intCast(g.n_hca))); + try launch.launch(g.f, .{ .grid = .{ .x = grid }, .block = .{ .x = threads } }, .{ .handle = stream, .d = g.d }, &a); + g.gathers += 1; + } + + /// An error if a wait timed out on the GPU or the proxy failed (the host checks after a synchronize). + pub fn check(g: *Gather) !void { + const c = g.ctrl(); + const bad = @atomicLoad(u32, &c[c_error], .acquire); + if (bad == 0 and !g.failed.load(.acquire)) return; + std.log.err("RoCE all-gather failed (sequence {d}, HCA {d}): {s}", .{ bad, c[c_hca], std.mem.sliceTo(&g.err, 0) }); + return error.ProxyFailed; + } + + fn teardown(g: *Gather) void { + if (g.running.swap(false, .acq_rel)) if (g.thread) |t| t.join(); + g.closeHcas(); + g.module.unload(); + g.counters.free(); + g.region.free(); + g.lib.close(); + } + + pub fn deinit(g: *Gather) void { + g.teardown(); + g.gpa.destroy(g); + } +}; + +/// The HCAs to stripe over, at most two: TF_ROCE_HCA, else NCCL_IB_HCA without "=^" prefixes or ":port" suffixes. +fn hcaNames(io: std.Io, gpa: std.mem.Allocator, names: *[max_hcas][64]u8, lens: *[max_hcas]usize) !u32 { + _ = io; + _ = gpa; + const raw = std.c.getenv("TF_ROCE_HCA") orelse std.c.getenv("NCCL_IB_HCA") orelse return error.NoHca; + var n: u32 = 0; + var it = std.mem.tokenizeScalar(u8, std.mem.span(raw), ','); + while (it.next()) |w| { + if (n == max_hcas) break; + var s = std.mem.trimStart(u8, std.mem.trim(u8, w, " "), "=^"); + if (std.mem.indexOfScalar(u8, s, ':')) |c| s = s[0..c]; + if (s.len == 0 or s.len > 63) continue; + @memcpy(names[n][0..s.len], s); + lens[n] = s.len; + n += 1; + } + if (n == 0) return error.NoHca; + return n; +} + +/// The HCA's RoCE v2 IPv4 GID (TF_ROCE_GID_INDEX overrides): its index differs between hosts' tables. +fn gidIndex(io: std.Io, gpa: std.mem.Allocator, name: []const u8) !c_int { + if (std.c.getenv("TF_ROCE_GID_INDEX")) |v| return std.fmt.parseInt(c_int, std.mem.span(v), 10) catch error.NoHca; + const cwd = std.Io.Dir.cwd(); + var i: c_int = 0; + while (i < 16) : (i += 1) { + var pbuf: [160]u8 = undefined; + const tpath = try std.fmt.bufPrint(&pbuf, "/sys/class/infiniband/{s}/ports/1/gid_attrs/types/{d}", .{ name, i }); + const t = cwd.readFileAlloc(io, tpath, gpa, .limited(64)) catch continue; + defer gpa.free(t); + if (!std.mem.startsWith(u8, t, "RoCE v2")) continue; + const gpath = try std.fmt.bufPrint(&pbuf, "/sys/class/infiniband/{s}/ports/1/gids/{d}", .{ name, i }); + const gid = cwd.readFileAlloc(io, gpath, gpa, .limited(64)) catch continue; + defer gpa.free(gid); + if (std.mem.startsWith(u8, gid, "0000:0000:0000:0000:0000:ffff:")) return i; + } + return error.NoHca; +} + +test "RoCE gather: both ranks agree one cap" { + // 511 KiB and 512 KiB round to the same slot; a 512 KiB shard must take the same path on both ranks + try std.testing.expectEqual(Layout.of(511 << 10).slot_bytes, Layout.of(512 << 10).slot_bytes); + const a = Gather.agreedCap(511 << 10, 512 << 10); + const b = Gather.agreedCap(512 << 10, 511 << 10); + try std.testing.expectEqual(a, b); + try std.testing.expectEqual(@as(usize, 511 << 10), a); +} diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 6c4e0991b6..81adb2529d 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -21,6 +21,7 @@ pub const Dim3 = launch.Dim3; pub const graph = @import("graph.zig"); pub const cublaslt = @import("cublaslt.zig"); pub const nccl = @import("nccl.zig"); +pub const roce = @import("roce.zig"); pub const triton = @import("triton.zig"); pub const aot = @import("aot.zig"); pub const kernels = @import("kernels.zig"); @@ -51,6 +52,7 @@ test { _ = @import("vit.zig"); _ = @import("experts.zig"); _ = @import("sample.zig"); + _ = @import("roce.zig"); _ = @import("rank_scatter.zig"); _ = launch; _ = abi; diff --git a/zig/src/cuda/tp_group.zig b/zig/src/cuda/tp_group.zig index 7f26c7c846..de2b39948f 100644 --- a/zig/src/cuda/tp_group.zig +++ b/zig/src/cuda/tp_group.zig @@ -1,9 +1,10 @@ -//! A tensor-parallel group over NCCL: NCCL_COMM_ID gives every rank the same id; rows meet in rank order. +//! A tensor-parallel group over NCCL (NCCL_COMM_ID; rows in rank order); two ranks add RoCE small gathers (roce.zig). const std = @import("std"); const nccl = @import("nccl.zig"); const driver = @import("driver.zig"); const stream_ = @import("stream.zig"); const memory = @import("memory.zig"); +const roce = @import("roce.zig"); const Stream = stream_.Stream; @@ -12,6 +13,7 @@ pub const Group = struct { comm: nccl.Comm, rank: usize, world: usize, + fast: ?*roce.Gather = null, // two ranks' small all-gathers by RoCE writes (openRoce), when both ranks brought it up /// Rank `rank` of `world` (NCCL opened at run time; NCCL_COMM_ID names rank 0's address, every rank the same). pub fn open(rank: usize, world: usize) !Group { @@ -26,13 +28,26 @@ pub const Group = struct { } pub fn close(g: *Group) void { + if (g.fast) |f| f.deinit(); _ = g.lib.api.ncclCommDestroy(g.comm); g.lib.close(); g.* = undefined; } + /// Two ranks: gathers up to `max_bytes` a rank go by RoCE writes if both ranks bring it up; else NCCL keeps them. + pub fn openRoce(g: *Group, gpa: std.mem.Allocator, io: std.Io, d: *const driver.Driver, s: Stream, max_bytes: usize) !void { + if (g.world != 2 or g.fast != null) return; + g.fast = try roce.Gather.init(gpa, io, d, g, s, roce.wanted(), max_bytes); + } + + /// A failed RoCE gather's error (its rows hold nothing valid): call it after a synchronize that reads a gather. + pub fn check(g: *const Group) !void { + if (g.fast) |f| try f.check(); + } + /// `bytes` at `mine` from every rank into `all` in rank order (world * bytes there), on `s`. pub fn allGather(g: *const Group, s: Stream, mine: u64, all: u64, bytes: usize) !void { + if (g.fast) |f| if (f.fits(bytes)) return f.gather(s.handle, mine, all, bytes); try g.lib.check(g.lib.api.ncclAllGather(mine, all, bytes, .u8, g.comm, s.handle), "ncclAllGather"); } @@ -68,6 +83,7 @@ pub const Group = struct { try dev.uploadAsync(g.rank * size, std.mem.asBytes(&mine), s.handle); try g.allGather(s, dev.ptr + g.rank * size, dev.ptr, size); try s.synchronize(); + try g.check(); try dev.download(0, std.mem.sliceAsBytes(out)); } }; diff --git a/zig/src/families/glm/config.zig b/zig/src/families/glm/config.zig index 2b23df7bc4..f7e9cd27c3 100644 --- a/zig/src/families/glm/config.zig +++ b/zig/src/families/glm/config.zig @@ -14,6 +14,7 @@ pub const Config = struct { run: u32 = 45, // backbone layers loaded and run: all, or the first few for a check that must fit (GLM_LAYERS) vocab: u32 = 154880, head_unit: u32 = 64, // the vocabulary's TP cut in whole units: 128 for EXL3, whose matrices split at 128 columns + exl3: bool = false, // routed experts in EXL3 (quant_method "exl3", Mia-AiLab's checkpoint), the rest bf16 eps: f32 = 1e-5, dense_layers: u32 = 3, dense_inter: u32 = 12288, // this Mac's dense MLP rows: all, or its half in TP2 @@ -151,6 +152,7 @@ pub fn parse(gpa: std.mem.Allocator, json: []const u8) !Config { c.vocab = @intCast(try int(t, "vocab_size")); if (root.get("quantization_config")) |q| if (q == .object) if (q.object.get("quant_method")) |m| if (m == .string and std.mem.eql(u8, m.string, "exl3")) { c.head_unit = 128; + c.exl3 = true; }; c.eps = @floatCast(try float(t, "rms_norm_eps")); c.dense_layers = @intCast(try int(t, "first_k_dense_replace")); diff --git a/zig/src/families/glm/cuda.zig b/zig/src/families/glm/cuda.zig index af9254de61..0c77e40d21 100644 --- a/zig/src/families/glm/cuda.zig +++ b/zig/src/families/glm/cuda.zig @@ -1,6 +1,13 @@ //! GLM-5.3-Flash on CUDA: the family's parts, ported from the Python family's CUDA path. pub const experts = @import("cuda_experts.zig"); +pub const exl3 = @import("cuda_exl3.zig"); +pub const vision = @import("cuda_vision.zig"); +pub const mm = @import("cuda_mm.zig"); +pub const vision_config = @import("vision_config.zig"); +pub const image_core = @import("image_core.zig"); +/// An image's bytes to the tower's patches, as the Metal family prepares them (core's decoder). +pub const image = image_core.For(@import("core").picture); pub const glue = @import("cuda_glue.zig"); pub const kda = @import("cuda_kda.zig"); pub const b16 = @import("cuda_b16.zig"); @@ -25,6 +32,10 @@ pub const native = @import("cuda_native.zig"); test { _ = experts; + _ = exl3; + _ = vision; + _ = mm; + _ = image_core; _ = glue; _ = kda; _ = b16; diff --git a/zig/src/families/glm/cuda_b16.zig b/zig/src/families/glm/cuda_b16.zig index d935b26a64..0509b459d2 100644 --- a/zig/src/families/glm/cuda_b16.zig +++ b/zig/src/families/glm/cuda_b16.zig @@ -3,6 +3,7 @@ const std = @import("std"); const cuda = @import("cuda"); const glue = @import("cuda_glue.zig"); +const mm = @import("cuda_mm.zig"); /// qmm.SPLIT_TARGET, B16_SPLIT_COLS, B16_BK, geometry.MLA_B16_SLICES and MLA_B16_ROWS_FROM. pub const split_target = 192; @@ -55,6 +56,7 @@ pub fn partBytes(m: usize, n: usize, k: usize) usize { /// qmm.matmul for a B16 weight: out [m][n] bf16 (fp32 with `wide`) = x . w; 128+ rows sum in registers, else `part`. pub fn matmul(set: glue.Set, stream: cuda.Stream, x: u64, x_stride: usize, w: u64, out: u64, part: u64, m: usize, n: usize, k: usize, wide: bool) !void { if (k % bk != 0) return error.KNotMultipleOf64; + if (set.mm) |mk| return mm.matmul(mk, stream, x, x_stride, .{ .w = w, .n = n, .k = k }, out, part, m, wide); // tiled weights const sk = splitK(n, k); var a: cuda.Args = .{}; a.add(x); diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index ee38e00b3e..f9e92556e0 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -11,11 +11,15 @@ const ranks = @import("cuda_ranks.zig"); const sampler = @import("cuda_sampler.zig"); const heat = @import("heat"); const modelopt = @import("core").modelopt; +const exl3 = @import("cuda_exl3.zig"); +const images_mod = @import("cuda_images.zig"); +const prof_mod = @import("cuda_prof.zig"); pub const Options = struct { capacity: usize, // tokens a sequence holds mtp: bool, // load the MTP head (drafts) settings: ranks.Settings, + vision: u32 = 0, // --vision's image tokens (0: none): every rank must hold the rows rank 0 sends }; /// Rank 0's own word on its request (cancelled or not); the other ranks have none and hear rank 0's. @@ -23,6 +27,8 @@ pub const Cancel = struct { ptr: ?*anyopaque, check: *const fn (?*anyopaque) boo /// The most drafts a chain holds (decode.DepthPolicy's top depth). pub const max_drafts = 3; +/// A greedy pick as it crosses ranks: (value, id) and padding, so the RoCE gather (16-byte packs) can carry it. +const pair_bytes = 16; /// A prompt pass's chunk rows (decode.prefill's). pub const prompt_rows = fw.prefill_rows; @@ -63,6 +69,10 @@ pub const Engine = struct { spans: []shard.Span, // every rank's vocabulary span lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, + group: ?*const cuda.tp_group.Group = null, // the main stream's all-gathers (RoCE on two ranks); null: lib and comm + images: ?*images_mod.Images = null, // --vision: a prompt pass's image rows (rank 0 encodes, the others receive) + profs: [fw.window_rows]?*prof_mod.Prof = @splat(null), // TF_GLM_PROF=1: verify sections' GPU time, by rows + prompt_profs: [4]?*prof_mod.Prof = @splat(null), // and prompt chunks': to 64, 256, 1024, 2048 rows settings: ranks.Settings, m: fw.Modules, stream: cuda.Stream, @@ -98,7 +108,8 @@ pub const Engine = struct { const cfg_path = try std.fs.path.join(a, &.{ checkpoint, "config.json" }); const text = try std.Io.Dir.cwd().readFileAlloc(io, cfg_path, a, .limited(1 << 24)); var why: modelopt.Why = .{}; - modelopt.admit(gpa, text, &why) catch |err| { // a quant_algo other than NVFP4 is refused before anything loads + const admitted = if (exl3.isExl3(gpa, text)) exl3.admit(gpa, text, &why) else modelopt.admit(gpa, text, &why); + admitted catch |err| { // anything but ModelOpt NVFP4 or EXL3 (4-bit mcg routed experts) is refused before loading std.log.err("{s}: {s}", .{ cfg_path, why.text() }); return err; }; @@ -107,7 +118,7 @@ pub const Engine = struct { errdefer stream.deinit(); const config_hash = try ranks.hashFile(a, io, cfg_path); const program_hash = try ranks.hashFile(a, io, "/proc/self/exe"); - const fp = ranks.Fingerprint.of( + var fp = ranks.Fingerprint.of( world, rank, o.settings, @@ -117,6 +128,7 @@ pub const Engine = struct { program_hash, checkpoint, ); + fp.vision = o.vision; try ranks.agree(d, lib, comm, stream, fp, try a.alloc(ranks.Fingerprint, world)); const share = try shard.Share.of(&c, rank, world); const spans = try gpa.alloc(shard.Span, world); @@ -131,6 +143,10 @@ pub const Engine = struct { for (layers, 0..) |*l, i| l.* = @intCast(i); e.* = .{ .gpa = gpa, .io = io, .heat_gate = .{ .bands = o.settings.heat, .root = heatRoot() }, .shared_buf = undefined, .d = d, .c = c, .share = share, .spans = spans, .lib = lib, .comm = comm, .settings = o.settings, .m = m, .stream = stream, .pieces = pieces, .w = undefined, .capacity = o.capacity, .bp = undefined, .bd = undefined, .bm = undefined, .hold = undefined, .pairs = undefined, .smp = undefined, .mtp = o.mtp }; e.w = try weights.load(gpa, io, d, stream, e.m.tools(), checkpoint, c, share, layers, o.mtp); + if (tiledLinears()) { // our bf16 linears (glm_mm) on their tiled weights; TF_GLM_MM=0 keeps the glue's + try weights.tileLinears(&e.w, &e.m.mm, stream, d); + e.m.tiled = true; + } errdefer e.w.deinit(); try e.bind(0); errdefer { @@ -147,9 +163,17 @@ pub const Engine = struct { errdefer if (e.bs) |*b| b.deinit(); e.hold = try cuda.DeviceBuffer.alloc(d, c.hidden * 2); errdefer e.hold.free(); - e.pairs = try cuda.DeviceBuffer.alloc(d, 8 * world); + e.pairs = try cuda.DeviceBuffer.alloc(d, pair_bytes * world); errdefer e.pairs.free(); + try d.check(d.api.cuMemsetD8_v2(e.pairs.ptr, 0, pair_bytes * world), "cuMemsetD8"); // each pair's padding stays zero e.smp = try sampler.Sampler.init(d, c.vocab, fw.window_rows); + if (prof_mod.Prof.wanted()) for (&e.profs) |*p| { + p.* = try prof_mod.Prof.init(d, io); + }; + if (prof_mod.Prof.wanted()) for (&e.prompt_profs) |*p| { + p.* = try prof_mod.Prof.init(d, io); + p.*.?.every = 4; + }; errdefer e.smp.deinit(); e.shared_buf = try cuda.DeviceBuffer.alloc(d, 8); try e.stream.synchronize(); @@ -182,12 +206,13 @@ pub const Engine = struct { } fn exchange(e: *const Engine) fw.Exchange { - return .{ .nccl = .{ .lib = e.lib, .comm = e.comm } }; + return .{ .nccl = .{ .lib = e.lib, .comm = e.comm, .group = e.group } }; } fn view(e: *Engine, b: *fw.Buffers, out: *[1]fw.Rank) fw.Forward { out.* = .{.{ .w = &e.w, .st = &e.cur.st, .b = b }}; - return .{ .d = e.d, .m = &e.m, .stream = e.stream, .ranks = out, .exchange = e.exchange(), .pieces = if (b == &e.bp) &e.pieces else null, .split = e.settings.split, .kda_wide = e.settings.kda_wide }; + const im: ?fw.ImageRows = if (e.images) |v| .{ .v = v, .slot = e.slot } else null; + return .{ .d = e.d, .m = &e.m, .stream = e.stream, .ranks = out, .exchange = e.exchange(), .pieces = if (b == &e.bp) &e.pieces else null, .split = e.settings.split, .kda_wide = e.settings.kda_wide, .images = im }; } /// Binds slot `slot`'s sequence for the steps that follow, making it (a zeroed state) on first use. @@ -228,6 +253,12 @@ pub const Engine = struct { return e.cur.st.pos; } + /// A follower's begin: rank 0's next request; no image rows from an earlier one splice into it. + pub fn beginFollowing(e: *Engine, sampling: ?sampler.Sampling) !void { + if (e.images) |v| v.clear(); + return e.begin(sampling); + } + /// A fresh sequence with `sampling` (null or temperature 0: greedy). pub fn begin(e: *Engine, sampling: ?sampler.Sampling) !void { e.cur.sampling = sampler.check(sampling); @@ -259,10 +290,15 @@ pub const Engine = struct { const rows = @min(len - start, fw.prefill_rows); if (part.len != @min(len - start, rows + 1)) return error.BadChunk; var one: [1]fw.Rank = undefined; - const f = e.view(&e.bp, &one); + var f = e.view(&e.bp, &one); const as_i32: []const i32 = @ptrCast(part); if (marks) |m| if (start > 0 and std.mem.indexOfScalar(u32, m.at, @intCast(start)) != null) m.call(m.ctx, @intCast(start)); + f.prof = e.prompt_profs[if (rows <= 64) 0 else if (rows <= 256) 1 else if (rows <= 1024) 2 else 3]; + const labels = [_][]const u8{ "prompt chunk to 64 rows", "prompt chunk to 256 rows", "prompt chunk to 1024 rows", "prompt chunk to 2048 rows" }; + if (f.prof) |p| try p.begin(e.stream, labels[if (rows <= 64) 0 else if (rows <= 256) 1 else if (rows <= 1024) 2 else 3]); + defer if (f.prof) |p| p.end(); try e.gate(cancel); + if (f.prof) |p| p.mark(e.stream, .gate); try f.run(as_i32[0..rows]); const last = start + rows == len; var first: ?u32 = null; @@ -271,15 +307,21 @@ pub const Engine = struct { try e.draw(e.bp.logits, 1, len, &t); first = t[0]; } + if (f.prof) |p| p.mark(e.stream, .draw); if (e.mtp) { if (last) try e.copy(e.hold.ptr, e.bp.fnormed + (rows - 1) * e.c.hidden * 2, e.c.hidden * 2); const nxt = as_i32[1..]; if (nxt.len > 0) { // decode._absorb_rows + const mp = f.prof; + f.prof = null; // the head's pass as one section try f.mtp(nxt, &.{e.bp.fnormed}); + f.prof = mp; try e.cur.st.setMtpLen(e.stream, e.cur.st.mtp_len + nxt.len); } } + if (f.prof) |p| p.mark(e.stream, .mtp); try f.commit(rows); + if (f.prof) |p| p.mark(e.stream, .commit); return first; } @@ -394,11 +436,19 @@ pub const Engine = struct { pub fn verify(e: *Engine, ids: []const u32, out: []u32) !void { if (ids.len == 0 or ids.len > fw.window_rows) return error.WindowTooWide; var one: [1]fw.Rank = undefined; - const f = e.view(&e.bd, &one); + var f = e.view(&e.bd, &one); const at = e.cur.st.pos; + f.prof = e.profs[ids.len - 1]; + if (f.prof) |p| try p.begin(e.stream, verifyLabel(ids.len)); try f.run(@ptrCast(ids)); e.cur.rows_at = e.bd.fnormed; try e.draw(e.bd.logits, ids.len, at + 1, out); + if (f.prof) |p| p.end(); + } + + fn verifyLabel(rows: usize) []const u8 { + const l = [_][]const u8{ "verify 1 row", "verify 2 rows", "verify 3 rows", "verify 4 rows", "verify 5 rows", "verify 6 rows", "verify 7 rows", "verify 8 rows" }; + return l[@min(rows, l.len) - 1]; } /// Several streams' windows in one forward: shared row ops, per-part attention and KDA. @@ -462,7 +512,10 @@ pub const Engine = struct { /// `rows` logits rows of this rank's vocabulary -> the tokens at positions first, first + 1, ... fn draw(e: *Engine, logits: u64, rows: usize, first: u64, out: []u32) !void { - if (e.cur.sampling) |s| return e.smp.draw(e.stream, e.exchange(), e.spans, &.{.{ .rank = e.share.rank, .logits = logits }}, rows, first, sampler.rule(s), out); + if (e.cur.sampling) |s| { + try e.smp.draw(e.stream, e.exchange(), e.spans, &.{.{ .rank = e.share.rank, .logits = logits }}, rows, first, sampler.rule(s), out); + return e.checkGathers(); // the draw read rows the forward's gathers made + } for (0..rows) |r| out[r] = try e.greedy(logits + r * e.share.vocab.len() * 2); } @@ -473,12 +526,14 @@ pub const Engine = struct { a.add(logits); a.add(@as(i64, @intCast(e.share.vocab.len()))); a.add(@as(u32, @intCast(e.share.vocab.lo))); - a.add(e.pairs.ptr + e.share.rank * 8); + a.add(e.pairs.ptr + e.share.rank * pair_bytes); try cuda.launch.launch(e.m.glue.top1, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1024 } }, e.stream, &a); - try e.lib.check(e.lib.api.ncclAllGather(e.pairs.ptr + e.share.rank * 8, e.pairs.ptr, 2, .u32, e.comm, e.stream.handle), "ncclAllGather"); + const mine = e.pairs.ptr + e.share.rank * pair_bytes; + if (e.group) |g| try g.allGather(e.stream, mine, e.pairs.ptr, pair_bytes) else try e.lib.check(e.lib.api.ncclAllGather(mine, e.pairs.ptr, pair_bytes, .u8, e.comm, e.stream.handle), "ncclAllGather"); try e.stream.synchronize(); - var all: [ranks.max_world][2]u32 = undefined; // Engine.init refuses a larger world - try e.d.check(e.d.api.cuMemcpyDtoH_v2(&all, e.pairs.ptr, world * 8), "cuMemcpyDtoH"); + try e.checkGathers(); + var all: [ranks.max_world][4]u32 = undefined; // Engine.init refuses a larger world + try e.d.check(e.d.api.cuMemcpyDtoH_v2(&all, e.pairs.ptr, world * pair_bytes), "cuMemcpyDtoH"); var pick = all[0]; for (all[1..world]) |p| { const v: f32 = @bitCast(p[0]); @@ -488,6 +543,17 @@ pub const Engine = struct { return pick[1]; } + /// A failed RoCE gather's error, after a synchronize whose result read gathered rows (they hold nothing then). + fn checkGathers(e: *const Engine) !void { + if (e.group) |g| try g.check(); + } + + /// TF_GLM_MM=0 keeps the bf16 linears on the glue's b16 kernels (the weights stay as the checkpoint stores them). + fn tiledLinears() bool { + const v = std.c.getenv("TF_GLM_MM") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); + } + fn heatRoot() []const u8 { const p = std.c.getenv("TF_HEAT_ROOT") orelse return heat.default_root; const s = std.mem.span(p); diff --git a/zig/src/families/glm/cuda_exl3.zig b/zig/src/families/glm/cuda_exl3.zig new file mode 100644 index 0000000000..05e380a7fa --- /dev/null +++ b/zig/src/families/glm/cuda_exl3.zig @@ -0,0 +1,200 @@ +//! GLM-5.3-Flash's EXL3 routed experts on CUDA (Mia-AiLab's checkpoint: 4-bit mcg trellis experts, the rest bf16). + +const std = @import("std"); +const modelopt = @import("core").modelopt; +const Span = @import("cuda_shard.zig").Span; +const cuda = @import("cuda"); +const experts = @import("cuda_experts.zig"); + +/// config.json's quantization_config (top level, else text_config's), when it has one. +fn block(root: std.json.Value) ?std.json.ObjectMap { + if (root != .object) return null; + if (root.object.get("quantization_config")) |q| if (q == .object) return q.object; + if (root.object.get("text_config")) |t| if (t == .object) if (t.object.get("quantization_config")) |q| if (q == .object) return q.object; + return null; +} + +/// Whether config.json names EXL3 (quant_method "exl3"): `admit` then says whether this family reads it. +pub fn isExl3(gpa: std.mem.Allocator, config_json: []const u8) bool { + const parsed = std.json.parseFromSlice(std.json.Value, gpa, config_json, .{}) catch return false; + defer parsed.deinit(); + const q = block(parsed.value) orelse return false; + const m = q.get("quant_method") orelse return false; + return m == .string and std.mem.eql(u8, m.string, "exl3"); +} + +/// The EXL3 checkpoints the kernels read: 4-bit routed experts on the "mcg" codebook (everything else bf16). +pub fn admit(gpa: std.mem.Allocator, config_json: []const u8, why: *modelopt.Why) !void { + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, config_json, .{}); + defer parsed.deinit(); + const q = block(parsed.value) orelse { + why.set("no quantization_config", .{}); + return error.UnsupportedQuant; + }; + const bits = q.get("bits") orelse std.json.Value{ .null = {} }; + if (bits != .integer or bits.integer != 4) { + why.set("EXL3 at {f} bits: the CUDA family reads 4-bit EXL3 experts", .{std.json.fmt(bits, .{})}); + return error.UnsupportedQuant; + } + const book = q.get("codebook") orelse std.json.Value{ .null = {} }; + if (book != .string or !std.mem.eql(u8, book.string, "mcg")) { + why.set("EXL3 codebook {f}: the CUDA family reads \"mcg\"", .{std.json.fmt(book, .{})}); + return error.UnsupportedQuant; + } +} + +/// One layer's EXL3 experts on a rank: trellis words [E][K/16][N/16][32] and fp16 scales (suh [E][K], svh [E][N]). +pub const Experts = struct { + gt: u64, + ut: u64, + dt: u64, + suh_g: u64, + suh_u: u64, + svh_g: u64, + svh_u: u64, + suh_d: u64, + svh_d: u64, + shx: bool, // every expert's gate and up share one input rotation: a prompt rotates a row once, not once a pick +}; + +const tile = 128; // a 16 x 16 tile's bytes at 4 bits + +/// Layer `i`'s EXL3 experts through `L` (cuda_weights' loader): gate, up by `cut`'s tile columns, down by its rows. +pub fn load(L: anytype, i: u32, count: usize, dims: usize, inter: usize, cut: Span) !Experts { + if (dims % 128 != 0 or cut.lo % 128 != 0 or cut.hi % 128 != 0) return error.NotExl3Mcg4; + const width = cut.len(); + var out: Experts = undefined; + out.shx = true; + const Mat = struct { k: usize, n: usize, by_cols: bool }; + const mats = [_]Mat{ + .{ .k = dims, .n = inter, .by_cols = true }, + .{ .k = dims, .n = inter, .by_cols = true }, + .{ .k = inter, .n = dims, .by_cols = false }, + }; + var first_suh: ?[]const u8 = null; + inline for (mats, .{ "gate_proj", "up_proj", "down_proj" }, .{ "gt", "ut", "dt" }, .{ "suh_g", "suh_u", "suh_d" }, .{ "svh_g", "svh_u", "svh_d" }) |m, proj, tw, su, sv| { + const kt = (if (m.by_cols) m.k else width) / 16; // the rank's tiles + const nt = (if (m.by_cols) width else m.n) / 16; + for (0..count) |e| { + const t = try L.get("layers.{d}.mlp.experts.{d}." ++ proj ++ ".trellis", .{ i, e }); + if (t.dtype != .i16 or t.rank != 3 or t.dim(0) != m.k / 16 or t.dim(1) != m.n / 16 or t.dim(2) != 64) return error.NotExl3Mcg4; + _ = L.get("layers.{d}.mlp.experts.{d}." ++ proj ++ ".mcg", .{ i, e }) catch return error.NotExl3Mcg4; + if (m.by_cols) { // every tile row, the rank's tile columns + const row = (m.n / 16) * tile; + for (0..kt) |r| try L.host.appendSlice(L.gpa, t.bytes[r * row + (cut.lo / 16) * tile ..][0 .. nt * tile]); + } else try L.host.appendSlice(L.gpa, t.bytes[(cut.lo / 16) * nt * tile ..][0 .. kt * nt * tile]); // its tile rows + } + @field(out, tw) = try L.put("layers.{d}.moe.exl3." ++ tw, .{i}); + for (0..count) |e| { // suh: the input's scales (gate, up: every input; down: the rank's) + const t = try L.get("layers.{d}.mlp.experts.{d}." ++ proj ++ ".suh", .{ i, e }); + if (t.dtype != .f16 or t.numel() != m.k) return error.NotExl3Mcg4; + const mine = if (m.by_cols) t.bytes else t.bytes[cut.lo * 2 .. cut.hi * 2]; + if (m.by_cols) { // one rotation shared by the layer's every gate and up input? + if (first_suh) |f| (if (!std.mem.eql(u8, f, mine)) { + out.shx = false; + }) else first_suh = mine; + } + try L.host.appendSlice(L.gpa, mine); + } + @field(out, su) = try L.put("layers.{d}.moe.exl3." ++ su, .{i}); + for (0..count) |e| { // svh: the output's scales (gate, up: the rank's; down: every output) + const t = try L.get("layers.{d}.mlp.experts.{d}." ++ proj ++ ".svh", .{ i, e }); + if (t.dtype != .f16 or t.numel() != m.n) return error.NotExl3Mcg4; + try L.host.appendSlice(L.gpa, if (m.by_cols) t.bytes[cut.lo * 2 .. cut.hi * 2] else t.bytes); + } + @field(out, sv) = try L.put("layers.{d}.moe.exl3." ++ sv, .{i}); + } + return out; +} + +/// glm_exl3's kernels. +pub const Kernels = struct { + rot_rows: cuda.Function, + rot_in: cuda.Function, + grouped: cuda.Function, + gateup: cuda.Function, + down: cuda.Function, + prompt_gu: cuda.Function, + prompt_dn: cuda.Function, + picks: cuda.Function, + + pub fn load(m: cuda.Module) !Kernels { + var k: Kernels = undefined; + inline for (@typeInfo(Kernels).@"struct".field_names) |name| @field(k, name) = try m.function("glm_exl3_" ++ name); + return k; + } +}; + +/// Scratch beside the MoE block's (its eact holds the down inputs): 8-slot picks, rotated inputs, K-split partials. +pub const Scratch = struct { picks: u64, xg: u64, xu: u64, z: u64 }; + +/// The picks a row holds (the kernels' pair = row * 8 + slot). +pub const top_k = 8; +/// Gate and up K slices in a decode window (down: 1); prompt chunks take the fused kernels, no partials. +pub const decode_splits = 4; + +/// Bytes of each Scratch buffer for windows of `rows` rows: (picks, xg, xu, z); `shx`: one input rotation a layer. +pub fn scratchBytes(rows: usize, dims: usize, width: usize, prefill: bool, shx: bool) [4]usize { + const pairs = rows * top_k; + const one = prefill and shx; // a prompt rotates each row once + return .{ + pairs * 4, + (if (one) rows else pairs) * dims * 2, + if (one) 16 else pairs * dims * 2, + if (prefill) 16 else @max(2 * decode_splits * pairs * width, pairs * dims) * 4, + }; +} + +/// One launch's arguments, then the launch. +fn go(f: cuda.Function, stream: cuda.Stream, grid: cuda.Dim3, threads: u32, args: anytype) !void { + var a: cuda.Args = .{}; + inline for (args) |v| a.add(v); + try cuda.launch.launch(f, .{ .grid = grid, .block = .{ .x = threads } }, stream, &a); +} + +/// Routed slots of `rows` rows (picks [rows][slots], shared last) into `ey` [rows][slots][dims] fp32, shared untouched. +pub fn run(k: *const Kernels, ek: *const experts.Kernels, stream: cuda.Stream, w: Experts, sc: Scratch, plan_buf: cuda.grouped.Plan, x: u64, x_stride: usize, picks: u64, slots: usize, rows: usize, count: usize, dims: usize, width: usize, limit: f32, prompt: bool, ey: u64, xd: u64) !void { + if (slots != top_k + 1) return error.UnsupportedTopK; + if (dims % 256 != 0 or width % 128 != 0) return error.NotExl3Mcg4; + const P = rows * top_k; + const D: i32 = @intCast(dims); + const N: i32 = @intCast(width); + const ys: i32 = @intCast(slots); + try go(k.picks, stream, .{ .x = @intCast((P + 255) / 256) }, 256, .{ picks, sc.picks, @as(i32, @intCast(rows)), @as(i32, @intCast(slots)) }); + const plan: experts.Plan = .{ .route = plan_buf, .slots = top_k, .experts = count, .prompt = prompt }; + try experts.route(ek, stream, sc.picks, P, plan); + const items: u32 = @intCast(experts.maxItems(P, count, if (prompt) experts.staged_tile else experts.tile)); + const r = plan_buf; + const one = prompt and w.shx; // the prompt kernels read row pair >> 3: one rotated row a token + if (one) { + try go(k.rot_rows, stream, .{ .x = @intCast(rows), .y = @intCast(dims / 128) }, 32, .{ x, @as(i32, @intCast(x_stride)), w.suh_g, sc.xg, D }); + } else { + try go(k.rot_in, stream, .{ .x = @intCast(P), .y = @intCast(dims / 128), .z = 2 }, 32, .{ x, @as(i32, @intCast(x_stride)), sc.picks, w.suh_g, w.suh_u, sc.xg, sc.xu, D }); + } + if (prompt) { // gate and up with their epilogue into xd; down with its epilogue into ey + try go(k.prompt_gu, stream, .{ .x = @intCast(width / 128), .y = items }, 256, .{ sc.xg, if (one) sc.xg else sc.xu, w.gt, w.ut, r.items, r.counts, r.members, w.svh_g, w.svh_u, w.suh_d, xd, D, N, @as(i32, if (one) 3 else 0), limit }); + return go(k.prompt_dn, stream, .{ .x = @intCast(dims / 256), .y = items }, 256, .{ xd, w.dt, r.items, r.counts, r.members, w.svh_d, ey, N, D, ys }); + } + const sk: i32 = decode_splits; + const nt = 8; // n tiles a block + try go(k.grouped, stream, .{ .x = items, .y = @intCast(width / (16 * nt)), .z = @intCast(2 * sk) }, 128, .{ sc.xg, sc.xu, w.gt, w.ut, r.items, r.counts, r.members, sc.z, D, N, @as(i32, @intCast(P)), sk }); + try go(k.gateup, stream, .{ .x = @intCast(P), .y = @intCast(width / 128) }, 32, .{ sc.z, sc.picks, w.svh_g, w.svh_u, w.suh_d, xd, @as(i32, @intCast(P)), N, sk, limit }); + try go(k.grouped, stream, .{ .x = items, .y = @intCast(dims / (16 * nt)), .z = 1 }, 128, .{ xd, xd, w.dt, w.dt, r.items, r.counts, r.members, sc.z, N, D, @as(i32, @intCast(P)), @as(i32, 1) }); + try go(k.down, stream, .{ .x = @intCast(P), .y = @intCast(dims / 128) }, 32, .{ sc.z, sc.picks, w.svh_d, ey, @as(i32, @intCast(P)), D, @as(i32, 1), ys }); +} + +test "EXL3 admission: Mia-AiLab's block passes; other widths and codebooks are refused by name" { + const gpa = std.testing.allocator; + const mia = + \\{"quantization_config": {"quant_method": "exl3", "bits": 4, "codebook": "mcg", "head_bits": 16}} + ; + try std.testing.expect(isExl3(gpa, mia)); + var why: modelopt.Why = .{}; + try admit(gpa, mia, &why); + try std.testing.expect(!isExl3(gpa, "{\"quantization_config\": {\"quant_method\": \"modelopt\"}}")); + try std.testing.expect(!isExl3(gpa, "{}")); + try std.testing.expectError(error.UnsupportedQuant, admit(gpa, "{\"quantization_config\": {\"quant_method\": \"exl3\", \"bits\": 3, \"codebook\": \"mcg\"}}", &why)); + try std.testing.expectEqualStrings("EXL3 at 3 bits: the CUDA family reads 4-bit EXL3 experts", why.text()); + try std.testing.expectError(error.UnsupportedQuant, admit(gpa, "{\"text_config\": {\"quantization_config\": {\"quant_method\": \"exl3\", \"bits\": 4, \"codebook\": \"mul1\"}}}", &why)); + try std.testing.expectEqualStrings("EXL3 codebook \"mul1\": the CUDA family reads \"mcg\"", why.text()); +} diff --git a/zig/src/families/glm/cuda_follow.zig b/zig/src/families/glm/cuda_follow.zig index 8072b7cfe1..1f31bad634 100644 --- a/zig/src/families/glm/cuda_follow.zig +++ b/zig/src/families/glm/cuda_follow.zig @@ -3,10 +3,12 @@ const std = @import("std"); const cuda = @import("cuda"); const engine = @import("cuda_engine.zig"); +const images_mod = @import("cuda_images.zig"); +const vcfg = @import("vision_config.zig"); const sampler = @import("cuda_sampler.zig"); const Engine = engine.Engine; -pub const Kind = enum(u32) { begin = 1, prefill, verify, keep, draft, end, drop, release, chunk, shared }; +pub const Kind = enum(u32) { begin = 1, prefill, verify, keep, draft, end, drop, release, chunk, shared, image }; /// One step as it crosses: 64 bytes, then `n` u32 ids. pub const Header = extern struct { @@ -88,6 +90,12 @@ pub const Link = struct { } } + /// `bytes` at `ptr` from rank 0 into every rank's `ptr` (image rows: the follower's buffer is the same size). + pub fn rows(l: *Link, ptr: u64, bytes: usize) !void { + if (bytes == 0) return; + try l.lib.check(l.lib.api.ncclBroadcast(ptr, ptr, bytes, .u8, 0, l.comm, l.stream.handle), "ncclBroadcast"); + } + /// Other ranks: the next step and its ids (valid until the next receive). pub fn receive(l: *Link) !struct { h: Header, ids: []const u32 } { try l.bcast(64); @@ -153,6 +161,16 @@ pub const Leader = struct { try x.e.begin(c); } + /// The prompt pass's image rows (rank 0 encoded them for the bound slot) to every rank: spans, then the rows. + pub fn images(x: Leader, v: *const images_mod.Images) !void { + try x.flush(); + var words: [3 * 64]u32 = undefined; + if (v.spans.items.len * 3 > words.len) return error.TooManyImages; + for (v.spans.items, 0..) |sp, i| words[3 * i ..][0..3].* = .{ sp.at, sp.n, sp.row }; + try x.send(.{ .kind = @backingInt(Kind.image) }, words[0 .. 3 * v.spans.items.len]); + try x.link.rows(v.rows.ptr, v.rowCount() * vcfg.out_hidden * 2); + } + /// `cancel` (rank 0's), asked before each chunk, decides for every rank; `from` every rank holds; `marks` rank 0's. pub fn prefill(x: Leader, ids: []const u32, from: ?engine.Resume, cancel: ?engine.Cancel, marks: ?engine.Marks) !u32 { try x.flush(); @@ -220,7 +238,7 @@ pub fn follow(e: *Engine, link: *Link) !void { const kind: Kind = @fromBackingInt(@intCast(h.kind)); if (kind != .end and kind != .drop and kind != .shared) try e.bind(h.slot); switch (kind) { - .begin => try e.begin(if (h.flags & 1 != 0) .{ .seed = h.seed, .temperature = h.temperature, .top_p = h.top_p, .min_p = h.min_p, .top_k = h.top_k } else null), + .begin => try e.beginFollowing(if (h.flags & 1 != 0) .{ .seed = h.seed, .temperature = h.temperature, .top_p = h.top_p, .min_p = h.min_p, .top_k = h.top_k } else null), .prefill => { const ids = try e.gpa.dupe(u32, step.ids); defer e.gpa.free(ids); @@ -262,6 +280,11 @@ pub fn follow(e: *Engine, link: *Link) !void { defer e.gpa.free(ids); try e.draft(h.flags & 1 != 0, ids, h.depth, &drafts); }, + .image => { + const v = e.images orelse return error.NoVision; // the ranks agreed on --vision (cuda_ranks) + try v.take(h.slot, step.ids); + try link.rows(v.rows.ptr, v.rowCount() * vcfg.out_hidden * 2); + }, .drop => for (step.ids) |id| e.dropState(id), .release => e.release(h.slot), .end => return, diff --git a/zig/src/families/glm/cuda_forward.zig b/zig/src/families/glm/cuda_forward.zig index 1e57e55a08..bbc5a92942 100644 --- a/zig/src/families/glm/cuda_forward.zig +++ b/zig/src/families/glm/cuda_forward.zig @@ -17,6 +17,10 @@ const moe = @import("cuda_moe.zig"); const quant = @import("cuda_quant.zig"); const q4 = @import("cuda_q4.zig"); const state = @import("cuda_state.zig"); +const mm_mod = @import("cuda_mm.zig"); +const prof_mod = @import("cuda_prof.zig"); +const images_mod = @import("cuda_images.zig"); +const exl3 = @import("cuda_exl3.zig"); const shared = @import("cuda_shared.zig"); pub const prefill_rows = state.prefill_rows; @@ -27,11 +31,14 @@ pub const Buffers = state.Buffers; /// The modules a forward launches from: our fatbins. pub const Modules = struct { glue: glue.Kernels, - mods: [7]cuda.Module, + mods: [9]cuda.Module, kda: kda.Kernels, lane: cuda.qmmf.Lane, // the NVFP4 linears' decode rows (qlinear) prompt: cuda.nvfp4.Prompt, // and their prompt rows experts: experts.Kernels, + exl3: exl3.Kernels, // an EXL3 checkpoint's routed experts + mm: mm_mod.Kernels, // glm_mm: the bf16 linears on tiled weights + tiled: bool = false, // the bf16 linears' weights are tiled (Engine.init, TF_GLM_MM unset or 1): glm_mm runs them copy_rows: cuda.Function, put_slot: cuda.Function, pack_dense: cuda.Function, @@ -43,7 +50,7 @@ pub const Modules = struct { var m: Modules = undefined; var n: usize = 0; errdefer for (m.mods[0..n]) |*x| x.unload(); - for (&m.mods, [_][]const u8{ k.glm_kda, k.glm_ops, k.experts, k.affine4_pack, k.glm_quant, k.qmm_group, k.glm_glue }) |*x, image| { + for (&m.mods, [_][]const u8{ k.glm_kda, k.glm_ops, k.experts, k.affine4_pack, k.glm_quant, k.qmm_group, k.glm_glue, k.glm_exl3, k.glm_mm }) |*x, image| { x.* = try cuda.Module.load(d, image); n += 1; } @@ -60,6 +67,9 @@ pub const Modules = struct { m.quant = try quant.Kernels.load(m.mods[4]); m.q4 = try q4.Kernels.load(m.mods[5], m.mods[3], ctx); m.glue = try glue.Kernels.load(m.mods[6]); + m.exl3 = try exl3.Kernels.load(m.mods[7]); + m.mm = try mm_mod.Kernels.load(m.mods[8]); + m.tiled = false; return m; } @@ -134,7 +144,8 @@ pub const Pieces = struct { /// How the ranks' partials meet: every rank in-process copying into the others' slots, or one rank over NCCL. pub const Exchange = union(enum) { local, - nccl: struct { lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm }, + // group: the main stream's all-gathers through it (RoCE writes for small ones on two ranks); side streams stay raw + nccl: struct { lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, group: ?*const cuda.tp_group.Group = null }, }; /// One rank of the forward: its weights, its sequence's caches, the window's buffers. @@ -146,6 +157,9 @@ pub const Rank = struct { segs: []const Seg = &.{}, }; +/// The engine's image rows and the sequence the forward runs on (they splice into that sequence's rows only). +pub const ImageRows = struct { v: *images_mod.Images, slot: u32 }; + /// One stream's rows of a shared window: rows row0..row0 + rows, on sequence `st` from its position. pub const Seg = struct { st: *State, row0: usize, rows: usize }; @@ -158,9 +172,11 @@ pub const Forward = struct { kda_wide: bool = true, // TF_GLM_KDA_DECODE_WIDE (the Python family's default): decode windows on KDA's three-kernel chain pieces: ?*const Pieces = null, // a prompt chunk's exchanges in row pieces on a second stream (null: whole) split: bool = true, // TF_GLM_PREFILL_REDUCE=split (the default with send/recv): each rank glues its own rows + images: ?ImageRows = null, // a prompt pass's image rows over their placeholders' embeddings (its sequence only) + prof: ?*prof_mod.Prof = null, // TF_GLM_PROF=1: events between the sections pub fn set(f: *const Forward) glue.Set { - return .{ .glue = &f.m.glue }; + return .{ .glue = &f.m.glue, .mm = if (f.m.tiled) &f.m.mm else null }; } /// reduce.sums' gather: every rank's partial of `rows` rows into every rank's gathered buffer, rank order. @@ -175,6 +191,7 @@ pub const Forward = struct { }, .nccl => |n| { const r = f.ranks[0]; + if (n.group) |g| return g.allGather(f.stream, r.b.gathered + r.w.share.rank * bytes, r.b.gathered, bytes); const count = rows * c.hidden; try n.lib.check(n.lib.api.ncclAllGather(r.b.gathered + r.w.share.rank * bytes, r.b.gathered, count, .f32, n.comm, f.stream.handle), "ncclAllGather"); }, @@ -248,7 +265,7 @@ pub const Forward = struct { } fn moeKernels(f: *const Forward) moe.Kernels { - return .{ .set = f.set(), .experts = &f.m.experts, .put_slot = f.m.put_slot }; + return .{ .set = f.set(), .experts = &f.m.experts, .exl3 = &f.m.exl3, .put_slot = f.m.put_slot, .prof = f.prof }; } const Site = enum { mixer, ffn }; @@ -451,11 +468,14 @@ pub const Forward = struct { } const c = &f.ranks[0].w.config; for (f.ranks) |r| try glue.embedB16(f.set(), f.stream, r.b.ids, r.w.embed, r.b.x, R, c.hidden, c.hc); + if (f.images) |im| for (f.ranks) |r| try im.v.splice(f.stream, im.slot, r.st.pos, r.st.pos + R, r.b.x, c.hc); + if (f.prof) |p| p.mark(f.stream, .embed); var cut: [max_pieces][2]usize = undefined; const pieces = if (f.pieces) |p| (if (f.ranks[0].b.prefill) ranges(R, p.count, &cut) else 0) else 0; // overlap.Overlap.begin: pieced with two pieces or more, and split takes a chunk of one piece too if (pieces > 1 or (pieces == 1 and f.splits(R))) try f.pieced(f.pieces.?, R, cut[0..pieces]) else try f.whole(R); try f.finish(R); + if (f.prof) |p| p.mark(f.stream, .head); } /// The unpieced loop over the layers (and the stream mean): each site's partials whole, then the glue. @@ -469,15 +489,19 @@ pub const Forward = struct { try f.pre(r, l.attn_hc, l.in_norm, R); if (r.segs.len > 0) try shared.mixer(f, r, l, li, di, R, slot(r, R)) else try f.mixer(r, l, li, di, R, @min(latent.chunksFor(r.st.pos + R), r.b.nch), slot(r, R)); } + if (f.prof) |p| p.mark(f.stream, if (f.ranks[0].w.layers[j].mixer == .kda) .kda else .mla); try f.gather(R); + if (f.prof) |p| p.mark(f.stream, .gather_a); for (f.ranks) |r| { const l = r.w.layers[j]; try f.post(r, R); try f.pre(r, l.ffn_hc, l.post_norm, R); try f.ffn(r, l, R, slot(r, R)); } + if (f.prof) |p| p.mark(f.stream, if (f.ranks[0].w.layers[j].ffn == .dense) .dense else .shared); // MoE: its fill try f.gather(R); for (f.ranks) |r| try f.post(r, R); + if (f.prof) |p| p.mark(f.stream, .gather_b); switch (f.ranks[0].w.layers[j].mixer) { .kda => li += 1, .mla => di += 1, @@ -512,6 +536,8 @@ pub const Forward = struct { const m = r.w.mtp.?; const b = r.b; try glue.embedB16(f.set(), f.stream, b.ids, r.w.embed, b.me, n, D, 1); + // the head reads each row's next token: an image's rows one place on + if (f.images) |im| try im.v.splice(f.stream, im.slot, r.st.mtp_len + 1, r.st.mtp_len + 1 + n, b.me, 1); if (r.st.mtp_len == 0) try f.d.check(f.d.api.cuMemsetD8Async(b.me, 0, D * 2, f.stream.handle), "cuMemsetD8Async"); // position 0's embedding is zeroed try glue.rmsnorm(f.set(), f.stream, b.me, D, m.enorm, b.mcat, 2 * D, null, c.eps, n, D); try glue.rmsnorm(f.set(), f.stream, b.hin, D, m.hnorm, b.mcat + D * 2, 2 * D, null, c.eps, n, D); diff --git a/zig/src/families/glm/cuda_glue.zig b/zig/src/families/glm/cuda_glue.zig index 733edcdb27..53414396d3 100644 --- a/zig/src/families/glm/cuda_glue.zig +++ b/zig/src/families/glm/cuda_glue.zig @@ -2,9 +2,10 @@ const std = @import("std"); const cuda = @import("cuda"); +const mm_mod = @import("cuda_mm.zig"); /// What the forward's kernels launch from: our glue's functions. -pub const Set = struct { glue: *const Kernels }; +pub const Set = struct { glue: *const Kernels, mm: ?*const mm_mod.Kernels = null }; // mm: bf16 linears tiled at load /// Our glm_glue.cu's functions. pub const Kernels = struct { diff --git a/zig/src/families/glm/cuda_images.zig b/zig/src/families/glm/cuda_images.zig new file mode 100644 index 0000000000..cfeb019864 --- /dev/null +++ b/zig/src/families/glm/cuda_images.zig @@ -0,0 +1,166 @@ +//! A prompt pass's image rows on every rank: rank 0's tower encodes them, the others receive them, all splice them. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const vision = @import("cuda_vision.zig"); +const vcfg = @import("vision_config.zig"); +const image_core = @import("image_core.zig"); +const image = image_core.For(@import("core").picture); + +/// What a request may hold: image tokens in all (the rows buffer), and images. +pub const Limits = struct { image_tokens: u32, max_images: u32 }; + +/// Prompt rows [at, at + n) take image rows [row, row + n). +pub const Span = struct { at: u32, n: u32, row: u32 }; + +pub const Images = struct { + gpa: std.mem.Allocator, + io: std.Io, + d: *const cuda.Driver, + module: cuda.Module, + run: vision.Run, + tower: ?vision.Tower = null, // rank 0's; the other ranks receive the rows + cfg: ?vcfg.Config = null, + rows: cuda.DeviceBuffer, // bf16 [limits.image_tokens][out_hidden] + map: cuda.DeviceBuffer, // i32 [max_rows]: a splice's feature row per prompt row + host_map: []i32, + limits: Limits, + spans: std.ArrayList(Span) = .empty, + slot: u32 = 0, // the sequence the spans belong to: another stream's rows at the same positions take none + marks: struct { image: u32, open: u32, close: u32 } = .{ .image = 0, .open = 0, .close = 0 }, + mark_ids: [3]u32 = @splat(0), // begin, image, end: text that spells one stays text + + /// The chat template writes each image's begin and end itself (the server only expands the image token). + pub const template_wraps = true; + + /// Rank 0 (`tower`) or another rank: the rows buffer, and rank 0's tower; null with `problem` when no tower. + pub fn init(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, stream: cuda.Stream, dir: []const u8, limits: Limits, max_rows: usize, tower: bool, problem: *[]const u8) !?*Images { + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const cwd = std.Io.Dir.cwd(); + const cfg_text = try cwd.readFileAlloc(io, try std.fs.path.join(a, &.{ dir, "config.json" }), a, .limited(1 << 24)); + const proc = cwd.readFileAlloc(io, try std.fs.path.join(a, &.{ dir, "processor_config.json" }), a, .limited(1 << 24)) catch null; + const cfg = (try vcfg.Config.parse(a, cfg_text, proc)) orelse { + problem.* = "this checkpoint has no vision tower: serve it without --vision"; + return null; + }; + const v = try gpa.create(Images); + errdefer gpa.destroy(v); + var module = try cuda.Module.load(d, cuda.kernels.glm_vision); + errdefer module.unload(); + v.* = .{ .gpa = gpa, .io = io, .d = d, .module = module, .run = .{ .stream = stream, .f = try vision.Funcs.load(module) }, .cfg = cfg, .rows = undefined, .map = undefined, .host_map = undefined, .limits = limits }; + v.rows = try cuda.DeviceBuffer.alloc(d, @as(usize, limits.image_tokens) * vcfg.out_hidden * 2); + errdefer v.rows.free(); + v.map = try cuda.DeviceBuffer.alloc(d, max_rows * 4); + errdefer v.map.free(); + v.host_map = try gpa.alloc(i32, max_rows); + errdefer gpa.free(v.host_map); + if (tower) v.tower = try vision.Tower.load(gpa, io, d, dir); + v.marks = .{ .image = cfg.image_token, .open = cfg.image_start, .close = cfg.image_end }; + v.mark_ids = .{ cfg.image_start, cfg.image_token, cfg.image_end }; + return v; + } + + pub fn deinit(v: *Images) void { + v.spans.deinit(v.gpa); + if (v.tower) |*t| t.deinit(); + v.gpa.free(v.host_map); + v.map.free(); + v.rows.free(); + v.module.unload(); + v.gpa.destroy(v); + } + + /// Device bytes the tower and the buffers hold. + pub fn bytes(v: *const Images) u64 { + return (if (v.tower) |t| t.bytes else 0) + v.rows.len + v.map.len; + } + + /// No rows splice into any pass until another begins (each request's prompt pass starts here). + pub fn clear(v: *Images) void { + v.spans.clearRetainingCapacity(); + } + + /// Rank 0: a prompt pass on `slot` from row `from`: the images whose rows end past it, through the tower. + pub fn begin(v: *Images, slot: u32, images: []const lanes.stream.Image, from: usize) !void { + v.clear(); + if (images.len > v.limits.max_images) return error.TooManyImages; + const t = &(v.tower orelse return error.NoVision); + var row: u32 = 0; + for (images) |img| { + if (img.gh * img.gw != img.tokens * vcfg.patches_per_token) return error.ImageOutOfStep; + if (img.at + img.tokens <= from) continue; // in the kept state already + if (row + img.tokens > v.limits.image_tokens) return error.ImageRows; + try t.encode(&v.run, img.pixels, .{ 1, img.gh, img.gw }, v.rows.ptr + @as(u64, row) * vcfg.out_hidden * 2); + try v.spans.append(v.gpa, .{ .at = img.at, .n = img.tokens, .row = row }); + row += img.tokens; + } + v.slot = slot; + } + + /// The other ranks: rank 0's spans (as `words`: at, n, row a span) on `slot`; the rows arrive into `rows` next. + pub fn take(v: *Images, slot: u32, words: []const u32) !void { + v.clear(); + if (words.len % 3 != 0) return error.BadStep; + var i: usize = 0; + while (i < words.len) : (i += 3) { + const sp: Span = .{ .at = words[i], .n = words[i + 1], .row = words[i + 2] }; + if (sp.row + sp.n > v.limits.image_tokens) return error.ImageRows; + try v.spans.append(v.gpa, sp); + } + v.slot = slot; + } + + /// Image rows the spans hold (what crosses to the other ranks). + pub fn rowCount(v: *const Images) usize { + var n: usize = 0; + for (v.spans.items) |sp| n = @max(n, sp.row + sp.n); + return n; + } + + /// Prompt rows [lo, hi) of `slot`, embedded at `emb` in `copies` streams, take the image rows they cross. + pub fn splice(v: *Images, s: cuda.Stream, slot: u32, lo: usize, hi: usize, emb: u64, copies: usize) !void { + if (v.spans.items.len == 0 or slot != v.slot) return; + const n = hi - lo; + if (n > v.host_map.len) return error.BadWindow; + var any = false; + for (v.host_map[0..n], lo..) |*m, p| { + m.* = -1; + for (v.spans.items) |sp| if (p >= sp.at and p < sp.at + sp.n) { + m.* = @intCast(sp.row + (p - sp.at)); + any = true; + }; + } + if (!any) return; + try v.map.upload(0, std.mem.sliceAsBytes(v.host_map[0..n])); // waits: the host map is reused next splice + var a: cuda.Args = .{}; + a.add(v.rows.ptr); + a.add(emb); + a.add(v.map.ptr); + a.add(@as(c_int, @intCast(n))); + a.add(@as(c_int, @intCast(vcfg.out_hidden))); + a.add(@as(c_int, @intCast(copies))); + try cuda.launch.launch(v.run.f.inject, .{ .grid = .{ .x = @intCast(n) }, .block = .{ .x = 256 } }, s, &a); + } + + /// The words for a prepare error the server cannot word (a JPEG library this machine lacks). + pub const explain = image.decoderProblem; + + /// An image's bytes as the tower reads them, at most `max_tokens` tokens (the server's prepare). + pub fn prepare(v: *const Images, a: std.mem.Allocator, data: []const u8, max_tokens: u32) !image_core.Prepared { + return image.prepare(v.io, a, data, v.cfg.?, max_tokens); + } +}; + +test "a splice's map: rows inside a span take its rows, others keep their embedding" { + var host: [8]i32 = undefined; + const spans = [_]Span{ .{ .at = 3, .n = 2, .row = 0 }, .{ .at = 6, .n = 4, .row = 2 } }; + for (&host, 2..) |*m, p| { + m.* = -1; + for (spans) |sp| if (p >= sp.at and p < sp.at + sp.n) { + m.* = @intCast(sp.row + (p - sp.at)); + }; + } + try std.testing.expectEqualSlices(i32, &.{ -1, 0, 1, -1, 2, 3, 4, 5 }, &host); +} diff --git a/zig/src/families/glm/cuda_lanes.zig b/zig/src/families/glm/cuda_lanes.zig index 0106c383b0..c257fca5ab 100644 --- a/zig/src/families/glm/cuda_lanes.zig +++ b/zig/src/families/glm/cuda_lanes.zig @@ -11,8 +11,9 @@ const be = lanes.backend; /// First tokens a handle names. const ring = 1024; -/// The widest verify window: the pending token and a full draft chain (glm-decode-check's, the AOT row variants). -pub const max_rows = 1 + engine.max_drafts; +/// The widest verify window: the pending token and a copy (the lane core proposes them), as the decode buffers hold. +pub const max_rows = fw.window_rows; +const mtp_rows = 1 + engine.max_drafts; /// A stream's place in the backend: its sequence's slot, the rows its last verify left uncommitted, its held drafts. const Lane = struct { @@ -56,12 +57,13 @@ pub const Lanes = struct { break :vt v; }; - /// Round loop facts: windows up to max_rows, drafts with the head loaded, shared windows unless TF_GLM_SHARED=0. + /// Round loop facts: copies up to max_rows, head chains to mtp_rows, shared windows unless TF_GLM_SHARED=0. pub fn facts(self: *const Lanes) lanes.Model { const drafting = self.x.e.mtp; const sharing = drafting and !self.alone; return .{ .exact_width = if (drafting) max_rows else 1, + .first_copy_rows = mtp_rows, // a first copy as wide as a head round; it grows as its tokens land .gpu_tokens = false, .mtp = drafting, .speculate = drafting, @@ -121,6 +123,7 @@ pub const Lanes = struct { fn prefillFn(ptr: *anyopaque, s: *lanes.Stream) anyerror!void { const self = of(ptr); const from = try self.opening(s); + try self.images(s, from); _ = self.take(try self.x.prefill(s.prompt(), from, .{ .ptr = s, .check = cancelled }, marksOf(s))); } @@ -129,6 +132,7 @@ pub const Lanes = struct { const self = of(ptr); const going: ?usize = if (self.lanes.getPtr(s)) |l| l.fill else null; const from = if (going == null) try self.opening(s) else null; + if (going == null) try self.images(s, from); const l = self.lanes.getPtr(s).?; if (going != null) try self.x.bind(l.slot); const at = going orelse if (from) |r| r.at else 0; @@ -152,11 +156,20 @@ pub const Lanes = struct { gop.value_ptr.held_n = 0; gop.value_ptr.fill = null; try self.x.bind(gop.value_ptr.slot); + if (self.x.e.images) |v| v.clear(); // no earlier request's rows splice into this one try self.x.begin(s.sampling); s.cached = 0; return self.resumeFrom(s, ids.len); } + /// A prompt pass with images: rank 0's tower encodes those past `from`, then every rank holds their rows. + fn images(self: *Lanes, s: *lanes.Stream, from: ?engine.Resume) !void { + if (s.images.len == 0) return; + const v = self.x.e.images orelse return error.NoVision; + try v.begin(self.x.e.slot, s.images, if (from) |r| r.at else 0); // a refusal here leaves every rank as it was + try self.x.images(v); + } + fn marksOf(s: *lanes.Stream) ?engine.Marks { return if (s.reuse.marks.len == 0) null else .{ .at = s.reuse.marks, .ctx = s, .call = keptAt }; } diff --git a/zig/src/families/glm/cuda_mlp.zig b/zig/src/families/glm/cuda_mlp.zig index 8e8d5eba11..a9e55254b9 100644 --- a/zig/src/families/glm/cuda_mlp.zig +++ b/zig/src/families/glm/cuda_mlp.zig @@ -3,12 +3,19 @@ const std = @import("std"); const cuda = @import("cuda"); const glue = @import("cuda_glue.zig"); +const b16 = @import("cuda_b16.zig"); -/// One dense layer on a rank: gate and up [width][dims] and down [dims][width], NVFP4 as Fp4Linear holds them. -pub const Layer = struct { gate: cuda.qmmf.Weight, up: cuda.qmmf.Weight, down: cuda.qmmf.Weight, width: usize, dims: usize, limit: f32 }; +/// A dense layer's weights: NVFP4 as Fp4Linear holds them, or the checkpoint's bf16 (an EXL3 checkpoint's). +pub const Weights = union(enum) { + nvfp4: struct { gate: cuda.qmmf.Weight, up: cuda.qmmf.Weight, down: cuda.qmmf.Weight }, + bf16: struct { gu: u64, down: u64 }, // [gate | up] rows [2 width][dims], down [dims][width] +}; -/// The block's scratch for up to `rows` rows: one projection's rows, the [gate | up] rows, the activation and its sums. -pub const Scratch = struct { tmp: u64, gu: u64, act: u64, xs_act: u64 }; +/// One dense layer on a rank: gate and up [width][dims] and down [dims][width]. +pub const Layer = struct { w: Weights, width: usize, dims: usize, limit: f32 }; + +/// The block's scratch for `rows` rows: a projection's rows, [gate | up], the activation, its sums, split-K partials. +pub const Scratch = struct { tmp: u64, gu: u64, act: u64, xs_act: u64, mm_part: u64 }; pub const Kernels = struct { set: glue.Set, linear: cuda.qlinear.Linear, copy_rows: cuda.Function }; @@ -31,13 +38,23 @@ pub fn block(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_st /// forward.out_proj's fill for the dense MLP: down of rows lo..hi (s.act) into out [hi - lo][dims] fp32. pub fn fill(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, lo: usize, hi: usize, prompt: bool, out: u64) !void { - try mm(k, stream, l.down, s.act + lo * l.width * 2, l.width, hi - lo, out, true, prompt); + switch (l.w) { + .nvfp4 => |w| try mm(k, stream, w.down, s.act + lo * l.width * 2, l.width, hi - lo, out, true, prompt), + .bf16 => |w| try b16.matmul(k.set, stream, s.act + lo * l.width * 2, l.width, w.down, out, s.mm_part, hi - lo, l.dims, l.width, true), + } } /// The block up to its down projection: [gate | up] and SwiGLU into s.act. pub fn body(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_stride: usize, rows: usize, prompt: bool) !void { const w = l.width; - for ([_]cuda.qmmf.Weight{ l.gate, l.up }, 0..) |proj, j| { + const q = switch (l.w) { + .nvfp4 => |q| q, + .bf16 => |b| { // [gate | up] in one matmul, as the shared expert's + try b16.matmul(k.set, stream, x, x_stride, b.gu, s.gu, s.mm_part, rows, 2 * w, l.dims, false); + return glue.swiglu(k.set, stream, s.gu, s.act, s.xs_act, l.limit, rows, w); + }, + }; + for ([_]cuda.qmmf.Weight{ q.gate, q.up }, 0..) |proj, j| { // Fp4Linear(out=a column slice): the rows computed contiguous, then copied into their half of each row try mm(k, stream, proj, x, x_stride, rows, s.tmp, false, prompt); var a: cuda.Args = .{}; diff --git a/zig/src/families/glm/cuda_mm.zig b/zig/src/families/glm/cuda_mm.zig new file mode 100644 index 0000000000..0abdfd7e39 --- /dev/null +++ b/zig/src/families/glm/cuda_mm.zig @@ -0,0 +1,109 @@ +//! GLM's bf16 linears on tensor cores (glm_mm): weights tiled once at load, any row count, row-stable K slices. +const std = @import("std"); +const cuda = @import("cuda"); + +/// A linear's tiled weights: [n / 8][k / 64][32 lanes][8 words] bf16 pairs (glm_mm_tile16's layout). +pub const Tiled = struct { w: u64, n: usize, k: usize }; + +/// glm_mm's bf16 kernels by row bucket, the K-slice reduce and the tiling. +pub const Kernels = struct { + by_rows: [4]cuda.Function, // 16, 32, 64, 128 rows a block + prompt: cuda.Function, // 128-row blocks, row blocks fastest: a prompt's blocks share each weight tile in L2 + reduce: cuda.Function, + tile: cuda.Function, + pdl: bool, // dependent launches: a matmul's weights stream in while the kernel before it finishes + + pub fn load(m: cuda.Module) !Kernels { + var k: Kernels = .{ .by_rows = undefined, .prompt = try m.function("glm_mm16_rows"), .reduce = try m.function("glm_mm_reduce"), .tile = try m.function("glm_mm_tile16"), .pdl = true }; + inline for (.{ "glm_mm16_16", "glm_mm16_32", "glm_mm16_64", "glm_mm16_128" }, 0..) |name, i| { + k.by_rows[i] = try m.function(name); + try k.by_rows[i].allowDynamicShared(smem(bucket(i))); + } + try k.prompt.allowDynamicShared(smem(128)); + if (std.c.getenv("TF_GLM_PDL")) |v| k.pdl = !std.mem.eql(u8, std.mem.span(v), "0"); + return k; + } +}; + +const split_target = 192; // column tiles times K slices a launch aims at (the GB10's SMs, a few blocks each) + +fn bucket(i: usize) usize { + return @as(usize, 16) << @intCast(i); +} + +fn smem(bm: usize) u32 { + return @intCast(3 * (bm * 128 + 8 * 32 * 8 * 4 + 128)); +} + +/// K slices fixed by the shape: column tiles times slices reach the target, each slice at least 4 groups of 64. +pub fn splitK(n: usize, k: usize) usize { + const tiles = (n + 63) / 64; + const steps = k / 64; + var sk: usize = 1; + while (sk < 8 and tiles * sk < split_target and steps % (sk * 2) == 0 and steps / (sk * 2) >= 4) sk *= 2; + return sk; +} + +/// Bytes of fp32 partials a call of `m` rows needs (0: one slice; windows past 64 rows go 64 at a time). +pub fn partBytes(m: usize, w: Tiled) usize { + const sk = splitK(w.n, w.k); + return if (sk > 1) sk * @min(m, 64) * w.n * 4 else 0; +} + +/// x [m][k] (rows `ldx` apart) by `w` -> out [m][n], bf16 or fp32 (`wide`); `part`: partBytes(m, w). +pub fn matmul(k: *const Kernels, s: cuda.Stream, x: u64, ldx: usize, w: Tiled, out: u64, part: u64, m: usize, wide: bool) !void { + if (w.k % 64 != 0 or w.n % 8 != 0 or x % 16 != 0 or ldx % 8 != 0) return error.BadShape; + var done: usize = 0; + while (done < m) { // windows to 64 rows a launch with K slices; prompts in one launch of 128-row blocks, one slice + const big = m - done >= 128; + const rows = if (big) m - done else @min(m - done, 64); + const i: usize = if (big) 3 else if (rows <= 16) 0 else if (rows <= 32) 1 else 2; + const bm = bucket(i); + const sk = if (big) 1 else splitK(w.n, w.k); // a prompt's blocks share each weight tile in L2: no slices + const o = out + done * w.n * @as(usize, if (wide) 4 else 2); + var a: cuda.Args = .{}; + a.add(x + done * ldx * 2); + a.add(@as(c_longlong, @intCast(ldx))); + a.add(w.w); + a.add(@as(u64, 0)); + a.add(o); + a.add(part); + a.add(@as(c_int, @intFromBool(wide))); + a.add(@as(c_int, @intCast(rows))); + a.add(@as(c_int, @intCast(w.n))); + a.add(@as(c_int, @intCast(w.k))); + a.add(@as(c_int, @intCast(sk))); + const cols: u32 = @intCast((w.n + 63) / 64); + const row_blocks: u32 = @intCast((rows + bm - 1) / bm); + if (big) { + try cuda.launch.launch(k.prompt, .{ .grid = .{ .x = row_blocks, .y = cols }, .block = .{ .x = 256 }, .shared = smem(128), .pdl = k.pdl }, s, &a); + } else try cuda.launch.launch(k.by_rows[i], .{ .grid = .{ .x = cols, .y = row_blocks, .z = @intCast(sk) }, .block = .{ .x = if (bm <= 32) 128 else 256 }, .shared = smem(bm), .pdl = k.pdl }, s, &a); + if (sk > 1) { + const total = rows * w.n; + var r: cuda.Args = .{}; + r.add(part); + r.add(o); + r.add(@as(c_longlong, @intCast(total))); + r.add(@as(c_int, @intCast(sk))); + r.add(@as(c_int, @intFromBool(wide))); + try cuda.launch.launch(k.reduce, .{ .grid = .{ .x = @intCast((total + 255) / 256) }, .block = .{ .x = 256 } }, s, &r); + } + done += rows; + } +} + +/// A row-major bf16 [n][k] linear at `src` tiled into `dst` (as many bytes). +pub fn tile(k: *const Kernels, s: cuda.Stream, src: u64, dst: u64, n: usize, kk: usize) !void { + const chunks = n / 8 * (kk / 64) * 32; + var a: cuda.Args = .{}; + a.add(src); + a.add(dst); + a.add(@as(c_int, @intCast(n))); + a.add(@as(c_int, @intCast(kk))); + try cuda.launch.launch(k.tile, .{ .grid = .{ .x = @intCast((chunks + 255) / 256) }, .block = .{ .x = 256 } }, s, &a); +} + +test "K slices as the shapes take them" { + try std.testing.expectEqual(@as(usize, 1), splitK(12288, 4096)); // 192 column tiles already + try std.testing.expectEqual(@as(usize, 4), splitK(4096, 2048)); // 64 tiles, 32 groups: 4 slices of 8 +} diff --git a/zig/src/families/glm/cuda_moe.zig b/zig/src/families/glm/cuda_moe.zig index ec860a188e..72132399d2 100644 --- a/zig/src/families/glm/cuda_moe.zig +++ b/zig/src/families/glm/cuda_moe.zig @@ -5,12 +5,15 @@ const cuda = @import("cuda"); const glue = @import("cuda_glue.zig"); const experts = @import("cuda_experts.zig"); const b16 = @import("cuda_b16.zig"); +const exl3 = @import("cuda_exl3.zig"); +const prof_mod = @import("cuda_prof.zig"); /// One MoE layer on a rank: router (bf16 [count][dims]) and fp32 bias, routed experts, shared [gate | up] and down. pub const Layer = struct { router: u64, bias: u64, - routed: cuda.experts.Layer, + routed: cuda.experts.Layer, // NVFP4 experts, or only the shape (width, dims, limit) when `exl3` holds them + exl3: ?exl3.Experts = null, count: usize, shared_gu: u64, shared_down: u64, @@ -34,10 +37,11 @@ pub const Scratch = struct { sxs: u64, // fp32 [rows][shared / 64] sy: u64, // fp32 [rows][dims] mm_part: u64, // fp32 split-K partials of the shared expert's matmuls + x3: exl3.Scratch = .{ .picks = 0, .xg = 0, .xu = 0, .z = 0 }, // EXL3 experts (eact holds their down inputs) }; /// The kernels the block runs: the captured Triton set, the grouped NVFP4 experts, the plan, and glm_ops. -pub const Kernels = struct { set: glue.Set, experts: *const experts.Kernels, put_slot: cuda.Function }; +pub const Kernels = struct { set: glue.Set, experts: *const experts.Kernels, exl3: *const exl3.Kernels, put_slot: cuda.Function, prof: ?*prof_mod.Prof = null }; /// forward.moe_block: out [rows][dims] fp32 partial; `prompt`: 64-pair items, the prompt kernels, bf16 slot rows. pub fn block(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_stride: usize, rows: usize, prompt: bool, out: u64) !void { @@ -52,6 +56,9 @@ pub fn body(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_str const shared_id: i32 = @intCast(l.count); try glue.router(k.set, stream, x, x_stride, l.router, s.mlog, s.router_part, rows, d, l.count); try glue.select(k.set, stream, s.mlog, l.bias, s.pick, s.wts, rows, l.top_k, l.count, l.scale, l.norm); + if (k.prof) |p| p.mark(stream, .router); + defer if (k.prof) |p| p.mark(stream, .routed); + if (l.exl3) |w| return exl3.run(k.exl3, k.experts, stream, w, s.x3, s.plan, x, x_stride, s.pick, slots, rows, l.count, d, l.routed.width, l.routed.limit, prompt, s.ey, s.eact); const plan: experts.Plan = .{ .route = s.plan, .slots = slots, .experts = l.count + 1, .prompt = prompt }; try experts.route(k.experts, stream, s.pick, rows * slots, plan); try experts.run(k.experts, stream, .gate_up, l.routed, plan, x, x_stride, rows, s.eact, shared_id); @@ -63,14 +70,15 @@ pub fn fill(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_str const d = l.routed.dims; const slots = l.top_k + 1; const n = hi - lo; - const eb: usize = if (prompt) 2 else 4; // the slot rows: bf16 in a prompt chunk + const bf16_slots = prompt and l.exl3 == null; // the slot rows: bf16 in an NVFP4 prompt chunk, else fp32 + const eb: usize = if (bf16_slots) 2 else 4; try b16.matmul(k.set, stream, x + lo * x_stride * 2, x_stride, l.shared_gu, s.sgu + lo * 2 * l.shared * 2, s.mm_part, n, 2 * l.shared, d, false); try glue.swiglu(k.set, stream, s.sgu + lo * 2 * l.shared * 2, s.sact + lo * l.shared * 2, s.sxs + lo * (l.shared / 64) * 4, l.routed.limit, n, l.shared); try b16.matmul(k.set, stream, s.sact + lo * l.shared * 2, l.shared, l.shared_down, s.sy + lo * d * 4, s.mm_part, n, d, l.shared, true); var a: cuda.Args = .{}; a.add(s.sy + lo * d * 4); a.add(s.ey + lo * slots * d * eb); - for ([_]usize{ n, d, slots, l.top_k, @intFromBool(prompt) }) |v| a.add(@as(c_int, @intCast(v))); + for ([_]usize{ n, d, slots, l.top_k, @intFromBool(bf16_slots) }) |v| a.add(@as(c_int, @intCast(v))); try cuda.launch.launch(k.put_slot, .{ .grid = .{ .x = @intCast((n * d + 255) / 256) }, .block = .{ .x = 256 } }, stream, &a); - try glue.combine(k.set, stream, s.ey + lo * slots * d * eb, !prompt, s.wts + lo * slots * 4, out, n, d, slots); + try glue.combine(k.set, stream, s.ey + lo * slots * d * eb, !bf16_slots, s.wts + lo * slots * 4, out, n, d, slots); } diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig index e85cb25ad7..a894a5bbb5 100644 --- a/zig/src/families/glm/cuda_native.zig +++ b/zig/src/families/glm/cuda_native.zig @@ -5,6 +5,7 @@ const cuda = @import("cuda"); const lanes = @import("lanes"); const engine = @import("cuda_engine.zig"); const follow = @import("cuda_follow.zig"); +const images_mod = @import("cuda_images.zig"); const glm_lanes = @import("cuda_lanes.zig"); const ranks = @import("cuda_ranks.zig"); const fw = @import("cuda_forward.zig"); @@ -31,6 +32,7 @@ pub const Options = struct { segments: usize = 1, slide: bool = false, room: ?Room = null, + vision: ?images_mod.Limits = null, // --vision: rank 0's tower, every rank's image rows }; /// The smallest room among the other ranks and whose it is. pub const Peer = struct { bytes: u64, rank: usize }; @@ -49,6 +51,7 @@ pub const Loaded = struct { lone: ?LoneRun = null, target: *engine.Engine, // the prompt cache keeps this engine's states (snap_*) peer_room: ?Peer = null, // rank 0: the server admits streams and kept states by this when it is smaller + vision: ?*images_mod.Images = null, // rank 0's tower for the server's image offer }; /// This rank's place: TF_GLM_RANK of TF_GLM_WORLD (both, or neither for a refusal that names them). @@ -85,13 +88,37 @@ pub fn capacity(context: usize) usize { const Owned = struct { gpa: std.mem.Allocator, - lib: cuda.nccl.Library, - comm: cuda.nccl.Comm, + group: cuda.tp_group.Group, // the ranks' communicator; two ranks add RoCE writes for the main stream's small gathers e: *engine.Engine, link: follow.Link, lanes: glm_lanes.Lanes, }; +extern "c" fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: c_int) c_int; + +/// Every peer's send/recv connected now, before the weights fill memory (connected on first use, ibv_reg_mr can fail). +fn warmPeers(d: *const cuda.Driver, g: *const cuda.tp_group.Group) !void { + const bytes = 4 << 20; // a prompt piece's exchange: every channel a large send/recv takes + var buf = try cuda.DeviceBuffer.alloc(d, 2 * bytes * g.world); + defer buf.free(); + var s = try cuda.Stream.init(d, false); + defer s.deinit(); + try g.lib.check(g.lib.api.ncclGroupStart(), "ncclGroupStart"); + const posted: anyerror!void = post: { + for (0..g.world) |peer| if (peer != g.rank) { + g.send(s, buf.ptr + peer * bytes, bytes, peer) catch |err| break :post err; + g.recv(s, buf.ptr + (g.world + peer) * bytes, bytes, peer) catch |err| break :post err; + }; + break :post {}; + }; + try g.lib.check(g.lib.api.ncclGroupEnd(), "ncclGroupEnd"); // the group closes on every path + try posted; + try s.synchronize(); +} + +/// The largest shard a rank sends by RoCE writes: a decode window's partials (8 rows of fp32) fit. +const roce_cap = 512 << 10; + pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { _ = kernels; // our own kernels: a captured set is not used const p = try place(); @@ -101,35 +128,54 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ std.log.err("{s}", .{problem}); return err; }; + if (p.world > 1) try cuda.tp_watch.start(io, p.rank, p.world); // rank 0's line before NCCL: a failure stops every rank const own = try gpa.create(Owned); errdefer gpa.destroy(own); own.gpa = gpa; - own.lib = try cuda.nccl.Library.open(); - errdefer own.lib.close(); + // every connection registered at init, while memory is free (lazily, a full GB10's ibv_reg_mr can refuse it) + if (std.c.getenv("NCCL_RUNTIME_CONNECT") == null) _ = setenv("NCCL_RUNTIME_CONNECT", "0", 0); + var lib = try cuda.nccl.Library.open(); + errdefer lib.close(); var id: cuda.nccl.UniqueId = undefined; - try own.lib.check(own.lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); // NCCL_COMM_ID: the same id on every rank - own.comm = null; - try own.lib.check(own.lib.api.ncclCommInitRank(&own.comm, @intCast(p.world), id, @intCast(p.rank)), "ncclCommInitRank"); - errdefer _ = own.lib.api.ncclCommDestroy(own.comm); + try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); // NCCL_COMM_ID: the same id on every rank + var comm: cuda.nccl.Comm = null; + try lib.check(lib.api.ncclCommInitRank(&comm, @intCast(p.world), id, @intCast(p.rank)), "ncclCommInitRank"); + own.group = .{ .lib = lib, .comm = comm, .rank = p.rank, .world = p.world }; + errdefer _ = own.group.lib.api.ncclCommDestroy(own.group.comm); + if (p.world > 1) try warmPeers(ctx.d, &own.group); const cap = capacity(o.context); - own.e = try engine.Engine.init(gpa, io, ctx.d, ctx, dir, &own.lib, own.comm, p.rank, p.world, .{ .capacity = cap, .mtp = o.drafts, .settings = set }); + own.e = try engine.Engine.init(gpa, io, ctx.d, ctx, dir, &own.group.lib, own.group.comm, p.rank, p.world, .{ .capacity = cap, .mtp = o.drafts, .settings = set, .vision = if (o.vision) |l| l.image_tokens else 0 }); errdefer own.e.deinit(); - own.link = try follow.Link.init(gpa, ctx.d, &own.lib, own.comm, own.e.stream, cap); + if (o.vision) |lim| { // before the rooms: the tower's bytes count against rank 0's + own.e.images = (try images_mod.Images.init(gpa, io, ctx.d, own.e.stream, dir, lim, fw.prefill_rows + 1, p.rank == 0, &problem)) orelse { + std.log.err("{s}", .{problem}); + return error.NoVisionTower; + }; + } + errdefer if (own.e.images) |v| v.deinit(); + try own.group.openRoce(gpa, io, ctx.d, own.e.stream, roce_cap); // both ranks bring it up or neither does + errdefer if (own.group.fast) |f| f.deinit(); + own.e.group = &own.group; + own.link = try follow.Link.init(gpa, ctx.d, &own.group.lib, own.group.comm, own.e.stream, cap); errdefer own.link.deinit(); // Rank 0 admits by the smallest room of the ranks; its sizes are the largest (the first ranks take the remainder). const rooms = try gpa.alloc(u64, p.world); defer gpa.free(rooms); const mine: u64 = if (o.room) |r| r.read(r.ctx) orelse 0 else std.math.maxInt(u64); - try ranks.gather(u64, ctx.d, &own.lib, own.comm, own.e.stream, p.rank, mine, rooms); + try ranks.gather(u64, ctx.d, &own.group.lib, own.group.comm, own.e.stream, p.rank, mine, rooms); var peer: ?Peer = null; for (rooms[1..], 1..) |r, q| if (peer == null or r < peer.?.bytes) { peer = .{ .bytes = r, .rank = q }; }; + // every rank loaded and agreed (cuda_ranks) before the first step; the digest names what they agreed on + if (p.world > 1) try cuda.tp_watch.ready(io, ranks.hashBytes(std.fs.path.basename(std.mem.trimEnd(u8, dir, "/")))); if (p.rank != 0) { // a follower never serves: it runs rank 0's steps, then leaves the process follow.follow(own.e, &own.link) catch |err| { std.log.err("rank {d}: following rank 0 stopped ({s})", .{ p.rank, @errorName(err) }); - std.process.exit(1); + cuda.tp_watch.fail(@errorName(err)); // rank 0 stops with it instead of waiting on this rank + std.c._exit(1); // the stream may hold kernels waiting on a gone rank: a normal exit would join them }; + cuda.tp_watch.ended(); std.log.info("rank {d}: rank 0 ended; leaving", .{p.rank}); std.process.exit(0); } @@ -150,6 +196,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ .deinit = close, .target = own.e, .peer_room = peer, + .vision = own.e.images, }; } @@ -190,12 +237,13 @@ pub fn snap_drop(ptr: *anyopaque, saved: *anyopaque) void { fn close(ptr: *anyopaque) void { const own: *Owned = @ptrCast(@alignCast(ptr)); + cuda.tp_watch.ended(); // the other ranks' lines close on purpose from here own.lanes.x.end() catch {}; own.lanes.deinit(); own.link.deinit(); + if (own.e.images) |v| v.deinit(); own.e.deinit(); - _ = own.lib.api.ncclCommDestroy(own.comm); - own.lib.close(); + own.group.close(); own.gpa.destroy(own); } diff --git a/zig/src/families/glm/cuda_prof.zig b/zig/src/families/glm/cuda_prof.zig new file mode 100644 index 0000000000..a4bb46cdf1 --- /dev/null +++ b/zig/src/families/glm/cuda_prof.zig @@ -0,0 +1,77 @@ +//! TF_GLM_PROF=1: where a forward's GPU time goes (events between its sections) against the call's wall time. +const std = @import("std"); +const cuda = @import("cuda"); + +pub const Section = enum { gate, embed, kda, mla, gather_a, dense, router, routed, shared, gather_b, head, draw, mtp, commit }; +const names = [_][]const u8{ "gate", "embed", "kda", "mla", "gather_a", "dense", "router", "routed", "shared", "gather_b", "head", "draw", "mtp", "commit" }; +const max_marks = 512; + +pub const Prof = struct { + d: *const cuda.Driver, + events: [max_marks + 1]cuda.Event, + ends: [max_marks]Section = undefined, + n: usize = 0, // marks this call + open: bool = false, + t0: i96 = 0, + io: std.Io, + sections: [names.len]f64 = @splat(0), // ms each section took on the GPU, summed over calls + gpu_ms: f64 = 0, + wall_ms: f64 = 0, + calls: u64 = 0, + label: []const u8 = "", + every: u64 = 64, // calls a report sums + + pub fn init(d: *const cuda.Driver, io: std.Io) !*Prof { + const p = try std.heap.page_allocator.create(Prof); + p.* = .{ .d = d, .events = undefined, .io = io }; + for (&p.events) |*e| e.* = try cuda.Event.init(d, true); + return p; + } + + /// TF_GLM_PROF=1 asks for it. + pub fn wanted() bool { + const v = std.c.getenv("TF_GLM_PROF") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); + } + + /// A call begins on `s` (its wall clock and its first event). + pub fn begin(p: *Prof, s: cuda.Stream, label: []const u8) !void { + p.n = 0; + p.open = true; + p.label = label; + p.t0 = std.Io.Clock.awake.now(p.io).toNanoseconds(); + try p.events[0].record(s); + } + + /// The work queued on `s` since the last mark belongs to `section`. + pub fn mark(p: *Prof, s: cuda.Stream, section: Section) void { + if (!p.open or p.n >= max_marks) return; + p.events[p.n + 1].record(s) catch return; + p.ends[p.n] = section; + p.n += 1; + } + + /// The call ended (its results read: the stream is idle): sections summed, a report every 64 calls. + pub fn end(p: *Prof) void { + if (!p.open) return; + p.open = false; + const wall: f64 = @as(f64, @floatFromInt(std.Io.Clock.awake.now(p.io).toNanoseconds() - p.t0)) / 1e6; + if (p.n == 0) return; + p.events[p.n].synchronize() catch return; + for (0..p.n) |i| p.sections[@backingInt(p.ends[i])] += cuda.Event.elapsedMs(p.events[i], p.events[i + 1]) catch 0; + p.gpu_ms += cuda.Event.elapsedMs(p.events[0], p.events[p.n]) catch 0; + p.wall_ms += wall; + p.calls += 1; + if (p.calls % p.every != 0) return; + const k: f64 = @floatFromInt(p.calls); + var buf: [512]u8 = undefined; + var w: std.Io.Writer = .fixed(&buf); + w.print("glm prof {s} ({d} calls): wall {d:.2} ms, gpu {d:.2} ms a call;", .{ p.label, p.calls, p.wall_ms / k, p.gpu_ms / k }) catch {}; + for (names, p.sections) |nm, ms| if (ms > 0) w.print(" {s} {d:.2}", .{ nm, ms / k }) catch {}; + std.log.info("{s}", .{w.buffered()}); + p.sections = @splat(0); + p.gpu_ms = 0; + p.wall_ms = 0; + p.calls = 0; + } +}; diff --git a/zig/src/families/glm/cuda_ranks.zig b/zig/src/families/glm/cuda_ranks.zig index ab5b01b770..5d4d79f0c5 100644 --- a/zig/src/families/glm/cuda_ranks.zig +++ b/zig/src/families/glm/cuda_ranks.zig @@ -60,7 +60,7 @@ pub fn environment(name: [:0]const u8) ?[]const u8 { /// What every rank must agree on, as bytes the communicator carries. pub const Fingerprint = extern struct { - layout: u32 = 5, // this struct's version + layout: u32 = 6, // this struct's version world: u32, rank: u32, split: u32, @@ -68,6 +68,7 @@ pub const Fingerprint = extern struct { pieces: u32, layers: u32, // TF_GLM_LAYERS (0: all) mtp: u32, // 1: the MTP head loads and runs (drafts on): every prompt pass adds its collectives + vision: u32 = 0, // --vision's image tokens (0: off): rank 0 sends image rows every rank must hold capacity: u64, // tokens a sequence holds heat_high: f64, // the heat bands in degrees (0 and 0: off): every rank must wait alike heat_low: f64, @@ -84,6 +85,7 @@ pub const Fingerprint = extern struct { pieces, layers, mtp, + vision, capacity, heat_high, heat_low, @@ -138,6 +140,7 @@ pub fn mismatch(all: []const Fingerprint) ?Mismatch { "pieces", "layers", "mtp", + "vision", "capacity", "heat_high", "heat_low", @@ -213,6 +216,7 @@ fn explain(f: Fingerprint.Field) []const u8 { .pieces => "TF_GLM_OVERLAP_PIECES", .layers => "TF_GLM_LAYERS", .mtp => "--no-drafts given to some ranks and not others", + .vision => "--vision (or its image tokens) given to some ranks and not others", .capacity => "the sequence capacity (the context)", .heat_high => "TF_HEAT_HIGH", .heat_low => "TF_HEAT_LOW", @@ -264,4 +268,7 @@ test "fingerprints: agreeing ranks pass; the first differing field and rank are try std.testing.expectEqual(Fingerprint.Field.checkpoint, mismatch(&all).?.field); all[1] = Fingerprint.of(2, 1, .{}, 8192, false, h0, h0, "/hub/snapshots/423acf37"); // --no-drafts on rank 1 only try std.testing.expectEqual(Mismatch{ .field = .mtp, .rank = 1 }, mismatch(&all).?); + all[1] = Fingerprint.of(2, 1, .{}, 8192, true, h0, h0, "/hub/snapshots/423acf37"); + all[1].vision = 4096; // --vision on rank 1 only + try std.testing.expectEqual(Mismatch{ .field = .vision, .rank = 1 }, mismatch(&all).?); } diff --git a/zig/src/families/glm/cuda_shard.zig b/zig/src/families/glm/cuda_shard.zig index 6b00faac54..812cfb508a 100644 --- a/zig/src/families/glm/cuda_shard.zig +++ b/zig/src/families/glm/cuda_shard.zig @@ -35,7 +35,7 @@ pub const Share = struct { .kda = try span(c.kda_heads, 1, rank, world), .mla = try span(c.mla_heads, 1, rank, world), .dense = try span(c.dense_inter, unit, rank, world), - .moe = try span(c.moe_inter, unit, rank, world), + .moe = try span(c.moe_inter, if (c.exl3) 128 else unit, rank, world), // EXL3 tiles and rotations: 128 .shared = try span(c.shared_inter, unit, rank, world), .vocab = try span(c.vocab, c.head_unit, rank, world), }; diff --git a/zig/src/families/glm/cuda_state.zig b/zig/src/families/glm/cuda_state.zig index 32559505c3..b90123aa15 100644 --- a/zig/src/families/glm/cuda_state.zig +++ b/zig/src/families/glm/cuda_state.zig @@ -2,6 +2,7 @@ const std = @import("std"); const cuda = @import("cuda"); +const exl3 = @import("cuda_exl3.zig"); const weights = @import("cuda_weights.zig"); const kda = @import("cuda_kda.zig"); const b16 = @import("cuda_b16.zig"); @@ -46,7 +47,7 @@ fn arena(d: *const cuda.Driver, sizes: []const usize, out: []u64) !Arena { } /// A rank's B16 matmul shapes ([n, k]) whose split-K partials the buffers hold. -fn b16Shapes(w: *const weights.Weights) [10][2]usize { +fn b16Shapes(w: *const weights.Weights) [12][2]usize { const c = &w.config; const s = &w.share; const n = s.kda.len() * c.kda_dim; @@ -56,6 +57,8 @@ fn b16Shapes(w: *const weights.Weights) [10][2]usize { .{ c.q_lora + c.kv_lora, c.hidden }, .{ s.mla.len() * c.nope, c.q_lora }, .{ c.i_dim + c.i_heads, c.hidden }, .{ c.i_heads * c.i_dim, c.q_lora }, .{ c.hidden, s.mla.len() * c.v_dim }, .{ 2 * s.shared.len(), c.hidden }, .{ c.hidden, 2 * c.hidden }, // the MTP head's eh_proj + if (c.exl3) .{ 2 * s.dense.len(), c.hidden } else .{ c.hidden, 2 * c.hidden }, // bf16 dense layers: [gate | up] + if (c.exl3) .{ c.hidden, s.dense.len() } else .{ c.hidden, 2 * c.hidden }, // and down }; } @@ -123,6 +126,7 @@ pub const Buffers = struct { var mm: usize = 4; // windows under rows_from rows keep their K slices apart (a prompt's tail chunk on the prefill buffers too) for (b16Shapes(w)) |sh_| mm = @max(mm, b16.partBytes(@min(rows, b16.rows_from - 1), sh_[0], sh_[1])); + const x3: [4]usize = if (c.exl3) exl3.scratchBytes(rows, D, ew, prefill, w.exl3_shx) else @splat(4); const hist_len = if (pairs > cuda.grouped.small) (pairs + cuda.grouped.small - 1) / cuda.grouped.small * (count + 1) else 1; const sizes = [_]usize{ rows * 4, rows * S * D * 2, rows * D * 2, rows * (D / 64) * 4, @@ -160,10 +164,11 @@ pub const Buffers = struct { rows * (D / 64) * 4, (if (prefill) 1 else rows) * s.vocab.len() * 2, rows * D * 2, rows * D * 2, rows * 2 * D * 2, rows * D * 2, // hin, me, mcat, mx rows * D * 4, // part: this rank's fp32 partial rows a split exchange sends from + x3[0], x3[1], x3[2], x3[3], // EXL3 experts: 8-slot picks, rotated inputs, a decode window's partials }; var p: [sizes.len]u64 = undefined; const a = try arena(d, &sizes, &p); - return .{ + var out: Buffers = .{ .arena = a, .rows = rows, .prefill = prefill, @@ -201,7 +206,7 @@ pub const Buffers = struct { .mm_part = p[42], .at = p[43], }, - .dense = .{ .tmp = p[44], .gu = p[45], .act = p[46], .xs_act = p[47] }, + .dense = .{ .tmp = p[44], .gu = p[45], .act = p[46], .xs_act = p[47], .mm_part = p[42] }, .moe = .{ .mlog = p[48], .router_part = p[26], .pick = p[49], .wts = p[50], .plan = .{ .members = p[51], .items = p[52], .counts = p[53], .rank = p[54], .hist = p[55] }, .eact = p[56], .ey = p[57], .sgu = p[58], .sact = p[59], .sxs = p[60], .sy = p[61], .mm_part = p[42] }, .hidden = p[62], .fnormed = p[63], @@ -213,6 +218,8 @@ pub const Buffers = struct { .mx = p[69], .part = p[70], }; + out.moe.x3 = .{ .picks = p[71], .xg = p[72], .xu = p[73], .z = p[74] }; + return out; } pub fn deinit(b: *Buffers) void { diff --git a/zig/src/families/glm/cuda_vision.zig b/zig/src/families/glm/cuda_vision.zig new file mode 100644 index 0000000000..6cf2825221 --- /dev/null +++ b/zig/src/families/glm/cuda_vision.zig @@ -0,0 +1,403 @@ +//! GLM-5.3-Flash's image tower on CUDA as glm5_next/vision.py: patches to the language model's rows (rank 0 encodes). + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); + +const H: i64 = 1024; // hidden +const HEADS: i64 = 16; +const HD: i64 = 64; // a head's width +const MLP: i64 = 4096; +const PATCH: i64 = 1176; // 3 channels x 2 frames x 14 x 14 +const MERGE: i64 = 2; +const OUT: i64 = 4096; // the merged rows' width (the language model's) +const PROJ: i64 = 10240; // the merger's SwiGLU width +pub const DEPTH = 24; +const CHUNK: i64 = 2048; // query rows a scores product covers +const ROWS: i64 = 4096; // rows a linear's fp32 product covers (the scratch's bound) +const EPS: f32 = 1e-5; +const LIMIT: f32 = 10.0; + +const Block = struct { n1: u64, n2: u64, qkv_w: u64, qkv_b: u64, qn: u64, kn: u64, proj_w: u64, proj_b: u64, g_w: u64, g_b: u64, u_w: u64, u_b: u64, d_w: u64, d_b: u64 }; + +const LinKey = struct { m: u64, n: u64, k: u64, f32_out: bool }; + +/// glm_vision's kernels and the stream they run on. +pub const Run = struct { + stream: cuda.Stream, + f: Funcs, +}; + +pub const Funcs = struct { + to_bf16: cuda.Function, + bias: cuda.Function, + embed32: cuda.Function, + rmsnorm: cuda.Function, + qkv: cuda.Function, + softmax: cuda.Function, + head_out: cuda.Function, + swiglu: cuda.Function, + ln_gelu: cuda.Function, + inject: cuda.Function, + + pub fn load(m: cuda.Module) !Funcs { + var f: Funcs = undefined; + inline for (@typeInfo(Funcs).@"struct".field_names) |name| @field(f, name) = try m.function("glm_vis_" ++ name); + return f; + } +}; + +fn launch(r: *const Run, f: cuda.Function, grid: cuda.Dim3, threads: u32, a: *cuda.Args) !void { + try cuda.launch.launch(f, .{ .grid = grid, .block = .{ .x = threads } }, r.stream, a); +} + +pub const Tower = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + lt: cuda.cublaslt.Library, + mem: cuda.DeviceBuffer, + blocks: [DEPTH]Block = undefined, + patch_w: u64 = 0, + patch_b: u64 = 0, + post: u64 = 0, + down_w: u64 = 0, // [4096, 2, 2, 1024]: the conv's taps in the merge block's (row, col) order, then the channels + down_b: u64 = 0, + m_proj: u64 = 0, + m_ln_w: u64 = 0, + m_ln_b: u64 = 0, + m_gate: u64 = 0, + m_up: u64 = 0, + m_down: u64 = 0, + inv_freq: cuda.DeviceBuffer, + linears: std.AutoHashMap(LinKey, cuda.cublaslt.Linear), + ws: cuda.DeviceBuffer, + ws_len: usize = 64 << 20, + bytes: u64 = 0, + stages: ?[]f32 = null, // checks: the fp32 stream after the patch embedding and after every block, read back + + /// The tower's tensors (model.visual.*, bf16) from the checkpoint's shards in ``dir``, in one allocation. + pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, dir: []const u8) !Tower { + var ck = try core.Checkpoint.openModelPrefix(gpa, io, dir, "model.visual."); + defer ck.close(); + var total: u64 = 0; + for (ck.files.items) |f| for (f.names.keys()) |name| if (std.mem.startsWith(u8, name, "model.visual.")) { + total += std.mem.alignForward(u64, (try ck.get(name)).bytes.len, 256); + }; + var t: Tower = .{ .gpa = gpa, .d = d, .lt = try cuda.cublaslt.Library.open(), .mem = try cuda.DeviceBuffer.alloc(d, total), .inv_freq = try cuda.DeviceBuffer.alloc(d, 16 * 4), .linears = .init(gpa), .ws = undefined, .bytes = total }; + errdefer t.deinit(); + t.ws = try cuda.DeviceBuffer.alloc(d, t.ws_len); + var at: u64 = 0; + var names = std.StringHashMap(u64).init(gpa); + defer names.deinit(); + for (ck.files.items) |f| for (f.names.keys()) |name| if (std.mem.startsWith(u8, name, "model.visual.")) { + const ten = try ck.get(name); + if (ten.dtype != .bf16) return error.UnsupportedVisionDType; + var bytes = ten.bytes; + var perm: []u16 = &.{}; + defer gpa.free(perm); + if (std.mem.eql(u8, name, "model.visual.downsample.weight")) { // [o, c, 2, 2] -> [o, 2, 2, c] + if (ten.rank != 4 or ten.dim(0) != OUT or ten.dim(1) != H) return error.BadVisionShape; + const src = std.mem.bytesAsSlice(u16, @as([]align(1) const u8, bytes)); + perm = try gpa.alloc(u16, src.len); + const c_n: usize = @intCast(H); + for (0..@intCast(OUT)) |o| for (0..c_n) |c| for (0..2) |kh| for (0..2) |kw| { + perm[((o * 2 + kh) * 2 + kw) * c_n + c] = src[((o * c_n + c) * 2 + kh) * 2 + kw]; + }; + bytes = std.mem.sliceAsBytes(perm); + } + const ptr = t.mem.ptr + at; + try d.check(d.api.cuMemcpyHtoD_v2(ptr, bytes.ptr, bytes.len), "vision weights"); + try names.put(name, ptr); + at += std.mem.alignForward(u64, bytes.len, 256); + }; + const get = struct { + fn f(m: *const std.StringHashMap(u64), comptime fmt: []const u8, args: anytype) !u64 { + var nb: [128]u8 = undefined; + const name = try std.fmt.bufPrint(&nb, "model.visual." ++ fmt, args); + return m.get(name) orelse { + std.log.err("glm vision: no tensor {s}", .{name}); + return error.MissingVisionTensor; + }; + } + }.f; + t.patch_w = try get(&names, "patch_embed.proj.weight", .{}); + t.patch_b = try get(&names, "patch_embed.proj.bias", .{}); + t.post = try get(&names, "post_layernorm.weight", .{}); + t.down_w = try get(&names, "downsample.weight", .{}); + t.down_b = try get(&names, "downsample.bias", .{}); + t.m_proj = try get(&names, "merger.proj.weight", .{}); + t.m_ln_w = try get(&names, "merger.post_projection_norm.weight", .{}); + t.m_ln_b = try get(&names, "merger.post_projection_norm.bias", .{}); + t.m_gate = try get(&names, "merger.gate_proj.weight", .{}); + t.m_up = try get(&names, "merger.up_proj.weight", .{}); + t.m_down = try get(&names, "merger.down_proj.weight", .{}); + for (&t.blocks, 0..) |*bl, i| { + bl.* = .{ + .n1 = try get(&names, "blocks.{d}.norm1.weight", .{i}), + .n2 = try get(&names, "blocks.{d}.norm2.weight", .{i}), + .qkv_w = try get(&names, "blocks.{d}.attn.qkv.weight", .{i}), + .qkv_b = try get(&names, "blocks.{d}.attn.qkv.bias", .{i}), + .qn = try get(&names, "blocks.{d}.attn.q_norm.weight", .{i}), + .kn = try get(&names, "blocks.{d}.attn.k_norm.weight", .{i}), + .proj_w = try get(&names, "blocks.{d}.attn.proj.weight", .{i}), + .proj_b = try get(&names, "blocks.{d}.attn.proj.bias", .{i}), + .g_w = try get(&names, "blocks.{d}.mlp.gate_proj.weight", .{i}), + .g_b = try get(&names, "blocks.{d}.mlp.gate_proj.bias", .{i}), + .u_w = try get(&names, "blocks.{d}.mlp.up_proj.weight", .{i}), + .u_b = try get(&names, "blocks.{d}.mlp.up_proj.bias", .{i}), + .d_w = try get(&names, "blocks.{d}.mlp.down_proj.weight", .{i}), + .d_b = try get(&names, "blocks.{d}.mlp.down_proj.bias", .{i}), + }; + } + // the rotary's frequencies: 1 / 10000^(i / 16), i = 0, 2 .. 30 over 32 (head width / 2) + var inv: [16]f32 = undefined; + for (&inv, 0..) |*v, i| v.* = 1.0 / std.math.pow(f32, 10000.0, @as(f32, @floatFromInt(2 * i)) / 32.0); + try t.inv_freq.upload(0, std.mem.sliceAsBytes(&inv)); + return t; + } + + pub fn deinit(t: *Tower) void { + var it = t.linears.valueIterator(); + while (it.next()) |l| l.deinit(); + t.linears.deinit(); + t.mem.free(); + t.inv_freq.free(); + t.ws.free(); + t.lt.close(); + } + + /// Stage `i`'s fp32 stream (N rows) into `dst` (checks only: it waits for the GPU). + fn stage(t: *Tower, ru: *const Run, dst: []f32, i: usize, x32: u64, N: i64) !void { + const n: usize = @intCast(N * H); + if ((i + 1) * n > dst.len) return; + try ru.stream.synchronize(); + try t.d.check(t.d.api.cuMemcpyDtoH_v2(@ptrCast(dst[i * n ..].ptr), x32, n * 4), "vision stage"); + } + + /// D[m, n] = X[m, k] . W[n, k]^T (fp32 or bf16 out), a plan per shape kept for the next image. + fn linear(t: *Tower, ru: *const Run, x: u64, w: u64, d: u64, m: i64, n: i64, kk: i64, f32_out: bool) !void { + const key: LinKey = .{ .m = @intCast(m), .n = @intCast(n), .k = @intCast(kk), .f32_out = f32_out }; + const gop = try t.linears.getOrPut(key); + if (!gop.found_existing) { + gop.value_ptr.* = cuda.cublaslt.Linear.init(&t.lt, key.m, key.n, key.k, if (f32_out) .f32 else .bf16, t.ws_len) catch |e| { + _ = t.linears.remove(key); + return e; + }; + } + try gop.value_ptr.run(x, w, d, t.ws.ptr, t.ws_len, ru.stream.handle); + } + + fn grid1(n: i64) cuda.Dim3 { + return .{ .x = @intCast(@min(std.math.divCeil(i64, n, 256) catch unreachable, 65535)) }; + } + + /// out = bf16(x . W^T + bias) (bias 0: none), in row chunks through ``acc``; with ``x32``, x32 += that instead. + fn dense(t: *Tower, ru: *const Run, x: u64, w: u64, bias: u64, out: u64, x32: u64, acc: u64, m: i64, n: i64, kk: i64) !void { + var r: i64 = 0; + while (r < m) : (r += ROWS) { + const rows = @min(ROWS, m - r); + try t.linear(ru, x + @as(u64, @intCast(r * kk * 2)), w, acc, rows, n, kk, true); + var a: cuda.Args = .{}; + a.add(acc); + a.add(bias); + a.add(if (out != 0) out + @as(u64, @intCast(r * n * 2)) else 0); + a.add(if (x32 != 0) x32 + @as(u64, @intCast(r * n * 4)) else 0); + a.add(@as(i64, rows * n)); + a.add(@as(i32, @intCast(n))); + try launch(ru, ru.f.bias, grid1(rows * n), 256, &a); + } + } + + /// The limited SwiGLU of two linears (biases 0: none) over x [m, kk] -> out [m, n] bf16, in row chunks. + fn swiglu(t: *Tower, ru: *const Run, x: u64, gw: u64, gb: u64, uw: u64, ub: u64, out: u64, acc: u64, acc2: u64, m: i64, n: i64, kk: i64) !void { + var r: i64 = 0; + const step = @divFloor(ROWS * MLP, n); // the fp32 products' rows within the scratch + while (r < m) : (r += step) { + const rows = @min(step, m - r); + const xr = x + @as(u64, @intCast(r * kk * 2)); + try t.linear(ru, xr, gw, acc, rows, n, kk, true); + try t.linear(ru, xr, uw, acc2, rows, n, kk, true); + var a: cuda.Args = .{}; + a.add(acc); + a.add(acc2); + a.add(gb); + a.add(ub); + a.add(out + @as(u64, @intCast(r * n * 2))); + a.add(@as(i64, rows * n)); + a.add(@as(i32, @intCast(n))); + a.add(LIMIT); + try launch(ru, ru.f.swiglu, grid1(rows * n), 256, &a); + } + } + + fn rmsnorm(ru: *const Run, x32: u64, w: u64, y: u64, rows: i64, cols: i64) !void { + var a: cuda.Args = .{}; + a.add(x32); + a.add(w); + a.add(y); + a.add(@as(i32, @intCast(cols))); + a.add(EPS); + try launch(ru, ru.f.rmsnorm, .{ .x = @intCast(rows) }, 256, &a); + } + + /// Scratch bytes ``encode`` takes for ``patches`` patches in segments of ``seg``. + pub fn scratchBytes(patches: i64, seg: i64) u64 { + const N = patches; + const q = @min(CHUNK, seg); + return @intCast(N * PATCH * 4 + N * H * 4 + N * H * 2 * 5 + N * MLP * 2 + 2 * ROWS * MLP * 4 + q * seg * 6 + q * HD * 4 + q * 4 + N * 8 + 16 * 1024); + } + + /// Features of an image or video (``grid`` [t, h, w]): fp32 patches [t h w, 1176] -> ``out`` [t h w / 4, 4096]. + pub fn encode(t: *Tower, ru: *const Run, patches: []const f32, grid: [3]i64, out: u64) !void { + const d = t.d; + const gt = grid[0]; + const gh = grid[1]; + const gw = grid[2]; + if (gt < 1 or gh < MERGE or gw < MERGE or @mod(gh, MERGE) != 0 or @mod(gw, MERGE) != 0) return error.BadImageGrid; + const seg = gh * gw; + const N = gt * seg; + if (@as(i64, @intCast(patches.len)) != N * PATCH) return error.BadImagePatches; + // each patch's (row, col) in merge-block order, every frame group alike + const pos = try t.gpa.alloc(i32, @intCast(N * 2)); + defer t.gpa.free(pos); + var i: i64 = 0; + while (i < N) : (i += 1) { + const j = @mod(i, seg); + const in_col = @mod(j, MERGE); + const in_row = @mod(@divTrunc(j, MERGE), MERGE); + const block_col = @mod(@divTrunc(j, MERGE * MERGE), @divTrunc(gw, MERGE)); + const block_row = @divTrunc(j, MERGE * MERGE * @divTrunc(gw, MERGE)); + pos[@intCast(i * 2)] = @intCast(block_row * MERGE + in_row); + pos[@intCast(i * 2 + 1)] = @intCast(block_col * MERGE + in_col); + } + var scratch = try cuda.DeviceBuffer.alloc(d, scratchBytes(N, seg)); + defer scratch.free(); + var at: u64 = scratch.ptr; + const take = struct { + fn f(p: *u64, n: i64) u64 { + const r = p.*; + p.* += std.mem.alignForward(u64, @intCast(n), 256); + return r; + } + }.f; + const q_rows = @min(CHUNK, seg); + const p32 = take(&at, N * PATCH * 4); + const x32 = take(&at, N * H * 4); + const xn = take(&at, N * H * 2); + const qh = take(&at, N * H * 2); + const kh = take(&at, N * H * 2); + const vt = take(&at, N * H * 2); + const att = take(&at, N * H * 2); + const big = take(&at, N * MLP * 2); + const acc = take(&at, ROWS * MLP * 4); + const acc2 = take(&at, ROWS * MLP * 4); + const sc = take(&at, q_rows * seg * 4); + const pr = take(&at, q_rows * seg * 2); + const oh = take(&at, q_rows * HD * 4); + const sums = take(&at, q_rows * 4); + const dpos = take(&at, N * 8); + try d.check(d.api.cuMemcpyHtoDAsync_v2(dpos, pos.ptr, pos.len * 4, ru.stream.handle), "vision positions"); + try d.check(d.api.cuMemcpyHtoDAsync_v2(p32, patches.ptr, patches.len * 4, ru.stream.handle), "vision patches"); + { // the patch embedding: the patches as bf16, then the conv as a linear (with its bias) into the fp32 stream + var a: cuda.Args = .{}; + a.add(p32); + a.add(big); + a.add(@as(i64, N * PATCH)); + try launch(ru, ru.f.to_bf16, grid1(N * PATCH), 256, &a); + var r: i64 = 0; + while (r < N) : (r += ROWS) { + const rows = @min(ROWS, N - r); + try t.linear(ru, big + @as(u64, @intCast(r * PATCH * 2)), t.patch_w, acc, rows, H, PATCH, true); + var b: cuda.Args = .{}; + b.add(acc); + b.add(t.patch_b); + b.add(x32 + @as(u64, @intCast(r * H * 4))); + b.add(@as(i64, rows * H)); + b.add(@as(i32, @intCast(H))); + try launch(ru, ru.f.embed32, grid1(rows * H), 256, &b); + } + } + const scale: f32 = 1.0 / @sqrt(@as(f32, HD)); + if (t.stages) |dst| try t.stage(ru, dst, 0, x32, N); + for (t.blocks, 1..) |bl, bi| { + defer if (t.stages) |dst| t.stage(ru, dst, bi, x32, N) catch {}; + try rmsnorm(ru, x32, bl.n1, xn, N, H); + try t.dense(ru, xn, bl.qkv_w, bl.qkv_b, big, 0, acc, N, 3 * H, H); + { + var a: cuda.Args = .{}; + a.add(big); + a.add(bl.qn); + a.add(bl.kn); + a.add(dpos); + a.add(t.inv_freq.ptr); + a.add(qh); + a.add(kh); + a.add(vt); + a.add(@as(i32, @intCast(N))); + a.add(@as(i32, @intCast(seg))); + a.add(@as(i32, @intCast(HEADS))); + a.add(EPS); + try launch(ru, ru.f.qkv, .{ .x = @intCast(std.math.divCeil(i64, N * HEADS, 8) catch unreachable) }, 256, &a); + } + // each frame group attends within itself, per head, in query chunks (q, k [t, H, seg, 64]; v transposed) + var g: i64 = 0; + while (g < gt) : (g += 1) { + var h: i64 = 0; + while (h < HEADS) : (h += 1) { + const blk: u64 = @intCast((g * HEADS + h) * seg * HD * 2); + var c: i64 = 0; + while (c < seg) : (c += CHUNK) { + const rows = @min(CHUNK, seg - c); + try t.linear(ru, qh + blk + @as(u64, @intCast(c * HD * 2)), kh + blk, sc, rows, seg, HD, true); + var a: cuda.Args = .{}; + a.add(sc); + a.add(pr); + a.add(sums); + a.add(@as(i32, @intCast(seg))); + a.add(scale); + try launch(ru, ru.f.softmax, .{ .x = @intCast(rows) }, 256, &a); + try t.linear(ru, pr, vt + blk, oh, rows, HD, seg, true); + var b: cuda.Args = .{}; + b.add(oh); + b.add(sums); + b.add(att + @as(u64, @intCast((g * seg + c) * H * 2))); + b.add(@as(i32, @intCast(rows))); + b.add(@as(i32, @intCast(HEADS))); + b.add(@as(i32, @intCast(h))); + try launch(ru, ru.f.head_out, grid1(rows * HD), 256, &b); + } + } + } + try t.dense(ru, att, bl.proj_w, bl.proj_b, 0, x32, acc, N, H, H); + try rmsnorm(ru, x32, bl.n2, xn, N, H); + try t.swiglu(ru, xn, bl.g_w, bl.g_b, bl.u_w, bl.u_b, big, acc, acc2, N, MLP, H); + try t.dense(ru, big, bl.d_w, bl.d_b, 0, x32, acc, N, H, MLP); + } + // the post-norm, then each merge block's four patches (adjacent rows) through the 2 x 2 conv as one linear + try rmsnorm(ru, x32, t.post, xn, N, H); + const M = @divExact(N, MERGE * MERGE); + try t.dense(ru, xn, t.down_w, t.down_b, att, 0, acc, M, OUT, 4 * H); + // the merger: projection, LayerNorm + GELU, the limited SwiGLU to 10240 and back + var r: i64 = 0; + while (r < M) : (r += ROWS) { + const rows = @min(ROWS, M - r); + try t.linear(ru, att + @as(u64, @intCast(r * OUT * 2)), t.m_proj, acc, rows, OUT, OUT, true); + var a: cuda.Args = .{}; + a.add(acc); + a.add(t.m_ln_w); + a.add(t.m_ln_b); + a.add(xn + @as(u64, @intCast(r * OUT * 2))); + a.add(@as(i32, @intCast(OUT))); + a.add(EPS); + try launch(ru, ru.f.ln_gelu, .{ .x = @intCast(rows) }, 256, &a); + } + r = 0; + const step = @divFloor(ROWS * MLP, PROJ); + while (r < M) : (r += step) { + const rows = @min(step, M - r); + try t.swiglu(ru, xn + @as(u64, @intCast(r * OUT * 2)), t.m_gate, 0, t.m_up, 0, big, acc, acc2, rows, PROJ, OUT); + try t.linear(ru, big, t.m_down, out + @as(u64, @intCast(r * OUT * 2)), rows, OUT, PROJ, false); + } + try ru.stream.synchronize(); // the scratch is freed on return + } +}; diff --git a/zig/src/families/glm/cuda_weights.zig b/zig/src/families/glm/cuda_weights.zig index 80bc44a8e0..fb9fbee479 100644 --- a/zig/src/families/glm/cuda_weights.zig +++ b/zig/src/families/glm/cuda_weights.zig @@ -12,6 +12,8 @@ const mlp = @import("cuda_mlp.zig"); const moe = @import("cuda_moe.zig"); const quant = @import("cuda_quant.zig"); const q4 = @import("cuda_q4.zig"); +const exl3 = @import("cuda_exl3.zig"); +const mm_mod = @import("cuda_mm.zig"); const Tensor = core.checkpoint.Tensor; const prefix = "model.language_model."; @@ -42,6 +44,7 @@ pub const Weights = struct { norm: u64 = 0, head: u64 = 0, // bf16 [share.vocab][dims] mtp: ?Mtp = null, + exl3_shx: bool = true, // EXL3: every layer's gate and up share one input rotation (a prompt rotates each row once) buffers: std.ArrayList(cuda.DeviceBuffer) = .empty, named: std.ArrayList(Named) = .empty, @@ -79,13 +82,13 @@ const Loader = struct { host: std.ArrayList(u8) = .empty, name_buf: [192]u8 = undefined, - fn get(L: *Loader, comptime fmt: []const u8, args: anytype) !Tensor { + pub fn get(L: *Loader, comptime fmt: []const u8, args: anytype) !Tensor { const name = try std.fmt.bufPrint(&L.name_buf, prefix ++ fmt, args); return L.ck.get(name); } /// The host bytes gathered so far uploaded as one tensor named `fmt`. - fn put(L: *Loader, comptime fmt: []const u8, args: anytype) !u64 { + pub fn put(L: *Loader, comptime fmt: []const u8, args: anytype) !u64 { const b = try cuda.DeviceBuffer.fromHost(L.d, L.host.items); try L.w.buffers.append(L.gpa, b); try L.w.named.append(L.gpa, .{ .name = try std.fmt.allocPrint(L.gpa, fmt, args), .ptr = b.ptr, .len = b.len }); @@ -258,10 +261,18 @@ const Loader = struct { const c = &L.w.config; const s = L.w.share.dense; const d = c.hidden; - return .{ + const w: mlp.Weights = if (c.exl3) bf16: { // an EXL3 checkpoint keeps its dense layers bf16 + inline for (.{ "gate_proj", "up_proj" }) |x| try L.rows(try L.get("layers.{d}.mlp." ++ x ++ ".weight", .{i}), s.lo, s.hi, c.dense_inter, .bf16); + const gu = try L.put("layers.{d}.mlp.gu", .{i}); + try L.cols(try L.get("layers.{d}.mlp.down_proj.weight", .{i}), s.lo, s.hi, .bf16); + break :bf16 .{ .bf16 = .{ .gu = gu, .down = try L.put("layers.{d}.mlp.down", .{i}) } }; + } else .{ .nvfp4 = .{ .gate = try L.fp4("layers.{d}.mlp.gate_proj", i, s.lo, s.hi, 0, d), .up = try L.fp4("layers.{d}.mlp.up_proj", i, s.lo, s.hi, 0, d), .down = try L.fp4("layers.{d}.mlp.down_proj", i, 0, d, s.lo, s.hi), + } }; + return .{ + .w = w, .width = s.len(), .dims = d, .limit = c.swiglu_limit, @@ -344,6 +355,39 @@ const Loader = struct { try L.whole(try L.get("layers.{d}.mlp.gate.e_score_correction_bias", .{i}), .f32); const bias = try L.put("layers.{d}.moe.bias", .{i}); const count: usize = c.experts; + var routed: cuda.experts.Layer = .{ .format = .nvfp4, .up = 0, .down = 0, .width = ni, .dims = d, .experts = count, .limit = c.swiglu_limit }; + var x3: ?exl3.Experts = null; + if (c.exl3) { // the MTP layer's experts too + x3 = try exl3.load(L, i, count, d, c.moe_inter, cut); + if (!x3.?.shx) L.w.exl3_shx = false; + } else routed = try L.nvfp4Experts(i); + const sh = L.w.share.shared; + inline for (.{ "gate_proj", "up_proj" }) |x| try L.rows(try L.get("layers.{d}.mlp.shared_experts." ++ x ++ ".weight", .{i}), sh.lo, sh.hi, c.shared_inter, .bf16); + const gu = try L.put("layers.{d}.moe.shared.gu", .{i}); + try L.cols(try L.get("layers.{d}.mlp.shared_experts.down_proj.weight", .{i}), sh.lo, sh.hi, .bf16); + const sd = try L.put("layers.{d}.moe.shared.down", .{i}); + return .{ + .router = router, + .bias = bias, + .routed = routed, + .exl3 = x3, + .count = count, + .shared_gu = gu, + .shared_down = sd, + .shared = sh.len(), + .top_k = c.topk, + .scale = c.routed_scale, + .norm = c.norm_topk, + }; + } + + /// Layer `i`'s routed NVFP4 experts packed as Experts4 (the MTP layer's quantized from its BF16 ones). + fn nvfp4Experts(L: *Loader, i: u32) !cuda.experts.Layer { + const c = &L.w.config; + const d: usize = c.hidden; + const cut = L.w.share.moe; + const ni = cut.len(); + const count: usize = c.experts; const up_len = X.bytes(1, ni, d, 2) / 4; const down_len = X.bytes(1, d, ni, 1) / 4; const up = try L.gpa.alloc(u32, count * up_len); @@ -383,23 +427,7 @@ const Loader = struct { try L.host.appendSlice(L.gpa, bytes); p.* = try L.put("layers.{d}.moe.experts.{s}", .{ i, what }); } - const sh = L.w.share.shared; - inline for (.{ "gate_proj", "up_proj" }) |x| try L.rows(try L.get("layers.{d}.mlp.shared_experts." ++ x ++ ".weight", .{i}), sh.lo, sh.hi, c.shared_inter, .bf16); - const gu = try L.put("layers.{d}.moe.shared.gu", .{i}); - try L.cols(try L.get("layers.{d}.mlp.shared_experts.down_proj.weight", .{i}), sh.lo, sh.hi, .bf16); - const sd = try L.put("layers.{d}.moe.shared.down", .{i}); - return .{ - .router = router, - .bias = bias, - .routed = .{ .format = .nvfp4, .up = ptrs[0], .down = ptrs[1], .up_scale = ptrs[2], .down_scale = ptrs[3], .width = ni, .dims = d, .experts = count, .limit = c.swiglu_limit }, - .count = count, - .shared_gu = gu, - .shared_down = sd, - .shared = sh.len(), - .top_k = c.topk, - .scale = c.routed_scale, - .norm = c.norm_topk, - }; + return .{ .format = .nvfp4, .up = ptrs[0], .down = ptrs[1], .up_scale = ptrs[2], .down_scale = ptrs[3], .width = ni, .dims = d, .experts = count, .limit = c.swiglu_limit }; } /// weights.load's MTPW: layer `layers` as plain DSA + MoE, norms, eh_proj, draft head from the rank's head rows. @@ -497,3 +525,63 @@ test "bf16 rounding of fp32 is to nearest even" { try std.testing.expectEqual(@as(u16, 0x3f82), toBf16(@bitCast(@as(u32, 0x3f818000)))); // a tie to even (up) try std.testing.expectEqual(@as(u16, 0x3f81), toBf16(@bitCast(@as(u32, 0x3f808001)))); } + +/// Every bf16 linear b16.matmul reads, retiled in place for glm_mm (the q4 draft head was made from the head already). +pub fn tileLinears(w: *const Weights, mk: *const mm_mod.Kernels, s: cuda.Stream, d: *const cuda.Driver) !void { + var list: std.ArrayList([3]usize) = .empty; // (pointer, n, k) + defer list.deinit(w.gpa); + const add = struct { + fn f(l: *std.ArrayList([3]usize), gpa: std.mem.Allocator, ptr: u64, n: usize, k: usize) !void { + if (ptr != 0) try l.append(gpa, .{ ptr, n, k }); + } + }.f; + const kd = @import("cuda_kda.zig").dk; + var layers: std.ArrayList(Layer) = .empty; + defer layers.deinit(w.gpa); + try layers.appendSlice(w.gpa, w.layers); + if (w.mtp) |m| { + try layers.append(w.gpa, m.layer); + try add(&list, w.gpa, m.eh, w.config.hidden, 2 * w.config.hidden); + } + for (layers.items) |l| { + switch (l.mixer) { + .kda => |x| { + const n = x.heads * kd; + try add(&list, w.gpa, x.proj, x.width(), x.dims); + try add(&list, w.gpa, x.fb, n, kd); + try add(&list, w.gpa, x.gb, n, kd); + try add(&list, w.gpa, x.o, x.dims, n); + }, + .mla => |x| { + try add(&list, w.gpa, x.proj, x.q_lora + x.lw, x.dims); + try add(&list, w.gpa, x.q_b, x.heads * x.qk, x.q_lora); + try add(&list, w.gpa, x.index.kw, x.index.d + x.index.heads, x.dims); + try add(&list, w.gpa, x.index.qb, x.index.heads * x.index.d, x.q_lora); + try add(&list, w.gpa, x.o, x.dims, x.heads * x.dv); + }, + } + switch (l.ffn) { + .dense => |x| switch (x.w) { + .bf16 => |b| { + try add(&list, w.gpa, b.gu, 2 * x.width, x.dims); + try add(&list, w.gpa, b.down, x.dims, x.width); + }, + .nvfp4 => {}, + }, + .moe => |x| { + try add(&list, w.gpa, x.shared_gu, 2 * x.shared, x.routed.dims); + try add(&list, w.gpa, x.shared_down, x.routed.dims, x.shared); + }, + } + } + try add(&list, w.gpa, w.head, w.share.vocab.len(), w.config.hidden); + var most: usize = 0; + for (list.items) |e| most = @max(most, e[1] * e[2] * 2); + var tmp = try cuda.DeviceBuffer.alloc(d, most); + defer tmp.free(); + for (list.items) |e| { + try mm_mod.tile(mk, s, e[0], tmp.ptr, e[1], e[2]); + try d.check(d.api.cuMemcpyDtoDAsync_v2(e[0], tmp.ptr, e[1] * e[2] * 2, s.handle), "cuMemcpyDtoDAsync"); + } + try s.synchronize(); +} diff --git a/zig/src/families/glm/image.zig b/zig/src/families/glm/image.zig index b22d34de11..4d0ee62758 100644 --- a/zig/src/families/glm/image.zig +++ b/zig/src/families/glm/image.zig @@ -1,116 +1,18 @@ //! An image's bytes to GLM5-Next vision patches as TensorFold 0.6.6's --vision made them (Pillow, then mlx-vlm). const std = @import("std"); const picture = @import("../../core/picture/picture.zig"); -const vision = @import("vision.zig"); - -/// One image ready for the tower: its patches, its patch grid and its tokens (gh * gw / 4). -pub const Prepared = struct { - pixels: []f32, // [gh * gw, 1176] - gh: u32, - gw: u32, - tokens: u32, - hash: u64, // its identity (core/picture/identity.zig): bytes and grid, so another budget is another image -}; - -/// The shared decoder's refusals, and a token budget too small for an image. -pub const Error = picture.Error || error{ImageTooSmall}; - -/// The words for a refusal the server cannot word: the JPEG library this machine lacks, or a CMYK JPEG. -pub const decoderProblem = picture.decoderProblem; - -// ---- the processor's geometry ---- - -pub const Geometry = struct { target_h: u32, target_w: u32, content_h: u32, content_w: u32 }; - -fn alignUp(v: u64, f: u64) u64 { - return (v + f - 1) / f * f; -} - -/// smart_resize and _resize_geometry: the padded canvas under `max_tokens` (at least `min_tokens`), the content size. -pub fn geometry(h: u32, w: u32, min_tokens: u32, max_tokens: u32) Error!Geometry { - const factor: u64 = vision.patch * vision.merge; - const frames: u64 = vision.temporal; - const per_token = frames * factor * factor; - const min_px = @as(u64, min_tokens) * per_token; - const max_px = @as(u64, max_tokens) * per_token; - if (min_tokens == 0 or max_tokens < min_tokens) return error.ImageTooSmall; - const H: u64 = h; - const W: u64 = w; - var ah = alignUp(H, factor); - var aw = alignUp(W, factor); - var pixels = frames * ah * aw; - if (pixels < min_px) { - const scale = @sqrt(@as(f64, @floatFromInt(min_px)) / @as(f64, @floatFromInt(frames * H * W))); - ah = alignUp(@max(1, @as(u64, @intFromFloat(@ceil(@as(f64, @floatFromInt(H)) * scale)))), factor); - aw = alignUp(@max(1, @as(u64, @intFromFloat(@ceil(@as(f64, @floatFromInt(W)) * scale)))), factor); - pixels = frames * ah * aw; - } - if (pixels > max_px) { // the tallest content that fits, by bisection on its height - if (max_px < frames * factor * factor) return error.ImageTooSmall; - var best: [2]u64 = .{ factor, factor }; - var low: i64 = 1; - var high: i64 = @intCast(H); - while (low <= high) { - const ch: u64 = @intCast(@divFloor(low + high, 2)); - const cw: u64 = @max(1, W * ch / H); // math.floor(width * content_h / height) - const cand: [2]u64 = .{ alignUp(ch, factor), alignUp(cw, factor) }; - if (frames * cand[0] * cand[1] <= max_px) { - best = cand; - low = @as(i64, @intCast(ch)) + 1; - } else high = @as(i64, @intCast(ch)) - 1; - } - ah = best[0]; - aw = best[1]; - } - var scale = @min(@as(f64, @floatFromInt(ah)) / @as(f64, @floatFromInt(H)), @as(f64, @floatFromInt(aw)) / @as(f64, @floatFromInt(W))); - if (frames * H * W >= min_px) scale = @min(1.0, scale); - const chf: u64 = @intFromFloat(@floor(@as(f64, @floatFromInt(H)) * scale)); - const cwf: u64 = @intFromFloat(@floor(@as(f64, @floatFromInt(W)) * scale)); - return .{ .target_h = @intCast(ah), .target_w = @intCast(aw), .content_h = @intCast(@max(1, @min(ah, chf))), .content_w = @intCast(@max(1, @min(aw, cwf))) }; -} - -/// The resized content padded with black, scaled, normalized and cut into 2x2 windows of 14x14 patches, row-major. -pub fn patchify(gpa: std.mem.Allocator, content: picture.Rgb, g: Geometry, mean: [3]f32, sd: [3]f32) Error![]f32 { - const P = vision.patch; - const M = vision.merge; - const gh = g.target_h / P; - const gw = g.target_w / P; - const out = try gpa.alloc(f32, @as(usize, gh) * gw * vision.patch_width); - const rescale: f32 = @floatCast(@as(f64, 1.0) / 255.0); - var row: usize = 0; - for (0..gh / M) |wy| for (0..gw / M) |wx| for (0..M) |my| for (0..M) |mx| { - const py0 = (wy * M + my) * P; - const px0 = (wx * M + mx) * P; - const dst = out[row * vision.patch_width ..][0..vision.patch_width]; - for (0..3) |ch| for (0..vision.temporal) |t| for (0..P) |dy| for (0..P) |dx| { - const y = py0 + dy; - const x = px0 + dx; - const v: u8 = if (y < content.height and x < content.width) content.pixels[(y * content.width + x) * 3 + ch] else 0; - const f = @as(f32, @floatFromInt(v)) * rescale; - dst[((ch * vision.temporal + t) * P + dy) * P + dx] = (f - mean[ch]) / sd[ch]; - }; - row += 1; - }; - return out; -} - -/// Bytes to patches: decoded as 0.6.6 decoded them (core's picture), sized for `max_tokens`, resized, padded and cut. -pub fn prepare(io: std.Io, gpa: std.mem.Allocator, bytes: []const u8, c: vision.Config, max_tokens: u32) Error!Prepared { - var img = try picture.decode(io, gpa, bytes); - defer img.deinit(gpa); - return prepareRgb(gpa, img, picture.identity.ofBytes(0x1a6e, bytes), c, max_tokens); -} - -/// RGB to patches (Pillow's bicubic resize, the padded canvas cut); the hash covers the bytes' hash and the grid. -pub fn prepareRgb(gpa: std.mem.Allocator, img: picture.Rgb, bytes_hash: u64, c: vision.Config, max_tokens: u32) Error!Prepared { - const g = try geometry(@intCast(img.height), @intCast(img.width), @min(c.min_tokens, max_tokens), max_tokens); - const resized = try picture.resample.resize(gpa, .pillow, .of(img.pixels, img.width, img.height), g.content_w, g.content_h); - defer gpa.free(resized); - const content: picture.Rgb = .{ .width = g.content_w, .height = g.content_h, .pixels = resized }; - const pixels = try patchify(gpa, content, g, c.mean, c.std); - const grid = [2]u32{ g.target_h / vision.patch, g.target_w / vision.patch }; - return .{ .pixels = pixels, .gh = grid[0], .gw = grid[1], .tokens = grid[0] * grid[1] / vision.patches_per_token, .hash = picture.identity.withGrid(bytes_hash, grid) }; -} +const vision = @import("vision_config.zig"); +const core = @import("image_core.zig"); +const steps = core.For(picture); + +pub const Prepared = core.Prepared; +pub const Geometry = core.Geometry; +pub const geometry = core.geometry; +pub const Error = steps.Error; +pub const decoderProblem = steps.decoderProblem; +pub const patchify = steps.patchify; +pub const prepare = steps.prepare; +pub const prepareRgb = steps.prepareRgb; test "geometry follows smart_resize" { // the reference dumps' grids: 517x389 pads to 532x392 unscaled; 96x96 grows to 112; 20x15 to the 16-token floor diff --git a/zig/src/families/glm/image_core.zig b/zig/src/families/glm/image_core.zig new file mode 100644 index 0000000000..3f0a087177 --- /dev/null +++ b/zig/src/families/glm/image_core.zig @@ -0,0 +1,117 @@ +//! An image's bytes to GLM5-Next vision patches as TensorFold 0.6.6's --vision made them, over the caller's decoder. +const std = @import("std"); +const vision = @import("vision_config.zig"); + +/// One image ready for the tower: its patches, its patch grid and its tokens (gh * gw / 4). +pub const Prepared = struct { + pixels: []f32, // [gh * gw, 1176] + gh: u32, + gw: u32, + tokens: u32, + hash: u64, // its identity (core/picture/identity.zig): bytes and grid, so another budget is another image +}; + +// ---- the processor's geometry ---- + +pub const Geometry = struct { target_h: u32, target_w: u32, content_h: u32, content_w: u32 }; + +fn alignUp(v: u64, f: u64) u64 { + return (v + f - 1) / f * f; +} + +/// smart_resize and _resize_geometry: the padded canvas under `max_tokens` (at least `min_tokens`), the content size. +pub fn geometry(h: u32, w: u32, min_tokens: u32, max_tokens: u32) error{ImageTooSmall}!Geometry { + const factor: u64 = vision.patch * vision.merge; + const frames: u64 = vision.temporal; + const per_token = frames * factor * factor; + const min_px = @as(u64, min_tokens) * per_token; + const max_px = @as(u64, max_tokens) * per_token; + if (min_tokens == 0 or max_tokens < min_tokens) return error.ImageTooSmall; + const H: u64 = h; + const W: u64 = w; + var ah = alignUp(H, factor); + var aw = alignUp(W, factor); + var pixels = frames * ah * aw; + if (pixels < min_px) { + const scale = @sqrt(@as(f64, @floatFromInt(min_px)) / @as(f64, @floatFromInt(frames * H * W))); + ah = alignUp(@max(1, @as(u64, @intFromFloat(@ceil(@as(f64, @floatFromInt(H)) * scale)))), factor); + aw = alignUp(@max(1, @as(u64, @intFromFloat(@ceil(@as(f64, @floatFromInt(W)) * scale)))), factor); + pixels = frames * ah * aw; + } + if (pixels > max_px) { // the tallest content that fits, by bisection on its height + if (max_px < frames * factor * factor) return error.ImageTooSmall; + var best: [2]u64 = .{ factor, factor }; + var low: i64 = 1; + var high: i64 = @intCast(H); + while (low <= high) { + const ch: u64 = @intCast(@divFloor(low + high, 2)); + const cw: u64 = @max(1, W * ch / H); // math.floor(width * content_h / height) + const cand: [2]u64 = .{ alignUp(ch, factor), alignUp(cw, factor) }; + if (frames * cand[0] * cand[1] <= max_px) { + best = cand; + low = @as(i64, @intCast(ch)) + 1; + } else high = @as(i64, @intCast(ch)) - 1; + } + ah = best[0]; + aw = best[1]; + } + var scale = @min(@as(f64, @floatFromInt(ah)) / @as(f64, @floatFromInt(H)), @as(f64, @floatFromInt(aw)) / @as(f64, @floatFromInt(W))); + if (frames * H * W >= min_px) scale = @min(1.0, scale); + const chf: u64 = @intFromFloat(@floor(@as(f64, @floatFromInt(H)) * scale)); + const cwf: u64 = @intFromFloat(@floor(@as(f64, @floatFromInt(W)) * scale)); + return .{ .target_h = @intCast(ah), .target_w = @intCast(aw), .content_h = @intCast(@max(1, @min(ah, chf))), .content_w = @intCast(@max(1, @min(aw, cwf))) }; +} + +/// The steps that read pixels, over `picture` (core/picture: the Metal build's relative import, a CUDA build's core). +pub fn For(comptime picture: type) type { + return struct { + /// The shared decoder's refusals, and a token budget too small for an image. + pub const Error = picture.Error || error{ImageTooSmall}; + + /// The words for a refusal the server cannot word: the JPEG library this machine lacks, or a CMYK JPEG. + pub const decoderProblem = picture.decoderProblem; + + /// The resized content padded black, scaled, normalized, cut into 2x2 windows of 14x14 patches, row-major. + pub fn patchify(gpa: std.mem.Allocator, content: picture.Rgb, g: Geometry, mean: [3]f32, sd: [3]f32) Error![]f32 { + const P = vision.patch; + const M = vision.merge; + const gh = g.target_h / P; + const gw = g.target_w / P; + const out = try gpa.alloc(f32, @as(usize, gh) * gw * vision.patch_width); + const rescale: f32 = @floatCast(@as(f64, 1.0) / 255.0); + var row: usize = 0; + for (0..gh / M) |wy| for (0..gw / M) |wx| for (0..M) |my| for (0..M) |mx| { + const py0 = (wy * M + my) * P; + const px0 = (wx * M + mx) * P; + const dst = out[row * vision.patch_width ..][0..vision.patch_width]; + for (0..3) |ch| for (0..vision.temporal) |t| for (0..P) |dy| for (0..P) |dx| { + const y = py0 + dy; + const x = px0 + dx; + const v: u8 = if (y < content.height and x < content.width) content.pixels[(y * content.width + x) * 3 + ch] else 0; + const f = @as(f32, @floatFromInt(v)) * rescale; + dst[((ch * vision.temporal + t) * P + dy) * P + dx] = (f - mean[ch]) / sd[ch]; + }; + row += 1; + }; + return out; + } + + /// Bytes to patches: decoded as 0.6.6 decoded them, sized for `max_tokens`, resized, padded and cut. + pub fn prepare(io: std.Io, gpa: std.mem.Allocator, bytes: []const u8, c: vision.Config, max_tokens: u32) Error!Prepared { + var img = try picture.decode(io, gpa, bytes); + defer img.deinit(gpa); + return prepareRgb(gpa, img, picture.identity.ofBytes(0x1a6e, bytes), c, max_tokens); + } + + /// RGB to patches (Pillow's bicubic resize, the padded canvas cut); the hash covers the bytes and the grid. + pub fn prepareRgb(gpa: std.mem.Allocator, img: picture.Rgb, bytes_hash: u64, c: vision.Config, max_tokens: u32) Error!Prepared { + const g = try geometry(@intCast(img.height), @intCast(img.width), @min(c.min_tokens, max_tokens), max_tokens); + const resized = try picture.resample.resize(gpa, .pillow, .of(img.pixels, img.width, img.height), g.content_w, g.content_h); + defer gpa.free(resized); + const content: picture.Rgb = .{ .width = g.content_w, .height = g.content_h, .pixels = resized }; + const pixels = try patchify(gpa, content, g, c.mean, c.std); + const grid = [2]u32{ g.target_h / vision.patch, g.target_w / vision.patch }; + return .{ .pixels = pixels, .gh = grid[0], .gw = grid[1], .tokens = grid[0] * grid[1] / vision.patches_per_token, .hash = picture.identity.withGrid(bytes_hash, grid) }; + } + }; +} diff --git a/zig/src/families/glm/vision.zig b/zig/src/families/glm/vision.zig index 3665d63f2d..6e6c8e9c86 100644 --- a/zig/src/families/glm/vision.zig +++ b/zig/src/families/glm/vision.zig @@ -10,86 +10,21 @@ const readAll = @import("weights_fix.zig").readAll; const Engine = @import("engine.zig").Engine; const Ref = wts.Ref; -// The tower this file reads; `Config.parse` refuses a checkpoint whose vision_config differs. -pub const hidden = 1024; -pub const heads = 16; -pub const head_dim = 64; -pub const depth = 24; -pub const mlp_inter = 4096; -pub const out_hidden = 4096; // the language model's width -pub const merger_inter = 10240; -pub const patch = 14; -pub const temporal = 2; -pub const merge = 2; -pub const patch_width = 3 * temporal * patch * patch; // 1176: a patch's values, channel by frame by row by column -pub const patches_per_token = merge * merge; - -/// --vision's limits: image tokens a prompt may hold in all (scratch and rows are sized for them), and images. -pub const Limits = struct { image_tokens: u32 = 4096, max_images: u32 = 4 }; - -pub const Config = struct { - eps: f32 = 1e-5, - limit: f32 = 10, // swiglu_limit - image_token: u32, - image_start: u32, - image_end: u32, - video: ?[3]u32 = null, // video_start_token_id, video_token_id, video_end_token_id: the checkpoint reads videos - mean: [3]f32, - std: [3]f32, - min_tokens: u32, - max_tokens: u32, - - /// config.json's vision_config and image token ids, processor_config.json's image processor; null: no tower. - pub fn parse(arena: std.mem.Allocator, config_json: []const u8, processor_json: ?[]const u8) !?Config { - const doc = try std.json.parseFromSliceLeaky(std.json.Value, arena, config_json, .{}); - const vc = doc.object.get("vision_config") orelse return null; - if (vc != .object) return null; - const want = .{ .{ "depth", depth }, .{ "hidden_size", hidden }, .{ "num_heads", heads }, .{ "intermediate_size", mlp_inter }, .{ "out_hidden_size", out_hidden }, .{ "projection_intermediate_size", merger_inter }, .{ "patch_size", patch }, .{ "temporal_patch_size", temporal }, .{ "spatial_merge_size", merge }, .{ "in_channels", 3 } }; - inline for (want) |kv| { - const v = vc.object.get(kv[0]) orelse return error.VisionConfig; - if (v != .integer or v.integer != kv[1]) { - std.log.err("glm vision: vision_config.{s} is {f}; this tower reads {d}", .{ kv[0], std.json.fmt(v, .{}), kv[1] }); - return error.VisionConfig; - } - } - const act = vc.object.get("hidden_act") orelse return error.VisionConfig; - if (act != .string or !std.mem.eql(u8, act.string, "silu")) return error.VisionConfig; - var c: Config = .{ .image_token = 0, .image_start = 0, .image_end = 0, .mean = .{ 0.48145466, 0.4578275, 0.40821073 }, .std = .{ 0.26862954, 0.26130258, 0.27577711 }, .min_tokens = 16, .max_tokens = 8000 }; - if (vc.object.get("rms_norm_eps")) |v| c.eps = @floatCast(number(v) orelse return error.VisionConfig); - if (vc.object.get("swiglu_limit")) |v| c.limit = @floatCast(number(v) orelse return error.VisionConfig); - inline for (.{ .{ "image_token_id", "image_token" }, .{ "image_start_token_id", "image_start" }, .{ "image_end_token_id", "image_end" } }) |kv| { - const v = doc.object.get(kv[0]) orelse return error.VisionConfig; - if (v != .integer) return error.VisionConfig; - @field(c, kv[1]) = @intCast(v.integer); - } - if (doc.object.get("video_start_token_id")) |s| if (doc.object.get("video_token_id")) |t| if (doc.object.get("video_end_token_id")) |e| { - if (s == .integer and t == .integer and e == .integer) c.video = .{ @intCast(s.integer), @intCast(t.integer), @intCast(e.integer) }; - }; - const pj = processor_json orelse return c; - const pdoc = try std.json.parseFromSliceLeaky(std.json.Value, arena, pj, .{}); - const ip = pdoc.object.get("image_processor") orelse return c; - if (ip != .object) return error.VisionConfig; - inline for (.{ .{ "patch_size", patch }, .{ "temporal_patch_size", temporal }, .{ "merge_size", merge } }) |kv| { - if (ip.object.get(kv[0])) |v| if (v != .integer or v.integer != kv[1]) return error.VisionConfig; - } - if (ip.object.get("patch_expand_factor")) |v| if (v != .integer or v.integer != 1) return error.VisionConfig; - inline for (.{ "image_mean", "image_std" }, .{ &c.mean, &c.std }) |key, dst| if (ip.object.get(key)) |v| { - if (v != .array or v.array.items.len != 3) return error.VisionConfig; - for (v.array.items, 0..) |x, i| dst[i] = @floatCast(number(x) orelse return error.VisionConfig); - }; - if (ip.object.get("min_image_tokens")) |v| c.min_tokens = @intCast(v.integer); - if (ip.object.get("max_image_tokens")) |v| c.max_tokens = @intCast(v.integer); - return c; - } - - fn number(v: std.json.Value) ?f64 { - return switch (v) { - .float => |f| f, - .integer => |i| @floatFromInt(i), - else => null, - }; - } -}; +const vcfg = @import("vision_config.zig"); +pub const hidden = vcfg.hidden; +pub const heads = vcfg.heads; +pub const head_dim = vcfg.head_dim; +pub const depth = vcfg.depth; +pub const mlp_inter = vcfg.mlp_inter; +pub const out_hidden = vcfg.out_hidden; +pub const merger_inter = vcfg.merger_inter; +pub const patch = vcfg.patch; +pub const temporal = vcfg.temporal; +pub const merge = vcfg.merge; +pub const patch_width = vcfg.patch_width; +pub const patches_per_token = vcfg.patches_per_token; +pub const Limits = vcfg.Limits; +pub const Config = vcfg.Config; const Kernels = struct { lib: mtl.Library, diff --git a/zig/src/families/glm/vision_config.zig b/zig/src/families/glm/vision_config.zig new file mode 100644 index 0000000000..6b9ce215c1 --- /dev/null +++ b/zig/src/families/glm/vision_config.zig @@ -0,0 +1,83 @@ +//! GLM-5.3-Flash's image tower: its shape and config reading, Metal-free (every backend's image.zig reads it). +const std = @import("std"); + +// The tower this file reads; `Config.parse` refuses a checkpoint whose vision_config differs. +pub const hidden = 1024; +pub const heads = 16; +pub const head_dim = 64; +pub const depth = 24; +pub const mlp_inter = 4096; +pub const out_hidden = 4096; // the language model's width +pub const merger_inter = 10240; +pub const patch = 14; +pub const temporal = 2; +pub const merge = 2; +pub const patch_width = 3 * temporal * patch * patch; // 1176: a patch's values, channel by frame by row by column +pub const patches_per_token = merge * merge; + +/// --vision's limits: image tokens a prompt may hold in all (scratch and rows are sized for them), and images. +pub const Limits = struct { image_tokens: u32 = 4096, max_images: u32 = 4 }; + +pub const Config = struct { + eps: f32 = 1e-5, + limit: f32 = 10, // swiglu_limit + image_token: u32, + image_start: u32, + image_end: u32, + video: ?[3]u32 = null, // video_start_token_id, video_token_id, video_end_token_id: the checkpoint reads videos + mean: [3]f32, + std: [3]f32, + min_tokens: u32, + max_tokens: u32, + + /// config.json's vision_config and image token ids, processor_config.json's image processor; null: no tower. + pub fn parse(arena: std.mem.Allocator, config_json: []const u8, processor_json: ?[]const u8) !?Config { + const doc = try std.json.parseFromSliceLeaky(std.json.Value, arena, config_json, .{}); + const vc = doc.object.get("vision_config") orelse return null; + if (vc != .object) return null; + const want = .{ .{ "depth", depth }, .{ "hidden_size", hidden }, .{ "num_heads", heads }, .{ "intermediate_size", mlp_inter }, .{ "out_hidden_size", out_hidden }, .{ "projection_intermediate_size", merger_inter }, .{ "patch_size", patch }, .{ "temporal_patch_size", temporal }, .{ "spatial_merge_size", merge }, .{ "in_channels", 3 } }; + inline for (want) |kv| { + const v = vc.object.get(kv[0]) orelse return error.VisionConfig; + if (v != .integer or v.integer != kv[1]) { + std.log.err("glm vision: vision_config.{s} is {f}; this tower reads {d}", .{ kv[0], std.json.fmt(v, .{}), kv[1] }); + return error.VisionConfig; + } + } + const act = vc.object.get("hidden_act") orelse return error.VisionConfig; + if (act != .string or !std.mem.eql(u8, act.string, "silu")) return error.VisionConfig; + var c: Config = .{ .image_token = 0, .image_start = 0, .image_end = 0, .mean = .{ 0.48145466, 0.4578275, 0.40821073 }, .std = .{ 0.26862954, 0.26130258, 0.27577711 }, .min_tokens = 16, .max_tokens = 8000 }; + if (vc.object.get("rms_norm_eps")) |v| c.eps = @floatCast(number(v) orelse return error.VisionConfig); + if (vc.object.get("swiglu_limit")) |v| c.limit = @floatCast(number(v) orelse return error.VisionConfig); + inline for (.{ .{ "image_token_id", "image_token" }, .{ "image_start_token_id", "image_start" }, .{ "image_end_token_id", "image_end" } }) |kv| { + const v = doc.object.get(kv[0]) orelse return error.VisionConfig; + if (v != .integer) return error.VisionConfig; + @field(c, kv[1]) = @intCast(v.integer); + } + if (doc.object.get("video_start_token_id")) |s| if (doc.object.get("video_token_id")) |t| if (doc.object.get("video_end_token_id")) |e| { + if (s == .integer and t == .integer and e == .integer) c.video = .{ @intCast(s.integer), @intCast(t.integer), @intCast(e.integer) }; + }; + const pj = processor_json orelse return c; + const pdoc = try std.json.parseFromSliceLeaky(std.json.Value, arena, pj, .{}); + const ip = pdoc.object.get("image_processor") orelse return c; + if (ip != .object) return error.VisionConfig; + inline for (.{ .{ "patch_size", patch }, .{ "temporal_patch_size", temporal }, .{ "merge_size", merge } }) |kv| { + if (ip.object.get(kv[0])) |v| if (v != .integer or v.integer != kv[1]) return error.VisionConfig; + } + if (ip.object.get("patch_expand_factor")) |v| if (v != .integer or v.integer != 1) return error.VisionConfig; + inline for (.{ "image_mean", "image_std" }, .{ &c.mean, &c.std }) |key, dst| if (ip.object.get(key)) |v| { + if (v != .array or v.array.items.len != 3) return error.VisionConfig; + for (v.array.items, 0..) |x, i| dst[i] = @floatCast(number(x) orelse return error.VisionConfig); + }; + if (ip.object.get("min_image_tokens")) |v| c.min_tokens = @intCast(v.integer); + if (ip.object.get("max_image_tokens")) |v| c.max_tokens = @intCast(v.integer); + return c; + } + + fn number(v: std.json.Value) ?f64 { + return switch (v) { + .float => |f| f, + .integer => |i| @floatFromInt(i), + else => null, + }; + } +}; diff --git a/zig/src/native/cuda_bound.zig b/zig/src/native/cuda_bound.zig index 2fd155bceb..c825d65f4b 100644 --- a/zig/src/native/cuda_bound.zig +++ b/zig/src/native/cuda_bound.zig @@ -18,6 +18,12 @@ pub fn vtable(comptime H: type, v: be.Backend.VTable) be.Backend.VTable { return x.inner.vtable.prefill_step.?(x.inner.ptr, s); } }.f else null, + .prefill_many = if (v.prefill_many != null) struct { + fn f(p: *anyopaque, ss: []const *lanes.Stream) anyerror!bool { + const x = H.bind(p); + return x.inner.vtable.prefill_many.?(x.inner.ptr, ss); + } + }.f else null, .first = struct { fn f(p: *anyopaque, s: *lanes.Stream, position: u64) anyerror!u64 { const x = H.bind(p); diff --git a/zig/src/native/cuda_vision_offer.zig b/zig/src/native/cuda_vision_offer.zig index 56b8ac6d21..6b4439e8d4 100644 --- a/zig/src/native/cuda_vision_offer.zig +++ b/zig/src/native/cuda_vision_offer.zig @@ -1,18 +1,20 @@ //! What a CUDA family with an image tower offers the server: its processor, image token, markers and limits. const std = @import("std"); const api = @import("engine_api"); -const nemotron = @import("nemotron"); const Allocator = std.mem.Allocator; -/// Info.vision for a Nemotron Omni engine's tower; null without --vision. -pub fn offer(v: ?*nemotron.native.vision.Vision) ?api.Vision { +/// Info.vision for a family's tower (Nemotron Omni's, GLM-5.3-Flash's); null without --vision. +pub fn offer(v: anytype) ?api.Vision { const x = v orelse return null; + const T = @typeInfo(@TypeOf(x)).pointer.child; const prepare = struct { fn f(ctx: *anyopaque, a: Allocator, bytes: []const u8, max_tokens: u32) anyerror!api.PreparedImage { - const t: *nemotron.native.vision.Vision = @ptrCast(@alignCast(ctx)); + const t: *T = @ptrCast(@alignCast(ctx)); const p = try t.prepare(a, bytes, max_tokens); return .{ .pixels = p.pixels, .gh = p.gh, .gw = p.gw, .tokens = p.tokens, .hash = p.hash }; } }.f; - return .{ .ctx = x, .prepare = prepare, .image_token = x.marks.image, .markers = &x.mark_ids, .wrap = .{ x.marks.open, x.marks.close }, .image_tokens = x.limits.image_tokens, .max_images = x.limits.max_images, .explain = nemotron.native.vision.Vision.explain }; + // a template that writes each image's begin and end itself: the server expands the image token alone + const wrap: ?[2]u32 = if (@hasDecl(T, "template_wraps")) null else .{ x.marks.open, x.marks.close }; + return .{ .ctx = x, .prepare = prepare, .image_token = x.marks.image, .markers = &x.mark_ids, .wrap = wrap, .image_tokens = x.limits.image_tokens, .max_images = x.limits.max_images, .explain = T.explain }; } diff --git a/zig/tests/cuda/glm_exl3_tests.zig b/zig/tests/cuda/glm_exl3_tests.zig new file mode 100644 index 0000000000..59b1fc37bd --- /dev/null +++ b/zig/tests/cuda/glm_exl3_tests.zig @@ -0,0 +1,238 @@ +//! GLM's EXL3 routed experts (cuda_exl3.run, decode and prompt) against float64 matrices from core's EXL3 decoder. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const glm = @import("glm_cuda"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); +const Gpu = check.Gpu; + +const fmt = core.exl3_format; +const h = qr.helpers; +const Rig = qr.RigT; + +const E = 12; // routed experts (the shared one is id E, skipped) +const D = 512; // dims +const N = 256; // the rank's width of each expert +const slots = glm.exl3.top_k + 1; +const limit: f32 = 7.0; + +/// One weight set: every expert's trellis words and scales as the loader lays them out, and their float64 matrices. +const Set = struct { + gt: u64, + ut: u64, + dt: u64, + suh_g: u64, + suh_u: u64, + svh_g: u64, + svh_u: u64, + suh_d: u64, + svh_d: u64, + shx: bool, + wg: []f64, // [E][D][N] + wu: []f64, + wd: []f64, // [E][N][D] + + fn experts(s: Set) glm.exl3.Experts { + return .{ .gt = s.gt, .ut = s.ut, .dt = s.dt, .suh_g = s.suh_g, .suh_u = s.suh_u, .svh_g = s.svh_g, .svh_u = s.svh_u, .suh_d = s.suh_d, .svh_d = s.svh_d, .shx = s.shx }; + } +}; + +fn signs(r: *Rig, n: usize, scale: f16) ![]f16 { + const out = try r.a.alloc(f16, n); + for (out) |*v| v.* = if (r.rng.boolean()) scale else -scale; + return out; +} + +/// `E` random K x N matrices: trellis words, suh (shared by every expert when `one`), svh; device copies and float64. +fn matrices(r: *Rig, k: usize, n: usize, shared_suh: ?[]const f16, out_suh: *u64, out_svh: *u64, w: []f64) !u64 { + const words = k / 16 * (n / 16) * 64; + const tr = try r.a.alloc(i16, E * words); + for (tr) |*v| v.* = r.rng.int(i16); + const suh = try r.a.alloc(f16, E * k); + const svh = try r.a.alloc(f16, E * n); + const scratch = try r.a.alloc(f64, 2 * k * n); + for (0..E) |e| { + if (shared_suh) |s| @memcpy(suh[e * k ..][0..k], s) else @memcpy(suh[e * k ..][0..k], try signs(r, k, 1.0)); + @memcpy(svh[e * n ..][0..n], try signs(r, n, 0.02)); + fmt.dequantize(.{ .halves = 8 }, .mcg, tr[e * words ..][0..words], k / 16, n / 16, suh[e * k ..][0..k], svh[e * n ..][0..n], w[e * k * n ..][0 .. k * n], scratch); + } + out_suh.* = try r.dev(std.mem.sliceAsBytes(suh)); + out_svh.* = try r.dev(std.mem.sliceAsBytes(svh)); + return r.dev(std.mem.sliceAsBytes(tr)); +} + +fn weights(r: *Rig, one: bool) !Set { + var s: Set = undefined; + s.shx = one; + s.wg = try r.a.alloc(f64, E * D * N); + s.wu = try r.a.alloc(f64, E * D * N); + s.wd = try r.a.alloc(f64, E * N * D); + const rot = if (one) try signs(r, D, 1.0) else null; // one input rotation for every gate and up + s.gt = try matrices(r, D, N, rot, &s.suh_g, &s.svh_g, s.wg); + s.ut = try matrices(r, D, N, rot, &s.suh_u, &s.svh_u, s.wu); + s.dt = try matrices(r, N, D, null, &s.suh_d, &s.svh_d, s.wd); + return s; +} + +/// Each row's 8 distinct routed picks, then the shared expert's id. +fn picks(r: *Rig, rows: usize) ![]i32 { + const out = try r.a.alloc(i32, rows * slots); + for (0..rows) |row| { + var seen: [E]bool = @splat(false); + for (0..slots - 1) |j| { + var e = r.rng.uintLessThan(usize, E); + while (seen[e]) e = r.rng.uintLessThan(usize, E); + seen[e] = true; + out[row * slots + j] = @intCast(e); + } + out[row * slots + slots - 1] = E; + } + return out; +} + +const Run = struct { ey: u64, rows: usize }; + +/// The family's dispatch on `x` (bf16 [rows][D]) and `pk`: ey [rows][slots][D] fp32, the shared slot left as filled. +fn run(r: *Rig, k: *const glm.exl3.Kernels, ek: *const glm.experts.Kernels, s: Set, x: []const u16, pk: []const i32, rows: usize, prompt: bool) !Run { + const P = rows * glm.exl3.top_k; + const tile: usize = if (prompt) glm.experts.staged_tile else glm.experts.tile; + const wide = P > cuda.grouped.small; + const plan: cuda.grouped.Plan = .{ + .members = try r.zeros(P * 4), + .items = try r.zeros(cuda.grouped.maxItems(P, E, tile) * 12), + .counts = try r.zeros(8), + .rank = try r.zeros((if (wide) P else 1) * 4), + .hist = try r.zeros((if (wide) (P + 1023) / 1024 * E else 1) * 4), + }; + const sz = glm.exl3.scratchBytes(rows, D, N, prompt, s.shx); + const sc: glm.exl3.Scratch = .{ .picks = try r.zeros(sz[0]), .xg = try r.zeros(sz[1]), .xu = try r.zeros(sz[2]), .z = try r.zeros(sz[3]) }; + const ey = try r.zeros(rows * slots * D * 4); + const xd = try r.zeros(P * N * 2); + try glm.exl3.run(k, ek, r.s, s.experts(), sc, plan, try r.dev(std.mem.sliceAsBytes(x)), D, try r.dev(std.mem.sliceAsBytes(pk)), slots, rows, E, D, N, limit, prompt, ey, xd); + try r.s.synchronize(); + return .{ .ey = ey, .rows = rows }; +} + +/// The kernels' SwiGLU: gate and up rounded to bf16 and limited, silu to bf16, the product to bf16. +fn act(g: f64, u: f64) f64 { + const gg = @min(h.bfOf(h.bf16Of(@floatCast(g))), limit); + const uu = @max(@min(h.bfOf(h.bf16Of(@floatCast(u))), limit), -limit); + return h.bfOf(h.bf16Of(@floatCast(h.bfOf(h.bf16Of(@floatCast(gg / (1 + @exp(-gg))))) * uu))); +} + +/// Float64 slot rows for routed slots (the shared slot's want is NaN: checked apart) and their tolerances. +fn reference(r: *Rig, s: Set, x: []const u16, pk: []const i32, rows: usize, want: []f64, tol: []f64) !void { + const a = try r.a.alloc(f64, N); + for (0..rows) |row| for (0..slots - 1) |j| { + const e: usize = @intCast(pk[row * slots + j]); + for (0..N) |n| { + var g: f64 = 0; + var u: f64 = 0; + for (0..D) |kk| { + const xv = h.bfOf(x[row * D + kk]); + g += xv * s.wg[(e * D + kk) * N + n]; + u += xv * s.wu[(e * D + kk) * N + n]; + } + a[n] = act(g, u); + } + for (0..D) |n| { + var y: f64 = 0; + var mag: f64 = 0; + for (0..N) |kk| { + y += a[kk] * s.wd[(e * N + kk) * D + n]; + mag += @abs(a[kk] * s.wd[(e * N + kk) * D + n]); + } + const at = (row * slots + j) * D + n; + want[at] = y; + tol[at] = mag * (1.0 / 64.0) + 1e-6; // fp16 rotated inputs, bf16 activations, fp16 down inputs + } + }; +} + +/// Routed slots within tolerance; the shared slot untouched (the fill's bytes). +fn expect(r: *Rig, name: []const u8, out: Run, s: Set, x: []const u16, pk: []const i32) !void { + const n = out.rows * slots * D; + const want = try r.a.alloc(f64, n); + const tol = try r.a.alloc(f64, n); + try reference(r, s, x, pk, out.rows, want, tol); + const raw = try r.back(out.ey, n * 4); + var most: f64 = 0; + var at: usize = 0; + for (0..out.rows) |row| for (0..slots) |j| for (0..D) |c| { + const i = (row * slots + j) * D + c; + const bits = std.mem.readInt(u32, raw[i * 4 ..][0..4], .little); + if (j == slots - 1) { // the shared slot: the fill's bytes + try check.expect(bits == 0xffffffff, "{s}: the shared slot left alone", .{name}); + continue; + } + const g: f64 = @as(f32, @bitCast(bits)); + const q = @abs(g - want[i]) / tol[i]; + if (!(q <= most)) { + most = q; + at = i; + } + }; + r.worst = @max(r.worst, most); + std.debug.print("RESULT {s} ({d} rows): largest difference {d:.3} of its tolerance (element {d})\n", .{ name, out.rows, most, at }); + try check.expect(most <= 1, "{s} within its float64 tolerance", .{name}); +} + +pub fn run_(gpu: Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var stream = try cuda.Stream.init(gpu.d, false); + defer stream.deinit(); + var prng = std.Random.DefaultPrng.init(7); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = stream, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + var xm = try cuda.Module.load(gpu.d, cuda.kernels.glm_exl3); + defer xm.unload(); + var pm = try cuda.Module.load(gpu.d, cuda.kernels.experts); + defer pm.unload(); + const k = try glm.exl3.Kernels.load(xm); + var ek = try glm.experts.Kernels.load(gpu.d, pm, @intCast(try gpu.ctx.attribute(.multiprocessor_count))); + defer ek.unload(); + for ([_]bool{ true, false }) |one| { + const s = try weights(&r, one); + for ([_]usize{ 1, 3, 8 }) |rows| { // decode windows + const x = try r.bfs(rows * D, 1.0); + const pk = try picks(&r, rows); + const got = try run(&r, &k, &ek, s, x, pk, rows, false); + try expect(&r, if (one) "glm-exl3 decode, one rotation" else "glm-exl3 decode, a rotation an expert", got, s, x, pk); + if (rows > 1) { // every row's bits as alone + const all = try r.back(got.ey, rows * slots * D * 4); + for (0..rows) |row| { + const alone = try run(&r, &k, &ek, s, x[row * D ..][0..D], pk[row * slots ..][0..slots], 1, false); + const mine = try r.back(alone.ey, slots * D * 4); + try check.expect(std.mem.eql(u8, mine[0 .. (slots - 1) * D * 4], all[row * slots * D * 4 ..][0 .. (slots - 1) * D * 4]), "glm-exl3 decode: row {d} of {d} equals the row alone", .{ row, rows }); + } + } + } + for ([_]usize{ 1, 40, 200 }) |rows| { // prompt chunks + const x = try r.bfs(rows * D, 1.0); + const pk = try picks(&r, rows); + const got = try run(&r, &k, &ek, s, x, pk, rows, true); + try expect(&r, if (one) "glm-exl3 prompt, one rotation" else "glm-exl3 prompt, a rotation an expert", got, s, x, pk); + } + } + { // a NaN in one row's input: that row's routed slots NaN, every other row finite + const s = try weights(&r, true); + const rows = 4; + const x = try r.bfs(rows * D, 1.0); + x[2 * D + 17] = 0x7fc0; + const pk = try picks(&r, rows); + for ([_]bool{ false, true }) |prompt| { + const got = try run(&r, &k, &ek, s, x, pk, rows, prompt); + const raw = try r.back(got.ey, rows * slots * D * 4); + for (0..rows) |row| for (0..slots - 1) |j| for (0..D) |c| { + const v: f32 = @bitCast(std.mem.readInt(u32, raw[((row * slots + j) * D + c) * 4 ..][0..4], .little)); + try check.expect(if (row == 2) std.math.isNan(v) else std.math.isFinite(v), "glm-exl3 {s}: a NaN stays in its row ({d})", .{ if (prompt) "prompt" else "decode", row }); + }; + } + } + check.pass("glm-exl3: decode and prompt slot rows within {d:.3} of their float64 tolerance, rows as alone, NaN kept to its row", .{r.worst}); +} diff --git a/zig/tests/cuda/glm_mm_bench.zig b/zig/tests/cuda/glm_mm_bench.zig new file mode 100644 index 0000000000..b8532a4d8f --- /dev/null +++ b/zig/tests/cuda/glm_mm_bench.zig @@ -0,0 +1,86 @@ +//! GLM's bf16 linears on one GPU: the family's b16.matmul against glm_mm (tiled weights) on a rank's shapes, by rows. +const std = @import("std"); +const cuda = @import("cuda"); +const glm = @import("glm_cuda"); +const check = @import("check.zig"); +const Gpu = check.Gpu; + +/// A rank's linears of GLM-5.3-Flash at two ranks (n, k): KDA projections and output, MLA, shared expert, dense layer. +const shapes = [_][2]usize{ .{ 6416, 4096 }, .{ 4096, 2048 }, .{ 2048, 4096 }, .{ 4096, 1024 }, .{ 12288, 4096 }, .{ 4096, 6144 }, .{ 2560, 4096 }, .{ 4096, 4096 } }; +const reps = 30; + +pub fn run(gpu: Gpu) !void { + const a = gpu.gpa; + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var gm = try cuda.Module.load(gpu.d, cuda.kernels.glm_glue); + defer gm.unload(); + var mm_mod = try cuda.Module.load(gpu.d, cuda.kernels.glm_mm); + defer mm_mod.unload(); + const glue = try glm.glue.Kernels.load(gm); + const set: glm.glue.Set = .{ .glue = &glue }; + const mk = try glm.mm.Kernels.load(mm_mod); + var prng = std.Random.DefaultPrng.init(3); + const rng = prng.random(); + var t0 = try cuda.Event.init(gpu.d, true); + defer t0.deinit(); + var t1 = try cuda.Event.init(gpu.d, true); + defer t1.deinit(); + var total = [2][7]f64{ @splat(0), @splat(0) }; + for (shapes) |sh| { + const n = sh[0]; + const k = sh[1]; + const host_w = try a.alloc(u16, n * k); + defer a.free(host_w); + for (host_w) |*v| v.* = @truncate(@as(u32, @bitCast((rng.float(f32) * 2 - 1) * 0.02)) >> 16); + var w = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(host_w)); + defer w.free(); + var wt = try cuda.DeviceBuffer.alloc(gpu.d, n * k * 2); + defer wt.free(); + try glm.mm.tile(&mk, s, w.ptr, wt.ptr, n, k); + const tw: glm.mm.Tiled = .{ .w = wt.ptr, .n = n, .k = k }; + const host_x = try a.alloc(u16, 512 * k); + defer a.free(host_x); + for (host_x) |*v| v.* = @truncate(@as(u32, @bitCast(rng.float(f32) * 2 - 1)) >> 16); + var x = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(host_x)); + defer x.free(); + var out_a = try cuda.DeviceBuffer.alloc(gpu.d, 512 * n * 4); + defer out_a.free(); + var out_b = try cuda.DeviceBuffer.alloc(gpu.d, 512 * n * 4); + defer out_b.free(); + var part = try cuda.DeviceBuffer.alloc(gpu.d, @max(16, @max(glm.b16.partBytes(127, n, k), glm.mm.partBytes(128, tw)))); + defer part.free(); + var ms: [2][7]f64 = undefined; + for ([_]usize{ 1, 2, 4, 8, 64, 128, 512 }, 0..) |rows, ri| { + for (0..2) |which| { + for (0..2) |round| { + try t0.record(s); + for (0..reps) |_| { + if (which == 0) try glm.b16.matmul(set, s, x.ptr, k, w.ptr, out_a.ptr, part.ptr, rows, n, k, true) else try glm.mm.matmul(&mk, s, x.ptr, k, tw, out_b.ptr, part.ptr, rows, true); + } + try t1.record(s); + try t1.synchronize(); + if (round == 1) ms[which][ri] = @as(f64, try cuda.Event.elapsedMs(t0, t1)) / reps; + } + total[which][ri] += ms[which][ri]; + } + // the two products agree (different summation orders: a few fp32 ulps of the row's magnitude) + const ga = try a.alloc(f32, rows * n); + defer a.free(ga); + const gb = try a.alloc(f32, rows * n); + defer a.free(gb); + try out_a.download(0, std.mem.sliceAsBytes(ga)); + try out_b.download(0, std.mem.sliceAsBytes(gb)); + var most: f32 = 0; + var big: f32 = 0; + for (ga, gb) |u, v| { + most = @max(most, @abs(u - v)); + big = @max(big, @abs(u)); + } + try check.expect(most <= big * 1e-3, "glm-mm-bench: {d} x {d} at {d} rows: the products agree", .{ n, k, rows }); + } + std.debug.print("RESULT {d} x {d} (rows 1/2/4/8/64/128/512):\n b16 {d:.3}/{d:.3}/{d:.3}/{d:.3}/{d:.3}/{d:.3}/{d:.3} ms\n glm_mm {d:.3}/{d:.3}/{d:.3}/{d:.3}/{d:.3}/{d:.3}/{d:.3} ms\n", .{ n, k, ms[0][0], ms[0][1], ms[0][2], ms[0][3], ms[0][4], ms[0][5], ms[0][6], ms[1][0], ms[1][1], ms[1][2], ms[1][3], ms[1][4], ms[1][5], ms[1][6] }); + } + std.debug.print("RESULT all shapes once (rows 1/2/4/8/64/128/512):\n b16 {d:.2}/{d:.2}/{d:.2}/{d:.2}/{d:.2}/{d:.2}/{d:.2} ms\n glm_mm {d:.2}/{d:.2}/{d:.2}/{d:.2}/{d:.2}/{d:.2}/{d:.2} ms\n", .{ total[0][0], total[0][1], total[0][2], total[0][3], total[0][4], total[0][5], total[0][6], total[1][0], total[1][1], total[1][2], total[1][3], total[1][4], total[1][5], total[1][6] }); + check.pass("glm-mm-bench: both products agree on every shape and row count", .{}); +} diff --git a/zig/tests/cuda/glm_vision_tests.zig b/zig/tests/cuda/glm_vision_tests.zig new file mode 100644 index 0000000000..f7d5968b54 --- /dev/null +++ b/zig/tests/cuda/glm_vision_tests.zig @@ -0,0 +1,104 @@ +//! GLM's image tower on CUDA against vision_ref.py: the patches, and features as near float32 as torch's bf16 run. +const std = @import("std"); +const cuda = @import("cuda"); +const glm = @import("glm_cuda"); +const check = @import("check.zig"); +const Gpu = check.Gpu; + +fn readF32(gpu: Gpu, dir: []const u8, name: []const u8) ![]f32 { + const path = try std.fs.path.join(gpu.gpa, &.{ dir, name }); + defer gpu.gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, gpu.gpa, .limited(1 << 30)); + defer gpu.gpa.free(bytes); + const out = try gpu.gpa.alloc(f32, bytes.len / 4); + @memcpy(std.mem.sliceAsBytes(out), bytes[0 .. out.len * 4]); + return out; +} + +/// Root-mean-square of `want - got` over the rms of `want`, and the largest difference over that rms. +fn errors(want: []const f32, got: []const f32) [2]f64 { + var d2: f64 = 0; + var w2: f64 = 0; + var most: f64 = 0; + for (want, got) |w, g| { + const d: f64 = @as(f64, g) - w; + d2 += d * d; + w2 += @as(f64, w) * w; + most = @max(most, @abs(d)); + } + const rms = @sqrt(w2 / @as(f64, @floatFromInt(want.len))); + return .{ @sqrt(d2 / @as(f64, @floatFromInt(want.len))) / rms, most / rms }; +} + +pub fn run(gpu: Gpu, model: []const u8, case: []const u8, bf16_case: []const u8, image: []const u8) !void { + const a = gpu.gpa; + var arena: std.heap.ArenaAllocator = .init(a); + defer arena.deinit(); + const cwd = std.Io.Dir.cwd(); + const cfg_path = try std.fs.path.join(arena.allocator(), &.{ model, "config.json" }); + const proc_path = try std.fs.path.join(arena.allocator(), &.{ model, "processor_config.json" }); + const cfg = try cwd.readFileAlloc(gpu.io, cfg_path, arena.allocator(), .limited(1 << 24)); + const proc = try cwd.readFileAlloc(gpu.io, proc_path, arena.allocator(), .limited(1 << 24)); + const c = (try glm.vision_config.Config.parse(arena.allocator(), cfg, proc)) orelse return error.NoVisionTower; + const want_patches = try readF32(gpu, case, "patches.f32"); + defer a.free(want_patches); + const want = try readF32(gpu, case, "features.f32"); + defer a.free(want); + const own = try readF32(gpu, bf16_case, "features.f32"); // the torch tower in bf16: its own distance from float32 + defer a.free(own); + const bound = errors(want, own); + + // preparation: core's decoder and resize, the processor's canvas and merge order + const bytes = try cwd.readFileAlloc(gpu.io, image, arena.allocator(), .limited(64 << 20)); + const p = try glm.image.prepare(gpu.io, a, bytes, c, c.max_tokens); + defer a.free(p.pixels); + try check.expect(p.pixels.len == want_patches.len, "glm-vision: {d} patch values, the reference {d}", .{ p.pixels.len, want_patches.len }); + const pe = errors(want_patches, p.pixels); + std.debug.print("RESULT glm-vision patches {d} x {d} ({d} tokens): rms difference {e:.2} of their rms, largest {e:.2}\n", .{ p.gh, p.gw, p.tokens, pe[0], pe[1] }); + try check.expect(pe[1] < 0.02, "glm-vision: the patches as the processor makes them", .{}); + + // the tower on the reference's patches + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.glm_vision); + defer mod.unload(); + const ru: glm.vision.Run = .{ .stream = s, .f = try glm.vision.Funcs.load(mod) }; + var t = try glm.vision.Tower.load(a, gpu.io, gpu.d, model); + defer t.deinit(); + var out = try cuda.DeviceBuffer.alloc(gpu.d, @as(usize, p.tokens) * glm.vision_config.out_hidden * 2); + defer out.free(); + const grid = [3]i64{ 1, p.gh, p.gw }; + { // where the stream parts from the reference's: after the embedding, then each block + const all = readF32(gpu, case, "all.f32") catch null; + const embed = readF32(gpu, case, "embed.f32") catch null; + if (all != null and embed != null) { + defer a.free(all.?); + defer a.free(embed.?); + const n = embed.?.len; + const got = try a.alloc(f32, (glm.vision.DEPTH + 1) * n); + defer a.free(got); + t.stages = got; + try t.encode(&ru, want_patches, grid, out.ptr); + t.stages = null; + const e0 = errors(embed.?, got[0..n]); + std.debug.print("STAGE embed: rms {d:.4}\n", .{e0[0]}); + for (0..glm.vision.DEPTH) |b| { + const e = errors(all.?[b * n ..][0..n], got[(b + 1) * n ..][0..n]); + std.debug.print("STAGE block {d}: rms {d:.4}, largest {d:.3}\n", .{ b, e[0], e[1] }); + } + } + } + for ([_][]const f32{ want_patches, p.pixels }, [_][]const u8{ "the reference's patches", "its own patches" }) |patches, what| { + try t.encode(&ru, patches, grid, out.ptr); + const raw = try a.alloc(u16, @as(usize, p.tokens) * glm.vision_config.out_hidden); + defer a.free(raw); + try out.download(0, std.mem.sliceAsBytes(raw)); + const got = try a.alloc(f32, raw.len); + defer a.free(got); + for (got, raw) |*g, r| g.* = @bitCast(@as(u32, r) << 16); + const fe = errors(want, got); + std.debug.print("RESULT glm-vision features from {s}: rms difference {d:.4} of their rms ({d:.3}x the torch bf16 run's {d:.4}), largest {d:.3}\n", .{ what, fe[0], fe[0] / bound[0], bound[0], fe[1] }); + try check.expect(fe[0] <= 1.1 * bound[0], "glm-vision: features as close to float32 as the torch tower's bf16 run", .{}); + } + check.pass("glm-vision: the processor's patches as the reference's, the features as near float32 as torch's bf16 run", .{}); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index fe98682c86..89343d4be7 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -8,6 +8,7 @@ const bench = @import("bench.zig"); const oracle_tests = @import("oracle_tests.zig"); const libs_tests = @import("libs_tests.zig"); const rank_scatter_tests = @import("rank_scatter_tests.zig"); +const roce_gather_tests = @import("roce_gather_tests.zig"); const sample_tests = @import("sample_tests.zig"); const glue_tests = @import("glue_tests.zig"); const window_profile = @import("window_profile.zig"); @@ -21,6 +22,9 @@ const fp8_experts_tests = @import("fp8_experts_tests.zig"); const nvfp4_experts_tests = @import("nvfp4_experts_tests.zig"); const quant_ref_tests = @import("quant_ref_tests.zig"); const glm_ref_tests = @import("glm_ref_tests.zig"); +const glm_exl3_tests = @import("glm_exl3_tests.zig"); +const glm_vision_tests = @import("glm_vision_tests.zig"); +const glm_mm_bench = @import("glm_mm_bench.zig"); const carveout_tests = @import("carveout_tests.zig"); const nvfp4_tests = @import("nvfp4_tests.zig"); const modelopt_tests = @import("modelopt_tests.zig"); @@ -52,6 +56,7 @@ const usage = \\ cublaslt bf16 GEMM through cuBLASLt against an fp64 reference \\ nccl one-rank NCCL all-reduce and all-gather \\ rank-scatter [id] rank-order reduce-scatter (NCCL_COMM_ID or an id file) against an all-gather + \\ roce-gather two ranks (NCCL_COMM_ID): the RoCE all-gather against NCCL's, graph replays, time a call \\ rank-sum the reduce-scatter's sum kernel alone, 2 to 5 ranks as each rank (one GPU, no NCCL) \\ nccl-pair one rank of a multi-host communicator: send/recv with every peer, all-reduce, all-gather \\ gdn-replay replay_kernel bits against the Python oracle's fixture @@ -90,6 +95,9 @@ const usage = \\ train-pack MODEL EXAMPLES [COUNT] training passes holding several sequences against one a pass: differences, time \\ kernel-contract the registered kernels through the contract's gates, each planted control red (needs -Dnvcc) \\ kernel-contract-limits DIR each registered launch against sm_80/86/89/120/121 from the fatbins in DIR + \\ glm-exl3 GLM's EXL3 routed experts, decode and prompt, against float64 (synthetic) + \\ glm-vision MODEL F32 BF16 IMG GLM's image tower and preparation against vision_ref.py's torch runs + \\ glm-mm-bench GLM's bf16 linears: b16.matmul against glm_mm, by rows (one GPU) \\ glm-ref GLM projections (matmuls, router, attention, indexer, hc) against float64 (synthetic) \\ glm-lanes-cover copy drafts, cancels, discards, eos and max_new on one lanes.Engine across requests \\ glm-lanes-concurrent concurrent streams on one lanes.Engine, each against its solo run @@ -146,6 +154,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { const world = try std.fmt.parseInt(usize, try arg(rest, 1), 10); return rank_scatter_tests.run(gpu, rank, world, if (rest.len > 2) rest[2] else null); } + if (std.mem.eql(u8, cmd, "roce-gather")) return roce_gather_tests.run(gpu, try std.fmt.parseInt(usize, try arg(rest, 0), 10)); if (std.mem.eql(u8, cmd, "rank-sum")) return rank_scatter_tests.sum(gpu); if (std.mem.eql(u8, cmd, "nccl-pair")) return libs_tests.ncclPair(gpu, try std.fmt.parseInt(c_int, try arg(rest, 0), 10), try std.fmt.parseInt(c_int, try arg(rest, 1), 10), try arg(rest, 2)); if (std.mem.eql(u8, cmd, "gdn-replay")) return oracle_tests.gdnReplay(gpu, try arg(rest, 0)); @@ -157,6 +166,9 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "nvfp4-experts")) return nvfp4_experts_tests.experts(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); if (std.mem.eql(u8, cmd, "glm-ref")) return glm_ref_tests.run(gpu); + if (std.mem.eql(u8, cmd, "glm-exl3")) return glm_exl3_tests.run_(gpu); + if (std.mem.eql(u8, cmd, "glm-mm-bench")) return glm_mm_bench.run(gpu); + if (std.mem.eql(u8, cmd, "glm-vision")) return glm_vision_tests.run(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8")) return exlr8_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8-dense")) return exlr8_dense_tests.run(gpu, if (rest.len > 0) rest[0] else null); diff --git a/zig/tests/cuda/roce_gather_tests.zig b/zig/tests/cuda/roce_gather_tests.zig new file mode 100644 index 0000000000..18809d55e2 --- /dev/null +++ b/zig/tests/cuda/roce_gather_tests.zig @@ -0,0 +1,135 @@ +//! Two ranks' RoCE all-gather (Group.openRoce) against NCCL's byte for byte, graph replays, time a call (NCCL_COMM_ID). +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const Gpu = check.Gpu; + +const cap = 512 << 10; +/// Shards a rank: tiny, a few decode rows, the cap, and past it or not a multiple of 16 (those stay on NCCL). +const sizes = [_]usize{ 16, 48, 4096, 8192, 32768, 65536, 262144, cap, cap + 16, 100 }; +const seeds = [_]u64{ 1, 2, 3 }; +const graph_gathers = 4; +const graph_replays = 64; +const reps = 2000; + +/// Rank `rank`'s shard for a case: bytes that differ by rank, size and seed. +fn shard(out: []u8, seed: u64, rank: usize) void { + var rng = std.Random.DefaultPrng.init(seed *% 1000003 +% rank *% 7919 +% out.len); + rng.random().bytes(out); +} + +pub fn run(gpu: Gpu, rank: usize) !void { + const a = gpu.gpa; + try check.expect(std.c.getenv("NCCL_COMM_ID") != null, "roce-gather: NCCL_COMM_ID names rank 0's address", .{}); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var g = try cuda.tp_group.Group.open(rank, 2); + defer g.close(); + try g.openRoce(a, gpu.io, gpu.d, s, cap); + try check.expect(g.fast != null, "roce-gather: the RoCE path did not come up on both ranks (see the warning above)", .{}); + + const most = sizes[std.mem.indexOfMax(usize, &sizes)]; + var mine = try cuda.DeviceBuffer.alloc(gpu.d, most); + defer mine.free(); + var by_group = try cuda.DeviceBuffer.alloc(gpu.d, 2 * most); + defer by_group.free(); + var by_nccl = try cuda.DeviceBuffer.alloc(gpu.d, 2 * most); + defer by_nccl.free(); + const host = try a.alloc(u8, most); + defer a.free(host); + const peer = try a.alloc(u8, most); + defer a.free(peer); + const got = try a.alloc(u8, 2 * most); + defer a.free(got); + const want = try a.alloc(u8, 2 * most); + defer a.free(want); + + var by_roce: usize = 0; + for (sizes) |n| for (seeds) |seed| { + shard(host[0..n], seed, rank); + shard(peer[0..n], seed, 1 - rank); + try mine.upload(0, host[0..n]); + try g.allGather(s, mine.ptr, by_group.ptr, n); + try g.lib.check(g.lib.api.ncclAllGather(mine.ptr, by_nccl.ptr, n, .u8, g.comm, s.handle), "ncclAllGather"); + try s.synchronize(); + try g.check(); + @memcpy(want[rank * n ..][0..n], host[0..n]); + @memcpy(want[(1 - rank) * n ..][0..n], peer[0..n]); + try by_group.download(0, got[0 .. 2 * n]); + try check.sameBytes("group all-gather against both ranks' shards", got[0 .. 2 * n], want[0 .. 2 * n]); + try by_nccl.download(0, got[0 .. 2 * n]); + try check.sameBytes("NCCL all-gather against both ranks' shards", got[0 .. 2 * n], want[0 .. 2 * n]); + if (g.fast.?.fits(n)) by_roce += 1; + }; + check.pass("roce-gather rank {d}: shards of {any} bytes (seeds 1-3) equal NCCL's all-gather; {d} of {d} cases by RoCE", .{ rank, sizes, by_roce, sizes.len * seeds.len }); + + try replays(gpu, &g, s, rank); + try timing(gpu, &g, s, rank); +} + +/// A graph of `graph_gathers` RoCE gathers replayed with new shards each time: each replay's rows are its own. +fn replays(gpu: Gpu, g: *cuda.tp_group.Group, s: cuda.Stream, rank: usize) !void { + const a = gpu.gpa; + const n = 8192; + var mine = try cuda.DeviceBuffer.alloc(gpu.d, graph_gathers * n); + defer mine.free(); + var all = try cuda.DeviceBuffer.alloc(gpu.d, graph_gathers * 2 * n); + defer all.free(); + try cuda.graph.beginCapture(s, .thread_local); + for (0..graph_gathers) |i| try g.allGather(s, mine.ptr + i * n, all.ptr + i * 2 * n, n); + var captured = try cuda.graph.endCapture(s); + defer captured.deinit(); + var exec = try captured.instantiate(); + defer exec.deinit(); + const host = try a.alloc(u8, graph_gathers * n); + defer a.free(host); + const peer = try a.alloc(u8, graph_gathers * n); + defer a.free(peer); + const got = try a.alloc(u8, graph_gathers * 2 * n); + defer a.free(got); + const want = try a.alloc(u8, graph_gathers * 2 * n); + defer a.free(want); + for (0..graph_replays) |r| { + for (0..graph_gathers) |i| { + shard(host[i * n ..][0..n], 100 + r * graph_gathers + i, rank); + shard(peer[i * n ..][0..n], 100 + r * graph_gathers + i, 1 - rank); + @memcpy(want[i * 2 * n + rank * n ..][0..n], host[i * n ..][0..n]); + @memcpy(want[i * 2 * n + (1 - rank) * n ..][0..n], peer[i * n ..][0..n]); + } + try mine.upload(0, host); + try exec.launchOn(s); + try s.synchronize(); + try g.check(); + try all.download(0, got); + try check.sameBytes("a replayed graph's gathers", got, want); + } + check.pass("roce-gather rank {d}: a captured graph of {d} gathers of {d} bytes replayed {d} times, each replay's rows its own", .{ rank, graph_gathers, n, graph_replays }); +} + +/// Median microseconds a call of `reps` back-to-back gathers, RoCE against NCCL, at decode-sized shards. +fn timing(gpu: Gpu, g: *cuda.tp_group.Group, s: cuda.Stream, rank: usize) !void { + var t0 = try cuda.Event.init(gpu.d, true); + defer t0.deinit(); + var t1 = try cuda.Event.init(gpu.d, true); + defer t1.deinit(); + var mine = try cuda.DeviceBuffer.alloc(gpu.d, 65536); + defer mine.free(); + var all = try cuda.DeviceBuffer.alloc(gpu.d, 2 * 65536); + defer all.free(); + for ([_]usize{ 4096, 8192, 32768, 65536 }) |n| { + var us: [2]f64 = undefined; + for (0..2) |which| { + for (0..2) |round| { // the first round warms + try t0.record(s); + for (0..reps) |_| { + if (which == 0) try g.allGather(s, mine.ptr, all.ptr, n) else try g.lib.check(g.lib.api.ncclAllGather(mine.ptr, all.ptr, n, .u8, g.comm, s.handle), "ncclAllGather"); + } + try t1.record(s); + try t1.synchronize(); + if (round == 1) us[which] = @as(f64, try cuda.Event.elapsedMs(t0, t1)) * 1000.0 / reps; + } + } + try g.check(); + check.pass("roce-gather rank {d}: {d} bytes a rank, {d} calls: RoCE {d:.1} us a call, NCCL {d:.1} us", .{ rank, n, reps, us[0], us[1] }); + } +}