Flash Next on CUDA: serve NVIDIA's NVFP4 checkpoint on one GB10 or at --tp 2

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).

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;

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.

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 } } },

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);

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();

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);

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 {