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
…s together (on ashhart#1069) Depends on ashhart#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 + ashhart#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.
…raft.zig The lean gate's new findings on this branch: multi-line comments folded to one line within 120 columns, and draftBatched, depthLog and batchHeads moved unchanged to cuda_lanes_draft.zig (as cuda_lanes_verify.zig holds the shared verify), which brings cuda_lanes.zig back to 587 lines; value() is pub for it, as lane() and ops() are. lean_new.py: 0 new.
|
Thanks @jschmied for batching the greedy streams' MTP heads, and for checking the batched rounds and accepted counts against - Synthema, Ash's AI Assistant. |
|
Please reserve and price the batched maximum at load. Growth after admission is exactly what the paged-KV work tries to rule out, and bounding the batch by the start-up reservation would quietly cap batching whenever the chunk is below 256 rows, which is where small-memory setups sit. The test that fails on any growth after admission fits that. Not in the middle of a change here. |
…slice counts Carried onto main: the review's recorded compatibility child, which the PR's own base did not need. Main's lane returns the projection's actual slice count and its single-stream head keeps it: cuda_mtp_ops.zig's forwardHead holds the embedding branch's count (:148) and the hidden branch's (:152), and the attention writeback takes the pass's own out_slices (:159). The batched head dropped all three. Its lane's return was discarded and the static fc_sk used instead, but the lane's count is not the static one: a wide projection folds its slices and returns one (cuda_tiles.zig's matmul returns z, one where the fold holds), and the batched head's rows are the sum over the lanes, so the fold can hold where a single stream's would not. The attention writeback took the static out_sk rather than the pass's out_slices, and the head projection's return was not explicitly discarded, as main's own projection does. The batched head now keeps the embedding count and the hidden count, takes the pass's out_slices for the attention writeback, and discards the head projection's return.
…e admitted capacity Carried onto main: the review's first finding, that the batched head grows the shared scratch after startup priced it. Startup reserves the pass's rows and the block scores from min(context, chunk) (cuda_lanes_load.zig), and the engine's reserve and reserveScores only ever grow (cuda_prompt.zig), so a request past the admission reallocates the shared scratch mid-round. The batched head asked for max(32, 2 * total) pass rows and for piece_rows (256) score rows, with total the sum of the lanes' rows. A 64-row chunk with two one-row heads and four 16-row paths asks for 132 pass rows and 256 score rows against the 64 admitted. The single-stream head asks the same engine for the same scratch but derives its requests from the capacity it was admitted (cuda_mtp_ops.zig's forwardHead: the piece is the capacity's half, and the score rows are priced to it). The batched head now does the same. The plan (cuda_mtp_plan.zig) takes the admitted capacity and the rows to draft: the pass's rows never past the capacity, the score rows are the piece, and the piece is the capacity's half. The lanes then go in groups whose rows fit the piece, each group one pass, so no group's rows are past what the pass holds. The engine reports its admitted rows (cuda_prompt.zig's admittedRows). The host test (test-flashnext-mtp-plan-host) checks the review's shape: the requests stay within the admitted capacity, and the lanes' groups fit the piece. Against the pre-fix arithmetic it fails on the pass rows (132 against the 64 admitted).
Carried onto main: the review's recorded context resolution for the first message's conflict, which the PR's own base did not need. Main's draws takes the request masks (allowed tokens a row may take) beside the sampling settings (cuda_mtp_ops.zig's draws), and the grouped windows' draw forwarded only the sampling settings. The grouped draw now forwards each window's masks (cuda_lanes_verify.zig), so a structured output's rows keep their allowed tokens in a shared pass as they do in a lone one.
Carried onto main: the review's second finding, that nothing tested the batched head against the head it replaces. The shared verify has its own regression (the PR's flashnext_windows_tests), but the head's shared pass had none: forwardHeadMany and draftMany over several streams were never compared with the single-stream forwardHead and draft they must equal. The new test (tf-cuda-test flashnext-mtp-parity) drafts several streams twice: each stream's head alone (forwardHead and draft, one stream a pass) and the same heads in one shared pass (forwardHeadMany and draftMany). It compares the head's logits byte for byte, the proposals, the probabilities and the confidence stop, over two one-row heads, four 16-row paths, and both together, with and without a confidence stop. The planted control breaks one stream's gather: one lane's gathered row lands on another's, which is what a gather that reads the wrong row leaves, and the comparison must go red. The test also checks that the streams' heads differ before the control, so the control cannot pass vacuously. The Linux compile census compiles the checker (compile-flashnext-mtp-parity-host): the test's own file, the runner's entry and the Flash Next CUDA family, for Linux, without images or execution. The test itself runs on a device, which is the lead's cell.
Carried onto main: the stack's three files grew past the engine's 600-line rule, which the review names as the third finding. The helpers move out; the code and the launch arguments do not change. cuda_batch.zig: the head's projection goes to cuda_batch_head.zig, which the head's write-back already lived in, and the MoE's write-back bookkeeping to cuda_batch_moe.zig. cuda_batch_state.zig: the lanes' admission checks and the rollback's shape go to cuda_batch_lanes.zig. cuda_lanes.zig: the checkpoint's device bytes go to cuda_lanes_bytes.zig and the engine's refusals in words to cuda_lanes_explain.zig. The build file's Linux census steps move to cuda_census.zig: the qwen35 rounds checker and the new Flash Next MTP parity checker, compiled for Linux without images or execution.
#1155's first record and main's #1217 both change gdnBlock in zig/src/families/flashnext/cuda_batch.zig: the record keeps each lane's rows on its own conv window and state, and #1217 runs a lone pass's prompt DeltaNet in one kernel. The record carries the shared pass; this commit carries #1217's block onto it. zig/src/families/flashnext/cuda_batch.zig:272-276: a lone pass keeps its projection and conv window before the fused branch, so the one-block path fills the rollback's keep as main does. zig/src/families/flashnext/cuda_batch.zig:278-283: the fused launch for a lone pass of 64 rows or more. zig/src/families/flashnext/cuda_batch.zig:305: the gated norm through gdn_fused.norm, the same kernel and grid as the inline launch it replaces. Refs #1155
|
Landed on main as 5e77d73 and 85770f4 on top of #1069, with five maintainer fix-ups and a carry onto #1217's fused GDN block (through 7faf345), thank you @jschmied. A GB10 parity run passes 6/6: a shared MTP head over 2, 4 and 6 streams (1- and 16-row mixes) equals each stream's own pass in logits, proposals, probabilities and the confidence stop. The served sweep is +35/+75/+116% at 2/4/8 streams with identical tokens, and prefill is flat. - Synthema, Ash's AI Assistant. |
|
This shipped in TensorFold 1.1.4: https://github.com/ashhart/TensorFold/releases/tag/v1.1.4 Thank you @jschmied for the contribution. - Synthema, Ash's AI Assistant. |
Depends on #1069 (drafted windows of several streams verified in one pass; this uses its lanes with row spans). The branch carries #1069's two commits (
3ced86d9,71fc0971) under its own one (5cc536ce), so only the last commit is this change; I'll rebase once #1069 lands.What this changes
The recipe's "Not yet" list says the MTP head drafts one stream at a time. With #1069 the verify windows of a round
share one pass, but the head still ran each stream's absorb pass and chain steps alone: on a GB10 at four streams a
round spent ~29 ms in head passes, against ~5 ms at one stream (per-pass timing: a verify pass is ~42 ms + 2.6 ms a
row, the c=4 chat round 110 ms).
facts()sets.draft_streams, so the lanes core asks a shared round's drafts in one call (as Nemotron's and GLM'sbackends do).
draftBatchedtakes the call's greedy streams (no trees, no media): rollbacks as before, thenOps.draftMany: eachstream's kept rows absorbed in one head pass (
forwardHeadMany, a lane each on its own head cache and position),then one pass a chain step for the streams still drafting, every active stream's best read in one download. Each
stream picks and stops by the same rule as alone (confidence, count, chain limit).
forwardHeadruns on it (row-wise kernels, the per-laneattention of Flash Next CUDA: drafted windows of several streams verified in one pass #1069), so each stream gets exactly the drafts it gets alone.
TF_FLASHNEXT_BATCH_HEADS=0restores it for everystream;
TF_FLASHNEXT_DEPTH_LOG=1prints each stream's chain length by round width every 200 calls.Receipt
Environment: TensorFold
maind1018b31+ Flash Next CUDA: drafted windows of several streams verified in one pass #1069 + this branch; one DGX Spark (GB10, sm_121, CUDA 13.0);nvidia/Qwen3.8-Flash-Next-FP4-Local@b9ce18a6,TF_FLASHNEXT_DRAFT_HEAD=affine4,--parallel 4 --context 40960.The numbers below were taken on a local stack that also carries Flash Next CUDA: NVFP4 experts on prompt tiles with the decode form's sums, bit for bit #1143 and Flash Next CUDA: no zero fills the one-rank ModelOpt pass reads back unchanged #1147 (prompt-side, no effect on decode).
Exactness: with the batched head and with
TF_FLASHNEXT_BATCH_HEADS=0every request's rounds and accepted/draftedcounts are equal (four-stream chat: 100 rounds, 155/256 each; code: 56 rounds, 199/229), i.e. the same drafts;
bench_concurrent.py --alone --serial: concurrent == alone 4/4, drafted =="draft": false1/1.Decode speed (
bench_concurrent.py, 256 tokens, greedy, aggregate tok/s):The four-stream chat round falls from ~110 to ~98 ms.
Chain lengths at four streams (
TF_FLASHNEXT_DEPTH_LOG, length: stream-rounds): 0:3 1:229 2:60 3:157 4:43 5:1006:144.
Kernel contract: no kernel changes.
Tests run:
zig build test --summary all(built-Dnvcc): 211/211 steps, 2,959 of 2,973 tests passed, 14 skipped.Not run, and why: sampled streams (they keep the one-at-a-time path); two ranks.
Gotchas
mainincludes its commits until it merges.Checklist
"draft": false(8/8), concurrent equals solo (8/8), and a resumed two-turn prompt equals a fresh one (7,147 of 7,184 tokens restored), each equal tomain's output (MTP drafts on,--prompt-cache-gib 4; measured at5cc536ce, before the later lean fix, which only folds comments and moves code). Each stream gets the same drafts as alone.TF_FLASHNEXT_BATCH_HEADS=0.