feat(mtp): MTP verify-decode infrastructure - #6
Open
sufubao wants to merge 56 commits into
Open
Conversation
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.
Restore blank lines that were stripped from pre-existing definitions (black-induced reformatting of upstream code that this PR didn't functionally change). Keeps the diff focused on the MTP feature; fixing historical formatting is out of scope for this PR.
Scope this branch to Qwen3.5 MTP support only by rolling back the EAGLE-mode draft optimization. The draft model again runs the full (mtp_step+1)-expanded verify layout instead of being narrowed to the single accepted row per request. - dp/chunked _draft_decode_eagle: restore full-layout draft (copy.copy + b_num_accepted_tokens=None so it routes to the (bs, False) graph); drop the per-rank padding helpers and accepted-row narrowing. - base_backend: remove _build_eagle_accepted_draft_input / _scatter_accepted_next_token_ids. - cuda_graph: the draft runs at multiples of (mtp_step+1) again, so collapse the dual batch-size sets to one and delete the now-redundant _get_graph_batch_sizes routing. Keep the (bs, is_mtp_verify_decode) graph key + verify-layout warmup (core GDN verify support, not the optimization). - static benchmark: eagle path now measures the full-layout draft cost. - tests: drop the two narrowed-draft tests; rewrite the dual-set tests to the single-set model (still cover the verify/normal key distinction).
Drop the remaining draft-side divergence from upstream so this branch is scoped to Qwen3.5 MTP support only. The draft decode no longer clears b_num_accepted_tokens to force a flat/normal layout; it reuses the main model_input (still copy.copy'd to isolate per-step input_ids/b_seq_len/ mem_indexes mutations) and runs the same (mtp_step+1)-grouped verify decode layout as the main model — exactly as upstream does. For the pure-full-attention draft (qwen3_5_mtp: full_attention_interval=1, no GDN) grouped and flat are numerically identical: each position k sees KV [0, s+k) either way, same page-table entries, same RoPE positions; the main verify forward already uses this geometry and is the validated path. The earlier flat-draft only added an unnecessary (bs, False) cudagraph layout + b_num_accepted_tokens gating; nothing the draft computes needs it. - chunked_prefill/dp_backend: 6 draft fns (vanilla/eagle + dp overlap variants) stop clearing b_num_accepted_tokens. - cuda_graph: draft warms up the verify graph key too (mtp_step>0 -> verify for both main and draft); delete the now-dead _is_mtp_draft_model. - tests: rewrite the warmup-layout test (main+draft both verify; mtp_step==0 -> normal) and drop the stale "draft uses normal layout" framing. Keep is_mtp_verify_decode (main-model GDN verify still needs it) and the committed fp8.py causal=True fix. Verified live (QW35-122B-A10B, eagle_with_att, mtp_step=1, tp4): GSM8K acc 0.964 / Invalid 0.000, accept 1.956/2.0 — matches pre-revert baseline (no regression). Codex independent pass concurred (high confidence).
…e plumbing - is_mtp_verify: drop the redundant `b_num_accepted_tokens is not None` clause (post grouped-revert it's implied by mtp_step>0 ∧ ¬prefill). - Replace the per-step host round-trip for b_num_accepted_tokens with a GPU-resident ReqManager.req_to_accept_len: a triton scatter_mtp_accept_len after verify + a GDN-only gather in init_mtp_verify_extra_state. Removes the gen_from_list H2D rebuild, the phase-2 req.mtp_accept_len writeback, and the host attr (linear-att offload + resets now read/write the buffer). - Drop the redundant `if mtp_step>0` guard inside decode_mtp/decode_overlap_mtp. - config_objs: inline the mtp draft-layer count, dropping the _mtp_added_layer_num helper (kept get_added_mtp_kv_layer_num inlined in envs_utils). - cpu_cache_meta: don't bump layer_num for linear-att models (the draft full-att slots are already in LinearAttCacheConfig.get_cpu_cache_big_page_bytes()). Static checks pass (ast, flake8). The req_to_accept_len refactor is not yet runtime-verified; pending a hybrid GSM8K + cudagraph-ON parity run.
Keep streamed_args_for_tool a byte-exact prefix of json.dumps(arguments) so the serving-layer stop-time reconciliation can't produce duplicated/invalid JSON. String values stream incrementally; non-string values are emitted only once </parameter> arrives (a partial number/array/bool isn't a prefix of its json.dumps form). Detect "value still streaming" via the terminator rather than match position ($ matches before the template's trailing newline). Also: - don't crash (IndexError) on >=2 <function=> in one <tool_call> block, and emit a name head for each; flush via _ensure_qwen3_stream_state. - an undefined first function no longer discards the whole block (a valid function after it is still emitted). - treat </tool_call> as an implicit function close so the closing '}' rides in the same args delta (no separate trailing-'}' delta to double-count). - drop dead _build_partial_arguments_json; dedup newline-strip / param-type helpers. Add unit_tests/server/test_qwen3_coder_stream_fc.py covering these across chunk boundaries.
Squash of logging and router-stats work: - Logging: colored levels, windowed cache stats, quieter per-request logs; restore config access logs and add debug status details. - Httpserver: reject oversized prompts and translate ValueError to 400. - Router stats: restore router_statics.update() on completion, output TPS via per-req deltas, skip aborted reqs, sweep deltas once per print interval, defensive group_request_id init. - Detokenization: keep req registration failures fatal. - Add tests for access log and prompt length guard. Claude-Session: https://claude.ai/code/session_0195z54gpFLwSyX82Hj5ooQf
Squash of 20 commits adding MTP / speculative-decode support: - New MTP model packages: qwen3_5_mtp, qwen3_5_moe_mtp (pre/post layer weights, transformer layer weights, model wiring) + mtp_model_factory. - qwen3next: GDN decode conv+pack kernel (gdn_decode_pack), causal conv1d spec kernel, shared-expert-gate kernel; verify-decode plumbing and GPU-resident req_to_accept_len. - basemodel: mtp_verify_extra_state, cuda_graph MTP-aware warmup decode input builder; backend changes (base/chunked_prefill/dp, generic post-process) for verify/accept handling. - distributed: flashinfer all-reduce + communication_op additions. - sampling/api: mtp_step CLI/sampling params plumbing. Conflicts resolved: - cuda_graph: kept HEAD capture progress bar calling the MTP _build_warmup_decode_model_input helper. - qwen3next GDN decode: dropped HEAD's strided-decode optimization (_rearrange_mixed_qkv + per-token-stride decode-only fused_recurrent) in favor of the branch's conv_pack decode + verify path. The two are mutually exclusive; HEAD's decode-only kernel asserted cu_seqlens is None, which crashed the MTP verify path. Reverted fused_recurrent.py to the branch version (cu_seqlens-capable) and removed the now-orphaned test_fused_recurrent_strided.py so the whole qwen3next GDN subsystem matches the tested MTP branch. Claude-Session: https://claude.ai/code/session_0195z54gpFLwSyX82Hj5ooQf
Squash of 20 commits adding visual-server memory reservation: - visualserver/model_infer: new mem_reserve.py + worst_case_reserve.py; model_rpc / __init__ wiring to reserve ViT activation memory up front. - vit/model.py and qwen2_vl / qwen2_5_vl / qwen3_vl visual models: report ViT activation delta above weights (not total reserved), worst-case grid rounds up so it never under-reserves the cap. - basemodel: supporting changes; api_cli: new reservation flag(s). Clean auto-merge with the MTP work (no conflicts). Claude-Session: https://claude.ai/code/session_0195z54gpFLwSyX82Hj5ooQf
Adds TritonFP8w8a8PerTensorQuantizationMethod (triton-fp8w8a8-pertensor / -pt) with deferred staging for fused multi-split and per-expert weights: - grouped MoE matmul: WEIGHT_SCALE_PER_TENSOR support. - per-token scaled_mm kernel: B_SCALE_IS_TENSOR support. - w8a8.py: new per-tensor method; api_cli --quant_type help updated. Clean auto-merge (single upstream commit; no conflicts). Claude-Session: https://claude.ai/code/session_0195z54gpFLwSyX82Hj5ooQf
sample() returns next_token_logprobs=None whenever logprobs aren't requested (flashinfer fast-path / not need_logprobs), so next_token_logprobs_cpu is None for ordinary requests. The MTP decode paths subscripted it unconditionally (next_token_logprobs_cpu[select_mask]) -> TypeError: 'NoneType' object is not subscriptable. The non-MTP path already passes None through and _post_handle handles it; mirror that in chunked_prefill.decode_mtp and dp_backend MTP decode. Claude-Session: https://claude.ai/code/session_0195z54gpFLwSyX82Hj5ooQf
TritonFP8w8a8PerTensorQuantizationMethod.apply() quantized the activation with vllm's scaled_fp8_quant, which is only bound under `if HAS_VLLM`. On images without vllm this raised `NameError: name 'scaled_fp8_quant' is not defined` at first prefill, crashing model init (router saw only EOFError). Use the pure-triton per-token kernel instead. Call lightllm_per_token_group_quant_fp8 directly rather than the per_token_group_quant_fp8 wrapper: the wrapper dispatches to the sgl kernel when present, which rejects group_size == k (per-token). group_size == k matches vllm's use_per_token_if_dynamic semantics. Verified on H200: per-token scale exact vs max(|x|)/448; mm cosine 0.9993 vs bf16 ref (> repo's 0.99 gate) across M=1/16/47/128/7 and non-aligned K/N. Claude-Session: https://claude.ai/code/session_0195z54gpFLwSyX82Hj5ooQf
get_scheduler_metadata was imported on the same line as flash_attn_varlen_func from sgl_kernel.flash_attn. Older sgl_kernel builds ship fa3 but not get_scheduler_metadata, so the whole import raised and the except nulled out flash_attn_varlen_func too -> fa3 reported unavailable -> flashinfer fallback (incl. a ~54s JIT compile at startup), despite fa3/fp.py already guarding get_scheduler_metadata is None. Import it in its own try/except. Claude-Session: https://claude.ai/code/session_01R9kHHvsKkjkniVZwUSpTvA
Brings in PR ModelTC#1349 (perf(qwen3next): drop q/k/v/a/b contiguous copies in GDN fused_recurrent decode). Semantic reconciliation in fused_recurrent.py (git auto-merged it silently since this branch never textually touched the file, but the auto-resolution would have broken MTP verify): ModelTC#1349 makes the fused recurrent kernel decode-only by asserting `cu_seqlens is None`. This branch's MTP `_gdn_verify_kernel` drives the same kernel with `cu_seqlens` (variable-length verify chunks), so the bare auto-merge would crash verify-decode on that assert. Resolution keeps BOTH: ModelTC#1349's per-token strided no-copy decode path AND the MTP verify varlen path. The strided kernel arithmetic is already general (bos * stride_tok), so only the host wrapper needed fixing: - drop the two `assert cu_seqlens is None` guards - restore `N = B if cu_seqlens is None else len(cu_seqlens) - 1` - generalize `_ensure_qkv_token_strided` to accept the verify layout [1, tokens, head, dim] (token dim = dim 1) in addition to the decode layout [tokens, 1, head, dim] (token dim = dim 0); both are contiguous-tail column views, so no copy is needed in either case. unit_tests/.../test_fused_recurrent_strided.py: dropped the `test_cu_seqlens_is_not_supported` negative test (it asserted the lifted decode-only contract); kept the decode strided-views equivalence test. Varlen verify correctness is covered E2E by the MTP GSM8K accuracy check. Claude-Session: https://claude.ai/code/session_01J3SiYM55DJg8ht4dmEPmUx
…t-in)
Add --gdn_prefill_backend {fla,flashqla} (default fla, no behavior change).
FlashQLA is Qwen's TileLang GDN chunked-prefill kernel; it replaces only the
chunk_gated_delta_rule prefill call (decode / MTP-verify fused_recurrent are
untouched). A memoized resolver selects it only on SM90+ with flash_qla
importable, otherwise falls back to the vendored FLA triton kernel with a
warning.
Hopper H200, Qwen3.5-27B TP4, A/B vs FLA:
- accuracy lossless (GSM8K 0.895 vs 0.875; kernel bit-exact, cos=1.0)
- long-context (ISL 8192) total throughput +8-10% across concurrency 1-128,
TTFT ~9% faster; multi-turn +1-3% (prefix/state cache absorbs most prefill)
FlashQLA's l2norm/gemm require last-dim-contiguous q/k/v, which the rearranged
mixed-qkv are not, so the flashqla path makes them contiguous.
Claude-Session: https://claude.ai/code/session_019aksHwfDGt51aMRbny6YQn
The triton scaled_mm_per_token kernel underperforms on the attention-projection shapes (qkv N=8192, o N=5120,K=6144) at decode batch sizes: it runs ~0.4-0.6x of vLLM's Cutlass fp8 scaled-mm and even slower than bf16, because its grid (cdiv(M,BM)*cdiv(N,BN), no split-K) can't fill the GPU at small N + small M. A split-K prototype confirmed the diagnosis but never matched Cutlass, which holds a ~2x structural lead. This capped end-to-end fp8 speedup at ~1.10x over bf16 vs vLLM's ~1.20x on Qwen3.5-27B / H200. Select the fp8 GEMM backend via LIGHTLLM_FP8_GEMM (auto|cutlass|sgl|triton); auto prefers vllm cutlass_scaled_mm, then sgl_kernel.fp8_scaled_mm, then the triton kernel, so images without a Cutlass-class kernel fall back cleanly. Lifts end-to-end fp8 throughput to ~1.19x over bf16 (matches vLLM); GSM8K lossless (cutlass 0.850 / sgl 0.860 / triton 0.825 / bf16 0.860). sgl_kernel rejects a per-tensor weight scale, so the scalar is expanded to per-channel [N] and cached on the weight_pack. Claude-Session: https://claude.ai/code/session_013NAh5zZG2KvrtmZHpKH3AY
# Conflicts: # lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py
一个 mtp decode step 原来是 4 次独立 graph replay (主模型 verify + 3 次 draft), 之间夹着大量 eager glue (inferstate 构建/gen_decode_params/page_table 拷贝/ copy_for_cuda_graph 输入拷贝/argmax/mem_indexes 旋转) 与 4 次 ~390us 的 cudaGraphLaunch。现在捕获为两个 graph: - verify graph: input gather + 主模型 forward + flashinfer 采样 (philox seed/offset 驻显存, 图内自增) + mtp verify + accept 更新 - draft graph: 3 次 draft forward + argmax 链 + 候选 scatter, 在 verify 结果回传后异步执行, 掩盖下一 step 的 cpu 准备 温度/top_k/top_p 为逐行静态 buffer, 支持 batch 内异构参数; 惩罚项/logprobs/ 自定义种子/logit_bias/约束解码等请求自动回退旧路径。多模态 decode 的 position delta 走静态显存 buffer 保证捕获安全。 Qwen3.5-27B H200 TP1 c1: ITL 7.67ms -> 6.38ms, 输出吞吐 ~107 -> ~137 tok/s, accept length 无回归。LIGHTLLM_DISABLE_MTP_FUSED_GRAPH=1 可关闭。
- mtp draft input hiddens now use post-final-norm capture for qwen3.5 (GEMMA +1 norm makes pre/post-norm materially different), matching the trained convention; default stays pre-norm for deepseek-style models - eagle draft chain steps>=1 write kv to static scratch slots instead of clobbering verify slots, and the verify position mapping is restored after the chain: next-round draft prefix kv keeps step0 content (sampled token, main-model hidden), aligning with sglang draft re-extend semantics; fixes p2/p3 chain accept decay - the chain scratch is reserved once at startup and reused: per-round allocation at high kv watermark could fail radix eviction (all tree tokens refed) and kill the infer thread, and the fused graph now captures the slot tensor as a constant (drops eagle staging entirely) - mtp_avg_token_per_step no longer counts the prefill-produced token as a verify step (metric was understated ~2%)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.