Skip to content

LFM2.5-24B-A2B on Metal: drafts with LiquidAI's DSpark drafter (proposal and gate) #1374

Open

Description

@BlivionIaG

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": false on 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), like gemma4/dflash.zig. It reads the target's rows
    after layers 2, 10, 18, 26 and 34 through prepare_features and features, 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 8 is an explicit option that quantizes the drafter's linears at load to 8-bit codes. It is named in
    --help and 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 --drafter nothing 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 --serial at
    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 SUMMARY lines 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

  1. docs/recipes/SKILL.md says 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?
  2. --drafter-bits 8 is 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?
  3. 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?
  4. 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):

Lane-core series, each behind a flag that is off by default:

Drafter family PRs, each on the one before it:

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions