Skip to content

Flash Next CUDA: drafted windows of several streams verified in one pass - #1069

Closed
Mirrdhyn wants to merge 2 commits into
ashhart:mainfrom
hygurlabs:flashnext-cuda-shared-verify
Closed

Mirrdhyn wants to merge 2 commits into
ashhart:mainfrom
hygurlabs:flashnext-cuda-shared-verify

Conversation

@Mirrdhyn

Copy link
Copy Markdown
Contributor

What this changes

Draft, as agreed: the MTP verify windows batched across streams in the lanes engine, with the same bit-for-bit contract. Is the "CUDA per-stream drafting precision" work on #286 touching cuda_mtp_ops.zig / verifyFn too? If so, I'll rebase this onto it rather than collide.

verifyFn ran one batched pass per drafting stream, so a round over N streams read the weights N times and the aggregate stayed flat (81.8, 70.3 and 64.0 tok/s at 1, 2 and 4 streams on a GB10, code prompt, --parallel 4). Now the drafted windows of a round share one pass, up to the round's 64 rows (TF_FLASHNEXT_WINDOW_ROWS, 0 keeps each window alone; never past TF_FLASHNEXT_CHUNK).

  • batch.Lane gains a row span and the lane's own rollback (keep), hidden-row (streams) and score (logits) destinations. rows = 1 and none of them is the one-row round as before, which keeps its launches. A lane's keep may hold a window's DeltaNet rollback only (rollbackOnly); a training pass's kept tensors are refused, as the round already refuses moe_add.
  • Per lane, with the row count and launches its lone pass makes: n-gram gather, media inject, PLE convolution and tail, DeltaNet front, chain (to the lane's kept state) and conv shift, attention (rope, cache writes, pooling, indexer scores, selection, chunks).
  • Once over all rows: the matmuls, hyper-connection blocks, router, experts and head. An untaught stream's rows read no pocket, as in round.
  • Engine.windows (cuda_prompt_windows.zig) runs the pass. Ops.verifyMany refuses every window before any stream's state moves, then leaves each stream as its verify leaves it (saved tail and PLE rows, host scores row, fed). verifyFn groups the drafted windows (cuda_lanes_verify.zig). Plain one-row rounds and two ranks keep their paths (two ranks would need one wire command for the batch).

2 commits, 16 files, +376/-48, all in zig/src/families/flashnext/ (cuda_batch*.zig, cuda_prompt*.zig, cuda_mtp_ops.zig, cuda_lanes*.zig, cuda_round_graph.zig, which refuses a window it never captures, cuda_media.zig and two test row types) and zig/tests/cuda/ (flashnext_windows_tests.zig, main.zig).

Receipt

  • Environment: TensorFold main d1018b31 + this branch; NVIDIA DGX Spark, GB10 sm_121, driver 595.84, nvcc 13.2.86; checkpoint Flash Next MLX 4-bit with MTP (qwen4_exp), drafts on (depth 6, confidence 0.6), --context 32768 --parallel 4, chunk 16384.

  • Exactness:

    • new tf-cuda-test flashnext-windows MODEL IDS: several streams' forced windows in one shared pass against each window alone, on a 100-token and a 2,600-token prompt (the indexer's sparse selection), windows of {1,1,1,1}, {5,5,5,5}, {1,16,3,7}, {16,16} and {16,16,16,16} rows (4 to 64 in all; past 16 rows the router takes its many-row kernel, which a lone window never reaches): each window's scores, hidden rows and last host row, then after a partial rollback the position, n-gram history and next row, byte for byte: 10 of 10 pass. With every lane reading the first lane's position it fails (124152 of 124160 words differ).
    • bench_concurrent.py --alone --serial, levels 1/2/4, T=0 (2 reps) and T=1: 0 unequal, 0 failed; at T=0 the same token_sha256 per prompt and level as with TF_FLASHNEXT_WINDOW_ROWS=0. At --parallel 2 --context 131072, levels 1/4, 3 reps: 0 unequal.
  • Aggregate tok/s, 512 tokens each, same binary with TF_FLASHNEXT_WINDOW_ROWS=0 (each window alone, main's path) and the default:

    prompt T 1 stream, alone / shared 4 streams, alone 4 streams, shared
    code 0 80.9 / 80.5 79.9 134.0 x1.68
    chat 0 66.1 / 65.6 68.1 122.4 x1.80
    French prose 0 55.0 / 55.0 57.1 105.9 x1.85
    JSON list 0 82.5 / 82.7 82.6 140.3 x1.70
    code 1 71.3 / 71.6 75.0 120.1 x1.60
    chat 1 60.2 / 60.5 62.4 104.0 x1.67
    French prose 1 71.7 / 71.8 69.0 109.8 x1.59
    JSON list 1 82.2 / 82.2 80.5 130.6 x1.62

    Tokens per round are unchanged (the drafts and their acceptance are the same). bench_openai.py at --parallel 2 --context 131072, one stream: fibonacci-raw 81.7, gpu-chat-no-think 69.4 tok/s.

  • Prompt speed: n/a, prompt chunks do not go through the shared verify.

  • Tests run: zig fmt --check, native, test, test-kernel-contract, test-golden, fuzz-cortheon, test-cuda-host pass; lean check 444 problems (main's count); tf-cuda-test smoke, symbols, route-rows, experts-rows, flashnext-hc, flashnext-attn, flashnext-dense-hc, flashnext-draft-affine, flashnext-shared, flashnext-nvfp4-tp, norms, norm-rows, sample, projection-f32, mxfp8, quant-ref, modelopt, glue, glm-ref and flashnext-windows pass; kernel-contract SUMMARY 85 kernels: 85 pass gates 1-4, 0 fail, 0 known findings; 7 controls: 7 red as planted, 0 missed; 899 limit problems, the same 899 lines as main at d1018b31 (no kernel changes here).

  • Not run, and why:

    • Two ranks: unchanged path; Flash Next's two-rank CUDA loader takes its ranks from --tp, which the CUDA build does not serve.
    • A taught pocket (Living Weights) with a shared verify: no taught checkpoint here; the shared pass follows round's per-row pocket scales.
    • Metal: no Mac here, no Metal file touched.

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 (prompt passes are untouched), and concurrent equals solo.
  • No precision traded for speed. The bits do not change.
  • No GPU kernel is new or changed (the shared pass launches the existing ones).
  • Prompt processing is no slower than the last release (not reached).
  • CUDA GB10 run; Metal not run (no Mac, no Metal file touched).
  • New tests fail before the change, pass after it, and skip cleanly without their dependency: the shared verify does not exist before it, so the receipt shows the test failing on a planted per-lane position bug; flashnext-windows needs the checkpoint, like flashnext-window-rank.
  • 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.

…against each alone, byte for byte, dense and sparse, up to 64 rows
ashhart pushed a commit that referenced this pull request Oct 24, 2026
…s together (on #1069)

Depends on #1069 (lanes with row spans). The head drafted one stream at a time: at four streams a
round spent ~29 ms in head passes against ~5 ms at one. Now the core asks a shared round's drafts
in one call (.draft_streams) and draftBatched runs them together: each greedy stream's kept rows
absorbed in one head pass (a lane each, on its own head cache), then one pass a chain step for the
streams still drafting, their bests read in one download. Every row runs the ops a lone head pass
runs on it, so each stream gets the drafts it gets alone (same confidence stop, same count).
Sampled streams, trees and media keep the one-at-a-time path. TF_FLASHNEXT_BATCH_HEADS=0 restores
it; TF_FLASHNEXT_DEPTH_LOG=1 prints each stream's chain length by round width.

FP4-Local, GB10, main + #1069 + local stack, affine4 head, four streams: code 130.2 -> 150.7,
chat 88.9 -> 99.8 tok/s aggregate; one stream unchanged (70.7 / 44.2). Every request's rounds and
accepted/drafted equal the one-at-a-time run's; concurrent == alone 4/4, drafted == serial 1/1.
@ashhart

ashhart commented Oct 24, 2026

Copy link
Copy Markdown
Owner

Landed on main together with #1155 (cb60c86 and 4c78a76, under 9801a84), thank you @Mirrdhyn. On a GB10 with Qwen3.8 Flash Next (MLX 4-bit, MTP), the served sweep against main goes from 65.7 to 88.8 tok/s with 2 streams (+35%), 66.0 to 115.7 with 4 (+75%) and 66.0 to 142.4 with 8 (+116%). Single-stream and prefill (2k to 64k) are flat, and all 20 token rows are identical to main.

- Synthema, Ash's AI Assistant.

@ashhart ashhart closed this Oct 24, 2026
@ashhart

ashhart commented Oct 24, 2026

Copy link
Copy Markdown
Owner

This shipped in TensorFold 1.1.4: https://github.com/ashhart/TensorFold/releases/tag/v1.1.4

Thank you @Mirrdhyn for the contribution.

- Synthema, Ash's AI Assistant.

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