Repository navigation
lfm2: keep a drafter's tapped rows on every forward - #1390
Open
BlivionIaG wants to merge 15 commits into
Open
BlivionIaG wants to merge 15 commits into
BlivionIaG wants to merge 15 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.
The lane core can put an external drafter in front of a target. The drafter reads the target's hidden rows after some layers (its taps). LFM2-MoE had no taps, so it could draft only with the core's own heads. prepare_features names the layers. From then on, each forward copies each row's output of those layers into the scratch (lfm2_tap, a plain copy of bf16 values). features() gives a stream's rows of the latest verify or one-token round. prefill_features hands each prompt chunk's rows to the drafter before the chunk's end is kept for reuse. The prompt's last row is kept beside the stream's caches for its first draft. A drafting stream resumes only a kept state that also holds its drafter's context (Snap.draft). Without a drafter nothing changes: no taps are kept and every forward runs as before. lfm2_tap registers with the kernel contract (metal_lfm2_tap.zig): the third of five taps' columns at 1, 10, 32 and 2,048 rows must hold the layer's bits and the other columns what the call found. Gate 5 checks its launch against every Apple chip we serve and this GPU. test-lfm2-dspark checks the same copy through the family's own kernels.
This was referenced Oct 26, 2026
ashhart
requested changes
Oct 26, 2026
ashhart
left a comment
ashhart
left a comment
Owner
There was a problem hiding this comment.
Thanks @BlivionIaG. The tap operator has the same contract gap as #1377 (no planted controls or recorded launch limits), and this stacks on #1377.
…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.
…planted controls Merges the shared-round contract records. lfm2_tap's launch is in lfm2_records.zig with the limits --offline compiled, not this GPU's. Two controls edit the source: rows packed at the hidden size, and a column of another tap overwritten. Each goes red at the shapes 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.
Contributor
Author
BlivionIaG
marked this pull request as ready for review
October 26, 2026 21:37
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.
Branch
series/dspark-taps, head3df2d0569. Sits on #1377 (lfm2/plain-decode,70f74ca34). Part of #1374.What this changes
LFM2's Metal backend can keep the rows of named layers for an external drafter (
prepare_features,features,prefill_features), copied by a new kernel,lfm2_tap. A drafting stream resumes only a kept state that also holdsthe drafter's context. Without a drafter nothing changes: no taps are kept and every forward runs as before. This PR
has no user of its own; the drafter PR (
series/dspark-drafter) is the first.lfm2_tapregisters with the kernel contract in its own file,zig/tests/contract/metal_lfm2_tap.zig. Its launch has recorded limits (M1 and M2 from offline compiles) and two planted controls, rows packed at the hidden size and another tap's column overwritten, both red as planted.Receipts at this head
tf-lfm2-exactat the DSpark stack head2e5966f68(it holds this PR's kernels), M5 Ultra (Mac17,15, macOS 27.0 build 26A428):--only lfm2SUMMARY 11 kernels: 11 pass gates 1-4, 0 fail, 0 known findings; 17 controls: 17 red as planted, 0 missed; 0 limit problems.--only tile_qmm6 pass, 0 fail.tf-lfm2-exact(real weights, LFM2.5-24B-A2B-MLX-4bit667408d4): 72 forward/cache checks passed, 0 failed, 0 skipped.zig build test test-lfm2 test-lfm2-dsparkat the stack head (M4 Pro, main8c36359d5merged in): 534 of 534 steps passed;zig fmt --checkclean;tools/zig/lean_check.py428 problems, the same as main.lfm2_tapreadingfrom[i.x == 0 ? 1 : i.x](a wrong source column), thefamily test
lfm2 dspark: the tap kernel copies each row's hidden values into its tap's columns, the rest untouchedfails, and the contract says
FAIL shapes metal lfm2 tap copy: at 1x10240x1: an output differs from the copy's bits,SUMMARY 1 kernels: 0 pass gates 1-4, 1 fail. With the kernel as committed both pass.--drafternothing changes: no taps are kept. On the stack head without the drafter, M5 Ultra, one stream: code 289.8 tok/s, chat 267.9; cold time to first token 0.149 s at 2,048 tokens and 7.60 s at 65,536 (medians of 3).Not measured
Related PRs
Sits on #1377. The drafter PR (#1391) sits on this PR.