Skip to content

fix: speculative verify with --mtp-draft > 2 commits wrong tokens#358

Merged
antirez merged 1 commit into
antirez:mainfrom
pandysp:fix-mtp-verify-topk
Jun 9, 2026
Merged

fix: speculative verify with --mtp-draft > 2 commits wrong tokens#358
antirez merged 1 commit into
antirez:mainfrom
pandysp:fix-mtp-verify-topk

Conversation

@pandysp

@pandysp pandysp commented Jun 8, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Summary

In metal_graph_verify_suffix_tops, the multi-row branch calls ds4_gpu_indexer_topk_tensor with n_tokens and top_k swapped: it asks for the top-top_rows tokens of a single row instead of the top-1 of each of top_rows rows. The signature is (selected, scores, n_comp, n_tokens, top_k), so row_tops[i>0] come back as row 0's runner-ups. At draft depth > 2 the verifier then accepts or rejects drafts against the wrong rows and commits tokens the model never produced.

It only fires above the shipped depth: at --mtp-draft 2, top_rows == 1 takes the dedicated argmax path just above this branch, so the swapped call is never reached. So this is latent (it does not affect default output) but raising the draft depth gives silently wrong tokens (a spurious EOS or divergence), and nothing in the suite covered the multi-row verify path.

Fix

Swap the last two arguments back (signature is (selected, scores, n_comp, n_tokens, top_k)):

- ds4_gpu_indexer_topk_tensor(g->comp_selected, g->spec_logits, DS4_N_VOCAB, 1, top_rows)
+ ds4_gpu_indexer_topk_tensor(g->comp_selected, g->spec_logits, DS4_N_VOCAB, top_rows, 1)

Test

Adds tests/ds4_test.c --mtp-verify-depth. It runs greedy speculative decode at draft 4 over a verbatim-copy task, then teacher-forces the committed tokens back through plain decode and requires each to be a (near-)argmax — the invariant speculative verify must preserve. Unlike comparing whole token streams, that tolerates the benign tie divergences of near-greedy speculation.

The MTP head is not in the default test model set, so it self-skips unless DS4_TEST_MTP points at an MTP GGUF:

make ds4_test
DS4_TEST_MODEL=<base.gguf> DS4_TEST_MTP=<mtp.gguf> ./ds4_test --mtp-verify-depth

Without the fix a committed token sits ~21 logits below the argmax (worst gap 20.96 at token 145); with the fix every committed token is the argmax (gap 0.00).

Notes

  • Verified on M4 Max / Metal, DeepSeek-V4-Flash q2-q4 imatrix base + MTP-Q4K-Q8_0 head. Ran alongside --server --metal-tensor-equivalence --long-context; all pass.
  • The call site is backend-shared, so CUDA is affected and fixed identically, but I have not verified it on a CUDA machine.
  • The other four ds4_gpu_indexer_topk_tensor call sites use the correct argument order; this was the only swapped one.

@pandysp
fix: speculative verify with --mtp-draft > 2 commits wrong tokens
a1661e0 
metal_graph_verify_suffix_tops() asked ds4_gpu_indexer_topk_tensor() for the
top-`top_rows` tokens of a single logits row (n_tokens=1, top_k=top_rows)
instead of the top-1 of each of `top_rows` rows. The signature is
(selected, scores, n_comp, n_tokens, top_k), so row_tops[i>0] came back as
row 0 runner-ups and the verifier accepted or rejected drafts against the
wrong rows, committing tokens the model never produced. Visible only at draft
depth > 2: at depth 2 top_rows == 1 and the code takes the dedicated argmax
path above this branch, so the swapped call is never reached.

Adds tests/ds4_test.c --mtp-verify-depth: runs greedy speculative decode at
draft 4 over a verbatim-copy task, then teacher-forces the committed tokens
back through plain decode and asserts each is a (near-)argmax at its position.
That is the invariant speculative verify must preserve, and unlike comparing
whole token streams it tolerates the benign tie divergences of near-greedy
speculation. Self-skips unless DS4_TEST_MTP points at an MTP GGUF. Without the
fix a committed token sits ~21 logits below the argmax (worst gap 20.96 at
token 145); with the fix every committed token is the argmax (gap 0.00).

The shared test engine loads the MTP head only when DS4_TEST_MTP is set, and
only on the fast engine, so the default suite and the quality engine are
unaffected.

Verified on M4 Max / Metal, DeepSeek-V4-Flash q2-q4 imatrix base + MTP-Q4K-Q8_0
head:
  make ds4_test
  DS4_TEST_MODEL=<base> DS4_TEST_MTP=<mtp> ./ds4_test --mtp-verify-depth
Ran alongside --server --metal-tensor-equivalence --long-context: all pass.
The swapped call site is backend-shared (ds4.c), so CUDA is affected and fixed
identically, but this was not verified on a CUDA machine.
@antirez
antirez merged commit e76c4ec into antirez:main Jun 9, 2026
@antirez

antirez commented Jun 9, 2026

Copy link
Copy Markdown
Owner

Thank you, merged.

@pandysp
pandysp deleted the fix-mtp-verify-topk branch June 9, 2026 17:27
@pandysp pandysp mentioned this pull request Jun 9, 2026
Open
TrevorS added a commit to TrevorS/ds4 that referenced this pull request Jun 17, 2026
@TrevorS
fork(gb10+mtp): GB10 speed/accuracy keep-set + MTP + batch-decode on …
ad44ba1 
…upstream

Single consolidated fork commit over antirez/ds4, tuned for DGX Spark (GB10,
sm_121, unified memory, ~236 GB/s). Collapses the prior land/steering/python
stack plus this session's batch-decode + managed-KV work into one commit.

Engine (speed/accuracy, GB10):
- MTP speculative decode: combined-forward verify, adaptive K=1/2 cascade
  (opt-in), deterministic verify default + DS4_CUDA_FAST_VERIFY fast path.
- DS4_BATCH_DECODE (opt-in): route plain ds4_session_eval through the n=1
  combined-verify batch forward (verify kernel set: expert-tile MoE, heads8
  attention, cuBLAS f16). +21% plain decode @4k and @65k, token-identical in
  greedy + seeded-sampled testing (near-tie caveat, same as MTP combined).
- managed-KV policy fix: ds4_gpu_should_use_managed_kv_cache always consults
  cudaMemGetInfo free memory instead of the context_bytes < 8 GiB blind spot
  that OOM-hung the box at 131k ctx with the model resident.
- MoE expert-tile/rowspan kernels, atomic-free MoE down, Q8 residency cache,
  HBM hot-span device-copy with the _exps. filter, host-register, window attn.

Product:
- Multi-profile directional-steering cache + /v1/steering server admin API.
- In-process Python ctypes bindings (libds4.so) exposing engine/session/steering.

Validation + tooling:
- CUDA correctness gate (tensor-equivalence, cuda-ppl, cpu-cuda-ppl) + primer.
- gamut perf suite, longctx sweep harness, MTP retrain harness (v2 head),
  eugr N-runs mean+/-std methodology, ds4-bench --batch-cost/--batch-check,
  cascade A/B harness, thermal/memory-watchdog soak driver.

Gates green: token-diff identity (greedy+sampled, 4k+65k), mtp-correctness
worst_rms=0, mtp-selfconsistency maxabs=0, agent 86.5% accept / 21.9 t/s, no
BOS spam.

Rebased onto upstream cafc134 ("Fix server const warning"; +25 commits over
base 59d9bc7: CUDA/ROCm SSD streaming + mixed-precision routed experts + loader
refactor, MTP verify-topk fix antirez#358, agent raw-mode/edit fixes, server hardening).
Conflict resolutions (4 files, 6 hunks):
- ds4.c verify MoE: kept our verify_ordered_moe block, composed upstream's new
  DS4_ROCM_BUILD async-finish guard ahead of it.
- ds4.c verify top-k (antirez#358): kept our per-row argmax loop over upstream's
  indexer_topk arg-swap fix; ours is correct + faster (avoids the ~17ms
  single-thread O(n_vocab) scan). Validated by upstream's own antirez#358 regression
  (test_mtp_verify_depth): worst_argmax_gap=0.000, max_chunk=5.
- ds4_cuda.cu / tests / Makefile: union merges (both sides' additions kept).
Post-rebase re-validation on GB10: build clean; fork-fixes intact (device-copy
8.20 GiB via _exps. filter, host-register 80.76 GiB, 344 dense f16 prewarms,
managed-KV fix firing under tight free-mem); mtp-verify-depth/mtp-correctness
(worst_rms=0)/mtp-selfconsistency (maxabs=0, bit-exact) all OK.
TrevorS added a commit to TrevorS/ds4 that referenced this pull request Jun 17, 2026
@TrevorS
fork(gb10+mtp): GB10 speed/accuracy keep-set + MTP + batch-decode on …
2daa202 
…upstream

Single consolidated fork commit over antirez/ds4, tuned for DGX Spark (GB10,
sm_121, unified memory, ~236 GB/s). Collapses the prior land/steering/python
stack plus this session's batch-decode + managed-KV work into one commit.

Engine (speed/accuracy, GB10):
- MTP speculative decode: combined-forward verify, adaptive K=1/2 cascade
  (opt-in), deterministic verify default + DS4_CUDA_FAST_VERIFY fast path.
- DS4_BATCH_DECODE (opt-in): route plain ds4_session_eval through the n=1
  combined-verify batch forward (verify kernel set: expert-tile MoE, heads8
  attention, cuBLAS f16). +21% plain decode @4k and @65k, token-identical in
  greedy + seeded-sampled testing (near-tie caveat, same as MTP combined).
- managed-KV policy fix: ds4_gpu_should_use_managed_kv_cache always consults
  cudaMemGetInfo free memory instead of the context_bytes < 8 GiB blind spot
  that OOM-hung the box at 131k ctx with the model resident.
- MoE expert-tile/rowspan kernels, atomic-free MoE down, Q8 residency cache,
  HBM hot-span device-copy with the _exps. filter, host-register, window attn.

Product:
- Multi-profile directional-steering cache + /v1/steering server admin API.
- In-process Python ctypes bindings (libds4.so) exposing engine/session/steering.

Validation + tooling:
- CUDA correctness gate (tensor-equivalence, cuda-ppl, cpu-cuda-ppl) + primer.
- gamut perf suite, longctx sweep harness, MTP retrain harness (v2 head),
  eugr N-runs mean+/-std methodology, ds4-bench --batch-cost/--batch-check,
  cascade A/B harness, thermal/memory-watchdog soak driver.

Gates green: token-diff identity (greedy+sampled, 4k+65k), mtp-correctness
worst_rms=0, mtp-selfconsistency maxabs=0, agent 86.5% accept / 21.9 t/s, no
BOS spam.

Rebased onto upstream cafc134 ("Fix server const warning"; +25 commits over
base 59d9bc7: CUDA/ROCm SSD streaming + mixed-precision routed experts + loader
refactor, MTP verify-topk fix antirez#358, agent raw-mode/edit fixes, server hardening).
Conflict resolutions (4 files, 6 hunks):
- ds4.c verify MoE: kept our verify_ordered_moe block, composed upstream's new
  DS4_ROCM_BUILD async-finish guard ahead of it.
- ds4.c verify top-k (antirez#358): kept our per-row argmax loop over upstream's
  indexer_topk arg-swap fix; ours is correct + faster (avoids the ~17ms
  single-thread O(n_vocab) scan). Validated by upstream's own antirez#358 regression
  (test_mtp_verify_depth): worst_argmax_gap=0.000, max_chunk=5.
- ds4_cuda.cu / tests / Makefile: union merges (both sides' additions kept).
Also pruned dead cuda_model_range_is_cached (our own pre-rebase orphan; base had
3 refs, the loader rework left ours unreferenced) -> build warning-clean.
Post-rebase re-validation on GB10: build clean; fork-fixes intact (device-copy
8.20 GiB via _exps. filter, host-register 80.76 GiB, 344 dense f16 prewarms,
managed-KV fix firing under tight free-mem). All gates green:
mtp-verify-depth (worst_argmax_gap=0.000, max_chunk=5), mtp-correctness
(worst_rms=0), mtp-selfconsistency (maxabs=0, bit-exact), cuda-ppl
(ppl=3.886775, delta=0.000000 vs committed baseline), cuda-tensor-equivalence
(worst_rms=0.0052, no layer fails, nonfinite=0).
TrevorS added a commit to TrevorS/ds4 that referenced this pull request Jun 17, 2026
@TrevorS
fork(gb10+mtp): GB10 speed/accuracy keep-set + MTP + batch-decode on …
542a9a0 
…upstream

Single consolidated fork commit over antirez/ds4, tuned for DGX Spark (GB10,
sm_121, unified memory, ~236 GB/s). Collapses the prior land/steering/python
stack plus this session's batch-decode + managed-KV work into one commit.

Engine (speed/accuracy, GB10):
- MTP speculative decode: combined-forward verify, adaptive K=1/2 cascade
  (opt-in), deterministic verify default + DS4_CUDA_FAST_VERIFY fast path.
- DS4_BATCH_DECODE (opt-in): route plain ds4_session_eval through the n=1
  combined-verify batch forward (verify kernel set: expert-tile MoE, heads8
  attention, cuBLAS f16). +21% plain decode @4k and @65k, token-identical in
  greedy + seeded-sampled testing (near-tie caveat, same as MTP combined).
- managed-KV policy fix: ds4_gpu_should_use_managed_kv_cache always consults
  cudaMemGetInfo free memory instead of the context_bytes < 8 GiB blind spot
  that OOM-hung the box at 131k ctx with the model resident.
- MoE expert-tile/rowspan kernels, atomic-free MoE down, Q8 residency cache,
  HBM hot-span device-copy with the _exps. filter, host-register, window attn.

Product:
- Multi-profile directional-steering cache + /v1/steering server admin API.
- In-process Python ctypes bindings (libds4.so) exposing engine/session/steering.

Validation + tooling:
- CUDA correctness gate (tensor-equivalence, cuda-ppl, cpu-cuda-ppl) + primer.
- gamut perf suite, longctx sweep harness, MTP retrain harness (v2 head),
  eugr N-runs mean+/-std methodology, ds4-bench --batch-cost/--batch-check,
  cascade A/B harness, thermal/memory-watchdog soak driver.

Gates green: token-diff identity (greedy+sampled, 4k+65k), mtp-correctness
worst_rms=0, mtp-selfconsistency maxabs=0, agent 86.5% accept / 21.9 t/s, no
BOS spam.

Rebased onto upstream cafc134 ("Fix server const warning"; +25 commits over
base 59d9bc7: CUDA/ROCm SSD streaming + mixed-precision routed experts + loader
refactor, MTP verify-topk fix antirez#358, agent raw-mode/edit fixes, server hardening).
Conflict resolutions (4 files, 6 hunks):
- ds4.c verify MoE: kept our verify_ordered_moe block, composed upstream's new
  DS4_ROCM_BUILD async-finish guard ahead of it.
- ds4.c verify top-k (antirez#358): kept our per-row argmax loop over upstream's
  indexer_topk arg-swap fix; ours is correct + faster (avoids the ~17ms
  single-thread O(n_vocab) scan). Validated by upstream's own antirez#358 regression
  (test_mtp_verify_depth): worst_argmax_gap=0.000, max_chunk=5.
- ds4_cuda.cu / tests / Makefile: union merges (both sides' additions kept).
Also pruned dead cuda_model_range_is_cached (our own pre-rebase orphan; base had
3 refs, the loader rework left ours unreferenced) -> build warning-clean.
Post-rebase re-validation on GB10: build clean; fork-fixes intact (device-copy
8.20 GiB via _exps. filter, host-register 80.76 GiB, 344 dense f16 prewarms,
managed-KV fix firing under tight free-mem). All gates green:
mtp-verify-depth (worst_argmax_gap=0.000, max_chunk=5), mtp-correctness
(worst_rms=0), mtp-selfconsistency (maxabs=0, bit-exact), cuda-ppl
(ppl=3.886775, delta=0.000000 vs committed baseline), cuda-tensor-equivalence
(worst_rms=0.0052, no layer fails, nonfinite=0).
redhunt07 pushed a commit to redhunt07/ds4-DSpark-GB10 that referenced this pull request Jul 11, 2026
@TrevorS @redhunt07
fork(gb10+mtp): GB10 speed/accuracy keep-set + MTP + batch-decode on …
ea912cd 
…upstream

Single consolidated fork commit over antirez/ds4, tuned for DGX Spark (GB10,
sm_121, unified memory, ~236 GB/s). Collapses the prior land/steering/python
stack plus this session's batch-decode + managed-KV work into one commit.

Engine (speed/accuracy, GB10):
- MTP speculative decode: combined-forward verify, adaptive K=1/2 cascade
  (opt-in), deterministic verify default + DS4_CUDA_FAST_VERIFY fast path.
- DS4_BATCH_DECODE (opt-in): route plain ds4_session_eval through the n=1
  combined-verify batch forward (verify kernel set: expert-tile MoE, heads8
  attention, cuBLAS f16). +21% plain decode @4k and @65k, token-identical in
  greedy + seeded-sampled testing (near-tie caveat, same as MTP combined).
- managed-KV policy fix: ds4_gpu_should_use_managed_kv_cache always consults
  cudaMemGetInfo free memory instead of the context_bytes < 8 GiB blind spot
  that OOM-hung the box at 131k ctx with the model resident.
- MoE expert-tile/rowspan kernels, atomic-free MoE down, Q8 residency cache,
  HBM hot-span device-copy with the _exps. filter, host-register, window attn.

Product:
- Multi-profile directional-steering cache + /v1/steering server admin API.
- In-process Python ctypes bindings (libds4.so) exposing engine/session/steering.

Validation + tooling:
- CUDA correctness gate (tensor-equivalence, cuda-ppl, cpu-cuda-ppl) + primer.
- gamut perf suite, longctx sweep harness, MTP retrain harness (v2 head),
  eugr N-runs mean+/-std methodology, ds4-bench --batch-cost/--batch-check,
  cascade A/B harness, thermal/memory-watchdog soak driver.

Gates green: token-diff identity (greedy+sampled, 4k+65k), mtp-correctness
worst_rms=0, mtp-selfconsistency maxabs=0, agent 86.5% accept / 21.9 t/s, no
BOS spam.

Rebased onto upstream cafc134 ("Fix server const warning"; +25 commits over
base 59d9bc7: CUDA/ROCm SSD streaming + mixed-precision routed experts + loader
refactor, MTP verify-topk fix antirez#358, agent raw-mode/edit fixes, server hardening).
Conflict resolutions (4 files, 6 hunks):
- ds4.c verify MoE: kept our verify_ordered_moe block, composed upstream's new
  DS4_ROCM_BUILD async-finish guard ahead of it.
- ds4.c verify top-k (antirez#358): kept our per-row argmax loop over upstream's
  indexer_topk arg-swap fix; ours is correct + faster (avoids the ~17ms
  single-thread O(n_vocab) scan). Validated by upstream's own antirez#358 regression
  (test_mtp_verify_depth): worst_argmax_gap=0.000, max_chunk=5.
- ds4_cuda.cu / tests / Makefile: union merges (both sides' additions kept).
Also pruned dead cuda_model_range_is_cached (our own pre-rebase orphan; base had
3 refs, the loader rework left ours unreferenced) -> build warning-clean.
Post-rebase re-validation on GB10: build clean; fork-fixes intact (device-copy
8.20 GiB via _exps. filter, host-register 80.76 GiB, 344 dense f16 prewarms,
managed-KV fix firing under tight free-mem). All gates green:
mtp-verify-depth (worst_argmax_gap=0.000, max_chunk=5), mtp-correctness
(worst_rms=0), mtp-selfconsistency (maxabs=0, bit-exact), cuda-ppl
(ppl=3.886775, delta=0.000000 vs committed baseline), cuda-tensor-equivalence
(worst_rms=0.0052, no layer fails, nonfinite=0).
iSevenDays pushed a commit to iSevenDays/ds4 that referenced this pull request Jul 17, 2026
@antirez
Merge pull request antirez#358 from pandysp/fix-mtp-verify-topk
7d26091 
fix: speculative verify with --mtp-draft > 2 commits wrong tokens
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@pandysp @antirez