diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index d4075393bc..574cae5df6 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -36,6 +36,8 @@ const kernels = [_]Kernel{ .{ .name = "nemotron_mamba", .flags = glue }, .{ .name = "nemotron_attention", .flags = glue }, .{ .name = "nemotron_keyed", .flags = glue }, + .{ .name = "norms", .flags = &.{"-O3"} }, // ours: RMS norm rows for every family (cuda/norms.zig) + .{ .name = "embed_rows", .flags = &.{"-O3"} }, // ours: 4-bit affine and bf16 token rows (cuda/embed_rows.zig) .{ .name = "qmmf", .flags = &.{"-O3"} }, // nvfp4/qmmf.cu's device code (FP8G, FP4), tensorfold_nvfp4_v3 .{ .name = "prompt16", .flags = &.{"-O3"} }, // nvfp4/prompt.cu's FP4 and FP8 tile 4, tensorfold_nvfp4_prompt_v1 .{ .name = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts @@ -125,6 +127,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu b.installArtifact(b.addExecutable(.{ .name = "tensorfold", .root_module = cli })); const runner = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/main.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip }); runner.addImport("cuda", cuda); + runner.addImport("core", mods.core); runner.addImport("lanes", mods.lanes); runner.addImport("nemotron", mods.nemotron); const exlr8 = exlr8_build.coreAlias(b, target, optimize, mods.core); diff --git a/zig/kernels/cuda/embed_rows.cu b/zig/kernels/cuda/embed_rows.cu new file mode 100644 index 0000000000..2ed32c9665 --- /dev/null +++ b/zig/kernels/cuda/embed_rows.cu @@ -0,0 +1,44 @@ +// Token rows for every CUDA family (float64 reference: core/act_ref.zig's affine4): 4-bit affine or bf16 tables. + +#include +#include + +namespace tf_embed { + +constexpr int NT = 256; + +} // namespace tf_embed + +using namespace tf_embed; + +// out[r, i] = bf16(q * s + b) of token ids[r]: q 4 bits (8 a word, low first), s and b bf16 a group of 64; d % 64 == 0. +extern "C" __global__ void __launch_bounds__(NT) tf_embed_rows_affine4(const int* __restrict__ ids, + const uint32_t* __restrict__ w, + const __nv_bfloat16* __restrict__ s, + const __nv_bfloat16* __restrict__ b, + __nv_bfloat16* __restrict__ out, int d) { + const int64_t r = blockIdx.x; + const int c = blockIdx.y * NT + threadIdx.x; // 8 values, one word + if (8 * c >= d) return; + const int64_t tok = ids[r]; + const uint32_t word = w[tok * (d / 8) + c]; + const int64_t g = tok * (d / 64) + c / 8; + const float sc = __bfloat162float(s[g]), bi = __bfloat162float(b[g]); + uint4 o; + __nv_bfloat162* h = reinterpret_cast<__nv_bfloat162*>(&o); +#pragma unroll + for (int i = 0; i < 4; ++i) + h[i] = __floats2bfloat162_rn(__fmaf_rn(static_cast((word >> (8 * i)) & 0xFu), sc, bi), + __fmaf_rn(static_cast((word >> (8 * i + 4)) & 0xFu), sc, bi)); + *reinterpret_cast(out + r * d + 8 * c) = o; +} + +// out[r] = w[ids[r]] for a bf16 table, 16 bytes a thread (d % 8 == 0). +extern "C" __global__ void __launch_bounds__(NT) tf_embed_rows_bf16(const int* __restrict__ ids, + const __nv_bfloat16* __restrict__ w, + __nv_bfloat16* __restrict__ out, int d) { + const int64_t r = blockIdx.x; + const int c = blockIdx.y * NT + threadIdx.x; + if (8 * c >= d) return; + *reinterpret_cast(out + r * d + 8 * c) = *reinterpret_cast(w + static_cast(ids[r]) * d + 8 * c); +} diff --git a/zig/kernels/cuda/nemotron_norms.cu b/zig/kernels/cuda/nemotron_norms.cu index d49387026b..b6171c685d 100644 --- a/zig/kernels/cuda/nemotron_norms.cu +++ b/zig/kernels/cuda/nemotron_norms.cu @@ -52,19 +52,6 @@ __device__ __forceinline__ void scale_store(const float (&x)[MAXE], int D, float using namespace tf_glue; -// MLX affine 4-bit token rows: out = bf16(fma(q, scale, bias)), D / 8 words and D / 64 scales a token. -extern "C" __global__ void __launch_bounds__(NT) tf_nemo_embed(const int* __restrict__ ids, const uint32_t* __restrict__ W, - const __nv_bfloat16* __restrict__ S, const __nv_bfloat16* __restrict__ B, - __nv_bfloat16* __restrict__ out, int D) { - const int64_t tok = ids[blockIdx.x]; - for (int i = threadIdx.x; i < D; i += NT) { - const uint32_t word = W[tok * (D / 8) + i / 8]; - const float q = static_cast((word >> (4 * (i % 8))) & 0xFu); - const int64_t g = tok * (D / 64) + i / 64; - out[static_cast(blockIdx.x) * D + i] = __float2bfloat16_rn(__fmaf_rn(q, bf(S[g]), bf(B[g]))); - } -} - // A bf16 token table's rows (ModelOpt checkpoints keep the embedding unquantized): out = W[ids]. extern "C" __global__ void __launch_bounds__(NT) tf_nemo_embed16(const int* __restrict__ ids, const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ out, int D) { diff --git a/zig/kernels/cuda/norms.cu b/zig/kernels/cuda/norms.cu new file mode 100644 index 0000000000..8c2e4a245d --- /dev/null +++ b/zig/kernels/cuda/norms.cu @@ -0,0 +1,476 @@ +// RMS norm rows for every CUDA family (float64 references: core/norms_ref.zig), a 256-thread block a row. + +// A row's sums run in one fixed order, so its bits never depend on the row count; kinds 0 fp32, 1 fp16, 2 bf16. + +#include +#include +#include + +namespace tf_norms { + +constexpr int NT = 1024; + +template +__device__ __forceinline__ float ld(const void* p, int64_t i) { + if constexpr (K == 0) return static_cast(p)[i]; + else if constexpr (K == 1) return __half2float(static_cast(p)[i]); + else return __bfloat162float(static_cast(p)[i]); +} + +template +__device__ __forceinline__ void st(void* p, int64_t i, float v) { + if constexpr (K == 0) static_cast(p)[i] = v; + else if constexpr (K == 1) static_cast<__half*>(p)[i] = __float2half_rn(v); + else static_cast<__nv_bfloat16*>(p)[i] = __float2bfloat16_rn(v); +} + +__device__ __forceinline__ float load(const void* p, int kind, int64_t i) { + if (kind == 0) return static_cast(p)[i]; + if (kind == 1) return __half2float(static_cast(p)[i]); + return __bfloat162float(static_cast(p)[i]); +} + +__device__ __forceinline__ void store(void* p, int kind, int64_t i, float v) { + if (kind == 0) static_cast(p)[i] = v; + else if (kind == 1) static_cast<__half*>(p)[i] = __float2half_rn(v); + else static_cast<__nv_bfloat16*>(p)[i] = __float2bfloat16_rn(v); +} + +// rsqrt(sum / width + eps) of a block's per-thread partial sums, the same order for every row. +__device__ __forceinline__ float inv_rms(float part, int width, float eps) { + __shared__ float warps[NT / 32]; + __shared__ float inv; + for (int m = 16; m > 0; m >>= 1) part = __fadd_rn(part, __shfl_xor_sync(0xFFFFFFFFu, part, m)); + if ((threadIdx.x & 31) == 0) warps[threadIdx.x >> 5] = part; + __syncthreads(); + if (threadIdx.x < 32) { + float s = warps[threadIdx.x]; + for (int m = 16; m > 0; m >>= 1) s = __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, m)); + if (threadIdx.x == 0) inv = rsqrtf(__fadd_rn(__fdiv_rn(s, static_cast(width)), eps)); + } + __syncthreads(); + const float r = inv; + __syncthreads(); // `warps` and `inv` are reused by a second norm in the same block + return r; +} + +} // namespace tf_norms + +using namespace tf_norms; + +namespace tf_norms { + +constexpr int CACHED = 8; // a thread keeps its values in registers up to CACHED * NT = 8,192 wide + +// rms rows; WIDE re-reads the row from memory (wider than CACHED * NT), the same arithmetic in the same order. +template +__device__ __forceinline__ void rms_row(const void* __restrict__ x, const float* __restrict__ w, void* __restrict__ y, + int width, float eps) { + const int64_t at = static_cast(blockIdx.x) * width; + float part = 0.f; + if (WIDE) { + for (int i = threadIdx.x; i < width; i += NT) { + const float a = ld(x, at + i); + part = __fmaf_rn(a, a, part); + } + const float inv = inv_rms(part, width, eps); + for (int i = threadIdx.x; i < width; i += NT) { + float a = __fmul_rn(ld(x, at + i), inv); + if (w != nullptr) a = __fmul_rn(a, w[i]); + st(y, at + i, a); + } + return; + } + float v[CACHED], wv[CACHED]; +#pragma unroll + for (int j = 0; j < CACHED; ++j) { // every load issued before any is used + const int i = threadIdx.x + j * NT; + v[j] = i < width ? ld(x, at + i) : 0.f; + wv[j] = (w != nullptr && i < width) ? w[i] : 1.f; + } +#pragma unroll + for (int j = 0; j < CACHED; ++j) part = __fmaf_rn(v[j], v[j], part); // zeros past the row add nothing + const float inv = inv_rms(part, width, eps); +#pragma unroll + for (int j = 0; j < CACHED; ++j) { + const int i = threadIdx.x + j * NT; + if (i < width) st(y, at + i, w != nullptr ? __fmul_rn(__fmul_rn(v[j], inv), wv[j]) : __fmul_rn(v[j], inv)); + } +} + +template +__device__ __forceinline__ void add_rms_row(const void* __restrict__ x, float* __restrict__ res, const float* __restrict__ w1, + const float* __restrict__ w2, void* __restrict__ y, int width, float eps) { + const int64_t at = static_cast(blockIdx.x) * width; + if (WIDE) { + float inv1 = 1.f; + if (w1 != nullptr) { + float part = 0.f; + for (int i = threadIdx.x; i < width; i += NT) { + const float a = ld(x, at + i); + part = __fmaf_rn(a, a, part); + } + inv1 = inv_rms(part, width, eps); + } + float part = 0.f; + for (int i = threadIdx.x; i < width; i += NT) { + float a = ld(x, at + i); + if (w1 != nullptr) a = __fmul_rn(__fmul_rn(a, inv1), w1[i]); + const float r = __fadd_rn(res[at + i], a); + res[at + i] = r; + part = __fmaf_rn(r, r, part); + } + const float inv2 = inv_rms(part, width, eps); + for (int i = threadIdx.x; i < width; i += NT) st(y, at + i, __fmul_rn(__fmul_rn(res[at + i], inv2), w2[i])); + return; + } + float v[CACHED], rv[CACHED], a1[CACHED], a2[CACHED]; +#pragma unroll + for (int j = 0; j < CACHED; ++j) { // every load issued before any is used + const int i = threadIdx.x + j * NT; + const bool in = i < width; + v[j] = in ? ld(x, at + i) : 0.f; + rv[j] = in ? res[at + i] : 0.f; + a1[j] = (w1 != nullptr && in) ? w1[i] : 0.f; + a2[j] = in ? w2[i] : 0.f; + } + float inv1 = 1.f; + if (w1 != nullptr) { + float part = 0.f; +#pragma unroll + for (int j = 0; j < CACHED; ++j) part = __fmaf_rn(v[j], v[j], part); + inv1 = inv_rms(part, width, eps); + } + float part = 0.f; +#pragma unroll + for (int j = 0; j < CACHED; ++j) { + const float a = w1 != nullptr ? __fmul_rn(__fmul_rn(v[j], inv1), a1[j]) : v[j]; + rv[j] = __fadd_rn(rv[j], a); + part = __fmaf_rn(rv[j], rv[j], part); + } + const float inv2 = inv_rms(part, width, eps); +#pragma unroll + for (int j = 0; j < CACHED; ++j) { + const int i = threadIdx.x + j * NT; + if (i < width) { + res[at + i] = rv[j]; + st(y, at + i, __fmul_rn(__fmul_rn(rv[j], inv2), a2[j])); + } + } +} + +} // namespace tf_norms + +// y = x * rsqrt(mean(x^2) + eps) * w (w null: no weight), in fp32, rounded once to y's kind. +extern "C" __global__ void __launch_bounds__(NT) tf_rms_rows(const void* x, int x_kind, const float* w, void* y, + int y_kind, int width, float eps) { +#define TF_RMS(XK, YK) \ + if (x_kind == XK && y_kind == YK) { \ + if (width > CACHED * NT) rms_row(x, w, y, width, eps); \ + else rms_row(x, w, y, width, eps); \ + return; \ + } + TF_RMS(0, 0) TF_RMS(0, 1) TF_RMS(0, 2) TF_RMS(1, 0) TF_RMS(1, 1) TF_RMS(1, 2) TF_RMS(2, 0) TF_RMS(2, 1) TF_RMS(2, 2) +#undef TF_RMS +} + +// res (fp32, in place) += rms(x) * w1 (a sandwich norm; null: res += x), then y = rms(res) * w2 in y's kind. +extern "C" __global__ void __launch_bounds__(NT) tf_add_rms_rows(const void* x, int x_kind, float* res, const float* w1, + const float* w2, void* y, int y_kind, int width, + float eps) { +#define TF_ADD(XK, YK) \ + if (x_kind == XK && y_kind == YK) { \ + if (width > CACHED * NT) add_rms_row(x, res, w1, w2, y, width, eps); \ + else add_rms_row(x, res, w1, w2, y, width, eps); \ + return; \ + } + TF_ADD(0, 0) TF_ADD(0, 1) TF_ADD(0, 2) TF_ADD(1, 0) TF_ADD(1, 1) TF_ADD(1, 2) TF_ADD(2, 0) TF_ADD(2, 1) TF_ADD(2, 2) +#undef TF_ADD +} + +// bf16 rows with 16-byte loads (width % 8, at most 8 * VEC_MOST * VNT, 16 B aligned): a short straight line. +namespace tf_norms { + +constexpr int VNT = 256; +constexpr int VEC_MOST = 4; // 8,192 wide + +__device__ __forceinline__ void unpack8(const uint4& u, float* f) { + const __nv_bfloat162* h = reinterpret_cast(&u); +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float2 t = __bfloat1622float2(h[i]); + f[2 * i] = t.x; + f[2 * i + 1] = t.y; + } +} + +__device__ __forceinline__ uint4 pack8(const float* f) { + uint4 u; + __nv_bfloat162* h = reinterpret_cast<__nv_bfloat162*>(&u); +#pragma unroll + for (int i = 0; i < 4; ++i) h[i] = __floats2bfloat162_rn(f[2 * i], f[2 * i + 1]); + return u; +} + +__device__ __forceinline__ void load8f(const float* p, float* f) { + const float4 a = reinterpret_cast(p)[0]; + const float4 b = reinterpret_cast(p)[1]; + f[0] = a.x, f[1] = a.y, f[2] = a.z, f[3] = a.w, f[4] = b.x, f[5] = b.y, f[6] = b.z, f[7] = b.w; +} + +__device__ __forceinline__ void store8f(float* p, const float* f) { + reinterpret_cast(p)[0] = make_float4(f[0], f[1], f[2], f[3]); + reinterpret_cast(p)[1] = make_float4(f[4], f[5], f[6], f[7]); +} + +// rsqrt(sum / width + eps) over a VNT-thread block: warp trees, then warp 0's tree over the 8 warp sums. +__device__ __forceinline__ float inv_rms_v(float part, int width, float eps) { + __shared__ float warps[VNT / 32]; + __shared__ float inv; + for (int m = 16; m > 0; m >>= 1) part = __fadd_rn(part, __shfl_xor_sync(0xFFFFFFFFu, part, m)); + if ((threadIdx.x & 31) == 0) warps[threadIdx.x >> 5] = part; + __syncthreads(); + if (threadIdx.x < 32) { + float s = threadIdx.x < VNT / 32 ? warps[threadIdx.x] : 0.f; + for (int m = 4; m > 0; m >>= 1) s = __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, m)); + if (threadIdx.x == 0) inv = rsqrtf(__fadd_rn(__fdiv_rn(s, static_cast(width)), eps)); + } + __syncthreads(); + const float r = inv; + __syncthreads(); + return r; +} + +} // namespace tf_norms + +// tf_rms_rows for bf16 x and y (w fp32, null: none), 16-byte loads. +extern "C" __global__ void __launch_bounds__(VNT) tf_rms_rows_bf16(const __nv_bfloat16* __restrict__ x, + const float* __restrict__ w, + __nv_bfloat16* __restrict__ y, int width, float eps) { + const int64_t at = static_cast(blockIdx.x) * width; + const int chunks = width / 8; + float v[VEC_MOST][8], wv[VEC_MOST][8]; + float part = 0.f; +#pragma unroll + for (int j = 0; j < VEC_MOST; ++j) { + const int c = threadIdx.x + j * VNT; + if (c < chunks) { + unpack8(reinterpret_cast(x + at)[c], v[j]); + if (w != nullptr) load8f(w + 8 * c, wv[j]); + } + } +#pragma unroll + for (int j = 0; j < VEC_MOST; ++j) + if (threadIdx.x + j * VNT < chunks) +#pragma unroll + for (int i = 0; i < 8; ++i) part = __fmaf_rn(v[j][i], v[j][i], part); + const float inv = inv_rms_v(part, width, eps); +#pragma unroll + for (int j = 0; j < VEC_MOST; ++j) { + const int c = threadIdx.x + j * VNT; + if (c < chunks) { + float o[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) o[i] = w != nullptr ? __fmul_rn(__fmul_rn(v[j][i], inv), wv[j][i]) : __fmul_rn(v[j][i], inv); + reinterpret_cast(y + at)[c] = pack8(o); + } + } +} + +// tf_add_rms_rows for bf16 x and y, 16-byte loads. +extern "C" __global__ void __launch_bounds__(VNT) tf_add_rms_rows_bf16(const __nv_bfloat16* __restrict__ x, + float* __restrict__ res, + const float* __restrict__ w1, + const float* __restrict__ w2, + __nv_bfloat16* __restrict__ y, int width, + float eps) { + const int64_t at = static_cast(blockIdx.x) * width; + const int chunks = width / 8; + float v[VEC_MOST][8], rv[VEC_MOST][8], a1[VEC_MOST][8], a2[VEC_MOST][8]; +#pragma unroll + for (int j = 0; j < VEC_MOST; ++j) { + const int c = threadIdx.x + j * VNT; + if (c < chunks) { + unpack8(reinterpret_cast(x + at)[c], v[j]); + load8f(res + at + 8 * c, rv[j]); + if (w1 != nullptr) load8f(w1 + 8 * c, a1[j]); + load8f(w2 + 8 * c, a2[j]); + } + } + float inv1 = 1.f; + if (w1 != nullptr) { + float part = 0.f; +#pragma unroll + for (int j = 0; j < VEC_MOST; ++j) + if (threadIdx.x + j * VNT < chunks) +#pragma unroll + for (int i = 0; i < 8; ++i) part = __fmaf_rn(v[j][i], v[j][i], part); + inv1 = inv_rms_v(part, width, eps); + } + float part = 0.f; +#pragma unroll + for (int j = 0; j < VEC_MOST; ++j) + if (threadIdx.x + j * VNT < chunks) +#pragma unroll + for (int i = 0; i < 8; ++i) { + const float a = w1 != nullptr ? __fmul_rn(__fmul_rn(v[j][i], inv1), a1[j][i]) : v[j][i]; + rv[j][i] = __fadd_rn(rv[j][i], a); + part = __fmaf_rn(rv[j][i], rv[j][i], part); + } + const float inv2 = inv_rms_v(part, width, eps); +#pragma unroll + for (int j = 0; j < VEC_MOST; ++j) { + const int c = threadIdx.x + j * VNT; + if (c < chunks) { + float o[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) o[i] = __fmul_rn(__fmul_rn(rv[j][i], inv2), a2[j][i]); + store8f(res + at + 8 * c, rv[j]); + reinterpret_cast(y + at)[c] = pack8(o); + } + } +} + +// Rows that hand on a bf16 residual and fp32 sums over each 64 outputs (the next affine matmul's group sums). +namespace tf_norms { + +__device__ __forceinline__ void load8w(const float* p, float* f) { load8f(p, f); } +__device__ __forceinline__ void load8w(const __nv_bfloat16* p, float* f) { unpack8(*reinterpret_cast(p), f); } + +// A thread's 8 stored outputs summed in order, then the group's 8 lanes in a fixed tree (every lane of the warp calls). +__device__ __forceinline__ float group_sum8(const uint4& packed, bool live) { + float f[8]; + unpack8(packed, f); + float s = 0.f; + if (live) +#pragma unroll + for (int i = 0; i < 8; ++i) s = __fadd_rn(s, f[i]); + s = __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, 4)); + s = __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, 2)); + return __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, 1)); +} + +// h = bf16(x + r) (r null: h = x), y = bf16(rms(h) * w), xs y's 64-value sums; width % 64, 16-byte loads. +template +__device__ __forceinline__ void add_rms_h(const __nv_bfloat16* __restrict__ x, const __nv_bfloat16* __restrict__ r, + const WT* __restrict__ w, __nv_bfloat16* __restrict__ h, + __nv_bfloat16* __restrict__ y, float* __restrict__ xs, int width, float eps) { + const int64_t at = static_cast(blockIdx.x) * width; + const int chunks = width / 8; + float v[VEC_MOST][8], wv[VEC_MOST][8]; +#pragma unroll + for (int j = 0; j < VEC_MOST; ++j) { + const int c = threadIdx.x + j * VNT; + if (c < chunks) { + unpack8(reinterpret_cast(x + at)[c], v[j]); + load8w(w + 8 * c, wv[j]); + if (r != nullptr) { + float rv[8]; + unpack8(reinterpret_cast(r + at)[c], rv); +#pragma unroll + for (int i = 0; i < 8; ++i) v[j][i] = __fadd_rn(v[j][i], rv[i]); + const uint4 hp = pack8(v[j]); + reinterpret_cast(h + at)[c] = hp; + unpack8(hp, v[j]); // the norm reads h as stored + } + } + } + float part = 0.f; +#pragma unroll + for (int j = 0; j < VEC_MOST; ++j) + if (threadIdx.x + j * VNT < chunks) +#pragma unroll + for (int i = 0; i < 8; ++i) part = __fmaf_rn(v[j][i], v[j][i], part); + const float inv = inv_rms_v(part, width, eps); +#pragma unroll + for (int j = 0; j < VEC_MOST; ++j) { + const int c = threadIdx.x + j * VNT; + const bool live = c < chunks; + float o[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) o[i] = __fmul_rn(__fmul_rn(v[j][i], inv), wv[j][i]); + const uint4 yp = pack8(o); + if (live) reinterpret_cast(y + at)[c] = yp; + if (xs != nullptr && j * VNT < chunks) { // whole warps of whole groups: width % 64 == 0 + const float s = group_sum8(yp, live); + if (live && (c & 7) == 0) xs[static_cast(blockIdx.x) * (width / 64) + c / 8] = s; + } + } +} + +} // namespace tf_norms + +extern "C" __global__ void __launch_bounds__(VNT) tf_add_rms_rows_bf16h(const __nv_bfloat16* x, const __nv_bfloat16* r, + const __nv_bfloat16* w, __nv_bfloat16* h, + __nv_bfloat16* y, float* xs, int width, float eps) { + add_rms_h(x, r, w, h, y, xs, width, eps); +} + +extern "C" __global__ void __launch_bounds__(VNT) tf_add_rms_rows_bf16h_f32w(const __nv_bfloat16* x, + const __nv_bfloat16* r, const float* w, + __nv_bfloat16* h, __nv_bfloat16* y, + float* xs, int width, float eps) { + add_rms_h(x, r, w, h, y, xs, width, eps); +} + +// Per head of dv = 8 L: out = bf16(silu(z) * rms(y) * w), y and z rows at strides y_ld, z_ld; xs its 64-value sums. +template +__device__ __forceinline__ void gated_rms(const __nv_bfloat16* __restrict__ y, long long y_ld, + const __nv_bfloat16* __restrict__ z, long long z_ld, + const __nv_bfloat16* __restrict__ w, __nv_bfloat16* __restrict__ out, + float* __restrict__ xs, int heads, float eps) { + constexpr int DV = 8 * L; + const int64_t r = blockIdx.x; + const int lane = threadIdx.x & 31, sub = lane % L; + const int hd = (blockIdx.y * blockDim.x + threadIdx.x) / L; + const bool live = hd < heads; + float yv[8], zv[8], wv[8]; + if (live) { + unpack8(*reinterpret_cast(y + r * y_ld + hd * DV + 8 * sub), yv); + unpack8(*reinterpret_cast(z + r * z_ld + hd * DV + 8 * sub), zv); + unpack8(*reinterpret_cast(w + 8 * sub), wv); + } else { +#pragma unroll + for (int i = 0; i < 8; ++i) yv[i] = zv[i] = wv[i] = 0.f; + } + float part = 0.f; +#pragma unroll + for (int i = 0; i < 8; ++i) part = __fmaf_rn(yv[i], yv[i], part); +#pragma unroll + for (int m = L / 2; m > 0; m >>= 1) part = __fadd_rn(part, __shfl_xor_sync(0xFFFFFFFFu, part, m)); + const float inv = rsqrtf(__fadd_rn(__fdiv_rn(part, static_cast(DV)), eps)); + float o[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const float sg = __fdiv_rn(1.f, __fadd_rn(1.f, expf(-zv[i]))); + o[i] = __fmul_rn(__fmul_rn(zv[i], sg), __fmul_rn(__fmul_rn(yv[i], inv), wv[i])); + } + const uint4 op = pack8(o); + const int64_t at = (r * heads + hd) * DV + 8 * sub; + if (live) *reinterpret_cast(out + at) = op; + if (xs != nullptr) { + const float s = group_sum8(op, live); + if (live && (sub & 7) == 0) xs[at / 64] = s; + } +} + +extern "C" __global__ void __launch_bounds__(256) tf_gated_rms_rows_64(const __nv_bfloat16* y, long long y_ld, + const __nv_bfloat16* z, long long z_ld, + const __nv_bfloat16* w, __nv_bfloat16* out, + float* xs, int heads, float eps) { + gated_rms<8>(y, y_ld, z, z_ld, w, out, xs, heads, eps); +} + +extern "C" __global__ void __launch_bounds__(256) tf_gated_rms_rows_128(const __nv_bfloat16* y, long long y_ld, + const __nv_bfloat16* z, long long z_ld, + const __nv_bfloat16* w, __nv_bfloat16* out, + float* xs, int heads, float eps) { + gated_rms<16>(y, y_ld, z, z_ld, w, out, xs, heads, eps); +} + +extern "C" __global__ void __launch_bounds__(256) tf_gated_rms_rows_256(const __nv_bfloat16* y, long long y_ld, + const __nv_bfloat16* z, long long z_ld, + const __nv_bfloat16* w, __nv_bfloat16* out, + float* xs, int heads, float eps) { + gated_rms<32>(y, y_ld, z, z_ld, w, out, xs, heads, eps); +} diff --git a/zig/src/core/act_ref.zig b/zig/src/core/act_ref.zig new file mode 100644 index 0000000000..ccab2c254e --- /dev/null +++ b/zig/src/core/act_ref.zig @@ -0,0 +1,22 @@ +//! Activation rows in float64, the reference every backend's activation kernels are checked against. +const std = @import("std"); + +pub fn sigmoid(x: f64) f64 { + return 1 / (1 + @exp(-x)); +} + +/// A 4-bit affine table's value: q * scale + bias. +pub fn affine4(q: u4, scale: f64, bias: f64) f64 { + return @as(f64, @floatFromInt(q)) * scale + bias; +} + +/// A bf16 output of fp32 arithmetic through expf: half a bf16 ulp plus 2^-18 of |want|. +pub fn tolerance(want: f64) f64 { + return @abs(want) * (std.math.pow(f64, 2, -8) + std.math.pow(f64, 2, -18)) + 1e-30; +} + +test "activations against hand values" { + try std.testing.expectEqual(@as(f64, 0.5), sigmoid(0)); + try std.testing.expectEqual(@as(f64, 7 * 0.5 - 1), affine4(7, 0.5, -1)); + try std.testing.expect(tolerance(1) > tolerance(0.5)); +} diff --git a/zig/src/core/norms_ref.zig b/zig/src/core/norms_ref.zig new file mode 100644 index 0000000000..ef2a2d14e2 --- /dev/null +++ b/zig/src/core/norms_ref.zig @@ -0,0 +1,83 @@ +//! RMS norm rows in float64, the reference every backend's norm kernels are checked against (CUDA: cuda/norms.zig). +const std = @import("std"); + +/// y = x / sqrt(mean(x^2) + eps) * w for one row (w null: no weight). +pub fn rms(x: []const f64, w: ?[]const f64, eps: f64, y: []f64) void { + var ss: f64 = 0; + for (x) |v| ss += v * v; + const inv = 1 / @sqrt(ss / @as(f64, @floatFromInt(x.len)) + eps); + for (x, y, 0..) |v, *o, i| o.* = v * inv * (if (w) |ww| ww[i] else 1); +} + +/// The residual step for one row: res += rms(x) * w1 (w1 null: res += x), then y = rms(res) * w2. +pub fn addRms(x: []const f64, res: []f64, w1: ?[]const f64, w2: []const f64, eps: f64, y: []f64) void { + if (w1) |w| { + rms(x, w, eps, y); // y as scratch: rms(x) * w1 + for (res, y) |*r, v| r.* += v; + } else for (res, x) |*r, v| { + r.* += v; + } + rms(res, w2, eps, y); +} + +/// One once-rounded output's tolerance: half an ulp of its kind at |want| plus 2^-20 of |want| for fp32 arithmetic. +pub fn tolerance(want: f64, kind: Kind) f64 { + const half_ulp: f64 = switch (kind) { + .f32 => std.math.pow(f64, 2, -24), + .f16 => std.math.pow(f64, 2, -11), + .bf16 => std.math.pow(f64, 2, -8), + }; + return @abs(want) * (half_ulp + std.math.pow(f64, 2, -20)) + 1e-30; +} + +/// The residual step's bounds: res_i 2^-20 of |res0_i| + |add_i| (add: what it adds), y_i that error carried on. +pub fn addRmsBounds(res0: []const f64, add: []const f64, w2: []const f64, y: []const f64, res: []const f64, kind: Kind, res_tol: []f64, y_tol: []f64) void { + var ss: f64 = 0; + for (res) |v| ss += v * v; + const inv = 1 / @sqrt(ss / @as(f64, @floatFromInt(res.len))); + for (res0, add, w2, y, res_tol, y_tol) |r0, a, w, yy, *rt, *yt| { + rt.* = (@abs(r0) + @abs(a)) * std.math.pow(f64, 2, -20) + 1e-30; + yt.* = tolerance(yy, kind) + rt.* * inv * @abs(w); + } +} + +/// y's sums over each 64 values. +pub fn groupSums(y: []const f64, xs: []f64) void { + for (xs, 0..) |*o, j| { + var t: f64 = 0; + for (y[64 * j ..][0..64]) |v| t += v; + o.* = t; + } +} + +/// A sum of 64 fp32 values' tolerance: 64 roundings of at most 2^-24 of the running |sum| <= sum |y|. +pub fn groupSumTolerance(y: []const f64) f64 { + var t: f64 = 0; + for (y) |v| t += @abs(v); + return t * 64 * std.math.pow(f64, 2, -24) + 1e-30; +} + +/// silu(z) * rms(y) * w for one head. +pub fn gatedRms(y: []const f64, z: []const f64, w: []const f64, eps: f64, out: []f64) void { + rms(y, w, eps, out); + for (out, z) |*o, zz| o.* *= zz / (1 + @exp(-zz)); +} + +/// The element kinds the kernels read and write, numbered as they are (hip/ops/common.hpp: 0 fp32, 1 fp16, 2 bf16). +pub const Kind = enum(u8) { f32 = 0, f16 = 1, bf16 = 2 }; + +test "a row's norm and the residual step against hand values" { + var y: [4]f64 = undefined; + rms(&.{ 1, -1, 1, -1 }, &.{ 2, 2, 2, 2 }, 0, &y); + try std.testing.expectEqualSlices(f64, &.{ 2, -2, 2, -2 }, &y); + rms(&.{ 3, 4 }, null, 0, y[0..2]); + try std.testing.expectApproxEqRel(@as(f64, 3.0 / @sqrt(12.5)), y[0], 1e-15); + var res = [_]f64{ 1, 1, 1, 1 }; + addRms(&.{ 1, -1, 1, -1 }, &res, null, &.{ 1, 1, 1, 1 }, 0, &y); + try std.testing.expectEqualSlices(f64, &.{ 2, 0, 2, 0 }, &res); + try std.testing.expectApproxEqRel(@as(f64, @sqrt(2.0)), y[0], 1e-15); + res = .{ 0, 0, 0, 0 }; + addRms(&.{ 2, -2, 2, -2 }, &res, &.{ 3, 3, 3, 3 }, &.{ 1, 1, 1, 1 }, 0, &y); // res = rms(x) * 3 = +-3 + try std.testing.expectEqualSlices(f64, &.{ 3, -3, 3, -3 }, &res); + try std.testing.expect(tolerance(1, .bf16) > tolerance(1, .f32)); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 3f4467205f..736248a93d 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -16,6 +16,8 @@ pub const shard_edit = @import("shard_edit.zig"); pub const quant = @import("quant/quant.zig"); pub const registry = @import("registry.zig"); pub const tuning = @import("tuning.zig"); +pub const norms_ref = @import("norms_ref.zig"); +pub const act_ref = @import("act_ref.zig"); test { _ = safetensors; @@ -25,6 +27,8 @@ test { _ = quant; _ = registry; _ = tuning; + _ = norms_ref; + _ = act_ref; _ = exl3_format; _ = exlr8; _ = @import("exl3_rect_test.zig"); diff --git a/zig/src/cuda/embed_rows.zig b/zig/src/cuda/embed_rows.zig new file mode 100644 index 0000000000..a8d4a67b4c --- /dev/null +++ b/zig/src/cuda/embed_rows.zig @@ -0,0 +1,50 @@ +//! embed_rows.cu's token rows for every CUDA family: 4-bit affine (MLX layout, groups of 64) or bf16 tables. +const std = @import("std"); +const kernels = @import("kernels.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); + +const Driver = @import("driver.zig").Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +const threads = 256; + +pub const EmbedRows = struct { + mod: Module, + affine4_fn: Function, + bf16_fn: Function, + + pub fn load(d: *const Driver) !EmbedRows { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.embed_rows); + errdefer mod.unload(); + return .{ .mod = mod, .affine4_fn = try mod.function("tf_embed_rows_affine4"), .bf16_fn = try mod.function("tf_embed_rows_bf16") }; + } + + pub fn unload(g: *EmbedRows) void { + g.mod.unload(); + } + + /// out [rows, d] bf16 from ids (i32 [rows]): w u32 [vocab, d / 8], scales and biases bf16 [vocab, d / 64]. + pub fn affine4(g: *const EmbedRows, s: Stream, ids: u64, w: u64, scales: u64, biases: u64, out: u64, rows: usize, d: usize) !void { + if (d % 64 != 0 or out % 16 != 0) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + for ([_]u64{ ids, w, scales, biases, out }) |v| a.add(v); + a.add(@as(i32, @intCast(d))); + try launch_.launch(g.affine4_fn, .{ .grid = .{ .x = @intCast(rows), .y = @intCast((d / 8 + threads - 1) / threads), .z = 1 }, .block = .{ .x = threads } }, s, &a); + } + + /// out [rows, d] = w [vocab, d] bf16 rows at ids. + pub fn bf16(g: *const EmbedRows, s: Stream, ids: u64, w: u64, out: u64, rows: usize, d: usize) !void { + if (d % 8 != 0 or out % 16 != 0 or w % 16 != 0) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + for ([_]u64{ ids, w, out }) |v| a.add(v); + a.add(@as(i32, @intCast(d))); + try launch_.launch(g.bf16_fn, .{ .grid = .{ .x = @intCast(rows), .y = @intCast((d / 8 + threads - 1) / threads), .z = 1 }, .block = .{ .x = threads } }, s, &a); + } +}; diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 63201f48d3..0fa8bf6e28 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -44,6 +44,8 @@ pub const nemotron_route: []const u8 = if (available) &Blob("fatbin_nemotron_rou pub const nemotron_mamba: []const u8 = if (available) &Blob("fatbin_nemotron_mamba").bytes else &.{}; pub const nemotron_attention: []const u8 = if (available) &Blob("fatbin_nemotron_attention").bytes else &.{}; pub const nemotron_keyed: []const u8 = if (available) &Blob("fatbin_nemotron_keyed").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 &.{}; /// 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/norms.zig b/zig/src/cuda/norms.zig new file mode 100644 index 0000000000..e746777704 --- /dev/null +++ b/zig/src/cuda/norms.zig @@ -0,0 +1,146 @@ +//! norms.cu's RMS norm rows for every CUDA family: a block a row, its bits independent of the row count. +const std = @import("std"); +const kernels = @import("kernels.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); + +const Driver = @import("driver.zig").Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +/// Element kinds, numbered as the kernels (and hip/ops/common.hpp) number them. +pub const Kind = enum(u8) { f32 = 0, f16 = 1, bf16 = 2 }; + +const threads = 1024; +/// norms.cu's 16-byte forms: VNT threads, VEC_MOST loads of 8 a thread. +const vec_threads = 256; +const vec_most = 4; + +/// The 16-byte forms take the row: width a multiple of 8 within VEC_MOST * 8 * VNT, every address 16-byte aligned. +fn vectorOk(width: usize, ptrs: []const u64) bool { + if (width % 8 != 0 or width > vec_most * 8 * vec_threads) return false; + for (ptrs) |p| if (p % 16 != 0) return false; + return true; +} + +pub const Norms = struct { + mod: Module, + rms_fn: Function, + add_rms_fn: Function, + rms_bf16_fn: Function, + add_rms_bf16_fn: Function, + add_rms_h_fn: Function, + add_rms_h_f32w_fn: Function, + gated_fns: [3]Function, // head dim 64, 128, 256 + + pub fn load(d: *const Driver) !Norms { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.norms); + errdefer mod.unload(); + return .{ + .mod = mod, + .rms_fn = try mod.function("tf_rms_rows"), + .add_rms_fn = try mod.function("tf_add_rms_rows"), + .rms_bf16_fn = try mod.function("tf_rms_rows_bf16"), + .add_rms_bf16_fn = try mod.function("tf_add_rms_rows_bf16"), + .add_rms_h_fn = try mod.function("tf_add_rms_rows_bf16h"), + .add_rms_h_f32w_fn = try mod.function("tf_add_rms_rows_bf16h_f32w"), + .gated_fns = .{ try mod.function("tf_gated_rms_rows_64"), try mod.function("tf_gated_rms_rows_128"), try mod.function("tf_gated_rms_rows_256") }, + }; + } + + pub fn unload(n: *Norms) void { + n.mod.unload(); + } + + /// `rows` rows of `width`: y = x * rsqrt(mean(x^2) + eps) * w (w null: no weight), fp32 weights. + pub fn rms(n: *const Norms, s: Stream, x: u64, x_kind: Kind, w: ?u64, y: u64, y_kind: Kind, rows: usize, width: usize, eps: f32) !void { + if (rows == 0) return; + if (x_kind == .bf16 and y_kind == .bf16 and vectorOk(width, &.{ x, y, w orelse 0 })) { + var a: launch_.Args = .{}; + a.add(x); + a.add(w orelse 0); + a.add(y); + a.add(@as(i32, @intCast(width))); + a.add(eps); + return launch_.launch(n.rms_bf16_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = vec_threads } }, s, &a); + } + var a: launch_.Args = .{}; + a.add(x); + a.add(@as(i32, @backingInt(x_kind))); + a.add(w orelse 0); + a.add(y); + a.add(@as(i32, @backingInt(y_kind))); + a.add(@as(i32, @intCast(width))); + a.add(eps); + try launch_.launch(n.rms_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = threads } }, s, &a); + } + + /// The residual step: res (fp32, in place) += rms(x) * w1 (w1 null: res += x), then y = rms(res) * w2. + pub fn addRms(n: *const Norms, s: Stream, x: u64, x_kind: Kind, res: u64, w1: ?u64, w2: u64, y: u64, y_kind: Kind, rows: usize, width: usize, eps: f32) !void { + if (rows == 0) return; + if (x_kind == .bf16 and y_kind == .bf16 and vectorOk(width, &.{ x, y, res, w1 orelse 0, w2 })) { + var a: launch_.Args = .{}; + a.add(x); + a.add(res); + a.add(w1 orelse 0); + a.add(w2); + a.add(y); + a.add(@as(i32, @intCast(width))); + a.add(eps); + return launch_.launch(n.add_rms_bf16_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = vec_threads } }, s, &a); + } + var a: launch_.Args = .{}; + a.add(x); + a.add(@as(i32, @backingInt(x_kind))); + a.add(res); + a.add(w1 orelse 0); + a.add(w2); + a.add(y); + a.add(@as(i32, @backingInt(y_kind))); + a.add(@as(i32, @intCast(width))); + a.add(eps); + try launch_.launch(n.add_rms_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = threads } }, s, &a); + } + + /// bf16 residual: h = bf16(x + r) (r null: h = x), y = bf16(rms(h) * w); xs (or null): y's 64-value sums. + pub fn addRmsH(n: *const Norms, s: Stream, x: u64, r: ?u64, w: u64, w_kind: Kind, h: u64, y: u64, xs: ?u64, rows: usize, width: usize, eps: f32) !void { + if (width % 64 != 0 or !vectorOk(width, &.{ x, r orelse 0, w, h, y }) or w_kind == .f16) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + for ([_]u64{ x, r orelse 0, w, h, y, xs orelse 0 }) |v| a.add(v); + a.add(@as(i32, @intCast(width))); + a.add(eps); + const f = if (w_kind == .bf16) n.add_rms_h_fn else n.add_rms_h_f32w_fn; + try launch_.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = vec_threads } }, s, &a); + } + + /// Per head (dv 64/128/256): out = bf16(silu(z) * rms(y) * w), y/z rows at strides y_ld/z_ld; xs: 64-value sums. + pub fn gatedRms(n: *const Norms, s: Stream, y: u64, y_ld: usize, z: u64, z_ld: usize, w: u64, out: u64, xs: ?u64, rows: usize, heads: usize, dv: usize, eps: f32) !void { + const which: usize = switch (dv) { + 64 => 0, + 128 => 1, + 256 => 2, + else => return error.RowShape, + }; + if (y_ld % 8 != 0 or z_ld % 8 != 0 or !vectorOk(8, &.{ y, z, w, out })) return error.RowShape; + if (rows == 0 or heads == 0) return; + var a: launch_.Args = .{}; + a.add(y); + a.add(@as(i64, @intCast(y_ld))); + a.add(z); + a.add(@as(i64, @intCast(z_ld))); + for ([_]u64{ w, out, xs orelse 0 }) |v| a.add(v); + a.add(@as(i32, @intCast(heads))); + a.add(eps); + const per_block = 256 / (dv / 8); + const grid: launch_.Dim3 = .{ .x = @intCast(rows), .y = @intCast((heads + per_block - 1) / per_block), .z = 1 }; + try launch_.launch(n.gated_fns[which], .{ .grid = grid, .block = .{ .x = 256 } }, s, &a); + } +}; + +test "kinds keep the kernels' numbering" { + try std.testing.expectEqual(@as(u8, 2), @backingInt(Kind.bf16)); +} diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 6ad88717e7..3053c09a6b 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -31,6 +31,8 @@ pub const nvfp4 = @import("nvfp4.zig"); pub const segments = @import("segments.zig"); pub const grouped = @import("grouped.zig"); pub const qlinear = @import("qlinear.zig"); +pub const norms = @import("norms.zig"); +pub const embed_rows = @import("embed_rows.zig"); pub const experts = @import("experts.zig"); pub const sample = @import("sample.zig"); pub const modelopt = @import("modelopt.zig"); @@ -39,6 +41,8 @@ test { _ = @import("memory.zig"); _ = @import("grouped.zig"); _ = @import("qlinear.zig"); + _ = @import("norms.zig"); + _ = @import("embed_rows.zig"); _ = @import("experts.zig"); _ = @import("sample.zig"); _ = launch; diff --git a/zig/src/families/nemotron/cuda_forward.zig b/zig/src/families/nemotron/cuda_forward.zig index 52efe6a449..5e3b34d0a3 100644 --- a/zig/src/families/nemotron/cuda_forward.zig +++ b/zig/src/families/nemotron/cuda_forward.zig @@ -62,7 +62,7 @@ pub const Forward = struct { captures: []const kept.Mark = &.{}, // marks inside the prompt chunk: their Mamba states are copied out as it runs pub fn init(c: Config, w: *const weights.Weights, b: *const state.Buffers, ops: kern.Ops, max_len: usize, nch: usize, sampled: bool) Forward { - return .{ .c = c, .w = w, .b = b, .ops = ops, .glue = .{ .set = if (ops.k.triton) |*t| t else null, .f = &ops.k.glue, .s = ops.s }, .max_len = max_len, .nch = nch, .sampled = sampled }; + return .{ .c = c, .w = w, .b = b, .ops = ops, .glue = .{ .set = if (ops.k.triton) |*t| t else null, .f = &ops.k.glue, .s = ops.s, .emb = &ops.k.emb }, .max_len = max_len, .nch = nch, .sampled = sampled }; } /// Rows of logits draw their next tokens into `out`, row r at position META[0] + r + 1, by `rule` when sampled. diff --git a/zig/src/families/nemotron/cuda_glue.zig b/zig/src/families/nemotron/cuda_glue.zig index 4c313e9331..47c1525901 100644 --- a/zig/src/families/nemotron/cuda_glue.zig +++ b/zig/src/families/nemotron/cuda_glue.zig @@ -16,7 +16,6 @@ pub const chunk_blocks = 64; /// The native kernels, by their extern "C" names. pub const Fns = struct { - embed: cuda.Function, embed16: cuda.Function, add_rmsnorm: cuda.Function, add_moe_norm: cuda.Function, @@ -36,7 +35,6 @@ pub const Fns = struct { /// The modules of nemotron_norms, _route, _mamba, _attention and _keyed, in that order. pub fn resolve(m: []const cuda.Module) !Fns { return .{ - .embed = try m[0].function("tf_nemo_embed"), .embed16 = try m[0].function("tf_nemo_embed16"), .add_rmsnorm = try m[0].function("tf_nemo_add_rmsnorm"), .add_moe_norm = try m[0].function("tf_nemo_add_moe_norm"), @@ -75,6 +73,7 @@ pub const Glue = struct { set: ?*const cuda.aot.Set, // a captured Triton set (GB10's qualified one), else our own kernels f: *const Fns, s: cuda.Stream, + emb: *const cuda.embed_rows.EmbedRows, // the shared token rows (embed_rows.cu) fn tri(g: Glue) ?triton.Tri { return if (g.set) |x| .{ .set = x, .s = g.s } else null; @@ -96,10 +95,7 @@ pub const Glue = struct { /// MLX 4-bit token rows, dequantized. pub fn embed(g: Glue, ids: u64, w: u64, s: u64, b: u64, out: u64, rows: usize, d: usize) !void { if (g.tri()) |t| return t.embed(ids, w, s, b, out, rows, d); - var a: cuda.Args = .{}; - ptrs(&a, &.{ ids, w, s, b, out }); - a.add(int(d)); - try g.go(g.f.embed, .{ rows, 1, 1 }, 256, 0, &a); + try g.emb.affine4(g.s, ids, w, s, b, out, rows, d); } /// Rows of a bf16 table [n, d] (ours with or without a captured set, whose embed reads MLX words). diff --git a/zig/src/families/nemotron/cuda_kernels.zig b/zig/src/families/nemotron/cuda_kernels.zig index caa2ff7801..d23e589224 100644 --- a/zig/src/families/nemotron/cuda_kernels.zig +++ b/zig/src/families/nemotron/cuda_kernels.zig @@ -41,6 +41,7 @@ pub const Kernels = struct { mods: [22]cuda.Module, triton: ?cuda.aot.Set, // a captured Triton set for the glue (GB10's qualified one); null: our own glue kernels glue: glue.Fns, + emb: cuda.embed_rows.EmbedRows, affine: cuda.qlinear.Affine4Kernels, // the 4-bit projections: qmm_group, lane_gemv and qmm_prefill router: cuda.grouped.Router, experts: cuda.experts.Affine4Experts, // the 4-bit grouped experts: experts, experts_prefill and experts_pack @@ -86,6 +87,8 @@ pub const Kernels = struct { k.torch = try torch_ops.Functions.resolve(k.mods[8..14]); k.draws = try cuda.sample.Draws.resolve(k.mods[14]); k.glue = try glue.Fns.resolve(k.mods[16..21]); + k.emb = try cuda.embed_rows.EmbedRows.load(d); + errdefer k.emb.unload(); k.grammar_mask = try k.mods[14].function("tf_grammar_mask"); k.modelopt = null; k.triton = if (triton_dir) |dir| try cuda.aot.Set.load(gpa, io, d, ctx.device, dir) else null; @@ -104,6 +107,7 @@ pub const Kernels = struct { } pub fn deinit(k: *Kernels) void { + k.emb.unload(); if (k.modelopt) |*m| { m.lane.unload(); m.gemm.unload(); diff --git a/zig/src/families/nemotron/glue_ref.zig b/zig/src/families/nemotron/glue_ref.zig index eadcbb2bb9..408e6f47d7 100644 --- a/zig/src/families/nemotron/glue_ref.zig +++ b/zig/src/families/nemotron/glue_ref.zig @@ -47,7 +47,7 @@ fn scaleStore(x: []const f32, inv: f32, w: []const u16, out: []u16, xs: []f32) v } } -/// tf_nemo_embed: MLX 4-bit rows, fma(q, scale, bias) rounded to bf16. +/// The 4-bit embed (embed_rows.cu): MLX 4-bit rows, fma(q, scale, bias) rounded to bf16. pub fn embed(ids: []const i32, w: []const u32, s: []const u16, b: []const u16, out: []u16, d: usize) void { for (ids, 0..) |id, r| { const tok: usize = @intCast(id); diff --git a/zig/tests/cuda/glue_tests.zig b/zig/tests/cuda/glue_tests.zig index 022e3f8d27..31a72a1956 100644 --- a/zig/tests/cuda/glue_tests.zig +++ b/zig/tests/cuda/glue_tests.zig @@ -30,6 +30,7 @@ const Rig = struct { a: std.mem.Allocator, // host data, freed together at the end s: cuda.Stream, f: glue.Fns, + emb: cuda.embed_rows.EmbedRows, rng: std.Random, bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, bad: usize = 0, @@ -37,7 +38,7 @@ const Rig = struct { sparsest: f64 = 1, // the smallest share of nonzero elements in any reference: none may be all zeros fn g(r: *Rig) glue.Glue { - return .{ .set = null, .f = &r.f, .s = r.s }; + return .{ .set = null, .f = &r.f, .s = r.s, .emb = &r.emb }; } /// A device copy of host values, landed before any kernel on the rig's stream reads it. @@ -315,8 +316,9 @@ pub fn run(gpu: Gpu) !void { loaded += 1; } var prng = std.Random.DefaultPrng.init(0x5eed_9a7e); - var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = try cuda.Stream.init(gpu.d, true), .f = try glue.Fns.resolve(&mods), .rng = prng.random() }; + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = try cuda.Stream.init(gpu.d, true), .f = try glue.Fns.resolve(&mods), .emb = try cuda.embed_rows.EmbedRows.load(gpu.d), .rng = prng.random() }; defer { + r.emb.unload(); r.s.deinit(); for (r.bufs.items) |*b| b.free(); r.bufs.deinit(gpu.gpa); diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 7ab32f226c..8aa42dead0 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -24,6 +24,8 @@ const modelopt_tests = @import("modelopt_tests.zig"); const train_pack = @import("train_pack.zig"); const exlr8_tests = @import("exlr8_tests.zig"); const exlr8_dense_tests = @import("exlr8_dense_tests.zig"); +const norms_tests = @import("norms_tests.zig"); +const norm_rows_tests = @import("norm_rows_tests.zig"); const usage = \\usage: tf-cuda-test @@ -40,6 +42,8 @@ const usage = \\ gdn-replay replay_kernel bits against the Python oracle's fixture \\ gdn-tree tree_kernel bits against the Python oracle's fixture \\ triton a Triton cubin's bits against the Python oracle's fixture + \\ norms RMS norm and the residual step against float64 references (synthetic) + \\ norm-rows bf16-residual and gated norms, token rows against float64 (synthetic) \\ sample sample.cu's keyed draws against the Metal rule's host references (synthetic rows) \\ glue Nemotron's glue kernels against their host references at its shapes (synthetic rows) \\ window-profile MODEL WIDTHS IDS_FILE a decode window's GPU ms by kernel class (WIDTHS like 1,4,16), beside its graph @@ -118,6 +122,8 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "exlr8-dense")) return exlr8_dense_tests.run(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exlr8-bench")) return exlr8_tests.bench(gpu, try arg(rest, 0), if (rest.len > 1) rest[1] else "home", if (rest.len > 2) try std.fmt.parseInt(u16, rest[2], 10) else 40, if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 4); if (std.mem.eql(u8, cmd, "sample")) return sample_tests.draws(gpu); + if (std.mem.eql(u8, cmd, "norms")) return norms_tests.run(gpu); + if (std.mem.eql(u8, cmd, "norm-rows")) return norm_rows_tests.run(gpu); if (std.mem.eql(u8, cmd, "glue")) return glue_tests.run(gpu); if (std.mem.eql(u8, cmd, "window-profile")) return window_profile.run(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2)); if (std.mem.eql(u8, cmd, "tree-accept")) return tree_accept.run(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), try arg(rest, 3)); diff --git a/zig/tests/cuda/norm_rows_tests.zig b/zig/tests/cuda/norm_rows_tests.zig new file mode 100644 index 0000000000..50f0e4117b --- /dev/null +++ b/zig/tests/cuda/norm_rows_tests.zig @@ -0,0 +1,226 @@ +//! norms.cu's bf16-residual and gated rows and embed_rows.cu's token rows against float64, bits at every row count. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); + +const nref = core.norms_ref; +const aref = core.act_ref; +const counts = [_]usize{ 1, 3, 17, 64 }; +const most = counts[counts.len - 1]; +const eps: f32 = 1e-6; + +fn bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +fn bf(h: u16) f64 { + return @as(f32, @bitCast(@as(u32, h) << 16)); +} + +const Rig = struct { + gpu: check.Gpu, + a: std.mem.Allocator, + s: cuda.Stream, + rng: std.Random, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + worst: f64 = 0, + + fn dev(r: *Rig, bytes: []const u8) !u64 { + const b = try cuda.DeviceBuffer.fromHost(r.gpu.d, bytes); + try r.bufs.append(r.gpu.gpa, b); + return b.ptr; + } + + fn out(r: *Rig, bytes: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(r.gpu.d, bytes); + try r.bufs.append(r.gpu.gpa, b); + try b.fill8(0xff, r.s.handle); // a value the kernel misses reads NaN + return b.ptr; + } + + fn back(r: *Rig, comptime T: type, ptr: u64, n: usize) ![]T { + const o = try r.a.alloc(T, n); + try r.s.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = ptr, .len = n * @sizeOf(T) }; + try b.download(0, std.mem.sliceAsBytes(o)); + return o; + } + + /// bf16 values: mostly within [-scale, scale), one in 64 twenty times larger. + fn bfs(r: *Rig, n: usize, scale: f32) []u16 { + const v = r.a.alloc(u16, n) catch unreachable; + for (v, 0..) |*x, i| x.* = bf16((r.rng.float(f32) * 2 - 1) * scale * (if (i % 64 == 7) @as(f32, 20) else 1)); + return v; + } + + /// Weights near 1. + fn weights(r: *Rig, n: usize) []f32 { + const v = r.a.alloc(f32, n) catch unreachable; + for (v) |*x| x.* = 0.5 + r.rng.float(f32); + return v; + } + + fn within(r: *Rig, what: []const u8, got: f64, want: f64, tol: f64) !void { + const q = @abs(got - want) / tol; + if (!(q <= 1)) { + try check.expect(false, "{s}: within float64 tolerance (got {d}, want {d}, {d:.3} tolerances)", .{ what, got, want, q }); + } + r.worst = @max(r.worst, q); + } + + /// Group sums `xs` of the stored bf16 rows `y` ([rows, k]). + fn sums(r: *Rig, what: []const u8, y: []const u16, xs: []const f32, rows: usize, k: usize) !void { + var row: [64]f64 = undefined; + for (0..rows) |i| for (0..k / 64) |j| { + for (&row, 0..) |*v, e| v.* = bf(y[i * k + 64 * j + e]); + var want: [1]f64 = undefined; + nref.groupSums(&row, &want); + try r.within(what, xs[i * (k / 64) + j], want[0], nref.groupSumTolerance(&row)); + }; + } +}; + +fn addRmsH(r: *Rig, n: *const cuda.norms.Norms, width: usize, residual: bool, wk: cuda.norms.Kind) !void { + const x = r.bfs(most * width, 1); + const rv = r.bfs(most * width, 1); + const wf = r.weights(width); + const wb = try r.a.alloc(u16, width); + for (wb, wf) |*o, v| o.* = bf16(v); + const wvals = try r.a.alloc(f64, width); + for (wvals, wf, wb) |*o, f, b| o.* = if (wk == .bf16) bf(b) else f; + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const dr = try r.dev(std.mem.sliceAsBytes(rv)); + const dw = try r.dev(if (wk == .bf16) std.mem.sliceAsBytes(wb) else std.mem.sliceAsBytes(wf)); + var hs: [counts.len]u64 = undefined; + var ys: [counts.len]u64 = undefined; + var xss: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + hs[c] = try r.out(m * width * 2); + ys[c] = try r.out(m * width * 2); + xss[c] = try r.out(m * width / 64 * 4); + try n.addRmsH(r.s, dx, if (residual) dr else null, dw, wk, hs[c], ys[c], xss[c], m, width, eps); + } + var name: [96]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "add_rms_h width {d}{s}, {t} weights", .{ width, if (residual) "" else ", no residual", wk }); + const h = try r.back(u16, hs[counts.len - 1], most * width); + const y = try r.back(u16, ys[counts.len - 1], most * width); + const xs = try r.back(f32, xss[counts.len - 1], most * width / 64); + const hv = try r.a.alloc(f64, width); + const want = try r.a.alloc(f64, width); + for (0..most) |row| { + for (hv, 0..) |*v, i| { + const at = row * width + i; + if (residual) { + const sum = bf(x[at]) + bf(rv[at]); + try r.within(label, bf(h[at]), sum, nref.tolerance(sum, .bf16)); + v.* = bf(h[at]); + } else v.* = bf(x[at]); + } + nref.rms(hv, wvals, eps, want); + for (want, 0..) |w, i| try r.within(label, bf(y[row * width + i]), w, nref.tolerance(w, .bf16)); + } + try r.sums(label, y, xs, most, width); + for (counts[0 .. counts.len - 1], 0..) |m, c| { + try check.sameBytes(label, std.mem.sliceAsBytes(try r.back(u16, ys[c], m * width)), std.mem.sliceAsBytes(y[0 .. m * width])); + try check.sameBytes(label, std.mem.sliceAsBytes(try r.back(f32, xss[c], m * width / 64)), std.mem.sliceAsBytes(xs[0 .. m * width / 64])); + if (residual) try check.sameBytes(label, std.mem.sliceAsBytes(try r.back(u16, hs[c], m * width)), std.mem.sliceAsBytes(h[0 .. m * width])); + } +} + +fn gatedRms(r: *Rig, n: *const cuda.norms.Norms, heads: usize, dv: usize) !void { + const width = heads * dv; + const y_ld = width + 64; // a stacked row with something after the heads + const z_ld = 2 * width; + const y = r.bfs(most * y_ld, 1); + const z = r.bfs(most * z_ld, 3); + const wb = try r.a.alloc(u16, dv); + for (wb, r.weights(dv)) |*o, v| o.* = bf16(v); + const wvals = try r.a.alloc(f64, dv); + for (wvals, wb) |*o, b| o.* = bf(b); + const dy = try r.dev(std.mem.sliceAsBytes(y)); + const dz = try r.dev(std.mem.sliceAsBytes(z)); + const dw = try r.dev(std.mem.sliceAsBytes(wb)); + var os: [counts.len]u64 = undefined; + var xss: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + os[c] = try r.out(m * width * 2); + xss[c] = try r.out(m * width / 64 * 4); + try n.gatedRms(r.s, dy, y_ld, dz, z_ld, dw, os[c], xss[c], m, heads, dv, eps); + } + var name: [96]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "gated rms {d} heads of {d}", .{ heads, dv }); + const o = try r.back(u16, os[counts.len - 1], most * width); + const xs = try r.back(f32, xss[counts.len - 1], most * width / 64); + const yv = try r.a.alloc(f64, dv); + const zv = try r.a.alloc(f64, dv); + const want = try r.a.alloc(f64, dv); + for (0..most) |row| for (0..heads) |hd| { + for (yv, zv, 0..) |*a, *b, e| { + a.* = bf(y[row * y_ld + hd * dv + e]); + b.* = bf(z[row * z_ld + hd * dv + e]); + } + nref.gatedRms(yv, zv, wvals, eps, want); + for (want, 0..) |w, e| try r.within(label, bf(o[(row * heads + hd) * dv + e]), w, aref.tolerance(w)); + }; + try r.sums(label, o, xs, most, width); + for (counts[0 .. counts.len - 1], 0..) |m, c| { + try check.sameBytes(label, std.mem.sliceAsBytes(try r.back(u16, os[c], m * width)), std.mem.sliceAsBytes(o[0 .. m * width])); + try check.sameBytes(label, std.mem.sliceAsBytes(try r.back(f32, xss[c], m * width / 64)), std.mem.sliceAsBytes(xs[0 .. m * width / 64])); + } +} + +fn embed(r: *Rig, d: usize) !void { + var e = try cuda.embed_rows.EmbedRows.load(r.gpu.d); + defer e.unload(); + const vocab = 300; + const w = try r.a.alloc(u32, vocab * d / 8); + for (w) |*x| x.* = r.rng.int(u32); + const sc = r.bfs(vocab * d / 64, 0.1); + const bi = r.bfs(vocab * d / 64, 0.5); + const tab = r.bfs(vocab * d, 1); + const ids = try r.a.alloc(i32, most); + for (ids) |*x| x.* = @intCast(r.rng.uintLessThan(usize, vocab)); + const did = try r.dev(std.mem.sliceAsBytes(ids)); + const o4 = try r.out(most * d * 2); + const o16 = try r.out(most * d * 2); + try e.affine4(r.s, did, try r.dev(std.mem.sliceAsBytes(w)), try r.dev(std.mem.sliceAsBytes(sc)), try r.dev(std.mem.sliceAsBytes(bi)), o4, most, d); + try e.bf16(r.s, did, try r.dev(std.mem.sliceAsBytes(tab)), o16, most, d); + const g4 = try r.back(u16, o4, most * d); + const g16 = try r.back(u16, o16, most * d); + for (0..most) |row| { + const t: usize = @intCast(ids[row]); + for (0..d) |i| { + const q: u4 = @truncate(w[t * d / 8 + i / 8] >> @intCast(4 * (i % 8))); + const want = aref.affine4(q, bf(sc[t * d / 64 + i / 64]), bf(bi[t * d / 64 + i / 64])); + try r.within("embed affine4", bf(g4[row * d + i]), want, nref.tolerance(want, .bf16)); + } + try check.sameBytes("embed bf16", std.mem.sliceAsBytes(g16[row * d ..][0..d]), std.mem.sliceAsBytes(tab[t * d ..][0..d])); + } +} + +pub fn run(gpu: check.Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var prng = std.Random.DefaultPrng.init(27); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = try cuda.Stream.init(gpu.d, false), .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + r.s.deinit(); + } + var n = try cuda.norms.Norms.load(gpu.d); + defer n.unload(); + try addRmsH(&r, &n, 5120, true, .bf16); // Qwen3.8-27B's residual + try addRmsH(&r, &n, 5120, false, .bf16); // the first layer + try addRmsH(&r, &n, 2048, true, .f32); + try addRmsH(&r, &n, 2688, true, .bf16); // Nemotron 3.5 Lightning's hidden + try gatedRms(&r, &n, 48, 128); // the 27B's GDN output + try gatedRms(&r, &n, 5, 64); + try gatedRms(&r, &n, 3, 256); + try embed(&r, 5120); // the 27B's hidden + try embed(&r, 2688); + try embed(&r, 192); + check.pass("norm rows: add_rms_h, gated rms, embed rows within float64 tolerance (worst {d:.3}), group sums of the stored outputs, every row count's bits equal", .{r.worst}); +} diff --git a/zig/tests/cuda/norms_tests.zig b/zig/tests/cuda/norms_tests.zig new file mode 100644 index 0000000000..20dceaa55b --- /dev/null +++ b/zig/tests/cuda/norms_tests.zig @@ -0,0 +1,196 @@ +//! norms.cu against core/norms_ref.zig's float64 rows, and every row's bits at every row count. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); + +const ref = core.norms_ref; +const Kind = cuda.norms.Kind; + +const widths = [_]usize{ 64, 2560, 5120, 7168, 9216 }; // 9,216: past 8,192, the re-reading path +const counts = [_]usize{ 1, 3, 17, 300 }; +const most = counts[counts.len - 1]; +const eps: f32 = 1e-6; + +const Rig = struct { + gpu: check.Gpu, + a: std.mem.Allocator, + s: cuda.Stream, + n: *const cuda.norms.Norms, + rng: std.Random, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + worst: f64 = 0, + + fn dev(r: *Rig, bytes: []const u8) !u64 { + const b = try cuda.DeviceBuffer.fromHost(r.gpu.d, bytes); + try r.bufs.append(r.gpu.gpa, b); + return b.ptr; + } + + fn out(r: *Rig, bytes: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(r.gpu.d, bytes); + try r.bufs.append(r.gpu.gpa, b); + try b.fill8(0xff, r.s.handle); // a row the kernel misses reads NaN + return b.ptr; + } + + fn back(r: *Rig, ptr: u64, len: usize) ![]u8 { + const o = try r.a.alloc(u8, len); + try r.s.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = ptr, .len = len }; + try b.download(0, o); + return o; + } + + /// Values a residual stream holds: mostly within [-1, 1), one channel in 64 twenty times larger. + fn values(r: *Rig, n: usize) []f32 { + const v = r.a.alloc(f32, n) catch unreachable; + for (v, 0..) |*x, i| x.* = (r.rng.float(f32) * 2 - 1) * (if (i % 64 == 7) @as(f32, 20) else 1); + return v; + } +}; + +fn size(k: Kind) usize { + return switch (k) { + .f32 => 4, + .f16, .bf16 => 2, + }; +} + +fn bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +/// `v` stored as `k` (bytes) and the values it then holds (f64). +fn stored(a: std.mem.Allocator, v: []const f32, k: Kind) !struct { bytes: []u8, vals: []f64 } { + const bytes = try a.alloc(u8, v.len * size(k)); + const vals = try a.alloc(f64, v.len); + for (v, 0..) |x, i| switch (k) { + .f32 => { + std.mem.writeInt(u32, bytes[4 * i ..][0..4], @bitCast(x), .little); + vals[i] = x; + }, + .bf16 => { + const h = bf16(x); + std.mem.writeInt(u16, bytes[2 * i ..][0..2], h, .little); + vals[i] = @as(f32, @bitCast(@as(u32, h) << 16)); + }, + .f16 => { + const h: f16 = @floatCast(x); + std.mem.writeInt(u16, bytes[2 * i ..][0..2], @bitCast(h), .little); + vals[i] = h; + }, + }; + return .{ .bytes = bytes, .vals = vals }; +} + +fn read(raw: []const u8, k: Kind, i: usize) f64 { + return switch (k) { + .f32 => @as(f32, @bitCast(std.mem.readInt(u32, raw[4 * i ..][0..4], .little))), + .bf16 => @as(f32, @bitCast(@as(u32, std.mem.readInt(u16, raw[2 * i ..][0..2], .little)) << 16)), + .f16 => @as(f16, @bitCast(std.mem.readInt(u16, raw[2 * i ..][0..2], .little))), + }; +} + +/// Every element of `raw` (kind `k`) within `tol` of `want` (null: core's tolerance of a once-rounded value). +fn close(r: *Rig, what: []const u8, raw: []const u8, k: Kind, want: []const f64, tol: ?[]const f64) !void { + var q_most: f64 = 0; + for (want, 0..) |w, i| { + const t = if (tol) |tt| tt[i] else ref.tolerance(w, @fromBackingInt(@intCast(@backingInt(k)))); + const q = @abs(read(raw, k, i) - w) / t; + if (!(q <= q_most)) q_most = q; + } + r.worst = @max(r.worst, q_most); + try check.expect(q_most <= 1, "{s}: within its float64 tolerance (worst {d:.3})", .{ what, q_most }); +} + +fn rmsCase(r: *Rig, width: usize, xk: Kind, yk: Kind, weighted: bool) !void { + const x = try stored(r.a, r.values(most * width), xk); + const wv = r.values(width); + for (wv) |*v| v.* = 0.5 + @abs(v.*) / 20; + const w = try stored(r.a, wv, .f32); + const dx = try r.dev(x.bytes); + const dw = try r.dev(w.bytes); + const want = try r.a.alloc(f64, most * width); + for (0..most) |row| ref.rms(x.vals[row * width ..][0..width], if (weighted) w.vals else null, eps, want[row * width ..][0..width]); + var name: [96]u8 = undefined; + const outs = try r.a.alloc(u64, counts.len); + for (counts, outs) |m, *o| { + o.* = try r.out(m * width * size(yk)); + try r.n.rms(r.s, dx, xk, if (weighted) dw else null, o.*, yk, m, width, eps); + } + const label = try std.fmt.bufPrint(&name, "rms {t} -> {t}, width {d}{s}", .{ xk, yk, width, if (weighted) "" else ", no weight" }); + const wide = try r.back(outs[counts.len - 1], most * width * size(yk)); + try close(r, label, wide, yk, want, null); + for (counts[0 .. counts.len - 1], outs[0 .. counts.len - 1]) |m, o| + try check.sameBytes(label, try r.back(o, m * width * size(yk)), wide[0 .. m * width * size(yk)]); +} + +fn addCase(r: *Rig, width: usize, sandwich: bool) !void { + const x = try stored(r.a, r.values(most * width), .bf16); + const res0 = try stored(r.a, r.values(most * width), .f32); + const w1v = r.values(width); + const w2v = r.values(width); + for (w1v, w2v) |*p, *q| { + p.* = 0.5 + @abs(p.*) / 20; + q.* = 0.5 + @abs(q.*) / 20; + } + const w1 = try stored(r.a, w1v, .f32); + const w2 = try stored(r.a, w2v, .f32); + const dx = try r.dev(x.bytes); + const dw1 = try r.dev(w1.bytes); + const dw2 = try r.dev(w2.bytes); + const want_res = try r.a.dupe(f64, res0.vals); + const want_y = try r.a.alloc(f64, most * width); + const added = try r.a.alloc(f64, most * width); + const res_tol = try r.a.alloc(f64, most * width); + const y_tol = try r.a.alloc(f64, most * width); + for (0..most) |row| { + const at = row * width; + const xs = x.vals[at..][0..width]; + if (sandwich) ref.rms(xs, w1.vals, eps, added[at..][0..width]) else @memcpy(added[at..][0..width], xs); + ref.addRms(xs, want_res[at..][0..width], if (sandwich) w1.vals else null, w2.vals, eps, want_y[at..][0..width]); + ref.addRmsBounds(res0.vals[at..][0..width], added[at..][0..width], w2.vals, want_y[at..][0..width], want_res[at..][0..width], .bf16, res_tol[at..][0..width], y_tol[at..][0..width]); + } + var name: [96]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "add_rms bf16 + fp32 residual, width {d}{s}", .{ width, if (sandwich) ", sandwich" else "" }); + const ys = try r.a.alloc(u64, counts.len); + const rs = try r.a.alloc(u64, counts.len); + for (counts, ys, rs) |m, *y, *rr| { + rr.* = try r.dev(res0.bytes[0 .. m * width * 4]); + y.* = try r.out(m * width * 2); + try r.n.addRms(r.s, dx, .bf16, rr.*, if (sandwich) dw1 else null, dw2, y.*, .bf16, m, width, eps); + } + const wide_y = try r.back(ys[counts.len - 1], most * width * 2); + const wide_r = try r.back(rs[counts.len - 1], most * width * 4); + try close(r, label, wide_y, .bf16, want_y, y_tol); + try close(r, label, wide_r, .f32, want_res, res_tol); + for (counts[0 .. counts.len - 1], ys[0 .. counts.len - 1], rs[0 .. counts.len - 1]) |m, y, rr| { + try check.sameBytes(label, try r.back(y, m * width * 2), wide_y[0 .. m * width * 2]); + try check.sameBytes(label, try r.back(rr, m * width * 4), wide_r[0 .. m * width * 4]); + } +} + +pub fn run(gpu: check.Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var n = try cuda.norms.Norms.load(gpu.d); + defer n.unload(); + var prng = std.Random.DefaultPrng.init(17); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .n = &n, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + for (widths) |w| { + try rmsCase(&r, w, .f32, .bf16, true); + try rmsCase(&r, w, .bf16, .bf16, true); + try rmsCase(&r, w, .f16, .f32, false); + try addCase(&r, w, true); + try addCase(&r, w, false); + } + check.pass("norms: rms and add_rms within float64 tolerance (worst {d:.3}), every row count's bits equal", .{r.worst}); +}