Skip to content

Flash Next BF16 prompts: larger CUDA tiles with identical output bytes - #1846

Open
laudney wants to merge 16 commits into
ashhart:mainfrom
laudney:pr-fn-bf16-prompt-ready
Open

laudney wants to merge 16 commits into
ashhart:mainfrom
laudney:pr-fn-bf16-prompt-ready

Conversation

@laudney

@laudney laudney commented Nov 1, 2026 •

Copy link
Copy Markdown
Contributor

Flash Next's large BF16 prompt projections use 64-row qmmf tiles. This PR adds a 128x128 tile on sm_12x. Standalone on PR1, its corrected 4218-row proxy is 819.752 → 502.403 ms (1.63x), and 64k cold TTFT is 45.166 → 40.650 s (10.00%). Those measurements precede the PR2-first stack; they are not an additive full-series gain.

Refs #1754. Depends on shared #1792, after PR2 #1784 and frozen PR1 #1771. One operator commit above #1792; head 9a91abef6. Incremental diff.

What

  • Add BF16/FP32 128x128 prompt entries with three K64 stages at eligible rows >=512 on sm_12x. Keep small-row and unsupported-device paths.
  • Retain pinned full-input K slices, ordered K16 MMA and output rounding. Keep one canonical qmmf body and scratch for the maximum eight slices.
  • Reuse one prompt eligibility predicate in qmmf and shared-slot selection. Direct shared output yields to eligible prompt projection, and remains whole-output only. Default split TP2 retains PR2's BF16 gather, full-K640 shared down and raw FP32 output-column copy; it never selects a half-output strided kernel. Forced-direct fixtures still test shared entries. The prompt default stays 1536 rows.

Gate (GB10, sm_121, CUDA 13.0.88)

Ash's latest order puts #1792 after #1784; this PR follows #1792. The issue requirements call for a separate replacement, matching bytes, contract records and a planted control. CI is checked separately on the current head.

  • Replacement bytes: qualified qmmf-prompt covers nine normal checkpoint projection profiles, MTP hidden rows, bounded full/sharded all-logit heads and nonfinite/subnormal/signed-zero inputs. The four parent contracts compare against the replaced 64-row path: 4 pass, 0 fail; 1 planted control red, 0 missed at gates 1–4. Repeats, unchanged inputs and guards pass.
  • Code generation/resources: qualified CUDA/qmmf blobs are unchanged. Baseline/shared K64 entries are preserved, and 20 prompt SM entries match the qualified image. CUDA 13.0 records cover five compiled SMs; loaded sm_121 entries match 95 registers (unsplit) or 255 (fused), zero user static shared bytes, 256-thread bounds and 98,432 dynamic shared bytes.
  • Current ordered host gate: 62/62 pass, including the explicitly discovered composition target (36/36), parent contract/scratch tests and shared tests. Dispatch boundaries 511/512/513, pinned slices and forced direct calls remain covered. The composition fixture checks the unique BF16 gather, two immediate interleave copies, full-K640 down and raw FP32 column copy; it accounts for 64-row reducers and initializes mock mixed input. The old filtered 30/200 totals did not discover that fixture.
  • Windows fixture compatibility: the four POSIX socket-pair tests now explicitly skip on Windows. The full FlashNext test module cross-compiles for x86_64-windows-gnu; Windows execution remains a separate CI check. The corrected shared target passes 36/36 on Linux, and the full final series passes 228/228 again. Runtime and CUDA code are unchanged.
  • Fixture sensitivity: the full FlashNext host module passes 228/228 at two seeds; the bad half-output direct-store guard control is red. This is host coverage, not a new device run.
  • Native/consumer scope: standalone-on-PR1 exact draft/plain, mixed/cancel, cold 2k/8k/32k/64k, concurrent 8k and non-drafting cache-reuse checks pass. Earlier qualified CUDA-runner/ModelOpt-startup compilation passed. At the ordered K32 tip, 15/15 consumer build steps pass, including tf-cuda-test fast native and tensorfold-native safe native. Inert images are compile-only; those binaries were not executed.

The ordered full-series source proof reports only two host-test-file changes against the qualified combined series; runtime and CUDA source are unchanged. The TP1 and TP2 split-on, plus TP2 split-off native receipts linked from PR2 #1784 are reused on that basis. The canonical qmmf image retains the 450-entry qualified union. Other physical SMs and later compilers remain unrun; focused gates do not establish a globally clean CUDA suite.

Numbers

The standalone-on-PR1 proxy alternates calls into one physical output buffer. Nine normal rank-0 shapes use source counts 36/12/48/96/48/96/96/1/1 (434 calls); MTP and all-logit heads are separate.

Standalone scope Before BF16 parent Change
4218-row weighted projection proxy 819.752 ms 502.403 ms 1.632x
Cold 64k median TTFT, three fresh prompts 45.166 s 40.650 s 10.00% lower

One buffer controls address effects; unchanged paths still show timing noise. These values do not measure this operator's incremental effect after PR2 or the net full-series gain.

Overlap

Order: #1771 → #1784 → #1792 → this PR → K32 #1848. HC #1847/#1849 is independent on PR2. GDN prefetch and NVFP4 K32-down stay local and off.

#1764's family-level prompt16 dispatcher recomputes split-K and can bypass qmmf. Choose one canonical BF16 dispatcher before combining the proposals, so pinned output cuts retain full-input sums and the K32 selector remains reachable.

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.
@laudney
laudney marked this pull request as ready for review November 1, 2026 14:43
@laudney
laudney force-pushed the pr-fn-bf16-prompt-ready branch from cafda94 to 1a8f2f7 Compare November 1, 2026 15:35
@laudney laudney changed the title Flash Next CUDA: add exact BF16 prompt projection tiles Flash Next BF16 prompts: larger CUDA tiles with identical output bytes Nov 1, 2026
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 force-pushed the pr-fn-bf16-prompt-ready branch from 1a8f2f7 to 9a91abe Compare November 1, 2026 16:21
Remove the whole shared expert's scatter by writing FP32 results
straight into slot 10. Retain PR2's full-input column copies for
split outputs, honor pinned K slices, and verify both routes with
replacement bytes, controls, records and a host trace. Skip the
POSIX socket fixture on Windows.
Use 128-row tiles for eligible BF16 prompt projections while keeping
the pinned K slices, ordered MMA chain and output rounding. Let the
shared-slot selector prefer this path and verify replacement bytes,
row boundaries, full MTP scratch, controls and measured resources.
@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 than #1764's tile (the 512-row sm_12x tiling), and post an A/B against main on both checkpoints. #1848 follows this PR.

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