DSpark: replay-free partial accepts + encoder-batched captures (+40% greedy, byte-lossless)#590
DSpark: replay-free partial accepts + encoder-batched captures (+40% greedy, byte-lossless)#590gilbert-barajas wants to merge 2 commits into
Conversation
The batched verifier steps the compressor token by token, so the mutable frontier after every verify row is captured in-stream into per-depth slot buffers (a few small blits per layer), generalizing the MTP prefix1 mechanism to whole blocks. A partial accept of k drafts now commits by copying slot k-1 back over the live frontiers and adopting the batch's own logits row, instead of restoring the pre-block snapshot and re-decoding each accepted token through the full target model (~24 ms per accepted token on an M5 Max). The restore+replay path survives as the fallback and as the TP-lockstep route; DS4_DSPARK_NOREPLAY=0 pins the old behavior. The commit path also retargets the DSpark drafter capture at the committed depth: batch row k of the verifier's hidden-state capture becomes the new anchor hc and the capture window is trimmed to the accepted prefix, so the proposer's initial-cache seeding stays healthy after partial accepts (the old path left the batch capture stale, costing the following cycle its draft). The aligned-chunk bulk compressor path has no per-row boundary to capture at, so the encoder takes the per-token path only while checkpoints are requested; prefill and MTP are unaffected. Measured overhead of the in-stream captures is ~4 ms per verify against ~24 ms per replayed token. Gates on an M5 Max (Q2_K DeepSeek-V4-Flash target, DSpark drafter, greedy, 512-token code generation, scheduler off, medians of 3 cooled runs): - conf 0.5: 50.4 -> 52.8 t/s (replay 846 ms -> 0; checkpoint commits cost 5 ms per run) - conf 0.3: 45.6 -> 53.4 t/s; replay-free makes lower confidence the new best operating point (avg_accept 4.11) - dspark-verify-depth: worst_argmax_gap=0.000, including a conf-0.5 run with partial accepts flowing through the new path - token-stream parity vs plain greedy on 6 prompts: the new path diverges on 2/6 at batch-verifier near-ties vs 3/6 for the old path (the acceptance fixture's redis case already fails byte-identity on this quant with the unmodified base -- a pre-existing near-tie class, not introduced here)
The per-depth frontier checkpoints ride the batched verify interleaved with the per-row compressor dispatches, and each ds4_gpu_tensor_copy there ended the cached batch compute encoder, opened a blit encoder for one copy, and forced the next dispatch to reopen the compute encoder -- four encoder transitions per captured tensor, several hundred per verify. ds4_gpu_tensor_copy_f32_inline encodes the same copy as a cpy_f32_f32 compute dispatch instead, which stays inside the open batch encoder (the buffers use default hazard tracking, the same ordering the row-to-row compressor chain already depends on). Metal-only in effect: CUDA/ROCm alias it to the plain device copy, and it falls back to the blit path outside a batch command buffer. Snapshot and commit copies keep the blit path -- they run in their own command buffers with no encoder churn. The committed token streams are byte-identical to the blit encoding at conf 0.3 and 0.5 (512-token generation) and across the 6-prompt parity set; dspark-verify-depth still reports worst_argmax_gap=0.000, including a conf-0.5 run with partial accepts through the checkpoint path. Measured on an M5 Max (Q2_K DeepSeek-V4-Flash target, DSpark drafter, greedy 512-token code generation, scheduler off, medians of 3 cooled runs): - conf 0.3: 53.4 -> 55.3 t/s - conf 0.5: 52.8 -> 54.0 t/s - verify phase at conf 0.5: 6048 -> 5880 ms per run against 5630 ms with captures disabled -- the checkpoint overhead drops from ~4 ms to ~2 ms per verify, now mostly the copies' own memory traffic.
83326be to
567e798
Compare
|
Independent datapoint from another M5 Max (128 GB, Metal, mixed Q2+Q4
Baseline sanity: base numbers here match yours (485 t/s prefill / ~39-40 t/s gen @2k — full sweep in #587). Happy to re-run the acceptance fixture and the A/B on this branch if useful. |
|
Appreciate the independent M5 Max datapoint — that ~24 ms/token replay match is exactly the confirmation this needed. Same thing I hit before writing the fix: baseline DSpark net-negative despite a high accept rate. And you've got the tie-break right — that greedy flip in Please do run the A/B on the branch — that's the number that matters. #590's whole job is killing that replay cost: partial accepts no longer replay, so your If it's easy: |
|
Ran the branch (
Takeaway from this datapoint: on mixed quants the default confidence gate may leave most of the branch's benefit on the table — possibly worth a note in the DSpark README section, or a quant-aware default down the line. |
2 participants
Two coupled improvements to the DSpark speculative-decode path (
fc9efd1), Metal, no backend regressions.What
replay=0.000).ds4_gpu_tensor_copy_f32_inline) rather than separate blit encoders, which were tearing down and rebuilding the encoder ~350–500× per verify. Bounds-checked blit fallback retained.Gains
Metal, DeepSeek-V4-Flash Q2_K, full residency, server decode-phase, cooled:
replay=0.000, accept_rate 94.5%confirms partial accepts commit through the checkpoint path.Testing (per CONTRIBUTING.md)
Machine: Apple M5 Max, 128 GB, Metal.
make clean && make.commdiff = 0 new); all are the macOS-27didModifyRange:SDK deprecation, not from this change.make cpubuilds clean.make dspark-verify-depth→worst_argmax_gap = 0.000— every committed greedy token (partial accepts included) is the exact plain-decode argmax../ds4_test --metal-kernels→ OK; compressor/router exactness allmax_ulp = 0/0/0.--dspark-off decode is byte-for-byte the same code. Base decode measured unchanged (server baseline 40.18 t/s;ds4-benchbase 485 prefill / 39 gen @ ctx 2048).--logprob-vectorswas skipped — the official continuation.vecisn't in the repo; glad to run it against the official vectors. A cleands4-benchbefore/after CSV is best captured on a thermally-settled box (this machine was saturated during testing).Notes for the reviewer
net_savedin the stats line reads understated with the scheduler disabled: the replay-free path skips the re-decode whose timing fedsaved_ms, sosaved_msno longer accrues in that mode andnet_savedshows a misleading negative. Wall-clock throughput is the ground truth (+40%). Glad to creditsaved_msfrom the verify-pass timing in a follow-up._inlinecopy + fallback).ds4_cuda.cugets a 9-lineextern "C"alias to the existing device copy. CPU / GLM / distributed / tensor-parallel sessions fall back to a plain one-token decode. C99 throughout the core.Diff: +363 / −10 across 4 files (
ds4.c,ds4_metal.m,ds4_gpu.h,ds4_cuda.cu).A follow-up PR adds opt-in
DS4_DSPARK_SAMPLE=1sampling-mode speculation at temperature > 0 via the rejection rule — distributionally lossless — and will be sent after this one merges.