Flash Next shared down: remove the scatter with direct slot output

CHANGELOG.md

@@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel
 
 ## Unreleased
 
+- 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.
 
 - The tokenizer supports NFKC and NFKD normalizers, as Nemotron Parse uses, and a normalized added token decodes as the checkpoint spells it, matching Hugging Face tokenizers. No model we serve decodes differently (#1931, part of #1927). Thanks to @laudney.

docs/kernel-contract.md

@@ -336,3 +336,13 @@ A third registration executes the real cut kernel against the full FP32 bound an
 The product fixture writes float output so BF16 rounding cannot conceal that control.
 Both modes still require exceptional-class, repeat, guard, shape, accumulator and separate launch checks.
 See [the mode's scope and device checks](recipes/m5-attention-precision.md); other families' rounding is unchanged.
+
+Flash Next's BF16 shared-down entries `tf_qmmf_shared16_half`, `tf_qmmf_shared16`, `tf_qmmf_shared32` and
+`tf_qmmf_shared64` preserve the contiguous projection's FP32 result and store it directly in expert slot 10.
+The baseline shared expert has all 640 intermediate channels and runs only on rank 0 at two ranks.
+The contract checks all four geometries, nonfinite inputs, output bytes and untouched routed slots.
+`moetile_records.zig` holds launch geometry and measured nvcc 13.0 resources on all five served SMs.
+Later compiler receipts remain pending. Check each build with `tf-cuda-test kernel-contract-limits zig-out/fatbin`.
+`tf-cuda-test moetile` compares the direct slot with the contiguous product and scatter at 1, 16, 17, 32, 64, 65,
+1536 and 4218 rows, five repeats with guards. `moetile-contract` runs the replacement gates and planted control.
+`moetile-bench` alternates baseline and direct-slot GPU event timings at 4218 rows; no current timing is claimed.

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/build/flashnext_glue_host.zig

@@ -2,6 +2,13 @@
 const std = @import("std");
 
 pub fn steps(b: *std.Build, cuda: *std.Build.Module, core: *std.Build.Module, family: *std.Build.Module, all: *std.Build.Step) void {
+    const compose = b.step("test-flashnext-compose-host", "Dense heads and shared expert composition without a GPU");
+    compose.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = family, .filters = &.{ "split out lane", "head split", "dense split", "full heads", "half the heads", "shared expert composition" } })).step);
+    const moe = b.step("test-moetile-host", "Shared down slot fixtures, launch dispatch and contract records without a GPU");
+    all.dependOn(moe);
+    moe.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = cuda, .filters = &.{"shared down stride"} })).step);
+    const fixtures = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/moetile_tests.zig"), .target = b.graph.host, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "flashnext", .module = family } } });
+    moe.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = fixtures })).step);
     const root = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/contract_flashnext_glue_test.zig"), .target = b.graph.host, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "flashnext", .module = family } } });
     const run = b.addRunArtifact(b.addTest(.{ .root_module = root }));
     all.dependOn(&run.step);

zig/kernels/cuda/flashnext_batch.cu

@@ -37,9 +37,16 @@ extern "C" __global__ void fn_shift_window(__nv_bfloat16* __restrict__ win, cons
 // ---- the DeltaNet chain over a chunk: flashnext_gdn.cu's chain cut into front, rows and norm, with its bits ----
 
 namespace fn_gdn {
-constexpr int NK = 16, NV = 48, DK = 128, DV = 128, TAPS = 4;
-constexpr int C = 2 * NK * DK + NV * DV;
-constexpr int PW = C + NV * DV + 2 * NV;
+constexpr int DK = 128, DV = 128, TAPS = 4;
+
+// A layer's key and value heads (all 16 and 48, or one rank's 8 and 24): its conv channels and projection row width.
+template <int NK, int NV> struct G {
+    static constexpr int C = 2 * NK * DK + NV * DV;
+    static constexpr int PW = C + NV * DV + 2 * NV;
+};
+
+// The whole layer's heads, for the pending-row kernels that have no _h8 twin (a template's own NK, NV hide these).
+constexpr int NK = 16, NV = 48;
 
 static __device__ __forceinline__ float bf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); }
 
@@ -49,8 +56,10 @@ static __device__ __forceinline__ float sigmoidf_(float x) { return 1.0f / (1.0f
 
 static __device__ __forceinline__ float softplusf_(float x) { return x > 20.0f ? x : log1pf(expf(x)); }
 
+template <int NK, int NV>
 static __device__ __forceinline__ float conv_act(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs,
                                                  const __nv_bfloat16* __restrict__ cw, int r, int c) {
+    constexpr int C = G<NK, NV>::C, PW = G<NK, NV>::PW;
     float acc = 0.0f;
 #pragma unroll
     for (int tap = 0; tap < TAPS; ++tap) {
@@ -71,18 +80,19 @@ static __device__ __forceinline__ const float* gbs(const float* log, int log_row
 }  // namespace fn_late
 
 // Program (row, unit), 256 threads: key head u's normed q and k (u < NK), else a value head's v, decay g and beta.
-template <bool LOG>  // LOG: k, v and the gates go into the log's row after the pending ones
+template <int NK, int NV, bool LOG>  // LOG: k, v and the gates go into the log's row after the pending ones
 static __device__ __forceinline__ void gdn_front(
         const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, const __nv_bfloat16* __restrict__ cw,
         const float* __restrict__ a_log, const float* __restrict__ dt_bias, float* __restrict__ qn,
         float* __restrict__ kn, float* __restrict__ v, float* __restrict__ gb, const unsigned* __restrict__ ctl, int log_rows) {
     using namespace fn_gdn;
+    constexpr int C = G<NK, NV>::C, PW = G<NK, NV>::PW;
     const int r = blockIdx.x, u = blockIdx.y, t = threadIdx.x, warp = t >> 5, lane = t & 31;
     const int kr = LOG ? (int)(((*ctl & 0xffffu) + (*ctl >> 16) + (unsigned)r) % (unsigned)log_rows) : r;
     if (u < NK) {
         __shared__ float x[2][DK];
         const int hk = u;
-        x[t / DK][t % DK] = conv_act(P, cs, cw, r, t < DK ? hk * DK + t : NK * DK + hk * DK + (t - DK));
+        x[t / DK][t % DK] = conv_act<NK, NV>(P, cs, cw, r, t < DK ? hk * DK + t : NK * DK + hk * DK + (t - DK));
         __syncthreads();
         if (warp < 2) {
             float v4[4], ss = 0.0f;
@@ -100,7 +110,7 @@ static __device__ __forceinline__ void gdn_front(
         }
     } else {
         const int hv = u - NK;
-        if (t < DV) v[(size_t)kr * NV * DV + hv * DV + t] = conv_act(P, cs, cw, r, 2 * NK * DK + hv * DV + t);
+        if (t < DV) v[(size_t)kr * NV * DV + hv * DV + t] = conv_act<NK, NV>(P, cs, cw, r, 2 * NK * DK + hv * DV + t);
         else if (t == DV) {
             const float b = __bfloat162float(P[(size_t)r * PW + C + NV * DV + hv]);
             const float a = __bfloat162float(P[(size_t)r * PW + C + NV * DV + NV + hv]);
@@ -114,20 +124,29 @@ extern "C" __global__ void __launch_bounds__(256) fn_gdn_front(
         const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, const __nv_bfloat16* __restrict__ cw,
         const float* __restrict__ a_log, const float* __restrict__ dt_bias, float* __restrict__ qn,
         float* __restrict__ kn, float* __restrict__ v, float* __restrict__ gb) {
-    gdn_front<false>(P, cs, cw, a_log, dt_bias, qn, kn, v, gb, nullptr, 1);
+    gdn_front<16, 48, false>(P, cs, cw, a_log, dt_bias, qn, kn, v, gb, nullptr, 1);
 }
 
-// fn_gdn_front with the window's k, v and gates into the stream's log after its pending rows (q as fn_gdn_front's).
+// fn_gdn_front (whole heads) with the window's k, v and gates into the log after its pending rows (q as there).
 extern "C" __global__ void __launch_bounds__(256) fn_gdn_front_log(
         const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, const __nv_bfloat16* __restrict__ cw,
         const float* __restrict__ a_log, const float* __restrict__ dt_bias, float* __restrict__ qn,
         float* __restrict__ log, const unsigned* __restrict__ ctl, int log_rows) {
     using namespace fn_gdn;
     float* vl = log + (size_t)log_rows * NK * DK;
-    gdn_front<true>(P, cs, cw, a_log, dt_bias, qn, log, vl, vl + (size_t)log_rows * NV * DV, ctl, log_rows);
+    gdn_front<NK, NV, true>(P, cs, cw, a_log, dt_bias, qn, log, vl, vl + (size_t)log_rows * NV * DV, ctl, log_rows);
+}
+
+// fn_gdn_front on one rank's 8 key and 24 value heads (a projection row 8,240 wide, 5,120 conv channels).
+extern "C" __global__ void __launch_bounds__(256) fn_gdn_front_h8(
+        const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, const __nv_bfloat16* __restrict__ cw,
+        const float* __restrict__ a_log, const float* __restrict__ dt_bias, float* __restrict__ qn,
+        float* __restrict__ kn, float* __restrict__ v, float* __restrict__ gb) {
+    gdn_front<8, 24, false>(P, cs, cw, a_log, dt_bias, qn, kn, v, gb, nullptr, 1);
 }
 
 // The chain through the chunk from state_in to state_out (in place, or apart for a rollback): a warp 4 state rows.
+template <int NK, int NV>
 static __device__ __forceinline__ void gdn_rows(
         const float* __restrict__ qn, const float* __restrict__ kn, const float* __restrict__ v,
         const float* __restrict__ gb, const float* state_in, float* state_out, int rows, float* __restrict__ ys_out) {
@@ -197,13 +216,13 @@ static __device__ __forceinline__ void gdn_rows(
 extern "C" __global__ void __launch_bounds__(128) fn_gdn_rows(
         const float* __restrict__ qn, const float* __restrict__ kn, const float* __restrict__ v,
         const float* __restrict__ gb, float* state, int rows, float* __restrict__ ys_out) {
-    gdn_rows(qn, kn, v, gb, state, state, rows, ys_out);
+    gdn_rows<16, 48>(qn, kn, v, gb, state, state, rows, ys_out);
 }
 
 extern "C" __global__ void __launch_bounds__(128) fn_gdn_rows_to(
         const float* __restrict__ qn, const float* __restrict__ kn, const float* __restrict__ v,
         const float* __restrict__ gb, const float* state_in, float* state_out, int rows, float* __restrict__ ys_out) {
-    gdn_rows(qn, kn, v, gb, state_in, state_out, rows, ys_out);
+    gdn_rows<16, 48>(qn, kn, v, gb, state_in, state_out, rows, ys_out);
 }
 
 // A window's chain after its stream's pending rows: those first (state only), then its rows read out; fn_gdn_rows' sums
@@ -300,10 +319,25 @@ extern "C" __global__ void fn_late_roll(const fn_roll::Jobs t, int keep) {
     }
 }
 
-extern "C" __global__ void __launch_bounds__(128) fn_gdn_norm(
+// fn_gdn_rows and fn_gdn_rows_to on one rank's 8 key and 24 value heads.
+extern "C" __global__ void __launch_bounds__(128) fn_gdn_rows_h8(
+        const float* __restrict__ qn, const float* __restrict__ kn, const float* __restrict__ v,
+        const float* __restrict__ gb, float* state, int rows, float* __restrict__ ys_out) {
+    gdn_rows<8, 24>(qn, kn, v, gb, state, state, rows, ys_out);
+}
+
+extern "C" __global__ void __launch_bounds__(128) fn_gdn_rows_to_h8(
+        const float* __restrict__ qn, const float* __restrict__ kn, const float* __restrict__ v,
+        const float* __restrict__ gb, const float* state_in, float* state_out, int rows, float* __restrict__ ys_out) {
+    gdn_rows<8, 24>(qn, kn, v, gb, state_in, state_out, rows, ys_out);
+}
+
+template <int NK, int NV>
+static __device__ __forceinline__ void gdn_norm(
         const float* __restrict__ ys_in, const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ norm_w,
         float eps, __nv_bfloat16* __restrict__ out, float* __restrict__ xs) {
     using namespace fn_gdn;
+    constexpr int C = G<NK, NV>::C, PW = G<NK, NV>::PW;
     const int r = blockIdx.x, hv = blockIdx.y, t = threadIdx.x, warp = t >> 5, lane = t & 31;
     __shared__ float ys[DV];
     __shared__ float rinv;
@@ -328,6 +362,19 @@ extern "C" __global__ void __launch_bounds__(128) fn_gdn_norm(
     if (lane == 0) xs[(size_t)r * (NV * DV / 32) + hv * (DV / 32) + warp] = gs;
 }
 
+extern "C" __global__ void __launch_bounds__(128) fn_gdn_norm(
+        const float* __restrict__ ys_in, const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ norm_w,
+        float eps, __nv_bfloat16* __restrict__ out, float* __restrict__ xs) {
+    gdn_norm<16, 48>(ys_in, P, norm_w, eps, out, xs);
+}
+
+// fn_gdn_norm on one rank's 24 value heads (z read from its 8,240-wide projection rows).
+extern "C" __global__ void __launch_bounds__(128) fn_gdn_norm_h8(
+        const float* __restrict__ ys_in, const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ norm_w,
+        float eps, __nv_bfloat16* __restrict__ out, float* __restrict__ xs) {
+    gdn_norm<8, 24>(ys_in, P, norm_w, eps, out, xs);
+}
+
 // ---- sparse attention (QSA): each row's blocks ----------------------------------------------------------------------
 
 // attention._select's order-preserving key of a block score: negative floats flipped, positive ones topped.
@@ -471,13 +518,28 @@ extern "C" __global__ void fn_copy_cols(const uint16_t* __restrict__ src, long l
     if (col < width) dst[r * dst_pitch + col] = src[r * src_pitch + col];
 }
 
+// A rank's picks as its experts' ids: first..first+local-1 become 0..local-1, any other id (the shared slot too) local.
+extern "C" __global__ void fn_ep_picks(const int* __restrict__ picks, int* __restrict__ out, int n, int first, int local) {
+    const int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= n) return;
+    const int p = picks[i];
+    out[i] = p >= first && p < first + local ? p - first : local;
+}
+
 #include "flashnext_gdn_tpr.cuh"
 
 // The prompt chain on two threads a state row (gdn_tpr): fn_gdn_rows_to's bits, state_in == state_out in place.
 extern "C" __global__ void __launch_bounds__(256) fn_gdn_chain_tpr(const float* qn, const float* kn, const float* v,
                                                                  const float* gb, const float* state_in, float* state_out,
                                                                  int rows, float* ys_out) {
-    gdn_tpr::chain<2, 16>(qn, kn, v, gb, state_in, state_out, rows, ys_out);
+    gdn_tpr::chain<16, 48, 2, 16>(qn, kn, v, gb, state_in, state_out, rows, ys_out);
+}
+
+// fn_gdn_chain_tpr on one rank's 8 key and 24 value heads.
+extern "C" __global__ void __launch_bounds__(256) fn_gdn_chain_tpr_h8(const float* qn, const float* kn, const float* v,
+                                                                    const float* gb, const float* state_in,
+                                                                    float* state_out, int rows, float* ys_out) {
+    gdn_tpr::chain<8, 24, 2, 16>(qn, kn, v, gb, state_in, state_out, rows, ys_out);
 }
 
 // The prompt GDN block in one kernel a value head (front, chain, gated norm; their bits), no fp32 rows through memory.