spec: MTP draft-head subset + even drafts / small-batch verify for qwen4exp

[gaetan-puleo]

Summary

Speculative-decoding (MTP) improvements for qwen4exp (Qwen3.8-Flash-Next UD-IQ4_XS) on gfx1151. Two kinds of change: the --spec-draft-mtp-vocab draft-head row subset (previously accepted only for qwen35/qwen35moe), and three fixes that make MTP actually pay off at several sequences and on small verify steps.

Measured on the same machine, server interleaved A/B, medians.

Changes

1. Enable the MTP draft-head row subset for qwen4exp (src/llama-model.cpp, src/models/qwen4exp.cpp). --spec-draft-mtp-vocab <n> scores drafted tokens against only the first n rows of the target LM head (a load-time row copy, no GGUF/weight change) and scatters them into a full-vocab -inf row. On a qwen4exp MTP sidecar the flag used to be silently ignored, leaving the draft on the full 248k-row Q6_K head. N=65536 is the sweet spot.
2. Even-length MTP drafts across sequences (STRIX_SPEC_EVEN_DRAFTS, default on; common/speculative.cpp, tools/server/server-context.cpp). The target's hybrid memory splits a batch into ubatches that need the same token count per sequence, so drafts of different lengths turned one verify into extra whole-model passes, each a fresh graph. All drafting sequences now pad to the longest confident run. Multi-sequence aggregate: 31.4 -> 36.3 t/s (adaptive).
3. Vector-kernel chunks for the small-batch MoE router and experts (STRIX_F32_VEC_CHUNK_MAX=32, STRIX_MOE_VEC_CHUNK=4 / _MAX_T=16; ggml/src/ggml-cuda/ggml-cuda.cu). A 5-16 token verify step's F32 router [2560x512] and its routed experts (past the vector limits: IQ3_S 4, IQ4_NL 6) fell to tiny GEMM/MMQ tiles; they now run the vector kernel over column chunks. Multi-sequence -np 3: 37.0 -> 42.7 t/s.
4. Small products of a verify step off one-workgroup kernels (STRIX_Q_VEC_CHUNK_MAX/ROWS, STRIX_CONCAT_T_MIN; ggml-cuda.cu, concat.cu). Quant weights of <=1024 output rows (hc_down [10240->320], attn k/v, shared-expert gate/up) chunked through the vector kernel, and the GDN conv concat uses the tiled transpose from 2 tokens. Multi-sequence np=3: +2.6 %.

Results

area	before	after	delta
MTP multi-sequence (-np 3)	37.05 t/s	44.3 t/s	+19.6 % (now level with plain batched decode, 45.2)
MTP multi-sequence, adaptive drafts	31.4 t/s	36.3 t/s	+16 %
MTP single-stream, --spec-draft-mtp-vocab 65536	38.0 t/s	40.0 t/s	+5.9 % (measured on the same harness)

Correctness

• The row-subset change is numerically inert on the default path (no subset) and byte-identical output with the subset on 3 prompts (same content sha1), acceptance unchanged at N=65536.
• The three speculative fixes change which kernel small batches take, not the arithmetic; identical content sha1 on single-sequence, and on the measured multi-sequence outputs.
• test-backend-ops test -b ROCm0 -o MUL_MAT,MUL_MAT_ID: 14577/14577.

Not included

The adaptive/cost-aware draft path (#102) is intentionally not part of this change.

Note

This branch was split out cleanly: the non-speculative decode kernels (Q8_0 deep-prefetch, HIP graph keys, IQ3_S vec-dot) live on the decode PR #106, not here.

[dzannotti]

Tested on Strix Halo (gfx1151, HIP), base ec01a7dfc vs this PR. Our Flash-Next GGUF has no MTP head, so the speculative part was measured on Qwen3.6-35B-A3B (MTP).

3-slot MTP serving (-np 3 --spec-type draft-mtp --spec-draft-n-max 3, llama-benchy pp512/tg256, concurrency 3, ABBA):

	base	PR	Δ
decode throughput	109.2 t/s	125.6 t/s	+15.0%
prefill throughput	1174 t/s	1134 t/s	−3.4%
TTFT	1317 ms	1386 ms	+5.3%

Non-speculative batches (no draft model; -r 8 ABBA, 16 samples per arm, slowest quarter dropped as warm-up):

tokens	Qwen3-1.7B Q8_0 (dense)	gpt-oss-20b MXFP4 (MoE)
4	–	167.8 → 167.1 (−0.4%)
8	–	228.7 → 225.1 (−1.6%)
9	759.0 → 749.4 (−1.3%)	–
12	–	360.5 → 294.1 (−18.4%)
16	1334.6 → 1287.6 (−3.5%)	452.6 → 327.2 (−27.7%)
24	1458.6 → 1432.5 (−1.8%)	–
32	1910.6 → 1818.3 (−4.8%)	–

Regressions:

• Chunked dispatch, not gated: the chunked MMVQ / F32 vector / MoE vector dispatch in ggml-cuda.cu applies to every CUDA/HIP device (only the concat change checks RDNA3.5). It slows models it was not tuned for, including a 28% loss on gpt-oss at 16 tokens, a common multi-slot and verify width.
• Even drafts: default-on in both common/speculative.cpp and the server, for every spec type and architecture.

Please gate these to the measured target, and make even drafts opt-in.

Minor: 7 new STRIX_* env vars are parsed with unvalidated atoi, and concat.cu calls getenv on every concat.

---

Comment generated via automated review of clanker Claude Code.

[Nathanw1014]

Nice one. One catch we hit: with the shared Flash-Next MTP sidecar (Unsloth's mtp-...-shared), --spec-draft-mtp-vocab quietly does nothing. That file has no output.weight of its own and borrows the target's LM head, so the subset never gets built or used.

The fix is small: take the subset rows from the target's head when the draft has none, and use the subset whenever the MTP block scores with the model head. With it, decode goes from 32.5 to 34.3 t/s on a Strix Halo (Vulkan), with the same acceptance.

It also needs a one-line rebase after #111, which renamed the batch size accessor in common/speculative.cpp:

            if (batch.size() == 0 || (even && n_confident == 0)) {

I have both on a branch (your four commits rebased on master, plus the fix as a fifth). Happy to push it here if that's easiest, or I can leave suggestions instead.

Assisted-by: Claude Opus 5.5