Skip to content

Flash Next TP=2: exact output-column splits, 21% faster plain decode - #1784

Open
laudney wants to merge 15 commits into
ashhart:mainfrom
laudney:pr-fn-tp2-split
Open

laudney wants to merge 15 commits into
ashhart:mainfrom
laudney:pr-fn-tp2-split

Conversation

@laudney

@laudney laudney commented Oct 31, 2026 •

Copy link
Copy Markdown
Contributor

Refs #1754. Depends only on #1771, frozen at 6d2fc0361. This standalone PR2 is two commits on that base: the output-column split and its frozen-PR1 oracle. Incremental diff. The five operator PRs follow PR2.

Why

At --tp 2, running every BF16 dense projection whole on both ranks leaves much of the checkpoint's work duplicated. Splitting output columns keeps one GPU's sums and makes the two-GB10 setup faster: the standalone PR2 comparison gives 21.02% higher plain decode and 7.98% lower cold 64k TTFT against corrected PR1, with the default 1,536-row chunk unchanged.

What

  1. Split outputs, keep full dot products. Gather attention and DeltaNet head outputs and the shared expert's intermediate in BF16. Each rank computes its output columns over the full input, with K slices pinned to the whole weight's plan. Gather disjoint FP32 output columns as raw bits, without summing partial dot products across ranks.
  2. Keep state with its heads. The loader cuts the corresponding projection, convolution, attention and recurrent-state geometry. Rank agreement and admission check that layout; the MTP head retains its full geometry. PR1's whole 640-wide routed experts and ordered expert-slot combination remain intact. Split shared outputs keep full-K640 projections and raw FP32 column copies, with peer columns left zero.
  3. Check the split against frozen bytes. Frozen-PR1 GDN oracles cover full and half heads. Projection and launch-record tests check the head ownership, BF16 input gathers, pinned slices and raw FP32 output-column copies.

Gate

Ash's landing order is PR1, PR2, then #1792. All five operator PRs follow this standalone PR2. The arithmetic requirement is to keep one GPU's sum order and reply bytes, then show that output splitting pays against unsplit dense layers.

Standalone qualification uses one or two DGX Spark GB10s, sm_121, CUDA 13.0.88, and nvidia/Qwen3.8-Flash-Next-NVFP4 at revision fc694b54fb0174e0913e6adf86691ef85a4ead47.

  • Bytes: 720 frozen-PR1 GDN full/half output comparisons and 326 real-checkpoint projection comparisons pass, including mixed non-finite inputs and repeated runs.
  • Engine: six prose/code draft/plain and cross-TP equality cases pass. Eight concurrency cells at each TP, mixed requests, cancellation, non-drafting cache resume and cold 2k–64k checks pass.
  • Complete standalone CUDA contract: 546 kernels pass gates 1–4, with six declared known findings; all 349 planted controls are red. The command exits 1 for 181 resource problems reproduced on PR1. All seven new live entries pass; this is not a globally clean full-contract result.
  • Fresh full host gate on the identical source tree at 34704a314, Zig 0.17.0: zig build test test-kernel-contract test-golden fuzz-cortheon -j6 --summary all --test-timeout 10m exits 0. All 486 build steps succeed; 6,346 of 6,375 tests pass and 29 are skipped. Goldens give 322/335 equal with 13 declared known differences and zero unexpected differences, then 55/55 equal with zero unexpected differences. Cortheon fuzz passes 19/19.
  • Limits: physical-device qualification is CUDA 13.0.88 on sm_121. Later-compiler qualification and a runtime peer-kill test remain unrun. TP2 drafting uses cache zero; drafting with a positive or auto cache remains refused, as in PR1. CI is reported separately on the PR.

The later operator series has separate combined-runtime receipts: 16 affected contract cases pass, all seven affected planted controls are red, and the 326 checkpoint head/output comparisons pass again. Its 450 qmmf SM entries match the qualified standalone SASS union. That later runtime also passes TP1/TP2 draft/plain and canonical-token checks, mixed requests, cancellation, cold 2k/8k/64k prefixes, four concurrent 8k requests and non-drafting cache reuse; its TP2 split-disabled control passes. These receipts apply to the later kernel composition.

Numbers

Repeated standalone PR2 comparison against corrected PR1, TP2, safe native builds, same checkpoint and 1,536-row chunk:

Metric Corrected PR1 Standalone PR2 Change
Plain decode 18.971 tok/s 22.959 tok/s +21.02%
Cold 64k TTFT 73.782 s 67.893 s −7.98%

Cold 2k/8k/32k TTFT falls by 7.13%/8.24%/8.23%. The separate same-branch off/on/off control gives a 21.12% plain-decode gain. The standalone 64k profile has three fresh prompts per setting. TP1 retains higher cold throughput in the recorded reference runs.

Later operator series: absolute standard spark-docker benchmark of the combined 64d499e8f runtime. This run includes the later operators, with MTP and copy drafts enabled, prompt cache zero, a four-active-request cap, 4,218–4,219 input tokens and 512 output tokens per request:

Client concurrency Requests Aggregate output tok/s Mean request decode tok/s p50 TTFT (s)
1 4 33.0 45.9 4.264
2 4 41.4 28.2 5.056
4 8 55.9 19.5 6.637
8 16 56.0 18.0 40.360

All 32 requests complete without failures: 16,384 output tokens, zero cache hits, and 70.37% draft acceptance (11,548 accepted of 16,410 verified draft tokens). At client concurrency eight, requests queue behind the four-active cap. These are absolute later-series results, not a standalone PR2 measurement or an additive estimate of the gains above. The later composition's 64k exactness checks use one fresh prompt per setting.

Overlap

Shared down #1792, BF16 prompt #1846, K32 staging #1848, HC normalization #1847 and HC stream mix #1849 are separate follow-ups to this PR. Review PR2 relative to pr-fn-nvfp4.

BF16 #1846 and HC #1847 overlap #1764 and need a maintainer choice of canonical implementations before both alternatives are combined. PR2's output cuts require pinned full-weight K slices. #1764's family dispatcher bypasses qmmf and recomputes splitK, so a later combination must preserve those pinned plans.

@laudney laudney changed the title Flash Next on CUDA: split the dense layers across two ranks with exact rank sums Flash Next CUDA: split dense output columns with exact gathers Nov 1, 2026
@laudney
laudney force-pushed the pr-fn-tp2-split branch 2 times, most recently from 64d499e to 5970dab Compare November 1, 2026 15:35
@laudney laudney changed the title Flash Next CUDA: split dense output columns with exact gathers Flash Next TP=2: exact output-column splits, 21% faster plain decode Nov 1, 2026
@laudney

laudney commented Nov 2, 2026

Copy link
Copy Markdown
Contributor Author

@ashhart #1784 also conflicts with main (4de3944ad), in ten files: CHANGELOG.md, zig/kernels/cuda/flashnext_batch.cu, zig/src/core/contract/registry.zig, and seven files under zig/src/families/flashnext/ plus zig/tests/cuda/contract_family_sets.zig. It still carries PR 1 as it was at 6d2fc0361.

Do you want us to rebase it now, or wait until #1771 lands and rebase it once onto main then? Unless you prefer otherwise, we'll wait, so it gets one rebase onto the landed PR 1.

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.
At load, --tp 4 or 8 failed with a bare BadRank, and --tp 2 on an
MLX checkpoint or on captured kernels with a bare
UnsupportedQuantization, an error other one-GPU load paths also
raise. Each --tp refusal now has its own error, and Flash Next sets
explains_open so cuda.zig asks explain() at load: with no engine it
words only load refusals, so a request's words (out-of-memory, a
long prompt) never describe a failed load. The CHANGELOG entry
spells Qwen3.8-Flash-Next and credits ashhart#1771 for ashhart#1754.
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.
@laudney

laudney commented Nov 2, 2026

Copy link
Copy Markdown
Contributor Author

@ashhart Rebased onto #1771 as it now stands (db7e9f547, on main 933b969fa) and force-pushed. The branch is #1771's 13 commits plus this PR's two, so the old PR 1 copy at 6d2fc0361 is gone. Against current main (089a18dff), only CHANGELOG.md conflicts.

Before Now
Head 9ec8e7edb 1077d5fbf
Output-column split 348bbaee6 1bfd6cce6
Frozen GDN oracle test 9ec8e7edb 1077d5fbf
PR 1 under it 6d2fc0361 (frozen copy) db7e9f547 (#1771)

Main changed a lot under this PR since its base. The port:

  • Pending DeltaNet rows (fn_gdn_front_log, fn_gdn_rows_late, fn_late_roll): main runs these at --tp 1 only and refuses TF_FLASHNEXT_KEEP_LATE=1 at --tp, so they stay whole-head. gdn_front takes LOG next to <NK, NV>.
  • flashnext cuda: a shared round's attention in one launch a kernel #1739's one-launch shared-round attention: whole heads only. At the head split, a shared round runs its attention lane by lane, as before.
  • Packed lane parts, the prompt-cache snapshot module and the draft layout take the rank's heads. The side-stream shared expert still runs on one GPU, and cuda_shared_split.run covers --tp 2 and prompt fills.
  • To keep cuda_mtp_ops.zig under 600 lines, the rewind passes its kernel arguments as tuples. Same order and types.

Checks on 1077d5fbf: zig fmt and the lean check (0 new problems), zig build test test-kernel-contract test-golden fuzz-cortheon exits 0, flashnext_batch_serve_test.py 8/8, and compile-cuda-tests test-cuda-host pass. The nvcc 13.0 limits check on all five SMs gives the same 209 problems as #1771's tree, all from main. All seven new entries are current, and main's existing kernels keep their register counts.

I did not rerun anything on a GPU after this rebase. The TP=2 byte and speed numbers in the description were measured on the pre-rebase commits.

@ashhart

ashhart commented Nov 2, 2026

Copy link
Copy Markdown
Owner

Thanks @laudney. We reviewed this independently, and it lands after our TP=2 run on two GB10s. No dot product is split across ranks: every production cut is by output rows over the full K, and cut weights pin sk to the whole weight's splitK, checked before any read. Heads stay with their state, joins are copies not sums, and the shared slot uses #1771's one-rank-nonzero combine. One GPU is unchanged, uneven splits are refused at load, and admission prices the half. Two facts we'll note: the split is on by default at TP=2 and only for BF16 dense checkpoints, and the +21.02% decode and −7.98% cold-64K TTFT were measured before the rebase onto #1771 as it landed, so our run will re-measure them. Also, #1764 (landing now) must keep the pinned plans through its dispatcher, and we'll check that at rebase.

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