Repository navigation
Conversation
Qwen4Exp never produced the auxiliary features a DFlash drafter consumes, and Qwen4ExpForConditionalGenerationConfig rejected every spec method except mtp and ngram. Add the capture: Qwen4ExpModel inherits EagleModelMixin, both entry classes inherit SupportsEagle3, and the layer loop materializes the completed output of the drafted layers (their pending hyper-connection combine), averages it over the hc_count streams, and returns it with the sampled stream. DFlash2's fc.weight is [2560, 12800], so the 5 features named by target_layer_ids [5, 15, 25, 35, 45] are what the checkpoint expects. Both get_mtp_target_hidden_states() call sites in the V2 model runner sliced the hook result unconditionally. It returns None whenever the draft is not the native MTP layer, so a DFlash draft crashed profile_run with 'NoneType' object is not subscriptable. Verified on 2x A100 (TP2, expert parallel) against a Qwen3.8-Flash-Next AWQ W4A16 checkpoint with incoai/Qwen3.8-Flash-Next-DFlash2 at its trained block size 8: mean acceptance length 4.08 and 224.9 output tok/s at concurrency 1, against 2.98 and 180.1 with MTP-3 over the same 10 prompts. Chat, qwen3 reasoning parsing and qwen3_xml tool calls all work. Known gap: prefix reuse falls back to whole-block granularity while a DFlash draft is loaded, so the fine-grained --prefix-match-unit path needs a follow-up. Co-Authored-By: pi (Qwen3.8-Flash-Next) <noreply@github.com>
|
Heads-up from a production GLM-5.3-Flash deployment: the DFlash2 drafter failure here (KV-cache planning declining the drafter's sliding-window layers) has the same root cause as our #120 — the fix is a pure-Python patch to Patches plus a CPU-only repro (prints |
TL;DR
Qwen3.8-Flash-Next can now use a DFlash / DFlash2 draft model for speculative decoding. On 2x A100-80GB this lifts single-stream output from 180.1 to 224.9 tok/s (+25%) over the native MTP-3 drafter, with mean acceptance length going from 2.98 to 4.08.
The draft-model side (
qwen3_dflash2.py, the DFlash2 speculator, the registry entry) was already in this repo. What was missing was on the target side, so a DFlash draft could never actually run against Qwen4Exp. Three small changes, no new abstractions.This is opt-in. Nothing changes unless you pass
--speculative-configwith"method":"dflash".What was missing
1. The target never produced the features the drafter was trained on. DFlash drafts condition on the target's hidden states at specific layers (
dflash_config.target_layer_ids, here[5, 15, 25, 35, 45]), concatenated. The draft'sfc.weightis[2560, 12800], i.e. 5 features of width 2560. Qwen4Exp had no capture hooks at all.Qwen4ExpModelnow inheritsEagleModelMixin, both entry classes inheritSupportsEagle3, and the layer loop captures the completed output of the requested layers, averaged over thehc_counthyper-connection streams. "Completed" means the pending hyper-connection combine is materialized first, the same way the layer body and the existing pipeline-parallel and deepstack paths materialize it, so the captured tensor is the real layer output rather than a half-mixed stream.2. The config check rejected it.
Qwen4ExpForConditionalGenerationConfig.verify_and_update_configraisedNotImplementedErrorfor any speculative method other thanmtp,ngram,ngram_gpu.dflashis now allowed.3. A latent crash in the V2 model runner. Both call sites of
get_mtp_target_hidden_states()sliced the return value unconditionally. The hook is declared-> torch.Tensor | Noneand returnsNonefor every non-MTP drafter, soprofile_rundied with'NoneType' object is not subscriptable. Now guarded. MTP behaviour is unchanged.Benchmarks
Target
/data/calclout_Qwen3.8-Flash-Next-Abliterated-AWQ-W4A16(AWQ W4A16,qwen4_exp), draftincoai/Qwen3.8-Flash-Next-DFlash2. 2x A100-80GB, TP2,--enable-expert-parallel,--mamba-cache-mode=align,VLLM_PLE_CPU_OFFLOAD=1, MTP comparison uses the documented MTP-3 recipe. 10 fixed math/code/chat prompts, greedy, warmed server, reported as output tokens per wall second. Numbers in parentheses are mean acceptance length.Acceptance detail. Mean acceptance length is accepted tokens per verification step; drafted-token rate is accepted tokens over drafted tokens.
Per-drafted-token acceptance drops as the block grows while accepted tokens per step keep rising, which is where the throughput comes from. Concurrency-4 throughput has roughly 10% run-to-run spread on this stack (block 8 measured 519, 589 and 634 tok/s on three runs of the same bench), so blocks 6 to 8 are tied there and acceptance length is the better signal.
Per-position acceptance at block 8, single stream, 1,592 verification steps:
Smooth decay rather than a collapse, which is what tells you the captured features are the right ones.
KV cache. Available KV tokens: 2,138,244 with the MTP draft, 2,377,232 with DFlash block 4 falling to 2,161,700 at block 8. Every DFlash config leaves more KV cache than the MTP draft.
Functional checks. Chat completion,
--reasoning-parser=qwen3, and--tool-call-parser=qwen3_xml(returnsfinish_reason: tool_callswith valid arguments) all work at every block size tested.Block size you have to pass
Not a code issue, but a real usability trap for anyone trying DFlash2 on this model. QSA ring capacity is
compress_ratio * cdiv(compress_ratio + num_speculative_tokens, compress_ratio)(vllm/models/qwen4_exp/common/qsa_cache.py:847) and must divide the attention block size. This checkpoint hasindexer_compress_ratio = 4, so capacity jumps to 12 fornum_speculative_tokens5 to 8, and 12 does not divide the block size that--mamba-cache-mode=alignrequires (1,616 to 1,648). Startup fails withQSA ring capacity 12 must divide the attention block size 1616.--block-size 1632--block-size 1632--block-size 1680Known limitation, and why it is not fixed here
Prefix-cache reuse falls back to whole-block granularity whenever a DFlash draft is loaded, so
--prefix-match-unit=16and--enable-mamba-fine-grained-prefix-cachestop helping:It is the draft group, not the larger block size: block 5 runs at the automatic block size and still reuses only whole blocks, so anything shorter than two blocks gets no hit. The decision lives in
resolve_kv_cache_block_sizes(vllm/v1/core/kv_cache_utils.py) and the partial-mamba tail handling in the scheduler, which is adjacent to what #120 changes. I kept it out of this PR rather than destabilise Qwen4Exp's KV grouping; this PR is what makes the drafter run at all, the reuse granularity is the follow-up.How to test
Start the model with a DFlash draft instead of the MTP config, everything else unchanged:
--speculative-config '{"method":"dflash","model":"incoai/Qwen3.8-Flash-Next-DFlash2","num_speculative_tokens":7}' \ --block-size 1680Startup must log
Using Eagle3 auxiliary layers from config: (6, 16, 26, 36, 46)(the+1offsetget_eagle3_aux_layers_from_configapplies to DFlash layer ids) and then reachApplication startup complete./metricsthen reportsvllm:spec_decode_*counters, which is howthe tables above were read. The baseline column is the documented MTP-3 recipe for this model.
Not covered
tests/evals/qwen4_exp/). Speculative decoding is verified by rejection sampling and output looked correct, but nobody has scored it.Relation to #120
Same genre, disjoint code. #120 gives GLM-5.3 a KV group for sliding-window draft layers and adds the equivalent
EagleModelMixin/SupportsEagle3plumbing toglm5next. This PR does the Qwen4Exp side (its own hyper-connection capture semantics and the MTP-buffer guard) and touches no file #120 touches.Notes for reviewers
qwen4_expaux hidden-state capture or DFlash for this model. [GLM5][Spec Decode] Group sliding-window draft layers separately in the GLM-5.3 KV builder (DFlash2 on A100) #120 is GLM-5.3 only.pre-commit runpasses on the three changed files, includingmypy-3.10. The# type: ignore[union-attr]on the MTP-buffer slice is pre-existing and covershidden_states, not the guarded hook.