Repository navigation
flashnext cuda: pending DeltaNet rows, reused lane buffers and the shared expert on a side stream - #1763
Conversation
|
Thanks @Arminova. We're starting our own device checks while your receipt finishes: the kernel contract (your three new kernels against the ones they replace) on one GB10, and served cells on |
c065ddf to
7b37e1d
Compare
|
Kernel contract on this change (GB10, nvcc 13.3.73; run at c065ddf, before the rebase onto
The limit problems are |
|
Our device checks for this PR are done, @Arminova, and they're clean:
When your own receipt is in and you mark it ready, we'll gate it and land it. TP2 stays untested on our side, as in #1762. |
7b37e1d to
890c4e6
Compare
|
Thanks for running it. Our receipt is in the description: it ran on |
…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).
890c4e6 to
ac213a3
Compare
|
Thank you @Arminova. 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 @Arminova! This is in TensorFold 1.2.2, out now: https://github.com/ashhart/TensorFold/releases/tag/v1.2.2 |
What this changes
Three changes to a shared pass's rounds on one GPU. Each keeps the bits of
mainand has its own switch:TF_FLASHNEXT_KEEP_LATE=0replays as before).TF_FLASHNEXT_LANE_PARTS=0).TF_FLASHNEXT_SHARED_SIDE=0).Three new kernels (
fn_gdn_front_log,fn_gdn_rows_late, andfn_late_rollfor a rollback's windows in one launch), each registered with the kernel contract as a bit-equal replacement of the kernels it stands for; the set joins the family sets, socontract_tests.zigis unchanged. Off or refused with tensor parallelism.We had also batched every stream's DeltaNet steps into one launch each, pending rows included. We took that part out of this PR because #1740 does the same and was agreed first; we're glad to compare notes on it or review #1740.
Part of #1762.
Receipt
Environment: TensorFold
main94b90c9dband this change (one commit); DGX Spark (GB10, 128 GB), Ubuntu 24.04 aarch64, 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. Since then rebased ontomain20705d604; the rebase moved only the contract set's registration.Exactness:
bench_concurrent --levels 1,4 --alone --serial: every reply equals its solo run and"draft": false. Every reply'stoken_sha256equalsmain's at 1, 4 and 16 streams, greedy and seeded sampled, code and chat, at default settings and withTF_FLASHNEXT_CHUNK=8192 TF_FLASHNEXT_DRAFT_HEAD=affine4;bench_cacheon cold 512 / 2,048 / 16,384-token prompts (and 8K / 32K / 128K at the second setting) givesmain's replies, and repeated (resumed) and concurrent requests equal the first.Decode speed at default settings,
main→ this change, tok/s, mean of two alternated rounds (main, this change, flashnext cuda: prompt passes with the same bits and less work, and a prompt scratch sized by use #1764, flashnext cuda: prompt passes with the same bits and less work, and a prompt scratch sized by use #1764, this change,main):bench_openaicode / chat (256 tokens, 3 reps)bench_concurrentcode x1 / x4bench_concurrentchat x1 / x4--context 32768 --parallel 16) code / chatWith
TF_FLASHNEXT_CHUNK=8192 TF_FLASHNEXT_DRAFT_HEAD=affine4(one round):bench_concurrentcode x1 / x4 74.6 → 75.8 / 131.1 → 132.1, chat 45.8 → 46.8 / 83.4 → 85.7.Prompt speed: this change does not touch prompt passes; cold prefill at 2,048 / 16,384 tokens 1,569 → 1,590 / 1,677 → 1,679 tok/s (
bench_cache --require-cold).Memory: the prompt cache gets 7,233 → 6,583 MiB at default settings (the pending-row logs, 54.6 MiB a stream, and the reused lane buffers), 1,006 → 1,009 MiB with 8,192-row chunks.
Tests:
zig build test445 / 445 steps, 0 failures (2,431 tests run, 14 skipped; cached steps not counted); after the rebasetest-cuda-host274 / 274 steps, 3,911 passed, 8 skipped.Kernel contract (nvcc 13.3.73, GB10):
kernel-contract-limitson this change: none of its records stale, no launch over a limit. The full run'sSUMMARYline follows in a comment.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. (The CUDASUMMARYfollows in a comment.)