diff --git a/CHANGELOG.md b/CHANGELOG.md index 6952ad3769..a52b030100 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -67,6 +67,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - `serve --profile stable` names today's defaults explicitly, and startup reports the profile and any registered switch set away from it; seven CUDA KV switches (`TF_QWEN_KV_PAGES`, `TF_QWEN_KV_SHARE`, `TF_GLM_KV_SHARE`, `TF_KV_SHARE`, `TF_FLASHNEXT_KV_PAGES` and two `_ALL` checks) are read through one registry, with environment values winning byte for byte; `fast-exact` and `fast` are refused until qualified (#1944, part of #1660). Thanks to @garvrao80. - Compaction keeps a memory note sent as text parts instead of summarizing the conversation again, and a note or remembered file path that arrived with carriage returns no longer carries them into the next note; Cortheon's checks read a shell command wherever it sits among the call's fields, its search fallback stays within its 120-byte cap, and its text cuts never end in a partial character (#1939). Thanks to @ppardus. +- 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 equal their solo runs in 8 of 8 cells (#1771, for #1754). Thanks to @laudney. - A model family can declare that its replies never carry a reasoning block, so the whole reply is answer content whatever `enable_thinking` the chat template renders; the OCR models in #1927 are the first to need it (#1933). Thanks to @laudney. - GLM-5.3-Flash on Macs with `--learn-conversations`: a prompt state restored from disk is charged at its full size (its attention rows as well as its KDA states) before it is chosen, prefetched or recalled, so the prompt cache no longer grows past its budget by about 2.5 GB a restored state at 130,000 tokens (part of #1572). diff --git a/docs/recipes/flash-next-cuda.md b/docs/recipes/flash-next-cuda.md index a422cbc8d2..0554230046 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/kernels/cuda/flashnext_batch.cu b/zig/kernels/cuda/flashnext_batch.cu index 8722a528f9..a2cf17ea72 100644 --- a/zig/kernels/cuda/flashnext_batch.cu +++ b/zig/kernels/cuda/flashnext_batch.cu @@ -525,6 +525,14 @@ 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. diff --git a/zig/src/core/contract/flashnext_cuda_records.zig b/zig/src/core/contract/flashnext_cuda_records.zig index a1d63d1c23..4987349d3b 100644 --- a/zig/src/core/contract/flashnext_cuda_records.zig +++ b/zig/src/core/contract/flashnext_cuda_records.zig @@ -31,6 +31,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/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 48353c2eb3..ee28d09152 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -9,12 +9,11 @@ const wt = @import("cuda_prompt_wt.zig"); const load = @import("cuda_prompt_load.zig"); const parts = @import("cuda_batch_state.zig"); const late = @import("cuda_batch_late.zig"); -const shared = @import("cuda_batch_side.zig"); const hc = @import("cuda_batch_hc.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"); @@ -444,22 +443,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); - const side = try shared.fork(self, native, n); - 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); - try shared.join(self, side, native, n); - 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 96d818387d..a17d43c8a1 100644 --- a/zig/src/families/flashnext/cuda_batch_kernels.zig +++ b/zig/src/families/flashnext/cuda_batch_kernels.zig @@ -23,6 +23,8 @@ pub const Kernels = struct { prompt: ?cuda.nvfp4.Prompt, prompt_rows: usize, 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, @@ -39,6 +41,7 @@ pub const Kernels = struct { 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, @@ -74,11 +77,14 @@ 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.select_lanes = try k.batch.function("fn_qsa_select_lanes"); @@ -104,6 +110,7 @@ pub const Kernels = struct { pub fn deinit(k: *Kernels) void { if (k.expert_tile) |*tile| tile.unload(); + k.fp8.unload(); k.nvfp4.unload(); if (k.prompt) |*p| p.unload(); k.native_lane.unload(); diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index e43d92fd6d..f98fbf5057 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -270,6 +270,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, @@ -372,7 +373,7 @@ 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 + lanes_bytes; - 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; @@ -433,6 +434,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_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_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_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index abfaae8903..5c9f2e8bbd 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -117,10 +117,12 @@ pub fn Namespace(comptime Tri: type) type { return local; } - /// A request or load this engine refuses, in words; null: none of its own. - pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { - return @import("cuda_lanes_explain.zig").explain(err); + /// A request this engine refuses, in words; with no engine (cuda.zig's open), a load it refuses. + pub fn explain(ctx: ?*anyopaque, err: anyerror) ?[]const u8 { + const words = @import("cuda_lanes_explain.zig"); + return if (ctx == null) words.open(err) else words.explain(err); } + pub const explains_open = true; /// A stream's caches; with the head, its cache, a window's buffers and the drafts it holds for the next round. const Lane = struct { @@ -210,7 +212,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 }; } 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..386bd56044 100644 --- a/zig/src/families/flashnext/cuda_lanes_explain.zig +++ b/zig/src/families/flashnext/cuda_lanes_explain.zig @@ -1,5 +1,17 @@ //! The Flash Next CUDA engine's refusals in words. +/// A load this engine refuses before it serves, in words (cuda.zig asks with no engine: `explains_open`). +pub fn open(err: anyerror) ?[]const u8 { + return switch (err) { + error.FlashNextRanksOverTwo => "Flash Next on CUDA runs on one GPU or at --tp 2 (two ranks): serve with --tp 1 or --tp 2", + error.FlashNextRanksNeedModelopt => "--tp 2 splits Flash Next's ModelOpt checkpoints only; serve this MLX checkpoint on one GPU, without --tp", + error.FlashNextRanksOnCapturedGlue => "--tp 2 runs Flash Next's own glue, not a captured set: set TENSORFOLD_CUDA_KERNELS=native, or serve on one GPU", + error.BadFlashNextKvPages, error.BadFlashNextPageTest, error.FlashNextPagesUnsupported, error.UnsupportedCheckpoint => explain(err), + error.UnsupportedLearningFormat, error.TrainingOnCapturedGlue => explain(err), + else => null, + }; +} + /// A request or load this engine refuses, in words; null: none of its own. pub fn explain(err: anyerror) ?[]const u8 { return switch (err) { @@ -16,3 +28,13 @@ pub fn explain(err: anyerror) ?[]const u8 { error.TrainingOnCapturedGlue => "--slide trains on Flash Next's own glue: serve it without TENSORFOLD_CUDA_KERNELS", }; } + +test "--tp refusals at load are worded, and a request's words stay out of a load's" { + const t = @import("std").testing; + for ([_]anyerror{ error.FlashNextRanksOverTwo, error.FlashNextRanksNeedModelopt, error.FlashNextRanksOnCapturedGlue }) |e| + try t.expect(open(e) != null); + try t.expectEqualStrings(explain(error.UnsupportedLearningFormat).?, open(error.UnsupportedLearningFormat).?); + // generic errors other load paths raise at one GPU too, and a request's out-of-memory, keep cuda.zig's own words + for ([_]anyerror{ error.UnsupportedQuantization, error.BadRank, error.OutOfDeviceMemory, error.PromptTooLong }) |e| + try t.expect(open(e) == null); +} diff --git a/zig/src/families/flashnext/cuda_lanes_load.zig b/zig/src/families/flashnext/cuda_lanes_load.zig index 4eb3a6aa9e..f3f570bc0b 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -14,12 +14,14 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu if (o.context > prompt.max_tokens) return error.ContextOverWindow; if (o.segments != 1) return error.UnsupportedOption; - if ((o.tp != 1 and o.tp != 2) or o.rank >= o.tp) return error.BadRank; + if (o.tp != 1 and o.tp != 2) return error.FlashNextRanksOverTwo; + if (o.rank >= o.tp) return error.BadRank; // refused before the stream, the glue or a captured set loads const chunk = try @import("cuda_chunk.zig").soloRows(try prompt.chunkRows()); // the largest pass, solo or not _ = 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; + if (o.tp != 1 and mode != .modelopt) return error.FlashNextRanksNeedModelopt; + if (o.tp != 1 and captured != null) return error.FlashNextRanksOnCapturedGlue; 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(); @@ -50,7 +52,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(); diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index e50399fae2..4301a1eacf 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -10,6 +10,8 @@ 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"); @@ -449,6 +451,7 @@ test { _ = @import("cuda_pages.zig"); _ = @import("cuda_kv_lanes.zig"); _ = @import("cuda_config.zig"); + _ = @import("cuda_lanes_explain.zig"); _ = @import("cuda_projection_load.zig"); _ = @import("cuda_rank_engine.zig"); _ = @import("cuda_round_graph.zig"); 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 519b4de746..3554a35cc8 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..7de0cc72cf --- /dev/null +++ b/zig/src/families/flashnext/cuda_moe_native.zig @@ -0,0 +1,62 @@ +//! 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 shared = @import("cuda_batch_side.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"); + const side = try shared.fork(self, native, n); + 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 shared.join(self, side, 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_ple.zig b/zig/src/families/flashnext/cuda_ple.zig index 333f8da38f..094d61e575 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; @@ -204,7 +205,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 }) @@ -330,6 +331,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, @@ -339,8 +343,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, @@ -370,13 +375,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| { @@ -396,15 +402,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); @@ -438,8 +445,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), } @@ -478,7 +487,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); @@ -536,5 +545,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_rank_engine.zig b/zig/src/families/flashnext/cuda_rank_engine.zig index 35a1650541..6821e6d3bf 100644 --- a/zig/src/families/flashnext/cuda_rank_engine.zig +++ b/zig/src/families/flashnext/cuda_rank_engine.zig @@ -30,11 +30,14 @@ 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" }); @@ -44,6 +47,11 @@ pub fn identity(a: std.mem.Allocator, io: std.Io, dir: []const u8, context: usiz 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()) }; } +/// 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 { gpa: std.mem.Allocator, rank: usize, @@ -59,12 +67,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) { @@ -490,3 +501,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..e65906b5d5 100644 --- a/zig/src/families/flashnext/cuda_rank_memory.zig +++ b/zig/src/families/flashnext/cuda_rank_memory.zig @@ -37,7 +37,7 @@ fn tensorBytes(name: []const u8, e: Entry, world: usize, rank: usize) !u64 { 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; 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/native/cuda.zig b/zig/src/native/cuda.zig index d62fae7eb3..39220e4d53 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 f73e1fc60d..84da81e3d4 100644 --- a/zig/tests/cuda/contract_family_sets.zig +++ b/zig/tests/cuda/contract_family_sets.zig @@ -1,4 +1,4 @@ -//! The contract's family sets: Nemotron's pocket, MiniMax-M3, MOSS-Image, GB10's grouped affine-4, Flash Next, EXL3. +//! The contract's family sets: pocket, MiniMax-M3, MOSS-Image, Flash Next NVFP4, grouped affine-4, Flash Next, EXL3. const std = @import("std"); const cuda = @import("cuda"); const check = @import("check.zig"); @@ -11,6 +11,7 @@ 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, flashnext: *@import("contract_flashnext_more.zig").Set, exl3: *@import("contract_exl3.zig").Set, @@ -23,6 +24,8 @@ 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); @@ -42,6 +45,7 @@ pub const Sets = struct { .pocket = pocket, .minimax = minimax, .moss = moss, + .modelopt = modelopt, .grouped = grouped, .flashnext = flashnext, .exl3 = exl3, @@ -54,6 +58,7 @@ pub const Sets = struct { sets.flashnext.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_modelopt.zig b/zig/tests/cuda/contract_flashnext_modelopt.zig new file mode 100644 index 0000000000..d93d885c29 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_modelopt.zig @@ -0,0 +1,407 @@ +//! 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 = 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/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_model_rank_tests.zig b/zig/tests/cuda/flashnext_model_rank_tests.zig index ebb937ae71..80f80fbd51 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(); 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_window_tests.zig b/zig/tests/cuda/flashnext_window_tests.zig index 999d773cd0..48673199cf 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(); diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index fd2376f18a..229afb681b 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -96,6 +96,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"); @@ -137,6 +138,7 @@ 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 \\ 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 @@ -246,6 +248,7 @@ const usage = \\ 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) @@ -348,6 +351,7 @@ 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, "nvfp4-fact")) return nvfp4_fact_tests.run(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); @@ -483,6 +487,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);