Skip to content

flashnext cuda: pending DeltaNet rows, reused lane buffers and the shared expert on a side stream - #1763

Merged
ashhart merged 2 commits into
ashhart:mainfrom
Arminova:flashnext/decode-streams
Nov 2, 2026
Merged

ashhart merged 2 commits into
ashhart:mainfrom
Arminova:flashnext/decode-streams

Conversation

@Arminova

@Arminova Arminova commented Oct 31, 2026 •

Copy link
Copy Markdown
Contributor

What this changes

Three changes to a shared pass's rounds on one GPU. Each keeps the bits of main and has its own switch:

  • A partly kept window's DeltaNet rows stay pending in a per-stream log instead of being replayed; the next window's chain applies them first, so a rollback only moves the conv windows and the n-gram tail (TF_FLASHNEXT_KEEP_LATE=0 replays as before).
  • A contiguous lane's fixed-size buffers are reused across streams instead of allocated per request (TF_FLASHNEXT_LANE_PARTS=0).
  • A decode or verify pass's shared expert runs on a side stream beside the routed experts (TF_FLASHNEXT_SHARED_SIDE=0).

Three new kernels (fn_gdn_front_log, fn_gdn_rows_late, and fn_late_roll for a rollback's windows in one launch), each registered with the kernel contract as a bit-equal replacement of the kernels it stands for; the set joins the family sets, so contract_tests.zig is unchanged. Off or refused with tensor parallelism.

We had also batched every stream's DeltaNet steps into one launch each, pending rows included. We took that part out of this PR because #1740 does the same and was agreed first; we're glad to compare notes on it or review #1740.

Part of #1762.

Receipt

  • Environment: TensorFold main 94b90c9db and this change (one commit); DGX Spark (GB10, 128 GB), Ubuntu 24.04 aarch64, driver 580.173.02, nvcc 13.3.73; nvidia/Qwen3.8-Flash-Next-FP4-Local at b9ce18a6; --context 262144 --parallel 3, bf16 KV, prompt cache at its default size; a fresh server for every run. Since then rebased onto main 20705d604; the rebase moved only the contract set's registration.

  • Exactness: bench_concurrent --levels 1,4 --alone --serial: every reply equals its solo run and "draft": false. Every reply's token_sha256 equals main's at 1, 4 and 16 streams, greedy and seeded sampled, code and chat, at default settings and with TF_FLASHNEXT_CHUNK=8192 TF_FLASHNEXT_DRAFT_HEAD=affine4; bench_cache on cold 512 / 2,048 / 16,384-token prompts (and 8K / 32K / 128K at the second setting) gives main's replies, and repeated (resumed) and concurrent requests equal the first.

  • Decode speed at default settings, main → this change, tok/s, mean of two alternated rounds (main, this change, flashnext cuda: prompt passes with the same bits and less work, and a prompt scratch sized by use #1764, flashnext cuda: prompt passes with the same bits and less work, and a prompt scratch sized by use #1764, this change, main):

    greedy seeded sampled
    bench_openai code / chat (256 tokens, 3 reps) 66.1 → 67.3 (+2 %) / 41.4 → 42.4 (+2 %)
    bench_concurrent code x1 / x4 66.1 → 67.2 (+2 %) / 117.8 → 121.1 (+3 %) 64.7 → 65.2 (+1 %) / 89.3 → 91.2 (+2 %)
    bench_concurrent chat x1 / x4 41.4 → 42.3 (+2 %) / 76.3 → 78.9 (+3 %) 40.2 → 40.3 (0 %) / 64.5 → 65.8 (+2 %)
    16 streams (--context 32768 --parallel 16) code / chat 279.5 → 282.9 (+1 %) / 231.6 → 245.9 (+6 %) 195.3 → 200.3 (+3 %) / 159.1 → 167.1 (+5 %)
    decode after a cold 512 / 2,048 / 16,384-token prompt 68.2 → 69.3 / 38.8 → 39.3 / 45.5 → 46.1 (+1.5 %)

    With TF_FLASHNEXT_CHUNK=8192 TF_FLASHNEXT_DRAFT_HEAD=affine4 (one round): bench_concurrent code x1 / x4 74.6 → 75.8 / 131.1 → 132.1, chat 45.8 → 46.8 / 83.4 → 85.7.

  • Prompt speed: this change does not touch prompt passes; cold prefill at 2,048 / 16,384 tokens 1,569 → 1,590 / 1,677 → 1,679 tok/s (bench_cache --require-cold).

  • Memory: the prompt cache gets 7,233 → 6,583 MiB at default settings (the pending-row logs, 54.6 MiB a stream, and the reused lane buffers), 1,006 → 1,009 MiB with 8,192-row chunks.

  • Tests: zig build test 445 / 445 steps, 0 failures (2,431 tests run, 14 skipped; cached steps not counted); after the rebase test-cuda-host 274 / 274 steps, 3,911 passed, 8 skipped.

  • Kernel contract (nvcc 13.3.73, GB10): kernel-contract-limits on this change: none of its records stale, no launch over a limit. The full run's SUMMARY line follows in a comment.

  • Not run: TP2 (one GB10; every 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. (The CUDA SUMMARY follows in a comment.)
  • 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.

@ashhart

ashhart commented Oct 31, 2026

Copy link
Copy Markdown
Owner

Thanks @Arminova. We're starting our own device checks while your receipt finishes: the kernel contract (your three new kernels against the ones they replace) on one GB10, and served cells on Qwen3.8-Flash-Next-FP4-Local against main on another (prompt sizes, the 1/4/16-stream sweep, every reply compared). TP2 stays untested on our side too, as in #1762. #1764 shows a conflict with main and builds on this one, so we'll take it after #1763; no rush on that rebase until this one's through.

@Arminova
Arminova force-pushed the flashnext/decode-streams branch from c065ddf to 7b37e1d Compare October 31, 2026 20:06
@Arminova

Copy link
Copy Markdown
Contributor Author

Kernel contract on this change (GB10, nvcc 13.3.73; run at c065ddf, before the rebase onto 20705d604, which moved only the set's registration):

SUMMARY 466 kernels: 466 pass gates 1-4, 0 fail, 6 known findings; 324 controls: 324 red as planted, 0 missed; 223 limit problems

The limit problems are main's own records on this nvcc build (kernel-contract-limits on main's images lists the same 251 stale records); none is one of this PR's, and no launch is over a limit. The three new kernels' gate lines all pass. Marking this ready.

@ashhart

ashhart commented Nov 1, 2026

Copy link
Copy Markdown
Owner

Our device checks for this PR are done, @Arminova, and they're clean:

  • Kernel contract on a GB10: your three new kernels (fn_gdn_front_log, fn_gdn_rows_late, fn_late_roll) are bit-equal to the kernels they replace at every shape, repeat byte-identical over five runs and keep the stated non-finite rules. The full contract is 512/512, with all 342 controls red.
  • Served on Qwen3.8-Flash-Next MLX 4-bit, against main: tokens equal on all 20 prompt rows. The first run's long-prompt drop turned out to be our box's: with the order reversed, it's main's middle runs that drop at 64k, so it follows position, not your change. Four-stream greedy chat was +6.4% in the first order.

When your own receipt is in and you mark it ready, we'll gate it and land it. TP2 stays untested on our side, as in #1762.

@Arminova

Arminova commented Nov 1, 2026

Copy link
Copy Markdown
Contributor Author

Thanks for running it. Our receipt is in the description: it ran on main 94b90c9db plus this change, before the rebase onto 20705d604, which moved only the contract set's registration. The kernel contract's SUMMARY is in the comment above, and the PR is marked ready. If you'd like the receipt rerun on the rebased head 7b37e1d3d, say so and we'll add it.

Arminova and others added 2 commits November 2, 2026 01:39
…ared expert on a side stream

Three changes to a shared pass's rounds on one GPU. Each keeps the bits of main and has its own switch:

- A partly kept window's DeltaNet rows stay pending in a per-stream log instead of being replayed; the next window's
  chain applies them first, so a rollback only moves the conv windows and the n-gram tail.
  TF_FLASHNEXT_KEEP_LATE=0 replays as before.
- A contiguous lane's fixed-size buffers are reused across streams instead of allocated for every request.
  TF_FLASHNEXT_LANE_PARTS=0 allocates them as before.
- A decode or verify pass's shared expert runs on a side stream beside the routed experts.
  TF_FLASHNEXT_SHARED_SIDE=0 keeps it in line.

Three new kernels (fn_gdn_front_log, fn_gdn_rows_late, and fn_late_roll for a rollback's windows in one launch), each
registered with the kernel contract as a bit-equal replacement of the kernels it stands for
(zig/tests/cuda/contract_flashnext_late.zig, contract_flashnext_roll.zig; one set in contract_flashnext_more.zig,
registered with the family sets).
One launch for all streams' DeltaNet steps is left to ashhart#1737 / ashhart#1740.

Off or refused with tensor parallelism. Code moved into new files where a file would pass 600 lines
(cuda_batch_late, cuda_batch_side, cuda_lane_parts, cuda_seq_parts, cuda_lanes_snap,
cuda_mtp_roll).
@ashhart
ashhart force-pushed the flashnext/decode-streams branch from 890c4e6 to ac213a3 Compare November 2, 2026 01:39
@ashhart
ashhart merged commit 65686d7 into ashhart:main Nov 2, 2026
3 of 9 checks passed
@ashhart

ashhart commented Nov 2, 2026

Copy link
Copy Markdown
Owner

Thank you @Arminova. This is merged into main with your commits under your own name, so you're listed as a TensorFold contributor. It ships in the next release, credited to you in the changelog and the release notes.

@ashhart

ashhart commented Nov 3, 2026

Copy link
Copy Markdown
Owner

Thanks @Arminova! This is in TensorFold 1.2.2, out now: https://github.com/ashhart/TensorFold/releases/tag/v1.2.2

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