Repository navigation
Conversation
…nd the writer that places rows window_rows.cuh's decode partials, merge and prompt tiles take a Pg (a stream's row of page ids, elements a page, positions a page) as a template switch. With the switch off every kernel in window_rows.cu and contract_controls.cu compiles to the same SASS on sm_80/86/89/120/121 as before; the paged head dim 256 kernels are the new image window_pages.cu: - tf_wrows_part_256_pg, tf_wrows_merge_256_pg, tf_wrows_prompt_256_pg: a decode chunk and a prompt tile lie in one page (pages hold whole 256-key chunks), so a page id is read once a chunk or tile. - tf_wrows_put: k and v rows into the pages at their positions, a block a row, 16 bytes a thread. WindowAttention.loadPages loads the image on demand (only Dory asks); Window.pg selects the paged kernels, and the unpaged calls are unchanged. All four are registered with the kernel contract: tf_wrows_put places the caches of each case in a shuffled arena bit for bit, then the attention kernels read them through the table, each output within float64's range, at pages of 256 and 1,024 positions. Two controls (two streams' tables swapped, a stale page id) go red at gate 3. Launch records from nvcc 13.0's sm_80/86/89/120/121 images.
…am leasing, a window set aside for the first At --loops 2 and 256K a stream's 14 full-attention caches take 14 GiB; the 42 sliding rings are a fixed 3,072 keys (504 MiB) whatever the context, so only the full caches page. A page holds 256 positions of every full cache (K then V each), 14 MiB at 2 loops, 24 MiB at 4; the sliding rings stay one buffer a slot, because they do not grow and their writes wrap. - Loaded.layout with engine_paged, as GLM's: the first stream's rows are pages too, and the plan sets one window of them aside, so one stream can always reach the context. Every stream leases its pages from kv_pool.Pool at its prompt pass (prompt, reply and one verify window) and gives them back when it ends; a request that cannot lease waits (LaneMemory), and F.fits is the host's admission. --parallel counts a slot's rings and table row and a page each; the prompt cache gets the rest (kept states are copies of varying size: snap_bytes 0). - One table row a slot, so a decode round mixes streams in one launch. Model.attach makes the arena, the tables and the rings of the slots the plan names; forward refuses rows past a slot's pages. Positions are written by tf_wrows_put and read through the table by the paged window kernels; the bits do not change. - Prompt-cache states copy page runs in the contiguous layout, so a kept state moves between the two forms. - TENSORFOLD_DORY_KV_PAGES=0 keeps every stream's caches contiguous and allocated at open (what `tensorfold run` and the checks' contiguous side use). tf-cuda-test dory-pages runs one model contiguous and one paged over four slots on shuffled pages (prompts to 3,100 rows, page edges at 255/256/257, 64 rounds of 1 to 16 rows) and then the lane core over 1/3/8 streams, greedy, sampled and drafted: logits bit-equal, replies equal, the pool idle. DORY_PAGED=1 runs dory-serve and dory-states on pages.
|
Thanks @jschmied. We ran this on a GB10 with the Dory 4B NVFP4 checkpoint against main, in alternating base/new runs. Exact, with one small cost.
Our rule is that a merge doesn't regress anything, however small. Could the full-attention rows keep main's contiguous path wherever paging isn't needed (the first stream's window, or any stream whose rows fit its set-aside window), and go through the page table only past that? That's how the Qwen paging slice ended up flat. Once 2k is level with main we'll land it straight away. - Synthema, Ash's AI Assistant. |
…of a step while one is free findAligned finds the highest such run; claimRun's body moves to claimAt, unchanged. A lease that already holds pages, or finds no free aligned run, falls back to cover, so nothing waits on it.
…s path, the table only past that Review of ashhart#1284: 2k prefill read -0.6 to -1 % against main. The arena is now cache-major (each full cache's K over every page, then V), so page ids p..p+n-1 are rows p*256.. of one buffer, and a lease taken as one run from a multiple of pagesOf(context) is exactly main's slot w of a contiguous cache. Streams lease that way while a run is free (kv_pool.coverRun), the first stream always. - A batch whose every stream sits in such a window runs main's kernels on the window index as the slot: the fused row writes, tf_wrows_prompt_256 and the contiguous decode, no tf_wrows_put, its own graphs (flat_graphs). Any stream on scattered pages sends the whole batch through the table, which is valid for the windowed streams too: one launch set, no split, graphs keyed by rows as before. - Prompt-cache copies coalesce consecutive page ids into one copy. dory-pages maps slots 0, 1 on windows and 2, 3 on shuffled pages and adds rounds of a lone stream and a windowed pair; the lane-core replies run scattered and windowed. DORY_SCATTER=0 runs dory-serve and dory-states on windows.
|
The new paged prompt cases include |
…s a later page At 256 rows and no prefix the table holds one page, so ids[1] read past it: a Debug panic, undefined in ReleaseFast. With one page the control is planted from the next served shape on, where it stays red.
|
Thanks, done in three commits on top (no rebase): 9f00cc6, 9564058, 79ea7b2. What changed. The arena is now cache-major (each full cache's K over every page, then V), so page ids p..p+n-1 are one contiguous run of rows. A stream leases its pages as one run starting at a multiple of pagesOf(context) whenever one is free ( Where the 2k cost was. nsys on the PR as it was, 2k prefill, L=2, per pass: paged attention 48.1 ms vs contiguous 47.95 ms. Served, alternating base/new, 3+3 (dory-nvfp4 L=2, --parallel 8, --context 32768; prefill = median tok/s of 8 cold prompts; decode = aggregate tok/s, 256 tokens, temp 0):
Two of the three new 2k runs fall inside main's range (4820–4835). The third (4793) is low, and its 8 samples run 4751–4851. Main's own 2k runs on this box today span 4726–4835. Gates (L=1/2/4, all pass):
@Maxrio9 thanks, with #1316 fixed on main (3cb8cb5) the new zero-prefix cases now run; results are in the merge line above. One more Debug-only crash, on main and not in this PR: |
…s path, the table only past that Review of #1284: 2k prefill read -0.6 to -1 % against main. The arena is now cache-major (each full cache's K over every page, then V), so page ids p..p+n-1 are rows p*256.. of one buffer, and a lease taken as one run from a multiple of pagesOf(context) is exactly main's slot w of a contiguous cache. Streams lease that way while a run is free (kv_pool.coverRun), the first stream always. - A batch whose every stream sits in such a window runs main's kernels on the window index as the slot: the fused row writes, tf_wrows_prompt_256 and the contiguous decode, no tf_wrows_put, its own graphs (flat_graphs). Any stream on scattered pages sends the whole batch through the table, which is valid for the windowed streams too: one launch set, no split, graphs keyed by rows as before. - Prompt-cache copies coalesce consecutive page ids into one copy. dory-pages maps slots 0, 1 on windows and 2, 3 on shuffled pages and adds rounds of a lone stream and a windowed pair; the lane-core replies run scattered and windowed. DORY_SCATTER=0 runs dory-serve and dory-states on windows.
… caches stay the default (#1284) Paging changed Dory's default open, memory plan and admission; new behaviour is off by default until device checks qualify it. Unset or 0 keeps main's contiguous caches; any other value takes the pool's pages. Mirrors the device tests' DORY_PAGED switch.
|
Landed on main, thanks @jschmied, with two fix-ups on top: paging is opt-in with |
…differently, five-SM RECORD lines tf-cuda-test kernel-contract-limits on nvcc 13.3's fatbins for sm_80/86/89/120/121 found 539 stale records on 31a864f (543 with #1284's window_pages launches): 206 launches whose register counts (and a few static shared sizes) differ from the 13.2 or 13.0 records they were checked against. The 156 ordinary ones go into a 13.3 set (nvcc_133_records.zig), the 50 pending ones into a 13.3 receipts file the pending lookup reads (measurements/pending_133.zig), as qwen_lanes_133.zig's decode receipts do; the tests that pinned those launches as unmeasured under 13.3 now expect the receipts. The tool's RECORD line always carries five entries: an SM a launch does not apply to keeps the record the compiler is checked against, so the line pastes as an Override (it held one or two entries for cluster and GB10-only launches).
…refuses a context off the page size, window_pages only with paging (ashhart#1284)
…refuses a context off the page size, window_pages only with paging (#1284)
Follows #1233 (Dory, landed in #1256-#1261): the full-attention caches move into pages of the core KV pool. Two commits on main
2ff019107:b5eadd6a1(thewindow_pageskernels, contract registrations and records) and09852121d(Dory on pages, with tests).What this changes
At 2 loops a Dory stream's 14 full-attention caches take 56 KiB a token (14 GiB at 256K; 24 GiB at 4 loops).
Only these page. The 42 sliding rings are a fixed 3,072 keys (504 MiB a stream), whatever the context, and their writes
wrap, so they stay one buffer a slot and count as the layout's fixed bytes.
window_rows.cu's decode chunk, so a chunk and aprompt tile never straddle a page.
Loaded.layoutwithengine_paged(as GLM's): every stream, the first included, leases its pages fromkv_pool.Poolat itsprompt pass (prompt, reply, one verify window) and returns them when it ends. The plan sets one window of pages aside, so one
stream can always reach the context. A request that cannot lease waits (
LaneMemory);fitsis the host's admission.--parallelcounts a slot's rings and table row plus a page each; the prompt cache gets the rest.
window_rows.cuhtakes aPgtemplate switch; the head dim 256 decode partials, merge and prompt tiles are new kernels ina new image
window_pages.cu(tf_wrows_part_256_pg,_merge_256_pg,_prompt_256_pg), andtf_wrows_putwrites k/v rows intothe pages. Only addresses change; no sum changes order. With the switch off, every kernel in
window_rows.cuandcontract_controls.cuhas identical SASS on sm_80/86/89/120/121.WindowAttention.loadPagesloads the image on demand (Dory only).tensorfold-native;TENSORFOLD_DORY_KV_PAGES=0keeps every stream contiguous and allocated at open.tensorfold runstays contiguous.Receipt
Environment: one DGX Spark (GB10, sm_121, CUDA 13.0, nvcc 13.0.88),
Dory-1.0-NVFP4(ModelOpt), greedy unless noted.tf-cuda-test dory-pages(new), L = 1, 2, 4: a contiguous and a paged model over 4 slots on shuffled pages (prompts 255..3,100 rows,page edges at 255/256/257, 64 rounds of 1-16 rows): every row's fp32 logits bit-equal. Then the lane core over 1/3/8 streams,
greedy, sampled and drafted: replies equal, pool idle.
On pages (
DORY_PAGED=1), L = 1, 2, 4:dory-serve: eager == graphs, drafted == plain, concurrent == solo at 2/4/8 (with and withoutdrafts).
dory-states: resumed == fresh at 8,000, 8,192 and 32,001 into caches wiped to 0xFF, greedy and sampled.Served (
tensorfold-native, 2 loops, 24 GB budget, 128K window,--parallel 4): contiguous refuses ("needs 30.63 GiB ... leaves 21.02GiB: 2 streams fit"); paged starts (3 more taking 1,360 pages, 18.59 GiB, enough for one 131,072-token window). An 85k-token prompt
with three of 15-19k concurrently, and four 50-58k prompts (which cannot all fit; they wait for pages), all equal their solo replies.
Chat resume: 18,503 of 18,943 prompt tokens cached, reply equal to a fresh server's (0.40 s against 4.7 s).
Decode (
dory-bench, 2 loops, no drafts, tok/s together, 3 ABAB rounds):Means 43.7 vs 43.7 and 326.5 vs 324.2 (-0.7 %); contiguous alone moves 0.5 %.
Prefill (
dory-prefill, 2 loops, prompt time alone, best of three passes per run, 3 alternating runs, ms):Within +-0.3 % at every length (16k is 0.8 % faster paged), inside the contiguous arm's own spread; the extra spill (56-byte
stack) costs nothing measurable.
Kernel contract: the four new launches pass gates 1-4 at pages of 256 and 1,024 positions (decode and prompt).
tf_wrows_putplacesthe caches in a shuffled arena bit for bit before the attention kernels read them. Two controls (two streams' tables swapped, a stale
page id) go red at gate 3. Full run on this branch: 340 kernels, 338 pass, 91/91 controls red; the two failures are w4a8's
planted-non-finite cases on
main, and the 143 limit problems are the launchesmainkeeps pending (none is a window_pages launch).Launch records are from nvcc 13.0.
Tests:
zig build test -Dnvcc=/usr/local/cuda/bin/nvccfrom the repo root exit 0;zig fmt --checkclean; lean check 0 new.Gotchas
(table above), so no split was made.
Checklist