Repository navigation
perf(qwen3_5 cuda): kept DeltaNet rows replay in the next slot scan instead of a full-state keep (+1 to +8 %, same bits) - #1824
Conversation
|
Thanks @jrb-ai. Folding the kept rows into the scan that loads the state anyway is a neat saving, and +7% to +8% at 16
We'll run it on our GB10s with the NVFP4 27B and the 35B-A3B too, and refresh the nvcc 13.2 records ourselves. |
|
Hey Ash, yep, totally reasonable. Here are the details you requested. 1. A resumed prompt with rows pending. Our earlier prompts were too short to resume. This family keeps a prompt state only at a multiple of 4,096 rows inside the prompt, and the host store plans none for a raw The test now uses two chat turns:
Each turn is sent twice, greedy and then seeded sampled, drafts on, on six servers of this branch:
Each request has one Rows cannot be pending at a served save or restore, though:
So the turns above show that resumes stay exact with the late keep on, but they cannot put rows pending at those points. To cover that directly, With J = 3, 16-row tree windows and 160 tokens of the 3,001-token That check found one gap. 2. A contract case for the replay.
On GB10 it gives 37 PASS lines in about two minutes and ends: It follows 3. The credit. It is code, not only the approach. From fkiene's
That commit changes TensorFold files vendored in that repository. Its LICENSE puts the harness, benchmarks, tests and docs under MIT and says the vendored engine keeps its own licence; THIRD_PARTY_NOTICES names that licence as Apache-2.0. The commit itself is the challenge's validation commit, co-authored by fkiene. Ours: the cache-owned pending buffer (so a tree's path waits, not only a chain prefix), the replay without write-back, the batched loads, Thanks for running it on your GB10s with the NVFP4 27B and Qwen3.6-35B-A3B, and for refreshing the 13.2 records. |
…of a full-state replay keep() copies each kept row's DeltaNet heads (decay, write, key) from the window's log into the cache's own pending buffer in path order, so a tree's path replays as a chain. Each recorded slot scan replays the pending rows into the state it loads with tf_q35k_apply's fp32 steps in the same order, and writes the state back at a one-row window or at 24 pending rows. Prompt passes and prefix or reuse snapshots land the pending rows first (tf_q35k_apply_pend). TF_QWEN_KEEP_LATE=0 replays every keep at once, as before. The replay inside the slot scan is adapted from fkiene's submission to the cuda.fast Qwen3.8-27B challenge (Layr-Labs/cudafast-qwen38-27b-tf-engine d7c540f).
saveCache lands the cache's pending kept rows before it copies, and restoreCache drops the rows pending in the state it replaces, so the restore a prompt pass opens with no longer relies on a cache just reset; the reuse vtable's save and restore lose their flushes. tf-qwen35-cuda's tree mode takes TF_Q35_RESUME_TREE=J: a kept copy after window J, restored at the end into the same cache with rows pending again, and the windows after J run again with equal logits.
Each tf_q35_gdn_slots_* scan replaying pending rows, written back and not, and tf_q35k_apply_pend, bit-equal to tf_q35k_apply then the scan at the 27B's linear widths, two layers and two streams a launch; the last pending row dropped and the last two swapped are its controls. It runs in kernel-contract and alone as tf-cuda-test qwen-pending.
…replay The pending kept rows carry its code (Apache-2.0, as part of the TensorFold engine vendored in that repository).
58d41e2 to
40733f0
Compare
|
Thank you @jrb-ai. 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. |
Thanks Ash! Appreciate your review and support. |
|
@jrb-ai, a heads-up on #1824, which landed this morning. On a GB10, Qwen3.8-27B MLX 4-bit served chats fail on main: 8 of 64 greedy chat requests each time, in two separate runs, at prompts of 524 to 32,780 tokens. Each fails mid-reply with |
|
The Served on a GB10 with
Authored by JRB AI in a private capacity. This contribution does not represent the views or positions of any organisation. |
|
Thanks @jrb-ai! This is in TensorFold 1.2.2, out now: https://github.com/ashhart/TensorFold/releases/tag/v1.2.2 |
What this changes
On CUDA, the Qwen3.5 family keeps a verified window with
tf_q35k_apply, which reads and writes every linear layer's recurrent state once a stream every round. This change leaves the kept rows pending and lets the next round's slot scan, which loads that state anyway, apply them. Every output keeps its bits.TF_QWEN_KEEP_LATE=0keeps today's keep.keep()copies each kept row's DeltaNet heads (decay, write, key) from the window's log into the cache's own pending buffer, in path order, so a tree's path replays as a chain. The conv tails and the attention compaction run as before.tf_q35k_apply's fp32 steps in the same order, and writes the state back only at a one-row window or once 24 rows are pending.saveCache) land the pending rows first with one new kernel,tf_q35k_apply_pend(tf_q35k_apply's loop over the pending buffer). A restore (restoreCache) drops the rows pending in the state it replaces.streamBytes), so pending rows never wait in shared recurrent scratch between rounds.Seggains the pending count and buffer (pn,plog); the four contract tests that mirrorSegfollow.This is the Qwen3.5-family counterpart of #1763, which keeps Flash Next's pending DeltaNet rows in a per-stream log.
Receipt
Environment: TensorFold
main72cc3ea89and this change (one commit; since rebased ontomaina8c02eef, whose 34 commits touch none of its files); DGX Spark (GB10, 128 GB), Ubuntu 24.04 aarch64, kernel 7.0.0-1019-nvidia, driver 580.178.04; both built innvcr.io/nvidia/pytorch:26.07-py3(nvcc 13.3.73) with Zig 0.17.0,-Doptimize=fast;Vontra/Qwen3.8-27B-MLX-4bitat70ae7facwithz-lab/Qwen3.8-27B-DFlash2at50307d4c(default drafter bits);serve --context 16384 --parallel 16 --no-thinking, prompt cache at its default; a fresh server for every run, in the ordermain, this change, this change,main. The receipts marked "after review" ran on this branch's head (a8c02eefand four commits), built the same way.Exactness:
bench_concurrent --levels 1,4,16 --alone --serial, greedy and seeded sampled, code and chat: every reply equals its solo run and"draft": falsein all four runs (1,008 solo and 408 serial comparisons, none unequal). Every completion of every benchmark here (342 a run) has the sametoken_shaonmainand this change. The served receipts ran this change before its contract records and the four test mirrors were added; those leave every CUDA image byte-identical and no engine source changed.bench_cache --require-coldon the 3,001- and 12,001-token prompts ofzig/tests/cuda/flashnext/lr7: repeated and concurrent requests equal the first on both.Resumed prompts (after review): this family keeps a prompt state only at a multiple of 4,096 rows inside a prompt, and the host store plans none for a raw
/v1/completionsprompt, so the 33- to 3,001-token prompts tried first never resumed. A first chat turn of 9,452 tokens and a second of 9,540 (the first, its 64-token reply, a follow-up), each sent twice, greedy and then seeded sampled, ran on six servers: the host store and thePrefixstore (TF_QWEN_KV_PAGES=256 TF_QWEN_KV_SHARE=64), each withTF_QWEN_KEEP_LATEunset,=0, and no prompt cache (--prompt-cache-gib 0). Where a store was on, every send after the first resumed at 8,192 tokens; each request (two turns, greedy and sampled) has onetoken_sha256across all six servers and both of its sends. On thePrefixstore,bench_cache --require-cold --require-reuseon the 12,001-token prompt qualifies with the switch unset and=0(repeated and concurrent requests resume at 8,192), and its replies' hash equals a server's without a store.Rows pending at a kept copy and a restore (after review): served traffic has none at either, because a copy is taken at a mark inside a prompt pass, which lands them first, and a restore starts a request on a cache just reset. So
tf-qwen35-cuda treetakesTF_Q35_RESUME_TREE=J: a kept copy after window J, restored at the end into the same cache with rows pending again, and the windows after J run again. With J = 3, 16-row tree windows and 160 tokens, the copy lands 24 pending rows and the 136 rows after it repeat their logits byte for byte, with the switch unset and=0; all four runs' logits (with and without the copy, both settings) are equal. With this change's twocuda_reuse.ziglines reverted it fails (ResumeMismatch).Contract case (after review):
tf-cuda-test qwen-pending(zig/tests/cuda/contract_qwen_pending.zig, also run bykernel-contract) on GB10. Each slot scan's replay, written back and not, andtf_q35k_apply_pend, bit-equal totf_q35k_applyand then the scan: 30 cases (a decode row, a 4-row chain, two 16-row trees and a 64-row chain at 1, 7, 8, 9, 24 and 48 pending rows, on every scan its slot peak allows) and the six counts oftf_q35k_apply_pend, endingPASS qwen35 pending: controls RED as planted (a dropped pending row, two swapped) on every scan and tf_q35k_apply_pend.Decode speed,
main-> this change, tok/s, mean of the two alternated rounds:bench_openaicode / chat (256 tokens, 3 reps)bench_concurrentcode x1 / x4 / x16bench_concurrentchat x1 / x4 / x16The round gap (p50, chat greedy) at x1 / x4 / x16: 81.5 -> 80.6 / 111.0 -> 107.2 / 171.6 -> 156.2 ms. The two rounds of each binary agree within 0.5 % at x1.
Prompt speed: the cold first request's
ttft_s(bench_cache --require-cold) at 3,001 / 12,001 tokens:main2.27, 2.24 / 9.79, 9.02 s; this change 2.21, 2.23 / 9.01, 9.02 s. A prompt pass does more work only when rows are pending (a later turn in the same stream), and then lands them in one launch.Memory: 1.90 -> 1.96 GiB a stream at a 16,384-token window (the pending buffer); 16 streams admitted either way.
Tests:
zig build test-cuda-hostonmain: 4,174 of 4,186 pass, 10 skipped, 2 failed (bothdrama_voice_toolcleanup tests, directory permissions in this container); on this change the same two fail and nothing else (pending_132_test's fixture now reads the one-slot scan's new 13.0 register count).test-kernel-contractandtest-delta-slots-contract-host: 21 / 21 steps, 126 / 126 tests.zig fmt --checkclean on every changed file.tools/zig/lean_check.py: 424 problems, the same asmain. After review, on this branch's head: 4,120 of 4,132 pass, 10 skipped, and the same twodrama_voice_tooltests fail;zig fmt --checkclean;lean_check.py424, none in a changed file.Kernel contract (GB10, nvcc 13.3.73),
zig build kernel-contracton this change:SUMMARY 527 kernels: 527 pass gates 1-4, 0 fail, 6 known findings; 344 controls: 344 red as planted, 0 missed; 229 limit problems. The four contract tests that mirrorSegtake its two new fields; without them the paged-attention runs lost the context. That run predates the review changes; the new case ran on its own (above), and no kernel has changed since.Kernel limits:
tf_q35k_apply_pendregisters as a pending lane launch besidetf_q35k_apply; the fivetf_q35_gdn_slots_*take a few more registers (64 -> 77 at sm_80 for one slot, for example).kernel-contract-limitson five-SM fatbins (sm_80, 86, 89, 120, 121) with nvcc 13.3.73: every changed launch's record current and the same 274 inherited problems asmain, launch for launch (the records oftf_q35_attn_partial8,_pagedand_legacy, whose registers moved withSeg, refreshed). With nvcc 13.0.88: the new launch's and the slot scans' base records current, and 197 problems on bothmainand this change, the same launches (the base records of those three attention launches were already stale onmain; their counts move by a few registers and are left for a refresh). Not re-derived: the 13.2 records of the slot scans (measurements/qwen_lanes_132.zig), with no 13.2 compiler here;tf_q35k_apply_pendstays pending under 13.2.Not run: Metal and HIP (CUDA-only change; their DeltaNet keeps are untouched); CUDA chips other than GB10 (launch records only); tensor parallelism (one GB10); the family's other checkpoints and MoE members (NVFP4 27B, 35B-A3B), which take the same keep.
Credit
This change carries code, not only the approach, from fkiene's submission to Layr Labs' cuda.fast Qwen3.8-27B challenge (Layr-Labs/cudafast-qwen38-27b-tf-engine
d7c540f, the challenge's validation commit, co-authored by fkiene), which deferred a kept chain's replay into the next scan through a second log half. From it, nearly as written:Seg'spnandplog, the scan'sbeforewrite after the replay and the write-back block at its end, theflushandfoldcall sites in the forward, and the flush before a reuse save; rewritten on the same shape:Cache.pendwithflushandfold, and the scan's replay loop. That commit changes TensorFold files vendored in that repository; its LICENSE puts the harness, benchmarks, tests and docs under MIT and leaves the vendored engine under its own licence, which THIRD_PARTY_NOTICES names as Apache-2.0. This change keeps the rows in the cache's own buffer instead, so a tree's path folds the same way and pending rows outlive the recurrent scratch lease, and adds the replay without write-back, the batched loads,tf_q35k_apply_pend, the prefix-cache landings andTF_QWEN_KEEP_LATE.docs/credits.mdnames it.Checklist
"draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.tf_q35k_apply's fp32 steps in the same order; no sum changes.SUMMARYline is in the receipt. (CUDA only.)contract_qwen_pending.zigneeds this change's kernel and replay;TF_Q35_RESUME_TREEfails with this change's twocuda_reuse.ziglines reverted. Both are GPU runs oftf-cuda-testandtf-qwen35-cuda, which have no skip path.)If this lands, a credit to
@JRB_AIon X would be appreciated: https://x.com/JRB_AIAuthored by JRB AI in a private capacity. This contribution does not represent the views or positions of any organisation.