feat(qwen3_5_mtp): Qwen3.5 / Qwen3.5-MoE MTP speculative decoding#1338
feat(qwen3_5_mtp): Qwen3.5 / Qwen3.5-MoE MTP speculative decoding#1338sufubao wants to merge 5 commits into
Conversation
There was a problem hiding this comment.
Code Review
This pull request implements comprehensive multi-token prediction (MTP) and speculative decoding support for Qwen3.5 and Qwen3Next models, including updates to attention backends, CUDA graph warmup layouts, and Triton kernels for spec-decode updates. Feedback on the changes highlights a performance concern regarding synchronous device-to-host transfers when validating b_num_accepted_tokens on the GPU, as well as a potential division-by-zero error in the causal conv1d update kernel if the batch size is empty.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| assert int(b_num_accepted_tokens.min()) >= 1 and int(b_num_accepted_tokens.max()) <= mtp_step + 1, ( | ||
| f"b_num_accepted_tokens out of range [1, {mtp_step + 1}]: " | ||
| f"min={int(b_num_accepted_tokens.min())} max={int(b_num_accepted_tokens.max())}" | ||
| ) |
There was a problem hiding this comment.
Calling .min() and .max() on the GPU tensor b_num_accepted_tokens and casting to int causes a synchronous device-to-host (D2H) transfer. Since this function is called on the eager decode hot path, this synchronization will stall the CPU and degrade inference performance. Consider performing this validation on the CPU side (e.g., on the list of mtp_accept_len in infer_batch.py before moving it to CUDA) to keep the execution fully asynchronous.
| assert conv_state_indices is not None | ||
| batch = conv_state_indices.size(0) | ||
| dim = x.size(1) |
There was a problem hiding this comment.
If batch (the size of conv_state_indices) is 0, calculating seqlen = x.size(0) // batch will raise a ZeroDivisionError. To make the function robust against empty batches (which can occur in edge cases or testing), add a defensive guard to return x immediately if batch == 0.
assert conv_state_indices is not None\n batch = conv_state_indices.size(0)\n if batch == 0:\n return x\n dim = x.size(1)493ce5e to
ed22f8e
Compare
Model-agnostic verify-decode machinery: MTP-verify dispatch in TpPartBaseModel, dedicated decode CUDA-graph capture/replay for the (mtp_step+1)-expanded verify layout, a shared mtp_verify_extra_state block on infer_struct/batch_objs, fa3 decode attention narrowed to the verify layout (b_att_seq_len + causal) for fp/fp8/mla, and env/kv-cache helpers for MTP added-layer accounting.
Self-contained dense (qwen3_5_mtp) and MoE (qwen3_5_moe_mtp) MTP draft packages: each carries its own draft wiring (reuse the main model's req/mem managers + rope caches, is_mtp_draft_model marker) and shares a weight-retarget mixin (mtp.* head, embeddings shared with the main model) plus the MTP pre-layer fuse. No shared model base class.
Gated-delta-net (linear attention) speculative-decode verify path for qwen3next: a per-sequence spec causal_conv1d kernel; a widened conv working slot split from the committed (narrow) persisted slot; MTP draft full-attn KV-slot accounting across the linear-att cache config, mem operator and req manager; and removal of the dead gen_b_req_mtp_start_loc kernel.
Wire the verify path through the inference backends: a single draft-model factory keyed on (model_type, mtp_mode); build the (mtp_step+1)-expanded verify decode batch; run the eagle + vanilla draft decode; verify accepted tokens; and thread per-request accept-lengths (b_num_accepted_tokens) from the chunked-prefill and dp backends into the model verify forward.
Behavioural/CUDA coverage for the subtle MTP paths: verify-extra-state metadata, decode CUDA-graph verify layouts, fa3 fp8 verify narrowing, GDN verify equivalence, the spec causal_conv1d kernel and its prefill->decode roundtrip, and the linear-att conv/SSM widened-slot split + snapshot + CPU-cache persistence. Also extends the static-inference MTP benchmark and anchors the .gitignore benchmark-output rule to /benchmark.
What
Adds Qwen3.5 / Qwen3.5-MoE Multi-Token Prediction (MTP / speculative decoding) end-to-end, together with the model-agnostic verify-decode machinery it needs. Builds on the
BaseMTPModelrefactor in #1337 (merged).Commits
TpPartBaseModel, dedicated decode CUDA-graph capture/replay for the(mtp_step+1)-expanded verify layout, a sharedmtp_verify_extra_stateblock, and fa3 decode attention narrowed to the verify layout (b_att_seq_len+ causal) for fp / fp8 / mla.mtp.*head, embeddings shared with the main model) and the MTP pre-layer fuse.causal_conv1dkernel, a widened conv working slot split from the committed (narrow) persisted slot, and MTP draft full-attn KV-slot accounting across the linear-att cache config / mem operator / req manager.(model_type, mtp_mode), the verify decode batch, eagle + vanilla draft decode, and per-request accept-length (b_num_accepted_tokens) transport through the chunked-prefill and dp backends..gitignorebenchmark-output rule anchored to/benchmark.Testing
pre-commit(black 21.12b0 + flake8 6.1.0) clean.unit_tests/cover verify-extra-state, decode CUDA-graph layouts, fa3 narrowing, GDN verify equivalence, linear-att conv/SSM split + CPU-cache persistence, and the draft-model factory.