Skip to content

flashnext cuda: a shared round's DeltaNet steps in one launch each over its lanes - #1942

Open
Arminova wants to merge 2 commits into
ashhart:mainfrom
Arminova:flashnext/deltanet-lanes
Open

Arminova wants to merge 2 commits into
ashhart:mainfrom
Arminova:flashnext/deltanet-lanes

Conversation

@Arminova

@Arminova Arminova commented Nov 2, 2026

Copy link
Copy Markdown
Contributor

What this changes

A shared round (decode or verify) ran each DeltaNet step once per stream: the rollback copies, the front (or, with pending rows, the front into the stream's log), the chain, and the conv window's shift. Now each step is one launch over all the round's lanes, each lane running the lone launch's program on its own buffers, so the bits are main's.

On by default; TF_FLASHNEXT_GDN_LANES=0 loops the lanes as before. A round falls back to the loop above 16 lanes, for a lane of 64 rows or more (that chain takes fn_gdn_chain_tpr), or for buffers not 16-byte aligned. Off with tensor parallelism. New file cuda_batch_gdn_lanes.zig.

Four new kernels (fn_gdn_lanes_keep, fn_gdn_lanes_front, fn_gdn_lanes_chain, fn_gdn_lanes_shift), each registered with the kernel contract as a bit-equal replacement of the per-lane launches it stands for (zig/tests/cuda/contract_flashnext_lanes.zig, in the Flash Next family set): 1 to 16 lanes of 1 to 32 rows each, mixing lanes with no rollback, with a rollback copy, and with pending rows (none, at the store threshold, past it, the log wrapping).

This is the batching we took out of #1763 for #1740, which optimisme has closed in its favour. It does not need #1739's lane table. Part of #1762.

Receipt

  • Environment: as flashnext cuda: pending DeltaNet rows, reused lane buffers and the shared expert on a side stream #1763 (GB10, driver 580.173.02, nvcc 13.3.73, nvidia/Qwen3.8-Flash-Next-FP4-Local at b9ce18a6, a fresh server for every run), on main 088cfbd35. Base: main; this change fdb221871. The head 574d965c6 is that code with the four kernels' measured launch records and the lane loop's else folded into its condition (so cuda_batch.zig stays within 600 lines next to flashnext cuda: a shared round's attention in one launch a kernel #1739), in one commit; served once more on the head: every reply equals main's at 1, 4 and 16 streams, alone and serial all equal. Mean of two alternated rounds (base, this change, this change, base).

  • Exactness: every reply's token_sha256 equals main's in every cell below: 1, 4 and 16 streams, greedy and seeded sampled, code and chat, and each of the 16 long prompts. bench_concurrent --alone --serial at 1 and 4 streams (--context 262144 --parallel 3) and at 16 streams (--context 32768 --parallel 16): 336 of 336 concurrent replies equal their solo run and 216 of 216 equal "draft": false. An 8,192-token prompt sent twice on 4 streams: the 5 replies resumed from 8,191 cached tokens equal the 2 fresh ones.

  • Speed, tok/s:

    bench_concurrent, 256 tokens greedy seeded sampled
    code x1 / x4 67.1 → 67.1 (+0.1 %) / 117.9 → 122.7 (+4.1 %) 66.4 → 66.2 (-0.4 %) / 91.9 → 93.7 (+2.0 %)
    chat x1 / x4 42.2 → 42.1 (-0.2 %) / 78.7 → 79.7 (+1.3 %) 41.7 → 41.7 (0.0 %) / 67.6 → 68.3 (+1.0 %)
    16 streams (--context 32768 --parallel 16) code / chat 283.0 → 307.2 (+8.6 %) / 248.7 → 270.2 (+8.6 %) 207.9 → 217.1 (+4.4 %) / 178.3 → 189.0 (+6.0 %)

    One stream takes the batched path too (a shared round of one lane) and ties.

  • Long contexts (flashnext cuda: a shared round's DeltaNet in one launch a kernel, behind TF_FLASHNEXT_GDN_LANES #1740's case, prompt cache off): --context 45056 --parallel 16 --prompt-cache-gib 0, 16 distinct prompts of about 30K tokens, --tokens 128: prompts of 25,152 to 32,208 tokens, --levels 4,16, greedy and seeded sampled. End-to-end tok/s (every token over the whole run, prefill included): 4 streams 6.69 → 6.70 / 6.63 → 6.64; 16 streams 6.78 → 6.86 (all 16 done in 302.2 → 298.4 s) / 6.88 → 6.91 (297.7 → 296.2 s). No cell is slower, and every reply equals main's. (aggregate_tps in these cells spans the first stream's first token to the last stream's last, so it mostly measures when the first of 16 staggered prefills ends; it moved 12.8 to 14.1 between runs of the same build.)

  • Tests on the head 574d965c6: zig build test 465 / 465 steps and test-cuda-host 285 / 285 steps, 0 failures; compile-cuda-tests builds. tools/zig/lean_check.py: 424 problems, the same as main (0 new).

  • Kernel contract (nvcc 13.3.73, GB10), on a4e1c3c6e (the head differs only by that one line in cuda_batch.zig): every gate line of the four kernels passes (sums: none declared).

    SUMMARY 534 kernels: 534 pass gates 1-4, 0 fail, 6 known findings; 344 controls: 344 red as planted, 0 missed; 229 limit problems

    kernel-contract-limits on fatbins for sm_80 / 86 / 89 / 120 / 121: none of the four kernels' records stale, no launch over a limit; the 274 stale records and the 229 limit problems are main's own on this nvcc build. A planted mutation (fn_gdn_lanes_shift moving the window one row too far) fails gates 1 and 3 from byte 0.

  • Not run: TP2 (one GB10; the change keeps today's path with tensor parallelism), Metal and HIP (CUDA-only change), CUDA chips other than GB10 (launch records only).

Checklist

  • Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.
  • No precision traded for speed. If the bits change, the description says which sums change and why.
  • A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt.
  • Prompt processing is no slower than the last release.
  • The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
  • New tests fail before the change, pass after it, and skip cleanly without their dependency. (The new kernels do not exist before this change.)
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths. No AI attribution lines.

…er its lanes

A shared verify or decode pass ran each DeltaNet step once per lane: the rollback copies, the front (or, with pending
rows, the front into the stream's log), the chain, and the conv window's shift. Now each step is one launch over all the
pass's lanes (fn_gdn_lanes_keep / _front / _chain / _shift), every lane running its lone launch's program on its own
buffers, so the bits are the per-lane path's.

On by default; TF_FLASHNEXT_GDN_LANES=0 loops the lanes as before. Falls back to the loop above 16 lanes, for a lane
of 64 rows or more, for unaligned row blocks, and with tensor parallelism. New code in cuda_batch_gdn_lanes.zig.

The four kernels are registered with the kernel contract as bit-equal replacements of the per-lane launches
(zig/tests/cuda/contract_flashnext_lanes.zig, in the Flash Next family set): 1 to 16 lanes of 1 to 32 rows, mixing
lanes with no rollback, with a rollback copy and with pending rows.
@ashhart

ashhart commented Nov 2, 2026

Copy link
Copy Markdown
Owner

Thanks @Arminova. We reviewed this independently. The four batched kernels call the same inlined gdn_front/gdn_rows/gdn_rows_late/shift_window programs as the lone launches, on each lane's own buffers, so the bits are main's. One ask, in the test: the contract's served shapes are 2, 5 and 16 lanes, so one lane (what a single stream takes) never runs, and nothing tests the 17-lane fallback. Could you add 1 to the served shapes? And a host test that once() returns false, with nothing launched, at 17 lanes, at a 64-row lane, on a misaligned pointer and with TF_FLASHNEXT_GDN_LANES=0? Our GB10 A/Bs on both Flash Next checkpoints are queued, and we'll post them here.

…t test of the rounds once() leaves to the loop

The contract's served shapes add one lane (a single stream's round). cuda_batch_gdn_lanes_test.zig runs once() with a
driver that counts launches: one, two and sixteen lanes launch three kernels (four with rollback copies); seventeen
lanes, a 64-row lane, a misaligned projection or keep buffer, tensor parallelism and TF_FLASHNEXT_GDN_LANES=0 return
false with nothing launched.
@Arminova

Arminova commented Nov 2, 2026

Copy link
Copy Markdown
Contributor Author

Thanks, Ash. Both are in 898839eb7, a new commit on top:

  • The contract's served shapes are now 1, 2, 5 and 16 lanes. On our GB10 the four kernels pass gates 1-3 at all four shapes, bit-equal to the per-lane launches (sums: none declared).
  • cuda_batch_gdn_lanes_test.zig runs once() with a driver that counts launches. As controls, 1, 2 and 16 lanes launch three kernels (four with rollback copies). 17 lanes, a 64-row lane, a misaligned projection pointer, a misaligned keep buffer, tensor parallelism and TF_FLASHNEXT_GDN_LANES=0 each return false with nothing launched. Changed to expect one launch at 17 lanes, zig build test fails there, so the test runs.

On the head: zig build test 465 / 465 steps, test-cuda-host 285 / 285, compile-cuda-tests builds, lean_check.py 0 new, zig fmt --check clean.

@ashhart

ashhart commented Nov 3, 2026

Copy link
Copy Markdown
Owner

Thanks @Arminova. Here's our GB10 A/B of 898839eb7 against main: Flash Next FP4, one GB10, order base-new-new-base.

Exact: 0 of 20 prompt rows differ, and every concurrent stream equals its run alone. Both 64K long runs pass on both arms.

Decode gains with more streams:

prompt, T 2 streams 4 streams 8 streams
code, 0 +2.0% +3.5% +5.8% (330.6 → 349.9)
chat, 0 +2.0% +4.1% +6.9% (217.2 → 232.1)
code, 1 +1.3% +2.4% +3.7%
chat, 1 +1.9% +3.0% +4.5%

Prompt speed is unchanged everywhere (2K to 64K, all within ±0.3%).

One stream is slightly slower, and steadily so:

  • code at T 0: 69.1 → 68.9 tok/s (-0.3%), both new runs below both base runs (69.1/69.1 against 68.9/68.9);
  • chat at T 0: 43.3 → 43.2 (-0.2%), likewise;
  • code at T 1: -0.4%, though there the runs overlap.

That's inside our noise threshold, but one stream is the number we're pushing hardest, and we hold to "never slower" there. Could the merged launch fall back to the per-lane launches when a round has one lane? The kernels are already bit-equal to them, so it changes no bytes, and every stream count keeps its gain. Once that's in, I'll rerun the same cells and land it on a clean one-stream result.

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