DSpark B2 rejection sampling + adaptive block sizing#482
Conversation
Implements Chen et al. (2023) / Leviathan et al. (2023) rejection sampling for the DSpark speculative decode path. At temp=0 (default), behavior is unchanged (greedy argmax matching). At temp>0, accepts draft tokens with probability min(1, target_prob/draft_prob) and samples correction tokens from the residual distribution. Key properties: - Distribution-exact: output is drawn from EXACTLY the target model's distribution regardless of draft quality - Token-identical to non-spec decode at temp=0 (greedy) - Numerically stable: all computations in log-probability space (handles 129K vocab without overflow) - Zero overhead when disabled: buffer only allocated when DS4_SPEC_TEMP is set Activation: DS4_SPEC_TEMP=0.6 DS4_SPEC_RNG_SEED=42 ./ds4 ... Based on the community analysis from ds4#468: @lobanov proved +30-48% speedup with component measurements; this wires the acceptance algorithm into the decode loop for end-to-end lossless spec decode. Co-Authored-By: Tang Feng (Audrey) <audreyt@audreyt.org> Co-Authored-By: lobanov <lobanov@users.noreply.github.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
ARDD adversarial review found two critical bugs: 1. OFF-BY-ONE (CRITICAL): metal_graph_verify_suffix_tops row[i] is the target distribution AFTER processing drafts[i], predicting drafts[i+1]. B2 was using row[i] to verify drafts[i] — wrong distribution entirely. Fix: prepend s->logits (from previous target eval) as row 0, shift verify output by one. Now drafts[0] verifies against s->logits and drafts[j>0] verifies against verify_row[j-1]. 2. RNG RESEEDED PER CALL (HIGH): b2_rng was stack-local, reseeded from time(NULL) on every speculative eval call. At >1 tok/s, multiple calls within the same second got identical seeds, producing correlated random sequences. Fix: persist RNG state in ds4_session struct, seed once. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
The skip-logits optimization (NULL logits on intermediate replay tokens) saves the lm_head matmul but doesn't meaningfully reduce replay cost — the 43-layer forward (~26ms/tok) dominates. Key finding from end-to-end testing: - Baseline: 35.04 tok/s (short ctx), 37.59 tok/s (long ctx) - DSpark with replay: 23-26 tok/s (SLOWER — replay ~26ms/tok eats margin) - DSpark fast-partial (no restore): 34.63 tok/s (AT PAR but corrupts compressed KV — phantom entries cause output divergence) The gap between lobanov's component measurements (+30-48%) and end-to-end is EXACTLY the replay cost. His calculation assumed zero replay, which only holds for full 5/5 commits (~16-21% of cycles). Path to speedup: multi-point checkpointing during verify — capture compressor state (layer_n_comp, attn_state_kv) at each draft position. On partial commit at K, restore from checkpoint K instead of frontier (position 0). This eliminates the O(K × 43-layer) replay entirely. This requires modifying metal_graph_verify_suffix_tops to snapshot intermediate compressor state — a change to ds4's batch verify internals best driven by the upstream maintainer. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Reverts hybrid/fast-partial experiments (corrupted compressed KV). V4-Flash compresses 41/43 layers (ratio=4 even, ratio=128 odd, only layers 0-1 uncompressed). The fast-partial approach cannot work without per-token verify checkpointing across all compressed layers — a deep architectural change. The correct B2 replay path is: - Correctness: PASS (token-identical at temp=0) - Performance: replay cost = baseline cost per token → net negative - Fix requires: per-token checkpointing in metal_graph_verify_suffix_tops
Conservative-then-aggressive strategy: start at block=2 (near-baseline), escalate to full block after seeing full commits (structured output detected), drop back on partial commits. Tracks previous cycle's acceptance via dspark_prev_accepted/drafted fields in session struct. Opt-in via DS4_DSPARK_ADAPTIVE=1 env var. Measured results (M5 Max 128GB, Q2 base + Q4K DSpark, clean bandwidth): JSON structured: +8.5% speedup (42.27 vs 38.95 tok/s) Markdown tables: +6.3% speedup (45.21 vs 42.54 tok/s) Code generation: -17% (needs higher acceptance — HyperDFlash path) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
|
Got the DSpark block speculative path running on CUDA (NVIDIA GB10 / DGX Spark) on top of this branch, in case the numbers are useful. The CUDA backend currently ships Measured on GB10 (Q2 V4-Flash base + Q4_K DSpark drafter, greedy,
Adaptive block sizing is doing the heavy lifting. With a fixed block, partial-commit replay dominates and the whole thing goes net-negative, which matches the #468 analysis. One caveat worth flagging: on free-form text the CUDA output is not strictly token-identical to non-spec decode. I traced one divergence to a target batch-verify vs plain-decode floating-point near-tie — the top two candidates were 0.29 logits apart, and about 10 of 80 positions in that prose sample were within a 0.3-logit margin. The noncausal kernel is not the cause: committed tokens are always the target argmax, so the drafter only changes the accepted length, never the committed value. The CUDA batch-verify just accumulates slightly differently from single-token decode; aligning the two would make it strictly token-identical. I couldn't check whether the Metal path is bit-identical here, so this may be part of why dspark is Metal-gated for now. Happy to share the patch or open a PR against the branch if it's useful. |
|
I cannot reproduce the speedup on my M5 Max 128GB using the code form this branch @machiabeli Here's detailed report from GLM-5.2/Pi, which is hopefully useful: PR #482 (DSpark B2 rejection sampling + adaptive block sizing) — Device Test ReportIndependent on-device test of GitHub PR #482 (antirez/ds4). 0. PR identity
1. PR's claimed results (from the PR body)
Key PR features: (1) B2 rejection sampling 2. Test environment
3. BuildCommands (in the PR worktree): Result: builds clean, exit 0, 0 warnings / 0 errors. Binary The PR loads the DSpark drafter via 4. Invocation conventions (env vars)
5. Test matrix + resultsAll runs: 5.1 Code prompt (creative, low acceptance) — n=256Prompt:
5.2 Repetitive JSON prompt (the PR's claimed +36% regime) — n=400Prompt: `'Generate a JSON array of 30 user objects. Each object must have exactly these keys in this order: "id" (incrementing integer starting at 1), "name", "email", "age". Use the same template for every object. Start: [{"id":1,'
The PR's headline +36% on repetitive JSON does NOT reproduce on this device — measured −24% (31.07 vs 38.72). The PR is slower than baseline on every prompt/configuration tested. 6. Root-cause: the PR runs with KV replay on every cyclePer-cycle timing ( Example PR cycles (repetitive JSON, Aggregate over 55 cycles: acceptance 63.9% (genuinely high — the PR should win here), but 100% of cycles replayed, avg replay 74.9ms. 19 of 19 full-accept cycles paid the replay. A full-accept cycle = verify 64ms + replay 132ms = 196ms / 5 commits = 39ms/tok ≈ baseline. The replay erases the speculation win. 7. Correctness claims — verified, do NOT hold on this device7.1 "Token-identical greedy at temp=0" — FAILSCommands (n=200, repetitive JSON prompt, temp=0.0): Hashes:
Neither PR mode reproduces plain greedy at temp=0. Root cause: 7.2 "Lossless B2 at any temperature" — NOT lossless as implementedThe PR's 8. ReproducibilityExact commands (from §5, with All raw outputs captured during testing: |
|
Hi @machiabeli, Unfortunately I cannot reproduce the speedup on my M5 Max 128GB using the code form this branch. Here's detailed report from GLM-5.2/Pi, which is hopefully useful: PR #482 (DSpark B2 rejection sampling + adaptive block sizing) — Device Test ReportIndependent on-device test of GitHub PR #482 (antirez/ds4). 0. PR identity
1. PR's claimed results (from the PR body)
Key PR features: (1) B2 rejection sampling 2. Test environment
3. BuildCommands (in the PR worktree): Result: builds clean, exit 0, 0 warnings / 0 errors. Binary The PR loads the DSpark drafter via 4. Invocation conventions (env vars)
5. Test matrix + resultsAll runs: 5.1 Code prompt (creative, low acceptance) — n=256Prompt:
5.2 Repetitive JSON prompt (the PR's claimed +36% regime) — n=400Prompt: `'Generate a JSON array of 30 user objects. Each object must have exactly these keys in this order: "id" (incrementing integer starting at 1), "name", "email", "age". Use the same template for every object. Start: [{"id":1,'
The PR's headline +36% on repetitive JSON does NOT reproduce on this device — measured −24% (31.07 vs 38.72). The PR is slower than baseline on every prompt/configuration tested. 6. Root-cause: the PR runs with KV replay on every cyclePer-cycle timing ( Example PR cycles (repetitive JSON, Aggregate over 55 cycles: acceptance 63.9% (genuinely high — the PR should win here), but 100% of cycles replayed, avg replay 74.9ms. 19 of 19 full-accept cycles paid the replay. A full-accept cycle = verify 64ms + replay 132ms = 196ms / 5 commits = 39ms/tok ≈ baseline. The replay erases the speculation win. 7. Correctness claims — verified, do NOT hold on this device7.1 "Token-identical greedy at temp=0" — FAILSCommands (n=200, repetitive JSON prompt, temp=0.0): Hashes:
Neither PR mode reproduces plain greedy at temp=0. Root cause: 7.2 "Lossless B2 at any temperature" — NOT lossless as implementedThe PR's 8. ReproducibilityExact commands (from §5, with All raw outputs captured during testing: |
|
@lobanov your §7.2 finding is right, and I tracked it down to a specific bug: Fix + PR against the branch: machiabeli#1. Summary:
Evidence the fix actually changes proposal sampling (not just the accept/reject coin flip): two runs with different Also added a CPU-only synthetic statistical test on my fork (not in the minimal PR, happy to add on request): for a single draft token, feeding a real sample from a synthetic drafter distribution into B2's accept/reject step reproduces the target distribution within noise ( Scope: this addresses your §7.2 (proposal not drawn from q). It does not touch §7.1 — the batch-verify-vs-single-decode floating point divergence that makes even the greedy DSpark path diverge from plain non-speculative decode at temp=0. That's a separate, pre-existing issue independent of B2 and still open; I haven't dug into On speed: this fix is about correctness, not throughput, and doesn't change the picture there. On my M5 Max 128GB with the Q4K Markov drafter, B2 is still slower than baseline on most workloads I tried after the fix — partial/correction-cycle replay is still the dominant cost, matching the PR's own "Key findings" on creative-prompt replay overhead. Repetitive structured output remains the one clear net win. |
|
Follow-up on my comment above: I added the actual losslessness proof to the PR, not just behavioral evidence that sampling changed.
That's what actually proves losslessness (the earlier seed-divergence evidence only proved the proposal was sampled, not that the output distribution is correct). PR updated: machiabeli#1. |
|
Quick follow-up on the §7.1 greedy divergence, from the CUDA side — hopefully a useful lead. This looks target-side, not the drafter: in the greedy path It only shows up at near-ties: the one divergence I traced had the top-2 target logits 0.29 apart, and ~10/80 positions in that prose sample were within a 0.3-logit margin. Structured/repetitive output stayed token-identical; free-form prose is where it flips. If it's worth closing the gap rather than documenting it as near-lossless: the N=2 MTP path already has an exact decode-order verifier — applying that same ordering to the block verify (instead of the fast batch path) might be the cleanest fix. I'm on CUDA so I can't confirm the Metal path is bit-identical, but the code structure lines up. |
4 participants
Summary
Adds lossless B2 rejection sampling (Chen/Leviathan 2023) to the DSpark speculative decode path, plus adaptive block sizing for workload-aware performance.
Builds on PR #480 (@audreyt) which provides the DSpark loader, Metal drafter, converter, and test infrastructure.
Verified speedups (M5 Max 128GB, Q2 base + Q4K drafter)
All outputs are token-identical to non-speculative greedy decode at temp=0. Lossless by construction at any temperature.
What's in this PR
B2 rejection sampling (
b2_rejection_sample, 170 lines) — log-space stable for 129K vocab. Accept with probability min(1, target/draft), sample correction from residual on reject. Activated viaDS4_SPEC_TEMP=0.6env var; greedy path unchanged when unset.Off-by-one fix —
metal_graph_verify_suffix_topsrow[i] predicts drafts[i+1], not drafts[i]. B2 target logits shifted by prependings->logitsas row 0.RNG state persistence — moved from stack-local (reseeded per call from
time(NULL)) tods4_sessionstruct. Prevents correlated random sequences at >1 tok/s.Adaptive block sizing (
DS4_DSPARK_ADAPTIVE=1) — starts conservative (block=2), escalates to full block after full commits, drops back on partial commits.Key findings
metal_graph_verify_suffix_tops(DeepSeek-V4-Flash-DSpark #468)Credits
Co-Authored-By: Tang Feng (Audrey) audreyt@audreyt.org
Co-Authored-By: lobanov lobanov@users.noreply.github.com