Repository navigation
LFM2.5-24B-A2B on Metal: drafts with LiquidAI's DSpark drafter (proposal and gate) #1374
Description
What
LiquidAI publishes a DSpark drafter for LFM2.5-24B-A2B: LiquidAI/LFM2.5-24B-A2B-DSpark at revision
e052e5b1dd103c38cdd4f8ddc068d7dea4586f55 (327.7M bf16 parameters, 655 MB). Its weights are under the LFM Open
License v1.0. TensorFold would not ship them. This issue proposes serving it on Metal with --drafter DIR for
lfm2_moe, and agreeing the gates first. The lane core verifies every draft, so a drafted reply equals its
"draft": false reply.
Why
DSpark is a one-stream accelerator on the Macs we have. The numbers below are from the M5 Ultra (three alternated
rounds, the drafter against the same binary without it, no other model process running) at the top of the proposed stack
(2e5966f68). Each PR carries its own receipts on its own commit and replaces them. On the M4 Pro the decode change (#1377) and the correctness checks are measured and in the PRs; the drafter's speed there has no clean measurement.
- One stream, bf16 drafter (the default), medians of 9 replies: code 355.9 tok/s against 289.8 without it (+22.8%), chat
313.1 against 267.9 (+16.9%). With--drafter-bits 8: 370.2 (+27.8%) and 324.3 (+21.1%). On 38 short prompts, drafted
against"draft": falseon the same server: bf16 364.5 against 262.6, 8-bit 389.7 against 262.1; every reply token-equal. - 4, 8 and 16 streams, a different prompt on each, aggregate tok/s against a server without the drafter: bf16 +0.7%,
+0.0%, -1.2%; 8-bit +1.0%, +0.3%, -0.7%. The guards keep most shared rounds plain there (about 97% of rounds at 8
streams in one trace). - Accepted tokens per step on 38 short prompts (greedy, 256 tokens, M5, 4-bit target): code and math 4.3 to 5.5,
multi-turn chat 2.6. Accepted tokens per round are about equal for bf16 and 8-bit (2.827 and 2.859). LiquidAI's model
card reports 5.16 on other datasets, hardware and target precision, so the figures do not compare directly. - Peak footprint (
--parallel 16 --context 131072), warm: 14.67 GiB (bf16) and 14.41 GiB (8-bit) against 13.90 GiB
without the drafter; after one cold 65,536-token request: 16.52, 16.26 and 15.14 GiB. Drafter weights 625 MiB
(bf16) and 361 MiB (8-bit). - Cold prefill, time to first token (median of 3), 2,048 tokens: 0.156 s (bf16), 0.152 s (8-bit), 0.149 s without; 65,536
tokens: 7.84, 7.69 and 7.60 s. - Against mlx-vlm 0.7.10's DSpark path (the same 4-bit target and 40 prompts, one stream, greedy): this drafter 398
to 399 tok/s, mlx-vlm with its drafter 229 to 235, mlx-vlm without it 197. Its replies differ from ours from the
first tokens (0 of 120 equal), so this is not a speed comparison on identical output.
Long context
The drafter's acceptance falls as the context grows, and drafting then decodes slower than plain decode.
- Controlled test (M5, same last 1,536 tokens after a filler of 4k to 130,816 tokens, 4 tails, drafting forced every
round): accepted drafts over proposed drafts are 0.550, 0.529, 0.148 and 0.070 at 4,096, 16,384, 65,536 and 130,816
tokens. Decode against plain decode: -10% at 65,536 and -17% at 130,816. - It is not a defect in our path. The GPU drafts equal the fp32 host reference's at 8,192, 32,768, 65,536 and 130,816
tokens (9 of 9, first-row logit cosine at least 0.9999). mlx-vlm 0.7.10's DSpark, on the same token ids, accepts
1.23 tokens a step at 130,816 against 4.63 at 4,096. - The model card gives no trained context length. We did not test why acceptance falls.
- The lone-stream guard (lanes: a lone stream holds drafts only while its drafted rounds measure faster #1382) re-tested drafting on every new request. A request that follows a losing one in the
same context band now starts plain, and a probe is spaced so that it costs at most 3% of plain time. One stream,
drafter against none, cold 130,816-token prompts: -11.6%, -15.5% and -5.6% before; -9.0%, -0.7% and -0.7% after.
The first request after a server start still loses 5% to 9% at 64k and above. - An experiment, not proposed: a drafter that reads only the last 4,096 keys had 0.419 accepted over proposed at
130,816 tokens. It changes what the released drafter reads.
Work that came out of this and is proposed separately (no drafter needed): plain decode with several LFM2 streams
is much faster when conv and attention run once per layer over all streams (M5 Ultra, a different prompt on each
stream, three alternated rounds: 8 streams 551.6-552.6 to 734.1-734.8 tok/s, 16 streams 661.0-661.7 to
921.5-923.3 tok/s). See #1377.
Design
- The drafter is family code (
families/lfm2_moe/dspark*.zig), likegemma4/dflash.zig. It reads the target's rows
after layers 2, 10, 18, 26 and 34 throughprepare_featuresandfeatures, and uses the target's embedding and
head. It runs the released bf16 matrices as they are. New kernels register with the kernel contract. --drafter-bits 8is an explicit option that quantizes the drafter's linears at load to 8-bit codes. It is named in
--helpand the startup log. Its drafts match the fp32 reference's. 4 is refused.- Verification and scheduling stay in the lane core. No second scheduler: the confidence head's survivals go to the
existing window allocator as each draft's chance. A checkpoint without the head uses the core's depth rule. There is
no verify flag. - This follows the paper. DSpark's prefix scheduler (Section 3.2.2, Algorithm 1) starts each request at zero verified
tokens and admits tokens while the estimated throughput rises, and Section 5.4 reports that the average verification
length decreases with load. The guards apply that rule to round times measured live, because our cost table is
timed on a short history and does not see that long contexts cost more. - Core additions are opt-in and off by default, each in its own PR from a lane-core series:
Drafted.held_on_target,
the shared-round guard, chain chances with granted depths and hold prices, the lone-stream guard, and per-stream
copy rates. Other families are not changed. - Without
--drafternothing changes for the target. - No environment switches. Tree verify and a lag-2 budget were built and measured no gain; they are not proposed.
Gates we propose
- Drafted equals
"draft": false, token for token (exact.py's 8 prompts,bench_concurrent --alone --serialat
1, 4, 8 and 16 streams, same and distinct prompts). - A resumed prompt equals a fresh one with the drafter on; a cancelled stream leaves the next request equal to its
solo run; each concurrent reply equals its solo run. - The drafter against an fp32 host reference of the released checkpoint (
tf-lfm2-dspark): drafts and first-row
logits, every pick's margin printed. - Kernel contract
SUMMARYlines on M5 Ultra and M4 Pro, with planted controls and recorded launch limits. M1 and M2
limits come from offline compiles. An M3 Pro has run the stack head (receipt in lfm2: a cheaper DSpark hold (one concurrent encoder) #1394); no run covers M1 or M2 yet (a test plan for LFM2: shared-round decode #1377 is in that PR). - Without
--drafter: replies, decode, cold prefill (2k, 8k, 32k,--prompt-cache-gib 0) and peak memory equal to
or within noise of the parent build. - With
--drafter: three alternated rounds on one machine, code and chat, 1/4/8/16 streams, a different prompt a
stream with the same-prompt cells labelled; cold prefill with the drafter; peak memory.
Questions
docs/recipes/SKILL.mdsays drafting stays on for every stream a family admits, and that a drafts-off or serial
fallback at some stream count is a bug. For this drafter, with a different prompt on each stream, the guards run
most shared rounds plain at 4 or more streams, and plain at very long contexts, because drafted rounds measure
slower there. They keep re-measuring drafting and turn it on again when it wins. DSpark's own scheduler also
verifies fewer tokens as load rises (see Design). Which do you prefer: (a) keep the guards; (b) draft every round
and accept the loss; (c) admit the drafter only up to a stream count and refuse above it?--drafter-bits 8is a lossy copy of the released weights made at load, off by default. Is an explicit option for
it acceptable, or should the drafter run bf16 only?- The released drafter loses acceptance at long context (see Long context). Is documenting that enough, or do you want
an option that limits the keys the drafter reads, which would differ from the released model? - Do you want one shared DSpark drafter across families later, instead of one drafter per family?
Plan (stack)
All of these are drafts. Each carries its own receipts, measured on its own commits. Numbers in this issue are replaced
by the PR receipts as they arrive.
Independent PRs (on main):
- server: capabilities --json's Engines struct moves to cli_capabilities.zig #1375 capabilities split (landed).
- tools: bench clients report TTFT, prefill tok/s, ITL and round gaps; bench_concurrent --prompts #1376 bench metrics and
bench_concurrent --prompts(landed). - LFM2: shared-round decode #1377 LFM2 shared-round decode: conv and attention once per layer over all streams, a GPU argmax. No drafter needed.
Lane-core series, each behind a flag that is off by default:
- lanes: Drafted passes held windows to a target that reads them itself #1378
Drafted.held_on_target. - lanes: a shared round drafts only while drafting measures faster (Model.shared_guard) #1379 the shared-round guard.
- lanes: a copy window may be priced by its stream's own copies (Model.own_copy_rates) #1380 per-stream copy rates (
Model.own_copy_rates). - lanes: an external drafter's chain chances price verify rows; holds follow the grants and pay for their pass #1381 chain chances, grants and hold prices (on lanes: a shared round drafts only while drafting measures faster (Model.shared_guard) #1379).
- lanes: a lone stream holds drafts only while its drafted rounds measure faster #1382 the lone-stream guard (on lanes: an external drafter's chain chances price verify rows; holds follow the grants and pay for their pass #1381).
Drafter family PRs, each on the one before it:
- lfm2: keep a drafter's tapped rows on every forward #1390 taps: the target keeps the drafter's tapped rows.
- lfm2: draft with LiquidAI's DSpark drafter (--drafter) #1391 the drafter with
--drafter(sits on lfm2: keep a drafter's tapped rows on every forward #1390, lanes: Drafted passes held windows to a target that reads them itself #1378, lanes: a shared round drafts only while drafting measures faster (Model.shared_guard) #1379). - lfm2: the DSpark confidence head schedules the verify #1392 scheduled verify (sits on lfm2: draft with LiquidAI's DSpark drafter (--drafter) #1391, lanes: a lone stream holds drafts only while its drafted rounds measure faster #1382, lanes: a copy window may be priced by its stream's own copies (Model.own_copy_rates) #1380).
- lfm2: the DSpark drafter loads one matrix at a time and keeps a lane in one buffer #1393 drafter memory: the loader and one buffer a lane.
- lfm2: a cheaper DSpark hold (one concurrent encoder) #1394 a cheaper hold.
Open points: the cross-family A/B on current main per head, and final-head receipts on the M5 and the M4, are being added to the PRs.
Activity