Repository navigation
lfm2_moe CUDA: attention pages from the shared KV pool, and the prompt cache (stacked on #1492) - #1503
Draft
jschmied wants to merge 17 commits into
Draft
lfm2_moe CUDA: attention pages from the shared KV pool, and the prompt cache (stacked on #1492)#1503jschmied wants to merge 17 commits into
jschmied wants to merge 17 commits into
Conversation
The host drew every verify row by scanning its 128,000 bf16 logits, after
the forward and with the GPU idle. lfm2_argmax picks each head row's first
largest logit on the GPU and flags a row with a nonfinite logit. verifyFn
takes the pick for greedy rows without a mask; sampled, masked and flagged
rows draw on the host as before.
The kernel is registered with the contract ("metal lfm2 argmax": the host
draw's pick, ties to the lower id, nonfinite rows flagged), and
tf-lfm2-exact checks every window's picks against the host draw.
…ctions lfm2_conv, lfm2_qk, lfm2_attn_mma and lfm2_attn_merge each call an inline body (lfm2_conv_row, lfm2_qk_row, lfm2_attn_rows, lfm2_merge_row) with the same arguments, so another kernel can run a row's arithmetic unchanged. The kernels' bits do not change.
The family's tests (config, draws, state) ran only inside `zig build test`. test-lfm2 runs the engine's test binary filtered to lfm2_moe, so a change to the family can run them without the whole suite.
…ll its streams A recorded call of several streams launched each stream's own short conv (30 layers) and its qk, attention and merge (10 layers): at 16 one-row streams 960 launches a forward for little work each. lfm2_segments.metal runs each of them once a layer over every row, from a table of the streams' cache addresses, row offsets and key counts (segments.zig); each row's arithmetic is the per-stream kernel's, called through the same inline bodies. A stream of more than 8 rows folds a chunk's sub-spans in its 8-row tiles, as its own 16-row tiles do. Prompt calls and one-stream windows keep their launches. Every cache buffer the table reaches is declared to the encoder: a new stream's caches join the server's residency set only at its next flush, so a kernel that reaches them through the table alone may read them before they are resident. The scratch's attention partials now hold each stream's apart (partialCount: a 32-row scratch's grow fourfold; the served scratch of 2,048 rows already holds them). The contract registers "metal lfm2 stream attention" and "metal lfm2 stream conv" as bit-equal replacements of the per-stream kernels, with float64 attention and a host replay of the conv steps. test-lfm2 checks the table, the row map, the tiles and the declared buffers, and the partials' count at 1, 32 and 2,048 rows.
… more From 8 rows the dense twin took its MMA window kernel, a threadgroup for 16 outputs over every row. A 2,048-wide matrix then has 128 threadgroups and runs slower than the scalar twin's 4-row passes; shared rounds of 8 to 32 rows call such matrices at every layer. The 6,144-wide short conv input and the head keep the window. Both kernels give the tiles' bits (the load check runs each matrix's twin at every width); the contract's twin now also serves 16 rows of 6,144 from 2,048.
… receipt README pointed at no native recipe for LFM2.5-24B-A2B. The page gives the checkpoint, the serve command, the shared-round kernels and their tests, and the measured M5 Ultra decode rates before and after them.
…mits and four planted controls The five launches are in lfm2_records.zig, with the limits that --offline compiled, not this GPU's. lfm2_attn_segs reserves its 128 threads in its source, as M1 and M2 cannot build it for another count. Four controls edit the source: a tie to the highest id, a nonfinite logit not flagged, and every stream reading or writing the first stream's conv state or kv cache. Each goes red at its gate.
M1 and M2 cap a pipeline's threads as its registers rise, and no run covers them. The server refuses them by name at open. An M3 is served, with one log line that it is not yet tested with shared rounds.
The shared-round kernels carry recorded M1 and M2 launch limits in the kernel contract, and no run covers those chips. The server no longer refuses them.
The shared-window chunk and merge bodies move to paged_attn_segs.cuh, included twice by nemotron_attention.cu: Nemotron's head-128 kernels keep their names and SASS, and head 64 adds tf_paged_attn_chunk_segs_64 and tf_paged_attn_merge_segs_64. prefill_attention.cu instantiates pattn_kernel<64, 8, 4, 8> over paged caches. sassid on sm_80/86/89/120/121: 0 functions changed, 15 added.
lfm2_moe.cu holds what the shared operators do not: the short conv over each stream's state with every window row's snapshot, the commit of an accepted row's conv state, and the router (fp32 weights, top-4 of sigmoid + bias, ties to the lower id, a NaN ranked as -inf). Each has a planted control. Launch records cover these kernels and the head-64 paged attention on nvcc 13.0, with one 13.2 override.
The MLX 4-bit checkpoint (embedding at 4, 6 or 8 bits) on the shared affine-4 projections and experts, with prompt chunks of 1024 rows, one-row decode and copy-drafted windows over --parallel streams. The ten attention layers keep their keys and values in 512-position pages, and each stream has its own conv state. There is no drafter and no prompt cache yet.
…ount The gated short conv over each stream's state and the conv-state commit move out of lfm2_moe.cuh into liquid_conv.cuh, templated on the tap count, so LFM3 (7 taps) can take the same bodies as LFM2.5 (3). lfm2_moe.cu keeps the router and instantiates the 3-tap pair as tf_liquid_conv_3 and tf_liquid_commit_3; the controls become tf_ctl_liquid_conv_3 and tf_ctl_liquid_commit_3. cuobjdump -sass on sm_80/86/89/120/121, nvcc 13.0 and 13.2: the moved bodies are identical under their new names, every other function unchanged.
…m leasing The kernels already read keys and values through a slot's page table, so the pool changes only what the table holds and how large the arena is; no kernel changes. A page is 512 positions of the ten attention layers' keys and values (10 MiB, 20 KiB a token); the conv states stay a fixed part a slot (240 KiB). - Loaded.layout with engine_paged, as Dory's: the plan sets a window of pages aside for the first stream; each stream leases its prompt, its reply and one verify window at its prompt pass and gives them back when it ends. A request that cannot lease waits (LaneMemory); fits is the host's admission. - Model.attach makes the arena and the tables and conv states of the slots the plan names; forward refuses rows past a slot's pages. The idle server logs the pool's census. - TENSORFOLD_LFM2_KV_PAGES=0 keeps every stream's own window of pages, made at open. Served on the GB10 at --parallel 8, 16k: drafted == plain and concurrent == solo, also with 8 long prompts on a 52-page pool where the last waits for pages; the replies equal the fixed pages' ones.
…the pages' rows through its position A state at `at` is the slot's conv states and its keys and values at positions [0, at) of the ten attention layers: one copy a run of consecutive whole pages, then the last page's rows a layer's keys and values. It is copied behind the chunk that ends at the mark and restored into the new stream's leased pages before its next chunk, so a later turn's pass starts at its history. Kept states are device copies within the plan's cache budget (no snapshot slots), with lookahead 0 as the other headless families. Served on the GB10: eight conversations (turns 2 to 4, prompts 444 to 6,335 tokens, greedy, sampled and drafted) reply alike with the cache, without it, and with fixed pages; 24 of 32 requests resume (cached 440 to 6,243 tokens). Through a 0.1 GiB cache that keeps one state, 6 evictions and 1 hit reply as without a cache.
jschmied
force-pushed
the
lfm2-cuda-pool
branch
from
October 27, 2026 22:45
fcd8014 to
f31a350
Compare
Owner
|
Thanks @jschmied. Leasing from the core pool and keeping states at chat marks is the shape we want, and checking the census after every drain is a nice touch. With #1492's commits on main, both of your commits apply cleanly. One thing must change before this leaves draft. A cold chat turn that keeps a state is 50 ms slower than on #1492 (1.27 s against 1.22 s), and the prompt cache is on by default. Each kept state is a new device allocation ( To your questions:
Then we'll serve it on a GB10:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Part of #1488. Stacked on #1492 (itself on #1377): only the last two commits are new here. Draft until #1492 is ready, then I'll rebase onto main.
What
LFM2.5-24B-A2B on CUDA, step 2: attention pages from the shared CUDA KV pool, and the prompt cache.
21602413aPool pages. The 10 attention layers' K/V live in pages of the corekv_pool(512 positions, 10 MiB a page), Dory's shape. A stream leases pages for prompt + reply + 16 at its prompt pass, waits when the pool is short, and returns them at release.forwardrefuses rows past a slot's pages. The kernels already read through a page table, so no kernel or build file changes, and SASS is unchanged by construction.TENSORFOLD_LFM2_KV_PAGES=0keeps lfm2_moe: serve LFM2.5-24B-A2B on CUDA (plain serve, shared kernels; stacked on #1377) #1492's fixed pages per stream.fcd801437Prompt cache. A kept state is the slot's conv states plus K/V rows[0, at)of the 10 layers, copied run by run with whole pages coalesced. It is restored into the new stream's leased pages. Like the other CUDA families, it keeps at chat history marks and copies its state.6 files, +453 / −31 against #1492.
Verification (GB10, real weights, served, greedy and sampled)
lean_check0 new against lfm2_moe: serve LFM2.5-24B-A2B on CUDA (plain serve, shared kernels; stacked on #1377) #1492's head;zig build testpasses.Speed (GB10, against #1492)
Paging itself costs nothing: with the cache off, the cold turn takes 1.205 s. Keeping a state costs about 50 ms on a cold 10k turn. I haven't pinned that down yet; the candidates are the ~200 MB device allocation per state and the mark's extra chunk.
Questions
TENSORFOLD_DORY_KV_PAGES(dory: full-attention KV in pages of the core pool #1284)?warm_quiet). Should LFM2 get it in this PR or a follow-up?3c5ba9dbeand the two commits after it). The pool's API is unchanged, so I expect a clean rebase.