Reconcile DSpark runtime (#480) + B2 rejection sampling/adaptive block sizing (#482)

[stephenlthorn]

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.

[stephenlthorn]

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.

[stephenlthorn]

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

1. 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.
2. 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]

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:

1. 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.
2. 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.