Metal: Reconcile DSpark runtime (#480) + B2 rejection sampling/adaptive block sizing (#482) by stephenlthorn · Pull Request #502 · antirez/ds4

stephenlthorn · 2026-07-05T09:19:35Z

Summary

This reconciles the two open DSpark PRs from issue #468 into one branch that builds cleanly and passes the model-free test suite:

Add DSpark speculative draft runtime #480 (@audreyt) - DSpark Metal runtime, converter, prefix-checkpoint partial commits
DSpark B2 rejection sampling + adaptive block sizing #482 (@machiabeli) - B2 rejection sampling (lossless temp>0) + adaptive block sizing (DS4_DSPARK_ADAPTIVE)

The two PRs share an earlier common scaffold but diverged afterward, so this isn't a plain merge - I cherry-picked both commit chains onto main (skipping #482's skip-logits-replay commit + its own revert, which cancel out to a no-op) and hand-resolved one real conflict.

Fix included

The conflict was in the partial-commit path of the speculative loop: audreyt's prefix-checkpoint fast-commit (skip replay when the batch verify already captured a KV row for each committed position) and machiabeli's B2 correction token (sampled fresh during rejection, never processed by the verify pass) interact badly - a correction token has no captured KV row, so routing it through the fast-commit path would silently commit against a stale/wrong cache row.

Fixed by gating the fast-commit path on commit_has_correction: any commit that includes a B2 correction token always takes the frontier-restore + replay path, same as before #480's optimization existed. Also fixed a small leak of b2_target_logits on the prefix-commit early return.

Testing

On Apple Silicon (Metal), no GPU-heavy model runs yet on this machine:

make (Metal): clean, 0 warnings
make cpu: clean
./ds4_test --dspark-binder --dspark-markov-bf16 --dspark-runtime --server: all pass
Model-dependent tests (--dspark-speculative-block, --mtp-verify-depth) self-skip without DS4_TEST_DSPARK/DS4_TEST_MTP - not yet run against a real DSpark GGUF

I plan to run the real acceptance/speed benchmarks from #468 (block=2 vs block=5 vs adaptive, JSON/code/creative prompt mix) on a machine with the model loaded and will report back with numbers. Flagging this now as a reviewable integration point rather than waiting, since #480 and #482 are both stalled on being combined.

Thanks to @audreyt and @machiabeli for the actual implementation work - this PR is integration/reconciliation only, no new algorithmic contribution beyond the correction-token fix above.

@lobanov

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>

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>

stephenlthorn · 2026-07-05T11:17:31Z

On-hardware validation (M4 Max 128GB, q2-imatrix base + Q4_K DSpark drafter)

Correctness: PASS

Test	Result	Detail
`--mtp-verify-depth`	OK	nspec=206 max_chunk=5 worst_argmax_gap=0.000 at=-1
`--dspark-speculative-block`	OK	nspec=96 max_chunk=6 worst_argmax_gap=0.000 at=-1

The worst_argmax_gap invariant holds at 0.000 on both: every committed speculative token is exactly the target model's argmax on replay.

Generation tok/s by prompt type (`--nothink -n 128`)

Prompt	baseline	legacy_mtp	dspark_b2	dspark_adaptive	dspark_b5	Best DSpark vs baseline
json_repetitive	28.40	31.44	31.54	37.44	38.10	+34.2% (b5)
md_table	31.53	31.57	30.55	33.75	31.91	+7.0% (adaptive)
code_gen	31.35	31.42	30.16	28.81	30.00	-3.6% to 0% (wash)
prose	31.31	31.29	26.81	21.68	21.21	-31% to -32% (regression)

This confirms @lobanov's original acceptance-rate problem (issue #468: "average acceptance 1.5-2 tokens, 5% slowdown") is resolved by this branch's combination of the prefix-checkpoint path (#480) and the B2 correction fix (#482) - we're seeing +34% on repetitive JSON with the same Q4_K drafter class on the same q2 base quant.

temp=0 identicality vs baseline

json/table/code: byte-identical across all 5 configs. Prose diverges on all 3 DSpark configs (legacy_mtp stays identical). Root cause: a near-tied target logit gets flipped by the drafter's slight numeric bias, and greedy decoding cascades the divergence from there. The worst_argmax_gap invariant still holds (every committed token is still argmax-correct against the target model at the moment it's verified) - so this is "target-argmax-replay lossless," not "byte-identical," on ties. Worth calling out explicitly in the README's DSpark section since it currently reads like byte-identical is the guarantee.

Findings

block=5 regresses on prose (-32%): low acceptance means more rejected drafts to verify plus partial-commit replay cost. block=2 nearly halves the loss (-14%) but still trails baseline.
Adaptive escalates to b5 under low acceptance and pays the same penalty as fixed b5 on prose. A ceiling (e.g. cap escalation at b3 unless recent acceptance is consistently full) would likely recover most of the structured-output win without the prose regression - flagging as a possible follow-up, happy to try it if useful.
Prefill throughput is ~3% lower across the board with DSpark loaded (extra resident weights + verify overhead) - minor, not blocking.
Single-sample-per-cell on an M4 Max; didn't control for thermal throttling. The b5 vs b2 prose gap (-32% vs -14%) should be re-run 2-3x before treating it as a stable number rather than noise.
DS4_SPEC_TEMP (B2 stochastic path) not yet exercised - all runs above are temp=0. That's a separate correctness surface (rejection sampling at temp>0) worth a dedicated follow-up pass.

Raw logs/output and a re-runnable bench script are on hand if useful for someone else to reproduce.

On-hardware validation on M4 Max (see PR antirez#502 comment) confirmed adaptive mode's single-step 2->full jump regresses low-acceptance workloads: one incidental full commit at block=2 escalated straight to block=5, and the next low-acceptance cycle then paid the full block=5 partial-commit replay cost before dropping back down (-31% on prose vs baseline). Stage the escalation instead: require 2 consecutive full commits before using the full block, one intermediate step at block=3. Drop back to block=2 immediately on any partial commit, same as before. Also document in the README that DSpark's temp=0 correctness guarantee is "every committed token is the target model's argmax on verify", not byte- identical output to a non-speculative run: near-tied logits can flip which side of the tie a drafter-influenced pass commits, which is expected and not a correctness bug, but was previously undocumented and looked like a stronger guarantee than the code actually provides.

stephenlthorn · 2026-07-05T11:26:56Z

Pushed a follow-up commit addressing the two items from the benchmark comment above:

Staged adaptive escalation (2 → 3 → full, requiring 2 consecutive full commits before using the full block, instead of jumping straight from 2 to full on the first full commit). This targets the prose regression: previously one incidental full commit at block=2 would escalate to block=5, and the very next low-acceptance cycle paid the full block=5 partial-commit replay cost before dropping back down.
README clarification: documented that DSpark's temp=0 guarantee is "every committed token is the target model's argmax at verify time," not byte-identical output to a non-speculative run — near-tied logits can flip which side of a tie gets committed under speculation, which is expected, not a bug.

Not yet re-benchmarked on real hardware — the person with the GPU/model loaded is queued to re-run the prose case (and ideally 2-3x to control for the thermal noise flagged earlier) against this staged version. Will report back with numbers once that's done.

stephenlthorn · 2026-07-05T18:50:29Z

Re-benchmark: staged escalation does not fix the prose regression

Follow-up to the earlier validation comment. Re-ran only the prose case (3x, to rule out thermal noise) against the staged-escalation commit (7ad4ea2).

	tok/s	vs baseline
baseline	31.31	—
adaptive, single-jump (2→full)	21.68	-30.8%
adaptive, staged (2→3→full, this commit)	22.47 (mean of 22.61/22.47/22.32, spread ~1.3%)	-28.2%

Correctness unaffected: --dspark-speculative-block still reports worst_argmax_gap=0.000.

Diagnosis: this was the wrong fix for the problem

The +0.8 t/s improvement is inside measurement noise, not a real fix. Staging the escalation only slows down how fast adaptive reaches block=5, it doesn't change the fact that block=5 (or even block=2) costs more than it saves once you're there:

Prose drafter acceptance at temp=0 is below the break-even point identified in DSpark B2 rejection sampling + adaptive block sizing #482's writeup (~26ms partial-commit replay penalty per rejected token) even at block=2. The staged guard ("2 consecutive full commits before escalating") essentially never fires on prose, so this version just spends longer at block=2 before giving up and trying block=5 anyway - which is why it loses slightly less, not because escalation got smarter.
31.31 t/s baseline is effectively "DSpark off." On this prompt class, drafter verify overhead exceeds what accepted tokens save, independent of block size.

Recommendation

Don't present staged escalation as a prose fix - it isn't one, though it's still worth keeping as protection against future flaky single-commit escalation. Two actual paths forward, and I'd defer to whoever picks this up next on which:

Scope adaptive down: cap DS4_DSPARK_ADAPTIVE at block=3 by default instead of escalating to the full block, and describe it in docs as "block=2..3," not "block=2..5." This trades away the structured-output ceiling (+34% on repetitive JSON was measured at fixed block=5, not adaptive) for a policy that can't regress prose as hard.
Accept the workload-dependence and make it explicit: keep block=5 available but not the adaptive default, and document that DSpark is a net win only on structured/repetitive output at this drafter quality - creative text should run with --mtp-draft unset (or explicitly 2) until a higher-quality drafter (Q8, ~9GB extra resident) is validated to raise prose acceptance above break-even.

Not picking one of these myself since it's a product tradeoff (peak structured-output speedup vs. worst-case prose safety) rather than a bug fix. Happy to implement whichever is preferred.

OPS-NeoRetro

If you have the time, please implement DSpark for CUDA, ROCm and the CPU debug path. This is a Metal-only PR. You also did some improvements to DS4, but for this PR, this is a good start for you. Please rename this to "Metal: Reconcile DSpark runtime (#480) + B2 rejection sampling/adaptive block sizing (#482)".

OPS-NeoRetro · 2026-07-09T05:55:42Z

+extern "C" int ds4_gpu_attention_decode_raw_batch_heads_noncausal_tensor(
+        ds4_gpu_tensor       *heads,
+        const void             *model_map,
+        uint64_t                model_size,
+        uint64_t                sinks_offset,
+        const ds4_gpu_tensor *q,
+        const ds4_gpu_tensor *raw_kv,
+        uint32_t                n_tokens,
+        uint32_t                n_raw,
+        uint32_t                raw_cap,
+        uint32_t                raw_start,
+        uint32_t                n_head,
+        uint32_t                head_dim) {
+    (void)heads; (void)model_map; (void)model_size; (void)sinks_offset;
+    (void)q; (void)raw_kv; (void)n_tokens; (void)n_raw; (void)raw_cap;
+    (void)raw_start; (void)n_head; (void)head_dim;
+    return 0;
+}
+


Why did you write a stub?

OPS-NeoRetro · 2026-07-09T05:56:35Z

+extern "C" int ds4_gpu_attention_decode_raw_batch_heads_noncausal_tensor(
+        ds4_gpu_tensor       *heads,
+        const void             *model_map,
+        uint64_t                model_size,
+        uint64_t                sinks_offset,
+        const ds4_gpu_tensor *q,
+        const ds4_gpu_tensor *raw_kv,
+        uint32_t                n_tokens,
+        uint32_t                n_raw,
+        uint32_t                raw_cap,
+        uint32_t                raw_start,
+        uint32_t                n_head,
+        uint32_t                head_dim) {
+    (void)heads; (void)model_map; (void)model_size; (void)sinks_offset;
+    (void)q; (void)raw_kv; (void)n_tokens; (void)n_raw; (void)raw_cap;
+    (void)raw_start; (void)n_head; (void)head_dim;
+    return 0;
+}
+


Why did you also write a stub here?

OPS-NeoRetro · 2026-07-09T06:09:44Z

+    if (fread(magic, 1, sizeof(magic), fp) != sizeof(magic) || memcmp(magic, "GGUF", 4) != 0) {
+        die("bad GGUF self-test file");
+    }
+    (void)read_u32_le_fp(fp, "GGUF version");


Don't forget to add a GGUF version check in deepseek4-quantize

OPS-NeoRetro · 2026-07-09T10:48:28Z

+    default:
+        return "speculative draft disabled";
+    }
+}


Don't forget to end your files with a newline

OPS-NeoRetro · 2026-07-09T10:48:43Z

+
+const char *ds4_dspark_spec_gate_reason(ds4_dspark_spec_gate gate);
+
+#endif


Don't forget to end your new files with a newline

OPS-NeoRetro · 2026-07-09T10:53:06Z

+#ifndef DS4_GIT_COMMIT
+#define DS4_GIT_COMMIT "unknown"
+#endif


Good addition to DS4

OPS-NeoRetro · 2026-07-09T10:54:35Z

+    (void)e;
+    (void)dataset_path;
+    (void)output_dir;
+    (void)target_model_name_or_path;
+    (void)chat_template;
+    (void)ctx_size;
+    (void)max_prompts;
+    (void)max_tokens;
+    fprintf(stderr, "ds4: DSpark target cache export requires a graph backend build\n");
+    return 1;


Please add a CPU path

gsrunion · 2026-07-09T11:53:01Z

ROCm/gfx1151 port + result: DSpark runs correctly on ROCm but is net-negative on bandwidth-bound Strix Halo — the batched verify doesn't pay off here.

Ported the block-draft path to the ROCm backend and ran it on DeepSeek-V4-Flash IQ2XXS (81GB) on a Strix Halo box (Radeon 8060S, 128GB, ROCm 7.2.4). Drafter converted from deepseek-ai/DeepSeek-V4-Flash-DSpark shards 46–48 via deepseek4-quantize --dspark-only → 11.5GB Markov DSpark GGUF (kind=dspark, draft=5, runtime_mtp=yes).

Port (happy to open as a PR if useful): the only backend gap was ds4_gpu_attention_decode_raw_batch_heads_noncausal_tensor (stubbed on both CUDA and ROCm). Implemented it on ROCm by adding a noncausal flag to attention_decode_mixed_kernel (extends the existing single_all all-visible mode → every query in the block sees all raw KV, visible_comp=n_comp), threaded through attention_decode_batch_launch, and forced the general kernel for the noncausal path (the head8-online fast paths bake causal bounds). Plus the ds4_mtp_draft_runtime_supported gate, which hard-rejects the dspark family off Metal.

Correctness: good. Temp-0 output byte-identical to plain; acceptance healthy (drafted=5 committed=5 on repetitive spans). So the kernel is right.

Performance: net-negative on this hardware. Agent/code-edit prompt, temp 0, matched:

config	decode t/s
plain (no spec)	15.3
DSpark (draft 5 / 8 / adaptive)	11.6–12.3
prompt-lookup (#396)	18.7

Prose is worse (7–9 t/s). Draft-depth sweep (repetitive prompt): verify ≈ 253 ms @draft2, 330 ms @Draft5, 329 ms @draft8 — the block verify of 5–8 tokens costs ~5× a single-token decode (~65 ms), while committed plateaus at ~4.4 tokens regardless of depth. So we pay ~5 decodes' worth of bandwidth to commit ~4.4 tokens ⇒ slower than plain.

Why (hypothesis): on UMA Strix Halo, batch-1 decode is memory-bandwidth-bound, and the routed MoE experts don't amortize across a verify block — each of the 5 block tokens routes to its own 6/256 experts, so expert bytes read scale ~linearly with block size even in a batched verify (only the dense/attention Q8 stream, ~6.5 GB, amortizes). Once dense is amortized, experts dominate and the block verify approaches the sequential cost. DSpark's 60–85% figures are on compute-bound datacenter GPUs where batched verify is nearly free; that condition doesn't hold on a 212 GB/s iGPU. (Consistent with the per-kernel bandwidth accounting I posted in #16.)

Takeaway for the ROCm/UMA target: prompt-lookup (#396) beats DSpark here because it has zero drafter cost and near-100% acceptance on echo-heavy work. DSpark's ceiling on this class of hardware is bounded by expert-gather bandwidth in the verify, not by the drafter or the port. Full logs / the noncausal patch / conversion recipe available if you want them — and this box (your QA hardware twin) is free nights US-Eastern for any block-verify batching experiments.

stephenlthorn · 2026-07-09T12:22:14Z

@OPS-NeoRetro Renamed as suggested - it is indeed a Metal-only PR, and the title now says so.

On implementing CUDA/ROCm/CPU here, I'd argue against expanding this PR's scope, and @gsrunion's excellent data two comments up is most of the reason:

ROCm: the port already exists (gsrunion's), it's correct, and it's net-negative on Strix Halo - the project's primary ROCm target - because routed-MoE expert reads don't amortize across the verify block on bandwidth-bound UMA hardware. Enabling the runtime there today would ship a slowdown. Their noncausal-kernel patch has standalone value and deserves to be its own PR (with the runtime gate kept off, or opt-in) rather than folded into this one.
CUDA: the primary CUDA target is DGX Spark / GB10, which is also unified-memory at ~273 GB/s - the same bandwidth-bound regime as Strix Halo, so the same economics very likely apply. I also can't test CUDA changes from this machine, and per AGENT.md/CONTRIBUTING.md untested backend changes shouldn't ship.
CPU: per AGENT.md the CPU path is reference/debug only (and large CPU runs can trigger a macOS kernel VM bug), so a speculative-decode fast path there has no production value.

The current state - Metal runtime + explicit ds4_mtp_draft_runtime_supported gate with non-Metal backends loading DSpark GGUFs inertly, documented in the README - keeps every backend correct while limiting the speedup claim to the one backend where it's been validated. That gating pattern came from #480 and I think it's the right call until someone demonstrates a backend where the verify economics work (which per gsrunion is a batching/bandwidth question, not a porting question).

@gsrunion this is a great writeup - the expert-gather-doesn't-amortize analysis cleanly explains why the datacenter DSpark numbers don't transfer to UMA. Please do open the noncausal kernel port as its own PR; correctness-verified backend parity is valuable even with the runtime gated off, and it gives anyone experimenting on Strix Halo a working baseline. The prompt-lookup comparison (#396 winning on echo-heavy work with zero drafter cost) is also a useful datapoint for the docs.

gsrunion · 2026-07-09T16:41:22Z

Thanks — agreed on all three points, and keeping #502 Metal-only with the #480 gate is the right call. The verify-economics ceiling is a batching/bandwidth question, not a porting one, so gating off-by-default until someone demonstrates a backend where a block verify amortizes is exactly right.

I'll open the noncausal ROCm kernel as its own PR as you suggested. It fills the ROCm stub of ds4_gpu_attention_decode_raw_batch_heads_noncausal_tensor with a real implementation (a noncausal flag on attention_decode_mixed_kernel, extending the existing single_all all-visible path, threaded through attention_decode_batch_launch), keeps the runtime gated off by default with a DS4_DSPARK_ROCM=1 opt-in for experimenters, and is validated on gfx1151 (Radeon 8060S / Strix Halo): temp-0 output byte-identical to plain decode, drafted=5 committed=5 acceptance.

Since it depends on this PR's declaration + Metal impl + ROCm stub, I'll base it so it's a clean 3-file delta once #502 lands (rather than stacking a noisy diff now). Happy to keep the DS4_DSPARK_ROCM opt-in or drop to pure gated-off parity — whichever you prefer.

audreyt and others added 9 commits July 5, 2026 04:16

Add guarded DSpark loader plumbing

ddcc812

Add DeepSpec nonseq draft-head schema gate

06ab968

Add DeepSpec DSpark scaffold

2764efa

Add DSpark speculative draft runtime

8d005ab

Make DSpark nonseq runtime ready

7a51e42

Commit DSpark partial accepts from prefix checkpoints

75fdeee

audreyt mentioned this pull request Jul 5, 2026

Add DSpark speculative draft runtime #480

Closed

OPS-NeoRetro suggested changes Jul 9, 2026

View reviewed changes

stephenlthorn changed the title ~~Reconcile DSpark runtime (#480) + B2 rejection sampling/adaptive block sizing (#482)~~ Metal: Reconcile DSpark runtime (#480) + B2 rejection sampling/adaptive block sizing (#482) Jul 9, 2026


		const char *ds4_dspark_spec_gate_reason(ds4_dspark_spec_gate gate);

		#endif No newline at end of file

Conversation

stephenlthorn commented Jul 5, 2026

Summary

Fix included

Testing

Uh oh!

stephenlthorn commented Jul 5, 2026

On-hardware validation (M4 Max 128GB, q2-imatrix base + Q4_K DSpark drafter)

Correctness: PASS

Generation tok/s by prompt type (--nothink -n 128)

temp=0 identicality vs baseline

Findings

Uh oh!

stephenlthorn commented Jul 5, 2026

Uh oh!

stephenlthorn commented Jul 5, 2026

Re-benchmark: staged escalation does not fix the prose regression

Diagnosis: this was the wrong fix for the problem

Recommendation

Uh oh!

OPS-NeoRetro left a comment

Choose a reason for hiding this comment

Uh oh!

OPS-NeoRetro Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

OPS-NeoRetro Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

OPS-NeoRetro Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

OPS-NeoRetro Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

OPS-NeoRetro Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

OPS-NeoRetro Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

OPS-NeoRetro Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

gsrunion commented Jul 9, 2026

Uh oh!

stephenlthorn commented Jul 9, 2026

Uh oh!

gsrunion commented Jul 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Generation tok/s by prompt type (`--nothink -n 128`)