N-gram self-speculative decoding for hybrid GDN/SA models

[dahai80]

Motivation

Single-token decode on Apple Silicon is bandwidth-bound: each step streams the full model weights for one token. Speculative decoding can beat this ceiling by verifying multiple draft tokens per weight load — but standard speculative decoding needs a separate draft model (extra weights, extra memory).

N-gram self-speculative decoding generates draft tokens on the CPU from an n-gram table built from the prompt + generated tokens, with zero GPU overhead. The model's own forward pass verifies the draft. On a hit, one weight load yields 2 tokens; on a miss, we fall back to 1 token with cache rollback. No draft model required.

This is especially effective for hybrid GDN/SA architectures (e.g. Qwen3.5/3.6), where the GatedDeltaNet (GDN) layers have O(1) recurrent state that must be checkpointed and rolled back on draft rejection — KV-cache-only speculative decode would corrupt GDN state.

Proposal

Add an opt-in n-gram self-speculative decode path:

1. NgramDraftTable (generate.py) — CPU-side n-gram table (trigram → bigram → unigram fallback). O(1) hash probe, zero GPU cost. Built from the prompt, rolling-updated on each accepted token.

2. ngram_speculative_generate_step() (generate.py) — self-speculative decode generator. S=2 verification: forward pass produces 1–2 tokens per step. On acceptance: yield 2; on rejection: yield 1 + rollback cache. CLI flags --ngram-spec and --ngram-n (default n=3).

3. ArraysCache.checkpoint() / rollback() / trim() (cache.py) — generic cache state save/restore for GDN recurrent state during speculative verification. trim is a no-op for ArraysCache (state-based, not offset-based like KVCache); rollback is used instead. ~18 lines, no behavior change for existing caches.

4. GatedDeltaNet._conv1d_decode_multi() (qwen3_5.py) — vectorized S>1 decode path for the conv1d in GDN layers, using mx.stack + batched multiply instead of a sequential Python loop. Dispatched only when S>1 and a cache exists (the speculative verification step).

Correctness

• S=2 verification: the model's own logits verify the draft token, so output is identical to greedy decoding (verified: ngram-spec output matches baseline token-for-token).
• On draft rejection, ArraysCache.rollback() restores GDN recurrent state + KV cache to the pre-draft checkpoint, so the next step proceeds correctly from the rejected position.
• The fast path is opt-in (--ngram-spec); without it, behavior is unchanged.

Performance

Measured on Qwen3.6-27B (hybrid: 48 GDN + 16 SA layers) on M2 Ultra 137GB:

Mode	Speed	Notes
Baseline (S=1)	44.6 tok/s	bandwidth-bound
N-gram spec (1 draft)	52.1 tok/s	+17%, 44% draft hit rate
N-gram spec (2 drafts)	34.4 tok/s	0.77× — S=3 overhead exceeds gain

The 1-draft (S=2) configuration gives a reliable speedup; higher draft counts are counterproductive due to quadratic multi-token forward overhead on bandwidth-bound models.

Note on quantization: the benchmark above used quant2-all (an experimental 2-bit mixed recipe from #1466). The +17% is architecture-level — it reduces the number of GPU forward passes per token by verifying CPU-generated drafts, which is independent of weight precision. On standard 4-bit / bf16, the speedup should reproduce (the bandwidth ceiling is the same or lower-precision models are more bandwidth-bound, so the relative gain should hold or grow). I have not yet re-measured on standard upstream quant formats due to local environment constraints; happy to rerun on request.

Relationship to existing issues

• n-gram hashing for speculative decoding #851 (n-gram hashing for speculative decoding) — related but distinct: n-gram hashing for speculative decoding #851 proposes n-gram hashing; this issue proposes full self-speculative verification with GDN recurrent-state rollback for hybrid architectures, where the hard part is safely checkpointing/rolling back GDN O(1) state (not just KV cache).
• Decode performance optimizations for MLX-LM inference #1450 (decode performance omnibus) — this is one of the items tracked there, split into a focused issue for a clean PR trail.

A PR implementing the above will follow this issue.

Refs #851. Refs #1450.