diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index a7975e5bfe..d2aa4323a2 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -10,8 +10,10 @@ const torch_ops = &[_][]const u8{ "-O3", "--fmad=false", "--ftz=false" }; const kernels = [_]Kernel{ .{ .name = "gdn", .flags = &.{ "-O3", "--fmad=false" } }, // cuda/kernels/gdn.py, tensorfold_gdn_v2 + .{ .name = "flashnext_gdn", .flags = &.{ "-O3", "--fmad=false" } }, // qwen4_exp/cuda/gdn.cu chain, one row .{ .name = "probe", .flags = &.{"-O3"} }, .{ .name = "qmm_group", .flags = &.{"-O3"} }, // cuda/kernels/qmm.py, tensorfold_qmm_v5 + .{ .name = "qmm", .flags = &.{"-O3"} }, // cuda/kernels/qmm.py lane matmul, groups of 32 .{ .name = "qmm_prefill", .flags = &.{"-O3"} }, .{ .name = "experts", .flags = &.{"-O3"} }, // cuda/experts.py, tensorfold_experts_v7 .{ .name = "experts_prefill", .flags = &.{"-O3"} }, @@ -55,7 +57,7 @@ fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti return cuda; } -fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, draft_ids: *std.Build.Module) struct { core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module } { +fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, draft_ids: *std.Build.Module) struct { core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, flashnext: *std.Build.Module, tokenizer: *std.Build.Module } { const tokenizer = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = optimize, .link_libc = true }); const core = b.createModule(.{ .root_source_file = b.path("zig/src/core/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); core.addImport("tokenizer", tokenizer); @@ -65,7 +67,11 @@ fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin nemotron.addImport("core", core); nemotron.addImport("lanes", lanes); nemotron.addImport("nemotron_draft_ids", draft_ids); - return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .tokenizer = tokenizer }; + const flashnext = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_launch.zig"), .target = target, .optimize = optimize, .link_libc = true }); + flashnext.addImport("cuda", cuda); + flashnext.addImport("core", core); + flashnext.addImport("lanes", lanes); + return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .flashnext = flashnext, .tokenizer = tokenizer }; } /// Linux targets: fatbins (-Dnvcc builds them, -Dfatbins embeds prebuilt ones), `tensorfold` and `tf-cuda-test`. @@ -100,25 +106,45 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu const runner = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/main.zig"), .target = target, .optimize = optimize, .link_libc = true }); runner.addImport("cuda", cuda); b.installArtifact(b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner })); - nativeServer(b, target, optimize, cuda, mods.lanes, mods.nemotron, mods.tokenizer); + // not installed: one embedding row through the captured _embed cubin + const embed_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_embed_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + embed_check.addImport("cuda", cuda); + embed_check.addImport("core", mods.core); + embed_check.addImport("flashnext", mods.flashnext); + const embed_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-embed", .root_module = embed_check })); + const embed_model = b.option([]const u8, "flashnext_model", "Flash Next checkpoint directory for the embedding check"); + const embed_kernels = b.option([]const u8, "flashnext_kernels", "captured Triton pack for the embedding check"); + const embed_token = b.option([]const u8, "flashnext_token", "token id for the embedding check") orelse "7"; + if (embed_model != null and embed_kernels != null) embed_run.addArgs(&.{ embed_model.?, embed_kernels.?, embed_token }); + b.step("flashnext-embed", "Upload the Flash Next embedding table and run _embed for one token").dependOn(&embed_run.step); + const prompt_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_prompt_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + prompt_check.addImport("cuda", cuda); + prompt_check.addImport("core", mods.core); + prompt_check.addImport("flashnext", mods.flashnext); + const prompt_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-prompt", .root_module = prompt_check })); + const prompt_ids = b.option([]const u8, "flashnext_ids", "file of prompt token ids for the Flash Next prompt check"); + const prompt_count = b.option([]const u8, "flashnext_count", "greedy tokens to emit") orelse "8"; + if (embed_model != null and embed_kernels != null and prompt_ids != null) prompt_run.addArgs(&.{ embed_model.?, embed_kernels.?, prompt_ids.?, prompt_count }); + b.step("flashnext-prompt", "Run one Flash Next prompt on the captured CUDA kernels").dependOn(&prompt_run.step); + nativeServer(b, target, optimize, cuda, mods.lanes, mods.nemotron, mods.flashnext, mods.tokenizer); } /// The CUDA engines a native server opens (native/cuda.zig), over the given runtime and families. -fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module) struct { api: *std.Build.Module, engines: *std.Build.Module } { +fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, flashnext: *std.Build.Module) struct { api: *std.Build.Module, engines: *std.Build.Module } { const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); const mod = b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, - .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "engine_api", .module = api }, .{ .name = "lanes", .module = lanes }, .{ .name = "nemotron", .module = nemotron } }, + .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "engine_api", .module = api }, .{ .name = "lanes", .module = lanes }, .{ .name = "nemotron", .module = nemotron }, .{ .name = "flashnext", .module = flashnext } }, }); return .{ .api = api, .engines = mod }; } /// `zig build native`: tensorfold-native with the CUDA engines into zig-out/native/bin, as the Metal build makes it. -fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module) void { - const m = engines(b, target, optimize, cuda, lanes, nemotron); +fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, flashnext: *std.Build.Module, tokenizer: *std.Build.Module) void { + const m = engines(b, target, optimize, cuda, lanes, nemotron, flashnext); // the HTTP side keeps its safety checks; the engine below it runs at `optimize` (the tokenizer is the family's) const template = b.createModule(.{ .root_source_file = b.path("zig/src/core/template/template.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true }); const exe = b.addExecutable(.{ .name = "tensorfold-native", .root_module = b.createModule(.{ @@ -137,8 +163,8 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, step: *std.Build.S const host = b.graph.host; const cuda = runtime(b, host, .debug, &.{}); const mods = family(b, host, .debug, cuda, draft_ids); - const native = engines(b, host, .debug, cuda, mods.lanes, mods.nemotron).engines; - for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); + const native = engines(b, host, .debug, cuda, mods.lanes, mods.nemotron, mods.flashnext).engines; + for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.flashnext, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); const cli = b.createModule(.{ .root_source_file = b.path("zig/src/cli/cuda_main.zig"), .target = host, .optimize = .debug, .link_libc = true }); cli.addImport("cuda", cuda); cli.addImport("core", mods.core); diff --git a/zig/kernels/cuda/experts.cu b/zig/kernels/cuda/experts.cu index edd83834c1..9e46b57a19 100644 --- a/zig/kernels/cuda/experts.cu +++ b/zig/kernels/cuda/experts.cu @@ -280,3 +280,8 @@ __global__ void plan_scatter(const int* __restrict__ picks, int P, int E, const int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float); TF_EXPERT(1) TF_EXPERT(0) +// Groups of 32: SwiGLU gate and up together (epilogue 2), and the down projection to fp32. +template __global__ void tf_experts::expert_kernel<32, 2, 2, 2, 4>(const __nv_bfloat16*, \ + int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float); +template __global__ void tf_experts::expert_kernel<32, 1, 0, 2, 4>(const __nv_bfloat16*, \ + int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float); diff --git a/zig/kernels/cuda/experts_pack.cu b/zig/kernels/cuda/experts_pack.cu index 6aa7413793..327517f5b4 100644 --- a/zig/kernels/cuda/experts_pack.cu +++ b/zig/kernels/cuda/experts_pack.cu @@ -37,6 +37,8 @@ pack_kernel(const uint32_t* __restrict__ words, const uint16_t* __restrict__ sca } // namespace tf_experts_pack -// Groups of 64 inputs. +// Groups of 64 inputs, then groups of 32. template __global__ void tf_experts_pack::pack_kernel<2>(const uint32_t*, const uint16_t*, const uint16_t*, uint32_t*, int, int, int, int); +template __global__ void tf_experts_pack::pack_kernel<1>(const uint32_t*, const uint16_t*, const uint16_t*, + uint32_t*, int, int, int, int); diff --git a/zig/kernels/cuda/flashnext_gdn.cu b/zig/kernels/cuda/flashnext_gdn.cu new file mode 100644 index 0000000000..5ce476b7df --- /dev/null +++ b/zig/kernels/cuda/flashnext_gdn.cu @@ -0,0 +1,143 @@ +// One value head per block: depthwise conv, the delta-rule step, then the gated RMSNorm. 16 key heads, 48 value heads. +// Same operations as the one-row chain, compiled without FMA contraction. + +#include +#include + +constexpr int NK = 16, NV = 48, DK = 128, DV = 128, TAPS = 4; + +static __device__ __forceinline__ float bf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +static __device__ __forceinline__ float warp_sum(float x) { + for (int o = 16; o; o >>= 1) x += __shfl_xor_sync(0xffffffffu, x, o); + return x; +} + +static __device__ __forceinline__ float sigmoidf_(float x) { return 1.0f / (1.0f + expf(-x)); } + +static __device__ __forceinline__ float softplusf_(float x) { return x > 20.0f ? x : log1pf(expf(x)); } + +static __device__ __forceinline__ void update(float (&s)[4][4], const float (&kk)[4], const float* vrow, int warp, float g, + float beta) { +#pragma unroll + for (int j = 0; j < 4; ++j) { + float kv = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + s[j][i] = s[j][i] * g; + kv = kv + s[j][i] * kk[i]; + } + kv = warp_sum(kv); + const float delta = (vrow[warp * 4 + j] - kv) * beta; +#pragma unroll + for (int i = 0; i < 4; ++i) s[j][i] = s[j][i] + kk[i] * delta; + } +} + +extern "C" __global__ void __launch_bounds__(1024) flashnext_gdn_chain( + const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, + const __nv_bfloat16* __restrict__ cw, const float* __restrict__ state_in, const float* __restrict__ a_log, + const float* __restrict__ dt_bias, const __nv_bfloat16* __restrict__ norm_w, float eps, int rows, + __nv_bfloat16* __restrict__ out, float* __restrict__ xs, float* __restrict__ state_out) { + constexpr int C = 2 * NK * DK + NV * DV; + constexpr int PW = C + NV * DV + 2 * NV; + const int hv = blockIdx.x, hk = hv / (NV / NK); + const int t = threadIdx.x, warp = t >> 5, lane = t & 31; + __shared__ float qs[DK], ks[DK], vs[DV], ys[DV]; + __shared__ float gates[2], rinv; + int c = -1; + if (t < DK) c = hk * DK + t; + else if (t < 2 * DK) c = NK * DK + hk * DK + (t - DK); + else if (t < 2 * DK + DV) c = 2 * NK * DK + hv * DV + (t - 2 * DK); + float s[4][4]; + const size_t sbase = (size_t)hv * DV * DK; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) s[j][i] = state_in[sbase + (size_t)(warp * 4 + j) * DK + lane * 4 + i]; + const __nv_bfloat16* pz = P + C + hv * DV + (t < DV ? t : 0); + const __nv_bfloat16* pb = P + C + NV * DV + hv; + const __nv_bfloat16* pa = pb + NV; + for (int r = 0; r < rows; ++r) { + if (c >= 0) { + float acc = 0.0f; +#pragma unroll + for (int tap = 0; tap < TAPS; ++tap) { + const int at = r + tap; + const float x = at < TAPS - 1 ? __bfloat162float(cs[at * C + c]) + : __bfloat162float(P[(size_t)(at - (TAPS - 1)) * PW + c]); + acc = acc + __bfloat162float(cw[c * TAPS + tap]) * x; + } + const float act = bf(acc / (1.0f + expf(-acc))); + if (t < DK) qs[t] = act; + else if (t < 2 * DK) ks[t - DK] = act; + else vs[t - 2 * DK] = act; + } + __syncthreads(); + if (warp < 2) { + float* x = warp == 0 ? qs : ks; + float v4[4], ss = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + v4[i] = x[lane * 4 + i]; + ss = ss + v4[i] * v4[i]; + } + ss = warp_sum(ss); + float inv = 1.0f / sqrtf(ss + 1e-6f); + if (warp == 0) inv = inv * (1.0f / sqrtf((float)DK)); + __syncwarp(); +#pragma unroll + for (int i = 0; i < 4; ++i) x[lane * 4 + i] = v4[i] * inv; + } else if (warp == 2 && lane == 0) { + const float b = __bfloat162float(pb[(size_t)r * PW]); + const float a = __bfloat162float(pa[(size_t)r * PW]); + gates[0] = expf(-expf(a_log[hv]) * softplusf_(a + dt_bias[hv])); + gates[1] = bf(sigmoidf_(b)); + } + __syncthreads(); + const float g = gates[0], beta = gates[1]; + float kk[4], qq[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + kk[i] = ks[lane * 4 + i]; + qq[i] = qs[lane * 4 + i]; + } + update(s, kk, vs, warp, g, beta); +#pragma unroll + for (int j = 0; j < 4; ++j) { + float o = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) o = o + s[j][i] * qq[i]; + o = warp_sum(o); + if (lane == 0) ys[warp * 4 + j] = bf(o); + } + __syncthreads(); + if (warp == 0) { + float ss = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float y = ys[lane * 4 + i]; + ss = ss + y * y; + } + ss = warp_sum(ss); + if (lane == 0) rinv = 1.0f / sqrtf(ss / (float)DV + eps); + } + __syncthreads(); + if (t < DV) { + const float yn = bf(bf(ys[t] * rinv) * __bfloat162float(norm_w[t])); + const float z = __bfloat162float(pz[(size_t)r * PW]); + const float o = bf(yn * sigmoidf_(z)); + out[(size_t)r * NV * DV + hv * DV + t] = __float2bfloat16_rn(o); + const float gs = warp_sum(o); + if (lane == 0) xs[(size_t)r * (NV * DV / 32) + hv * (DV / 32) + warp] = gs; + } + __syncthreads(); + } + if (state_out != nullptr) { +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + state_out[sbase + (size_t)(warp * 4 + j) * DK + lane * 4 + i] = s[j][i]; + } +} diff --git a/zig/kernels/cuda/qmm.cu b/zig/kernels/cuda/qmm.cu new file mode 100644 index 0000000000..773d1d5926 --- /dev/null +++ b/zig/kernels/cuda/qmm.cu @@ -0,0 +1,230 @@ +// Device code of src/tensorfold/cuda/kernels/qmm.cu (lines 1-230, comments and ATen includes dropped), checked by zig/tests/cuda/copies.py. + +#include +#include +#include +#include + +#include "qmm_frag.cuh" + +namespace tf_qmm { + +using namespace qmm_frag; + +template +__global__ void __launch_bounds__(WM * WN * 32) qmm_kernel( + const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, const uint32_t* __restrict__ w, + const __nv_bfloat16* __restrict__ scales, const __nv_bfloat16* __restrict__ biases, + void* __restrict__ out, float* __restrict__ part, int M, int N, int K, int SK, int npad, int ldx, int group) { + using T = LaneTile; + extern __shared__ __align__(128) unsigned char buf[]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int wm = warp / WN, wn = warp % WN; + const int KG = K / GS, per = KG / SK; + const int2 at = tile_of(blockIdx.x, M, N, BM, BN, group); + const int m0 = at.x, n0 = at.y, slice = blockIdx.z, g0 = slice * per; + + auto stage = [&](int s) { return buf + s * T::STAGE; }; + auto load = [&](int s, int g, int m0) { + unsigned char* p = stage(s); + 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 + static_cast(row) * ldx + g * GS + ch * 8, + m0 + r < M); + } + unsigned char* pw = p + T::X; + constexpr int TILE_BYTES = 64 * GS / 2; // one stored 64-column tile's group block + for (int c = tid; c < T::W / 16; c += T::THREADS) { + const int t = c / (TILE_BYTES / 16), off = c % (TILE_BYTES / 16); + const size_t tile = static_cast(n0 / 64 + t) * KG + g; + cp16(pw + c * 16, reinterpret_cast(w) + tile * TILE_BYTES + off * 16); + } + unsigned char* ps = pw + T::W; + for (int c = tid; c < 2 * (T::S / 16); c += T::THREADS) { + const int which = c / (T::S / 16), off = c % (T::S / 16); + const __nv_bfloat16* src = (which ? biases : scales) + static_cast(g) * npad + n0 + off * 8; + cp16(ps + which * T::S + off * 16, src); + } + float* px = reinterpret_cast(ps + 2 * T::S); + for (int r = tid; r < BM; r += T::THREADS) + cp4(px + r, xs + static_cast(min(m0 + r, M - 1)) * KG + g); + }; + + { + 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; + auto epilogue = [&](float (&acc_)[T::MT][T::NT][4], const float (&d_)[T::MT][T::NT][4], + const float (&s_)[T::NT][2], const float (&b_)[T::NT][2], const float (&x_)[T::MT][2]) { +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#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(x_[i][e >> 1], b_[j][e & 1], __fmaf_rn(d_[i][j][e], s_[j][e & 1], + acc_[i][j][e])); + }; + float dp[T::MT][T::NT][4], sp[T::NT][2], bp[T::NT][2], xp[T::MT][2]; +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < per) load(s, g0 + s, m0); + commit(); + } + for (int it = 0; it < per; ++it) { + wait(); + __syncthreads(); + const int next = it + STAGES - 1; + if (next < per) load(next % STAGES, g0 + next, m0); + commit(); + const unsigned char* p = stage(it % STAGES); + const uint32_t* pw = reinterpret_cast(p + T::X); + const __nv_bfloat16* ps = reinterpret_cast(p + T::X + T::W); + const float* px = reinterpret_cast(p + T::X + T::W + 2 * T::S); + uint32_t words[T::NT][GS / 32]; +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int v = 0; v < GS / 32; ++v) + words[j][v] = pw[((wn * T::NT + j) * 32 + lane) * (GS / 32) + v]; + float d[T::MT][T::NT][4]; +#pragma unroll + for (int kt = 0; kt < GS / 16; ++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) { + const uint32_t b0 = pair(words[j][kt / 2], (kt & 1) * 8); + const uint32_t b1 = pair(words[j][kt / 2], (kt & 1) * 8 + 4); +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + if (kt == 0) mma0(d[i][j], a[i], b0, b1); + else mma(d[i][j], a[i], b0, b1); + } + } + if (PIPE && kt == 0 && it > 0) epilogue(acc, dp, sp, bp, xp); // the previous group's, between MMAs + } + float sv[T::NT][2], bv[T::NT][2], xv[T::MT][2]; +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const int col = wn * (BN / WN) + j * 8 + (lane & 3) * 2; + const __nv_bfloat162 s2 = *reinterpret_cast(ps + col); + const __nv_bfloat162 b2 = *reinterpret_cast(ps + BN + col); + sv[j][0] = __low2float(s2); sv[j][1] = __high2float(s2); + bv[j][0] = __low2float(b2); bv[j][1] = __high2float(b2); + } +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + const int row = wm * (BM / WM) + i * 16 + (lane >> 2); + xv[i][0] = px[row]; + xv[i][1] = px[row + 8]; + } + if constexpr (PIPE) { +#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) dp[i][j][e] = d[i][j][e]; + xp[i][0] = xv[i][0]; xp[i][1] = xv[i][1]; + } +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + sp[j][0] = sv[j][0]; sp[j][1] = sv[j][1]; bp[j][0] = bv[j][0]; bp[j][1] = bv[j][1]; + } + } else { + epilogue(acc, d, sv, bv, xv); + } + } + if (PIPE && per > 0) epilogue(acc, dp, sp, bp, xp); + wait<0>(); + __syncthreads(); + if constexpr (CLUSTER) { +#if __CUDA_ARCH__ < 900 + __trap(); // no clusters before sm_90: the host never launches this +#else + constexpr int E = T::MT * T::NT * 4; + auto cluster = cooperative_groups::this_cluster(); + float* mine = reinterpret_cast(buf); + if (slice != 0) { +#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) mine[((i * T::NT + j) * 4 + e) * T::THREADS + tid] = acc[i][j][e]; + } + cluster.sync(); + if (slice == 0) { + for (int peer = 1; peer < SK; ++peer) { + const float* theirs = cluster.map_shared_rank(mine, peer); +#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] = acc[i][j][e] + theirs[((i * T::NT + j) * 4 + e) * T::THREADS + tid]; + } + } + cluster.sync(); // peers keep their memory until slice 0 has read it + if (slice != 0) return; +#endif + } +#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 * (BN / WN) + j * 8 + (lane & 3) * 2; +#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 && !CLUSTER) { + float* dst = part + (static_cast(slice) * M + row) * N + col; + if (col < N) dst[0] = v0; + if (col + 1 < N) dst[1] = v1; + } else if (F32) { + float* dst = reinterpret_cast(out) + static_cast(row) * N + col; + if (col < N) dst[0] = v0; + if (col + 1 < N) dst[1] = v1; + } else { + __nv_bfloat16* dst = reinterpret_cast<__nv_bfloat16*>(out) + static_cast(row) * N + col; + if (col + 1 < N && (N & 1) == 0) + *reinterpret_cast<__nv_bfloat162*>(dst) = __floats2bfloat162_rn(v0, v1); + else { + if (col < N) dst[0] = __float2bfloat16_rn(v0); + if (col + 1 < N) dst[1] = __float2bfloat16_rn(v1); + } + } + } + } + } +} + +template +__global__ void reduce_kernel(const float* __restrict__ part, void* __restrict__ out, long long total, int SK) { + const long long i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= total) return; + float acc = part[i]; + for (int s = 1; s < SK; ++s) acc = acc + part[s * total + i]; + if (F32) reinterpret_cast(out)[i] = acc; + else reinterpret_cast<__nv_bfloat16*>(out)[i] = __float2bfloat16_rn(acc); +} +} // namespace tf_qmm + +// Decode lane: groups of 32, 16 rows, 64 columns, one by four warps, four stages, bf16 out. +template __global__ void tf_qmm::qmm_kernel<32, 16, 64, 1, 4, 4, false, false, false>( + const __nv_bfloat16*, const float*, const uint32_t*, const __nv_bfloat16*, const __nv_bfloat16*, + void*, float*, int, int, int, int, int, int, int); diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 7881d97b11..80363d43f2 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -13,8 +13,10 @@ fn Blob(comptime import_name: []const u8) type { pub const available = options.with_kernels; pub const gdn: []const u8 = if (available) &Blob("fatbin_gdn").bytes else &.{}; +pub const flashnext_gdn: []const u8 = if (available) &Blob("fatbin_flashnext_gdn").bytes else &.{}; pub const probe: []const u8 = if (available) &Blob("fatbin_probe").bytes else &.{}; pub const qmm_group: []const u8 = if (available) &Blob("fatbin_qmm_group").bytes else &.{}; +pub const qmm: []const u8 = if (available) &Blob("fatbin_qmm").bytes else &.{}; pub const qmm_prefill: []const u8 = if (available) &Blob("fatbin_qmm_prefill").bytes else &.{}; pub const experts: []const u8 = if (available) &Blob("fatbin_experts").bytes else &.{}; pub const experts_prefill: []const u8 = if (available) &Blob("fatbin_experts_prefill").bytes else &.{}; @@ -37,7 +39,7 @@ pub const gdn_symbols = struct { /// Every instantiation gdn.cu exports; tests resolve each one so a wrong name cannot hide behind an unused path. pub const gdn_variants = [_]TreeVariant{ - tree(0, 8, 4, true), tree(1, 8, 4, false), tree(2, 8, 2, false), tree(2, 4, 4, false), + tree(0, 8, 4, true), tree(1, 8, 4, false), tree(2, 8, 2, false), tree(2, 4, 4, false), tree(4, 2, 4, false), tree(8, 2, 4, false), tree(16, 2, 2, false), tree(32, 2, 1, false), }; diff --git a/zig/src/families/flashnext/cuda_attn.zig b/zig/src/families/flashnext/cuda_attn.zig new file mode 100644 index 0000000000..47d823fc1a --- /dev/null +++ b/zig/src/families/flashnext/cuda_attn.zig @@ -0,0 +1,442 @@ +//! Layer 47's full attention for one decode token at position 0: project, prepare, attend, gate, project out. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const hc = @import("cuda_hc.zig"); +const qmm = @import("cuda_qmm.zig"); + +const dims: usize = 2560; +const q_heads: usize = 24; +const kv_heads: usize = 2; +const head_dim: usize = 256; +const index_heads: usize = 4; +const index_dim: usize = 128; +const half: usize = 32; +const eps: f32 = 1e-6; +const attn_scale: f32 = 0.0625; +const proj_n: usize = q_heads * 2 * head_dim + 2 * kv_heads * head_dim + (index_heads + 1) * index_dim; +const out_k: usize = q_heads * head_dim; +const nch: usize = 3; +const prefix = "language_model.model.layers.47.self_attn."; + +const Gap = struct { off: usize, steps: u32, at: usize }; + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +fn readBf(dst: []u16, raw: []const u8) void { + for (dst, 0..) |*o, i| o.* = std.mem.readInt(u16, raw[2 * i ..][0..2], .little); +} + +fn treeSum(buf: []f32) f32 { + var n = buf.len; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + return buf[0]; +} + +fn sumSq(x: []const u16) f32 { + var buf: [head_dim]f32 = undefined; + for (x, 0..) |bits, i| { + const v = promote(bits); + buf[i] = v * v; + } + return treeSum(buf[0..x.len]); +} + +/// glue._attn_prep at position 0: RMSNorm, round, rotate-half. A zero angle stores the normalized row. +fn normRot(x: []const u16, scale_w: []const f32, inv: []const f32, pos: f32, out: []u16) void { + const width = x.len; + const rinv = 1.0 / @sqrt(sumSq(x) / @as(f32, @floatFromInt(width)) + eps); + var xn: [head_dim]f32 = undefined; + for (0..width) |d| xn[d] = promote(toBf16(promote(x[d]) * rinv * scale_w[d])); + for (0..width) |d| { + const partner = if (d < half) d + half else if (d < 2 * half) d - half else d; + const axis_i = if (d < half) d else if (d < 2 * half) d - half else 0; + const ang = pos * inv[axis_i]; + const c = @cos(ang); + const s = @sin(ang); + const a = xn[d]; + const b = xn[partner]; + const rot = if (d < half) a * c - b * s else if (d < 2 * half) b * s + a * c else a; + out[d] = toBf16(rot); + } +} + +fn copyRow(src: []const u16, dst: []u16) void { + @memcpy(dst, src); +} + +/// One key: the softmax weight is 1, so each query head copies its KV head's value. +fn attend(q: []const u16, k: []const u16, v: []const u16, out: []u16) void { + _ = q; + _ = k; + copyRow(v, out); +} + +fn applyGate(o: []const u16, g: []const u16, out: []u16, xs: []f32) void { + for (o, g, out) |ov, gv, *y| y.* = toBf16(promote(ov) / (1.0 + @exp(-promote(gv)))); + for (xs, 0..) |*s, i| { + var buf: [32]f32 = undefined; + for (0..32) |j| buf[j] = promote(out[i * 32 + j]); + s.* = treeSum(buf[0..]); + } +} + +fn gap(got: []const u16, want: []const u16) Gap { + var off: usize = 0; + var steps: u32 = 0; + var at: usize = 0; + for (got, want, 0..) |a, b, i| { + const d = hc.mixedSteps(a, b); + if (d != 0) off += 1; + if (d > steps) { + steps = d; + at = i; + } + } + return .{ .off = off, .steps = steps, .at = at }; +} + +fn shardPath(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) ![]u8 { + const index = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); + defer gpa.free(index); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const file = parsed.value.object.get("weight_map").?.object.get(prefix ++ "q_proj.weight").?.string; + return std.fs.path.join(gpa, &.{ dir, file }); +} + +const Triple = struct { w: core.safetensors.Tensor, s: core.safetensors.Tensor, b: core.safetensors.Tensor }; + +fn triple(file: *const core.safetensors.File, name: []const u8, rows: usize, k8: usize) !Triple { + var buf: [180]u8 = undefined; + const w = file.get(try std.fmt.bufPrint(&buf, "{s}.weight", .{name})) orelse return error.MissingTensor; + const s = file.get(try std.fmt.bufPrint(&buf, "{s}.scales", .{name})) orelse return error.MissingTensor; + const b = file.get(try std.fmt.bufPrint(&buf, "{s}.biases", .{name})) orelse return error.MissingTensor; + if (!w.is(.u32, &.{ rows, k8 }) or !s.is(.bf16, &.{ rows, k8 / 4 }) or !b.is(.bf16, &.{ rows, k8 / 4 })) return error.UnexpectedTensor; + return .{ .w = w, .s = s, .b = b }; +} + +fn stitch(gpa: std.mem.Allocator, parts: []const Triple) !struct { words: []u8, scales: []u8, biases: []u8 } { + var wb: usize = 0; + var sb: usize = 0; + for (parts) |p| { + wb += p.w.bytes.len; + sb += p.s.bytes.len; + } + const words = try gpa.alloc(u8, wb); + errdefer gpa.free(words); + const scales = try gpa.alloc(u8, sb); + errdefer gpa.free(scales); + const biases = try gpa.alloc(u8, sb); + var wo: usize = 0; + var so: usize = 0; + for (parts) |p| { + @memcpy(words[wo..][0..p.w.bytes.len], p.w.bytes); + wo += p.w.bytes.len; + @memcpy(scales[so..][0..p.s.bytes.len], p.s.bytes); + @memcpy(biases[so..][0..p.b.bytes.len], p.b.bytes); + so += p.s.bytes.len; + } + return .{ .words = words, .scales = scales, .biases = biases }; +} + +fn promoteScale(bytes: []const u8, out: []f32) void { + for (out, 0..) |*o, i| o.* = promote(std.mem.readInt(u16, bytes[2 * i ..][0..2], .little)); +} + +fn invFreq(out: []f32) void { + const n: f64 = @floatFromInt(out.len); + for (out, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / n)); +} + +fn prepare(pa: []const u16, q_scale: []const f32, k_scale: []const f32, i_scale: []const f32, inv: []const f32, q: []u16, k: []u16, v: []u16, iq: []u16, ik: []u16) void { + const q_pairs = q_heads * 2 * head_dim; + const kv = kv_heads * head_dim; + for (0..q_heads) |h| normRot(pa[h * 2 * head_dim ..][0..head_dim], q_scale, inv, 0, q[h * head_dim ..][0..head_dim]); + for (0..kv_heads) |h| { + normRot(pa[q_pairs + h * head_dim ..][0..head_dim], k_scale, inv, 0, k[h * head_dim ..][0..head_dim]); + copyRow(pa[q_pairs + kv + h * head_dim ..][0..head_dim], v[h * head_dim ..][0..head_dim]); + } + const iq0 = q_pairs + 2 * kv; + for (0..index_heads) |h| normRot(pa[iq0 + h * index_dim ..][0..index_dim], i_scale, inv, 0, iq[h * index_dim ..][0..index_dim]); + copyRow(pa[iq0 + index_heads * index_dim ..][0..index_dim], ik); +} + +fn hostAttend(q: []const u16, k: []const u16, v: []const u16, out: []u16) void { + const g = q_heads / kv_heads; + for (0..q_heads) |h| { + const hk = h / g; + attend(q[h * head_dim ..][0..head_dim], k[hk * head_dim ..][0..head_dim], v[hk * head_dim ..][0..head_dim], out[h * head_dim ..][0..head_dim]); + } +} + +fn gatesOf(pa: []const u16, out: []u16) void { + for (0..q_heads) |h| copyRow(pa[h * 2 * head_dim + head_dim ..][0..head_dim], out[h * head_dim ..][0..head_dim]); +} + +/// The projection, prep, one-key attention, gate and output projection, against the host of the same mixed row. +pub fn decode(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, io: std.Io, model_dir: []const u8, tri: Tri, mixed: u64, xs: u64) !u32 { + const path = try shardPath(gpa, io, model_dir); + defer gpa.free(path); + var file = try core.safetensors.File.open(gpa, io, path); + defer file.close(io); + const parts = [_]Triple{ + try triple(&file, prefix ++ "q_proj", q_heads * 2 * head_dim, dims / 8), + try triple(&file, prefix ++ "k_proj", kv_heads * head_dim, dims / 8), + try triple(&file, prefix ++ "v_proj", kv_heads * head_dim, dims / 8), + try triple(&file, prefix ++ "indexer.index_qk_proj", (index_heads + 1) * index_dim, dims / 8), + }; + const out_w = try triple(&file, prefix ++ "o_proj", dims, out_k / 8); + const qn = file.get(prefix ++ "q_norm.weight") orelse return error.MissingTensor; + const kn = file.get(prefix ++ "k_norm.weight") orelse return error.MissingTensor; + const iqn = file.get(prefix ++ "indexer.q_layernorm.weight") orelse return error.MissingTensor; + if (!qn.is(.bf16, &.{head_dim}) or !kn.is(.bf16, &.{head_dim}) or !iqn.is(.bf16, &.{index_dim})) return error.UnexpectedTensor; + if (proj_n != 13952) return error.UnexpectedTensor; + + const stacked = try stitch(gpa, &parts); + defer gpa.free(stacked.words); + defer gpa.free(stacked.scales); + defer gpa.free(stacked.biases); + var lane = try qmm.pack(gpa, stacked.words, stacked.scales, stacked.biases, proj_n, dims / 8); + defer lane.deinit(gpa); + var ow = try qmm.pack(gpa, out_w.w.bytes, out_w.s.bytes, out_w.b.bytes, dims, out_k / 8); + defer ow.deinit(gpa); + + var q_scale: [head_dim]f32 = undefined; + var k_scale: [head_dim]f32 = undefined; + var i_scale: [index_dim]f32 = undefined; + var inv: [half]f32 = undefined; + promoteScale(qn.bytes, &q_scale); + promoteScale(kn.bytes, &k_scale); + promoteScale(iqn.bytes, &i_scale); + invFreq(&inv); + + var w_b = try cuda.DeviceBuffer.fromHost(driver, lane.weight); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, lane.scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, lane.biases); + defer b_b.free(); + var pa_b = try cuda.DeviceBuffer.alloc(driver, proj_n * 2); + defer pa_b.free(); + try qmm.matmul(driver, stream.*, mixed, xs, w_b.ptr, s_b.ptr, b_b.ptr, pa_b.ptr, 1, proj_n, dims); + + var qw_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&q_scale)); + defer qw_b.free(); + var kw_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&k_scale)); + defer kw_b.free(); + var iw_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&i_scale)); + defer iw_b.free(); + var inv_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&inv)); + defer inv_b.free(); + var pos_host: i32 = 0; + var pos_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.asBytes(&pos_host)); + defer pos_b.free(); + var q_b = try cuda.DeviceBuffer.alloc(driver, q_heads * head_dim * 2); + defer q_b.free(); + var kc_b = try cuda.DeviceBuffer.alloc(driver, kv_heads * head_dim * 2); + defer kc_b.free(); + var vc_b = try cuda.DeviceBuffer.alloc(driver, kv_heads * head_dim * 2); + defer vc_b.free(); + var iq_b = try cuda.DeviceBuffer.alloc(driver, index_heads * index_dim * 2); + defer iq_b.free(); + var ik_b = try cuda.DeviceBuffer.alloc(driver, index_dim * 2); + defer ik_b.free(); + var dummy = try cuda.DeviceBuffer.alloc(driver, 256); + defer dummy.free(); + try tri.attnPrep(pa_b.ptr, pos_b.ptr, qw_b.ptr, kw_b.ptr, iw_b.ptr, inv_b.ptr, q_b.ptr, kc_b.ptr, vc_b.ptr, dummy.ptr, dummy.ptr, iq_b.ptr, ik_b.ptr, dummy.ptr, dummy.ptr, eps, 1, 0); + + var po_b = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * head_dim * 4); + defer po_b.free(); + var pm_b = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * 4); + defer pm_b.free(); + var pl_b = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * 4); + defer pl_b.free(); + var o_b = try cuda.DeviceBuffer.alloc(driver, out_k * 2); + defer o_b.free(); + try tri.attnChunks(q_b.ptr, kc_b.ptr, vc_b.ptr, dummy.ptr, dummy.ptr, pos_b.ptr, po_b.ptr, pm_b.ptr, pl_b.ptr, dummy.ptr, dummy.ptr, dummy.ptr, 1, kv_heads, 1); + try tri.attnMerge(po_b.ptr, pm_b.ptr, pl_b.ptr, pos_b.ptr, o_b.ptr, dummy.ptr, dummy.ptr, 1, kv_heads); + + var gated_b = try cuda.DeviceBuffer.alloc(driver, out_k * 2); + defer gated_b.free(); + var gxs_b = try cuda.DeviceBuffer.alloc(driver, (out_k / 32) * 4); + defer gxs_b.free(); + try tri.attnGate(o_b.ptr, pa_b.ptr, gated_b.ptr, gxs_b.ptr, 1, proj_n, q_heads, head_dim); + + var ow_b = try cuda.DeviceBuffer.fromHost(driver, ow.weight); + defer ow_b.free(); + var os_b = try cuda.DeviceBuffer.fromHost(driver, ow.scales); + defer os_b.free(); + var ob_b = try cuda.DeviceBuffer.fromHost(driver, ow.biases); + defer ob_b.free(); + var branch_b = try cuda.DeviceBuffer.alloc(driver, dims * 2); + defer branch_b.free(); + try qmm.matmul(driver, stream.*, gated_b.ptr, gxs_b.ptr, ow_b.ptr, os_b.ptr, ob_b.ptr, branch_b.ptr, 1, dims, out_k); + + const row_raw = try gpa.alloc(u8, dims * 2); + defer gpa.free(row_raw); + const xs_raw = try gpa.alloc(u8, (dims / 32) * 4); + defer gpa.free(xs_raw); + const pa_raw = try gpa.alloc(u8, proj_n * 2); + defer gpa.free(pa_raw); + const q_raw = try gpa.alloc(u8, q_heads * head_dim * 2); + defer gpa.free(q_raw); + const k_raw = try gpa.alloc(u8, kv_heads * head_dim * 2); + defer gpa.free(k_raw); + const v_raw = try gpa.alloc(u8, kv_heads * head_dim * 2); + defer gpa.free(v_raw); + const iq_raw = try gpa.alloc(u8, index_heads * index_dim * 2); + defer gpa.free(iq_raw); + const ik_raw = try gpa.alloc(u8, index_dim * 2); + defer gpa.free(ik_raw); + const o_raw = try gpa.alloc(u8, out_k * 2); + defer gpa.free(o_raw); + const gated_raw = try gpa.alloc(u8, out_k * 2); + defer gpa.free(gated_raw); + const gxs_raw = try gpa.alloc(u8, (out_k / 32) * 4); + defer gpa.free(gxs_raw); + const branch_raw = try gpa.alloc(u8, dims * 2); + defer gpa.free(branch_raw); + try stream.synchronize(); + var mixed_view: cuda.DeviceBuffer = .{ .d = driver, .ptr = mixed, .len = row_raw.len }; + var xs_view: cuda.DeviceBuffer = .{ .d = driver, .ptr = xs, .len = xs_raw.len }; + try mixed_view.download(0, row_raw); + try xs_view.download(0, xs_raw); + try pa_b.download(0, pa_raw); + try q_b.download(0, q_raw); + try kc_b.download(0, k_raw); + try vc_b.download(0, v_raw); + try iq_b.download(0, iq_raw); + try ik_b.download(0, ik_raw); + try o_b.download(0, o_raw); + try gated_b.download(0, gated_raw); + try gxs_b.download(0, gxs_raw); + try branch_b.download(0, branch_raw); + + const row = try gpa.alloc(u16, dims); + defer gpa.free(row); + readBf(row, row_raw); + const xs_f = try gpa.alloc(f32, dims / 32); + defer gpa.free(xs_f); + for (xs_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xs_raw[4 * i ..][0..4], .little)); + const host_pa = try gpa.alloc(u16, proj_n); + defer gpa.free(host_pa); + try qmm.dotRow(row, xs_f, stacked.words, stacked.scales, stacked.biases, proj_n, dims, host_pa); + const gpu_pa = try gpa.alloc(u16, proj_n); + defer gpa.free(gpu_pa); + readBf(gpu_pa, pa_raw); + const proj = gap(gpu_pa, host_pa); + + const host_q = try gpa.alloc(u16, q_heads * head_dim); + defer gpa.free(host_q); + const host_k = try gpa.alloc(u16, kv_heads * head_dim); + defer gpa.free(host_k); + const host_v = try gpa.alloc(u16, kv_heads * head_dim); + defer gpa.free(host_v); + const host_iq = try gpa.alloc(u16, index_heads * index_dim); + defer gpa.free(host_iq); + const host_ik = try gpa.alloc(u16, index_dim); + defer gpa.free(host_ik); + prepare(gpu_pa, &q_scale, &k_scale, &i_scale, &inv, host_q, host_k, host_v, host_iq, host_ik); + const gpu_q = try gpa.alloc(u16, host_q.len); + defer gpa.free(gpu_q); + const gpu_k = try gpa.alloc(u16, host_k.len); + defer gpa.free(gpu_k); + const gpu_v = try gpa.alloc(u16, host_v.len); + defer gpa.free(gpu_v); + const gpu_iq = try gpa.alloc(u16, host_iq.len); + defer gpa.free(gpu_iq); + const gpu_ik = try gpa.alloc(u16, host_ik.len); + defer gpa.free(gpu_ik); + readBf(gpu_q, q_raw); + readBf(gpu_k, k_raw); + readBf(gpu_v, v_raw); + readBf(gpu_iq, iq_raw); + readBf(gpu_ik, ik_raw); + const qg = gap(gpu_q, host_q); + const kg = gap(gpu_k, host_k); + const vg = gap(gpu_v, host_v); + const iqg = gap(gpu_iq, host_iq); + const ikg = gap(gpu_ik, host_ik); + + const host_o = try gpa.alloc(u16, out_k); + defer gpa.free(host_o); + hostAttend(gpu_q, gpu_k, gpu_v, host_o); + const gpu_o = try gpa.alloc(u16, out_k); + defer gpa.free(gpu_o); + readBf(gpu_o, o_raw); + const og = gap(gpu_o, host_o); + + const gate_row = try gpa.alloc(u16, out_k); + defer gpa.free(gate_row); + gatesOf(gpu_pa, gate_row); + const host_gated = try gpa.alloc(u16, out_k); + defer gpa.free(host_gated); + const host_gxs = try gpa.alloc(f32, out_k / 32); + defer gpa.free(host_gxs); + applyGate(gpu_o, gate_row, host_gated, host_gxs); + const gpu_gated = try gpa.alloc(u16, out_k); + defer gpa.free(gpu_gated); + readBf(gpu_gated, gated_raw); + const gg = gap(gpu_gated, host_gated); + var xs_ulp: u32 = 0; + for (host_gxs, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, gxs_raw[4 * i ..][0..4], .little)); + const dist = hc.ulps(got, want); + if (dist > xs_ulp) xs_ulp = dist; + } + + const gxs_f = try gpa.alloc(f32, out_k / 32); + defer gpa.free(gxs_f); + for (gxs_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, gxs_raw[4 * i ..][0..4], .little)); + const host_branch = try gpa.alloc(u16, dims); + defer gpa.free(host_branch); + try qmm.dotRow(gpu_gated, gxs_f, out_w.w.bytes, out_w.s.bytes, out_w.b.bytes, dims, out_k, host_branch); + const gpu_branch = try gpa.alloc(u16, dims); + defer gpa.free(gpu_branch); + readBf(gpu_branch, branch_raw); + const bg = gap(gpu_branch, host_branch); + + std.debug.print("attn proj n {d} off {d} max_steps {d} at {d} got {x:0>4} host {x:0>4} head {x:0>4}\n", .{ proj_n, proj.off, proj.steps, proj.at, gpu_pa[proj.at], host_pa[proj.at], gpu_pa[0] }); + std.debug.print("attn prep q {d}/{d} k {d}/{d} v {d}/{d} iq {d}/{d} ik {d}/{d} head {x:0>4} host {x:0>4}\n", .{ qg.off, qg.steps, kg.off, kg.steps, vg.off, vg.steps, iqg.off, iqg.steps, ikg.off, ikg.steps, gpu_q[0], host_q[0] }); + std.debug.print("attn out off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ og.off, og.steps, gpu_o[0], host_o[0] }); + std.debug.print("attn gate off {d} max_steps {d} xs_ulp {d} head {x:0>4} host {x:0>4}\n", .{ gg.off, gg.steps, xs_ulp, gpu_gated[0], host_gated[0] }); + std.debug.print("attn oproj n {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ dims, bg.off, bg.steps, gpu_branch[0], host_branch[0] }); + if (proj.off > 2 or proj.steps > 1 or qg.steps != 0 or kg.off > 1 or kg.steps > 1 or vg.steps != 0 or iqg.steps != 0 or ikg.steps != 0 or og.steps != 0 or gg.off > 1 or gg.steps > 1 or xs_ulp > 8192 or bg.steps != 0) return 1; + return 0; +} + +test "position zero stores the normalized query" { + var x: [head_dim]u16 = @splat(0); + x[0] = 0x3f80; + var scale_w: [head_dim]f32 = @splat(1); + var inv: [half]f32 = @splat(1); + var out: [head_dim]u16 = undefined; + normRot(&x, &scale_w, &inv, 0, &out); + try std.testing.expectEqual(@as(u16, 0x4180), out[0]); + try std.testing.expectEqual(@as(u16, 0), out[half]); +} + +test "the gate rounds one times sigmoid of zero to one half" { + var o: [head_dim]u16 = @splat(0); + var g: [head_dim]u16 = @splat(0); + o[0] = 0x3f80; + var out: [head_dim]u16 = undefined; + var xs: [head_dim / 32]f32 = undefined; + applyGate(&o, &g, &out, &xs); + try std.testing.expectEqual(@as(u16, 0x3f00), out[0]); +} diff --git a/zig/src/families/flashnext/cuda_embed.zig b/zig/src/families/flashnext/cuda_embed.zig new file mode 100644 index 0000000000..0b31f38a69 --- /dev/null +++ b/zig/src/families/flashnext/cuda_embed.zig @@ -0,0 +1,171 @@ +//! The MLX 4-bit group-32 token table: host dequant matches glue._embed, and the same bytes go to the GPU. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); + +const Tensor = core.safetensors.Tensor; + +pub const weight_name = "language_model.model.embed_tokens.weight"; +pub const scales_name = "language_model.model.embed_tokens.scales"; +pub const biases_name = "language_model.model.embed_tokens.biases"; + +const Triple = struct { weight: Tensor, scales: Tensor, biases: Tensor }; + +fn tensorAt(obj: std.json.ObjectMap, data: []const u8, name: []const u8, dtype: core.safetensors.DType) !Tensor { + const o = (obj.get(name) orelse return error.MissingTensor).object; + const got = core.safetensors.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; + if (got != dtype) return error.UnexpectedTensor; + const shape = o.get("shape").?.array.items; + if (shape.len == 0 or shape.len > 4) return error.UnexpectedTensor; + const offs = o.get("data_offsets").?.array.items; + const begin: usize = @intCast(offs[0].integer); + const end: usize = @intCast(offs[1].integer); + if (end < begin or end > data.len) return error.BadSafetensors; + var t: Tensor = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .bytes = data[begin..end] }; + for (shape, 0..) |d, i| t.shape[i] = @intCast(d.integer); + return t; +} + +/// The three embedding tensors inside one safetensors header. Other tensors, including rank above 4, are ignored. +pub fn readTriple(gpa: std.mem.Allocator, header: []const u8, data: []const u8) !Triple { + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, header, .{}); + defer parsed.deinit(); + const obj = parsed.value.object; + return .{ + .weight = try tensorAt(obj, data, weight_name, .u32), + .scales = try tensorAt(obj, data, scales_name, .bf16), + .biases = try tensorAt(obj, data, biases_name, .bf16), + }; +} + +/// The shard that holds the embedding table, mapped until `close`. +pub const Mapped = struct { + io: std.Io, + file: std.Io.File, + map: std.Io.File.MemoryMap, + weight: Tensor, + scales: Tensor, + biases: Tensor, + + pub fn open(gpa: std.mem.Allocator, io: std.Io, path: []const u8) !Mapped { + var file = try std.Io.Dir.cwd().openFile(io, path, .{}); + errdefer file.close(io); + const len: usize = @intCast(try file.length(io)); + if (len < 8) return error.BadSafetensors; + var map = try std.Io.File.MemoryMap.create(io, file, .{ .len = len, .protection = .{ .read = true, .write = false }, .populate = false }); + errdefer map.destroy(io); + const header_len: usize = @intCast(std.mem.readInt(u64, map.memory[0..8], .little)); + if (header_len > len - 8) return error.BadSafetensors; + const triple = try readTriple(gpa, map.memory[8..][0..header_len], map.memory[8 + header_len ..]); + return .{ .io = io, .file = file, .map = map, .weight = triple.weight, .scales = triple.scales, .biases = triple.biases }; + } + + /// One named tensor in this shard. Rank 1 to 4. + pub fn lookup(self: *const Mapped, gpa: std.mem.Allocator, name: []const u8, dtype: core.safetensors.DType) !Tensor { + const header_len: usize = @intCast(std.mem.readInt(u64, self.map.memory[0..8], .little)); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, self.map.memory[8..][0..header_len], .{}); + defer parsed.deinit(); + return tensorAt(parsed.value.object, self.map.memory[8 + header_len ..], name, dtype); + } + + pub fn close(self: *Mapped) void { + self.map.destroy(self.io); + self.file.close(self.io); + self.* = undefined; + } +}; + +pub const Table = struct { + w: cuda.DeviceBuffer, + s: cuda.DeviceBuffer, + b: cuda.DeviceBuffer, + rows: usize, + dims: usize, + + /// Copies the three tensors. `dims` is hidden size (eight values per stored word). + pub fn upload(d: *const cuda.Driver, weight: Tensor, scales: Tensor, biases: Tensor) !Table { + if (weight.dtype != .u32 or scales.dtype != .bf16 or biases.dtype != .bf16) return error.UnsupportedQuantization; + if (weight.rank != 2 or scales.rank != 2 or biases.rank != 2) return error.UnexpectedTensor; + const rows = weight.dim(0); + const dims = weight.dim(1) * 8; + if (dims == 0 or dims % 32 != 0 or scales.dim(0) != rows or biases.dim(0) != rows) return error.UnexpectedTensor; + if (scales.dim(1) != dims / 32 or biases.dim(1) != dims / 32) return error.UnexpectedTensor; + var w = try cuda.DeviceBuffer.fromHost(d, weight.bytes); + errdefer w.free(); + var s = try cuda.DeviceBuffer.fromHost(d, scales.bytes); + errdefer s.free(); + var b = try cuda.DeviceBuffer.fromHost(d, biases.bytes); + errdefer b.free(); + return .{ .w = w, .s = s, .b = b, .rows = rows, .dims = dims }; + } + + pub fn deinit(self: *Table) void { + self.w.free(); + self.s.free(); + self.b.free(); + self.* = undefined; + } +}; + +/// fp32 to bf16, round to nearest even, the rounding Triton uses. +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn f32FromBf16(b: []const u8, i: usize) f32 { + return @bitCast(@as(u32, std.mem.readInt(u16, b[2 * i ..][0..2], .little)) << 16); +} + +/// One token row: each group of 32 is `nibble * scale + bias`, stored bf16. Same order as glue._embed. +pub fn dequant(weight: []const u8, scales: []const u8, biases: []const u8, dims: usize, token: usize, out: []u16) !void { + if (dims == 0 or dims % 32 != 0 or out.len != dims) return error.UnexpectedTensor; + const groups = dims / 32; + const words = dims / 8; + const woff = token * words * 4; + const soff = token * groups * 2; + if (woff + words * 4 > weight.len or soff + groups * 2 > scales.len or soff + groups * 2 > biases.len) return error.UnexpectedTensor; + for (0..groups) |g| { + const scale = f32FromBf16(scales[soff..], g); + const bias = f32FromBf16(biases[soff..], g); + for (0..4) |i| { + const word = std.mem.readInt(u32, weight[woff + (g * 4 + i) * 4 ..][0..4], .little); + for (0..8) |j| { + const q: f32 = @floatFromInt((word >> @intCast(j * 4)) & 0xF); + out[g * 32 + i * 8 + j] = toBf16(q * scale + bias); + } + } + } +} + +test "embedding bytes are read when another tensor has rank above 4" { + const header = + \\{"high":{"dtype":"U8","shape":[1,1,1,1,1],"data_offsets":[0,1]}, + \\ "language_model.model.embed_tokens.weight":{"dtype":"U32","shape":[1,4],"data_offsets":[1,17]}, + \\ "language_model.model.embed_tokens.scales":{"dtype":"BF16","shape":[1,1],"data_offsets":[17,19]}, + \\ "language_model.model.embed_tokens.biases":{"dtype":"BF16","shape":[1,1],"data_offsets":[19,21]}} + ; + var data: [21]u8 = @splat(0); + data[1] = 0x21; + std.mem.writeInt(u16, data[17..19], 0x3f80, .little); + const triple = try readTriple(std.testing.allocator, header, &data); + try std.testing.expectEqual(@as(usize, 4), triple.weight.dim(1)); + var out: [32]u16 = undefined; + try dequant(triple.weight.bytes, triple.scales.bytes, triple.biases.bytes, 32, 0, &out); + try std.testing.expectEqual(@as(u16, 0x3f80), out[0]); +} + +test "one embedding group dequants nibbles with scale and bias" { + var weight: [16]u8 = @splat(0); + std.mem.writeInt(u32, weight[0..4], 0x21, .little); + var scales: [2]u8 = undefined; + std.mem.writeInt(u16, &scales, 0x3f80, .little); + const biases: [2]u8 = @splat(0); + var out: [32]u16 = undefined; + try dequant(&weight, &scales, &biases, 32, 0, &out); + try std.testing.expectEqual(@as(u16, 0x3f80), out[0]); + try std.testing.expectEqual(@as(u16, 0x4000), out[1]); + try std.testing.expectEqual(@as(u16, 0), out[2]); +} diff --git a/zig/src/families/flashnext/cuda_embed_main.zig b/zig/src/families/flashnext/cuda_embed_main.zig new file mode 100644 index 0000000000..294e0e6b69 --- /dev/null +++ b/zig/src/families/flashnext/cuda_embed_main.zig @@ -0,0 +1,393 @@ +//! One token through the captured _embed cubin, compared with the host dequant of the same row. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); + +const embed = flash.embed; + +fn shardPath(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) ![]u8 { + const index = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); + defer gpa.free(index); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const file = parsed.value.object.get("weight_map").?.object.get(embed.weight_name).?.string; + return std.fs.path.join(gpa, &.{ dir, file }); +} + +pub fn main(init: std.process.Init) !u8 { + const gpa = init.gpa; + const io = init.io; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 4) { + std.debug.print("usage: flashnext-embed \n", .{}); + return 2; + } + const token = try std.fmt.parseInt(u32, args[3], 10); + const path = try shardPath(gpa, io, args[1]); + defer gpa.free(path); + var mapped = try embed.Mapped.open(gpa, io, path); + defer mapped.close(); + const weight = mapped.weight; + const scales = mapped.scales; + const biases = mapped.biases; + + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + var table = try embed.Table.upload(&driver, weight, scales, biases); + defer table.deinit(); + if (token >= table.rows) return error.UnexpectedTensor; + const host = try gpa.alloc(u16, table.dims); + defer gpa.free(host); + try embed.dequant(weight.bytes, scales.bytes, biases.bytes, table.dims, token, host); + + var set = try cuda.aot.Set.load(gpa, io, &driver, ctx.device, args[2]); + defer set.deinit(); + var stream = try cuda.Stream.init(&driver, true); + defer stream.deinit(); + var ids = try cuda.DeviceBuffer.fromHost(&driver, std.mem.asBytes(&token)); + defer ids.free(); + var out = try cuda.DeviceBuffer.alloc(&driver, table.dims * 2); + defer out.free(); + try flash.Tri.embed(.{ .set = &set, .s = stream }, ids.ptr, table.w.ptr, table.s.ptr, table.b.ptr, out.ptr, table.dims, 1, 1); + + const gpu = try gpa.alloc(u8, table.dims * 2); + defer gpa.free(gpu); + try stream.synchronize(); + try out.download(0, gpu); + + var mismatch: usize = 0; + for (host, 0..) |want, i| { + const got = std.mem.readInt(u16, gpu[2 * i ..][0..2], .little); + if (got != want) mismatch += 1; + } + std.debug.print("token {d} rows {d} dims {d} mismatch {d} head", .{ token, table.rows, table.dims, mismatch }); + for (0..@min(8, table.dims)) |i| std.debug.print(" {x:0>4}", .{std.mem.readInt(u16, gpu[2 * i ..][0..2], .little)}); + std.debug.print("\n", .{}); + if (mismatch != 0) return 1; + + const streams: usize = 4; + var wide = try cuda.DeviceBuffer.alloc(&driver, streams * table.dims * 2); + defer wide.free(); + try flash.Tri.embed(.{ .set = &set, .s = stream }, ids.ptr, table.w.ptr, table.s.ptr, table.b.ptr, wide.ptr, table.dims, streams, 1); + const wide_bytes = try gpa.alloc(u8, streams * table.dims * 2); + defer gpa.free(wide_bytes); + try stream.synchronize(); + try wide.download(0, wide_bytes); + const wide_u16 = try gpa.alloc(u16, streams * table.dims); + defer gpa.free(wide_u16); + for (wide_u16, 0..) |*v, i| v.* = std.mem.readInt(u16, wide_bytes[2 * i ..][0..2], .little); + + var copies_differ = false; + for (1..streams) |s| { + if (!std.mem.eql(u16, wide_u16[0..table.dims], wide_u16[s * table.dims ..][0..table.dims])) copies_differ = true; + } + std.debug.print("embed copies {d} identical {}\n", .{ streams, !copies_differ }); + + const nc = table.dims / 256; + var pss = try cuda.DeviceBuffer.alloc(&driver, nc * streams * 4); + defer pss.free(); + const tri: flash.Tri = .{ .set = &set, .s = stream }; + try tri.hcWriteback(wide.ptr, wide.ptr, pss.ptr, wide.ptr, "*bf16", wide.ptr, wide.ptr, "*bf16", wide.ptr, "*bf16", table.dims, 1, table.dims, streams, 0, 1, 1, 1); + const pss_bytes = try gpa.alloc(u8, nc * streams * 4); + defer gpa.free(pss_bytes); + try stream.synchronize(); + try pss.download(0, pss_bytes); + const want = try gpa.alloc(f32, nc * streams); + defer gpa.free(want); + try flash.hc.partialSums(wide_u16, table.dims, streams, want); + var max_ulp: u32 = 0; + var streams_differ = false; + for (0..nc) |c| { + const first = std.mem.readInt(u32, pss_bytes[4 * c * streams ..][0..4], .little); + for (0..streams) |s| { + const bits = std.mem.readInt(u32, pss_bytes[4 * (c * streams + s) ..][0..4], .little); + if (bits != first) streams_differ = true; + const got: f32 = @bitCast(bits); + const dist = flash.hc.ulps(got, want[c * streams + s]); + if (dist > max_ulp) max_ulp = dist; + } + } + std.debug.print("writeback streams identical {} max_ulp {d}\n", .{ !streams_differ, max_ulp }); + if (copies_differ or streams_differ or max_ulp > 4) return 1; + + const norm_mismatch = try runDown(gpa, io, args[1], &driver, &stream, &set, &mapped, wide.ptr, pss.ptr, wide_u16, pss_bytes, table.dims, streams, nc); + if (norm_mismatch != 0) return 1; + const ngram = try flash.ple.table(flash.Tri, gpa, io, args[1], &driver, &stream, tri, token); + const scored = try flash.mix.scores(flash.Tri, gpa, io, args[1], &driver, &stream, tri, wide_u16); + return if (ngram == 0 and scored == 0) 0 else 1; +} + +fn runDown(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, driver: *cuda.Driver, stream: *cuda.Stream, set: *cuda.aot.Set, mapped: *const embed.Mapped, h: u64, pss: u64, wide_u16: []const u16, pss_bytes: []const u8, dims: usize, streams: usize, nc: usize) !usize { + const prefix = "language_model.model.layers.0.attn_hyper_connection."; + var name: [180]u8 = undefined; + const down_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_down.weight", .{prefix}), .u32); + const down_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_down.scales", .{prefix}), .bf16); + const down_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_down.biases", .{prefix}), .bf16); + const inj_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}block_inject_weight.weight", .{prefix}), .u32); + const inj_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}block_inject_weight.scales", .{prefix}), .bf16); + const inj_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}block_inject_weight.biases", .{prefix}), .bf16); + const gamma = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}hc_norm.weight", .{prefix}), .bf16); + const up_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_up.weight", .{prefix}), .u32); + const up_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_up.scales", .{prefix}), .bf16); + const up_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_up.biases", .{prefix}), .bf16); + const n = down_w.dim(0) + inj_w.dim(0); + const k8 = down_w.dim(1); + const kg = k8 / 4; + if (n != 324 or k8 != 1280 or gamma.dim(0) != streams * dims) return error.UnexpectedTensor; + + const stacked_w = try gpa.alloc(u8, down_w.bytes.len + inj_w.bytes.len); + defer gpa.free(stacked_w); + @memcpy(stacked_w[0..down_w.bytes.len], down_w.bytes); + @memcpy(stacked_w[down_w.bytes.len..], inj_w.bytes); + const stacked_s = try gpa.alloc(u8, down_s.bytes.len + inj_s.bytes.len); + defer gpa.free(stacked_s); + @memcpy(stacked_s[0..down_s.bytes.len], down_s.bytes); + @memcpy(stacked_s[down_s.bytes.len..], inj_s.bytes); + const stacked_b = try gpa.alloc(u8, down_b.bytes.len + inj_b.bytes.len); + defer gpa.free(stacked_b); + @memcpy(stacked_b[0..down_b.bytes.len], down_b.bytes); + @memcpy(stacked_b[down_b.bytes.len..], inj_b.bytes); + const tiled = try flash.hc.tileWords(gpa, stacked_w, n, k8); + defer gpa.free(tiled); + const scales_t = try flash.hc.transposeBf16(gpa, stacked_s, n, kg); + defer gpa.free(scales_t); + const biases_t = try flash.hc.transposeBf16(gpa, stacked_b, n, kg); + defer gpa.free(biases_t); + const scale = try gpa.alloc(f32, gamma.dim(0)); + defer gpa.free(scale); + for (scale, 0..) |*o, i| o.* = @bitCast(@as(u32, std.mem.readInt(u16, gamma.bytes[2 * i ..][0..2], .little)) << 16); + + var w_b = try cuda.DeviceBuffer.fromHost(driver, tiled); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, scales_t); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, biases_t); + defer b_b.free(); + var scale_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(scale)); + defer scale_b.free(); + var norm_b = try cuda.DeviceBuffer.alloc(driver, streams * dims * 2); + defer norm_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, n * 2); + defer out_b.free(); + const sk: usize = 32; + var part_b = try cuda.DeviceBuffer.alloc(driver, sk * n * 4); + defer part_b.free(); + const tri: flash.Tri = .{ .set = set, .s = stream.* }; + try tri.hcDown(h, pss, scale_b.ptr, norm_b.ptr, w_b.ptr, s_b.ptr, b_b.ptr, out_b.ptr, part_b.ptr, 1e-6, 1, n, k8 * 8, dims, nc, streams, sk); + const low: usize = 320; + var act_b = try cuda.DeviceBuffer.alloc(driver, low * 2); + defer act_b.free(); + var xs_b = try cuda.DeviceBuffer.alloc(driver, (low / 32) * 4); + defer xs_b.free(); + var inj_out = try cuda.DeviceBuffer.alloc(driver, streams * 2); + defer inj_out.free(); + try tri.hcReduceAct(part_b.ptr, act_b.ptr, xs_b.ptr, inj_out.ptr, sk, 1, streams, low, n, 1); + const up_n = up_w.dim(0); + const up_k8 = up_w.dim(1); + if (up_n != streams * dims or up_k8 != 40) return error.UnexpectedTensor; + const up_tiled = try flash.hc.tileWords(gpa, up_w.bytes, up_n, up_k8); + defer gpa.free(up_tiled); + const up_scales = try flash.hc.transposeBf16(gpa, up_s.bytes, up_n, up_k8 / 4); + defer gpa.free(up_scales); + const up_biases = try flash.hc.transposeBf16(gpa, up_b.bytes, up_n, up_k8 / 4); + defer gpa.free(up_biases); + var up_wb = try cuda.DeviceBuffer.fromHost(driver, up_tiled); + defer up_wb.free(); + var up_sb = try cuda.DeviceBuffer.fromHost(driver, up_scales); + defer up_sb.free(); + var up_bb = try cuda.DeviceBuffer.fromHost(driver, up_biases); + defer up_bb.free(); + var mixed_b = try cuda.DeviceBuffer.alloc(driver, dims * 2); + defer mixed_b.free(); + var xsm_b = try cuda.DeviceBuffer.alloc(driver, (dims / 32) * 4); + defer xsm_b.free(); + try tri.hcUpmix(act_b.ptr, xs_b.ptr, up_wb.ptr, up_sb.ptr, up_bb.ptr, norm_b.ptr, mixed_b.ptr, xsm_b.ptr, 1, up_n, up_k8 * 8, dims, streams); + + const got = try gpa.alloc(u8, streams * dims * 2); + defer gpa.free(got); + try stream.synchronize(); + try norm_b.download(0, got); + const pss_f = try gpa.alloc(f32, nc * streams); + defer gpa.free(pss_f); + for (pss_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, pss_bytes[4 * i ..][0..4], .little)); + const want = try gpa.alloc(u16, streams * dims); + defer gpa.free(want); + try flash.hc.normed(wide_u16, pss_f, scale, dims, streams, 1e-6, want); + var mismatch: usize = 0; + for (want, 0..) |w, i| { + if (std.mem.readInt(u16, got[2 * i ..][0..2], .little) != w) mismatch += 1; + } + var part4: [4]u8 = undefined; + try part_b.download(0, &part4); + const gpu_part: f32 = @bitCast(std.mem.readInt(u32, &part4, .little)); + const gpu_norm = try gpa.alloc(u16, 320); + defer gpa.free(gpu_norm); + for (gpu_norm, 0..) |*o, i| o.* = std.mem.readInt(u16, got[2 * i ..][0..2], .little); + const host_part = try flash.hc.firstPartial(gpu_norm, stacked_w[0 .. k8 * 4], stacked_s[0 .. kg * 2], stacked_b[0 .. kg * 2], 10); + const col_ulp = flash.hc.ulps(gpu_part, host_part); + std.debug.print("hcdown normed mismatch {d} col0 gpu {d} host {d} ulp {d}\n", .{ mismatch, gpu_part, host_part, col_ulp }); + if (col_ulp > 8) return 1; + + const part_bytes = try gpa.alloc(u8, sk * n * 4); + defer gpa.free(part_bytes); + try part_b.download(0, part_bytes); + const part_f = try gpa.alloc(f32, sk * n); + defer gpa.free(part_f); + for (part_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, part_bytes[4 * i ..][0..4], .little)); + const act_want = try gpa.alloc(u16, low); + defer gpa.free(act_want); + const inj_want = try gpa.alloc(u16, streams); + defer gpa.free(inj_want); + const xs_want = try gpa.alloc(f32, low / 32); + defer gpa.free(xs_want); + try flash.hc.reduceAct(part_f, sk, n, low, streams, act_want, inj_want, xs_want); + const act_bytes = try gpa.alloc(u8, low * 2); + defer gpa.free(act_bytes); + const inj_bytes = try gpa.alloc(u8, streams * 2); + defer gpa.free(inj_bytes); + const xs_bytes = try gpa.alloc(u8, (low / 32) * 4); + defer gpa.free(xs_bytes); + try act_b.download(0, act_bytes); + try inj_out.download(0, inj_bytes); + try xs_b.download(0, xs_bytes); + var act_mismatch: usize = 0; + for (act_want, 0..) |w, i| { + if (std.mem.readInt(u16, act_bytes[2 * i ..][0..2], .little) != w) act_mismatch += 1; + } + var inj_mismatch: usize = 0; + for (inj_want, 0..) |w, i| { + if (std.mem.readInt(u16, inj_bytes[2 * i ..][0..2], .little) != w) inj_mismatch += 1; + } + var xs_ulp: u32 = 0; + for (xs_want, 0..) |w, i| { + const got_xs: f32 = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const dist = flash.hc.ulps(got_xs, w); + if (dist > xs_ulp) xs_ulp = dist; + } + std.debug.print("reduce act mismatch {d} inject mismatch {d} xs max_ulp {d}\n", .{ act_mismatch, inj_mismatch, xs_ulp }); + if (act_mismatch != 0 or inj_mismatch != 0 or xs_ulp > 4) return 1; + + const mixed_bytes = try gpa.alloc(u8, dims * 2); + defer gpa.free(mixed_bytes); + try mixed_b.download(0, mixed_bytes); + const act_u16 = try gpa.alloc(u16, low); + defer gpa.free(act_u16); + for (act_u16, 0..) |*o, i| o.* = std.mem.readInt(u16, act_bytes[2 * i ..][0..2], .little); + const xs_f = try gpa.alloc(f32, low / 32); + defer gpa.free(xs_f); + for (xs_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const norm_u16 = try gpa.alloc(u16, streams * dims); + defer gpa.free(norm_u16); + for (norm_u16, 0..) |*o, i| o.* = std.mem.readInt(u16, got[2 * i ..][0..2], .little); + const host_mixed = try gpa.alloc(u16, dims); + defer gpa.free(host_mixed); + try flash.hc.mix(act_u16, xs_f, norm_u16, up_w.bytes, up_s.bytes, up_b.bytes, up_n, up_k8, dims, streams, host_mixed); + var mixed_off: usize = 0; + var max_steps: u32 = 0; + for (host_mixed, 0..) |host_m, i| { + const got_m = std.mem.readInt(u16, mixed_bytes[2 * i ..][0..2], .little); + const dist = flash.hc.mixedSteps(got_m, host_m); + if (dist != 0) mixed_off += 1; + if (dist > max_steps) max_steps = dist; + } + std.debug.print("upmix dims {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ dims, mixed_off, max_steps, std.mem.readInt(u16, mixed_bytes[0..2], .little), host_mixed[0] }); + if (max_steps > 2) return 1; + const rest = try runProj(gpa, io, model_dir, driver, stream, mapped, mixed_b, xsm_b, tri, h, inj_out.ptr, wide_u16); + if (rest != 0) return 1; + return mismatch; +} + +const proj_names = [_][]const u8{ "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a" }; +const proj_rows = [_]usize{ 10240, 6144, 48, 48 }; + +fn runProj(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, driver: *cuda.Driver, stream: *cuda.Stream, mapped: *const embed.Mapped, mixed: cuda.DeviceBuffer, xs: cuda.DeviceBuffer, tri: flash.Tri, h: u64, inj: u64, hidden: []const u16) !u32 { + const prefix = "language_model.model.layers.0.linear_attn."; + const k8: usize = 320; + const kg: usize = 80; + var name: [160]u8 = undefined; + var n: usize = 0; + var words_len: usize = 0; + var scale_len: usize = 0; + var ws: [4]core.safetensors.Tensor = undefined; + var ss: [4]core.safetensors.Tensor = undefined; + var bs: [4]core.safetensors.Tensor = undefined; + for (proj_names, proj_rows, 0..) |part, rows, i| { + ws[i] = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}{s}.weight", .{ prefix, part }), .u32); + ss[i] = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}{s}.scales", .{ prefix, part }), .bf16); + bs[i] = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}{s}.biases", .{ prefix, part }), .bf16); + if (ws[i].dim(0) != rows or ws[i].dim(1) != k8 or ss[i].dim(0) != rows or ss[i].dim(1) != kg) return error.UnexpectedTensor; + n += rows; + words_len += ws[i].bytes.len; + scale_len += ss[i].bytes.len; + } + if (n != 16480) return error.UnexpectedTensor; + const k = k8 * 8; + const row = try gpa.alloc(u16, k); + defer gpa.free(row); + const words = try gpa.alloc(u8, words_len); + defer gpa.free(words); + const scales = try gpa.alloc(u8, scale_len); + defer gpa.free(scales); + const biases = try gpa.alloc(u8, scale_len); + defer gpa.free(biases); + var wo: usize = 0; + var so: usize = 0; + for (ws, ss, bs) |w, s, b| { + @memcpy(words[wo..][0..w.bytes.len], w.bytes); + wo += w.bytes.len; + @memcpy(scales[so..][0..s.bytes.len], s.bytes); + @memcpy(biases[so..][0..b.bytes.len], b.bytes); + so += s.bytes.len; + } + var lane = try flash.qmm.pack(gpa, words, scales, biases, n, k8); + defer lane.deinit(gpa); + var w_b = try cuda.DeviceBuffer.fromHost(driver, lane.weight); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, lane.scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, lane.biases); + defer b_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, n * 2); + defer out_b.free(); + try flash.qmm.matmul(driver, stream.*, mixed.ptr, xs.ptr, w_b.ptr, s_b.ptr, b_b.ptr, out_b.ptr, 1, n, k); + const row_bytes = try gpa.alloc(u8, k * 2); + defer gpa.free(row_bytes); + const xs_bytes = try gpa.alloc(u8, kg * 4); + defer gpa.free(xs_bytes); + const got_bytes = try gpa.alloc(u8, n * 2); + defer gpa.free(got_bytes); + try stream.synchronize(); + try mixed.download(0, row_bytes); + try xs.download(0, xs_bytes); + try out_b.download(0, got_bytes); + for (row, 0..) |*o, i| o.* = std.mem.readInt(u16, row_bytes[2 * i ..][0..2], .little); + const xs_f = try gpa.alloc(f32, kg); + defer gpa.free(xs_f); + for (xs_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const host = try gpa.alloc(u16, n); + defer gpa.free(host); + try flash.qmm.dotRow(row, xs_f, words, scales, biases, n, k, host); + var off: usize = 0; + var max_steps: u32 = 0; + for (host, 0..) |want, i| { + const got = std.mem.readInt(u16, got_bytes[2 * i ..][0..2], .little); + const dist = flash.hc.mixedSteps(got, want); + if (dist != 0) off += 1; + if (dist > max_steps) max_steps = dist; + } + std.debug.print("proj n {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ n, off, max_steps, std.mem.readInt(u16, got_bytes[0..2], .little), host[0] }); + if (max_steps != 0) return max_steps; + const slices = try flash.gdn.memory(gpa, driver, stream, mapped, out_b, got_bytes) orelse return 1; + defer gpa.free(slices); + const wrote = try flash.mlp.connect(flash.Tri, gpa, driver, stream, mapped, tri, h, inj, hidden, slices); + if (wrote != 0) return wrote; + const routed = try flash.moe.experts(flash.Tri, gpa, driver, stream, mapped, tri, mixed.ptr); + if (routed != 0) return routed; + return flash.attn.decode(flash.Tri, gpa, driver, stream, io, model_dir, tri, mixed.ptr, xs.ptr); +} diff --git a/zig/src/families/flashnext/cuda_gdn.zig b/zig/src/families/flashnext/cuda_gdn.zig new file mode 100644 index 0000000000..b53fb5d50d --- /dev/null +++ b/zig/src/families/flashnext/cuda_gdn.zig @@ -0,0 +1,409 @@ +//! One-row Gated DeltaNet: depthwise conv, the delta-rule state, the gated norm, then the output projection. + +const std = @import("std"); +const cuda = @import("cuda"); +const embed = @import("cuda_embed.zig"); +const hc = @import("cuda_hc.zig"); +const qmm = @import("cuda_qmm.zig"); + +pub const nk: usize = 16; +pub const nv: usize = 48; +pub const dk: usize = 128; +pub const dv: usize = 128; +pub const taps: usize = 4; +pub const conv_dim: usize = 2 * nk * dk + nv * dv; +pub const proj_width: usize = conv_dim + nv * dv + 2 * nv; +pub const value_dim: usize = nv * dv; +pub const groups: usize = value_dim / qmm.group; + +const symbol: [:0]const u8 = "flashnext_gdn_chain"; + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +fn bf(v: f32) f32 { + return promote(toBf16(v)); +} + +fn sigmoid(x: f32) f32 { + return 1.0 / (1.0 + @exp(-x)); +} + +fn softplus(x: f32) f32 { + return if (x > 20) x else @log(1.0 + @exp(x)); +} + +fn warpSum(v: *[32]f32) void { + var span: usize = 16; + while (span != 0) : (span >>= 1) { + var next: [32]f32 = undefined; + for (0..32) |lane| next[lane] = v[lane] + v[lane ^ span]; + v.* = next; + } +} + +fn convAct(proj: []const u16, cs: []const u16, cw: []const u16, row: usize, channel: usize) f32 { + var acc: f32 = 0; + for (0..taps) |tap| { + const at = row + tap; + const x: f32 = if (at < taps - 1) promote(cs[at * conv_dim + channel]) else promote(proj[(at - (taps - 1)) * proj_width + channel]); + acc = acc + promote(cw[channel * taps + tap]) * x; + } + return bf(acc / (1.0 + @exp(-acc))); +} + +fn normPair(x: []f32, scale_dk: bool) void { + var part: [32]f32 = undefined; + for (0..32) |lane| { + var ss: f32 = 0; + for (0..4) |i| { + const v = x[lane * 4 + i]; + ss = ss + v * v; + } + part[lane] = ss; + } + warpSum(&part); + var inv = 1.0 / @sqrt(part[0] + 1e-6); + if (scale_dk) inv = inv * (1.0 / @sqrt(@as(f32, @floatFromInt(dk)))); + for (0..dk) |i| x[i] = x[i] * inv; +} + +fn update(state: []f32, base: usize, ks: []const f32, vs: []const f32, g: f32, beta: f32) void { + for (0..32) |warp| { + var kk: [32][4]f32 = undefined; + for (0..32) |lane| { + for (0..4) |i| kk[lane][i] = ks[lane * 4 + i]; + } + for (0..4) |j| { + var kv: [32]f32 = undefined; + const row = warp * 4 + j; + for (0..32) |lane| { + var acc: f32 = 0; + for (0..4) |i| { + const at = base + row * dk + lane * 4 + i; + state[at] = state[at] * g; + acc = acc + state[at] * kk[lane][i]; + } + kv[lane] = acc; + } + warpSum(&kv); + const delta = (vs[row] - kv[0]) * beta; + for (0..32) |lane| { + for (0..4) |i| { + const at = base + row * dk + lane * 4 + i; + state[at] = state[at] + kk[lane][i] * delta; + } + } + } + } +} + +fn readout(state: []const f32, base: usize, qs: []const f32, ys: []f32) void { + for (0..32) |warp| { + for (0..4) |j| { + var part: [32]f32 = undefined; + const row = warp * 4 + j; + for (0..32) |lane| { + var acc: f32 = 0; + for (0..4) |i| acc = acc + state[base + row * dk + lane * 4 + i] * qs[lane * 4 + i]; + part[lane] = acc; + } + warpSum(&part); + ys[row] = bf(part[0]); + } + } +} + +/// One fresh or continued row. `state` is read and written in place, layout [value head, row, column]. +pub fn chain(proj: []const u16, cs: []const u16, cw: []const u16, state: []f32, a_log: []const f32, dt_bias: []const f32, norm_w: []const u16, eps: f32, rows: usize, out: []u16, xs: []f32) !void { + if (rows == 0 or proj.len < rows * proj_width or cs.len < (taps - 1) * conv_dim or cw.len < conv_dim * taps) return error.UnexpectedTensor; + if (state.len < nv * dv * dk or a_log.len < nv or dt_bias.len < nv or norm_w.len < dv) return error.UnexpectedTensor; + if (out.len < rows * value_dim or xs.len < rows * groups) return error.UnexpectedTensor; + for (0..rows) |r| { + for (0..nv) |hv| { + const hk = hv / (nv / nk); + var qs: [dk]f32 = undefined; + var ks: [dk]f32 = undefined; + var vs: [dv]f32 = undefined; + for (0..dk) |t| qs[t] = convAct(proj, cs, cw, r, hk * dk + t); + for (0..dk) |t| ks[t] = convAct(proj, cs, cw, r, nk * dk + hk * dk + t); + for (0..dv) |t| vs[t] = convAct(proj, cs, cw, r, 2 * nk * dk + hv * dv + t); + normPair(&qs, true); + normPair(&ks, false); + const b = promote(proj[r * proj_width + conv_dim + nv * dv + hv]); + const a = promote(proj[r * proj_width + conv_dim + nv * dv + nv + hv]); + const g = @exp(-@exp(a_log[hv]) * softplus(a + dt_bias[hv])); + const beta = bf(sigmoid(b)); + const base = hv * dv * dk; + update(state, base, &ks, &vs, g, beta); + var ys: [dv]f32 = undefined; + readout(state, base, &qs, &ys); + var part: [32]f32 = undefined; + for (0..32) |lane| { + var ss: f32 = 0; + for (0..4) |i| { + const y = ys[lane * 4 + i]; + ss = ss + y * y; + } + part[lane] = ss; + } + warpSum(&part); + const rinv = 1.0 / @sqrt(part[0] / @as(f32, @floatFromInt(dv)) + eps); + var gated: [dv]f32 = undefined; + for (0..dv) |t| { + const yn = bf(bf(ys[t] * rinv) * promote(norm_w[t])); + const z = promote(proj[r * proj_width + conv_dim + hv * dv + t]); + gated[t] = bf(yn * sigmoid(z)); + out[r * value_dim + hv * dv + t] = toBf16(gated[t]); + } + for (0..dv / 32) |warp| { + for (0..32) |lane| part[lane] = gated[warp * 32 + lane]; + warpSum(&part); + xs[r * groups + hv * (dv / 32) + warp] = part[0]; + } + } + } +} + +fn orderedUlps(a: f32, b: f32) u32 { + const ai: i32 = @bitCast(a); + const bi: i32 = @bitCast(b); + const ao: i32 = if (ai < 0) std.math.minInt(i32) - ai else ai; + const bo: i32 = if (bi < 0) std.math.minInt(i32) - bi else bi; + const d = ao - bo; + return @intCast(if (d < 0) -d else d); +} + +/// The one-row kernel. Conv state is [3, conv channels], recurrence state is [heads, 128, 128]. +pub fn launch(d: *const cuda.Driver, stream: cuda.Stream, proj: u64, cs: u64, cw: u64, state_in: u64, a_log: u64, dt_bias: u64, norm_w: u64, eps: f32, out: u64, xs: u64, state_out: u64) !void { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var module = try cuda.Module.load(d, cuda.kernels.flashnext_gdn); + defer module.unload(); + const f = try module.function(symbol); + var args: cuda.Args = .{}; + args.add(proj); + args.add(cs); + args.add(cw); + args.add(state_in); + args.add(a_log); + args.add(dt_bias); + args.add(norm_w); + args.add(eps); + args.add(@as(c_int, 1)); + args.add(out); + args.add(xs); + args.add(state_out); + try cuda.launch.launch(f, .{ .grid = .{ .x = nv }, .block = .{ .x = 1024 } }, stream, &args); +} + +fn f32FromBf16(bytes: []const u8) f32 { + return promote(std.mem.readInt(u16, bytes[0..2], .little)); +} + +/// Layer 0's conv, recurrence and output projection. A match also returns the output face's eight unreduced K slices. +pub fn memory(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, mapped: *const embed.Mapped, proj_dev: cuda.DeviceBuffer, proj_bytes: []const u8) !?[]f32 { + const prefix = "language_model.model.layers.0.linear_attn."; + var name: [160]u8 = undefined; + const conv = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}conv1d.weight", .{prefix}), .bf16); + const a_bits = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}A_log", .{prefix}), .bf16); + const dt_bits = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}dt_bias", .{prefix}), .bf16); + const norm = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}norm.weight", .{prefix}), .bf16); + const out_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}out_proj.weight", .{prefix}), .u32); + const out_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}out_proj.scales", .{prefix}), .bf16); + const out_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}out_proj.biases", .{prefix}), .bf16); + if (conv.dim(0) != conv_dim or conv.dim(1) != taps or norm.dim(0) != dv) return error.UnexpectedTensor; + if (a_bits.dim(0) != nv or dt_bits.dim(0) != nv) return error.UnexpectedTensor; + if (out_w.dim(0) != 2560 or out_w.dim(1) != value_dim / 8) return error.UnexpectedTensor; + if (proj_bytes.len != proj_width * 2) return error.UnexpectedTensor; + + const proj = try gpa.alloc(u16, proj_width); + defer gpa.free(proj); + for (proj, 0..) |*o, i| o.* = std.mem.readInt(u16, proj_bytes[2 * i ..][0..2], .little); + const cs_bytes = try gpa.alloc(u8, (taps - 1) * conv_dim * 2); + defer gpa.free(cs_bytes); + @memset(cs_bytes, 0); + const a_log = try gpa.alloc(f32, nv); + defer gpa.free(a_log); + const dt_bias = try gpa.alloc(f32, nv); + defer gpa.free(dt_bias); + for (0..nv) |i| { + a_log[i] = f32FromBf16(a_bits.bytes[2 * i ..]); + dt_bias[i] = f32FromBf16(dt_bits.bytes[2 * i ..]); + } + const state_n = nv * dv * dk; + const state = try gpa.alloc(f32, state_n); + defer gpa.free(state); + @memset(state, 0); + var st_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(state)); + defer st_b.free(); + const cs = try gpa.alloc(u16, (taps - 1) * conv_dim); + defer gpa.free(cs); + @memset(cs, 0); + const cw = try gpa.alloc(u16, conv_dim * taps); + defer gpa.free(cw); + for (cw, 0..) |*o, i| o.* = std.mem.readInt(u16, conv.bytes[2 * i ..][0..2], .little); + const norm_w = try gpa.alloc(u16, dv); + defer gpa.free(norm_w); + for (norm_w, 0..) |*o, i| o.* = std.mem.readInt(u16, norm.bytes[2 * i ..][0..2], .little); + const host_out = try gpa.alloc(u16, value_dim); + defer gpa.free(host_out); + const host_xs = try gpa.alloc(f32, groups); + defer gpa.free(host_xs); + const eps: f32 = 1e-6; + try chain(proj, cs, cw, state, a_log, dt_bias, norm_w, eps, 1, host_out, host_xs); + + var cs_b = try cuda.DeviceBuffer.fromHost(driver, cs_bytes); + defer cs_b.free(); + var cw_b = try cuda.DeviceBuffer.fromHost(driver, conv.bytes); + defer cw_b.free(); + var a_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(a_log)); + defer a_b.free(); + var dt_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(dt_bias)); + defer dt_b.free(); + var norm_b = try cuda.DeviceBuffer.fromHost(driver, norm.bytes); + defer norm_b.free(); + var st_out = try cuda.DeviceBuffer.alloc(driver, state_n * 4); + defer st_out.free(); + var gout = try cuda.DeviceBuffer.alloc(driver, value_dim * 2); + defer gout.free(); + var gxs = try cuda.DeviceBuffer.alloc(driver, groups * 4); + defer gxs.free(); + try launch(driver, stream.*, proj_dev.ptr, cs_b.ptr, cw_b.ptr, st_b.ptr, a_b.ptr, dt_b.ptr, norm_b.ptr, eps, gout.ptr, gxs.ptr, st_out.ptr); + + const got_bytes = try gpa.alloc(u8, value_dim * 2); + defer gpa.free(got_bytes); + const xs_bytes = try gpa.alloc(u8, groups * 4); + defer gpa.free(xs_bytes); + const st_bytes = try gpa.alloc(u8, state_n * 4); + defer gpa.free(st_bytes); + try stream.synchronize(); + try gout.download(0, got_bytes); + try gxs.download(0, xs_bytes); + try st_out.download(0, st_bytes); + + var out_off: usize = 0; + var out_steps: u32 = 0; + for (host_out, 0..) |want, i| { + const got = std.mem.readInt(u16, got_bytes[2 * i ..][0..2], .little); + const dist = hc.mixedSteps(got, want); + if (dist != 0) out_off += 1; + if (dist > out_steps) out_steps = dist; + } + var xs_off: usize = 0; + var xs_ulp: u32 = 0; + for (host_xs, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const dist = orderedUlps(got, want); + if (dist != 0) xs_off += 1; + if (dist > xs_ulp) xs_ulp = dist; + } + var st_off: usize = 0; + var st_ulp: u32 = 0; + for (state, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, st_bytes[4 * i ..][0..4], .little)); + const dist = orderedUlps(got, want); + if (dist != 0) st_off += 1; + if (dist > st_ulp) st_ulp = dist; + } + const got_head = std.mem.readInt(u16, got_bytes[0..2], .little); + std.debug.print("gdn out {d} off {d} max_steps {d} xs_off {d} xs_ulp {d} state_off {d} state_ulp {d} head {x:0>4} host {x:0>4}\n", .{ value_dim, out_off, out_steps, xs_off, xs_ulp, st_off, st_ulp, got_head, host_out[0] }); + + const got_u16 = try gpa.alloc(u16, value_dim); + defer gpa.free(got_u16); + for (got_u16, 0..) |*o, i| o.* = std.mem.readInt(u16, got_bytes[2 * i ..][0..2], .little); + const got_xs = try gpa.alloc(f32, groups); + defer gpa.free(got_xs); + for (got_xs, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const host_proj = try gpa.alloc(u16, 2560); + defer gpa.free(host_proj); + try qmm.dotRow(got_u16, got_xs, out_w.bytes, out_s.bytes, out_b.bytes, 2560, value_dim, host_proj); + var lane = try qmm.pack(gpa, out_w.bytes, out_s.bytes, out_b.bytes, 2560, value_dim / 8); + defer lane.deinit(gpa); + var w_b = try cuda.DeviceBuffer.fromHost(driver, lane.weight); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, lane.scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, lane.biases); + defer b_b.free(); + var branch = try cuda.DeviceBuffer.alloc(driver, 2560 * 2); + defer branch.free(); + try qmm.matmul(driver, stream.*, gout.ptr, gxs.ptr, w_b.ptr, s_b.ptr, b_b.ptr, branch.ptr, 1, 2560, value_dim); + const branch_bytes = try gpa.alloc(u8, 2560 * 2); + defer gpa.free(branch_bytes); + try stream.synchronize(); + try branch.download(0, branch_bytes); + var proj_off: usize = 0; + var proj_steps: u32 = 0; + for (host_proj, 0..) |want, i| { + const got = std.mem.readInt(u16, branch_bytes[2 * i ..][0..2], .little); + const dist = hc.mixedSteps(got, want); + if (dist != 0) proj_off += 1; + if (dist > proj_steps) proj_steps = dist; + } + std.debug.print("gdn out_proj n {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ 2560, proj_off, proj_steps, std.mem.readInt(u16, branch_bytes[0..2], .little), host_proj[0] }); + if (out_off != 0 or xs_off != 0 or st_off != 0 or proj_off != 0) return null; + // Eight K slices, left unreduced for the write-back. + const sk: usize = 8; + const n: usize = 2560; + var part_b = try cuda.DeviceBuffer.alloc(driver, sk * n * 4); + defer part_b.free(); + try qmm.partials(driver, stream.*, gout.ptr, gxs.ptr, w_b.ptr, s_b.ptr, b_b.ptr, branch.ptr, part_b.ptr, 1, n, value_dim, sk); + const raw = try gpa.alloc(u8, sk * n * 4); + defer gpa.free(raw); + try stream.synchronize(); + try part_b.download(0, raw); + const slices = try gpa.alloc(f32, sk * n); + errdefer gpa.free(slices); + for (slices, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, raw[4 * i ..][0..4], .little)); + return slices; +} + +test "warp sum is the sum of the lanes" { + var v: [32]f32 = @splat(0); + v[0] = 1; + v[31] = 3; + warpSum(&v); + try std.testing.expectEqual(@as(f32, 4), v[0]); + try std.testing.expectEqual(@as(f32, 4), v[17]); +} + +test "a zero projection and zero weights stay a zero row" { + const gpa = std.testing.allocator; + const proj = try gpa.alloc(u16, proj_width); + defer gpa.free(proj); + @memset(proj, 0); + const cs = try gpa.alloc(u16, (taps - 1) * conv_dim); + defer gpa.free(cs); + @memset(cs, 0); + const cw = try gpa.alloc(u16, conv_dim * taps); + defer gpa.free(cw); + @memset(cw, 0); + const state = try gpa.alloc(f32, nv * dv * dk); + defer gpa.free(state); + @memset(state, 0); + const a_log = try gpa.alloc(f32, nv); + defer gpa.free(a_log); + @memset(a_log, 0); + const dt = try gpa.alloc(f32, nv); + defer gpa.free(dt); + @memset(dt, 0); + const norm = try gpa.alloc(u16, dv); + defer gpa.free(norm); + @memset(norm, 0); + const out = try gpa.alloc(u16, value_dim); + defer gpa.free(out); + const xs = try gpa.alloc(f32, groups); + defer gpa.free(xs); + try chain(proj, cs, cw, state, a_log, dt, norm, 1e-6, 1, out, xs); + try std.testing.expectEqual(@as(u16, 0), out[0]); + try std.testing.expectEqual(@as(u16, 0), out[value_dim - 1]); + try std.testing.expectEqual(@as(f32, 0), xs[0]); + try std.testing.expectEqual(@as(f32, 0), state[0]); +} diff --git a/zig/src/families/flashnext/cuda_hc.zig b/zig/src/families/flashnext/cuda_hc.zig new file mode 100644 index 0000000000..65b661904b --- /dev/null +++ b/zig/src/families/flashnext/cuda_hc.zig @@ -0,0 +1,307 @@ +//! Layer-0 hyper-connection host checks: mode 0 writes each 256-wide chunk's sum of squares. + +const std = @import("std"); + +/// glue._hc_writeback mode 0: pss[c, s] is a pairwise sum of squares of stream s, chunk c. +pub fn partialSums(h: []const u16, dims: usize, streams: usize, out: []f32) !void { + const block: usize = 256; + if (dims == 0 or dims % block != 0 or h.len != streams * dims) return error.UnexpectedTensor; + const nc = dims / block; + if (out.len != nc * streams) return error.UnexpectedTensor; + for (0..nc) |c| { + for (0..streams) |s| { + var buf: [256]f32 = undefined; + for (0..block) |i| { + const bits: u32 = h[s * dims + c * block + i]; + const v: f32 = @bitCast(bits << 16); + buf[i] = v * v; + } + var n: usize = block; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + out[c * streams + s] = buf[0]; + } + } +} + +test "mode 0 partial sums are the squares of one 256-wide chunk" { + var h: [256]u16 = @splat(0); + h[0] = 0x3f80; + var out: [1]f32 = undefined; + try partialSums(&h, 256, 1, &out); + try std.testing.expectEqual(@as(f32, 1), out[0]); +} + +/// Steps between two positive finite floats. Triton's block sum lands a couple of steps from a pairwise sum. +pub fn ulps(a: f32, b: f32) u32 { + const xa: i32 = @bitCast(a); + const xb: i32 = @bitCast(b); + const d = xa - xb; + return @intCast(if (d < 0) -d else d); +} + +test "ulps counts steps between positive floats" { + const a: f32 = 1; + const b: f32 = @bitCast(@as(u32, @bitCast(a)) + 2); + try std.testing.expectEqual(@as(u32, 2), ulps(a, b)); +} + +pub const tile_bn: usize = 64; + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +/// MLX (n, k8) words to [n/64][k/32][64][4], n padded up to 64. qmm.tile_words. +pub fn tileWords(gpa: std.mem.Allocator, words: []const u8, n: usize, k8: usize) ![]u8 { + if (k8 == 0 or k8 % 4 != 0 or words.len != n * k8 * 4) return error.UnexpectedTensor; + const npad = (n + tile_bn - 1) / tile_bn * tile_bn; + const kg = k8 / 4; + const out = try gpa.alloc(u8, (npad / tile_bn) * kg * tile_bn * 4 * 4); + @memset(out, 0); + for (0..n) |r| { + const tile = r / tile_bn; + const local = r % tile_bn; + for (0..kg) |g| { + for (0..4) |w| { + const src = ((r * k8) + g * 4 + w) * 4; + const dst = (((tile * kg + g) * tile_bn + local) * 4 + w) * 4; + @memcpy(out[dst..][0..4], words[src..][0..4]); + } + } + } + return out; +} + +/// (rows, cols) bf16 to (cols, rows), the tiled Q4 scale layout. +pub fn transposeBf16(gpa: std.mem.Allocator, src: []const u8, rows: usize, cols: usize) ![]u8 { + if (src.len != rows * cols * 2) return error.UnexpectedTensor; + const out = try gpa.alloc(u8, src.len); + for (0..rows) |r| { + for (0..cols) |c| { + const from = (r * cols + c) * 2; + const to = (c * rows + r) * 2; + @memcpy(out[to..][0..2], src[from..][0..2]); + } + } + return out; +} + +/// The normed row _qmm_hcdown writes: bf16(h * rinv * scale), one rinv per stream from that stream's partial sums. +pub fn normed(h: []const u16, pss: []const f32, scale: []const f32, dims: usize, streams: usize, eps: f32, out: []u16) !void { + const nc = dims / 256; + if (dims == 0 or dims % 256 != 0 or h.len != streams * dims or pss.len != nc * streams or scale.len != h.len or out.len != h.len) return error.UnexpectedTensor; + for (0..streams) |s| { + var total: f32 = 0; + for (0..nc) |c| total += pss[c * streams + s]; + const rinv = 1.0 / @sqrt(total / @as(f32, @floatFromInt(dims)) + eps); + for (0..dims) |i| { + const hv: f32 = @bitCast(@as(u32, h[s * dims + i]) << 16); + out[s * dims + i] = toBf16(hv * rinv * scale[s * dims + i]); + } + } +} + +/// fp32 dot of the first K-slice against output column 0, groups of 32, from the MLX row (not the tiled copy). +pub fn firstPartial(row_normed: []const u16, row_words: []const u8, scale_row: []const u8, bias_row: []const u8, groups: usize) !f32 { + if (row_normed.len < groups * 32 or row_words.len < groups * 16 or scale_row.len < groups * 2 or bias_row.len < groups * 2) return error.UnexpectedTensor; + var acc: f32 = 0; + for (0..groups) |g| { + var buf: [32]f32 = undefined; + var p: f32 = 0; + for (0..4) |i| { + const word = std.mem.readInt(u32, row_words[(g * 4 + i) * 4 ..][0..4], .little); + for (0..8) |j| { + const q: f32 = @floatFromInt((word >> @intCast(j * 4)) & 0xF); + const xv: f32 = @bitCast(@as(u32, row_normed[g * 32 + i * 8 + j]) << 16); + buf[i * 8 + j] = xv; + p += xv * q; + } + } + var n: usize = 32; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + const s: f32 = @bitCast(@as(u32, std.mem.readInt(u16, scale_row[g * 2 ..][0..2], .little)) << 16); + const b: f32 = @bitCast(@as(u32, std.mem.readInt(u16, bias_row[g * 2 ..][0..2], .little)) << 16); + acc += p * s + buf[0] * b; + } + return acc; +} + +test "tiled words keep row 0 in the first tile's local row 0" { + var words: [16]u8 = @splat(0); + std.mem.writeInt(u32, words[0..4], 0x21, .little); + const out = try tileWords(std.testing.allocator, &words, 1, 4); + defer std.testing.allocator.free(out); + try std.testing.expectEqual(@as(u32, 0x21), std.mem.readInt(u32, out[0..4], .little)); + try std.testing.expectEqual(@as(u32, 0), std.mem.readInt(u32, out[16..20], .little)); +} + +test "the down norm writes bf16 of h times rinv times scale" { + var h: [256]u16 = @splat(0); + h[0] = 0x3f80; + const pss = [_]f32{1}; + var scale: [256]f32 = @splat(1); + var out: [256]u16 = undefined; + try normed(&h, &pss, &scale, 256, 1, 1e-6, &out); + const rinv = 1.0 / @sqrt(1.0 / 256.0 + 1e-6); + try std.testing.expectEqual(toBf16(rinv), out[0]); + try std.testing.expectEqual(@as(u16, 0), out[1]); +} + +fn roundBf16(v: f32) f32 { + return @bitCast(@as(u32, toBf16(v)) << 16); +} + +/// glue._hc_reduce_act: sum the K slices in order, round, divide by the stream count, SiLU, and the inject gates. +pub fn reduceAct(part: []const f32, sk: usize, ndn: usize, low: usize, streams: usize, act: []u16, inj: []u16, xs: []f32) !void { + if (low == 0 or low % 32 != 0 or ndn < low + streams or part.len != sk * ndn or act.len != low or inj.len != streams or xs.len != low / 32) return error.UnexpectedTensor; + for (0..low) |i| { + var sum: f32 = 0; + for (0..sk) |s| sum += part[s * ndn + i]; + const y = roundBf16(roundBf16(sum) / @as(f32, @floatFromInt(streams))); + act[i] = toBf16(y / (1.0 + @exp(-y))); + } + var buf: [32]f32 = undefined; + for (0..low / 32) |g| { + for (0..32) |j| buf[j] = @bitCast(@as(u32, act[g * 32 + j]) << 16); + var n: usize = 32; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + xs[g] = buf[0]; + } + for (0..streams) |s| { + var sum: f32 = 0; + for (0..sk) |k| sum += part[k * ndn + low + s]; + const y = roundBf16(roundBf16(sum) / @as(f32, @floatFromInt(streams))); + const sig = roundBf16(1.0 / (1.0 + @exp(-y))); + inj[s] = toBf16(2.0 * sig); + } +} + +test "a zero down-projection sum is SiLU zero and an inject gate of one" { + var part: [33]f32 = @splat(0); + var act: [32]u16 = undefined; + var inj: [1]u16 = undefined; + var xs: [1]f32 = undefined; + try reduceAct(&part, 1, 33, 32, 1, &act, &inj, &xs); + try std.testing.expectEqual(@as(u16, 0), act[0]); + try std.testing.expectEqual(@as(u16, 0x3f80), inj[0]); + try std.testing.expectEqual(@as(f32, 0), xs[0]); +} + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +/// One mixed hidden row: up-dot each stream, sigmoid-gate it by that stream's normed value, then average. +pub fn mix(act: []const u16, xs: []const f32, normed_row: []const u16, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k8: usize, dims: usize, streams: usize, out: []u16) !void { + const kg = k8 / 4; + if (dims == 0 or k8 == 0 or k8 % 4 != 0 or act.len < kg * 32 or xs.len < kg or normed_row.len != streams * dims or out.len != dims) return error.UnexpectedTensor; + if (words.len != n * k8 * 4 or scales.len != n * kg * 2 or biases.len != scales.len) return error.UnexpectedTensor; + for (0..dims) |col| { + var total: f32 = 0; + for (0..streams) |s| { + const row = s * dims + col; + var acc: f32 = 0; + const word_row = words[row * k8 * 4 ..][0 .. k8 * 4]; + const scale_row = scales[row * kg * 2 ..][0 .. kg * 2]; + const bias_row = biases[row * kg * 2 ..][0 .. kg * 2]; + for (0..kg) |g| { + var p: f32 = 0; + for (0..4) |i| { + const word = std.mem.readInt(u32, word_row[(g * 4 + i) * 4 ..][0..4], .little); + for (0..8) |j| { + const q: f32 = @floatFromInt((word >> @intCast(j * 4)) & 0xF); + p += promote(act[g * 32 + i * 8 + j]) * q; + } + } + const sc = promote(std.mem.readInt(u16, scale_row[g * 2 ..][0..2], .little)); + const bias = promote(std.mem.readInt(u16, bias_row[g * 2 ..][0..2], .little)); + acc += p * sc + xs[g] * bias; + } + const up = roundBf16(acc); + const sig = roundBf16(1.0 / (1.0 + @exp(-up))); + total += roundBf16(sig * promote(normed_row[row])); + } + out[col] = toBf16(total / @as(f32, @floatFromInt(streams))); + } +} + +fn bf16Steps(a: u16, b: u16) u32 { + const oa: i32 = if (a >= 0x8000) ~@as(i32, a) else a; + const ob: i32 = if (b >= 0x8000) ~@as(i32, b) else b; + const d = oa - ob; + return @intCast(if (d < 0) -d else d); +} + +pub fn mixedSteps(gpu: u16, host: u16) u32 { + return bf16Steps(gpu, host); +} + +/// Mode 4's branch: the K slices added in slice order, then one bf16 round. +pub fn sumSlices(part: []const f32, sk: usize, n: usize, out: []u16) !void { + if (sk == 0 or n == 0 or part.len != sk * n or out.len != n) return error.UnexpectedTensor; + for (0..n) |i| { + var acc = part[i]; + for (1..sk) |k| acc += part[k * n + i]; + out[i] = toBf16(acc); + } +} + +/// Mode 4's hidden update: one bf16 rounding of hidden + branch * inject, the write-back's fused multiply-add. +pub fn applyBranch(h: []const u16, branch: []const u16, inj: []const u16, dims: usize, streams: usize, out: []u16) !void { + if (dims == 0 or h.len != streams * dims or branch.len != dims or inj.len != streams or out.len != h.len) return error.UnexpectedTensor; + for (0..streams) |s| { + const gate = promote(inj[s]); + for (0..dims) |i| { + out[s * dims + i] = toBf16(promote(h[s * dims + i]) + promote(branch[i]) * gate); + } + } +} + +test "slices add in order and round once" { + const part = [_]f32{ 1, 0, 0.5, 0 }; + var out: [2]u16 = undefined; + try sumSlices(&part, 2, 2, &out); + try std.testing.expectEqual(toBf16(1.5), out[0]); + try std.testing.expectEqual(@as(u16, 0), out[1]); +} + +test "a branch of one times an inject gate of one adds one" { + const h = [_]u16{0}; + const branch = [_]u16{0x3f80}; + const inj = [_]u16{0x3f80}; + var out: [1]u16 = undefined; + try applyBranch(&h, &branch, &inj, 1, 1, &out); + try std.testing.expectEqual(@as(u16, 0x3f80), out[0]); +} + +test "the fused product is rounded once with the hidden value" { + const h = [_]u16{0x3b34}; + const branch = [_]u16{0xbe16}; + const inj = [_]u16{0x3c9a}; + var out: [1]u16 = undefined; + try applyBranch(&h, &branch, &inj, 1, 1, &out); + try std.testing.expectEqual(@as(u16, 0xb6f0), out[0]); +} + +test "a zero activation mixes to a zero hidden value" { + var act: [32]u16 = @splat(0); + const xs = [_]f32{0}; + var row: [1]u16 = @splat(0); + var words: [16]u8 = @splat(0); + var scale: [2]u8 = @splat(0); + var bias: [2]u8 = @splat(0); + var mixed: [1]u16 = undefined; + try mix(&act, &xs, &row, &words, &scale, &bias, 1, 4, 1, 1, &mixed); + try std.testing.expectEqual(@as(u16, 0), mixed[0]); +} diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig new file mode 100644 index 0000000000..7502108299 --- /dev/null +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -0,0 +1,215 @@ +//! Flash Next on the native CUDA server: one stream's prompt, then one token at a time. No draft head. + +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const prompt = @import("cuda_prompt.zig"); + +const be = lanes.backend; +const ring = 1024; + +/// `Tri` is the captured-kernel launcher. `open` loads the checkpoint and serves greedy or keyed draws. +pub fn Namespace(comptime Tri: type) type { + return struct { + pub const model_type = "qwen4_exp"; + pub const formats: []const []const u8 = &.{"mlx-q4g32"}; + pub const default_context: i64 = 262144; + pub const prefill_step: u32 = 0; + + pub const Options = struct { context: usize, drafts: bool }; + + pub const Loaded = struct { + backend: be.Backend, + facts: lanes.Model, + rows: u32, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + }; + + const Lane = struct { seq: prompt.Seq }; + + const Cuda = struct { + gpa: std.mem.Allocator, + context: usize, + engine: *prompt.Engine, + set: cuda.aot.Set, + stream: cuda.Stream, + lanes: std.AutoHashMapUnmanaged(*lanes.Stream, Lane) = .empty, + drawn: [ring]u32 = undefined, + next: u64 = 0, + + fn tri(self: *Cuda) Tri { + return .{ .set = &self.set, .s = self.stream }; + } + + fn backend(self: *Cuda) be.Backend { + return .{ .ptr = self, .vtable = &.{ + .prefill = prefillFn, + .first = firstFn, + .queue = queueFn, + .read = readFn, + .verify = verifyFn, + .keep = keepFn, + .draft = draftFn, + .release = releaseFn, + } }; + } + + fn facts() lanes.Model { + return .{ + .exact_width = 1, + .gpu_tokens = false, + .mtp = false, + .speculate = false, + .drafts = 0, + .batch_rows = 1, + .max_streams = 1, + }; + } + + fn take(self: *Cuda, token: u32) u64 { + const h = self.next; + self.drawn[h % ring] = token; + self.next += 1; + return h; + } + + fn tokenAt(self: *Cuda, seq: *prompt.Seq, s: *lanes.Stream, position: u64) !u32 { + if (s.sampling) |sm| { + const n = prompt.vocabulary; + const vals = try self.gpa.alloc(f64, n); + defer self.gpa.free(vals); + const ids = try self.gpa.alloc(u64, n); + defer self.gpa.free(ids); + for (vals, ids, 0..) |*v, *id, i| { + const bits = std.mem.readInt(u16, seq.logits[2 * i ..][0..2], .little); + v.* = @floatCast(@as(f32, @bitCast(@as(u32, bits) << 16))); + id.* = @intCast(i); + } + return @intCast(try lanes.sampling.choose(self.gpa, vals, ids, position, sm)); + } + return prompt.argmax(seq.logits); + } + + fn lane(self: *Cuda, s: *lanes.Stream) !*Lane { + return self.lanes.getPtr(s) orelse return error.UnknownStream; + } + }; + + pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: []const u8, o: Options) !Loaded { + _ = o.drafts; // this model is served one token at a time; there is no draft head + const driver: *cuda.Driver = @constCast(ctx.d); + const own = try gpa.create(Cuda); + errdefer gpa.destroy(own); + own.gpa = gpa; + own.context = o.context; + own.stream = try cuda.Stream.init(driver, true); + errdefer own.stream.deinit(); + own.set = try cuda.aot.Set.load(gpa, io, driver, ctx.device, kernels); + errdefer own.set.deinit(); + own.engine = try prompt.Engine.init(gpa, io, driver, &own.stream, dir); + own.lanes = .empty; + own.next = 0; + return .{ + .backend = own.backend(), + .facts = Cuda.facts(), + .rows = 1, + .ctx = own, + .deinit = release, + }; + } + + fn release(p: *anyopaque) void { + const own: *Cuda = @ptrCast(@alignCast(p)); + var it = own.lanes.valueIterator(); + while (it.next()) |l| l.seq.deinit(own.gpa); + own.lanes.deinit(own.gpa); + own.engine.deinit(); + own.set.deinit(); + own.stream.deinit(); + own.gpa.destroy(own); + } + + fn of(ptr: *anyopaque) *Cuda { + return @ptrCast(@alignCast(ptr)); + } + + fn prefillFn(ptr: *anyopaque, s: *lanes.Stream) anyerror!void { + const self = of(ptr); + const ids = s.prompt(); + if (ids.len == 0 or ids.len + s.max_new > self.context) return error.PromptTooLong; + const gop = try self.lanes.getOrPut(self.gpa, s); + if (gop.found_existing) gop.value_ptr.seq.deinit(self.gpa); + gop.value_ptr.* = .{ .seq = self.engine.newSeq(ids.len + s.max_new) catch |err| { + self.lanes.removeByPtr(gop.key_ptr); + return err; + } }; + self.engine.predict(Tri, self.tri(), &gop.value_ptr.seq, ids) catch |err| { + gop.value_ptr.seq.deinit(self.gpa); + self.lanes.removeByPtr(gop.key_ptr); + return err; + }; + } + + fn firstFn(ptr: *anyopaque, s: *lanes.Stream, position: u64) anyerror!u64 { + const self = of(ptr); + const l = try self.lane(s); + if (position != l.seq.fed) return error.PositionMismatch; + return self.take(try self.tokenAt(&l.seq, s, position)); + } + + fn queueFn(ptr: *anyopaque, s: *lanes.Stream, feed: be.Feed, position: u64) anyerror!u64 { + _ = ptr; + _ = s; + _ = feed; + _ = position; + return error.NotPipelined; + } + + fn readFn(ptr: *anyopaque, handle: u64) anyerror!u32 { + const self = of(ptr); + if (handle >= self.next or self.next - handle > ring) return error.NoSuchToken; + return self.drawn[handle % ring]; + } + + fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + const self = of(ptr); + for (windows, out) |w, *o| { + if (w.held != 0 or w.parents != null) return error.HeldDraftsNotBuilt; + const l = try self.lane(w.stream); + const rows = w.rows(); + var input = w.pending; + for (0..rows) |r| { + if (w.positions[r] != l.seq.fed + 1) return error.PositionMismatch; + var one = [_]u32{input}; + try self.engine.predict(Tri, self.tri(), &l.seq, &one); + o.sampled[r] = try self.tokenAt(&l.seq, w.stream, w.positions[r]); + if (r + 1 < rows) input = w.tokens[r]; + } + if (w.tokens.len > 0) @memcpy(o.drafts[0..w.tokens.len], w.tokens); + } + } + + fn keepFn(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) anyerror!void { + _ = ptr; + for (windows, paths) |w, path| { + if (path.len != w.rows()) return error.CannotRewind; + for (path, 0..) |r, i| if (r != i) return error.CannotRewind; + } + } + + fn draftFn(ptr: *anyopaque, requests: []const be.DraftRequest) anyerror!void { + _ = ptr; + _ = requests; + return error.NoDraftHead; + } + + fn releaseFn(ptr: *anyopaque, s: *lanes.Stream) void { + const self = of(ptr); + if (self.lanes.fetchRemove(s)) |kv| { + var l = kv.value; + l.seq.deinit(self.gpa); + } + } + }; +} diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig new file mode 100644 index 0000000000..75339636e9 --- /dev/null +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -0,0 +1,270 @@ +//! Flash Next's Triton launches from the captured cubins: each grid and constexpr matches the Python wrapper. + +const std = @import("std"); +const cuda = @import("cuda"); +const aot = cuda.aot; + +pub const embed = @import("cuda_embed.zig"); +pub const hc = @import("cuda_hc.zig"); +pub const qmm = @import("cuda_qmm.zig"); +pub const gdn = @import("cuda_gdn.zig"); +pub const mlp = @import("cuda_mlp.zig"); +pub const moe = @import("cuda_moe.zig"); +pub const attn = @import("cuda_attn.zig"); +pub const ple = @import("cuda_ple.zig"); +pub const mix = @import("cuda_mix.zig"); +pub const prompt = @import("cuda_prompt.zig"); +pub const native = @import("cuda_lanes.zig").Namespace(Tri); + +const p = aot.ptr; + +fn int(name: []const u8, v: usize) aot.Arg { + return aot.int(name, @intCast(v)); +} + +fn ci(name: []const u8, v: usize) aot.Const { + return aot.ci(name, @intCast(v)); +} + +fn u(x: usize) u32 { + return @intCast(x); +} + +fn cdiv(a: usize, b: usize) u32 { + return u((a + b - 1) / b); +} + +/// One program per token row and per group of 32 hidden features. +pub fn embedGrid(rows: usize, dims: usize) [3]u32 { + return .{ u(rows), u(dims / 32), 1 }; +} + +/// RMSNorm's BLOCK is the next power of two at or above the group width. +pub fn rmsBlock(group: usize) usize { + return std.math.ceilPowerOfTwo(usize, group) catch unreachable; +} + +pub const Tri = struct { + set: *const aot.Set, + s: cuda.Stream, + + fn run(t: Tri, name: []const u8, grid: [3]u32, args: []const aot.Arg, consts: []const aot.Const) !void { + try t.set.run(t.s, name, grid, args, consts); + } + + /// glue.embed: ids (R,) int32 into (R, copies * dims) bf16, MLX 4-bit group 32. + pub fn embed(t: Tri, ids: u64, weight: u64, scales: u64, biases: u64, out: u64, dims: usize, copies: usize, rows: usize) !void { + try t.run("_embed", embedGrid(rows, dims), &.{ + p("IDS", "*i32", ids), + p("W", "*i32", weight), + p("S", "*bf16", scales), + p("B", "*bf16", biases), + p("OUT", "*bf16", out), + }, &.{ ci("D", dims), ci("S_COPIES", copies) }); + } + + /// glue.attn_gate: one program per row and per query head. + pub fn attnGate(t: Tri, o: u64, proj: u64, out: u64, xs: u64, rows: usize, pair_width: usize, q_heads: usize, head_dim: usize) !void { + try t.run("_attn_gate", .{ u(rows), u(q_heads), 1 }, &.{ + p("O", "*bf16", o), + p("P", "*bf16", proj), + p("OUT", "*bf16", out), + p("XS", "*fp32", xs), + }, &.{ ci("PW", pair_width), ci("NQ", q_heads), ci("HD", head_dim) }); + } + + /// glue.attn_prep for text positions. The grid includes the indexer key. `length` is unused when MODE is 0. + pub fn attnPrep(t: Tri, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, kc: u64, vc: u64, ks: u64, vs: u64, iq: u64, ikc: u64, rope: u64, delta: u64, eps: f32, rows: usize, length: usize) !void { + try t.run("_attn_prep", .{ u(rows), u(24 + 2 + 4 + 1), 1 }, &.{ + p("P", "*bf16", proj), + p("POS0", "*i32", pos), + p("QW", "*fp32", q_scale), + p("KW", "*fp32", k_scale), + p("IW", "*fp32", i_scale), + p("INV", "*fp32", inv), + p("Q", "*bf16", q), + p("KC", "*bf16", kc), + p("VC", "*bf16", vc), + p("KS", "*fp16", ks), + p("VS", "*fp16", vs), + p("IQ", "*bf16", iq), + p("IKC", "*bf16", ikc), + p("ROPE", "*i32", rope), + p("DELTA", "*i32", delta), + int("length", length), + aot.float("eps", eps), + }, &.{ ci("PW", 13952), ci("NQ", 24), ci("NKV", 2), ci("HD", 256), ci("NI", 4), ci("IHD", 128), ci("HALF", 32), ci("BITS", 0), ci("MODE", 0), ci("S1", 11), ci("S2", 10) }); + } + + /// attention._chunks: one decode row, dense, the captured three-chunk scratch, one chunk launched. + pub fn attnChunks(t: Tri, q: u64, kc: u64, vc: u64, ks: u64, vs: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, chunks: usize) !void { + try t.run("_chunks", .{ u(rows), u(kv), u(chunks) }, &.{ + p("Q", "*bf16", q), + p("KC", "*bf16", kc), + p("VC", "*bf16", vc), + p("KSC", "*fp16", ks), + p("VSC", "*fp16", vs), + p("POS0", "*i32", pos), + p("PO", "*fp32", po), + p("PM", "*fp32", pm), + p("PL", "*fp32", pl), + p("IDS", "*i32", ids), + p("NKR", "*i32", nk), + p("SPR", "*i32", sparse), + }, &.{ ci("H", 24), ci("HK", 2), ci("D", 256), ci("G", 12), ci("CH", 512), ci("NCH", 3), aot.cf("SCALE", 0.0625), ci("IDW", 2052), ci("QSA", 0), ci("BITS", 0) }); + } + + /// attention._merge: one program per row and per KV head. + pub fn attnMerge(t: Tri, po: u64, pm: u64, pl: u64, pos: u64, out: u64, nk: u64, sparse: u64, rows: usize, kv: usize) !void { + try t.run("_merge", .{ u(rows), u(kv), 1 }, &.{ + p("PO", "*fp32", po), + p("PM", "*fp32", pm), + p("PL", "*fp32", pl), + p("POS0", "*i32", pos), + p("OUT", "*bf16", out), + p("NKR", "*i32", nk), + p("SPR", "*i32", sparse), + }, &.{ ci("H", 24), ci("HK", 2), ci("D", 256), ci("G", 12), ci("CH", 512), ci("NCH", 3), ci("QSA", 0), ci("BITS", 0) }); + } + + /// glue.hc_reduce_act: split-K sum fused with hc_act. M folds to a constexpr when it is 1. + pub fn hcReduceAct(t: Tri, part: u64, act: u64, xs: u64, inj: u64, sk: usize, rows: usize, streams: usize, low: usize, ndn: usize, has_inj: usize) !void { + try t.run("_hc_reduce_act", .{ u(rows), 1, 1 }, &.{ + p("PART", "*fp32", part), + p("ACT", "*bf16", act), + p("XS", "*fp32", xs), + p("INJ", "*bf16", inj), + int("M", rows), + }, &.{ ci("SK", sk), ci("S", streams), ci("LOW", low), ci("LOWP", rmsBlock(low)), ci("HAS_INJ", has_inj), ci("NDN", ndn) }); + } + + /// qmm.hc_upmix: the up projection and the stream mix. N is the wide row, K is the low rank. M folds when it is 1. + pub fn hcUpmix(t: Tri, act: u64, xs: u64, weight: u64, scales: u64, biases: u64, normed: u64, mixed: u64, xs_mixed: u64, rows: usize, n: usize, k: usize, dims: usize, streams: usize) !void { + try t.run("_qmm_upmix", .{ cdiv(rows, 16), u(dims / 32), 1 }, &.{ + p("X", "*bf16", act), + p("XS", "*fp32", xs), + p("W", "*i32", weight), + p("S", "*bf16", scales), + p("B", "*bf16", biases), + p("NORMED", "*bf16", normed), + p("MIXED", "*bf16", mixed), + p("XSM", "*fp32", xs_mixed), + int("M", rows), + }, &.{ ci("N", n), ci("K", k), ci("D", dims), ci("SS", streams), ci("BM", 16), ci("DB", 32), ci("GPI", 2), ci("SBN", 64) }); + } + + /// qmm.hc_down for one decode row. N=324 is the layer down stacked with inject. M folds when it is 1. + pub fn hcDown(t: Tri, h: u64, pss: u64, scale: u64, normed: u64, weight: u64, scales: u64, biases: u64, out: u64, part: u64, eps: f32, rows: usize, n: usize, k: usize, dims: usize, nc: usize, streams: usize, sk: usize) !void { + try t.run("_qmm_hcdown", .{ cdiv(rows, 16), cdiv(n, 64), u(sk) }, &.{ + p("H", "*bf16", h), + p("PSS", "*fp32", pss), + p("SCALE", "*fp32", scale), + p("NORMED", "*bf16", normed), + p("W", "*i32", weight), + p("S", "*bf16", scales), + p("B", "*bf16", biases), + p("OUT", "*bf16", out), + p("PART", "*fp32", part), + int("M", rows), + aot.float("eps", eps), + }, &.{ ci("N", n), ci("K", k), ci("D", dims), ci("NC", nc), ci("SS", streams), ci("SK", sk), ci("BM", 16), ci("BLOCK_N", 64), ci("GPI", 2), ci("SBN", 64) }); + } + + /// glue.hc_writeback. Pointer types follow the tensor dtypes Triton specialized (mode 0..4). + pub fn hcWriteback(t: Tri, h: u64, hout: u64, pss: u64, branch: u64, branch_ty: []const u8, inj: u64, y: u64, y_ty: []const u8, wts: u64, wts_ty: []const u8, rs: usize, rows: usize, dims: usize, streams: usize, mode: usize, topk: usize, slots: usize, world: usize) !void { + try t.run("_hc_writeback", .{ u(rows), u(dims / 256), 1 }, &.{ + p("H", "*bf16", h), + p("HOUT", "*bf16", hout), + p("PSS", "*fp32", pss), + p("BR", branch_ty, branch), + p("INJ", "*bf16", inj), + p("Y", y_ty, y), + p("WTS", wts_ty, wts), + int("RS", rs), + }, &.{ ci("D", dims), ci("S", streams), ci("MODE", mode), ci("TOPK", topk), ci("SLOTS", slots), ci("BLOCK", 256), ci("WORLD", world) }); + } + + /// glue.ple_embed: gathered n-gram rows, MLX 4-bit group 32. `scale` is a constexpr. + pub fn pleEmbed(t: Tri, weight: u64, scales: u64, biases: u64, out: u64, xs: u64, rows: usize, heads: usize, dh: usize, scale: f32) !void { + try t.run("_ple_embed", .{ u(rows), u(heads), 1 }, &.{ + p("W", "*i32", weight), + p("S", "*bf16", scales), + p("B", "*bf16", biases), + p("OUT", "*bf16", out), + p("XS", "*fp32", xs), + }, &.{ ci("HEADS", heads), ci("DH", dh), aot.cf("SCALE", scale) }); + } + + /// glue.ple_gate: one program per row. + pub fn pleGate(t: Tri, keys: u64, vals: u64, h: u64, norm_key: u64, norm_query: u64, gated: u64, pss: u64, eps: f32, rows: usize, dims: usize, streams: usize) !void { + try t.run("_ple_gate", .{ u(rows), 1, 1 }, &.{ + p("KEYS", "*bf16", keys), + p("VALS", "*bf16", vals), + p("H", "*bf16", h), + p("NK", "*fp32", norm_key), + p("NQ", "*fp32", norm_query), + p("GATED", "*bf16", gated), + p("PSS", "*fp32", pss), + aot.float("eps", eps), + }, &.{ ci("D", dims), ci("S", streams), ci("BLOCK", 512) }); + } + + /// glue.ple_conv. R folds to a constexpr when it is 1. + pub fn pleConv(t: Tri, gated: u64, pss: u64, norm_conv: u64, tail: u64, conv_w: u64, h: u64, hout: u64, nrow: u64, eps: f32, rows: usize, dims: usize, streams: usize, taps: usize, dilation: usize) !void { + try t.run("_ple_conv", .{ u(rows), u(streams * dims / 512), 1 }, &.{ + p("GATED", "*bf16", gated), + p("PSS", "*fp32", pss), + p("NC", "*fp32", norm_conv), + p("TAIL", "*bf16", tail), + p("CW", "*bf16", conv_w), + p("H", "*bf16", h), + p("HOUT", "*bf16", hout), + p("NROW", "*bf16", nrow), + aot.float("eps", eps), + int("R", rows), + }, &.{ ci("D", dims), ci("S", streams), ci("TAPS", taps), ci("DIL", dilation), ci("BLOCK", 512) }); + } + + /// moe.router for one decode row. M folds to a constexpr. 513 rows: the experts, then the shared gate. + pub fn router(t: Tri, x: u64, w: u64, out: u64, x_stride: usize) !void { + try t.run("_router", .{ 1, cdiv(513, 32), 1 }, &.{ + p("X", "*bf16", x), + p("W", "*bf16", w), + p("OUT", "*fp32", out), + int("M", 1), + int("x_stride", x_stride), + }, &.{ ci("D", 2560), ci("NE", 513), ci("BM", 16), ci("BLOCK_E", 32), ci("BK", 256) }); + } + + /// moe._topk_rows: one program, the row's 10 experts and the shared expert. + pub fn topkRows(t: Tri, logits: u64, pick: u64, wts: u64) !void { + try t.run("_topk_rows", .{ 1, 1, 1 }, &.{ + p("L", "*fp32", logits), + p("PICK", "*i32", pick), + p("WTS", "*fp32", wts), + }, &.{ ci("NE", 512), ci("NL", 513), ci("TOPK", 10), ci("SLOTS", 11), ci("BLOCK", 1024), ci("SLOTP", 16) }); + } +}; + +test { + _ = embed; + _ = hc; + _ = qmm; + _ = gdn; + _ = mlp; + _ = moe; + _ = attn; + _ = ple; + _ = mix; + _ = prompt; + _ = native; +} + +test "embed grid is one program per row and per group of 32" { + try std.testing.expectEqual([3]u32{ 4, 80, 1 }, embedGrid(4, 2560)); +} + +test "rmsnorm block is the next power of two at or above the group" { + try std.testing.expectEqual(@as(usize, 4096), rmsBlock(2560)); + try std.testing.expectEqual(@as(usize, 16384), rmsBlock(10240)); +} diff --git a/zig/src/families/flashnext/cuda_mix.zig b/zig/src/families/flashnext/cuda_mix.zig new file mode 100644 index 0000000000..ce050ef0af --- /dev/null +++ b/zig/src/families/flashnext/cuda_mix.zig @@ -0,0 +1,289 @@ +//! The final mixer and the word scores for one decode token: four embedding copies, no inject gate. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const hc = @import("cuda_hc.zig"); +const qmm = @import("cuda_qmm.zig"); + +const dims: usize = 2560; +const streams: usize = 4; +const low: usize = 320; +const wide: usize = streams * dims; +const sk: usize = 32; +const per: usize = 10; +const eps: f32 = 1e-6; +const vocab: usize = 248320; +const prefix = "language_model.model.hyper_connection_mixer."; + +const Tensor = core.safetensors.Tensor; +const Triple = struct { w: Tensor, s: Tensor, b: Tensor }; + +const Gap = struct { off: usize = 0, steps: u32 = 0, at: usize = 0 }; + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +fn roundBf16(v: f32) f32 { + return promote(toBf16(v)); +} + +fn tree32(buf: *[32]f32) f32 { + var n: usize = 32; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + return buf[0]; +} + +fn readU16(raw: []const u8, i: usize) u16 { + return std.mem.readInt(u16, raw[2 * i ..][0..2], .little); +} + +fn shardPath(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) ![]u8 { + const index = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); + defer gpa.free(index); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const file = parsed.value.object.get("weight_map").?.object.get(prefix ++ "hc_norm.weight").?.string; + return std.fs.path.join(gpa, &.{ dir, file }); +} + +fn triple(file: *const core.safetensors.File, name: []const u8, rows: usize, k8: usize) !Triple { + var buf: [180]u8 = undefined; + const w = file.get(try std.fmt.bufPrint(&buf, "{s}.weight", .{name})) orelse return error.MissingTensor; + const s = file.get(try std.fmt.bufPrint(&buf, "{s}.scales", .{name})) orelse return error.MissingTensor; + const b = file.get(try std.fmt.bufPrint(&buf, "{s}.biases", .{name})) orelse return error.MissingTensor; + if (!w.is(.u32, &.{ rows, k8 }) or !s.is(.bf16, &.{ rows, k8 / 4 }) or !b.is(.bf16, &.{ rows, k8 / 4 })) return error.UnexpectedTensor; + return .{ .w = w, .s = s, .b = b }; +} + +/// Split-K sum, one bf16 round, divide by the stream count, SiLU. No inject columns. +fn siluRow(part: []const f32, slices: usize, n: usize, nstreams: usize, act: []u16, xs: []f32) !void { + if (n == 0 or n % 32 != 0 or slices == 0 or part.len != slices * n or act.len != n or xs.len != n / 32) return error.UnexpectedTensor; + for (0..n) |i| { + var sum = part[i]; + for (1..slices) |s| sum += part[s * n + i]; + const y = roundBf16(roundBf16(sum) / @as(f32, @floatFromInt(nstreams))); + act[i] = toBf16(y / (1.0 + @exp(-y))); + } + for (xs, 0..) |*o, g| { + var buf: [32]f32 = undefined; + for (0..32) |j| buf[j] = promote(act[g * 32 + j]); + o.* = tree32(&buf); + } +} + +fn gap(got: []const u16, want: []const u16) Gap { + var g = Gap{}; + for (got, want, 0..) |a, b, i| { + const d = hc.mixedSteps(a, b); + if (d != 0) g.off += 1; + if (d > g.steps) { + g.steps = d; + g.at = i; + } + } + return g; +} + +fn u16s(dst: []u16, raw: []const u8) void { + for (dst, 0..) |*o, i| o.* = readU16(raw, i); +} + +test "a zero down sum without an inject gate is SiLU zero" { + var part: [32]f32 = @splat(0); + var act: [32]u16 = undefined; + var xs: [1]f32 = undefined; + try siluRow(&part, 1, 32, 1, &act, &xs); + try std.testing.expectEqual(@as(u16, 0), act[0]); + try std.testing.expectEqual(@as(f32, 0), xs[0]); +} + +/// Mixer read-out of `hidden` (four stream copies) and the head's scores of that mixed row. +pub fn scores(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, hidden: []const u16) !u32 { + if (hidden.len != wide) return error.UnexpectedTensor; + const path = try shardPath(gpa, io, model_dir); + defer gpa.free(path); + var file = try core.safetensors.File.open(gpa, io, path); + defer file.close(io); + const gamma = file.get(prefix ++ "hc_norm.weight") orelse return error.MissingTensor; + if (!gamma.is(.bf16, &.{wide})) return error.UnexpectedTensor; + const down = try triple(&file, prefix ++ "input_mix_weight_down", low, wide / 8); + const up = try triple(&file, prefix ++ "input_mix_weight_up", wide, low / 8); + const head = try triple(&file, "language_model.lm_head", vocab, dims / 8); + + const scale = try gpa.alloc(f32, wide); + defer gpa.free(scale); + for (scale, 0..) |*o, i| o.* = promote(readU16(gamma.bytes, i)); + const pss = try gpa.alloc(f32, (dims / 256) * streams); + defer gpa.free(pss); + try hc.partialSums(hidden, dims, streams, pss); + const host_norm = try gpa.alloc(u16, wide); + defer gpa.free(host_norm); + try hc.normed(hidden, pss, scale, dims, streams, eps, host_norm); + + var h_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(hidden)); + defer h_b.free(); + var pss_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(pss)); + defer pss_b.free(); + var scale_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(scale)); + defer scale_b.free(); + const down_w = try hc.tileWords(gpa, down.w.bytes, low, wide / 8); + defer gpa.free(down_w); + const down_s = try hc.transposeBf16(gpa, down.s.bytes, low, wide / 32); + defer gpa.free(down_s); + const down_b = try hc.transposeBf16(gpa, down.b.bytes, low, wide / 32); + defer gpa.free(down_b); + var dw = try cuda.DeviceBuffer.fromHost(driver, down_w); + defer dw.free(); + var ds = try cuda.DeviceBuffer.fromHost(driver, down_s); + defer ds.free(); + var db = try cuda.DeviceBuffer.fromHost(driver, down_b); + defer db.free(); + var norm_b = try cuda.DeviceBuffer.alloc(driver, wide * 2); + defer norm_b.free(); + var dn_b = try cuda.DeviceBuffer.alloc(driver, low * 2); + defer dn_b.free(); + var part_b = try cuda.DeviceBuffer.alloc(driver, sk * low * 4); + defer part_b.free(); + try tri.hcDown(h_b.ptr, pss_b.ptr, scale_b.ptr, norm_b.ptr, dw.ptr, ds.ptr, db.ptr, dn_b.ptr, part_b.ptr, eps, 1, low, wide, dims, dims / 256, streams, sk); + + var act_b = try cuda.DeviceBuffer.alloc(driver, low * 2); + defer act_b.free(); + var xs_b = try cuda.DeviceBuffer.alloc(driver, (low / 32) * 4); + defer xs_b.free(); + try tri.hcReduceAct(part_b.ptr, act_b.ptr, xs_b.ptr, act_b.ptr, sk, 1, streams, low, low, 0); + + const up_w = try hc.tileWords(gpa, up.w.bytes, wide, low / 8); + defer gpa.free(up_w); + const up_s = try hc.transposeBf16(gpa, up.s.bytes, wide, low / 32); + defer gpa.free(up_s); + const up_b = try hc.transposeBf16(gpa, up.b.bytes, wide, low / 32); + defer gpa.free(up_b); + var uw = try cuda.DeviceBuffer.fromHost(driver, up_w); + defer uw.free(); + var us = try cuda.DeviceBuffer.fromHost(driver, up_s); + defer us.free(); + var ub = try cuda.DeviceBuffer.fromHost(driver, up_b); + defer ub.free(); + var mixed_b = try cuda.DeviceBuffer.alloc(driver, dims * 2); + defer mixed_b.free(); + var xsm_b = try cuda.DeviceBuffer.alloc(driver, (dims / 32) * 4); + defer xsm_b.free(); + try tri.hcUpmix(act_b.ptr, xs_b.ptr, uw.ptr, us.ptr, ub.ptr, norm_b.ptr, mixed_b.ptr, xsm_b.ptr, 1, wide, low, dims, streams); + + const norm_raw = try gpa.alloc(u8, wide * 2); + defer gpa.free(norm_raw); + const part_raw = try gpa.alloc(u8, sk * low * 4); + defer gpa.free(part_raw); + const act_raw = try gpa.alloc(u8, low * 2); + defer gpa.free(act_raw); + const xs_raw = try gpa.alloc(u8, (low / 32) * 4); + defer gpa.free(xs_raw); + const mixed_raw = try gpa.alloc(u8, dims * 2); + defer gpa.free(mixed_raw); + const xsm_raw = try gpa.alloc(u8, (dims / 32) * 4); + defer gpa.free(xsm_raw); + try stream.synchronize(); + try norm_b.download(0, norm_raw); + try part_b.download(0, part_raw); + try act_b.download(0, act_raw); + try xs_b.download(0, xs_raw); + try mixed_b.download(0, mixed_raw); + try xsm_b.download(0, xsm_raw); + + var norm_off: usize = 0; + for (host_norm, 0..) |want, i| { + if (readU16(norm_raw, i) != want) norm_off += 1; + } + const gpu_norm = try gpa.alloc(u16, per * 32); + defer gpa.free(gpu_norm); + u16s(gpu_norm, norm_raw); + const host_part = try hc.firstPartial(gpu_norm, down.w.bytes[0 .. per * 16], down.s.bytes[0 .. per * 2], down.b.bytes[0 .. per * 2], per); + const gpu_part: f32 = @bitCast(std.mem.readInt(u32, part_raw[0..4], .little)); + const col_ulp = hc.ulps(gpu_part, host_part); + + const part_f = try gpa.alloc(f32, sk * low); + defer gpa.free(part_f); + for (part_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, part_raw[4 * i ..][0..4], .little)); + const act_want = try gpa.alloc(u16, low); + defer gpa.free(act_want); + const xs_want = try gpa.alloc(f32, low / 32); + defer gpa.free(xs_want); + try siluRow(part_f, sk, low, streams, act_want, xs_want); + var act_off: usize = 0; + for (act_want, 0..) |want, i| { + if (readU16(act_raw, i) != want) act_off += 1; + } + var xs_ulp: u32 = 0; + for (xs_want, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, xs_raw[4 * i ..][0..4], .little)); + const dist = hc.ulps(got, want); + if (dist > xs_ulp) xs_ulp = dist; + } + + const act_u = try gpa.alloc(u16, low); + defer gpa.free(act_u); + u16s(act_u, act_raw); + const xs_f = try gpa.alloc(f32, low / 32); + defer gpa.free(xs_f); + for (xs_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xs_raw[4 * i ..][0..4], .little)); + const norm_u = try gpa.alloc(u16, wide); + defer gpa.free(norm_u); + u16s(norm_u, norm_raw); + const host_mixed = try gpa.alloc(u16, dims); + defer gpa.free(host_mixed); + try hc.mix(act_u, xs_f, norm_u, up.w.bytes, up.s.bytes, up.b.bytes, wide, low / 8, dims, streams, host_mixed); + const got_mixed = try gpa.alloc(u16, dims); + defer gpa.free(got_mixed); + u16s(got_mixed, mixed_raw); + const mixed_gap = gap(got_mixed, host_mixed); + std.debug.print("mix norm {d} col_ulp {d} act {d} xs_ulp {d} off {d} max_steps {d} at {d} got {x:0>4} host {x:0>4} head {x:0>4}\n", .{ norm_off, col_ulp, act_off, xs_ulp, mixed_gap.off, mixed_gap.steps, mixed_gap.at, got_mixed[mixed_gap.at], host_mixed[mixed_gap.at], got_mixed[0] }); + + var lane = try qmm.pack(gpa, head.w.bytes, head.s.bytes, head.b.bytes, vocab, dims / 8); + defer lane.deinit(gpa); + var hw = try cuda.DeviceBuffer.fromHost(driver, lane.weight); + defer hw.free(); + var hs = try cuda.DeviceBuffer.fromHost(driver, lane.scales); + defer hs.free(); + var hb = try cuda.DeviceBuffer.fromHost(driver, lane.biases); + defer hb.free(); + var logits_b = try cuda.DeviceBuffer.alloc(driver, vocab * 2); + defer logits_b.free(); + try qmm.matmul(driver, stream.*, mixed_b.ptr, xsm_b.ptr, hw.ptr, hs.ptr, hb.ptr, logits_b.ptr, 1, vocab, dims); + const xsm_f = try gpa.alloc(f32, dims / 32); + defer gpa.free(xsm_f); + for (xsm_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xsm_raw[4 * i ..][0..4], .little)); + const host_logits = try gpa.alloc(u16, vocab); + defer gpa.free(host_logits); + try qmm.dotRow(got_mixed, xsm_f, head.w.bytes, head.s.bytes, head.b.bytes, vocab, dims, host_logits); + const logits_raw = try gpa.alloc(u8, vocab * 2); + defer gpa.free(logits_raw); + try stream.synchronize(); + try logits_b.download(0, logits_raw); + const got_logits = try gpa.alloc(u16, vocab); + defer gpa.free(got_logits); + u16s(got_logits, logits_raw); + const score_gap = gap(got_logits, host_logits); + std.debug.print("scores n {d} off {d} max_steps {d} at {d} got {x:0>4} host {x:0>4} head {x:0>4}\n", .{ vocab, score_gap.off, score_gap.steps, score_gap.at, got_logits[score_gap.at], host_logits[score_gap.at], got_logits[0] }); + if (score_gap.off > 0 and score_gap.off <= 8) { + for (got_logits, host_logits, 0..) |a, b, i| { + const d = hc.mixedSteps(a, b); + if (d != 0) std.debug.print("score {d} gpu {x:0>4} host {x:0>4} steps {d}\n", .{ i, a, b, d }); + } + } + + if (norm_off != 0 or col_ulp > 8 or act_off != 0 or xs_ulp > 4 or mixed_gap.steps > 2 or score_gap.steps > 1 or score_gap.off > 3) return 1; + return 0; +} diff --git a/zig/src/families/flashnext/cuda_mlp.zig b/zig/src/families/flashnext/cuda_mlp.zig new file mode 100644 index 0000000000..bd7e875d72 --- /dev/null +++ b/zig/src/families/flashnext/cuda_mlp.zig @@ -0,0 +1,255 @@ +//! Layer 0's second hyper-connection: the mixer branch written back, then the same read-out on the mlp weights. + +const std = @import("std"); +const cuda = @import("cuda"); +const embed = @import("cuda_embed.zig"); +const hc = @import("cuda_hc.zig"); + +const branch_sk: usize = 8; + +/// `slices` are the output face's unreduced K slices. `h` is the four streams, `inj` the first connection's gates. +pub fn connect(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, mapped: *const embed.Mapped, tri: Tri, h: u64, inj: u64, hidden: []const u16, slices: []const f32) !u32 { + if (hidden.len == 0 or hidden.len % 4 != 0) return error.UnexpectedTensor; + const streams: usize = 4; + const dims = hidden.len / streams; + if (slices.len != branch_sk * dims) return error.UnexpectedTensor; + const inj_bytes = try gpa.alloc(u8, streams * 2); + defer gpa.free(inj_bytes); + try stream.synchronize(); + var inj_buf: cuda.DeviceBuffer = .{ .d = driver, .ptr = inj, .len = inj_bytes.len }; + try inj_buf.download(0, inj_bytes); + const inj_u16 = try gpa.alloc(u16, streams); + defer gpa.free(inj_u16); + for (inj_u16, 0..) |*o, i| o.* = std.mem.readInt(u16, inj_bytes[2 * i ..][0..2], .little); + const branch = try gpa.alloc(u16, dims); + defer gpa.free(branch); + try hc.sumSlices(slices, branch_sk, dims, branch); + const host_h = try gpa.alloc(u16, hidden.len); + defer gpa.free(host_h); + try hc.applyBranch(hidden, branch, inj_u16, dims, streams, host_h); + + var part = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(slices)); + defer part.free(); + const nc = dims / 256; + var pss_b = try cuda.DeviceBuffer.alloc(driver, nc * streams * 4); + defer pss_b.free(); + try tri.hcWriteback(h, h, pss_b.ptr, part.ptr, "*fp32", inj, h, "*bf16", h, "*bf16", dims, 1, dims, streams, 4, 1, 1, branch_sk); + const got_h = try gpa.alloc(u8, hidden.len * 2); + defer gpa.free(got_h); + const pss_bytes = try gpa.alloc(u8, nc * streams * 4); + defer gpa.free(pss_bytes); + try stream.synchronize(); + var h_buf: cuda.DeviceBuffer = .{ .d = driver, .ptr = h, .len = got_h.len }; + try h_buf.download(0, got_h); + try pss_b.download(0, pss_bytes); + var h_off: usize = 0; + var h_steps: u32 = 0; + for (host_h, 0..) |want, i| { + const got = std.mem.readInt(u16, got_h[2 * i ..][0..2], .little); + const dist = hc.mixedSteps(got, want); + if (dist != 0) h_off += 1; + if (dist > h_steps) h_steps = dist; + } + const want_pss = try gpa.alloc(f32, nc * streams); + defer gpa.free(want_pss); + try hc.partialSums(host_h, dims, streams, want_pss); + var pss_ulp: u32 = 0; + for (want_pss, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, pss_bytes[4 * i ..][0..4], .little)); + const dist = hc.ulps(got, want); + if (dist > pss_ulp) pss_ulp = dist; + } + std.debug.print("mlp writeback off {d} max_steps {d} pss_ulp {d} head {x:0>4} host {x:0>4} branch {x:0>4}\n", .{ h_off, h_steps, pss_ulp, std.mem.readInt(u16, got_h[0..2], .little), host_h[0], branch[0] }); + if (h_steps != 0 or pss_ulp > 4) return 1; + + const prefix = "language_model.model.layers.0.mlp_hyper_connection."; + var name: [180]u8 = undefined; + const down_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_down.weight", .{prefix}), .u32); + const down_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_down.scales", .{prefix}), .bf16); + const down_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_down.biases", .{prefix}), .bf16); + const inj_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}block_inject_weight.weight", .{prefix}), .u32); + const inj_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}block_inject_weight.scales", .{prefix}), .bf16); + const inj_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}block_inject_weight.biases", .{prefix}), .bf16); + const gamma = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}hc_norm.weight", .{prefix}), .bf16); + const up_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_up.weight", .{prefix}), .u32); + const up_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_up.scales", .{prefix}), .bf16); + const up_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_up.biases", .{prefix}), .bf16); + const n = down_w.dim(0) + inj_w.dim(0); + const k8 = down_w.dim(1); + const kg = k8 / 4; + if (n != 324 or k8 != 1280 or gamma.dim(0) != streams * dims) return error.UnexpectedTensor; + const stacked_w = try gpa.alloc(u8, down_w.bytes.len + inj_w.bytes.len); + defer gpa.free(stacked_w); + @memcpy(stacked_w[0..down_w.bytes.len], down_w.bytes); + @memcpy(stacked_w[down_w.bytes.len..], inj_w.bytes); + const stacked_s = try gpa.alloc(u8, down_s.bytes.len + inj_s.bytes.len); + defer gpa.free(stacked_s); + @memcpy(stacked_s[0..down_s.bytes.len], down_s.bytes); + @memcpy(stacked_s[down_s.bytes.len..], inj_s.bytes); + const stacked_b = try gpa.alloc(u8, down_b.bytes.len + inj_b.bytes.len); + defer gpa.free(stacked_b); + @memcpy(stacked_b[0..down_b.bytes.len], down_b.bytes); + @memcpy(stacked_b[down_b.bytes.len..], inj_b.bytes); + const tiled = try hc.tileWords(gpa, stacked_w, n, k8); + defer gpa.free(tiled); + const scales_t = try hc.transposeBf16(gpa, stacked_s, n, kg); + defer gpa.free(scales_t); + const biases_t = try hc.transposeBf16(gpa, stacked_b, n, kg); + defer gpa.free(biases_t); + const scale = try gpa.alloc(f32, gamma.dim(0)); + defer gpa.free(scale); + for (scale, 0..) |*o, i| o.* = @bitCast(@as(u32, std.mem.readInt(u16, gamma.bytes[2 * i ..][0..2], .little)) << 16); + var w_b = try cuda.DeviceBuffer.fromHost(driver, tiled); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, scales_t); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, biases_t); + defer b_b.free(); + var scale_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(scale)); + defer scale_b.free(); + var norm_b = try cuda.DeviceBuffer.alloc(driver, streams * dims * 2); + defer norm_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, n * 2); + defer out_b.free(); + const down_sk: usize = 32; + var down_part = try cuda.DeviceBuffer.alloc(driver, down_sk * n * 4); + defer down_part.free(); + try tri.hcDown(h, pss_b.ptr, scale_b.ptr, norm_b.ptr, w_b.ptr, s_b.ptr, b_b.ptr, out_b.ptr, down_part.ptr, 1e-6, 1, n, k8 * 8, dims, nc, streams, down_sk); + const got_norm = try gpa.alloc(u8, streams * dims * 2); + defer gpa.free(got_norm); + try stream.synchronize(); + try norm_b.download(0, got_norm); + const pss_f = try gpa.alloc(f32, nc * streams); + defer gpa.free(pss_f); + for (pss_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, pss_bytes[4 * i ..][0..4], .little)); + const gpu_h = try gpa.alloc(u16, hidden.len); + defer gpa.free(gpu_h); + for (gpu_h, 0..) |*o, i| o.* = std.mem.readInt(u16, got_h[2 * i ..][0..2], .little); + const want_norm = try gpa.alloc(u16, hidden.len); + defer gpa.free(want_norm); + try hc.normed(gpu_h, pss_f, scale, dims, streams, 1e-6, want_norm); + var norm_off: usize = 0; + var norm_steps: u32 = 0; + var norm_at: usize = 0; + for (want_norm, 0..) |want, i| { + const got = std.mem.readInt(u16, got_norm[2 * i ..][0..2], .little); + if (got != want) { + const dist = hc.mixedSteps(got, want); + if (norm_off == 0 or dist > norm_steps) { + norm_steps = dist; + norm_at = i; + } + norm_off += 1; + } + } + var part4: [4]u8 = undefined; + try down_part.download(0, &part4); + const gpu_part: f32 = @bitCast(std.mem.readInt(u32, &part4, .little)); + const gpu_norm = try gpa.alloc(u16, 320); + defer gpa.free(gpu_norm); + for (gpu_norm, 0..) |*o, i| o.* = std.mem.readInt(u16, got_norm[2 * i ..][0..2], .little); + const host_part = try hc.firstPartial(gpu_norm, stacked_w[0 .. k8 * 4], stacked_s[0 .. kg * 2], stacked_b[0 .. kg * 2], 10); + const col_ulp = hc.ulps(gpu_part, host_part); + std.debug.print("mlp hcdown normed off {d} max_steps {d} at {d} col0 gpu {d} host {d} ulp {d}\n", .{ norm_off, norm_steps, norm_at, gpu_part, host_part, col_ulp }); + if (norm_steps > 1 or col_ulp > 8) return 1; + + const low: usize = 320; + var act_b = try cuda.DeviceBuffer.alloc(driver, low * 2); + defer act_b.free(); + var xs_b = try cuda.DeviceBuffer.alloc(driver, (low / 32) * 4); + defer xs_b.free(); + var inj_out = try cuda.DeviceBuffer.alloc(driver, streams * 2); + defer inj_out.free(); + try tri.hcReduceAct(down_part.ptr, act_b.ptr, xs_b.ptr, inj_out.ptr, down_sk, 1, streams, low, n, 1); + const part_bytes = try gpa.alloc(u8, down_sk * n * 4); + defer gpa.free(part_bytes); + try stream.synchronize(); + try down_part.download(0, part_bytes); + const part_f = try gpa.alloc(f32, down_sk * n); + defer gpa.free(part_f); + for (part_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, part_bytes[4 * i ..][0..4], .little)); + const act_want = try gpa.alloc(u16, low); + defer gpa.free(act_want); + const inj_want = try gpa.alloc(u16, streams); + defer gpa.free(inj_want); + const xs_want = try gpa.alloc(f32, low / 32); + defer gpa.free(xs_want); + try hc.reduceAct(part_f, down_sk, n, low, streams, act_want, inj_want, xs_want); + const act_bytes = try gpa.alloc(u8, low * 2); + defer gpa.free(act_bytes); + const inj_out_bytes = try gpa.alloc(u8, streams * 2); + defer gpa.free(inj_out_bytes); + const xs_bytes = try gpa.alloc(u8, (low / 32) * 4); + defer gpa.free(xs_bytes); + try act_b.download(0, act_bytes); + try inj_out.download(0, inj_out_bytes); + try xs_b.download(0, xs_bytes); + var act_off: usize = 0; + for (act_want, 0..) |want, i| { + if (std.mem.readInt(u16, act_bytes[2 * i ..][0..2], .little) != want) act_off += 1; + } + var inj_off: usize = 0; + for (inj_want, 0..) |want, i| { + if (std.mem.readInt(u16, inj_out_bytes[2 * i ..][0..2], .little) != want) inj_off += 1; + } + var xs_ulp: u32 = 0; + for (xs_want, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const dist = hc.ulps(got, want); + if (dist > xs_ulp) xs_ulp = dist; + } + std.debug.print("mlp reduce act mismatch {d} inject mismatch {d} xs max_ulp {d}\n", .{ act_off, inj_off, xs_ulp }); + if (act_off != 0 or inj_off != 0 or xs_ulp > 4) return 1; + + const up_n = up_w.dim(0); + const up_k8 = up_w.dim(1); + if (up_n != streams * dims or up_k8 != 40) return error.UnexpectedTensor; + const up_tiled = try hc.tileWords(gpa, up_w.bytes, up_n, up_k8); + defer gpa.free(up_tiled); + const up_scales = try hc.transposeBf16(gpa, up_s.bytes, up_n, up_k8 / 4); + defer gpa.free(up_scales); + const up_biases = try hc.transposeBf16(gpa, up_b.bytes, up_n, up_k8 / 4); + defer gpa.free(up_biases); + var up_wb = try cuda.DeviceBuffer.fromHost(driver, up_tiled); + defer up_wb.free(); + var up_sb = try cuda.DeviceBuffer.fromHost(driver, up_scales); + defer up_sb.free(); + var up_bb = try cuda.DeviceBuffer.fromHost(driver, up_biases); + defer up_bb.free(); + var mixed_b = try cuda.DeviceBuffer.alloc(driver, dims * 2); + defer mixed_b.free(); + var xsm_b = try cuda.DeviceBuffer.alloc(driver, (dims / 32) * 4); + defer xsm_b.free(); + try tri.hcUpmix(act_b.ptr, xs_b.ptr, up_wb.ptr, up_sb.ptr, up_bb.ptr, norm_b.ptr, mixed_b.ptr, xsm_b.ptr, 1, up_n, up_k8 * 8, dims, streams); + const mixed_bytes = try gpa.alloc(u8, dims * 2); + defer gpa.free(mixed_bytes); + try stream.synchronize(); + try mixed_b.download(0, mixed_bytes); + const act_u16 = try gpa.alloc(u16, low); + defer gpa.free(act_u16); + for (act_u16, 0..) |*o, i| o.* = std.mem.readInt(u16, act_bytes[2 * i ..][0..2], .little); + const xs_f = try gpa.alloc(f32, low / 32); + defer gpa.free(xs_f); + for (xs_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const norm_u16 = try gpa.alloc(u16, hidden.len); + defer gpa.free(norm_u16); + for (norm_u16, 0..) |*o, i| o.* = std.mem.readInt(u16, got_norm[2 * i ..][0..2], .little); + const host_mixed = try gpa.alloc(u16, dims); + defer gpa.free(host_mixed); + try hc.mix(act_u16, xs_f, norm_u16, up_w.bytes, up_s.bytes, up_b.bytes, up_n, up_k8, dims, streams, host_mixed); + var mixed_off: usize = 0; + var max_steps: u32 = 0; + var mixed_at: usize = 0; + for (host_mixed, 0..) |want, i| { + const got = std.mem.readInt(u16, mixed_bytes[2 * i ..][0..2], .little); + const dist = hc.mixedSteps(got, want); + if (dist != 0) mixed_off += 1; + if (dist > max_steps) { + max_steps = dist; + mixed_at = i; + } + } + std.debug.print("mlp upmix dims {d} off {d} max_steps {d} at {d} head {x:0>4} host {x:0>4}\n", .{ dims, mixed_off, max_steps, mixed_at, std.mem.readInt(u16, mixed_bytes[0..2], .little), host_mixed[0] }); + if (max_steps > 2) return 1; + return 0; +} diff --git a/zig/src/families/flashnext/cuda_moe.zig b/zig/src/families/flashnext/cuda_moe.zig new file mode 100644 index 0000000000..5a0030dfc0 --- /dev/null +++ b/zig/src/families/flashnext/cuda_moe.zig @@ -0,0 +1,512 @@ +//! Layer 0's experts on one mixed row: router, top 10 plus the shared expert, then each slot's gate, up and down. + +const std = @import("std"); +const cuda = @import("cuda"); +const embed = @import("cuda_embed.zig"); +const hc = @import("cuda_hc.zig"); + +const count: usize = 512; +const top_k: usize = 10; +const slots: usize = top_k + 1; +const dims: usize = 2560; +const width: usize = 640; +const routed_rows: usize = count + 1; +const tile: usize = 16; +const block_bytes: usize = 160 * 4; + +const pack_symbol: [:0]const u8 = "_ZN15tf_experts_pack11pack_kernelILi1EEEvPKjPKtS4_Pjiiii"; +const plan_symbol: [:0]const u8 = "_ZN10tf_experts11plan_kernelEPKiiiiPiS2_S2_"; +const up_symbol: [:0]const u8 = "_ZN10tf_experts13expert_kernelILi32ELi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; +const down_symbol: [:0]const u8 = "_ZN10tf_experts13expert_kernelILi32ELi1ELi0ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; + +const Proj = struct { + routed_w: []const u8, + routed_s: []const u8, + routed_b: []const u8, + shared_w: []const u8, + shared_s: []const u8, + shared_b: []const u8, + n: usize, + k: usize, +}; + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +fn roundTrip(v: f32) f32 { + return promote(toBf16(v)); +} + +/// The expert kernel's group sum: each lane adds its eight inputs, then the quad adds as (0+1)+(2+3). +fn groupSum(x: []const u16, g: usize) f32 { + var pair: [2]f32 = .{ 0, 0 }; + for (0..4) |t| { + var s: f32 = 0; + const base = g * 32 + t * 8; + for (0..8) |j| s += promote(x[base + j]); + if (t % 2 == 0) pair[t / 2] = s else pair[t / 2] += s; + } + return pair[0] + pair[1]; +} + +/// One expert matrix: groups in order, `p * scale + xs * bias`, fp32. `xs` is `groupSum` of `x`. +fn project(x: []const u16, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k: usize, out: []f32) void { + const kg = k / 32; + const k8 = k / 8; + var xs: [80]f32 = undefined; + for (0..kg) |g| xs[g] = groupSum(x, g); + for (0..n) |col| { + var acc: f32 = 0; + const wb = col * k8 * 4; + const sb = col * kg * 2; + for (0..kg) |g| { + var p: f32 = 0; + for (0..4) |i| { + const word = std.mem.readInt(u32, words[wb + (g * 4 + i) * 4 ..][0..4], .little); + const base = g * 32 + i * 8; + inline for (0..8) |j| { + const q: f32 = @floatFromInt((word >> (j * 4)) & 0xF); + p += promote(x[base + j]) * q; + } + } + const s = promote(std.mem.readInt(u16, scales[sb + g * 2 ..][0..2], .little)); + const b = promote(std.mem.readInt(u16, biases[sb + g * 2 ..][0..2], .little)); + acc = @mulAdd(f32, xs[g], b, @mulAdd(f32, p, s, acc)); + } + out[col] = acc; + } +} + +/// bf16(SiLU(bf16(gate)) * bf16(up)), the SwiGLU epilogue with no clip. +fn swiglu(gate: f32, up: f32) u16 { + const gv = roundTrip(gate); + const uv = roundTrip(up); + const silu = roundTrip(gv / (1 + @exp(-gv))); + return toBf16(silu * uv); +} + +/// Router rows in K tiles of 256, the same tile the captured program uses, added in tile order. +fn routerDots(x: []const u16, w: []const u16, out: []f32) void { + const k = x.len; + for (out, 0..) |*o, e| { + var acc: f32 = 0; + const row = w[e * k ..][0..k]; + var t: usize = 0; + while (t < k) : (t += 256) { + var part: f32 = 0; + for (0..256) |i| part += promote(x[t + i]) * promote(row[t + i]); + acc += part; + } + o.* = acc; + } +} + +/// Largest `k` of the first `n` logits (lower id on a tie), weights exp(l - l0) / sum rounded to bf16, then the shared gate. +fn choose(logits: []const f32, n: usize, k: usize, picks: []i32, wts: []f32) void { + var scratch: [count]f32 = undefined; + for (0..n) |i| scratch[i] = logits[i]; + var top = scratch[0]; + for (scratch[1..n]) |v| top = @max(top, v); + var total: f32 = 0; + var exs: [16]f32 = undefined; + for (0..k) |slot| { + var m = scratch[0]; + var idx: usize = 0; + for (scratch[1..n], 1..) |v, i| if (v > m) { + m = v; + idx = i; + }; + const ex = @exp(m - top); + picks[slot] = @intCast(idx); + exs[slot] = ex; + total += ex; + scratch[idx] = -std.math.inf(f32); + } + for (0..k) |slot| wts[slot] = roundTrip(exs[slot] / total); + const shared = roundTrip(logits[n]); + picks[k] = @intCast(n); + wts[k] = roundTrip(1 / (1 + @exp(-shared))); +} + +fn cint(v: usize) c_int { + return @intCast(v); +} + +fn loadProj(mapped: *const embed.Mapped, gpa: std.mem.Allocator, name: []const u8, n: usize, k: usize) !Proj { + const prefix = "language_model.model.layers.0.mlp."; + var buf: [180]u8 = undefined; + const k8 = k / 8; + const kg = k / 32; + const rw = try mapped.lookup(gpa, try std.fmt.bufPrint(&buf, "{s}switch_mlp.{s}.weight", .{ prefix, name }), .u32); + const rs = try mapped.lookup(gpa, try std.fmt.bufPrint(&buf, "{s}switch_mlp.{s}.scales", .{ prefix, name }), .bf16); + const rb = try mapped.lookup(gpa, try std.fmt.bufPrint(&buf, "{s}switch_mlp.{s}.biases", .{ prefix, name }), .bf16); + const sw = try mapped.lookup(gpa, try std.fmt.bufPrint(&buf, "{s}shared_expert.{s}.weight", .{ prefix, name }), .u32); + const ss = try mapped.lookup(gpa, try std.fmt.bufPrint(&buf, "{s}shared_expert.{s}.scales", .{ prefix, name }), .bf16); + const sb = try mapped.lookup(gpa, try std.fmt.bufPrint(&buf, "{s}shared_expert.{s}.biases", .{ prefix, name }), .bf16); + if (rw.dim(0) != count or rw.dim(1) != n or rw.dim(2) != k8) return error.UnexpectedTensor; + if (rs.dim(0) != count or rs.dim(1) != n or rs.dim(2) != kg) return error.UnexpectedTensor; + if (rb.dim(0) != count or rb.dim(1) != n or rb.dim(2) != kg) return error.UnexpectedTensor; + if (sw.dim(0) != n or sw.dim(1) != k8 or ss.dim(0) != n or ss.dim(1) != kg) return error.UnexpectedTensor; + if (sb.dim(0) != n or sb.dim(1) != kg) return error.UnexpectedTensor; + return .{ .routed_w = rw.bytes, .routed_s = rs.bytes, .routed_b = rb.bytes, .shared_w = sw.bytes, .shared_s = ss.bytes, .shared_b = sb.bytes, .n = n, .k = k }; +} + +fn copyExpert(dst_w: []u8, dst_s: []u8, dst_b: []u8, p: Proj, id: i32, slot: usize) void { + const k8 = p.k / 8; + const kg = p.k / 32; + const wb = p.n * k8 * 4; + const sb = p.n * kg * 2; + const routed = id != count; + const wsrc = if (routed) p.routed_w[@as(usize, @intCast(id)) * wb ..][0..wb] else p.shared_w; + const ssrc = if (routed) p.routed_s[@as(usize, @intCast(id)) * sb ..][0..sb] else p.shared_s; + const bsrc = if (routed) p.routed_b[@as(usize, @intCast(id)) * sb ..][0..sb] else p.shared_b; + @memcpy(dst_w[slot * wb ..][0..wb], wsrc); + @memcpy(dst_s[slot * sb ..][0..sb], ssrc); + @memcpy(dst_b[slot * sb ..][0..sb], bsrc); +} + +fn packDownload(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: cuda.Stream, pack_fn: cuda.Function, words: []const u8, scales: []const u8, biases: []const u8, e: usize, n: usize, k: usize) ![]u8 { + const kg = k / 32; + const nb = n / 32; + const out_len = e * nb * kg * block_bytes; + var w_b = try cuda.DeviceBuffer.fromHost(driver, words); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, biases); + defer b_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, out_len); + defer out_b.free(); + var a: cuda.Args = .{}; + a.add(w_b.ptr); + a.add(s_b.ptr); + a.add(b_b.ptr); + a.add(out_b.ptr); + a.add(cint(n)); + a.add(cint(k / 8)); + a.add(cint(kg)); + a.add(cint(nb)); + try cuda.launch.launch(pack_fn, .{ .grid = .{ .x = @intCast(kg), .y = @intCast(nb), .z = @intCast(e) }, .block = .{ .x = 160 } }, stream, &a); + const out = try gpa.alloc(u8, out_len); + errdefer gpa.free(out); + try stream.synchronize(); + try out_b.download(0, out); + return out; +} + +fn stackHalves(gpa: std.mem.Allocator, gate: []const u8, up: []const u8) ![]u8 { + if (gate.len != up.len or gate.len % block_bytes != 0) return error.UnexpectedTensor; + const nblk = gate.len / block_bytes; + const out = try gpa.alloc(u8, gate.len + up.len); + for (0..nblk) |b| { + @memcpy(out[b * 2 * block_bytes ..][0..block_bytes], gate[b * block_bytes ..][0..block_bytes]); + @memcpy(out[b * 2 * block_bytes + block_bytes ..][0..block_bytes], up[b * block_bytes ..][0..block_bytes]); + } + return out; +} + +fn maxUlp(got: []const f32, want: []const f32) struct { off: usize, ulp: u32 } { + var off: usize = 0; + var ulp: u32 = 0; + for (got, want) |g, w| { + const dist = hc.ulps(g, w); + if (dist != 0) off += 1; + if (dist > ulp) ulp = dist; + } + return .{ .off = off, .ulp = ulp }; +} + +/// The 512 routed experts and the shared expert, on the mixed row already at `mixed`. +pub fn experts(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, mapped: *const embed.Mapped, tri: Tri, mixed: u64) !u32 { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + const prefix = "language_model.model.layers.0.mlp."; + var name: [180]u8 = undefined; + const gate = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}gate.weight", .{prefix}), .bf16); + const sg_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}shared_expert_gate.weight", .{prefix}), .u32); + const sg_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}shared_expert_gate.scales", .{prefix}), .bf16); + const sg_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}shared_expert_gate.biases", .{prefix}), .bf16); + if (gate.dim(0) != count or gate.dim(1) != dims) return error.UnexpectedTensor; + if (sg_w.dim(0) != 1 or sg_w.dim(1) != dims / 8) return error.UnexpectedTensor; + + const row_raw = try gpa.alloc(u8, dims * 2); + defer gpa.free(row_raw); + try stream.synchronize(); + var mixed_buf: cuda.DeviceBuffer = .{ .d = driver, .ptr = mixed, .len = row_raw.len }; + try mixed_buf.download(0, row_raw); + const row = try gpa.alloc(u16, dims); + defer gpa.free(row); + for (row, 0..) |*o, i| o.* = std.mem.readInt(u16, row_raw[2 * i ..][0..2], .little); + + const router = try gpa.alloc(u16, routed_rows * dims); + defer gpa.free(router); + @memcpy(std.mem.sliceAsBytes(router[0 .. count * dims]), gate.bytes); + try embed.dequant(sg_w.bytes, sg_s.bytes, sg_b.bytes, dims, 0, router[count * dims ..][0..dims]); + const host_l = try gpa.alloc(f32, routed_rows); + defer gpa.free(host_l); + routerDots(row, router, host_l); + + var w_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(router)); + defer w_b.free(); + var logits_b = try cuda.DeviceBuffer.alloc(driver, routed_rows * 4); + defer logits_b.free(); + var pick_b = try cuda.DeviceBuffer.alloc(driver, slots * 4); + defer pick_b.free(); + var wts_b = try cuda.DeviceBuffer.alloc(driver, slots * 4); + defer wts_b.free(); + try tri.router(mixed, w_b.ptr, logits_b.ptr, dims); + try tri.topkRows(logits_b.ptr, pick_b.ptr, wts_b.ptr); + + const logit_raw = try gpa.alloc(u8, routed_rows * 4); + defer gpa.free(logit_raw); + const pick_raw = try gpa.alloc(u8, slots * 4); + defer gpa.free(pick_raw); + const wts_raw = try gpa.alloc(u8, slots * 4); + defer gpa.free(wts_raw); + try stream.synchronize(); + try logits_b.download(0, logit_raw); + try pick_b.download(0, pick_raw); + try wts_b.download(0, wts_raw); + const gpu_l = try gpa.alloc(f32, routed_rows); + defer gpa.free(gpu_l); + for (gpu_l, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, logit_raw[4 * i ..][0..4], .little)); + var gpu_pick: [slots]i32 = undefined; + var gpu_w: [slots]f32 = undefined; + for (0..slots) |i| { + gpu_pick[i] = std.mem.readInt(i32, pick_raw[4 * i ..][0..4], .little); + gpu_w[i] = @bitCast(std.mem.readInt(u32, wts_raw[4 * i ..][0..4], .little)); + } + const route = maxUlp(gpu_l, host_l); + var from_gpu: [slots]i32 = undefined; + var w_gpu: [slots]f32 = undefined; + choose(gpu_l, count, top_k, &from_gpu, &w_gpu); + var from_host: [slots]i32 = undefined; + var w_host: [slots]f32 = undefined; + choose(host_l, count, top_k, &from_host, &w_host); + var pick_off: usize = 0; + var host_pick_off: usize = 0; + for (gpu_pick, from_gpu, from_host) |g, a, b| { + if (g != a) pick_off += 1; + if (g != b) host_pick_off += 1; + } + const wts_same = maxUlp(&gpu_w, &w_gpu); + var host_wts_off: usize = 0; + var host_wts_steps: u32 = 0; + for (gpu_w, w_host) |g, h| { + const steps = hc.mixedSteps(toBf16(g), toBf16(h)); + if (steps != 0) host_wts_off += 1; + if (steps > host_wts_steps) host_wts_steps = steps; + } + std.debug.print("router off {d} max_ulp {d} head {x:0>8} host {x:0>8}\n", .{ route.off, route.ulp, @as(u32, @bitCast(gpu_l[0])), @as(u32, @bitCast(host_l[0])) }); + std.debug.print("topk pick_off {d} host_pick_off {d} wts_ulp {d} host_wts_off {d} host_wts_steps {d} picks", .{ pick_off, host_pick_off, wts_same.ulp, host_wts_off, host_wts_steps }); + for (gpu_pick) |id| std.debug.print(" {d}", .{id}); + std.debug.print("\n", .{}); + + var packed_ids: [slots]i32 = undefined; + var npack: usize = 0; + var local: [slots]i32 = undefined; + for (gpu_pick, 0..) |id, s| { + if (id < 0 or id > count) return error.UnexpectedTensor; + var found: ?usize = null; + for (packed_ids[0..npack], 0..) |prev, i| if (prev == id) { + found = i; + }; + if (found) |i| { + local[s] = @intCast(i); + } else { + packed_ids[npack] = id; + local[s] = @intCast(npack); + npack += 1; + } + } + + const gate_p = try loadProj(mapped, gpa, "gate_proj", width, dims); + const up_p = try loadProj(mapped, gpa, "up_proj", width, dims); + const down_p = try loadProj(mapped, gpa, "down_proj", dims, width); + const gathered_w = try gpa.alloc(u8, npack * width * (dims / 8) * 4); + defer gpa.free(gathered_w); + const gathered_s = try gpa.alloc(u8, npack * width * (dims / 32) * 2); + defer gpa.free(gathered_s); + const gathered_b = try gpa.alloc(u8, gathered_s.len); + defer gpa.free(gathered_b); + const up_w = try gpa.alloc(u8, gathered_w.len); + defer gpa.free(up_w); + const up_s = try gpa.alloc(u8, gathered_s.len); + defer gpa.free(up_s); + const up_b = try gpa.alloc(u8, gathered_s.len); + defer gpa.free(up_b); + const down_w = try gpa.alloc(u8, npack * dims * (width / 8) * 4); + defer gpa.free(down_w); + const down_s = try gpa.alloc(u8, npack * dims * (width / 32) * 2); + defer gpa.free(down_s); + const down_b = try gpa.alloc(u8, down_s.len); + defer gpa.free(down_b); + for (packed_ids[0..npack], 0..) |id, slot| { + copyExpert(gathered_w, gathered_s, gathered_b, gate_p, id, slot); + copyExpert(up_w, up_s, up_b, up_p, id, slot); + copyExpert(down_w, down_s, down_b, down_p, id, slot); + } + + var pack_mod = try cuda.Module.load(driver, cuda.kernels.experts_pack); + defer pack_mod.unload(); + var exp_mod = try cuda.Module.load(driver, cuda.kernels.experts); + defer exp_mod.unload(); + const pack_fn = try pack_mod.function(pack_symbol); + const packed_g = try packDownload(gpa, driver, stream.*, pack_fn, gathered_w, gathered_s, gathered_b, npack, width, dims); + defer gpa.free(packed_g); + const packed_u = try packDownload(gpa, driver, stream.*, pack_fn, up_w, up_s, up_b, npack, width, dims); + defer gpa.free(packed_u); + const packed_d = try packDownload(gpa, driver, stream.*, pack_fn, down_w, down_s, down_b, npack, dims, width); + defer gpa.free(packed_d); + const stacked = try stackHalves(gpa, packed_g, packed_u); + defer gpa.free(stacked); + + var up_wb = try cuda.DeviceBuffer.fromHost(driver, stacked); + defer up_wb.free(); + var down_wb = try cuda.DeviceBuffer.fromHost(driver, packed_d); + defer down_wb.free(); + var local_pick: [slots]i32 = undefined; + for (0..slots) |s| local_pick[s] = local[s]; + var picks_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&local_pick)); + defer picks_b.free(); + var members_b = try cuda.DeviceBuffer.alloc(driver, slots * 4); + defer members_b.free(); + var items_b = try cuda.DeviceBuffer.alloc(driver, slots * 3 * 4); + defer items_b.free(); + var counts_b = try cuda.DeviceBuffer.alloc(driver, 8); + defer counts_b.free(); + var act_b = try cuda.DeviceBuffer.alloc(driver, slots * width * 2); + defer act_b.free(); + var y_b = try cuda.DeviceBuffer.alloc(driver, slots * dims * 4); + defer y_b.free(); + const plan_fn = try exp_mod.function(plan_symbol); + var plan_a: cuda.Args = .{}; + plan_a.add(picks_b.ptr); + plan_a.add(cint(slots)); + plan_a.add(cint(npack)); + plan_a.add(cint(tile)); + plan_a.add(members_b.ptr); + plan_a.add(items_b.ptr); + plan_a.add(counts_b.ptr); + try cuda.launch.launch(plan_fn, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1024 } }, stream.*, &plan_a); + + const up_fn = try exp_mod.function(up_symbol); + const down_fn = try exp_mod.function(down_symbol); + try launchExpert(up_fn, stream.*, mixed, dims, slots, up_wb.ptr, dims / 32, width / 32, items_b.ptr, counts_b.ptr, members_b.ptr, act_b.ptr, width, slots * (width / 32)); + try launchExpert(down_fn, stream.*, act_b.ptr, width, 0, down_wb.ptr, width / 32, dims / 32, items_b.ptr, counts_b.ptr, members_b.ptr, y_b.ptr, dims, slots * (dims / 32)); + + const act_raw = try gpa.alloc(u8, slots * width * 2); + defer gpa.free(act_raw); + const y_raw = try gpa.alloc(u8, slots * dims * 4); + defer gpa.free(y_raw); + try stream.synchronize(); + try act_b.download(0, act_raw); + try y_b.download(0, y_raw); + + var act_off: usize = 0; + var act_steps: u32 = 0; + var y_off: usize = 0; + var y_ulp: u32 = 0; + var y_steps: u32 = 0; + var head_g: u32 = 0; + var head_h: u32 = 0; + const gate_col = try gpa.alloc(f32, width); + defer gpa.free(gate_col); + const up_col = try gpa.alloc(f32, width); + defer gpa.free(up_col); + const act_row = try gpa.alloc(u16, width); + defer gpa.free(act_row); + const down_col = try gpa.alloc(f32, dims); + defer gpa.free(down_col); + const gw = width * (dims / 8) * 4; + const gs = width * (dims / 32) * 2; + const dw = dims * (width / 8) * 4; + const ds = dims * (width / 32) * 2; + for (0..slots) |s| { + const src: usize = @intCast(local[s]); + project(row, gathered_w[src * gw ..][0..gw], gathered_s[src * gs ..][0..gs], gathered_b[src * gs ..][0..gs], width, dims, gate_col); + project(row, up_w[src * gw ..][0..gw], up_s[src * gs ..][0..gs], up_b[src * gs ..][0..gs], width, dims, up_col); + for (act_row, gate_col, up_col) |*o, gv, uv| o.* = swiglu(gv, uv); + for (0..width) |c| { + const got = std.mem.readInt(u16, act_raw[(s * width + c) * 2 ..][0..2], .little); + const dist = hc.mixedSteps(got, act_row[c]); + if (dist != 0) act_off += 1; + if (dist > act_steps) act_steps = dist; + } + project(act_row, down_w[src * dw ..][0..dw], down_s[src * ds ..][0..ds], down_b[src * ds ..][0..ds], dims, width, down_col); + for (0..dims) |c| { + const got: f32 = @bitCast(std.mem.readInt(u32, y_raw[(s * dims + c) * 4 ..][0..4], .little)); + const dist = hc.ulps(got, down_col[c]); + const steps = hc.mixedSteps(toBf16(got), toBf16(down_col[c])); + if (dist != 0) y_off += 1; + if (dist > y_ulp) y_ulp = dist; + if (steps > y_steps) y_steps = steps; + if (s == 0 and c == 0) { + head_g = @bitCast(got); + head_h = @bitCast(down_col[c]); + } + } + } + std.debug.print("experts slots {d} act_off {d} act_steps {d} y_off {d} y_ulp {d} y_steps {d} head {x:0>8} host {x:0>8}\n", .{ slots, act_off, act_steps, y_off, y_ulp, y_steps, head_g, head_h }); + if (pick_off != 0 or host_pick_off != 0 or wts_same.ulp != 0 or host_wts_steps > 1 or act_steps != 0 or y_steps != 0) return 1; + return 0; +} + +fn launchExpert(f: cuda.Function, stream: cuda.Stream, x: u64, x_stride: usize, slot_n: usize, w: u64, kg: usize, nb: usize, items: u64, counts: u64, members: u64, out: u64, n: usize, units: usize) !void { + var a: cuda.Args = .{}; + a.add(x); + a.add(cint(x_stride)); + a.add(cint(slot_n)); + a.add(w); + a.add(cint(kg)); + a.add(cint(nb)); + a.add(items); + a.add(counts); + a.add(members); + a.add(out); + a.add(cint(n)); + a.add(@as(f32, 0)); + const grid = (units + 3) / 4; + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(grid) }, .block = .{ .x = 128 } }, stream, &a); +} + +test "a tied logit keeps the lower expert id and the shared gate rounds to one half" { + const logits = [_]f32{ 1, 3, 3, 0 }; + var picks: [3]i32 = undefined; + var wts: [3]f32 = undefined; + choose(&logits, 3, 2, &picks, &wts); + try std.testing.expectEqual([3]i32{ 1, 2, 3 }, picks); + try std.testing.expectEqual(@as(f32, 0.5), wts[0]); + try std.testing.expectEqual(@as(f32, 0.5), wts[1]); + try std.testing.expectEqual(@as(f32, 0.5), wts[2]); +} + +test "each lane of the quad contributes to the group sum" { + var x: [32]u16 = @splat(0); + x[0] = 0x3f80; + x[8] = 0x3f80; + x[16] = 0x3f80; + x[24] = 0x3f80; + try std.testing.expectEqual(@as(f32, 4), groupSum(&x, 0)); +} + +test "scale times the nibble dot plus the group bias is five" { + var words: [16]u8 = @splat(0); + std.mem.writeInt(u32, words[0..4], 0x21, .little); + var scale: [2]u8 = undefined; + std.mem.writeInt(u16, &scale, 0x3f80, .little); + var bias: [2]u8 = undefined; + std.mem.writeInt(u16, &bias, 0x3f80, .little); + var x: [32]u16 = @splat(0); + x[0] = 0x3f80; + x[1] = 0x3f80; + var out: [1]f32 = undefined; + project(&x, &words, &scale, &bias, 1, 32, &out); + try std.testing.expectEqual(@as(f32, 5), out[0]); +} diff --git a/zig/src/families/flashnext/cuda_ple.zig b/zig/src/families/flashnext/cuda_ple.zig new file mode 100644 index 0000000000..751f8a8c75 --- /dev/null +++ b/zig/src/families/flashnext/cuda_ple.zig @@ -0,0 +1,589 @@ +//! Layer 2's n-gram table for one decode token: hash the row ids, gather those rows, dequant like glue._ple_embed. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const embed = @import("cuda_embed.zig"); +const hc = @import("cuda_hc.zig"); + +const golden: u64 = 0x9E3779B97F4A7C15; +const mix1: u64 = 0xBF58476D1CE4E5B9; +const mix2: u64 = 0x94D049BB133111EB; +const prime_step: u64 = 10007; +const max_n: usize = 8; +const max_heads: usize = 64; +const max_shards: usize = 128; +const layer: usize = 1; + +const Tensor = core.safetensors.Tensor; + +const Hash = struct { + n: usize, + per: usize, + heads: usize, + eos: i64, + dh: usize, + shards: usize, + sizes: [max_heads]i64 = undefined, + offsets: [max_heads]i64 = undefined, + mult: [max_n]i64 = undefined, + starts: [max_shards + 1]i64 = undefined, + + fn init(args: struct { + vocab: u64, + ngram: usize, + per: usize, + base: u64, + divisor: u64, + shards: usize, + seed: u64, + eos: i64, + ple_dim: usize, + ple_index: usize, + }) !Hash { + if (args.ngram < 2 or args.ngram > max_n or args.per == 0 or args.shards == 0 or args.shards > max_shards) return error.UnexpectedTensor; + if (args.base < 2 or args.divisor == 0) return error.UnexpectedTensor; + const heads = (args.ngram - 1) * args.per; + if (heads > max_heads or args.ple_dim == 0 or args.ple_dim % heads != 0) return error.UnexpectedTensor; + var h = Hash{ .n = args.ngram, .per = args.per, .heads = heads, .eos = args.eos, .dh = args.ple_dim / heads, .shards = args.shards }; + var total: i64 = 0; + for (0..heads) |head| { + const size: i64 = @intCast(nthPrimeAfter(args.base - 1, args.ple_index * heads + head + 1)); + h.sizes[head] = size; + h.offsets[head] = total; + total += size; + } + const divisor: i64 = @intCast(args.divisor); + const shard_n: i64 = @intCast(args.shards); + const padded = divCeil(total, divisor) * divisor; + const base: i64 = @divFloor(padded, shard_n); + const extra: usize = @intCast(@mod(padded, shard_n)); + var acc: i64 = 0; + for (0..args.shards) |i| { + h.starts[i] = acc; + acc += base + @as(i64, @intFromBool(i < extra)); + } + h.starts[args.shards] = acc; + const half: u64 = @max(1, (((1 << 63) - 1) / @max(args.vocab, 1)) / 2); + const seed = args.seed +% prime_step *% @as(u64, @intCast(args.ple_index)); + for (0..args.ngram) |i| { + const mixed = splitmix64(seed +% golden *% @as(u64, @intCast(i + 1))); + h.mult[i] = @intCast(2 * (mixed % half) + 1); + } + return h; + } + + fn locate(h: Hash, row: i64) !struct { shard: usize, local: usize } { + if (row < 0 or row >= h.starts[h.shards]) return error.UnexpectedTensor; + var lo: usize = 0; + var hi = h.shards; + while (lo + 1 < hi) { + const mid = (lo + hi) / 2; + if (h.starts[mid] <= row) lo = mid else hi = mid; + } + return .{ .shard = lo, .local = @intCast(row - h.starts[lo]) }; + } +}; + +fn splitmix64(value: u64) u64 { + var v = value +% golden; + v = (v ^ (v >> 30)) *% mix1; + v = (v ^ (v >> 27)) *% mix2; + return v ^ (v >> 31); +} + +fn isPrime(value: u64) bool { + if (value < 2) return false; + if (value % 2 == 0) return value == 2; + var d: u64 = 3; + while (d <= value / d) : (d += 2) if (value % d == 0) return false; + return true; +} + +fn nthPrimeAfter(start: u64, count: usize) u64 { + var prime = start; + for (0..count) |_| { + prime += 1; + while (!isPrime(prime)) prime += 1; + } + return prime; +} + +fn divCeil(a: i64, b: i64) i64 { + return @divFloor(a + b - 1, b); +} + +/// Row ids of one new token after `history` (n - 1 ids). A fresh decode window is n - 1 copies of eos. +fn lastIds(h: Hash, history: []const i64, token: i64, out: []i64) !void { + if (history.len != h.n - 1 or out.len != h.heads) return error.UnexpectedTensor; + var seq: [max_n]i64 = undefined; + for (history, 0..) |v, i| seq[i] = v; + seq[history.len] = token; + const width = h.n; + var eos_at: [max_n]i64 = undefined; + for (0..width) |i| eos_at[i] = if (seq[i] == h.eos) @intCast(i) else -1; + var acc: [max_n]i64 = undefined; + acc[0] = eos_at[0]; + for (1..width) |i| acc[i] = @max(acc[i - 1], eos_at[i]); + var before: [max_n]i64 = undefined; + before[0] = -1; + for (1..width) |i| before[i] = acc[i - 1]; + var shifted: [max_n]i64 = undefined; + const pos: i64 = @intCast(width - 1); + const in_segment = pos - (before[width - 1] + 1); + for (0..h.n) |shift| { + const source = pos - @as(i64, @intCast(shift)); + const taken = seq[@intCast(@max(source, 0))]; + const keep = in_segment >= @as(i64, @intCast(shift)) and source >= 0; + shifted[shift] = if (keep) taken else h.eos; + } + var at: usize = 0; + var ngram: usize = 2; + while (ngram <= h.n) : (ngram += 1) { + var mixed = shifted[0] *% h.mult[0]; + for (1..ngram) |p| mixed = mixed ^ (shifted[p] *% h.mult[p]); + for (0..h.per) |_| { + out[at] = @mod(mixed, h.sizes[at]) + h.offsets[at]; + at += 1; + } + } +} + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn sumGroups(row: []const u16, out: []f32) void { + const groups = row.len / 32; + for (0..groups) |g| { + var buf: [32]f32 = undefined; + for (0..32) |i| buf[i] = @bitCast(@as(u32, row[g * 32 + i]) << 16); + var n: usize = 32; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + out[g] = buf[0]; + } +} + +const Open = struct { + io: std.Io, + path: []u8, + file: std.Io.File, + map: std.Io.File.MemoryMap, + parsed: std.json.Parsed(std.json.Value), + at: usize, + + fn load(gpa: std.mem.Allocator, io: std.Io, path: []const u8) !Open { + const owned = try gpa.dupe(u8, path); + errdefer gpa.free(owned); + var file = try std.Io.Dir.cwd().openFile(io, owned, .{}); + errdefer file.close(io); + const len: usize = @intCast(try file.length(io)); + if (len < 8) return error.BadSafetensors; + var map = try std.Io.File.MemoryMap.create(io, file, .{ .len = len, .protection = .{ .read = true, .write = false }, .populate = false }); + errdefer map.destroy(io); + const header_len: usize = @intCast(std.mem.readInt(u64, map.memory[0..8], .little)); + if (header_len > len - 8) return error.BadSafetensors; + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, map.memory[8..][0..header_len], .{}); + return .{ .io = io, .path = owned, .file = file, .map = map, .parsed = parsed, .at = 8 + header_len }; + } + + fn close(self: *Open, gpa: std.mem.Allocator) void { + self.parsed.deinit(); + self.map.destroy(self.io); + self.file.close(self.io); + gpa.free(self.path); + self.* = undefined; + } + + fn tensor(self: *const Open, name: []const u8, dtype: core.safetensors.DType) !Tensor { + const o = (self.parsed.value.object.get(name) orelse return error.MissingTensor).object; + const got = core.safetensors.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; + if (got != dtype) return error.UnexpectedTensor; + const shape = o.get("shape").?.array.items; + if (shape.len == 0 or shape.len > 4) return error.UnexpectedTensor; + const offs = o.get("data_offsets").?.array.items; + const begin: usize = @intCast(offs[0].integer); + const end: usize = @intCast(offs[1].integer); + const data = self.map.memory[self.at..]; + if (end < begin or end > data.len) return error.BadSafetensors; + var t: Tensor = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .bytes = data[begin..end] }; + for (shape, 0..) |d, i| t.shape[i] = @intCast(d.integer); + return t; + } +}; + +fn readI64(t: Tensor, out: []i64) !void { + if (t.dtype != .i64 or t.rank != 1 or t.dim(0) != out.len or t.bytes.len != out.len * 8) return error.UnexpectedTensor; + for (out, 0..) |*o, i| o.* = std.mem.readInt(i64, t.bytes[8 * i ..][0..8], .little); +} + +fn weightName(buf: []u8, shard: usize, nested: bool) ![]const u8 { + return if (nested) + std.fmt.bufPrint(buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.shards.{d}.weight", .{ layer, shard }) + else + std.fmt.bufPrint(buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.shard_{d}.weight", .{ layer, shard }); +} + +fn partName(buf: []u8, shard: usize, nested: bool, part: []const u8) ![]const u8 { + return if (nested) + std.fmt.bufPrint(buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.shards.{d}.{s}", .{ layer, shard, part }) + else + std.fmt.bufPrint(buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.shard_{d}.{s}", .{ layer, shard, part }); +} + +fn findOpen(opened: []Open, path: []const u8) ?*Open { + for (opened) |*o| if (std.mem.eql(u8, o.path, path)) return o; + return null; +} + +fn copyRow(file: *const Open, shard: usize, nested: bool, local: usize, dh: usize, words: []u8, scales: []u8, biases: []u8) !void { + var name: [180]u8 = undefined; + const w = try file.tensor(try partName(&name, shard, nested, "weight"), .u32); + const s = try file.tensor(try partName(&name, shard, nested, "scales"), .bf16); + const b = try file.tensor(try partName(&name, shard, nested, "biases"), .bf16); + const words_w = dh / 8; + const groups = dh / 32; + if (w.dim(1) != words_w or s.dim(1) != groups or b.dim(1) != groups or local >= w.dim(0)) return error.UnexpectedTensor; + const wb = words_w * 4; + const gb = groups * 2; + @memcpy(words, w.bytes[local * wb ..][0..wb]); + @memcpy(scales, s.bytes[local * gb ..][0..gb]); + @memcpy(biases, b.bytes[local * gb ..][0..gb]); +} + +fn tableScale(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, wm: std.json.ObjectMap) !f32 { + const name = "language_model.model.layers.1.ple.ple_embedding.ngram_embedding.weight_scale"; + const file = wm.get(name) orelse return 1; + const path = try std.fs.path.join(gpa, &.{ dir, file.string }); + defer gpa.free(path); + var opened = try Open.load(gpa, io, path); + defer opened.close(gpa); + const t = try opened.tensor(name, blk: { + const o = opened.parsed.value.object.get(name).?.object; + break :blk core.safetensors.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; + }); + if (t.numel() != 1) return error.UnexpectedTensor; + return switch (t.dtype) { + .f32 => @bitCast(std.mem.readInt(u32, t.bytes[0..4], .little)), + .f64 => @floatCast(@as(f64, @bitCast(std.mem.readInt(u64, t.bytes[0..8], .little)))), + .bf16 => @bitCast(@as(u32, std.mem.readInt(u16, t.bytes[0..2], .little)) << 16), + else => error.UnexpectedTensor, + }; +} + +fn jsonU64(v: std.json.Value) !u64 { + if (v != .integer or v.integer < 0) return error.UnexpectedTensor; + return @intCast(v.integer); +} + +fn optU64(o: std.json.ObjectMap, name: []const u8, default: u64) !u64 { + const v = o.get(name) orelse return default; + if (v == .null) return default; + return jsonU64(v); +} + +const NgramSpec = struct { + vocab: u64, + ngram: usize, + per: usize, + base: u64, + divisor: u64, + shards: usize, + seed: u64, + eos: i64, + ple_dim: usize, + ple_index: usize, +}; + +/// `ple_layer_ids` value 2 is decoder layer 1. The eos id is the text config's, which the hash mixes in. +fn readSpec(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !NgramSpec { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 24)); + defer gpa.free(bytes); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer parsed.deinit(); + const root = parsed.value.object; + const text = if (root.get("text_config")) |v| if (v == .object and v.object.count() > 0) v.object else root else root; + const layers = try jsonU64(text.get("num_hidden_layers") orelse return error.MissingTensor); + if (layers <= layer) return error.UnexpectedTensor; + const hidden = try jsonU64(text.get("hidden_size") orelse return error.MissingTensor); + const ids = text.get("ple_layer_ids") orelse return error.MissingTensor; + if (ids != .array) return error.UnexpectedTensor; + var seen: [512]bool = @splat(false); + var ple_index: usize = 0; + var found = false; + for (ids.array.items) |item| { + const id = try jsonU64(item); + if (id == 0 or id > layers or id > seen.len) return error.UnexpectedTensor; + if (id == layer + 1) found = true; + if (!seen[id - 1] and id - 1 < layer) ple_index += 1; + seen[id - 1] = true; + } + if (!found) return error.UnexpectedTensor; + const eos: i64 = if (text.get("eos_token_id")) |v| blk: { + if (v == .null) break :blk 0; + if (v == .array) { + if (v.array.items.len == 0) return error.UnexpectedTensor; + break :blk @intCast(try jsonU64(v.array.items[0])); + } + break :blk @intCast(try jsonU64(v)); + } else 0; + return .{ + .vocab = try jsonU64(text.get("vocab_size") orelse return error.MissingTensor), + .ngram = @intCast(try optU64(text, "ngram_size", 3)), + .per = @intCast(try optU64(text, "heads_per_ngram", 8)), + .base = try optU64(text, "ngram_vocab_size_base", 20_000_000), + .divisor = try optU64(text, "make_ngram_vocab_size_divisible_by", 128), + .shards = @intCast(try optU64(text, "split_ngram_parts", 128)), + .seed = try optU64(text, "seed", 1234), + .eos = eos, + .ple_dim = @intCast(try optU64(text, "ple_embed_dim", hidden)), + .ple_index = ple_index, + }; +} + +/// The text config's eos id, which a fresh n-gram window repeats. +pub fn eosOf(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8) !i64 { + const cfg = try readSpec(gpa, io, model_dir); + return cfg.eos; +} + +/// Host dequant of one token's sixteen n-gram rows. `history` is the previous n - 1 ids. +pub fn embedding(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, history: []const i64, token: u32, out: []u16, xs: []f32) !void { + const cfg = try readSpec(gpa, io, model_dir); + const hash = try Hash.init(.{ + .vocab = cfg.vocab, + .ngram = cfg.ngram, + .per = cfg.per, + .base = cfg.base, + .divisor = cfg.divisor, + .shards = cfg.shards, + .seed = cfg.seed, + .eos = cfg.eos, + .ple_dim = cfg.ple_dim, + .ple_index = cfg.ple_index, + }); + if (history.len != hash.n - 1 or out.len != hash.heads * hash.dh or xs.len != out.len / 32) return error.UnexpectedTensor; + var ids: [max_heads]i64 = undefined; + try lastIds(hash, history, @intCast(token), ids[0..hash.heads]); + const index_path = try std.fs.path.join(gpa, &.{ model_dir, "model.safetensors.index.json" }); + defer gpa.free(index_path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index_path, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const wm = parsed.value.object.get("weight_map").?.object; + const scale = try tableScale(gpa, io, model_dir, wm); + const heads = hash.heads; + const dh = hash.dh; + const words = try gpa.alloc(u8, heads * (dh / 8) * 4); + defer gpa.free(words); + const scales = try gpa.alloc(u8, heads * (dh / 32) * 2); + defer gpa.free(scales); + const biases = try gpa.alloc(u8, scales.len); + defer gpa.free(biases); + var opened: [max_heads]Open = undefined; + var nopen: usize = 0; + defer for (opened[0..nopen]) |*o| o.close(gpa); + var key: [180]u8 = undefined; + for (0..heads) |h| { + const loc = try hash.locate(ids[h]); + var nested = false; + const mapped = if (wm.get(try weightName(&key, loc.shard, false))) |v| + v.string + else blk: { + nested = true; + break :blk (wm.get(try weightName(&key, loc.shard, true)) orelse return error.MissingTensor).string; + }; + const path = try std.fs.path.join(gpa, &.{ model_dir, mapped }); + defer gpa.free(path); + const file = findOpen(opened[0..nopen], path) orelse file: { + opened[nopen] = try Open.load(gpa, io, path); + nopen += 1; + break :file &opened[nopen - 1]; + }; + const wb = (dh / 8) * 4; + const gb = (dh / 32) * 2; + try copyRow(file, loc.shard, nested, loc.local, dh, words[h * wb ..][0..wb], scales[h * gb ..][0..gb], biases[h * gb ..][0..gb]); + } + for (0..heads) |h| try embed.dequant(words, scales, biases, dh, h, out[h * dh ..][0..dh]); + if (scale != 1) { + for (out) |*v| { + const f: f32 = @bitCast(@as(u32, v.*) << 16); + v.* = toBf16(f * scale); + } + } + sumGroups(out, xs); +} + +/// One decode token. `ple_layer_ids` `[2]` is this layer. Returns 0 when the gathered rows match the host dequant. +pub fn table(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, token: u32) !u32 { + const cfg = try readSpec(gpa, io, model_dir); + const hash = try Hash.init(.{ + .vocab = cfg.vocab, + .ngram = cfg.ngram, + .per = cfg.per, + .base = cfg.base, + .divisor = cfg.divisor, + .shards = cfg.shards, + .seed = cfg.seed, + .eos = cfg.eos, + .ple_dim = cfg.ple_dim, + .ple_index = cfg.ple_index, + }); + var history: [max_n]i64 = undefined; + for (0..hash.n - 1) |i| history[i] = hash.eos; + var ids: [max_heads]i64 = undefined; + try lastIds(hash, history[0 .. hash.n - 1], token, ids[0..hash.heads]); + + const index_path = try std.fs.path.join(gpa, &.{ model_dir, "model.safetensors.index.json" }); + defer gpa.free(index_path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index_path, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const wm = parsed.value.object.get("weight_map").?.object; + + var cname: [160]u8 = undefined; + const mult_name = try std.fmt.bufPrint(&cname, "language_model.model.layers.{d}.ple.ple_embedding.layer_multipliers", .{layer}); + const const_file = wm.get(mult_name) orelse return error.MissingTensor; + const const_path = try std.fs.path.join(gpa, &.{ model_dir, const_file.string }); + defer gpa.free(const_path); + var constants = try Open.load(gpa, io, const_path); + defer constants.close(gpa); + var got_m: [max_n]i64 = undefined; + var got_o: [max_heads]i64 = undefined; + var got_s: [max_heads]i64 = undefined; + try readI64(try constants.tensor(mult_name, .i64), got_m[0..hash.n]); + const off_name = try std.fmt.bufPrint(&cname, "language_model.model.layers.{d}.ple.ple_embedding.ngram_heads_offsets", .{layer}); + try readI64(try constants.tensor(off_name, .i64), got_o[0..hash.heads]); + const size_name = try std.fmt.bufPrint(&cname, "language_model.model.layers.{d}.ple.ple_embedding.ngram_heads_vocab_sizes", .{layer}); + try readI64(try constants.tensor(size_name, .i64), got_s[0..hash.heads]); + if (!std.mem.eql(i64, got_m[0..hash.n], hash.mult[0..hash.n]) or !std.mem.eql(i64, got_o[0..hash.heads], hash.offsets[0..hash.heads]) or !std.mem.eql(i64, got_s[0..hash.heads], hash.sizes[0..hash.heads])) { + std.debug.print("ngram constants differ\n", .{}); + return 1; + } + const scale = try tableScale(gpa, io, model_dir, wm); + + const heads = hash.heads; + const dh = hash.dh; + const words = try gpa.alloc(u8, heads * (dh / 8) * 4); + defer gpa.free(words); + const scales = try gpa.alloc(u8, heads * (dh / 32) * 2); + defer gpa.free(scales); + const biases = try gpa.alloc(u8, scales.len); + defer gpa.free(biases); + var opened: [max_heads]Open = undefined; + var nopen: usize = 0; + defer for (opened[0..nopen]) |*o| o.close(gpa); + var key: [180]u8 = undefined; + for (0..heads) |h| { + const loc = try hash.locate(ids[h]); + var nested = false; + const mapped = if (wm.get(try weightName(&key, loc.shard, false))) |v| + v.string + else blk: { + nested = true; + break :blk (wm.get(try weightName(&key, loc.shard, true)) orelse return error.MissingTensor).string; + }; + const path = try std.fs.path.join(gpa, &.{ model_dir, mapped }); + defer gpa.free(path); + const file = findOpen(opened[0..nopen], path) orelse file: { + opened[nopen] = try Open.load(gpa, io, path); + nopen += 1; + break :file &opened[nopen - 1]; + }; + const wb = (dh / 8) * 4; + const gb = (dh / 32) * 2; + try copyRow(file, loc.shard, nested, loc.local, dh, words[h * wb ..][0..wb], scales[h * gb ..][0..gb], biases[h * gb ..][0..gb]); + } + + const host = try gpa.alloc(u16, heads * dh); + defer gpa.free(host); + for (0..heads) |h| try embed.dequant(words, scales, biases, dh, h, host[h * dh ..][0..dh]); + if (scale != 1) { + for (host) |*v| { + const f: f32 = @bitCast(@as(u32, v.*) << 16); + v.* = toBf16(f * scale); + } + } + const host_xs = try gpa.alloc(f32, heads * dh / 32); + defer gpa.free(host_xs); + sumGroups(host, host_xs); + + var w_b = try cuda.DeviceBuffer.fromHost(driver, words); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, biases); + defer b_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, host.len * 2); + defer out_b.free(); + var xs_b = try cuda.DeviceBuffer.alloc(driver, host_xs.len * 4); + defer xs_b.free(); + try tri.pleEmbed(w_b.ptr, s_b.ptr, b_b.ptr, out_b.ptr, xs_b.ptr, 1, heads, dh, scale); + const got_bytes = try gpa.alloc(u8, host.len * 2); + defer gpa.free(got_bytes); + const xs_bytes = try gpa.alloc(u8, host_xs.len * 4); + defer gpa.free(xs_bytes); + try stream.synchronize(); + try out_b.download(0, got_bytes); + try xs_b.download(0, xs_bytes); + const gpu = try gpa.alloc(u16, host.len); + defer gpa.free(gpu); + for (gpu, 0..) |*o, i| o.* = std.mem.readInt(u16, got_bytes[2 * i ..][0..2], .little); + var off: usize = 0; + var steps: u32 = 0; + var at: usize = 0; + for (gpu, host, 0..) |a, b, i| { + const d = hc.mixedSteps(a, b); + if (d != 0) off += 1; + if (d > steps) { + steps = d; + at = i; + } + } + var xs_ulp: u32 = 0; + var xs_at: usize = 0; + for (host_xs, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const dist = hc.ulps(got, want); + if (dist > xs_ulp) { + xs_ulp = dist; + xs_at = i; + } + } + std.debug.print("ngram token {d} heads {d} values {d} off {d} max_steps {d} at {d} got {x:0>4} host {x:0>4} xs_ulp {d} xs_at {d} head {x:0>4}\n", .{ token, heads, host.len, off, steps, at, gpu[at], host[at], xs_ulp, xs_at, gpu[0] }); + if (off != 0 or steps != 0 or xs_ulp != 0) return 1; + return 0; +} + +test "a fresh window hashes token 7 into sixteen rows" { + const h = try Hash.init(.{ + .vocab = 248320, + .ngram = 3, + .per = 8, + .base = 20_000_000, + .divisor = 128, + .shards = 128, + .seed = 1234, + .eos = 248044, + .ple_dim = 2560, + .ple_index = 0, + }); + try std.testing.expectEqual(@as(i64, 23703573157769), h.mult[0]); + const history = [_]i64{ 248044, 248044 }; + var ids: [16]i64 = undefined; + try lastIds(h, &history, 7, &ids); + const want = [_]i64{ 2927653, 34980843, 54748278, 66612378, 97814964, 109013870, 126560393, 151352333, 167888935, 182235580, 215170017, 237467519, 247510681, 278779700, 296141806, 304994522 }; + try std.testing.expectEqualSlices(i64, &want, &ids); + const loc = try h.locate(ids[0]); + try std.testing.expectEqual(@as(usize, 1), loc.shard); + try std.testing.expectEqual(@as(usize, 427641), loc.local); +} diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig new file mode 100644 index 0000000000..5dad8c369b --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -0,0 +1,403 @@ +//! One prompt, drafts off: each token is one fused decode row, layers in order, then the argmax. + +const std = @import("std"); +const cuda = @import("cuda"); +const embed = @import("cuda_embed.zig"); +const qmm = @import("cuda_qmm.zig"); +const gdn = @import("cuda_gdn.zig"); +const ple = @import("cuda_ple.zig"); +const weights = @import("cuda_prompt_wt.zig"); +const step = @import("cuda_prompt_step.zig"); +const load = @import("cuda_prompt_load.zig"); + +const layers_n = weights.layers_n; +const linear_n = weights.linear_n; +const attn_n = weights.attn_n; +const dims = weights.dims; +const streams = weights.streams; +const wide = weights.wide; +const vocab = weights.vocab; +const experts_n = weights.experts_n; +const slots = weights.slots; +const moe_w = weights.moe_w; +const head_dim = weights.head_dim; +const index_dim = weights.index_dim; +const proj_n = weights.proj_n; +const out_k = weights.out_k; +const ple_layer = weights.ple_layer; +const inj_stride = weights.inj_stride; +const wts_stride = weights.wts_stride; +const ple_tail_n = weights.ple_tail_n; +const state_n = weights.state_n; +const kv_heads = weights.kv_heads; +const Store = weights.Store; +const Face = weights.Face; +const Q4 = weights.Q4; +const HcW = weights.HcW; +const linearLayer = weights.linearLayer; +const linearIndex = weights.linearIndex; +const attnIndex = weights.attnIndex; +const promote = weights.promote; +const asF32 = weights.asF32; +const uploadQ4 = weights.uploadQ4; +const join3 = weights.join3; +const loadHc = weights.loadHc; +const stitch = weights.stitch; +const Scratch = step.Scratch; +const writeback = step.writeback; +const readout = step.readout; +const gdnStep = step.gdnStep; +const attnStep = step.attnStep; +const moeStep = step.moeStep; +const pleStep = step.pleStep; +const Layer = load.Layer; +const loadLayer = load.loadLayer; + +pub fn argmax(raw: []const u8) u32 { + var best = promote(std.mem.readInt(u16, raw[0..2], .little)); + var at: usize = 0; + var i: usize = 1; + while (i < vocab) : (i += 1) { + const v = promote(std.mem.readInt(u16, raw[2 * i ..][0..2], .little)); + if (v > best) { + best = v; + at = i; + } + } + return @intCast(at); +} + +const Pass = struct { + gpa: std.mem.Allocator, + io: std.Io, + driver: *cuda.Driver, + stream: *cuda.Stream, + store: *Store, + model_dir: []const u8, + s: *Scratch, + table: embed.Table, + rows: usize, + mode: []u8, + history: *[2]i64, + conv: []cuda.DeviceBuffer, + state: []cuda.DeviceBuffer, + kc: []cuda.DeviceBuffer, + vc: []cuda.DeviceBuffer, + ikc: []cuda.DeviceBuffer, + tail: cuda.DeviceBuffer, + mixer: HcW, + head_q: Q4, + logits: []u8, + plan_fn: cuda.Function, + up_fn: cuda.Function, + down_fn: cuda.Function, +}; + +fn embedRow(comptime Tri: type, p: *Pass, tri: Tri, token: u32, h: u64) !void { + const id: i32 = @intCast(token); + try p.s.ids.upload(0, std.mem.asBytes(&id)); + try tri.embed(p.s.ids.ptr, p.table.w.ptr, p.table.s.ptr, p.table.b.ptr, h, dims, streams, 1); +} + +fn runLayer(comptime Tri: type, p: *Pass, tri: Tri, lw: *Layer, layer: usize, ids: []const u32, pos0: usize) !void { + for (ids, 0..) |token, t| { + const h = p.s.h.ptr + @as(u64, t) * wide * 2; + const inj = p.s.inj.ptr + @as(u64, t) * inj_stride; + const y = p.s.y.ptr + @as(u64, t) * slots * dims * 4; + const wt = p.s.wts.ptr + @as(u64, t) * wts_stride; + if (layer == ple_layer and p.mode[t] == 2) { + try writeback(tri, p.s, h, 2, inj, y, wt); + p.mode[t] = 0; + } + if (layer == ple_layer) try pleStep(Tri, p.gpa, p.io, p.driver, p.stream, tri, p.s, p.model_dir, p.history, token, h, lw.key, lw.val, lw.nk.ptr, lw.nq.ptr, lw.nc.ptr, lw.ple_cw.ptr, &p.tail); + if (p.mode[t] == 2) { + try writeback(tri, p.s, h, 2, inj, y, wt); + p.mode[t] = 0; + } else try writeback(tri, p.s, h, 0, inj, y, wt); + try readout(tri, p.s, lw.attn_hc, h, inj); + if (lw.linear) { + const li = linearIndex(layer); + try gdnStep(p.gpa, p.driver, p.stream, p.s, lw.proj, lw.out, p.conv[li], lw.cw.ptr, lw.a_log.ptr, lw.dt.ptr, lw.norm.ptr, &p.state[li]); + } else { + const ai = attnIndex(layer); + try attnStep(Tri, p.driver, p.stream, tri, p.s, lw.proj, lw.out, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, p.kc[ai].ptr, p.vc[ai].ptr, p.ikc[ai].ptr, @intCast(pos0 + t)); + } + try writeback(tri, p.s, h, 4, inj, y, wt); + try readout(tri, p.s, lw.mlp_hc, h, inj); + try moeStep(Tri, tri, p.s, p.stream, p.plan_fn, p.up_fn, p.down_fn, lw.router.ptr, lw.experts_up.ptr, lw.experts_down.ptr, y, wt); + p.mode[t] = 2; + } +} + +fn zeros(gpa: std.mem.Allocator, driver: *cuda.Driver, n: usize, len: usize) ![]cuda.DeviceBuffer { + const out = try gpa.alloc(cuda.DeviceBuffer, n); + var i: usize = 0; + errdefer { + for (out[0..i]) |*b| b.free(); + gpa.free(out); + } + while (i < n) : (i += 1) { + out[i] = try cuda.DeviceBuffer.alloc(driver, len); + try out[i].fill8(0, null); + } + return out; +} + +fn head(comptime Tri: type, p: *Pass, tri: Tri, row: usize) !void { + if (p.mode[row] != 2) return error.UnexpectedTensor; + const h = p.s.h.ptr + @as(u64, row) * wide * 2; + const inj = p.s.inj.ptr + @as(u64, row) * inj_stride; + const y = p.s.y.ptr + @as(u64, row) * slots * dims * 4; + const wt = p.s.wts.ptr + @as(u64, row) * wts_stride; + try writeback(tri, p.s, h, 2, inj, y, wt); + p.mode[row] = 0; + try readout(tri, p.s, p.mixer, h, inj); + try qmm.matmul(p.driver, p.stream.*, p.s.mixed.ptr, p.s.xsm.ptr, p.head_q.w.ptr, p.head_q.s.ptr, p.head_q.b.ptr, p.s.logits.ptr, 1, vocab, dims); + try p.stream.synchronize(); + try p.s.logits.download(0, p.logits); +} + +fn forward(comptime Tri: type, p: *Pass, tri: Tri, layers: []Layer, ids: []const u32, pos0: usize) !void { + @memset(p.mode, 0); + for (ids, 0..) |token, t| try embedRow(Tri, p, tri, token, p.s.h.ptr + @as(u64, t) * wide * 2); + for (layers, 0..) |*lw, layer| try runLayer(Tri, p, tri, lw, layer, ids, pos0); + try head(Tri, p, tri, ids.len - 1); +} + +pub const vocabulary: usize = vocab; + +/// One loaded checkpoint: the embedding table, the mixer, the head, and a scratch row the forwards share. +pub const Engine = struct { + gpa: std.mem.Allocator, + io: std.Io, + driver: *cuda.Driver, + stream: *cuda.Stream, + store: Store, + model_dir: []u8, + table: embed.Table, + mixer: HcW, + head_q: Q4, + scratch: Scratch, + scratch_rows: usize, + mode: []u8, + eos: i64, + experts: cuda.Module, + plan_fn: cuda.Function, + up_fn: cuda.Function, + down_fn: cuda.Function, + layers: [layers_n]Layer, + n_layers: usize, + + pub fn init(gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, model_dir: []const u8) !*Engine { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + const e = try gpa.create(Engine); + errdefer gpa.destroy(e); + e.gpa = gpa; + e.io = io; + e.driver = driver; + e.stream = stream; + e.n_layers = 0; + errdefer for (e.layers[0..e.n_layers]) |*l| l.free(gpa); + e.model_dir = try gpa.dupe(u8, model_dir); + errdefer gpa.free(e.model_dir); + e.store = try Store.open(gpa, io, e.model_dir); + errdefer e.store.deinit(); + const weight = try e.store.tensor(embed.weight_name); + const scales = try e.store.tensor(embed.scales_name); + const biases = try e.store.tensor(embed.biases_name); + e.table = try embed.Table.upload(driver, weight, scales, biases); + errdefer e.table.deinit(); + e.eos = try ple.eosOf(gpa, io, e.model_dir); + e.scratch = try Scratch.alloc(driver, 1); + errdefer e.scratch.free(); + try e.scratch.dummy.fill8(0, null); + e.scratch_rows = 1; + e.mode = try gpa.alloc(u8, 1); + errdefer gpa.free(e.mode); + e.mixer = try loadHc(gpa, driver, &e.store, "language_model.model.hyper_connection_mixer", false); + errdefer e.mixer.free(); + var buf: [80]u8 = undefined; + const head_w = try join3(gpa, &e.store, &buf, "language_model.lm_head", ""); + defer gpa.free(head_w.w); + defer gpa.free(head_w.s); + defer gpa.free(head_w.b); + if (head_w.w.len != vocab * (dims / 8) * 4) return error.UnexpectedTensor; + e.head_q = try uploadQ4(gpa, driver, head_w.w, head_w.s, head_w.b, vocab, dims / 8); + errdefer e.head_q.free(); + e.experts = try cuda.Module.load(driver, cuda.kernels.experts); + errdefer e.experts.unload(); + e.plan_fn = try e.experts.function(weights.plan_symbol); + e.up_fn = try e.experts.function(weights.up_symbol); + e.down_fn = try e.experts.function(weights.down_symbol); + var pack_mod = try cuda.Module.load(driver, cuda.kernels.experts_pack); + defer pack_mod.unload(); + const pack_fn = try pack_mod.function(weights.pack_symbol); + for (0..layers_n) |i| { + e.layers[i] = try loadLayer(gpa, driver, stream, &e.store, pack_fn, i); + e.n_layers += 1; + } + e.store.release(); + e.store.closeFile(); + return e; + } + + pub fn deinit(e: *Engine) void { + e.experts.unload(); + for (e.layers[0..e.n_layers]) |*l| l.free(e.gpa); + e.head_q.free(); + e.mixer.free(); + e.gpa.free(e.mode); + e.scratch.free(); + e.table.deinit(); + e.store.deinit(); + e.gpa.free(e.model_dir); + e.gpa.destroy(e); + } + + fn grow(e: *Engine, rows: usize) !void { + if (rows <= e.scratch_rows) return; + var next = try Scratch.alloc(e.driver, rows); + errdefer next.free(); + try next.dummy.fill8(0, null); + const mode = try e.gpa.alloc(u8, rows); + e.scratch.free(); + e.gpa.free(e.mode); + e.scratch = next; + e.mode = mode; + e.scratch_rows = rows; + } + + /// A fresh cache. `cap` is how many tokens it can hold, prompt plus generated. + pub fn newSeq(e: *Engine, cap: usize) !Seq { + if (cap == 0) return error.UnexpectedTensor; + const logits = try e.gpa.alloc(u8, vocab * 2); + errdefer e.gpa.free(logits); + const conv = try zeros(e.gpa, e.driver, linear_n, (gdn.taps - 1) * gdn.conv_dim * 2); + errdefer freeOwned(e.gpa, conv); + const state = try zeros(e.gpa, e.driver, linear_n, state_n * 4); + errdefer freeOwned(e.gpa, state); + const kc = try zeros(e.gpa, e.driver, attn_n, cap * kv_heads * head_dim * 2); + errdefer freeOwned(e.gpa, kc); + const vc = try zeros(e.gpa, e.driver, attn_n, cap * kv_heads * head_dim * 2); + errdefer freeOwned(e.gpa, vc); + const ikc = try zeros(e.gpa, e.driver, attn_n, cap * index_dim * 2); + errdefer freeOwned(e.gpa, ikc); + var tail = try cuda.DeviceBuffer.alloc(e.driver, ple_tail_n * wide * 2); + errdefer tail.free(); + try tail.fill8(0, null); + return .{ + .conv = conv, + .state = state, + .kc = kc, + .vc = vc, + .ikc = ikc, + .tail = tail, + .history = .{ e.eos, e.eos }, + .fed = 0, + .logits = logits, + .cap = cap, + }; + } + + /// Run `ids` on from `seq.fed` and leave the next token's scores in `seq.logits`. + pub fn predict(e: *Engine, comptime Tri: type, tri: Tri, seq: *Seq, ids: []const u32) !void { + if (ids.len == 0 or seq.fed + ids.len > seq.cap) return error.UnexpectedTensor; + try e.grow(ids.len); + var pass = Pass{ + .gpa = e.gpa, + .io = e.io, + .driver = e.driver, + .stream = e.stream, + .store = &e.store, + .model_dir = e.model_dir, + .s = &e.scratch, + .table = e.table, + .rows = ids.len, + .mode = e.mode, + .history = &seq.history, + .conv = seq.conv, + .state = seq.state, + .kc = seq.kc, + .vc = seq.vc, + .ikc = seq.ikc, + .tail = seq.tail, + .mixer = e.mixer, + .head_q = e.head_q, + .logits = seq.logits, + .plan_fn = e.plan_fn, + .up_fn = e.up_fn, + .down_fn = e.down_fn, + }; + try forward(Tri, &pass, tri, &e.layers, ids, seq.fed); + seq.fed += ids.len; + } +}; + +/// One stream's caches. `fed` tokens have already gone through the model. +pub const Seq = struct { + conv: []cuda.DeviceBuffer, + state: []cuda.DeviceBuffer, + kc: []cuda.DeviceBuffer, + vc: []cuda.DeviceBuffer, + ikc: []cuda.DeviceBuffer, + tail: cuda.DeviceBuffer, + history: [2]i64, + fed: usize, + logits: []u8, + cap: usize, + + pub fn deinit(s: *Seq, gpa: std.mem.Allocator) void { + freeBufs(s.conv); + freeBufs(s.state); + freeBufs(s.kc); + freeBufs(s.vc); + freeBufs(s.ikc); + gpa.free(s.conv); + gpa.free(s.state); + gpa.free(s.kc); + gpa.free(s.vc); + gpa.free(s.ikc); + if (s.tail.ptr != 0) s.tail.free(); + gpa.free(s.logits); + } +}; + +fn freeBufs(bs: []cuda.DeviceBuffer) void { + for (bs) |*b| b.free(); +} + +fn freeOwned(gpa: std.mem.Allocator, bs: []cuda.DeviceBuffer) void { + freeBufs(bs); + gpa.free(bs); +} + +/// `out` receives `new_tokens` greedy ids. The first is the prompt's argmax. The rest are fed back one at a time. +pub fn run(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, model_dir: []const u8, tokens: []const u32, new_tokens: usize, out: []u32) !void { + if (tokens.len == 0 or new_tokens == 0 or out.len < new_tokens or tokens.len + new_tokens > 256) return error.UnexpectedTensor; + var e = try Engine.init(gpa, io, driver, stream, model_dir); + defer e.deinit(); + var seq = try e.newSeq(tokens.len + new_tokens); + defer seq.deinit(gpa); + try e.predict(Tri, tri, &seq, tokens); + out[0] = argmax(seq.logits); + std.debug.print("token {d}\n", .{out[0]}); + var tok = out[0]; + for (1..new_tokens) |i| { + var one = [_]u32{tok}; + try e.predict(Tri, tri, &seq, &one); + tok = argmax(seq.logits); + out[i] = tok; + std.debug.print("token {d}\n", .{tok}); + } +} + +test "linear layers skip every fourth index and attention starts at 3" { + try std.testing.expect(linearLayer(0)); + try std.testing.expect(!linearLayer(3)); + try std.testing.expect(!linearLayer(47)); + try std.testing.expectEqual(@as(usize, 0), linearIndex(0)); + try std.testing.expectEqual(@as(usize, 2), linearIndex(2)); + try std.testing.expectEqual(@as(usize, 35), linearIndex(46)); + try std.testing.expectEqual(@as(usize, 0), attnIndex(3)); + try std.testing.expectEqual(@as(usize, 11), attnIndex(47)); +} diff --git a/zig/src/families/flashnext/cuda_prompt_load.zig b/zig/src/families/flashnext/cuda_prompt_load.zig new file mode 100644 index 0000000000..f7192ce768 --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_load.zig @@ -0,0 +1,389 @@ +//! One resident copy of every decoder layer. The prompt forward reads these, it does not reload them. + +const std = @import("std"); +const cuda = @import("cuda"); +const embed = @import("cuda_embed.zig"); +const gdn = @import("cuda_gdn.zig"); +const wt = @import("cuda_prompt_wt.zig"); + +const layers_n = wt.layers_n; +const dims = wt.dims; +const wide = wt.wide; +const experts_n = wt.experts_n; +const moe_w = wt.moe_w; +const head_dim = wt.head_dim; +const index_dim = wt.index_dim; +const proj_n = wt.proj_n; +const out_k = wt.out_k; +const ple_layer = wt.ple_layer; +const block_bytes = wt.block_bytes; +const cint = wt.cint; +const Store = wt.Store; +const Face = wt.Face; +const Q4 = wt.Q4; +const HcW = wt.HcW; +const linearLayer = wt.linearLayer; +const asF32 = wt.asF32; +const uploadQ4 = wt.uploadQ4; +const join3 = wt.join3; +const loadHc = wt.loadHc; +const stitch = wt.stitch; + +const OwnedFace = struct { + face: Face, + w: []u8, + s: []u8, + b: []u8, + + fn free(self: *OwnedFace, gpa: std.mem.Allocator) void { + gpa.free(self.w); + gpa.free(self.s); + gpa.free(self.b); + } +}; + +pub fn faceOf(store: *Store, buf: []u8, prefix: []const u8, name: []const u8, routed: bool, n: usize, k: usize) !Face { + const w = try store.hold(try std.fmt.bufPrint(buf, "{s}{s}.weight", .{ prefix, name })); + const sc = try store.hold(try std.fmt.bufPrint(buf, "{s}{s}.scales", .{ prefix, name })); + const b = try store.hold(try std.fmt.bufPrint(buf, "{s}{s}.biases", .{ prefix, name })); + const k8 = k / 8; + const kg = k / 32; + if (routed) { + if (!w.is(.u32, &.{ experts_n, n, k8 }) or !sc.is(.bf16, &.{ experts_n, n, kg }) or !b.is(.bf16, &.{ experts_n, n, kg })) return error.UnexpectedTensor; + } else if (!w.is(.u32, &.{ n, k8 }) or !sc.is(.bf16, &.{ n, kg }) or !b.is(.bf16, &.{ n, kg })) return error.UnexpectedTensor; + return .{ .w = w.bytes, .s = sc.bytes, .b = b.bytes, .n = n, .k = k, .routed = routed }; +} + +pub const Layer = struct { + attn_hc: HcW, + mlp_hc: HcW, + proj: Q4, + out: Q4, + router: cuda.DeviceBuffer, + linear: bool, + cw: cuda.DeviceBuffer, + a_log: cuda.DeviceBuffer, + dt: cuda.DeviceBuffer, + norm: cuda.DeviceBuffer, + q_scale: cuda.DeviceBuffer, + k_scale: cuda.DeviceBuffer, + i_scale: cuda.DeviceBuffer, + key: Q4, + val: Q4, + nk: cuda.DeviceBuffer, + nq: cuda.DeviceBuffer, + nc: cuda.DeviceBuffer, + ple_cw: cuda.DeviceBuffer, + /// Packed gate and up for every routed expert plus the shared expert at index 512. + experts_up: cuda.DeviceBuffer, + experts_down: cuda.DeviceBuffer, + + pub fn free(self: *Layer, gpa: std.mem.Allocator) void { + _ = gpa; + self.attn_hc.free(); + self.mlp_hc.free(); + self.proj.free(); + self.out.free(); + self.router.free(); + self.cw.free(); + self.a_log.free(); + self.dt.free(); + self.norm.free(); + self.q_scale.free(); + self.k_scale.free(); + self.i_scale.free(); + self.key.free(); + self.val.free(); + self.nk.free(); + self.nq.free(); + self.nc.free(); + self.ple_cw.free(); + self.experts_up.free(); + self.experts_down.free(); + } +}; + +fn deadBuf(d: *cuda.Driver) cuda.DeviceBuffer { + return .{ .d = d, .ptr = 0, .len = 0 }; +} + +fn blankFace() OwnedFace { + return .{ .face = .{ .w = &.{}, .s = &.{}, .b = &.{}, .n = 0, .k = 0, .routed = false }, .w = &.{}, .s = &.{}, .b = &.{} }; +} + +/// Copy one expert tensor out of the checkpoint so the shard can be closed. +fn keepFace(gpa: std.mem.Allocator, store: *Store, buf: []u8, prefix: []const u8, name: []const u8, routed: bool, n: usize, k: usize) !OwnedFace { + const view = try faceOf(store, buf, prefix, name, routed, n, k); + const w = try gpa.dupe(u8, view.w); + errdefer gpa.free(w); + const s = try gpa.dupe(u8, view.s); + errdefer gpa.free(s); + const b = try gpa.dupe(u8, view.b); + return .{ .face = .{ .w = w, .s = s, .b = b, .n = n, .k = k, .routed = routed }, .w = w, .s = s, .b = b }; +} + +fn packDevice(driver: *cuda.Driver, stream: *cuda.Stream, pack_fn: cuda.Function, words: []const u8, scales: []const u8, biases: []const u8, e: usize, n: usize, k: usize) !cuda.DeviceBuffer { + const kg = k / 32; + const nb = n / 32; + var w_b = try cuda.DeviceBuffer.fromHost(driver, words); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, biases); + defer b_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, e * nb * kg * block_bytes); + errdefer out_b.free(); + var a: cuda.Args = .{}; + a.add(w_b.ptr); + a.add(s_b.ptr); + a.add(b_b.ptr); + a.add(out_b.ptr); + a.add(cint(n)); + a.add(cint(k / 8)); + a.add(cint(kg)); + a.add(cint(nb)); + try cuda.launch.launch(pack_fn, .{ .grid = .{ .x = @intCast(kg), .y = @intCast(nb), .z = @intCast(e) }, .block = .{ .x = 160 } }, stream.*, &a); + try stream.synchronize(); + return out_b; +} + +/// Routed experts, then the shared expert as index `experts_n`, in the kernel's packed layout. +fn packAll(driver: *cuda.Driver, stream: *cuda.Stream, pack_fn: cuda.Function, routed: OwnedFace, shared: OwnedFace) !cuda.DeviceBuffer { + const n = routed.face.n; + const k = routed.face.k; + var bulk = try packDevice(driver, stream, pack_fn, routed.w, routed.s, routed.b, experts_n, n, k); + defer bulk.free(); + var one = try packDevice(driver, stream, pack_fn, shared.w, shared.s, shared.b, 1, n, k); + defer one.free(); + if (bulk.len != experts_n * one.len) return error.UnexpectedTensor; + var out = try cuda.DeviceBuffer.alloc(driver, bulk.len + one.len); + errdefer out.free(); + try out.copyFrom(0, bulk.ptr, bulk.len, stream.handle); + try out.copyFrom(bulk.len, one.ptr, one.len, stream.handle); + try stream.synchronize(); + return out; +} + +fn stackHalves(gpa: std.mem.Allocator, gate: []const u8, up: []const u8) ![]u8 { + if (gate.len != up.len or gate.len % block_bytes != 0) return error.UnexpectedTensor; + const nblk = gate.len / block_bytes; + const out = try gpa.alloc(u8, gate.len + up.len); + for (0..nblk) |b| { + @memcpy(out[b * 2 * block_bytes ..][0..block_bytes], gate[b * block_bytes ..][0..block_bytes]); + @memcpy(out[b * 2 * block_bytes + block_bytes ..][0..block_bytes], up[b * block_bytes ..][0..block_bytes]); + } + return out; +} + +/// Gate and up blocks side by side, which is the SwiGLU weight the expert kernel reads. +fn stackUp(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, gate: cuda.DeviceBuffer, up: cuda.DeviceBuffer) !cuda.DeviceBuffer { + const graw = try gpa.alloc(u8, gate.len); + defer gpa.free(graw); + const uraw = try gpa.alloc(u8, up.len); + defer gpa.free(uraw); + try stream.synchronize(); + try gate.download(0, graw); + try up.download(0, uraw); + const stacked = try stackHalves(gpa, graw, uraw); + defer gpa.free(stacked); + return cuda.DeviceBuffer.fromHost(driver, stacked); +} + +pub fn loadLayer(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, store: *Store, pack_fn: cuda.Function, layer: usize) !Layer { + var buf: [200]u8 = undefined; + const base = try std.fmt.bufPrint(&buf, "language_model.model.layers.{d}", .{layer}); + var base_buf: [80]u8 = undefined; + @memcpy(base_buf[0..base.len], base); + const root = base_buf[0..base.len]; + var attn_hc = try loadHc(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.attn_hyper_connection", .{root}), true); + errdefer attn_hc.free(); + var mlp_hc = try loadHc(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.mlp_hyper_connection", .{root}), true); + errdefer mlp_hc.free(); + const mlp = try std.fmt.bufPrint(&buf, "{s}.mlp", .{root}); + var mlp_buf: [96]u8 = undefined; + @memcpy(mlp_buf[0..mlp.len], mlp); + const mlp_root = mlp_buf[0..mlp.len]; + const gate = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.gate.weight", .{mlp_root})); + if (!gate.is(.bf16, &.{ experts_n, dims })) return error.UnexpectedTensor; + const gate_copy = try gpa.alloc(u8, gate.bytes.len); + defer gpa.free(gate_copy); + @memcpy(gate_copy, gate.bytes); + const sg = try join3(gpa, store, &buf, mlp_root, ".shared_expert_gate"); + defer gpa.free(sg.w); + defer gpa.free(sg.s); + defer gpa.free(sg.b); + const router = try gpa.alloc(u16, (experts_n + 1) * dims); + defer gpa.free(router); + @memcpy(std.mem.sliceAsBytes(router[0 .. experts_n * dims]), gate_copy); + try embed.dequant(sg.w, sg.s, sg.b, dims, 0, router[experts_n * dims ..][0..dims]); + var router_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(router)); + errdefer router_b.free(); + + const linear = linearLayer(layer); + var proj_q = Q4.empty(driver); + var out_q = Q4.empty(driver); + errdefer proj_q.free(); + errdefer out_q.free(); + var cw_b = deadBuf(driver); + var a_b = deadBuf(driver); + var dt_b = deadBuf(driver); + var nw_b = deadBuf(driver); + var qs_b = deadBuf(driver); + var ks_b = deadBuf(driver); + var is_b = deadBuf(driver); + var key_q = Q4.empty(driver); + var val_q = Q4.empty(driver); + var nk_b = deadBuf(driver); + var nq_b = deadBuf(driver); + var nc_b = deadBuf(driver); + var pcw_b = deadBuf(driver); + errdefer cw_b.free(); + errdefer a_b.free(); + errdefer dt_b.free(); + errdefer nw_b.free(); + errdefer qs_b.free(); + errdefer ks_b.free(); + errdefer is_b.free(); + errdefer nk_b.free(); + errdefer nq_b.free(); + errdefer nc_b.free(); + errdefer pcw_b.free(); + if (linear) { + const pre = try std.fmt.bufPrint(&buf, "{s}.linear_attn", .{root}); + var pre_buf: [120]u8 = undefined; + @memcpy(pre_buf[0..pre.len], pre); + const lp = pre_buf[0..pre.len]; + const parts = [_][]const u8{ ".in_proj_qkv", ".in_proj_z", ".in_proj_b", ".in_proj_a" }; + const stitched = try stitch(gpa, store, lp, &parts); + defer gpa.free(stitched.w); + defer gpa.free(stitched.s); + defer gpa.free(stitched.b); + if (stitched.w.len != gdn.proj_width * (dims / 8) * 4) return error.UnexpectedTensor; + proj_q = try uploadQ4(gpa, driver, stitched.w, stitched.s, stitched.b, gdn.proj_width, dims / 8); + const out = try join3(gpa, store, &buf, lp, ".out_proj"); + defer gpa.free(out.w); + defer gpa.free(out.s); + defer gpa.free(out.b); + out_q = try uploadQ4(gpa, driver, out.w, out.s, out.b, dims, gdn.value_dim / 8); + const conv = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.conv1d.weight", .{lp})); + if (conv.bytes.len != gdn.conv_dim * gdn.taps * 2) return error.UnexpectedTensor; + cw_b = try cuda.DeviceBuffer.fromHost(driver, conv.bytes); + const a_log = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.A_log", .{lp}))); + defer gpa.free(a_log); + const dt = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.dt_bias", .{lp}))); + defer gpa.free(dt); + if (a_log.len != gdn.nv or dt.len != gdn.nv) return error.UnexpectedTensor; + a_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(a_log)); + dt_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(dt)); + const norm = try store.copyOf(try std.fmt.bufPrint(&buf, "{s}.norm.weight", .{lp})); + defer gpa.free(norm); + nw_b = try cuda.DeviceBuffer.fromHost(driver, norm); + } else { + const pre = try std.fmt.bufPrint(&buf, "{s}.self_attn", .{root}); + var pre_buf: [120]u8 = undefined; + @memcpy(pre_buf[0..pre.len], pre); + const ap = pre_buf[0..pre.len]; + const parts = [_][]const u8{ ".q_proj", ".k_proj", ".v_proj", ".indexer.index_qk_proj" }; + const stitched = try stitch(gpa, store, ap, &parts); + defer gpa.free(stitched.w); + defer gpa.free(stitched.s); + defer gpa.free(stitched.b); + if (stitched.w.len != proj_n * (dims / 8) * 4) return error.UnexpectedTensor; + proj_q = try uploadQ4(gpa, driver, stitched.w, stitched.s, stitched.b, proj_n, dims / 8); + const out = try join3(gpa, store, &buf, ap, ".o_proj"); + defer gpa.free(out.w); + defer gpa.free(out.s); + defer gpa.free(out.b); + out_q = try uploadQ4(gpa, driver, out.w, out.s, out.b, dims, out_k / 8); + const qn = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.q_norm.weight", .{ap}))); + defer gpa.free(qn); + const kn = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.k_norm.weight", .{ap}))); + defer gpa.free(kn); + const iqn = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.indexer.q_layernorm.weight", .{ap}))); + defer gpa.free(iqn); + if (qn.len != head_dim or kn.len != head_dim or iqn.len != index_dim) return error.UnexpectedTensor; + qs_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(qn)); + ks_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(kn)); + is_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(iqn)); + } + if (layer == ple_layer) { + const pre = try std.fmt.bufPrint(&buf, "{s}.ple", .{root}); + var pre_buf: [96]u8 = undefined; + @memcpy(pre_buf[0..pre.len], pre); + const pp = pre_buf[0..pre.len]; + const key = try join3(gpa, store, &buf, pp, ".key_proj"); + defer gpa.free(key.w); + defer gpa.free(key.s); + defer gpa.free(key.b); + const val = try join3(gpa, store, &buf, pp, ".value_proj"); + defer gpa.free(val.w); + defer gpa.free(val.s); + defer gpa.free(val.b); + key_q = try uploadQ4(gpa, driver, key.w, key.s, key.b, wide, dims / 8); + val_q = try uploadQ4(gpa, driver, val.w, val.s, val.b, dims, dims / 8); + const nk = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_key.weight", .{pp}))); + defer gpa.free(nk); + const nq = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_query.weight", .{pp}))); + defer gpa.free(nq); + const nc = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_conv.weight", .{pp}))); + defer gpa.free(nc); + if (nk.len != wide or nq.len != wide or nc.len != wide) return error.UnexpectedTensor; + nk_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(nk)); + nq_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(nq)); + nc_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(nc)); + const conv = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.conv1d.weight", .{pp})); + if (conv.bytes.len != wide * 4 * 2) return error.UnexpectedTensor; + pcw_b = try cuda.DeviceBuffer.fromHost(driver, conv.bytes); + } + errdefer key_q.free(); + errdefer val_q.free(); + + var shared: [3]OwnedFace = .{ blankFace(), blankFace(), blankFace() }; + errdefer for (&shared) |*f| f.free(gpa); + shared[0] = try keepFace(gpa, store, &buf, mlp_root, ".shared_expert.gate_proj", false, moe_w, dims); + shared[1] = try keepFace(gpa, store, &buf, mlp_root, ".shared_expert.up_proj", false, moe_w, dims); + shared[2] = try keepFace(gpa, store, &buf, mlp_root, ".shared_expert.down_proj", false, dims, moe_w); + var routed: [3]OwnedFace = .{ blankFace(), blankFace(), blankFace() }; + errdefer for (&routed) |*f| f.free(gpa); + routed[0] = try keepFace(gpa, store, &buf, mlp_root, ".switch_mlp.gate_proj", true, moe_w, dims); + routed[1] = try keepFace(gpa, store, &buf, mlp_root, ".switch_mlp.up_proj", true, moe_w, dims); + routed[2] = try keepFace(gpa, store, &buf, mlp_root, ".switch_mlp.down_proj", true, dims, moe_w); + var gate_p = try packAll(driver, stream, pack_fn, routed[0], shared[0]); + errdefer gate_p.free(); + var up_p = try packAll(driver, stream, pack_fn, routed[1], shared[1]); + errdefer up_p.free(); + var down_p = try packAll(driver, stream, pack_fn, routed[2], shared[2]); + errdefer down_p.free(); + for (&shared) |*f| f.free(gpa); + for (&routed) |*f| f.free(gpa); + shared = .{ blankFace(), blankFace(), blankFace() }; + routed = .{ blankFace(), blankFace(), blankFace() }; + var experts_up = try stackUp(gpa, driver, stream, gate_p, up_p); + errdefer experts_up.free(); + gate_p.free(); + up_p.free(); + store.release(); + return .{ + .attn_hc = attn_hc, + .mlp_hc = mlp_hc, + .proj = proj_q, + .out = out_q, + .router = router_b, + .linear = linear, + .cw = cw_b, + .a_log = a_b, + .dt = dt_b, + .norm = nw_b, + .q_scale = qs_b, + .k_scale = ks_b, + .i_scale = is_b, + .key = key_q, + .val = val_q, + .nk = nk_b, + .nq = nq_b, + .nc = nc_b, + .ple_cw = pcw_b, + .experts_up = experts_up, + .experts_down = down_p, + }; +} diff --git a/zig/src/families/flashnext/cuda_prompt_main.zig b/zig/src/families/flashnext/cuda_prompt_main.zig new file mode 100644 index 0000000000..bb84cafefc --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_main.zig @@ -0,0 +1,50 @@ +//! Prompt tokens in, greedy token ids out. The ids file is whitespace or comma separated. + +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); + +const want = [_]u32{ 1596, 1144, 310, 5707, 310, 1156, 25, 328 }; + +pub fn main(init: std.process.Init) !u8 { + const gpa = init.gpa; + const io = init.io; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 5) { + std.debug.print("usage: flashnext-prompt \n", .{}); + return 2; + } + const count = try std.fmt.parseInt(usize, args[4], 10); + const text = try std.Io.Dir.cwd().readFileAlloc(io, args[3], gpa, .limited(1 << 20)); + defer gpa.free(text); + var ids: std.ArrayList(u32) = .empty; + defer ids.deinit(gpa); + var it = std.mem.tokenizeAny(u8, text, ", \t\r\n"); + while (it.next()) |part| try ids.append(gpa, try std.fmt.parseInt(u32, part, 10)); + if (ids.items.len == 0 or count == 0) return 2; + + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + var set = try cuda.aot.Set.load(gpa, io, &driver, ctx.device, args[2]); + defer set.deinit(); + var stream = try cuda.Stream.init(&driver, true); + defer stream.deinit(); + const tri: flash.Tri = .{ .set = &set, .s = stream }; + const got = try gpa.alloc(u32, count); + defer gpa.free(got); + try flash.prompt.run(flash.Tri, gpa, io, &driver, &stream, tri, args[1], ids.items, count, got); + std.debug.print("tokens", .{}); + for (got) |id| std.debug.print(" {d}", .{id}); + std.debug.print("\n", .{}); + if (count == want.len and ids.items.len == 58) { + var same = true; + for (got, &want) |g, w| if (g != w) { + same = false; + }; + std.debug.print("{s}\n", .{if (same) "match" else "differ"}); + if (!same) return 1; + } + return 0; +} diff --git a/zig/src/families/flashnext/cuda_prompt_step.zig b/zig/src/families/flashnext/cuda_prompt_step.zig new file mode 100644 index 0000000000..de02eb2eef --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_step.zig @@ -0,0 +1,316 @@ +//! One fused decode row: hyper-connection, linear or full attention, experts, and the n-gram step. + +const std = @import("std"); +const cuda = @import("cuda"); +const qmm = @import("cuda_qmm.zig"); +const gdn = @import("cuda_gdn.zig"); +const ple = @import("cuda_ple.zig"); +const wt = @import("cuda_prompt_wt.zig"); + +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const low = wt.low; +const eps = wt.eps; +const experts_n = wt.experts_n; +const slots = wt.slots; +const moe_w = wt.moe_w; +const block_bytes = wt.block_bytes; +const tile = wt.tile; +const q_heads = wt.q_heads; +const kv_heads = wt.kv_heads; +const head_dim = wt.head_dim; +const index_heads = wt.index_heads; +const index_dim = wt.index_dim; +const half = wt.half; +const proj_n = wt.proj_n; +const out_k = wt.out_k; +const nch = wt.nch; +const ple_tail_n = wt.ple_tail_n; +const state_n = wt.state_n; +const hc_sk = wt.hc_sk; +const out_sk = wt.out_sk; +const inj_stride = wt.inj_stride; +const wts_stride = wt.wts_stride; +const vocab = wt.vocab; +const cint = wt.cint; +const toBf16 = wt.toBf16; +const pack_symbol = wt.pack_symbol; +const plan_symbol = wt.plan_symbol; +const up_symbol = wt.up_symbol; +const down_symbol = wt.down_symbol; +const Face = wt.Face; +const Q4 = wt.Q4; +const HcW = wt.HcW; + +pub const Scratch = struct { + h: cuda.DeviceBuffer, + y: cuda.DeviceBuffer, + wts: cuda.DeviceBuffer, + inj: cuda.DeviceBuffer, + pss: cuda.DeviceBuffer, + norm: cuda.DeviceBuffer, + dn: cuda.DeviceBuffer, + hc_part: cuda.DeviceBuffer, + act: cuda.DeviceBuffer, + xs_act: cuda.DeviceBuffer, + mixed: cuda.DeviceBuffer, + xsm: cuda.DeviceBuffer, + proj: cuda.DeviceBuffer, + part: cuda.DeviceBuffer, + gout: cuda.DeviceBuffer, + gxs: cuda.DeviceBuffer, + st: cuda.DeviceBuffer, + ple_emb: cuda.DeviceBuffer, + ple_xs: cuda.DeviceBuffer, + ple_keys: cuda.DeviceBuffer, + ple_vals: cuda.DeviceBuffer, + ple_gated: cuda.DeviceBuffer, + ple_pss: cuda.DeviceBuffer, + ple_nrow: cuda.DeviceBuffer, + ple_part: cuda.DeviceBuffer, + pa: cuda.DeviceBuffer, + q: cuda.DeviceBuffer, + iq: cuda.DeviceBuffer, + po: cuda.DeviceBuffer, + pm: cuda.DeviceBuffer, + pl: cuda.DeviceBuffer, + ao: cuda.DeviceBuffer, + gated: cuda.DeviceBuffer, + gxs_attn: cuda.DeviceBuffer, + dummy: cuda.DeviceBuffer, + inv: cuda.DeviceBuffer, + pos: cuda.DeviceBuffer, + ids: cuda.DeviceBuffer, + logits: cuda.DeviceBuffer, + router_l: cuda.DeviceBuffer, + picks: cuda.DeviceBuffer, + moe_members: cuda.DeviceBuffer, + moe_items: cuda.DeviceBuffer, + moe_counts: cuda.DeviceBuffer, + moe_act: cuda.DeviceBuffer, + + pub fn alloc(driver: *cuda.Driver, rows: usize) !Scratch { + var s: Scratch = undefined; + s.h = try cuda.DeviceBuffer.alloc(driver, rows * wide * 2); + s.y = try cuda.DeviceBuffer.alloc(driver, rows * slots * dims * 4); + s.wts = try cuda.DeviceBuffer.alloc(driver, rows * wts_stride); + s.inj = try cuda.DeviceBuffer.alloc(driver, rows * inj_stride); + s.pss = try cuda.DeviceBuffer.alloc(driver, (dims / 256) * streams * 4); + s.norm = try cuda.DeviceBuffer.alloc(driver, wide * 2); + s.dn = try cuda.DeviceBuffer.alloc(driver, (low + 4) * 2); + s.hc_part = try cuda.DeviceBuffer.alloc(driver, hc_sk * (low + 4) * 4); + s.act = try cuda.DeviceBuffer.alloc(driver, low * 2); + s.xs_act = try cuda.DeviceBuffer.alloc(driver, (low / 32) * 4); + s.mixed = try cuda.DeviceBuffer.alloc(driver, dims * 2); + s.xsm = try cuda.DeviceBuffer.alloc(driver, (dims / 32) * 4); + s.proj = try cuda.DeviceBuffer.alloc(driver, gdn.proj_width * 2); + s.part = try cuda.DeviceBuffer.alloc(driver, out_sk * dims * 4); + s.gout = try cuda.DeviceBuffer.alloc(driver, gdn.value_dim * 2); + s.gxs = try cuda.DeviceBuffer.alloc(driver, (gdn.value_dim / 32) * 4); + s.st = try cuda.DeviceBuffer.alloc(driver, state_n * 4); + s.ple_emb = try cuda.DeviceBuffer.alloc(driver, dims * 2); + s.ple_xs = try cuda.DeviceBuffer.alloc(driver, (dims / 32) * 4); + s.ple_keys = try cuda.DeviceBuffer.alloc(driver, wide * 2); + s.ple_vals = try cuda.DeviceBuffer.alloc(driver, dims * 2); + s.ple_gated = try cuda.DeviceBuffer.alloc(driver, wide * 2); + s.ple_pss = try cuda.DeviceBuffer.alloc(driver, streams * 4); + s.ple_nrow = try cuda.DeviceBuffer.alloc(driver, wide * 2); + s.ple_part = try cuda.DeviceBuffer.alloc(driver, 4 * dims * 4); + s.pa = try cuda.DeviceBuffer.alloc(driver, proj_n * 2); + s.q = try cuda.DeviceBuffer.alloc(driver, q_heads * head_dim * 2); + s.iq = try cuda.DeviceBuffer.alloc(driver, index_heads * index_dim * 2); + s.po = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * head_dim * 4); + s.pm = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * 4); + s.pl = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * 4); + s.ao = try cuda.DeviceBuffer.alloc(driver, out_k * 2); + s.gated = try cuda.DeviceBuffer.alloc(driver, out_k * 2); + s.gxs_attn = try cuda.DeviceBuffer.alloc(driver, (out_k / 32) * 4); + s.dummy = try cuda.DeviceBuffer.alloc(driver, 4096); + var inv: [half]f32 = undefined; + for (&inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / @as(f64, half))); + s.inv = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&inv)); + s.pos = try cuda.DeviceBuffer.alloc(driver, 4); + s.ids = try cuda.DeviceBuffer.alloc(driver, 4); + s.logits = try cuda.DeviceBuffer.alloc(driver, vocab * 2); + s.router_l = try cuda.DeviceBuffer.alloc(driver, (experts_n + 1) * 4); + s.picks = try cuda.DeviceBuffer.alloc(driver, slots * 4); + s.moe_members = try cuda.DeviceBuffer.alloc(driver, slots * 4); + s.moe_items = try cuda.DeviceBuffer.alloc(driver, slots * 3 * 4); + s.moe_counts = try cuda.DeviceBuffer.alloc(driver, 8); + s.moe_act = try cuda.DeviceBuffer.alloc(driver, slots * moe_w * 2); + return s; + } + + pub fn free(self: *Scratch) void { + self.h.free(); + self.y.free(); + self.wts.free(); + self.inj.free(); + self.pss.free(); + self.norm.free(); + self.dn.free(); + self.hc_part.free(); + self.act.free(); + self.xs_act.free(); + self.mixed.free(); + self.xsm.free(); + self.proj.free(); + self.part.free(); + self.gout.free(); + self.gxs.free(); + self.st.free(); + self.ple_emb.free(); + self.ple_xs.free(); + self.ple_keys.free(); + self.ple_vals.free(); + self.ple_gated.free(); + self.ple_pss.free(); + self.ple_nrow.free(); + self.ple_part.free(); + self.pa.free(); + self.q.free(); + self.iq.free(); + self.po.free(); + self.pm.free(); + self.pl.free(); + self.ao.free(); + self.gated.free(); + self.gxs_attn.free(); + self.dummy.free(); + self.inv.free(); + self.pos.free(); + self.ids.free(); + self.logits.free(); + self.router_l.free(); + self.picks.free(); + self.moe_members.free(); + self.moe_items.free(); + self.moe_counts.free(); + self.moe_act.free(); + } +}; + +pub fn writeback(tri: anytype, s: *Scratch, h: u64, mode: usize, inj: u64, y: u64, wts: u64) !void { + const dummy = s.dummy.ptr; + if (mode == 0) { + try tri.hcWriteback(h, h, s.pss.ptr, dummy, "*bf16", dummy, dummy, "*bf16", dummy, "*bf16", dims, 1, dims, streams, 0, 1, 1, 1); + } else if (mode == 2) { + try tri.hcWriteback(h, h, s.pss.ptr, dummy, "*bf16", inj, y, "*fp32", wts, "*fp32", dims, 1, dims, streams, 2, 10, 11, 1); + } else if (mode == 4) { + try tri.hcWriteback(h, h, s.pss.ptr, s.part.ptr, "*fp32", inj, dummy, "*bf16", dummy, "*bf16", dims, 1, dims, streams, 4, 1, 1, out_sk); + } else return error.UnexpectedTensor; +} + +pub fn readout(tri: anytype, s: *Scratch, hc_w: HcW, h: u64, inj: u64) !void { + try tri.hcDown(h, s.pss.ptr, hc_w.scale.ptr, s.norm.ptr, hc_w.down.w.ptr, hc_w.down.s.ptr, hc_w.down.b.ptr, s.dn.ptr, s.hc_part.ptr, eps, 1, hc_w.n_down, wide, dims, dims / 256, streams, hc_sk); + const inj_ptr = if (hc_w.has_inj == 1) inj else s.act.ptr; + try tri.hcReduceAct(s.hc_part.ptr, s.act.ptr, s.xs_act.ptr, inj_ptr, hc_sk, 1, streams, low, hc_w.n_down, hc_w.has_inj); + try tri.hcUpmix(s.act.ptr, s.xs_act.ptr, hc_w.up.w.ptr, hc_w.up.s.ptr, hc_w.up.b.ptr, s.norm.ptr, s.mixed.ptr, s.xsm.ptr, 1, wide, low, dims, streams); +} + +fn shiftTail(gpa: std.mem.Allocator, stream: *cuda.Stream, window: cuda.DeviceBuffer, taps_n: usize, row_bytes: usize, fresh: []const u8) !void { + const bytes = taps_n * row_bytes; + const raw = try gpa.alloc(u8, bytes); + defer gpa.free(raw); + try stream.synchronize(); + try window.download(0, raw); + const next = try gpa.alloc(u8, bytes); + defer gpa.free(next); + const keep = taps_n - 1; + @memcpy(next[0 .. keep * row_bytes], raw[row_bytes..]); + @memcpy(next[keep * row_bytes ..], fresh); + try window.upload(0, next); +} + +pub fn gdnStep(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, s: *Scratch, proj_q: Q4, out_q: Q4, conv: cuda.DeviceBuffer, cw: u64, a_log: u64, dt: u64, norm: u64, state: *cuda.DeviceBuffer) !void { + try qmm.matmul(driver, stream.*, s.mixed.ptr, s.xsm.ptr, proj_q.w.ptr, proj_q.s.ptr, proj_q.b.ptr, s.proj.ptr, 1, proj_q.n, proj_q.k); + try gdn.launch(driver, stream.*, s.proj.ptr, conv.ptr, cw, state.ptr, a_log, dt, norm, eps, s.gout.ptr, s.gxs.ptr, s.st.ptr); + try state.copyFrom(0, s.st.ptr, state_n * 4, stream.handle); + const row = try gpa.alloc(u8, gdn.conv_dim * 2); + defer gpa.free(row); + try stream.synchronize(); + try s.proj.download(0, row); + try shiftTail(gpa, stream, conv, gdn.taps - 1, gdn.conv_dim * 2, row); + try qmm.partials(driver, stream.*, s.gout.ptr, s.gxs.ptr, out_q.w.ptr, out_q.s.ptr, out_q.b.ptr, s.dummy.ptr, s.part.ptr, 1, dims, gdn.value_dim, out_sk); +} + +pub fn attnStep(comptime Tri: type, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, s: *Scratch, proj_q: Q4, out_q: Q4, q_scale: u64, k_scale: u64, i_scale: u64, kc: u64, vc: u64, ikc: u64, pos: i32) !void { + try qmm.matmul(driver, stream.*, s.mixed.ptr, s.xsm.ptr, proj_q.w.ptr, proj_q.s.ptr, proj_q.b.ptr, s.pa.ptr, 1, proj_n, dims); + try s.pos.upload(0, std.mem.asBytes(&pos)); + try tri.attnPrep(s.pa.ptr, s.pos.ptr, q_scale, k_scale, i_scale, s.inv.ptr, s.q.ptr, kc, vc, s.dummy.ptr, s.dummy.ptr, s.iq.ptr, ikc, s.dummy.ptr, s.dummy.ptr, eps, 1, 0); + try tri.attnChunks(s.q.ptr, kc, vc, s.dummy.ptr, s.dummy.ptr, s.pos.ptr, s.po.ptr, s.pm.ptr, s.pl.ptr, s.dummy.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads, 1); + try tri.attnMerge(s.po.ptr, s.pm.ptr, s.pl.ptr, s.pos.ptr, s.ao.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads); + try tri.attnGate(s.ao.ptr, s.pa.ptr, s.gated.ptr, s.gxs_attn.ptr, 1, proj_n, q_heads, head_dim); + try qmm.partials(driver, stream.*, s.gated.ptr, s.gxs_attn.ptr, out_q.w.ptr, out_q.s.ptr, out_q.b.ptr, s.dummy.ptr, s.part.ptr, 1, dims, out_k, out_sk); +} + +fn launchExpert(f: cuda.Function, stream: cuda.Stream, x: u64, x_stride: usize, slot_n: usize, w: u64, kg: usize, nb: usize, items: u64, counts: u64, members: u64, out: u64, n: usize, units: usize) !void { + var a: cuda.Args = .{}; + a.add(x); + a.add(cint(x_stride)); + a.add(cint(slot_n)); + a.add(w); + a.add(cint(kg)); + a.add(cint(nb)); + a.add(items); + a.add(counts); + a.add(members); + a.add(out); + a.add(cint(n)); + a.add(@as(f32, 0)); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast((units + 3) / 4) }, .block = .{ .x = 128 } }, stream, &a); +} + +pub fn moeStep(comptime Tri: type, tri: Tri, s: *Scratch, stream: *cuda.Stream, plan_fn: cuda.Function, up_fn: cuda.Function, down_fn: cuda.Function, router_w: u64, up: u64, down: u64, y: u64, wts: u64) !void { + try tri.router(s.mixed.ptr, router_w, s.router_l.ptr, dims); + try tri.topkRows(s.router_l.ptr, s.picks.ptr, wts); + var plan_a: cuda.Args = .{}; + plan_a.add(s.picks.ptr); + plan_a.add(cint(slots)); + plan_a.add(cint(experts_n + 1)); + plan_a.add(cint(tile)); + plan_a.add(s.moe_members.ptr); + plan_a.add(s.moe_items.ptr); + plan_a.add(s.moe_counts.ptr); + try cuda.launch.launch(plan_fn, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1024 } }, stream.*, &plan_a); + try launchExpert(up_fn, stream.*, s.mixed.ptr, dims, slots, up, dims / 32, moe_w / 32, s.moe_items.ptr, s.moe_counts.ptr, s.moe_members.ptr, s.moe_act.ptr, moe_w, slots * (moe_w / 32)); + try launchExpert(down_fn, stream.*, s.moe_act.ptr, moe_w, 0, down, moe_w / 32, dims / 32, s.moe_items.ptr, s.moe_counts.ptr, s.moe_members.ptr, y, dims, slots * (dims / 32)); +} + +fn reduceBf16(gpa: std.mem.Allocator, stream: *cuda.Stream, part: cuda.DeviceBuffer, sk: usize, n: usize, dst: cuda.DeviceBuffer) !void { + const raw = try gpa.alloc(u8, sk * n * 4); + defer gpa.free(raw); + try stream.synchronize(); + try part.download(0, raw); + const row = try gpa.alloc(u16, n); + defer gpa.free(row); + for (0..n) |i| { + var acc: f32 = @bitCast(std.mem.readInt(u32, raw[4 * i ..][0..4], .little)); + for (1..sk) |s| acc += @as(f32, @bitCast(std.mem.readInt(u32, raw[4 * (s * n + i) ..][0..4], .little))); + row[i] = toBf16(acc); + } + try dst.upload(0, std.mem.sliceAsBytes(row)); +} + +pub fn pleStep(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, s: *Scratch, model_dir: []const u8, history: *[2]i64, token: u32, h: u64, key_q: Q4, val_q: Q4, nk: u64, nq: u64, nc: u64, cw: u64, tail: *cuda.DeviceBuffer) !void { + const emb = try gpa.alloc(u16, dims); + defer gpa.free(emb); + const xs = try gpa.alloc(f32, dims / 32); + defer gpa.free(xs); + try ple.embedding(gpa, io, model_dir, history, token, emb, xs); + try s.ple_emb.upload(0, std.mem.sliceAsBytes(emb)); + try s.ple_xs.upload(0, std.mem.sliceAsBytes(xs)); + try qmm.matmul(driver, stream.*, s.ple_emb.ptr, s.ple_xs.ptr, key_q.w.ptr, key_q.s.ptr, key_q.b.ptr, s.ple_keys.ptr, 1, wide, dims); + try qmm.partials(driver, stream.*, s.ple_emb.ptr, s.ple_xs.ptr, val_q.w.ptr, val_q.s.ptr, val_q.b.ptr, s.ple_vals.ptr, s.ple_part.ptr, 1, dims, dims, 4); + try reduceBf16(gpa, stream, s.ple_part, 4, dims, s.ple_vals); + try tri.pleGate(s.ple_keys.ptr, s.ple_vals.ptr, h, nk, nq, s.ple_gated.ptr, s.ple_pss.ptr, eps, 1, dims, streams); + try tri.pleConv(s.ple_gated.ptr, s.ple_pss.ptr, nc, tail.ptr, cw, h, h, s.ple_nrow.ptr, eps, 1, dims, streams, 4, 3); + const nrow = try gpa.alloc(u8, wide * 2); + defer gpa.free(nrow); + try stream.synchronize(); + try s.ple_nrow.download(0, nrow); + try shiftTail(gpa, stream, tail.*, ple_tail_n, wide * 2, nrow); + history[0] = history[1]; + history[1] = token; +} diff --git a/zig/src/families/flashnext/cuda_prompt_wt.zig b/zig/src/families/flashnext/cuda_prompt_wt.zig new file mode 100644 index 0000000000..a2598efd7d --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_wt.zig @@ -0,0 +1,398 @@ +//! Checkpoint bytes for one drafts-off prompt: the safetensors reader and the quantized linears. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const hc = @import("cuda_hc.zig"); +const qmm = @import("cuda_qmm.zig"); +const gdn = @import("cuda_gdn.zig"); + +pub const layers_n: usize = 48; +pub const linear_n: usize = 36; +pub const attn_n: usize = 12; +pub const dims: usize = 2560; +pub const streams: usize = 4; +pub const wide: usize = streams * dims; +pub const low: usize = 320; +pub const eps: f32 = 1e-6; +pub const vocab: usize = 248320; +pub const experts_n: usize = 512; +pub const slots: usize = 11; +pub const moe_w: usize = 640; +pub const block_bytes: usize = 160 * 4; +pub const tile: usize = 16; +pub const q_heads: usize = 24; +pub const kv_heads: usize = 2; +pub const head_dim: usize = 256; +pub const index_heads: usize = 4; +pub const index_dim: usize = 128; +pub const half: usize = 32; +pub const proj_n: usize = 13952; +pub const out_k: usize = 6144; +pub const nch: usize = 3; +pub const ple_layer: usize = 1; +pub const inj_stride: usize = 16; +pub const wts_stride: usize = 48; +pub const ple_tail_n: usize = 9; +pub const state_n: usize = gdn.nv * gdn.dv * gdn.dk; +pub const hc_sk: usize = 32; +pub const out_sk: usize = 8; + +pub const pack_symbol: [:0]const u8 = "_ZN15tf_experts_pack11pack_kernelILi1EEEvPKjPKtS4_Pjiiii"; +pub const plan_symbol: [:0]const u8 = "_ZN10tf_experts11plan_kernelEPKiiiiPiS2_S2_"; +pub const up_symbol: [:0]const u8 = "_ZN10tf_experts13expert_kernelILi32ELi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; +pub const down_symbol: [:0]const u8 = "_ZN10tf_experts13expert_kernelILi32ELi1ELi0ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; + +const Tensor = core.safetensors.Tensor; + +const Held = struct { rel: []u8, file: Shard }; + +const Shard = struct { + io: std.Io, + file: std.Io.File, + map: std.Io.File.MemoryMap, + at: usize, + arena: std.heap.ArenaAllocator, + names: std.StringHashMapUnmanaged(core.safetensors.Entry), + + fn load(gpa: std.mem.Allocator, io: std.Io, path: []const u8) !Shard { + var file = try std.Io.Dir.cwd().openFile(io, path, .{}); + errdefer file.close(io); + const len: usize = @intCast(try file.length(io)); + if (len < 8) return error.BadSafetensors; + var map = try std.Io.File.MemoryMap.create(io, file, .{ .len = len, .protection = .{ .read = true, .write = false }, .populate = false }); + errdefer map.destroy(io); + const header_len: usize = @intCast(std.mem.readInt(u64, map.memory[0..8], .little)); + if (header_len > len - 8) return error.BadSafetensors; + var arena = std.heap.ArenaAllocator.init(gpa); + errdefer arena.deinit(); + var names: std.StringHashMapUnmanaged(core.safetensors.Entry) = .empty; + const parsed = try std.json.parseFromSlice(std.json.Value, arena.allocator(), map.memory[8..][0..header_len], .{}); + defer parsed.deinit(); + var it = parsed.value.object.iterator(); + while (it.next()) |kv| { + if (std.mem.eql(u8, kv.key_ptr.*, "__metadata__")) continue; + const o = kv.value_ptr.object; + const dtype = core.safetensors.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; + const shape = o.get("shape").?.array.items; + if (shape.len > core.safetensors.max_rank) continue; + var e: core.safetensors.Entry = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .begin = 0, .end = 0 }; + var n: usize = dtype.size(); + for (shape, 0..) |d, i| { + e.shape[i] = @intCast(d.integer); + n *= e.shape[i]; + } + const offs = o.get("data_offsets").?.array.items; + e.begin = @intCast(offs[0].integer); + e.end = @intCast(offs[1].integer); + if (e.end < e.begin or e.end - e.begin != n or e.end > len - 8 - header_len) return error.BadSafetensors; + try names.put(arena.allocator(), try arena.allocator().dupe(u8, kv.key_ptr.*), e); + } + return .{ .io = io, .file = file, .map = map, .at = 8 + header_len, .arena = arena, .names = names }; + } + + fn close(self: *Shard) void { + self.map.destroy(self.io); + self.file.close(self.io); + self.arena.deinit(); + self.* = undefined; + } + + fn get(self: *const Shard, name: []const u8) ?Tensor { + const e = self.names.get(name) orelse return null; + return .{ .dtype = e.dtype, .rank = e.rank, .shape = e.shape, .bytes = self.map.memory[self.at + e.begin .. self.at + e.end] }; + } +}; + +pub const Store = struct { + gpa: std.mem.Allocator, + io: std.Io, + dir: []const u8, + parsed: std.json.Parsed(std.json.Value), + path: []u8 = &.{}, + rel: []u8 = &.{}, + file: ?Shard = null, + held: [12]Held = undefined, + nheld: usize = 0, + + pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Store { + const index = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); + defer gpa.free(index); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index, gpa, .limited(1 << 26)); + defer gpa.free(text); + return .{ .gpa = gpa, .io = io, .dir = dir, .parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}) }; + } + + pub fn deinit(self: *Store) void { + self.release(); + self.closeFile(); + self.parsed.deinit(); + } + + pub fn release(self: *Store) void { + for (self.held[0..self.nheld]) |*h| { + h.file.close(); + self.gpa.free(h.rel); + } + self.nheld = 0; + } + + pub fn hold(self: *Store, key: []const u8) !Tensor { + const file = try self.shard(key); + for (self.held[0..self.nheld]) |*h| { + if (std.mem.eql(u8, h.rel, file)) return h.file.get(key) orelse return error.MissingTensor; + } + if (self.nheld == self.held.len) return error.UnexpectedTensor; + const rel = try self.gpa.dupe(u8, file); + errdefer self.gpa.free(rel); + const path = try std.fs.path.join(self.gpa, &.{ self.dir, file }); + defer self.gpa.free(path); + self.held[self.nheld] = .{ .rel = rel, .file = try Shard.load(self.gpa, self.io, path) }; + self.nheld += 1; + return self.held[self.nheld - 1].file.get(key) orelse { + std.debug.print("absent {s}\n", .{key}); + return error.MissingTensor; + }; + } + + pub fn closeFile(self: *Store) void { + if (self.file) |*f| f.close(); + self.file = null; + if (self.path.len != 0) self.gpa.free(self.path); + self.path = &.{}; + if (self.rel.len != 0) self.gpa.free(self.rel); + self.rel = &.{}; + } + + fn shard(self: *Store, key: []const u8) ![]const u8 { + const wm = self.parsed.value.object.get("weight_map") orelse return error.MissingTensor; + const v = wm.object.get(key) orelse { + std.debug.print("missing {s}\n", .{key}); + return error.MissingTensor; + }; + return v.string; + } + + fn use(self: *Store, key: []const u8) !void { + const file = try self.shard(key); + if (self.file != null and std.mem.eql(u8, file, self.rel)) return; + self.closeFile(); + self.rel = try self.gpa.dupe(u8, file); + self.path = try std.fs.path.join(self.gpa, &.{ self.dir, file }); + self.file = try Shard.load(self.gpa, self.io, self.path); + } + + pub fn tensor(self: *Store, key: []const u8) !Tensor { + try self.use(key); + return self.file.?.get(key) orelse { + std.debug.print("absent {s}\n", .{key}); + return error.MissingTensor; + }; + } + + pub fn copyOf(self: *Store, key: []const u8) ![]u8 { + const t = try self.tensor(key); + const out = try self.gpa.alloc(u8, t.bytes.len); + @memcpy(out, t.bytes); + return out; + } +}; + +pub const Face = struct { + w: []const u8, + s: []const u8, + b: []const u8, + n: usize, + k: usize, + routed: bool, +}; + +pub const Q4 = struct { + w: cuda.DeviceBuffer, + s: cuda.DeviceBuffer, + b: cuda.DeviceBuffer, + n: usize, + k: usize, + + pub fn free(self: *Q4) void { + self.w.free(); + self.s.free(); + self.b.free(); + } + + pub fn empty(d: *cuda.Driver) Q4 { + const z: cuda.DeviceBuffer = .{ .d = d, .ptr = 0, .len = 0 }; + return .{ .w = z, .s = z, .b = z, .n = 0, .k = 0 }; + } +}; + +pub const HcW = struct { + scale: cuda.DeviceBuffer, + down: Q4, + up: Q4, + n_down: usize, + has_inj: usize, + + pub fn free(self: *HcW) void { + self.scale.free(); + self.down.free(); + self.up.free(); + } +}; + +pub fn linearLayer(layer: usize) bool { + return layer % 4 != 3; +} + +pub fn linearIndex(layer: usize) usize { + return layer - layer / 4; +} + +pub fn attnIndex(layer: usize) usize { + return layer / 4; +} + +pub fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +pub fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +pub fn cint(v: usize) c_int { + return @intCast(v); +} + +pub fn cat(gpa: std.mem.Allocator, parts: []const []const u8) ![]u8 { + var n: usize = 0; + for (parts) |p| n += p.len; + const out = try gpa.alloc(u8, n); + var at: usize = 0; + for (parts) |p| { + @memcpy(out[at..][0..p.len], p); + at += p.len; + } + return out; +} + +pub fn asF32(gpa: std.mem.Allocator, t: Tensor) ![]f32 { + if (t.rank != 1) return error.UnexpectedTensor; + const out = try gpa.alloc(f32, t.dim(0)); + if (t.dtype == .f32 and t.bytes.len == out.len * 4) { + for (out, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, t.bytes[4 * i ..][0..4], .little)); + } else if (t.dtype == .bf16 and t.bytes.len == out.len * 2) { + for (out, 0..) |*o, i| o.* = promote(std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little)); + } else return error.UnexpectedTensor; + return out; +} + +pub fn uploadQ4(gpa: std.mem.Allocator, driver: *cuda.Driver, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k8: usize) !Q4 { + var lane = try qmm.pack(gpa, words, scales, biases, n, k8); + defer lane.deinit(gpa); + var w = try cuda.DeviceBuffer.fromHost(driver, lane.weight); + errdefer w.free(); + var s = try cuda.DeviceBuffer.fromHost(driver, lane.scales); + errdefer s.free(); + const b = try cuda.DeviceBuffer.fromHost(driver, lane.biases); + return .{ .w = w, .s = s, .b = b, .n = n, .k = k8 * 8 }; +} + +pub fn uploadTiled(gpa: std.mem.Allocator, driver: *cuda.Driver, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k8: usize) !Q4 { + const kg = k8 / 4; + const tiled = try hc.tileWords(gpa, words, n, k8); + defer gpa.free(tiled); + const ts = try hc.transposeBf16(gpa, scales, n, kg); + defer gpa.free(ts); + const tb = try hc.transposeBf16(gpa, biases, n, kg); + defer gpa.free(tb); + var w = try cuda.DeviceBuffer.fromHost(driver, tiled); + errdefer w.free(); + var s = try cuda.DeviceBuffer.fromHost(driver, ts); + errdefer s.free(); + const b = try cuda.DeviceBuffer.fromHost(driver, tb); + return .{ .w = w, .s = s, .b = b, .n = n, .k = k8 * 8 }; +} + +pub fn join3(gpa: std.mem.Allocator, store: *Store, buf: []u8, prefix: []const u8, name: []const u8) !struct { w: []u8, s: []u8, b: []u8 } { + const w = try store.copyOf(try std.fmt.bufPrint(buf, "{s}{s}.weight", .{ prefix, name })); + errdefer gpa.free(w); + const s = try store.copyOf(try std.fmt.bufPrint(buf, "{s}{s}.scales", .{ prefix, name })); + errdefer gpa.free(s); + const b = try store.copyOf(try std.fmt.bufPrint(buf, "{s}{s}.biases", .{ prefix, name })); + return .{ .w = w, .s = s, .b = b }; +} + +pub fn loadHc(gpa: std.mem.Allocator, driver: *cuda.Driver, store: *Store, prefix: []const u8, inject: bool) !HcW { + var buf: [180]u8 = undefined; + const down = try join3(gpa, store, &buf, prefix, ".input_mix_weight_down"); + defer gpa.free(down.w); + defer gpa.free(down.s); + defer gpa.free(down.b); + const up = try join3(gpa, store, &buf, prefix, ".input_mix_weight_up"); + defer gpa.free(up.w); + defer gpa.free(up.s); + defer gpa.free(up.b); + const gamma = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.hc_norm.weight", .{prefix})); + const scale = try asF32(gpa, gamma); + defer gpa.free(scale); + if (scale.len != wide) return error.UnexpectedTensor; + var n: usize = low; + var words: []u8 = undefined; + var scales: []u8 = undefined; + var biases: []u8 = undefined; + var inj_w: []u8 = &.{}; + var inj_s: []u8 = &.{}; + var inj_b: []u8 = &.{}; + if (inject) { + const inj = try join3(gpa, store, &buf, prefix, ".block_inject_weight"); + inj_w = inj.w; + inj_s = inj.s; + inj_b = inj.b; + n = low + 4; + words = try cat(gpa, &.{ down.w, inj.w }); + scales = try cat(gpa, &.{ down.s, inj.s }); + biases = try cat(gpa, &.{ down.b, inj.b }); + } else { + words = try gpa.dupe(u8, down.w); + scales = try gpa.dupe(u8, down.s); + biases = try gpa.dupe(u8, down.b); + } + defer gpa.free(words); + defer gpa.free(scales); + defer gpa.free(biases); + defer if (inject) gpa.free(inj_w); + defer if (inject) gpa.free(inj_s); + defer if (inject) gpa.free(inj_b); + const k8 = wide / 8; + if (words.len != n * k8 * 4) return error.UnexpectedTensor; + var down_q = try uploadTiled(gpa, driver, words, scales, biases, n, k8); + errdefer down_q.free(); + var up_q = try uploadTiled(gpa, driver, up.w, up.s, up.b, wide, low / 8); + errdefer up_q.free(); + const scale_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(scale)); + return .{ .scale = scale_b, .down = down_q, .up = up_q, .n_down = n, .has_inj = @intFromBool(inject) }; +} + +pub fn stitch(gpa: std.mem.Allocator, store: *Store, prefix: []const u8, names: []const []const u8) !struct { w: []u8, s: []u8, b: []u8, n: usize } { + var buf: [180]u8 = undefined; + var ws: [6][]u8 = undefined; + var ss: [6][]u8 = undefined; + var bs: [6][]u8 = undefined; + var n: usize = 0; + for (names, 0..) |part, i| { + const got = try join3(gpa, store, &buf, prefix, part); + ws[i] = got.w; + ss[i] = got.s; + bs[i] = got.b; + n += 1; + } + defer for (0..n) |i| { + gpa.free(ws[i]); + gpa.free(ss[i]); + gpa.free(bs[i]); + }; + return .{ .w = try cat(gpa, ws[0..n]), .s = try cat(gpa, ss[0..n]), .b = try cat(gpa, bs[0..n]), .n = n }; +} diff --git a/zig/src/families/flashnext/cuda_qmm.zig b/zig/src/families/flashnext/cuda_qmm.zig new file mode 100644 index 0000000000..50df60a80f --- /dev/null +++ b/zig/src/families/flashnext/cuda_qmm.zig @@ -0,0 +1,211 @@ +//! Group-32 lane matmul: MLX words packed as qmm.pack, then the shared CUDA kernel's 16-row tile. + +const std = @import("std"); +const cuda = @import("cuda"); + +pub const group: usize = 32; + +const offsets = [8]u32{ 0, 8, 16, 24, 1, 9, 17, 25 }; + +/// LaneTile<32, 16, 64, 1, 4, 4>::SMEM. One by four warps, four stages. +pub const smem: u32 = 9472; +pub const threads: u32 = 128; + +const symbol: [:0]const u8 = "_ZN6tf_qmm10qmm_kernelILi32ELi16ELi64ELi1ELi4ELi4ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii"; + +pub const Packed = struct { + weight: []u8, + scales: []u8, + biases: []u8, + n: usize, + k: usize, + npad: usize, + + pub fn deinit(self: *Packed, gpa: std.mem.Allocator) void { + gpa.free(self.weight); + gpa.free(self.scales); + gpa.free(self.biases); + self.* = undefined; + } +}; + +fn nibble(words: []const u8, k8: usize, col: usize, pos: usize) u32 { + const word_i = pos / 8; + const shift: u5 = @intCast((pos % 8) * 4); + const off = (col * k8 + word_i) * 4; + const word = std.mem.readInt(u32, words[off..][0..4], .little); + return (word >> shift) & 0xF; +} + +/// MLX (n, k/8) words and (n, k/32) scales and biases -> the lane layout. n is padded to 128. +pub fn pack(gpa: std.mem.Allocator, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k8: usize) !Packed { + if (k8 == 0 or k8 % 4 != 0 or words.len != n * k8 * 4) return error.UnexpectedTensor; + const k = k8 * 8; + if (k % group != 0) return error.UnexpectedTensor; + const kg = k / group; + if (scales.len != n * kg * 2 or biases.len != scales.len) return error.UnexpectedTensor; + const npad = (n + 127) / 128 * 128; + const tiles = npad / 64; + const weight = try gpa.alloc(u8, tiles * kg * 8 * 32 * 4); + errdefer gpa.free(weight); + @memset(weight, 0); + for (0..n) |col| { + const tile = col / 64; + const inner = col % 64; + const j = inner / 8; + const r = inner % 8; + for (0..kg) |g| { + for (0..4) |c| { + var bits: u32 = 0; + for (0..8) |p| { + const pos = g * group + 2 * c + offsets[p]; + bits |= nibble(words, k8, col, pos) << @intCast(p * 4); + } + const idx = r * 4 + c; + const dst = (((tile * kg + g) * 8 + j) * 32 + idx) * 4; + std.mem.writeInt(u32, weight[dst..][0..4], bits, .little); + } + } + } + const scales_out = try gpa.alloc(u8, kg * npad * 2); + errdefer gpa.free(scales_out); + const biases_out = try gpa.alloc(u8, kg * npad * 2); + errdefer gpa.free(biases_out); + @memset(scales_out, 0); + @memset(biases_out, 0); + for (0..kg) |g| { + for (0..n) |col| { + const src = (col * kg + g) * 2; + const dst = (g * npad + col) * 2; + @memcpy(scales_out[dst..][0..2], scales[src..][0..2]); + @memcpy(biases_out[dst..][0..2], biases[src..][0..2]); + } + } + return .{ .weight = weight, .scales = scales_out, .biases = biases_out, .n = n, .k = k, .npad = npad }; +} + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +/// One row: each group is `dot(x, q) * scale + xs * bias`, groups added in order, then rounded to bf16. +pub fn dotRow(x: []const u16, xs: []const f32, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k: usize, out: []u16) !void { + if (k == 0 or k % group != 0 or x.len < k or xs.len < k / group or out.len < n) return error.UnexpectedTensor; + const kg = k / group; + const k8 = k / 8; + if (words.len != n * k8 * 4 or scales.len != n * kg * 2 or biases.len != scales.len) return error.UnexpectedTensor; + for (0..n) |col| { + var acc: f32 = 0; + const word_row = words[col * k8 * 4 ..][0 .. k8 * 4]; + const scale_row = scales[col * kg * 2 ..][0 .. kg * 2]; + const bias_row = biases[col * kg * 2 ..][0 .. kg * 2]; + for (0..kg) |g| { + var p: f32 = 0; + for (0..4) |i| { + const word = std.mem.readInt(u32, word_row[(g * 4 + i) * 4 ..][0..4], .little); + for (0..8) |j| { + const q: f32 = @floatFromInt((word >> @intCast(j * 4)) & 0xF); + p += promote(x[g * group + i * 8 + j]) * q; + } + } + const s = promote(std.mem.readInt(u16, scale_row[g * 2 ..][0..2], .little)); + const b = promote(std.mem.readInt(u16, bias_row[g * 2 ..][0..2], .little)); + acc = @mulAdd(f32, xs[g], b, @mulAdd(f32, p, s, acc)); + } + out[col] = toBf16(acc); + } +} + +/// qmm.cu's 16-row launch for one K slice: x (m, k) bf16 and xs (m, k/32) fp32 into out (m, n) bf16. +pub fn matmul(d: *const cuda.Driver, stream: cuda.Stream, x: u64, xs: u64, weight: u64, scales: u64, biases: u64, out: u64, m: usize, n: usize, k: usize) !void { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + if (m == 0 or m > 16 or k == 0 or k % group != 0) return error.UnexpectedTensor; + var module = try cuda.Module.load(d, cuda.kernels.qmm); + defer module.unload(); + const f = try module.function(symbol); + try f.allowDynamicShared(smem); + const bm: usize = 16; + const bn: usize = 64; + const rows_t = (m + bm - 1) / bm; + const cols_t = (n + bn - 1) / bn; + const sweep = @max(@as(usize, 1), @min(rows_t, (12 << 20) / (bm * k * 2))); + const npad = (n + 127) / 128 * 128; + const ldx = k; + var args: cuda.Args = .{}; + args.add(x); + args.add(xs); + args.add(weight); + args.add(scales); + args.add(biases); + args.add(out); + args.add(@as(u64, 0)); + for ([_]usize{ m, n, k, 1, npad, ldx, sweep }) |v| args.add(@as(c_int, @intCast(v))); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows_t * cols_t) }, .block = .{ .x = threads }, .shared = smem }, stream, &args); +} + +/// The same kernel as `matmul` with one block per K slice. `part` receives the unreduced [sk, m, n] fp32 slices. +pub fn partials(d: *const cuda.Driver, stream: cuda.Stream, x: u64, xs: u64, weight: u64, scales: u64, biases: u64, out: u64, part: u64, m: usize, n: usize, k: usize, sk: usize) !void { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + if (m == 0 or m > 16 or sk == 0 or k == 0 or k % group != 0 or (k / group) % sk != 0) return error.UnexpectedTensor; + var module = try cuda.Module.load(d, cuda.kernels.qmm); + defer module.unload(); + const f = try module.function(symbol); + try f.allowDynamicShared(smem); + const bm: usize = 16; + const bn: usize = 64; + const rows_t = (m + bm - 1) / bm; + const cols_t = (n + bn - 1) / bn; + const sweep = @max(@as(usize, 1), @min(rows_t, (12 << 20) / (bm * k * 2))); + const npad = (n + 127) / 128 * 128; + const ldx = k; + var args: cuda.Args = .{}; + args.add(x); + args.add(xs); + args.add(weight); + args.add(scales); + args.add(biases); + args.add(out); + args.add(part); + for ([_]usize{ m, n, k, sk, npad, ldx, sweep }) |v| args.add(@as(c_int, @intCast(v))); + try cuda.launch.launch(f, .{ + .grid = .{ .x = @intCast(rows_t * cols_t), .z = @intCast(sk) }, + .block = .{ .x = threads }, + .shared = smem, + }, stream, &args); +} + +test "one group's first nibble pair lands in the lane word" { + var words: [16]u8 = @splat(0); + std.mem.writeInt(u32, words[0..4], 0x21, .little); + var scale: [2]u8 = undefined; + std.mem.writeInt(u16, &scale, 0x3f80, .little); + const bias: [2]u8 = @splat(0); + var got = try pack(std.testing.allocator, &words, &scale, &bias, 1, 4); + defer got.deinit(std.testing.allocator); + try std.testing.expectEqual(@as(usize, 128), got.npad); + try std.testing.expectEqual(@as(u32, 0x20001), std.mem.readInt(u32, got.weight[0..4], .little)); + try std.testing.expectEqual(@as(u32, 0), std.mem.readInt(u32, got.weight[4..8], .little)); + try std.testing.expectEqual(@as(u16, 0x3f80), std.mem.readInt(u16, got.scales[0..2], .little)); +} + +test "scale times the nibble dot plus the group bias rounds to five" { + var words: [16]u8 = @splat(0); + std.mem.writeInt(u32, words[0..4], 0x21, .little); + var scale: [2]u8 = undefined; + std.mem.writeInt(u16, &scale, 0x3f80, .little); + var bias: [2]u8 = undefined; + std.mem.writeInt(u16, &bias, 0x3f80, .little); + var x: [32]u16 = @splat(0); + x[0] = 0x3f80; + x[1] = 0x3f80; + const xs = [_]f32{2}; + var out: [1]u16 = undefined; + try dotRow(&x, &xs, &words, &scale, &bias, 1, 32, &out); + try std.testing.expectEqual(@as(u16, 0x40a0), out[0]); +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index b066798721..b3cd5bf120 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -5,10 +5,11 @@ const cuda = @import("cuda"); const api = @import("engine_api"); const lanes = @import("lanes"); const nemotron = @import("nemotron"); +const flashnext = @import("flashnext"); const Allocator = std.mem.Allocator; /// The CUDA families: namespaces with `model_type`, `formats`, `default_context`, `prefill_step` and `open`. -const registry = .{nemotron.native}; +const registry = .{ nemotron.native, flashnext.native }; pub const backends: []const []const u8 = &.{"cuda"}; pub const families: []const api.Family = blk: { @@ -230,4 +231,6 @@ test "chip classes name the compute capability" { test "every registered family is listed for capabilities" { try std.testing.expectEqual(@as(usize, registry.len), families.len); try std.testing.expectEqualStrings("nemotron_h", families[0].model_type); + try std.testing.expectEqualStrings("qwen4_exp", families[1].model_type); + try std.testing.expectEqualStrings("mlx-q4g32", families[1].formats[0]); } diff --git a/zig/tests/cuda/copies.py b/zig/tests/cuda/copies.py index 8626189454..9ddf0f1491 100644 --- a/zig/tests/cuda/copies.py +++ b/zig/tests/cuda/copies.py @@ -33,6 +33,13 @@ "gdn.cu": ("src/tensorfold/cuda/kernels/gdn.cu", 1, 319, (3, 4), (10, "tf_gdn"), GDN_FOOTER), "qmm_frag.cuh": ("src/tensorfold/cuda/kernels/qmm_frag.cuh", 1, 118, (), None, ""), "experts.cuh": ("src/tensorfold/cuda/experts.cuh", 1, 105, (), None, ""), + "qmm.cu": ( + "src/tensorfold/cuda/kernels/qmm.cu", 1, 230, (3, 5), (12, "tf_qmm"), + "// Decode lane: groups of 32, 16 rows, 64 columns, one by four warps, four stages, bf16 out.\n" + "template __global__ void tf_qmm::qmm_kernel<32, 16, 64, 1, 4, 4, false, false, false>(\n" + " const __nv_bfloat16*, const float*, const uint32_t*, const __nv_bfloat16*, const __nv_bfloat16*,\n" + " void*, float*, int, int, int, int, int, int, int);\n", + ), "qmm_group.cu": ( "src/tensorfold/cuda/kernels/qmm_group.cu", 1, 348, (3, 5), (14, "tf_qmm_group"), "// The instantiations Nemotron's windows launch on sm_121 (tile 2: rows <= 16, bf16 out).\n" @@ -52,7 +59,12 @@ "#define TF_EXPERT(EPI) template __global__ void tf_experts::expert_kernel<64, 1, EPI, 2, 4>(const __nv_bfloat16*, \\\n" " int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float);\n" "TF_EXPERT(1)\n" - "TF_EXPERT(0)\n", + "TF_EXPERT(0)\n" + "// Groups of 32: SwiGLU gate and up together (epilogue 2), and the down projection to fp32.\n" + "template __global__ void tf_experts::expert_kernel<32, 2, 2, 2, 4>(const __nv_bfloat16*, \\\n" + " int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float);\n" + "template __global__ void tf_experts::expert_kernel<32, 1, 0, 2, 4>(const __nv_bfloat16*, \\\n" + " int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float);\n", ), "experts_prefill.cu": ( "src/tensorfold/cuda/experts_prefill.cu", 1, 153, (3, 4, 7), (11, "tf_experts_prefill"), @@ -64,8 +76,10 @@ ), "experts_pack.cu": ( "src/tensorfold/cuda/experts_pack.cu", 1, 42, (3, 4, 6), (8, "tf_experts_pack"), - "// Groups of 64 inputs.\n" + "// Groups of 64 inputs, then groups of 32.\n" "template __global__ void tf_experts_pack::pack_kernel<2>(const uint32_t*, const uint16_t*, const uint16_t*,\n" + " uint32_t*, int, int, int, int);\n" + "template __global__ void tf_experts_pack::pack_kernel<1>(const uint32_t*, const uint16_t*, const uint16_t*,\n" " uint32_t*, int, int, int, int);\n", ), "prefill_attention.cu": (