DSpark under --temp > 0: a decode/verify consistency bug, a working stochastic-verify implementation with no net win on M5 Max, and an unresolved kernel issue
- Lingua principale
- C
- Stelle
- 22.3k
- Fork
- 2.1k
- Merge medio
- 1g 3h
- PR unite (30g)
- 4
Descrizione
Hi, I spent some time (with Claude Code as pair-programmer) investigating why DSpark's speculative decoding currently does nothing under real sampling (`--temp > 0`), and whether the batched-verify infrastructure could be made to pay off if fixed. Sharing what came out of it in case any of it is useful. I ended up reverting all the code changes locally (didn't want to ship something I couldn't fully validate), so this is a text report, no PR/diff attached — happy to share the actual commits (currently only on a local, un-pushed branch) if useful.
**Setup**: M5 Max, Metal backend, DeepSeek-V4-Flash (43 layers, 256 experts / 6 active, MLA+SWA), `--warm-weights`, single GPU, `--batched-session` off.
## 1. A real correctness bug: indexer sparse-threshold inconsistency between decode and batched verify
`metal_graph_encode_decode_layer` (single-token decode) gates the indexer query chain behind `metal_graph_decode_indexer_sparse_threshold()` (1024 compressed rows by default) before falling back to dense mixed attention -- the comment there explains this is a genuine perf tradeoff (indexer setup dominates below ~2K context).
The batched path (`metal_graph_encode_layer_attention_batch`, used by both real prompt prefill and DSpark's `metal_graph_verify_suffix_tops`) instead gates purely on `n_comp > DS4_N_INDEXER_TOP_K` (512). So for `n_comp` in `(512, 1024]`, ordinary decode attends *densely* over all compressed rows, but DSpark's verifier attends over the *indexed top-512* -- a different candidate set than what decode would actually produce for the same position. Not just a perf discrepancy: the verifier isn't exactly reproducing what decode does in that window.
I only fixed this for small (`n_tokens <= 8`, i.e. actual verify-shaped) batches, gated so real large prefill chunks keep the original threshold -- decode's amortization argument for going sparse earlier doesn't obviously transfer to genuine multi-thousand-token prefill, and I didn't want to touch that without separately validating it.
## 2. Stochastic (temp>0) verify: implemented, correct, but no net win on this hardware
`ds4_session_eval_dspark_speculative_argmax` only ever worked for greedy decoding -- it compares drafted tokens to the target's argmax via exact match, and the CLI/server call sites gate the whole speculative path to `--temp <= 0` accordingly. Under real production sampling this path is currently dead code.
I implemented the standard Leviathan/Chen stochastic accept/reject (accept a drafted token with probability `min(1, p(x)/q(x))`, residual-resample from `max(0, p-q)/Z` on rejection) and validated the residual-resample math with a 300k-trial statistical test confirming the marginal distribution of the emitted token matches the target distribution exactly, not just approximately -- happy to share that test file.
Three concrete perf fixes along the way (each verified by direct measurement, not assumed):
- reused the existing `spec_frontier_commit_prefix` fast-commit path for stochastic partial-accepts instead of falling back to a full sequential replay on every rejection
- routed `draft_n==1` cycles through a plain decode instead of the batched verifier -- same required output (post-token logits), ~1/3 the measured cost, since the batched verifier pays for a full 43-layer batch dispatch even at `n_tokens==1`
- the indexer threshold fix above
Net effect at the shipped default config (`--dspark-confidence 0.6`, scheduler on): went from ~35-45% slower than not using DSpark at all, down to ~4-6% slower. Tested across confidence 0.01/0.3/0.4/0.5/0.6/0.7/0.8 and scheduler on/off; 0.6 + scheduler-on (today's default) remained the best configuration found. Did not reach a net positive gain on this specific hardware/model combination.
## 3. Attempted fusion kernel: correct in isolation, broken live, root cause not found
Decode has ~8 hand-fused kernels (e.g. `kernel_dsv4_qkv_pair_quad_compressor_store_q8_0`, doing attn_q_a + attn_kv + attn/indexer-compressor kv+gate in one dispatch) that the generic batch path lacks, presumably because decode is always exactly one row and never needed batching. I wrote a row-batched sibling (attn_q_a + attn_kv + indexer-compressor kv/gate only -- DSpark's batch path never needs attn_compressor, that's zero_prefix-only) and validated it against the currently-shipped unfused reference on real model weights under `MTL_SHADER_VALIDATION=1 MTL_DEBUG_LAYER=1`: 126/126 exact matches across all compressed layers for `n_tokens` in {1,2,5,6,8} (found and fixed one real bug this way -- a missing per-branch output-stride override, invisible in the original single-row kernel since its row index is always 0).
Wired into `metal_graph_encode_layer_attention_batch` behind an explicit opt-in flag (set only by the DSpark verify call site, cleared right after), it produces wrong generated text -- and does so *reproducibly differently* depending on whether `MTL_SHADER_VALIDATION` is enabled: correct without it, garbled with it, same seed, deterministic either way. That's consistent with a synchronization/integration issue somewhere in how the shared per-layer command buffer/encoder is used, but I wasn't able to pin it down (ruled out an insufficient trigger heuristic, and `g_batch_encoder_concurrent`'s FFN-overlap window, as causes). Left it disabled by default. If anyone wants to pick this up, an Xcode GPU Frame Capture on the verify call would probably show the actual dispatch ordering directly and make this fast to root-cause.
---
Happy to answer questions or share the actual commits/diffs if that would help anyone dig into any of this further.
Guida per i contributori
Apri la guida per i contributori
Valutazione
Questa issue non è ancora stata valutata.