dspark: fold first-token forward into verify + bit-exact M5 Metal decode wins - #778
Open
polymorf wants to merge 3 commits into
Open
dspark: fold first-token forward into verify + bit-exact M5 Metal decode wins#778polymorf wants to merge 3 commits into
polymorf wants to merge 3 commits into
Conversation
Metal DSpark decode reworked around the DFlash block layout and a set of regime signals, all opt-in and deterministic at temp 0: - DS4_DSPARK_FOLD=1 verifies the committed token as batch row zero instead of paying a standalone one-token forward per cycle (~22 ms). The existing accept loop, prefix slots, replay fallback and hidden capture apply to the folded block unchanged. Folding re-arms only after a full accept; any miss or partial drops back to the standalone path, which refreshes the draft conditioning and returns control to the scheduler. - DS4_DSPARK_CONFIDENCE_GRID=a,b,c grades the confidence pruning threshold by draft position (stale-conditioned folds over-prune depth with the flat default; 0.55,0.6,0.65 measured optimal under fold). - DS4_DSPARK_SCHEDULER_BACKOFF=1 doubles pause lengths per consecutive unproductive cycle (bounded x8) and raises the position-0 engagement bar (+0.05/step, bounded +0.20); one accepted draft resets both. - DS4_DSPARK_MARGIN_GATE=<logits> skips standalone proposing when the target's own top1-top2 logit margin is below the bar: a hesitant target is content the weaker draft cannot predict, and the margin is free on the host. Fold-chain proposing is never gated, a full accept just validated the chain. - The per-row batch compressor update loops take the packed ratio-4 pool (bit-identical per the in-tree pack gate), and a fused sequential rows kernel with bit-exact tests covers large non-aligned chunks (n_tokens >= 16; small verify batches measured faster unfused). - mv_ext gains r1_6..8 row tiles so 6-8 row verify blocks stop reading every dense weight twice; the speculative head keeps its eight-row padding on CUDA only. A Metal port of the CUDA fused markov argmax ships behind DS4_DSPARK_GPU_MARKOV (per-position submit latency outweighs the CPU matvec it replaces on Apple; kept as the building block for a batched chain). DS4_METAL_GAP_PROFILE reports GPU-timeline holes at token boundaries. speed-bench/ngram_similarity.py characterizes workload self-similarity for speculation triage. With --dspark --mtp <support.gguf> plus FOLD, GRID=0.55,0.6,0.65, BACKOFF=1 and MARGIN_GATE=1, greedy CLI throughput on M5 Max vs plain decode: structured payload 46.2 -> 62.8 t/s (+36%), codegen 46.8 -> 48.2 (+2.9%), Italian prose 44.8 -> 46.6 (+4.0%); no measured workload regresses. The default bit-exact path is untouched: official bench replicas keep the canonical frontier logits SHA-256 (eb794f497861d7d9e373665f6e115d8ebe4e17c13c431aa7e318282f16a0f21d) and full make test plus the focused Metal kernel suites pass. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The staged B tile stores element (n, k) at sb[n*NK + k]: dim0 is K (stride 1), dim1 is N (stride NK). The tB tensor declared its extents as (NR1, NK), which only described the layout correctly because NR1 == NK == 32; any future instantiation with NR1 != NK would compute against a transposed view of the staged tile. Spell the extents (NK, NR1). Extents and strides are identical while NR1 == 32, so the shipped kernel is unchanged: verified bit-exact on full-vocabulary logit frontiers, and the Metal kernel tests pass. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- HC producer spread: give every comb threadgroup one live NR0=2 cluster instead of packing both clusters on groups 2..5, raising the producer grid from 6 to 10 threadgroups. Per-row lane traversal, reduce tree and epilogue are identical across both shapes, so the outputs are bit-identical; the wider grid measured ~+0.12% decode on M5 Max (40-core). Default on M5-class GPUs, kill switch DS4_METAL_DISABLE_M5_HC_SPREAD. - FFN island into HC expand: at the parallel-FFN join, publish the routed/shared outputs with a memory barrier inside the still-open concurrent encoder and let the single HC expansion dispatch land in the same encoder, saving one encoder transition per layer. Kill switch DS4_METAL_DISABLE_M5_FFN_ISLAND_INTO_HC. Both changes verified bit-exact against the previous defaults on full-vocabulary logit frontiers (8 runs each, abort-on-diff harness); official bench keeps the canonical frontier hash with prefill 786-788 t/s and steady decode 45.50-45.58 t/s on M5 Max. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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
1 participant
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.
Three commits, measured on M5 Max 40-core / 128 GB with the ds4f-q2 model. The default (non-speculative) decode path stays bit-exact throughout: every change was gated by an abort-on-diff A/B harness comparing full-vocabulary logit frontiers, and the official bench keeps the canonical frontier hash (prefill 786-788 t/s, steady decode 45.50-45.58 t/s).
dspark: fold the first token's forward into the verify batch
With
--dspark, each accept/miss cycle previously paid a standalone target forward for the first token before the verify batch. This folds that forward into the verify dispatch as row 0 (the DFlash protocol shape), so one batched pass covers first token + verification. The fold is self-regulating: it re-arms only on a full accept; any partial accept or miss falls back to the standalone path, which resynchronizes draft conditioning.Also in this commit, all opt-in via env:
DS4_DSPARK_CONFIDENCE_GRID— per-position draft confidence thresholds (best:0.55,0.6,0.65)DS4_DSPARK_SCHEDULER_BACKOFF— adaptive selectivity + regime pauses on consecutive miss streaksDS4_DSPARK_MARGIN_GATE— skip proposing when the target's top1-top2 logit margin says it is hesitating (the draft has no chance there); applied at the standalone site onlyWith
DS4_DSPARK_FOLD=1 DS4_DSPARK_CONFIDENCE_GRID=0.55,0.6,0.65 DS4_DSPARK_SCHEDULER_BACKOFF=1 DS4_DSPARK_MARGIN_GATE=1(temp 0,-n 512), against plain decode on the same channel:No channel regresses — the margin gate is what turned prose from a loss into a gain (the speculative path only works its genuinely predictable spans, at ~87% acceptance). The speculative flow is deterministic run-to-run.
metal: fix staged B-tile tensor extents in the mm_id mpp kernel
Latent-bug fix, also submitted standalone as #777 (kept as its own commit here so either merge order works): the
tBextents were declared transposed and only accidentally correct becauseNR1 == NK == 32. Bit-identical today; required for any futureNR1 != NKinstantiation.metal: two M5 decode encoder/occupancy wins, both bit-exact
DS4_METAL_DISABLE_M5_HC_SPREAD.DS4_METAL_DISABLE_M5_FFN_ISLAND_INTO_HC.make testpasses,--metal-kernelsgreen, ROCm/CUDA paths untouched apart from the already-merged DSpark support hooks.🤖 Generated with Claude Code