DSpark speculative decode breaks --temp 0 greedy identity (batch-verify vs single-token KV numerics)

[nexus-cw]

Summary

DSpark speculative decode (--mtp <support.gguf> --dspark) does not preserve
byte-identical greedy (--temp 0) output vs. non-speculative decode on longer
generations, despite the accept/verify logic itself correctly matching argmax
(independently code-reviewed).

Root cause

The compressed-KV "compressor frontier" state for every accepted token
(including the always-accepted bonus token) is computed by the verify
batch's own multi-token GEMM, ds4_gpu_matmul_f16_pair_tensor
(ds4.c:27454 and the equivalent decode-batch site around ds4.c:28685-28753
on current main). This is a numerically distinct code path from ordinary
single-token greedy decode, which computes the identical logical quantity via
a different, fused single-vector kernel,
ds4_gpu_matmul_f16_pair_compressor_store_tensor (ds4.c:22993-23011).

These two kernels are not required to be, and empirically are not,
bit-identical: a GEMM computing one row of a [n_tokens x width] batched
matmul and a fused single-vector kernel computing the same logical 1 x width row use different tiling/accumulation order internally (standard
FP16/BF16 GEMM non-associativity across kernel implementations). So every
token committed through DSpark's accept path gets a compressor-frontier
state that is numerically close to, but not identical to, what sequential
single-token decode would have produced for the exact same token. This state
feeds every subsequent layer's attention, so the drift accumulates silently
until it flips an argmax tie -- consistent with the divergence we observed
being a late (tens of tokens in), paraphrase-level flip rather than an
immediate/garbled break.

This is why spec_frontier_snapshot/_restore (rollback/restore) is not the
culprit: it faithfully snapshots and restores exactly what the batch-verify
pass computed. The bug is upstream of it, in what values get written into
the frontier tensors by the accept-commit path in the first place.
Specifically, ds4_session_eval_dspark_speculative_argmax's accept-commit
logic had a fast path for full accepts (commit_drafts == draft_n) that
committed the batch-verify-derived KV/compressor-frontier state directly,
instead of re-deriving it via the fused single-token kernel the way plain
decode does.

Reproduction

Resident IQ2XXS quant of DeepSeek-V4-Flash + a DSpark support GGUF
(lineage-matched drafter, dspark.target_layer_ids=[40,41,42]), on CUDA
(NVIDIA GB10 / sm_121):

./ds4 --cuda -m <IQ2XXS-model.gguf> --temp 0 -n 400 \
  -p "Explain the TCP three-way handshake."
./ds4 --cuda -m <IQ2XXS-model.gguf> --mtp <DSpark-support.gguf> --dspark \
  --temp 0 -n 400 -p "Explain the TCP three-way handshake."

We confirmed this directly against unpatched main (54b36ed, current tip at
the time of filing): the drafter-enabled run diverges from the no-drafter
baseline within the first couple of sentences of the response (word/phrasing
choices differ from there on, e.g. "...agree on initial sequence numbers and
parameters..." vs. "...agree on initial sequence numbers and synchronize
their connection state...") -- a near-tied paraphrase-level divergence, not
garbled output, consistent with an FP-drift-driven argmax flip rather than a
logic bug. This is a --temp 0 contract violation, since both runs are
greedy decodes of the same model/prompt and should be byte-identical. (Exact
divergence onset is prompt/pairing-dependent; on our downstream fork's build
of this same drafter pairing the same prompt instead diverged ~55-65 tokens
in -- the divergence is a drift effect, not a fixed offset.)

Fix

PR #<n> fixes this. The DSpark accept-commit fast path for full accepts
is removed; every accept (full or partial) now falls through to the
rollback+replay path that already existed for partial accepts, which rolls
the KV/compressor-frontier state back to the pre-verify snapshot and
re-derives each accepted token's state one token at a time through the same
single-token decode kernel plain greedy decode uses
(metal_graph_eval_token_raw_swa). This guarantees the post-accept
state -- and the logits read off it -- are bit-identical to what pure decode
would have produced for the same tokens.

Verified byte-identical drafter-vs-no-drafter output on multiple prompts at
both 400 and 800 generated tokens, including the reproducer above extended
past its original divergence point, on both our downstream fork and this
issue's target branch built from current main.

Cost: this adds one extra fused single-token forward pass per accepted token,
even on full accepts (previously only partial accepts paid this cost). On our
hardware/pairing this makes an already fairly marginal drafter throughput
case measurably worse (see below) -- correctness first, and the PR states
this plainly rather than hiding it.

Supporting measurements (informational, not part of the correctness fix)

These are separate, previously-gathered measurements on the same
hardware/pairing (NVIDIA GB10, resident IQ2XXS DeepSeek-V4-Flash + the DSpark
drafter above), offered as data points rather than requests:

• Throughput, resident model: no-drafter baseline ~16.87 t/s generation.
  Pre-fix drafter-enabled: ~16.07 t/s default confidence (-4.7%), ~10.81 t/s
  forced-accept. Post-fix (this issue's PR), drafter-enabled generation is
  further reduced to roughly 14.0-14.6 t/s across several prompts (-14% to
  -17% vs. no-drafter) -- the correctness fix's replay-per-accepted-token
  cost applied to full accepts as well as partial ones.
• Throughput, streamed model (SSD-streaming, different pairing):
  no-drafter baseline ~5.25 t/s, drafter-enabled ~3.85 t/s default confidence,
  ~1.62 t/s forced-accept.
• Accepted-depth distribution: with default confidence (0.9), acceptance
  depth is heavily concentrated at depth 1 (bonus-token-only accepts), with a
  long tail out to depth 8-9 in occasional bursts, rather than a flat/typical
  distribution. We traced this to the drafter checkpoint's own
  dspark.noise_token_id metadata deliberately seeding the non-bonus draft
  positions' Markov chain with noise/mask tokens rather than chaining true
  per-position argmax embeddings -- this appears to be intentional per the
  checkpoint format (matching the "greedy argmax-only path" framing in the
  original DSpark work this format is based on), not a wiring gap on our
  side, but it does bound how much the drafter can help on this particular
  drafter/target pairing. Flagging in case it's useful context for anyone
  tuning drafter checkpoints or expected acceptance-depth behavior against
  this codebase.

Happy to share full logs/scripts if useful.

🤖 Generated with Claude Code

[nexus-cw]

Opened a fix for this: #659

[nexus-cw]

While assessing rebase impact of #659 against ds4f-mxfp4, we analyzed 74e8f11 ("Commit partial DSpark verifier blocks without replay"). We believe the no-replay commit path there inherits the same greedy-identity mechanism described in this issue. All line refs below are on the current ds4f-mxfp4 tip.

Code path:

• The verify batch projects compressor KV/score for the whole block with the batched matrix-matrix GEMM ds4_gpu_matmul_f16_tensor over n_tokens rows (ds4.c:27473 and ds4.c:27485).
• 74e8f11 adds !g->spec_capture_prefixes to the aligned_chunk condition (ds4.c:27601), forcing the per-token compressor stepping loop — but that loop still consumes the batch-GEMM-projected rows via row views of metal_graph_batch_comp_kv/sc into ds4_gpu_compressor_update_tensor (ds4.c:27724), with emitted rows FP8-quantized at ds4.c:27752. Stepping token-by-token restores the update order of decode, but not its input numerics.
• Plain greedy decode computes the same logical rows via the fused single-vector kernel ds4_gpu_matmul_f16_pair_compressor_store_tensor (ds4.c:22062). These two projection paths differ in tiling/accumulation order — the comment at ds4.c:~26490 in this same tree acknowledges exactly this ("mixing those two accumulation orders changes a few FP8 rounding decisions").
• The prefix snapshots captured at ds4.c:27770 / ds4.c:28066 therefore hold verify-GEMM-derived frontier state, and the new partial-accept block (ds4.c:61391-61417, via spec_frontier_commit_prefix at ds4.c:49739) commits that state directly with no target-model replay.

So on a partial accept, the committed compressor frontiers (and any FP8 rows emitted during the verify block) can differ in low bits from what non-speculative decode would have produced, and later tokens can diverge at --temp 0. Same mechanism as this issue, just reached through the new prefix-slot path instead of the old fast paths.

Repro: the temp-0 A/B from the top of this issue should reproduce it on ds4f-mxfp4 — long greedy generations with --mtp <support.gguf> --dspark vs. plain decode, diff the token streams. Divergence is prompt/length-dependent since it needs an FP8 rounding decision to flip, so longer runs make it deterministic to observe.

This is a static analysis (we can't run Metal here), but the kernel-path argument is the same one that was empirically confirmed on main for #658, so confidence is high.

Also a heads-up: #659's replay approach rewrites the same accept path 74e8f11 touches, so it will conflict with the ds4f-mxfp4 rewrite. Happy to re-derive the fix on top of the new prefix-slot verifier if that's the direction that lands — the prefix-slot machinery would actually pair nicely with a replay (or a fused-kernel re-projection of just the accepted tokens), which could keep partial accepts cheaper than our current full replay.