diff --git a/CHANGELOG.md b/CHANGELOG.md index d2c7a05a4d..af715d2d96 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -129,6 +129,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Flash Next's MTP head norms its 10,240-value rows with tiled launches on CUDA sm_120 and sm_121 and split wave launches on HIP gfx1201, storing the same bytes as before: the norm kernel runs 1.4x to 2.5x faster on a GB10 and 4x to 5x on a Radeon AI PRO R9700 at 1 to 16 rows (#1702, fixes #1689). Thanks to @laudney. +- Qwen3.8 Flash Next on CUDA serves `nvidia/Qwen3.8-Flash-Next-NVFP4` at `--tp 2` on two GPUs, one rank a host: its BF16 attention, DeltaNet and shared-expert projections, its n-gram table of E4M3 rows with one BF16 scale, and its MTP head's E4M3 experts with 128x128 block scales load as stored. The CUDA server now serves `--tp` (`--rank`, `--master`, `--master-port`); families that open without ranks refuse it by name, and `--idle-unload` is refused with it. The ranks compute and compare a cached digest of every checkpoint byte, so `TENSORFOLD_CHECKPOINT_SHA256` is no longer required (it still overrides the digest). A crash on the first prompt at `--tp 2`, a vtable built on the stack, is fixed. At `--tp 2` each rank owns half the routed experts whole, so replies are one GPU's bytes; on two GB10s drafted replies equal `"draft": false` (6 of 6) and concurrent replies their solo runs (8 of 8 cells). + ## 1.2.0 - CUDA EXL3: the dense projections zero their split-K counters only when the buffer grows, on the calls' stream, instead of before every 128-row block of every call; each block's last program already resets its own counter. Replies are token-identical, and on a GB10 with the Qwen3.8-27B 3.50 bpw pack decode is 0.4% to 0.9% faster (#1475). Thanks to @diwakergupta. diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 09b68a97bc..27730483c6 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -310,3 +310,15 @@ The sum fixtures use integers with exact FP32 partials below 2^24, including poo FP16 and BF16 simulations must fail in each of the three accumulator orders. Pooling declares its nine-term reduction separately from the physical grid width, so the small stress grid remains a valid launch. NaN and both infinity signs have explicit class controls. These registrations still require the device gates. + +Flash Next's BF16 shared-down entries `tf_qmmf_shared16_half`, `tf_qmmf_shared16`, `tf_qmmf_shared32` and +`tf_qmmf_shared64` preserve the contiguous projection's FP32 result and store it directly in expert slot 10. +The baseline shared expert has all 640 intermediate channels and runs only on rank 0 at two ranks. +The contract checks all four geometries, nonfinite inputs, output bytes and untouched routed slots. +`moetile_records.zig` holds launch geometry and measured nvcc 13.0 resources on all five served SMs. +Later compiler receipts remain pending. Check each build with `tf-cuda-test kernel-contract-limits zig-out/fatbin`. +`tf-cuda-test moetile` compares the direct slot with the contiguous product and scatter at 1, 16, 17, 32, 64, 65, +1536 and 4218 rows, five repeats with guards. `moetile-contract` runs the replacement gates and planted control. +`moetile-bench` alternates baseline and direct-slot GPU event timings at 4218 rows; no current timing is claimed. +Automatic shared-down dispatch yields to an eligible 128-row prompt projection from 512 rows. +The direct entry remains available for forced byte fixtures; automatic selection and direct support are checked separately. diff --git a/docs/recipes/flash-next-cuda.md b/docs/recipes/flash-next-cuda.md index 10c6634eb5..cdb7613037 100644 --- a/docs/recipes/flash-next-cuda.md +++ b/docs/recipes/flash-next-cuda.md @@ -19,6 +19,28 @@ contract receipts. Compile success and host tests do not establish checkpoint fi limits must be measured with the exact checkpoint, binary, device topology and request settings. Prefix retention with MTP is outside this path's qualified scope; evaluate with `--prompt-cache-gib 0`. +### NVIDIA's NVFP4 checkpoint on two GB10s + +`nvidia/Qwen3.8-Flash-Next-NVFP4` (the checkpoint SGLang's Flash Next recipes serve) loads as stored: NVFP4 routed +experts with plain scales, BF16 attention, DeltaNet and shared-expert projections, its n-gram table as E4M3 rows +times one BF16 scale (`bf16(e4m3 * scale)`, as SGLang computes it), and its MTP head's routed experts as E4M3 with +BF16 128x128 block scales on the block-FP8 expert kernels (BF16 activations). Two ranks, one a host, rank 1 first: + +```bash +# on rank 1 +tensorfold-native serve nvidia/Qwen3.8-Flash-Next-NVFP4 --tp 2 --rank 1 --master 192.0.2.1 --context 65536 --prompt-cache-gib 0 +# on rank 0 (192.0.2.1), which serves HTTP +tensorfold-native serve nvidia/Qwen3.8-Flash-Next-NVFP4 --tp 2 --rank 0 --master 192.0.2.1 --context 65536 --prompt-cache-gib 0 +``` + +- Both ranks run the same executable (they compare it) and need NCCL 2.30 or later; set `NCCL_SOCKET_IFNAME` and + `NCCL_IB_HCA` for the link between the hosts. `--master-port` (default 29551) names rank 0's control port. +- Each rank sends a digest of every checkpoint byte, cached after the first read (about 2 minutes for this + checkpoint's 124 GiB on a GB10), and both wait up to 30 minutes for the other. `TENSORFOLD_CHECKPOINT_SHA256`, when + set the same on both ranks, replaces the digest and keeps the wait at a minute. +- `--prompt-cache-gib 0` is required with drafts at `--tp 2`, and `--idle-unload` is refused with `--tp`. +- At `--tp 2` each stream's window verifies alone, so concurrent streams do not yet share a verify pass. + ## Experimental K32 prompt experts `TF_FLASHNEXT_PROMPT_EXPERT_TILE=1` selects the K32 staging pipeline for NVFP4 prompt gate-up; diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 15ac6ed602..719edfeff7 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -220,6 +220,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st step.dependOn(&omni_slide.step); const skip_drafts = b.step("test-drafter-skip-cuda-host", "Counted shared draft skips without devices"); for ([_]*std.Build.Module{ mods.nemotron, mods.qwen35, flashnextModule(b, host, .debug, cuda, mods.core, mods.lanes, mods.qwen_vision, mods.api) }) |m| skip_drafts.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m, .filters = &.{"zero grant:"} })).step); + @import("qmmf_prompt_host.zig").steps(b, cuda, mods.core, step); @import("qwen_contract.zig").host(b, mods.core, step); @import("qwen_vision_weights_host.zig").host(b, mods.core, step); @import("inc4_host.zig").round(b, cuda, mods.core); diff --git a/zig/build/flashnext_glue_host.zig b/zig/build/flashnext_glue_host.zig index 98b2657ac2..baabbc355c 100644 --- a/zig/build/flashnext_glue_host.zig +++ b/zig/build/flashnext_glue_host.zig @@ -2,6 +2,13 @@ const std = @import("std"); pub fn steps(b: *std.Build, cuda: *std.Build.Module, core: *std.Build.Module, family: *std.Build.Module, all: *std.Build.Step) void { + const compose = b.step("test-flashnext-compose-host", "Dense heads and shared expert composition without a GPU"); + compose.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = family, .filters = &.{ "split out lane", "head split", "dense split", "full heads", "half the heads", "shared expert composition" } })).step); + const moe = b.step("test-moetile-host", "Shared down slot fixtures, launch dispatch and contract records without a GPU"); + all.dependOn(moe); + moe.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = cuda, .filters = &.{"shared down stride"} })).step); + const fixtures = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/moetile_tests.zig"), .target = b.graph.host, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "flashnext", .module = family } } }); + moe.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = fixtures })).step); const root = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/contract_flashnext_glue_test.zig"), .target = b.graph.host, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "flashnext", .module = family } } }); const run = b.addRunArtifact(b.addTest(.{ .root_module = root })); all.dependOn(&run.step); diff --git a/zig/build/qmmf_prompt_host.zig b/zig/build/qmmf_prompt_host.zig new file mode 100644 index 0000000000..f58c421f7b --- /dev/null +++ b/zig/build/qmmf_prompt_host.zig @@ -0,0 +1,27 @@ +//! BF16 prompt scratch envelopes and contract registrations without a GPU. +const std = @import("std"); +const runtime = @import("cuda.zig").runtime; + +pub fn steps(b: *std.Build, cuda: *std.Build.Module, core: *std.Build.Module, step: *std.Build.Step) void { + const host = b.graph.host; + const prompt_scratch = b.step("test-qmmf-prompt-scratch-host", "BF16 prompt scratch bounds in Debug and ReleaseFast without a GPU"); + for ([_]std.builtin.OptimizeMode{ .debug, .fast }) |optimize| { + const tests = b.createModule(.{ + .root_source_file = b.path("zig/tests/cuda/qmmf_prompt_tests.zig"), + .target = host, + .optimize = optimize, + .link_libc = true, + .imports = &.{.{ .name = "cuda", .module = runtime(b, host, optimize, &.{}) }}, + }); + prompt_scratch.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = tests, .filters = &.{"qmmf prompt scratch"} })).step); + } + step.dependOn(prompt_scratch); + const qmmf_prompt_contract = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/cuda/contract_qmmf_prompt.zig"), + .target = host, + .link_libc = true, + .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core } }, + }), .filters = &.{"qmmf prompt contract"} })); + b.step("test-qmmf-prompt-contract-host", "Four BF16 prompt instances and their contract/limit coverage without a GPU").dependOn(&qmmf_prompt_contract.step); + step.dependOn(&qmmf_prompt_contract.step); +} diff --git a/zig/kernels/cuda/flashnext_batch.cu b/zig/kernels/cuda/flashnext_batch.cu index 8b66e200aa..21713f791f 100644 --- a/zig/kernels/cuda/flashnext_batch.cu +++ b/zig/kernels/cuda/flashnext_batch.cu @@ -32,9 +32,13 @@ extern "C" __global__ void fn_shift_window(__nv_bfloat16* __restrict__ win, cons // ---- the DeltaNet chain over a chunk: flashnext_gdn.cu's chain cut into front, rows and norm, with its bits ---- namespace fn_gdn { -constexpr int NK = 16, NV = 48, DK = 128, DV = 128, TAPS = 4; -constexpr int C = 2 * NK * DK + NV * DV; -constexpr int PW = C + NV * DV + 2 * NV; +constexpr int DK = 128, DV = 128, TAPS = 4; + +// A layer's key and value heads (all 16 and 48, or one rank's 8 and 24): its conv channels and projection row width. +template struct G { + static constexpr int C = 2 * NK * DK + NV * DV; + static constexpr int PW = C + NV * DV + 2 * NV; +}; static __device__ __forceinline__ float bf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } @@ -44,8 +48,10 @@ static __device__ __forceinline__ float sigmoidf_(float x) { return 1.0f / (1.0f static __device__ __forceinline__ float softplusf_(float x) { return x > 20.0f ? x : log1pf(expf(x)); } +template static __device__ __forceinline__ float conv_act(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, const __nv_bfloat16* __restrict__ cw, int r, int c) { + constexpr int C = G::C, PW = G::PW; float acc = 0.0f; #pragma unroll for (int tap = 0; tap < TAPS; ++tap) { @@ -59,16 +65,18 @@ static __device__ __forceinline__ float conv_act(const __nv_bfloat16* __restrict } // namespace fn_gdn // Program (row, unit), 256 threads: key head u's normed q and k (u < NK), else a value head's v, decay g and beta. -extern "C" __global__ void __launch_bounds__(256) fn_gdn_front( +template +static __device__ __forceinline__ void gdn_front( const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, const __nv_bfloat16* __restrict__ cw, const float* __restrict__ a_log, const float* __restrict__ dt_bias, float* __restrict__ qn, float* __restrict__ kn, float* __restrict__ v, float* __restrict__ gb) { using namespace fn_gdn; + constexpr int C = G::C, PW = G::PW; const int r = blockIdx.x, u = blockIdx.y, t = threadIdx.x, warp = t >> 5, lane = t & 31; if (u < NK) { __shared__ float x[2][DK]; const int hk = u; - x[t / DK][t % DK] = conv_act(P, cs, cw, r, t < DK ? hk * DK + t : NK * DK + hk * DK + (t - DK)); + x[t / DK][t % DK] = conv_act(P, cs, cw, r, t < DK ? hk * DK + t : NK * DK + hk * DK + (t - DK)); __syncthreads(); if (warp < 2) { float v4[4], ss = 0.0f; @@ -86,7 +94,7 @@ extern "C" __global__ void __launch_bounds__(256) fn_gdn_front( } } else { const int hv = u - NK; - if (t < DV) v[(size_t)r * NV * DV + hv * DV + t] = conv_act(P, cs, cw, r, 2 * NK * DK + hv * DV + t); + if (t < DV) v[(size_t)r * NV * DV + hv * DV + t] = conv_act(P, cs, cw, r, 2 * NK * DK + hv * DV + t); else if (t == DV) { const float b = __bfloat162float(P[(size_t)r * PW + C + NV * DV + hv]); const float a = __bfloat162float(P[(size_t)r * PW + C + NV * DV + NV + hv]); @@ -96,7 +104,23 @@ extern "C" __global__ void __launch_bounds__(256) fn_gdn_front( } } +extern "C" __global__ void __launch_bounds__(256) fn_gdn_front( + const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, const __nv_bfloat16* __restrict__ cw, + const float* __restrict__ a_log, const float* __restrict__ dt_bias, float* __restrict__ qn, + float* __restrict__ kn, float* __restrict__ v, float* __restrict__ gb) { + gdn_front<16, 48>(P, cs, cw, a_log, dt_bias, qn, kn, v, gb); +} + +// fn_gdn_front on one rank's 8 key and 24 value heads (a projection row 8,240 wide, 5,120 conv channels). +extern "C" __global__ void __launch_bounds__(256) fn_gdn_front_h8( + const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, const __nv_bfloat16* __restrict__ cw, + const float* __restrict__ a_log, const float* __restrict__ dt_bias, float* __restrict__ qn, + float* __restrict__ kn, float* __restrict__ v, float* __restrict__ gb) { + gdn_front<8, 24>(P, cs, cw, a_log, dt_bias, qn, kn, v, gb); +} + // The chain through the chunk from state_in to state_out (in place, or apart for a rollback): a warp 4 state rows. +template static __device__ __forceinline__ void gdn_rows( const float* __restrict__ qn, const float* __restrict__ kn, const float* __restrict__ v, const float* __restrict__ gb, const float* state_in, float* state_out, int rows, float* __restrict__ ys_out) { @@ -166,19 +190,34 @@ static __device__ __forceinline__ void gdn_rows( extern "C" __global__ void __launch_bounds__(128) fn_gdn_rows( const float* __restrict__ qn, const float* __restrict__ kn, const float* __restrict__ v, const float* __restrict__ gb, float* state, int rows, float* __restrict__ ys_out) { - gdn_rows(qn, kn, v, gb, state, state, rows, ys_out); + gdn_rows<16, 48>(qn, kn, v, gb, state, state, rows, ys_out); } extern "C" __global__ void __launch_bounds__(128) fn_gdn_rows_to( const float* __restrict__ qn, const float* __restrict__ kn, const float* __restrict__ v, const float* __restrict__ gb, const float* state_in, float* state_out, int rows, float* __restrict__ ys_out) { - gdn_rows(qn, kn, v, gb, state_in, state_out, rows, ys_out); + gdn_rows<16, 48>(qn, kn, v, gb, state_in, state_out, rows, ys_out); } -extern "C" __global__ void __launch_bounds__(128) fn_gdn_norm( +// fn_gdn_rows and fn_gdn_rows_to on one rank's 8 key and 24 value heads. +extern "C" __global__ void __launch_bounds__(128) fn_gdn_rows_h8( + const float* __restrict__ qn, const float* __restrict__ kn, const float* __restrict__ v, + const float* __restrict__ gb, float* state, int rows, float* __restrict__ ys_out) { + gdn_rows<8, 24>(qn, kn, v, gb, state, state, rows, ys_out); +} + +extern "C" __global__ void __launch_bounds__(128) fn_gdn_rows_to_h8( + const float* __restrict__ qn, const float* __restrict__ kn, const float* __restrict__ v, + const float* __restrict__ gb, const float* state_in, float* state_out, int rows, float* __restrict__ ys_out) { + gdn_rows<8, 24>(qn, kn, v, gb, state_in, state_out, rows, ys_out); +} + +template +static __device__ __forceinline__ void gdn_norm( const float* __restrict__ ys_in, const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ norm_w, float eps, __nv_bfloat16* __restrict__ out, float* __restrict__ xs) { using namespace fn_gdn; + constexpr int C = G::C, PW = G::PW; const int r = blockIdx.x, hv = blockIdx.y, t = threadIdx.x, warp = t >> 5, lane = t & 31; __shared__ float ys[DV]; __shared__ float rinv; @@ -203,6 +242,19 @@ extern "C" __global__ void __launch_bounds__(128) fn_gdn_norm( if (lane == 0) xs[(size_t)r * (NV * DV / 32) + hv * (DV / 32) + warp] = gs; } +extern "C" __global__ void __launch_bounds__(128) fn_gdn_norm( + const float* __restrict__ ys_in, const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ norm_w, + float eps, __nv_bfloat16* __restrict__ out, float* __restrict__ xs) { + gdn_norm<16, 48>(ys_in, P, norm_w, eps, out, xs); +} + +// fn_gdn_norm on one rank's 24 value heads (z read from its 8,240-wide projection rows). +extern "C" __global__ void __launch_bounds__(128) fn_gdn_norm_h8( + const float* __restrict__ ys_in, const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ norm_w, + float eps, __nv_bfloat16* __restrict__ out, float* __restrict__ xs) { + gdn_norm<8, 24>(ys_in, P, norm_w, eps, out, xs); +} + // ---- sparse attention (QSA): each row's blocks ---------------------------------------------------------------------- // attention._select's order-preserving key of a block score: negative floats flipped, positive ones topped. @@ -331,13 +383,28 @@ extern "C" __global__ void fn_copy_cols(const uint16_t* __restrict__ src, long l if (col < width) dst[r * dst_pitch + col] = src[r * src_pitch + col]; } +// A rank's picks as its experts' ids: first..first+local-1 become 0..local-1, any other id (the shared slot too) local. +extern "C" __global__ void fn_ep_picks(const int* __restrict__ picks, int* __restrict__ out, int n, int first, int local) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + const int p = picks[i]; + out[i] = p >= first && p < first + local ? p - first : local; +} + #include "flashnext_gdn_tpr.cuh" // The prompt chain on two threads a state row (gdn_tpr): fn_gdn_rows_to's bits, state_in == state_out in place. extern "C" __global__ void __launch_bounds__(256) fn_gdn_chain_tpr(const float* qn, const float* kn, const float* v, const float* gb, const float* state_in, float* state_out, int rows, float* ys_out) { - gdn_tpr::chain<2, 16>(qn, kn, v, gb, state_in, state_out, rows, ys_out); + gdn_tpr::chain<16, 48, 2, 16>(qn, kn, v, gb, state_in, state_out, rows, ys_out); +} + +// fn_gdn_chain_tpr on one rank's 8 key and 24 value heads. +extern "C" __global__ void __launch_bounds__(256) fn_gdn_chain_tpr_h8(const float* qn, const float* kn, const float* v, + const float* gb, const float* state_in, + float* state_out, int rows, float* ys_out) { + gdn_tpr::chain<8, 24, 2, 16>(qn, kn, v, gb, state_in, state_out, rows, ys_out); } // The prompt GDN block in one kernel a value head (front, chain, gated norm; their bits), no fp32 rows through memory. diff --git a/zig/kernels/cuda/flashnext_gdn_fused.cuh b/zig/kernels/cuda/flashnext_gdn_fused.cuh index 3b15c092e0..2501229f5d 100644 --- a/zig/kernels/cuda/flashnext_gdn_fused.cuh +++ b/zig/kernels/cuda/flashnext_gdn_fused.cuh @@ -9,8 +9,10 @@ constexpr int SMEM = 2 * STAGE + 2 * STEPS * DV * 4 + STEPS * 2 * DK * 4 + STEPS __device__ __forceinline__ void pbar() { asm volatile("bar.sync 1, 128;\n" ::: "memory"); } // conv_act for rows r0 .. r0 + n - 1 of channel c, each P / cs row loaded once (conv_act's taps, order and rounding) +template __device__ __forceinline__ void conv_rows(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, const __nv_bfloat16* __restrict__ cw, int r0, int n, int c, float (&out)[STEPS]) { + constexpr int C = G::C, PW = G::PW; float w[TAPS]; #pragma unroll for (int tap = 0; tap < TAPS; ++tap) w[tap] = __bfloat162float(cw[c * TAPS + tap]); @@ -30,18 +32,20 @@ __device__ __forceinline__ void conv_rows(const __nv_bfloat16* __restrict__ P, c } // producers: rows r0 .. r0 + n - 1's q, k (normed as fn_gdn_front), v, g, beta into the stage +template __device__ __forceinline__ void front_batch(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, const __nv_bfloat16* __restrict__ cw, const float* __restrict__ a_log, const float* __restrict__ dt_bias, int r0, int n, int hk, int hv, int p, unsigned char* base, float (*x)[2][DK]) { + constexpr int C = G::C, PW = G::PW; float o[STEPS]; - conv_rows(P, cs, cw, r0, n, hk * DK + p, o); // fn_gdn_front unit hk, t = p: q's channel + conv_rows(P, cs, cw, r0, n, hk * DK + p, o); // fn_gdn_front unit hk, t = p: q's channel #pragma unroll for (int j = 0; j < STEPS; ++j) if (j < n) x[j][0][p] = o[j]; - conv_rows(P, cs, cw, r0, n, NK * DK + hk * DK + p, o); // t = p + DK: k's channel + conv_rows(P, cs, cw, r0, n, NK * DK + hk * DK + p, o); // t = p + DK: k's channel #pragma unroll for (int j = 0; j < STEPS; ++j) if (j < n) x[j][1][p] = o[j]; - conv_rows(P, cs, cw, r0, n, 2 * NK * DK + hv * DV + p, o); // unit NK + hv, t = p: v + conv_rows(P, cs, cw, r0, n, 2 * NK * DK + hv * DV + p, o); // unit NK + hv, t = p: v #pragma unroll for (int j = 0; j < STEPS; ++j) if (j < n) reinterpret_cast(base + j * ROW + 2 * DK * 4)[p] = o[j]; if (p < n) { // unit NK + hv, t == DV: row r0 + p's g and beta @@ -72,9 +76,11 @@ __device__ __forceinline__ void front_batch(const __nv_bfloat16* __restrict__ P, } // producers: fn_gdn_norm for n rows from row r0 (readouts ys[j][dv]) +template __device__ __forceinline__ void norm_rows(const float (*ys)[DV], const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ norm_w, float eps, __nv_bfloat16* __restrict__ out, float* __restrict__ xs, int r0, int n, int hv, int p, float* rinv) { + constexpr int C = G::C, PW = G::PW; const int warp = p >> 5, lane = p & 31; for (int j = warp; j < n; j += PR / 32) { float ss = 0.0f; @@ -105,6 +111,7 @@ __device__ __forceinline__ void norm_rows(const float (*ys)[DV], const __nv_bflo pbar(); } +template __device__ __forceinline__ void run(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, const __nv_bfloat16* __restrict__ cw, const float* __restrict__ a_log, const float* __restrict__ dt_bias, const float* state_in, float* state_out, int rows, @@ -121,9 +128,9 @@ __device__ __forceinline__ void run(const __nv_bfloat16* __restrict__ P, const _ const int batches = (rows + STEPS - 1) / STEPS; auto fill = [&](int bi) { const int r0 = bi * STEPS, n = min(STEPS, rows - r0); - front_batch(P, cs, cw, a_log, dt_bias, r0, n, hk, hv, p, stage + (bi & 1) * STAGE, xb); + front_batch(P, cs, cw, a_log, dt_bias, r0, n, hk, hv, p, stage + (bi & 1) * STAGE, xb); }; - // the chain's registers (gdn_tpr::chain<2, 16>'s) + // the chain's registers (gdn_tpr::chain's) constexpr int L = 32 / TPR; const int row = threadIdx.x / TPR, t = threadIdx.x % TPR; const size_t sbase = (size_t)hv * DV * DK + (size_t)row * DK; @@ -139,7 +146,7 @@ __device__ __forceinline__ void run(const __nv_bfloat16* __restrict__ P, const _ const int r0 = bi * STEPS, n = min(STEPS, rows - r0); if (producer) { if (bi + 1 < batches) fill(bi + 1); - if (bi >= 1) norm_rows(ysb[(bi - 1) & 1], P, norm_w, eps, out, xs, (bi - 1) * STEPS, STEPS, hv, p, rinv); + if (bi >= 1) norm_rows(ysb[(bi - 1) & 1], P, norm_w, eps, out, xs, (bi - 1) * STEPS, STEPS, hv, p, rinv); } else { const unsigned char* base = stage + (bi & 1) * STAGE; for (int rr = 0; rr < n; ++rr) { @@ -185,7 +192,7 @@ __device__ __forceinline__ void run(const __nv_bfloat16* __restrict__ P, const _ __syncthreads(); } if (producer) { - if (batches > 0) norm_rows(ysb[(batches - 1) & 1], P, norm_w, eps, out, xs, (batches - 1) * STEPS, rows - (batches - 1) * STEPS, hv, p, rinv); + if (batches > 0) norm_rows(ysb[(batches - 1) & 1], P, norm_w, eps, out, xs, (batches - 1) * STEPS, rows - (batches - 1) * STEPS, hv, p, rinv); } else { #pragma unroll for (int l = 0; l < L; ++l) @@ -200,5 +207,14 @@ extern "C" __global__ void __launch_bounds__(gdn_fused::THREADS) fn_gdn_fused( const float* __restrict__ a_log, const float* __restrict__ dt_bias, const float* state_in, float* state_out, int rows, const __nv_bfloat16* __restrict__ norm_w, float eps, __nv_bfloat16* __restrict__ out, float* __restrict__ xs, float* __restrict__ ys_out) { - gdn_fused::run(P, cs, cw, a_log, dt_bias, state_in, state_out, rows, norm_w, eps, out, xs, ys_out); + gdn_fused::run<16, 48>(P, cs, cw, a_log, dt_bias, state_in, state_out, rows, norm_w, eps, out, xs, ys_out); +} + +// fn_gdn_fused on one rank's 8 key and 24 value heads (8,240-wide projection rows, a 5,120-channel conv window). +extern "C" __global__ void __launch_bounds__(gdn_fused::THREADS) fn_gdn_fused_h8( + const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, const __nv_bfloat16* __restrict__ cw, + const float* __restrict__ a_log, const float* __restrict__ dt_bias, const float* state_in, float* state_out, int rows, + const __nv_bfloat16* __restrict__ norm_w, float eps, __nv_bfloat16* __restrict__ out, float* __restrict__ xs, + float* __restrict__ ys_out) { + gdn_fused::run<8, 24>(P, cs, cw, a_log, dt_bias, state_in, state_out, rows, norm_w, eps, out, xs, ys_out); } diff --git a/zig/kernels/cuda/flashnext_gdn_tpr.cuh b/zig/kernels/cuda/flashnext_gdn_tpr.cuh index 013359c7d7..d73f4fcb58 100644 --- a/zig/kernels/cuda/flashnext_gdn_tpr.cuh +++ b/zig/kernels/cuda/flashnext_gdn_tpr.cuh @@ -3,7 +3,7 @@ #define KREG 0 #endif namespace gdn_tpr { -constexpr int NK = 16, NV = 48, DK = 128, DV = 128; +constexpr int DK = 128, DV = 128; __device__ __forceinline__ float bf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } __device__ __forceinline__ void cp16(void* dst, const void* src) { asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(static_cast(__cvta_generic_to_shared(dst))), "l"(src)); @@ -27,7 +27,8 @@ __device__ __forceinline__ float tree(float (&x)[32 / TPR]) { for (int m = TPR / 2; m >= 1; m >>= 1) r = r + __shfl_xor_sync(0xffffffffu, r, m); return r; } -template +// NK key and NV value heads (a layer's 16 and 48, or one rank's 8 and 24), TPR threads a state row, STEPS rows a stage +template __device__ __forceinline__ void chain(const float* __restrict__ qn, const float* __restrict__ kn, const float* __restrict__ v, const float* __restrict__ gb, const float* state_in, float* state_out, int rows, float* __restrict__ ys_out) { diff --git a/zig/kernels/cuda/qmmf.cu b/zig/kernels/cuda/qmmf.cu index c18084c111..92e3ae0708 100644 --- a/zig/kernels/cuda/qmmf.cu +++ b/zig/kernels/cuda/qmmf.cu @@ -47,248 +47,42 @@ struct Tile { static constexpr int S = MODE == FP4 || MODE == FP8G ? BN * 4 : MODE == MXFP8 ? BN * 2 : 0; static constexpr int STAGE = (X + W + S + 127) / 128 * 128; static constexpr int PARTIALS = MT * NT * 4 * THREADS * 4; - static constexpr int SMEM = STAGES * STAGE > PARTIALS ? STAGES * STAGE : PARTIALS; + static constexpr int BARRIERS = MODE == BF16 && BN > 64 ? 128 : 0; + static constexpr int SMEM = (STAGES * STAGE > PARTIALS ? STAGES * STAGE : PARTIALS) + BARRIERS; }; +template +__device__ __forceinline__ void qmmf_body( + const __nv_bfloat16* __restrict__ x, const unsigned char* __restrict__ w, const uint8_t* __restrict__ bs, + float scale, void* __restrict__ out, float* __restrict__ part, int M, int N, int K, int SK, int npad, int ldx, + int group, int ldy) { +#define TF_QMMF_STRIDED +#include "qmmf_body.cuh" +#undef TF_QMMF_STRIDED +} + + template __global__ void __launch_bounds__(WM * WN * 32) qmmf_kernel( const __nv_bfloat16* __restrict__ x, const unsigned char* __restrict__ w, const uint8_t* __restrict__ bs, float scale, void* __restrict__ out, float* __restrict__ part, int M, int N, int K, int SK, int npad, int ldx, int group) { - using T = Tile; - static_assert(BN == 64 || BN == 32, "a block reads one 64-column tile of words and scales, or its half"); - 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, slice_groups = KG / SK, per = FUSE ? KG : slice_groups; - const int2 at = tile_of(blockIdx.x, M, N, BM, BN, group); - const int m0 = at.x, n0 = at.y, slice = FUSE ? 0 : blockIdx.z, g0 = slice * per; - - auto stage = [&](int s) { return buf + s * T::STAGE; }; - auto load = [&](int s, int g) { - 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 = MODE == FP4 ? 64 * GS / 2 : MODE == BF16 ? 64 * GS * 2 : 64 * GS; - if constexpr (BN == 32) { - // half a packed 64-column tile: its n8 tiles 0-3 or 4-7, contiguous in every format's word order - const int half = (n0 % 64) / 32; - for (int c = tid; c < T::W / 16; c += T::THREADS) - cp16(pw + c * 16, w + (static_cast(n0 / 64) * KG + g) * TILE_BYTES + half * (TILE_BYTES / 2) + c * 16); - if constexpr (T::S > 0) { - unsigned char* ps = pw + T::W; - for (int c = tid; c < T::S / 16; c += T::THREADS) - cp16(ps + c * 16, bs + (static_cast(n0 / 64) * KG + g) * (2 * T::S) + half * T::S + c * 16); - } - } else { - for (int c = tid; c < T::W / 16; c += T::THREADS) { - const int t = c / (TILE_BYTES / 16), off = c % (TILE_BYTES / 16); - cp16(pw + c * 16, w + (static_cast(n0 / 64 + t) * KG + g) * TILE_BYTES + off * 16); - } - if constexpr (T::S > 0) { - unsigned char* ps = pw + T::W; - for (int c = tid; c < T::S / 16; c += T::THREADS) - cp16(ps + c * 16, bs + (static_cast(n0 / 64) * KG + g) * T::S + c * 16); - } - } - }; +#include "qmmf_body.cuh" +} - float acc[T::MT][T::NT][4]; - float tot[FUSE ? T::MT : 1][FUSE ? T::NT : 1][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; -#pragma unroll - for (int s = 0; s < STAGES - 1; ++s) { - if (s < per) load(s, g0 + s); - commit(); - } - for (int it = 0; it < per; ++it) { - if constexpr (FUSE) { - if (it > 0 && it % slice_groups == 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) { - tot[i][j][e] = it == slice_groups ? acc[i][j][e] : tot[i][j][e] + acc[i][j][e]; - acc[i][j][e] = 0.0f; - } - } - } - wait(); - __syncthreads(); - const int next = it + STAGES - 1; - if (next < per) load(next % STAGES, g0 + next); - commit(); - const unsigned char* p = stage(it % STAGES); - const unsigned char* pw = p + T::X; - const uint8_t* ps = p + T::X + T::W; - float d[T::MT][T::NT][4]; - uint4 wq[T::NT]; - uint32_t sq[T::NT][2]; -#pragma unroll - for (int j = 0; j < T::NT; ++j) { - const int jj = wn * T::NT + j; - const int col = wn * (BN / WN) + j * 8 + (lane & 3) * 2; - if constexpr (MODE == FP4 || MODE == FP8G) { - const uint2 v = *reinterpret_cast(ps + col * 4); - sq[j][0] = v.x; - sq[j][1] = v.y; - } else if constexpr (MODE == MXFP8) { - const uint32_t v = *reinterpret_cast(ps + col * 2); - sq[j][0] = v & 0xFFFFu; - sq[j][1] = v >> 16; - } - if constexpr (MODE == FP4) { - const uint2 u = reinterpret_cast(pw)[jj * 32 + lane]; - wq[j] = make_uint4(u.x, u.y, 0u, 0u); - } else if constexpr (MODE != BF16) { - wq[j] = reinterpret_cast(pw)[jj * 32 + lane]; - } - } -#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) { - uint32_t b0, b1; - if constexpr (MODE == FP4) { - const uint32_t word = kt < 2 ? wq[j].x : wq[j].y; - b0 = fp4pair(word, (kt & 1) * 8); - b1 = fp4pair(word, (kt & 1) * 8 + 4); - } else if constexpr (MODE == BF16) { - const uint2 v = reinterpret_cast(pw)[((wn * T::NT + j) * 4 + kt) * 32 + lane]; - b0 = v.x; - b1 = v.y; - } else { - const uint32_t word = comp(wq[j], kt); - b0 = fp8pair(word & 0xFFFFu); - b1 = fp8pair(word >> 16); - } -#pragma unroll - for (int i = 0; i < T::MT; ++i) { - if constexpr (MODE == FP8 || MODE == BF16) mma(acc[i][j], a[i], b0, b1); - else if (MODE == FP4 || (MODE == MXFP8 && (kt & 1) == 0) || (MODE == FP8G && kt == 0)) - mma0(d[i][j], a[i], b0, b1); - else mma(d[i][j], a[i], b0, b1); - } - } - if constexpr (MODE == FP4 || MODE == MXFP8 || MODE == FP8G) { - if (MODE == FP4 || (MODE == MXFP8 && (kt & 1)) || (MODE == FP8G && kt == GS / 16 - 1)) { - const int blk = MODE == FP4 ? kt : kt / 2; -#pragma unroll - for (int j = 0; j < T::NT; ++j) { - float s0, s1; - if constexpr (MODE == FP8G) { - s0 = __uint_as_float(sq[j][0]); - s1 = __uint_as_float(sq[j][1]); - } else if constexpr (MODE == FP4) { - s0 = e4m3f(static_cast(sq[j][0] >> (8 * blk))); - s1 = e4m3f(static_cast(sq[j][1] >> (8 * blk))); - } else { - s0 = e8m0f(static_cast(sq[j][0] >> (8 * blk))); - s1 = e8m0f(static_cast(sq[j][1] >> (8 * blk))); - } -#pragma unroll - for (int i = 0; i < T::MT; ++i) -#pragma unroll - for (int e = 0; e < 4; ++e) acc[i][j][e] = __fmaf_rn(d[i][j][e], (e & 1) ? s1 : s0, acc[i][j][e]); - } - } - } - } - } - wait<0>(); - __syncthreads(); - if constexpr (FUSE) { - if (SK > 1) { -#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] = tot[i][j][e] + acc[i][j][e]; - } - } - if constexpr (CLUSTER) { -#if __CUDA_ARCH__ < 900 - __trap(); -#else - 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(); - 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; - if (SK > 1 && !CLUSTER && !FUSE) { - float* dst = part + (static_cast(slice) * M + row) * N + col; - if (col < N) dst[0] = acc[i][j][2 * h]; - if (col + 1 < N) dst[1] = acc[i][j][2 * h + 1]; - continue; - } - const float v0 = acc[i][j][2 * h] * scale, v1 = acc[i][j][2 * h + 1] * scale; - 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); - } - } - } - } +// The shared down has one K slice; only the output address changes from the existing BF16-weight product. +#define TF_SHARED_DOWN(NAME, BM, BN) \ +extern "C" __global__ void __launch_bounds__(128) NAME( \ + const __nv_bfloat16* x, const unsigned char* w, float scale, float* out, \ + int M, int N, int K, int npad, int ldx, int group, int ldy) { \ + qmmf_body( \ + x, w, nullptr, scale, out, nullptr, M, N, K, 1, npad, ldx, group, ldy); \ } +TF_SHARED_DOWN(tf_qmmf_shared16_half, 16, 32) +TF_SHARED_DOWN(tf_qmmf_shared16, 16, 64) +TF_SHARED_DOWN(tf_qmmf_shared32, 32, 64) +TF_SHARED_DOWN(tf_qmmf_shared64, 64, 64) +#undef TF_SHARED_DOWN template __global__ void reduce_kernel(const float* __restrict__ part, void* __restrict__ out, long long total, int SK, @@ -406,6 +200,17 @@ template __global__ void qmmf_kernel(c template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +// Prompt tiles retain the same m16n8k16 chains and fused slice boundaries as the 64x64 instances. +#define BF16_PROMPT(F32, FUSE) \ + template __global__ void qmmf_kernel( \ + const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, \ + void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +BF16_PROMPT(false, false) +BF16_PROMPT(false, true) +BF16_PROMPT(true, false) +BF16_PROMPT(true, true) +#undef BF16_PROMPT + template __global__ void reduce_kernel(const float* __restrict__, void* __restrict__, long long, int, float); template __global__ void reduce_kernel(const float* __restrict__, void* __restrict__, long long, int, float); diff --git a/zig/kernels/cuda/qmmf_body.cuh b/zig/kernels/cuda/qmmf_body.cuh new file mode 100644 index 0000000000..ad66314c56 --- /dev/null +++ b/zig/kernels/cuda/qmmf_body.cuh @@ -0,0 +1,282 @@ +// Strided stores are instantiated only by the shared-slot entries. + using T = Tile; +#ifdef TF_QMMF_STRIDED + static_assert(BN == 64 || BN == 32, "a block reads one 64-column tile of words and scales, or its half"); +#else + static_assert(BN == 64 || BN == 32 || (MODE == BF16 && BN == 128), "packed column tiles"); +#endif + 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, slice_groups = KG / SK, per = FUSE ? KG : slice_groups; + const int2 at = tile_of(blockIdx.x, M, N, BM, BN, group); + const int m0 = at.x, n0 = at.y, slice = FUSE ? 0 : blockIdx.z, g0 = slice * per; + +#ifndef TF_QMMF_STRIDED + uint64_t* ready = reinterpret_cast(buf + STAGES * T::STAGE); +#if __CUDA_ARCH__ >= 1200 + if constexpr (T::BARRIERS > 0) { + if (tid == 0) { +#pragma unroll + for (int s = 0; s < STAGES; ++s) + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" :: "r"(smem(ready + s)) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + __syncthreads(); + } +#endif +#endif + auto stage = [&](int s) { return buf + s * T::STAGE; }; + auto load = [&](int s, int g) { + 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 = MODE == FP4 ? 64 * GS / 2 : MODE == BF16 ? 64 * GS * 2 : 64 * GS; +#if __CUDA_ARCH__ >= 1200 && !defined(TF_QMMF_STRIDED) + if constexpr (T::BARRIERS > 0) { + if (tid == 0) { + const uint32_t bar = smem(ready + s); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" + :: "r"(bar), "r"(T::W) : "memory"); +#pragma unroll + for (int t = 0; t < BN / 64; ++t) { + const void* src = w + (static_cast(n0 / 64 + t) * KG + g) * TILE_BYTES; + asm volatile("cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" + :: "r"(smem(pw + t * TILE_BYTES)), "l"(src), "r"(TILE_BYTES), "r"(bar) : "memory"); + } + } + } else +#endif + if constexpr (BN == 32) { + // half a packed 64-column tile: its n8 tiles 0-3 or 4-7, contiguous in every format's word order + const int half = (n0 % 64) / 32; + for (int c = tid; c < T::W / 16; c += T::THREADS) + cp16(pw + c * 16, w + (static_cast(n0 / 64) * KG + g) * TILE_BYTES + half * (TILE_BYTES / 2) + c * 16); + if constexpr (T::S > 0) { + unsigned char* ps = pw + T::W; + for (int c = tid; c < T::S / 16; c += T::THREADS) + cp16(ps + c * 16, bs + (static_cast(n0 / 64) * KG + g) * (2 * T::S) + half * T::S + c * 16); + } + } else { + for (int c = tid; c < T::W / 16; c += T::THREADS) { + const int t = c / (TILE_BYTES / 16), off = c % (TILE_BYTES / 16); + cp16(pw + c * 16, w + (static_cast(n0 / 64 + t) * KG + g) * TILE_BYTES + off * 16); + } + if constexpr (T::S > 0) { + unsigned char* ps = pw + T::W; + for (int c = tid; c < T::S / 16; c += T::THREADS) + cp16(ps + c * 16, bs + (static_cast(n0 / 64) * KG + g) * T::S + c * 16); + } + } + }; + + float acc[T::MT][T::NT][4]; + float tot[FUSE ? T::MT : 1][FUSE ? T::NT : 1][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; +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < per) load(s, g0 + s); + commit(); + } + for (int it = 0; it < per; ++it) { + if constexpr (FUSE) { + if (it > 0 && it % slice_groups == 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) { + tot[i][j][e] = it == slice_groups ? acc[i][j][e] : tot[i][j][e] + acc[i][j][e]; + acc[i][j][e] = 0.0f; + } + } + } + wait(); +#if __CUDA_ARCH__ >= 1200 && !defined(TF_QMMF_STRIDED) + if constexpr (T::BARRIERS > 0) { + const uint32_t bar = smem(ready + it % STAGES), phase = (it / STAGES) & 1; + asm volatile("{ .reg .pred done; WAIT: mbarrier.try_wait.parity.shared::cta.b64 done, [%0], %1;" + " @!done bra WAIT; }" :: "r"(bar), "r"(phase) : "memory"); + } +#endif + __syncthreads(); + const int next = it + STAGES - 1; + if (next < per) load(next % STAGES, g0 + next); + commit(); + const unsigned char* p = stage(it % STAGES); + const unsigned char* pw = p + T::X; + const uint8_t* ps = p + T::X + T::W; + float d[T::MT][T::NT][4]; + uint4 wq[T::NT]; + uint32_t sq[T::NT][2]; +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const int jj = wn * T::NT + j; + const int col = wn * (BN / WN) + j * 8 + (lane & 3) * 2; + if constexpr (MODE == FP4 || MODE == FP8G) { + const uint2 v = *reinterpret_cast(ps + col * 4); + sq[j][0] = v.x; + sq[j][1] = v.y; + } else if constexpr (MODE == MXFP8) { + const uint32_t v = *reinterpret_cast(ps + col * 2); + sq[j][0] = v & 0xFFFFu; + sq[j][1] = v >> 16; + } + if constexpr (MODE == FP4) { + const uint2 u = reinterpret_cast(pw)[jj * 32 + lane]; + wq[j] = make_uint4(u.x, u.y, 0u, 0u); + } else if constexpr (MODE != BF16) { + wq[j] = reinterpret_cast(pw)[jj * 32 + lane]; + } + } +#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) { + uint32_t b0, b1; + if constexpr (MODE == FP4) { + const uint32_t word = kt < 2 ? wq[j].x : wq[j].y; + b0 = fp4pair(word, (kt & 1) * 8); + b1 = fp4pair(word, (kt & 1) * 8 + 4); + } else if constexpr (MODE == BF16) { + const uint2 v = reinterpret_cast(pw)[((wn * T::NT + j) * 4 + kt) * 32 + lane]; + b0 = v.x; + b1 = v.y; + } else { + const uint32_t word = comp(wq[j], kt); + b0 = fp8pair(word & 0xFFFFu); + b1 = fp8pair(word >> 16); + } +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + if constexpr (MODE == FP8 || MODE == BF16) mma(acc[i][j], a[i], b0, b1); + else if (MODE == FP4 || (MODE == MXFP8 && (kt & 1) == 0) || (MODE == FP8G && kt == 0)) + mma0(d[i][j], a[i], b0, b1); + else mma(d[i][j], a[i], b0, b1); + } + } + if constexpr (MODE == FP4 || MODE == MXFP8 || MODE == FP8G) { + if (MODE == FP4 || (MODE == MXFP8 && (kt & 1)) || (MODE == FP8G && kt == GS / 16 - 1)) { + const int blk = MODE == FP4 ? kt : kt / 2; +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + float s0, s1; + if constexpr (MODE == FP8G) { + s0 = __uint_as_float(sq[j][0]); + s1 = __uint_as_float(sq[j][1]); + } else if constexpr (MODE == FP4) { + s0 = e4m3f(static_cast(sq[j][0] >> (8 * blk))); + s1 = e4m3f(static_cast(sq[j][1] >> (8 * blk))); + } else { + s0 = e8m0f(static_cast(sq[j][0] >> (8 * blk))); + s1 = e8m0f(static_cast(sq[j][1] >> (8 * blk))); + } +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int e = 0; e < 4; ++e) acc[i][j][e] = __fmaf_rn(d[i][j][e], (e & 1) ? s1 : s0, acc[i][j][e]); + } + } + } + } + } + wait<0>(); + __syncthreads(); + if constexpr (FUSE) { + if (SK > 1) { +#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] = tot[i][j][e] + acc[i][j][e]; + } + } + if constexpr (CLUSTER) { +#if __CUDA_ARCH__ < 900 + __trap(); +#else + 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(); + 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; + if (SK > 1 && !CLUSTER && !FUSE) { + float* dst = part + (static_cast(slice) * M + row) * N + col; + if (col < N) dst[0] = acc[i][j][2 * h]; + if (col + 1 < N) dst[1] = acc[i][j][2 * h + 1]; + continue; + } + const float v0 = acc[i][j][2 * h] * scale, v1 = acc[i][j][2 * h + 1] * scale; + if (F32) { +#ifdef TF_QMMF_STRIDED + float* dst = reinterpret_cast(out) + static_cast(row) * ldy + col; +#else + float* dst = reinterpret_cast(out) + static_cast(row) * N + col; +#endif + if (col < N) dst[0] = v0; + if (col + 1 < N) dst[1] = v1; + } else { +#ifdef TF_QMMF_STRIDED + __nv_bfloat16* dst = reinterpret_cast<__nv_bfloat16*>(out) + static_cast(row) * ldy + col; +#else + __nv_bfloat16* dst = reinterpret_cast<__nv_bfloat16*>(out) + static_cast(row) * N + col; +#endif + 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); + } + } + } + } diff --git a/zig/src/core/contract/flashnext_cuda_records.zig b/zig/src/core/contract/flashnext_cuda_records.zig index a1d63d1c23..a562ae82ab 100644 --- a/zig/src/core/contract/flashnext_cuda_records.zig +++ b/zig/src/core/contract/flashnext_cuda_records.zig @@ -2,6 +2,7 @@ const CudaRecord = @import("registry.zig").CudaRecord; pub const cuda = [_]CudaRecord{ + .{ .launch = .{ .name = "flashnext_batch fn_copy_cols", .symbol = "fn_copy_cols", .image = "flashnext_batch", .threads = 256 }, .usage = @splat(.{ .regs = 12, .static = 0, .bound = 0 }), .later_pending = true }, .{ .launch = .{ .name = "flashnext_gdn flashnext_gdn_chain", .symbol = "flashnext_gdn_chain", .image = "flashnext_gdn", .threads = 1024 }, .usage = .{ .{ .regs = 51, .static = 2060, .bound = 1024 }, .{ .regs = 56, .static = 2060, .bound = 1024 }, .{ .regs = 56, .static = 2060, .bound = 1024 }, .{ .regs = 62, .static = 2060, .bound = 1024 }, .{ .regs = 62, .static = 2060, .bound = 1024 } } }, .{ .launch = .{ .name = "flashnext_hc fn_add_streams", .symbol = "fn_add_streams", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 35, .static = 0, .bound = 0 }, .{ .regs = 35, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks", .symbol = "fn_attn_chunks", .image = "flashnext_attn", .threads = 128, .dynamic = 43520, .opt_in = true }, .usage = .{ .{ .regs = 72, .static = 640, .bound = 128 }, .{ .regs = 72, .static = 640, .bound = 128 }, .{ .regs = 72, .static = 640, .bound = 128 }, .{ .regs = 72, .static = 640, .bound = 128 }, .{ .regs = 72, .static = 640, .bound = 128 } } }, @@ -31,6 +32,7 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "flashnext_mtp fn_row_best", .symbol = "fn_row_best", .image = "flashnext_mtp", .threads = 1024 }, .usage = .{ .{ .regs = 17, .static = 392, .bound = 1024 }, .{ .regs = 17, .static = 392, .bound = 1024 }, .{ .regs = 17, .static = 392, .bound = 1024 }, .{ .regs = 21, .static = 392, .bound = 1024 }, .{ .regs = 21, .static = 392, .bound = 1024 } } }, .{ .launch = .{ .name = "flashnext_hc fn_shared_scatter", .symbol = "fn_shared_scatter", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 12, .static = 0, .bound = 0 }, .{ .regs = 12, .static = 0, .bound = 0 }, .{ .regs = 12, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "flashnext_hc fn_shared_swiglu", .symbol = "fn_shared_swiglu", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 17, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_ep_picks", .symbol = "fn_ep_picks", .image = "flashnext_batch", .threads = 256 }, .usage = .{ .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "flashnext_batch fn_shift_window", .symbol = "fn_shift_window", .image = "flashnext_batch", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "flashnext_batch fn_sum_slices", .symbol = "fn_sum_slices", .image = "flashnext_batch", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 38, .static = 0, .bound = 0 }, .{ .regs = 38, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "norms tf_rms_rows_sums_tiles", .symbol = "tf_rms_rows_sums_tiles", .image = "norms", .threads = 256 }, .usage = .{ .{ .regs = 31, .static = 20612, .bound = 256 }, .{ .regs = 48, .static = 20612, .bound = 256 }, .{ .regs = 48, .static = 20612, .bound = 256 }, .{ .regs = 48, .static = 20612, .bound = 256 }, .{ .regs = 48, .static = 20612, .bound = 256 } } }, diff --git a/zig/src/core/contract/flashnext_heads_records.zig b/zig/src/core/contract/flashnext_heads_records.zig new file mode 100644 index 0000000000..c5f94d1855 --- /dev/null +++ b/zig/src/core/contract/flashnext_heads_records.zig @@ -0,0 +1,65 @@ +//! Flash Next's DeltaNet twins on one rank's 8 key and 24 value heads, nvcc 13.0's usage from every SM's image. +const registry = @import("registry.zig"); + +/// nvcc 13.0's usage (kernel-contract-limits on a GB10); later compilers' sets are pending. +pub const records = [_]registry.CudaRecord{ + .{ .launch = .{ .name = "flashnext_batch fn_gdn_front_h8", .symbol = "fn_gdn_front_h8", .image = "flashnext_batch", .threads = 256 }, .usage = .{ + .{ .regs = 26, .static = 1024, .bound = 256 }, + .{ .regs = 26, .static = 1024, .bound = 256 }, + .{ .regs = 26, .static = 1024, .bound = 256 }, + .{ .regs = 32, .static = 1024, .bound = 256 }, + .{ .regs = 32, .static = 1024, .bound = 256 }, + }, .later_pending = true }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_norm_h8", .symbol = "fn_gdn_norm_h8", .image = "flashnext_batch", .threads = 128 }, .usage = .{ + .{ .regs = 24, .static = 528, .bound = 128 }, + .{ .regs = 26, .static = 528, .bound = 128 }, + .{ .regs = 26, .static = 528, .bound = 128 }, + .{ .regs = 24, .static = 528, .bound = 128 }, + .{ .regs = 24, .static = 528, .bound = 128 }, + }, .later_pending = true }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_rows_h8", .symbol = "fn_gdn_rows_h8", .image = "flashnext_batch", .threads = 128 }, .usage = .{ + .{ .regs = 64, .static = 0, .bound = 128 }, + .{ .regs = 64, .static = 0, .bound = 128 }, + .{ .regs = 64, .static = 0, .bound = 128 }, + .{ .regs = 64, .static = 0, .bound = 128 }, + .{ .regs = 64, .static = 0, .bound = 128 }, + }, .later_pending = true }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_rows_to_h8", .symbol = "fn_gdn_rows_to_h8", .image = "flashnext_batch", .threads = 128 }, .usage = .{ + .{ .regs = 64, .static = 0, .bound = 128 }, + .{ .regs = 64, .static = 0, .bound = 128 }, + .{ .regs = 64, .static = 0, .bound = 128 }, + .{ .regs = 64, .static = 0, .bound = 128 }, + .{ .regs = 64, .static = 0, .bound = 128 }, + }, .later_pending = true }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_chain_tpr_h8", .symbol = "fn_gdn_chain_tpr_h8", .image = "flashnext_batch", .threads = 256, .dynamic = 49664, .opt_in = true }, .usage = .{ + .{ .regs = 129, .static = 0, .bound = 256 }, + .{ .regs = 129, .static = 0, .bound = 256 }, + .{ .regs = 129, .static = 0, .bound = 256 }, + .{ .regs = 154, .static = 0, .bound = 256 }, + .{ .regs = 154, .static = 0, .bound = 256 }, + }, .later_pending = true }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_fused_h8", .symbol = "fn_gdn_fused_h8", .image = "flashnext_batch", .threads = 384, .dynamic = 82496, .opt_in = true }, .usage = .{ + .{ .regs = 138, .static = 0, .bound = 384 }, + .{ .regs = 142, .static = 0, .bound = 384 }, + .{ .regs = 142, .static = 0, .bound = 384 }, + .{ .regs = 166, .static = 0, .bound = 384 }, + .{ .regs = 166, .static = 0, .bound = 384 }, + }, .later_pending = true }, +}; +test "Flash Next DeltaNet twins: measured on every SM, within its limits, launched as their full originals are" { + const std = @import("std"); + const t = std.testing; + try t.expectEqual(@as(usize, 6), records.len); + try t.expectEqual(@as(usize, 0), try registry.checkCuda(t.allocator, &records, false)); + for (records) |r| { + try t.expect(r.usage[0].regs != 0 and r.later_pending); + for (registry.cuda) |measured| try t.expect(!std.mem.eql(u8, r.launch.symbol, measured.launch.symbol)); + const twin = r.launch.symbol[0 .. r.launch.symbol.len - "_h8".len]; + const full = for (registry.cuda) |measured| { + if (std.mem.eql(u8, measured.launch.symbol, twin)) break measured.launch; + } else return error.NoFullRecord; + try t.expectEqual(full.threads, r.launch.threads); + try t.expectEqual(full.dynamic, r.launch.dynamic); + try t.expectEqual(full.opt_in, r.launch.opt_in); + } +} diff --git a/zig/src/core/contract/moetile_records.zig b/zig/src/core/contract/moetile_records.zig new file mode 100644 index 0000000000..be36d37656 --- /dev/null +++ b/zig/src/core/contract/moetile_records.zig @@ -0,0 +1,30 @@ +//! Shared-down launch geometry and nvcc 13.0 receipts; later compiler receipts remain pending. +const std = @import("std"); +const registry = @import("registry.zig"); +const limits = @import("limits.zig"); + +fn usage(regs: [5]u32) [5]limits.Usage { + var out: [5]limits.Usage = undefined; + for (&out, regs) |*u, count| u.* = .{ .regs = count, .static = 0, .bound = 128 }; + return out; +} +pub const records = [_]registry.CudaRecord{ + .{ .launch = .{ .name = "qmmf shared down half16 stride", .symbol = "tf_qmmf_shared16_half", .image = "qmmf", .threads = 128, .dynamic = 24576, .min_sm = 80 }, .usage = usage(.{ 64, 64, 64, 40, 40 }), .later_pending = true }, + .{ .launch = .{ .name = "qmmf shared down 16 stride", .symbol = "tf_qmmf_shared16", .image = "qmmf", .threads = 128, .dynamic = 40960, .min_sm = 80 }, .usage = usage(.{ 79, 79, 79, 38, 38 }), .later_pending = true }, + .{ .launch = .{ .name = "qmmf shared down 32 stride", .symbol = "tf_qmmf_shared32", .image = "qmmf", .threads = 128, .dynamic = 49152, .min_sm = 80 }, .usage = usage(.{ 93, 93, 93, 56, 56 }), .later_pending = true }, + .{ .launch = .{ .name = "qmmf shared down 64 stride", .symbol = "tf_qmmf_shared64", .image = "qmmf", .threads = 128, .dynamic = 65536, .opt_in = true, .min_sm = 80 }, .usage = usage(.{ 118, 118, 118, 72, 72 }), .later_pending = true }, +}; + +test "shared down records hold nvcc 13.0 resources and keep later compilers pending" { + try std.testing.expectEqual(@as(usize, 0), try registry.checkCuda(std.testing.allocator, &records, true)); + for (records) |record| { + try std.testing.expectEqual(record.usage, registry.usageFor(record, .{ .major = 13, .minor = 0 })); + for (record.usage) |u| { + try std.testing.expect(u.regs > 0); + try std.testing.expectEqual(@as(u32, 0), u.static); + try std.testing.expectEqual(@as(u32, 128), u.bound); + } + for ([_]limits.Nvcc{ .{ .major = 13, .minor = 2 }, .{ .major = 13, .minor = 3 } }) |release| + try std.testing.expectEqual(registry.none, registry.usageFor(record, release)); + } +} diff --git a/zig/src/core/contract/qmmf_prompt_records.zig b/zig/src/core/contract/qmmf_prompt_records.zig new file mode 100644 index 0000000000..48873411de --- /dev/null +++ b/zig/src/core/contract/qmmf_prompt_records.zig @@ -0,0 +1,40 @@ +//! BF16 prompt launches on sm_12x: nvcc 13.0's usage on sm_120 and sm_121; later releases remain pending. +const std = @import("std"); +const registry = @import("registry.zig"); +const limits = @import("limits.zig"); + +fn usage(comptime regs: u32) [5]limits.Usage { + const off: limits.Usage = .{ .regs = 0, .static = 0 }; + const on: limits.Usage = .{ .regs = regs, .static = 0, .bound = 256 }; + return .{ off, off, off, on, on }; +} + +fn record(comptime wide: bool, comptime fused: bool) registry.CudaRecord { + return .{ + .launch = .{ + .name = "qmmf BF16 prompt 128x128 " ++ (if (wide) "fp32" else "bf16") ++ (if (fused) " fused slices" else " unsplit"), + .symbol = std.fmt.comptimePrint("_ZN7tf_qmmf11qmmf_kernelILi4ELi128ELi128ELi2ELi4ELi3ELb{d}ELb0ELb{d}EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .{ @intFromBool(wide), @intFromBool(fused) }), + .image = "qmmf", + .threads = 256, + .dynamic = 98432, + .opt_in = true, + .min_sm = 120, + }, + .usage = usage(if (fused) 255 else 95), + .later_pending = true, + }; +} + +pub const records = [_]registry.CudaRecord{ record(false, false), record(false, true), record(true, false), record(true, true) }; + +test "qmmf prompt launch records hold nvcc 13.0's usage on sm_12x, later releases pending" { + try std.testing.expectEqual(@as(usize, 0), try registry.checkCuda(std.testing.allocator, &records, true)); + for (records) |r| { + try std.testing.expectEqual(@as(u32, 256), r.launch.threads); + try std.testing.expectEqual(@as(u32, 98432), r.launch.dynamic); + try std.testing.expect(r.launch.opt_in and !r.geometry_pending); + try std.testing.expect(registry.usageFor(r, .{ .major = 13, .minor = 0 })[4].regs > 0); + for ([_]limits.Nvcc{ .{ .major = 13, .minor = 2 }, .{ .major = 13, .minor = 3 } }) |rel| + try std.testing.expectEqual(registry.none, registry.usageFor(r, rel)); + } +} diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 8e8e10c965..3686283f3a 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -206,7 +206,6 @@ pub const cuda = @import("flashnext_cuda_records.zig").cuda ++ [_]CudaRecord{ .{ .launch = .{ .name = "experts3 tf_experts3_widen", .symbol = "tf_experts3_widen", .image = "experts3", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "experts3 tf_experts3_planes", .symbol = "tf_experts3_planes", .image = "experts3", .threads = 256 }, .usage = .{ .{ .regs = 26, .static = 0, .bound = 0 }, .{ .regs = 28, .static = 0, .bound = 0 }, .{ .regs = 28, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 } } }, } ++ @import("registry_media.zig").cuda ++ @import("nemotron_prompt_records.zig").cuda ++ @import("glm_cuda_records.zig").cuda ++ @import("registry_shared_attention.zig").records(CudaRecord) ++ @import("gemma_cuda_records.zig").cuda ++ @import("spark_cuda_records.zig").cuda ++ @import("row_cuda_records.zig").cuda ++ @import("flashnext_tile_records.zig").cuda ++ @import("window_rows_records.zig").cuda ++ @import("modelopt_lane_records.zig").cuda ++ @import("nvfp4_fact_records.zig").cuda ++ @import("lane_half_records.zig").cuda ++ @import("pocket_records.zig").cuda ++ @import("minimax_cuda_records.zig").cuda ++ @import("moss_image_cuda_records.zig").cuda; - /// The nvcc release `cuda`'s usage records come from: the compiler that builds the release fatbins. pub const cuda_release: limits.Nvcc = .{ .major = 13, .minor = 0 }; /// One launch's usage under a later nvcc, by the record's name. @@ -561,13 +560,11 @@ test "the registered launches fit every Apple family and every SM, and each plan test { _ = @import("registry_sets_test.zig"); _ = @import("nvfp4_tile_pending.zig"); -} -const paging = @import("paging_pending.zig"); -test { _ = paging; _ = @import("minimax_records.zig"); _ = @import("nemotron_ssd_records.zig"); } +const paging = @import("paging_pending.zig"); /// The Metal collector keeps measured slots and the restored vision declarations. pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("metal_pair_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal ++ @import("lfm2_records.zig").records ++ @import("nemotron_experts_records.zig").records ++ @import("nemotron_pages_records.zig").records ++ @import("nemotron_attention_records.zig").records ++ @import("nemotron_ssd_records.zig").records ++ @import("flashnext_attention_records.zig").records ++ @import("spark25_records.zig").records ++ @import("gemma_attention_records.zig").records ++ @import("moss_image_records.zig").records ++ @import("dsv41_metal_records.zig").records ++ @import("minimax_records.zig").records; @@ -575,11 +572,11 @@ const learning = @import("registry_learning.zig"); pub const gradient_pending = learning.gradient_pending; pub const qwen_learning_pending = learning.qwen_learning_pending; pub const glm_pending = @import("glm_sampling_pending.zig"); -pub const landed_cuda = cuda ++ @import("dory_cuda_records.zig").cuda ++ gradient_pending ++ qwen_learning_pending ++ paging.extra ++ @import("qwen_served_records.zig").records ++ cuda_pending ++ @import("qwen_decode_pending.zig").records ++ glm_pending.sampling ++ @import("glm_kv_pending.zig").records ++ @import("exl3_grouped_records.zig").records ++ @import("proposal_sampling_records.zig").records ++ @import("flashnext_pages_pending.zig").records ++ @import("glm_f32w_pending.zig").records ++ @import("penalty_sampling_records.zig").records; +pub const landed_cuda = cuda ++ @import("dory_cuda_records.zig").cuda ++ gradient_pending ++ qwen_learning_pending ++ paging.extra ++ @import("qwen_served_records.zig").records ++ cuda_pending ++ @import("qwen_decode_pending.zig").records ++ glm_pending.sampling ++ @import("glm_kv_pending.zig").records ++ @import("exl3_grouped_records.zig").records ++ @import("proposal_sampling_records.zig").records ++ @import("flashnext_pages_pending.zig").records ++ @import("flashnext_heads_records.zig").records ++ @import("glm_f32w_pending.zig").records ++ @import("penalty_sampling_records.zig").records; test { _ = learning; } -pub const all_cuda = @import("gb10_many_records.zig").records ++ @import("nvfp4_tile_pending.zig").records ++ landed_cuda ++ dsv41_pending ++ @import("inc4_pending.zig").trees ++ @import("inc4_pending.zig").rows ++ @import("affine_verify_records.zig").records; +pub const all_cuda = @import("qmmf_prompt_records.zig").records ++ @import("moetile_records.zig").records ++ @import("gb10_many_records.zig").records ++ @import("nvfp4_tile_pending.zig").records ++ landed_cuda ++ dsv41_pending ++ @import("inc4_pending.zig").trees ++ @import("inc4_pending.zig").rows ++ @import("affine_verify_records.zig").records; pub const dsv41_pending = @import("registry_union.zig").extra(landed_cuda, @import("dsv41_pending_1.zig").records ++ @import("dsv41_pending_2.zig").records); test { _ = @import("qwen4_cuda_receipts_test.zig"); @@ -589,6 +586,9 @@ test { _ = @import("glm_f32w_pending.zig"); _ = @import("registry_pending_test.zig"); _ = @import("flashnext_pages_pending.zig"); + _ = @import("moetile_records.zig"); + _ = @import("qmmf_prompt_records.zig"); + _ = @import("flashnext_heads_records.zig"); } test "unmeasured CUDA records cannot pass a compiled resource check" { for (cuda_pending) |r| { diff --git a/zig/src/cuda/qmmf.md b/zig/src/cuda/qmmf.md index 82f5c5a1ad..08fe1fa30d 100644 --- a/zig/src/cuda/qmmf.md +++ b/zig/src/cuda/qmmf.md @@ -17,3 +17,94 @@ row count, activation stride, partial output tiles, fused reduction, cluster reduction and explicit split reduction. The fixture is synthetic; model correctness and multi-device collective behavior require separate qualification. + +# BF16 prompt tiles + +On sm_12x, BF16 calls with at least 512 rows, N >= 128, K >= 512 and +128-column padded weights use a 128x128 CTA: WM=2, WN=4, 256 threads, +three K64 stages. Dynamic shared memory is 98,432 bytes, including the +stage barriers. Other formats, smaller grids and older devices use the +existing tiles. This threshold is a conservative screening choice; the +GB10 timing receipt must confirm that it pays. + +Each stage copies two independent packed 64-column weight tiles with +`cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes`. Their +source stride remains `(K / 64) * 8192` bytes. One thread issues the bulk +copies and accounts for 16,384 bytes at that stage's barrier. Swizzled +activation rows use the existing coalesced 16-byte `cp.async` loads, +including zero fill past M and the caller's activation stride. Each warp +waits for both copy paths before it reads the stage. The CTA barrier +before the next prefetch also prevents reuse while another warp still +reads the old stage. Barrier parity changes each time a stage is reused. +The larger instances retain a 16-byte copy fallback in older SM images; +the launcher loads and selects them only on sm_12x. + +## Byte contract + +The BF16 accumulator starts at positive FP32 zero for each K slice. +Groups visit K in ascending 64-element steps. Each group executes the +same four `mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32` +instructions, at K offsets 0, 16, 32 and 48, on the same operand fragments. +The hardware's internal K16 order is unchanged because the instruction +and fragment mapping are unchanged. Only independent output fragments +move between warps and CTAs; the stages change operand arrival time. + +`Weight.sk` still overrides `splitK(N, K)`. It does not depend on M or +the CTA dimensions. A fused call resets the accumulator after exactly +`K / (64 * SK)` groups, retains slice zero, and adds each later slice in +ascending order with FP32 additions. This matches the old fused, +cluster and explicit reduction paths. The epilogue multiplies the final +FP32 sum by `Weight.scale` once, then either stores FP32 or applies one +round-to-nearest-even BF16 conversion. There is no new intermediate +rounding, split or reduction tree. Whole decoder layers retain their +shape-driven K slices; explicit pins let the comparison test exercise +the same slices on each tile geometry. + +## Device checks and timing + +Rebuild the qmmf image and runner together. For an existing complete +fatbin directory `FATBINS`, compile only the changed image, then embed it: + +```sh +nvcc -fatbin -O3 -std=c++20 --expt-relaxed-constexpr \ + -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ \ + -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ \ + -gencode=arch=compute_121,code=sm_121 -Izig/kernels/cuda \ + zig/kernels/cuda/qmmf.cu -o "$FATBINS/qmmf.fatbin" +zig build compile-cuda-tests -Dfatbins="$FATBINS" -Doptimize=ReleaseFast +``` + +The build step emits the runner without installing it. To install the +runner with the same images, use `zig build -Dfatbins="$FATBINS" +-Doptimize=ReleaseFast`, then run: + +```sh +zig-out/bin/tf-cuda-test qmmf-prompt +zig-out/bin/tf-cuda-test qmmf-prompt-bench +``` + +The device test compares all bytes against the unchanged lane kernels +for current whole-layer N/K = 16480/2560, 13952/2560, 2560/6144, +640/2560, 2560/640 and HC up 10240/320, both output types, SK=1 and the shape-driven split, +and M=1/16/64/65/127/128/129/511/512/1536/4218. It forces the new CTA +below the production threshold, checks production dispatch separately, +checks output guards and repeats the largest case five times. Small +SK=4 calls exercise both cluster and explicit reduction baselines. +Inputs use varied BF16 exponents and signs to expose a changed sum order. +A separate N=129 tail fixture plants NaNs, infinities, signed zeros and +subnormals in the inputs, and checks that inputs and weights stay unchanged. + +The benchmark checks equality first, prints loaded register/shared-memory +usage and occupancy, then measures three A-B-A rounds with CUDA timing +events, three warmups and twenty calls per sample. It prints median +old/new times for all three shapes and a weighted 36/12/48-layer total. +Run on an idle GB10. These warm microbenchmarks do not prove the full +served mixer time or the complete kernel contract. + +Each output tile now reuses operands over twice the rows and columns: +ideal input traffic per output halves, CTA count falls by about four, +and each resident CTA has eight warps instead of four. Bulk weight copies +also reduce load instructions. These are reasons to test a 1.6x speedup, +not a measured result. Register pressure, the three-stage shared-memory +allocation and hardware MMA throughput can limit the gain. Device byte +checks, timings and compile-limit receipts remain required before landing. diff --git a/zig/src/cuda/qmmf.zig b/zig/src/cuda/qmmf.zig index 0b4774e55d..d95eb4403b 100644 --- a/zig/src/cuda/qmmf.zig +++ b/zig/src/cuda/qmmf.zig @@ -53,13 +53,30 @@ pub fn bucket(m: usize) u32 { } /// A packed projection on the device: codes or words, scale tiles, the tensor scale, its shape, its input scale (act). -pub const Weight = struct { mode: Mode, codes: u64, scales: u64, scale: f32 = 1.0, n: u32, k: u32, npad: u32, act: f32 = 0 }; +pub const Weight = struct { mode: Mode, codes: u64, scales: u64, scale: f32 = 1.0, n: u32, k: u32, npad: u32, act: f32 = 0, sk: u32 = 0 }; -/// What one call launches, from the shapes alone: tile, slices, cluster or reduce (qmmf_cuda's choice). +/// `w` with its K slices pinned to `sk`: 1, 2, 4 or 8, dividing its groups and leaving at least 8 a slice. +pub fn pinned(w: Weight, sk: u32) error{BadSplit}!Weight { + const groups = w.k / group; + if (sk == 0 or sk > 8 or sk & (sk - 1) != 0 or groups % sk != 0 or (sk != 1 and groups / sk < 8)) return error.BadSplit; + var out = w; + out.sk = sk; + return out; +} + +/// The base tile and slice meeting from shapes alone; Lane may enlarge a BF16 prompt tile on sm_12x. pub const Plan = struct { bm: u32, sk: u32, fused: bool, cluster: bool, reduce: bool }; pub fn plan(m: usize, n: usize, k: usize, clusters: bool) Plan { - const sk = splitK(n, k); + return planSlices(m, splitK(n, k), clusters); +} + +/// plan() for a packed weight: its pinned K slices if any, else its shape's. +pub fn planFor(m: usize, w: Weight, clusters: bool) Plan { + return planSlices(m, if (w.sk != 0) w.sk else splitK(w.n, w.k), clusters); +} + +fn planSlices(m: usize, sk: u32, clusters: bool) Plan { const fused = sk > 1 and m >= fused_rows; const cluster = !fused and sk > 1 and sk <= 8 and clusters; return .{ .bm = if (fused) 64 else bucket(m), .sk = sk, .fused = fused, .cluster = cluster, .reduce = !fused and sk > 1 and !cluster }; @@ -131,6 +148,20 @@ fn loadTiles(mod: Module, comptime mode: Mode, comptime wide: bool) ![7]Function return set; } +/// BF16 prompt tile: 128x128, eight warps, three 32 KiB stages and one barrier per stage. +pub const prompt_shared: u32 = 3 * (128 + 128) * group * 2 + 128; +pub const prompt_rows: usize = 512; + +pub fn promptSymbol(comptime out: Out, comptime fused: bool) [:0]const u8 { + return std.fmt.comptimePrint("_ZN7tf_qmmf11qmmf_kernelILi4ELi128ELi128ELi2ELi4ELi3ELb{d}ELb0ELb{d}EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .{ @intFromBool(out == .f32), @intFromBool(fused) }); +} + +/// Only large BF16 grids use the prompt tile; the packed allocation contains complete 128-column tiles. +pub fn promptEligible(m: usize, w: Weight) bool { + return m >= prompt_rows and w.mode == .bf16 and w.n >= 128 and + (w.k >= 512 or (w.n == 10240 and w.k == 320)) and w.npad % 128 == 0; +} + pub const Lane = struct { mod: Module, fns: [modes.len][7]Function, // per mode: lane16, lane16c, lane32, lane32c, lane64, lane64c, fused @@ -138,6 +169,8 @@ pub const Lane = struct { reduce: [2]Function, // by Out half: [2][2][2]Function, // by halfModeIndex, Out, then plain / cluster half_below: usize = half_below_default, + shared_down: [4]Function = undefined, + prompt: ?[2][2]Function = null, // Out, then plain / fused; loaded only on sm_12x clusters: bool, // sm_90 on: slices meet in a cluster's shared memory pub fn load(d: *const Driver, capability_major: u32) !Lane { @@ -162,6 +195,19 @@ pub const Lane = struct { try f.allowDynamicShared(halfShared(mode)); l.half[mi][oi][ci] = f; }; + for (@import("qmmf_shared.zig").symbols, &l.shared_down) |name, *f| { + f.* = try mod.function(name); + try f.allowDynamicShared(sharedBytes(.bf16, 64)); + } + if (capability_major == 12) { + var prompt: [2][2]Function = undefined; + inline for (.{ Out.bf16, Out.f32 }, 0..) |out, oi| inline for (.{ false, true }, 0..) |fused, fi| { + const f = try mod.function(promptSymbol(out, fused)); + try f.allowDynamicShared(prompt_shared); + prompt[oi][fi] = f; + }; + l.prompt = prompt; + } l.half_below = halfBelow(); return l; } @@ -172,7 +218,7 @@ pub const Lane = struct { /// Slice partials a call needs (fp32 [sk][m][n]) when it reduces them, else 0. pub fn partBytes(l: Lane, m: usize, w: Weight) usize { - const p = plan(m, w.n, w.k, l.clusters); + const p = planFor(m, w, l.clusters); return if (p.reduce) @as(usize, p.sk) * m * w.n * 4 else 0; } @@ -186,18 +232,28 @@ pub const Lane = struct { return l.matmulOut(s, x, ldx, m, w, y, .f32, part); } + pub const matmulF32Strided = @import("qmmf_shared.zig").matmulF32Strided; + + /// The actual prompt selection, including loaded instances and the pinned slice meeting plan. + pub fn usesPrompt(l: Lane, m: usize, w: Weight) bool { + const p = planFor(m, w, l.clusters); + return promptEligible(m, w) and !p.cluster and !p.reduce and l.prompt != null; + } + /// matmul() writing y as `out` says: bf16, or fp32. pub fn matmulOut(l: Lane, s: Stream, x: u64, ldx: usize, m: usize, w: Weight, y: u64, out: Out, part: u64) !void { - const p = plan(m, w.n, w.k, l.clusters); + const p = planFor(m, w, l.clusters); if (p.reduce and part == 0) return error.MissingLanePart; const set = (if (out == .f32) l.f32_fns else l.fns)[modeIndex(w.mode)]; const ti: usize = if (p.bm == 16) 0 else if (p.bm == 32) 1 else 2; - const rows_t = (m + p.bm - 1) / p.bm; + const prompt = l.usesPrompt(m, w); + const bm: usize = if (prompt) 128 else p.bm; + const rows_t = (m + bm - 1) / bm; const hm = halfModeIndex(w.mode); const half = hm != null and p.bm == 16 and !p.fused and rows_t * ((w.n + bn - 1) / bn) * p.sk < l.half_below; - const f = if (half) l.half[hm.?][@backingInt(out)][@intFromBool(p.cluster)] else if (p.fused) set[6] else set[ti * 2 + @intFromBool(p.cluster)]; - const cols: usize = if (half) 32 else bn; - const band = @max(1, @min(rows_t, (12 << 20) / (@as(usize, p.bm) * w.k * 2))); + const f = if (prompt) l.prompt.?[@backingInt(out)][@intFromBool(p.fused)] else if (half) l.half[hm.?][@backingInt(out)][@intFromBool(p.cluster)] else if (p.fused) set[6] else set[ti * 2 + @intFromBool(p.cluster)]; + const cols: usize = if (prompt) 128 else if (half) 32 else bn; + const band = @max(1, @min(rows_t, (12 << 20) / (bm * w.k * 2))); var args: launch_.Args = .{}; args.add(x); args.add(w.codes); @@ -214,8 +270,8 @@ pub const Lane = struct { args.add(@as(i32, @intCast(band))); try launch_.launch(f, .{ .grid = .{ .x = @intCast(rows_t * ((w.n + cols - 1) / cols)), .y = 1, .z = if (p.fused) 1 else p.sk }, - .block = .{ .x = threads, .y = 1, .z = 1 }, - .shared = if (half) halfShared(w.mode) else sharedBytes(w.mode, p.bm), + .block = .{ .x = if (prompt) 256 else threads, .y = 1, .z = 1 }, + .shared = if (prompt) prompt_shared else if (half) halfShared(w.mode) else sharedBytes(w.mode, p.bm), .cluster = if (p.cluster) .{ .x = 1, .y = 1, .z = p.sk } else null, }, s, &args); if (p.reduce) { @@ -303,3 +359,102 @@ test "instance names are qmmf.cu's" { try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi0ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp4, 64, false, true, false)); try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi0ELi32ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp4, 32, true, false, true)); } + +test "large BF16 tiles preserve pinned slices and leave decode and other formats unchanged" { + const w = try pinned(Weight{ .mode = .bf16, .codes = 0, .scales = 0, .n = 8240, .k = 2560, .npad = 8320 }, 1); + for ([_]usize{ 1, 16, 64, 65, 511, 512, 1536, 4218 }) |m| { + try std.testing.expectEqual(m >= prompt_rows, promptEligible(m, w)); + try std.testing.expectEqual(@as(u32, 1), planFor(m, w, true).sk); + } + var other = w; + other.mode = .fp8; + try std.testing.expect(!promptEligible(4218, other)); + try std.testing.expectEqual(@as(u32, 98432), prompt_shared); + const hc_up: Weight = .{ .mode = .bf16, .codes = 0, .scales = 0, .n = 10240, .k = 320, .npad = 10240 }; + try std.testing.expect(promptEligible(512, hc_up)); + try std.testing.expect(!promptEligible(511, hc_up)); + var narrow_up = hc_up; + narrow_up.n = 5120; + narrow_up.npad = 5120; + try std.testing.expect(!promptEligible(512, narrow_up)); + const out = try pinned(Weight{ .mode = .bf16, .codes = 0, .scales = 0, .n = 2560, .k = 3072, .npad = 2560 }, 4); + const p = planFor(4218, out, true); + try std.testing.expect(promptEligible(4218, out) and p.fused and p.sk == 4 and !p.reduce and !p.cluster); +} + +test "prompt dispatch changes CTA geometry and output instance without changing slices" { + const abi = @import("abi.zig"); + const Trace = struct { + var calls: usize = 0; + var grid: u32 = 0; + var block: u32 = 0; + var shared: u32 = 0; + var id: usize = 0; + var slices: i32 = 0; + fn run(f: abi.Function, gx: c_uint, _: c_uint, gz: c_uint, bx: c_uint, _: c_uint, _: c_uint, sm: c_uint, _: abi.Stream, args: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) abi.Result { + calls += 1; + grid = gx; + block = bx; + shared = sm; + id = @intFromPtr(f); + slices = @as(*const i32, @ptrCast(@alignCast(args.?[9].?))).*; + return if (gz == 1) abi.success else 1; + } + }; + var d: Driver = .{ .lib = undefined, .api = undefined }; + d.api.cuLaunchKernel = Trace.run; + const base: Function = .{ .d = &d, .handle = @ptrFromInt(10) }; + const fast: Function = .{ .d = &d, .handle = @ptrFromInt(20) }; + const fused: Function = .{ .d = &d, .handle = @ptrFromInt(21) }; + const wide: Function = .{ .d = &d, .handle = @ptrFromInt(23) }; + var lane: Lane = .{ + .mod = undefined, + .fns = @splat(@splat(base)), + .f32_fns = @splat(@splat(base)), + .reduce = @splat(base), + .half = @splat(@splat(@splat(base))), + .half_below = 0, + .prompt = .{ .{ fast, fused }, .{ fast, wide } }, + .clusters = false, + }; + const s: Stream = .{ .d = &d, .handle = null }; + var w = try pinned(Weight{ .mode = .bf16, .codes = 0, .scales = 0, .n = 8240, .k = 2560, .npad = 8320 }, 1); + for ([_]usize{ 1, 65, 511, 512, 4218 }) |m| { + Trace.calls = 0; + try lane.matmul(s, 0, w.k, m, w, 0, 0); + const fast_path = m >= prompt_rows; + const bm: usize = if (fast_path) 128 else bucket(m); + const cols: usize = if (fast_path) 128 else 64; + try std.testing.expectEqual(@as(usize, 1), Trace.calls); + try std.testing.expectEqual(@as(usize, if (fast_path) 20 else 10), Trace.id); + try std.testing.expectEqual(@as(u32, if (fast_path) 256 else 128), Trace.block); + try std.testing.expectEqual(if (fast_path) prompt_shared else sharedBytes(.bf16, @intCast(bm)), Trace.shared); + try std.testing.expectEqual((m + bm - 1) / bm * ((w.n + cols - 1) / cols), Trace.grid); + try std.testing.expectEqual(@as(i32, 1), Trace.slices); + } + w = try pinned(Weight{ .mode = .bf16, .codes = 0, .scales = 0, .n = 2560, .k = 3072, .npad = 2560 }, 4); + try lane.matmulF32(s, 0, w.k, 4218, w, 0, 0); + try std.testing.expectEqual(@as(usize, 23), Trace.id); + try std.testing.expectEqual(@as(i32, 4), Trace.slices); + lane.prompt = null; + try lane.matmulF32(s, 0, w.k, 4218, w, 0, 0); + try std.testing.expectEqual(@as(usize, 10), Trace.id); + try std.testing.expectEqual(@as(i32, 4), Trace.slices); +} + +test "a cut weight pinned to the whole weight's K slices plans as the whole one at every row count" { + const whole: Weight = .{ .mode = .bf16, .codes = 0, .scales = 0, .n = 16480, .k = 2560, .npad = 16512 }; + var cut = whole; + cut.n = 8240; + cut.npad = 8320; + try std.testing.expectEqual(@as(u32, 1), splitK(whole.n, whole.k)); + try std.testing.expectEqual(@as(u32, 2), splitK(cut.n, cut.k)); // unpinned, the half would sum two slices + const kept = try pinned(cut, splitK(whole.n, whole.k)); + for ([_]usize{ 1, 16, 64, 1536 }) |m| for ([_]bool{ false, true }) |cl| { + try std.testing.expectEqual(plan(m, whole.n, whole.k, cl), planFor(m, kept, cl)); + try std.testing.expectEqual(plan(m, cut.n, cut.k, cl), planFor(m, cut, cl)); + }; + try std.testing.expect(planFor(1536, cut, false).sk != planFor(1536, kept, false).sk); + try std.testing.expectEqual(@as(u32, 4), (try pinned(cut, 4)).sk); // 40 groups: 10 a slice + for ([_]u32{ 0, 3, 8, 16 }) |bad| try std.testing.expectError(error.BadSplit, pinned(cut, bad)); // 8: 5 a slice +} diff --git a/zig/src/cuda/qmmf_shared.zig b/zig/src/cuda/qmmf_shared.zig new file mode 100644 index 0000000000..0e10faa5ae --- /dev/null +++ b/zig/src/cuda/qmmf_shared.zig @@ -0,0 +1,97 @@ +//! BF16-weight products with one K slice write FP32 rows directly into a caller's row stride. +const q = @import("qmmf.zig"); +const launch = @import("launch.zig"); +const Stream = @import("stream.zig").Stream; +pub const symbols = [4][:0]const u8{ "tf_qmmf_shared16_half", "tf_qmmf_shared16", "tf_qmmf_shared32", "tf_qmmf_shared64" }; + +/// A forced strided launch supports one pinned BF16 K slice. +pub fn supported(l: q.Lane, m: usize, w: q.Weight) bool { + return w.mode == .bf16 and q.planFor(m, w, l.clusters).sk == 1; +} + +/// Prefer an available 128-row prompt projection over the smaller strided tile. +pub fn eligible(l: q.Lane, m: usize, w: q.Weight) bool { + return supported(l, m, w) and !l.usesPrompt(m, w); +} + +pub fn matmulF32Strided(l: q.Lane, s: Stream, x: u64, ldx: usize, m: usize, w: q.Weight, y: u64, ldy: usize) !void { + if (!supported(l, m, w)) return error.StridedLaneNotSupported; + if (m == 0 or w.n == 0 or w.k == 0 or w.k % q.group != 0 or w.npad < w.n or w.npad % 128 != 0 or + ldy < w.n or ldx < w.k) return error.BadStridedLaneShape; + const bm: usize = q.bucket(m); + const rows = (m + bm - 1) / bm; + const half = bm == 16 and rows * ((w.n + 63) / 64) < l.half_below; + const index: usize = if (half) 0 else if (bm == 16) 1 else if (bm == 32) 2 else 3; + const cols: usize = if (half) 32 else 64; + const band = @max(1, @min(rows, (12 << 20) / (bm * w.k * 2))); + var a: launch.Args = .{}; + a.add(x); + a.add(w.codes); + a.add(w.scale); + a.add(y); + for ([_]usize{ m, w.n, w.k, w.npad, if (m == 1) w.k else ldx, band, ldy }) |v| a.add(@as(i32, @intCast(v))); + try launch.launch(l.shared_down[index], .{ + .grid = .{ .x = @intCast(rows * ((w.n + cols - 1) / cols)) }, + .block = .{ .x = 128 }, + .shared = if (half) q.halfShared(.bf16) else q.sharedBytes(.bf16, @intCast(bm)), + }, s, &a); +} + +test "shared down stride: geometry, output address and stride match the contiguous launcher" { + const std = @import("std"); + const abi = @import("abi.zig"); + const Mock = struct { + var id: usize = 0; + var stride: i32 = 0; + var dest: u64 = 0; + var smem: u32 = 0; + fn run(f: abi.Function, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, sm: c_uint, _: abi.Stream, args: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) abi.Result { + id = @intFromPtr(f.?); + stride = @as(*const i32, @ptrCast(@alignCast(args.?[10].?))).*; + dest = @as(*const u64, @ptrCast(@alignCast(args.?[3].?))).*; + smem = sm; + return 0; + } + }; + var d: @import("driver.zig").Driver = undefined; + d.api.cuLaunchKernel = Mock.run; + var l: q.Lane = undefined; + l.clusters = true; + l.prompt = null; + l.half_below = 96; + for (&l.shared_down, 0..) |*f, i| f.* = .{ .d = &d, .handle = @ptrFromInt(101 + i) }; + const s: Stream = .{ .d = &d, .handle = null }; + const w: q.Weight = .{ .mode = .bf16, .codes = 11, .scales = 0, .n = 2560, .k = 640, .npad = 2560 }; + for ([_]usize{ 1, 16, 17, 32, 64, 65, 4218 }, [_]usize{ 0, 0, 2, 2, 3, 3, 3 }) |rows, i| { + try l.matmulF32Strided(s, 12, 640, rows, w, 1000, 28160); + try std.testing.expectEqual(101 + i, Mock.id); + try std.testing.expectEqual(@as(i32, 28160), Mock.stride); + try std.testing.expectEqual(@as(u64, 1000), Mock.dest); + try std.testing.expectEqual(if (i == 0) q.halfShared(.bf16) else q.sharedBytes(.bf16, q.bucket(rows)), Mock.smem); + } + var other = w; + other.mode = .mxfp8; + try std.testing.expectError(error.StridedLaneNotSupported, l.matmulF32Strided(s, 12, 640, 64, other, 1000, 28160)); + other = w; + other.k = 6144; + try std.testing.expect(!eligible(l, 64, other)); + try std.testing.expectError(error.BadStridedLaneShape, l.matmulF32Strided(s, 12, 640, 64, w, 1000, 2559)); + for ([_]bool{ false, true }) |loaded| { + l.prompt = if (loaded) @splat(@splat(l.shared_down[0])) else null; + for ([_]usize{ 511, 512, 513 }) |m| { + try std.testing.expectEqual(loaded and m >= q.prompt_rows, l.usesPrompt(m, w)); + try std.testing.expectEqual(!loaded or m < q.prompt_rows, eligible(l, m, w)); + try std.testing.expect(supported(l, m, w)); + try l.matmulF32Strided(s, 12, 640, m, w, 1000, 28160); + try std.testing.expectEqual(@as(usize, 104), Mock.id); + } + } + other = w; + other.k = 6144; + other.sk = 1; + try std.testing.expect(supported(l, 64, other)); + try std.testing.expect(eligible(l, 511, other)); + try std.testing.expect(!eligible(l, 512, other)); + other.sk = 4; + try std.testing.expect(!supported(l, 64, other)); +} diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 5489b019b8..a9543cf8bb 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -31,6 +31,7 @@ pub const roce = @import("roce.zig"); pub const triton = @import("triton.zig"); pub const aot = @import("aot.zig"); pub const kernels = @import("kernels.zig"); +pub const qmmf_shared = @import("qmmf_shared.zig"); pub const qmmf = @import("qmmf.zig"); pub const fp8 = @import("fp8.zig"); pub const mxfp8 = @import("mxfp8.zig"); @@ -110,6 +111,7 @@ test { _ = aot; _ = segments; _ = @import("qmmf.zig"); + _ = @import("qmmf_shared.zig"); _ = @import("fp8.zig"); _ = @import("mxfp8.zig"); _ = carveout; diff --git a/zig/src/families/flashnext/cuda_attn.zig b/zig/src/families/flashnext/cuda_attn.zig index a5f94ff442..1bff2bdf10 100644 --- a/zig/src/families/flashnext/cuda_attn.zig +++ b/zig/src/families/flashnext/cuda_attn.zig @@ -259,7 +259,7 @@ pub fn decode(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, 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, iq_b.ptr, ik_b.ptr, dummy.ptr, dummy.ptr, eps, 1, 0, .{}); + 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, iq_b.ptr, ik_b.ptr, dummy.ptr, dummy.ptr, eps, 1, 0, .full, .{}); var po_b = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * head_dim * 4); defer po_b.free(); @@ -269,8 +269,8 @@ pub fn decode(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, 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, pos_b.ptr, po_b.ptr, pm_b.ptr, pl_b.ptr, dummy.ptr, dummy.ptr, dummy.ptr, 1, kv_heads, 1, false, .{}); - 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, false); + try tri.attnChunks(q_b.ptr, kc_b.ptr, vc_b.ptr, pos_b.ptr, po_b.ptr, pm_b.ptr, pl_b.ptr, dummy.ptr, dummy.ptr, dummy.ptr, 1, .full, 1, false, .{}); + try tri.attnMerge(po_b.ptr, pm_b.ptr, pl_b.ptr, pos_b.ptr, o_b.ptr, dummy.ptr, dummy.ptr, 1, .full, false); var gated_b = try cuda.DeviceBuffer.alloc(driver, out_k * 2); defer gated_b.free(); diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 6066b1ea63..cbd1d8c636 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -11,9 +11,10 @@ const parts = @import("cuda_batch_state.zig"); const batch_lanes = @import("cuda_batch_lanes.zig"); const tiles = @import("cuda_tiles.zig"); const tk = @import("cuda_tile_kernels.zig"); -const nvfp4_tile = @import("cuda_nvfp4_tile.zig"); +const moe_native = @import("cuda_moe_native.zig"); const zs = @import("cuda_zero_slice.zig"); const ops = @import("cuda_batch_ops.zig"); +const split_lane = @import("cuda_split_lane.zig"); const layers_n = wt.layers_n; const linear_n = wt.linear_n; @@ -27,14 +28,11 @@ const experts_n = wt.experts_n; const slots = wt.slots; const moe_w = wt.moe_w; 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_layer = wt.ple_layer; const ple_tail_n = wt.ple_tail_n; @@ -44,6 +42,7 @@ const out_sk = wt.out_sk; const Q4 = wt.Q4; const HcW = wt.HcW; const Layer = load.Layer; +const Heads = @import("cuda_heads.zig").Heads; /// Normed-stream partial sums a row: D / 256 blocks of each stream. pub const pss_n = (dims / 256) * streams; @@ -78,6 +77,7 @@ pub const MoeAdd = parts.MoeAdd; pub const pocket_stride = parts.pocket_stride; pub const Hooks = parts.Hooks; pub const Kernels = parts.Kernels; +const Gdn = @import("cuda_batch_kernels.zig").Gdn; pub const Scratch = parts.Scratch; pub const SeqState = parts.SeqState; pub const scoreStride = parts.scoreStride; @@ -254,47 +254,51 @@ pub fn Pass(comptime Tri: type) type { pub fn gdnBlock(self: *Self, lw: *const Layer, layer: usize, li: usize, keep: LayerKeep) !void { const s = self.s; const n = self.rows; + const h = lw.heads; + const g = try self.k.gdnSet(h); // before any launch: the full kernels, or the _h8 twins for a rank's half + const pw = h.projWidth(); + const cd = h.convDim(); _ = try self.lane(s.mixed, s.xsm, lw.proj, s.proj, s.dummy, n, 1); // a lone pass keeps its projection and conv window; a shared pass keeps each lane's own if (self.lanes.len == 0) { - try self.copy(keep.proj, s.proj, n * gdn.proj_width * 2); - try self.copy(keep.hist, self.seq.conv[li].ptr, (gdn.taps - 1) * gdn.conv_dim * 2); + try self.copy(keep.proj, s.proj, n * pw * 2); + try self.copy(keep.hist, self.seq.conv[li].ptr, (gdn.taps - 1) * cd * 2); } // every row's conv, norms and gates, the recurrence, then the gated norm: the one-block chain's bits const fused = self.lanes.len == 0 and n >= 64 and gdn_fused.on(); if (fused) { // the three in one kernel a value head: the same per-row op sequences, no fp32 rows through memory const st = self.seq.state[li].ptr; - const in: gdn_fused.In = .{ .proj = s.proj, .conv = self.seq.conv[li].ptr, .cw = lw.cw.ptr, .a_log = lw.a_log.ptr, .dt = lw.dt.ptr, .state = st, .state_out = if (keep.state_out != 0) keep.state_out else st, .norm = lw.norm.ptr }; - try gdn_fused.launch(self.k.gdn_fused, self.stream, in, n, s.gout, s.gxs); + const in: gdn_fused.In = .{ .proj = s.proj, .conv = self.seq.conv[li].ptr, .cw = lw.cw.ptr, .a_log = lw.a_log.ptr, .dt = lw.dt.ptr, .state = st, .state_out = if (keep.state_out != 0) keep.state_out else st, .norm = lw.norm.ptr, .heads = h }; + try gdn_fused.launch(g.fused, self.stream, in, n, s.gout, s.gxs); } else if (self.lanes.len == 0) { - try self.gdnFront(lw, s.proj, self.seq.conv[li].ptr, n, s.gq, s.gk, s.gv, s.ggb); - try self.gdnChain(s.gq, s.gk, s.gv, s.ggb, self.seq.state[li].ptr, keep.state_out, n, s.ys); + try self.gdnFront(g, lw, s.proj, self.seq.conv[li].ptr, n, s.gq, s.gk, s.gv, s.ggb); + try self.gdnChain(g, h, s.gq, s.gk, s.gv, s.ggb, self.seq.state[li].ptr, keep.state_out, n, s.ys); } else { // shared rounds: each lane's rows kept, fronted and chained on its own conv window and state var row0: usize = 0; for (self.lanes) |l| { const lk: LayerKeep = if (l.keep) |kp| kp.layers[layer] else .{}; - const proj = s.proj + row0 * gdn.proj_width * 2; - const qk = row0 * gdn.nk * gdn.dk * 4; - const v = row0 * gdn.value_dim * 4; - const gb = row0 * gdn.nv * 2 * 4; - try self.copy(lk.proj, proj, l.rows * gdn.proj_width * 2); - try self.copy(lk.hist, l.seq.conv[li].ptr, (gdn.taps - 1) * gdn.conv_dim * 2); - try self.gdnFront(lw, proj, l.seq.conv[li].ptr, l.rows, s.gq + qk, s.gk + qk, s.gv + v, s.ggb + gb); - try self.gdnChain(s.gq + qk, s.gk + qk, s.gv + v, s.ggb + gb, l.seq.state[li].ptr, lk.state_out, l.rows, s.ys + v); + const proj = s.proj + row0 * pw * 2; + const qk = row0 * h.nk * gdn.dk * 4; + const v = row0 * h.valueDim() * 4; + const gb = row0 * h.nv * 2 * 4; + try self.copy(lk.proj, proj, l.rows * pw * 2); + try self.copy(lk.hist, l.seq.conv[li].ptr, (gdn.taps - 1) * cd * 2); + try self.gdnFront(g, lw, proj, l.seq.conv[li].ptr, l.rows, s.gq + qk, s.gk + qk, s.gv + v, s.ggb + gb); + try self.gdnChain(g, h, s.gq + qk, s.gk + qk, s.gv + v, s.ggb + gb, l.seq.state[li].ptr, lk.state_out, l.rows, s.ys + v); // the lane's conv window moves on now: the gated norm below reads none of it - try self.shiftWindow(l.seq.conv[li].ptr, proj, gdn.proj_width, l.rows, gdn.taps - 1, gdn.conv_dim); + try self.shiftWindow(l.seq.conv[li].ptr, proj, pw, l.rows, gdn.taps - 1, cd); row0 += l.rows; } } - if (!fused) try gdn_fused.norm(self.k.gnorm, self.stream, s.ys, s.proj, lw.norm.ptr, s.gout, s.gxs, n); - if (self.lanes.len == 0) try self.shiftWindow(self.seq.conv[li].ptr, s.proj, gdn.proj_width, n, gdn.taps - 1, gdn.conv_dim); - self.out_slices = try self.lane(s.gout, s.gxs, lw.out, s.dummy, s.part, n, out_sk); + if (!fused) try gdn_fused.norm(g.norm, self.stream, s.ys, s.proj, lw.norm.ptr, s.gout, s.gxs, n, h); + if (self.lanes.len == 0) try self.shiftWindow(self.seq.conv[li].ptr, s.proj, pw, n, gdn.taps - 1, cd); + self.out_slices = try split_lane.out(self, s.gout, s.gxs, lw); } /// fn_gdn_front over `rows` projection rows after the conv window `conv`. - fn gdnFront(self: *Self, lw: *const Layer, proj: u64, conv: u64, rows: usize, gq: u64, gk: u64, gv: u64, ggb: u64) !void { + fn gdnFront(self: *Self, g: *const Gdn, lw: *const Layer, proj: u64, conv: u64, rows: usize, gq: u64, gk: u64, gv: u64, ggb: u64) !void { var f: cuda.Args = .{}; f.add(proj); f.add(conv); @@ -305,11 +309,11 @@ pub fn Pass(comptime Tri: type) type { f.add(gk); f.add(gv); f.add(ggb); - try cuda.launch.launch(self.k.front, .{ .grid = .{ .x = @intCast(rows), .y = gdn.nk + gdn.nv }, .block = .{ .x = 256 } }, self.stream, &f); + try cuda.launch.launch(g.front, .{ .grid = .{ .x = @intCast(rows), .y = lw.heads.nk + lw.heads.nv }, .block = .{ .x = 256 } }, self.stream, &f); } /// fn_gdn_rows over `rows` rows from `state`, in place or (state_out) to fn_gdn_rows_to's; readouts into ys. - fn gdnChain(self: *Self, gq: u64, gk: u64, gv: u64, ggb: u64, state: u64, state_out: u64, rows: usize, ys: u64) !void { + fn gdnChain(self: *Self, g: *const Gdn, h: Heads, gq: u64, gk: u64, gv: u64, ggb: u64, state: u64, state_out: u64, rows: usize, ys: u64) !void { var a: cuda.Args = .{}; a.add(gq); a.add(gk); @@ -321,12 +325,12 @@ pub fn Pass(comptime Tri: type) type { a.add(if (state_out != 0) state_out else state); a.add(cint(rows)); a.add(ys); - return cuda.launch.launch(self.k.chain_tpr, .{ .grid = .{ .x = gdn.nv }, .block = .{ .x = 256 }, .shared = parts.chain_tpr_smem }, self.stream, &a); + return cuda.launch.launch(g.chain_tpr, .{ .grid = .{ .x = h.nv }, .block = .{ .x = 256 }, .shared = parts.chain_tpr_smem }, self.stream, &a); } if (state_out != 0) a.add(state_out); a.add(cint(rows)); a.add(ys); - try cuda.launch.launch(if (state_out != 0) self.k.chain_to else self.k.chain, .{ .grid = .{ .x = gdn.nv, .y = gdn.dv / 16 }, .block = .{ .x = 128 } }, self.stream, &a); + try cuda.launch.launch(if (state_out != 0) g.chain_to else g.chain, .{ .grid = .{ .x = h.nv, .y = gdn.dv / 16 }, .block = .{ .x = 128 } }, self.stream, &a); } pub fn attnBlock(self: *Self, lw: *const Layer, ai: usize) !void { @@ -338,12 +342,14 @@ pub fn Pass(comptime Tri: type) type { const s = self.s; const n = self.rows; const pa = s.proj; + const h = lw.heads; + const pn = h.projN(); _ = try self.lane(s.mixed, s.xsm, lw.proj, pa, s.dummy, n, 1); - const qw2 = q_heads * 2 * head_dim; - const kw = kv_heads * head_dim; - try self.copyCols(pa, proj_n, keep.pq, qw2, qw2, n); - try self.copyCols(pa + qw2 * 2, proj_n, keep.pk, kw, kw, n); - try self.copyCols(pa + (qw2 + kw) * 2, proj_n, keep.pv, kw, kw, n); + const qw2 = h.q * 2 * head_dim; + const kw = h.kv * head_dim; + try self.copyCols(pa, pn, keep.pq, qw2, qw2, n); + try self.copyCols(pa + qw2 * 2, pn, keep.pk, kw, kw, n); + try self.copyCols(pa + (qw2 + kw) * 2, pn, keep.pv, kw, kw, n); if (self.lanes.len > 0) { // shared rounds: each lane's rope, cache writes, selection and attention on its own stream's caches const own_seq = self.seq; @@ -356,13 +362,13 @@ pub fn Pass(comptime Tri: type) type { for (self.lanes, 0..) |l, i| { self.seq = l.seq; self.pos0 = l.pos; - try self.attnRows(lw, ai, pa + row0 * proj_n * 2, l.rows, .{ .pos = s.pos + i * 16, .q = s.q + row0 * q_heads * head_dim * 2, .iq = s.iq + row0 * index_heads * index_dim * 2, .ao = s.ao + row0 * out_k * 2 }); + try self.attnRows(lw, ai, pa + row0 * pn * 2, l.rows, .{ .pos = s.pos + i * 16, .q = s.q + row0 * h.q * head_dim * 2, .iq = s.iq + row0 * index_heads * index_dim * 2, .ao = s.ao + row0 * h.q * head_dim * 2 }); row0 += l.rows; } } else try self.attnRows(lw, ai, pa, n, .{ .pos = s.pos, .q = s.q, .iq = s.iq, .ao = s.ao }); const gated = if (keep.gated != 0) keep.gated else s.gated; - try self.tri.attnGate(s.ao, pa, gated, s.gxs_attn, n, proj_n, q_heads, head_dim); - self.out_slices = try self.lane(gated, s.gxs_attn, lw.out, s.dummy, s.part, n, out_sk); + try self.tri.attnGate(s.ao, pa, gated, s.gxs_attn, n, pn, h.q, head_dim); + self.out_slices = try split_lane.out(self, gated, s.gxs_attn, lw); } /// Where attnRows reads its rows' positions (16 bytes apart) and writes queries, indexer queries and outputs. @@ -371,12 +377,14 @@ pub fn Pass(comptime Tri: type) type { /// `n` rows of self.seq from self.pos0: rope, cache writes and pooling, then scores, selection and attention. fn attnRows(self: *Self, lw: *const Layer, ai: usize, pa: u64, n: usize, at: RowsAt) !void { const s = self.s; + const h = lw.heads; + const qr = h.q * head_dim * 2; // a row's queries, and its attention output const kc = self.seq.kc[ai].ptr; const vc = self.seq.vc[ai].ptr; const ikc = self.seq.ikc[ai].ptr; const pooled = self.seq.pooled[ai].ptr; const pg = self.seq.pg; - if (self.seq.media) |m| try m.attn(self.stream, pa, at.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, at.q, kc, vc, at.iq, ikc, eps, n, pg) else try self.tri.attnPrep(pa, at.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, at.q, kc, vc, at.iq, ikc, s.dummy, s.dummy, eps, n, 0, pg); + if (self.seq.media) |m| try m.attn(self.stream, pa, at.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, at.q, kc, vc, at.iq, ikc, eps, n, h, pg) else try self.tri.attnPrep(pa, at.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, at.q, kc, vc, at.iq, ikc, s.dummy, s.dummy, eps, n, 0, h, pg); // attention.qsa_pool for the blocks these rows complete, then each block of rows' scores and attention if (self.seq.media) |m| try m.pooled(self.stream, ikc, pooled, at.pos, lw.ik_scale.ptr, s.inv, eps, n, pg) else try self.tri.qsaPool(ikc, pooled, at.pos, lw.ik_scale.ptr, s.inv, eps, n, pg); const nb = scoreStride(self.seq.cap); @@ -409,11 +417,11 @@ pub fn Pass(comptime Tri: type) type { try cuda.launch.launch(self.k.select, .{ .grid = .{ .x = @intCast(m) }, .block = .{ .x = 1024 } }, self.stream, &a); if (!fused) { const chunks = cdiv(@min(ends, wt.dense_keys), 512); - try self.tri.attnChunks(at.q + r0 * q_heads * head_dim * 2, kc, vc, pos, s.po, s.pm, s.pl, s.qids, s.qnk, s.qspr, m, kv_heads, chunks, true, pg); - try self.tri.attnMerge(s.po, s.pm, s.pl, pos, at.ao + r0 * out_k * 2, s.qnk, s.qspr, m, kv_heads, true); + try self.tri.attnChunks(at.q + r0 * qr, kc, vc, pos, s.po, s.pm, s.pl, s.qids, s.qnk, s.qspr, m, h, chunks, true, pg); + try self.tri.attnMerge(s.po, s.pm, s.pl, pos, at.ao + r0 * qr, s.qnk, s.qspr, m, h, true); } } - if (fused) try self.tri.attnFused(at.q + g0 * q_heads * head_dim * 2, kc, vc, at.pos + (g0 / att_rows) * 16, at.ao + g0 * out_k * 2, s.qids, s.qnk, s.qspr, gm, kv_heads, pg); + if (fused) try self.tri.attnFused(at.q + g0 * qr, kc, vc, at.pos + (g0 / att_rows) * 16, at.ao + g0 * qr, s.qids, s.qnk, s.qspr, gm, h, pg); } } @@ -434,28 +442,7 @@ pub fn Pass(comptime Tri: type) type { try self.copy(keep.moe_wts, s.wts, pairs * 4); const e = experts_n + 1; const plan: cuda.grouped.Plan = .{ .members = s.members, .items = s.items, .counts = s.counts, .rank = s.rank, .hist = s.hist }; - if (lw.native_experts) |*native| { - const owned = &native.experts; - if (zs.clearsY(native.world, native.rank)) try self.stream.d.check(self.stream.d.api.cuMemsetD8Async(s.y, 0, pairs * dims * 4, self.stream.handle), "cuMemsetD8Async"); - const pt = nvfp4_tile.size(n); - try self.k.router.route(self.stream, s.picks, pairs, e, pt, plan); - const items = cuda.grouped.maxItems(pairs, e, pt); - try nvfp4_tile.run(self.k, pt, self.stream, .up, .{ .x = s.mixed, .stride = dims, .slots = slots }, owned.view(), plan, s.moe_act, items, experts_n); - try nvfp4_tile.run(self.k, pt, self.stream, .down_f32, .{ .x = s.moe_act, .stride = owned.width }, owned.view(), plan, s.y, items, experts_n); - if (native.rank == 0) { - const glue = self.tri.own orelse return error.NativeGlueRequired; - try self.k.native_lane.matmul(self.stream, s.mixed, dims, n, native.shared[0].native.?, s.shared_gate, s.native_part); - try self.k.native_lane.matmul(self.stream, s.mixed, dims, n, native.shared[1].native.?, s.shared_up, s.native_part); - try glue.sharedSwiGLU(self.stream, s.shared_gate, s.shared_up, s.shared_act, n, moe_w); - try self.k.native_lane.matmulF32(self.stream, s.shared_act, moe_w, n, native.shared[2].native.?, s.shared_out, s.native_part); - try glue.sharedScatter(self.stream, s.shared_out, s.y, n, dims, slots); - } - if (native.world != 1) { - const collective = self.m.collective orelse return error.TensorParallelUnavailable; - try collective.sum(self.stream, s.y, pairs * dims); - } - return; - } + if (lw.native_experts) |*native| return moe_native.run(self, native, n, pairs, plan); if (n >= 64) { // prompt rows on 64-pair factored tiles, expert_rows' bits try self.k.router.route(self.stream, s.picks, pairs, e, tk.fexp_tile, plan); diff --git a/zig/src/families/flashnext/cuda_batch_kernels.zig b/zig/src/families/flashnext/cuda_batch_kernels.zig index 1443cf1398..f1b6d8affd 100644 --- a/zig/src/families/flashnext/cuda_batch_kernels.zig +++ b/zig/src/families/flashnext/cuda_batch_kernels.zig @@ -7,6 +7,33 @@ const cubin_qmm = state.cubin_qmm; const qmm_smem = state.qmm_smem; const chain_tpr_smem = state.chain_tpr_smem; const gdn_fused_smem = state.gdn_fused_smem; +const Heads = @import("cuda_heads.zig").Heads; + +/// The DeltaNet kernels of one head geometry: the full heads', or their `_h8` twins' (8 key and 24 value heads). +pub const Gdn = struct { + front: cuda.Function, + chain: cuda.Function, + chain_to: cuda.Function, + chain_tpr: cuda.Function, + /// front + chain + gated norm in one kernel a value head (prompt passes of 64 rows or more). + fused: cuda.Function, + norm: cuda.Function, + + /// The entry point of `field` for the full heads (suffix "") or a twin set (suffix "_h8"). + pub fn symbol(comptime field: []const u8, comptime suffix: [:0]const u8) [:0]const u8 { + const names = .{ .front = "fn_gdn_front", .chain = "fn_gdn_rows", .chain_to = "fn_gdn_rows_to", .chain_tpr = "fn_gdn_chain_tpr", .fused = "fn_gdn_fused", .norm = "fn_gdn_norm" }; + return @field(names, field) ++ suffix; + } + + fn load(m: cuda.Module, comptime suffix: [:0]const u8) !Gdn { + var g: Gdn = undefined; + inline for (@typeInfo(Gdn).@"struct".field_names) |name| @field(g, name) = try m.function(symbol(name, suffix)); + try g.chain_tpr.allowDynamicShared(chain_tpr_smem); + try g.fused.allowDynamicShared(gdn_fused_smem); + return g; + } +}; + pub const Kernels = struct { qmm: cuda.Module, experts: cuda.Module, @@ -20,17 +47,16 @@ pub const Kernels = struct { tiles: tk.TileKernels, native_lane: cuda.qmmf.Lane, nvfp4: cuda.experts.Nvfp4Experts, + /// Block-FP8 routed experts (the MTP head of NVIDIA's NVFP4 checkpoint; cuda_moe_native.zig). + fp8: cuda.experts.Fp8Experts, expert_tile: ?cuda.experts_tile.Tile = null, - front: cuda.Function, - chain: cuda.Function, - chain_to: cuda.Function, - chain_tpr: cuda.Function, - /// front + chain + gated norm in one kernel a value head (prompt passes of 64 rows or more). - gdn_fused: cuda.Function, - gnorm: cuda.Function, + /// The DeltaNet kernels for the whole layer's heads and for a rank's half under the head split (gdnSet). + gdn: Gdn, + gdn_h8: Gdn, sum: cuda.Function, moe_sum: cuda.Function, copy_cols: cuda.Function, + ep_picks: cuda.Function, shift: cuda.Function, /// Each sparse row's top blocks (attention._select / _select_tiles), native. select: cuda.Function, @@ -62,26 +88,31 @@ pub const Kernels = struct { try d.check(d.api.cuDeviceGetAttribute(&sms, .multiprocessor_count, device), "cuDeviceGetAttribute"); k.nvfp4 = try cuda.experts.Nvfp4Experts.load(d, @intCast(sms)); errdefer k.nvfp4.unload(); + k.fp8 = try cuda.experts.Fp8Experts.load(d, @intCast(sms)); + errdefer k.fp8.unload(); if (tile_mode != .off) k.expert_tile = try cuda.experts_tile.Tile.load(d, @intCast(sms), tile_mode); errdefer if (k.expert_tile) |*tile| tile.unload(); k.sum = try k.batch.function("fn_sum_slices"); k.moe_sum = try k.batch.function("fn_moe_sum"); k.copy_cols = try k.batch.function("fn_copy_cols"); + k.ep_picks = try k.batch.function("fn_ep_picks"); k.shift = try k.batch.function("fn_shift_window"); k.select = try k.batch.function("fn_qsa_select"); - k.front = try k.batch.function("fn_gdn_front"); - k.chain = try k.batch.function("fn_gdn_rows"); - k.chain_to = try k.batch.function("fn_gdn_rows_to"); - k.chain_tpr = try k.batch.function("fn_gdn_chain_tpr"); - try k.chain_tpr.allowDynamicShared(chain_tpr_smem); - k.gdn_fused = try k.batch.function("fn_gdn_fused"); - try k.gdn_fused.allowDynamicShared(gdn_fused_smem); - k.gnorm = try k.batch.function("fn_gdn_norm"); + k.gdn = try Gdn.load(k.batch, ""); + k.gdn_h8 = try Gdn.load(k.batch, "_h8"); return k; } + /// The DeltaNet kernels for a layer's heads `h`: the full set, or the `_h8` twins for half; any other is refused. + pub fn gdnSet(k: *const Kernels, h: Heads) error{RowShape}!*const Gdn { + if (h.eql(.full)) return &k.gdn; + if (h.eql(.half)) return &k.gdn_h8; + return error.RowShape; + } + pub fn deinit(k: *Kernels) void { if (k.expert_tile) |*tile| tile.unload(); + k.fp8.unload(); k.nvfp4.unload(); k.native_lane.unload(); k.batch.unload(); @@ -91,3 +122,26 @@ pub const Kernels = struct { k.qmm.unload(); } }; + +test "a split layer's DeltaNet launches, the rollback replay's front and chain too, take the _h8 twins; a whole one today's" { + const want = [_][:0]const u8{ "fn_gdn_front_h8", "fn_gdn_rows_h8", "fn_gdn_rows_to_h8", "fn_gdn_chain_tpr_h8", "fn_gdn_fused_h8", "fn_gdn_norm_h8" }; + const names = @typeInfo(Gdn).@"struct".field_names; + inline for (names, want) |field, name| { + try std.testing.expectEqualStrings(name, Gdn.symbol(field, "_h8")); + try std.testing.expectEqualStrings(name[0 .. name.len - 3], Gdn.symbol(field, "")); + } + var driver: cuda.Driver = undefined; + var k: Kernels = undefined; + inline for (names, 0..) |field, i| { + @field(k.gdn, field) = .{ .d = &driver, .handle = @ptrFromInt(0x100 + i) }; + @field(k.gdn_h8, field) = .{ .d = &driver, .handle = @ptrFromInt(0x200 + i) }; + } + const half = try k.gdnSet(.half); + const full = try k.gdnSet(.full); + try std.testing.expectEqual(k.gdn_h8.front.handle, half.front.handle); + try std.testing.expectEqual(k.gdn_h8.chain.handle, half.chain.handle); + try std.testing.expectEqual(k.gdn.front.handle, full.front.handle); + try std.testing.expectEqual(k.gdn.chain.handle, full.chain.handle); + inline for (names) |field| try std.testing.expect(@field(half, field).handle != @field(full, field).handle); + try std.testing.expectError(error.RowShape, k.gdnSet(.{ .nk = 4, .nv = 12, .q = 6, .kv = 1 })); +} diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index 98faec84be..4cbe5e50b8 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -212,6 +212,10 @@ pub const Scratch = struct { xsm: u64 = 0, proj: u64 = 0, part: u64 = 0, + rank_input: u64 = 0, + rank_gather: u64 = 0, + rank_output: u64 = 0, + rank_result: u64 = 0, /// part's second cap-stride slice is zeros (native lanes write only the first): a writeback adds it once. part_zero: bool = false, gout: u64 = 0, @@ -235,6 +239,7 @@ pub const Scratch = struct { gxs_attn: u64 = 0, router_l: u64 = 0, picks: u64 = 0, + picks_ep: u64 = 0, // the rank's own experts' picks (cuda_moe_native.zig) at --tp wts: u64 = 0, members: u64 = 0, items: u64 = 0, @@ -327,12 +332,13 @@ pub const Scratch = struct { n += cap * q_heads * head_dim * 2 + cap * index_heads * index_dim * 2; n += att_rows * nch * q_heads * (head_dim + 2) * 4 + cap * out_k * 2 * 2 + cap * (out_k / 32) * 4; n += fused_rows * (idw + 2) * 4; - n += cap * (experts_n + 1) * 4 + pairs * 4 * 2 + pairs * 4 + items * 3 * 4 + 8 + pairs * 4 + nblk * (experts_n + 1) * 4; + n += cap * (experts_n + 1) * 4 + pairs * 4 * 2 + pairs * 4 + items * 3 * 4 + 8 + pairs * 4 + nblk * (experts_n + 1) * 4 + pairs * 4; n += cap * moe_w * 6 + cap * dims * 4; n += pairs * moe_w * 2 + pairs * dims * 4 + cap * dims * 2; n += cap * dims * 2 + cap * (dims / 32) * 4 + cap * wide * 2 + cap * dims * 2 + 4 * cap * dims * 4; n += cap * wide * 2 * 2 + cap * streams * 4 + vocab * 2 + cap * vocab + half * 4 + 4096; n += cap * dims * 4 + cap * pocket_stride * 4; + n += cap * out_k * 4 + cap * dims * 6; return n; } @@ -362,6 +368,10 @@ pub const Scratch = struct { s.mixed = try s.add(gpa, d, cap * dims * 2); s.xsm = try s.add(gpa, d, cap * (dims / 32) * 4); s.proj = try s.add(gpa, d, cap * proj_row * 2); + s.rank_input = try s.add(gpa, d, cap * out_k * 2); + s.rank_gather = try s.add(gpa, d, cap * out_k * 2); + s.rank_output = try s.add(gpa, d, cap * dims * 2); + s.rank_result = try s.add(gpa, d, cap * dims * 4); s.part = try s.add(gpa, d, out_sk * cap * dims * 4); try s.bufs.items[s.bufs.items.len - 1].fill8(0, stream.handle); s.part_zero = true; @@ -385,6 +395,7 @@ pub const Scratch = struct { s.gxs_attn = try s.add(gpa, d, cap * (out_k / 32) * 4); s.router_l = try s.add(gpa, d, cap * (experts_n + 1) * 4); s.picks = try s.add(gpa, d, pairs * 4); + s.picks_ep = try s.add(gpa, d, pairs * 4); s.wts = try s.add(gpa, d, pairs * 4); s.members = try s.add(gpa, d, pairs * 4); s.items = try s.add(gpa, d, items * 3 * 4); diff --git a/zig/src/families/flashnext/cuda_collective.zig b/zig/src/families/flashnext/cuda_collective.zig index fd6d930836..6f2e30dc9a 100644 --- a/zig/src/families/flashnext/cuda_collective.zig +++ b/zig/src/families/flashnext/cuda_collective.zig @@ -16,12 +16,13 @@ pub const Collective = struct { comm: nccl.Comm = null, link: wire.Link, timeout_ms: u64, + rank: usize = 0, failed: bool = false, /// `link` has completed the source/settings handshake; it stays owned by the caller. pub fn init(link: wire.Link, rank: usize, timeout_ms: u64) !Collective { if (rank > 1 or timeout_ms == 0) return error.BadRank; - var c: Collective = .{ .lib = try nccl.Library.open(), .api = undefined, .link = link, .timeout_ms = timeout_ms }; + var c: Collective = .{ .lib = try nccl.Library.open(), .api = undefined, .link = link, .timeout_ms = timeout_ms, .rank = rank }; errdefer c.lib.close(); try c.lib.requireSupportedVersion(); const info = @typeInfo(Advanced).@"struct"; @@ -84,13 +85,22 @@ pub const Collective = struct { try c.wait(stream, wire.deadline(c.timeout_ms)); } - /// Disjoint BF16 vocabulary columns are concatenated without an arithmetic reduction. + /// BF16 rank blocks are concatenated without an arithmetic reduction. pub fn gather(c: *Collective, stream: cuda.Stream, local: u64, out: u64, count: usize) !void { + return c.gatherType(stream, local, out, count, .bf16); + } + + /// FP32 rank blocks are concatenated without rounding or reduction. + pub fn gatherF32(c: *Collective, stream: cuda.Stream, local: u64, out: u64, count: usize) !void { + return c.gatherType(stream, local, out, count, .f32); + } + + fn gatherType(c: *Collective, stream: cuda.Stream, local: u64, out: u64, count: usize, dtype: nccl.DataType) !void { if (c.failed or c.comm == null) return error.RankDisconnected; if (count == 0) return; errdefer c.abort(); try c.link.checkPeer(); - try c.submitted(c.lib.api.ncclAllGather(local, out, count, .bf16, c.comm, stream.handle)); + try c.submitted(c.lib.api.ncclAllGather(local, out, count, dtype, c.comm, stream.handle)); try c.wait(null, wire.deadline(c.timeout_ms)); } diff --git a/zig/src/families/flashnext/cuda_config.zig b/zig/src/families/flashnext/cuda_config.zig index 41b4099220..1054083b32 100644 --- a/zig/src/families/flashnext/cuda_config.zig +++ b/zig/src/families/flashnext/cuda_config.zig @@ -130,19 +130,32 @@ fn modeloptMismatch(text: std.json.ObjectMap, quant: std.json.ObjectMap) ?[]cons if (!projection(entries, routed, "NVFP4", 16)) return "quantization_config.routed_experts"; const p = object(entries, routed).?; if (given(p, "weight_scale_encoding")) |encoding| if (!holds(encoding, .{ .word = "csf" })) return "quantization_config.weight_scale_encoding"; + // each projection MXFP8/32 (FP4-Local) or unlisted (BF16, NVIDIA's NVFP4); a stitched group one or the other for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }) |part| { - const name = std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.mlp.shared_expert.{s}", .{ layer, part }) catch unreachable; - if (!projection(entries, name, "MXFP8", 32)) return "quantization_config.shared_expert"; - } - const parts: []const []const u8 = if (wt.linearLayer(layer)) &.{ "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a", "out_proj" } else &.{ "q_proj", "k_proj", "v_proj", "indexer.index_qk_proj", "o_proj" }; - for (parts) |part| { - const name = std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.{s}.{s}", .{ layer, if (wt.linearLayer(layer)) "linear_attn" else "self_attn", part }) catch unreachable; - if (!projection(entries, name, "MXFP8", 32)) return "quantization_config.attention_projection"; + if (declared(entries, layer, "mlp.shared_expert", &.{part}) == null) return "quantization_config.shared_expert"; } + const linear = wt.linearLayer(layer); + const prefix = if (linear) "linear_attn" else "self_attn"; + const stitched: []const []const u8 = if (linear) &.{ "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a" } else &.{ "q_proj", "k_proj", "v_proj", "indexer.index_qk_proj" }; + if (declared(entries, layer, prefix, stitched) == null) return "quantization_config.attention_projection"; + if (declared(entries, layer, prefix, &.{if (linear) "out_proj" else "o_proj"}) == null) return "quantization_config.attention_projection"; } return null; } +/// One loaded projection's parts: all MXFP8/32, or none listed (BF16); null when mixed or another format. +fn declared(entries: std.json.ObjectMap, layer: usize, prefix: []const u8, parts: []const []const u8) ?enum { mxfp8, bf16 } { + var listed: usize = 0; + var buf: [256]u8 = undefined; + for (parts) |part| { + const name = std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.{s}.{s}", .{ layer, prefix, part }) catch unreachable; + if (entries.get(name) == null) continue; + if (!projection(entries, name, "MXFP8", 32)) return null; + listed += 1; + } + return if (listed == 0) .bf16 else if (listed == parts.len) .mxfp8 else null; +} + /// The neurons a Living Weights pocket adds to each shared expert (living's Keys.pocket). const learned_key = "tensorfold_learned_neurons"; @@ -413,9 +426,49 @@ test "mixed precision requires every routed, shared and attention declaration an try std.testing.expectEqualStrings("quantization_config.kv_cache_quant_algo", modeloptMismatch(text, quant).?); try quant.put(a, "kv_cache_quant_algo", .null); try std.testing.expect(modeloptMismatch(text, quant) == null); + // a projection loaded alone may be BF16 (unlisted); a stitched group may not mix MXFP8 and BF16 parts _ = entries.swapRemove("model.language_model.layers.47.self_attn.o_proj"); + _ = entries.swapRemove("model.language_model.layers.3.mlp.shared_expert.up_proj"); + try quant.put(a, "quantized_layers", .{ .object = entries }); + try std.testing.expect(modeloptMismatch(text, quant) == null); + _ = entries.swapRemove("model.language_model.layers.47.self_attn.k_proj"); try quant.put(a, "quantized_layers", .{ .object = entries }); try std.testing.expectEqualStrings("quantization_config.attention_projection", modeloptMismatch(text, quant).?); + try addMxfp8(a, &entries, 47, "self_attn", "k_proj"); + _ = entries.swapRemove("model.language_model.layers.0.linear_attn.in_proj_a"); + try quant.put(a, "quantized_layers", .{ .object = entries }); + try std.testing.expectEqualStrings("quantization_config.attention_projection", modeloptMismatch(text, quant).?); +} + +test "NVIDIA's NVFP4 layout passes: routed experts NVFP4 with plain scales, every other projection BF16 (unlisted)" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var text: std.json.ObjectMap = .empty; + var rope_params: std.json.ObjectMap = .empty; + var mtp: std.json.ObjectMap = .empty; + var kinds: std.json.Array = .init(a); + var ple: std.json.Array = .init(a); + try served(a, &text, &rope_params, &mtp, &kinds, &ple); + var quant: std.json.ObjectMap = .empty; + try quant.put(a, "quant_method", .{ .string = "modelopt" }); + try quant.put(a, "quant_algo", .{ .string = "MIXED_PRECISION" }); + var entries: std.json.ObjectMap = .empty; + for (0..wt.layers_n) |layer| { + var nv: std.json.ObjectMap = .empty; + try nv.put(a, "quant_algo", .{ .string = "NVFP4" }); + try nv.put(a, "group_size", .{ .integer = 16 }); + try entries.put(a, try std.fmt.allocPrint(a, "model.language_model.layers.{d}.mlp.experts", .{layer}), .{ .object = nv }); + } + try quant.put(a, "quantized_layers", .{ .object = entries }); + try std.testing.expect(modeloptMismatch(text, quant) == null); + // one shared projection declared in another format is still refused + var nv: std.json.ObjectMap = .empty; + try nv.put(a, "quant_algo", .{ .string = "NVFP4" }); + try nv.put(a, "group_size", .{ .integer = 16 }); + try entries.put(a, "model.language_model.layers.5.mlp.shared_expert.gate_proj", .{ .object = nv }); + try quant.put(a, "quantized_layers", .{ .object = entries }); + try std.testing.expectEqualStrings("quantization_config.shared_expert", modeloptMismatch(text, quant).?); } fn addMxfp8(a: std.mem.Allocator, entries: *std.json.ObjectMap, layer: usize, prefix: []const u8, part: []const u8) !void { diff --git a/zig/src/families/flashnext/cuda_dense_split.zig b/zig/src/families/flashnext/cuda_dense_split.zig new file mode 100644 index 0000000000..c08b7a236c --- /dev/null +++ b/zig/src/families/flashnext/cuda_dense_split.zig @@ -0,0 +1,210 @@ +//! Whether each of two ranks keeps half of the BF16 dense projections, and what a rank's half holds on the device. +const std = @import("std"); +const cuda = @import("cuda"); +const safetensors = @import("core").safetensors; +const Tensor = safetensors.Tensor; +const wt = @import("cuda_prompt_wt.zig"); +const gdn = @import("cuda_gdn.zig"); +const Axis = @import("cuda_projection_load.zig").Axis; +const Cut = @import("cuda_modelopt_experts.zig").Cut; +const Heads = @import("cuda_heads.zig").Heads; +const heads_cut = @import("heads_cut.zig"); +const Entry = safetensors.Entry; + +comptime { + // DeltaNet's out_proj and attention's o_proj take the same input width, so one cut serves both + std.debug.assert(gdn.value_dim == wt.out_k); +} + +/// The tensors the split is decided from: layer 0's DeltaNet out projection and its shared expert's down projection. +const deciders = [_][]const u8{ + "language_model.model.layers.0.linear_attn.out_proj.weight", + "language_model.model.layers.0.mlp.shared_expert.down_proj.weight", +}; + +/// TF_FLASHNEXT_SPLIT_DENSE: unset or 1 splits, 0 keeps every dense projection whole. +pub fn setting(value: ?[]const u8) !bool { + const v = value orelse return true; + if (std.mem.eql(u8, v, "0")) return false; + if (std.mem.eql(u8, v, "1")) return true; + return error.BadDenseSplit; +} + +pub fn enabled() !bool { + return setting(if (std.c.getenv("TF_FLASHNEXT_SPLIT_DENSE")) |v| std.mem.span(v) else null); +} + +/// At two ranks with the setting on, a checkpoint whose deciders are BF16 splits; world 1 never reads the setting. +pub fn applies(store: anytype, world: usize) !bool { + if (world != 2 or !try enabled()) return false; + for (deciders) |name| if (!store.has(name)) return false; + for (deciders) |name| if ((try store.tensor(name)).dtype != .bf16) return false; + return true; +} + +/// The heads a rank computes of every model layer: half of each under the split, else the whole layer. +pub fn heads(store: anytype, world: usize) !Heads { + return if (try applies(store, world)) .half else .full; +} + +pub fn appliesDir(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize) !bool { + if (world != 2) return false; + var store = try wt.Store.open(gpa, io, dir); + defer store.deinit(); + return applies(&store, world); +} + +/// How the split cuts a checkpoint projection by its name; null: kept whole on every rank. +pub fn axisOf(name: []const u8) ?Axis { + for ([_][]const u8{ ".linear_attn.out_proj.weight", ".self_attn.o_proj.weight", ".mlp.shared_expert.gate_proj.weight", ".mlp.shared_expert.up_proj.weight", ".mlp.shared_expert.down_proj.weight" }) |suffix| { + if (std.mem.endsWith(u8, name, suffix)) return .rows; + } + return null; +} + +/// Device bytes of one rank's cut of a BF16 [n, k] projection, its rows padded as packBf16 pads them. +pub fn residentBytes(n: usize, k: usize, axis: Axis, world: usize) !u64 { + const part = (try Cut.init(if (axis == .rows) n else k, world, 0)).width(); + const rows = if (axis == .rows) part else n; + const cols = if (axis == .rows) k else part; + return try std.math.mul(u64, try std.math.mul(u64, cuda.nvfp4.padded(rows), cols), 2); +} + +/// One rank's device bytes of a tensor the head split cuts (in-projection parts, conv1d, A_log, dt_bias); null: none. +pub fn headBytes(name: []const u8, e: Entry, world: usize) !?u64 { + const h = try Heads.split(world); + const bytes: u64 = e.end - e.begin; + inline for (.{ .{ ".linear_attn", heads_cut.delta_parts }, .{ ".self_attn", heads_cut.attention_parts } }) |class| { + inline for (class[1], 0..) |part, i| if (std.mem.endsWith(u8, name, class[0] ++ part ++ ".weight")) { + if (e.dtype != .bf16 or e.rank != 2) return null; + const s = if (comptime std.mem.eql(u8, class[0], ".linear_attn")) try heads_cut.delta(h, 0) else try heads_cut.attention(h, 0); + const k: u64 = e.dim(1); + // the stitched weight packs as one: its rows past the last part's padded once, counted with the first + const pad: u64 = if (i == 0) cuda.nvfp4.padded(s.n) - s.n else 0; + return (s.ranges[i].rows() + pad) * k * 2; + }; + } + if (std.mem.endsWith(u8, name, ".linear_attn.conv1d.weight")) { + if (e.rank == 0 or e.dim(0) != Heads.full.convDim() or bytes % e.dim(0) != 0) return error.UnexpectedTensor; + return bytes / e.dim(0) * h.convDim(); + } + for ([_][]const u8{ ".linear_attn.A_log", ".linear_attn.dt_bias" }) |suffix| if (std.mem.endsWith(u8, name, suffix)) { + if (e.rank != 1 or e.dim(0) != Heads.full.nv or bytes % e.dim(0) != 0) return error.UnexpectedTensor; + return bytes / e.dim(0) * h.nv; + }; + return null; +} + +/// A checkpoint of two tensors, each of one dtype: what `applies` reads. +const Fake = struct { + dtypes: [2]safetensors.DType, + present: bool = true, + + pub fn has(f: *const Fake, name: []const u8) bool { + _ = name; + return f.present; + } + + pub fn tensor(f: *const Fake, name: []const u8) !Tensor { + const i: usize = if (std.mem.eql(u8, name, deciders[0])) 0 else 1; + return .{ .dtype = f.dtypes[i], .rank = 2, .shape = .{ 2560, 64, 1, 1, 1 }, .bytes = &.{} }; + } +}; + +test "the dense split setting is on unless 0 and refuses any other value" { + try std.testing.expect(try setting(null)); + try std.testing.expect(try setting("1")); + try std.testing.expect(!try setting("0")); + try std.testing.expectError(error.BadDenseSplit, setting("2")); + try std.testing.expectError(error.BadDenseSplit, setting("")); +} + +test "the dense split applies at two ranks to BF16 layers only, and world 1 never reads the setting" { + const env = @import("sys").env; + defer _ = env.unsetenv("TF_FLASHNEXT_SPLIT_DENSE"); + _ = env.unsetenv("TF_FLASHNEXT_SPLIT_DENSE"); + const bf16: Fake = .{ .dtypes = .{ .bf16, .bf16 } }; + try std.testing.expect(try applies(&bf16, 2)); + try std.testing.expect(!try applies(&bf16, 1)); + try std.testing.expect(!try applies(&Fake{ .dtypes = .{ .f8_e4m3, .f8_e4m3 } }, 2)); + try std.testing.expect(!try applies(&Fake{ .dtypes = .{ .bf16, .f8_e4m3 } }, 2)); + try std.testing.expect(!try applies(&Fake{ .dtypes = .{ .f8_e4m3, .bf16 } }, 2)); + try std.testing.expect(!try applies(&Fake{ .dtypes = .{ .bf16, .bf16 }, .present = false }, 2)); + _ = env.setenv("TF_FLASHNEXT_SPLIT_DENSE", "0", 1); + try std.testing.expect(!try applies(&bf16, 2)); + _ = env.setenv("TF_FLASHNEXT_SPLIT_DENSE", "1", 1); + try std.testing.expect(try applies(&bf16, 2)); + _ = env.setenv("TF_FLASHNEXT_SPLIT_DENSE", "on", 1); + try std.testing.expectError(error.BadDenseSplit, applies(&bf16, 2)); + try std.testing.expect(!try applies(&bf16, 1)); + try std.testing.expect(!try appliesDir(std.testing.allocator, std.testing.io, "/nonexistent", 1)); +} + +test "every model layer's heads: whole at one rank (the MTP head's too), half each at two under the split, whole without it" { + const env = @import("sys").env; + defer _ = env.unsetenv("TF_FLASHNEXT_SPLIT_DENSE"); + _ = env.unsetenv("TF_FLASHNEXT_SPLIT_DENSE"); + const bf16: Fake = .{ .dtypes = .{ .bf16, .bf16 } }; + // the MTP head's layer loads at world 1 (cuda_prompt_load.zig loadLayerAt), so its heads are the whole layer's + try std.testing.expect((try heads(&bf16, 1)).eql(.full)); + try std.testing.expect((try heads(&bf16, 2)).eql(.half)); + try std.testing.expect((try heads(&Fake{ .dtypes = .{ .f8_e4m3, .f8_e4m3 } }, 2)).eql(.full)); + _ = env.setenv("TF_FLASHNEXT_SPLIT_DENSE", "0", 1); + try std.testing.expect((try heads(&bf16, 2)).eql(.full)); + try std.testing.expect((try heads(&bf16, 1)).eql(.full)); +} + +test "the head split's per-rank bytes: stitched parts as packed once, half the conv channels, half the gate values" { + const k = 2560; + const bf = struct { + fn entry(n: usize, cols: usize) Entry { + return .{ .dtype = .bf16, .rank = 2, .shape = .{ n, cols, 1, 1, 1 }, .begin = 0, .end = n * cols * 2 }; + } + }; + const root = "model.language_model.layers.0.linear_attn"; + // 5,120 + 3,072 + 24 + 24 = 8,240 rows, padded once to 8,320: the 80 pad rows go with in_proj_qkv's 5,120 + try std.testing.expectEqual(@as(?u64, (5120 + 80) * k * 2), try headBytes(root ++ ".in_proj_qkv.weight", bf.entry(10240, k), 2)); + try std.testing.expectEqual(@as(?u64, 3072 * k * 2), try headBytes(root ++ ".in_proj_z.weight", bf.entry(6144, k), 2)); + try std.testing.expectEqual(@as(?u64, 24 * k * 2), try headBytes(root ++ ".in_proj_b.weight", bf.entry(48, k), 2)); + try std.testing.expectEqual(@as(?u64, 24 * k * 2), try headBytes(root ++ ".in_proj_a.weight", bf.entry(48, k), 2)); + const attn = "model.language_model.layers.3.self_attn"; + // 6,144 + 256 + 256 + 640 = 7,296 rows, a multiple of 128: no pad + try std.testing.expectEqual(@as(?u64, 6144 * k * 2), try headBytes(attn ++ ".q_proj.weight", bf.entry(12288, k), 2)); + try std.testing.expectEqual(@as(?u64, 256 * k * 2), try headBytes(attn ++ ".k_proj.weight", bf.entry(512, k), 2)); + try std.testing.expectEqual(@as(?u64, 256 * k * 2), try headBytes(attn ++ ".v_proj.weight", bf.entry(512, k), 2)); + try std.testing.expectEqual(@as(?u64, 640 * k * 2), try headBytes(attn ++ ".indexer.index_qk_proj.weight", bf.entry(640, k), 2)); + const conv: Entry = .{ .dtype = .bf16, .rank = 3, .shape = .{ 10240, 1, 4, 1, 1 }, .begin = 0, .end = 10240 * 4 * 2 }; + try std.testing.expectEqual(@as(?u64, 5120 * 4 * 2), try headBytes(root ++ ".conv1d.weight", conv, 2)); + const gates: Entry = .{ .dtype = .f32, .rank = 1, .shape = .{ 48, 1, 1, 1, 1 }, .begin = 0, .end = 48 * 4 }; + try std.testing.expectEqual(@as(?u64, 24 * 4), try headBytes(root ++ ".A_log", gates, 2)); + try std.testing.expectEqual(@as(?u64, 24 * 4), try headBytes(root ++ ".dt_bias", gates, 2)); + // at one rank every share is the whole tensor + try std.testing.expectEqual(@as(?u64, 10240 * 4 * 2), try headBytes(root ++ ".conv1d.weight", conv, 1)); + try std.testing.expectEqual(@as(?u64, 48 * 4), try headBytes(root ++ ".A_log", gates, 1)); + // whole on both ranks, or not BF16: not the head split's to price + for ([_][]const u8{ root ++ ".norm.weight", attn ++ ".q_norm.weight", attn ++ ".indexer.k_layernorm.weight", attn ++ ".o_proj.weight", root ++ ".out_proj.weight" }) |name| + try std.testing.expectEqual(@as(?u64, null), try headBytes(name, bf.entry(256, k), 2)); + var fp8 = bf.entry(10240, k); + fp8.dtype = .f8_e4m3; + try std.testing.expectEqual(@as(?u64, null), try headBytes(root ++ ".in_proj_qkv.weight", fp8, 2)); + try std.testing.expectError(error.UnexpectedTensor, headBytes(root ++ ".conv1d.weight", bf.entry(64, 4), 2)); + try std.testing.expectError(error.UnevenHeadSplit, headBytes(root ++ ".A_log", gates, 3)); +} + +test "the split cuts output rows for mixer and shared projections" { + try std.testing.expectEqual(Axis.rows, axisOf("model.language_model.layers.0.linear_attn.out_proj.weight").?); + try std.testing.expectEqual(Axis.rows, axisOf("model.language_model.layers.3.self_attn.o_proj.weight").?); + try std.testing.expectEqual(@as(?Axis, null), axisOf("model.language_model.layers.0.linear_attn.in_proj_qkv.weight")); + try std.testing.expectEqual(@as(?Axis, null), axisOf("model.language_model.layers.0.linear_attn.out_proj.weight_scale")); + try std.testing.expectEqual(@as(?Axis, null), axisOf("model.language_model.layers.0.mlp.gate.weight")); + try std.testing.expectEqual(@as(u64, 1280 * 6144 * 2), try residentBytes(2560, 6144, .rows, 2)); + for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }) |name| { + var buf: [160]u8 = undefined; + try std.testing.expectEqual(Axis.rows, axisOf(try std.fmt.bufPrint(&buf, "model.language_model.layers.0.mlp.shared_expert.{s}.weight", .{name})).?); + } + try std.testing.expectEqual(@as(u64, 384 * 2560 * 2), try residentBytes(640, 2560, .rows, 2)); + try std.testing.expectEqual(@as(u64, 2560 * 320 * 2), try residentBytes(2560, 640, .cols, 2)); + try std.testing.expectEqual(@as(u64, 2560 * 6144 * 2), try residentBytes(2560, 6144, .cols, 1)); + try std.testing.expectError(error.UnevenExpertSplit, residentBytes(640, 2560, .rows, 3)); + try std.testing.expectError(error.UnevenExpertSplit, residentBytes(2560, 6144, .cols, 0)); +} diff --git a/zig/src/families/flashnext/cuda_disk.zig b/zig/src/families/flashnext/cuda_disk.zig index ab8bcda0be..dc2e032e2c 100644 --- a/zig/src/families/flashnext/cuda_disk.zig +++ b/zig/src/families/flashnext/cuda_disk.zig @@ -88,14 +88,14 @@ pub fn Kept(comptime F: type) type { const head = c.head != null; const st = try c.gpa.create(snap.Snap); errdefer c.gpa.destroy(st); - st.* = .{ .buf = try cuda.DeviceBuffer.alloc(c.engine.driver, snap.bytes(at, head)), .at = at, .history = .{ 0, 0 }, .head = head }; + st.* = .{ .buf = try cuda.DeviceBuffer.alloc(c.engine.driver, snap.bytes(at, head, .full)), .at = at, .history = .{ 0, 0 }, .head = head }; return @ptrCast(st); } /// A state from a stream that did not draft lacks the head's rows this server's parts name: it stays off disk. fn whole(ptr: *anyopaque, saved: api.prompt_cache.Saved) bool { const st: *const snap.Snap = @ptrCast(@alignCast(saved)); - return st.head == (of(ptr).head != null) and st.rank_id == 0; + return st.head == (of(ptr).head != null) and st.rank_id == 0 and st.heads.eql(.full); } /// A fixed prompt's state at a mark inside its last chunk, hashed, and the prompt rate it took. @@ -149,7 +149,7 @@ test "the image covers a state's device bytes once, its rows split into 256-toke defer parts.deinit(a); try partsOf(at, head, &parts, a); try std.testing.expect(rows.image.valid(parts.items, at, rows.block_tokens)); - try std.testing.expectEqual(words + snap.bytes(at, head), rows.image.imageBytes(parts.items)); + try std.testing.expectEqual(words + snap.bytes(at, head, .full), rows.image.imageBytes(parts.items)); }; } @@ -160,7 +160,7 @@ test "the image's parts are cuda_snap.move's pieces in its order: sizes, offsets defer parts.deinit(a); try partsOf(at, head, &parts, a); try std.testing.expectEqual(@as(u64, words), parts.items[0].bytes); - const o = snap.order(at, head); + const o = snap.order(at, head, .full); var i: usize = 1; var off: u64 = words; // the image's offset; move's buffer offset is this less the words var moved: u64 = 0; @@ -185,7 +185,7 @@ test "the image's parts are cuda_snap.move's pieces in its order: sizes, offsets } } try std.testing.expectEqual(parts.items.len, i); - try std.testing.expectEqual(@as(u64, snap.bytes(at, head)), moved); + try std.testing.expectEqual(@as(u64, snap.bytes(at, head, .full)), moved); try std.testing.expectEqual(rows.image.imageBytes(parts.items), off); }; } diff --git a/zig/src/families/flashnext/cuda_fp8_experts_load.zig b/zig/src/families/flashnext/cuda_fp8_experts_load.zig new file mode 100644 index 0000000000..3c79c3e52f --- /dev/null +++ b/zig/src/families/flashnext/cuda_fp8_experts_load.zig @@ -0,0 +1,94 @@ +//! The MTP head's block-FP8 routed experts: E4M3 weights, BF16 128x128 weight_scale_inv, packed for Fp8Experts (W8A16). +const std = @import("std"); +const cuda = @import("cuda"); +const experts = @import("cuda_modelopt_experts.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const fx = cuda.experts.Fp8Experts; + +const blk = fx.block; + +/// One expert's share of each buffer: gate|up words, down words, f32 scales (gate then up), down scales. +pub const Expert = struct { up: []u32, down: []u32, up_scale: []f32, down_scale: []f32 }; + +/// The words and scales one expert takes in each buffer, `width` intermediate channels by `dims` model columns. +pub const Shape = struct { + up: usize, + down: usize, + up_scale: usize, + down_scale: usize, + + pub fn of(dims: usize, width: usize) !Shape { + if (dims == 0 or width == 0 or dims % blk != 0 or width % blk != 0) return error.UnexpectedTensor; + const weights = std.math.mul(usize, dims, width) catch return error.UnexpectedTensor; + const blocks = (dims / blk) * (width / blk); + return .{ .up = weights / 2, .down = weights / 4, .up_scale = 2 * blocks, .down_scale = blocks }; + } +}; + +/// A layer of `count` experts on one GPU; a rank cut would split the 128-wide scale blocks, so world > 1 is refused. +pub fn load(a: std.mem.Allocator, d: *const cuda.Driver, store: anytype, prefix: []const u8, world: usize, dims: usize, width: usize, count: usize) !experts.Owned { + defer store.release(); // loadShape's format probe borrowed the first gate weight + if (world != 1) return error.UnevenExpertSplit; + const one = try Shape.of(dims, width); + if (count == 0) return error.UnexpectedTensor; + const up_total = std.math.mul(usize, count, one.up * 4) catch return error.UnexpectedTensor; + const down_total = std.math.mul(usize, count, one.down * 4) catch return error.UnexpectedTensor; + var up = try cuda.DeviceBuffer.alloc(d, up_total); + errdefer up.free(); + var down = try cuda.DeviceBuffer.alloc(d, down_total); + errdefer down.free(); + const up_scale = try a.alloc(f32, count * one.up_scale); + defer a.free(up_scale); + const down_scale = try a.alloc(f32, count * one.down_scale); + defer a.free(down_scale); + const up_words = try a.alloc(u32, one.up); + defer a.free(up_words); + const down_words = try a.alloc(u32, one.down); + defer a.free(down_words); + for (0..count) |i| { + const at: Expert = .{ .up = up_words, .down = down_words, .up_scale = up_scale[i * one.up_scale ..][0..one.up_scale], .down_scale = down_scale[i * one.down_scale ..][0..one.down_scale] }; + try packExpert(a, store, prefix, i, dims, width, at); + // a failed wait may leave a copy writing `up` or `down`: drain before the errdefers free them (fromHost) + up.upload(i * one.up * 4, std.mem.sliceAsBytes(up_words)) catch |e| return drained(d, e); + down.upload(i * one.down * 4, std.mem.sliceAsBytes(down_words)) catch |e| return drained(d, e); + } + var us = try cuda.DeviceBuffer.fromHost(d, std.mem.sliceAsBytes(up_scale)); + errdefer us.free(); + const ds = try cuda.DeviceBuffer.fromHost(d, std.mem.sliceAsBytes(down_scale)); + return .{ .up = up, .down = down, .up_scale = us, .down_scale = ds, .width = width, .dims = dims, .count = count, .format = .fp8g }; +} + +fn drained(d: *const cuda.Driver, e: anyerror) anyerror { + cuda.completion.drain(d, .{ .d = d, .handle = null }, "an FP8 expert upload's default stream and context did not drain"); + return e; +} + +/// Expert `i` checked and packed as Fp8Experts.packGateUp / packOne lay it out, its scales widened to f32. +pub fn packExpert(a: std.mem.Allocator, store: anytype, prefix: []const u8, i: usize, dims: usize, width: usize, out: Expert) !void { + defer store.release(); + const one = try Shape.of(dims, width); + if (out.up.len != one.up or out.down.len != one.down or out.up_scale.len != one.up_scale or out.down_scale.len != one.down_scale) return error.UnexpectedTensor; + const half = one.up_scale / 2; + const gate = try part(store, prefix, i, "gate_proj", width, dims, out.up_scale[0..half]); + const up = try part(store, prefix, i, "up_proj", width, dims, out.up_scale[half..]); + try fx.packGateUp(a, out.up, gate, up, width, dims); + fx.packOne(out.down, try part(store, prefix, i, "down_proj", dims, width, out.down_scale), dims, width); +} + +/// One projection's E4M3 weight [n, k]; its BF16 weight_scale_inv [n/128, k/128] widened (exactly) into `scale`. +fn part(store: anytype, prefix: []const u8, i: usize, name: []const u8, n: usize, k: usize, scale: []f32) ![]const u8 { + var buf: [512]u8 = undefined; + const w = try store.get(try std.fmt.bufPrint(&buf, "{s}.experts.{d}.{s}.weight", .{ prefix, i, name })); + if (!w.is(.f8_e4m3, &.{ n, k }) or w.bytes.len != n * k) return error.UnexpectedTensor; + const s = try store.get(try std.fmt.bufPrint(&buf, "{s}.experts.{d}.{s}.weight_scale_inv", .{ prefix, i, name })); + if (!s.is(.bf16, &.{ n / blk, k / blk }) or s.bytes.len != 2 * scale.len) return error.UnexpectedTensor; + for (scale, 0..) |*v, j| { + v.* = wt.promote(std.mem.readInt(u16, s.bytes[2 * j ..][0..2], .little)); + if (!std.math.isFinite(v.*)) return error.UnsupportedQuantization; + } + return w.bytes; +} + +test { + _ = @import("cuda_fp8_experts_load_test.zig"); +} diff --git a/zig/src/families/flashnext/cuda_fp8_experts_load_test.zig b/zig/src/families/flashnext/cuda_fp8_experts_load_test.zig new file mode 100644 index 0000000000..01311a3477 --- /dev/null +++ b/zig/src/families/flashnext/cuda_fp8_experts_load_test.zig @@ -0,0 +1,306 @@ +//! Block-FP8 MTP experts on the host: scale widening and packing, world 2 refused, the weight store's load, headFp8. +const std = @import("std"); +const core = @import("core"); +const cuda = @import("cuda"); +const fp8 = @import("cuda_fp8_experts_load.zig"); +const experts = @import("cuda_modelopt_experts.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const lanes_bytes = @import("cuda_lanes_bytes.zig"); +const Fake = @import("cuda_modelopt_threads_test.zig").Fake; +const env = @import("sys").env; +const fx = cuda.experts.Fp8Experts; +const ta = std.testing.allocator; +const io = std.testing.io; + +const count = 2; +const width = 128; // NI: one scale block +const dims = 256; // D: two scale blocks +const projs = [_][]const u8{ "gate_proj", "up_proj", "down_proj" }; + +/// Two experts as the checkpoint stores them: E4M3 codes [expert][gate, up, down] and their BF16 scale_inv values. +const Set = struct { + codes: [count][3][width * dims]u8, + scales: [count][3][2]u16, // gate / up [1, 2], down [2, 1] + values: [count][3][2]f32, // the scales' values, built in f32 and exact in BF16 + releases: usize = 0, + + fn init(s: *Set, seed: u64) void { + var prng = std.Random.DefaultPrng.init(seed); + const r = prng.random(); + s.releases = 0; + for (0..count) |e| for (0..3) |p| { + r.bytes(&s.codes[e][p]); + for (&s.values[e][p], &s.scales[e][p]) |*v, *bits| { + // 2^-9 to 2^-6 with all seven mantissa bits used, so a widening that drops any shows + const m: f32 = @floatFromInt(r.uintLessThan(u8, 128)); + v.* = std.math.ldexp(1 + m / 128, -9 + @as(i32, r.uintLessThan(u8, 4))); + bits.* = @truncate(@as(u32, @bitCast(v.*)) >> 16); + } + }; + } + + fn shape(p: usize, scale: bool) [2]usize { + const n: usize = if (p == 2) dims else width; + const k: usize = if (p == 2) width else dims; + return if (scale) .{ n / 128, k / 128 } else .{ n, k }; + } + + pub fn has(_: *Set, name: []const u8) bool { + return std.mem.indexOf(u8, name, ".experts.") != null; + } + + pub fn get(s: *Set, name: []const u8) !core.safetensors.Tensor { + const marker = std.mem.indexOf(u8, name, ".experts.") orelse return error.MissingTensor; + const tail = name[marker + ".experts.".len ..]; + const dot = std.mem.indexOfScalar(u8, tail, '.') orelse return error.MissingTensor; + const e = try std.fmt.parseInt(usize, tail[0..dot], 10); + if (e >= count) return error.MissingTensor; + for (projs, 0..) |proj, p| { + const rest = tail[dot + 1 ..]; + if (!std.mem.startsWith(u8, rest, proj)) continue; + if (std.mem.eql(u8, rest[proj.len..], ".weight")) return tensor(.f8_e4m3, shape(p, false), &s.codes[e][p]); + if (std.mem.eql(u8, rest[proj.len..], ".weight_scale_inv")) return tensor(.bf16, shape(p, true), std.mem.sliceAsBytes(&s.scales[e][p])); + } + return error.MissingTensor; + } + + pub fn release(s: *Set) void { + s.releases += 1; + } +}; + +fn tensor(dtype: core.safetensors.DType, shape: [2]usize, bytes: []const u8) core.safetensors.Tensor { + var t: core.safetensors.Tensor = .{ .dtype = dtype, .rank = 2, .shape = @splat(1), .bytes = bytes }; + t.shape[0] = shape[0]; + t.shape[1] = shape[1]; + return t; +} + +const one = fp8.Shape.of(dims, width) catch unreachable; + +/// The layer's words and scales as Fp8Experts lays them out, from the set's codes and f32 values. +const Want = struct { + up: []u32, + down: []u32, + up_scale: []f32, + down_scale: []f32, + + fn init(s: *const Set) !Want { + const w: Want = .{ .up = try ta.alloc(u32, count * one.up), .down = try ta.alloc(u32, count * one.down), .up_scale = try ta.alloc(f32, count * 4), .down_scale = try ta.alloc(f32, count * 2) }; + for (0..count) |e| { + try fx.packGateUp(ta, w.up[e * one.up ..][0..one.up], &s.codes[e][0], &s.codes[e][1], width, dims); + fx.packOne(w.down[e * one.down ..][0..one.down], &s.codes[e][2], dims, width); + // [E][2][NI/128][D/128]: gate's two blocks, then up's; [E][1][D/128][NI/128]: down's two + @memcpy(w.up_scale[e * 4 ..][0..2], &s.values[e][0]); + @memcpy(w.up_scale[e * 4 + 2 ..][0..2], &s.values[e][1]); + @memcpy(w.down_scale[e * 2 ..][0..2], &s.values[e][2]); + } + return w; + } + + fn deinit(w: Want) void { + ta.free(w.up); + ta.free(w.down); + ta.free(w.up_scale); + ta.free(w.down_scale); + } + + fn expectOwned(w: Want, o: experts.Owned) !void { + try std.testing.expectEqual(cuda.experts.Format.fp8g, o.view().format); + try std.testing.expectEqual(@as(usize, width), o.view().width); + try std.testing.expectEqual(@as(usize, dims), o.view().dims); + try std.testing.expectEqual(@as(usize, count), o.view().experts); + try std.testing.expectEqual(@as(f32, 0), o.view().limit); + for ([_]struct { []const u8, cuda.DeviceBuffer }{ .{ std.mem.sliceAsBytes(w.up), o.up }, .{ std.mem.sliceAsBytes(w.down), o.down }, .{ std.mem.sliceAsBytes(w.up_scale), o.up_scale }, .{ std.mem.sliceAsBytes(w.down_scale), o.down_scale } }) |pair| { + try std.testing.expectEqualSlices(u8, pair[0], @as([*]const u8, @ptrFromInt(pair[1].ptr))[0..pair[1].len]); + } + } +}; + +test "Flash Next expert loader: FP8 MTP experts: a hand-built 2-expert set packs as Fp8Experts lays it out, BF16 block scales widened exactly" { + const s = try ta.create(Set); + defer ta.destroy(s); + s.init(0xf8); + const want = try Want.init(s); + defer want.deinit(); + const up = try ta.alloc(u32, one.up); + defer ta.free(up); + const down = try ta.alloc(u32, one.down); + defer ta.free(down); + var us: [4]f32 = undefined; + var ds: [2]f32 = undefined; + for (0..count) |e| { + try fp8.packExpert(ta, s, "language_model.mtp.layers.0.mlp", e, dims, width, .{ .up = up, .down = down, .up_scale = &us, .down_scale = &ds }); + try std.testing.expectEqualSlices(u32, want.up[e * one.up ..][0..one.up], up); + try std.testing.expectEqualSlices(u32, want.down[e * one.down ..][0..one.down], down); + try std.testing.expectEqualSlices(f32, want.up_scale[e * 4 ..][0..4], &us); + try std.testing.expectEqualSlices(f32, want.down_scale[e * 2 ..][0..2], &ds); + } + try std.testing.expectEqual(@as(usize, count), s.releases); + // the whole layer through loadShape from a plain reader, on the counting fake driver + Fake.reset(); + var d = Fake.driver(); + var owned = try experts.loadShape(ta, &d, s, "language_model.mtp.layers.0.mlp", 1, 0, dims, width, count); + try want.expectOwned(owned); + owned.free(); + try std.testing.expectEqual(@as(usize, 0), Fake.live_device); + // a NaN scale is refused, a F32 scale (as Drama 3 stores them) is not this layout; nothing stays on the device + s.scales[1][2][1] = 0x7fc0; + try std.testing.expectError(error.UnsupportedQuantization, experts.loadShape(ta, &d, s, "language_model.mtp.layers.0.mlp", 1, 0, dims, width, count)); + try std.testing.expectEqual(@as(usize, 0), Fake.live_device); + try std.testing.expectError(error.UnexpectedTensor, fp8.packExpert(ta, s, "p", 0, dims, width, .{ .up = up, .down = down[1..], .up_scale = &us, .down_scale = &ds })); +} + +test "Flash Next expert loader: FP8 MTP experts: world 2 is refused before the driver is touched, the format probe's borrow released" { + const s = try ta.create(Set); + defer ta.destroy(s); + s.init(1); + Fake.reset(); + var d = Fake.driver(); + // 128 channels cut in two are 32-column whole (Cut accepts them), but each half splits a 128-wide scale block + try std.testing.expectError(error.UnevenExpertSplit, experts.loadShape(ta, &d, s, "language_model.mtp.layers.0.mlp", 2, 0, dims, width, count)); + try std.testing.expectError(error.UnevenExpertSplit, experts.loadShape(ta, &d, s, "language_model.mtp.layers.0.mlp", 2, 1, dims, width, count)); + try std.testing.expectEqual(@as(usize, 2), s.releases); + for (Fake.calls) |n| try std.testing.expectEqual(@as(usize, 0), n); +} + +/// One tensor of a fixture shard: its checkpoint name, dtype, [rows, cols] and bytes. +const Entry = struct { name: []const u8, dtype: []const u8, shape: [2]usize, data: []const u8 }; + +/// `entries` in one shard, model.safetensors, with the index that names it. +fn fixture(dir: std.Io.Dir, entries: []const Entry) !void { + var header: std.Io.Writer.Allocating = .init(ta); + defer header.deinit(); + var index: std.Io.Writer.Allocating = .init(ta); + defer index.deinit(); + var data: std.ArrayList(u8) = .empty; + defer data.deinit(ta); + try header.writer.writeByte('{'); + try index.writer.writeAll("{\"weight_map\":{"); + for (entries, 0..) |e, i| { + if (i > 0) { + try header.writer.writeByte(','); + try index.writer.writeByte(','); + } + const at = data.items.len; + try data.appendSlice(ta, e.data); + try header.writer.print("\"{s}\":{{\"dtype\":\"{s}\",\"shape\":[{d},{d}],\"data_offsets\":[{d},{d}]}}", .{ e.name, e.dtype, e.shape[0], e.shape[1], at, data.items.len }); + try index.writer.print("\"{s}\":\"model.safetensors\"", .{e.name}); + } + try header.writer.writeByte('}'); + try index.writer.writeAll("}}"); + var image: std.Io.Writer.Allocating = .init(ta); + defer image.deinit(); + try image.writer.writeAll(&std.mem.toBytes(std.mem.nativeToLittle(u64, header.written().len))); + try image.writer.writeAll(header.written()); + try image.writer.writeAll(data.items); + try dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = image.written() }); + try dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = index.written() }); +} + +/// The set under the checkpoint's own names (mtp.layers.0...), which the store reaches as language_model.mtp... +fn setFixture(dir: std.Io.Dir, s: *const Set, names: std.mem.Allocator) !void { + var entries: [count * 3 * 2]Entry = undefined; + for (0..count) |e| for (projs, 0..) |proj, p| { + const at = (e * 3 + p) * 2; + entries[at] = .{ .name = try std.fmt.allocPrint(names, "mtp.layers.0.mlp.experts.{d}.{s}.weight", .{ e, proj }), .dtype = "F8_E4M3", .shape = Set.shape(p, false), .data = &s.codes[e][p] }; + entries[at + 1] = .{ .name = try std.fmt.allocPrint(names, "mtp.layers.0.mlp.experts.{d}.{s}.weight_scale_inv", .{ e, proj }), .dtype = "BF16", .shape = Set.shape(p, true), .data = std.mem.sliceAsBytes(&s.scales[e][p]) }; + }; + try fixture(dir, &entries); +} + +test "Flash Next expert loader: FP8 MTP experts: the weight store with default load threads takes the FP8 loop, not the threaded NVFP4 loader" { + _ = env.unsetenv("TF_FLASHNEXT_LOAD_THREADS"); // default: 8 threads, which an NVFP4 layer at world 1 takes + const s = try ta.create(Set); + defer ta.destroy(s); + s.init(7); + var arena = std.heap.ArenaAllocator.init(ta); + defer arena.deinit(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try setFixture(tmp.dir, s, arena.allocator()); + const path = try tmp.dir.realPathFileAlloc(io, ".", ta); + defer ta.free(path); + var store = try wt.Store.open(ta, io, path); + defer store.deinit(); + const want = try Want.init(s); + defer want.deinit(); + Fake.reset(); + var d = Fake.driver(); + var owned = try experts.loadShape(ta, &d, &store, "language_model.mtp.layers.0.mlp", 1, 0, dims, width, count); + try want.expectOwned(owned); + try std.testing.expectEqual(@as(usize, 0), Fake.calls[@backingInt(Fake.Kind.pinned)]); // no NVFP4 staging + try std.testing.expect(store.staging == null); + try std.testing.expectEqual(@as(usize, 0), store.held.items.len); // every borrow released + owned.free(); + try std.testing.expectEqual(@as(usize, 0), Fake.live_device); +} + +test "Flash Next expert loader: FP8 MTP experts: headFp8 reads the head's first gate weight dtype; U8 or absent is false" { + const s = try ta.create(Set); + defer ta.destroy(s); + s.init(3); + var arena = std.heap.ArenaAllocator.init(ta); + defer arena.deinit(); + const codes: [8]u8 = @splat(0); + const cases = [_]struct { entries: []const Entry, fp8: bool }{ + .{ .entries = &.{.{ .name = "mtp.layers.0.mlp.experts.0.gate_proj.weight", .dtype = "U8", .shape = .{ 2, 4 }, .data = &codes }}, .fp8 = false }, + .{ .entries = &.{.{ .name = "language_model.mtp.layers.0.mlp.experts.0.gate_proj.weight", .dtype = "F8_E4M3", .shape = .{ 2, 4 }, .data = &codes }}, .fp8 = true }, + .{ .entries = &.{.{ .name = "model.language_model.embed_tokens.weight", .dtype = "U8", .shape = .{ 2, 4 }, .data = &codes }}, .fp8 = false }, + }; + for (cases) |c| { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try fixture(tmp.dir, c.entries); + const path = try tmp.dir.realPathFileAlloc(io, ".", ta); + defer ta.free(path); + try std.testing.expectEqual(c.fp8, try lanes_bytes.headFp8(ta, io, path)); + } + // the set's own layout, under the checkpoint's mtp.layers.0 names + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try setFixture(tmp.dir, s, arena.allocator()); + const path = try tmp.dir.realPathFileAlloc(io, ".", ta); + defer ta.free(path); + try std.testing.expect(try lanes_bytes.headFp8(ta, io, path)); +} + +/// A wait that fails while `fail` is set; the context's own wait succeeds and is counted, each free sees the count. +const Waits = struct { + var fail = false; + var ctx_waits: usize = 0; + var waits_at_free: usize = std.math.maxInt(usize); + fn stream(_: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + return if (fail) cuda.abi.error_not_ready else cuda.abi.success; // a wait that has not finished, unlogged + } + fn context() callconv(.c) cuda.abi.Result { + ctx_waits += 1; + return cuda.abi.success; + } + fn free(p: cuda.abi.DevicePtr) callconv(.c) cuda.abi.Result { + waits_at_free = @min(waits_at_free, ctx_waits); + return Fake.devFree(p); + } +}; + +test "Flash Next expert loader: FP8 MTP experts: a failed upload wait drains the context before the buffers are freed" { + const s = try ta.create(Set); + defer ta.destroy(s); + s.init(3); + Fake.reset(); + var d = Fake.driver(); + d.api.cuStreamSynchronize = Waits.stream; + d.api.cuCtxSynchronize = Waits.context; + d.api.cuMemFree_v2 = Waits.free; + cuda.copy_source.setCurrent(@ptrFromInt(11)); + Waits.fail = true; + Waits.ctx_waits = 0; + Waits.waits_at_free = std.math.maxInt(usize); + defer { + Waits.fail = false; + cuda.copy_source.synchronizeContext(&d) catch unreachable; + } + try std.testing.expect(std.meta.isError(experts.loadShape(ta, &d, s, "language_model.mtp.layers.0.mlp", 1, 0, dims, width, count))); + try std.testing.expect(Waits.waits_at_free != std.math.maxInt(usize) and Waits.waits_at_free >= 1); + try std.testing.expectEqual(@as(usize, 0), Fake.live_device); +} diff --git a/zig/src/families/flashnext/cuda_gdn_fused.zig b/zig/src/families/flashnext/cuda_gdn_fused.zig index fdec57027b..4534fecb1a 100644 --- a/zig/src/families/flashnext/cuda_gdn_fused.zig +++ b/zig/src/families/flashnext/cuda_gdn_fused.zig @@ -2,6 +2,7 @@ const std = @import("std"); const cuda = @import("cuda"); const gdn = @import("cuda_gdn.zig"); +const Heads = @import("cuda_heads.zig").Heads; const parts = @import("cuda_batch_state.zig"); const eps = @import("cuda_prompt_wt.zig").eps; @@ -12,7 +13,7 @@ pub fn on() bool { } /// fn_gdn_fused's inputs: projection rows, conv window and taps, gate parameters, the state and where it goes. -pub const In = struct { proj: u64, conv: u64, cw: u64, a_log: u64, dt: u64, state: u64, state_out: u64, norm: u64 }; +pub const In = struct { proj: u64, conv: u64, cw: u64, a_log: u64, dt: u64, state: u64, state_out: u64, norm: u64, heads: Heads }; /// fn_gdn_front, fn_gdn_chain_tpr and fn_gdn_norm over `rows` rows in one launch (their bits): out and xs. pub fn launch(f: cuda.Function, stream: cuda.Stream, in: In, rows: usize, out: u64, xs: u64) !void { @@ -22,14 +23,14 @@ pub fn launch(f: cuda.Function, stream: cuda.Stream, in: In, rows: usize, out: u a.add(in.norm); a.add(eps); for ([_]u64{ out, xs, 0 }) |v| a.add(v); // no fp32 readouts kept - try cuda.launch.launch(f, .{ .grid = .{ .x = gdn.nv }, .block = .{ .x = parts.gdn_fused_threads }, .shared = parts.gdn_fused_smem }, stream, &a); + try cuda.launch.launch(f, .{ .grid = .{ .x = in.heads.nv }, .block = .{ .x = parts.gdn_fused_threads }, .shared = parts.gdn_fused_smem }, stream, &a); } /// fn_gdn_norm over `rows` readouts ys: the gated norm into out, its sums into xs. -pub fn norm(f: cuda.Function, stream: cuda.Stream, ys: u64, proj: u64, norm_w: u64, out: u64, xs: u64, rows: usize) !void { +pub fn norm(f: cuda.Function, stream: cuda.Stream, ys: u64, proj: u64, norm_w: u64, out: u64, xs: u64, rows: usize, heads: Heads) !void { var b: cuda.Args = .{}; for ([_]u64{ ys, proj, norm_w }) |v| b.add(v); b.add(eps); for ([_]u64{ out, xs }) |v| b.add(v); - try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = gdn.nv }, .block = .{ .x = gdn.dv } }, stream, &b); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = heads.nv }, .block = .{ .x = gdn.dv } }, stream, &b); } diff --git a/zig/src/families/flashnext/cuda_glue.zig b/zig/src/families/flashnext/cuda_glue.zig index 14a17dc451..aa9d9e6ea0 100644 --- a/zig/src/families/flashnext/cuda_glue.zig +++ b/zig/src/families/flashnext/cuda_glue.zig @@ -3,6 +3,7 @@ const std = @import("std"); const cuda = @import("cuda"); const wt = @import("cuda_prompt_wt.zig"); +const Heads = @import("cuda_heads.zig").Heads; /// The attention tiles' head width and query heads a KV head (flashnext_attn.cu's HD and QR). const tile_head_dim = 256; @@ -219,25 +220,26 @@ pub const Glue = struct { try go(g.f.gate, s, .{ rows, q_heads, 1 }, tile_head_dim, 0, &a); } - /// Text positions, a bf16 cache: the model's heads (cuda_prompt_wt.zig) normed, roped and stored. - pub fn attnPrep(g: *const Glue, s: cuda.Stream, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, eps: f32, rows: usize, pg: Pg) !void { + /// Text positions, a bf16 cache: the layer's heads `h` and the whole indexer normed, roped and stored. + pub fn attnPrep(g: *const Glue, s: cuda.Stream, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, eps: f32, rows: usize, h: Heads, pg: Pg) !void { if (wt.head_dim != tile_head_dim) return error.RowShape; var a: cuda.Args = .{}; ptrs(&a, &.{ proj, pos, q_scale, k_scale, i_scale, inv, q, kc, vc, iq, ikc }); a.add(eps); - ints(&a, &.{ wt.proj_n, wt.q_heads, wt.kv_heads, wt.index_heads, wt.index_dim, wt.half }); + ints(&a, &.{ h.projN(), h.q, h.kv, wt.index_heads, wt.index_dim, wt.half }); if (pg.pt != 0) a.add(pg); const f = if (pg.pt != 0) g.f.prep_paged else g.f.prep; - try go(f, s, .{ rows, wt.q_heads + wt.kv_heads + wt.index_heads + 1, 1 }, tile_head_dim, 0, &a); + try go(f, s, .{ rows, h.q + h.kv + wt.index_heads + 1, 1 }, tile_head_dim, 0, &a); } /// Attention over 512-key chunks; `qsa` the sparse form (NCH 5, rows' key lists), else dense (NCH 3). - pub fn attnChunks(g: *const Glue, s: cuda.Stream, q: u64, kc: u64, vc: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, chunks: usize, qsa: bool, pg: Pg) !void { - const group = wt.q_heads / kv; + pub fn attnChunks(g: *const Glue, s: cuda.Stream, q: u64, kc: u64, vc: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, h: Heads, chunks: usize, qsa: bool, pg: Pg) !void { + const kv: usize = h.kv; + const group = h.q / kv; if (group > tile_heads) return error.RowShape; var a: cuda.Args = .{}; ptrs(&a, &.{ q, kc, vc, pos, po, pm, pl, ids, nk, sparse }); - ints(&a, &.{ wt.q_heads, kv, group, 512, if (qsa) wt.nch else wt.dense_nch }); + ints(&a, &.{ h.q, kv, group, 512, if (qsa) wt.nch else wt.dense_nch }); a.add(@as(f32, 0.0625)); ints(&a, &.{ wt.index_idw, @intFromBool(qsa) }); if (pg.pt != 0) a.add(pg); @@ -248,23 +250,24 @@ pub const Glue = struct { } /// Prompt rows' attention with every chunk merged in the program (fn_attn_fused): out bf16, attnChunks then - pub fn attnFused(g: *const Glue, s: cuda.Stream, q: u64, kc: u64, vc: u64, pos: u64, out: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, pg: Pg) !void { - const group = wt.q_heads / kv; + pub fn attnFused(g: *const Glue, s: cuda.Stream, q: u64, kc: u64, vc: u64, pos: u64, out: u64, ids: u64, nk: u64, sparse: u64, rows: usize, h: Heads, pg: Pg) !void { + const kv: usize = h.kv; + const group = h.q / kv; if (group > tile_heads) return error.RowShape; var a: cuda.Args = .{}; ptrs(&a, &.{ q, kc, vc, pos, out, ids, nk, sparse }); - ints(&a, &.{ wt.q_heads, kv, group, 512 }); + ints(&a, &.{ h.q, kv, group, 512 }); a.add(@as(f32, 0.0625)); ints(&a, &.{ wt.index_idw, 1 }); if (pg.pt != 0) a.add(pg); try go(if (pg.pt != 0) g.f.fused_paged else g.f.fused, s, .{ rows, kv, 1 }, 128, oneSmem(), &a); } - pub fn attnMerge(g: *const Glue, s: cuda.Stream, po: u64, pm: u64, pl: u64, pos: u64, out: u64, nk: u64, sparse: u64, rows: usize, kv: usize, qsa: bool) !void { + pub fn attnMerge(g: *const Glue, s: cuda.Stream, po: u64, pm: u64, pl: u64, pos: u64, out: u64, nk: u64, sparse: u64, rows: usize, h: Heads, qsa: bool) !void { var a: cuda.Args = .{}; ptrs(&a, &.{ po, pm, pl, pos, out, nk, sparse }); - ints(&a, &.{ wt.q_heads, kv, wt.q_heads / kv, 512, if (qsa) wt.nch else wt.dense_nch, @intFromBool(qsa) }); - try go(g.f.merge, s, .{ rows, wt.q_heads, 1 }, tile_head_dim, 0, &a); + ints(&a, &.{ h.q, h.kv, h.q / h.kv, 512, if (qsa) wt.nch else wt.dense_nch, @intFromBool(qsa) }); + try go(g.f.merge, s, .{ rows, h.q, 1 }, tile_head_dim, 0, &a); } pub fn qsaPool(g: *const Glue, s: cuda.Stream, ikc: u64, pooled: u64, pos: u64, k_scale: u64, inv: u64, eps: f32, rows: usize, pg: Pg) !void { diff --git a/zig/src/families/flashnext/cuda_heads.zig b/zig/src/families/flashnext/cuda_heads.zig new file mode 100644 index 0000000000..79846c1643 --- /dev/null +++ b/zig/src/families/flashnext/cuda_heads.zig @@ -0,0 +1,91 @@ +//! A layer's head geometry: its DeltaNet key and value heads, attention query and KV heads, and the widths they set. + +const std = @import("std"); +const gdn = @import("cuda_gdn.zig"); +const wt = @import("cuda_prompt_wt.zig"); + +pub const Heads = struct { + nk: u32, + nv: u32, + q: u32, + kv: u32, + + /// The checkpoint's whole layer: what one rank holds today, and the MTP head always. + pub const full: Heads = .{ .nk = 16, .nv = 48, .q = 24, .kv = 2 }; + /// One rank's half under the head split: key and value heads halved (1:3 kept), one KV head and its 12 queries. + pub const half: Heads = .{ .nk = 8, .nv = 24, .q = 12, .kv = 1 }; + + pub fn eql(a: Heads, b: Heads) bool { + return a.nk == b.nk and a.nv == b.nv and a.q == b.q and a.kv == b.kv; + } + + /// The DeltaNet conv channels: q and k of every key head, v of every value head. + pub fn convDim(h: Heads) usize { + return 2 * h.nk * gdn.dk + h.nv * gdn.dv; + } + + /// A DeltaNet projection row: the conv channels, then z, b and a. + pub fn projWidth(h: Heads) usize { + return h.convDim() + h.nv * gdn.dv + 2 * h.nv; + } + + /// A DeltaNet output row, the out projection's input (an attention row's q * head_dim is the same width). + pub fn valueDim(h: Heads) usize { + return h.nv * gdn.dv; + } + + /// An attention projection row: each query head's q and gate, k and v of each KV head, then the whole indexer. + pub fn projN(h: Heads) usize { + return (2 * h.q + 2 * h.kv) * wt.head_dim + wt.index_heads * wt.index_dim + wt.index_dim; + } + + /// Bytes of one position's keys (or values) in the bf16 attention cache. + pub fn kvRow(h: Heads) usize { + return h.kv * wt.head_dim * 2; + } + + /// What each of `world` ranks computes of a layer the head split cuts: the whole layer at one rank, half at two. + pub fn split(world: usize) error{UnevenHeadSplit}!Heads { + return switch (world) { + 1 => .full, + 2 => .half, + else => error.UnevenHeadSplit, + }; + } +}; + +comptime { + const f = Heads.full; + std.debug.assert(f.nk == gdn.nk and f.nv == gdn.nv and f.q == wt.q_heads and f.kv == wt.kv_heads); + std.debug.assert(f.convDim() == gdn.conv_dim and f.projWidth() == gdn.proj_width and f.valueDim() == gdn.value_dim); + std.debug.assert(f.projN() == wt.proj_n and f.kvRow() == wt.kv_heads * wt.head_dim * 2); + for ([_]Heads{ f, Heads.half }) |h| std.debug.assert(h.q * wt.head_dim == h.valueDim() and h.nv == 3 * h.nk); +} + +test "the full heads' widths are the checkpoint constants they replace" { + const f = Heads.full; + try std.testing.expectEqual(gdn.conv_dim, f.convDim()); + try std.testing.expectEqual(gdn.proj_width, f.projWidth()); + try std.testing.expectEqual(gdn.value_dim, f.valueDim()); + try std.testing.expectEqual(wt.proj_n, f.projN()); + try std.testing.expectEqual(wt.out_k, f.q * wt.head_dim); + try std.testing.expectEqual(@as(usize, wt.kv_heads * wt.head_dim * 2), f.kvRow()); + try std.testing.expectEqual(@as(usize, 1024), f.kvRow()); +} + +test "half the heads: 5,120 conv channels, 8,240-wide DeltaNet rows, 3,072 values, 7,296-wide attention rows" { + const h = Heads.half; + try std.testing.expectEqual(@as(usize, 5120), h.convDim()); + try std.testing.expectEqual(@as(usize, 8240), h.projWidth()); + try std.testing.expectEqual(@as(usize, 3072), h.valueDim()); + try std.testing.expectEqual(@as(usize, 7296), h.projN()); + try std.testing.expectEqual(@as(usize, 512), h.kvRow()); + try std.testing.expectEqual(h.q / h.kv, Heads.full.q / Heads.full.kv); +} + +test "one rank holds the whole layer, each of two half of it, and no other world splits" { + try std.testing.expect((try Heads.split(1)).eql(.full)); + try std.testing.expect((try Heads.split(2)).eql(.half)); + for ([_]usize{ 0, 3, 4 }) |w| try std.testing.expectError(error.UnevenHeadSplit, Heads.split(w)); + try std.testing.expect(!Heads.half.eql(.full)); +} diff --git a/zig/src/families/flashnext/cuda_images_test.zig b/zig/src/families/flashnext/cuda_images_test.zig index a4a5b4a823..3c7d901526 100644 --- a/zig/src/families/flashnext/cuda_images_test.zig +++ b/zig/src/families/flashnext/cuda_images_test.zig @@ -42,6 +42,7 @@ fn run(comptime Cuda: type, comptime open: anytype, comptime drop: anytype, draf engine.chunk = 16; engine.eos = 7; engine.round_graph = .{}; + engine.heads = .full; var ids = [_]u32{1}; var head: mtp.Head = undefined; head.ids = &ids; @@ -140,6 +141,7 @@ fn release(comptime Cuda: type, comptime open: anytype) !void { engine.chunk = 16; engine.eos = 7; engine.round_graph = .{}; + engine.heads = .full; var self: Cuda = undefined; self.gpa = a; self.context = 64; diff --git a/zig/src/families/flashnext/cuda_kv_lanes.zig b/zig/src/families/flashnext/cuda_kv_lanes.zig index 296a396be0..d2711aa5a8 100644 --- a/zig/src/families/flashnext/cuda_kv_lanes.zig +++ b/zig/src/families/flashnext/cuda_kv_lanes.zig @@ -50,7 +50,7 @@ pub const Head = struct { draft_n: usize, rows: usize }; pub fn layout(capacity: usize, tokens: usize, head: ?Head, extra: usize) kv.Layout { const layers = wt.attn_n + @intFromBool(head != null); const run = seqRunBytes() + if (head) |h| mtp.Draft.runBytes(h.draft_n, h.rows) else 0; - return .{ .fixed_bytes = mostPages(capacity, tokens) * 4 + extra, .token_bytes = pages.pageBytesOf(tokens, layers) / tokens, .page_tokens = @intCast(tokens), .snap_bytes = snap.bytes(0, head != null), .run_bytes = run }; + return .{ .fixed_bytes = mostPages(capacity, tokens) * 4 + extra, .token_bytes = pages.pageBytesOf(tokens, layers) / tokens, .page_tokens = @intCast(tokens), .snap_bytes = snap.bytes(0, head != null, .full), .run_bytes = run }; } /// Where a stream's caches came from: no pool, the one stream on its own caches beside a pool, or a pool lane. diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 69ac0f595d..5e7fa114e5 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -202,7 +202,7 @@ pub fn Namespace(comptime Tri: type) type { .score = @import("cuda_score.zig").Hook(Cuda, Tri).call, .first_masked = firstMaskedFn, }; - const ranked_table = @import("cuda_score.zig").withoutScore(scoring_table); + const ranked_table = comptime @import("cuda_score.zig").withoutScore(scoring_table); return .{ .ptr = self, .proposal_head = if (self.head) |h| &h.service else null, .vtable = if (self.ranks == null) &scoring_table else &ranked_table }; } @@ -281,7 +281,7 @@ pub fn Namespace(comptime Tri: type) type { dropLane(self, l); self.lanes.removeByPtr(gop.key_ptr); } - if (drafting) l.draft = if (from == .pool) try self.kv.pool.?.draft(&l.seq, cap) else try mtp.Draft.init(self.gpa, self.engine.driver, self.stream, cap, self.head.?.ids.len, self.engine.chunk); + if (drafting) l.draft = if (from == .pool) try self.kv.pool.?.draft(&l.seq, cap) else try mtp.Draft.init(self.gpa, self.engine.driver, self.stream, cap, self.head.?.ids.len, self.engine.chunk, self.engine.heads); if (s.images.len != 0) { const v = self.vision orelse return error.VisionNotEnabled; l.seq.media = try v.state(ids, s.images); @@ -564,7 +564,7 @@ pub fn Namespace(comptime Tri: type) type { // -- the prompt cache's copies ------------------------------------------------------------------------ fn snapBytes(ptr: *anyopaque, at: u32) u64 { - return snap.bytes(at, of(ptr).head != null); + return snap.bytes(at, of(ptr).head != null, of(ptr).engine.heads); } /// The owner stream's state, its prompt pass standing at `at` (a drafting stream's with the head's rows). diff --git a/zig/src/families/flashnext/cuda_lanes_admission_test.zig b/zig/src/families/flashnext/cuda_lanes_admission_test.zig index 898c0bd9c6..6265ebd214 100644 --- a/zig/src/families/flashnext/cuda_lanes_admission_test.zig +++ b/zig/src/families/flashnext/cuda_lanes_admission_test.zig @@ -17,6 +17,7 @@ fn run(comptime Cuda: type, comptime open: anytype, comptime drop: anytype, coun engine.eos = 7; engine.chunk = 16; engine.round_graph = .{}; + engine.heads = .full; var ids = [_]u32{1}; var head: mtp.Head = undefined; head.ids = &ids; diff --git a/zig/src/families/flashnext/cuda_lanes_bytes.zig b/zig/src/families/flashnext/cuda_lanes_bytes.zig index bb959edc89..ba1a163a96 100644 --- a/zig/src/families/flashnext/cuda_lanes_bytes.zig +++ b/zig/src/families/flashnext/cuda_lanes_bytes.zig @@ -7,7 +7,7 @@ const config = @import("cuda_config.zig"); /// Device bytes of the checkpoint: the n-gram tables stay mapped, the MTP head counts; 0 if unread. pub fn deviceWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) u64 { const main = weights.deviceBytes(gpa, io, dir) catch return 0; - return main + mtp.headBytes(); + return main + mtp.headBytes(headFp8(gpa, io, dir) catch return 0); } pub fn placementWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize, rank: usize, drafts: bool) !u64 { @@ -16,11 +16,21 @@ pub fn placementWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, if (mode == .modelopt) { const main = try @import("cuda_rank_memory.zig").estimate(gpa, io, dir, world, rank); // Dense draft projections and vocabulary rows exceed the affine head's estimate. - return main + if (drafts and rank == 0) mtp.headBytes() + (1 << 30) else @as(u64, 0); + return main + if (drafts and rank == 0) mtp.headBytes(try headFp8(gpa, io, dir)) + (1 << 30) else @as(u64, 0); } if (world != 1) return error.UnsupportedQuantization; const main = try weights.deviceBytes(gpa, io, dir); - return main + if (drafts) mtp.headBytes() else @as(u64, 0); + return main + if (drafts) mtp.headBytes(try headFp8(gpa, io, dir)) else @as(u64, 0); +} + +/// Whether the MTP head's routed experts are block FP8: its first expert's gate weight is F8_E4M3 (absent: false). +pub fn headFp8(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !bool { + var store = try weights.Store.open(gpa, io, dir); + defer store.deinit(); + // the store's aliases take the checkpoint's mtp.layers.0... as well as language_model.mtp.layers.0... + const name = "language_model.mtp.layers.0.mlp.experts.0.gate_proj.weight"; + if (!store.has(name)) return false; + return (try store.tensor(name)).dtype == .f8_e4m3; } /// The tower's weights and one call's scratch, which rank zero holds beside the text weights with --vision. diff --git a/zig/src/families/flashnext/cuda_lanes_explain.zig b/zig/src/families/flashnext/cuda_lanes_explain.zig index 4459d56c04..dd44c86ed4 100644 --- a/zig/src/families/flashnext/cuda_lanes_explain.zig +++ b/zig/src/families/flashnext/cuda_lanes_explain.zig @@ -8,6 +8,8 @@ pub fn explain(err: anyerror) ?[]const u8 { error.BadFlashNextKvPages, error.BadFlashNextPageTest => "TF_FLASHNEXT_KV_PAGES and _PAGE_TEST take positions a page, a multiple of 4", error.FlashNextPagesUnsupported => "TF_FLASHNEXT_KV_PAGES needs one rank, our own glue and no --slide: unset it", error.UnsupportedCheckpoint => "the checkpoint differs from the supported Flash Next CUDA layout; the log names its config field", + error.BadDenseSplit => "TF_FLASHNEXT_SPLIT_DENSE takes 1 (split the BF16 dense projections at --tp 2, the default) or 0", + error.DecodeGraphWithDenseSplit => "TF_FLASHNEXT_DECODE_GRAPH is not qualified with the dense split: unset it or set TF_FLASHNEXT_SPLIT_DENSE=0", error.TensorRankDraftSnapshotUnavailable => "tensor-rank MTP snapshots are unavailable: disable the prompt cache for this configuration", error.RankDisconnected, error.RankTimeout, error.RemoteRank, error.BadRankFrame => "a tensor-parallel rank failed; restart both ranks from the same checkpoint before retrying", else => @import("cuda_images.zig").Vision.explain(err), diff --git a/zig/src/families/flashnext/cuda_lanes_load.zig b/zig/src/families/flashnext/cuda_lanes_load.zig index 5830e8705f..6c7d4c2c4d 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -20,6 +20,8 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu _ = try @import("cuda_modelopt_threads.zig").configured(); // TF_FLASHNEXT_LOAD_THREADS, refused before any load const mode = try config.mode(gpa, io, dir); if (o.tp != 1 and (mode != .modelopt or captured != null)) return error.UnsupportedQuantization; + // the decode graph captures each layer's collectives; the dense split adds its own until a later stage qualifies it + if (@import("cuda_round_graph.zig").enabled() and try @import("cuda_dense_split.zig").appliesDir(gpa, io, dir, o.tp)) return error.DecodeGraphWithDenseSplit; if (o.slide and (mode != .affine or o.tp != 1)) return error.UnsupportedLearningFormat; if (o.slide and captured != null) return error.TrainingOnCapturedGlue; const page_tokens = try @import("cuda_kv_lanes.zig").paging(); @@ -43,7 +45,7 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu id.image_tokens = v.image_tokens; id.max_images = v.max_images; } - own.ranks = try ranks.Session.init(gpa, io, o.master orelse "127.0.0.1", o.master_port, o.rank, id); + own.ranks = try ranks.Session.init(gpa, io, o.master orelse "127.0.0.1", o.master_port, o.rank, id, ranks.linkWait()); } errdefer if (own.ranks) |r| { r.poison(); @@ -81,16 +83,16 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu own.lanes = .empty; own.next = 0; if (o.slide) own.lay = try @import("cuda_slide.zig").layoutOf(gpa, io, dir); - const own_seq_bytes = try prompt.seqBytesPaged(window, try @import("cuda_pages.zig").pageTest()); + const own_seq_bytes = try prompt.seqBytesPaged(window, try @import("cuda_pages.zig").pageTest(), own.engine.heads); var stream_bytes = own_seq_bytes; if (own.head) |h| { - stream_bytes += mtp.Draft.bytes(window, h.ids.len, own.engine.chunk); - } else if (o.drafts and o.tp == 2) stream_bytes += mtp.Draft.bytes(window, 1, own.engine.chunk); + stream_bytes += mtp.Draft.bytes(window, h.ids.len, own.engine.chunk, own.engine.heads); + } else if (o.drafts and o.tp == 2) stream_bytes += mtp.Draft.bytes(window, 1, own.engine.chunk, own.engine.heads); if (o.vision) |v| stream_bytes += @as(usize, v.image_tokens) * @import("cuda_prompt_wt.zig").dims * 2 + 16 * o.context; if (page_tokens > 0) { // the lane's head runs at the engine's chunk: its cache on the stream's pages, its buffers in the run own.kv.head = if (own.head) |h| .{ .draft_n = h.ids.len, .rows = own.engine.chunk } else null; - const head_bytes = if (own.head) |h| mtp.Draft.bytes(window, h.ids.len, own.engine.chunk) else 0; + const head_bytes = if (own.head) |h| mtp.Draft.bytes(window, h.ids.len, own.engine.chunk, own.engine.heads) else 0; own.kv.layout = @import("cuda_kv_lanes.zig").layout(o.context, page_tokens, own.kv.head, stream_bytes - own_seq_bytes - head_bytes); } return .{ diff --git a/zig/src/families/flashnext/cuda_lanes_restore_test.zig b/zig/src/families/flashnext/cuda_lanes_restore_test.zig index aa441f6166..9c8921452d 100644 --- a/zig/src/families/flashnext/cuda_lanes_restore_test.zig +++ b/zig/src/families/flashnext/cuda_lanes_restore_test.zig @@ -62,6 +62,7 @@ pub fn run(comptime Cuda: type, comptime open: anytype, comptime restore: anytyp engine.eos = 7; engine.chunk = 16; engine.round_graph = .{}; + engine.heads = .full; var pair: [2]std.posix.socket_t = undefined; if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; defer _ = std.c.close(pair[0]); diff --git a/zig/src/families/flashnext/cuda_lanes_verify.zig b/zig/src/families/flashnext/cuda_lanes_verify.zig index d921a2a5a5..e6f2554ff1 100644 --- a/zig/src/families/flashnext/cuda_lanes_verify.zig +++ b/zig/src/families/flashnext/cuda_lanes_verify.zig @@ -7,7 +7,7 @@ const be = lanes.backend; /// Windows a shared pass takes at most (a round's streams). const round_streams = 32; -/// The drafted windows verified in shared passes of up to `cap` rows (alone on two ranks), then drawn. +/// The drafted windows verified in shared passes of up to `cap` rows (both ranks' at TP=2), then drawn. pub fn windows(self: anytype, wins: []const be.Window, out: []be.Verified, drafted: []const usize, row_ids: []const [mtp.max_rows]u32, cap: usize) !void { var seqs: [round_streams]*prompt.Seq = undefined; var ds: [round_streams]*mtp.Draft = undefined; @@ -25,10 +25,16 @@ pub fn windows(self: anytype, wins: []const be.Window, out: []be.Verified, draft idss[n] = row_ids[drafted[at + n]][0..w.rows()]; rows += w.rows(); } - if (n == 1 or self.ranks != null or cap <= 1) { + if (n == 1 or cap <= 1) { for (seqs[0..n], ds[0..n], idss[0..n]) |q, d, ids| try self.ops().verify(q, d, ids); + } else if (self.ranks) |r| { + try r.verifyMany(@TypeOf(self.tri()), self.tri(), seqs[0..n], ds[0..n], idss[0..n], self.logits.ptr); } else try self.ops().verifyMany(seqs[0..n], ds[0..n], idss[0..n], self.logits.ptr); for (drafted[at..][0..n], ds[0..n]) |wi, d| try self.ops().draws(d, wins[wi].stream.sampling, wins[wi].masks, out[wi].sampled[0..wins[wi].rows()]); at += n; } } + +test { + _ = @import("cuda_lanes_verify_test.zig"); +} diff --git a/zig/src/families/flashnext/cuda_lanes_verify_test.zig b/zig/src/families/flashnext/cuda_lanes_verify_test.zig new file mode 100644 index 0000000000..038d9dd7ca --- /dev/null +++ b/zig/src/families/flashnext/cuda_lanes_verify_test.zig @@ -0,0 +1,95 @@ +//! Shared verify dispatch with a recording owner: TP=2 groups go to the ranks, a lone window verifies alone. +const std = @import("std"); +const lanes = @import("lanes"); +const prompt = @import("cuda_prompt.zig"); +const mtp = @import("cuda_mtp.zig"); +const verify = @import("cuda_lanes_verify.zig"); +const be = lanes.backend; +const t = std.testing; + +/// Each call a pass made: 0 a lone verify, else a shared pass of that many windows (100 + n on the ranks). +const Log = struct { + calls: [16]usize = undefined, + n: usize = 0, + fn add(l: *Log, v: usize) void { + l.calls[l.n] = v; + l.n += 1; + } +}; + +const Owner = struct { + const Lane = struct { seq: prompt.Seq = undefined, draft: ?mtp.Draft = @as(mtp.Draft, undefined) }; // a set optional: windows() unwraps it + const Ranks = struct { + log: *Log, + pub fn verifyMany(r: *Ranks, comptime Tri: type, _: Tri, seqs: []const *prompt.Seq, ds: []const *mtp.Draft, idss: []const []const u32, _: u64) !void { + try t.expect(seqs.len == ds.len and seqs.len == idss.len); + r.log.add(100 + seqs.len); + } + }; + const Ops = struct { + log: *Log, + pub fn verify(o: Ops, _: *prompt.Seq, _: *mtp.Draft, _: []const u32) !void { + o.log.add(0); + } + pub fn verifyMany(o: Ops, seqs: []const *prompt.Seq, _: []const *mtp.Draft, _: []const []const u32, _: u64) !void { + o.log.add(seqs.len); + } + pub fn draws(_: Ops, _: *mtp.Draft, _: anytype, _: []const u32, _: []u32) !void {} + }; + log: Log = .{}, + lanes_: [4]Lane = .{ .{}, .{}, .{}, .{} }, + streams: [4]lanes.Stream = undefined, + ranks: ?*Ranks = null, + logits: struct { ptr: u64 = 0x1000 } = .{}, + + pub fn lane(o: *Owner, s: *lanes.Stream) !*Lane { + for (&o.streams, 0..) |*x, i| if (x == s) return &o.lanes_[i]; + return error.UnknownStream; + } + pub fn tri(_: *Owner) u8 { + return 0; + } + pub fn ops(o: *Owner) Ops { + return .{ .log = &o.log }; + } +}; + +fn run(o: *Owner, rows: []const usize, cap: usize) !void { + var wins: [4]be.Window = undefined; + var out: [4]be.Verified = undefined; + var sampled: [4][mtp.max_rows]u32 = undefined; + var drafted: [4]usize = undefined; + var row_ids: [4][mtp.max_rows]u32 = undefined; + for (rows, 0..) |r, i| { + o.streams[i].sampling = null; + wins[i] = .{ .stream = &o.streams[i], .pending = 0, .held = @intCast(r - 1), .tokens = &.{}, .parents = null, .positions = &.{} }; + out[i] = .{ .sampled = &sampled[i], .drafts = &.{} }; + drafted[i] = i; + row_ids[i] = @splat(1); + } + o.log = .{}; + try verify.windows(o, wins[0..rows.len], out[0..rows.len], drafted[0..rows.len], row_ids[0..rows.len], cap); +} + +test "shared verify dispatch: at TP=2 a group goes to the ranks, a lone window and a split group verify alone" { + var log: Log = .{}; + var ranks: Owner.Ranks = .{ .log = &log }; + var o: Owner = .{}; + o.ranks = &ranks; + try run(&o, &.{ 2, 3, 2 }, 64); + try t.expectEqual(@as(usize, 1), log.n); + try t.expectEqual(@as(usize, 103), log.calls[0]); + try t.expectEqual(@as(usize, 0), o.log.n); + log = .{}; + try run(&o, &.{4}, 64); + try t.expectEqual(@as(usize, 0), log.n); + try t.expectEqualSlices(usize, &.{0}, o.log.calls[0..o.log.n]); + // a cap of 4 rows splits two 3-row windows into two lone verifies + try run(&o, &.{ 3, 3 }, 4); + try t.expectEqual(@as(usize, 0), log.n); + try t.expectEqualSlices(usize, &.{ 0, 0 }, o.log.calls[0..o.log.n]); + // one rank: the shared pass stays on this engine + o.ranks = null; + try run(&o, &.{ 2, 3, 2 }, 64); + try t.expectEqualSlices(usize, &.{3}, o.log.calls[0..o.log.n]); +} diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 27d78bba79..8e8f0ddf94 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -6,10 +6,13 @@ const aot = cuda.aot; pub const embed = @import("cuda_embed.zig"); pub const projection = @import("cuda_projection.zig"); +pub const projection_load = @import("cuda_projection_load.zig"); pub const rank_engine = @import("cuda_rank_engine.zig"); pub const collective = @import("cuda_collective.zig"); pub const rank_wire = @import("cuda_rank_wire.zig"); pub const modelopt_experts = @import("cuda_modelopt_experts.zig"); +/// The MoE pass's native routed experts (NVFP4, block FP8), whose FP8 call a device test takes. +pub const moe_native = @import("cuda_moe_native.zig"); pub const hc = @import("cuda_hc.zig"); pub const draft_affine = @import("cuda_draft_affine.zig"); pub const qmm = @import("cuda_qmm.zig"); @@ -39,6 +42,10 @@ pub const native = @import("cuda_lanes.zig").Namespace(Tri); pub const Glue = glue_.Glue; /// A stream's caches through a page table, as the cache launches take them (TF_FLASHNEXT_KV_PAGES) pub const Pg = glue_.Pg; +/// A layer's head geometry, which the attention launches take (the whole layer: Heads.full) +pub const Heads = @import("cuda_heads.zig").Heads; +/// What one rank keeps of a layer under the head split, which the head split's device test cuts by +pub const heads_cut = @import("heads_cut.zig"); const glue_ = @import("cuda_glue.zig"); const p = aot.ptr; @@ -148,9 +155,10 @@ pub const Tri = struct { } /// glue.attn_prep for text positions, bf16 KV (BITS 0, the scale pointers unread); `length` unused at MODE 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, iq: u64, ikc: u64, rope: u64, delta: u64, eps: f32, rows: usize, length: usize, pg: Pg) !void { - if (t.glue("_attn_prep")) |g| return g.attnPrep(t.s, proj, pos, q_scale, k_scale, i_scale, inv, q, kc, vc, iq, ikc, eps, rows, pg); + 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, iq: u64, ikc: u64, rope: u64, delta: u64, eps: f32, rows: usize, length: usize, h: Heads, pg: Pg) !void { + if (t.glue("_attn_prep")) |g| return g.attnPrep(t.s, proj, pos, q_scale, k_scale, i_scale, inv, q, kc, vc, iq, ikc, eps, rows, h, pg); if (pg.pt != 0) return error.PagedCapturedGlue; // a captured entry reads whole windows + if (!h.eql(.full)) return error.RowShape; // and bakes in the whole layer's heads try t.run("_attn_prep", .{ u(rows), u(24 + 2 + 4 + 1), 1 }, &.{ p("P", "*bf16", proj), p("POS0", "*i32", pos), @@ -173,10 +181,11 @@ pub const Tri = struct { } /// attention._chunks over 512-key chunks, bf16 KV; `qsa` the sparse form (NCH 5, QSA 1), else dense (NCH 3, QSA 0). - pub fn attnChunks(t: Tri, q: u64, kc: u64, vc: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, chunks: usize, qsa: bool, pg: Pg) !void { - if (t.glue("_chunks")) |g| return g.attnChunks(t.s, q, kc, vc, pos, po, pm, pl, ids, nk, sparse, rows, kv, chunks, qsa, pg); + pub fn attnChunks(t: Tri, q: u64, kc: u64, vc: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, h: Heads, chunks: usize, qsa: bool, pg: Pg) !void { + if (t.glue("_chunks")) |g| return g.attnChunks(t.s, q, kc, vc, pos, po, pm, pl, ids, nk, sparse, rows, h, chunks, qsa, pg); if (pg.pt != 0) return error.PagedCapturedGlue; - try t.run("_chunks", .{ u(rows), u(kv), u(chunks) }, &.{ + if (!h.eql(.full)) return error.RowShape; + try t.run("_chunks", .{ u(rows), h.kv, u(chunks) }, &.{ p("Q", "*bf16", q), p("KC", "*bf16", kc), p("VC", "*bf16", vc), @@ -193,9 +202,10 @@ pub const Tri = struct { } /// attention._merge: one program per row and per KV head, in attnChunks' form. - pub fn attnMerge(t: Tri, po: u64, pm: u64, pl: u64, pos: u64, out: u64, nk: u64, sparse: u64, rows: usize, kv: usize, qsa: bool) !void { - if (t.glue("_merge")) |g| return g.attnMerge(t.s, po, pm, pl, pos, out, nk, sparse, rows, kv, qsa); - try t.run("_merge", .{ u(rows), u(kv), 1 }, &.{ + pub fn attnMerge(t: Tri, po: u64, pm: u64, pl: u64, pos: u64, out: u64, nk: u64, sparse: u64, rows: usize, h: Heads, qsa: bool) !void { + if (t.glue("_merge")) |g| return g.attnMerge(t.s, po, pm, pl, pos, out, nk, sparse, rows, h, qsa); + if (!h.eql(.full)) return error.RowShape; + try t.run("_merge", .{ u(rows), h.kv, 1 }, &.{ p("PO", "*fp32", po), p("PM", "*fp32", pm), p("PL", "*fp32", pl), @@ -292,9 +302,9 @@ pub const Tri = struct { } /// Prompt rows' attention, chunks merged in the program, ours only. - pub fn attnFused(t: Tri, q: u64, kc: u64, vc: u64, pos: u64, out: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, pg: Pg) !void { + pub fn attnFused(t: Tri, q: u64, kc: u64, vc: u64, pos: u64, out: u64, ids: u64, nk: u64, sparse: u64, rows: usize, h: Heads, pg: Pg) !void { const g = t.glue("_chunks") orelse return error.NoGlue; - return g.attnFused(t.s, q, kc, vc, pos, out, ids, nk, sparse, rows, kv, pg); + return g.attnFused(t.s, q, kc, vc, pos, out, ids, nk, sparse, rows, h, pg); } /// Whether every launch that reads or writes the caches is ours, so a stream's caches may be paged. @@ -413,6 +423,7 @@ pub const Tri = struct { }; test { + _ = @import("cuda_split_lane_test.zig"); _ = embed; _ = hc; _ = qmm; @@ -433,6 +444,7 @@ test { _ = @import("cuda_config.zig"); _ = @import("cuda_projection_load.zig"); _ = @import("cuda_rank_engine.zig"); + _ = @import("cuda_rank_verify.zig"); _ = @import("cuda_round_graph.zig"); _ = @import("cuda_rank_memory.zig"); _ = @import("cuda_rank_wire.zig"); @@ -442,6 +454,9 @@ test { _ = @import("cuda_mtp_check.zig"); _ = @import("cuda_draft_affine.zig"); _ = @import("cuda_head.zig"); + _ = @import("cuda_heads.zig"); + _ = @import("cuda_batch_kernels.zig"); + _ = @import("cuda_mtp_draft.zig"); _ = fn_back; _ = pocket; _ = train; @@ -502,11 +517,19 @@ test "a paged launch never reaches a captured entry" { var set: aot.Set = undefined; const t: Tri = .{ .set = &set, .s = undefined }; const pg: Pg = .{ .pt = 64, .ps = 1, .ptok = 32 }; - try std.testing.expectError(error.PagedCapturedGlue, t.attnPrep(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1e-6, 1, 0, pg)); + try std.testing.expectError(error.PagedCapturedGlue, t.attnPrep(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1e-6, 1, 0, .full, pg)); try std.testing.expectError(error.PagedCapturedGlue, t.qsaPool(0, 0, 0, 0, 0, 1e-6, 1, pg)); try std.testing.expectError(error.PagedCapturedGlue, t.qsaScores(0, 0, 0, 0, 16, 1, 1, pg)); - try std.testing.expectError(error.PagedCapturedGlue, t.attnChunks(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 1, true, pg)); - try std.testing.expectError(error.NoGlue, t.attnFused(0, 0, 0, 0, 0, 0, 0, 0, 1, 2, pg)); + try std.testing.expectError(error.PagedCapturedGlue, t.attnChunks(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, .full, 1, true, pg)); + try std.testing.expectError(error.NoGlue, t.attnFused(0, 0, 0, 0, 0, 0, 0, 0, 1, .full, pg)); +} + +test "a captured entry takes the whole layer's heads only" { + var set: aot.Set = undefined; + const t: Tri = .{ .set = &set, .s = undefined }; + try std.testing.expectError(error.RowShape, t.attnPrep(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1e-6, 1, 0, .half, .{})); + try std.testing.expectError(error.RowShape, t.attnChunks(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, .half, 1, true, .{})); + try std.testing.expectError(error.RowShape, t.attnMerge(0, 0, 0, 0, 0, 0, 0, 1, .half, true)); } test { diff --git a/zig/src/families/flashnext/cuda_media.zig b/zig/src/families/flashnext/cuda_media.zig index d8eac19a5a..72bf2c8e30 100644 --- a/zig/src/families/flashnext/cuda_media.zig +++ b/zig/src/families/flashnext/cuda_media.zig @@ -3,6 +3,7 @@ const cuda = @import("cuda"); const shared = @import("qwen_vision"); const wt = @import("cuda_prompt_wt.zig"); const Pg = @import("cuda_glue.zig").Pg; +const Heads = @import("cuda_heads.zig").Heads; pub const State = struct { image: shared.media.Media, prep: cuda.Function, @@ -23,17 +24,17 @@ pub const View = struct { pool: cuda.Function, prep_paged: cuda.Function, pool_paged: cuda.Function, - pub fn attn(v: View, stream: cuda.Stream, proj: u64, pos: u64, qw: u64, kw: u64, iw: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, eps: f32, rows: usize, pg: Pg) !void { + pub fn attn(v: View, stream: cuda.Stream, proj: u64, pos: u64, qw: u64, kw: u64, iw: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, eps: f32, rows: usize, h: Heads, pg: Pg) !void { var a: cuda.Args = .{}; for ([_]u64{ proj, pos, qw, kw, iw, inv, q, kc, vc, iq, ikc }) |x| a.add(x); a.add(eps); - for ([_]usize{ wt.proj_n, wt.q_heads, wt.kv_heads, wt.index_heads, wt.index_dim, wt.half }) |x| a.add(@as(i32, @intCast(x))); + for ([_]usize{ h.projN(), h.q, h.kv, wt.index_heads, wt.index_dim, wt.half }) |x| a.add(@as(i32, @intCast(x))); a.add(v.image.rope); a.add(@as(i32, @intCast(v.image.n))); a.add(v.image.delta); if (pg.pt != 0) a.add(pg); const f = if (pg.pt != 0) v.prep_paged else v.prep; - try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = wt.q_heads + wt.kv_heads + wt.index_heads + 1, .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, stream, &a); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(h.q + h.kv + wt.index_heads + 1), .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, stream, &a); } pub fn pooled(v: View, stream: cuda.Stream, ikc: u64, out: u64, pos: u64, w: u64, inv: u64, eps: f32, rows: usize, pg: Pg) !void { var a: cuda.Args = .{}; diff --git a/zig/src/families/flashnext/cuda_media_test.zig b/zig/src/families/flashnext/cuda_media_test.zig index 2f1b1547ce..bce139085a 100644 --- a/zig/src/families/flashnext/cuda_media_test.zig +++ b/zig/src/families/flashnext/cuda_media_test.zig @@ -45,7 +45,7 @@ test "the splice, attention and pooled indexer launches get every pointer, scala try m.image.inject(s, 4000, 2, 3, 4); try std.testing.expectEqualSlices(u64, &.{ 1000, 2000, 4000, 2, 20, 2560, 4 }, Capture.values[0][0..7]); try std.testing.expectEqual([7]u32{ 3, 40, 1, 256, 1, 1, 0 }, Capture.configs[0]); - try m.attn(s, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1e-6, 3, .{}); + try m.attn(s, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1e-6, 3, .full, .{}); const eps: u32 = @bitCast(@as(f32, 1e-6)); const delta: u32 = @bitCast(@as(i32, -2)); try std.testing.expectEqualSlices(u64, &.{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, eps, wt.proj_n, wt.q_heads, wt.kv_heads, wt.index_heads, wt.index_dim, wt.half, 3000, 20, delta }, Capture.values[1][0..21]); @@ -74,7 +74,7 @@ test "a paged stream's media rows take the paged launches: the same arguments, i Capture.calls = 0; const m = view(&d, 1000); const pg: @import("cuda_glue.zig").Pg = .{ .pt = 7000, .ps = 64 * 30784 / 2, .ptok = 64 }; - try m.attn(s, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1e-6, 3, pg); + try m.attn(s, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1e-6, 3, .full, pg); try m.pooled(s, 1, 2, 3, 4, 5, 1e-6, 8, pg); const eps: u32 = @bitCast(@as(f32, 1e-6)); const delta: u32 = @bitCast(@as(i32, -2)); diff --git a/zig/src/families/flashnext/cuda_mixed_experts.zig b/zig/src/families/flashnext/cuda_mixed_experts.zig index b26225731d..bf6efb0f1f 100644 --- a/zig/src/families/flashnext/cuda_mixed_experts.zig +++ b/zig/src/families/flashnext/cuda_mixed_experts.zig @@ -1,4 +1,4 @@ -//! Routed NVFP4 intermediate partitions and a separately owned MXFP8 shared expert. +//! Whole routed experts per rank and a shared expert with optional exact output-column partitioning. const std = @import("std"); const cuda = @import("cuda"); const wt = @import("cuda_prompt_wt.zig"); @@ -10,6 +10,7 @@ pub const Owned = struct { shared: [3]wt.Q4, rank: usize, world: usize, + split_shared: bool = false, pub fn free(self: *Owned) void { self.experts.free(); @@ -17,20 +18,21 @@ pub const Owned = struct { } }; -pub fn load(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, world: usize, rank: usize) !Owned { +pub fn load(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, world: usize, rank: usize, split: bool) !Owned { var e = try routed.load(a, d, store, prefix, world, rank); errdefer e.free(); var shared: [3]wt.Q4 = @splat(wt.Q4.empty(d)); errdefer for (&shared) |*p| p.free(); // The shared branch participates exactly once in the sum across ranks. - if (rank == 0) { + if (rank == 0 or split) { var buf: [256]u8 = undefined; for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }, 0..) |name, i| { const n = if (i == 2) wt.dims else wt.moe_w; const k = if (i == 2) wt.moe_w else wt.dims; - shared[i] = try projections.load(a, d, store, try std.fmt.bufPrint(&buf, "{s}.shared_expert.{s}", .{ prefix, name }), n, k); + const key = try std.fmt.bufPrint(&buf, "{s}.shared_expert.{s}", .{ prefix, name }); + shared[i] = if (split) try projections.loadCut(a, d, store, key, n, k, try routed.Cut.init(n, world, rank), .rows) else try projections.load(a, d, store, key, n, k); if (shared[i].native == null) return error.UnsupportedQuantization; } } - return .{ .experts = e, .shared = shared, .rank = rank, .world = world }; + return .{ .experts = e, .shared = shared, .rank = rank, .world = world, .split_shared = split }; } diff --git a/zig/src/families/flashnext/cuda_modelopt_experts.zig b/zig/src/families/flashnext/cuda_modelopt_experts.zig index 7ed5be798c..7bae50f6cf 100644 --- a/zig/src/families/flashnext/cuda_modelopt_experts.zig +++ b/zig/src/families/flashnext/cuda_modelopt_experts.zig @@ -1,4 +1,4 @@ -//! Rank-local intermediate-channel TP for NVFP4 routed experts; the shared expert retains its own projection format. +//! Routed experts: NVFP4 (whole experts a rank at --tp), or block FP8 at world 1; the shared expert keeps its format. const std = @import("std"); const core = @import("core"); const cuda = @import("cuda"); @@ -6,6 +6,7 @@ const wt = @import("cuda_prompt_wt.zig"); const ex = cuda.experts.Nvfp4Experts; const slice = core.nvfp4_slice; const threads = @import("cuda_modelopt_threads.zig"); +const fp8 = @import("cuda_fp8_experts_load.zig"); pub const Cut = struct { from: usize, @@ -23,6 +24,17 @@ pub const Cut = struct { } }; +/// A rank's routed experts at --tp: `count / world` whole ones from `first`, so every pair keeps one GPU's sums. +pub const Experts = struct { + first: usize, + count: usize, + + pub fn init(total: usize, world: usize, rank: usize) !Experts { + if (world == 0 or rank >= world or total == 0 or total % world != 0) return error.UnevenExpertSplit; + return .{ .first = rank * (total / world), .count = total / world }; + } +}; + pub const Owned = struct { up: cuda.DeviceBuffer, down: cuda.DeviceBuffer, @@ -31,6 +43,10 @@ pub const Owned = struct { width: usize, dims: usize, count: usize, + /// The global id of the rank's first expert (0 at world 1); `count` experts from it are this rank's. + first: usize = 0, + /// .nvfp4 (global scales a projection) or .fp8g (f32 128x128 block scales, cuda_fp8_experts_load.zig). + format: cuda.experts.Format = .nvfp4, pub fn free(self: *Owned) void { self.up.free(); @@ -41,7 +57,7 @@ pub const Owned = struct { } pub fn view(self: *const Owned) cuda.experts.Layer { - return .{ .format = .nvfp4, .up = self.up.ptr, .down = self.down.ptr, .up_scale = self.up_scale.ptr, .down_scale = self.down_scale.ptr, .width = self.width, .dims = self.dims, .experts = self.count }; + return .{ .format = self.format, .up = self.up.ptr, .down = self.down.ptr, .up_scale = self.up_scale.ptr, .down_scale = self.down_scale.ptr, .width = self.width, .dims = self.dims, .experts = self.count }; } }; @@ -65,8 +81,11 @@ pub fn load(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []c } pub fn loadShape(a: std.mem.Allocator, d: *const cuda.Driver, store: anytype, prefix: []const u8, world: usize, rank: usize, dims: usize, full_width: usize, count: usize) !Owned { - const cut = try Cut.init(full_width, world, rank); + const cut = try Cut.init(full_width, 1, 0); // every expert whole: --tp splits experts, not their channels + const own = try Experts.init(count, world, rank); if (dims == 0 or dims % 32 != 0 or count == 0) return error.UnexpectedTensor; + // F8_E4M3 weights load in their own loop; U8 (NVFP4) goes on to the threaded or one-expert paths below + if (try formatOf(store, prefix) == .fp8g) return fp8.load(a, d, store, prefix, world, dims, full_width, count); // one GPU and the weight store: packed on threads (cuda_modelopt_threads.zig); other readers or 1 thread: this loop if (comptime @typeInfo(@TypeOf(store)).pointer.child == wt.Store) { const n = try threads.configured(); @@ -75,32 +94,33 @@ pub fn loadShape(a: std.mem.Allocator, d: *const cuda.Driver, store: anytype, pr const width = cut.width(); const up_bytes = try packedBytes(width, dims, 2); const down_bytes = try packedBytes(dims, width, 1); - const up_total = std.math.mul(usize, count, up_bytes) catch return error.UnexpectedTensor; - const down_total = std.math.mul(usize, count, down_bytes) catch return error.UnexpectedTensor; - const globals_count = std.math.mul(usize, count, 2) catch return error.UnexpectedTensor; + const up_total = std.math.mul(usize, own.count, up_bytes) catch return error.UnexpectedTensor; + const down_total = std.math.mul(usize, own.count, down_bytes) catch return error.UnexpectedTensor; + const globals_count = std.math.mul(usize, own.count, 2) catch return error.UnexpectedTensor; var up = try cuda.DeviceBuffer.alloc(d, up_total); errdefer up.free(); var down = try cuda.DeviceBuffer.alloc(d, down_total); errdefer down.free(); const up_globals = try a.alloc(f32, globals_count); defer a.free(up_globals); - const down_globals = try a.alloc(f32, count); + const down_globals = try a.alloc(f32, own.count); defer a.free(down_globals); const packed_up = try a.alloc(u32, up_bytes / 4); defer a.free(packed_up); const packed_down = try a.alloc(u32, down_bytes / 4); defer a.free(packed_down); var buf: [512]u8 = undefined; - for (0..count) |i| { - var gate = try readPart(a, store, try std.fmt.bufPrint(&buf, "{s}.experts.{d}.gate_proj", .{ prefix, i }), cut, false, dims, full_width); + for (0..own.count) |i| { + const e = own.first + i; + var gate = try readPart(a, store, try std.fmt.bufPrint(&buf, "{s}.experts.{d}.gate_proj", .{ prefix, e }), cut, false, dims, full_width); defer gate.deinit(a); - var value = try readPart(a, store, try std.fmt.bufPrint(&buf, "{s}.experts.{d}.up_proj", .{ prefix, i }), cut, false, dims, full_width); + var value = try readPart(a, store, try std.fmt.bufPrint(&buf, "{s}.experts.{d}.up_proj", .{ prefix, e }), cut, false, dims, full_width); defer value.deinit(a); ex.packGateUp(packed_up, .{ gate.codes, gate.scales }, .{ value.codes, value.scales }, width, dims); try up.upload(i * up_bytes, std.mem.sliceAsBytes(packed_up)); up_globals[i * 2] = gate.global; up_globals[i * 2 + 1] = value.global; - var output = try readPart(a, store, try std.fmt.bufPrint(&buf, "{s}.experts.{d}.down_proj", .{ prefix, i }), cut, true, dims, full_width); + var output = try readPart(a, store, try std.fmt.bufPrint(&buf, "{s}.experts.{d}.down_proj", .{ prefix, e }), cut, true, dims, full_width); defer output.deinit(a); ex.packOne(packed_down, output.codes, output.scales, dims, width); try down.upload(i * down_bytes, std.mem.sliceAsBytes(packed_down)); @@ -109,7 +129,30 @@ pub fn loadShape(a: std.mem.Allocator, d: *const cuda.Driver, store: anytype, pr var us = try cuda.DeviceBuffer.fromHost(d, std.mem.sliceAsBytes(up_globals)); errdefer us.free(); const ds = try cuda.DeviceBuffer.fromHost(d, std.mem.sliceAsBytes(down_globals)); - return .{ .up = up, .down = down, .up_scale = us, .down_scale = ds, .width = width, .dims = dims, .count = count }; + return .{ .up = up, .down = down, .up_scale = us, .down_scale = ds, .width = width, .dims = dims, .count = own.count, .first = own.first }; +} + +/// The layer's format from its first expert's gate weight dtype; the loader's next release returns the borrow. +fn formatOf(store: anytype, prefix: []const u8) !cuda.experts.Format { + var buf: [512]u8 = undefined; + const w = try store.get(try std.fmt.bufPrint(&buf, "{s}.experts.0.gate_proj.weight", .{prefix})); + errdefer store.release(); // a refused dtype returns the borrow here; the loaders return it otherwise + return switch (w.dtype) { + .u8 => .nvfp4, + .f8_e4m3 => .fp8g, + else => error.UnsupportedQuantization, + }; +} + +test "expert-parallel ranges give each rank whole experts and cover every expert once" { + const a = try Experts.init(512, 2, 0); + const b = try Experts.init(512, 2, 1); + try std.testing.expectEqual(@as(usize, 0), a.first); + try std.testing.expectEqual(a.count, b.first); + try std.testing.expectEqual(@as(usize, 512), b.first + b.count); + try std.testing.expectEqual(Experts{ .first = 0, .count = 512 }, try Experts.init(512, 1, 0)); + try std.testing.expectError(error.UnevenExpertSplit, Experts.init(512, 2, 2)); + try std.testing.expectError(error.UnevenExpertSplit, Experts.init(511, 2, 0)); } test "intermediate-channel cuts cover the source once and preserve 32-input kernel groups" { diff --git a/zig/src/families/flashnext/cuda_modelopt_threads_test.zig b/zig/src/families/flashnext/cuda_modelopt_threads_test.zig index f9e230906d..1dd5b965e8 100644 --- a/zig/src/families/flashnext/cuda_modelopt_threads_test.zig +++ b/zig/src/families/flashnext/cuda_modelopt_threads_test.zig @@ -68,15 +68,15 @@ const ta = std.testing.allocator; const env = @import("sys").env; /// A driver whose device memory is host memory, counting calls and live allocations; cuda.memory's limit refuses. -const Fake = struct { - const Kind = enum { pinned, device, copy }; - var calls: [3]usize = @splat(0); - var live_pinned: usize = 0; - var live_device: usize = 0; +pub const Fake = struct { + pub const Kind = enum { pinned, device, copy }; + pub var calls: [3]usize = @splat(0); + pub var live_pinned: usize = 0; + pub var live_device: usize = 0; var spans: [256][]u8 = undefined; var nspans: usize = 0; - fn reset() void { + pub fn reset() void { calls = @splat(0); } fn tally(k: Kind) void { @@ -117,7 +117,7 @@ const Fake = struct { live_device += 1; return abi.success; } - fn devFree(p: abi.DevicePtr) callconv(.c) abi.Result { + pub fn devFree(p: abi.DevicePtr) callconv(.c) abi.Result { give(p); live_device -= 1; return abi.success; @@ -137,7 +137,7 @@ const Fake = struct { s.* = null; return 1; } - fn driver() cuda.Driver { + pub fn driver() cuda.Driver { var d: cuda.Driver = undefined; d.func_name = null; d.api.cuMemHostAlloc = hostAlloc; diff --git a/zig/src/families/flashnext/cuda_moe_native.zig b/zig/src/families/flashnext/cuda_moe_native.zig new file mode 100644 index 0000000000..086f0d92b9 --- /dev/null +++ b/zig/src/families/flashnext/cuda_moe_native.zig @@ -0,0 +1,60 @@ +//! The MoE pass's native experts: NVFP4 or block-FP8 pairs (whole experts a rank at --tp), shared on rank 0, rank sum. +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const mixed = @import("cuda_mixed_experts.zig"); +const nvfp4_tile = @import("cuda_nvfp4_tile.zig"); +const zs = @import("cuda_zero_slice.zig"); +const Kernels = @import("cuda_batch_kernels.zig").Kernels; + +const dims = wt.dims; +const experts_n = wt.experts_n; +const slots = wt.slots; + +/// moeBlockKeep's native layer: each routed pair's fp32 row into y [pairs, dims], the shared slot's on rank 0. +pub fn run(self: anytype, native: *const mixed.Owned, n: usize, pairs: usize, plan: cuda.grouped.Plan) !void { + const s = self.s; + const owned = &native.experts; + // the rank's experts and one skip id past them (the shared slot at world 1; every other rank's pick at --tp) + const e = owned.count + 1; + const skip: i32 = @intCast(owned.count); + const picks = if (native.world == 1) s.picks else try ownPicks(self, s.picks, s.picks_ep, pairs, owned.first, owned.count); + if (zs.clearsY(native.world, native.rank)) try self.stream.d.check(self.stream.d.api.cuMemsetD8Async(s.y, 0, pairs * dims * 4, self.stream.handle), "cuMemsetD8Async"); + switch (owned.format) { + .nvfp4 => { + const pt = nvfp4_tile.size(n); + try self.k.router.route(self.stream, picks, pairs, e, pt, plan); + const items = cuda.grouped.maxItems(pairs, e, pt); + try nvfp4_tile.run(self.k, pt, self.stream, .up, .{ .x = s.mixed, .stride = dims, .slots = slots }, owned.view(), plan, s.moe_act, items, skip); + try nvfp4_tile.run(self.k, pt, self.stream, .down_f32, .{ .x = s.moe_act, .stride = owned.width }, owned.view(), plan, s.y, items, skip); + }, + .fp8g => try fp8(self.k, self.stream, s.picks, s.mixed, owned.view(), plan, pairs, n, s.y, s.moe_act), + else => return error.UnsupportedQuantization, + } + try @import("cuda_shared_split.zig").run(self, native, n); + if (native.world != 1) { + const collective = self.m.collective orelse return error.TensorParallelUnavailable; + try collective.sum(self.stream, s.y, pairs * dims); + } +} + +/// `picks` as this rank's expert ids into `out`: its own experts from 0, every other id (and the shared slot) `count`. +fn ownPicks(self: anytype, picks: u64, out: u64, pairs: usize, first: usize, count: usize) !u64 { + var a: cuda.Args = .{}; + a.add(picks); + a.add(out); + for ([_]usize{ pairs, first, count }) |v| a.add(@as(i32, @intCast(v))); + try cuda.launch.launch(self.k.ep_picks, .{ .grid = .{ .x = @intCast((pairs + 255) / 256) }, .block = .{ .x = 256 } }, self.stream, &a); + return out; +} + +/// Block-FP8 pairs (W8A16) at NVFP4's decode/prompt tile: `act` bf16, `out` fp32; the shared slot is skipped. +pub fn fp8(k: *const Kernels, stream: cuda.Stream, picks: u64, x: u64, layer: cuda.experts.Layer, plan: cuda.grouped.Plan, pairs: usize, n: usize, out: u64, act: u64) !void { + if (layer.format != .fp8g) return error.UnsupportedQuantization; + const e = experts_n + 1; + const prompt = nvfp4_tile.size(n) != 16; + const tile: usize = if (prompt) cuda.experts.promptTile(.fp8g) else 16; + try k.router.route(stream, picks, pairs, e, tile, plan); + const items = cuda.grouped.maxItems(pairs, e, tile); + try k.fp8.run(stream, .up, .{ .x = x, .stride = layer.dims, .slots = slots }, layer, plan, prompt, act, items, experts_n); + try k.fp8.run(stream, .down_f32, .{ .x = act, .stride = layer.width }, layer, plan, prompt, out, items, experts_n); +} diff --git a/zig/src/families/flashnext/cuda_mtp.zig b/zig/src/families/flashnext/cuda_mtp.zig index 8fe177ec9a..f1191cb48e 100644 --- a/zig/src/families/flashnext/cuda_mtp.zig +++ b/zig/src/families/flashnext/cuda_mtp.zig @@ -258,9 +258,12 @@ pub const Head = struct { } }; -/// Device bytes the head's weights take beyond the main model's (for the server's admission estimate). -pub fn headBytes() usize { - const experts = (wt.experts_n + 1) * 3 * wt.moe_w * dims / 2 * 5 / 4; +/// Device bytes the head's weights take beyond the main model's; `fp8`: block-FP8 routed experts (headFp8). +pub fn headBytes(fp8: bool) usize { + const weights = 3 * wt.moe_w * dims; + const nvfp4 = weights / 2 * 5 / 4; + const routed = if (fp8) weights + 3 * (wt.moe_w / 128) * (dims / 128) * 4 else nvfp4; + const experts = wt.experts_n * routed + nvfp4; const attn = (wt.proj_n + dims) * dims / 2 * 5 / 4; const draft = 79_591 * dims / 2 * 5 / 4; return experts + attn + draft + 2 * dims * dims / 2 * 5 / 4; @@ -272,6 +275,13 @@ pub const Draft = @import("cuda_mtp_draft.zig").Draft; pub const Ops = @import("cuda_mtp_ops.zig").Ops; pub const check = @import("cuda_mtp_check.zig").check; +test "the head's bytes: NVFP4 experts as before, FP8 routed experts a byte a weight and their f32 block scales" { + const before = (wt.experts_n + 1) * 3 * wt.moe_w * dims / 2 * 5 / 4 + (wt.proj_n + dims) * dims / 2 * 5 / 4 + 79_591 * dims / 2 * 5 / 4 + 2 * dims * dims / 2 * 5 / 4; + try std.testing.expectEqual(before, headBytes(false)); + // an expert's three projections: 640 x 2560 bytes and 5 x 20 scales each, against 5/8 of a byte a weight + try std.testing.expectEqual(wt.experts_n * (3 * (640 * 2560 + 5 * 20 * 4) - 3 * 640 * 2560 * 5 / 8), headBytes(true) - headBytes(false)); +} + test "the draft vocabulary is 79,591 sorted ids under the vocabulary" { const ids = try draftIds(std.testing.allocator); defer std.testing.allocator.free(ids); diff --git a/zig/src/families/flashnext/cuda_mtp_check.zig b/zig/src/families/flashnext/cuda_mtp_check.zig index 53e3b2d057..3cef40b1ee 100644 --- a/zig/src/families/flashnext/cuda_mtp_check.zig +++ b/zig/src/families/flashnext/cuda_mtp_check.zig @@ -117,7 +117,7 @@ pub fn check(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cu const cap = ids.len + count + max_rows; var seq = try e.newSeq(cap); defer seq.deinit(gpa); - var d = try Draft.init(gpa, driver, stream.*, cap, h.ids.len, e.chunk); + var d = try Draft.init(gpa, driver, stream.*, cap, h.ids.len, e.chunk, e.heads); defer d.deinit(); for (0..repeat) |rep| { // serial: the plain prompt pass, then one row a forward diff --git a/zig/src/families/flashnext/cuda_mtp_draft.zig b/zig/src/families/flashnext/cuda_mtp_draft.zig index c2f23d0598..3ea3294370 100644 --- a/zig/src/families/flashnext/cuda_mtp_draft.zig +++ b/zig/src/families/flashnext/cuda_mtp_draft.zig @@ -16,18 +16,14 @@ const vocab = wt.vocab; const eps = wt.eps; const layers_n = wt.layers_n; const linear_n = wt.linear_n; -const state_n = wt.state_n; const out_sk = wt.out_sk; -const kv_heads = wt.kv_heads; -const head_dim = wt.head_dim; -const index_dim = wt.index_dim; const ple_tail_n = wt.ple_tail_n; const Q4 = wt.Q4; const HcW = wt.HcW; const max_chain = mtp.max_chain; const max_rows = mtp.max_rows; -const conv_bytes = mtp.conv_bytes; -const proj_bytes = mtp.proj_bytes; +const sizes = @import("cuda_seq_bytes.zig"); +const Heads = @import("cuda_heads.zig").Heads; /// One draft and its temperature-1 probability under the head. pub const Pick = struct { token: u32, prob: f64 }; @@ -74,6 +70,8 @@ pub const Draft = struct { saved: cuda.DeviceBuffer, /// A window's DeltaNet end states: swapped in when every row is kept, else the kept rows replay on the sequence's. spare: [linear_n]cuda.DeviceBuffer = undefined, + /// Where `saved` holds each piece and how large the spares are, at the model's heads (its sequence's). + layout: Layout, /// A window ran with the states written to `spare` and is not yet kept or rolled back. open: bool = false, keep: batch.Keep = .{}, @@ -86,28 +84,49 @@ pub const Draft = struct { host: []u16, best_host: [max_rows * 4]f32 = undefined, - const hist_at: usize = 0; - const proj_at: usize = hist_at + linear_n * conv_bytes; - pub const tail_at: usize = proj_at + linear_n * max_rows * proj_bytes; - /// The window's n-gram rows (the pass's ple_nrow), saved for a rollback to shift the kept ones into the tail. - pub const nrow_at: usize = tail_at + ple_tail_n * wide * 2; - const saved_bytes: usize = nrow_at + max_rows * wide * 2; + /// `saved`'s pieces and the spares' size at the model's heads: conv windows, projection rows, tail, n-gram rows. + pub const Layout = struct { + heads: Heads, + conv: usize, + proj: usize, + hist_at: usize = 0, + proj_at: usize, + tail_at: usize, + /// The window's n-gram rows (the pass's ple_nrow), saved for a rollback to shift the kept ones into the tail. + nrow_at: usize, + saved: usize, + state: usize, - /// Device bytes `init(cap, draft_n, rows)` allocates. - pub fn bytes(cap: usize, draft_n: usize, rows: usize) usize { - return cap * (2 * prompt.kv_row_bytes + index_dim * 2) + prompt.pooledBytes(cap) + @max(rows, max_rows) * wide * 2 + wide * 2 + max_rows * vocab * 2 + draft_n * 2 + max_rows * 16 + saved_bytes + linear_n * state_n * 4; + pub fn of(h: Heads) Layout { + const conv = sizes.convBytes(h); + const proj = h.projWidth() * 2; + const proj_at = linear_n * conv; + const tail_at = proj_at + linear_n * max_rows * proj; + const nrow_at = tail_at + ple_tail_n * wide * 2; + return .{ .heads = h, .conv = conv, .proj = proj, .proj_at = proj_at, .tail_at = tail_at, .nrow_at = nrow_at, .saved = nrow_at + max_rows * wide * 2, .state = sizes.stateBytes(h) }; + } + }; + + /// The head's own cache: its layer is whole on every rank. + const head_heads: Heads = .full; + + /// Device bytes `init(cap, draft_n, rows, h)` allocates: the head's cache at full KV rows, the rest at `h`. + pub fn bytes(cap: usize, draft_n: usize, rows: usize, h: Heads) usize { + const l = Layout.of(h); + const head = 2 * sizes.kvBytes(cap, head_heads) + sizes.indexBytes(cap) + sizes.pooledBytes(cap); + return head + @max(rows, max_rows) * wide * 2 + wide * 2 + max_rows * vocab * 2 + draft_n * 2 + max_rows * 16 + l.saved + linear_n * l.state; } - /// A head cache of `cap` entries; `rows` the widest prompt chunk whose streams it keeps. - pub fn init(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: cuda.Stream, cap: usize, draft_n: usize, rows: usize) !Draft { - var d: Draft = .{ .stream = stream, .gpa = gpa, .cap = cap, .rows_cap = @max(rows, max_rows), .kc = undefined, .vc = undefined, .ikc = undefined, .pooled = undefined, .streams = undefined, .out = undefined, .logits = undefined, .dlogits = undefined, .best = undefined, .saved = undefined, .host = undefined }; - d.kc = try cuda.DeviceBuffer.alloc(driver, cap * kv_heads * head_dim * 2); + /// A head cache of `cap` entries; `rows` the widest prompt chunk whose streams it keeps; `h` the model's heads. + pub fn init(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: cuda.Stream, cap: usize, draft_n: usize, rows: usize, h: Heads) !Draft { + var d: Draft = .{ .stream = stream, .gpa = gpa, .cap = cap, .rows_cap = @max(rows, max_rows), .kc = undefined, .vc = undefined, .ikc = undefined, .pooled = undefined, .streams = undefined, .out = undefined, .logits = undefined, .dlogits = undefined, .best = undefined, .saved = undefined, .host = undefined, .layout = .of(h) }; + d.kc = try cuda.DeviceBuffer.alloc(driver, sizes.kvBytes(cap, head_heads)); errdefer d.kc.free(); - d.vc = try cuda.DeviceBuffer.alloc(driver, cap * kv_heads * head_dim * 2); + d.vc = try cuda.DeviceBuffer.alloc(driver, sizes.kvBytes(cap, head_heads)); errdefer d.vc.free(); - d.ikc = try cuda.DeviceBuffer.alloc(driver, cap * index_dim * 2); + d.ikc = try cuda.DeviceBuffer.alloc(driver, sizes.indexBytes(cap)); errdefer d.ikc.free(); - d.pooled = try cuda.DeviceBuffer.alloc(driver, prompt.pooledBytes(cap)); + d.pooled = try cuda.DeviceBuffer.alloc(driver, sizes.pooledBytes(cap)); errdefer d.pooled.free(); d.streams = try cuda.DeviceBuffer.alloc(driver, d.rows_cap * wide * 2); errdefer d.streams.free(); @@ -119,11 +138,11 @@ pub const Draft = struct { errdefer d.dlogits.free(); d.best = try cuda.DeviceBuffer.alloc(driver, max_rows * 16); errdefer d.best.free(); - d.saved = try cuda.DeviceBuffer.alloc(driver, saved_bytes); + d.saved = try cuda.DeviceBuffer.alloc(driver, d.layout.saved); errdefer d.saved.free(); var made: usize = 0; errdefer for (d.spare[0..made]) |*b| b.free(); - while (made < linear_n) : (made += 1) d.spare[made] = try cuda.DeviceBuffer.alloc(driver, state_n * 4); + while (made < linear_n) : (made += 1) d.spare[made] = try cuda.DeviceBuffer.alloc(driver, d.layout.state); d.host = try gpa.alloc(u16, @max(max_rows * vocab, draft_n)); errdefer gpa.free(d.host); errdefer cuda.completion.drain(driver, stream, "Flash Next draft fills did not drain before rollback"); @@ -138,8 +157,8 @@ pub const Draft = struct { var li: usize = 0; for (&d.keep.layers, 0..) |*lk, layer| { if (!wt.linearLayer(layer)) continue; - lk.hist = d.saved.ptr + hist_at + li * conv_bytes; - lk.proj = d.saved.ptr + proj_at + li * max_rows * proj_bytes; + lk.hist = d.saved.ptr + d.layout.hist_at + li * d.layout.conv; + lk.proj = d.saved.ptr + d.layout.proj_at + li * max_rows * d.layout.proj; li += 1; } d.aim(); @@ -147,9 +166,12 @@ pub const Draft = struct { /// Device bytes `initIn` takes from a run: every buffer but the head's cache, each on a 256-byte boundary. pub fn runBytes(draft_n: usize, rows: usize) usize { - return bytes(0, draft_n, rows) + (6 + linear_n) * 256; + return bytes(0, draft_n, rows, pool_heads) + (6 + linear_n) * 256; } + /// A pool lane's model heads: the pages hold the whole layer (one rank; cuda_lanes_load.zig refuses pages at --tp). + const pool_heads: Heads = .full; + /// A pool lane's head: its cache on `paged`'s pages after the model's layers, its buffers from `base`'s run. pub fn initIn(gpa: std.mem.Allocator, driver: *cuda.Driver, cap: usize, draft_n: usize, rows: usize, paged: *const pages.Paged, base: u64, len: usize) !Draft { if (len < runBytes(draft_n, rows)) return error.FlashNextRunTooShort; @@ -165,14 +187,14 @@ pub const Draft = struct { } }; const head = wt.attn_n; - var d: Draft = .{ .gpa = gpa, .cap = cap, .rows_cap = @max(rows, max_rows), .kc = Carve.page(driver, paged.at(head, .key)), .vc = Carve.page(driver, paged.at(head, .value)), .ikc = Carve.page(driver, paged.at(head, .index)), .pooled = Carve.page(driver, paged.at(head, .pooled)), .pg = paged.pg(), .streams = undefined, .out = undefined, .logits = undefined, .dlogits = undefined, .best = undefined, .saved = undefined, .host = undefined }; + var d: Draft = .{ .gpa = gpa, .cap = cap, .rows_cap = @max(rows, max_rows), .kc = Carve.page(driver, paged.at(head, .key)), .vc = Carve.page(driver, paged.at(head, .value)), .ikc = Carve.page(driver, paged.at(head, .index)), .pooled = Carve.page(driver, paged.at(head, .pooled)), .pg = paged.pg(), .streams = undefined, .out = undefined, .logits = undefined, .dlogits = undefined, .best = undefined, .saved = undefined, .host = undefined, .layout = .of(pool_heads) }; d.streams = Carve.take(driver, &at, d.rows_cap * wide * 2); d.out = Carve.take(driver, &at, wide * 2); d.logits = Carve.take(driver, &at, max_rows * vocab * 2); d.dlogits = Carve.take(driver, &at, draft_n * 2); d.best = Carve.take(driver, &at, max_rows * 16); - d.saved = Carve.take(driver, &at, saved_bytes); - for (&d.spare) |*b| b.* = Carve.take(driver, &at, state_n * 4); + d.saved = Carve.take(driver, &at, d.layout.saved); + for (&d.spare) |*b| b.* = Carve.take(driver, &at, d.layout.state); std.debug.assert(at - base <= len); d.host = try gpa.alloc(u16, @max(max_rows * vocab, draft_n)); d.point(); @@ -220,3 +242,37 @@ pub const Draft = struct { d.rows = 0; } }; + +test "a draft at 65,536 tokens and 1,536-row chunks: the head's cache whole, saved rows and spares at the model's heads" { + const cap = 65536; + const rows = 1536; + const draft_n = 32768; + // the head's cache: two KV heads of 256, 128 indexer keys a position, a pooled row a block of 4 + const head: usize = cap * (2 * 2 * 256 * 2 + 128 * 2) + cap / 4 * 128 * 2; + // streams rows, the head's output row, a window's scores, the draft head's, best values + const buffers: usize = rows * 4 * 2560 * 2 + 4 * 2560 * 2 + 32 * 248320 * 2 + draft_n * 2 + 32 * 16; + // saved: 36 conv windows, 36 layers' 32 projection rows, the n-gram tail and 32 n-gram rows; 36 spare states + const tail: usize = 9 * 4 * 2560 * 2 + 32 * 4 * 2560 * 2; + const full: usize = 36 * 3 * 10240 * 2 + 36 * 32 * 16480 * 2 + tail + 36 * 48 * 128 * 128 * 4; + const half: usize = 36 * 3 * 5120 * 2 + 36 * 32 * 8240 * 2 + tail + 36 * 24 * 128 * 128 * 4; + try std.testing.expectEqual(head + buffers + full, Draft.bytes(cap, draft_n, rows, .full)); + try std.testing.expectEqual(head + buffers + half, Draft.bytes(cap, draft_n, rows, .half)); + try std.testing.expectEqual(@as(usize, 356_893_184), Draft.bytes(cap, draft_n, rows, .full)); + try std.testing.expectEqual(@as(usize, 280_179_200), Draft.bytes(cap, draft_n, rows, .half)); + // with the stream's own caches at half the heads + try std.testing.expectEqual(@as(usize, 1_395_057_152), sizes.seqBytes(cap, .half) + Draft.bytes(cap, draft_n, rows, .half)); +} + +test "the saved pieces follow one another at the model's heads, and the spares hold its states" { + for ([_]Heads{ .full, .half }) |h| { + const l = Draft.Layout.of(h); + try std.testing.expectEqual(@as(usize, 36 * 3 * h.convDim() * 2), l.proj_at); + try std.testing.expectEqual(l.proj_at + 36 * 32 * h.projWidth() * 2, l.tail_at); + try std.testing.expectEqual(l.tail_at + 9 * 4 * 2560 * 2, l.nrow_at); + try std.testing.expectEqual(l.nrow_at + 32 * 4 * 2560 * 2, l.saved); + try std.testing.expectEqual(sizes.stateBytes(h), l.state); + try std.testing.expect(l.heads.eql(h)); + } + try std.testing.expectEqual(@as(usize, 41_021_440), Draft.Layout.of(.full).saved); + try std.testing.expectEqual(@as(usize, 20_930_560), Draft.Layout.of(.half).saved); +} diff --git a/zig/src/families/flashnext/cuda_mtp_ops.zig b/zig/src/families/flashnext/cuda_mtp_ops.zig index 993e5328ce..9ac82e4d80 100644 --- a/zig/src/families/flashnext/cuda_mtp_ops.zig +++ b/zig/src/families/flashnext/cuda_mtp_ops.zig @@ -34,7 +34,6 @@ const Q4 = wt.Q4; const HcW = wt.HcW; const max_chain = mtp.max_chain; const max_rows = mtp.max_rows; -const conv_bytes = mtp.conv_bytes; const proj_bytes = mtp.proj_bytes; const fc_sk = mtp.fc_sk; const piece_rows = mtp.piece_rows; @@ -123,7 +122,7 @@ pub fn Ops(comptime Tri: type) type { const h = o.h; if (next.len == 0) return; if (d.len + next.len > d.cap) return error.PromptTooLong; - const want = @max(@as(usize, 32), 2 * @min(next.len, piece_rows)); + const want = plan.reserveRows(next.len); const s = try e.reserve(want); const piece = @max(@as(usize, 1), @min(piece_rows, s.cap / 2)); // the engine's block scores (the passes share one stream), for rows past 2,048 keys of the head's cache @@ -417,15 +416,16 @@ pub fn Ops(comptime Tri: type) type { const n = ids.len; if (n == 0 or n > max_rows) return error.WindowTooWide; if (seq.fed + n > seq.cap) return error.PromptTooLong; + if (!d.layout.heads.eql(seq.heads)) return error.DraftHeads; // its spares swap with the sequence's states d.settle(seq); d.fed0 = seq.fed; d.rows = n; d.history = seq.history; @memcpy(d.ids[0..n], ids); const keeps = n > 1; - if (keeps) try copyAsync(o.stream(), d.saved.ptr + Draft.tail_at, seq.tail.ptr, ple_tail_n * wide * 2); + if (keeps) try copyAsync(o.stream(), d.saved.ptr + d.layout.tail_at, seq.tail.ptr, ple_tail_n * wide * 2); try o.e.forward(Tri, o.tri, seq, ids, .{ .logits = .all, .logits_out = d.logits.ptr, .keep = if (keeps) &d.keep else null, .streams = d.streams.ptr }); - if (keeps) try copyAsync(o.stream(), d.saved.ptr + Draft.nrow_at, o.e.batch_s.?.ple_nrow, n * wide * 2); + if (keeps) try copyAsync(o.stream(), d.saved.ptr + d.layout.nrow_at, o.e.batch_s.?.ple_nrow, n * wide * 2); d.open = keeps; } @@ -433,9 +433,10 @@ pub fn Ops(comptime Tri: type) type { pub fn verifyMany(o: Self, seqs: []const *prompt.Seq, ds: []const *Draft, idss: []const []const u32, out: u64) !void { if (o.main != null) return error.UnsupportedRankOperation; // every window refused before any stream's state moves - for (seqs, idss) |seq, ids| { + for (seqs, idss, ds) |seq, ids, d| { if (ids.len == 0 or ids.len > max_rows) return error.WindowTooWide; if (seq.fed + ids.len > seq.cap) return error.PromptTooLong; + if (!d.layout.heads.eql(seq.heads)) return error.DraftHeads; } o.e.round_graph.clear(); const wins = try o.e.gpa.alloc(prompt.Engine.Window, seqs.len); @@ -448,14 +449,14 @@ pub fn Ops(comptime Tri: type) type { d.history = seq.history; @memcpy(d.ids[0..n], ids); const keeps = n > 1; - if (keeps) try copyAsync(o.stream(), d.saved.ptr + Draft.tail_at, seq.tail.ptr, ple_tail_n * wide * 2); + if (keeps) try copyAsync(o.stream(), d.saved.ptr + d.layout.tail_at, seq.tail.ptr, ple_tail_n * wide * 2); w.* = .{ .seq = seq, .ids = ids, .keep = if (keeps) &d.keep else null, .streams = d.streams.ptr, .logits = d.logits.ptr }; } try o.e.windows(Tri, o.tri, wins, out); var row0: usize = 0; for (ds, idss) |d, ids| { const keeps = ids.len > 1; - if (keeps) try copyAsync(o.stream(), d.saved.ptr + Draft.nrow_at, o.e.batch_s.?.ple_nrow + row0 * wide * 2, ids.len * wide * 2); + if (keeps) try copyAsync(o.stream(), d.saved.ptr + d.layout.nrow_at, o.e.batch_s.?.ple_nrow + row0 * wide * 2, ids.len * wide * 2); d.open = keeps; row0 += ids.len; } @@ -492,19 +493,27 @@ pub fn Ops(comptime Tri: type) type { const st = o.stream(); const s = if (e.batch_s) |*bs| bs else return error.CannotRewind; const k = &e.batch_k; + // refused before any change: kernels for every layer's heads, saved rows at the sequence's + if (!d.layout.heads.eql(seq.heads)) return error.DraftHeads; + for (e.layers[0..e.n_layers]) |*lw| if (lw.linear) { + _ = try k.gdnSet(lw.heads); + if (!lw.heads.eql(seq.heads)) return error.DraftHeads; + }; seq.fed = d.fed0 + keep; var hist = d.history; for (d.ids[0..keep]) |t| hist = .{ hist[1], t }; seq.history = hist; - try copyAsync(st, seq.tail.ptr, d.saved.ptr + Draft.tail_at, ple_tail_n * wide * 2); - try shift(k, st, seq.tail.ptr, d.saved.ptr + Draft.nrow_at, wide, keep, ple_tail_n, wide); + try copyAsync(st, seq.tail.ptr, d.saved.ptr + d.layout.tail_at, ple_tail_n * wide * 2); + try shift(k, st, seq.tail.ptr, d.saved.ptr + d.layout.nrow_at, wide, keep, ple_tail_n, wide); var li: usize = 0; for (e.layers[0..e.n_layers], 0..) |*lw, layer| { if (!lw.linear) continue; defer li += 1; // the sequence's state is still the window's start (the pass wrote its end state to the spare) const lk = d.keep.layers[layer]; - try copyAsync(st, seq.conv[li].ptr, lk.hist, conv_bytes); + const h = lw.heads; + const g = try k.gdnSet(h); + try copyAsync(st, seq.conv[li].ptr, lk.hist, (gdn.taps - 1) * h.convDim() * 2); var f: cuda.Args = .{}; f.add(lk.proj); f.add(seq.conv[li].ptr); @@ -515,7 +524,7 @@ pub fn Ops(comptime Tri: type) type { f.add(s.gk); f.add(s.gv); f.add(s.ggb); - try cuda.launch.launch(k.front, .{ .grid = .{ .x = @intCast(keep), .y = gdn.nk + gdn.nv }, .block = .{ .x = 256 } }, st, &f); + try cuda.launch.launch(g.front, .{ .grid = .{ .x = @intCast(keep), .y = h.nk + h.nv }, .block = .{ .x = 256 } }, st, &f); var a: cuda.Args = .{}; a.add(s.gq); a.add(s.gk); @@ -524,8 +533,8 @@ pub fn Ops(comptime Tri: type) type { a.add(seq.state[li].ptr); a.add(cint(keep)); a.add(s.ys); - try cuda.launch.launch(k.chain, .{ .grid = .{ .x = gdn.nv, .y = gdn.dv / 16 }, .block = .{ .x = 128 } }, st, &a); - try shift(k, st, seq.conv[li].ptr, lk.proj, gdn.proj_width, keep, gdn.taps - 1, gdn.conv_dim); + try cuda.launch.launch(g.chain, .{ .grid = .{ .x = h.nv, .y = gdn.dv / 16 }, .block = .{ .x = 128 } }, st, &a); + try shift(k, st, seq.conv[li].ptr, lk.proj, h.projWidth(), keep, gdn.taps - 1, h.convDim()); } d.rows = keep; d.open = false; diff --git a/zig/src/families/flashnext/cuda_mtp_plan.zig b/zig/src/families/flashnext/cuda_mtp_plan.zig index a6494ced2d..cdc98cffaf 100644 --- a/zig/src/families/flashnext/cuda_mtp_plan.zig +++ b/zig/src/families/flashnext/cuda_mtp_plan.zig @@ -13,11 +13,17 @@ pub const Plan = struct { piece: usize, }; +/// Two scratch rows per token, with the bounded count widened before multiplication. +pub fn reserveRows(total: usize) usize { + const rows: usize = @min(total, piece_rows); + return @max(@as(usize, 32), 2 * rows); +} + /// The head's plan from the admitted capacity and the rows to draft. pub fn plan(admitted: usize, total: usize) Plan { // the pass's rows as they stand: the head asks for at least 32, as the single-stream head does const cap = @max(admitted, 32); - const want = @max(@as(usize, 32), 2 * @min(total, piece_rows)); + const want = reserveRows(total); return .{ .pass_rows = @min(want, cap), .score_rows = @max(@as(usize, 1), @min(piece_rows, cap / 2)), diff --git a/zig/src/families/flashnext/cuda_mtp_plan_test.zig b/zig/src/families/flashnext/cuda_mtp_plan_test.zig index 08e89b149c..2d38ba950b 100644 --- a/zig/src/families/flashnext/cuda_mtp_plan_test.zig +++ b/zig/src/families/flashnext/cuda_mtp_plan_test.zig @@ -32,3 +32,18 @@ test "Flash Next head plan: one lane's rows never exceed the piece" { try std.testing.expectEqual(@as(usize, 1), plan.groups(&rows, p.piece, groups[0..])); try std.testing.expectEqual(@as(usize, 1), groups[0]); } + +test "Flash Next head scratch: reserve two rows per token through the 256-token boundary" { + try std.testing.expectEqual(@as(usize, 512), plan.reserveRows(256)); + try std.testing.expectEqual(@as(usize, 512), plan.reserveRows(16384)); + for ([_]usize{ 32, 64, 512, 1536, 16384 }) |admitted| { + for (1..16385) |total| { + const p = plan.plan(admitted, total); + const used: usize = @min(total, p.piece); + try std.testing.expect(p.pass_rows <= admitted); + try std.testing.expect(2 * used <= p.pass_rows); + const bounded: usize = if (total < 256) total else 256; + try std.testing.expectEqual(@max(@as(usize, 32), 2 * bounded), plan.reserveRows(total)); + } + } +} diff --git a/zig/src/families/flashnext/cuda_page_test_safety.zig b/zig/src/families/flashnext/cuda_page_test_safety.zig index 8fdef20995..20dafbb9a6 100644 --- a/zig/src/families/flashnext/cuda_page_test_safety.zig +++ b/zig/src/families/flashnext/cuda_page_test_safety.zig @@ -128,8 +128,8 @@ test "diagnostic pages: admission includes full final pages and the device table try std.testing.expectEqual(@as(usize, 2), l.pages); try std.testing.expectEqual(2 * pages.pageBytes(48), l.arena); try std.testing.expectEqual(@as(usize, 8), l.table); - const contiguous = wt.attn_n * (cap * (2 * prompt.kv_row_bytes + wt.index_dim * 2) + prompt.pooledBytes(cap)); - try std.testing.expectEqual(prompt.seqBytes(cap) - contiguous + l.total, try prompt.seqBytesPaged(cap, 48)); - try std.testing.expect(try prompt.seqBytesPaged(cap, 48) > prompt.seqBytes(cap)); - try std.testing.expectEqual(prompt.seqBytes(cap), try prompt.seqBytesPaged(cap, 0)); + const contiguous = wt.attn_n * (cap * (2 * wt.kv_heads * wt.head_dim * 2 + wt.index_dim * 2) + prompt.pooledBytes(cap)); + try std.testing.expectEqual(prompt.seqBytes(cap, .full) - contiguous + l.total, try prompt.seqBytesPaged(cap, 48, .full)); + try std.testing.expect(try prompt.seqBytesPaged(cap, 48, .full) > prompt.seqBytes(cap, .full)); + try std.testing.expectEqual(prompt.seqBytes(cap, .full), try prompt.seqBytesPaged(cap, 0, .full)); } diff --git a/zig/src/families/flashnext/cuda_paging_batch_test.zig b/zig/src/families/flashnext/cuda_paging_batch_test.zig index 356cf096ac..dc236fac97 100644 --- a/zig/src/families/flashnext/cuda_paging_batch_test.zig +++ b/zig/src/families/flashnext/cuda_paging_batch_test.zig @@ -45,7 +45,7 @@ const Fake = struct { pub fn fusedAttentionOwned(_: @This()) bool { return C.tri().fusedAttentionOwned(); } - pub fn attnPrep(_: @This(), _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, pg: flash.Pg) !void { + pub fn attnPrep(_: @This(), _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, pg: flash.Pg) !void { try std.testing.expectEqual(C.expected_pg, pg); } pub fn qsaPool(_: @This(), _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, _: anytype, pg: flash.Pg) !void { @@ -120,6 +120,7 @@ test "paged grouped prompts: every current attention route receives the sequence s.gxs_attn = 0x1200000000; s.part = 0x1300000000; var lw: Layer = undefined; + lw.heads = .full; lw.proj.native = null; lw.proj.w.ptr = 0x1000; lw.proj.s.ptr = 0x2000; @@ -132,6 +133,7 @@ test "paged grouped prompts: every current attention route receives the sequence lw.out.b.ptr = 0x6000; lw.out.k = 6144; lw.out.n = 7680; + lw.dense_cut = null; lw.q_scale.ptr = 0x7000; lw.k_scale.ptr = 0x8000; lw.i_scale.ptr = 0x9000; diff --git a/zig/src/families/flashnext/cuda_paging_launch_test.zig b/zig/src/families/flashnext/cuda_paging_launch_test.zig index f012392f6c..c9bc52bb90 100644 --- a/zig/src/families/flashnext/cuda_paging_launch_test.zig +++ b/zig/src/families/flashnext/cuda_paging_launch_test.zig @@ -51,12 +51,12 @@ test "paged launches: prep, pool, tiled QSA, ahead, one-buffer and fused select Capture.count = 0; Capture.symbols = @splat(0); Capture.pages = @splat(.{}); - try g.attnPrep(stream, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1e-6, 1, pg); + try g.attnPrep(stream, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1e-6, 1, .full, pg); try g.qsaPool(stream, 1, 2, 3, 4, 5, 1e-6, 1, pg); try g.qsaScores(stream, 1, 2, 3, 4, 1024, 1, 600, pg); - try g.attnChunks(stream, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 16, 2, 5, true, pg); - try g.attnChunks(stream, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 17, 2, 5, true, pg); - try g.attnFused(stream, 1, 2, 3, 4, 5, 6, 7, 8, 300, 2, pg); + try g.attnChunks(stream, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 16, .full, 5, true, pg); + try g.attnChunks(stream, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 17, .full, 5, true, pg); + try g.attnFused(stream, 1, 2, 3, 4, 5, 6, 7, 8, 300, .full, pg); try std.testing.expectEqual(@as(usize, 6), Capture.count); for (Capture.symbols, Capture.pages, 0..) |id, actual, i| { try std.testing.expectEqual(2 * i + 1 + @intFromBool(pg.pt != 0), id); diff --git a/zig/src/families/flashnext/cuda_paging_snap_test.zig b/zig/src/families/flashnext/cuda_paging_snap_test.zig index 352ce074af..19ca305be7 100644 --- a/zig/src/families/flashnext/cuda_paging_snap_test.zig +++ b/zig/src/families/flashnext/cuda_paging_snap_test.zig @@ -91,7 +91,7 @@ test "paged snapshots: save and restore keep contiguous order across partial shu try std.testing.expectEqual(kept.buf.ptr + off, copy[0]); off += copy[2]; } - try std.testing.expectEqual(@as(u64, snap.bytes(150, false)), off); + try std.testing.expectEqual(@as(u64, snap.bytes(150, false, .full)), off); const p = f.paged(&d); var i: usize = 2 * wt.linear_n; for (std.enums.values(pages.Part)) |part| for (0..wt.attn_n) |layer| { @@ -187,7 +187,7 @@ test "paged snapshots: the pooled MTP head occupies the thirteenth layer with th try std.testing.expectEqual(kept.buf.ptr + off, copy[0]); off += copy[2]; } - try std.testing.expectEqual(@as(u64, snap.bytes(150, true)), off); + try std.testing.expectEqual(@as(u64, snap.bytes(150, true, .full)), off); Fake.n = 0; seq.fed = 0; try snap.restore(stream, &seq, &head, kept); diff --git a/zig/src/families/flashnext/cuda_ple.zig b/zig/src/families/flashnext/cuda_ple.zig index e4ba3ca438..42f2cef6d3 100644 --- a/zig/src/families/flashnext/cuda_ple.zig +++ b/zig/src/families/flashnext/cuda_ple.zig @@ -7,6 +7,7 @@ const embed = @import("cuda_embed.zig"); const hc = @import("cuda_hc.zig"); const ple_pool = @import("cuda_ple_pool.zig"); const ple_reader = @import("cuda_ple_reader.zig"); +const ple_fp8 = @import("cuda_ple_fp8.zig"); const golden: u64 = 0x9E3779B97F4A7C15; const mix1: u64 = 0xBF58476D1CE4E5B9; @@ -208,7 +209,7 @@ pub fn partName(buf: []u8, shard: usize, nested: bool, part: []const u8) ![]cons return partNameFor(buf, shard, nested, part, false); } -fn partNameFor(buf: []u8, shard: usize, nested: bool, part: []const u8, modelopt: bool) ![]const u8 { +pub fn partNameFor(buf: []u8, shard: usize, nested: bool, part: []const u8, modelopt: bool) ![]const u8 { const prefix = if (modelopt) "model.language_model" else "language_model.model"; return if (nested) std.fmt.bufPrint(buf, "{s}.layers.{d}.ple.ple_embedding.ngram_embedding.shards.{d}.{s}", .{ prefix, layer, shard, part }) @@ -334,6 +335,9 @@ pub fn eosOf(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8) !i64 { return cfg.eos; } +/// A table's row format: MLX affine words, NVFP4 codes with E4M3 group scales, or FP8 rows with one table scale. +pub const Format = enum { mlx, nvfp4, fp8 }; + /// The n-gram table held open for an engine's life: the hash, the scale, and each shard's slices of its mapped file. pub const Ngram = struct { gpa: std.mem.Allocator, @@ -343,8 +347,9 @@ pub const Ngram = struct { w: [max_shards][]const u8 = undefined, s: [max_shards][]const u8 = undefined, b: [max_shards][]const u8 = undefined, - format: enum { mlx, nvfp4 } = .mlx, + format: Format = .mlx, tables: [max_shards]core.nvfp4_rows.Table = undefined, + tables8: [max_shards]ple_fp8.Table = undefined, readers: [max_shards]ple_reader.Reader = undefined, read_mode: ple_reader.Mode = .mmap, warm_ahead: bool = false, @@ -374,13 +379,14 @@ pub const Ngram = struct { const wm = try core.safetensors.index.weightMap(parsed.value); const first = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding.shard_0.weight"; const modelopt = wm.contains(first); - g.format = if (modelopt) .nvfp4 else .mlx; + g.format = if (modelopt) try ple_fp8.format(gpa, io, model_dir, wm, first) else .mlx; if (modelopt) g.read_mode = try ple_reader.mode(if (std.c.getenv("TF_FLASHNEXT_PLE_IO")) |v| std.mem.span(v) else null); g.warm_ahead = @import("cuda_ple_ahead.zig").enabled(if (std.c.getenv("TF_FLASHNEXT_WARM_AHEAD")) |v| std.mem.span(v) else null); - g.scale = if (modelopt) - try readScale(gpa, io, model_dir, wm, "model.language_model.layers.1.ple.ple_embedding.ngram_embedding.weight_scale_2", true) - else - try tableScale(gpa, io, model_dir, wm); + g.scale = switch (g.format) { + .mlx => try tableScale(gpa, io, model_dir, wm), + .nvfp4 => try readScale(gpa, io, model_dir, wm, "model.language_model.layers.1.ple.ple_embedding.ngram_embedding.weight_scale_2", true), + .fp8 => try readScale(gpa, io, model_dir, wm, "model.language_model.layers.1.ple.ple_embedding.ngram_embedding.weight_scale", true), + }; const dh = g.hash.dh; var key: [180]u8 = undefined; for (0..g.hash.shards) |shard| { @@ -400,15 +406,16 @@ pub const Ngram = struct { try g.files.append(gpa, o); break :file &g.files.items[g.files.items.len - 1]; }; - const w = try file.tensor(try partNameFor(&key, shard, nested, "weight", modelopt), if (modelopt) .u8 else .u32); + const dtype: core.safetensors.DType = switch (g.format) { + .mlx => .u32, + .nvfp4 => .u8, + .fp8 => .f8_e4m3, + }; + const w = try file.tensor(try partNameFor(&key, shard, nested, "weight", modelopt), dtype); const expected_rows: usize = @intCast(g.hash.starts[shard + 1] - g.hash.starts[shard]); if (w.rank != 2 or w.dim(0) != expected_rows) return error.UnexpectedTensor; if (modelopt) { - const sc = try file.tensor(try partNameFor(&key, shard, nested, "weight_scale", true), .f8_e4m3); - const table = try core.nvfp4_rows.Table.init(w, sc, g.scale); - if (table.width != dh) return error.UnexpectedTensor; - g.tables[shard] = table; - g.readers[shard] = ple_reader.Reader.init(&file.file, io, table); + try ple_fp8.openShard(&g, file, shard, nested, w); } else { const sc = try file.tensor(try partNameFor(&key, shard, nested, "scales", false), .bf16); const bias = try file.tensor(try partNameFor(&key, shard, nested, "biases", false), .bf16); @@ -442,8 +449,10 @@ pub const Ngram = struct { const row = out[i * dh ..][0..dh]; switch (g.format) { .mlx => try embed.dequant(g.w[loc.shard], g.s[loc.shard], g.b[loc.shard], dh, loc.local, row), - .nvfp4 => if (g.read_mode == .pread) - try g.readers[loc.shard].read(g.gpa, g.tables[loc.shard], loc.local, row) + .nvfp4, .fp8 => if (g.read_mode == .pread) + try g.readers[loc.shard].read(g.gpa, loc.local, row) + else if (g.format == .fp8) + try g.tables8[loc.shard].gather(&.{loc.local}, row) else try g.tables[loc.shard].gather(&.{loc.local}, row), } @@ -482,7 +491,7 @@ pub const Ngram = struct { hist[history.len - 1] = token; } const n_threads = @min(ple_pool.max_parts, items / 4); - if (g.format == .nvfp4 and items > 1) { + if (g.format != .mlx and items > 1) { try @import("cuda_ple_gather.zig").run(g, ids, out); } else if (n_threads <= 1 or g.pool == null) { try g.fill(ids, out, 0, items); @@ -540,5 +549,6 @@ test { test { _ = ple_reader; + _ = ple_fp8; _ = @import("cuda_ple_ahead.zig"); } diff --git a/zig/src/families/flashnext/cuda_ple_ahead.zig b/zig/src/families/flashnext/cuda_ple_ahead.zig index 90bfdb7007..eca7dda5ac 100644 --- a/zig/src/families/flashnext/cuda_ple_ahead.zig +++ b/zig/src/families/flashnext/cuda_ple_ahead.zig @@ -6,7 +6,7 @@ pub fn enabled(value: ?[]const u8) bool { } pub fn warm(g: anytype, tokens: []const u32, from: usize, to: usize) void { - if (!g.warm_ahead or g.format != .nvfp4 or from >= to or to > tokens.len or from < g.hash.n - 1) return; + if (!g.warm_ahead or (g.format != .nvfp4 and g.format != .fp8) or from >= to or to > tokens.len or from < g.hash.n - 1) return; var history: [ple.max_n]i64 = undefined; const hist = history[0 .. g.hash.n - 1]; for (tokens[from - hist.len .. from], hist) |token, *h| h.* = token; @@ -15,7 +15,7 @@ pub fn warm(g: anytype, tokens: []const u32, from: usize, to: usize) void { ple.lastIds(g.hash, hist, token, ids[0..g.hash.heads]) catch return; for (ids[0..g.hash.heads]) |id| { const at = g.hash.locate(id) catch return; - g.readers[at.shard].hint(g.tables[at.shard], at.local); + g.readers[at.shard].hint(at.local); } for (0..hist.len - 1) |i| hist[i] = hist[i + 1]; hist[hist.len - 1] = token; @@ -27,13 +27,13 @@ test "PLE read-ahead: future hashes equal a full token history and flags leave n var ids: [128]i64 = undefined; var count: usize = 0; start: i64, - pub fn hint(r: @This(), _: u8, row: usize) void { + pub fn hint(r: @This(), row: usize) void { ids[count] = r.start + @as(i64, @intCast(row)); count += 1; } }; const hash = try ple.Hash.init(.{ .vocab = 100, .ngram = 3, .per = 2, .base = 13, .divisor = 2, .shards = 2, .seed = 1234, .eos = 99, .ple_dim = 64, .ple_index = 0 }); - var g: struct { hash: ple.Hash, warm_ahead: bool = false, format: enum { nvfp4, mlx } = .nvfp4, readers: [2]Recorder, tables: [2]u8 = .{ 0, 0 } } = .{ .hash = hash, .readers = .{ .{ .start = hash.starts[0] }, .{ .start = hash.starts[1] } } }; + var g: struct { hash: ple.Hash, warm_ahead: bool = false, format: enum { nvfp4, mlx, fp8 } = .nvfp4, readers: [2]Recorder } = .{ .hash = hash, .readers = .{ .{ .start = hash.starts[0] }, .{ .start = hash.starts[1] } } }; const tokens = [_]u32{ 7, 3, 9, 22, 4, 7, 9, 3 }; const defaults: ple.Ngram = .{ .gpa = std.testing.allocator, .hash = hash, .scale = 1 }; try std.testing.expect(!defaults.warm_ahead); @@ -68,4 +68,9 @@ test "PLE read-ahead: future hashes equal a full token history and flags leave n warm(&g, &tokens, 4, tokens.len + 1); warm(&g, &tokens, tokens.len, tokens.len); try std.testing.expectEqual(before, Recorder.count); + const nvfp4 = Recorder.ids; + Recorder.count = 0; + g.format = .fp8; + warm(&g, &tokens, 4, tokens.len); + try std.testing.expectEqualSlices(i64, nvfp4[0..before], Recorder.ids[0..Recorder.count]); } diff --git a/zig/src/families/flashnext/cuda_ple_fp8.zig b/zig/src/families/flashnext/cuda_ple_fp8.zig new file mode 100644 index 0000000000..e88e458790 --- /dev/null +++ b/zig/src/families/flashnext/cuda_ple_fp8.zig @@ -0,0 +1,83 @@ +//! Borrowed FP8 n-gram rows (E4M3 codes times one table scale, rounded once to BF16), and a ModelOpt shard's open. +const std = @import("std"); +const core = @import("core"); +const ple = @import("cuda_ple.zig"); +const ple_reader = @import("cuda_ple_reader.zig"); +const quant = core.quant; +const Tensor = core.safetensors.Tensor; + +pub const Error = error{ BadFp8Table, BadFp8Rows, AliasedFp8Rows }; + +pub const Table = struct { + codes: []const u8, + rows: usize, + width: usize, + scale: f32, + + /// Borrow exact row-major E4M3 codes and the table's scale; the caller retains the map. + pub fn init(weight: Tensor, scale: f32) Error!Table { + if (weight.rank != 2 or weight.dtype != .f8_e4m3 or weight.dim(0) == 0 or weight.dim(1) == 0) + return error.BadFp8Table; + const size = std.math.mul(usize, weight.dim(0), weight.dim(1)) catch return error.BadFp8Table; + if (weight.bytes.len != size or !std.math.isFinite(scale)) return error.BadFp8Table; + return .{ .codes = weight.bytes, .rows = weight.dim(0), .width = weight.dim(1), .scale = scale }; + } + + /// Token-major requested rows, with FP32 products rounded once to BF16; malformed requests write nothing. + pub fn gather(t: Table, ids: []const usize, out: []u16) Error!void { + const count = std.math.mul(usize, ids.len, t.width) catch return error.BadFp8Rows; + if (out.len != count) return error.BadFp8Rows; + for (ids) |id| if (id >= t.rows) return error.BadFp8Rows; + const bytes = std.mem.sliceAsBytes(out); + if (overlaps(bytes, t.codes) or overlaps(bytes, std.mem.sliceAsBytes(ids))) return error.AliasedFp8Rows; + for (ids, 0..) |id, row| { + for (t.codes[id * t.width ..][0..t.width], out[row * t.width ..][0..t.width]) |code, *o| { + const v: f32 = @floatCast(quant.fp8.e4m3(code)); + o.* = quant.convert.bf16(v * t.scale); + } + } + } +}; + +fn overlaps(a: []const u8, b: []const u8) bool { + if (a.len == 0 or b.len == 0) return false; + const x = @intFromPtr(a.ptr); + const y = @intFromPtr(b.ptr); + return if (x <= y) y - x < a.len else x - y < b.len; +} + +/// A ModelOpt table's row format, from shard 0's weight dtype (U8: NVFP4, F8_E4M3: FP8). +pub fn format(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, wm: std.json.ObjectMap, name: []const u8) !ple.Format { + const file = wm.get(name) orelse return error.MissingTensor; + const path = try std.fs.path.join(gpa, &.{ dir, file.string }); + defer gpa.free(path); + var opened = try ple.Open.load(gpa, io, path); + defer opened.close(gpa); + const t = opened.file.get(name) orelse return error.MissingTensor; + return switch (t.dtype) { + .u8 => .nvfp4, + .f8_e4m3 => .fp8, + else => error.UnexpectedTensor, + }; +} + +/// One shard's ModelOpt table (NVFP4 with E4M3 scales a 16 values, or FP8 rows) and its positional reader. +pub fn openShard(g: *ple.Ngram, file: *const ple.Open, shard: usize, nested: bool, w: Tensor) !void { + if (g.format == .fp8) { + const table = try Table.init(w, g.scale); + if (table.width != g.hash.dh) return error.UnexpectedTensor; + g.tables8[shard] = table; + g.readers[shard] = ple_reader.Reader.initFp8(&file.file, file.io, table); + return; + } + var key: [180]u8 = undefined; + const sc = try file.tensor(try ple.partNameFor(&key, shard, nested, "weight_scale", true), .f8_e4m3); + const table = try core.nvfp4_rows.Table.init(w, sc, g.scale); + if (table.width != g.hash.dh) return error.UnexpectedTensor; + g.tables[shard] = table; + g.readers[shard] = ple_reader.Reader.init(&file.file, file.io, table); +} + +test { + _ = @import("cuda_ple_fp8_test.zig"); +} diff --git a/zig/src/families/flashnext/cuda_ple_fp8_test.zig b/zig/src/families/flashnext/cuda_ple_fp8_test.zig new file mode 100644 index 0000000000..c6e5265dd2 --- /dev/null +++ b/zig/src/families/flashnext/cuda_ple_fp8_test.zig @@ -0,0 +1,171 @@ +//! FP8 n-gram rows: every E4M3 code against an f64 reference, refusals, and a synthetic FP8 checkpoint through Ngram. +const std = @import("std"); +const core = @import("core"); +const ple = @import("cuda_ple.zig"); +const Table = @import("cuda_ple_fp8.zig").Table; +const a = std.testing.allocator; +const io = std.testing.io; +const base = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding"; + +fn tensor(dtype: core.safetensors.DType, rows: usize, cols: usize, bytes: []const u8) core.safetensors.Tensor { + return .{ .dtype = dtype, .rank = 2, .shape = .{ rows, cols, 1, 1, 1 }, .bytes = bytes }; +} + +fn e4m3(byte: u8) f64 { + const mag = byte & 127; + if (mag == 127) return std.math.nan(f64); + const e: i32 = @as(i32, mag >> 3) - 7; + const value = if (mag < 8) @as(f64, @floatFromInt(mag)) / 512 else std.math.ldexp(1 + @as(f64, @floatFromInt(mag & 7)) / 8, e); + return if (byte & 128 != 0) -value else value; +} + +fn rounded(value: f64) u16 { + const v: f32 = @floatCast(value); + const bits: u32 = @bitCast(v); + const high = bits >> 16; + const low = bits & 65535; + return @intCast(high + @intFromBool(low > 32768 or (low == 32768 and high % 2 == 1))); +} + +fn expectValue(code: u8, scale: f32, actual: u16) !void { + if (code & 127 == 127) { + try std.testing.expect(actual & 0x7fc0 == 0x7fc0); + } else try std.testing.expectEqual(rounded(e4m3(code) * scale), actual); +} + +test "FP8 rows: every E4M3 code times a BF16 scale rounds once to BF16 in reordered repeated gathers" { + var codes: [256]u8 = undefined; + for (&codes, 0..) |*c, i| c.* = @intCast(i); + for ([_]u32{ 0xBE53, 0x3B8D, 0x4241 }) |bits| { + const scale: f32 = @bitCast(bits << 16); + const t = try Table.init(tensor(.f8_e4m3, 16, 16, &codes), scale); + var ids: [18]usize = undefined; + for (0..16) |i| ids[i] = 15 - i; + ids[16] = 3; + ids[17] = 3; + var out: [18 * 16]u16 = undefined; + try t.gather(&ids, &out); + for (ids, 0..) |id, row| for (0..16) |col| try expectValue(codes[id * 16 + col], scale, out[row * 16 + col]); + } +} + +test "FP8 rows: a bad row or destination refuses before writing and maps are never aliased" { + var storage: [64]u8 align(8) = @splat(0x38); + const t = try Table.init(tensor(.f8_e4m3, 2, 8, storage[0..16]), 0.5); + var out: [16]u16 = @splat(0xbeef); + try std.testing.expectError(error.BadFp8Rows, t.gather(&.{ 0, 2 }, &out)); + for (out) |v| try std.testing.expectEqual(@as(u16, 0xbeef), v); + try std.testing.expectError(error.BadFp8Rows, t.gather(&.{0}, &out)); + try std.testing.expectError(error.AliasedFp8Rows, t.gather(&.{0}, std.mem.bytesAsSlice(u16, storage[0..16]))); + var ids: [2]usize = .{ 0, 0 }; + try std.testing.expectError(error.AliasedFp8Rows, t.gather(ids[0..1], std.mem.bytesAsSlice(u16, std.mem.sliceAsBytes(&ids)))); + try t.gather(&.{}, &.{}); +} + +test "FP8 rows: invalid dtype shape extent and scale cannot create a borrowed view" { + const codes: [16]u8 = @splat(0); + try std.testing.expectError(error.BadFp8Table, Table.init(tensor(.u8, 2, 8, &codes), 1)); + try std.testing.expectError(error.BadFp8Table, Table.init(tensor(.f8_e4m3, 2, 7, &codes), 1)); + try std.testing.expectError(error.BadFp8Table, Table.init(tensor(.f8_e4m3, 0, 16, codes[0..0]), 1)); + try std.testing.expectError(error.BadFp8Table, Table.init(tensor(.f8_e4m3, 2, 8, &codes), std.math.nan(f32))); + try std.testing.expectError(error.BadFp8Table, Table.init(tensor(.f8_e4m3, std.math.maxInt(usize), 8, &codes), 1)); +} + +const scale_bits: u16 = 0xBE53; + +fn codeAt(id: usize, col: usize) u8 { + return @intCast((id * 7 + col * 13) % 256); +} + +fn fixture(dir: std.Io.Dir, include_scale: bool) !void { + const config = + \\{"text_config":{"num_hidden_layers":2,"hidden_size":32,"vocab_size":100,"ple_layer_ids":[2], + \\ "ngram_size":2,"heads_per_ngram":2,"ngram_vocab_size_base":13,"make_ngram_vocab_size_divisible_by":2, + \\ "split_ngram_parts":2,"ple_embed_dim":32,"eos_token_id":99}} + ; + try dir.writeFile(io, .{ .sub_path = "config.json", .data = config }); + const h = try ple.Hash.init(.{ .vocab = 100, .ngram = 2, .per = 2, .base = 13, .divisor = 2, .shards = 2, .seed = 1234, .eos = 99, .ple_dim = 32, .ple_index = 0 }); + var header: std.Io.Writer.Allocating = .init(a); + defer header.deinit(); + var index: std.Io.Writer.Allocating = .init(a); + defer index.deinit(); + var data: std.ArrayList(u8) = .empty; + defer data.deinit(a); + try header.writer.writeByte('{'); + try index.writer.writeAll("{\"weight_map\":{"); + for (0..2) |shard| { + const rows: usize = @intCast(h.starts[shard + 1] - h.starts[shard]); + const begin = data.items.len; + for (0..rows) |row| { + const id: usize = @intCast(h.starts[shard] + @as(i64, @intCast(row))); + for (0..16) |col| try data.append(a, codeAt(id, col)); + } + if (shard > 0) { + try header.writer.writeByte(','); + try index.writer.writeByte(','); + } + try header.writer.print("\"{s}.shard_{d}.weight\":{{\"dtype\":\"F8_E4M3\",\"shape\":[{d},16],\"data_offsets\":[{d},{d}]}}", .{ base, shard, rows, begin, data.items.len }); + try index.writer.print("\"{s}.shard_{d}.weight\":\"table.safetensors\"", .{ base, shard }); + } + if (include_scale) { + const begin = data.items.len; + try data.appendSlice(a, &std.mem.toBytes(std.mem.nativeToLittle(u16, scale_bits))); + try header.writer.print(",\"{s}.weight_scale\":{{\"dtype\":\"BF16\",\"shape\":[1],\"data_offsets\":[{d},{d}]}}", .{ base, begin, data.items.len }); + try index.writer.print(",\"{s}.weight_scale\":\"table.safetensors\"", .{base}); + } + try header.writer.writeByte('}'); + try index.writer.writeAll("}}"); + var image: std.Io.Writer.Allocating = .init(a); + defer image.deinit(); + try image.writer.writeAll(&std.mem.toBytes(std.mem.nativeToLittle(u64, header.written().len))); + try image.writer.writeAll(header.written()); + try image.writer.writeAll(data.items); + try dir.writeFile(io, .{ .sub_path = "table.safetensors", .data = image.written() }); + try dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = index.written() }); +} + +test "FP8 PLE rows: mapped and positional reads equal the reference dequantization" { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try fixture(tmp.dir, true); + const path = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + defer a.free(path); + var g = try ple.Ngram.init(a, io, path); + defer g.deinit(); + try std.testing.expectEqual(ple.Format.fp8, g.format); + const scale: f32 = @bitCast(@as(u32, scale_bits) << 16); + try std.testing.expectEqual(scale, g.scale); + const tokens = [_]u32{ 7, 7, 12, 7, 3, 9, 12, 7, 1, 44, 63, 98 }; + var mapped: [tokens.len * 32]u16 = undefined; + var positional: [tokens.len * 32]u16 = undefined; + var mapped_xs: [tokens.len]f32 = undefined; + var positional_xs: [tokens.len]f32 = undefined; + var first = [_]i64{99}; + var second = first; + g.read_mode = .mmap; + try g.rows(&first, &tokens, &mapped, &mapped_xs); + g.read_mode = .pread; + try g.rows(&second, tokens[0..5], positional[0 .. 5 * 32], positional_xs[0..5]); + try g.rows(&second, tokens[5..], positional[5 * 32 ..], positional_xs[5..]); + try std.testing.expectEqualSlices(u16, &mapped, &positional); + try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(&mapped_xs), std.mem.sliceAsBytes(&positional_xs)); + try std.testing.expectEqualSlices(i64, &first, &second); + var hist = [_]i64{99}; + for (tokens, 0..) |token, t| { + var ids: [2]i64 = undefined; + try ple.lastIds(g.hash, &hist, token, &ids); + hist = .{token}; + for (ids, 0..) |id, head| for (0..16) |col| { + try expectValue(codeAt(@intCast(id), col), scale, mapped[t * 32 + head * 16 + col]); + }; + } +} + +test "FP8 PLE requires its table scale" { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try fixture(tmp.dir, false); + const path = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + defer a.free(path); + try std.testing.expectError(error.MissingTensor, ple.Ngram.init(a, io, path)); +} diff --git a/zig/src/families/flashnext/cuda_ple_reader.zig b/zig/src/families/flashnext/cuda_ple_reader.zig index ec71436a61..1746447c0b 100644 --- a/zig/src/families/flashnext/cuda_ple_reader.zig +++ b/zig/src/families/flashnext/cuda_ple_reader.zig @@ -1,39 +1,77 @@ -//! Positional NVFP4 row reads borrow the mapped shard's open file and use the shared exact dequantizer. +//! Positional NVFP4 and FP8 row reads borrow the mapped shard's open file and use the shared exact dequantizers. const std = @import("std"); const core = @import("core"); +const fp8 = @import("cuda_ple_fp8.zig"); pub const Mode = enum { mmap, pread }; pub fn mode(text: ?[]const u8) !Mode { const value = text orelse return if (@import("builtin").os.tag == .linux) .pread else .mmap; return std.meta.stringToEnum(Mode, value) orelse error.BadPleIoMode; } +/// A row's bytes in the shard file. +pub const Span = struct { at: u64, len: u64 }; pub const Reader = struct { file: std.Io.File, io: std.Io, codes_at: u64, scales_at: u64, + table: union(enum) { nvfp4: core.nvfp4_rows.Table, fp8: fp8.Table }, pub fn init(file: *const core.safetensors.File, io: std.Io, table: core.nvfp4_rows.Table) Reader { const base = @intFromPtr(file.map.memory.ptr); - return .{ .file = file.file, .io = io, .codes_at = @intFromPtr(table.codes.ptr) - base, .scales_at = @intFromPtr(table.scales.ptr) - base }; + return .{ .file = file.file, .io = io, .codes_at = @intFromPtr(table.codes.ptr) - base, .scales_at = @intFromPtr(table.scales.ptr) - base, .table = .{ .nvfp4 = table } }; } - pub fn read(r: Reader, a: std.mem.Allocator, table: core.nvfp4_rows.Table, row: usize, out: []u16) !void { - if (row >= table.rows or out.len != table.width) return error.BadNvfp4Rows; - if (overlap(std.mem.sliceAsBytes(out), table.codes) or overlap(std.mem.sliceAsBytes(out), table.scales)) return error.AliasedNvfp4Rows; + /// An FP8 table has no scale bytes a row. + pub fn initFp8(file: *const core.safetensors.File, io: std.Io, table: fp8.Table) Reader { + return .{ .file = file.file, .io = io, .codes_at = @intFromPtr(table.codes.ptr) - @intFromPtr(file.map.memory.ptr), .scales_at = 0, .table = .{ .fp8 = table } }; + } + /// A row's codes, then its scales (empty for FP8). + pub fn spans(r: Reader, row: usize) [2]Span { + const codes: u64, const scales: u64 = switch (r.table) { + .nvfp4 => |t| .{ t.width / 2, t.width / 16 }, + .fp8 => |t| .{ t.width, 0 }, + }; + return .{ .{ .at = r.codes_at + row * codes, .len = codes }, .{ .at = r.scales_at + row * scales, .len = scales } }; + } + pub fn read(r: Reader, a: std.mem.Allocator, row: usize, out: []u16) !void { + const dest = std.mem.sliceAsBytes(out); + switch (r.table) { + .nvfp4 => |t| { + if (row >= t.rows or out.len != t.width) return error.BadNvfp4Rows; + if (overlap(dest, t.codes) or overlap(dest, t.scales)) return error.AliasedNvfp4Rows; + }, + .fp8 => |t| { + if (row >= t.rows or out.len != t.width) return error.BadFp8Rows; + if (overlap(dest, t.codes)) return error.AliasedFp8Rows; + }, + } + const at = r.spans(row); var stack: [512]u8 = undefined; - const size = table.width / 2 + table.width / 16; + const size: usize = @intCast(at[0].len + at[1].len); const bytes = if (size <= stack.len) stack[0..size] else try a.alloc(u8, size); defer if (size > stack.len) a.free(bytes); - const codes = bytes[0 .. table.width / 2]; + const codes = bytes[0..@intCast(at[0].len)]; const scales = bytes[codes.len..]; - if (try r.file.readPositionalAll(r.io, codes, r.codes_at + row * codes.len) != codes.len) return error.PleShortRead; - if (try r.file.readPositionalAll(r.io, scales, r.scales_at + row * scales.len) != scales.len) return error.PleShortRead; - const local: core.nvfp4_rows.Table = .{ .codes = codes, .scales = scales, .width = table.width, .rows = 1, .global = table.global }; - try local.gather(&.{0}, out); + if (try r.file.readPositionalAll(r.io, codes, at[0].at) != codes.len) return error.PleShortRead; + if (scales.len != 0 and try r.file.readPositionalAll(r.io, scales, at[1].at) != scales.len) return error.PleShortRead; + switch (r.table) { + .nvfp4 => |t| { + const local: core.nvfp4_rows.Table = .{ .codes = codes, .scales = scales, .width = t.width, .rows = 1, .global = t.global }; + try local.gather(&.{0}, out); + }, + .fp8 => |t| { + const local: fp8.Table = .{ .codes = codes, .rows = 1, .width = t.width, .scale = t.scale }; + try local.gather(&.{0}, out); + }, + } } - pub fn hint(r: Reader, table: core.nvfp4_rows.Table, row: usize) void { - if (@import("builtin").os.tag != .linux or row >= table.rows) return; + pub fn hint(r: Reader, row: usize) void { + const rows = switch (r.table) { + inline else => |t| t.rows, + }; + if (@import("builtin").os.tag != .linux or row >= rows) return; const linux = std.os.linux; - _ = linux.fadvise(r.file.handle, @intCast(r.codes_at + row * (table.width / 2)), @intCast(table.width / 2), linux.POSIX_FADV.WILLNEED); - _ = linux.fadvise(r.file.handle, @intCast(r.scales_at + row * (table.width / 16)), @intCast(table.width / 16), linux.POSIX_FADV.WILLNEED); + for (r.spans(row)) |s| { + if (s.len != 0) _ = linux.fadvise(r.file.handle, @intCast(s.at), @intCast(s.len), linux.POSIX_FADV.WILLNEED); + } } }; @@ -67,17 +105,50 @@ test "PLE positional reader: every code and scale encoding equals mapped dequant try tmp.dir.writeFile(t.io, .{ .sub_path = "table", .data = &data }); const file = try tmp.dir.openFile(t.io, "table", .{}); defer file.close(t.io); - const reader: Reader = .{ .file = file, .io = t.io, .codes_at = 7, .scales_at = 7 + codes.len }; const table: core.nvfp4_rows.Table = .{ .codes = &codes, .scales = &scales, .rows = 256, .width = 16, .global = 1.375 }; + const reader: Reader = .{ .file = file, .io = t.io, .codes_at = 7, .scales_at = 7 + codes.len, .table = .{ .nvfp4 = table } }; for (0..256) |row| { var want: [16]u16 = undefined; var got: [16]u16 = undefined; try table.gather(&.{row}, &want); - try reader.read(t.allocator, table, row, &got); + try reader.read(t.allocator, row, &got); + try t.expectEqualSlices(u16, &want, &got); + } + try t.expectError(error.AliasedNvfp4Rows, reader.read(t.allocator, 0, std.mem.bytesAsSlice(u16, codes[0..32]))); + var bad: [16]u16 = @splat(0xbeef); + try t.expectError(error.BadNvfp4Rows, reader.read(t.allocator, 256, &bad)); + for (bad) |v| try t.expectEqual(@as(u16, 0xbeef), v); +} + +test "PLE positional reader: every FP8 code equals mapped dequantization and hints its codes alone" { + const t = std.testing; + var tmp = t.tmpDir(.{}); + defer tmp.cleanup(); + var codes: [64 * 16]u8 align(2) = undefined; + for (&codes, 0..) |*c, i| c.* = @intCast(i % 256); + var data: [5 + codes.len]u8 = undefined; + @memset(data[0..5], 0x42); + @memcpy(data[5..], &codes); + try tmp.dir.writeFile(t.io, .{ .sub_path = "table", .data = &data }); + const file = try tmp.dir.openFile(t.io, "table", .{}); + defer file.close(t.io); + const table: fp8.Table = .{ .codes = &codes, .rows = 64, .width = 16, .scale = @bitCast(@as(u32, 0xBE53_0000)) }; + const reader: Reader = .{ .file = file, .io = t.io, .codes_at = 5, .scales_at = 0, .table = .{ .fp8 = table } }; + for (0..64) |row| { + var want: [16]u16 = undefined; + var got: [16]u16 = undefined; + try table.gather(&.{row}, &want); + try reader.read(t.allocator, row, &got); try t.expectEqualSlices(u16, &want, &got); } - try t.expectError(error.AliasedNvfp4Rows, reader.read(t.allocator, table, 0, std.mem.bytesAsSlice(u16, codes[0..32]))); + try t.expectError(error.AliasedFp8Rows, reader.read(t.allocator, 0, std.mem.bytesAsSlice(u16, codes[0..32]))); var bad: [16]u16 = @splat(0xbeef); - try t.expectError(error.BadNvfp4Rows, reader.read(t.allocator, table, 256, &bad)); + try t.expectError(error.BadFp8Rows, reader.read(t.allocator, 64, &bad)); + try t.expectError(error.BadFp8Rows, reader.read(t.allocator, 0, bad[0..15])); for (bad) |v| try t.expectEqual(@as(u16, 0xbeef), v); + const at = reader.spans(9); + try t.expectEqual(Span{ .at = 5 + 9 * 16, .len = 16 }, at[0]); + try t.expectEqual(@as(u64, 0), at[1].len); + const nvfp4: Reader = .{ .file = file, .io = t.io, .codes_at = 5, .scales_at = 900, .table = .{ .nvfp4 = .{ .codes = &codes, .scales = codes[0..64], .rows = 64, .width = 32, .global = 1 } } }; + try t.expectEqualSlices(Span, &.{ .{ .at = 5 + 9 * 16, .len = 16 }, .{ .at = 900 + 9 * 2, .len = 2 } }, &nvfp4.spans(9)); } diff --git a/zig/src/families/flashnext/cuda_projection_load.zig b/zig/src/families/flashnext/cuda_projection_load.zig index c0c37f1b79..90336a2d02 100644 --- a/zig/src/families/flashnext/cuda_projection_load.zig +++ b/zig/src/families/flashnext/cuda_projection_load.zig @@ -1,8 +1,10 @@ //! Projection loading preserves dense BF16 bits and MXFP8 codes/scales; affine layouts keep their existing packer. const std = @import("std"); const cuda = @import("cuda"); +const core = @import("core"); const wt = @import("cuda_prompt_wt.zig"); const projection = @import("cuda_projection.zig"); +const Cut = @import("cuda_modelopt_experts.zig").Cut; pub const Micro = struct { n: usize, @@ -127,6 +129,108 @@ pub fn loadStitched(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, pre return wt.uploadQ4(a, d, planes.w, planes.s, planes.b, n, k / 8); } +/// What a rank's cut of a projection keeps: output rows, or each row's input columns. +pub const Axis = enum { rows, cols }; + +/// A rank's own BF16 [cut, k] rows or [n, cut] input columns, copied before the store moves on. +pub fn cutBytes(a: std.mem.Allocator, t: core.safetensors.Tensor, cut: Cut, axis: Axis) ![]u8 { + if (t.dtype != .bf16 or t.rank != 2 or t.dim(0) == 0 or t.dim(1) == 0) return error.UnexpectedTensor; + const n = t.dim(0); + const k = t.dim(1); + if (t.bytes.len != try std.math.mul(usize, try std.math.mul(usize, n, k), 2)) return error.UnexpectedTensor; + if (cut.from >= cut.to or cut.to > (if (axis == .rows) n else k)) return error.UnexpectedTensor; + const w = cut.width(); + switch (axis) { + .rows => return a.dupe(u8, t.bytes[cut.from * k * 2 ..][0 .. w * k * 2]), + .cols => { + if (w % 64 != 0) return error.UnexpectedTensor; + const out = try a.alloc(u8, n * w * 2); + for (0..n) |row| @memcpy(out[row * w * 2 ..][0 .. w * 2], t.bytes[(row * k + cut.from) * 2 ..][0 .. w * 2]); + return out; + }, + } +} + +/// A BF16 [n, k] projection's cut on this rank, packed as `load` packs a whole one; other formats are refused. +pub fn loadCut(a: std.mem.Allocator, d: *const cuda.Driver, store: anytype, prefix: []const u8, n: usize, k: usize, cut: Cut, axis: Axis) !wt.Q4 { + var buf: [512]u8 = undefined; + const t = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.weight", .{prefix})); + if (t.dtype != .bf16) return error.UnsupportedQuantization; + if (!t.is(.bf16, &.{ n, k })) return error.UnexpectedTensor; + const bytes = try cutBytes(a, t, cut, axis); + defer a.free(bytes); + const rows = if (axis == .rows) cut.width() else n; + const cols = if (axis == .rows) k else cut.width(); + var value = try projection.dense(a, d, .{ .dtype = .bf16, .rank = 2, .shape = .{ rows, cols, 1, 1, 1 }, .bytes = bytes }); + if (axis == .rows) value.native.?.sk = cuda.qmmf.splitK(n, k); + return value; +} + +/// Up to four [from, to) output-row ranges kept of one stitched part, copied in order. +pub const Ranges = struct { + cuts: [4]Cut = undefined, + len: usize = 0, + + pub fn of(cuts: []const Cut) Ranges { + std.debug.assert(cuts.len <= 4); + var out: Ranges = .{ .len = cuts.len }; + @memcpy(out.cuts[0..cuts.len], cuts); + return out; + } + + pub fn items(self: *const Ranges) []const Cut { + return self.cuts[0..self.len]; + } + + pub fn rows(self: Ranges) usize { + var n: usize = 0; + for (self.items()) |c| n += c.width(); + return n; + } +}; + +/// Copy `kept` rows of a row-major `src` into the front of `dst`, in order; returns the bytes written. +pub fn copyRanges(dst: []u8, src: []const u8, row_bytes: usize, kept: Ranges) !usize { + var at: usize = 0; + for (kept.items()) |c| { + if (c.from >= c.to or try std.math.mul(usize, c.to, row_bytes) > src.len) return error.UnexpectedTensor; + const len = c.width() * row_bytes; + if (len > dst.len - at) return error.UnexpectedTensor; + @memcpy(dst[at..][0..len], src[c.from * row_bytes ..][0..len]); + at += len; + } + return at; +} + +/// One rank's BF16 [n, k] rows of a stitched projection: each part's ranges in order, copied before the store moves on. +pub fn stitchCut(a: std.mem.Allocator, store: anytype, prefix: []const u8, parts: []const []const u8, ranges: []const Ranges, n: usize, k: usize) ![]u8 { + if (parts.len == 0 or parts.len != ranges.len) return error.UnexpectedTensor; + const out = try a.alloc(u8, try std.math.mul(usize, try std.math.mul(usize, n, k), 2)); + errdefer a.free(out); + var buf: [512]u8 = undefined; + var at: usize = 0; + for (parts, ranges) |name, kept| { + const t = try store.tensor(try std.fmt.bufPrint(&buf, "{s}{s}.weight", .{ prefix, name })); + if (t.dtype != .bf16) return error.UnsupportedQuantization; + if (t.rank != 2 or t.dim(0) == 0 or t.dim(1) != k or t.bytes.len != try std.math.mul(usize, t.dim(0) * k, 2)) + return error.UnexpectedTensor; + at += try copyRanges(out[at..], t.bytes, k * 2, kept); + } + if (at != out.len) return error.UnexpectedTensor; + return out; +} + +/// `stitchCut` packed as `loadStitched` packs, its K slices pinned to `sk` (the whole stitched shape's) for equal sums. +pub fn loadStitchedCut(a: std.mem.Allocator, d: *const cuda.Driver, store: anytype, prefix: []const u8, parts: []const []const u8, ranges: []const Ranges, n: usize, k: usize, sk: u32) !wt.Q4 { + // the pin is checked on the shape before anything is read, so a bad one never reaches the driver + const view = try cuda.qmmf.pinned(cuda.modelopt.bf16Weight(0, n, k), sk); + const bytes = try stitchCut(a, store, prefix, parts, ranges, n, k); + defer a.free(bytes); + var p = try projection.dense(a, d, .{ .dtype = .bf16, .rank = 2, .shape = .{ n, k, 1, 1, 1 }, .bytes = bytes }); + p.native.?.sk = view.sk; + return p; +} + test { _ = @import("cuda_projection_load_test.zig"); } diff --git a/zig/src/families/flashnext/cuda_projection_load_test.zig b/zig/src/families/flashnext/cuda_projection_load_test.zig index a76e08c304..432b2bbca9 100644 --- a/zig/src/families/flashnext/cuda_projection_load_test.zig +++ b/zig/src/families/flashnext/cuda_projection_load_test.zig @@ -1,6 +1,8 @@ -//! A store that invalidates its previous view on every read, like advancing mapped shards. +//! MXFP8 stitching through a store that moves on every read, like mapped shards; BF16 cuts by rows or input columns. const std = @import("std"); +const cuda = @import("cuda"); const core = @import("core"); +const Cut = @import("cuda_modelopt_experts.zig").Cut; const source = @import("cuda_projection_load.zig"); const Tensor = core.safetensors.Tensor; @@ -61,3 +63,260 @@ test "MXFP8 source refuses incompatible shapes, dtypes, extents and missing scal try std.testing.expectError(error.UnexpectedTensor, source.readStitched(a, &fake, "root", &.{})); try std.testing.expectError(error.UnexpectedTensor, source.readStitched(a, &fake, "root", &.{ ".a", ".a", ".a", ".a", ".a", ".a", ".a" })); } + +const rows6 = 6; +const cols256 = 256; + +/// A [6, 256] BF16 tensor whose every element holds a distinct value. +fn dense(bytes: *[rows6 * cols256 * 2]u8) Tensor { + for (0..rows6 * cols256) |i| std.mem.writeInt(u16, bytes[i * 2 ..][0..2], @intCast(i + 1), .little); + return .{ .dtype = .bf16, .rank = 2, .shape = .{ rows6, cols256, 1, 1, 1 }, .bytes = bytes }; +} + +fn at(bytes: []const u8, i: usize) u16 { + return std.mem.readInt(u16, bytes[i * 2 ..][0..2], .little); +} + +test "BF16 row and column cuts equal the source, and two ranks cover every element once" { + const a = std.testing.allocator; + var bytes: [rows6 * cols256 * 2]u8 = undefined; + const t = dense(&bytes); + var seen: [rows6 * cols256]u8 = @splat(0); + for (0..2) |r| { + const cut = try Cut.init(cols256, 2, r); + const cols = try source.cutBytes(a, t, cut, .cols); + defer a.free(cols); + try std.testing.expectEqual(rows6 * 128 * 2, cols.len); + for (0..rows6) |row| for (0..128) |c| { + const v = at(cols, row * 128 + c); + try std.testing.expectEqual(at(t.bytes, row * cols256 + cut.from + c), v); + seen[v - 1] += 1; + }; + const half: Cut = .{ .from = 3 * r, .to = 3 * r + 3 }; + const rows = try source.cutBytes(a, t, half, .rows); + defer a.free(rows); + try std.testing.expectEqualSlices(u8, t.bytes[half.from * cols256 * 2 ..][0 .. 3 * cols256 * 2], rows); + for (0..rows.len / 2) |i| seen[at(rows, i) - 1] += 1; + } + for (seen) |count| try std.testing.expectEqual(@as(u8, 2), count); +} + +fn cutTrial(a: std.mem.Allocator) !void { + var bytes: [rows6 * cols256 * 2]u8 = undefined; + const t = dense(&bytes); + for ([_]source.Axis{ .rows, .cols }) |axis| { + const out = try source.cutBytes(a, t, .{ .from = 2, .to = if (axis == .rows) 4 else 66 }, axis); + a.free(out); + } +} + +test "BF16 cuts free every allocation on failure" { + try std.testing.checkAllAllocationFailures(std.testing.allocator, cutTrial, .{}); +} + +const Plain = struct { + t: Tensor, + + pub fn tensor(self: *Plain, _: []const u8) !Tensor { + return self.t; + } +}; + +test "BF16 cuts refuse other formats, a wrong extent and a column width not a multiple of 64" { + const a = std.testing.allocator; + var bytes: [rows6 * cols256 * 2]u8 = undefined; + const t = dense(&bytes); + const half: Cut = .{ .from = 128, .to = 256 }; + var other = t; + other.dtype = .f8_e4m3; + try std.testing.expectError(error.UnexpectedTensor, source.cutBytes(a, other, half, .cols)); + var short = t; + short.bytes = t.bytes[0 .. t.bytes.len - 2]; + try std.testing.expectError(error.UnexpectedTensor, source.cutBytes(a, short, half, .cols)); + try std.testing.expectError(error.UnexpectedTensor, source.cutBytes(a, t, .{ .from = 128, .to = 320 }, .cols)); + try std.testing.expectError(error.UnexpectedTensor, source.cutBytes(a, t, .{ .from = 3, .to = 7 }, .rows)); + try std.testing.expectError(error.UnexpectedTensor, source.cutBytes(a, t, .{ .from = 4, .to = 4 }, .rows)); + try std.testing.expectError(error.UnexpectedTensor, source.cutBytes(a, t, try Cut.init(cols256, 8, 1), .cols)); + // the refusals return before the driver is touched + const d: *const cuda.Driver = undefined; + var store: Plain = .{ .t = other }; + try std.testing.expectError(error.UnsupportedQuantization, source.loadCut(a, d, &store, "p", rows6, cols256, half, .cols)); + store.t = t; + try std.testing.expectError(error.UnexpectedTensor, source.loadCut(a, d, &store, "p", rows6, 512, half, .cols)); +} + +const heads_cut = @import("heads_cut.zig"); +const Heads = @import("cuda_heads.zig").Heads; +const Ranges = source.Ranges; +const qmmf = cuda.qmmf; +const k64 = 64; + +/// The whole checkpoint parts at 64 inputs, each read into one buffer that the next read overwrites. +const Whole = struct { + scratch: []u8, + dtype: @FieldType(Tensor, "dtype") = .bf16, + + const names = [_][]const u8{ "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a", "q_proj", "k_proj", "v_proj", "indexer.index_qk_proj" }; + const counts = [_]usize{ 10240, 6144, 48, 48, 12288, 512, 512, 640 }; + + fn part(name: []const u8) !usize { + if (!std.mem.startsWith(u8, name, "root.") or !std.mem.endsWith(u8, name, ".weight")) return error.MissingTensor; + const stem = name["root.".len .. name.len - ".weight".len]; + for (names, 0..) |n, i| if (std.mem.eql(u8, n, stem)) return i; + return error.MissingTensor; + } + + fn value(p: usize, i: usize) u16 { + const x: u32 = @truncate((i + p * 0x1000003) *% 0x9E3779B1); + return @truncate(x >> 16); + } + + pub fn tensor(self: *Whole, name: []const u8) !Tensor { + const p = try part(name); + const n = counts[p]; + const len = switch (self.dtype) { + .f8_e4m3 => n * k64, + .u32 => n * (k64 / 8) * 4, + else => n * k64 * 2, + }; + @memset(self.scratch, 0xee); + for (0..n * k64) |i| std.mem.writeInt(u16, self.scratch[i * 2 ..][0..2], value(p, i), .little); + return .{ .dtype = self.dtype, .rank = 2, .shape = .{ n, if (self.dtype == .u32) k64 / 8 else k64, 1, 1, 1 }, .bytes = self.scratch[0..len] }; + } +}; + +fn wholeStore(a: std.mem.Allocator) !Whole { + return .{ .scratch = try a.alloc(u8, 12288 * k64 * 2) }; +} + +test "the head split keeps each stitched part's ranged rows in order, and two ranks cover every row" { + const a = std.testing.allocator; + var fake = try wholeStore(a); + defer a.free(fake.scratch); + for ([_]bool{ true, false }) |linear| { + const parts = if (linear) heads_cut.delta_parts else heads_cut.attention_parts; + const first: usize = if (linear) 0 else 4; + var seen: [4][12288]u8 = @splat(@splat(0)); + for (0..2) |r| { + const s = if (linear) try heads_cut.delta(.half, r) else try heads_cut.attention(.half, r); + try std.testing.expectEqual(@as(usize, if (linear) 8240 else 7296), s.n); + const bytes = try source.stitchCut(a, &fake, "root", &parts, &s.ranges, s.n, k64); + defer a.free(bytes); + try std.testing.expectEqual(s.n * k64 * 2, bytes.len); + var pos: usize = 0; + for (parts, s.ranges, 0..) |name, kept, i| { + var name_buf: [64]u8 = undefined; + const whole = try a.dupe(u8, (try fake.tensor(try std.fmt.bufPrint(&name_buf, "root{s}.weight", .{name}))).bytes); + defer a.free(whole); + try std.testing.expectEqual(Whole.counts[first + i] * k64 * 2, whole.len); + for (kept.items()) |c| { + const len = c.width() * k64 * 2; + try std.testing.expectEqualSlices(u8, whole[c.from * k64 * 2 ..][0..len], bytes[pos..][0..len]); + for (c.from..c.to) |row| seen[i][row] += 1; + pos += len; + } + } + try std.testing.expectEqual(bytes.len, pos); + } + for (0..4) |i| for (seen[i][0..Whole.counts[first + i]]) |count| { + try std.testing.expectEqual(@as(u8, if (!linear and i == 3) 2 else 1), count); + }; + } +} + +fn expectRanges(want: []const [2]usize, got: Ranges) !void { + try std.testing.expectEqual(want.len, got.len); + for (want, got.items()) |w, c| { + try std.testing.expectEqual(w[0], c.from); + try std.testing.expectEqual(w[1], c.to); + } +} + +test "each rank's ranges are the design's: q, k, v, z, b, a, the attention heads, the indexer whole" { + for (0..2) |r| { + const d = try heads_cut.delta(.half, r); + const qkv = [_][2]usize{ .{ 1024 * r, 1024 * r + 1024 }, .{ 2048 + 1024 * r, 3072 + 1024 * r }, .{ 4096 + 3072 * r, 7168 + 3072 * r } }; + try expectRanges(&qkv, d.ranges[0]); + try expectRanges(&.{.{ 3072 * r, 3072 * r + 3072 }}, d.ranges[1]); + for (d.ranges[2..]) |kept| try expectRanges(&.{.{ 24 * r, 24 * r + 24 }}, kept); + try expectRanges(&qkv, try heads_cut.convChannels(.half, r)); + const t = try heads_cut.attention(.half, r); + try expectRanges(&.{.{ 6144 * r, 6144 * r + 6144 }}, t.ranges[0]); + for (t.ranges[1..3]) |kept| try expectRanges(&.{.{ 256 * r, 256 * r + 256 }}, kept); + try expectRanges(&.{.{ 0, 640 }}, t.ranges[3]); + } + const whole = try heads_cut.delta(.full, 0); + try std.testing.expectEqual(@as(usize, 16480), whole.n); + try expectRanges(&.{.{ 0, 6144 }}, whole.ranges[1]); + try std.testing.expectEqual(@as(usize, 13952), (try heads_cut.attention(.full, 0)).n); + try std.testing.expectError(error.UnevenHeadSplit, heads_cut.delta(.full, 1)); + try std.testing.expectError(error.UnevenHeadSplit, heads_cut.attention(.half, 2)); + try std.testing.expectError(error.UnevenHeadSplit, heads_cut.delta(.{ .nk = 8, .nv = 24, .q = 24, .kv = 2 }, 0)); + try std.testing.expectError(error.UnevenHeadSplit, heads_cut.valueHeads(.{ .nk = 0, .nv = 24, .q = 12, .kv = 1 }, 0)); + try std.testing.expectError(error.UnevenHeadSplit, heads_cut.valueHeads(.{ .nk = 8, .nv = 0x80000018, .q = 12, .kv = 1 }, 0)); +} + +test "a cut projection pins the whole stitched shape's K slices, so each kept row plans as the whole one does" { + for ([_]bool{ true, false }) |linear| { + const whole_n: usize = if (linear) 16480 else 13952; + for (0..2) |r| { + const s = if (linear) try heads_cut.delta(.half, r) else try heads_cut.attention(.half, r); + try std.testing.expectEqual(qmmf.splitK(whole_n, 2560), s.sk); + try std.testing.expect(qmmf.splitK(s.n, 2560) != s.sk); + const w = try qmmf.pinned(cuda.modelopt.bf16Weight(0, s.n, 2560), s.sk); + for ([_]bool{ true, false }) |cl| for ([_]usize{ 1, 16, 64, 1536 }) |m| { + try std.testing.expectEqual(qmmf.plan(m, whole_n, 2560, cl), qmmf.planFor(m, w, cl)); + }; + } + } +} + +test "conv1d keeps the rank's q, k and v channels, and A_log and dt_bias its 24 value heads" { + const a = std.testing.allocator; + const whole = try a.alloc(u8, 10240 * 4 * 2); + defer a.free(whole); + for (0..10240 * 4) |i| std.mem.writeInt(u16, whole[i * 2 ..][0..2], Whole.value(9, i), .little); + const t: Tensor = .{ .dtype = .bf16, .rank = 3, .shape = .{ 10240, 1, 4, 1, 1 }, .bytes = whole }; + var values: [48]f32 = undefined; + for (&values, 0..) |*v, i| v.* = @floatFromInt(i); + for (0..2) |r| { + const local = try heads_cut.convRows(a, t, .half, r); + defer a.free(local); + try std.testing.expectEqual(@as(usize, 5120 * 4 * 2), local.len); + for (0..5120) |row| { + const channel = if (row < 1024) 1024 * r + row else if (row < 2048) 2048 + 1024 * r + row - 1024 else 4096 + 3072 * r + row - 2048; + try std.testing.expectEqualSlices(u8, whole[channel * 8 ..][0..8], local[row * 8 ..][0..8]); + } + const kept = try heads_cut.valueSlice(&values, .half, r); + try std.testing.expectEqual(@as(usize, 24), kept.len); + for (kept, 0..) |v, i| try std.testing.expectEqual(@as(f32, @floatFromInt(24 * r + i)), v); + } + var short = t; + short.bytes = whole[0 .. whole.len - 2]; + try std.testing.expectError(error.UnexpectedTensor, heads_cut.convRows(a, short, .half, 0)); + var other = t; + other.dtype = .f32; + try std.testing.expectError(error.UnexpectedTensor, heads_cut.convRows(a, other, .half, 0)); + try std.testing.expectError(error.UnexpectedTensor, heads_cut.valueSlice(values[0..24], .half, 0)); +} + +test "stitched cuts refuse MXFP8 and affine parts, other widths and a wrong local row count" { + const a = std.testing.allocator; + var fake = try wholeStore(a); + defer a.free(fake.scratch); + const s = try heads_cut.delta(.half, 0); + // the refusals return before the driver is touched + const d: *const cuda.Driver = undefined; + for ([_]@FieldType(Tensor, "dtype"){ .f8_e4m3, .u32 }) |dtype| { + fake.dtype = dtype; + try std.testing.expectError(error.UnsupportedQuantization, source.loadStitchedCut(a, d, &fake, "root", &heads_cut.delta_parts, &s.ranges, s.n, k64, s.sk)); + } + fake.dtype = .bf16; + try std.testing.expectError(error.UnexpectedTensor, source.loadStitchedCut(a, d, &fake, "root", &heads_cut.delta_parts, &s.ranges, s.n, 128, s.sk)); + try std.testing.expectError(error.UnexpectedTensor, source.loadStitchedCut(a, d, &fake, "root", &heads_cut.delta_parts, &s.ranges, s.n + 1, k64, s.sk)); + try std.testing.expectError(error.UnexpectedTensor, source.loadStitchedCut(a, d, &fake, "root", heads_cut.delta_parts[0..3], &s.ranges, s.n, k64, s.sk)); + try std.testing.expectError(error.BadSplit, source.loadStitchedCut(a, d, &fake, "root", &heads_cut.delta_parts, &s.ranges, s.n, k64, 3)); + var past = s.ranges; + past[2] = Ranges.of(&.{.{ .from = 40, .to = 64 }}); + try std.testing.expectError(error.UnexpectedTensor, source.stitchCut(a, &fake, "root", &heads_cut.delta_parts, &past, s.n, k64)); +} diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index fc70793092..248e6fe8bf 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -2,7 +2,6 @@ const std = @import("std"); const cuda = @import("cuda"); const embed = @import("cuda_embed.zig"); -const gdn = @import("cuda_gdn.zig"); const ple = @import("cuda_ple.zig"); const weights = @import("cuda_prompt_wt.zig"); const load = @import("cuda_prompt_load.zig"); @@ -20,16 +19,12 @@ 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; @@ -95,26 +90,12 @@ fn made(comptime B: type, gpa: std.mem.Allocator, n: usize, src: anytype) ![]B { pub const vocabulary: usize = vocab; /// Prompt plus reply tokens a sequence can hold: the model's window. pub const max_tokens: usize = weights.window; -/// One position's keys and values in one attention layer, bf16. -pub const kv_row_bytes: usize = kv_heads * head_dim * 2; -/// Device bytes `Engine.newSeq(cap)` allocates: conv tails, recurrent states, KV with indexer keys, the n-gram tail. -pub fn seqBytes(cap: usize) usize { - const linear = linear_n * ((gdn.taps - 1) * gdn.conv_dim * 2 + state_n * 4); - const attn = attn_n * (cap * (2 * kv_row_bytes + index_dim * 2) + pooledBytes(cap)); - return linear + attn + ple_tail_n * wide * 2; -} -/// Contiguous storage or the page diagnostic's rounded arena and table, with every live allocation priced. -pub fn seqBytesPaged(cap: usize, tokens: usize) !usize { - if (cap == 0 or cap > max_tokens) return error.BadFlashNextPageTest; - if (tokens == 0) return seqBytes(cap); - const p = try pages.checkedLayout(cap, tokens); - const attention = attn_n * (cap * (2 * kv_row_bytes + index_dim * 2) + pooledBytes(cap)); - return std.math.add(usize, seqBytes(cap) - attention, p.total) catch error.BadFlashNextPageTest; -} -/// One attention layer's pooled indexer keys for a cache of `cap` tokens: a bf16 row of 128 a block of 4. -pub fn pooledBytes(cap: usize) usize { - return (cap + weights.index_ratio - 1) / weights.index_ratio * index_dim * 2; -} +/// A stream's device bytes at a rank's heads (cuda_seq_bytes.zig). +const sizes = @import("cuda_seq_bytes.zig"); +pub const seqBytes = sizes.seqBytes; +pub const seqBytesPaged = sizes.seqBytesPaged; +pub const pooledBytes = sizes.pooledBytes; +const Heads = @import("cuda_heads.zig").Heads; const round_graph = @import("cuda_round_graph.zig"); const pages = @import("cuda_pages.zig"); /// One loaded checkpoint: the embedding table, every layer, the mixer and head, and the batched pass's scratch. @@ -142,6 +123,8 @@ pub const Engine = struct { chunk: usize, /// Rows a lane's prompt pass takes while another stream decodes: TF_FLASHNEXT_CHUNK. busy_chunk: usize = batch.default_chunk, + /// The heads every model layer holds on this rank, which each stream's caches take (the MTP head's are full). + heads: Heads = .full, round_graph: round_graph.Cache = .{}, image_factory: ?*@import("cuda_images.zig").Vision = null, collective: ?*@import("cuda_collective.zig").Collective = null, @@ -220,6 +203,8 @@ pub const Engine = struct { try e.pockets.set(driver, stream, i, pk.gate[0 .. pk.n * dims], pk.up[0 .. pk.n * dims], pk.down[0 .. pk.n * dims], used); learned += used; } + e.heads = e.layers[0].heads; + for (e.layers[0..e.n_layers]) |*l| if (!l.heads.eql(e.heads)) return error.UnexpectedTensor; if (learned > 0) std.log.info("flashnext: the folder's Living Weights pockets hold {d} learned neurons beside the shared experts", .{learned}); e.store.loaded(); // every layer is on the device: no pinned staging, borrowed plane or open shard while serving e.ngram = try ple.Ngram.init(gpa, io, e.model_dir); @@ -302,9 +287,10 @@ pub const Engine = struct { errdefer if (paged) |*p| p.deinit(e.gpa); e.round_graph.clear(); if (cap == 0 or cap > max_tokens) return error.UnexpectedTensor; - _ = try seqBytesPaged(cap, if (paged) |p| p.tokens else 0); + _ = try seqBytesPaged(cap, if (paged) |p| p.tokens else 0, e.heads); var seq: Seq = .{ .stream = e.stream.*, + .heads = e.heads, .paged = paged, .conv = &.{}, .state = &.{}, @@ -322,11 +308,11 @@ pub const Engine = struct { // One rollback owns every published array, and drains before releasing any destination. errdefer seq.deinit(e.gpa); seq.logits = try e.gpa.alloc(u8, vocab * 2); - seq.conv = try zeros(e.gpa, e.driver, e.stream.*, linear_n, (gdn.taps - 1) * gdn.conv_dim * 2); - seq.state = try zeros(e.gpa, e.driver, e.stream.*, linear_n, state_n * 4); - seq.kc = if (seq.paged) |*p| try p.views(e.gpa, e.driver, .key) else try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); - seq.vc = if (seq.paged) |*p| try p.views(e.gpa, e.driver, .value) else try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); - seq.ikc = if (seq.paged) |*p| try p.views(e.gpa, e.driver, .index) else try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * index_dim * 2); + seq.conv = try zeros(e.gpa, e.driver, e.stream.*, linear_n, sizes.convBytes(e.heads)); + seq.state = try zeros(e.gpa, e.driver, e.stream.*, linear_n, sizes.stateBytes(e.heads)); + seq.kc = if (seq.paged) |*p| try p.views(e.gpa, e.driver, .key) else try zeros(e.gpa, e.driver, e.stream.*, attn_n, sizes.kvBytes(cap, e.heads)); + seq.vc = if (seq.paged) |*p| try p.views(e.gpa, e.driver, .value) else try zeros(e.gpa, e.driver, e.stream.*, attn_n, sizes.kvBytes(cap, e.heads)); + seq.ikc = if (seq.paged) |*p| try p.views(e.gpa, e.driver, .index) else try zeros(e.gpa, e.driver, e.stream.*, attn_n, sizes.indexBytes(cap)); seq.pooled = if (seq.paged) |*p| try p.views(e.gpa, e.driver, .pooled) else try zeros(e.gpa, e.driver, e.stream.*, attn_n, pooledBytes(cap)); seq.tail = try cuda.DeviceBuffer.alloc(e.driver, ple_tail_n * wide * 2); try seq.tail.fill8(0, e.stream.handle); @@ -459,7 +445,7 @@ pub const Engine = struct { const has_media = for (seqs) |q| { if (q.media != null) break true; } else false; - if (e.collective != null and round_graph.enabled() and !has_media) { + if (e.collective != null and round_graph.enabled() and !has_media and e.layers[0].dense_cut == null) { const key = try round_graph.Key.init(seqs, s.ids, scores, out); try e.round_graph.run(key, &p, ids); } else try p.chunk(ids, 0, null); @@ -478,6 +464,8 @@ pub const Seq = struct { stream: ?cuda.Stream = null, media: ?@import("cuda_media.zig").State = null, rank_id: u64 = 0, + /// The heads its DeltaNet states and KV caches hold (its engine's), which a snapshot's save and restore read. + heads: Heads = .full, conv: []cuda.DeviceBuffer, state: []cuda.DeviceBuffer, kc: []cuda.DeviceBuffer, @@ -540,13 +528,6 @@ test "linear layers skip every fourth index and attention starts at 3" { try std.testing.expectEqual(@as(usize, 11), attnIndex(47)); } -test "a sequence's device bytes grow by the twelve attention layers' keys, values, indexer keys and pooled blocks" { - const per_token = attn_n * (2 * kv_heads * head_dim * 2 + index_dim * 2); - try std.testing.expectEqual(per_token, seqBytes(6) - seqBytes(5)); - try std.testing.expectEqual(per_token + attn_n * index_dim * 2, seqBytes(5) - seqBytes(4)); - try std.testing.expect(seqBytes(max_tokens) > max_tokens * per_token); -} - test "a failed fill frees the buffer it was filling and every one before it" { const Fake = struct { live: *usize, @@ -586,4 +567,5 @@ pub const chunkRows = chunks.chunkRows; test { _ = @import("cuda_sequence_completion_test.zig"); _ = chunks; + _ = sizes; } diff --git a/zig/src/families/flashnext/cuda_prompt_load.zig b/zig/src/families/flashnext/cuda_prompt_load.zig index 45623e98cc..56158bd162 100644 --- a/zig/src/families/flashnext/cuda_prompt_load.zig +++ b/zig/src/families/flashnext/cuda_prompt_load.zig @@ -7,6 +7,10 @@ const gdn = @import("cuda_gdn.zig"); const wt = @import("cuda_prompt_wt.zig"); const mixed = @import("cuda_mixed_experts.zig"); const projections = @import("cuda_projection_load.zig"); +const dense_split = @import("cuda_dense_split.zig"); +const Cut = @import("cuda_modelopt_experts.zig").Cut; +const Heads = @import("cuda_heads.zig").Heads; +const heads_cut = @import("heads_cut.zig"); const layers_n = wt.layers_n; const dims = wt.dims; @@ -80,6 +84,10 @@ pub const Layer = struct { experts_up: cuda.DeviceBuffer, experts_down: cuda.DeviceBuffer, native_experts: ?mixed.Owned = null, + /// The input columns of `out` this rank keeps under the dense split (null: the whole projection). + dense_cut: ?Cut = null, + /// The heads this rank computes of the layer: the whole layer unless the head split cuts it. + heads: Heads = .full, pub fn free(self: *Layer, gpa: std.mem.Allocator) void { _ = gpa; @@ -268,6 +276,11 @@ pub fn loadLayerAt(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S } fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, store: *Store, pack_fn: cuda.Function, base: []const u8, linear: bool, with_ple: bool, world: usize, rank: usize, pocket: ?*Pocket) !Layer { + // decided before any view is held: deciding maps layer 0's shard; the split halves the heads and cuts the outs + const heads = try dense_split.heads(store, world); + const cut: ?Cut = if (!heads.eql(.full)) try Cut.init(dims, world, rank) else null; + // the rank's place among the heads' holders: 0 when one rank holds them all + const at: usize = if (cut != null) rank else 0; var buf: [200]u8 = undefined; var base_buf: [80]u8 = undefined; @memcpy(base_buf[0..base.len], base); @@ -339,19 +352,26 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S 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" }; - proj_q = try projections.loadStitched(gpa, driver, store, lp, &parts, gdn.proj_width, dims); - out_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.out_proj", .{lp}), dims, gdn.value_dim); + if (cut != null) { + const st = try heads_cut.delta(heads, at); + proj_q = try projections.loadStitchedCut(gpa, driver, store, lp, &heads_cut.delta_parts, &st.ranges, st.n, dims, st.sk); + } else proj_q = try projections.loadStitched(gpa, driver, store, lp, &heads_cut.delta_parts, gdn.proj_width, dims); + const out_name = try std.fmt.bufPrint(&buf, "{s}.out_proj", .{lp}); + out_q = if (cut) |c| try projections.loadCut(gpa, driver, store, out_name, dims, gdn.value_dim, c, .rows) else try projections.load(gpa, driver, store, out_name, dims, gdn.value_dim); 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); + if (cut != null) { + const kept = try heads_cut.convRows(gpa, conv, heads, at); + defer gpa.free(kept); + cw_b = try cuda.DeviceBuffer.fromHost(driver, kept); + } else 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)); + a_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(try heads_cut.valueSlice(a_log, heads, at))); + dt_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(try heads_cut.valueSlice(dt, heads, at))); 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); @@ -360,9 +380,12 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S 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" }; - proj_q = try projections.loadStitched(gpa, driver, store, ap, &parts, proj_n, dims); - out_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.o_proj", .{ap}), dims, out_k); + if (cut != null) { + const st = try heads_cut.attention(heads, at); + proj_q = try projections.loadStitchedCut(gpa, driver, store, ap, &heads_cut.attention_parts, &st.ranges, st.n, dims, st.sk); + } else proj_q = try projections.loadStitched(gpa, driver, store, ap, &heads_cut.attention_parts, proj_n, dims); + const out_name = try std.fmt.bufPrint(&buf, "{s}.o_proj", .{ap}); + out_q = if (cut) |c| try projections.loadCut(gpa, driver, store, out_name, dims, out_k, c, .rows) else try projections.load(gpa, driver, store, out_name, dims, out_k); const qn = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.q_norm.weight", .{ap}))); defer gpa.free(qn); const kn = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.k_norm.weight", .{ap}))); @@ -406,7 +429,7 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S var experts_down = deadBuf(driver); errdefer experts_down.free(); if (store.has(try std.fmt.bufPrint(&buf, "{s}.experts.0.gate_proj.weight", .{mlp_root}))) { - native_experts = try mixed.load(gpa, driver, store, mlp_root, world, rank); + native_experts = try mixed.load(gpa, driver, store, mlp_root, world, rank, cut != null); } else { if (world != 1 or rank != 0) return error.UnsupportedQuantization; const legacy = legacy_block: { @@ -468,5 +491,7 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S .experts_up = experts_up, .experts_down = experts_down, .native_experts = native_experts, + .dense_cut = cut, + .heads = heads, }; } diff --git a/zig/src/families/flashnext/cuda_prompt_step.zig b/zig/src/families/flashnext/cuda_prompt_step.zig index 58f76a512f..ddef6e6fdd 100644 --- a/zig/src/families/flashnext/cuda_prompt_step.zig +++ b/zig/src/families/flashnext/cuda_prompt_step.zig @@ -18,7 +18,6 @@ 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; @@ -240,12 +239,12 @@ pub fn attnStep(comptime Tri: type, driver: *cuda.Driver, stream: *cuda.Stream, // stream-ordered (queued kernels still read s.pos), then waited for: `pos` lives on this frame try s.pos.uploadAsync(0, std.mem.asBytes(&pos), stream.handle); try stream.synchronize(); - try tri.attnPrep(s.pa.ptr, s.pos.ptr, q_scale, k_scale, i_scale, s.inv.ptr, s.q.ptr, kc, vc, s.iq.ptr, ikc, s.dummy.ptr, s.dummy.ptr, eps, 1, 0, pg); + try tri.attnPrep(s.pa.ptr, s.pos.ptr, q_scale, k_scale, i_scale, s.inv.ptr, s.q.ptr, kc, vc, s.iq.ptr, ikc, s.dummy.ptr, s.dummy.ptr, eps, 1, 0, .full, pg); // attention.attention, dense: chunks of 512 keys up to the row's position (the row-by-row path stays dense) const keys: usize = @as(usize, @intCast(pos)) + 1; if (keys > wt.dense_nch * 512) return error.UnexpectedTensor; - try tri.attnChunks(s.q.ptr, kc, vc, s.pos.ptr, s.po.ptr, s.pm.ptr, s.pl.ptr, s.dummy.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads, (keys + 511) / 512, false, pg); - 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, false); + try tri.attnChunks(s.q.ptr, kc, vc, s.pos.ptr, s.po.ptr, s.pm.ptr, s.pl.ptr, s.dummy.ptr, s.dummy.ptr, s.dummy.ptr, 1, .full, (keys + 511) / 512, false, pg); + 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, .full, false); 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); } diff --git a/zig/src/families/flashnext/cuda_rank_engine.zig b/zig/src/families/flashnext/cuda_rank_engine.zig index 35a1650541..59d9cb2f97 100644 --- a/zig/src/families/flashnext/cuda_rank_engine.zig +++ b/zig/src/families/flashnext/cuda_rank_engine.zig @@ -7,6 +7,7 @@ const wire = @import("cuda_rank_wire.zig"); const Collective = @import("cuda_collective.zig").Collective; const wt = @import("cuda_prompt_wt.zig"); const mtp = @import("cuda_mtp.zig"); +const verify_rows = @import("cuda_rank_verify.zig"); pub const Identity = extern struct { config: [32]u8, @@ -20,6 +21,7 @@ pub const Identity = extern struct { head_sharded: u64 = 0, image_tokens: u64 = 0, max_images: u64 = 0, + dense_split: u64 = 0, }; fn fileHash(a: std.mem.Allocator, io: std.Io, path: []const u8) ![32]u8 { @@ -30,18 +32,26 @@ fn fileHash(a: std.mem.Allocator, io: std.Io, path: []const u8) ![32]u8 { return out; } +/// The checkpoint's bytes as both ranks see them: TENSORFOLD_CHECKPOINT_SHA256 when set, else the cached full digest. pub fn identity(a: std.mem.Allocator, io: std.Io, dir: []const u8, context: usize, chunk: usize) !Identity { - const digest = std.mem.span(std.c.getenv("TENSORFOLD_CHECKPOINT_SHA256") orelse return error.NoCheckpointIdentity); - if (digest.len != 64) return error.NoCheckpointIdentity; var checkpoint: [32]u8 = undefined; - _ = std.fmt.hexToBytes(&checkpoint, digest) catch return error.NoCheckpointIdentity; + if (std.c.getenv("TENSORFOLD_CHECKPOINT_SHA256")) |given| { + const digest = std.mem.span(given); + if (digest.len != 64) return error.NoCheckpointIdentity; + _ = std.fmt.hexToBytes(&checkpoint, digest) catch return error.NoCheckpointIdentity; + } else checkpoint = try @import("core").checkpoint_identity.fullIdentity(a, io, dir); const config = try std.fs.path.join(a, &.{ dir, "config.json" }); defer a.free(config); const index = try std.fs.path.join(a, &.{ dir, "model.safetensors.index.json" }); defer a.free(index); const program = try std.process.executablePathAlloc(io, a); defer a.free(program); - return .{ .config = try fileHash(a, io, config), .index = try fileHash(a, io, index), .program = try fileHash(a, io, program), .checkpoint = checkpoint, .context = context, .chunk = chunk, .decode_graph = @intFromBool(@import("cuda_round_graph.zig").enabled()), .head_sharded = @intFromBool(try @import("cuda_head.zig").enabled()) }; + return .{ .config = try fileHash(a, io, config), .index = try fileHash(a, io, index), .program = try fileHash(a, io, program), .checkpoint = checkpoint, .context = context, .chunk = chunk, .decode_graph = @intFromBool(@import("cuda_round_graph.zig").enabled()), .head_sharded = @intFromBool(try @import("cuda_head.zig").enabled()), .dense_split = @intFromBool(try @import("cuda_dense_split.zig").appliesDir(a, io, dir, 2)) }; +} + +/// How long a rank waits for its peer: a minute with a given digest, 30 when the peer may be hashing every byte. +pub fn linkWait() u64 { + return if (std.c.getenv("TENSORFOLD_CHECKPOINT_SHA256") != null) 60_000 else 30 * 60_000; } pub const Session = struct { @@ -59,12 +69,15 @@ pub const Session = struct { stopped: bool = false, available: std.atomic.Value(bool) = .init(true), - pub fn init(a: std.mem.Allocator, io: std.Io, host: []const u8, port: u16, rank: usize, mine: Identity) !*Session { + /// `wait_ms` bounds the link and the hello: the peer may still be reading every byte for its digest (`linkWait`). + pub fn init(a: std.mem.Allocator, io: std.Io, host: []const u8, port: u16, rank: usize, mine: Identity, wait_ms: u64) !*Session { + if (rank > 1) return error.BadRank; const s = try a.create(Session); errdefer a.destroy(s); - const link = try wire.Link.open(io, host, port, rank, 60_000); + if (wait_ms > 60_000) std.log.info("flashnext: rank {d} waits up to {d} min for rank {d} and its checkpoint digest", .{ rank, wait_ms / 60_000, 1 - rank }); + const link = try wire.Link.open(io, host, port, rank, wait_ms); errdefer link.close(); - const until = wire.deadline(60_000); + const until = wire.deadline(wait_ms); var peer: Identity = undefined; const hello: wire.Header = .{ .kind = .hello, .arg = rank, .bytes = @sizeOf(Identity) }; if (rank == 0) { @@ -221,6 +234,8 @@ pub const Session = struct { try s.ack(h); } + pub const verifyMany = verify_rows.send; + pub fn keep(s: *Session, comptime Tri: type, tri: Tri, q: *prompt.Seq, d: *mtp.Draft, count: usize) !void { if (!s.drafting or count == 0 or count > d.rows) return error.CannotRewind; const h = try s.command(.{ .kind = .keep, .object = q.rank_id, .arg = count }, &.{}); @@ -235,7 +250,7 @@ pub const Session = struct { const item = try s.windows.getOrPut(s.gpa, q.rank_id); if (!item.found_existing) { errdefer _ = s.windows.remove(q.rank_id); - item.value_ptr.* = try mtp.Draft.init(s.gpa, s.engine.driver, s.engine.stream.*, q.cap, 1, s.engine.chunk); + item.value_ptr.* = try verify_rows.makeWindow(s, q); } return item.value_ptr; } @@ -339,6 +354,7 @@ pub const Session = struct { const ops: mtp.Ops(Tri) = .{ .e = s.engine, .tri = tri, .h = undefined }; try ops.verify(q, try s.window(q), ids[0..n]); }, + .verify_many => try verify_rows.receive(s, Tri, tri, h, payload, logits), .keep => { if (!s.drafting) return error.UnsupportedRankOperation; const q = s.seqs.get(h.object) orelse return error.UnknownRankSequence; @@ -409,6 +425,7 @@ pub fn shape(h: wire.Header) !void { .forward => if (h.object == 0 or h.arg != 0 or h.bytes == 0 or h.bytes % 4 != 0 or (h.flags != @backingInt(prompt.batch.Logits.last) and h.flags != @backingInt(prompt.batch.Logits.none))) return error.BadRankFrame, .verify => if (h.object == 0 or h.arg != 0 or h.flags != 0 or h.bytes == 0 or h.bytes % 4 != 0 or h.bytes / 4 > mtp.max_rows) return error.BadRankFrame, .keep => if (h.object == 0 or h.arg == 0 or h.arg > mtp.max_rows or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, + .verify_many => if (h.object != 0 or h.flags != 0 or h.arg == 0 or h.arg > verify_rows.max_windows or h.bytes < 16 * h.arg or h.bytes > verify_rows.max_bytes) return error.BadRankFrame, .round => if (h.object != 0 or h.arg != 0 or h.flags != 0 or h.bytes == 0 or h.bytes % 12 != 0 or h.bytes / 12 > 32) return error.BadRankFrame, .save, .restore => if (h.object == 0 or h.arg == 0 or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, .drop_seq, .drop_snap => if (h.object == 0 or h.arg != 0 or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, @@ -448,7 +465,15 @@ test "rank operations refuse invalid geometry and unsupported hooks before execu try shape(.{ .kind = .round, .bytes = 24 }); try shape(.{ .kind = .verify, .object = 1, .bytes = mtp.max_rows * 4 }); try shape(.{ .kind = .keep, .object = 1, .arg = mtp.max_rows }); + try shape(.{ .kind = .verify_many, .arg = 2, .bytes = 32 }); + try shape(.{ .kind = .verify_many, .arg = verify_rows.max_windows, .bytes = verify_rows.max_bytes }); for ([_]wire.Header{ + .{ .kind = .verify_many, .bytes = 16 }, + .{ .kind = .verify_many, .arg = verify_rows.max_windows + 1, .bytes = verify_rows.max_bytes }, + .{ .kind = .verify_many, .arg = 2, .bytes = 31 }, + .{ .kind = .verify_many, .arg = 1, .bytes = verify_rows.max_bytes + 4 }, + .{ .kind = .verify_many, .object = 1, .arg = 1, .bytes = 16 }, + .{ .kind = .verify_many, .arg = 1, .bytes = 16, .flags = 1 }, .{ .kind = .verify, .object = 1, .bytes = (mtp.max_rows + 1) * 4 }, .{ .kind = .verify, .object = 1, .bytes = 3 }, .{ .kind = .verify, .bytes = 4 }, @@ -490,3 +515,30 @@ test "tensor ranks agree on image admission limits" { try std.testing.expectError(error.BadRankFrame, shape(.{ .kind = .media, .object = 1, .arg = 4096, .bytes = 65 * 12 })); _ = @import("cuda_rank_media.zig"); } + +test "rank identity: the cached full digest by default, a given digest when set, and the wait each allows" { + const t = std.testing; + const env = @import("sys").env; + var tmp = t.tmpDir(.{}); + defer tmp.cleanup(); + const header = "{\"w\":{\"dtype\":\"U8\",\"shape\":[4],\"data_offsets\":[0,4]}}"; + var shard: [8 + header.len + 4]u8 = undefined; + std.mem.writeInt(u64, shard[0..8], header.len, .little); + @memcpy(shard[8..][0..header.len], header); + @memcpy(shard[8 + header.len ..], "abcd"); + try tmp.dir.writeFile(t.io, .{ .sub_path = "model.safetensors", .data = &shard }); + try tmp.dir.writeFile(t.io, .{ .sub_path = "config.json", .data = "{\"hidden_size\":8}" }); + try tmp.dir.writeFile(t.io, .{ .sub_path = "model.safetensors.index.json", .data = "{\"weight_map\":{\"w\":\"model.safetensors\"}}" }); + const dir = try tmp.dir.realPathFileAlloc(t.io, ".", t.allocator); + defer t.allocator.free(dir); + _ = env.unsetenv("TENSORFOLD_CHECKPOINT_SHA256"); + try t.expectEqual(@as(u64, 30 * 60_000), linkWait()); + const computed = try identity(t.allocator, t.io, dir, 4096, 64); + try t.expectEqual(try @import("core").checkpoint_identity.fullIdentity(t.allocator, t.io, dir), computed.checkpoint); + _ = env.setenv("TENSORFOLD_CHECKPOINT_SHA256", "abababababababababababababababababababababababababababababababab", 1); + defer _ = env.unsetenv("TENSORFOLD_CHECKPOINT_SHA256"); + try t.expectEqual(@as(u64, 60_000), linkWait()); + try t.expectEqual(@as([32]u8, @splat(0xab)), (try identity(t.allocator, t.io, dir, 4096, 64)).checkpoint); + _ = env.setenv("TENSORFOLD_CHECKPOINT_SHA256", "ab", 1); + try t.expectError(error.NoCheckpointIdentity, identity(t.allocator, t.io, dir, 4096, 64)); +} diff --git a/zig/src/families/flashnext/cuda_rank_memory.zig b/zig/src/families/flashnext/cuda_rank_memory.zig index 642510a680..bad76b127c 100644 --- a/zig/src/families/flashnext/cuda_rank_memory.zig +++ b/zig/src/families/flashnext/cuda_rank_memory.zig @@ -3,12 +3,14 @@ const std = @import("std"); const core = @import("core"); const cuda = @import("cuda"); const wt = @import("cuda_prompt_wt.zig"); +const dense_split = @import("cuda_dense_split.zig"); const Entry = core.safetensors.Entry; pub fn estimate(a: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize, rank: usize) !u64 { if ((world != 1 and world != 2) or rank >= world) return error.BadRank; var store = try wt.Store.open(a, io, dir); defer store.deinit(); + const split = try dense_split.applies(&store, world); const wm = store.parsed.value.object.get("weight_map") orelse return error.MissingTensor; var files: std.StringHashMapUnmanaged(void) = .empty; defer files.deinit(a); @@ -23,21 +25,23 @@ pub fn estimate(a: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize, defer file.close(io); var tensors = file.names.iterator(); while (tensors.next()) |tensor| { - total = try std.math.add(u64, total, try tensorBytes(tensor.key_ptr.*, tensor.value_ptr.*, world, rank)); + total = try std.math.add(u64, total, try tensorBytes(tensor.key_ptr.*, tensor.value_ptr.*, world, rank, split)); } } return total; } -fn tensorBytes(name: []const u8, e: Entry, world: usize, rank: usize) !u64 { +fn tensorBytes(name: []const u8, e: Entry, world: usize, rank: usize, split: bool) !u64 { const bytes: u64 = e.end - e.begin; if (wt.hostOnly(name)) return 0; + if (split) if (try dense_split.headBytes(name, e, world)) |cut| return cut; + if (split and e.dtype == .bf16 and e.rank == 2) if (dense_split.axisOf(name)) |axis| return dense_split.residentBytes(e.dim(0), e.dim(1), axis, world); if (rank != 0 and std.mem.indexOf(u8, name, ".mlp.shared_expert.") != null) return 0; if (std.mem.indexOf(u8, name, ".mlp.experts.") != null) { if (std.mem.endsWith(u8, name, ".weight_scale_2")) return bytes; if (!std.mem.endsWith(u8, name, ".weight")) return 0; if (e.dtype != .u8 or e.rank != 2 or bytes % (8 * world) != 0) return error.UnexpectedTensor; - // NVFP4 device blocks hold 128 code words and 16 scale words; global scales stay replicated. + // NVFP4 device blocks hold 128 code words and 16 scale words; a rank's whole experts average bytes / world. return try std.math.mul(u64, bytes / (8 * world), 9); } if (std.mem.indexOf(u8, name, ".weight_scale") != null or std.mem.endsWith(u8, name, ".input_scale")) return 0; @@ -63,16 +67,71 @@ fn centered(name: []const u8) bool { test "rank estimates partition routed storage while retaining replicated and rank-zero shared operators" { const e: Entry = .{ .dtype = .u8, .rank = 2, .shape = .{ 640, 1280, 1, 1, 1 }, .begin = 0, .end = 640 * 1280 }; const name = "model.language_model.layers.0.mlp.experts.0.gate_proj.weight"; - try std.testing.expectEqual(@as(u64, cuda.experts.Nvfp4Experts.bytes(1, 320, 2560, 1)), try tensorBytes(name, e, 2, 0)); - try std.testing.expectEqual(try tensorBytes(name, e, 2, 0), try tensorBytes(name, e, 2, 1)); + try std.testing.expectEqual(@as(u64, cuda.experts.Nvfp4Experts.bytes(1, 320, 2560, 1)), try tensorBytes(name, e, 2, 0, false)); + try std.testing.expectEqual(try tensorBytes(name, e, 2, 0, false), try tensorBytes(name, e, 2, 1, false)); var shared = e; shared.dtype = .f8_e4m3; shared.shape = .{ 640, 2560, 1, 1, 1 }; shared.end = 640 * 2560; const shared_name = "model.language_model.layers.0.mlp.shared_expert.gate_proj.weight"; - try std.testing.expectEqual(@as(u64, 640 * 2560 * 33 / 32), try tensorBytes(shared_name, shared, 2, 0)); - try std.testing.expectEqual(@as(u64, 0), try tensorBytes(shared_name, shared, 2, 1)); + try std.testing.expectEqual(@as(u64, 640 * 2560 * 33 / 32), try tensorBytes(shared_name, shared, 2, 0, false)); + try std.testing.expectEqual(@as(u64, 0), try tensorBytes(shared_name, shared, 2, 1, false)); const norm: Entry = .{ .dtype = .bf16, .rank = 1, .shape = .{ 256, 1, 1, 1, 1 }, .begin = 0, .end = 512 }; - try std.testing.expectEqual(@as(u64, 1024), try tensorBytes("model.language_model.layers.3.self_attn.q_norm.weight", norm, 2, 1)); - try std.testing.expectEqual(@as(u64, 512), try tensorBytes("model.language_model.layers.0.linear_attn.norm.weight", norm, 2, 1)); + try std.testing.expectEqual(@as(u64, 1024), try tensorBytes("model.language_model.layers.3.self_attn.q_norm.weight", norm, 2, 1, false)); + try std.testing.expectEqual(@as(u64, 512), try tensorBytes("model.language_model.layers.0.linear_attn.norm.weight", norm, 2, 1, false)); +} + +test "rank estimates under the dense split count each rank's out-projection output rows, padded as BF16 packing pads" { + const out: Entry = .{ .dtype = .bf16, .rank = 2, .shape = .{ 2560, 6144, 1, 1, 1 }, .begin = 0, .end = 2560 * 6144 * 2 }; + const names = [_][]const u8{ "model.language_model.layers.0.linear_attn.out_proj.weight", "model.language_model.layers.3.self_attn.o_proj.weight" }; + for (names) |name| { + for (0..2) |rank| try std.testing.expectEqual(@as(u64, 2560 * 3072 * 2), try tensorBytes(name, out, 2, rank, true)); + try std.testing.expectEqual(@as(u64, 2560 * 6144 * 2), try tensorBytes(name, out, 2, 0, false)); + try std.testing.expectEqual(@as(u64, 2560 * 6144 * 2), try tensorBytes(name, out, 1, 0, false)); + } + try std.testing.expectEqual(@as(u64, 0), try tensorBytes("mtp.layers.0.self_attn.o_proj.weight", out, 2, 0, true)); +} + +test "rank estimates under the head split count each rank's in-projection rows, conv channels and gate values" { + const k = 2560; + const Part = struct { name: []const u8, n: usize, kept: usize }; + const delta = [_]Part{ .{ .name = "in_proj_qkv", .n = 10240, .kept = 5120 }, .{ .name = "in_proj_z", .n = 6144, .kept = 3072 }, .{ .name = "in_proj_b", .n = 48, .kept = 24 }, .{ .name = "in_proj_a", .n = 48, .kept = 24 } }; + const attn = [_]Part{ .{ .name = "q_proj", .n = 12288, .kept = 6144 }, .{ .name = "k_proj", .n = 512, .kept = 256 }, .{ .name = "v_proj", .n = 512, .kept = 256 }, .{ .name = "indexer.index_qk_proj", .n = 640, .kept = 640 } }; + for ([_][]const Part{ &delta, &attn }, [_][]const u8{ "0.linear_attn.", "3.self_attn." }, [_]usize{ 8320, 7296 }) |parts, class, packed_rows| { + var split: [2]u64 = .{ 0, 0 }; + var buf: [120]u8 = undefined; + for (parts) |p| { + const name = try std.fmt.bufPrint(&buf, "model.language_model.layers.{s}{s}.weight", .{ class, p.name }); + const e: Entry = .{ .dtype = .bf16, .rank = 2, .shape = .{ p.n, k, 1, 1, 1 }, .begin = 0, .end = p.n * k * 2 }; + for (&split, 0..) |*sum, rank| sum.* += try tensorBytes(name, e, 2, rank, true); + // unsplit, each part is the whole tensor on every rank, padded on its own as today + try std.testing.expectEqual(@as(u64, cuda.nvfp4.padded(p.n) * k * 2), try tensorBytes(name, e, 2, 1, false)); + try std.testing.expect(try tensorBytes(name, e, 2, 0, true) >= p.kept * k * 2); + } + // the parts' shares sum to the stitched cut as one BF16 pack pads it: 8,240 rows to 8,320, 7,296 to itself + for (split) |sum| try std.testing.expectEqual(@as(u64, packed_rows * k * 2), sum); + } + const conv: Entry = .{ .dtype = .bf16, .rank = 3, .shape = .{ 10240, 1, 4, 1, 1 }, .begin = 0, .end = 10240 * 4 * 2 }; + const conv_name = "model.language_model.layers.0.linear_attn.conv1d.weight"; + for (0..2) |rank| try std.testing.expectEqual(@as(u64, 5120 * 4 * 2), try tensorBytes(conv_name, conv, 2, rank, true)); + try std.testing.expectEqual(@as(u64, 10240 * 4 * 2), try tensorBytes(conv_name, conv, 2, 1, false)); + const gates: Entry = .{ .dtype = .f32, .rank = 1, .shape = .{ 48, 1, 1, 1, 1 }, .begin = 0, .end = 48 * 4 }; + for ([_][]const u8{ "model.language_model.layers.0.linear_attn.A_log", "model.language_model.layers.0.linear_attn.dt_bias" }) |name| { + for (0..2) |rank| try std.testing.expectEqual(@as(u64, 24 * 4), try tensorBytes(name, gates, 2, rank, true)); + try std.testing.expectEqual(@as(u64, 48 * 4), try tensorBytes(name, gates, 2, 0, false)); + } + // the norms stay whole on both ranks, and the MTP head's layer is the draft's to price + const norm: Entry = .{ .dtype = .bf16, .rank = 1, .shape = .{ 128, 1, 1, 1, 1 }, .begin = 0, .end = 256 }; + try std.testing.expectEqual(@as(u64, 256), try tensorBytes("model.language_model.layers.0.linear_attn.norm.weight", norm, 2, 1, true)); + const q: Entry = .{ .dtype = .bf16, .rank = 2, .shape = .{ 12288, k, 1, 1, 1 }, .begin = 0, .end = 12288 * k * 2 }; + try std.testing.expectEqual(@as(u64, 0), try tensorBytes("mtp.layers.0.self_attn.q_proj.weight", q, 2, 0, true)); +} + +test "shared output-row cuts include padded gate/up rows and full-input down rows on both ranks" { + const up: Entry = .{ .dtype = .bf16, .rank = 2, .shape = .{ 640, 2560, 1, 1, 1 }, .begin = 0, .end = 640 * 2560 * 2 }; + const down: Entry = .{ .dtype = .bf16, .rank = 2, .shape = .{ 2560, 640, 1, 1, 1 }, .begin = 0, .end = 2560 * 640 * 2 }; + for (0..2) |rank| { + try std.testing.expectEqual(@as(u64, 384 * 2560 * 2), try tensorBytes("model.language_model.layers.0.mlp.shared_expert.gate_proj.weight", up, 2, rank, true)); + try std.testing.expectEqual(@as(u64, 1280 * 640 * 2), try tensorBytes("model.language_model.layers.0.mlp.shared_expert.down_proj.weight", down, 2, rank, true)); + } } diff --git a/zig/src/families/flashnext/cuda_rank_verify.zig b/zig/src/families/flashnext/cuda_rank_verify.zig new file mode 100644 index 0000000000..a9621e1f69 --- /dev/null +++ b/zig/src/families/flashnext/cuda_rank_verify.zig @@ -0,0 +1,133 @@ +//! Control carries a shared verify pass: each window's sequence and rows, then their ids; both ranks run one pass. +const std = @import("std"); +const prompt = @import("cuda_prompt.zig"); +const wire = @import("cuda_rank_wire.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const mtp = @import("cuda_mtp.zig"); +const Session = @import("cuda_rank_engine.zig").Session; + +/// Windows a shared pass takes at most (a round's streams). +pub const max_windows = 32; +/// The bytes a frame can hold: every window's header, then a shared pass's rows of ids. +pub const max_bytes = 12 * max_windows + 4 * mtp.round_rows; + +/// A decoded frame: window w is sequence ids[w], rows[w] token ids from `tokens` in window order. +pub const Frame = struct { + count: usize, + ids: [max_windows]u64 = undefined, + rows: [max_windows]u32 = undefined, + tokens: []const u8, + + /// Window w's token ids, each read from the frame's little-endian words. + pub fn window(f: *const Frame, w: usize, out: *[mtp.max_rows]u32) []const u32 { + var at: usize = 0; + for (f.rows[0..w]) |r| at += r; + for (out[0..f.rows[w]], 0..) |*v, i| v.* = std.mem.readInt(u32, f.tokens[(at + i) * 4 ..][0..4], .little); + return out[0..f.rows[w]]; + } +}; + +/// `count` windows from `payload`: exact length, 1 to max_rows rows each, at most `limit` rows, distinct sequences. +pub fn decode(payload: []const u8, count: usize, limit: usize) !Frame { + if (count == 0 or count > max_windows or payload.len < 12 * count or payload.len > max_bytes) return error.BadRankFrame; + var f: Frame = .{ .count = count, .tokens = payload[12 * count ..] }; + var rows: usize = 0; + for (0..count) |w| { + f.ids[w] = std.mem.readInt(u64, payload[w * 12 ..][0..8], .little); + f.rows[w] = std.mem.readInt(u32, payload[w * 12 + 8 ..][0..4], .little); + if (f.ids[w] == 0 or f.rows[w] == 0 or f.rows[w] > mtp.max_rows) return error.BadRankFrame; + for (f.ids[0..w]) |prior| if (prior == f.ids[w]) return error.DuplicateRankSequence; + rows += f.rows[w]; + } + if (rows > limit or f.tokens.len != 4 * rows) return error.BadRankFrame; + for (0..rows) |i| if (std.mem.readInt(u32, f.tokens[i * 4 ..][0..4], .little) >= wt.vocab) return error.InvalidToken; + return f; +} + +/// The frame for `ids[w]`'s window `idss[w]` into `out`; refuses what `decode` refuses at the same `limit`. +pub fn encode(ids: []const u64, idss: []const []const u32, limit: usize, out: []u8) ![]const u8 { + if (ids.len != idss.len) return error.BadRankFrame; + var rows: usize = 0; + for (idss) |w| rows += w.len; + const n = 12 * ids.len + 4 * rows; + if (n > out.len) return error.BadRankFrame; + var at = 12 * ids.len; + for (ids, idss, 0..) |id, w, i| { + std.mem.writeInt(u64, out[i * 12 ..][0..8], id, .little); + std.mem.writeInt(u32, out[i * 12 + 8 ..][0..4], @intCast(w.len), .little); + for (w) |token| { + std.mem.writeInt(u32, out[at..][0..4], token, .little); + at += 4; + } + } + _ = try decode(out[0..n], ids.len, limit); + return out[0..n]; +} + +/// Every key in `map`, missing ones by `make(ctx, i)`, before any getPtr; a failed make's key stays absent. +pub fn ensure(map: anytype, gpa: std.mem.Allocator, keys: []const u64, ctx: anytype, comptime make: anytype) !void { + for (keys, 0..) |key, i| { + const item = try map.getOrPut(gpa, key); + if (item.found_existing) continue; + errdefer _ = map.remove(key); + item.value_ptr.* = try make(ctx, i); + } +} + +/// Rank 1's window state for a sequence: one draft row, the head's absorb rows (Session.window's form). +pub fn makeWindow(s: *Session, q: *const prompt.Seq) !mtp.Draft { + return mtp.Draft.init(s.gpa, s.engine.driver, s.engine.stream.*, q.cap, 1, s.engine.chunk, s.engine.heads); +} + +fn limitOf(s: *const Session) usize { + return @min(mtp.round_rows, s.engine.chunk); +} + +/// Rank 0: the frame to rank 1, the same pass here, then rank 1's ack; a failure after the send poisons the session. +pub fn send(s: *Session, comptime Tri: type, tri: Tri, seqs: []const *prompt.Seq, ds: []const *mtp.Draft, idss: []const []const u32, out: u64) !void { + if (!s.drafting or seqs.len == 0 or seqs.len > max_windows or seqs.len != ds.len or seqs.len != idss.len) return error.UnsupportedRankOperation; + var ids: [max_windows]u64 = undefined; + for (seqs, idss, 0..) |q, w, i| { + if (q.rank_id == 0 or q.fed + w.len > q.cap) return error.UnsupportedRankOperation; + ids[i] = q.rank_id; + } + var raw: [max_bytes]u8 = undefined; + const bytes = try encode(ids[0..seqs.len], idss, limitOf(s), &raw); + const h = try s.command(.{ .kind = .verify_many, .arg = seqs.len, .bytes = @intCast(bytes.len) }, bytes); + errdefer s.poison(); + const ops: mtp.Ops(Tri) = .{ .e = s.engine, .tri = tri, .h = undefined }; + try ops.verifyMany(seqs, ds, idss, out); + try s.ack(h); +} + +/// Rank 1: every check before any state moves, every window draft made before any is borrowed, then the pass. +pub fn receive(s: *Session, comptime Tri: type, tri: Tri, h: wire.Header, payload: []const u8, logits: u64) !void { + if (!s.drafting) return error.UnsupportedRankOperation; + const f = try decode(payload, @intCast(h.arg), limitOf(s)); + var seqs: [max_windows]*prompt.Seq = undefined; + for (f.ids[0..f.count], f.rows[0..f.count], 0..) |id, rows, w| { + seqs[w] = s.seqs.get(id) orelse return error.UnknownRankSequence; + if (seqs[w].fed + rows > seqs[w].cap) return error.PromptTooLong; + } + const Make = struct { + s: *Session, + seqs: []const *prompt.Seq, + fn make(ctx: @This(), i: usize) !mtp.Draft { + return makeWindow(ctx.s, ctx.seqs[i]); + } + }; + try ensure(&s.windows, s.gpa, f.ids[0..f.count], Make{ .s = s, .seqs = seqs[0..f.count] }, Make.make); + var ds: [max_windows]*mtp.Draft = undefined; + var tokens: [max_windows][mtp.max_rows]u32 = undefined; + var idss: [max_windows][]const u32 = undefined; + for (0..f.count) |w| { + ds[w] = s.windows.getPtr(f.ids[w]).?; + idss[w] = f.window(w, &tokens[w]); + } + const ops: mtp.Ops(Tri) = .{ .e = s.engine, .tri = tri, .h = undefined }; + try ops.verifyMany(seqs[0..f.count], ds[0..f.count], idss[0..f.count], logits); +} + +test { + _ = @import("cuda_rank_verify_test.zig"); +} diff --git a/zig/src/families/flashnext/cuda_rank_verify_test.zig b/zig/src/families/flashnext/cuda_rank_verify_test.zig new file mode 100644 index 0000000000..34d4e32420 --- /dev/null +++ b/zig/src/families/flashnext/cuda_rank_verify_test.zig @@ -0,0 +1,77 @@ +//! The shared verify frame without ranks: round trips, every refusal, and two-phase window creation. +const std = @import("std"); +const v = @import("cuda_rank_verify.zig"); +const mtp = @import("cuda_mtp.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const t = std.testing; + +test "shared verify frame: one and 32 windows round trip in window order" { + var raw: [v.max_bytes]u8 = undefined; + const one = try v.encode(&.{7}, &.{&.{ 1, 2, 3 }}, 64, &raw); + const f1 = try v.decode(one, 1, 64); + var row: [mtp.max_rows]u32 = undefined; + try t.expectEqual(@as(u64, 7), f1.ids[0]); + try t.expectEqualSlices(u32, &.{ 1, 2, 3 }, f1.window(0, &row)); + var ids: [v.max_windows]u64 = undefined; + var rows: [v.max_windows][2]u32 = undefined; + var idss: [v.max_windows][]const u32 = undefined; + for (0..v.max_windows) |w| { + ids[w] = 100 + w; + rows[w] = .{ @intCast(w), @intCast(wt.vocab - 1 - w) }; + idss[w] = &rows[w]; + } + const all = try v.encode(&ids, &idss, 64, &raw); + try t.expectEqual(@as(usize, 12 * 32 + 4 * 64), all.len); + const f = try v.decode(all, v.max_windows, 64); + for (0..v.max_windows) |w| { + try t.expectEqual(ids[w], f.ids[w]); + try t.expectEqualSlices(u32, &rows[w], f.window(w, &row)); + } +} + +test "shared verify frame: refused before any state moves" { + var raw: [v.max_bytes]u8 = undefined; + const good = try v.encode(&.{ 1, 2 }, &.{ &.{5}, &.{ 6, 7 } }, 64, &raw); + try t.expectError(error.BadRankFrame, v.decode(good, 0, 64)); + try t.expectError(error.BadRankFrame, v.decode(good, v.max_windows + 1, 64)); + try t.expectError(error.BadRankFrame, v.decode(good[0 .. good.len - 4], 2, 64)); + var long: [v.max_bytes]u8 = undefined; + @memcpy(long[0..good.len], good); + try t.expectError(error.BadRankFrame, v.decode(long[0 .. good.len + 4], 2, 64)); + // the total row bound: three windows of max_rows past a 64-row pass + var wide: [mtp.max_rows]u32 = @splat(1); + try t.expectError(error.BadRankFrame, v.encode(&.{ 1, 2, 3 }, &.{ &wide, &wide, &wide }, 64, &raw)); + try t.expectError(error.BadRankFrame, v.encode(&.{1}, &.{&.{}}, 64, &raw)); + var too_wide: [mtp.max_rows + 1]u32 = @splat(1); + try t.expectError(error.BadRankFrame, v.encode(&.{1}, &.{&too_wide}, 1024, &raw)); + try t.expectError(error.DuplicateRankSequence, v.encode(&.{ 4, 4 }, &.{ &.{1}, &.{1} }, 64, &raw)); + try t.expectError(error.InvalidToken, v.encode(&.{1}, &.{&.{wt.vocab}}, 64, &raw)); + try t.expectError(error.BadRankFrame, v.encode(&.{0}, &.{&.{1}}, 64, &raw)); +} + +test "shared verify frame: ensure makes every missing window before any pointer is taken" { + var map: std.AutoHashMapUnmanaged(u64, [64]u8) = .empty; + defer map.deinit(t.allocator); + var keys: [40]u64 = undefined; + for (&keys, 0..) |*k, i| k.* = 1000 + i; + const Make = struct { + fn make(fail_at: usize, i: usize) ![64]u8 { + if (i == fail_at) return error.OutOfDeviceMemory; + return @splat(@intCast(i)); + } + }; + try v.ensure(&map, t.allocator, &keys, @as(usize, std.math.maxInt(usize)), Make.make); + var seen: [40]*[64]u8 = undefined; + for (keys, 0..) |k, i| { + seen[i] = map.getPtr(k).?; + try t.expectEqual(@as(u8, @intCast(i)), seen[i][0]); + for (seen[0..i]) |p| try t.expect(p != seen[i]); + } + // a failed make leaves its key absent and the four made before it in place + var fresh: std.AutoHashMapUnmanaged(u64, [64]u8) = .empty; + defer fresh.deinit(t.allocator); + try t.expectError(error.OutOfDeviceMemory, v.ensure(&fresh, t.allocator, keys[0..8], @as(usize, 4), Make.make)); + try t.expectEqual(@as(u32, 4), fresh.count()); + try t.expect(fresh.get(keys[4]) == null); + for (keys[0..4], 0..) |k, i| try t.expectEqual(@as(u8, @intCast(i)), fresh.get(k).?[0]); +} diff --git a/zig/src/families/flashnext/cuda_rank_wire.zig b/zig/src/families/flashnext/cuda_rank_wire.zig index 98606b2994..23fc60e548 100644 --- a/zig/src/families/flashnext/cuda_rank_wire.zig +++ b/zig/src/families/flashnext/cuda_rank_wire.zig @@ -6,7 +6,7 @@ const p = std.posix; const net = std.Io.net; pub const max_payload = 1 << 20; -pub const Kind = enum(u32) { hello = 1, ready, create, forward, round, drop_seq, save, restore, drop_snap, stop, ack, verify, keep, media }; +pub const Kind = enum(u32) { hello = 1, ready, create, forward, round, drop_seq, save, restore, drop_snap, stop, ack, verify, keep, media, verify_many }; pub const Header = struct { kind: Kind, sequence: u64 = 0, diff --git a/zig/src/families/flashnext/cuda_score.zig b/zig/src/families/flashnext/cuda_score.zig index 71b6e573cd..00043dc976 100644 --- a/zig/src/families/flashnext/cuda_score.zig +++ b/zig/src/families/flashnext/cuda_score.zig @@ -52,6 +52,33 @@ pub fn CapabilityCases(comptime Cuda: type) type { self.ranks = @ptrFromInt(16); try std.testing.expect(self.backend().vtable.score == null); } + + test "decision capability: the ranked table outlives backend(), as the scoring table does" { + var self: Cuda = undefined; + self.head = null; + for ([_]usize{ 0, 16 }) |r| { + self.ranks = if (r == 0) null else @ptrFromInt(r); + // a real call: inlined, a table built in backend() would sit in this frame and survive clobber() + const vt = @call(.never_inline, Cuda.backend, .{&self}).vtable; + var here: u8 = 0; + std.mem.doNotOptimizeAway(&here); + const apart = @max(@intFromPtr(vt), @intFromPtr(&here)) - @min(@intFromPtr(vt), @intFromPtr(&here)); + try std.testing.expect(apart > 1 << 20); // static data, not this thread's stack + clobber(); + // read before backend() runs again: a table in its frame would be rebuilt at the same address + const step = vt.prefill_step; + const scores = vt.score != null; + try std.testing.expect(step != null and step == self.backend().vtable.prefill_step); + try std.testing.expectEqual(r == 0, scores); + } + } + + /// Writes over the stack below the caller, where a table built inside backend() would have lived. + noinline fn clobber() void { + var junk: [4096]u8 = undefined; + @memset(&junk, 0xa5); + std.mem.doNotOptimizeAway(&junk); + } }; } diff --git a/zig/src/families/flashnext/cuda_seq_bytes.zig b/zig/src/families/flashnext/cuda_seq_bytes.zig new file mode 100644 index 0000000000..323cb3beda --- /dev/null +++ b/zig/src/families/flashnext/cuda_seq_bytes.zig @@ -0,0 +1,91 @@ +//! One stream's device bytes at a rank's heads: DeltaNet conv windows and states, attention caches, the n-gram tail. +const std = @import("std"); +const wt = @import("cuda_prompt_wt.zig"); +const gdn = @import("cuda_gdn.zig"); +const pages = @import("cuda_pages.zig"); +const Heads = @import("cuda_heads.zig").Heads; + +const linear_n = wt.linear_n; +const attn_n = wt.attn_n; +const index_dim = wt.index_dim; +const max_tokens = wt.window; + +/// One DeltaNet layer's conv window: the last taps - 1 rows of the heads' conv channels, bf16. +pub fn convBytes(h: Heads) usize { + return (gdn.taps - 1) * h.convDim() * 2; +} + +/// One DeltaNet layer's recurrent state: a dk x dv fp32 matrix for each of the heads' value heads. +pub fn stateBytes(h: Heads) usize { + return h.nv * gdn.dk * gdn.dv * 4; +} + +/// One attention layer's keys (or values) for `cap` positions of the heads' KV heads. +pub fn kvBytes(cap: usize, h: Heads) usize { + return cap * h.kvRow(); +} + +/// One attention layer's indexer keys for `cap` positions: whole on every rank. +pub fn indexBytes(cap: usize) usize { + return cap * index_dim * 2; +} + +/// One attention layer's pooled indexer keys for a cache of `cap` tokens: a bf16 row of 128 a block of 4. +pub fn pooledBytes(cap: usize) usize { + return (cap + wt.index_ratio - 1) / wt.index_ratio * index_dim * 2; +} + +/// Every attention layer's contiguous caches for `cap` positions: keys, values, indexer keys and pooled blocks. +fn attnBytes(cap: usize, h: Heads) usize { + return attn_n * (2 * kvBytes(cap, h) + indexBytes(cap) + pooledBytes(cap)); +} + +/// Device bytes `Engine.newSeq(cap)` allocates: conv tails, recurrent states, KV with indexer keys, the n-gram tail. +pub fn seqBytes(cap: usize, h: Heads) usize { + const linear = linear_n * (convBytes(h) + stateBytes(h)); + return linear + attnBytes(cap, h) + wt.ple_tail_n * wt.wide * 2; +} + +/// Contiguous storage or the page diagnostic's rounded arena and table, with every live allocation priced. +pub fn seqBytesPaged(cap: usize, tokens: usize, h: Heads) !usize { + if (cap == 0 or cap > max_tokens) return error.BadFlashNextPageTest; + if (tokens == 0) return seqBytes(cap, h); + // the pages hold the whole layer's KV rows (cuda_pages.zig): a rank's half heads take none + if (!h.eql(.full)) return error.FlashNextPagesUnsupported; + const p = try pages.checkedLayout(cap, tokens); + return std.math.add(usize, seqBytes(cap, h) - attnBytes(cap, h), p.total) catch error.BadFlashNextPageTest; +} + +test "a sequence's device bytes grow by the twelve attention layers' keys, values, indexer keys and pooled blocks" { + for ([_]Heads{ .full, .half }) |h| { + const per_token = attn_n * (2 * h.kv * wt.head_dim * 2 + index_dim * 2); + try std.testing.expectEqual(per_token, seqBytes(6, h) - seqBytes(5, h)); + try std.testing.expectEqual(per_token + attn_n * index_dim * 2, seqBytes(5, h) - seqBytes(4, h)); + try std.testing.expect(seqBytes(max_tokens, h) > max_tokens * per_token); + } +} + +test "a stream at 65,536 tokens: 1.84 GiB with the whole layer, 1.04 GiB with a rank's half of the heads" { + const cap = 65536; + // 36 DeltaNet layers: 3 conv rows of 10,240 or 5,120 bf16 channels, 48 or 24 fp32 states of 128 x 128 + const full_linear: usize = 36 * (3 * 10240 * 2 + 48 * 128 * 128 * 4); + const half_linear: usize = 36 * (3 * 5120 * 2 + 24 * 128 * 128 * 4); + // 12 attention layers: 2 or 1 KV heads of 256 bf16 keys and values, 128 indexer keys, a pooled row a block of 4 + const full_attn: usize = 12 * (cap * (2 * 2 * 256 * 2 + 128 * 2) + cap / 4 * 128 * 2); + const half_attn: usize = 12 * (cap * (2 * 1 * 256 * 2 + 128 * 2) + cap / 4 * 128 * 2); + const tail: usize = 9 * 4 * 2560 * 2; + try std.testing.expectEqual(full_linear + full_attn + tail, seqBytes(cap, .full)); + try std.testing.expectEqual(half_linear + half_attn + tail, seqBytes(cap, .half)); + try std.testing.expectEqual(@as(usize, 1_977_913_344), seqBytes(cap, .full)); + try std.testing.expectEqual(@as(usize, 1_114_877_952), seqBytes(cap, .half)); + try std.testing.expectEqual(@as(usize, 30720), convBytes(.half)); + try std.testing.expectEqual(@as(usize, 1_572_864), stateBytes(.half)); + try std.testing.expectEqual(@as(usize, cap * 512), kvBytes(cap, .half)); +} + +test "the page diagnostic prices the whole layer's pages and refuses a rank's half of the heads" { + try std.testing.expectEqual(seqBytes(64, .half), try seqBytesPaged(64, 0, .half)); + try std.testing.expectError(error.FlashNextPagesUnsupported, seqBytesPaged(64, 16, .half)); + try std.testing.expect(try seqBytesPaged(64, 16, .full) > 0); + try std.testing.expectError(error.BadFlashNextPageTest, seqBytesPaged(0, 0, .half)); +} diff --git a/zig/src/families/flashnext/cuda_sequence_completion_test.zig b/zig/src/families/flashnext/cuda_sequence_completion_test.zig index ddf4e69bfa..7362b7d298 100644 --- a/zig/src/families/flashnext/cuda_sequence_completion_test.zig +++ b/zig/src/families/flashnext/cuda_sequence_completion_test.zig @@ -98,6 +98,7 @@ test "Flash Next completion: allocation refusals drain fills before every constr e.stream = &stream; e.eos = 1; e.round_graph = .{}; + e.heads = .full; try t.expectError(error.OutOfDeviceMemory, e.newSeq(4)); _ = try cuda.memory.setAllocationLimit(null); try t.expectEqual(deny, Fake.allocations); @@ -119,6 +120,7 @@ test "Flash Next completion: failed asynchronous fill and an interrupted prefill e.stream = &stream; e.eos = 1; e.round_graph = .{}; + e.heads = .full; try t.expectError(error.NotReady, e.newSeq(4)); try t.expectEqual(Fake.allocations, Fake.frees); if (Fake.early != 0) std.log.warn("completion control: allocs={d}, fills={d}, frees={d}, syncs={d}", .{ Fake.allocations, Fake.fills, Fake.frees, Fake.syncs }); @@ -132,6 +134,7 @@ test "Flash Next completion: failed asynchronous fill and an interrupted prefill e.stream = &stream; e.eos = 1; e.round_graph = .{}; + e.heads = .full; var seq = try e.newSeq(4); // Model a queued prefill kernel writing every owned sequence buffer before request refusal. for (Fake.pending[1 .. Fake.allocations + 1]) |*pending| pending.* = true; @@ -145,13 +148,13 @@ test "Flash Next completion: draft constructor fill failures and queued window w for ([_]usize{ 1, 2, 3, 4 }) |fail_fill| { var d = Fake.init(0, fail_fill); const stream: cuda.Stream = .{ .d = &d, .handle = @ptrFromInt(328) }; - try t.expectError(error.NotReady, Draft.init(t.allocator, &d, stream, 4, 8, 1)); + try t.expectError(error.NotReady, Draft.init(t.allocator, &d, stream, 4, 8, 1, .full)); try t.expectEqual(Fake.allocations, Fake.frees); try t.expectEqual(@as(usize, 0), Fake.early); } var d = Fake.init(0, 0); const stream: cuda.Stream = .{ .d = &d, .handle = @ptrFromInt(329) }; - var draft = try Draft.init(t.allocator, &d, stream, 4, 8, 1); + var draft = try Draft.init(t.allocator, &d, stream, 4, 8, 1, .full); @memset(Fake.pending[1 .. Fake.allocations + 1], true); draft.deinit(); try t.expectEqual(Fake.allocations, Fake.frees); diff --git a/zig/src/families/flashnext/cuda_shared_split.zig b/zig/src/families/flashnext/cuda_shared_split.zig new file mode 100644 index 0000000000..5d4b6bfd70 --- /dev/null +++ b/zig/src/families/flashnext/cuda_shared_split.zig @@ -0,0 +1,34 @@ +//! A shared expert keeps full dot products: gather BF16 intermediate rows, then write disjoint output columns. +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const ops = @import("cuda_batch_ops.zig"); +const mixed = @import("cuda_mixed_experts.zig"); + +pub fn run(self: anytype, native: *const mixed.Owned, n: usize) !void { + if (!native.split_shared and native.rank != 0) return; + const s = self.s; + const dims = wt.dims; + const glue = self.tri.own orelse return error.NativeGlueRequired; + const width = if (native.split_shared) wt.moe_w / 2 else wt.moe_w; + try self.k.native_lane.matmul(self.stream, s.mixed, dims, n, native.shared[0].native.?, s.shared_gate, s.native_part); + try self.k.native_lane.matmul(self.stream, s.mixed, dims, n, native.shared[1].native.?, s.shared_up, s.native_part); + try glue.sharedSwiGLU(self.stream, s.shared_gate, s.shared_up, s.shared_act, n, width); + var input = s.shared_act; + if (native.split_shared) { + const c = self.m.collective orelse return error.TensorParallelUnavailable; + try c.gather(self.stream, input, s.rank_gather, n * width); + for (0..2) |rank| try ops.copyCols(self.k, self.stream, s.rank_gather + rank * n * width * 2, width, s.rank_input + rank * width * 2, wt.moe_w, width, n); + input = s.rank_input; + } + const down = native.shared[2].native.?; + if (!native.split_shared and down.n == dims and cuda.qmmf_shared.eligible(self.k.native_lane, n, down)) { + const target = s.y + (wt.slots - 1) * dims * 4; + try self.k.native_lane.matmulF32Strided(self.stream, input, wt.moe_w, n, down, target, wt.slots * dims); + } else { + try self.k.native_lane.matmulF32(self.stream, input, wt.moe_w, n, down, s.shared_out, s.native_part); + if (native.split_shared) { + // Every routed slot and the peer's output columns stay at the zeros set before the expert launches. + try ops.copyCols(self.k, self.stream, s.shared_out, dims, s.y + (wt.slots - 1) * dims * 4 + native.rank * dims * 2, wt.slots * dims * 2, dims, n); + } else try glue.sharedScatter(self.stream, s.shared_out, s.y, n, dims, wt.slots); + } +} diff --git a/zig/src/families/flashnext/cuda_snap.zig b/zig/src/families/flashnext/cuda_snap.zig index 7db8fe4043..0045d74fbd 100644 --- a/zig/src/families/flashnext/cuda_snap.zig +++ b/zig/src/families/flashnext/cuda_snap.zig @@ -4,26 +4,28 @@ const std = @import("std"); const cuda = @import("cuda"); const prompt = @import("cuda_prompt.zig"); const wt = @import("cuda_prompt_wt.zig"); -const gdn = @import("cuda_gdn.zig"); const Draft = @import("cuda_mtp_draft.zig").Draft; const pages = @import("cuda_pages.zig"); +const sizes = @import("cuda_seq_bytes.zig"); +const Heads = @import("cuda_heads.zig").Heads; -const conv_bytes = (gdn.taps - 1) * gdn.conv_dim * 2; -const state_bytes = wt.state_n * 4; const index_row = wt.index_dim * 2; const tail_bytes = wt.ple_tail_n * wt.wide * 2; +/// The MTP head's cache: its own layer, whole on every rank. +const head_heads: Heads = .full; -/// A state's shape for its disk image (cuda_disk.zig): layer counts and each piece's bytes a row. -pub const shape = .{ .linear = wt.linear_n, .attn = wt.attn_n, .conv = conv_bytes, .state = state_bytes, .kv_row = prompt.kv_row_bytes, .index_row = index_row, .ratio = wt.index_ratio, .tail = tail_bytes }; +/// A state's shape for its disk image (cuda_disk.zig, one rank, the whole layer): layer counts and bytes a row. +pub const shape = .{ .linear = wt.linear_n, .attn = wt.attn_n, .conv = sizes.convBytes(.full), .state = sizes.stateBytes(.full), .kv_row = Heads.full.kvRow(), .index_row = index_row, .ratio = wt.index_ratio, .tail = tail_bytes }; -/// One attention cache's first `at` key, value and indexer-key rows, and the pooled blocks they complete. -fn attnBytes(at: usize) usize { - return at * (2 * prompt.kv_row_bytes + index_row) + at / wt.index_ratio * index_row; +/// One attention cache's first `at` key, value and indexer-key rows at heads `h`, and the pooled blocks they complete. +fn attnBytes(at: usize, h: Heads) usize { + return at * (2 * h.kvRow() + index_row) + at / wt.index_ratio * index_row; } -/// Device bytes a state after `at` tokens takes: the DeltaNet's are fixed, each attention cache grows by the token. -pub fn bytes(at: usize, head: bool) usize { - return wt.linear_n * (conv_bytes + state_bytes) + (wt.attn_n + @intFromBool(head)) * attnBytes(at) + tail_bytes; +/// Device bytes a state at the model's heads `h` after `at` tokens takes: fixed DeltaNet pieces, attention by token. +pub fn bytes(at: usize, head: bool, h: Heads) usize { + const own = wt.linear_n * (sizes.convBytes(h) + sizes.stateBytes(h)) + wt.attn_n * attnBytes(at, h) + tail_bytes; + return own + if (head) attnBytes(at, head_heads) else 0; } /// One kept state; `head`: it holds the MTP head's cache over the same rows, which a drafting stream needs. @@ -33,6 +35,8 @@ pub const Snap = struct { at: usize, history: [2]i64, head: bool, + /// The model's heads its pieces hold (the saved sequence's): a sequence of other heads cannot take it. + heads: Heads = .full, }; const Way = enum { save, restore }; @@ -40,22 +44,22 @@ const Way = enum { save, restore }; /// A piece `move` copies: `layers` buffers of `n` bytes, rows of `row` bytes (0: whole) a `per`-position block each. pub const Piece = struct { layers: usize, n: usize, row: usize = 0, per: usize = 1 }; -/// `move`'s pieces in its order at `at` (the disk image names the same): DeltaNet, attention, PLE tail, `head`'s four. -pub fn order(at: usize, head: bool) struct { p: [11]Piece, len: usize } { +/// `move`'s pieces in order at `at` and the model's heads `h` (the disk image's at full): DeltaNet, attention, tail. +pub fn order(at: usize, head: bool, h: Heads) struct { p: [11]Piece, len: usize } { var o: [11]Piece = undefined; - o[0] = .{ .layers = wt.linear_n, .n = conv_bytes }; - o[1] = .{ .layers = wt.linear_n, .n = state_bytes }; - o[2..6].* = attnOrder(wt.attn_n, at); + o[0] = .{ .layers = wt.linear_n, .n = sizes.convBytes(h) }; + o[1] = .{ .layers = wt.linear_n, .n = sizes.stateBytes(h) }; + o[2..6].* = attnOrder(wt.attn_n, at, h); o[6] = .{ .layers = 1, .n = tail_bytes }; - if (head) o[7..11].* = attnOrder(1, at); + if (head) o[7..11].* = attnOrder(1, at, head_heads); return .{ .p = o, .len = if (head) 11 else 7 }; } /// An attention cache's pieces for `at` rows: complete pooled blocks only (the rest pool again as later rows land). -fn attnOrder(layers: usize, at: usize) [4]Piece { +fn attnOrder(layers: usize, at: usize, h: Heads) [4]Piece { return .{ - .{ .layers = layers, .n = at * prompt.kv_row_bytes, .row = prompt.kv_row_bytes }, - .{ .layers = layers, .n = at * prompt.kv_row_bytes, .row = prompt.kv_row_bytes }, + .{ .layers = layers, .n = at * h.kvRow(), .row = h.kvRow() }, + .{ .layers = layers, .n = at * h.kvRow(), .row = h.kvRow() }, .{ .layers = layers, .n = at * index_row, .row = index_row }, .{ .layers = layers, .n = at / wt.index_ratio * index_row, .row = index_row, .per = wt.index_ratio }, }; @@ -65,7 +69,7 @@ fn attnOrder(layers: usize, at: usize) [4]Piece { fn move(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, base: u64, at: usize, way: Way) !void { const dr = stream.d; var off: u64 = 0; - const o = order(at, d != null); + const o = order(at, d != null, seq.heads); var bufs: [11][]cuda.DeviceBuffer = undefined; bufs[0..7].* = .{ seq.conv, seq.state, seq.kc, seq.vc, seq.ikc, seq.pooled, (&seq.tail)[0..1] }; if (d) |h| bufs[7..11].* = .{ (&h.kc)[0..1], (&h.vc)[0..1], (&h.ikc)[0..1], (&h.pooled)[0..1] }; @@ -99,7 +103,7 @@ fn move(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, base: u64, at: usize, off += p.n; } } - std.debug.assert(off == bytes(at, d != null)); + std.debug.assert(off == bytes(at, d != null, seq.heads)); } /// A copy of `seq` (and of the head's cache, with `d`) as it stands after `seq.fed` tokens. @@ -108,7 +112,7 @@ pub fn save(gpa: std.mem.Allocator, driver: *const cuda.Driver, stream: cuda.Str if (d) |h| if (h.len != at or h.drafted != 0) return error.NotAtMark; const snap = try gpa.create(Snap); errdefer gpa.destroy(snap); - snap.* = .{ .buf = try cuda.DeviceBuffer.alloc(driver, bytes(at, d != null)), .at = at, .history = seq.history, .head = d != null }; + snap.* = .{ .buf = try cuda.DeviceBuffer.alloc(driver, bytes(at, d != null, seq.heads)), .at = at, .history = seq.history, .head = d != null, .heads = seq.heads }; errdefer snap.buf.free(); errdefer cuda.completion.drain(driver, stream, "a failed snapshot save did not drain"); try move(stream, seq, d, snap.buf.ptr, at, .save); @@ -119,6 +123,7 @@ pub fn save(gpa: std.mem.Allocator, driver: *const cuda.Driver, stream: cuda.Str /// `seq` (and the head's cache, with `d`) becomes `snap`'s state; refused before any copy when it cannot hold it. pub fn restore(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, snap: *const Snap) !void { if (snap.at > seq.cap) return error.SequenceTooShort; + if (!snap.heads.eql(seq.heads)) return error.SnapShape; if (d) |h| if (!snap.head or snap.at > h.cap) return error.NoHeadState; // a state with the head's rows serves a stream without drafts too: those rows stay in the buffer try move(stream, seq, d, snap.buf.ptr, snap.at, .restore); @@ -138,12 +143,40 @@ pub fn drop(gpa: std.mem.Allocator, stream: cuda.Stream, snap: *Snap) void { } test "a state's bytes: about 110 MiB of DeltaNet states and the n-gram tail, then 27 KiB a token" { - try std.testing.expectEqual(@as(usize, 12 * (2 * 2 * 256 * 2 + 128 * 2)), bytes(1, false) - bytes(0, false)); + try std.testing.expectEqual(@as(usize, 12 * (2 * 2 * 256 * 2 + 128 * 2)), bytes(1, false, .full) - bytes(0, false, .full)); // a complete block of 4 adds its pooled row - try std.testing.expectEqual(@as(usize, 12 * (4 * (2 * 2 * 256 * 2 + 128 * 2) + 128 * 2)), bytes(4, false) - bytes(0, false)); + try std.testing.expectEqual(@as(usize, 12 * (4 * (2 * 2 * 256 * 2 + 128 * 2) + 128 * 2)), bytes(4, false, .full) - bytes(0, false, .full)); // the head's cache is one more attention layer's rows - try std.testing.expectEqual(@as(usize, 13 * (2 * 2 * 256 * 2 + 128 * 2)), bytes(1, true) - bytes(0, true)); - try std.testing.expect(bytes(0, false) > 100 << 20 and bytes(0, false) < 120 << 20); + try std.testing.expectEqual(@as(usize, 13 * (2 * 2 * 256 * 2 + 128 * 2)), bytes(1, true, .full) - bytes(0, true, .full)); + try std.testing.expect(bytes(0, false, .full) > 100 << 20 and bytes(0, false, .full) < 120 << 20); +} + +test "a state at a rank's half of the heads: half the conv windows, states and KV rows, the indexer and the head's whole" { + const h = Heads.half; + try std.testing.expectEqual(@as(usize, 36 * (3 * 5120 * 2 + 24 * 128 * 128 * 4) + 9 * 4 * 2560 * 2), bytes(0, false, h)); + try std.testing.expectEqual(@as(usize, 12 * (2 * 256 * 2 + 128 * 2)), bytes(1, false, h) - bytes(0, false, h)); + // the head's cache keeps its two KV heads beside the model's one + try std.testing.expectEqual(@as(usize, 12 * (2 * 256 * 2 + 128 * 2) + (2 * 2 * 256 * 2 + 128 * 2)), bytes(1, true, h) - bytes(0, true, h)); + const o = order(8, true, h); + try std.testing.expectEqual(@as(usize, 3 * 5120 * 2), o.p[0].n); + try std.testing.expectEqual(@as(usize, 24 * 128 * 128 * 4), o.p[1].n); + for (o.p[2..4]) |p| try std.testing.expectEqual(@as(usize, 512), p.row); + for (o.p[7..9]) |p| try std.testing.expectEqual(@as(usize, 1024), p.row); + var sum: usize = 0; + for (o.p[0..o.len]) |p| sum += p.layers * p.n; + try std.testing.expectEqual(bytes(8, true, h), sum); +} + +test "a state of the whole layer restores into no sequence of half the heads, and the reverse" { + var driver: cuda.Driver = undefined; + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + for ([_][2]Heads{ .{ .full, .half }, .{ .half, .full } }) |pair| { + const kept: Snap = .{ .buf = .{ .d = &driver, .ptr = 0x100, .len = 1, .borrowed = true }, .at = 1, .history = .{ 0, 0 }, .head = false, .heads = pair[0] }; + var seq: prompt.Seq = .{ .heads = pair[1], .conv = &.{}, .state = &.{}, .kc = &.{}, .vc = &.{}, .ikc = &.{}, .pooled = &.{}, .tail = .{ .d = &driver, .ptr = 0, .len = 0 }, .history = .{ 0, 0 }, .fed = 0, .logits = &.{}, .cap = 8 }; + // refused before any copy: no driver call is made + try std.testing.expectError(error.SnapShape, restore(stream, &seq, null, &kept)); + try std.testing.expectEqual(@as(usize, 0), seq.fed); + } } test { diff --git a/zig/src/families/flashnext/cuda_split_lane.zig b/zig/src/families/flashnext/cuda_split_lane.zig new file mode 100644 index 0000000000..743134e4a4 --- /dev/null +++ b/zig/src/families/flashnext/cuda_split_lane.zig @@ -0,0 +1,32 @@ +//! Exact output-column projections: gather BF16 heads, compute full dot products, then gather FP32 outputs. +const wt = @import("cuda_prompt_wt.zig"); +const ops = @import("cuda_batch_ops.zig"); +const Layer = @import("cuda_prompt_load.zig").Layer; + +const dims = wt.dims; +const out_sk = wt.out_sk; + +/// Write full FP32 output rows into part without adding partial products across ranks. +pub fn out(self: anytype, x: u64, xs: u64, lw: *const Layer) !usize { + const s = self.s; + const n = self.rows; + const cut = lw.dense_cut orelse return self.lane(x, xs, lw.out, s.dummy, s.part, n, out_sk); + const native = lw.out.native orelse return error.UnexpectedTensor; + const c = self.m.collective orelse return error.TensorParallelUnavailable; + if (n == 0 or n > s.cap or !lw.heads.eql(.half) or cut.width() != dims / 2 or + cut.from != c.rank * (dims / 2) or native.n != dims / 2 or native.k != wt.out_k or + native.sk != @import("cuda").qmmf.splitK(dims, wt.out_k)) return error.UnexpectedTensor; + const width = lw.heads.valueDim(); + try c.gather(self.stream, x, s.rank_gather, n * width); + for (0..2) |rank| try ops.copyCols(self.k, self.stream, s.rank_gather + rank * n * width * 2, width, s.rank_input + rank * width * 2, 2 * width, width, n); + try self.k.native_lane.matmulF32(self.stream, s.rank_input, 2 * width, n, native, s.rank_output, s.native_part); + try c.gatherF32(self.stream, s.rank_output, s.rank_result, n * cut.width()); + // copyCols moves 16-bit words, so two words preserve every FP32 value's bits. + for (0..2) |rank| try ops.copyCols(self.k, self.stream, s.rank_result + rank * n * cut.width() * 4, cut.width() * 2, s.part + rank * cut.width() * 4, dims * 2, cut.width() * 2, n); + return 1; +} + +test { + _ = @import("cuda_dense_split.zig"); + if (@import("builtin").os.tag != .windows) _ = @import("cuda_split_lane_test.zig"); +} diff --git a/zig/src/families/flashnext/cuda_split_lane_test.zig b/zig/src/families/flashnext/cuda_split_lane_test.zig new file mode 100644 index 0000000000..cba09243e8 --- /dev/null +++ b/zig/src/families/flashnext/cuda_split_lane_test.zig @@ -0,0 +1,318 @@ +//! Exact rank gathers over recorded launches, with the unsplit lane as the control. +const std = @import("std"); +const builtin = @import("builtin"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const parts = @import("cuda_batch_state.zig"); +const batch = @import("cuda_batch.zig"); +const split_lane = @import("cuda_split_lane.zig"); +const Layer = @import("cuda_prompt_load.zig").Layer; +const Cut = @import("cuda_modelopt_experts.zig").Cut; +const Collective = @import("cuda_collective.zig").Collective; +const nccl = cuda.nccl; + +const dims = wt.dims; +const out_k = wt.out_k; +const qmmf_id = 0x100; +const reduce_id = 0x200; +const legacy_id = 0x300; +const copy_id = 0x400; +const strided_id = 0x500; +const x_at: u64 = 0x10_0000; +const part_at: u64 = 0x2000_0000; +const native_part_at: u64 = 0x4000_0000; + +/// One recorded call: a launch's function and the arguments that place it, a memset, or a rank sum. +const Call = struct { kind: enum { launch, memset, sum, gather }, f: usize = 0, a: [6]i64 = @splat(0) }; + +const Rec = struct { + var calls: [64]Call = undefined; + var n: usize = 0; + + fn arg(comptime T: type, args: ?[*]?*anyopaque, i: usize) i64 { + const v = @as(*const T, @ptrCast(@alignCast(args.?[i].?))).*; + return if (T == u64) @bitCast(v) else @intCast(v); + } + + fn push(c: Call) void { + calls[n] = c; + n += 1; + } + + fn launch(f: cuda.abi.Function, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: cuda.abi.Stream, args: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) c_int { + const id = @intFromPtr(f); + push(.{ + .kind = .launch, + .f = id, + .a = switch (id) { + // x, y, part, m, n and k as n * 10^6 + k, and the row stride a qmmf tile reads + qmmf_id => .{ arg(u64, args, 0), arg(u64, args, 4), arg(u64, args, 5), arg(i32, args, 6), arg(i32, args, 7) * 1_000_000 + arg(i32, args, 8), arg(i32, args, 11) }, + strided_id => .{ arg(u64, args, 0), arg(u64, args, 3), arg(i32, args, 4), arg(i32, args, 5) * 1_000_000 + arg(i32, args, 6), arg(i32, args, 10), arg(i32, args, 8) }, + copy_id => .{ arg(u64, args, 0), arg(i64, args, 1), arg(u64, args, 2), arg(i64, args, 3), arg(i32, args, 4), 0 }, + reduce_id => .{ arg(u64, args, 0), arg(u64, args, 1), arg(i64, args, 2), arg(i32, args, 3), 0, 0 }, + else => .{ arg(u64, args, 0), arg(u64, args, 6), arg(c_int, args, 7), arg(c_int, args, 8), arg(c_int, args, 9), arg(c_int, args, 10) }, + }, + }); + return 0; + } + + fn memset(ptr: u64, _: u8, count: usize, _: cuda.abi.Stream) callconv(.c) c_int { + push(.{ .kind = .memset, .a = .{ @bitCast(ptr), @intCast(count), 0, 0, 0, 0 } }); + return 0; + } + + fn allReduce(send: cuda.abi.DevicePtr, recv: cuda.abi.DevicePtr, count: usize, dt: nccl.DataType, op: nccl.RedOp, _: nccl.Comm, _: cuda.abi.Stream) callconv(.c) c_int { + push(.{ .kind = .sum, .a = .{ @bitCast(send), @bitCast(recv), @intCast(count), @backingInt(dt), @backingInt(op), 0 } }); + return 0; + } + + fn gather(send: u64, recv: u64, count: usize, dt: nccl.DataType, _: nccl.Comm, _: cuda.abi.Stream) callconv(.c) c_int { + push(.{ .kind = .gather, .a = .{ @bitCast(send), @bitCast(recv), @intCast(count), @backingInt(dt), 0, 0 } }); + return 0; + } + + fn asyncError(_: nccl.Comm, state: *c_int) callconv(.c) c_int { + state.* = 0; + return 0; + } + + fn abort(_: nccl.Comm) callconv(.c) c_int { + return 0; + } +}; + +const Rig = struct { + d: cuda.Driver, + k: parts.Kernels, + s: parts.Scratch, + c: Collective, + pair: [2]c_int, + + fn init(r: *Rig) !void { + r.d = undefined; + r.d.api.cuLaunchKernel = Rec.launch; + r.d.api.cuMemsetD8Async = Rec.memset; + const f = struct { + fn of(d: *const cuda.Driver, id: usize) cuda.Function { + return .{ .d = d, .handle = @ptrFromInt(id) }; + } + }.of; + r.k = undefined; + r.k.copy_cols = f(&r.d, copy_id); + r.k.lane = f(&r.d, legacy_id); + r.k.tiles.fact = f(&r.d, legacy_id + 1); + r.k.tiles.fact_fold = f(&r.d, legacy_id + 2); + @memset(std.mem.asBytes(&r.k.native_lane), 0); + r.k.native_lane.fns = @splat(@splat(f(&r.d, qmmf_id))); + r.k.native_lane.f32_fns = @splat(@splat(f(&r.d, qmmf_id))); + r.k.native_lane.reduce = @splat(f(&r.d, reduce_id)); + r.k.native_lane.clusters = false; + r.s = undefined; + r.s.cap = 256; + r.s.part = part_at; + r.s.native_part = native_part_at; + r.s.dummy = 0x99; + r.s.part_zero = true; + r.s.rank_input = 0x5000_0000; + r.s.rank_gather = 0x6000_0000; + r.s.rank_output = 0x7000_0000; + r.s.rank_result = 0x8000_0000; + // the peer's end of a socket pair stays open, so the sum's peer check passes + if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &r.pair) != 0) return error.SocketPair; + r.c = undefined; + r.c.rank = 0; + r.c.failed = false; + r.c.comm = @ptrFromInt(0x1000); + r.c.link = .{ .fd = r.pair[0] }; + r.c.timeout_ms = 1000; + r.c.lib.api.ncclAllReduce = Rec.allReduce; + r.c.lib.api.ncclAllGather = Rec.gather; + r.c.api.ncclCommGetAsyncError = Rec.asyncError; + r.c.api.ncclCommAbort = Rec.abort; + Rec.n = 0; + } + + fn deinit(r: *Rig) void { + for (r.pair) |fd| _ = std.c.close(fd); + } + + fn pass(r: *Rig, rows: usize, collective: bool) batch.Pass(void) { + const m: parts.Model = .{ .collective = if (collective) &r.c else null, .table_w = 0, .table_s = 0, .table_b = 0, .layers = &.{}, .mixer = undefined, .head = undefined, .ngram = undefined }; + return .{ .gpa = std.testing.allocator, .io = undefined, .stream = .{ .d = &r.d, .handle = null }, .tri = {}, .k = &r.k, .s = &r.s, .m = m, .seq = undefined, .hooks = undefined, .rows = rows }; + } + + fn tryPin(w: cuda.qmmf.Weight) cuda.qmmf.Weight { + return cuda.qmmf.pinned(w, cuda.qmmf.splitK(dims, out_k)) catch unreachable; + } + + fn layer(r: *Rig, k: usize, native: bool, cut: ?Cut) Layer { + const z: cuda.DeviceBuffer = .{ .d = &r.d, .ptr = 0, .len = 0 }; + var lw: Layer = undefined; + lw.heads = if (cut != null) .half else .full; + lw.out = .{ .w = z, .s = z, .b = z, .n = dims, .k = k, .native = if (native) cuda.modelopt.bf16Weight(0x2000, dims, k) else null }; + if (cut != null) { + lw.out.n = dims / 2; + if (native) lw.out.native = tryPin(cuda.modelopt.bf16Weight(0x2000, dims / 2, k)); + } + lw.dense_cut = cut; + return lw; + } +}; + +fn recorded() []const Call { + return Rec.calls[0..Rec.n]; +} + +test "split out lane gathers full BF16 inputs and disjoint FP32 outputs without a reduction" { + if (builtin.os.tag == .windows) return error.SkipZigTest; + var r: Rig = undefined; + try r.init(); + defer r.deinit(); + for ([_]usize{ 1, 4, 17, 33, 64, 256 }) |n| for (0..2) |rank| { + r.c.rank = rank; + const cut = try Cut.init(dims, 2, rank); + const lw = r.layer(out_k, true, cut); + var p = r.pass(n, true); + Rec.n = 0; + try std.testing.expectEqual(@as(usize, 1), try split_lane.out(&p, x_at, 0x20, &lw)); + const calls = recorded(); + try std.testing.expect(calls.len == 7 or calls.len == 8); + try std.testing.expectEqual(Call{ .kind = .gather, .a = .{ x_at, @intCast(r.s.rank_gather), @intCast(n * out_k / 2), @backingInt(nccl.DataType.bf16), 0, 0 } }, calls[0]); + for (0..2) |peer| { + const c = calls[1 + peer]; + try std.testing.expectEqual(copy_id, c.f); + try std.testing.expectEqualSlices(i64, &.{ @intCast(r.s.rank_gather + peer * n * out_k), out_k / 2, @intCast(r.s.rank_input + peer * out_k), out_k, out_k / 2, 0 }, &c.a); + } + const q = calls[3]; + try std.testing.expectEqual(qmmf_id, q.f); + try std.testing.expectEqual(@as(i64, @intCast(r.s.rank_input)), q.a[0]); + try std.testing.expectEqual(@as(i64, 1280 * 1_000_000 + 6144), q.a[4]); + try std.testing.expectEqual(@as(i64, out_k), q.a[5]); + const at = calls.len - 3; + try std.testing.expectEqual(Call{ .kind = .gather, .a = .{ @intCast(r.s.rank_output), @intCast(r.s.rank_result), @intCast(n * dims / 2), @backingInt(nccl.DataType.f32), 0, 0 } }, calls[at]); + for (0..2) |peer| { + const c = calls[at + 1 + peer]; + try std.testing.expectEqual(copy_id, c.f); + try std.testing.expectEqualSlices(i64, &.{ @intCast(r.s.rank_result + peer * n * dims * 2), dims, @intCast(part_at + peer * dims * 2), dims * 2, dims, 0 }, &c.a); + } + for (calls) |c| try std.testing.expect(c.kind != .sum and c.kind != .memset); + }; +} + +test "split out lane: an unsplit layer launches exactly what Pass.lane launches, with or without a collective" { + if (builtin.os.tag == .windows) return error.SkipZigTest; + var r: Rig = undefined; + try r.init(); + defer r.deinit(); + var want: [64]Call = undefined; + for ([_]usize{ 1, 4, 64, 256 }) |n| for ([_]bool{ true, false }) |native| for ([_]bool{ true, false }) |collective| { + const lw = r.layer(out_k, native, null); + var p = r.pass(n, collective); + Rec.n = 0; + const lane_slices = try p.lane(x_at, 0x20, lw.out, r.s.dummy, r.s.part, n, wt.out_sk); + const count = Rec.n; + @memcpy(want[0..count], recorded()); + Rec.n = 0; + var q = r.pass(n, collective); + try std.testing.expectEqual(lane_slices, try split_lane.out(&q, x_at, 0x20, &lw)); + try std.testing.expectEqual(count, Rec.n); + for (want[0..count], recorded()) |a, b| { + try std.testing.expectEqual(a.kind, b.kind); + try std.testing.expectEqual(a.f, b.f); + try std.testing.expectEqualSlices(i64, &a.a, &b.a); + try std.testing.expect(b.kind != .sum); + } + }; +} + +test "split out lane: a split layer refuses a missing collective or a non-native weight" { + if (builtin.os.tag == .windows) return error.SkipZigTest; + var r: Rig = undefined; + try r.init(); + defer r.deinit(); + const cut = try Cut.init(dims, 2, 0); + const native = r.layer(out_k, true, cut); + const legacy = r.layer(out_k, false, cut); + var p = r.pass(4, false); + try std.testing.expectError(error.TensorParallelUnavailable, split_lane.out(&p, x_at, 0x20, &native)); + var q = r.pass(4, true); + try std.testing.expectError(error.UnexpectedTensor, split_lane.out(&q, x_at, 0x20, &legacy)); + try std.testing.expectEqual(@as(usize, 0), Rec.n); +} + +const SharedGlue = struct { + pub fn sharedSwiGLU(_: *SharedGlue, _: cuda.Stream, _: u64, _: u64, _: u64, _: usize, _: usize) !void {} + pub fn sharedScatter(_: *SharedGlue, _: cuda.Stream, input: u64, output: u64, rows: usize, width: usize, slots: usize) !void { + Rec.push(.{ .kind = .launch, .f = 0x600, .a = .{ @intCast(input), @intCast(output), @intCast(rows), @intCast(width), @intCast(slots), 0 } }); + } +}; + +test "shared expert composition keeps whole direct slots and split full-K output copies" { + if (builtin.os.tag == .windows) return error.SkipZigTest; + var r: Rig = undefined; + try r.init(); + defer r.deinit(); + var glue: SharedGlue = .{}; + r.s.mixed = x_at; + r.s.shared_gate = 0x9000_0000; + r.s.shared_up = 0x9100_0000; + r.s.shared_act = 0x9200_0000; + r.s.shared_out = 0x9300_0000; + r.s.y = 0x9400_0000; + for (&r.k.native_lane.shared_down) |*f| f.* = .{ .d = &r.d, .handle = @ptrFromInt(strided_id) }; + const Self = struct { + k: *parts.Kernels, + s: *parts.Scratch, + m: parts.Model, + stream: cuda.Stream, + tri: struct { own: ?*SharedGlue }, + }; + for ([_]bool{ false, true }) |split| for ([_]bool{ false, true }) |prompt_loaded| for ([_]usize{ 64, 512 }) |n| for (0..2) |rank| { + r.c.rank = rank; + r.k.native_lane.prompt = if (prompt_loaded) @splat(@splat(cuda.Function{ .d = &r.d, .handle = @ptrFromInt(qmmf_id) })) else null; + var native: @import("cuda_mixed_experts.zig").Owned = undefined; + native.rank = rank; + native.split_shared = split; + const width: usize = if (split) wt.moe_w / 2 else wt.moe_w; + const columns: usize = if (split) dims / 2 else dims; + for (0..3) |i| native.shared[i].native = cuda.modelopt.bf16Weight(0x2000, if (i == 2) columns else width, if (i == 2) wt.moe_w else dims); + var self: Self = .{ .k = &r.k, .s = &r.s, .m = r.pass(n, true).m, .stream = .{ .d = &r.d, .handle = null }, .tri = .{ .own = &glue } }; + Rec.n = 0; + try @import("cuda_shared_split.zig").run(&self, &native, n); + if (!split and rank != 0) { + try std.testing.expectEqual(@as(usize, 0), Rec.n); + continue; + } + const direct = !split and !(prompt_loaded and n >= cuda.qmmf.prompt_rows); + const last = recorded()[Rec.n - 1]; + if (direct) { + try std.testing.expectEqual(strided_id, last.f); + try std.testing.expectEqualSlices(i64, &.{ @intCast(r.s.shared_act), @intCast(r.s.y + (wt.slots - 1) * dims * 4), @intCast(n), dims * 1_000_000 + wt.moe_w, wt.slots * dims, wt.moe_w }, &last.a); + } else { + const down = recorded()[Rec.n - 2]; + try std.testing.expectEqual(qmmf_id, down.f); + try std.testing.expectEqual(@as(i64, @intCast(columns * 1_000_000 + wt.moe_w)), down.a[4]); + try std.testing.expectEqual(@as(i64, wt.moe_w), down.a[5]); + if (split) { + try std.testing.expectEqual(copy_id, last.f); + try std.testing.expectEqualSlices(i64, &.{ @intCast(r.s.shared_out), dims, @intCast(r.s.y + (wt.slots - 1) * dims * 4 + rank * dims * 2), wt.slots * dims * 2, dims, 0 }, &last.a); + var gather_at: ?usize = null; + for (recorded(), 0..) |call, i| if (call.kind == .gather) { + try std.testing.expect(gather_at == null); + gather_at = i; + }; + const at = gather_at orelse return error.TestExpectedGather; + try std.testing.expectEqual(at + 5, Rec.n); + try std.testing.expectEqual(Call{ .kind = .gather, .a = .{ @intCast(r.s.shared_act), @intCast(r.s.rank_gather), @intCast(n * width), @backingInt(nccl.DataType.bf16), 0, 0 } }, recorded()[at]); + for (0..2) |peer| { + const copy = recorded()[at + 1 + peer]; + try std.testing.expectEqual(copy_id, copy.f); + try std.testing.expectEqualSlices(i64, &.{ @intCast(r.s.rank_gather + peer * n * width * 2), @intCast(width), @intCast(r.s.rank_input + peer * width * 2), wt.moe_w, @intCast(width), 0 }, ©.a); + } + try std.testing.expectEqual(@as(i64, @intCast(r.s.rank_input)), down.a[0]); + try std.testing.expectEqual(@as(i64, @intCast(r.s.shared_out)), down.a[1]); + for (recorded()) |c| try std.testing.expect(c.f != strided_id and c.kind != .sum); + } else try std.testing.expectEqual(@as(usize, 0x600), last.f); + } + }; +} diff --git a/zig/src/families/flashnext/cuda_train_buf.zig b/zig/src/families/flashnext/cuda_train_buf.zig index 6ea4416373..0600f6be73 100644 --- a/zig/src/families/flashnext/cuda_train_buf.zig +++ b/zig/src/families/flashnext/cuda_train_buf.zig @@ -48,7 +48,7 @@ pub fn memory() usize { n += pack_rows * vocab * 2; // every row's logits, then their gradient n += scratchBytes(); n += weightBytes(); - n += prompt.seqBytes(max_rows); + n += prompt.seqBytes(max_rows, .full); // training runs at one rank: the whole layer return n + (256 << 20); } diff --git a/zig/src/families/flashnext/heads_cut.zig b/zig/src/families/flashnext/heads_cut.zig new file mode 100644 index 0000000000..1da3348044 --- /dev/null +++ b/zig/src/families/flashnext/heads_cut.zig @@ -0,0 +1,102 @@ +//! What one rank keeps of a layer under the head split: stitched in-projection rows, conv channels and per-head values. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const gdn = @import("cuda_gdn.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const Heads = @import("cuda_heads.zig").Heads; +const load = @import("cuda_projection_load.zig"); +const Cut = @import("cuda_modelopt_experts.zig").Cut; +const Ranges = load.Ranges; + +/// The stitched parts of each class, in the order their rows are kept. +pub const delta_parts = [_][]const u8{ ".in_proj_qkv", ".in_proj_z", ".in_proj_b", ".in_proj_a" }; +pub const attention_parts = [_][]const u8{ ".q_proj", ".k_proj", ".v_proj", ".indexer.index_qk_proj" }; + +/// One rank's cut of a stitched in-projection: each part's kept rows, the local width, the whole shape's K slices. +pub const Stitch = struct { ranges: [4]Ranges, n: usize, sk: u32 }; + +/// `h` divides every count of the full heads by one world, and `rank` is one of its ranks. +fn check(h: Heads, rank: usize) error{UnevenHeadSplit}!void { + const f = Heads.full; + if (h.nk == 0 or f.nk % h.nk != 0) return error.UnevenHeadSplit; + const w: usize = f.nk / h.nk; + if (f.nv != w * h.nv or f.q != w * h.q or f.kv != w * h.kv or rank >= w) return error.UnevenHeadSplit; +} + +fn span(from: usize, width: usize) Cut { + return .{ .from = from, .to = from + width }; +} + +/// The q, k and v channels of the rank's key and value heads, in the layout `in_proj_qkv` and conv1d share. +pub fn convChannels(h: Heads, rank: usize) !Ranges { + try check(h, rank); + const qk = h.nk * gdn.dk; + const v = h.nv * gdn.dv; + const whole_qk = Heads.full.nk * gdn.dk; + return Ranges.of(&.{ span(rank * qk, qk), span(whole_qk + rank * qk, qk), span(2 * whole_qk + rank * v, v) }); +} + +/// The rank's value heads, the rows of `in_proj_b` and `in_proj_a` and the entries of `A_log` and `dt_bias`. +pub fn valueHeads(h: Heads, rank: usize) !Cut { + try check(h, rank); + return span(rank * h.nv, h.nv); +} + +/// The DeltaNet in-projection: q, k and v channels, z of each value head, then b and a. +pub fn delta(h: Heads, rank: usize) !Stitch { + const vh = try valueHeads(h, rank); + const z = h.nv * gdn.dv; + const s: Stitch = .{ + .ranges = .{ try convChannels(h, rank), Ranges.of(&.{span(rank * z, z)}), Ranges.of(&.{vh}), Ranges.of(&.{vh}) }, + .n = h.projWidth(), + .sk = cuda.qmmf.splitK(Heads.full.projWidth(), wt.dims), + }; + std.debug.assert(rows(s.ranges) == s.n); + return s; +} + +/// The attention in-projection: each query head's q and gate rows, k and v of the KV head, the indexer whole. +pub fn attention(h: Heads, rank: usize) !Stitch { + try check(h, rank); + const q = h.q * 2 * wt.head_dim; + const kv = h.kv * wt.head_dim; + const s: Stitch = .{ + .ranges = .{ + Ranges.of(&.{span(rank * q, q)}), + Ranges.of(&.{span(rank * kv, kv)}), + Ranges.of(&.{span(rank * kv, kv)}), + Ranges.of(&.{span(0, wt.index_heads * wt.index_dim + wt.index_dim)}), + }, + .n = h.projN(), + .sk = cuda.qmmf.splitK(Heads.full.projN(), wt.dims), + }; + std.debug.assert(rows(s.ranges) == s.n); + return s; +} + +fn rows(ranges: [4]Ranges) usize { + var n: usize = 0; + for (ranges) |r| n += r.rows(); + return n; +} + +/// The rank's channels of a BF16 conv1d weight (channel-major, `taps` per channel), copied before the store moves on. +pub fn convRows(a: std.mem.Allocator, t: core.safetensors.Tensor, h: Heads, rank: usize) ![]u8 { + const kept = try convChannels(h, rank); + const row = gdn.taps * 2; + if (t.dtype != .bf16 or t.rank == 0 or t.dim(0) != Heads.full.convDim() or t.bytes.len != t.dim(0) * row) + return error.UnexpectedTensor; + const out = try a.alloc(u8, h.convDim() * row); + errdefer a.free(out); + if (try load.copyRanges(out, t.bytes, row, kept) != out.len) return error.UnexpectedTensor; + return out; +} + +/// The rank's entries of a per-value-head vector (`A_log`, `dt_bias`) of the full heads. +pub fn valueSlice(values: []const f32, h: Heads, rank: usize) ![]const f32 { + if (values.len != Heads.full.nv) return error.UnexpectedTensor; + const c = try valueHeads(h, rank); + return values[c.from..c.to]; +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 1927709a92..7e721e36f3 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -34,7 +34,7 @@ const registry = .{ nemotron.native, nemotron.omni, glm.native, minimax.native, /// --record opens here (engines.zig refuses it for a backend without this). pub const records = true; pub const backends: []const []const u8 = &.{"cuda"}; -pub const serves: []const []const u8 = &.{ "--mtp-drafts", "--mtp-confidence", "--policy", "--checkpoint-slots", "--spill-gib" }; +pub const serves: []const []const u8 = &.{ "--tp", "--mtp-drafts", "--mtp-confidence", "--policy", "--checkpoint-slots", "--spill-gib" }; pub const families: []const api.Family = blk: { var out: [registry.len + 4]api.Family = undefined; for (registry, 0..) |F, i| out[i] = .{ .model_type = F.model_type, .formats = F.formats }; diff --git a/zig/src/native/cuda_refusal.zig b/zig/src/native/cuda_refusal.zig index d86d1bbfb2..72b58f9c8d 100644 --- a/zig/src/native/cuda_refusal.zig +++ b/zig/src/native/cuda_refusal.zig @@ -26,12 +26,22 @@ pub fn backend(o: api.Open) ?[]const u8 { if (o.learn != null and !o.learn_conversations) return "--learn learns shared prompt prefixes on Macs only; CUDA keeps conversation states on disk with --learn-conversations: serve with it, or without --learn, --learn-dir or --learn-gib"; if (o.speed_up != null) return "--speed-up links two Macs: serve without --speed-up on CUDA"; if (o.prompt_cache_over_cap) return "--prompt-cache-over-cap lifts a Mac's 70% cap: on CUDA the prompt cache takes at most the memory the weights and streams leave, so serve without it"; + // Fish and Qwen4 open outside family(), so their --tp is refused here + if (std.mem.eql(u8, o.model_type, "fish_omni") or std.mem.eql(u8, o.model_type, "qwen4") or std.mem.eql(u8, o.model_type, "qwen4_text")) + if (unranked(o)) |why| return why; return null; } /// LFM2 reads a checkpoint's pictures through its own tower without --vision: the flag is refused, not ignored. pub fn lfm2(o: api.Open) ?[]const u8 { - return if (o.vision != null) lfm2_vision else null; + return if (o.vision != null) lfm2_vision else unranked(o); +} + +pub const no_ranks = "this CUDA family has no tensor-parallel ranks: serve without --tp"; + +/// --tp for a family that opens outside `family()` (Fish, Qwen4, LFM2): refused, never ignored. +pub fn unranked(o: api.Open) ?[]const u8 { + return if (o.tp != 1) no_ranks else null; } const lfm2_vision = "--vision: LFM2 reads the checkpoint's images through its own tower without the flag; " ++ @@ -42,7 +52,7 @@ pub fn family(comptime F: type, o: api.Open) ?[]const u8 { if (o.vision != null and !(@hasDecl(F, "Options") and @hasField(F.Options, "vision"))) return "--vision: this CUDA family reads no images; serve it without --vision"; const ranked = @hasDecl(F, "Options") and @hasField(F.Options, "rank"); - if (o.tp != 1 and !ranked) return "this CUDA family has no tensor-parallel ranks: serve without --tp"; + if (o.tp != 1 and !ranked) return no_ranks; if (@hasDecl(F, "tp_draft_snapshots")) { if (!F.tp_draft_snapshots and o.tp > 1 and o.drafts and (o.prompt_cache_gib orelse 1) > 0) return "Flash Next tensor-rank MTP snapshots are unavailable: use --prompt-cache-gib 0 with --tp and drafting"; @@ -70,6 +80,23 @@ test "a family must consume rank options before tensor parallelism is admitted" }; try std.testing.expectStringStartsWith(family(Single, .{ .dir = "d", .model_type = "m", .tp = 2 }) orelse "", "this CUDA family has no tensor-parallel ranks"); try std.testing.expect(family(Ranked, .{ .dir = "d", .model_type = "m", .tp = 2 }) == null); + try std.testing.expectEqualStrings(no_ranks, family(Single, .{ .dir = "d", .model_type = "m", .tp = 2 }).?); + // the families past family(): one rank opens, two are refused by the same words + try std.testing.expect(unranked(.{ .dir = "d", .model_type = "qwen4", .tp = 1 }) == null); + try std.testing.expectEqualStrings(no_ranks, unranked(.{ .dir = "d", .model_type = "qwen4", .tp = 2 }).?); + try std.testing.expectEqualStrings(no_ranks, lfm2(.{ .dir = "d", .model_type = "lfm2", .tp = 2 }).?); + for ([_][]const u8{ "fish_omni", "qwen4", "qwen4_text" }) |m| { + try std.testing.expect(backend(.{ .dir = "d", .model_type = m, .tp = 1 }) == null); + try std.testing.expectEqualStrings(no_ranks, backend(.{ .dir = "d", .model_type = m, .tp = 2 }).?); + } +} + +test "the CUDA server serves --tp, and Flash Next's two ranks pass family admission" { + try std.testing.expect(for (@import("cuda.zig").serves) |x| { + if (std.mem.eql(u8, x, "--tp")) break true; + } else false); + const flashnext = @import("flashnext").native; + try std.testing.expect(family(flashnext, .{ .dir = "d", .model_type = flashnext.model_type, .tp = 2, .prompt_cache_gib = 0 }) == null); } test "--learn is refused whatever the family, and --slide or a learned-disk size alone is not" { diff --git a/zig/src/server/engines.zig b/zig/src/server/engines.zig index 00d3868e08..4993ddb59a 100644 --- a/zig/src/server/engines.zig +++ b/zig/src/server/engines.zig @@ -33,7 +33,7 @@ pub fn identity(a: Allocator, args: cli.Args) build_identity.Identity { pub fn open(a: Allocator, gpa: Allocator, io: std.Io, dir: []const u8, model_type: []const u8, args: cli.Args, problem: *[]const u8) !?Opened { if (args.idle_unload > 0) { const world = if (std.c.getenv("TF_GLM_WORLD")) |value| std.mem.span(value) else null; - if (@import("idle_unload_policy.zig").refusal(native.backends, model_type, args.idle_unload, world)) |why| { + if (@import("idle_unload_policy.zig").refusal(native.backends, model_type, args.idle_unload, world, args.tp)) |why| { problem.* = why; return null; } diff --git a/zig/src/server/idle_unload_policy.zig b/zig/src/server/idle_unload_policy.zig index 2746c326ff..37407c5c35 100644 --- a/zig/src/server/idle_unload_policy.zig +++ b/zig/src/server/idle_unload_policy.zig @@ -1,9 +1,10 @@ //! Refuse unload where closing the current engine also ends processes needed for reopening. const std = @import("std"); -pub fn refusal(backends: []const []const u8, model_type: []const u8, seconds: i64, world: ?[]const u8) ?[]const u8 { - if (seconds <= 0 or !std.mem.eql(u8, model_type, "glm5_next")) return null; - if (backends.len == 0 or !std.mem.eql(u8, backends[0], "cuda")) return null; +pub fn refusal(backends: []const []const u8, model_type: []const u8, seconds: i64, world: ?[]const u8, tp: u32) ?[]const u8 { + if (seconds <= 0 or backends.len == 0 or !std.mem.eql(u8, backends[0], "cuda")) return null; + if (tp >= 2) return "--idle-unload is not supported with --tp on CUDA: closing rank 0 ends its followers; serve with --idle-unload 0"; + if (!std.mem.eql(u8, model_type, "glm5_next")) return null; const ranks = std.fmt.parseInt(usize, world orelse return null, 10) catch return null; if (ranks < 2) return null; return "--idle-unload is not supported for GLM CUDA with TF_GLM_WORLD >= 2: closing ends its followers; serve with --idle-unload 0"; @@ -14,13 +15,23 @@ test "idle review: GLM CUDA env ranks refuse unload independently of the CLI top const args: @import("cli.zig").Args = .{ .idle_unload = 1 }; try t.expectEqual(@as(u32, 1), args.tp); for ([_][]const u8{ "2", "3", "4", "8", "16" }) |world| { - const why = refusal(&.{"cuda"}, "glm5_next", args.idle_unload, world); + const why = refusal(&.{"cuda"}, "glm5_next", args.idle_unload, world, args.tp); try t.expect(why != null); try t.expect(std.mem.indexOf(u8, why.?, "TF_GLM_WORLD") != null); } - for ([_]i64{ 0, -1 }) |seconds| try t.expectEqual(null, refusal(&.{"cuda"}, "glm5_next", seconds, "4")); - for ([_]?[]const u8{ null, "1", "invalid" }) |world| try t.expectEqual(null, refusal(&.{"cuda"}, "glm5_next", 1, world)); - try t.expectEqual(null, refusal(&.{"metal"}, "glm5_next", 1, "4")); - try t.expectEqual(null, refusal(&.{"cuda"}, "qwen3_5", 1, "4")); - try t.expectEqual(null, refusal(&.{}, "glm5_next", 1, "4")); + for ([_]i64{ 0, -1 }) |seconds| try t.expectEqual(null, refusal(&.{"cuda"}, "glm5_next", seconds, "4", 1)); + for ([_]?[]const u8{ null, "1", "invalid" }) |world| try t.expectEqual(null, refusal(&.{"cuda"}, "glm5_next", 1, world, 1)); + try t.expectEqual(null, refusal(&.{"metal"}, "glm5_next", 1, "4", 1)); + try t.expectEqual(null, refusal(&.{"cuda"}, "qwen3_5", 1, "4", 1)); + try t.expectEqual(null, refusal(&.{}, "glm5_next", 1, "4", 1)); +} + +test "idle review: CLI ranks on CUDA refuse unload for every family, and one rank or no unload opens" { + const t = std.testing; + for ([_][]const u8{ "qwen4_exp", "glm5_next", "qwen3_5" }) |model| { + try t.expect(std.mem.indexOf(u8, refusal(&.{"cuda"}, model, 1, null, 2).?, "--tp") != null); + try t.expectEqual(null, refusal(&.{"cuda"}, model, 0, null, 2)); + try t.expectEqual(null, refusal(&.{"metal"}, model, 1, null, 2)); + } + try t.expectEqual(null, refusal(&.{"cuda"}, "qwen4_exp", 1, null, 1)); } diff --git a/zig/tests/cuda/contract_family_sets.zig b/zig/tests/cuda/contract_family_sets.zig index a602175cb3..ab224f0442 100644 --- a/zig/tests/cuda/contract_family_sets.zig +++ b/zig/tests/cuda/contract_family_sets.zig @@ -1,9 +1,11 @@ -//! The family sets in the full kernel contract: Nemotron's pocket, MiniMax-M3's, MOSS-Image's, GB10's grouped affine-4. +//! Family contract sets: pocket, MiniMax-M3, MOSS-Image, grouped affine-4, Flash Next BF16 dense and FP8 expert downs. const std = @import("std"); const cuda = @import("cuda"); const check = @import("check.zig"); const g = @import("core").contract.gates; +const Prompt = @import("contract_qmmf_prompt.zig").Registrations; +const SharedDown = @import("contract_qmmf_shared.zig").Registrations; const Grouped = @import("contract_gb10_many.zig").Registrations; pub const Sets = struct { @@ -11,7 +13,10 @@ pub const Sets = struct { pocket: *@import("contract_pocket.zig").Set, minimax: *@import("contract_minimax.zig").Set, moss: *@import("contract_moss_image.zig").Set, + modelopt: *@import("contract_flashnext_modelopt.zig").Set, grouped: *Grouped, + shared_down: *SharedDown, + prompt: *Prompt, /// Each set's kernels appended to `list`, in this order. pub fn init(gpu: check.Gpu, s: cuda.Stream, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel), affine4: *cuda.qlinear.Affine4Kernels, cap: u32) !Sets { @@ -21,17 +26,30 @@ pub const Sets = struct { errdefer minimax.deinit(); const moss = try @import("contract_moss_image.zig").Set.init(gpu, s, a, list); errdefer moss.deinit(); + const modelopt = try @import("contract_flashnext_modelopt.zig").Set.init(gpu, s, a, list); + errdefer modelopt.deinit(); const grouped = try gpu.gpa.create(Grouped); errdefer gpu.gpa.destroy(grouped); grouped.* = Grouped.init(gpu, s, affine4); errdefer grouped.deinit(); try grouped.append(cap, a, list); - return .{ .gpa = gpu.gpa, .pocket = pocket, .minimax = minimax, .moss = moss, .grouped = grouped }; + const shared_down = try SharedDown.create(gpu, s, &modelopt.lane, a, list); + errdefer shared_down.deinit(); + const prompt = try gpu.gpa.create(Prompt); + errdefer gpu.gpa.destroy(prompt); + prompt.* = Prompt.init(gpu, s, &modelopt.lane); + errdefer prompt.deinit(); + try prompt.append(cap, a, list); + return .{ .prompt = prompt, .shared_down = shared_down, .gpa = gpu.gpa, .pocket = pocket, .minimax = minimax, .moss = moss, .modelopt = modelopt, .grouped = grouped }; } pub fn deinit(sets: Sets) void { + sets.prompt.deinit(); + sets.gpa.destroy(sets.prompt); + sets.shared_down.deinit(); sets.grouped.deinit(); sets.gpa.destroy(sets.grouped); + sets.modelopt.deinit(); sets.moss.deinit(); sets.minimax.deinit(); sets.pocket.deinit(); diff --git a/zig/tests/cuda/contract_flashnext_columns.zig b/zig/tests/cuda/contract_flashnext_columns.zig new file mode 100644 index 0000000000..d5deba87d3 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_columns.zig @@ -0,0 +1,94 @@ +//! Rank-block column copies preserve raw 16-bit words, including both halves of FP32 values and untouched slots. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const Rig = @import("contract_rig.zig").Rig; +const c = core.contract.spec; +const g = core.contract.gates; + +fn noSums(_: c.Shape) usize { + return 0; +} + +pub const Columns = struct { + rig: Rig, + f: cuda.Function, + slot: bool = false, + control: bool = false, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + input: std.ArrayList(u16) = .empty, + src: u64 = 0, + dst: u64 = 0, + msg: [160]u8 = undefined, + const vt = g.hooks(Columns); + + pub fn deinit(self: *Columns) void { + self.input.deinit(self.rig.gpu.gpa); + self.rig.deinit(); + } + + pub fn kernel(self: *Columns) g.Kernel { + return .{ .ctx = self, .vt = &vt, .control = if (self.control) .shapes else null, .spec = .{ + .name = if (self.control) "control: Flash Next columns read rank zero twice" else if (self.slot) + "cuda flashnext columns into the shared expert slot" + else + "cuda flashnext rank blocks into full input and output rows", + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = "every input word is copied without arithmetic; all other slots retain their bytes" }, + .tiles = .{ .{ .size = 16 }, .{ .size = 256, .step = 1, .min = 1 }, .{} }, + .served = &.{ .{ .rows = 1, .n = 320, .k = 1 }, .{ .rows = 17, .n = 2560, .k = 1 }, .{ .rows = 1536, .n = 3072, .k = 1 } }, + } }; + } + + pub fn prepare(self: *Columns, shape: c.Shape, _: g.Fill, seed: u64) !bool { + self.rig.reset(); + self.shape = shape; + if (shape.rows == 0 or shape.n == 0) return false; + try self.input.resize(self.rig.gpu.gpa, 2 * shape.rows * shape.n); + var rng = std.Random.DefaultPrng.init(seed); + for (self.input.items) |*v| v.* = rng.random().int(u16); + self.src = try self.rig.input(std.mem.sliceAsBytes(self.input.items)); + self.dst = try self.rig.output(self.count()); + return true; + } + + pub fn run(self: *Columns) !void { + const s = self.shape; + const slots: usize = if (self.slot) 11 else 1; + const pitch = 2 * s.n * slots; + for (0..2) |rank| { + var a: cuda.Args = .{}; + a.add(self.src + (if (self.control) 0 else rank) * s.rows * s.n * 2); + a.add(@as(i64, @intCast(s.n))); + a.add(self.dst + ((slots - 1) * 2 * s.n + rank * s.n) * 2); + a.add(@as(i64, @intCast(pitch))); + a.add(@as(i32, @intCast(s.n))); + try cuda.launch.launch(self.f, .{ .grid = .{ .x = @intCast((s.n + 255) / 256), .y = @intCast(s.rows) }, .block = .{ .x = 256 } }, self.rig.s, &a); + } + try self.rig.s.synchronize(); + } + + pub fn count(self: *Columns) usize { + return self.shape.rows * self.shape.n * 4 * @as(usize, if (self.slot) 11 else 1); + } + + pub fn reference(_: *Columns, want: []f64, mag: []f64) !void { + @memset(want, 0); + @memset(mag, 0); + } + + pub fn check(self: *Columns, _: g.Gate) !?[]const u8 { + const got = try self.rig.outputBytes(); + const s = self.shape; + const slots: usize = if (self.slot) 11 else 1; + for (0..s.rows) |row| for (0..slots) |slot| for (0..2) |rank| for (0..s.n) |col| { + const index = ((row * slots + slot) * 2 + rank) * s.n + col; + const want: u16 = if (slot == slots - 1) self.input.items[(rank * s.rows + row) * s.n + col] else 0xffff; + if (std.mem.readInt(u16, got[index * 2 ..][0..2], .little) != want) + return try std.fmt.bufPrint(&self.msg, "row {d} slot {d} rank {d} column {d} differs", .{ row, slot, rank, col }); + }; + return null; + } +}; diff --git a/zig/tests/cuda/contract_flashnext_glue.zig b/zig/tests/cuda/contract_flashnext_glue.zig index 885509c92b..5e796f50c9 100644 --- a/zig/tests/cuda/contract_flashnext_glue.zig +++ b/zig/tests/cuda/contract_flashnext_glue.zig @@ -13,8 +13,8 @@ const tk = flash.tile_kernels; pub const Kind = enum { chain, scores, readout, fused }; -const nk = 16; -const nv = 48; +const nk = flash.Heads.full.nk; +const nv = flash.Heads.full.nv; const dk = 128; const dv = 128; const di = 128; diff --git a/zig/tests/cuda/contract_flashnext_glue_test.zig b/zig/tests/cuda/contract_flashnext_glue_test.zig index 13e0d8be0c..72d74a9e5e 100644 --- a/zig/tests/cuda/contract_flashnext_glue_test.zig +++ b/zig/tests/cuda/contract_flashnext_glue_test.zig @@ -38,3 +38,83 @@ test "the fused block declares the state alias and binds one buffer for it" { try std.testing.expectEqual(@as(?u64, 64), a.integer(7)); } else return error.MissingFusedArgs; } + +const lay = @import("flashnext_gdn_heads.zig"); +const heads = @import("contract_flashnext_heads.zig"); + +test "a rank's DeltaNet heads keep the design's rows of the full layout, and the two ranks cover it once" { + const want = [_][2]usize{ .{ 0, 1024 }, .{ 2048, 1024 }, .{ 4096, 3072 }, .{ 10240, 3072 }, .{ 16384, 24 }, .{ 16432, 24 } }; + const step = [_]usize{ 1024, 1024, 3072, 3072, 24, 24 }; + for (0..2) |r| { + const cs = lay.cuts(.proj, .half, r); + try std.testing.expectEqual(want.len, cs.slice().len); + for (cs.slice(), want, step) |got, w, s| try std.testing.expectEqual(lay.Cut{ .from = w[0] + s * r, .len = w[1] }, got); + } + try std.testing.expectEqual(@as(usize, 8240), lay.width(.proj, .half)); + for (std.enums.values(lay.Part)) |part| { + const n = lay.width(part, .full); + const seen = try std.testing.allocator.alloc(u8, n); + defer std.testing.allocator.free(seen); + @memset(seen, 0); + var local: usize = 0; + for (0..2) |r| { + const cs = lay.cuts(part, .half, r); + for (cs.slice()) |x| for (seen[x.from..][0..x.len]) |*v| { + v.* += 1; + }; + for (cs.slice()) |x| local += x.len; + } + var at: usize = 0; // the full heads' one rank keeps every element in place + for (lay.cuts(part, .full, 0).slice()) |x| { + try std.testing.expectEqual(at, x.from); + at += x.len; + } + try std.testing.expectEqual(n, at); + try std.testing.expectEqual(2 * lay.width(part, .half), local); + for (seen) |v| try std.testing.expectEqual(@as(u8, 1), v); + } +} + +test "the gather copies each kept element from where fullIndex says it lies" { + const a = std.testing.allocator; + const rows = 3; + const n = rows * lay.width(.proj, .full); + const src = try a.alloc(u32, n); + defer a.free(src); + for (src, 0..) |*v, i| v.* = @intCast(i); + for (0..2) |r| { + const got = try lay.gather(a, std.mem.sliceAsBytes(src), 4, .proj, .half, r); + defer a.free(got); + const vals = std.mem.bytesAsSlice(u32, got); + const lw = lay.width(.proj, .half); + try std.testing.expectEqual(rows * lw, vals.len); + for (0..rows) |row| for (0..lw) |i| { + try std.testing.expectEqual(row * lay.width(.proj, .full) + lay.fullIndex(.proj, .half, r, i), vals[row * lw + i]); + }; + } + const same = try lay.gather(a, std.mem.sliceAsBytes(src), 4, .proj, .full, 0); + defer a.free(same); + try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(src), same); +} + +test "each DeltaNet twin is a candidate on both ranks, replacing its full original, at the served rows and seams" { + for (std.enums.values(heads.Kind)) |kind| for (0..2) |r| { + const s = heads.spec(kind, r); + try std.testing.expect(std.mem.indexOf(u8, s.name, lay.symbol(kind, true)) != null); + try std.testing.expect(std.mem.startsWith(u8, s.replaces.?, lay.symbol(kind, false))); + try std.testing.expectEqual(kind == .rows or kind == .fused, s.in_place); + var served: [3]bool = @splat(false); + for (s.served) |x| for ([_]usize{ 64, parts.default_chunk, parts.max_chunk }, 0..) |want, i| { + if (x.rows == want and x.n == lay.Heads.half.valueDim()) served[i] = true; + }; + for (served) |yes| try std.testing.expect(yes); + const list = try c.shapes(std.testing.allocator, s); + defer std.testing.allocator.free(list); + var seams: [4]bool = @splat(false); + for (list) |x| for ([_]usize{ 1, 15, 16, 17 }, 0..) |want, i| { + if (x.rows == want) seams[i] = true; + }; + for (seams) |yes| try std.testing.expect(yes); + }; + try std.testing.expect(!std.mem.eql(u8, heads.spec(.front, 0).name, heads.spec(.front, 1).name)); +} diff --git a/zig/tests/cuda/contract_flashnext_groups.zig b/zig/tests/cuda/contract_flashnext_groups.zig new file mode 100644 index 0000000000..8b8dbe34dc --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_groups.zig @@ -0,0 +1,70 @@ +//! The runner's Flash Next prompt groups as one set: the full layer's glue and the DeltaNet twins on each rank's heads. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const glue_k = @import("contract_flashnext_glue.zig"); +const heads_k = @import("contract_flashnext_heads.zig"); +const Columns = @import("contract_flashnext_columns.zig").Columns; +const lay = @import("flashnext_gdn_heads.zig"); + +const g = core.contract.gates; +const parts = flash.batch_state; +const tk = flash.tile_kernels; + +pub const Set = struct { + glue: ?flash.Glue = null, + fns: glue_k.Fns = undefined, + glues: [std.enums.values(glue_k.Kind).len]glue_k.Glue = undefined, + n_glues: usize = 0, + twins: lay.Fns = undefined, + ranks: [std.enums.values(heads_k.Kind).len * 2]heads_k.Rank = undefined, + n_ranks: usize = 0, + columns: [4]Columns = undefined, + n_columns: usize = 0, + + /// Every registration onto `kernels`, from the flashnext_batch, flashnext_attn and flashnext_hcf images. + pub fn add(self: *Set, gpu: check.Gpu, s: cuda.Stream, batch: cuda.Module, attn: cuda.Module, hcf: cuda.Module, kernels: *std.ArrayList(g.Kernel)) !void { + const a = gpu.gpa; + self.glue = try flash.Glue.load(gpu.d); + self.fns = .{ + .tpr = try batch.function("fn_gdn_chain_tpr"), + .rows_to = try batch.function("fn_gdn_rows_to"), + .tile = try attn.function("_ZN8qsa_tile6scoresILi8ELi4EEEvPK13__nv_bfloat16S3_PKiPfiiii"), + .scores = try attn.function("fn_qsa_scores"), + .hcf_down = try hcf.function(tk.cubin_hcf_down), + .hcf_up = try hcf.function(tk.cubin_hcf_up), + .front = try batch.function("fn_gdn_front"), + .norm = try batch.function("fn_gdn_norm"), + .fused = try batch.function("fn_gdn_fused"), + }; + try self.fns.fused.allowDynamicShared(parts.gdn_fused_smem); + try self.fns.tpr.allowDynamicShared(parts.chain_tpr_smem); + try self.fns.hcf_down.allowDynamicShared(tk.hcf_down_smem); + try self.fns.hcf_up.allowDynamicShared(tk.hcf_up_smem); + for (std.enums.values(glue_k.Kind), &self.glues) |kind, *m| { + m.* = .{ .rig = .{ .gpu = gpu, .s = s }, .kind = kind, .f = &self.fns, .glue = &self.glue.? }; + self.n_glues += 1; + } + for (&self.glues) |*m| try kernels.append(a, m.kernel()); + self.twins = try lay.Fns.load(batch); + for (std.enums.values(heads_k.Kind), 0..) |kind, i| for (0..2) |rank| { + self.ranks[i * 2 + rank] = .{ .rig = .{ .gpu = gpu, .s = s }, .kind = kind, .rank = rank, .f = &self.twins }; + self.n_ranks += 1; + }; + for (&self.ranks) |*m| try kernels.append(a, m.kernel()); + for (&self.columns, 0..) |*m, i| { + m.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = try batch.function("fn_copy_cols"), .slot = i % 2 == 1, .control = i >= 2 }; + self.n_columns += 1; + try kernels.append(a, m.kernel()); + } + } + + pub fn deinit(self: *Set) void { + for (self.columns[0..self.n_columns]) |*m| m.deinit(); + for (self.ranks[0..self.n_ranks]) |*m| m.deinit(); + for (self.glues[0..self.n_glues]) |*m| m.deinit(); + if (self.glue) |*x| x.deinit(); + } +}; diff --git a/zig/tests/cuda/contract_flashnext_heads.zig b/zig/tests/cuda/contract_flashnext_heads.zig new file mode 100644 index 0000000000..b4dcd0a889 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_heads.zig @@ -0,0 +1,282 @@ +//! Flash Next's DeltaNet twins on one rank's heads: each kernel's bytes equal its full original's on those heads. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const rig_ = @import("contract_rig.zig"); +const lay = @import("flashnext_gdn_heads.zig"); + +const c = core.contract.spec; +const g = core.contract.gates; +const parts = flash.batch_state; +const Heads = lay.Heads; +const Part = lay.Part; +pub const Kind = lay.Kind; +const dk = lay.dk; +const dv = lay.dv; + +fn noSums(_: c.Shape) usize { + return 0; +} + +/// The prompt pass's rows on one rank's heads (n: its value heads x dv): the fused seam, the default and the most. +const half_served = [_]c.Shape{ + .{ .rows = 64, .n = Heads.half.nv * dv, .k = dk }, + .{ .rows = parts.default_chunk, .n = Heads.half.nv * dv, .k = dk }, + .{ .rows = parts.max_chunk, .n = Heads.half.nv * dv, .k = dk }, +}; + +const ranks = 2; + +const names = blk: { + var n: [std.enums.values(Kind).len][ranks][]const u8 = undefined; + for (std.enums.values(Kind)) |k| for (0..ranks) |r| { + n[@backingInt(k)][r] = std.fmt.comptimePrint("cuda flashnext DeltaNet {s} on rank {d}'s 8 key and 24 value heads", .{ lay.symbol(k, true), r }); + }; + break :blk n; +}; + +const originals = blk: { + var n: [std.enums.values(Kind).len][]const u8 = undefined; + for (std.enums.values(Kind)) |k| n[@backingInt(k)] = std.fmt.comptimePrint("{s} on the same heads of the full layout", .{lay.symbol(k, false)}); + break :blk n; +}; + +/// rows = token rows; n = the candidate's value heads x dv; k = dk. +pub fn spec(kind: Kind, rank: usize) c.Spec { + return .{ .name = names[@backingInt(kind)][rank], .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = "every output byte, the state's included, equal to the full kernel's on these heads, non-finite inputs included" }, .replaces = originals[@backingInt(kind)], .tiles = .{ .{ .size = 16 }, .{}, .{} }, .in_place = kind == .rows or kind == .fused, .served = &half_served }; +} + +/// A buffer the kernel reads: what its rows hold (null: not cut, the norm weight), its element bytes and its values. +const In = struct { part: ?Part, elem: usize = 4, lo: f32, hi: f32, decay: bool = false }; +/// A buffer the kernel writes: what its rows hold, its element bytes, and whether it is the state read in place. +const Out = struct { part: Part, elem: usize = 4, state: bool = false }; + +const proj: In = .{ .part = .proj, .elem = 2, .lo = -1, .hi = 1 }; +const conv: In = .{ .part = .conv, .elem = 2, .lo = -1, .hi = 1 }; +const cw: In = .{ .part = .taps, .elem = 2, .lo = -0.5, .hi = 0.5 }; +const gates: In = .{ .part = .gates, .lo = -1, .hi = 1 }; +const norm_w: In = .{ .part = null, .elem = 2, .lo = 0.5, .hi = 1.5 }; +const keys: In = .{ .part = .keys, .lo = -0.1, .hi = 0.1 }; +const chain_in = [_]In{ keys, keys, .{ .part = .values, .lo = -1, .hi = 1 }, .{ .part = .gb, .lo = 0, .hi = 1, .decay = true } }; +const state: In = .{ .part = .state, .lo = -0.5, .hi = 0.5 }; + +fn inputs(kind: Kind) []const In { + return switch (kind) { + .front => &.{ proj, conv, cw, gates, gates }, + .norm => &.{ .{ .part = .values, .lo = -1, .hi = 1 }, proj, norm_w }, + .rows => &chain_in, + .rows_to, .tpr => &(chain_in ++ [_]In{state}), + .fused => &.{ proj, conv, cw, gates, gates, norm_w }, + }; +} + +fn outputs(kind: Kind) []const Out { + return switch (kind) { + .front => &.{ .{ .part = .keys }, .{ .part = .keys }, .{ .part = .values }, .{ .part = .gb } }, + .norm => &.{ .{ .part = .values, .elem = 2 }, .{ .part = .sums } }, + .rows => &.{ .{ .part = .values }, .{ .part = .state, .state = true } }, + .rows_to, .tpr => &.{ .{ .part = .values }, .{ .part = .state } }, + .fused => &.{ .{ .part = .values, .elem = 2 }, .{ .part = .sums }, .{ .part = .state, .state = true }, .{ .part = .values } }, + }; +} + +/// Bytes of a buffer in the layout of heads h: the norm weight (no part) is dv values whatever the heads. +fn bytesOf(part: ?Part, elem: usize, m: usize, h: Heads) usize { + const p = part orelse return dv * elem; + return lay.bytes(p, elem, m, h); +} + +fn put(host: []u8, in: In, at: usize, v: f32) void { + if (in.elem == 2) std.mem.bytesAsSlice(u16, host)[at] = c.toBf16(v) else std.mem.bytesAsSlice(f32, host)[at] = v; +} + +pub const Rank = struct { + rig: rig_.Rig, + kind: Kind, + /// The candidate's heads: one rank's (its twins), cut from the full layout. + heads: Heads = .half, + rank: usize, + f: *const lay.Fns, + rows: usize = 0, + /// The full-layout inputs and, by the same slot, the candidate's: rank r's heads gathered from them. + wide: [6]u64 = @splat(0), + ins: [6]u64 = @splat(0), + outs: [4]u64 = @splat(0), + nout: usize = 0, + /// The full outputs the original writes, back to back, poisoned before each baseline. + base: u64 = 0, + /// The candidate's state read and written in place (0: none), and the starting bytes of it and of its full twin. + state: u64 = 0, + fresh: ?[]u8 = null, + wide_fresh: ?[]u8 = null, + want: std.ArrayList(u8) = .empty, + seen: std.ArrayList(u8) = .empty, + msg: [160]u8 = undefined, + + pub fn deinit(self: *Rank) void { + self.drop(); + self.want.deinit(self.rig.gpu.gpa); + self.seen.deinit(self.rig.gpu.gpa); + self.rig.deinit(); + } + + fn drop(self: *Rank) void { + if (self.fresh) |f| self.rig.gpu.gpa.free(f); + if (self.wide_fresh) |f| self.rig.gpu.gpa.free(f); + self.fresh = null; + self.wide_fresh = null; + } + + pub fn kernel(self: *Rank) g.Kernel { + return .{ .spec = spec(self.kind, self.rank), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Rank); + + /// Full-layout values of a buffer, with gate 1's plants at known places of rank r's heads. + fn values(self: *Rank, in: In, r: std.Random, bad: bool) ![]u8 { + const m = self.rows; + const host = try self.rig.gpu.gpa.alloc(u8, bytesOf(in.part, in.elem, m, .full)); + const n = host.len / in.elem; + for (0..n) |i| { + var v = in.lo + r.float(f32) * (in.hi - in.lo); + if (in.decay and i % 2 == 0) v = -v; // as the full fixture: a log decay, then beta in (0, 1) + put(host, in, i, v); + } + const part = in.part orelse return host; + if (!bad) return host; + const h = self.heads; + switch (part) { + .proj => { + self.plant(host, in, 0, 3, std.math.nan(f32)); // q of the rank's first key head + self.plant(host, in, m / 2, 2 * h.nk * dk + 7, std.math.inf(f32)); // v + self.plant(host, in, m - 1, h.convDim() + 5, -std.math.inf(f32)); // z + }, + .keys => self.plant(host, in, 0, 3, std.math.nan(f32)), + .values => self.plant(host, in, m / 2, 5, std.math.inf(f32)), + .state => self.plant(host, in, 0, dv * dk + 7, -std.math.inf(f32)), + else => {}, + } + return host; + } + + /// Element i of row `row` of the rank's local layout, placed in the full-layout buffer. + fn plant(self: *Rank, host: []u8, in: In, row: usize, i: usize, v: f32) void { + const part = in.part.?; + put(host, in, row * lay.width(part, .full) + lay.fullIndex(part, self.heads, self.rank, i), v); + } + + fn local(self: *Rank, host: []const u8, part: ?Part, elem: usize) ![]u8 { + const a = self.rig.gpu.gpa; + const p = part orelse return a.dupe(u8, host); + return lay.gather(a, host, elem, p, self.heads, self.rank); + } + + pub fn prepare(self: *Rank, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + self.drop(); + self.state = 0; + self.nout = 0; + self.rows = shape.rows; + if (shape.rows == 0 or shape.n != self.heads.valueDim() or shape.k != dk) return false; + var prng = std.Random.DefaultPrng.init(seed ^ 0x6865616473); + const r = prng.random(); + const a = self.rig.gpu.gpa; + const bad = fill == .nonfinite; + for (inputs(self.kind), 0..) |in, i| { + const host = try self.values(in, r, bad); + defer a.free(host); + self.wide[i] = try self.rig.input(host); + const mine = try self.local(host, in.part, in.elem); + defer a.free(mine); + self.ins[i] = try self.rig.input(mine); + } + var total: usize = 0; + for (outputs(self.kind), 0..) |o, i| { + self.outs[i] = try self.rig.output(bytesOf(o.part, o.elem, self.rows, self.heads)); + self.nout += 1; + total += bytesOf(o.part, o.elem, self.rows, .full); + if (!o.state) continue; + // the state read and written where it lies: every launch of either kernel starts from these bytes + self.state = self.outs[i]; + self.wide_fresh = try self.values(state, r, bad); + self.fresh = try self.local(self.wide_fresh.?, .state, 4); + } + self.base = try self.rig.scratch(total); + return true; + } + + /// The kernel of this kind built for heads h, on inputs `in` into outputs `o`, as the pass binds it. + fn launch(self: *Rank, h: Heads, in: []const u64, o: []const u64) !void { + const f = self.f.of(self.kind, h); + const s = self.rig.s; + const m = self.rows; + switch (self.kind) { + .front => try lay.front(f, s, h, m, in[0], in[1], in[2], in[3], in[4], o[0], o[1], o[2], o[3]), + .norm => try lay.norm(f, s, h, m, in[0], in[1], in[2], o[0], o[1]), + .rows => try lay.chain(f, s, .rows, h, m, in[0], in[1], in[2], in[3], o[1], o[1], o[0]), + .rows_to, .tpr => try lay.chain(f, s, self.kind, h, m, in[0], in[1], in[2], in[3], in[4], o[1], o[0]), + .fused => try lay.fused(f, s, h, m, in[0], in[1], in[2], in[3], in[4], o[2], o[2], in[5], o[0], o[1], o[3]), + } + } + + fn upload(self: *Rank, ptr: u64, bytes: []const u8) !void { + const b: cuda.DeviceBuffer = .{ .d = self.rig.gpu.d, .ptr = ptr, .len = bytes.len }; + try b.upload(0, bytes); + } + + pub fn run(self: *Rank) !void { + if (self.fresh) |f| try self.upload(self.state, f); + try self.launch(self.heads, &self.ins, self.outs[0..self.nout]); + try self.rig.s.synchronize(); + } + + /// The full original on the full inputs into its own outputs, rank r's heads gathered in the candidate's order. + pub fn replaced(self: *Rank) ![]const u8 { + const a = self.rig.gpu.gpa; + if (self.fresh) |f| try self.upload(self.state, f); + const list = outputs(self.kind); + var o: [4]u64 = undefined; + var at = self.base; + for (list, 0..) |x, i| { + o[i] = at; + at += bytesOf(x.part, x.elem, self.rows, .full); + } + const ff = try a.alloc(u8, at - self.base); + defer a.free(ff); + @memset(ff, 0xff); + try self.upload(self.base, ff); + for (list, 0..) |x, i| if (x.state) try self.upload(o[i], self.wide_fresh.?); + try self.launch(.full, &self.wide, o[0..list.len]); + try self.rig.s.synchronize(); + self.want.clearRetainingCapacity(); + for (list, 0..) |x, i| { + const got = try self.rig.bytes(o[i], bytesOf(x.part, x.elem, self.rows, .full)); + const mine = try lay.gather(a, got, x.elem, x.part, self.heads, self.rank); + defer a.free(mine); + try self.want.appendSlice(a, mine); + } + return self.want.items; + } + + /// Gate 1's stated rule: the outputs on the planted inputs equal the original's on the same heads. + pub fn check(self: *Rank, gate: g.Gate) !?[]const u8 { + _ = gate; + self.seen.clearRetainingCapacity(); + try self.seen.appendSlice(self.rig.gpu.gpa, try self.rig.outputBytes()); + const want = try self.replaced(); + if (std.mem.eql(u8, self.seen.items, want)) return null; + return try std.fmt.bufPrint(&self.msg, "bytes differ from the full kernel's from byte {d}", .{std.mem.indexOfDiff(u8, self.seen.items, want).?}); + } + + pub fn count(self: *Rank) usize { + var total: usize = 0; + for (outputs(self.kind)) |x| total += bytesOf(x.part, x.elem, self.rows, self.heads); + return total; + } + + pub fn reference(_: *Rank, _: []f64, _: []f64) !void { + return error.NoFloat64Reference; + } +}; diff --git a/zig/tests/cuda/contract_flashnext_modelopt.zig b/zig/tests/cuda/contract_flashnext_modelopt.zig new file mode 100644 index 0000000000..a8c1abc1f2 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_modelopt.zig @@ -0,0 +1,410 @@ +//! NVIDIA's NVFP4 Flash Next checkpoint: BF16 dense and shared projections, FP8 MTP-head expert downs, --tp picks. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const Rig = @import("contract_rig.zig").Rig; +const c = core.contract.spec; +const g = core.contract.gates; +const Dense = core.contract.dense.Dense; +const q = cuda.qmmf; +const ex = cuda.experts; +const grouped = cuda.grouped; + +fn depth(s: c.Shape) usize { + return s.k / 16 + s.k / 128 + 24; +} + +/// Stress weights: integers 1 to 7 over the first half of each row and their negatives over the second. +fn stress(rng: std.Random, n: usize, k: usize, out: []f32) void { + for (0..n) |j| for (0..k / 2) |i| { + const v: f32 = @floatFromInt(1 + rng.uintLessThan(u32, 7)); + out[j * k + i] = v; + out[j * k + k / 2 + i] = -v; + }; +} + +// The dense and shared layers at decode rows: DeltaNet qkvz+ab, attention qkv+indexer, both outs, shared gate+up, down. +const bf16_served = [_]c.Shape{ + .{ .rows = 1, .n = 8240, .k = 2560 }, .{ .rows = 64, .n = 7296, .k = 2560 }, + .{ .rows = 1, .n = 1280, .k = 6144 }, .{ .rows = 64, .n = 1280, .k = 6144 }, + .{ .rows = 1, .n = 320, .k = 2560 }, .{ .rows = 64, .n = 1280, .k = 640 }, + .{ .rows = 1, .n = 16480, .k = 2560 }, .{ .rows = 16, .n = 16480, .k = 2560 }, + .{ .rows = 1, .n = 13952, .k = 2560 }, .{ .rows = 64, .n = 13952, .k = 2560 }, + .{ .rows = 1, .n = 2560, .k = 6144 }, .{ .rows = 16, .n = 2560, .k = 6144 }, + .{ .rows = 4, .n = 1280, .k = 2560 }, .{ .rows = 64, .n = 1280, .k = 2560 }, + .{ .rows = 1, .n = 2560, .k = 640 }, .{ .rows = 64, .n = 2560, .k = 640 }, +}; + +/// Flash Next's qmmf BF16-weight product: 16-, 32- and 64-row tiles, with the lane's K-slice plan. +pub const Bf16 = struct { + rig: Rig, + lane: *const q.Lane, + format: q.Out, + control: bool = false, + d: ?Dense = null, + w: q.Weight = undefined, + x: u64 = 0, + out: u64 = 0, + part: u64 = 0, + + const vt = g.hooks(Bf16); + + pub fn deinit(self: *Bf16) void { + if (self.d) |d| d.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Bf16) g.Kernel { + const name = if (self.control) "control: Flash Next qmmf BF16 weights scaled 1.5x" else switch (self.format) { + .bf16 => "cuda qmmf BF16 weights, bf16 out (Flash Next dense and shared layers)", + .f32 => "cuda qmmf BF16 weights, fp32 out (Flash Next dense and shared layers)", + }; + return .{ .ctx = self, .vt = &vt, .control = if (self.control) .shapes else null, .spec = .{ + .name = name, + .accum = .f32, + .depth = depth, + .ops = 1, + .out = if (self.format == .f32) .f32 else .bf16, + .tiles = .{ .{ .size = 16 }, .{ .size = 64 }, .{ .size = 64, .step = 64, .min = 64 } }, + .served = &bf16_served, + } }; + } + + pub fn prepare(self: *Bf16, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (self.d) |d| d.deinit(); + self.d = null; + self.rig.reset(); + if (shape.rows == 0 or shape.n == 0 or shape.k == 0 or shape.k % 64 != 0) return false; + const a = self.rig.gpu.gpa; + const d = try Dense.init(a, shape, 1); + self.d = d; + d.fillX(fill, seed, 256); + const vals = try a.alloc(f32, shape.n * shape.k); + defer a.free(vals); + var prng = std.Random.DefaultPrng.init(seed ^ 0xb16d); + const rng = prng.random(); + if (fill == .stress) stress(rng, shape.n, shape.k, vals) else for (vals) |*v| { + v.* = rng.float(f32) * 4 - 2; + } + const bits = try a.alloc(u16, vals.len); + defer a.free(bits); + for (bits, vals, d.w) |*b, v, *w| { + b.* = c.toBf16(v); + w.* = c.fromBf16(b.*); + } + const npad = (shape.n + 127) / 128 * 128; + const tiled = try a.alloc(u16, npad * shape.k); + defer a.free(tiled); + q.packBf16(tiled, bits, shape.n, shape.k); + self.w = cuda.modelopt.bf16Weight(try self.rig.input(std.mem.sliceAsBytes(tiled)), shape.n, shape.k); + self.x = try self.rig.input(std.mem.sliceAsBytes(d.x)); + self.out = try self.rig.output(d.outputs() * @as(usize, if (self.format == .f32) 4 else 2)); + const parts = self.lane.partBytes(shape.rows, self.w); + self.part = if (parts == 0) 0 else try self.rig.scratch(parts); + return true; + } + + pub fn run(self: *Bf16) !void { + const s = self.d.?.shape; + var w = self.w; + if (self.control) w.scale = 1.5; + try self.lane.matmulOut(self.rig.s, self.x, s.k, s.rows, w, self.out, self.format, self.part); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Bf16) usize { + return self.d.?.outputs(); + } + pub fn reference(self: *Bf16, want: []f64, mag: []f64) !void { + try self.d.?.reference(want, mag); + } + pub fn simulate(self: *Bf16, acc: c.Accum, order: g.Order, out: []f64) !void { + try self.d.?.simulate(acc, order, out); + } +}; + +const experts = 4; // routed experts in the fixture; every pair picks `target`, the rest hold NaN codes +const target = 2; + +// The MTP head's routed down (2,560 x 640): pairs are tokens times ten picks, at decode and prompt row counts. +const decode_served = [_]c.Shape{ .{ .rows = 10, .n = 2560, .k = 640 }, .{ .rows = 40, .n = 2560, .k = 640 }, .{ .rows = 160, .n = 2560, .k = 640 } }; +const prompt_served = [_]c.Shape{ .{ .rows = 160, .n = 2560, .k = 640 }, .{ .rows = 1280, .n = 2560, .k = 640 } }; + +/// The MTP head's fp8_experts down_f32: decode or prompt over one expert's grouped pairs. +pub const Fp8Down = struct { + rig: Rig, + set: *const Set, + prompt: bool, + control: bool = false, + d: ?Dense = null, + layer: ex.Layer = undefined, + plan: grouped.Plan = undefined, + picks: u64 = 0, + x: u64 = 0, + out: u64 = 0, + + const vt = g.hooks(Fp8Down); + + pub fn deinit(self: *Fp8Down) void { + if (self.d) |d| d.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Fp8Down) g.Kernel { + const name = if (self.control) "control: Flash Next FP8 expert down at 1.5x block scales" else if (self.prompt) + "cuda fp8_experts prompt down fp32 (Flash Next MTP head)" + else + "cuda fp8_experts decode down fp32 (Flash Next MTP head)"; + const rows: usize = if (self.prompt) ex.promptTile(.fp8g) else 16; + return .{ .ctx = self, .vt = &vt, .control = if (self.control) .shapes else null, .spec = .{ + .name = name, + .accum = .f32, + .depth = depth, + .ops = 1, + .out = .f32, + .tiles = .{ .{ .size = rows }, .{ .size = ex.Fp8Experts.cols, .step = 128, .min = 128 }, .{ .size = 128, .step = 128, .min = 128 } }, + .served = if (self.prompt) &prompt_served else &decode_served, + } }; + } + + pub fn prepare(self: *Fp8Down, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (self.d) |d| d.deinit(); + self.d = null; + self.rig.reset(); + const n = shape.n; + const k = shape.k; + if (shape.rows == 0 or n == 0 or k == 0 or n % 128 != 0 or k % 128 != 0) return false; + const a = self.rig.gpu.gpa; + const d = try Dense.init(a, shape, 1); + self.d = d; + d.fillX(fill, seed, 256); + var prng = std.Random.DefaultPrng.init(seed ^ 0xf8d0); + const rng = prng.random(); + // E4M3 codes [experts][n][k]: the target's values, every other expert's codes NaN so a misread goes red + const codes = try a.alloc(u8, experts * n * k); + defer a.free(codes); + @memset(codes, 0x7f); + const mine = codes[target * n * k ..][0 .. n * k]; + if (fill == .stress) { + const vals = try a.alloc(f32, n * k); + defer a.free(vals); + stress(rng, n, k, vals); + for (mine, vals) |*code, v| code.* = try codeOf(v); + } else for (mine) |*code| { + code.* = rng.int(u8); + if (code.* & 0x7f == 0x7f) code.* ^= 1; + } + // f32 block scales [experts][1][n/128][k/128]: 1 for sums the stress inputs keep exact, else 2^-8 to 2^-5 + const blocks = (n / 128) * (k / 128); + const scales = try a.alloc(f32, experts * blocks); + defer a.free(scales); + for (scales) |*s| s.* = if (fill == .stress) 1 else std.math.ldexp(@as(f32, 1), -5 - @as(i32, @intCast(rng.uintLessThan(u32, 4)))); + for (d.w, 0..) |*w, i| w.* = core.quant.fp8.e4m3(mine[i]) * scales[target * blocks + (i / k / 128) * (k / 128) + (i % k) / 128]; + if (self.control) for (scales) |*s| { + s.* *= 1.5; + }; + const packed_codes = try a.alloc(u32, experts * n * k / 4); + defer a.free(packed_codes); + for (0..experts) |e| ex.Fp8Experts.packOne(packed_codes[e * n * k / 4 ..][0 .. n * k / 4], codes[e * n * k ..][0 .. n * k], n, k); + const picks = try a.alloc(i32, shape.rows); + defer a.free(picks); + @memset(picks, target); + self.layer = .{ + .format = .fp8g, + .up = 0, + .down = try self.rig.input(std.mem.sliceAsBytes(packed_codes)), + .down_scale = try self.rig.input(std.mem.sliceAsBytes(scales)), + .width = k, + .dims = n, + .experts = experts, + }; + self.picks = try self.rig.input(std.mem.sliceAsBytes(picks)); + self.x = try self.rig.input(std.mem.sliceAsBytes(d.x)); + const pairs = shape.rows; + const wide = pairs > grouped.small; + const sizes = [5]usize{ pairs * 4, grouped.maxItems(pairs, experts, 16) * 12, 8, (if (wide) pairs else 1) * 4, (if (wide) (pairs + 1023) / 1024 * experts else 1) * 4 }; + var bufs: [5]u64 = undefined; + for (sizes, &bufs) |len, *b| b.* = try self.rig.scratch(len); + self.plan = .{ .members = bufs[0], .items = bufs[1], .counts = bufs[2], .rank = bufs[3], .hist = bufs[4] }; + self.out = try self.rig.output(d.outputs() * 4); + return true; + } + + /// The router's plan, then the down call: each pair's fp32 row of `n` outputs at its own row. + pub fn run(self: *Fp8Down) !void { + const pairs = self.d.?.shape.rows; + const tile: usize = if (self.prompt) ex.promptTile(.fp8g) else 16; + try self.set.router.route(self.rig.s, self.picks, pairs, experts, tile, self.plan); + const in: ex.Rows = .{ .x = self.x, .stride = self.layer.width }; + try self.set.fp8.run(self.rig.s, .down_f32, in, self.layer, self.plan, self.prompt, self.out, grouped.maxItems(pairs, experts, tile), -1); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Fp8Down) usize { + return self.d.?.outputs(); + } + pub fn reference(self: *Fp8Down, want: []f64, mag: []f64) !void { + try self.d.?.reference(want, mag); + } + pub fn simulate(self: *Fp8Down, acc: c.Accum, order: g.Order, out: []f64) !void { + try self.d.?.simulate(acc, order, out); + } +}; + +fn noSums(_: c.Shape) usize { + return 0; +} + +// Pairs a pass routes (tokens times eleven slots) at decode, a verify round and a 4,608-row prompt pass. +const picks_served = [_]c.Shape{ .{ .rows = 11, .n = 512, .k = 1 }, .{ .rows = 176, .n = 512, .k = 1 }, .{ .rows = 50688, .n = 512, .k = 1 } }; + +/// fn_ep_picks at --tp 2's second rank (first 256): its experts' picks to 0..255, every other id to 256. +pub const EpPicks = struct { + rig: Rig, + set: *const Set, + control: bool = false, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + picks: std.ArrayList(i32) = .empty, + dev: [2]u64 = @splat(0), + msg: [160]u8 = undefined, + + const vt = g.hooks(EpPicks); + const first = 256; + const local = 256; + + pub fn deinit(self: *EpPicks) void { + self.picks.deinit(self.rig.gpu.gpa); + self.rig.deinit(); + } + + pub fn kernel(self: *EpPicks) g.Kernel { + return .{ .ctx = self, .vt = &vt, .control = if (self.control) .shapes else null, .spec = .{ + .name = if (self.control) "control: fn_ep_picks one expert off" else "cuda flashnext fn_ep_picks (a rank's expert ids at --tp)", + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = "integer expert ids: an owned id less the rank's first, any other id the rank's count" }, + .tiles = .{ .{ .size = 256 }, .{}, .{} }, + .served = &picks_served, + } }; + } + + pub fn prepare(self: *EpPicks, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + _ = fill; + self.rig.reset(); + if (shape.rows == 0 or shape.n != 512) return false; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + self.shape = shape; + try self.picks.resize(self.rig.gpu.gpa, shape.rows); + for (self.picks.items, 0..) |*p, i| p.* = if (i % 11 == 10) 512 else @intCast(rng.uintLessThan(u32, 512)); // slot 10: shared + for ([_]i32{ 0, first - 1, first, first + local - 1 }, 0..) |edge, i| if (i < shape.rows) { + self.picks.items[i] = edge; + }; + self.dev[0] = try self.rig.input(std.mem.sliceAsBytes(self.picks.items)); + self.dev[1] = try self.rig.output(shape.rows * 4); + return true; + } + + pub fn run(self: *EpPicks) !void { + var a: cuda.Args = .{}; + a.add(self.dev[0]); + a.add(self.dev[1]); + const start: usize = if (self.control) first + 1 else first; + for ([_]usize{ self.shape.rows, start, local }) |v| a.add(@as(i32, @intCast(v))); + try cuda.launch.launch(self.set.ep_picks, .{ .grid = .{ .x = @intCast((self.shape.rows + 255) / 256) }, .block = .{ .x = 256 } }, self.rig.s, &a); + try self.rig.s.synchronize(); + } + + pub fn count(self: *EpPicks) usize { + return self.shape.rows * 4; + } + + pub fn reference(_: *EpPicks, want: []f64, mag: []f64) !void { + @memset(want, 0); + @memset(mag, 0); + } + + /// Every pick by the rule exactly. + pub fn check(self: *EpPicks, _: g.Gate) !?[]const u8 { + const got = std.mem.bytesAsSlice(i32, @as([]align(1) const u8, try self.rig.outputBytes())); + for (self.picks.items, 0..) |p, i| { + const want: i32 = if (p >= first and p < first + local) p - first else local; + if (got[i] != want) return try std.fmt.bufPrint(&self.msg, "pick {d} ({d}) mapped to {d}, not {d}", .{ i, p, got[i], want }); + } + return null; + } +}; + +/// The E4M3 code of an integer the stress weights use (1 to 7 and their negatives). +fn codeOf(v: f32) !u8 { + for (0..256) |i| if (core.quant.fp8.e4m3(@intCast(i)) == v) return @intCast(i); + return error.NotE4m3; +} + +/// The checkpoint's two kernel families, their kernels and controls appended to `list`. +pub const Set = struct { + gpa: std.mem.Allocator, + lane: q.Lane, + fp8: ex.Fp8Experts, + mod: cuda.Module, + batch: cuda.Module, + ep_picks: cuda.Function = undefined, + router: grouped.Router = undefined, + bf16: [3]Bf16 = undefined, + down: [3]Fp8Down = undefined, + picks: [2]EpPicks = undefined, + + pub fn init(gpu: check.Gpu, s: cuda.Stream, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel)) !*Set { + const set = try gpu.gpa.create(Set); + errdefer gpu.gpa.destroy(set); + const cap = try gpu.ctx.capability(); + const sms: usize = @intCast(try gpu.ctx.attribute(.multiprocessor_count)); + var lane = try q.Lane.load(gpu.d, cap / 10); + errdefer lane.unload(); + var fp8 = try ex.Fp8Experts.load(gpu.d, sms); + errdefer fp8.unload(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); + errdefer mod.unload(); + var batch = try cuda.Module.load(gpu.d, cuda.kernels.flashnext_batch); + errdefer batch.unload(); + set.* = .{ .gpa = gpu.gpa, .lane = lane, .fp8 = fp8, .mod = mod, .batch = batch }; + set.ep_picks = try set.batch.function("fn_ep_picks"); + set.router = try grouped.Router.resolve(set.mod); + const r: Rig = .{ .gpu = gpu, .s = s }; + set.bf16 = .{ .{ .rig = r, .lane = &set.lane, .format = .bf16 }, .{ .rig = r, .lane = &set.lane, .format = .f32 }, .{ .rig = r, .lane = &set.lane, .format = .bf16, .control = true } }; + set.down = .{ .{ .rig = r, .set = set, .prompt = false }, .{ .rig = r, .set = set, .prompt = true }, .{ .rig = r, .set = set, .prompt = false, .control = true } }; + for (&set.bf16) |*x| try list.append(a, x.kernel()); + for (&set.down) |*x| try list.append(a, x.kernel()); + set.picks = .{ .{ .rig = r, .set = set }, .{ .rig = r, .set = set, .control = true } }; + for (&set.picks) |*x| try list.append(a, x.kernel()); + return set; + } + + pub fn deinit(set: *Set) void { + for (&set.picks) |*x| x.deinit(); + for (&set.down) |*x| x.deinit(); + for (&set.bf16) |*x| x.deinit(); + set.batch.unload(); + set.mod.unload(); + set.fp8.unload(); + set.lane.unload(); + set.gpa.destroy(set); + } +}; + +/// These kernels' gates 1-4 and their controls alone (tf-cuda-test flashnext-modelopt-contract). +pub fn run(gpu: check.Gpu) !void { + const a = gpu.gpa; + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var list: std.ArrayList(g.Kernel) = .empty; + defer list.deinit(a); + const set = try Set.init(gpu, s, a, &list); + defer set.deinit(); + var t: g.Tally = .{}; + try g.runAll(a, list.items, 0x6b63, &t); + std.debug.print("SUMMARY {d} kernels: {d} pass gates 1-4, {d} fail; {d} controls: {d} red as planted, {d} missed\n", .{ t.kernels, t.passed, t.failed, t.controls, t.caught, t.missed }); + try check.expect(t.problems() == 0, "flashnext-modelopt-contract: {d} gate failures, {d} controls missed", .{ t.failed, t.missed }); +} diff --git a/zig/tests/cuda/contract_fn_pages.zig b/zig/tests/cuda/contract_fn_pages.zig index c0e10e2f06..513707ac29 100644 --- a/zig/tests/cuda/contract_fn_pages.zig +++ b/zig/tests/cuda/contract_fn_pages.zig @@ -228,10 +228,10 @@ const Bufs = struct { proj: u64, pos: u64, qw: u64, kw: u64, iw: u64, ikw: u64, /// cuda_batch.zig's attnRows over `n` rows from `pos0`: prep, pooling, each block's scores and selection, attention. fn pass(g: *const flash.Glue, s: cuda.Stream, select: cuda.Function, pos0: usize, n: usize, nb: usize, b: Bufs, pg: flash.Pg, image: ?media.View) !void { if (image) |im| { - try im.attn(s, b.proj, b.pos, b.qw, b.kw, b.iw, b.inv, b.q, b.kc, b.vc, b.iq, b.ikc, eps, n, pg); + try im.attn(s, b.proj, b.pos, b.qw, b.kw, b.iw, b.inv, b.q, b.kc, b.vc, b.iq, b.ikc, eps, n, .full, pg); try im.pooled(s, b.ikc, b.pooled, b.pos, b.ikw, b.inv, eps, n, pg); } else { - try g.attnPrep(s, b.proj, b.pos, b.qw, b.kw, b.iw, b.inv, b.q, b.kc, b.vc, b.iq, b.ikc, eps, n, pg); + try g.attnPrep(s, b.proj, b.pos, b.qw, b.kw, b.iw, b.inv, b.q, b.kc, b.vc, b.iq, b.ikc, eps, n, .full, pg); try g.qsaPool(s, b.ikc, b.pooled, b.pos, b.ikw, b.inv, eps, n, pg); } const fused = n > att_rows; @@ -251,11 +251,11 @@ fn pass(g: *const flash.Glue, s: cuda.Stream, select: cuda.Function, pos0: usize try cuda.launch.launch(select, .{ .grid = .{ .x = @intCast(m) }, .block = .{ .x = 1024 } }, s, &a); if (!fused) { const chunks = cdiv(@min(ends, dense_keys), 512); - try g.attnChunks(s, b.q + r0 * q_heads * head_dim * 2, b.kc, b.vc, pos, b.po, b.pm, b.pl, b.ids, b.nk, b.spr, m, kv_heads, chunks, true, pg); - try g.attnMerge(s, b.po, b.pm, b.pl, pos, b.ao + r0 * out_k * 2, b.nk, b.spr, m, kv_heads, true); + try g.attnChunks(s, b.q + r0 * q_heads * head_dim * 2, b.kc, b.vc, pos, b.po, b.pm, b.pl, b.ids, b.nk, b.spr, m, .full, chunks, true, pg); + try g.attnMerge(s, b.po, b.pm, b.pl, pos, b.ao + r0 * out_k * 2, b.nk, b.spr, m, .full, true); } } - if (fused) try g.attnFused(s, b.q + g0 * q_heads * head_dim * 2, b.kc, b.vc, b.pos + (g0 / att_rows) * 16, b.ao + g0 * out_k * 2, b.ids, b.nk, b.spr, gm, kv_heads, pg); + if (fused) try g.attnFused(s, b.q + g0 * q_heads * head_dim * 2, b.kc, b.vc, b.pos + (g0 / att_rows) * 16, b.ao + g0 * out_k * 2, b.ids, b.nk, b.spr, gm, .full, pg); } } diff --git a/zig/tests/cuda/contract_fn_pages_media.zig b/zig/tests/cuda/contract_fn_pages_media.zig index 306fa7c6be..0f8c31f049 100644 --- a/zig/tests/cuda/contract_fn_pages_media.zig +++ b/zig/tests/cuda/contract_fn_pages_media.zig @@ -99,7 +99,7 @@ test "paged media contract: production wrappers select both media symbols with f for (plans) |p| { Capture.calls = 0; const v = at(&d, f, 0x8000, p); - try v.attn(s, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1e-6, 16, pg); + try v.attn(s, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1e-6, 16, .full, pg); try v.pooled(s, 1, 2, 3, 4, 5, 1e-6, 16, pg); try std.testing.expectEqualSlices(usize, &.{ 3, 4 }, &Capture.symbols); for (0..2) |i| { diff --git a/zig/tests/cuda/contract_fn_pages_sums.zig b/zig/tests/cuda/contract_fn_pages_sums.zig index b5fbed9fc8..e7be092395 100644 --- a/zig/tests/cuda/contract_fn_pages_sums.zig +++ b/zig/tests/cuda/contract_fn_pages_sums.zig @@ -237,7 +237,7 @@ const Runner = struct { try self.glue.qsaScores(self.rig.s, p[1], pooled, p[0], p[8], blocks, self.rows, blocks, self.pg); } else if (self.isAttention()) { if (self.mode == .fused) { - try self.glue.attnFused(self.rig.s, p[1], kc, vc, p[0], p[8], p[2], p[3], p[4], self.rows, kh, self.pg); + try self.glue.attnFused(self.rig.s, p[1], kc, vc, p[0], p[8], p[2], p[3], p[4], self.rows, .full, self.pg); } else { var args: cuda.Args = .{}; for ([_]u64{ p[1], kc, vc, p[0], p[5], p[6], p[7], p[2], p[3], p[4] }) |v| args.add(v); @@ -249,7 +249,7 @@ const Runner = struct { const ahead = self.mode == .ahead; const shared: usize = (16 + (if (ahead) @as(usize, 128) else 32)) * (hd + 8) * 2 + 16 * 40 * 2; try self.launch(if (ahead) self.glue.f.chunks_ahead_paged else self.glue.f.chunks_one_paged, .{ self.rows, kh, 4 }, 128, shared, &args); - try self.glue.attnMerge(self.rig.s, p[5], p[6], p[7], p[0], p[8], p[3], p[4], self.rows, kh, true); + try self.glue.attnMerge(self.rig.s, p[5], p[6], p[7], p[0], p[8], p[3], p[4], self.rows, .full, true); } } else { var args: cuda.Args = .{}; diff --git a/zig/tests/cuda/contract_qmmf_prompt.zig b/zig/tests/cuda/contract_qmmf_prompt.zig new file mode 100644 index 0000000000..cd87eea157 --- /dev/null +++ b/zig/tests/cuda/contract_qmmf_prompt.zig @@ -0,0 +1,261 @@ +//! Four BF16 prompt instances through replacement, nonfinite, alias and FP32-sum gates. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const Rig = @import("contract_rig.zig").Rig; +const c = core.contract.spec; +const g = core.contract.gates; +const Dense = core.contract.dense.Dense; +const q = cuda.qmmf; +const cases = @import("qmmf_prompt_tests.zig"); + +const unsplit = [_]c.Shape{ + .{ .rows = 513, .n = 2560, .k = 640 }, + .{ .rows = 512, .n = 10240, .k = 320 }, + .{ .rows = 512, .n = 129, .k = 512 }, + .{ .rows = 512, .n = 129, .k = 2560 }, + .{ .rows = 513, .n = 8240, .k = 2560 }, + .{ .rows = 1536, .n = 7296, .k = 2560 }, + .{ .rows = 4218, .n = 8240, .k = 2560 }, + .{ .rows = 512, .n = 16480, .k = 2560 }, + .{ .rows = 512, .n = 13952, .k = 2560 }, +}; +const fused_shapes = [_]c.Shape{ + .{ .rows = 512, .n = 129, .k = 2048 }, + .{ .rows = 512, .n = 129, .k = 6144 }, + .{ .rows = 513, .n = 2560, .k = 3072 }, + .{ .rows = 4218, .n = 2560, .k = 3072 }, + .{ .rows = 512, .n = 2560, .k = 6144 }, + .{ .rows = 513, .n = 640, .k = 2560 }, + .{ .rows = 512, .n = 324, .k = 10240 }, + .{ .rows = 512, .n = 10240, .k = 2560 }, + .{ .rows = 512, .n = 2560, .k = 2560 }, +}; +const fused_fp32 = fused_shapes ++ [_]c.Shape{ + .{ .rows = 1024, .n = 2560, .k = 2560 }, // MTP fc_hidden: four streams of a 256-token piece + .{ .rows = 512, .n = 320, .k = 10240 }, + .{ .rows = 1536, .n = 320, .k = 10240 }, // final mixer under .all logits; padded N is 384 +}; + +fn slices(shape: c.Shape, fused: bool) u32 { + return if (fused) q.splitK(shape.n, shape.k) else 1; +} + +fn depth(shape: c.Shape) usize { + return shape.k / 16 + 24; +} + +pub fn spec(out: q.Out, fused: bool, control: bool) c.Spec { + return .{ + .name = if (control) "control: qmmf BF16 prompt scaled 1.5x" else switch (out) { + .bf16 => if (fused) "cuda qmmf BF16 prompt 128x128 bf16 fused" else "cuda qmmf BF16 prompt 128x128 bf16 unsplit", + .f32 => if (fused) "cuda qmmf BF16 prompt 128x128 fp32 fused" else "cuda qmmf BF16 prompt 128x128 fp32 unsplit", + }, + .accum = .f32, + .depth = depth, + .ops = 1, + .out = if (out == .f32) .f32 else .bf16, + .replaces = "qmmf BF16 64x64 with identical pinned K slices", + .tiles = .{ .{ .size = 128 }, .{ .size = 128, .min = 128 }, .{ .size = if (fused) 2048 else 512, .step = 256, .min = if (fused) 2048 else 512 } }, + .served = if (fused) (if (out == .f32) &fused_fp32 else &fused_shapes) else &unsplit, + }; +} + +pub const Product = struct { + rig: Rig, + lane: *const q.Lane, + format: q.Out, + fused: bool, + control: bool = false, + d: ?Dense = null, + w: q.Weight = undefined, + x: u64 = 0, + out: u64 = 0, + base: u64 = 0, + part: u64 = 0, + + pub fn deinit(self: *Product) void { + if (self.d) |d| d.deinit(); + self.rig.deinit(); + } + const vt = g.hooks(Product); + pub fn kernel(self: *Product) g.Kernel { + return .{ .spec = spec(self.format, self.fused, self.control), .ctx = self, .vt = &vt, .control = if (self.control) .shapes else null }; + } + fn width(self: *const Product) usize { + return if (self.format == .f32) 4 else 2; + } + pub fn prepare(self: *Product, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (self.d) |d| d.deinit(); + self.d = null; + self.rig.reset(); + if (shape.rows == 0 or shape.n < 128 or shape.k % 64 != 0 or shape.k < (if (self.fused) @as(usize, 2048) else 320)) return false; + const a = self.rig.gpu.gpa; + const d = try Dense.init(a, shape, 1); + self.d = d; + d.fillX(fill, seed, 16); + const weights = try a.alloc(u16, shape.n * shape.k); + defer a.free(weights); + var prng = std.Random.DefaultPrng.init(seed ^ 0xb16); + const random = prng.random(); + for (weights) |*v| v.* = c.toBf16(if (fill == .stress) @floatFromInt(1 + random.uintLessThan(u32, 7)) else random.float(f32) * 4 - 2); + if (fill == .stress) for (0..shape.n) |j| for (shape.k / 2..shape.k) |i| { + weights[j * shape.k + i] = weights[j * shape.k + i - shape.k / 2] | 0x8000; + }; + for (d.w, weights) |*o, v| o.* = c.fromBf16(v); + const npad = cuda.modelopt.padded(shape.n); + const packed_weights = try a.alloc(u16, npad * shape.k); + defer a.free(packed_weights); + q.packBf16(packed_weights, weights, shape.n, shape.k); + const codes = try self.rig.input(std.mem.sliceAsBytes(packed_weights)); + const sk = slices(shape, self.fused); + self.w = try q.pinned(cuda.modelopt.bf16Weight(codes, shape.n, shape.k), sk); + self.x = try self.rig.input(std.mem.sliceAsBytes(d.x)); + const bytes = self.count() * self.width(); + self.out = try self.rig.output(bytes); + self.base = try self.rig.scratch(bytes); + const partials = self.lane.partBytes(shape.rows, self.w); + self.part = if (partials == 0) 0 else try self.rig.scratch(partials); + if (shape.rows >= q.prompt_rows and (!q.promptEligible(shape.rows, self.w) or self.lane.prompt == null)) return error.PromptPathNotSelected; + return true; + } + pub fn run(self: *Product) !void { + const shape = self.d.?.shape; + var w = self.w; + if (self.control) w.scale = 1.5; + if (shape.rows >= q.prompt_rows) { + try self.lane.matmulOut(self.rig.s, self.x, shape.k, shape.rows, w, self.out, self.format, self.part); + } else { + // The harness shrinks single-shape gates; force the same instance instead of silently testing decode. + const band: usize = @max(1, @min((shape.rows + 127) / 128, (12 << 20) / (128 * shape.k * 2))); + var args: cuda.Args = .{}; + args.add(self.x); + args.add(w.codes); + args.add(w.scales); + args.add(w.scale); + args.add(self.out); + args.add(@as(u64, 0)); + for ([_]usize{ shape.rows, shape.n, shape.k, w.sk, w.npad, shape.k, band }) |v| args.add(@as(i32, @intCast(v))); + try cuda.launch.launch(self.lane.prompt.?[@backingInt(self.format)][@intFromBool(self.fused)], .{ + .grid = .{ .x = @intCast(((shape.rows + 127) / 128) * ((shape.n + 127) / 128)) }, + .block = .{ .x = 256 }, + .shared = q.prompt_shared, + }, self.rig.s, &args); + } + try self.rig.s.synchronize(); + } + pub fn replaced(self: *Product) ![]const u8 { + var lane = self.lane.*; + lane.prompt = null; + const shape = self.d.?.shape; + try lane.matmulOut(self.rig.s, self.x, shape.k, shape.rows, self.w, self.base, self.format, self.part); + return self.rig.bytes(self.base, self.count() * self.width()); + } + pub fn count(self: *Product) usize { + return self.d.?.outputs(); + } + pub fn reference(self: *Product, want: []f64, magnitude: []f64) !void { + try self.d.?.reference(want, magnitude); + } + pub fn simulate(self: *Product, acc: c.Accum, order: g.Order, out: []f64) !void { + try self.d.?.simulate(acc, order, out); + } +}; + +pub const Registrations = struct { + products: [5]Product, + pub fn init(gpu: check.Gpu, s: cuda.Stream, lane: *const q.Lane) Registrations { + var r: Registrations = undefined; + for ([_]q.Out{ .bf16, .f32 }, 0..) |out, oi| for ([_]bool{ false, true }, 0..) |fused, fi| { + r.products[oi * 2 + fi] = .{ .rig = .{ .gpu = gpu, .s = s }, .lane = lane, .format = out, .fused = fused }; + }; + r.products[4] = .{ .rig = .{ .gpu = gpu, .s = s }, .lane = lane, .format = .bf16, .fused = true, .control = true }; + return r; + } + pub fn deinit(self: *Registrations) void { + for (&self.products) |*p| p.deinit(); + } + pub fn append(self: *Registrations, cap: u32, a: std.mem.Allocator, kernels: *std.ArrayList(g.Kernel)) !void { + if (cap / 10 != 12) return; + for (&self.products) |*p| try kernels.append(a, p.kernel()); + } +}; + +pub fn run(gpu: check.Gpu) !void { + const cap = try gpu.ctx.capability(); + try check.expect(cap / 10 == 12, "qmmf prompt contract requires sm_12x", .{}); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var lane = try q.Lane.load(gpu.d, cap / 10); + defer lane.unload(); + var regs = Registrations.init(gpu, s, &lane); + defer regs.deinit(); + var kernels: std.ArrayList(g.Kernel) = .empty; + defer kernels.deinit(gpu.gpa); + try regs.append(cap, gpu.gpa, &kernels); + var tally: g.Tally = .{}; + try g.runAll(gpu.gpa, kernels.items, 0xb16, &tally); + std.debug.print("SUMMARY qmmf prompt gates 1-4: {d} pass, {d} fail; {d} controls red, {d} missed\n", .{ tally.passed, tally.failed, tally.caught, tally.missed }); + try check.expect(tally.problems() == 0, "qmmf prompt contract failed", .{}); +} + +test "qmmf prompt contract registers four production instances, served prompts and a replacement control" { + const a = std.testing.allocator; + var lane: q.Lane = undefined; + var regs = Registrations.init(undefined, undefined, &lane); + var kernels: std.ArrayList(g.Kernel) = .empty; + defer kernels.deinit(a); + try regs.append(89, a, &kernels); + try std.testing.expectEqual(@as(usize, 0), kernels.items.len); + try regs.append(121, a, &kernels); + try std.testing.expectEqual(@as(usize, 5), kernels.items.len); + inline for (.{ q.Out.bf16, q.Out.f32 }, 0..) |out, oi| inline for (.{ false, true }, 0..) |fused, fi| { + const k = kernels.items[oi * 2 + fi]; + try std.testing.expect(k.spec.replaces != null and k.vt.replaced != null); + try std.testing.expect(k.vt.simulate != null and k.control == null); + const symbol = q.promptSymbol(out, fused); + const record = for (core.contract.registry.all_cuda) |r| { + if (std.mem.eql(u8, r.launch.symbol, symbol)) break r; + } else return error.MissingPromptRecord; + try std.testing.expectEqual(q.prompt_shared, record.launch.dynamic); + try std.testing.expectEqual(@as(u32, 256), record.launch.threads); + try std.testing.expectEqual(@as(u32, 120), record.launch.min_sm); + try std.testing.expect(record.launch.opt_in); + for (k.spec.served) |shape| { + const sk = slices(shape, fused); + const w = try q.pinned(cuda.modelopt.bf16Weight(0, shape.n, shape.k), sk); + try std.testing.expect(shape.rows >= q.prompt_rows and q.promptEligible(shape.rows, w)); + try std.testing.expectEqual(fused, q.planFor(shape.rows, w, true).fused); + } + }; + try std.testing.expectEqual(c.Gate.shapes, kernels.items[4].control.?); +} + +test "qmmf prompt contract covers current normal operators, actual fused slices and MTP hidden rows" { + var seen: [9]bool = @splat(false); + for (cases.normal) |projection| { + const sk = q.splitK(projection.n, projection.k); + const s = spec(projection.out, sk > 1, false); + var found = false; + for (s.served) |shape| { + if (shape.n == projection.n and shape.k == projection.k) { + try std.testing.expectEqual(sk, slices(shape, sk > 1)); + found = true; + } + } + try std.testing.expect(found); + seen[sk] = true; + } + for ([_]usize{ 1, 2, 4, 8 }) |sk| try std.testing.expect(seen[sk]); + for ([_]usize{ 512, 1024 }) |rows_| { + var found = false; + for (spec(.f32, true, false).served) |shape| { + if (shape.rows == rows_ and shape.n == cases.mtp_hidden.n and shape.k == cases.mtp_hidden.k) found = true; + } + try std.testing.expect(found); + } + // Vocabulary probes stay in the bounded byte runner, rather than expanding FP64 or repeat gates to full heads. + for ([_]q.Out{ .bf16, .f32 }) |out| for (spec(out, false, false).served) |shape| + try std.testing.expect(shape.n < 124160); +} diff --git a/zig/tests/cuda/contract_qmmf_shared.zig b/zig/tests/cuda/contract_qmmf_shared.zig new file mode 100644 index 0000000000..827ef93918 --- /dev/null +++ b/zig/tests/cuda/contract_qmmf_shared.zig @@ -0,0 +1,159 @@ +//! Strided shared-down stores keep each BF16-weight product's FP32 bytes and leave every other slot unchanged. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const Rig = @import("contract_rig.zig").Rig; +const c = core.contract.spec; +const g = core.contract.gates; +const Dense = core.contract.dense.Dense; +const q = cuda.qmmf; + +fn depth(s: c.Shape) usize { + return s.k / 16 + 24; +} + +pub const Product = struct { + rig: Rig, + lane: *const q.Lane, + index: usize, + control: bool = false, + d: ?Dense = null, + w: q.Weight = undefined, + x: u64 = 0, + slots: u64 = 0, + out: u64 = 0, + base: u64 = 0, + + pub fn deinit(self: *Product) void { + if (self.d) |d| d.deinit(); + self.rig.deinit(); + } + const vt = g.hooks(Product); + + pub fn kernel(self: *Product) g.Kernel { + const names = [_][]const u8{ "cuda qmmf shared stride half16", "cuda qmmf shared stride 16", "cuda qmmf shared stride 32", "cuda qmmf shared stride 64" }; + const bm: usize = if (self.index < 2) 16 else if (self.index == 2) 32 else 64; + return .{ .ctx = self, .vt = &vt, .control = if (self.control) .shapes else null, .spec = .{ + .name = if (self.control) "control: shared down scaled 1.5x" else names[self.index], + .accum = .f32, + .depth = depth, + .ops = 1, + .out = .f32, + .replaces = "contiguous qmmf BF16-weight FP32 down before shared scatter", + .tiles = .{ .{ .size = bm }, .{ .size = if (self.index == 0) 32 else 64, .step = 32, .min = 32 }, .{ .size = 640, .step = 64, .min = 64 } }, + .served = switch (self.index) { + 0, 1 => &.{ .{ .rows = 1, .n = 2560, .k = 640 }, .{ .rows = 16, .n = 2560, .k = 640 } }, + 2 => &.{.{ .rows = 32, .n = 2560, .k = 640 }}, + else => &.{ .{ .rows = 65, .n = 2560, .k = 640 }, .{ .rows = 1536, .n = 2560, .k = 640 }, .{ .rows = 4218, .n = 2560, .k = 640 } }, + }, + } }; + } + pub fn prepare(self: *Product, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (self.d) |d| d.deinit(); + self.d = null; + self.rig.reset(); + if (shape.rows == 0 or shape.n == 0 or shape.k == 0 or shape.k % 64 != 0) return false; + const a = self.rig.gpu.gpa; + const d = try Dense.init(a, shape, 1); + self.d = d; + d.fillX(fill, seed, 256); + const weights = try a.alloc(u16, shape.n * shape.k); + defer a.free(weights); + var prng = std.Random.DefaultPrng.init(seed ^ 0x5107); + const rng = prng.random(); + for (weights) |*v| v.* = c.toBf16(if (fill == .stress) @floatFromInt(1 + rng.uintLessThan(u32, 7)) else rng.float(f32) - 0.5); + if (fill == .stress) for (0..shape.n) |j| for (shape.k / 2..shape.k) |i| { + weights[j * shape.k + i] = weights[j * shape.k + i - shape.k / 2] | 0x8000; + }; + for (d.w, weights) |*v, w| v.* = c.fromBf16(w); + const npad = cuda.modelopt.padded(shape.n); + const tiled = try a.alloc(u16, npad * shape.k); + defer a.free(tiled); + q.packBf16(tiled, weights, shape.n, shape.k); + self.w = cuda.modelopt.bf16Weight(try self.rig.input(std.mem.sliceAsBytes(tiled)), shape.n, shape.k); + self.x = try self.rig.input(std.mem.sliceAsBytes(d.x)); + self.slots = try self.rig.scratch(self.count() * 11 * 4); + self.out = try self.rig.output(self.count() * 4); + self.base = try self.rig.scratch(self.count() * 4); + return true; + } + pub fn run(self: *Product) !void { + const s = self.d.?.shape; + const bm: usize = if (self.index < 2) 16 else if (self.index == 2) 32 else 64; + const bn: usize = if (self.index == 0) 32 else 64; + const band = @max(1, @min((s.rows + bm - 1) / bm, (12 << 20) / (bm * s.k * 2))); + try self.rig.zero(self.slots, self.count() * 11 * 4); + var args: cuda.Args = .{}; + args.add(self.x); + args.add(self.w.codes); + args.add(if (self.control) @as(f32, 1.5) else self.w.scale); + args.add(self.slots + 10 * s.n * 4); + for ([_]usize{ s.rows, s.n, s.k, self.w.npad, s.k, band, 11 * s.n }) |v| args.add(@as(i32, @intCast(v))); + try cuda.launch.launch(self.lane.shared_down[self.index], .{ + .grid = .{ .x = @intCast(((s.rows + bm - 1) / bm) * ((s.n + bn - 1) / bn)) }, + .block = .{ .x = 128 }, + .shared = if (self.index == 0) q.halfShared(.bf16) else q.sharedBytes(.bf16, @intCast(bm)), + }, self.rig.s, &args); + const bytes = try self.rig.bytes(self.slots, self.count() * 11 * 4); + for (0..s.rows) |row| for (bytes[row * 11 * s.n * 4 ..][0 .. 10 * s.n * 4]) |v| if (v != 0) return error.SharedDownTouchedRoutedSlot; + for (0..s.rows) |row| try self.rig.gpu.d.check(self.rig.gpu.d.api.cuMemcpyDtoDAsync_v2( + self.out + row * s.n * 4, + self.slots + (row * 11 + 10) * s.n * 4, + s.n * 4, + self.rig.s.handle, + ), "cuMemcpyDtoDAsync"); + try self.rig.s.synchronize(); + } + pub fn replaced(self: *Product) ![]const u8 { + const lane = self.lane.*; + const s = self.d.?.shape; + try lane.matmulF32(self.rig.s, self.x, s.k, s.rows, self.w, self.base, 0); + return self.rig.bytes(self.base, self.count() * 4); + } + pub fn count(self: *Product) usize { + return self.d.?.outputs(); + } + pub fn reference(self: *Product, want: []f64, mag: []f64) !void { + try self.d.?.reference(want, mag); + } + pub fn simulate(self: *Product, acc: c.Accum, order: g.Order, out: []f64) !void { + try self.d.?.simulate(acc, order, out); + } +}; + +/// Owned registration storage keeps kernel contexts valid throughout the full contract run. +pub const Registrations = struct { + products: [5]Product, + + pub fn create(gpu: check.Gpu, stream: cuda.Stream, lane: *const q.Lane, a: std.mem.Allocator, kernels: *std.ArrayList(g.Kernel)) !*Registrations { + const self = try gpu.gpa.create(Registrations); + errdefer gpu.gpa.destroy(self); + for (&self.products, 0..) |*p, i| p.* = .{ .rig = .{ .gpu = gpu, .s = stream }, .lane = lane, .index = i % 4, .control = i == 4 }; + errdefer for (&self.products) |*p| p.deinit(); + for (&self.products) |*p| try kernels.append(a, p.kernel()); + return self; + } + + pub fn deinit(self: *Registrations) void { + const a = self.products[0].rig.gpu.gpa; + for (&self.products) |*p| p.deinit(); + a.destroy(self); + } +}; + +/// Complete gates 1-4 and planted control for the strided shared-down operator. +pub fn run(gpu: check.Gpu) !void { + var stream = try cuda.Stream.init(gpu.d, false); + defer stream.deinit(); + var lane = try q.Lane.load(gpu.d, (try gpu.ctx.capability()) / 10); + defer lane.unload(); + var kernels: std.ArrayList(g.Kernel) = .empty; + defer kernels.deinit(gpu.gpa); + const regs = try Registrations.create(gpu, stream, &lane, gpu.gpa, &kernels); + defer regs.deinit(); + var tally: g.Tally = .{}; + try g.runAll(gpu.gpa, kernels.items, 0x5107, &tally); + std.debug.print("SUMMARY shared down gates 1-4: {d} pass, {d} fail; {d} controls red, {d} missed\n", .{ tally.passed, tally.failed, tally.caught, tally.missed }); + try check.expect(tally.problems() == 0, "shared down contract failed", .{}); +} diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index b9158471f4..413575a471 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -57,7 +57,7 @@ const qkr_k = @import("contract_qk_round.zig"); const fact_k = @import("contract_flashnext_fact.zig"); const route_k = @import("contract_route_rows.zig"); const fexp_k = @import("contract_flashnext_fexp.zig"); -const glue_k = @import("contract_flashnext_glue.zig"); +const fn_groups_k = @import("contract_flashnext_groups.zig"); const q35_contiguous = @import("contract_q35_contiguous.zig"); const contract = core.contract; const g = contract.gates; @@ -484,27 +484,9 @@ pub fn run(gpu: check.Gpu) !void { defer attn.unload(); var hcf_mod = try cuda.Module.load(gpu.d, kk.flashnext_hcf); defer hcf_mod.unload(); - var fglue = try @import("flashnext").Glue.load(gpu.d); - defer fglue.deinit(); - const glue_fns: glue_k.Fns = .{ - .tpr = try batch.function("fn_gdn_chain_tpr"), - .rows_to = try batch.function("fn_gdn_rows_to"), - .tile = try attn.function("_ZN8qsa_tile6scoresILi8ELi4EEEvPK13__nv_bfloat16S3_PKiPfiiii"), - .scores = try attn.function("fn_qsa_scores"), - .hcf_down = try hcf_mod.function(ftk.cubin_hcf_down), - .hcf_up = try hcf_mod.function(ftk.cubin_hcf_up), - .front = try batch.function("fn_gdn_front"), - .norm = try batch.function("fn_gdn_norm"), - .fused = try batch.function("fn_gdn_fused"), - }; - try glue_fns.fused.allowDynamicShared(fparts.gdn_fused_smem); - try glue_fns.tpr.allowDynamicShared(fparts.chain_tpr_smem); - try glue_fns.hcf_down.allowDynamicShared(ftk.hcf_down_smem); - try glue_fns.hcf_up.allowDynamicShared(ftk.hcf_up_smem); - var glues: [4]glue_k.Glue = undefined; - for (std.enums.values(glue_k.Kind), &glues) |kind, *m| m.* = .{ .rig = .{ .gpu = gpu, .s = s }, .kind = kind, .f = &glue_fns, .glue = &fglue }; - defer for (&glues) |*m| m.deinit(); - for (&glues) |*m| try kernels.append(a, m.kernel()); + var fn_groups: fn_groups_k.Set = .{}; + defer fn_groups.deinit(); + try fn_groups.add(gpu, s, batch, attn, hcf_mod, &kernels); var wrows: [wrows_k.cfgs.len]wrows_k.Rows = undefined; for (&wrows, wrows_k.cfgs) |*x, f| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = f, .g = &window_fns, .ctl = mods[5], .arena = .init(a) }; defer for (&wrows) |*x| x.deinit(); diff --git a/zig/tests/cuda/flashnext_attn_tests.zig b/zig/tests/cuda/flashnext_attn_tests.zig index 8e27d470a3..c2f783ce49 100644 --- a/zig/tests/cuda/flashnext_attn_tests.zig +++ b/zig/tests/cuda/flashnext_attn_tests.zig @@ -77,7 +77,7 @@ fn prep(r: *Rig, g: *const flash.Glue) !void { kcs[c] = try r.out(cap * kv_heads * head_dim * 2); vcs[c] = try r.out(cap * kv_heads * head_dim * 2); ikcs[c] = try r.out(cap * index_dim * 2); - try g.attnPrep(r.s, dp, dpos, dqw, dkw, diw, dinv, qs[c], kcs[c], vcs[c], iqs[c], ikcs[c], eps, m, .{}); + try g.attnPrep(r.s, dp, dpos, dqw, dkw, diw, dinv, qs[c], kcs[c], vcs[c], iqs[c], ikcs[c], eps, m, .full, .{}); } const last = counts.len - 1; const q = try r.back(u16, qs[last], most * q_heads * head_dim); @@ -227,8 +227,8 @@ fn attention(r: *Rig, g: *const flash.Glue, sparse: bool) !void { for (counts, 0..) |m, c| { outs[c] = try r.out(m * q_heads * head_dim * 2); const chunks = if (sparse) (idw + 511) / 512 else (pos0 + m + 511) / 512; - try g.attnChunks(r.s, dq, dkc, dvc, dpos, po, pm, pl, dids, dnk, dspr, m, kv_heads, chunks, true, .{}); - try g.attnMerge(r.s, po, pm, pl, dpos, outs[c], dnk, dspr, m, kv_heads, true); + try g.attnChunks(r.s, dq, dkc, dvc, dpos, po, pm, pl, dids, dnk, dspr, m, .full, chunks, true, .{}); + try g.attnMerge(r.s, po, pm, pl, dpos, outs[c], dnk, dspr, m, .full, true); } const got = try r.back(u16, outs[counts.len - 1], rows * q_heads * head_dim); const name = if (sparse) "attention sparse" else "attention dense"; @@ -266,7 +266,7 @@ fn attention(r: *Rig, g: *const flash.Glue, sparse: bool) !void { for (counts[0 .. counts.len - 1], 0..) |m, c| try r.sameRows(name, u16, outs[c], got[0 .. m * q_heads * head_dim]); // prompt rows' fused launch (every chunk merged in the program): the chunks then merge's bytes const fused = try r.out(rows * q_heads * head_dim * 2); - try g.attnFused(r.s, dq, dkc, dvc, dpos, fused, dids, dnk, dspr, rows, kv_heads, .{}); + try g.attnFused(r.s, dq, dkc, dvc, dpos, fused, dids, dnk, dspr, rows, .full, .{}); try r.sameRows(if (sparse) "attention sparse fused" else "attention dense fused", u16, fused, got); } diff --git a/zig/tests/cuda/flashnext_fidelity_tests.zig b/zig/tests/cuda/flashnext_fidelity_tests.zig index 3f3c793ce1..d80456b940 100644 --- a/zig/tests/cuda/flashnext_fidelity_tests.zig +++ b/zig/tests/cuda/flashnext_fidelity_tests.zig @@ -31,7 +31,7 @@ pub fn run(gpu: check.Gpu, model: []const u8, rank: usize, port: u16, fixture: [ defer glue.deinit(); const tri: flash.Tri = .{ .s = stream, .own = &glue }; const identity = try flash.rank_engine.identity(gpu.gpa, gpu.io, model, context, chunk); - const session = try flash.rank_engine.Session.init(gpu.gpa, gpu.io, "127.0.0.1", port, rank, identity); + const session = try flash.rank_engine.Session.init(gpu.gpa, gpu.io, "127.0.0.1", port, rank, identity, flash.rank_engine.linkWait()); const e = flash.prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &stream, model, 2, rank) catch |err| { session.poison(); session.deinit(); diff --git a/zig/tests/cuda/flashnext_fp8_experts_tests.zig b/zig/tests/cuda/flashnext_fp8_experts_tests.zig new file mode 100644 index 0000000000..8ec27578b4 --- /dev/null +++ b/zig/tests/cuda/flashnext_fp8_experts_tests.zig @@ -0,0 +1,217 @@ +//! Flash Next's block-FP8 MTP experts at the head's shapes: loader and moe_native.fp8, decode and prompt, vs FP64. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const rig = @import("flashnext_rig.zig"); + +const experts = 512; // routed FP8 experts +const entries = experts + 1; // routing entries: the shared slot last, which the FP8 call skips +const NI = 640; // width +const D = 2560; // dims +const slots = 11; // ten routed picks a row, then the shared slot +const blk = 128; +const most = 200; +const projs = [_][]const u8{ "gate_proj", "up_proj", "down_proj" }; + +/// 512 experts' E4M3 codes [expert][gate, up, down][n][k] and BF16 block scales [expert][proj][n/128][k/128]. +const Source = struct { + codes: []u8, + scales: []u16, + + const per = NI * D; // codes a projection (gate / up [NI, D], down [D, NI]) + const blocks = (NI / blk) * (D / blk); + + fn init(a: std.mem.Allocator, r: std.Random) !Source { + const s: Source = .{ .codes = try a.alloc(u8, experts * 3 * per), .scales = try a.alloc(u16, experts * 3 * blocks) }; + r.bytes(s.codes); + for (s.codes) |*c| if (c.* & 0x7f == 0x7f) { + c.* ^= 1; // E4M3's NaN codes: the checkpoint holds none + }; + for (s.scales) |*v| v.* = rig.bf16(0.001 + r.float(f32) * 0.01); + return s; + } + + fn dims(p: usize) [2]usize { + return if (p == 2) .{ D, NI } else .{ NI, D }; + } + + /// The f64 value of projection `p` of expert `e` at (row, col). + fn value(s: *const Source, lut: *const [256]f64, e: usize, p: usize, row: usize, col: usize) f64 { + const nk = dims(p); + return lut[s.codes[(e * 3 + p) * per + row * nk[1] + col]] * rig.bf(s.scales[(e * 3 + p) * blocks + (row / blk) * (nk[1] / blk) + col / blk]); + } + + pub fn has(_: *Source, name: []const u8) bool { + return std.mem.indexOf(u8, name, ".experts.") != null; + } + + pub fn get(s: *Source, name: []const u8) !core.safetensors.Tensor { + const marker = std.mem.indexOf(u8, name, ".experts.") orelse return error.MissingTensor; + const tail = name[marker + ".experts.".len ..]; + const dot = std.mem.indexOfScalar(u8, tail, '.') orelse return error.MissingTensor; + const e = try std.fmt.parseInt(usize, tail[0..dot], 10); + if (e >= experts) return error.MissingTensor; + for (projs, 0..) |proj, p| { + const rest = tail[dot + 1 ..]; + if (!std.mem.startsWith(u8, rest, proj)) continue; + const nk = dims(p); + if (std.mem.eql(u8, rest[proj.len..], ".weight")) return tensor(.f8_e4m3, nk, s.codes[(e * 3 + p) * per ..][0..per]); + if (std.mem.eql(u8, rest[proj.len..], ".weight_scale_inv")) return tensor(.bf16, .{ nk[0] / blk, nk[1] / blk }, std.mem.sliceAsBytes(s.scales[(e * 3 + p) * blocks ..][0..blocks])); + } + return error.MissingTensor; + } + + pub fn release(_: *Source) void {} +}; + +fn tensor(dtype: core.safetensors.DType, shape: [2]usize, bytes: []const u8) core.safetensors.Tensor { + var t: core.safetensors.Tensor = .{ .dtype = dtype, .rank = 2, .shape = @splat(1), .bytes = bytes }; + t.shape[0] = shape[0]; + t.shape[1] = shape[1]; + return t; +} + +/// SwiGLU as the kernels round it: gate and up to bf16, silu to bf16, the product (no clamp). +fn swiglu(g: f64, u: f64) f64 { + const gv = rig.rbf(g); + const uv = rig.rbf(u); + return rig.rbf(gv / (1 + @exp(-gv))) * uv; +} + +/// One thread's pairs vs FP64: gate-up from x, down from the kernel's bf16 activations (quant_ref_experts' bounds). +const Ref = struct { + s: *const Source, + lut: *const [256]f64, + x: []const u16, + picks: []const i32, + act: []const u16, + y: []const f32, + lo: usize, + hi: usize, + worst: [2]f64 = .{ 0, 0 }, + bad: [2]usize = .{ 0, 0 }, + + fn note(f: *Ref, which: usize, got: f64, want: f64, tol: f64) void { + const q = @abs(got - want) / tol; + if (!(q <= 1)) f.bad[which] += 1; + if (q > f.worst[which] or std.math.isNan(q)) f.worst[which] = if (std.math.isNan(q)) std.math.inf(f64) else q; + } + + fn run(f: *Ref) void { + for (f.lo..f.hi) |p| { + const e: usize = @intCast(f.picks[p]); + if (e == experts) continue; + const xr = f.x[(p / slots) * D ..][0..D]; + for (0..NI) |j| { + var gs: f64 = 0; + var us: f64 = 0; + var abs: f64 = 0; + for (0..D) |i| { + const xv = rig.bf(xr[i]); + const g = xv * f.s.value(f.lut, e, 0, j, i); + const u = xv * f.s.value(f.lut, e, 1, j, i); + gs += g; + us += u; + abs += @abs(g) + @abs(u); + } + const a = swiglu(gs, us); + // a bf16 step either side of each of the three roundings, and fp32 sums + f.note(0, rig.bf(f.act[p * NI + j]), a, @abs(a) * std.math.pow(f64, 2, -6) + abs * std.math.pow(f64, 2, -12) + 1e-30); + } + for (0..D) |j| { + var sum: f64 = 0; + var abs: f64 = 0; + for (0..NI) |i| { + const t = rig.bf(f.act[p * NI + i]) * f.s.value(f.lut, e, 2, j, i); + sum += t; + abs += @abs(t); + } + f.note(1, f.y[p * D + j], sum, abs * std.math.pow(f64, 2, -12) + 1e-30); + } + } + } +}; + +/// Every row's ten distinct routed experts, then the shared slot. +fn picksOf(a: std.mem.Allocator, r: std.Random, rows: usize) ![]i32 { + const out = try a.alloc(i32, rows * slots); + for (0..rows) |row| { + for (0..slots - 1) |j| { + var e: i32 = @intCast(r.uintLessThan(usize, experts)); + while (std.mem.indexOfScalar(i32, out[row * slots ..][0..j], e) != null) e = @intCast(r.uintLessThan(usize, experts)); + out[row * slots + j] = e; + } + out[row * slots + slots - 1] = experts; + } + return out; +} + +pub fn run(gpu: check.Gpu) !void { + // decode rows take tile 16 and prompt rows tile 64 at the default fact rows (64) + if (std.c.getenv("TF_FLASHNEXT_FACT_ROWS") != null) return check.expect(false, "flashnext-fp8-experts runs with TF_FLASHNEXT_FACT_ROWS unset", .{}); + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var rng = std.Random.DefaultPrng.init(0xf8e4); + var r = try rig.Rig.init(gpu, a, rng.random()); + defer r.deinit(); + var source = try Source.init(a, rng.random()); + var lut: [256]f64 = undefined; + for (&lut, 0..) |*v, c| v.* = core.quant.fp8.e4m3(@intCast(c)); + // the production loader (block-FP8 branch) and dispatch: the batch pass's kernels, the layer's view + var owned = try flash.modelopt_experts.loadShape(gpu.gpa, gpu.d, &source, "language_model.mtp.layers.0.mlp", 1, 0, D, NI, experts); + defer owned.free(); + const layer = owned.view(); + try check.expect(layer.format == .fp8g and layer.width == NI and layer.dims == D and layer.experts == experts, "the head's experts load as block FP8, {d} x {d}, {d} experts", .{ NI, D, experts }); + var k = try flash.batch_state.Kernels.load(gpu.d); + defer k.deinit(); + const x = r.bfs(most * D, 1.0); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + var worst: [2]f64 = .{ 0, 0 }; + var checked: usize = 0; + for ([_]usize{ 1, 3, 16, 64, most }) |rows| { + const pairs = rows * slots; + const picks = try picksOf(a, r.rng, rows); + const plan: cuda.grouped.Plan = .{ + .members = try r.out(pairs * 4), + .items = try r.out(((pairs + 15) / 16 + entries) * 12), + .counts = try r.out(8), + .rank = try r.out(pairs * 4), + .hist = try r.out(((pairs + cuda.grouped.small - 1) / cuda.grouped.small) * entries * 4), + }; + // every byte 0xff: a pair the call writes changes, the shared slot's must not + const act = try r.out(pairs * NI * 2); + const y = try r.out(pairs * D * 4); + try flash.moe_native.fp8(&k, r.s, try r.dev(std.mem.sliceAsBytes(picks)), dx, layer, plan, pairs, rows, y, act); + const got_act = try r.back(u16, act, pairs * NI); + const got_y = try r.back(f32, y, pairs * D); + for (picks, 0..) |e, p| if (e == experts) { + for (got_act[p * NI ..][0..NI]) |v| try check.expect(v == 0xffff, "{d} rows: the shared slot's activations untouched", .{rows}); + for (std.mem.sliceAsBytes(got_y[p * D ..][0..D])) |v| try check.expect(v == 0xff, "{d} rows: the shared slot's down rows untouched", .{rows}); + }; + const threads = @min(32, @max(1, std.Thread.getCpuCount() catch 1)); + var refs: [32]Ref = undefined; + var handles: [32]?std.Thread = @splat(null); + const share = (pairs + threads - 1) / threads; + for (0..threads) |t| { + refs[t] = .{ .s = &source, .lut = &lut, .x = x, .picks = picks, .act = got_act, .y = got_y, .lo = @min(pairs, t * share), .hi = @min(pairs, (t + 1) * share) }; + handles[t] = std.Thread.spawn(.{}, Ref.run, .{&refs[t]}) catch blk: { + refs[t].run(); + break :blk null; + }; + } + for (handles[0..threads]) |h| if (h) |th| th.join(); + var bad: [2]usize = .{ 0, 0 }; + for (refs[0..threads]) |f| for (0..2) |i| { + bad[i] += f.bad[i]; + worst[i] = @max(worst[i], f.worst[i]); + }; + const path = if (rows >= 64) "prompt (tile 64)" else "decode (tile 16)"; + try check.expect(bad[0] == 0, "fp8 head experts {s}, {d} rows: gate-up within FP64 tolerance ({d} outside, worst {d:.3})", .{ path, rows, bad[0], worst[0] }); + try check.expect(bad[1] == 0, "fp8 head experts {s}, {d} rows: fp32 down within FP64 tolerance ({d} outside, worst {d:.3})", .{ path, rows, bad[1], worst[1] }); + checked += rows * (slots - 1); + } + check.pass("flashnext-fp8-experts: 512 block-FP8 experts at {d} x {d} through moe_native.fp8, decode rows 1/3/16 and prompt rows 64/200; {d} routed pairs within FP64 (worst gate-up {d:.3}, down {d:.3} tolerances), the shared slot untouched", .{ NI, D, checked, worst[0], worst[1] }); +} diff --git a/zig/tests/cuda/flashnext_gdn_heads.zig b/zig/tests/cuda/flashnext_gdn_heads.zig new file mode 100644 index 0000000000..0d08e9e9cf --- /dev/null +++ b/zig/tests/cuda/flashnext_gdn_heads.zig @@ -0,0 +1,272 @@ +//! One rank's DeltaNet heads inside the full layout: the ranges each buffer keeps, their gather, the kernels' launches. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); + +pub const Heads = flash.Heads; +const parts = flash.batch_state; +const full = Heads.full; + +pub const dk = 128; +pub const dv = 128; +pub const taps = 4; +pub const eps: f32 = 1e-6; + +/// A DeltaNet block buffer, by what one of its rows holds. +pub const Part = enum { + /// A projection row: q, k and v channels, then z, b and a. + proj, + /// A conv window row: the q, k and v channels. + conv, + /// The conv taps, channel-major: one row of channels x taps. + taps, + /// a_log or dt_bias: one value a value head. + gates, + /// Normed q or k: dk values a key head. + keys, + /// v, the readouts or the gated-norm rows: dv values a value head. + values, + /// Decay and beta: two values a value head. + gb, + /// The recurrent state, one row: dv x dk values a value head. + state, + /// The gated norm's 32-wide sums: dv / 32 values a value head. + sums, +}; + +/// Elements of one row of `part` for the heads h. +pub fn width(part: Part, h: Heads) usize { + return switch (part) { + .proj => h.projWidth(), + .conv => h.convDim(), + .taps => h.convDim() * taps, + .gates => h.nv, + .keys => h.nk * dk, + .values => h.nv * dv, + .gb => h.nv * 2, + .state => h.nv * dv * dk, + .sums => h.nv * dv / 32, + }; +} + +/// Rows of `part` in a pass over m token rows: the conv window's taps - 1, one for the taps, gates and state. +pub fn rowsIn(part: Part, m: usize) usize { + return switch (part) { + .conv => taps - 1, + .taps, .gates, .state => 1, + else => m, + }; +} + +/// Bytes of a whole buffer of `part` (elements `elem` bytes each) over m token rows, for the heads h. +pub fn bytes(part: Part, elem: usize, m: usize, h: Heads) usize { + return rowsIn(part, m) * width(part, h) * elem; +} + +pub const Cut = struct { from: usize, len: usize }; + +pub const Cuts = struct { + buf: [6]Cut = undefined, + len: usize = 0, + + fn add(c: *Cuts, from: usize, len: usize) void { + c.buf[c.len] = .{ .from = from, .len = len }; + c.len += 1; + } + + pub fn slice(c: *const Cuts) []const Cut { + return c.buf[0..c.len]; + } +}; + +/// Rank r's ranges of a full-layout row of `part`, in the order its local row (heads h) holds them. +pub fn cuts(part: Part, h: Heads, r: usize) Cuts { + var c: Cuts = .{}; + switch (part) { + .proj, .conv, .taps => { + const t: usize = if (part == .taps) taps else 1; + const qk = h.nk * dk; + const vw = h.nv * dv; + c.add(r * qk * t, qk * t); + c.add((full.nk * dk + r * qk) * t, qk * t); + c.add((2 * full.nk * dk + r * vw) * t, vw * t); + if (part == .proj) { + const z = full.convDim(); + c.add(z + r * vw, vw); + c.add(z + full.nv * dv + r * h.nv, h.nv); + c.add(z + full.nv * dv + full.nv + r * h.nv, h.nv); + } + }, + else => { + const w = width(part, h); + c.add(r * w, w); + }, + } + return c; +} + +/// The full-layout row index of element i of rank r's local row of `part`. +pub fn fullIndex(part: Part, h: Heads, r: usize, i: usize) usize { + var left = i; + const cs = cuts(part, h, r); + for (cs.slice()) |c| { + if (left < c.len) return c.from + left; + left -= c.len; + } + unreachable; +} + +/// Rank r's elements (`elem` bytes each) of a full-layout buffer of `part`, in the local layout of heads h. +pub fn gather(a: std.mem.Allocator, src: []const u8, elem: usize, part: Part, h: Heads, r: usize) ![]u8 { + const fw = width(part, full) * elem; + const lw = width(part, h) * elem; + if (src.len % fw != 0) return error.LayoutLength; + const n = src.len / fw; + const out = try a.alloc(u8, n * lw); + const cs = cuts(part, h, r); + for (0..n) |row| { + var at = row * lw; + for (cs.slice()) |c| { + @memcpy(out[at..][0 .. c.len * elem], src[row * fw + c.from * elem ..][0 .. c.len * elem]); + at += c.len * elem; + } + } + return out; +} + +pub const Kind = enum { front, norm, rows, rows_to, tpr, fused }; + +/// The kernel of `kind`: the full layer's, or (half) its twin on one rank's 8 key and 24 value heads. +pub fn symbol(kind: Kind, half: bool) [:0]const u8 { + return switch (kind) { + .front => if (half) "fn_gdn_front_h8" else "fn_gdn_front", + .norm => if (half) "fn_gdn_norm_h8" else "fn_gdn_norm", + .rows => if (half) "fn_gdn_rows_h8" else "fn_gdn_rows", + .rows_to => if (half) "fn_gdn_rows_to_h8" else "fn_gdn_rows_to", + .tpr => if (half) "fn_gdn_chain_tpr_h8" else "fn_gdn_chain_tpr", + .fused => if (half) "fn_gdn_fused_h8" else "fn_gdn_fused", + }; +} + +/// A full kernel and its twin, by kind, from the flashnext_batch image (dynamic shared bytes allowed). +pub const Fns = struct { + full: [6]cuda.Function, + half: [6]cuda.Function, + + pub fn load(batch: cuda.Module) !Fns { + var f: Fns = undefined; + for (std.enums.values(Kind)) |k| { + const i = @backingInt(k); + f.full[i] = try batch.function(symbol(k, false)); + f.half[i] = try batch.function(symbol(k, true)); + const smem: u32 = switch (k) { + .tpr => parts.chain_tpr_smem, + .fused => parts.gdn_fused_smem, + else => continue, + }; + try f.full[i].allowDynamicShared(smem); + try f.half[i].allowDynamicShared(smem); + } + return f; + } + + /// Replace only the full-head oracle; candidate half handles stay initialized and owned by their original module. + pub fn withFull(f: Fns, batch: cuda.Module) !Fns { + var out = f; + for (std.enums.values(Kind)) |k| { + const i = @backingInt(k); + out.full[i] = try batch.function(symbol(k, false)); + const smem: u32 = switch (k) { + .tpr => parts.chain_tpr_smem, + .fused => parts.gdn_fused_smem, + else => continue, + }; + try out.full[i].allowDynamicShared(smem); + } + return out; + } + + /// The kernel of `kind` built for heads h. + pub fn of(f: *const Fns, kind: Kind, h: Heads) cuda.Function { + return if (h.eql(.full)) f.full[@backingInt(kind)] else f.half[@backingInt(kind)]; + } +}; + +test "a frozen full-head oracle replaces no candidate half handle" { + const Fake = struct { + var shared: usize = 0; + fn get(out: *cuda.abi.Function, _: cuda.abi.Module, name: [*:0]const u8) callconv(.c) c_int { + for (std.enums.values(Kind)) |k| if (std.mem.eql(u8, std.mem.span(name), symbol(k, false))) { + out.* = @ptrFromInt(0x300 + @as(usize, @backingInt(k))); + return 0; + }; + return 1; + } + fn attribute(_: cuda.abi.Function, _: cuda.abi.FunctionAttribute, _: c_int) callconv(.c) c_int { + shared += 1; + return 0; + } + }; + var driver: cuda.Driver = undefined; + driver.api.cuModuleGetFunction = Fake.get; + driver.api.cuFuncSetAttribute = Fake.attribute; + var candidate: Fns = undefined; + for (&candidate.full, &candidate.half, 0..) |*f, *h, i| { + f.* = .{ .d = &driver, .handle = @ptrFromInt(0x100 + i) }; + h.* = .{ .d = &driver, .handle = @ptrFromInt(0x200 + i) }; + } + Fake.shared = 0; + const oracle = try candidate.withFull(.{ .d = &driver, .handle = @ptrFromInt(0x10) }); + for (candidate.full, candidate.half, oracle.full, oracle.half, 0..) |f, h, of, oh, i| { + try std.testing.expectEqual(@as(usize, 0x100 + i), @intFromPtr(f.handle)); + try std.testing.expectEqual(@as(usize, 0x300 + i), @intFromPtr(of.handle)); + try std.testing.expectEqual(h.handle, oh.handle); + } + try std.testing.expectEqual(@as(usize, 2), Fake.shared); +} + +fn cint(n: usize) c_int { + return @intCast(n); +} + +/// fn_gdn_front: projection rows p after the conv window cs into normed q and k, v, decay and beta. +pub fn front(f: cuda.Function, s: cuda.Stream, h: Heads, rows: usize, p: u64, cs: u64, cw: u64, a_log: u64, dt: u64, qn: u64, kn: u64, v: u64, gb: u64) !void { + var a: cuda.Args = .{}; + for ([_]u64{ p, cs, cw, a_log, dt, qn, kn, v, gb }) |x| a.add(x); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = h.nk + h.nv }, .block = .{ .x = 256 } }, s, &a); +} + +/// The chain over `rows` rows: in place (rows, st_out unused), to st_out (rows_to), or two threads a state row (tpr). +pub fn chain(f: cuda.Function, s: cuda.Stream, kind: Kind, h: Heads, rows: usize, qn: u64, kn: u64, v: u64, gb: u64, st_in: u64, st_out: u64, ys: u64) !void { + var a: cuda.Args = .{}; + for ([_]u64{ qn, kn, v, gb, st_in }) |x| a.add(x); + if (kind != .rows) a.add(st_out); + a.add(cint(rows)); + a.add(ys); + const cfg: cuda.launch.Config = switch (kind) { + .tpr => .{ .grid = .{ .x = h.nv }, .block = .{ .x = 256 }, .shared = parts.chain_tpr_smem }, + .rows, .rows_to => .{ .grid = .{ .x = h.nv, .y = dv / 16 }, .block = .{ .x = 128 } }, + else => return error.NotAChain, + }; + try cuda.launch.launch(f, cfg, s, &a); +} + +/// fn_gdn_norm: the readouts ys gated by z of the projection rows p, into out (bf16) and its sums xs. +pub fn norm(f: cuda.Function, s: cuda.Stream, h: Heads, rows: usize, ys: u64, p: u64, nw: u64, out: u64, xs: u64) !void { + var a: cuda.Args = .{}; + for ([_]u64{ ys, p, nw }) |x| a.add(x); + a.add(eps); + for ([_]u64{ out, xs }) |x| a.add(x); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = h.nv }, .block = .{ .x = dv } }, s, &a); +} + +/// fn_gdn_fused: front, chain and gated norm in one launch a value head; ys 0 keeps no fp32 readouts. +pub fn fused(f: cuda.Function, s: cuda.Stream, h: Heads, rows: usize, p: u64, cs: u64, cw: u64, a_log: u64, dt: u64, st_in: u64, st_out: u64, nw: u64, out: u64, xs: u64, ys: u64) !void { + var a: cuda.Args = .{}; + for ([_]u64{ p, cs, cw, a_log, dt, st_in, st_out }) |x| a.add(x); + a.add(cint(rows)); + a.add(nw); + a.add(eps); + for ([_]u64{ out, xs, ys }) |x| a.add(x); + try cuda.launch.launch(f, .{ .grid = .{ .x = h.nv }, .block = .{ .x = parts.gdn_fused_threads }, .shared = parts.gdn_fused_smem }, s, &a); +} diff --git a/zig/tests/cuda/flashnext_gdn_heads_tests.zig b/zig/tests/cuda/flashnext_gdn_heads_tests.zig new file mode 100644 index 0000000000..d84d4c53ad --- /dev/null +++ b/zig/tests/cuda/flashnext_gdn_heads_tests.zig @@ -0,0 +1,187 @@ +//! tf-cuda-test flashnext-gdn-heads: each DeltaNet twin on a rank's heads gives the full kernel's bytes for them. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const rig_ = @import("flashnext_rig.zig"); +const lay = @import("flashnext_gdn_heads.zig"); + +const Rig = rig_.Rig; +const Heads = lay.Heads; +const Part = lay.Part; + +/// Token rows: one, a few, a stage of 16, the TPR seam's either side, and the default chunk. +const counts = [_]usize{ 1, 4, 16, 63, 64, 1536 }; + +/// A pass's inputs, host bytes in one layout: projection rows, conv window and taps, gates, norm weight, state. +const In = struct { p: []const u8, cs: []const u8, cw: []const u8, a_log: []const u8, dt: []const u8, nw: []const u8, st: []const u8 }; + +const Out = struct { what: []const u8, part: Part, elem: usize = 4 }; + +/// Front's rows, then each form's readouts, state, gated-norm rows and their sums (plain, _to, TPR, fused). +const outs = blk: { + var o: [4 + 4 * 4]Out = undefined; + o[0..4].* = .{ .{ .what = "front q", .part = .keys }, .{ .what = "front k", .part = .keys }, .{ .what = "front v", .part = .values }, .{ .what = "front decay and beta", .part = .gb } }; + for ([_][]const u8{ "plain", "_to", "TPR", "fused" }, 0..) |form, i| { + o[4 + 4 * i] = .{ .what = form ++ " readouts", .part = .values }; + o[5 + 4 * i] = .{ .what = form ++ " state", .part = .state }; + o[6 + 4 * i] = .{ .what = form ++ " gated-norm rows", .part = .values, .elem = 2 }; + o[7 + 4 * i] = .{ .what = form ++ " sums", .part = .sums }; + } + break :blk o; +}; + +fn upload(r: *Rig, ptr: u64, bytes: []const u8) !void { + const b: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = ptr, .len = bytes.len }; + try b.upload(0, bytes); +} + +/// Every form's launches for the heads h over m rows, as the pass makes them; the outputs in `outs`' order. +fn pass(r: *Rig, f: *const lay.Fns, h: Heads, m: usize, in: In) ![outs.len]u64 { + const s = r.s; + var o: [outs.len]u64 = undefined; + for (outs, &o) |x, *p| p.* = try r.out(lay.bytes(x.part, x.elem, m, h)); + const p = try r.dev(in.p); + const cs = try r.dev(in.cs); + const cw = try r.dev(in.cw); + const a_log = try r.dev(in.a_log); + const dt = try r.dev(in.dt); + const nw = try r.dev(in.nw); + const st = try r.dev(in.st); + try lay.front(f.of(.front, h), s, h, m, p, cs, cw, a_log, dt, o[0], o[1], o[2], o[3]); + for ([_]lay.Kind{ .rows, .rows_to, .tpr }, 0..) |kind, i| { + const b = 4 + 4 * i; + if (kind == .rows) try upload(r, o[b + 1], in.st); // in place: the state starts where it lies + const from = if (kind == .rows) o[b + 1] else st; + try lay.chain(f.of(kind, h), s, kind, h, m, o[0], o[1], o[2], o[3], from, o[b + 1], o[b]); + try lay.norm(f.of(.norm, h), s, h, m, o[b], p, nw, o[b + 2], o[b + 3]); + } + try upload(r, o[17], in.st); + try lay.fused(f.of(.fused, h), s, h, m, p, cs, cw, a_log, dt, o[17], o[17], nw, o[18], o[19], o[16]); + try s.synchronize(); + return o; +} + +/// Rank r's inputs gathered from the full layout (the norm weight is whole on both ranks). +fn gathered(a: std.mem.Allocator, in: In, r: usize) !In { + return .{ + .p = try lay.gather(a, in.p, 2, .proj, .half, r), + .cs = try lay.gather(a, in.cs, 2, .conv, .half, r), + .cw = try lay.gather(a, in.cw, 2, .taps, .half, r), + .a_log = try lay.gather(a, in.a_log, 4, .gates, .half, r), + .dt = try lay.gather(a, in.dt, 4, .gates, .half, r), + .nw = in.nw, + .st = try lay.gather(a, in.st, 4, .state, .half, r), + }; +} + +/// Directed non-finites affect separate heads of both ranks; other heads retain the finite random inputs. +fn mixed(a: std.mem.Allocator, in: In, m: usize) !In { + const p = try a.dupe(u8, in.p); + const st = try a.dupe(u8, in.st); + const h = Heads.half; + for (0..2) |rank| { + const q = lay.fullIndex(.proj, h, rank, 3); + const v = (m / 2) * Heads.full.projWidth() + lay.fullIndex(.proj, h, rank, 2 * h.nk * lay.dk + 4 * lay.dv + 7); + const z = (m - 1) * Heads.full.projWidth() + lay.fullIndex(.proj, h, rank, h.convDim() + 7 * lay.dv + 5); + std.mem.writeInt(u16, p[q * 2 ..][0..2], 0x7fc1, .little); + std.mem.writeInt(u16, p[v * 2 ..][0..2], 0x7f80, .little); + std.mem.writeInt(u16, p[z * 2 ..][0..2], 0xff80, .little); + const state = lay.fullIndex(.state, h, rank, 15 * lay.dv * lay.dk + 7); + std.mem.writeInt(u32, st[state * 4 ..][0..4], 0xff800000, .little); + } + var out = in; + out.p = p; + out.st = st; + return out; +} + +/// Each rank must retain finite witnesses as well as observe the planted non-finites. +fn expectMixed(bytes: []const u8, elem: usize) !void { + var finite: usize = 0; + var nonfinite: usize = 0; + var at: usize = 0; + while (at < bytes.len) : (at += elem) { + const bad = if (elem == 2) + std.mem.readInt(u16, bytes[at..][0..2], .little) & 0x7f80 == 0x7f80 + else + std.mem.readInt(u32, bytes[at..][0..4], .little) & 0x7f800000 == 0x7f800000; + if (bad) nonfinite += 1 else finite += 1; + } + try check.expect(finite > 0 and nonfinite > 0, "mixed GDN fixture needs finite and non-finite outputs on each rank", .{}); +} + +fn compare(r: *Rig, f: *const lay.Fns, oracle: *const lay.Fns, m: usize, in: In, frozen: bool, bad: bool) !usize { + const a = r.a; + const want = try pass(r, oracle, .full, m, in); + var checks: usize = 0; + if (frozen) { + const got = try pass(r, f, .full, m, in); + for (outs, want, got) |x, w, mine| { + const len = lay.bytes(x.part, x.elem, m, .full); + const what = try std.fmt.allocPrint(a, "flashnext-gdn-heads frozen {s} {s}, full, {d} rows", .{ if (bad) "mixed" else "finite", x.what, m }); + try check.sameBytes(what, try r.back(u8, mine, len), try r.back(u8, w, len)); + checks += 1; + } + } + for (0..2) |rank| { + const got = try pass(r, f, .half, m, try gathered(a, in, rank)); + for (outs, want, got, 0..) |x, w, mine, i| { + const whole = try r.back(u8, w, lay.bytes(x.part, x.elem, m, .full)); + const theirs = try lay.gather(a, whole, x.elem, x.part, .half, rank); + if (bad and (i == 0 or i == 6)) try expectMixed(theirs, x.elem); + const what = try std.fmt.allocPrint(a, "flashnext-gdn-heads {s}{s}, rank {d}, {d} rows", .{ if (bad) "mixed " else "", x.what, rank, m }); + try check.sameBytes(what, try r.back(u8, mine, theirs.len), theirs); + checks += 1; + } + } + return checks; +} + +/// The existing embedded full-versus-half check; its finite inputs and 240 comparisons are unchanged. +pub fn run(gpu: check.Gpu) !void { + return runWith(gpu, null); +} + +/// Freeze the full-head oracle to a separate image while retaining the embedded candidate handles. +pub fn runBaseline(gpu: check.Gpu, path: []const u8) !void { + const raw = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, gpu.gpa, .limited(64 << 20)); + defer gpu.gpa.free(raw); + const image = try gpu.gpa.alignedAlloc(u8, .@"8", raw.len); + defer gpu.gpa.free(image); + @memcpy(image, raw); + var baseline = try cuda.Module.load(gpu.d, image); + defer baseline.unload(); + try runWith(gpu, baseline); +} + +fn runWith(gpu: check.Gpu, baseline: ?cuda.Module) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var prng = std.Random.DefaultPrng.init(61); + var r = try Rig.init(gpu, a, prng.random()); + defer r.deinit(); + var batch = try cuda.Module.load(gpu.d, cuda.kernels.flashnext_batch); + defer batch.unload(); + const f = try lay.Fns.load(batch); + const oracle = if (baseline) |b| try f.withFull(b) else f; + const full = Heads.full; + var checks: usize = 0; + for (counts) |m| { + const in: In = .{ + .p = std.mem.sliceAsBytes(r.bfs(m * full.projWidth(), 1)), + .cs = std.mem.sliceAsBytes(r.bfs((lay.taps - 1) * full.convDim(), 1)), + .cw = std.mem.sliceAsBytes(r.bfs(full.convDim() * lay.taps, 0.5)), + .a_log = std.mem.sliceAsBytes(r.f32s(full.nv, -1, 1)), + .dt = std.mem.sliceAsBytes(r.f32s(full.nv, -1, 1)), + .nw = std.mem.sliceAsBytes(r.bfs(lay.dv, 1.5)), + .st = std.mem.sliceAsBytes(r.f32s(full.nv * lay.dv * lay.dk, -0.5, 0.5)), + }; + checks += try compare(&r, &f, &oracle, m, in, baseline != null, false); + if (baseline != null) checks += try compare(&r, &f, &oracle, m, try mixed(a, in, m), true, true); + } + if (baseline != null) return check.pass("flashnext-gdn-heads frozen baseline: {d} finite and directed mixed outputs " ++ + "of the candidate full kernels and both half-head layouts equal the external image's full kernels at 1 to 1,536 rows", .{checks}); + check.pass("flashnext-gdn-heads: {d} outputs of the _h8 front, rows, _to, TPR, norm and fused kernels on each rank's heads " ++ + "equal the full kernels' bytes for those heads at 1 to 1,536 rows", .{checks}); +} diff --git a/zig/tests/cuda/flashnext_model_rank_tests.zig b/zig/tests/cuda/flashnext_model_rank_tests.zig index ebb937ae71..2674e03b9f 100644 --- a/zig/tests/cuda/flashnext_model_rank_tests.zig +++ b/zig/tests/cuda/flashnext_model_rank_tests.zig @@ -22,7 +22,7 @@ pub fn run(gpu: check.Gpu, dir: []const u8, rank: usize, port: u16, ids_path: [] defer glue.deinit(); const tri: flash.Tri = .{ .s = stream, .own = &glue }; const identity = try ranks.identity(gpu.gpa, gpu.io, dir, context, chunk); - const session = try ranks.Session.init(gpu.gpa, gpu.io, "127.0.0.1", port, rank, identity); + const session = try ranks.Session.init(gpu.gpa, gpu.io, "127.0.0.1", port, rank, identity, ranks.linkWait()); const engine = prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &stream, dir, 2, rank) catch |err| { session.poison(); session.deinit(); @@ -36,7 +36,7 @@ pub fn run(gpu: check.Gpu, dir: []const u8, rank: usize, port: u16, ids_path: [] engine.collective = &session.collective; if (engine.chunk != chunk) return error.RankIdentity; const room = try session.room((try gpu.ctx.memInfo()).free); - try check.expect(room > prompt.seqBytes(context) * 3, "both ranks have room for experimental caches", .{}); + try check.expect(room > prompt.seqBytes(context, engine.heads) * 3, "both ranks have room for experimental caches", .{}); var logits = try cuda.DeviceBuffer.alloc(gpu.d, 2 * prompt.vocabulary * 2); defer logits.free(); if (rank == 1) { diff --git a/zig/tests/cuda/flashnext_mtp_parity_tests.zig b/zig/tests/cuda/flashnext_mtp_parity_tests.zig index f534d5e03b..27cbafc886 100644 --- a/zig/tests/cuda/flashnext_mtp_parity_tests.zig +++ b/zig/tests/cuda/flashnext_mtp_parity_tests.zig @@ -102,7 +102,7 @@ fn shape(gpu: check.Gpu, e: *prompt.Engine, ops: mtp.Ops(flash.Tri), stream: *cu for (nexts[i][0..w], 0..) |*t, r| t.* = @intCast((37 + 101 * i + 17 * r) % 50_000); inline for (.{ &alone, &together }) |side| { side.seqs[i] = try e.newSeq(cap); - side.drafts[i] = mtp.Draft.init(gpa, e.driver, stream.*, cap, ops.h.ids.len, e.chunk) catch |err| { + side.drafts[i] = mtp.Draft.init(gpa, e.driver, stream.*, cap, ops.h.ids.len, e.chunk, e.heads) catch |err| { side.seqs[i].deinit(gpa); return err; }; diff --git a/zig/tests/cuda/flashnext_nvfp4_tp_tests.zig b/zig/tests/cuda/flashnext_nvfp4_tp_tests.zig index 7f9b81f1b4..08a9ad593b 100644 --- a/zig/tests/cuda/flashnext_nvfp4_tp_tests.zig +++ b/zig/tests/cuda/flashnext_nvfp4_tp_tests.zig @@ -1,4 +1,4 @@ -//! Intermediate-channel NVFP4 partitions: exact source cuts, owned upload and FP32 expert partials against FP64. +//! Expert-parallel NVFP4 ranks: whole experts a rank, fn_ep_picks, FP64 rows, and their sum equal to one GPU's bytes. const std = @import("std"); const cuda = @import("cuda"); const core = @import("core"); @@ -10,7 +10,6 @@ const width = 64; const experts = 2; const slots = 3; const most = 17; -const rank_width = width / 2; const fx = cuda.experts.Nvfp4Experts; const Source = struct { @@ -77,11 +76,14 @@ pub fn run(gpu: check.Gpu) !void { var r = try rig.Rig.init(gpu, a, rng.random()); defer r.deinit(); var source = Source.init(); - var left = try flash.modelopt_experts.loadShape(a, gpu.d, &source, "model.mlp", 2, 0, d, width, experts); - defer left.free(); - var right = try flash.modelopt_experts.loadShape(a, gpu.d, &source, "model.mlp", 2, 1, d, width, experts); - defer right.free(); - try check.expect(source.releases == 2 * experts * 3, "every source projection released after its owned rank cut", .{}); + var whole = try flash.modelopt_experts.loadShape(a, gpu.d, &source, "model.mlp", 1, 0, d, width, experts); + defer whole.free(); + const one_gpu_reads = source.releases; + var ranks: [2]flash.modelopt_experts.Owned = undefined; + for (&ranks, 0..) |*o, rank| o.* = try flash.modelopt_experts.loadShape(a, gpu.d, &source, "model.mlp", 2, rank, d, width, experts); + defer for (&ranks) |*o| o.free(); + try check.expect(source.releases - one_gpu_reads == one_gpu_reads, "the two ranks read every expert once between them", .{}); + for (ranks, 0..) |o, rank| try check.expect(o.first == rank and o.count == 1 and o.width == width, "rank {d} owns one whole expert", .{rank}); const x = r.bfs(most * d, 0.125); const dx = try r.dev(std.mem.sliceAsBytes(x)); const picks = try a.alloc(i32, most * slots); @@ -91,6 +93,7 @@ pub fn run(gpu: check.Gpu) !void { picks[row * slots + 2] = experts; // shared expert is a separate format, skipped here } const dp = try r.dev(std.mem.sliceAsBytes(picks)); + const local = try r.out(most * slots * 4); const plan: cuda.grouped.Plan = .{ .members = try r.out(most * slots * 4), .items = try r.out((most * slots + experts + 1) * 12), @@ -101,65 +104,71 @@ pub fn run(gpu: check.Gpu) !void { var mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); defer mod.unload(); const router = try cuda.grouped.Router.resolve(mod); - var prior: [2]?[]f32 = .{ null, null }; - for ([_]usize{ 4, 1 }) |warps| { - var kernels = try fx.loadWarps(gpu.d, @intCast(try gpu.ctx.attribute(.multiprocessor_count)), warps); - defer kernels.unload(); - const act = try r.out(most * slots * rank_width * 2); - const out = try r.out(most * slots * d * 4); - const combined = try a.alloc(f32, most * slots * d); - @memset(combined, 0); - const reference_full = try a.alloc(f64, combined.len); - @memset(reference_full, 0); - const reference_abs = try a.alloc(f64, combined.len); - @memset(reference_abs, 0); - for ([_]*const flash.modelopt_experts.Owned{ &left, &right }, 0..) |owned, rank| { - var all: []f32 = undefined; - for ([_]usize{ most, 1, 7 }) |rows| { - try gpu.d.check(gpu.d.api.cuMemsetD8Async(out, 0, most * slots * d * 4, r.s.handle), "cuMemsetD8Async"); - const items = cuda.grouped.maxItems(rows * slots, experts + 1, 16); - try router.route(r.s, dp, rows * slots, experts + 1, 16, plan); - try kernels.run(r.s, .up, .{ .x = dx, .stride = d, .slots = slots }, owned.view(), plan, act, items, experts); - try kernels.run(r.s, .down_f32, .{ .x = act, .stride = rank_width }, owned.view(), plan, out, items, experts); - if (rows == most) all = try r.back(f32, out, most * slots * d) else try r.sameRows("NVFP4 rank output rows", f32, out, all[0 .. rows * slots * d]); + var batch = try cuda.Module.load(gpu.d, cuda.kernels.flashnext_batch); + defer batch.unload(); + const ep_picks = try batch.function("fn_ep_picks"); + var kernels = try fx.loadWarps(gpu.d, @intCast(try gpu.ctx.attribute(.multiprocessor_count)), 4); + defer kernels.unload(); + const act = try r.out(most * slots * width * 2); + const out = try r.out(most * slots * d * 4); + const pairs = most * slots; + // one GPU: both experts, global picks, the shared slot skipped + try gpu.d.check(gpu.d.api.cuMemsetD8Async(out, 0, pairs * d * 4, r.s.handle), "cuMemsetD8Async"); + try router.route(r.s, dp, pairs, experts + 1, 16, plan); + try kernels.run(r.s, .up, .{ .x = dx, .stride = d, .slots = slots }, whole.view(), plan, act, cuda.grouped.maxItems(pairs, experts + 1, 16), experts); + try kernels.run(r.s, .down_f32, .{ .x = act, .stride = width }, whole.view(), plan, out, cuda.grouped.maxItems(pairs, experts + 1, 16), experts); + const one = try r.back(f32, out, pairs * d); + const summed = try a.alloc(f32, pairs * d); + @memset(summed, 0); + for (ranks, 0..) |o, rank| { + var args: cuda.Args = .{}; + args.add(dp); + args.add(local); + for ([_]usize{ pairs, o.first, o.count }) |v| args.add(@as(i32, @intCast(v))); + try cuda.launch.launch(ep_picks, .{ .grid = .{ .x = 1 }, .block = .{ .x = 256 } }, r.s, &args); + const got_picks = try r.back(i32, local, pairs); + for (picks, got_picks) |g, l| { + const want: i32 = if (g >= o.first and g < o.first + o.count) g - @as(i32, @intCast(o.first)) else @intCast(o.count); + try check.expect(l == want, "fn_ep_picks maps global pick {d} to {d} on rank {d}, not {d}", .{ g, want, rank, l }); + } + try gpu.d.check(gpu.d.api.cuMemsetD8Async(out, 0, pairs * d * 4, r.s.handle), "cuMemsetD8Async"); + const skip: i32 = @intCast(o.count); + try router.route(r.s, local, pairs, o.count + 1, 16, plan); + try kernels.run(r.s, .up, .{ .x = dx, .stride = d, .slots = slots }, o.view(), plan, act, cuda.grouped.maxItems(pairs, o.count + 1, 16), skip); + try kernels.run(r.s, .down_f32, .{ .x = act, .stride = width }, o.view(), plan, out, cuda.grouped.maxItems(pairs, o.count + 1, 16), skip); + const all = try r.back(f32, out, pairs * d); + for (0..most) |row| for (0..slots) |slot| { + const e: usize = @intCast(picks[row * slots + slot]); + const at = (row * slots + slot) * d; + if (e == experts or e < o.first or e >= o.first + o.count) { + for (all[at..][0..d]) |v| try check.expect(@as(u32, @bitCast(v)) == 0, "a pair rank {d} does not own stays +0", .{rank}); + continue; } - if (prior[rank]) |ref| try check.sameBytes("NVFP4 CTA grouping leaves every rank output byte-exact", std.mem.sliceAsBytes(all), std.mem.sliceAsBytes(ref)) else prior[rank] = all; - for (0..most) |row| for (0..slots) |slot| { - const e: usize = @intCast(picks[row * slots + slot]); - if (e == experts) { - for (all[(row * slots + slot) * d ..][0..d]) |v| try check.expect(v == 0, "shared slot untouched by NVFP4", .{}); - continue; - } - var activation: [rank_width]f64 = undefined; - for (&activation, 0..) |*v, j| { - const channel = rank * rank_width + j; - var gate: f64 = 0; - var up: f64 = 0; - for (0..d) |k| { - gate += rig.bf(x[row * d + k]) * weight(&source, e, 0, channel, k); - up += rig.bf(x[row * d + k]) * weight(&source, e, 1, channel, k); - } - gate = rig.rbf(gate); - up = rig.rbf(up); - v.* = rig.rbf(rig.rbf(gate / (1 + @exp(-gate))) * up); + var activation: [width]f64 = undefined; + for (&activation, 0..) |*v, j| { + var gate: f64 = 0; + var up: f64 = 0; + for (0..d) |k| { + gate += rig.bf(x[row * d + k]) * weight(&source, e, 0, j, k); + up += rig.bf(x[row * d + k]) * weight(&source, e, 1, j, k); } - for (0..d) |j| { - var want: f64 = 0; - var mag: f64 = 0; - for (activation, 0..) |v, k| { - const term = v * weight(&source, e, 2, j, rank * rank_width + k); - want += term; - mag += @abs(term); - } - const at = (row * slots + slot) * d + j; - try r.within("NVFP4 rank FP32 down", all[at], want, rig.sumTol(mag, rank_width) + 1e-6); - combined[at] += all[at]; - reference_full[at] += want; - reference_abs[at] += mag; + gate = rig.rbf(gate); + up = rig.rbf(up); + v.* = rig.rbf(rig.rbf(gate / (1 + @exp(-gate))) * up); + } + for (0..d) |j| { + var want: f64 = 0; + var mag: f64 = 0; + for (activation, 0..) |v, k| { + const term = v * weight(&source, e, 2, j, k); + want += term; + mag += @abs(term); } - }; - } - for (combined, reference_full, reference_abs) |got, want, mag| try r.within("NVFP4 summed rank partials", got, want, rig.sumTol(mag, width) + 2e-6); - r.report("NVFP4 TP cuts: both intermediate halves, exact source reader and packing, independent FP64 reference, FP32 partials and untouched shared slot; rows 1/7/17 equal"); + try r.within("NVFP4 whole-expert FP32 down", all[at + j], want, rig.sumTol(mag, width) + 1e-6); + } + }; + for (summed, all) |*s, v| s.* += v; // the all-reduce: one rank's row plus the other's +0 } + try check.sameBytes("two expert-parallel ranks summed equal one GPU's bytes", std.mem.sliceAsBytes(summed), std.mem.sliceAsBytes(one)); + r.report("NVFP4 expert parallel: whole experts a rank, fn_ep_picks' local ids, FP64 rows, untouched skipped pairs, and the ranks' sum byte-equal to one GPU"); } diff --git a/zig/tests/cuda/flashnext_tp_heads_tests.zig b/zig/tests/cuda/flashnext_tp_heads_tests.zig new file mode 100644 index 0000000000..90250d137f --- /dev/null +++ b/zig/tests/cuda/flashnext_tp_heads_tests.zig @@ -0,0 +1,346 @@ +//! tf-cuda-test flashnext-tp-heads MODEL: each rank's cut in-projections and attention heads, the whole run's bytes. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const rig = @import("flashnext_rig.zig"); + +const cuts = flash.projection_load; +const hc = flash.heads_cut; +const Heads = flash.Heads; +const Kernels = flash.prompt.batch.Kernels; +/// The family's checkpoint reader (cuda_prompt_wt.zig Store). +const Store = @FieldType(flash.prompt.Engine, "store"); +const Weight = cuda.qmmf.Weight; + +const D = 2560; // dims +const most = 1536; // the default chunk +/// Decode rows, a stage of 16, the prompt tile's first row count and the default chunk. +const counts = [_]usize{ 1, 4, 16, 64, most }; +const head_dim = 256; +const index_heads = 4; +const index_dim = 128; +const eps: f32 = 1e-6; +const cap = 4096; // the attention cache's positions +const idw = 2052; // a sparse row's key list (cuda_prompt_wt.zig index_idw) +const dense_keys = 2051; // keys a row reads densely at most (index_top + 1 blocks of 4, less one) +const nch = 5; // a sparse row's 512-key chunks +const att_rows = flash.batch_state.att_rows; +const fused_rows = flash.batch_state.fused_rows; + +const Ctx = struct { + r: *rig.Rig, + k: *const Kernels, + g: *const flash.Glue, + store: *Store, + part: cuda.DeviceBuffer, + checks: usize = 0, + + /// A packed projection's weight, its buffer freed with the rig's. + fn keep(c: *Ctx, p: flash.projection.Projection) !Weight { + try c.r.bufs.append(c.r.gpu.gpa, p.w); + return p.native.?; + } + + /// One lane call as the pass makes it (bf16 rows), or its fp32 rows, into a buffer of 0xff bytes. + fn lane(c: *Ctx, comptime T: type, x: u64, m: usize, w: Weight) ![]T { + const dst = try c.r.out(m * w.n * @sizeOf(T)); + try c.k.native_lane.matmulOut(c.r.s, x, w.k, m, w, dst, if (T == f32) .f32 else .bf16, c.part.ptr); + return c.r.back(T, dst, m * w.n); + } + + fn same(c: *Ctx, comptime fmt: []const u8, args: anytype, got: []const u8, want: []const u8) !void { + try check.sameBytes(try std.fmt.allocPrint(c.r.a, fmt, args), got, want); + c.checks += 1; + } +}; + +/// Rank `mine`'s kept columns of whole rows [m, whole.n]: each part's ranges in order, as the cut's rows hold them. +fn kept(comptime T: type, a: std.mem.Allocator, rows: []const T, m: usize, whole: hc.Stitch, mine: hc.Stitch) ![]T { + const out = try a.alloc(T, m * mine.n); + for (0..m) |row| { + var at: usize = 0; + var base: usize = 0; + for (0..4) |p| { + for (mine.ranges[p].items()) |c| { + @memcpy(out[row * mine.n + at ..][0..c.width()], rows[row * whole.n + base + c.from ..][0..c.width()]); + at += c.width(); + } + base += whole.ranges[p].rows(); + } + } + return out; +} + +/// Activations: BF16 of a unit normal. +fn inputs(r: *rig.Rig, len: usize) ![]u16 { + const v = try r.a.alloc(u16, len); + for (v) |*h| h.* = rig.bf16(r.rng.floatNorm(f32)); + return v; +} + +/// One class's in-projection: each rank's pinned cut against the whole one's kept rows; returns the whole's bf16 rows. +fn projection(c: *Ctx, comptime what: []const u8, prefix: []const u8, parts: []const []const u8, whole: hc.Stitch, mine: [2]hc.Stitch, x: u64) ![]const u16 { + const r = c.r; + const all = try c.keep(try cuts.loadStitched(r.gpu.gpa, @constCast(r.gpu.d), c.store, prefix, parts, whole.n, D)); + var ranks: [2]Weight = undefined; + var loose: [2]Weight = undefined; + for (mine, &ranks, &loose) |s, *w, *l| { + w.* = try c.keep(try cuts.loadStitchedCut(r.gpu.gpa, r.gpu.d, c.store, prefix, parts, &s.ranges, s.n, D, s.sk)); + try check.expect(w.sk == whole.sk and w.n == s.n, what ++ ": the cut keeps the whole shape's {d} K slices", .{whole.sk}); + // the witness: the same rows packed without the pin, so its K slices follow its own shape + const bytes = try cuts.stitchCut(r.a, c.store, prefix, parts, &s.ranges, s.n, D); + l.* = try c.keep(try flash.projection.dense(r.gpu.gpa, r.gpu.d, .{ .dtype = .bf16, .rank = 2, .shape = .{ s.n, D, 1, 1, 1 }, .bytes = bytes })); + try check.expect(cuda.qmmf.planFor(most, l.*, false).sk != cuda.qmmf.planFor(most, w.*, false).sk, what ++ ": the unpinned cut slices K its own way", .{}); + } + var rows: []const u16 = &.{}; + for (counts) |m| { + const want16 = try c.lane(u16, x, m, all); + const want32 = try c.lane(f32, x, m, all); + if (m == most) rows = want16; + for (mine, ranks, loose, 0..) |s, w, l, rank| { + const k16 = try kept(u16, r.a, want16, m, whole, s); + const k32 = try kept(f32, r.a, want32, m, whole, s); + try c.same(what ++ " rank {d}, {d} rows: bf16 rows", .{ rank, m }, std.mem.sliceAsBytes(try c.lane(u16, x, m, w)), std.mem.sliceAsBytes(k16)); + try c.same(what ++ " rank {d}, {d} rows: fp32 rows", .{ rank, m }, std.mem.sliceAsBytes(try c.lane(f32, x, m, w)), std.mem.sliceAsBytes(k32)); + if (m != most) continue; + const off = try c.lane(f32, x, m, l); + try check.expect(!std.mem.eql(f32, off, k32), what ++ " rank {d}: the unpinned witness equals the whole rows, so the test cannot fail", .{rank}); + } + } + check.pass("flashnext-tp-heads {s}: each rank's {d} rows equal the whole {d}-row projection's kept rows byte for byte " ++ + "(bf16 and fp32) at 1 to 1,536 rows; the unpinned cut differs", .{ what, mine[0].n, whole.n }); + return rows; +} + +/// RoPE's inverse frequencies as the pass uploads them (base 10,000,000 over 32 pairs). +fn inverse(r: *rig.Rig) ![]f32 { + const inv = try r.a.alloc(f32, 32); + for (inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / 32)); + return inv; +} + +/// Head `kv` of every position of a whole cache [cap, 2, 256]: a rank's own cache [cap, 1, 256]. +fn headOf(a: std.mem.Allocator, whole: []const u16, kv: usize) ![]u16 { + const out = try a.alloc(u16, cap * head_dim); + for (0..cap) |p| @memcpy(out[p * head_dim ..][0..head_dim], whole[(p * 2 + kv) * head_dim ..][0..head_dim]); + return out; +} + +/// Heads [from, from + n) of each of `m` rows of `per`-head rows `width` wide. +fn heads(a: std.mem.Allocator, comptime T: type, rows: []const T, m: usize, per: usize, from: usize, n: usize, width: usize) ![]T { + const out = try a.alloc(T, m * n * width); + for (0..m) |row| @memcpy(out[row * n * width ..][0 .. n * width], rows[(row * per + from) * width ..][0 .. n * width]); + return out; +} + +/// The layer's norm weights and RoPE table on the device, one cache's initial bytes, and the rows' key lists. +const Fixed = struct { qw: u64, kw: u64, iw: u64, inv: u64, kc: []const u16, vc: []const u16, ikc: []const u16, ids: u64, nk: u64, spr: u64 }; + +/// What one run leaves: queries, caches, indexer queries, merged and fused rows, gated rows and their sums. +const Run = struct { q: []u16, kc: []u16, vc: []u16, iq: []u16, ikc: []u16, merged: []u16, fused: []u16, gated: []u16, sums: []f32 }; + +/// The pass's attention for `m` rows at the end of the cache with heads `h`: prep, chunks then merge, fused, gate. +fn attend(c: *Ctx, f: Fixed, h: Heads, pa: []const u16, m: usize, kc0: []const u16, vc0: []const u16) !Run { + const r = c.r; + const g = c.g; + const s = r.s; + const pos0 = cap - m; + const qr = h.q * head_dim * 2; + // a 16-byte position record a block of att_rows rows, as the pass stages them + const blocks = (m + att_rows - 1) / att_rows; + const recs = try r.a.alloc(i32, 4 * blocks); + @memset(recs, 0); + for (0..blocks) |b| recs[4 * b] = @intCast(pos0 + b * att_rows); + const pos = try r.dev(std.mem.sliceAsBytes(recs)); + const dpa = try r.dev(std.mem.sliceAsBytes(pa)); + const q = try r.out(m * qr); + const iq = try r.out(m * index_heads * index_dim * 2); + const kc = try r.dev(std.mem.sliceAsBytes(kc0)); + const vc = try r.dev(std.mem.sliceAsBytes(vc0)); + const ikc = try r.dev(std.mem.sliceAsBytes(f.ikc)); + try g.attnPrep(s, dpa, pos, f.qw, f.kw, f.iw, f.inv, q, kc, vc, iq, ikc, eps, m, h, .{}); + const po = try r.out(att_rows * nch * h.q * head_dim * 4); + const pm = try r.out(att_rows * nch * h.q * 4); + const pl = try r.out(att_rows * nch * h.q * 4); + const merged = try r.out(m * qr); + var r0: usize = 0; + while (r0 < m) : (r0 += att_rows) { + const n = @min(att_rows, m - r0); + const at = pos + (r0 / att_rows) * 16; + const chunks = (@min(pos0 + r0 + n, dense_keys) + 511) / 512; + try g.attnChunks(s, q + r0 * qr, kc, vc, at, po, pm, pl, f.ids + r0 * idw * 4, f.nk + r0 * 4, f.spr + r0 * 4, n, h, chunks, true, .{}); + try g.attnMerge(s, po, pm, pl, at, merged + r0 * qr, f.nk + r0 * 4, f.spr + r0 * 4, n, h, true); + } + const fused = try r.out(m * qr); + var g0: usize = 0; + while (g0 < m) : (g0 += fused_rows) { + const n = @min(fused_rows, m - g0); + try g.attnFused(s, q + g0 * qr, kc, vc, pos + (g0 / att_rows) * 16, fused + g0 * qr, f.ids + g0 * idw * 4, f.nk + g0 * 4, f.spr + g0 * 4, n, h, .{}); + } + const gated = try r.out(m * qr); + const sums = try r.out(m * h.q * head_dim / 32 * 4); + try g.attnGate(s, merged, dpa, gated, sums, m, h.projN(), h.q, head_dim); + return .{ + .q = try r.back(u16, q, m * h.q * head_dim), + .kc = try r.back(u16, kc, cap * h.kv * head_dim), + .vc = try r.back(u16, vc, cap * h.kv * head_dim), + .iq = try r.back(u16, iq, m * index_heads * index_dim), + .ikc = try r.back(u16, ikc, cap * index_dim), + .merged = try r.back(u16, merged, m * h.q * head_dim), + .fused = try r.back(u16, fused, m * h.q * head_dim), + .gated = try r.back(u16, gated, m * h.q * head_dim), + .sums = try r.back(f32, sums, m * h.q * head_dim / 32), + }; +} + +/// fn_attn_prep as Glue.attnPrep launches it, but with the row width `pw` given (the witness's whole layout's). +fn prepAt(c: *Ctx, f: Fixed, h: Heads, pw: usize, pa: u64, pos: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, m: usize) !void { + var a: cuda.Args = .{}; + for ([_]u64{ pa, pos, f.qw, f.kw, f.iw, f.inv, q, kc, vc, iq, ikc }) |v| a.add(v); + a.add(eps); + for ([_]usize{ pw, h.q, h.kv, index_heads, index_dim, 32 }) |v| a.add(@as(c_int, @intCast(v))); + try cuda.launch.launch(c.g.f.prep, .{ .grid = .{ .x = @intCast(m), .y = h.q + h.kv + index_heads + 1 }, .block = .{ .x = head_dim } }, c.r.s, &a); +} + +/// Attention at 12 queries and one KV head a rank against the same heads of the whole layer's run. +fn attention(c: *Ctx, rows: []const u16, whole: hc.Stitch, mine: [2]hc.Stitch) !void { + const r = c.r; + var buf: [120]u8 = undefined; + const root = "language_model.model.layers.3.self_attn"; + var norms: [3]u64 = undefined; + for ([_][]const u8{ ".q_norm.weight", ".k_norm.weight", ".indexer.q_layernorm.weight" }, &norms) |name, *d| { + const w = try c.store.normScale(r.a, try c.store.tensor(try std.fmt.bufPrint(&buf, "{s}{s}", .{ root, name }))); + d.* = try r.dev(std.mem.sliceAsBytes(w)); + } + // every row's key list as the selection leaves it past 2,048 keys: one list both ranks read (the indexer is whole) + const ids = try r.a.alloc(i32, most * idw); + const nk = try r.a.alloc(i32, most); + const spr = try r.a.alloc(i32, most); + const kc = r.bfs(cap * 2 * head_dim, 1); + const vc = r.bfs(cap * 2 * head_dim, 1); + const f0: Fixed = .{ .qw = norms[0], .kw = norms[1], .iw = norms[2], .inv = try r.dev(std.mem.sliceAsBytes(try inverse(r))), .kc = kc, .vc = vc, .ikc = r.bfs(cap * index_dim, 1), .ids = 0, .nk = 0, .spr = 0 }; + const full = Heads.full; + for ([_]usize{ 1, most }) |m| { + const pos0 = cap - m; + for (0..m) |row| { + nk[row] = @intCast(idw - 1 - row % 3); + spr[row] = 1; + for (0..idw) |i| ids[row * idw + i] = @intCast((i * 1237 + row * 31) % (pos0 + row + 1)); + } + var f = f0; + f.ids = try r.dev(std.mem.sliceAsBytes(ids[0 .. m * idw])); + f.nk = try r.dev(std.mem.sliceAsBytes(nk[0..m])); + f.spr = try r.dev(std.mem.sliceAsBytes(spr[0..m])); + const pa = rows[0 .. m * whole.n]; + const want = try attend(c, f, full, pa, m, kc, vc); + for (mine, 0..) |s, rank| { + const h = Heads.half; + const got = try attend(c, f, h, try kept(u16, r.a, pa, m, whole, s), m, try headOf(r.a, kc, rank), try headOf(r.a, vc, rank)); + const fmt = "flashnext-tp-heads attention rank {d}, {d} rows: {s}"; + const theirs = [_][]const u16{ try heads(r.a, u16, want.q, m, full.q, 12 * rank, 12, head_dim), try heads(r.a, u16, want.merged, m, full.q, 12 * rank, 12, head_dim), try heads(r.a, u16, want.fused, m, full.q, 12 * rank, 12, head_dim), try heads(r.a, u16, want.gated, m, full.q, 12 * rank, 12, head_dim) }; + for ([_][]const u16{ got.q, got.merged, got.fused, got.gated }, theirs, [_][]const u8{ "queries", "merged rows", "fused rows", "gated rows" }) |mine_rows, w, what| + try c.same(fmt, .{ rank, m, what }, std.mem.sliceAsBytes(mine_rows), std.mem.sliceAsBytes(w)); + try c.same(fmt, .{ rank, m, "gate sums" }, std.mem.sliceAsBytes(got.sums), std.mem.sliceAsBytes(try heads(r.a, f32, want.sums, m, full.q, 12 * rank, 12, head_dim / 32))); + try c.same(fmt, .{ rank, m, "cache keys" }, std.mem.sliceAsBytes(got.kc), std.mem.sliceAsBytes(try headOf(r.a, want.kc, rank))); + try c.same(fmt, .{ rank, m, "cache values" }, std.mem.sliceAsBytes(got.vc), std.mem.sliceAsBytes(try headOf(r.a, want.vc, rank))); + try c.same(fmt, .{ rank, m, "indexer queries" }, std.mem.sliceAsBytes(got.iq), std.mem.sliceAsBytes(want.iq)); + try c.same(fmt, .{ rank, m, "indexer keys" }, std.mem.sliceAsBytes(got.ikc), std.mem.sliceAsBytes(want.ikc)); + } + if (m == 1) continue; + try witness(c, f, try kept(u16, r.a, pa[0 .. 16 * whole.n], 16, whole, mine[1]), whole.n, kc, vc); + } + check.pass("flashnext-tp-heads attention: queries, cache rows, merged, fused and gated rows of 12 queries and one KV head " ++ + "a rank equal the whole run's same heads byte for byte at 1 and 1,536 rows over a 4,096-token cache; the whole PW differs", .{}); +} + +/// Prep and gate at 16 rows with the whole layout's PW over a rank's rows in a buffer sized for it: they must differ. +fn witness(c: *Ctx, f: Fixed, local: []const u16, whole_pw: usize, kc: []const u16, vc: []const u16) !void { + const r = c.r; + const h = Heads.half; + const m = 16; + const right = try attend(c, f, h, local, m, try headOf(r.a, kc, 1), try headOf(r.a, vc, 1)); + const padded = try r.a.alloc(u16, m * whole_pw); + @memset(padded, 0); + @memcpy(padded[0..local.len], local); + const pa = try r.dev(std.mem.sliceAsBytes(padded)); + var recs = [_]i32{ cap - m, 0, 0, 0 }; + const pos = try r.dev(std.mem.sliceAsBytes(&recs)); + const q = try r.out(m * h.q * head_dim * 2); + const iq = try r.out(m * index_heads * index_dim * 2); + const kd = try r.dev(std.mem.sliceAsBytes(try headOf(r.a, kc, 1))); + const vd = try r.dev(std.mem.sliceAsBytes(try headOf(r.a, vc, 1))); + try prepAt(c, f, h, whole_pw, pa, pos, q, kd, vd, iq, try r.dev(std.mem.sliceAsBytes(f.ikc)), m); + const gated = try r.out(m * h.q * head_dim * 2); + const sums = try r.out(m * h.q * head_dim / 32 * 4); + try c.g.attnGate(r.s, try r.dev(std.mem.sliceAsBytes(right.merged)), pa, gated, sums, m, whole_pw, h.q, head_dim); + try r.s.synchronize(); + const wq = try r.back(u16, q, m * h.q * head_dim); + const wg = try r.back(u16, gated, m * h.q * head_dim); + try check.expect(!std.mem.eql(u16, wq, right.q), "flashnext-tp-heads witness: prep at the whole PW gives the rank's queries", .{}); + try check.expect(!std.mem.eql(u16, wg, right.gated), "flashnext-tp-heads witness: the gate at the whole PW gives the rank's rows", .{}); + // row 0 starts where the rank's own row 0 does, so only the rows after it move + try check.expect(std.mem.eql(u16, wq[0 .. h.q * head_dim], right.q[0 .. h.q * head_dim]), "flashnext-tp-heads witness: row 0's queries", .{}); +} + +pub fn run(gpu: check.Gpu, dir: []const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var rng = std.Random.DefaultPrng.init(0x7f17_4ead); + var r = try rig.Rig.init(gpu, a, rng.random()); + defer r.deinit(); + var k = try Kernels.load(gpu.d); + defer k.deinit(); + var g = try flash.Glue.load(gpu.d); + defer g.deinit(); + // the unpinned witness's two K slices of up to 8,240 fp32 columns a row + var part = try cuda.DeviceBuffer.alloc(gpu.d, 2 * most * 8240 * 4); + defer part.free(); + var store = try Store.open(gpu.gpa, gpu.io, dir); + defer store.deinit(); + try store.detectNorms(); + var c: Ctx = .{ .r = &r, .k = &k, .g = &g, .store = &store, .part = part }; + const x = try r.dev(std.mem.sliceAsBytes(try inputs(&r, most * D))); + const half = Heads.half; + _ = try projection(&c, "layer 0 DeltaNet in-projection", "language_model.model.layers.0.linear_attn", &hc.delta_parts, try hc.delta(.full, 0), .{ try hc.delta(half, 0), try hc.delta(half, 1) }, x); + const whole = try hc.attention(.full, 0); + const mine = [2]hc.Stitch{ try hc.attention(half, 0), try hc.attention(half, 1) }; + const rows = try projection(&c, "layer 3 attention in-projection", "language_model.model.layers.3.self_attn", &hc.attention_parts, whole, mine, x); + try attention(&c, rows, whole, mine); + try outputs(&c, "language_model.model.layers.0.linear_attn.out_proj", D, 6144); + try outputs(&c, "language_model.model.layers.3.self_attn.o_proj", D, 6144); + try outputs(&c, "language_model.model.layers.0.mlp.shared_expert.gate_proj", 640, D); + try outputs(&c, "language_model.model.layers.0.mlp.shared_expert.up_proj", 640, D); + try outputs(&c, "language_model.model.layers.0.mlp.shared_expert.down_proj", D, 640); + check.pass("flashnext-tp-heads: {d} byte-for-byte checks of each rank's heads against the whole layer's", .{c.checks}); +} + +/// Output rows keep the whole projection's input and K slices, including the single-row and prompt paths. +fn outputs(c: *Ctx, prefix: []const u8, outputs_n: usize, width: usize) !void { + const r = c.r; + const all = try c.keep(try cuts.load(r.gpu.gpa, @constCast(r.gpu.d), c.store, prefix, outputs_n, width)); + var ranks: [2]Weight = undefined; + for (&ranks, 0..) |*w, rank| { + const cut = try flash.modelopt_experts.Cut.init(outputs_n, 2, rank); + w.* = try c.keep(try cuts.loadCut(r.gpu.gpa, r.gpu.d, c.store, prefix, outputs_n, width, cut, .rows)); + try check.expect(w.sk == cuda.qmmf.splitK(outputs_n, width), "output row cut pins the full K split", .{}); + } + const values = try inputs(r, most * width); + for ([_]bool{ false, true }) |nonfinite| { + if (nonfinite) values[0..3].* = .{ 0x7fc0, 0x7f80, 0xff80 }; + const x = try r.dev(std.mem.sliceAsBytes(values)); + for (counts) |m| { + const want = try c.lane(f32, x, m, all); + const got = try r.a.alloc(f32, m * outputs_n); + for (0..5) |_| { + for (ranks, 0..) |w, rank| { + const mine = try c.lane(f32, x, m, w); + for (0..m) |row| @memcpy(got[row * outputs_n + rank * (outputs_n / 2) ..][0 .. outputs_n / 2], mine[row * (outputs_n / 2) ..][0 .. outputs_n / 2]); + } + try c.same("{s} output columns, {d} rows, nonfinite={any}", .{ prefix, m, nonfinite }, std.mem.sliceAsBytes(got), std.mem.sliceAsBytes(want)); + } + } + } + check.pass("{s}: each rank's full-input output columns equal one GPU's FP32 bytes, five repeats", .{prefix}); +} diff --git a/zig/tests/cuda/flashnext_window_tests.zig b/zig/tests/cuda/flashnext_window_tests.zig index 999d773cd0..45d1ba164b 100644 --- a/zig/tests/cuda/flashnext_window_tests.zig +++ b/zig/tests/cuda/flashnext_window_tests.zig @@ -22,7 +22,7 @@ pub fn run(gpu: check.Gpu, dir: []const u8, rank: usize, port: u16, ids_path: [] const tri: flash.Tri = .{ .s = stream, .own = &glue }; var identity = try flash.rank_engine.identity(gpu.gpa, gpu.io, dir, context, 32); identity.mtp = 1; - const session = try flash.rank_engine.Session.init(gpu.gpa, gpu.io, "127.0.0.1", port, rank, identity); + const session = try flash.rank_engine.Session.init(gpu.gpa, gpu.io, "127.0.0.1", port, rank, identity, flash.rank_engine.linkWait()); const e = prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &stream, dir, 2, rank) catch |err| { session.poison(); session.deinit(); @@ -57,7 +57,7 @@ pub fn run(gpu: check.Gpu, dir: []const u8, rank: usize, port: u16, ids_path: [] } var window = try session.create(context); defer session.drop(&window); - var draft = try mtp.Draft.init(gpu.gpa, e.driver, stream, context, 1, 32); + var draft = try mtp.Draft.init(gpu.gpa, e.driver, stream, context, 1, 32, e.heads); defer draft.deinit(); const rows = try gpu.gpa.alloc(u8, expected.len); defer gpu.gpa.free(rows); diff --git a/zig/tests/cuda/flashnext_windows_tests.zig b/zig/tests/cuda/flashnext_windows_tests.zig index 2d264197ee..ea15d52ef4 100644 --- a/zig/tests/cuda/flashnext_windows_tests.zig +++ b/zig/tests/cuda/flashnext_windows_tests.zig @@ -66,7 +66,7 @@ fn shape(gpu: check.Gpu, e: *prompt.Engine, ops: mtp.Ops(flash.Tri), stream: *cu for (windows[i][0..w], 0..) |*t, r| t.* = @intCast((37 + 101 * i + 17 * r) % 50_000); inline for (.{ &alone, &together }) |side| { side.seqs[i] = try e.newSeq(cap); - side.drafts[i] = mtp.Draft.init(gpa, e.driver, stream.*, cap, 1, e.chunk) catch |err| { + side.drafts[i] = mtp.Draft.init(gpa, e.driver, stream.*, cap, 1, e.chunk, e.heads) catch |err| { side.seqs[i].deinit(gpa); return err; }; diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index eff29cfbdf..634f7fda77 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -93,6 +93,7 @@ const flashnext_rank_tests = @import("flashnext_rank_tests.zig"); const flashnext_draft_affine_tests = @import("flashnext_draft_affine_tests.zig"); const flashnext_shared_tests = @import("flashnext_shared_tests.zig"); const flashnext_nvfp4_tp_tests = @import("flashnext_nvfp4_tp_tests.zig"); +const flashnext_fp8_experts_tests = @import("flashnext_fp8_experts_tests.zig"); const nvfp4_fact_tests = @import("nvfp4_fact_tests.zig"); const lane_half_tests = @import("lane_half_tests.zig"); const flashnext_dense_hc_tests = @import("flashnext_dense_hc_tests.zig"); @@ -134,11 +135,17 @@ const usage = \\ flashnext-draft-affine proposal-head affine projection against FP64 and row-count invariance \\ flashnext-shared separate MXFP8 shared expert and reserved-slot scatter against FP64 \\ flashnext-nvfp4-tp rank-local expert cuts and FP32 partials against FP64 (one GPU) + \\ flashnext-fp8-experts the MTP head's block-FP8 experts (512 x 640 x 2560) through moe_native.fp8 against FP64 + \\ moetile whole shared down writes only slot 10 with identical FP32 bytes + \\ moetile-contract shared down replacement gates and the planted control + \\ moetile-bench alternate shared-down plus scatter and direct-slot GPU timings \\ nvfp4-fact NVFP4 experts' prompt tiles with the decode form's sums: its bytes at every row, FP64 \\ lane-half lane half tiles (MXFP8, BF16): bytes equal to the 64-column tile, FP64, cold-weight time \\ flashnext-dense-hc dense connection weights and FP32 normalization against FP64 \\ flashnext-hc Flash Next's own hyper-connection and n-gram glue against float64 (synthetic rows) \\ flashnext-attn Flash Next's own attention glue, pooling and block scores against float64 (synthetic) + \\ flashnext-gdn-heads [BASELINE] DeltaNet full/half bytes, optionally against a frozen flashnext_batch.fatbin + \\ flashnext-tp-heads MODEL each rank's cut in-projections and attention heads equal the whole layer's bytes \\ flashnext-pages the attention glue's caches through a page table: paged == contiguous, float64, controls \\ info driver, device and library versions \\ smoke copies, fills, launches, argument packing, module globals @@ -232,12 +239,17 @@ const usage = \\ kernel-contract the registered kernels through the contract's gates, each planted control red (needs -Dnvcc) \\ qmm-prefill-many-contract the GB10 bit-exact multi-projection prompt launch against separate launches \\ qmm-prefill-tile22-contract the GB10 large-prompt 2x2 tile against the default tile + \\ qmmf-prompt BF16 128x128 prompt tiles: byte equality at Flash Next TP=2 shapes + \\ qmmf-prompt-contract gates 1-4 for all four BF16 prompt instances and a scaled control + \\ qmmf-prompt-repro isolate M=65 N=8240 K=2560 SK=4 BF16 without clusters + \\ qmmf-prompt-bench rank0 normal-main weighted A-B-A timings, plus separate MTP hidden timings \\ qmm-prefill-many-bench same-process A-B-A screening for the GB10 launch fusion \\ mtp-rows tf_q36m_attn_rows: gates 1-4 (byte-equal to tf_q36m_attn a segment) and float64 alone \\ qwen-pages qwen35.cu's attention cache through a page table: paged == contiguous, float64, red controls \\ rope-yarn YaRN rope in tf_q35_rope, tf_q35_rope8 and tf_q36m_norm_rope against float64, to 1M \\ kernel-contract-limits DIR each registered launch against sm_80/86/89/120/121 from the fatbins in DIR \\ pocket-contract Nemotron's pocket kernels (learned neurons) through gates 1-4, controls red + \\ flashnext-modelopt-contract Flash Next NVFP4's BF16 dense and FP8 MTP-expert downs through gates 1-4 \\ minimax-contract MiniMax-M3.1-Flash's indexer and block-sparse attention through gates 1-4, controls red \\ moss-contract MOSS-Image's own kernels and its full attention through gates 1-4, controls red \\ glm-exl3 GLM's EXL3 routed experts, decode and prompt, against float64 (synthetic) @@ -339,11 +351,25 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "flashnext-draft-affine")) return flashnext_draft_affine_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-shared")) return flashnext_shared_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-nvfp4-tp")) return flashnext_nvfp4_tp_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-fp8-experts")) return flashnext_fp8_experts_tests.run(gpu); + if (std.mem.eql(u8, cmd, "moetile")) return @import("moetile_tests.zig").run(gpu); + if (std.mem.eql(u8, cmd, "moetile-contract")) return @import("contract_qmmf_shared.zig").run(gpu); + if (std.mem.eql(u8, cmd, "moetile-bench")) return @import("moetile_tests.zig").bench(gpu); if (std.mem.eql(u8, cmd, "nvfp4-fact")) return nvfp4_fact_tests.run(gpu); + if (std.mem.eql(u8, cmd, "qmmf-prompt")) return @import("qmmf_prompt_tests.zig").run(gpu); + if (std.mem.eql(u8, cmd, "qmmf-prompt-contract")) return @import("contract_qmmf_prompt.zig").run(gpu); + if (std.mem.eql(u8, cmd, "qmmf-prompt-repro")) return @import("qmmf_prompt_tests.zig").repro(gpu); + if (std.mem.eql(u8, cmd, "qmmf-prompt-bench")) return @import("qmmf_prompt_tests.zig").bench(gpu); if (std.mem.eql(u8, cmd, "lane-half")) return lane_half_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-dense-hc")) return flashnext_dense_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-hc")) return flashnext_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-attn")) return flashnext_attn_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-gdn-heads")) { + if (rest.len > 1) return error.GdnHeadsArguments; + const heads = @import("flashnext_gdn_heads_tests.zig"); + return if (rest.len == 1) heads.runBaseline(gpu, rest[0]) else heads.run(gpu); + } + if (std.mem.eql(u8, cmd, "flashnext-tp-heads")) return @import("flashnext_tp_heads_tests.zig").run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "flashnext-pages")) return @import("contract_fn_pages.zig").run(gpu); if (std.mem.eql(u8, cmd, "info")) return info(gpu); if (std.mem.eql(u8, cmd, "smoke")) return runtime_tests.smoke(gpu); @@ -466,6 +492,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "rope-yarn")) return @import("contract_qwen_rope_yarn.zig").run(gpu); if (std.mem.eql(u8, cmd, "kernel-contract-limits")) return contract_tests.fatbins(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "pocket-contract")) return @import("contract_pocket.zig").run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-modelopt-contract")) return @import("contract_flashnext_modelopt.zig").run(gpu); if (std.mem.eql(u8, cmd, "moss-contract")) return @import("contract_moss_image.zig").run(gpu); if (std.mem.eql(u8, cmd, "tree")) return tree_tests.run(gpu); if (std.mem.eql(u8, cmd, "block")) return block_tests.run(gpu); diff --git a/zig/tests/cuda/moetile_tests.zig b/zig/tests/cuda/moetile_tests.zig new file mode 100644 index 0000000000..b4294d4a84 --- /dev/null +++ b/zig/tests/cuda/moetile_tests.zig @@ -0,0 +1,152 @@ +//! Shared-down FP32 stores into slot 10 equal the contiguous product followed by scatter, with GPU event timing. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const Rig = @import("contract_rig.zig").Rig; +const bf16 = @import("flashnext_rig.zig").bf16; +const d = 2560; +const width = 640; +const slots = 11; +const most = 4218; +const row_counts = [_]usize{ 1, 16, 17, 32, 64, 65, 1536, most }; + +const Fixture = struct { + r: Rig, + lane: cuda.qmmf.Lane, + glue: flash.Glue, + shared: cuda.qmmf.Weight, + shared_x: u64, + shared_out: u64, + shared_slots: [2]u64, + + fn init(gpu: check.Gpu) !Fixture { + const a = gpu.gpa; + const s = try cuda.Stream.init(gpu.d, false); + errdefer { + var stream = s; + stream.deinit(); + } + var r: Rig = .{ .gpu = gpu, .s = s }; + errdefer r.deinit(); + var lane = try cuda.qmmf.Lane.load(gpu.d, (try gpu.ctx.capability()) / 10); + errdefer lane.unload(); + var glue = try flash.Glue.load(gpu.d); + errdefer glue.deinit(); + var prng = std.Random.DefaultPrng.init(0x5a4ed); + const rng = prng.random(); + const x = try a.alloc(u16, most * width); + defer a.free(x); + for (x) |*v| v.* = bf16(rng.float(f32) - 0.5); + const weights = try a.alloc(u16, d * width); + defer a.free(weights); + const tiled = try a.alloc(u16, weights.len); + defer a.free(tiled); + for (weights) |*v| v.* = bf16((rng.float(f32) - 0.5) * 0.1); + cuda.qmmf.packBf16(tiled, weights, d, width); + var result: Fixture = .{ + .r = undefined, + .lane = lane, + .glue = glue, + .shared = .{ .mode = .bf16, .codes = try r.input(std.mem.sliceAsBytes(tiled)), .scales = 0, .scale = 0.8123, .n = d, .k = width, .npad = d }, + .shared_x = try r.input(std.mem.sliceAsBytes(x)), + .shared_out = try r.scratch(most * d * 4), + .shared_slots = .{ try r.output(most * slots * d * 4), try r.output(most * slots * d * 4) }, + }; + result.r = r; + return result; + } + + fn deinit(f: *Fixture) void { + f.r.deinit(); + f.r.s.deinit(); + f.glue.deinit(); + f.lane.unload(); + } + + fn sharedDown(f: *Fixture, rows: usize, wide: bool) !void { + if (wide) { + try f.lane.matmulF32Strided(f.r.s, f.shared_x, width, rows, f.shared, f.shared_slots[1] + 10 * d * 4, slots * d); + } else { + try f.lane.matmulF32(f.r.s, f.shared_x, width, rows, f.shared, f.shared_out, 0); + try f.glue.sharedScatter(f.r.s, f.shared_out, f.shared_slots[0], rows, d, slots); + } + } + + fn poison(f: *Fixture, ptr: u64, bytes: usize) !void { + const b: cuda.DeviceBuffer = .{ .d = f.r.gpu.d, .ptr = ptr, .len = bytes }; + try b.fill8(0xff, f.r.s.handle); + } + + fn compare(f: *Fixture, rows: usize) !void { + const a = f.r.gpu.gpa; + try f.poison(f.shared_slots[0], most * slots * d * 4); + try f.sharedDown(rows, false); + const want_shared = try a.dupe(u8, try f.r.bytes(f.shared_slots[0], most * slots * d * 4)); + defer a.free(want_shared); + for (0..5) |_| { + try f.poison(f.shared_slots[1], most * slots * d * 4); + try f.sharedDown(rows, true); + try check.sameBytes("slot 10 and untouched routed slots", try f.r.bytes(f.shared_slots[1], want_shared.len), want_shared); + } + try check.expect(try f.r.intact(), "Shared down buffer guards", .{}); + check.pass("moetile rows={d}: all slots byte equal, five repeats, guards intact", .{rows}); + } + + fn time(f: *Fixture, wide: bool) !f64 { + var start = try cuda.Event.init(f.r.gpu.d, true); + defer start.deinit(); + var end = try cuda.Event.init(f.r.gpu.d, true); + defer end.deinit(); + for (0..3) |_| { + try f.sharedDown(most, wide); + } + try start.record(f.r.s); + for (0..10) |_| { + try f.sharedDown(most, wide); + } + try end.record(f.r.s); + try end.synchronize(); + return @as(f64, try cuda.Event.elapsedMs(start, end)) / 10; + } +}; + +pub fn run(gpu: check.Gpu) !void { + var f = try Fixture.init(gpu); + defer f.deinit(); + for (row_counts) |rows| try f.compare(rows); +} + +pub fn bench(gpu: check.Gpu) !void { + var f = try Fixture.init(gpu); + defer f.deinit(); + try f.compare(most); + var totals = [_]f64{ 0, 0 }; + for (0..6) |round| for (0..2) |j| { + const i = (round + j) % 2; + totals[i] += try f.time(i == 1); + }; + const old = totals[0] / 6; + const new = totals[1] / 6; + check.pass("moetile shared-down+scatter rows=4218 old={d:.3}ms new={d:.3}ms saved={d:.3}ms/layer {d:.3}ms/48-layer-pass", .{ old, new, old - new, (old - new) * 48 }); +} + +test "Shared down registrations cover four exact geometries and replacement hooks" { + const records = @import("core").contract.registry.all_cuda; + const names = cuda.qmmf_shared.symbols; + const shared = [_]u32{ 24576, 40960, 49152, 65536 }; + for (names, shared) |name, smem| { + const record = for (records) |record| { + if (std.mem.eql(u8, name, record.launch.symbol)) break record; + } else return error.MissingMoETileRecord; + try std.testing.expectEqual(smem, record.launch.dynamic); + try std.testing.expectEqual(@as(u32, 128), record.launch.threads); + try std.testing.expect(!record.geometry_pending); + } + var lane: cuda.qmmf.Lane = undefined; + for (0..4) |i| { + var p: @import("contract_qmmf_shared.zig").Product = .{ .rig = undefined, .lane = &lane, .index = i }; + const kernel = p.kernel(); + try std.testing.expect(kernel.spec.replaces != null and kernel.vt.replaced != null); + } +} diff --git a/zig/tests/cuda/qmmf_prompt_tests.zig b/zig/tests/cuda/qmmf_prompt_tests.zig new file mode 100644 index 0000000000..3721833ca1 --- /dev/null +++ b/zig/tests/cuda/qmmf_prompt_tests.zig @@ -0,0 +1,465 @@ +//! BF16 prompt tiles against the unchanged lane kernels, including forced small tiles and GPU event timing. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const q = cuda.qmmf; +const rows = [_]usize{ 1, 16, 64, 65, 127, 128, 129, 511, 512, 1536, 4218 }; +pub const Projection = struct { name: []const u8, n: usize, k: usize, out: q.Out, calls: usize }; +/// BF16 projections in the serving snapshot fc694b54: one normal main-model chunk on TP1 or rank 0. +pub const normal = [_]Projection{ + .{ .name = "GDN input", .n = 16480, .k = 2560, .out = .bf16, .calls = 36 }, + .{ .name = "attention input", .n = 13952, .k = 2560, .out = .bf16, .calls = 12 }, + .{ .name = "mixer output", .n = 2560, .k = 6144, .out = .f32, .calls = 48 }, + .{ .name = "shared gate/up", .n = 640, .k = 2560, .out = .bf16, .calls = 96 }, + .{ .name = "shared down", .n = 2560, .k = 640, .out = .f32, .calls = 48 }, + .{ .name = "HC up", .n = 10240, .k = 320, .out = .bf16, .calls = 96 }, + .{ .name = "HC down/inject", .n = 324, .k = 10240, .out = .f32, .calls = 96 }, + .{ .name = "PLE key", .n = 10240, .k = 2560, .out = .bf16, .calls = 1 }, + .{ .name = "PLE value", .n = 2560, .k = 2560, .out = .f32, .calls = 1 }, +}; +/// .all logits uses these additional shapes; HC up is already in normal, and these do not enter its total. +pub const all_logits = [_]Projection{ + .{ .name = "final HC down", .n = 320, .k = 10240, .out = .f32, .calls = 1 }, + .{ .name = "full vocabulary", .n = 248320, .k = 2560, .out = .bf16, .calls = 1 }, + .{ .name = "sharded vocabulary", .n = 124160, .k = 2560, .out = .bf16, .calls = 1 }, +}; +pub const mtp_hidden: Projection = .{ .name = "MTP hidden", .n = 2560, .k = 2560, .out = .f32, .calls = 1 }; +const guard = 128; + +fn fixturePartBytes(m: usize, n: usize) usize { + const scratch_rows: usize = @min(m, q.fused_rows - 1); + return 8 * scratch_rows * n * @sizeOf(f32); +} + +const Fixture = struct { + gpu: check.Gpu, + s: cuda.Stream, + lane: *q.Lane, + w: q.Weight, + x: cuda.DeviceBuffer, + codes: cuda.DeviceBuffer, + old: cuda.DeviceBuffer, + new: cuda.DeviceBuffer, + part: cuda.DeviceBuffer, + ldx: usize, + + fn init(gpu: check.Gpu, s: cuda.Stream, lane: *q.Lane, n: usize, k: usize, m: usize) !Fixture { + return initSized(gpu, s, lane, n, k, m, fixturePartBytes(m, n)); + } + + // Bounded all-logit probes run at M>=512, so neither baseline nor candidate needs a split partial plane. + fn initBounded(gpu: check.Gpu, s: cuda.Stream, lane: *q.Lane, n: usize, k: usize) !Fixture { + return initSized(gpu, s, lane, n, k, 1536, 16); + } + + fn initSized(gpu: check.Gpu, s: cuda.Stream, lane: *q.Lane, n: usize, k: usize, m: usize, part_bytes: usize) !Fixture { + const a = gpu.gpa; + const npad = cuda.modelopt.padded(n); + const ldx = k + 64; + var prng = std.Random.DefaultPrng.init(0x4218_bf16 + n); + const rng = prng.random(); + const host_w = try a.alloc(u16, n * k); + defer a.free(host_w); + const tiled = try a.alloc(u16, npad * k); + defer a.free(tiled); + const host_x = try a.alloc(u16, m * ldx); + defer a.free(host_x); + for (host_w) |*v| v.* = finite(rng); + for (host_x) |*v| v.* = finite(rng); + q.packBf16(tiled, host_w, n, k); + var x = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(host_x)); + errdefer x.free(); + var codes = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(tiled)); + errdefer codes.free(); + var old = try cuda.DeviceBuffer.alloc(gpu.d, m * n * 4 + 2 * guard); + errdefer old.free(); + var new = try cuda.DeviceBuffer.alloc(gpu.d, old.len); + errdefer new.free(); + var part = try cuda.DeviceBuffer.alloc(gpu.d, part_bytes); + errdefer part.free(); + return .{ .gpu = gpu, .s = s, .lane = lane, .w = .{ .mode = .bf16, .codes = codes.ptr, .scales = 0, .n = @intCast(n), .k = @intCast(k), .npad = @intCast(npad), .scale = 0.8123 }, .x = x, .codes = codes, .old = old, .new = new, .part = part, .ldx = ldx }; + } + + fn deinit(f: *Fixture) void { + f.part.free(); + f.new.free(); + f.old.free(); + f.codes.free(); + f.x.free(); + } + + fn baseline(f: *Fixture, m: usize, out: q.Out) !void { + try f.product(m, out, false, f.old.ptr + guard); + } + + // Force the new tile even at one row, independently of the production threshold and slice meeting choice. + fn candidate(f: *Fixture, m: usize, out: q.Out) !void { + try f.product(m, out, true, f.new.ptr + guard); + } + + fn product(f: *Fixture, m: usize, out: q.Out, candidate_: bool, dest: u64) !void { + if (!candidate_) { + var old = f.lane.*; + old.prompt = null; + return old.matmulOut(f.s, f.x.ptr, f.ldx, m, f.w, dest, out, f.part.ptr); + } + const sk = f.w.sk; + const fused = sk > 1; + const fun = f.lane.prompt.?[@backingInt(out)][@intFromBool(fused)]; + const band = @max(1, @min((m + 127) / 128, (12 << 20) / (128 * f.w.k * 2))); + var args: cuda.Args = .{}; + args.add(f.x.ptr); + args.add(f.w.codes); + args.add(f.w.scales); + args.add(f.w.scale); + args.add(dest); + args.add(@as(u64, 0)); + args.add(@as(i32, @intCast(m))); + args.add(@as(i32, @intCast(f.w.n))); + args.add(@as(i32, @intCast(f.w.k))); + args.add(@as(i32, @intCast(sk))); + args.add(@as(i32, @intCast(f.w.npad))); + args.add(@as(i32, @intCast(if (m == 1) f.w.k else f.ldx))); + args.add(@as(i32, @intCast(band))); + try cuda.launch.launch(fun, .{ + .grid = .{ .x = @intCast(((m + 127) / 128) * ((f.w.n + 127) / 128)), .y = 1, .z = 1 }, + .block = .{ .x = 256, .y = 1, .z = 1 }, + .shared = q.prompt_shared, + }, f.s, &args); + } + + fn compare(f: *Fixture, m: usize, out: q.Out, repeats: usize) !void { + const bytes = m * f.w.n * @as(usize, if (out == .f32) 4 else 2); + try check.expect(f.x.len >= ((m - 1) * f.ldx + f.w.k) * 2, "prompt fixture input span", .{}); + try check.expect(f.codes.len >= @as(usize, f.w.npad) * f.w.k * 2, "prompt fixture weight span", .{}); + try check.expect(f.old.len >= bytes + 2 * guard and f.new.len >= bytes + 2 * guard, "prompt fixture output spans", .{}); + try check.expect(f.part.len >= f.lane.partBytes(m, f.w), "prompt fixture split scratch: have {d}, need {d}", .{ f.part.len, f.lane.partBytes(m, f.w) }); + const a = f.gpu.gpa; + const want = try a.alloc(u8, bytes); + defer a.free(want); + const got = try a.alloc(u8, bytes); + defer a.free(got); + try f.old.fill8(0xa5, f.s.handle); + try f.part.fill8(0xff, f.s.handle); + try f.baseline(m, out); + try f.s.synchronize(); + try f.old.download(guard, want); + for (0..repeats) |_| { + try f.new.fill8(0x5a, f.s.handle); + try f.candidate(m, out); + try f.s.synchronize(); + try f.new.download(guard, got); + try check.sameBytes("forced BF16 prompt tile", got, want); + try f.guards(bytes); + } + try f.new.fill8(0x5a, f.s.handle); + try f.part.fill8(0xff, f.s.handle); + try f.lane.matmulOut(f.s, f.x.ptr, f.ldx, m, f.w, f.new.ptr + guard, out, f.part.ptr); + try f.s.synchronize(); + try f.new.download(guard, got); + try check.sameBytes("production BF16 prompt dispatch", got, want); + try f.guards(bytes); + } + + fn guards(f: *Fixture, bytes: usize) !void { + var buf: [guard]u8 = undefined; + for ([_]usize{ 0, guard + bytes }) |at| { + try f.new.download(at, &buf); + for (buf) |v| try check.expect(v == 0x5a, "prompt tile output guard", .{}); + } + } + + fn timed(f: *Fixture, m: usize, prompt: bool, out: q.Out) !void { + try f.product(m, out, prompt, f.new.ptr + guard); + } + + fn time(f: *Fixture, m: usize, prompt: bool, out: q.Out) !f64 { + var t0 = try cuda.Event.init(f.gpu.d, true); + defer t0.deinit(); + var t1 = try cuda.Event.init(f.gpu.d, true); + defer t1.deinit(); + for (0..3) |_| { + try f.timed(m, prompt, out); + } + try t0.record(f.s); + for (0..20) |_| { + try f.timed(m, prompt, out); + } + try t1.record(f.s); + try t1.synchronize(); + return @as(f64, try cuda.Event.elapsedMs(t0, t1)) / 20; + } +}; + +fn finite(rng: std.Random) u16 { + const bits = rng.int(u16); + return (bits & 0x807f) | (@as(u16, 118 + rng.uintLessThan(u16, 16)) << 7); +} + +fn load(gpu: check.Gpu) !q.Lane { + const cap = try gpu.ctx.capability(); + try check.expect(cap / 10 == 12, "qmmf prompt tiles require sm_12x, got sm_{d}", .{cap}); + return q.Lane.load(gpu.d, cap / 10); +} + +pub fn run(gpu: check.Gpu) !void { + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var lane = try load(gpu); + defer lane.unload(); + for (normal) |shape| { + var f = try Fixture.init(gpu, s, &lane, shape.n, shape.k, 4218); + defer f.deinit(); + for ([_]u32{ 1, q.splitK(shape.n, shape.k) }, 0..) |sk, si| { + if (si == 1 and sk == 1) continue; + f.w = try q.pinned(f.w, sk); + for ([_]bool{ false, true }) |clusters| { + lane.clusters = clusters; + for ([_]q.Out{ .bf16, .f32 }) |out| for (rows) |m| { + try f.compare(m, out, if (m == 4218) 5 else 1); + check.pass("qmmf-prompt M={d} N={d} K={d} SK={d} out={t} clusters={}", .{ m, f.w.n, f.w.k, sk, out, clusters }); + }; + } + } + } + for ([_]usize{ 512, 1024 }) |m| { + var f = try Fixture.init(gpu, s, &lane, mtp_hidden.n, mtp_hidden.k, m); + defer f.deinit(); + f.w = try q.pinned(f.w, q.splitK(f.w.n, f.w.k)); + try f.compare(m, mtp_hidden.out, 5); + check.pass("qmmf-prompt MTP hidden M={d} N=2560 K=2560 SK=4 fp32", .{m}); + } + for (all_logits) |shape| { + var f = try Fixture.initBounded(gpu, s, &lane, shape.n, shape.k); + defer f.deinit(); + f.w = try q.pinned(f.w, q.splitK(shape.n, shape.k)); + for ([_]usize{ 512, 1536 }) |m| { + try f.compare(m, shape.out, 1); + check.pass("qmmf-prompt all-logit {s} M={d} N={d} K={d} SK={d} out={t}", .{ shape.name, m, shape.n, shape.k, f.w.sk, shape.out }); + } + } + try special(gpu, s, &lane); +} + +/// Reproduce the first reported fault without running the preceding matrix of cases. +pub fn repro(gpu: check.Gpu) !void { + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var lane = try load(gpu); + defer lane.unload(); + lane.clusters = false; + var f = try Fixture.init(gpu, s, &lane, 8240, 2560, 4218); + defer f.deinit(); + f.w = try q.pinned(f.w, 4); + try f.compare(65, .bf16, 1); +} + +fn special(gpu: check.Gpu, s: cuda.Stream, lane: *q.Lane) !void { + var f = try Fixture.init(gpu, s, lane, 129, 2048, 129); + defer f.deinit(); + const plants = [_]u16{ 0x7fc1, 0xffc2, 0x7f80, 0xff80, 0x0001, 0x8001, 0, 0x8000 }; + for (plants, 0..) |value, row| try f.x.upload(row * f.ldx * 2, std.mem.asBytes(&value)); + const infinities = [_]u16{ 0x7f80, 0xff80 }; + try f.x.upload(8 * f.ldx * 2, std.mem.sliceAsBytes(&infinities)); + const x_before = try check.download(gpu, f.x); + defer gpu.gpa.free(x_before); + const w_before = try check.download(gpu, f.codes); + defer gpu.gpa.free(w_before); + for ([_]u32{ 1, 4 }) |sk| { + f.w = try q.pinned(f.w, sk); + for ([_]q.Out{ .bf16, .f32 }) |out| try f.compare(129, out, 5); + } + const x_after = try check.download(gpu, f.x); + defer gpu.gpa.free(x_after); + const w_after = try check.download(gpu, f.codes); + defer gpu.gpa.free(w_after); + try check.sameBytes("prompt inputs unchanged", x_after, x_before); + try check.sameBytes("prompt weights unchanged", w_after, w_before); + check.pass("qmmf-prompt N=129: nonfinite, subnormal and signed-zero input bytes, immutable inputs", .{}); +} + +fn timings(f: *Fixture, m: usize, out: q.Out) ![2]f64 { + var samples: [3][3]f64 = undefined; + for (&samples) |*sample| { + sample[0] = try f.time(m, false, out); + sample[1] = try f.time(m, true, out); + sample[2] = try f.time(m, false, out); + } + var old: [3]f64 = undefined; + var new: [3]f64 = undefined; + for (samples, 0..) |sample, j| { + old[j] = (sample[0] + sample[2]) / 2; + new[j] = sample[1]; + } + return .{ check.median(&old), check.median(&new) }; +} + +pub fn bench(gpu: check.Gpu) !void { + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var lane = try load(gpu); + defer lane.unload(); + var totals = [2]f64{ 0, 0 }; + for (normal) |shape| { + var f = try Fixture.init(gpu, s, &lane, shape.n, shape.k, 4218); + defer f.deinit(); + f.w.scale = 1; + f.w = try q.pinned(f.w, q.splitK(shape.n, shape.k)); + const out = shape.out; + try f.compare(4218, out, 1); + const fun = lane.prompt.?[@backingInt(out)][@intFromBool(f.w.sk > 1)]; + std.debug.print("LIMITS qmmf-prompt regs={d} static_shared={d} blocks_per_sm={d}\n", .{ try fun.attribute(.num_regs), try fun.attribute(.shared_size_bytes), try fun.occupancy(256, q.prompt_shared) }); + const sample = try timings(&f, 4218, out); + const calls: f64 = @floatFromInt(shape.calls); + totals[0] += sample[0] * calls; + totals[1] += sample[1] * calls; + std.debug.print("BENCH qmmf-prompt rank0 normal-main {s} M=4218 N={d} K={d} SK={d} out={t} calls={d} old={d:.4} ms new={d:.4} ms speedup={d:.3}x\n", .{ shape.name, f.w.n, f.w.k, f.w.sk, out, shape.calls, sample[0], sample[1], sample[0] / sample[1] }); + } + std.debug.print("BENCH qmmf-prompt rank0 normal-main weighted 36/12/48/96/48/96/96/1/1: old={d:.3} ms new={d:.3} ms speedup={d:.3}x; excludes MTP and all-logit head\n", .{ totals[0], totals[1], totals[0] / totals[1] }); + for ([_]usize{ 512, 1024 }) |m| { + var f = try Fixture.init(gpu, s, &lane, mtp_hidden.n, mtp_hidden.k, m); + defer f.deinit(); + f.w.scale = 1; + f.w = try q.pinned(f.w, q.splitK(f.w.n, f.w.k)); + try f.compare(m, mtp_hidden.out, 1); + const sample = try timings(&f, m, mtp_hidden.out); + std.debug.print("BENCH qmmf-prompt MTP hidden M={d} N=2560 K=2560 SK=4 out=f32 old={d:.4} ms new={d:.4} ms speedup={d:.3}x; separate from normal-main\n", .{ m, sample[0], sample[1], sample[0] / sample[1] }); + } +} + +test "qmmf prompt scratch fixture uses usize before multiplication and covers every split row count" { + try std.testing.expectEqual(@as(usize, 17_139_200), fixturePartBytes(65, 8240)); + try std.testing.expectEqual(@as(usize, 67_238_400), fixturePartBytes(4218, 8240)); + var lane: q.Lane = undefined; + lane.clusters = false; + for (normal) |shape| { + const n = shape.n; + const w = try q.pinned(cuda.modelopt.bf16Weight(0, n, shape.k), q.splitK(n, shape.k)); + for (1..16385) |m| { + try std.testing.expect(fixturePartBytes(m, n) >= lane.partBytes(m, w)); + try std.testing.expect(fixturePartBytes(4218, n) >= lane.partBytes(m, w)); + } + } +} + +test "qmmf prompt scratch keeps the current normal weights and bounds all-logit probes" { + var calls: usize = 0; + for (normal) |projection| calls += projection.calls; + try std.testing.expectEqual(@as(usize, 434), calls); + const final = cuda.modelopt.bf16Weight(0, 320, 10240); + try std.testing.expectEqual(@as(u32, 384), final.npad); + try std.testing.expect(q.promptEligible(512, final)); + var lane: q.Lane = undefined; + lane.clusters = false; + for (all_logits) |projection| { + const w = cuda.modelopt.bf16Weight(0, projection.n, projection.k); + if (projection.n >= 124160) try std.testing.expectEqual(q.Out.bf16, projection.out); + for ([_]usize{ 512, 1536 }) |m| { + try std.testing.expect(q.promptEligible(m, w)); + try std.testing.expectEqual(@as(usize, 0), lane.partBytes(m, w)); + } + } + const hidden = cuda.modelopt.bf16Weight(0, mtp_hidden.n, mtp_hidden.k); + try std.testing.expectEqual(q.Out.f32, mtp_hidden.out); + try std.testing.expectEqual(@as(u32, 4), q.splitK(hidden.n, hidden.k)); + try std.testing.expect(!q.promptEligible(4 * 127, hidden)); + try std.testing.expect(q.promptEligible(4 * 128, hidden)); + try std.testing.expect(q.promptEligible(4 * 256, hidden)); +} + +test "qmmf prompt scratch timed launches share one output while byte comparisons keep separate outputs" { + const Trace = struct { + var calls: usize = 0; + var ids: [2]usize = undefined; + var args: [2][13]u64 = undefined; + var geometry: [2][7]u32 = undefined; + fn arg(comptime T: type, raw: ?[*]?*anyopaque, i: usize) T { + return @as(*const T, @ptrCast(@alignCast(raw.?[i].?))).*; + } + fn launch(f: cuda.abi.Function, gx: c_uint, gy: c_uint, gz: c_uint, bx: c_uint, by: c_uint, bz: c_uint, sm: c_uint, _: cuda.abi.Stream, raw: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) cuda.abi.Result { + const at = calls; + calls += 1; + ids[at] = @intFromPtr(f); + geometry[at] = .{ gx, gy, gz, bx, by, bz, sm }; + for (0..3) |i| args[at][i] = arg(u64, raw, i); + args[at][3] = @as(u32, @bitCast(arg(f32, raw, 3))); + for (4..6) |i| args[at][i] = arg(u64, raw, i); + for (6..13) |i| args[at][i] = @intCast(arg(i32, raw, i)); + return cuda.abi.success; + } + }; + var driver: cuda.Driver = undefined; + driver.api.cuLaunchKernel = Trace.launch; + const old: cuda.Function = .{ .d = &driver, .handle = @ptrFromInt(10) }; + const prompt: cuda.Function = .{ .d = &driver, .handle = @ptrFromInt(20) }; + var lane: q.Lane = .{ + .mod = undefined, + .fns = @splat(@splat(old)), + .f32_fns = @splat(@splat(old)), + .reduce = @splat(old), + .half = @splat(@splat(@splat(old))), + .half_below = 0, + .prompt = @splat(@splat(prompt)), + .clusters = false, + }; + var f: Fixture = .{ + .gpu = undefined, + .s = .{ .d = &driver, .handle = null }, + .lane = &lane, + .w = try q.pinned(cuda.modelopt.bf16Weight(0x2000, 2560, 2560), 4), + .x = .{ .d = &driver, .ptr = 0x1000, .len = 1 }, + .codes = .{ .d = &driver, .ptr = 0x2000, .len = 1 }, + .old = .{ .d = &driver, .ptr = 0x3000, .len = 1 }, + .new = .{ .d = &driver, .ptr = 0x4000, .len = 1 }, + .part = .{ .d = &driver, .ptr = 0x5000, .len = 1 }, + .ldx = 2560, + }; + for ([_]usize{ 512, 1024 }) |m| for ([_]q.Out{ .bf16, .f32 }) |out| { + Trace.calls = 0; + try f.timed(m, false, out); + try f.timed(m, true, out); + try std.testing.expectEqual(@as(usize, 2), Trace.calls); + try std.testing.expectEqualSlices(usize, &.{ 10, 20 }, &Trace.ids); + for (Trace.args) |a| try std.testing.expectEqual(f.new.ptr + guard, a[4]); + Trace.calls = 0; + try f.baseline(m, out); + try f.candidate(m, out); + try std.testing.expectEqual(f.old.ptr + guard, Trace.args[0][4]); + try std.testing.expectEqual(f.new.ptr + guard, Trace.args[1][4]); + try std.testing.expect(Trace.args[0][4] != Trace.args[1][4]); + }; +} + +test "qmmf prompt scratch production envelope covers TP1 and TP2 through the chunk cap" { + const dims: usize = 2560; + const chunk_cap: usize = 16384; + const weights = [_][3]u32{ + .{ 16480, 2560, 1 }, .{ 13952, 2560, 1 }, .{ 8240, 2560, 1 }, .{ 7296, 2560, 1 }, + .{ 2560, 6144, 0 }, .{ 2560, 3072, 0 }, .{ 10240, 2560, 0 }, .{ 2560, 2560, 0 }, + .{ 324, 10240, 0 }, .{ 320, 10240, 0 }, .{ 10240, 320, 0 }, .{ 640, 2560, 0 }, + .{ 320, 2560, 0 }, .{ 2560, 640, 0 }, .{ 2560, 320, 0 }, .{ 248320, 2560, 0 }, + .{ 124160, 2560, 0 }, .{ 79591, 2560, 0 }, + }; + var lane: q.Lane = undefined; + lane.clusters = false; + for (weights) |shape| { + const w: q.Weight = .{ .mode = .bf16, .codes = 0, .scales = 0, .n = shape[0], .k = shape[1], .npad = (shape[0] + 127) / 128 * 128, .sk = shape[2] }; + for (1..chunk_cap + 1) |m| { + const native_bytes: usize = 8 * m * dims * 4; + try std.testing.expect(lane.partBytes(m, w) <= native_bytes); + if (q.promptEligible(m, w)) try std.testing.expectEqual(@as(usize, 0), lane.partBytes(m, w)); + if (w.n == dims) { + const p = q.planFor(m, w, false); + try std.testing.expect(@as(usize, p.sk) * m * w.n * 4 <= native_bytes); + } + } + } + // MTP reserves at least two rows per token before fc_hidden projects all four streams. + const fc: q.Weight = .{ .mode = .bf16, .codes = 0, .scales = 0, .n = 2560, .k = 2560, .npad = 2560 }; + for (1..257) |m| { + const cap: usize = @max(32, 2 * m); + const native_bytes: usize = 8 * cap * dims * 4; + try std.testing.expect(lane.partBytes(4 * m, fc) <= native_bytes); + try std.testing.expect(4 * (4 * m) * dims * 4 <= native_bytes); + } +}