Repository navigation
Conversation
… a tile ahead fn_attn_chunks_ahead_paged formed a tile's key addresses at the start of that tile: the key id, then the page id it names, then a block barrier before the copies could start, so every 32-key tile waited on two dependent loads. The key and page id of the next tile now load while the current one computes (the key after it two tiles ahead), and the address is formed from them when the tile is staged. The addresses are the same, so the sums and the bits are too. The contiguous kernels and the one-buffer and fused paged kernels keep their code.
…s are first used The paged pool allocated its whole arena at load as one cuMemAlloc. On a GB10 that costs decode twice. Past about 5.4 GiB in one physical allocation the paged decode attention runs 1.6-1.8x the time of the contiguous kernel while executing the same instructions with the same L2 misses; at today's plan (both paged streams' whole windows, 2146 pages of 256 rows, 15.75 GiB at 262,144 x 3) it is 1.46x at 4K and 1.72x at 16K. And the 15.75 GiB leave the host at load, before any stream needs a page: MemAvailable falls from 42.7 to 26.6 GiB and the n-gram table that decode reads every token (26.8 GiB) keeps only 72-96 % of itself in the page cache. The DeviceBuffer the pool holds is now a reserved virtual range (CUDA virtual memory) whose 256 MiB physical pieces are mapped, and zeroed as the arena was, the first time a lane's page table or state run reaches into them; addresses, page ids and kernels stay as they were. A lane's pages are mapped when its table is uploaded at open, so a piece that cannot be mapped fails the open with LaneMemory before the request runs, as a contiguous stream's open fails when its allocation does. The buffer is counted as before; freeing it unmaps every mapped piece and the range. A pool made without a device (Pool.init), an arena of one piece or less, or a device without virtual memory keeps the single cuMemAlloc. The load context's device reaches the pool through Pool.initOn. Decode with every stream on pages (TF_FLASHNEXT_KV_PAGES=256, TF_FLASHNEXT_KV_PAGES_ALL=1) against the contiguous default, GB10, 262,144 x 3, fresh servers in mirrored order, five fresh prompts per cell, 512 greedy tokens, time per round, with the arena mapped whole at load: 4K +2.21 / +2.78 %, 16K +1.55 / +2.25 %; mapped as used: 4K -0.01 / +0.08 %, 16K -0.13 / -0.16 % (the default's own two runs +-0.11 %). The paged attention's 2-row launches take 34.8 us against 50.8 us contiguous at 16K. First-token time level. Every reply identical.
… not the installed one The arena got its reserved range by installing its hook as the process's DeviceBuffer hook for one allocation and putting the old one back, though the hook is meant to be installed once at boot: another engine or thread allocating meanwhile could have taken the range, or found its own hook replaced. DeviceBuffer.allocVia takes the hook for one call instead (alloc is allocVia with the installed one), with the same counting and the owner kept on the buffer, so the arena never touches the installed hook. Tests: the range comes through the arena's hook and the installed one is unchanged; a hook that declines (no memory for its flags) falls back to cuMemAlloc and keeps no range; without a device, or for one piece or less, the arena is one cuMemAlloc.
A piece was marked mapped before its zeroing was issued, so a zeroing that failed left a mapped piece marked as ready that was never zeroed. The flag is now set after the zeroing, and a piece whose zeroing fails is released and left unmarked, as a piece that would not map. A test fails the zeroing.
…ements Comment only.
…ches, at load and between requests With the arena mapped as it is used, the first request mapped about four 256 MiB pieces inside its first-token time (its lane's run, its first pages) and a request that walked into a fresh piece mapped it there too: on a GB10 each piece took 6-9 ms of cuMemCreate and cuMemSetAccess. The pool now maps those pieces before an open asks for them: at load, the lanes' runs (runs are claimed from the top of the arena) and a piece's worth of the pages the free stack hands out next; after each lane closes, again the next piece's worth of free pages, so the mapping happens between requests. The pieces, their zeroing and the addresses are the same, so the bits are too. A piece that cannot map then is mapped, or refused, by the open that needs it, as before. A test checks which pages come next from the free stack: page 0 first, then the ones past those a store holds.
…t, and looks two pieces ahead A lane's close gives its state run back before its pages, so the top of the free stack held the run's pages, already mapped, and the pieces mapped ahead were the ones the next request did not need: the next open claims that run again first and takes the fresh pages below it. With TF_KV_SHARE a keep holds pages past those too, and on a GB10 each request still mapped one 256 MiB piece inside its first-token time. The pages mapped ahead are now the ones the free stack hands an open out after the run it claims, two pieces' worth: a 2K request's pages and a keep's. Same pieces, zeroing and addresses, so the same bits. A test closes a lane with a run and a keep and checks the next pages are the freed and fresh ones, never the run.
…o no reply waits for a piece The core sends a stream's last tokens only after its close, and the close mapped the pieces the next open reaches on the engine thread. With TF_KV_SHARE a keep holds the pages it used, so the free pages past them are fresh and a request's last token waited for a 256 MiB piece to be created, mapped and zeroed. The arena now hands those pieces to one thread with the load context current, which maps them and zeroes each on its own stream before marking it. An ensure that needs a piece the thread has in hand waits for it, one still queued it maps itself, and one the thread failed to map goes back unmarked and maps on the caller as before. The load's map-ahead stays on the caller. Same pieces, zeroing and addresses, so the same bits. Host tests: a posted piece maps off the caller's thread and is waited for, never mapped twice; a piece the thread fails to map is mapped by the caller.
test-gradient-serving pins the bytes of the serving kernels and launch wrappers. This branch changes four of them: flashnext_attn.cu (page ids loaded a tile ahead), memory.zig (DeviceBuffer.allocVia), and cuda_lanes.zig / cuda_lanes_load.zig (the page pool's arena opens with the load's device and io). Each row takes its file's new hash, with the change named beside it.
|
Thanks @Arminova. This is careful work, and the receipt is thorough: 40 of 40 token hashes in each of the four comparisons, concurrent equal to solo, and the slow cells weighed against what two equal builds would show. What I checked:
What needs to change:
Once these are met, we'll run main against the branch on Qwen3.8-Flash-Next NVFP4 on a GB10.
|
…age-id look-ahead leaves this change The look-ahead in fn_attn_chunks_ahead_paged (0fd55aa) is a change of its own; it is undone here, with its two register records and its serving pin back to main's, so this change is the page pool's arena alone. The kernel's code, records and pin row equal main 6c5cef8's.
…where it is one cuMemAlloc, as before The pieces open libcuda.so.1 through std.DynLib, which has no open on Windows, so the Windows build stopped at Arena.alloc. The pieces path is now compiled for Linux only; on other systems the arena is one cuMemAlloc, as it was before this change.
…fails: LaneMemory, the pool idle Pool.open with an arena whose VMM refuses cuMemCreate returns LaneMemory; afterwards no lane is used, the pool's counts equal the ones before the open, and a lane, its pages and a run fit again. Arena.Pieces is public for it.
…thread to map its pieces An open whose run or pages reached pieces not mapped yet mapped them on the engine thread, so every running stream's round waited for it (about 9.5 ms a piece). Now, while other streams run, the first prompt step posts those pieces to the arena's thread and returns without progress; the other streams' rounds go on between steps, and the open comes once the pieces are mapped. Alone, the open maps them itself as before, since nothing waits on it. A piece the thread failed to map is mapped by the open itself, so a failure still ends as LaneMemory, not as a wait.
…ode-ahead # Conflicts: # zig/tests/fixtures/gradient-serving.tsv
Arminova
changed the title
|
Thanks @ashhart. All six points are in; the branch has new commits on top (no force-push):
Main ab78aef is merged in ( Full receipt on that head (
|
|
Thanks @Arminova. Removing the kernel look-ahead, keeping pieces Linux-only, correcting the compiler account, adding the failed-open test, moving shared-open maps off the engine thread and narrowing the CHANGELOG answer all six points. The added arena/pool tests pass. Four new lean findings remain in three comment blocks:
That should bring lean from 345 back to main's 341. Afterward, we'll check the resolved serving pins and schedule the final GB10 memory/concurrent-open run.
|
What this changes
Before: with TF_FLASHNEXT_KV_PAGES the page pool allocates its whole arena at load (15.75 GiB at 262,144 x 3 on a
GB10). After (Linux): the arena is a reserved range whose 256 MiB pieces are mapped and zeroed when a lane first uses
them. A helper thread maps the pieces the next open will reach between requests. When an open reaches pieces that are
not mapped while other streams run, it posts them to that thread and waits in the prompt queue (a step with no
progress); the other streams' rounds go on, and the open comes once the pieces are mapped. Alone, the open maps them
itself, since nothing waits on it. Kernels and replies are unchanged. Without TF_FLASHNEXT_KV_PAGES, and on other
systems, nothing changes (one cuMemAlloc, as before).
(478152b: the kernel, its records and its pin row equal main's), the pieces on Linux only (fee044b), a pool test
(5cffcbb), the open that waits for the helper (273363f), the re-pin of cuda_lanes.zig (f4b5b3b), main ab78aef merged in (b710d1a),
the CHANGELOG line (b35b173, 7c50545). Each commit compiles (compile-config-server) and passes zig fmt on its own.
EXL3 Flash Next: host-resident K/V on the tier (6/8) #1522). Main ab78aef is merged in (b710d1a): one conflict, the cuda_lanes.zig pin row beside core: a bounded ladder of intermediate prompt-cache checkpoints #1953's, resolved to the
merged file's hash with both notes.
Checks on this head
(134 of 134), test-gradient-serving, test-qwen-serving-host, compile-config-server, zig fmt --check build.zig zig,
git diff --check: all exit 0.
x86_64-windows in a scratch copy: before fee044b 18 errors (dynamic_library.zig: unsupported platform, no close,
no lookup, from arena.zig's open), after it and on the head none.
dropping the run's errdefer fails it); the arena's reach with a slow and a failing mapper; the pool's gate with a
50 ms VMM beside a running stream (returns in under 25 ms, the helper maps, the open's own maps after it run nothing
on the caller's thread; a gate that lets the open go early fails it).
Receipt
Environment: GB10 (sm_121), driver 580.173.02, kernels built by nvcc 13.3.73, nvidia/Qwen3.8-Flash-Next-FP4-Local
b9ce18a6, default reserve; main = ab78aef, this head = b710d1a (the CHANGELOG commit on top changes no code). Both
on pages (TF_FLASHNEXT_KV_PAGES=256, TF_FLASHNEXT_KV_PAGES_ALL=1).
Token equality
On pages (TF_FLASHNEXT_KV_PAGES=256, TF_FLASHNEXT_KV_PAGES_ALL=1), the 20 rows (bench_prompts' 16 and four chat
prompts of 283-17,182 tokens), greedy and seeded sampled, 128 tokens, 262,144 x 3: main == this head 40 of 40; this
head's drafted == "draft": false 40 of 40; 4 concurrent streams (code / chat, T 0 / 1): each equals its request alone
and serial. In the test sequence below, replies equal in every row.
An open beside a running stream (Ash's point 4)
Fresh server, pages, one stream decoding 3,072 tokens from a 2K prompt; 5 s in, a second stream sends a 64K prompt,
then on the same server a 120K one. nsys, CUDA API calls by thread:
(Traced on f4b5b3b, the code of this head before main was merged in.) Against main, three fresh servers a side,
alternating (main, head, head, main, main, head), untraced:
The first stream's reply is the same in all twelve runs.
Memory
MemAvailable after load at 262,144 x 3, the n-gram table warmed, four starts a side: main 26.5-27.0 GiB, this head
40.8-41.6 GiB. Of the 15.75 GiB arena, the three lanes' state runs (357 MiB each) and the next two pieces are
mapped at load, about 1.5 GiB; the rest as streams reach it. (The first receipt, on the earlier head with the look-ahead, measured what that buys: the 26.8 GiB
n-gram table stayed 100 % resident against main's 66-83 %, and the first request after start reached its first token 75 ms sooner with the table warmed; with it evicted,
-5 to +29 ms.)
Speed, main on pages against this head on pages
The test sequence (base, new, new, base; a fresh server each): 262,144 x 3 with the prompt cache off and the n-gram
table evicted, then 32,768 x 8 with the cache on and the table warm. No row missing, no slow first request. Every cell
within -1.0 to +1.3 %; in two single-stream decode cells both head runs are below both main runs, by 1.0 % (2K from
10,000: 71.0 / 70.8 against 70.5 / 69.9 tok/s, 14.08 / 14.12 against 14.18 / 14.31 ms a token; 8K: 39.3 / 39.3
against 38.8 / 39.0 tok/s, 25.45 / 25.45 against 25.77 / 25.64 ms), and the same cells in the second pass +0.2 % each.
In the two earlier runs of this sequence on this PR's arena, the same two cells read +0.3 / -0.4 % and -0.4 / -0.2 %,
with main's own two runs 0.5-0.9 tok/s apart in them (more than this run's gap). Decode at 4 and 8 streams +0.7 to +1.3 % in six of eight cells, -0.1 / +0.1 % in two.