Repository navigation
spark25 CUDA: the full layers' caches in pages of the core pool, every stream leasing; sliding layers keep their rings - #1653
Conversation
|
Thanks @jschmied. We ran this on a GB10 (sm_121) against main, with Spark-X2.6-4B MLX 8-bit, fresh servers per arm and alternating base/new rounds. A Dory run on a second GB10 was also clean. Tokens are equal to main (0 of 16 prompt-test rows differ), and decode and the stream sweep are flat (1, 4 and 8 streams within 0.2%). But cold prefill on the longer prompts is slower, and every run of this PR came in under every run of main:
The loss grows with the prompt's length, so it looks like the paged full-attention prompt path: reading committed keys through the page table on each prompt chunk, in place of the contiguous cache. We don't land a change that slows a default path, so we're parking this for now. A prompt path that reads the stream's own pages contiguously while it fills (or a per-chunk page lookup hoisted out of the inner loop) should close it. Then we'll rerun the same cells. The memory win is real, and we'd like this in once prefill is level. |
|
/atc exchange b5b665b442fe766a80ed676116b7ab74e0bfc7c86b511dd5a5cdc383d3d5132e |
|
Thanks for the numbers. It reproduced here on the GB10 (same model, It was the paged prompt kernel. Under nsys, two 15,000-token prompts spent 1,082 ms in The fix is one commit, 67cd827, in Server A/B after the fix, same cells (min / median / max of the per-server medians, 3 rounds):
Paired by round, the fix against main is -0.4, +1.2 and +1.0 % at 2,048. At 9,216 it is +0.6, +1.5 and +1.2 %, and at 15,000 +0.3, +0.7 and +1.0 %. That leaves only the first round at 2,048 below main. The old head was noisier this session than in the repro (from -1.9 to +0.9 % per round). Tokens: the 8-token greedy replies are equal across all three arms in all 9 cells. Checks on the fix:
The A/B summary is on the review exchange as evidence, along with a |
|
Thanks @jschmied, the prompt fix works. We re-ran your new head on a GB10 (sm_121) against main with Spark-X2.6-4B MLX 8-bit, fresh servers per arm, alternating base, new, new, base.
That's about -0.3%. It's small, but it's in every cell and it's on by default, and our rule is no default slowdowns. So we'll hold this until the 8-stream sweep is level. Our guess is the paged decode reads ( |
|
Thanks. It reproduced here too: -0.6 to -0.75 % at 8 streams in every cell (3 alternating rounds), with tokens equal. It was the paged decode merge, not the reads. nsys's node-level graph trace serialises the graph and hides it. Measured that way,
The fix is 7c0a80b, a new commit. In the paged merge, a non-finite head is replayed in a 8 streams after the fix (aggregate tok/s, per server in round order, ABBA):
Paired rounds 1 and 3 are +0.0 to +0.3 %. Main's round 2 came in low. Greedy replies are equal across arms. Checks:
The A/B summary and a |
… page table (_pg twin), SASS of the others unchanged window_pages.cu instantiates prompt::kernel<256, true, NONE, true> as tf_wrows_prompt_256_p32_pg, the twin Spark2.5's full-attention prompt rows need (head dim 256, P32). loadPages fills Paged.prompt[1] at 256 as it does at 512. The record is measured: nvcc -cubin -O3 per sm 80/86/89/120/121, cuobjdump -res-usage (static without the 1 KB reserve).
…nd 32), with a stale-page control
…y stream leasing; sliding layers keep their rings Spark2.5's full-attention layers (every fourth: 9 of 36) take their K/V from pages of the core pool, pages of 256 positions; the sliding layers keep their rings. TENSORFOLD_SPARK25_KV_PAGES=0 keeps the linear caches. One arena: per full layer a K region, then a V region. Loaded.layout with engine_paged, attach and fits exported as Gemma 4's; a stream's fixed bytes are its sliding rings and table row. Pages are promised at the prompt pass, taken as rows reach them and returned at dropLane; forward refuses rows past a slot's pages (PastPages). The rows kernel writes kr/vr as before and tf_wrows_put copies them into the pages. Reads: decode through the whole table, prompt rows through the slot's row (tf_wrows_prompt_256_p32_pg). Kept states stay copies; snap.each walks a paged layer through the slot's pages. Kernels.load calls loadPages, build gives slot 0 a first page so a forward before attach has its rows, and attach drops the captured decode graphs (graph_gen). tf-cuda-test spark25-pages MODEL TEXT N: linear vs paged at context 8192, two lanes, the pool exactly the two widest requests' pages and shuffled; replies and last-row logits byte-equal over three prompts, a sampled reply, shared rounds and a state kept at 1,000 and resumed. On Spark-X2.6-4B MLX 8-bit (GB10), 64 tokens: 15/15 same, resumed rows 1000 both ways. With TENSORFOLD_SPARK25_KV_PAGES=0 the paged arm opens linear caches and the check says so.
…iles at head dim 256 too (the chunk rule is the decode's)
…ion count at head dim 256 too
… time, the next page's id loaded a page early
… a noinline helper, so its fast path builds no page lookup
7c0a80b to
c337859
Compare
|
Thank you @jschmied. 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 @jschmied! Spark 2.5's full-attention K/V in core pool pages on CUDA is in 1.2.1: https://github.com/ashhart/TensorFold/releases/tag/v1.2.1 |
…and Nemotron's Metal paging The switch registry grows by four rows (ashhart#1660's next group, all measured identical): TENSORFOLD_DORY_KV_PAGES (opt-in, ashhart#1284), TENSORFOLD_GEMMA4_KV_PAGES and TENSORFOLD_SPARK25_KV_PAGES (default-on, exact 0 opts out, ashhart#1619, ashhart#1653), and TF_NEMOTRON_METAL_KV_PAGES, the first metal-backend row (exact 1, ashhart#1537). A new .not_zero syntax covers Dory's on-unless-0 reader; the 0-opt-out rows reuse report_values. The four readers go through the registry's environment-first helper, so an unset variable and every set value behave byte for byte as before.
…and Nemotron's Metal paging The switch registry grows by four rows (ashhart#1660's next group, all measured identical): TENSORFOLD_DORY_KV_PAGES (opt-in, ashhart#1284), TENSORFOLD_GEMMA4_KV_PAGES and TENSORFOLD_SPARK25_KV_PAGES (default-on, exact 0 opts out, ashhart#1619, ashhart#1653), and TF_NEMOTRON_METAL_KV_PAGES, the first metal-backend row (exact 1, ashhart#1537). A new .not_zero syntax covers Dory's on-unless-0 reader; the 0-opt-out rows reuse report_values. The four readers go through the registry's environment-first helper, so an unset variable and every set value behave byte for byte as before.
…and Nemotron's Metal paging The switch registry grows by four rows (ashhart#1660's next group, all measured identical): TENSORFOLD_DORY_KV_PAGES (opt-in, ashhart#1284), TENSORFOLD_GEMMA4_KV_PAGES and TENSORFOLD_SPARK25_KV_PAGES (default-on, exact 0 opts out, ashhart#1619, ashhart#1653), and TF_NEMOTRON_METAL_KV_PAGES, the first metal-backend row (exact 1, ashhart#1537). A new .not_zero syntax covers Dory's on-unless-0 reader; the 0-opt-out rows reuse report_values. The four readers go through the registry's environment-first helper, so an unset variable and every set value behave byte for byte as before.
Builds on #1618's paged kernels. Spark2.5's full-attention layers (every fourth, 9 of 36) take their K/V from pages of the core pool, every stream leasing as it grows; the sliding layers keep their rings.
TENSORFOLD_SPARK25_KV_PAGES=0keeps the linear caches.Loaded.layoutwithengine_paged,attachandfits; a stream's fixed bytes are its sliding rings and its table row. Pages are promised at the prompt pass, taken as rows reach them and returned atdropLane;forwardrefuses rows past a slot's pages (PastPages).kr/vras before, andtf_wrows_putcopies them into the pages.tf_wrows_part_256_pg/merge_256_pg, already on main). Prompt rows go through the slot's row via newtf_wrows_prompt_256_p32_pg, the_pgtwin ofprompt::kernel<256, true, NONE, true>(head dim 256, P32).snap.eachwalks a paged layer through the slot's pages.buildgives slot 0 a first page, so a forward beforeattachhas its rows.attachdrops the captured decode graphs.window_attention.zig: paged prompt rows take pages of whole 32-key tiles at head dim 256 too, andtf_wrows_putplaces rows in pages of any position count at 256. Spark2.5 itself serves pages of 256 and does not need them. The contract's pages-of-32 case does.Checks (GB10, nvcc 13.0, on main f350cd3)
window_pages.cuagainst main: 45 functions over sm_80/86/89/120/121, 0 changed, 5 added (the new kernel), 0 unstable.nvcc -cubin -O3andcuobjdump -res-usage: 251/252/252 registers on sm_80/86/89 and 254 on sm_120/121.kernel-contract-limitsover the 5-SM fatbins shows it current.tf_wrows_prompt_256_p32_pg, 16/4 heads, over pages of 256 and 32, passes repeat, shapes, sums and nonfinite. The control (first page id a later page's) goes red at shapes. As noted on cuda: tree attention and head dim 512's window rows read keys from a page table (_pg kernels), SASS of the others unchanged #1618, the runner's SUMMARY line is not usable: main itself stops atunknown image name: fp8_experts, and the FAIL names before it are the same 220 here and on main.tf-cuda-test spark25-pages MODEL TEXT 64: linear vs paged at context 8192 with two lanes, the pool exactly the two widest requests' pages, shuffled. Replies and last-row logits are byte-equal, 15/15, over three prompts, a sampled reply, shared rounds and a state kept at 1,000 rows and resumed (1,000 rows both ways). WithTENSORFOLD_SPARK25_KV_PAGES=0the paged arm opens linear caches, and the check says so.zig build test test-kernel-contract test-golden test-server-security test-pull-security fuzz-cortheon,git diff --check: pass; lean 0 new.Speed, paged vs
TENSORFOLD_SPARK25_KV_PAGES=0, Spark-X2.6-4B MLX 8-bit,--parallel 4 --context 16384(on main d33dcfc before the rebase; same harness as #1619: fresh server per arm, three alternating rounds, 256 greedy tokens, two prompts per case; min / median / max):Paged and linear agree within 0.6 %, inside the arms' spread (the range is the two prompts).