Repository navigation
LFM2: shared-round decode - #1377
BlivionIaG wants to merge 10 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.
|
A small request, @BlivionIaG: this description links more than 20 issues and PRs, which is over Agent Traffic Control's limit, so ATC drops the PR from its board and collision checks (and marks every check incomplete for everyone). Could you trim the related-work list to the issues this PR actually fixes or depends on (20 at most)? Thanks. |
ashhart
left a comment
ashhart
left a comment
There was a problem hiding this comment.
Thanks @BlivionIaG. LFM2 shared-round decode is the right direction. Before device checks: the new operators need canonical planted controls and recorded launch/compiler limits in the kernel contract (the summaries show 0 controls, and limit declarations use current-device checks), and the shared tile dispatch it touches needs a cross-family A/B. Please don't skip M1-M3 launch records: qualify those chips or refuse them explicitly.
…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.
|
Thanks @ashhart. Pushed
|
|
Thanks @BlivionIaG, this is careful work and the numbers are good. Since our last review the five shared-round launches and the greedy pick have records in We merged this head onto current main ( Before landing we will run our own cells on one Mac against main: |
Branch
lfm2/plain-decode, head70f74ca34, with main8c36359d5merged in. Commits by BlivionIaG, no trailers.What this changes
On LFM2.5-24B-A2B (Metal), a round of several streams now runs each layer's short conv and attention once over all
its rows, not once a stream. Plain decode at 4 to 16 streams is faster on the M5 Ultra; every reply is token for token
the same.
lfm2: a verify's greedy rows take the GPU's pick of their logits(471193c).lfm2_argmaxpicks each head row'sfirst largest logit and flags a row with a nonfinite logit. A greedy row without a mask takes that pick; sampled,
masked and flagged rows draw on the host as before.
lfm2: the conv, qk, attention and merge kernels' bodies as inline functions(886d417). No change to the bits.build: test-lfm2 runs the LFM2-MoE family's tests alone(793b68e).lfm2: a shared round runs each layer's conv and attention once over all its streams(807aa4d). At 16 one-rowstreams a forward launched 960 small kernels for the convs and attention.
lfm2_segments.metalruns each of themonce a layer, from a table of the streams' cache addresses, row offsets and key counts (
segments.zig). Each rowcalls the per-stream kernel's inline body. Prompt calls and one-stream windows keep their launches. Every cache
buffer the table reaches is declared to the encoder (
useResource): a new stream's caches join the residency setonly at its next flush. The scratch keeps each stream's attention partials apart (
state.partialCount).tile_qmm: the twin's MMA window only for matrices of 4,096 outputs or more(520d617). Narrower matrices keep thescalar twin's passes. Both kernels give the tiles' bits, and the load check still runs every matrix's twin at every
width.
tile_qmmis used only bylfm2_moe(grep ofzig/srcandzig/tests).docs: the LFM2.5-24B-A2B recipe, its shared rounds and their M5 Ultra receipt(9df2cd9). README saidthere was no native recipe for this model; this adds one.
No arithmetic changes: the new kernels are bit-equal replacements, and the twin's two kernels give the same bits.
kernel contract: LFM2's shared-round launches carry recorded M1/M2 limits and four planted controls(9f74d91).lfm2_attn_segsandlfm2_attn_blockreserve 128 threads in their source. The merges of main are plain merge commits.Receipts at this head
Decode, M5 Ultra (Mac17,15, macOS 27.0 build 26A428)
Main
8c36359d5against this head, LFM2.5-24B-A2B-MLX-4bit667408d4without a drafter,serve --parallel 16 --prompt-cache-gib 0, a fresh server per build and round, three rounds in the order main/this, this/main, main/this,bench_concurrent.py --tokens 256 --reps 3 --temperatures 0 --levels 1,4,8,16, aggregate tok/s, median of 3 reps. No cell dropped; the sampler saw one server and one client at a time and no Zig build.A different prompt on each stream (
tools/bench_prompts.json):The same prompt on every stream (the client's code and chat prompts), change of medians across rounds: code +4.4%, +17.6%, +34.9%, +31.0% and chat +2.1%, +16.5%, +36.3%, +31.9% at 1, 4, 8 and 16 streams. The per-round values are in the raw output (
codeat 16 streams: 633.4, 632.5, 633.4 against 828.2, 829.9, 830.7).token_sha256of every stream and rep is identical to main's on all 12 cells in every round. Alone: 87 equal, 0 unequal, 0 failed (distinct prompts) and 174 equal (code and chat) per build and round; serial 8 of 8 equal.--prompt-cache-gib 0, 3 reps): 2,048 tokens 0.147 s on both builds; 32,768 tokens 3.038 s on main and 3.039 s here. Output hashes equal./usr/bin/time -l,--parallel 16, warm after 1-stream and 4-stream runs): 14.93 GB on both builds, a difference under 1 MB.Decode, M4 Pro
Quiet machine (load under 2, no Zig build, one server at a time), a different prompt on each stream, greedy 256 tokens, 3 alternated rounds, both builds merged with main
8c36359d5:token_sha256equal to main in every cell, concurrent equals alone 12 of 12,"draft": falseequal 4 of 4. Nemotron-3.5-Lightning 4-bit, 3 clean rounds: 1 stream 120.4, 120.0, 120.6 on main and 119.9, 120.0, 120.6 here; 4 streams 153.0, 152.8, 152.5 and 152.4, 156.3, 152.5 (within 1.3%, ranges overlap, tokens equal).tile_qmmis used only bylfm2_moe. The M5 runs of other families (Qwen3.8-27B with and without DFlash2, Qwen3.5-2B, Nemotron) are on the lane-core head in #1382.Kernel contract and exactness (M5 Ultra)
Built on the M4 with
-Dcpu=apple_m1 -Dversion=1.1.4 -Dmetal-std=metal4.0. A fulltf-kernel-contractrun is not posted: it refuses on main and on this branch alike (UnrecordedRotaryCompiler; it needs atf-rotary-sweepreceipt on the machine).The new kernels: "metal lfm2 argmax" (the host draw's pick, ties to the lower id, nonfinite rows flagged), "metal lfm2 stream attention" and "metal lfm2 stream conv" (bit-equal to the per-stream kernels they replace, with float64 attention and a host replay of the conv steps as their checks). The four planted controls: a tie to the highest id, a nonfinite logit not flagged, every stream reading the first stream's state, every stream writing the first stream's cache. Launch limits are in
zig/src/core/contract/lfm2_records.zig, the M1 and M2 ones from--offline.tf-lfm2-exact(real weights): 72 forward/cache checks passed, 0 failed, 0 skipped, on main and on this head.Failures the tests catch
useResourceloop removed): on the M5, 0 of 12, 0 of 24 and 0 of 48concurrent replies equal their solo runs at 4, 8 and 16 distinct streams, in both of 2 rounds; with it, 12 of 12,
24 of 24 and 48 of 48.
test-lfm2fails on that variant ("a shared round's table reaches only cache buffers itdeclares") and passes on this branch.
partialCountwith@min(rows, batch_rows)left untyped (it narrows to u6, andmost * subswraps):test-lfm2fails, "expected 262144, found 65536". Main's sizing, every stream's partials in the same rows: the same failure.
Both pass on this branch.
Tests and checks (M4 Pro, this head)
zig build test test-lfm2: 531 of 531 steps passed.zig fmt --checkonzig/src: clean.tools/zig/lean_check.py: 428 problems, the same as main.Test plan for another machine
What it changes: with several LFM2 streams, conv and attention run once per layer over all streams, and the greedy
pick runs on the GPU. A reply must not change. Speed with 4 or more streams should go up.
Chips: M4 and M5 are run, and an M3 Pro ran the DSpark stack head. No run covers M1 or M2. Their launch limits are recorded in the kernel contract from
offline compiles. A run on either is the most useful receipt.
Who: anyone with an Apple Silicon Mac (M1 to M5) and enough memory for the 4-bit checkpoint (about 15 GB of it
is used while serving).
Needs: macOS with Xcode and its Metal compiler, Zig 0.17.0, Python 3.12, the checkpoint
LiquidAI/LFM2.5-24B-A2B-MLX-4bit(revision667408d494bde191c2cb9dd4f4549827816b9275).Rules: one model process at a time. Close other heavy apps. Do not build while a speed cell runs.
1. Two builds
mainis the before.pris the after. Two folders, so you can switch without rebuilding.Write down:
git -C tf-main rev-parse HEAD,git -C tf-pr rev-parse HEAD, your chip, memory,sw_vers.2. Kernel contract (does this chip run the new kernels?)
Send the
SUMMARYlines. If a chip fails or refuses here, say so: the PR does not support it yet.3. Exactness (replies equal the before)
Both must end with 0 failed. The PR run has more checks (it also checks the GPU pick against the host).
4. Served replies and speed
Alternate the builds: main, pr, pr, main, main, pr (three rounds, A B, B A, A B), so heat and load hit both.
For each cell: start the server in one terminal, run the client in another, stop the server, and wait until
pgrep -fl tensorfold-nativeprints nothing. ReplaceBUILDwithtf-mainortf-pr, andROUNDNwithr1,r2,r3../BUILD/zig-out/native/bin/tensorfold-native serve $MODEL --parallel 16 --port 8080 --context 40960(
--promptsgives each stream a different prompt. A same-prompt run flatters the numbers.)Check:
token_sha256is the same in themainandprfiles.--alonerun. The client reports this.On an M5 Ultra the PR measured about +3% (1 stream) up to +40% (16 streams). On an M4 Pro: about −2% to +10%.
Throw a cell away if another model process or a Zig build ran during it. Say which.
5. Prompt speed (must not change)
In process, no server. Replace
BUILDas above:The prompt rate must be equal within noise.
6. Memory
For one
mainand oneprserver run of section 4, start the server with/usr/bin/time -lin front. After Ctrl-C,send the
maximum resident set sizeandpeak memory footprintlines. They must be equal within noise.7. Send back
One message with: both commit hashes, chip, memory, macOS, the
SUMMARYlines, the end of eachtf-lfm2-exactrun, any hash that differs, the per-round tok/s table, the
tf-lfm2-benchlines, the memory lines, and the cellsyou threw away with the reason.
Not covered here: other model families (a separate A/B), CUDA and HIP.
Not measured
Related PRs
The DSpark drafter PRs (#1390 to #1394) sit on this PR. They add the same
docs/recipes/lfm2.5-24b-a2b.mdand touchforward.zig,state.zig,kernels.zig,lfm2.metal,metal_main.zigandmetal_targets.zig; they rebase on this.Checklist
"draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.Issue
CONTRIBUTING asks for an issue before a new family, backend or arithmetic change. This is none of them: the kernels
are bit-equal and no reduction order changes.