Skip to content

Add native DSpark speculative decoding - #1914

Draft
Blaizzy wants to merge 3 commits into
mainfrom
pc/dspark-qwen38
Draft

Add native DSpark speculative decoding#1914
Blaizzy wants to merge 3 commits into
mainfrom
pc/dspark-qwen38

Conversation

@Blaizzy

@Blaizzy Blaizzy commented Aug 15, 2026

Copy link
Copy Markdown
Owner

Summary

  • add native DSpark checkpoint loading and speculative decoding for RadixArk/Qwen3.8-27B-DSpark
  • auto-detect DSpark checkpoints and expose --draft-kind dspark plus load-time --draft-bits quantization
  • add acceptance-adaptive draft depth while keeping explicit block sizes fixed for reproducible tuning
  • make speculative prefill sizing depend on live Metal headroom instead of a specific Mac tier
  • document CLI and server usage for Qwen3.8

Root cause and memory fix

Qwen3.5/3.8 previously coupled hidden-layer capture with Gated DeltaNet rollback capture. DSpark prefill only needs five hidden-state taps, but a 1k prompt retained per-token rollback intermediates across all 48 linear-attention layers. This could require more than 150 GB before other activations and caused unchunked prefill to be killed.

This change separates hidden capture from rollback-state capture. Rollback intermediates are now retained only for short speculative verification rounds. It also explicitly collects and clears the dense BF16 drafter tensors after load-time quantization.

Impact

On an Apple M5 Max with 48 GB unified memory, using Qwen3.8-27B-4bit with a fixed three-proposal DSpark block:

  • 1,029-token unchunked prefill: 857.6 tok/s
  • 64-token decode: 43.57 tok/s
  • peak Metal memory: 19.54 GB
  • baseline decode: ~33.6 tok/s (about 30% slower)
  • the same unchunked prefill previously OOMed

The quantized drafter's live parameter footprint is about 0.765 GB after reclaiming the dense load-time copy.

Validation

  • 162 passed, 1 deselected in mlx_vlm/tests/test_speculative.py
  • 31 passed targeted server and prompt tests
  • targeted Ruff fatal/error checks passed
  • git diff --check passed
  • real-model 1k-input/64-output server benchmark passed

The deselected speculative test is the existing exact-BF16 quantized-linear equality check; it remains unrelated to DSpark and differs only in numerical accumulation.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant