Skip to content

Flash Next HC prompts: fuse writeback and normalization - #1847

Open
laudney wants to merge 15 commits into
ashhart:mainfrom
laudney:pr-fn-hc-norm-ready
Open

laudney wants to merge 15 commits into
ashhart:mainfrom
laudney:pr-fn-hc-norm-ready

Conversation

@laudney

@laudney laudney commented Nov 1, 2026 •

Copy link
Copy Markdown
Contributor

Flash Next writes each HC row back to BF16 and then launches normalization to read it again. This PR combines those passes for eligible BF16 prompt rows, while keeping H, the partial square sums and the normalized rows available to the next operators. Earlier standalone measurements at the default 1536-row chunk give a 12% reduction in the source-weighted main-model operator subset. Whole-model cold prefill remains flat; no served speedup is claimed.

Refs #1754. Depends on #1784, the standalone TP column split. One HC operator commit above that base. Incremental diff.

What

  • Add fn_hc_writeback_norm: one 256-thread block per row and stream, for four 2560-value streams. It handles the existing sums, MoE and FP32-slice writeback modes, then normalizes the stored BF16 values in the same launch.
  • Use it for eligible prompt rows of at least 64, with the checkpoint's BF16 HC projections and injection rows. Decode, shared rounds, captured kernel sets, pockets, PLE, intermediate keeps and unsupported layouts keep their existing launches.
  • Add replacement contracts, planted controls, resource records and guarded byte probes through the checkpoint's HC chain. The benchmark accepts MODEL [ROWS] and reports each mode separately, with native zero-slice spacing and the same final timing buffers for old and new.

Gate

The device receipts below are from the qualified standalone HC branch on GB10/sm_121 with CUDA 13.0.88 and checkpoint revision fc694b54fb0174e0913e6adf86691ef85a4ead47. This restack retains those HC kernel and adapter bytes and adds the PR2 base; its fresh checks are host tests and true consumer compilation. Prompt chunks remain at the default 1536 rows.

  • Bytes and arithmetic: H, partial square sums and normalized rows equal the original writeback-plus-normalize launches byte for byte. FP32 expert/slice addition order, BF16 rounding, square-sum reductions and normalization order are preserved. Synthetic and checkpoint-layer chains pass, including non-finite inputs, five repeats, row boundaries, guards and unchanged inputs.
  • Contract and controls: SUMMARY hc_writeback_norm: 4 pass, 0 fail; 2 controls red, 0 missed. This includes the three writeback modes and the existing normalization companion; each planted control fails its intended gate.
  • Resources: nvcc 13.0 records cover sm_80/86/89/120/121. The loaded sm_121 entry matches 48 registers, 976 user static shared bytes and a 256-thread bound. All 90 old HC SM/function entries retain identical SASS; the new kernel adds five entries.
  • Host: 23 focused HC host tests and the resource checks pass, as do native/runner compilation, format and diff checks. Lean remains at 424 inherited findings, zero new.
  • Native: the earlier standalone HC-on-PR1 runs pass draft/plain token equality, mixed requests, cancellation, cold 2k/8k/32k/64k prompts and four concurrent 8k requests. Prefix-cache reuse passes with --no-drafts. The combined TP1/TP2 receipt and split-disabled control linked in Flash Next TP=2: exact output-column splits, 21% faster plain decode #1784 are later-series composition evidence, not a new isolated native run of this restack.
  • Limits: this combines fresh affected-operator checks with reused unchanged gates. The full PR2 contract exit 1 and its 181 inherited resource diagnostics remain recorded in Flash Next TP=2: exact output-column splits, 21% faster plain decode #1784; this is not a globally clean full-suite claim. Other physical GPUs and later compiler sets remain unrun. CI is separate.

Numbers

Earlier standalone HC-on-PR1 measurements: checkpoint layer-0 inputs, GB10, alternating old/new order, event timing, nine measured samples after two warmups. These are operator timings with repeated buffers. The source predicts 94 eligible main-model calls per full chunk: 1 attention sums, 46 attention MoE and 47 MLP slice calls; MTP and PLE are excluded.

1536-row writeback + norm Calls Old per call New per call
Attention, sums 1 0.4482 ms 0.2896 ms
Attention, MoE 46 1.4025 ms 1.2479 ms
MLP, one partial plus zero slice 47 0.7565 ms 0.6541 ms
Source-weighted estimate 94 100.519 ms 88.436 ms

The weighted saving is 12.083 ms (12.0%) per modeled main-model chunk, not a measured whole-chunk or served saving. In that earlier standalone profile, native cold 64k time is 44.777 s before and 44.884 s after (+0.24%); the 2k–64k profile does not establish a stable served gain.

Overlap

#1764 also changes HC writeback and normalization. Its ABI and block layout differ, so Ash must choose the canonical operator before merge; the two implementations cannot land together unchanged. This PR contains only writeback plus normalization. The vector stream mix is a separate child PR.

backend() built the ranked table with a runtime call, so at TP=2 it
returned a pointer into its own stack frame and the first prompt
chunk on two GB10s jumped to address 0x25. The table is now built at
comptime like the scoring table, and a test reads it after a real
call and an overwritten stack, so a table in that frame fails it.
Flash Next's two-rank path was built and consumed --tp, but the CUDA
server refused the flag. It now serves --tp, Fish, Qwen4 and LFM2,
which open outside family admission, refuse it with the same words
family() uses, and --idle-unload is refused with --tp on CUDA, since
closing rank 0 ends its followers and a reload would have none.
Two Flash Next ranks only started when the operator set
TENSORFOLD_CHECKPOINT_SHA256, which nothing documented. Each rank now
sends the cached full digest GLM's ranks use, and the variable still
overrides it. A cold digest reads every byte, so the link and its
hello wait 30 minutes for the peer when the digest is computed.
NVIDIA's NVFP4 Flash Next checkpoint keeps its attention, DeltaNet
and shared-expert projections in BF16, so they are absent from
quantized_layers, and the check refused every one of them. Each
projection may now be MXFP8/32 or unlisted, and the parts the loader
stitches into one projection must all be one or the other.
NVIDIA's NVFP4 Flash Next checkpoint stores its n-gram table as E4M3
rows with one BF16 scale for the whole table, which the loader
refused. A third row format reads them as the reference does, each
code times the scale rounded once to BF16, through the mapped
gather, the positional reader, the distinct gather and warm-ahead.
The MTP head of NVIDIA's NVFP4 Flash Next checkpoint keeps its routed
experts in E4M3 with BF16 128x128 block scales, which only an NVFP4
loader read. They now pack for the existing FP8 expert kernels, the
MoE pass dispatches them from a file of their own, and the head's
memory estimate counts them at a byte a weight with their scales; a
device test runs that call at the head's 512 experts, 640 by 2560,
against an f64 reference, as reply equality cannot catch a wrong
draft that verification rejects.
The recipe gains the two-rank command for
nvidia/Qwen3.8-Flash-Next-NVFP4 with what each rank needs: the same
executable, NCCL 2.30 or later, --prompt-cache-gib 0 with drafts and
the digest's first-start cost, and states that concurrent streams do
not yet share a verify pass; the changelog records the branch.
cuda.zig went past the 600-line rule with main's latest lines, so
Fish's and Qwen4's --tp refusal moves into refusals.backend(), which
open() already calls before either family opens, and the test that
the server serves --tp and admits Flash Next's two ranks moves beside
the refusals' own rank tests; open() reads as it does on main.
The checkpoint's BF16 dense and shared layers run main's qmmf BF16
lane and its MTP head's routed experts run main's fp8_experts
kernels, which the contract had no gates 1-4 fixture for at Flash
Next's shapes. A new set runs the lane's BF16-weight product with
bf16 and fp32 outputs at the dense and shared shapes (decode rows),
and the FP8 expert down at decode and prompt pair counts against
the float64 reference, each with a planted 1.5x scale control; the
other experts hold NaN codes, so a misread expert goes red. It joins
the full contract and runs alone as flashnext-modelopt-contract.
At --tp 2 two of six replies differed from one GPU because each
routed NVFP4 expert was cut by intermediate channels and the ranks'
FP32 halves added in the all-reduce. Each rank now loads 256 whole
experts, fn_ep_picks maps the pass's picks to the rank's own ids
with every other id (and the shared slot) skipped, and the rows a
rank skips stay +0, so the all-reduce adds x + 0 and every pair
keeps one GPU's sums. One GPU and the expert kernels are unchanged;
the TP device test now checks the two ranks' sum equals one GPU's
bytes, and fn_ep_picks has its contract fixture, control and nvcc
13.0 records.
The NVFP4 checkpoint's entry now says that at --tp 2 each rank owns
half the routed experts whole, so replies are one GPU's bytes, which
the expert-parallel fix on this branch made true.
PR1 adds three comment warnings against the current upstream base.
Keep the contract descriptions on one line within 120 columns so
Lean accepts the branch without changing test or runtime behavior.
Gather full BF16 head and shared inputs so each rank computes its
own output columns with the original full-K slice order. Copy the
disjoint FP32 outputs as raw bits, preserving replies while reducing
duplicated TP2 work and retaining whole routed experts.
Load an independently built full-head image as the GDN byte oracle.
Compare full and half head layouts on finite and mixed non-finite
inputs at both ranks, with repeated outputs and explicit witnesses
so templated kernels cannot hide a shared arithmetic change.
Remove the extra HC normalization launch and H read for eligible
BF16 prompt rows while keeping writeback bytes, reduction order
and fallbacks intact. Include replacement gates, measured resources
and row-count benchmarks without changing the head split or gathers.
@ashhart

ashhart commented Nov 2, 2026

Copy link
Copy Markdown
Owner

Thanks @laudney, and sorry for the wait. #1764 (@Arminova) covers the same ground as this PR (the BF16 128 x 128 prompt tile in #1846, the hyper-connection writeback-and-normalize fuse in #1847), and it's GREEN on both Flash Next checkpoints on a GB10, with cold prompts +3.5 to +5.2% on FP4. It also runs on any SM behind its own switches. So it lands first. Your #1847 does the fuse in one kernel where #1764 uses three, and #1846 tiles at 512 rows on sm_12x. Once #1764 is on main, could you rebase onto it and keep whatever is still faster on top? We'll A/B any increment against main then, on both checkpoints, and land it with your name on it.

@laudney

laudney commented Nov 2, 2026

Copy link
Copy Markdown
Contributor Author

@ashhart Understood. This one waits for #1764. Once #1764 is on main, I'll rebase this onto it as you asked, keep only what is still faster (the one-kernel fuse where #1764 uses three), and post an A/B against main on both checkpoints. #1849 follows this PR.

@ashhart

ashhart commented Nov 3, 2026

Copy link
Copy Markdown
Owner

@laudney #1764 is on main now, so this is ready for the rebase you described. Keep only what is still faster than #1764's tile, then post the A/B against main on both checkpoints.

@laudney

laudney commented Nov 3, 2026

Copy link
Copy Markdown
Contributor Author

@ashhart #1764 is on main, so I measured this PR's fused kernel against #1764's. Nothing here is faster on a real pass, so I suggest we close this PR. If you agree, I will close it.

Mode Where a prompt pass uses it This PR / #1764, 1,536 rows (two GB10s) 4,096 rows
norm0 layer 0's read-out 0.953 / 0.987 0.964 / 0.963
moe layer start (MoE slots) 0.989 / 0.995 0.990 / 0.989
slices1 layer start (FP32 MoE) 0.942 / 0.940 0.945 / 0.943
slices2 every mixer writeback on one GPU (projection slice + zero slice) 1.043 / 1.029 1.049 / 1.044

Every layer runs one slices2 writeback, and this PR's kernel is slower there. With moe at the layer start, a layer comes out about 10 to 15 µs slower at 1,536 rows. With slices1, the two cancel. The cause is the grid: this PR uses one block for each (row, stream), so each stream block reads the row's branch again. #1764 reads it once a row for all four streams.

The second difference is in the read-out. This PR's read-out calls the qmmf lane directly. #1764's readoutNative goes through nativeMatmul, so from 1,024 rows its HC down and up use the prompt16 tile. A port of this PR would have to keep that too.

Method. One CUDA binary builds main's flashnext_hc.cu and this PR's flashnext_hc_fused.cuh side by side, and launches both on the same buffers: synthetic BF16 H, injections and scales, FP32 MoE slots, weights and slices, at 64 / 256 / 1,024 / 1,536 / 4,096 rows. H, PSS and NORM are equal byte for byte in all 20 cases on both GPUs. Times are the median of 7 repetitions of 100 launches (CUDA events, 10 warm-up launches). The test ran on GB10 (sm_121), nvcc 13.4.92 and driver 580.178.04.

I did not run a served A/B on the two checkpoints. These kernels have a fixed shape (2,560 × 4 streams), and they read no checkpoint values, so the kernel result is the same for NVFP4 and FP4-Local. The best case for this PR is a small loss, so a served run cannot show a gain.

#1849 does not depend on this fuse on main: its vector mix replaces fn_hc_dense_mix8, which #1764 leaves unchanged. I am restacking it onto main by itself, as you asked there.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants