Repository navigation
lfm2: a cheaper DSpark hold (one concurrent encoder) - #1394
Open
BlivionIaG wants to merge 43 commits into
Open
BlivionIaG wants to merge 43 commits into
BlivionIaG wants to merge 43 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.
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
Contributor
Author
|
@ashhart Following up on your last checklist: the cumulative fixes are now pushed here, at
Receipts, reproducer and limitations. The final allocation-only follow-up removes duplicate probability and held-grammar buffers without changing arithmetic or decoding policy. Its shared-lane/native checks pass (341 tests, three skipped), and the separate 64-cell M5 before/after comparison completed with matching token hashes and throughput changes between -0.45% and +0.68%. That comparison is additional evidence, not a repeat of the full hardware suite on the final commit. All nine CI checks are green on |
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.
Review fixes
Rebuilds the DSpark stack on main (
dd0a9cc5e), excluding the rejected shared/lone guard commits. Releases partially initialized drafter buffers, tests allocator failures, validates rank and theta, widens quantizer arithmetic, and refuses a drafter when qualification leaves only width one. BF16 remains the released default; 8-bit is opt-in and 4-bit is refused. Weight conversion handles one matrix at a time.Performance
Source
099455490measures shared 17–32-row costs using two independently qualified windows, anchored to the single-window timing. It also avoids querying a stateful copy proposer twice for one granted-depth decision. No stream-count fallback or new draft-disable gate is introduced.The M5 comparison comprises 128 cells across four rotated runs, code/prose prompts, 1/4/8/16 streams, 256 generated tokens, and main/plain, earlier PR/drafted, candidate/plain, candidate/drafted arms. Every request's token hash agrees.
At 16 streams, drafted throughput improves 10.96% code / 8.86% prose over the earlier PR, and 6.78% / 8.05% over main. At four streams, it loses 1.93% code / 0.42% prose against the earlier PR. The task owner explicitly accepted this tradeoff; this is not a universal no-regression claim. Candidate plain four-stream prose measures 2.04% below main.
Full results and limitations, request-level evidence, discarded experiments, reproducer.
Validation and remaining work
For source
099455490: scoped host checks 117/117 tests, default build 239/239 steps; M5 LFM kernel contracts 11 pass with 17 negative controls red; independent FP32 drafter 9/9; 20 steady rounds with zero host allocations/resizes; resumed-cache, sampled, cancellation, 8-bit, and 65,536/130,816-token equality checks pass. M4 kernel contracts and FP32 reference pass, but memory pressure prevented its new served/cache/concurrency qualification. Historical M4 evidence is retained and distinguished from this exact-source limit.The candidate's M5 startup peak was 22.7 MiB higher than the preceding PR in one measurement each; calibration adds about one second to warm startup. Qwen/DFlash2 arithmetic and kernel paths are unchanged; the generic granted-depth fix applies only when that capability is enabled.
Published allocation follow-up
98f5cc72breuses final probability buffers and reads held grammar tokens directly into their final row buffer. It preserves arithmetic and decoding policy, and passes 57/57 shared-lane/native host build steps (341 tests passed, three skipped). Its bounded-memory regression fails on the previous implementation and passes on the cleanup. The M5 native build passes; its 1/4/8/16-stream before/after comparison is queued behind another qualification. No additional throughput gain is claimed yet.CI is restarting on the new head, and the prior changes-requested review requires maintainer re-review. The performance and full hardware receipts above qualify source
099455490; the allocation follow-up has the narrower validation explicitly described here.