Repository navigation
lfm2: a cheaper DSpark hold (one concurrent encoder) - #1394
Open
BlivionIaG wants to merge 36 commits into
Open
BlivionIaG wants to merge 36 commits into
BlivionIaG wants to merge 36 commits into
Conversation
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. This builds on #1391, so it inherits that review's blockers (runtime lossy BF16 conversion, the 4-bit buffer mismatch, dropped async failure flags, the contract gaps). Its concurrent encoder and shared stage flush also need device race, resume and replacement-equality checks.
Contributor
Author
|
Thanks @ashhart. Device checks on an M4 Pro at the stack (details in the body, head |
BlivionIaG
marked this pull request as ready for review
October 26, 2026 21:37
7 of 9 tasks
Contributor
|
An M3 receipt for this head, since the body says no run covers M1, M2 or M3. M3 Pro (36 GB), macOS 26.7.1, Zig 0.17.0, head
|
Contributor
Author
Contributor
Author
Thanks for the run, now we have proof it works |
9 tasks done
Owner
|
Thanks @BlivionIaG, and thanks to @jschmied for the M3 Pro receipt. The concurrent hold encoder puts a barrier after This PR's own changes look ready. It is blocked only by what it inherits from #1391 and #1392: the guards, the |
Contributor
|
An M1 receipt for this head, for the load-time memory the body lists as not measured. M1 Pro (32 GB, MacBookPro18,1), macOS 27.0.1, Xcode 27.1 with Metal Toolchain 27A266a, Zig 0.17.0, head 2e5966f: zig build native -Dcpu=apple_m1: built. |
Contributor
|
An M1 receipt for this head: the load-time memory the body lists as not measured, and a slow-prompt finding on M1 that main has too. M1 Pro (32 GB, MacBookPro18,1), macOS 27.0.1, Xcode 27.1 with Metal Toolchain 27A266a, Zig 0.17.0, Load-time memory (head
Cold prompts are slow on M1, on this head and on main
|
This was referenced Oct 27, 2026
Owner
|
Thanks @Arminova for the M1 Pro receipt. The load peak (13.63 GiB without a drafter) is the number the body was missing. The slow cold prompts show on main too, so this stack didn't cause them. The long prompt that ignored SIGTERM looks like the same whole-prompt pass, which has no cancel point on main either. We'll look at it with #1377. @BlivionIaG, nothing has been pushed to #1390-#1394 since our last comments, so where things stand is unchanged. This PR's own changes are fine. What still blocks the stack:
Once that's in, we'll run these on the stack head: drafted against plain and concurrent against solo at 1 to 16 streams, the long-context cells, and alternated speed rounds against main. |
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.
engine.zig is 592 lines and the shared-round guard needs room in it. pipelinedRound and landInflight move unchanged to pipelined.zig as round and land; Engine.familyRound and Result become public for them. No behaviour change.
LiquidAI/LFM2.5-24B-A2B-DSpark is LiquidAI's speculative-decoding drafter for LFM2.5-24B-A2B. `tensorfold serve` now takes it with --drafter DIR for lfm2_moe checkpoints. The lane core verifies every draft, so a drafted reply equals its "draft": false reply. The drafter reads the target's rows after layers 2, 10, 18, 26 and 34 (the taps). fc and hidden_norm turn each row into the context keys and values of five Qwen3 layers. A block of the pending token and mask tokens runs the layers with non-causal attention, then the target's head. The Markov head adds W2 W1[previous token] to each row's logits and takes the argmax, one row after the other. A block is as long as the deepest chain its pass holds, at most the nine it was trained on. The drafter runs on the family's own kernels and core/tile_qmm (on GPUs with tensor units, core/affine_mm for absorb passes past the twin's rows): its bf16 linears are quantized to 8-bit MLX affine codes at load, and the q and k rows are reordered so that the target's half-split RoPE kernel gives the drafter's interleaved RoPE. New kernels: - lfm2_markov_scan and lfm2_markov_pick (lfm2_dspark.metal): a Markov step's W1 gather, W2 dot over 8-bit codes, bias add and argmax, ties to the lowest token, a NaN score flagged; - lfm2_attn_block and lfm2_attn_block_merge: the target's attention and merge kernels with a FULL template flag, every block row seeing the keys [0, len + 1). The target's own instantiations keep their causal count. Both register with the kernel contract (metal_lfm2_dspark.zig, metal_lfm2_attn.zig), gate 5 against every Apple chip we serve and this GPU. Rows reach the drafter in passes: a prompt chunk in one pass of up to 2,048 rows; a round's verified rows are staged on the host (128 a lane) and absorbed by the next hold, so a round without holds does no drafter GPU work. A hold pass takes as many streams as their blocks fit in 32 rows. The drafter sets Facts.shared_guard, so shared rounds draft only while drafted rounds measure faster than plain ones at that stream count; a lone stream keeps the depth rule's own plain guard. The drafter keeps its context with each kept prompt state, so a drafting stream resumes a prompt as before. Its lane's bytes join the streams the load limit admits. --drafter-bits 0 is refused for this drafter, and a refused open frees the model it loaded. Tests: test-lfm2-dspark checks the config's admission and refusals, the quantizer, the RoPE reordering through lfm2_qk, the Markov step and the block attention on the GPU. tf-lfm2-dspark runs the real drafter against an fp32 host reference of the released checkpoint and times its passes.
… the GPU A hold that fits in one pass signals a GPU event and returns without a wait. The next verify waits on that event on the GPU, and a new kernel, lfm2_draft_ids, copies each window's drafts from the drafter's picks into the verify's ids. The host reads the picks only after the verify, or when a grammar, an absorb, a keep or a release needs them. Multi-pass holds stay synchronous. The LFM2 backend takes held windows (Window.held) from a HeldSource the drafter gives, and Drafted.held_on_target passes windows without fillers to it as they are. lfm2_draft_ids registers with the kernel contract (metal lfm2 draft ids): each stream's first pick, or 0 for a token past the vocabulary, copied exactly; gate 5 against every Apple chip we serve and this GPU.
The recipe says where the drafter comes from (LiquidAI/LFM2.5-24B-A2B-DSpark at its revision), that its weights are under the LFM Open License v1.0 and not shipped with TensorFold, how to serve with --drafter, what the drafter runs and which kernels and checks cover it. The README row and the recipe index name the drafts.
… holds follow the grants A drafter that gives each held chain draft's chance (Drafter.VTable.chances) may set Facts.chances. Drafted.facts then sets Model.draft_probabilities and Model.granted_depth, so the window allocator prices each stream's verify rows by those chances, as it does a head's node chances. With Facts.confidence_floor it instead verifies the drafts at least that likely (a fixed window when 0). With granted_depth a stream's next hold is as deep as its last grant plus a slack: one more draft when the grant took every held draft, none when it stopped short of them (Stream.grant_short), and after a grant of none a probe of draft_slack drafts at the depth rule's probe interval. Facts.slack sets one slack for all three. A shared round's granted depths are cut so the drafting streams' blocks fit one hold pass of batch_rows. When the shared-round guard starts a drafted stretch, its first round grants every stream one draft, so the stretch measures drafted rounds rather than plain ones left by grants of none. All of it is behind Facts.chances and Model.granted_depth, off by default; windows.allocate records grant_short only for granted depths. Tests: chances price the chain's rows and drafted replies equal one-token decode alone and at 1, 4 and 8 shared streams, with the scheduled verify ragged across streams and verifying fewer rows than a fixed window; the facts reach the lane config; granted depths, their slacks and the one-pass cut; holds as wide as the grants.
…py, and holds kept only while they pay The depth rule asks one head draft for a stream with a copy ahead, and a shared round's granted holds were priced only by their verify rows. A drafter whose hold runs as a pass of its own (Facts.copy_first) pays that pass for the one draft, and the copy round then drops it; Model.copy_first now asks it for none. A drafter may also give its hold pass's timed cost by blocks and depth (Facts.hold_costs, priced by hold_price.zig). With granted depths a shared round then keeps granted holds, most expected drafts a row first, only while each raises the next round's expected tokens a ms over its verify rows and the pass; a probe after a grant of none keeps its hold. A stream's last grant records the drafts it expected to keep (Stream.grant_gain). Both are off by default (copy_first false, hold_costs empty). Tests: with copy_first no hold is asked for a copy round and replies equal one-token decode at 1 and 3 streams; a hold's cost between and past its timed passes; granted holds kept while they pay, none for a dear pass, and a probe's hold kept.
…ntract LiquidAI's DSpark drafter carries a confidence head: one linear over a draft's final block row (the row the target's head reads) and the W1 row of the token before it, 2,048 + 256 inputs, to one logit. lfm2_dspark_confidence (lfm2_dspark.metal) computes that logit for every row of a hold pass: 256 threads a row, fp32 products and sums, the eight simdgroups' partials summed in threadgroup memory. Nothing calls it yet. It registers with the kernel contract (metal_lfm2_confidence.zig): the served pass shapes against a float64 reference, the accumulation gate against reduced-precision sums, and gate 5 against every Apple chip we serve and this GPU.
…, --verify-costs) With the scheduled verify a hold also runs lfm2_dspark_confidence, and each lane keeps sigmoid(logit + bias) for its drafts. chances() gives their running product, the chance that a draft and every draft before it land, as Facts.chances: the lane core's window allocator then picks each stream's verify rows by those chances and the verify costs, and the next hold drafts as deep as the stream's grant. No second scheduler. While a hold still runs, its rows are priced by the lane's last settled confidences (draft j by the last hold's draft j, the deepest one repeated past it), so a verify does not wait on the GPU for them. A lane with none settled yet waits. Two serve flags on macOS, listed by capabilities --json and refused unless an lfm2_moe checkpoint is served with --drafter: - --draft-verify depth|fixed|scheduled: the core's depth rule, the whole block every round (the published fixed window), or the confidence schedule. scheduled is the default. - --verify-costs FILE: a table of verify forward ms by streams and rows a stream, written by tf-lfm2-verify-costs on the serving Mac, in place of the costs timed at load. Shared rounds keep the shared-round guard in every mode but fixed. Tests: test-lfm2-dspark reads a cost table, refuses malformed ones and checks the default mode; the server's parser tests take both flags and refuse an unknown mode. tf-lfm2-dspark checks the head against its fp32 reference.
…fied rows Each prompt chunk's rows were absorbed at once in a command buffer of their own, waited on before the first token, and the prompt's last row in a second one at the first draft request. The prefill pass is waited on before its rows are handed over, so a chunk that fits the lane's stage (128 rows) can be copied there on the host like a round's verified rows. The next hold, or a kept prompt state, absorbs the stage first, as before. Longer chunks are absorbed at once.
A DSpark hold is a pass of its own: the drafter runs its five layers and the target's head for every drafting stream before the verify. The lane core can now weigh that pass (Facts.copy_first, Facts.hold_costs, Model.own_copy_rates); the drafter gives it what it needs. - copy_first: a stream with a copy ahead verifies the copy, and the drafter is asked no hold for it. - hold_costs: after its warm-up the drafter times holds of 1, 2, 4, 8 and 16 blocks at one draft and at four (two at 16 blocks, one 32-row pass) on empty lanes, best of three. Each settled one-pass hold then moves the timed passes at its block counts a tenth of the way toward its own GPU ms, so the price follows the served contexts. - own_copy_rates: with a drafter, a stream's copy windows are priced by its own copies. Test: a 4-block hold at twice its price raises the 4-block passes by a tenth and leaves the others; one between 4 and 8 blocks moves both; a hold without GPU time teaches nothing.
…eal rows tf-lfm2-dspark times a shared round's hold over 2 to 16 streams, by the wall clock and on the GPU, and absorbs of 128 rows and of a whole absorb pass, so a hold's price can be checked against a verify's. tf-lfm2-verify-costs reads token id files after HISTORY: each of the 16 streams takes its history and then its verify rows from its own sixteenth of the files' ids (the history shrinks to fit), so a shared round's rows route the experts their own text routes. Without id files the rows are synthetic as before; the table's header names their kind.
The recipe names the confidence head, the three --draft-verify modes (scheduled by default), how the lane core prices a hold pass and a copy, the lone-stream guard, and --verify-costs with tf-lfm2-verify-costs.
…in one buffer Memory of the drafter, no change to its drafts: - The loader keeps the matrices in the safetensors file (Checkpoint.addFileSplit) and reads each into one reused stage, quantizes it and drops it. Load no longer holds the 655 MB bf16 file beside the 8-bit codes, only the largest matrix (W1 and W2). The codes are the same bytes. - A lane is one Metal buffer of laneBytes: its layers' keys and values, then its staged rows. Opening a stream allocated eleven buffers. Lane.kv gives a layer's offset; the rows and their layout are unchanged. - The absorb norms its context rows in place, so the rotated query heads go to the freed buffer and x, which only a hold reads, takes 32 rows instead of 2,048. - The waited passes end in one helper that commits, waits and reports a failed command buffer. - At load the drafter logs its weights, scratch and a stream's bytes. tf-lfm2-dspark counts host allocations: twenty rounds of hold, chances, held drafts and a staged absorb through the drafter's interface, after four warm ones, must take no allocation or resize of the drafter's allocator and one unwaited hold pass each.
The recipe says how the drafter loads and keeps its lanes, what it logs at load and that tf-lfm2-dspark counts host allocations over served rounds.
encodeTwin picks the window kernel from window_rows rows on, for matrices of window_n outputs or more, and the scalar twin below. A caller that measured its own cut (the DSpark drafter's hold projections) now passes it to encodeTwinFrom; encodeTwin keeps its choice and its callers are unchanged. The kernels are unchanged.
A hold ran every dispatch in a serial encoder, its blocks' attention one block after the other. A hold now encodes on a concurrent encoder and puts a barrier only where a step reads an earlier step's writes: - q, k and v, and gate and up, are projections of one input: one sums pass, then their products free to overlap (projectAll). The absorb's k and v do the same. - The attention partials hold a whole hold pass's query rows, so the blocks of a shared hold attend side by side in one wave, then merge in one. - A projection picks its kernel by rows and output width: the scalar twin up to 8 rows, the window kernel past them where the GPU has no tensor units, and the tensor units from 9 rows (the target's head, 128,000 outputs, from 5) where it has them; outputs of the hidden width or fewer stay on the twin to its rows. The kernels are the target's own and core's, unchanged; only their order and the rows each launch takes change. Drafts are checked by tf-lfm2-dspark against the fp32 reference, and every draft is still verified by the target.
In plain rounds each lane stages its tapped rows and absorbs them in a waited pass when its 128-row stage fills. With one prompt on every stream the stages fill in the same round; with a different prompt a stream they fill on different rounds, so 16 streams waited on 16 passes of up to 128 rows each stage's worth of rounds. A pass that a full stage starts now takes every other lane's staged rows too (up to 2,048 rows, 32 pieces a pass), so the lanes' stages fill together after the first one.
The recipe says how a hold encodes its passes and picks a kernel for each projection, and how a full stage absorbs every lane's rows.
…refused, and hidden_act must be silu The unwaited hold read its picks with nonfinite ignored and a bad id as 0; the GPU copy into a verify window did the same. Both now refuse as the synchronous path does (DraftNonfinite): readPicks errors, and lfm2_draft_ids sets a flag the verify reads. Planted NaN, Inf and invalid-id tests cover the host and the kernel. The config refuses a hidden_act other than silu by name.
… the checkpoint and joins its workers on every path
… 0, the default), no quantizing at load The drafter's linears and Markov W2 stay bf16: new rows and tile kernels for the linears and a bf16 scan. An explicit --drafter-bits other than 0 is refused by name. 8-bit codes stay behind Options.bits for tools.
…, the README, the recipe and the startup log --drafter-bits takes 0, 4 and 8. LFM2's drafter runs bf16 as released when it is not given or 0, 8-bit codes quantized at load on 8, and refuses 4. The other families refuse 8 by name.
…y gates 1-4, recorded M1/M2 limits and six planted controls The few-rows and 64x64 tile kernels and the scan over a bf16 W2 pass gates 1 to 4. Six controls edit the source: a lane skipping half of k, a group's last row, an untransposed W tile, a lost last column, a NaN score not flagged and every stream reading the first stream's previous token. Each goes red at its gate. The tile kernel compiles for 832 threads on M1 and M2.
…nces, else the depth rule --draft-verify and --verify-costs are gone with the fixed mode, the cost-table loader and tf-lfm2-verify-costs. A checkpoint with the confidence head verifies by its chances. One without it verifies by the depth rule.
BlivionIaG
force-pushed
the
series/dspark-hold
branch
from
October 28, 2026 08:57
2e5966f to
017919f
Compare
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.
The DSpark stack head, rebuilt on current main
d0c9dc9fbwithout the rejected shared and lone guards. Part of #1374; includes the lane pricing and drafter work reviewed in #1381 and #1390–#1393.What this changes
encodeTwinFromselects the twin's window threshold without changingencodeTwinor its callers.Review fixes
granted_depth,chances,copy_firstandhold_costs.Drafter.initcleanup so buffers are released before their list. Regression tests inject scratch and warm allocation failures, including in Debug.--drafterby name when the target's local wider-lane qualification leaves only one-row verification.hold_price.atDepth's description.Receipt
Apple M4 Pro (20 GPU cores, 48 GiB), macOS 26.5.1, Zig 0.17.0, Apple Metal compiler 32023.883;
-Dcpu=apple_m1.zig fmt --check zigandgit diff --checkpass. Lean: 428 findings on main and this source, zero new.667408d494bde191c2cb9dd4f4549827816b9275, both bf16 and 8-bit drafters passed all 12 same/distinct-prompt cells at 1, 4, 8 and 16 streams: concurrent equals solo and drafted equals"draft": false, zero unequal replies or failed requests. Public bench prompts, greedy, 64 tokens, seed 1234, one repetition, prompt cache off.1bb6e5db216,036,097,352 bytes (14.935 GiB); corrected stack 15,527,930,776 bytes (14.462 GiB)./memoryagrees withtime -l. Reduction: 484.6 MiB. This comparison includes the rebase onto current main, rather than isolating lfm2: the DSpark drafter loads one matrix at a time and keeps a lane in one buffer #1393 on an identical base.Current-head receipt retains the token-hash comparisons, exact memory counters, commands and remaining qualification.
Extended qualification and remaining performance issue
M4 Pro and M5 Ultra pass seeded sampled equality, real two-turn resumed/fresh equality, and cancellation/replacement stress at 6 and 16 slots with Metal API validation. M5 also passes plain/drafted equality through 130,816 prompt tokens and independent fp32 drafter references through 130,815 history tokens. Full request hashes, counters, alternating rounds and a reproduction harness are in the linked receipt.
Performance is not fully qualified: three alternating rounds show M5 drafted code decode +19.44% at one stream but up to -4.23% at wider concurrency; M4 shows losses up to -13.92%. Cold prefill is also slower with the drafter. Qwen/DFlash2 on M5 is within 0.7% of main, while M4 varies with independently measured build-specific cost records. No guards were restored and no inconclusive experimental pricing changes are included.
All nine GitHub checks passed on code revision
2dbbc0e8fbefore the receipt-only update. M1/M2/M3/CUDA/HIP were not rerun; focused Metal checks are not the full cross-family kernel-contract gate.