Repository navigation
Conversation
…hared pass, each row's bits its lone pass's
…against each alone, byte for byte, dense and sparse, up to 64 rows
This was referenced Oct 21, 2026
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.
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. |
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. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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/verifyFntoo? If so, I'll rebase this onto it rather than collide.verifyFnran 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 pastTF_FLASHNEXT_CHUNK).batch.Lanegains a row span and the lane's own rollback (keep), hidden-row (streams) and score (logits) destinations.rows = 1and none of them is the one-row round as before, which keeps its launches. A lane'skeepmay hold a window's DeltaNet rollback only (rollbackOnly); a training pass's kept tensors are refused, as the round already refusesmoe_add.round.Engine.windows(cuda_prompt_windows.zig) runs the pass.Ops.verifyManyrefuses every window before any stream's state moves, then leaves each stream as itsverifyleaves it (saved tail and PLE rows, host scores row,fed).verifyFngroups 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.zigand two test row types) andzig/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:
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 sametoken_sha256per prompt and level as withTF_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:Tokens per round are unchanged (the drafts and their acceptance are the same).
bench_openai.pyat--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-hostpass; lean check 444 problems (main's count);tf-cuda-testsmoke, 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-contractSUMMARY 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 atd1018b31(no kernel changes here).Not run, and why:
--tp, which the CUDA build does not serve.round's per-row pocket scales.Checklist
"draft": false, a resumed prompt equals a fresh one (prompt passes are untouched), and concurrent equals solo.flashnext-windowsneeds the checkpoint, likeflashnext-window-rank.