Repository navigation
flashnext cuda: prompt cache marks inside a chunk are kept on the way instead of cutting the pass - #1866
flashnext cuda: prompt cache marks inside a chunk are kept on the way instead of cutting the pass#1866Arminova wants to merge 21 commits into
Conversation
…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).
… prompt scratch sized by use Seven changes to prompt passes. Each keeps the bits of main: - BF16 projections' prompt passes of 1,024 rows or more run on prompt16's 128 x 128 tile with the lane's sums (TF_FLASHNEXT_PROMPT_ROWS sets the row count; TF_FLASHNEXT_PROMPT_BF16=0: the lane). - A layer's writeback and its read-out's normalize run in one launch (TF_FLASHNEXT_HC_FUSE=0: two launches). - A native projection asked for K slices returns its reduced slice and one zero slice (the old sums' bits, -0 to +0); the other slices are never zeroed. - tf_route_logits_wide walks the expert blocks first, so a row tile's input rows are read from DRAM once (its contract test launches the new grid). - fn_qsa_select keeps 8 keys in flight a pass and places its picks by warp runs: the same ids, ties included. - The fused prompt DeltaNet kernel runs its batch's gated norm on the chain warps. - The prompt scratch sizes head_local, native_part, hc_part and part by what reads them; the prompt cache gets the rest. Five new kernels (fn_hc_writeback_norm0 / _moe / _slices, tf_prompt_bf16 / _f32): the first three registered with the kernel contract as bit-equal replacements of the writeback and normalize launches (contract_flashnext_wbnorm.zig), the BF16 tile under the projection contract's FP64 reference (contract_projection_f32.zig). The default chunk stays 1,536. Off with tensor parallelism where noted in the code. The read-out helpers moved to cuda_batch_hc.zig and the scratch sizes to cuda_scratch_sizes.zig.
…ent, deeper greedy drafts, CUDA graphs Seven changes to drafted decoding. Replies keep the bits of main (drafts change, accepted tokens do not): - A drafted stream's first token is sent before its first drafts (TF_FLASHNEXT_FIRST_EARLY=0: after). - The n-gram table is read ahead into the page cache and loaded weight shards are dropped from it (TF_FLASHNEXT_TABLE_WARM=0, TF_FLASHNEXT_DROP_LOADED=0). - The draft chain stops on the running product of its probabilities; the core gets the chain's chances. - Greedy streams draft up to 8 deep; sampled streams keep 6. - Sampled draws read a row's top_k candidates from the GPU instead of the whole row (fn_row_topk, its candidates in column order, registered with the kernel contract against a host oracle; TF_FLASHNEXT_GPU_TOPK=0: the whole row). - A drafted verify window leaves its scores on the device (TF_FLASHNEXT_VERIFY_HOST_LOGITS=1: copied). - Drafted windows, draft-head pieces and rollbacks run as CUDA graphs on one GPU (TF_FLASHNEXT_WINDOW_GRAPH=0: plain launches). Off with tensor parallelism. New files cuda_mtp_draw, cuda_mtp_pieces, cuda_window_graph.
… instead of cutting the pass A prompt-cache mark inside a chunk cut the pass there, so a prompt ran as two passes, the second often a few rows. Now the states a mark keeps are taken inside the pass: the DeltaNet state from the fused kernel at the mark's row (a pass holding a mark runs the fused kernel whatever its row count; both kernels give the same bits), the conv and PLE tails copied after. Replies keep the bits of main, resumed equals fresh. TF_FLASHNEXT_KEEP_INPASS=0 cuts the pass as before. Off with tensor parallelism. New file cuda_batch_marks.zig.
…#1763 and ashhart#1868) Loaded shards now leave the page cache through main's rule (ashhart#1868), which keeps the shards that hold the n-gram table; this branch's own release and its TF_FLASHNEXT_DROP_LOADED switch are dropped. The n-gram table warm-up stays.
…so cuda_batch.zig stays within 600 lines The attention and MoE blocks' writeback and read-out (fused where cuda_batch_hc.fuses allows, else the two launches) move from cuda_batch.zig into cuda_batch_hc.attnReadout and mlpReadout: the same calls in the same order.
…attention runs in one launch With TF_FLASHNEXT_ATTN_LANES=1 (ashhart#1739) a shared round may take the one-launch attention or the per-lane one; the graph captured for one must not replay for the other, so SharedKey carries cuda_batch_lanes.useAttn's answer.
|
@Arminova, the FP4 check is in: |
…ASHNEXT_TABLE_WARM=1 With the table out of the page cache, the read-ahead it queued right after load delayed the first prompts' own table reads. Unset, 0 and other values now leave it off, as with TF_FLASHNEXT_WARM_AHEAD; a host test checks the default.
…rm-up off by default)
89a6600) Conflicts: cuda_batch.zig / cuda_batch_hc.zig take main's ashhart#1764 (its 600-line trim 9693706 instead of our move); cuda_lanes.zig keeps first_before_draft with main's comptime ranked vtable. b5c96c5 (the MTP head's experts on the call's last row only) moves into cuda_mtp_pieces.zig's Piece, which this branch runs forwardHead's pieces through: a piece knows whether it is the call's last (`last`, also in the graph's HeadKey) and only that one runs the experts, the write-back and the copy out, on row 0. 89a6600's FP8 table rows: the opt-in warm-up reads their codes.
…st in place A graph's key held the cache's geometry, which moves every 256 tokens past 2,048, and the sequence's number, new with every request, so each of them meant a new capture, and 80 % of a capture is the instantiation. The cache now holds a graph under its key without those two (slotOf); when only they differ, the pass is recorded again and cuGraphExecUpdate gives the held executable its values. A pass of other launches still gets a new executable. Replays need the whole key, as before. The host test counts recordings, instantiations, updates and launches on a driver of fakes; with the slot off it fails at the first new length.
|
Thanks @Arminova. Here's our GB10 A/B of Exact: 0 of 20 prompt rows differ, every concurrent stream equals its run alone, and both 64K long runs pass on both arms. Decode gains are large and steady:
After a prompt, decode is +6.7% at 2K and +2.1 to +2.6% at 32K. One cell loses, steadily: decode right after a 16K prompt.
8K and 64K are slightly down too (-0.7% and -0.8%), with the runs overlapping. This matches the 16K first-phase drop we saw on #1786 (-5.0%), so it may be the same mechanism: something about how a chunk's kept marks or the drafted round's state sit at that prompt length. Could you look at what changes right after a 16K prefill? For example: how many kept marks the last chunk leaves, whether the first drafted rounds re-read them, or a buffer sized to a 16K boundary. Once 16K is level, this lands on the gains above. |
…ndow graphs updated in place) ashhart#1786's head now merges main (ashhart#1764 as landed, b5c96c5, 89a6600) and keeps a held window graph across lengths and requests. cuda_lanes.zig would pass 600 lines with main's two: step()'s reuse hook call is one line.
…taNet chain without the per-row mark check The two cut rows in fn_gdn_fused's chain changed the whole kernel's code: 0.8 % more GPU time on every prompt pass of 64 rows or more (Nsight, cold 16K prompt: 470 vs 467 ms over 360 launches). The chain is now a template on whether it keeps marks; the kernel takes the copy without the check when both cut rows are 0, the one with it otherwise (464 ms). Same arguments and bits. The contract keeps two marks on even row counts and none on odd ones, so both copies run.
…e, a request or a length updates it in place A held graph's slot kept the sequence's buffers, so a request on another lane ran its windows eagerly, then captured them, on every cell's first request; a length crossing into the score grid (one more launch) failed the in-place update and instantiated a new executable, several times a request. Each key is now a shape (rows, outputs given or not, score grid on or off) and a hash of everything else (buffers, lengths, the sequence's number); the shape picks the held graph, the hash says whether it replays or is updated. The several-stream keys take the same path.
|
Thanks @Arminova. Here's our GB10 A/B of this head (
We can't land a loss on the long-prompt rows, even with the one-stream gain. The 32K and 64K decode drop looks like a per-round cost that grows with the kept context. Could you check what the fused DeltaNet chain without the per-row mark check does once a long prompt's state is kept? The 8-stream greedy code row may be the same cost at higher batch. Happy to rerun as soon as you push. |
…lashnext/prompt-cache
Conflicts only where ashhart#1786's squash met its own earlier commits: firstEarly now lives in cuda_mtp.zig (main's move); the load cache finishes before the n-gram warm-up (main's); this branch's in-pass marks (Hooks.marks, Mark) and Kernels.profile() kept.
|
@ashhart, this PR's own numbers on What it gains: the first token.
Your MLX rows above:
Your 16K question: decode right after a new prompt equals
Your matrix on The branch still merges cleanly with today's |
|
@Arminova, our rerun on
So your merge settled it, and our text agrees with yours. Thanks for chasing that down. |
What this changes
A prompt-cache mark inside a chunk cuts today's prompt pass there, so a prompt runs as two passes and the second is often a few rows (on a short chat prompt, a second pass over the rows after the history mark). This change keeps the mark's state inside the one pass instead:
fn_gdn_fusedtakes up to two marks and writes the state after each mark's row into the kept state's buffer). A pass holding a mark runs the fused kernel whatever its row count; the fused and the unfused kernels give the same bits.TF_FLASHNEXT_KEEP_INPASS=0cuts the pass at the mark as before. Off with tensor parallelism and for image prompts.Replies keep the bits of
main; a resumed prompt equals a fresh one. New filecuda_batch_marks.zig. The kernel contract'sfn_gdn_fusedentry now launches with two marks (a pass's middle row and its last) and checks each kept state againstfn_gdn_rows_toover that many rows.Depends on #1786 (#1763 is in
main). #1786's branch, withmain088cfbd35merged in, is merged into this one (merge commits on top, nothing rewritten). This PR's own change is commitfd3c0cc59. 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-Localatb9ce18a6,--context 262144 --parallel 3, bf16 KV, prompt cache at its default size, a fresh server for every run), onmainaf9d7cde7. Base: flashnext cuda: drafted rounds: first token early, n-gram table resident, deeper greedy drafts, CUDA graphs #1786's headbd11f9ab1; this changefd3c0cc59. Mean of two alternated rounds (base, this change, this change, base) unless noted.Exactness: every reply's
token_sha256equals the base's at 1, 4 and 16 streams, greedy and seeded sampled, code and chat, on the first-token prompts (new, resumed, drafted and"draft": false) and on every cold prompt below. A 32K prompt repeated (resumed from 32,767 cached tokens) and sent concurrently equals the first, on the base and on this change alike.bench_concurrent --alone --serialat 1 and 4 streams (3 reps) and at 16 streams (--context 32768 --parallel 16): every reply equals its solo and serial run. Drafted replies equal"draft": falseon four prompts (300 to 15K tokens).First token, ms (lower is better),
ttft2-style short chat prompt (~90 tokens), one stream: a new prompt 228.1 → 185.7 (-18.6 %); the same prompt with drafts off 219.5 → 174.8 (-20.4 %); the same prompt resumed from the cache 64.2 → 64.5 (level).Decode speed at default settings, base → this change, tok/s:
bench_openaicode / chat (256 tokens, 3 reps)bench_concurrentcode x1 / x4bench_concurrentchat x1 / x4--context 32768 --parallel 16) code / chatThe 16-stream chat cell was lower in both rounds (303.1 / 303.0 → 296.6 / 299.5); the other 16-stream cells overlap between rounds.
Prompt speed: cold 512 / 2,048 / 16,384-token prompts (
bench_cache --require-cold, prompt cache on) first token 368 → 363 / 1,151 → 1,153 / 8,991 → 9,006 ms. With--prompt-cache-gib 0(0 cached tokens, one warm-up prompt per server, exact token arrays:main'sdocs/**/*.mdtokenized by the server), 2,048 / 8,192 / 32,768 / 65,536 tokens 1,788 → 1,790 (+0.1 %) / 1,872 → 1,870 (-0.1 %) / 1,835 → 1,835 (0 %) / 1,798 → 1,797 (-0.1 %) tok/s: without the cache there are no marks, and plain prompt passes are unchanged.Memory: unchanged. Both builds hold the same fixed use within the log's 0.1 GiB rounding; a mark's kept state is the same buffer the cut pass kept before, and the prompt cache's size follows the host's free memory at start on both.
Tests on the head:
zig build test456 / 456 steps andtest-cuda-host276 / 276 steps, 0 failures.tools/zig/lean_check.py: 424 problems, the same asmainaf9d7cde7(0 new).Kernel contract (nvcc 13.3.73, GB10), on this change: every gate line of this change's kernels passes,
fn_gdn_fusedwith its two cache marks included (the kept states bit-equal tofn_gdn_rows_toover each mark's rows, non-finite inputs included; five runs byte-identical; 7 shapes).SUMMARY 536 kernels: 536 pass gates 1-4, 0 fail, 6 known findings; 344 controls: 344 red as planted, 0 missed; 229 limit problemskernel-contract-limitson fatbins for sm_80 / 86 / 89 / 120 / 121: none of this change's records stale (fn_gdn_fusedkeeps its registers with the marks), no launch over a limit; the 274 stale records and the 229 limit problems above aremain's own on this nvcc build.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
"draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.SUMMARYline is in the receipt.