From be4a904ba2fd574a134023233370fb7c84cc4d4a Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 15 Oct 2026 08:26:46 +0800 Subject: [PATCH] [Bugfix] Preserve singleton-first MTP fallback token dimensions Assisted-by: Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- benchmarks/flashnext_prefill/HANDOFF.md | 8 + tests/models/qwen4_exp/test_mtp_hc_prefill.py | 195 ++++++++++++++++++ vllm/models/qwen4_exp/nvidia/mtp.py | 118 ++++++++--- 3 files changed, 293 insertions(+), 28 deletions(-) diff --git a/benchmarks/flashnext_prefill/HANDOFF.md b/benchmarks/flashnext_prefill/HANDOFF.md index 6c6ffa526..2b07831fc 100644 --- a/benchmarks/flashnext_prefill/HANDOFF.md +++ b/benchmarks/flashnext_prefill/HANDOFF.md @@ -153,6 +153,14 @@ sharded HC、Q8_0 MTP experts、TP4/PP1/DCP1/PCP1、C1、BACKED 动态形状， 原归约，HC 行分片复用现有投影/通信/归一化实现，最终恢复原 token 顺序和 两路完整输出，padding 行不暴露给 proposer。 +启用该候选时，普通草稿回退模型使用独立的 UNBACKED token 维度编译配置， +并声明 ids、position、hidden、embedding 的行数关系。大 prefill 由新 view +处理后，原回退可能首次在 M1 编译；BACKED/AOT 关闭 guard 会将这一维度 +固化，后续 M6 的投影/position 错位。最小配对回归中，旧 BACKED 路径在 +M6 仍返回 M1 输出；UNBACKED 在首次 M1、随后 M6/M16/M2/M1/M64 及 AOT +缓存重新加载时保留正确行数。只有完整满足上述 admission 的模型复制 +这一配置；主模型、decode view 和未加速模型保留原配置和计算策略。 + `hc_mtp_sequence_runtime.py` 验证真实 MTP/decoder/HC adapter 的数值、行边界、 编译、变输入 graph replay、target hidden 不变和输出所有权；合成 producer 测试不能替代整模型质量或性能。候选尚需整模型 cold prefill、长文任务、 diff --git a/tests/models/qwen4_exp/test_mtp_hc_prefill.py b/tests/models/qwen4_exp/test_mtp_hc_prefill.py index 1b07eae1e..eed7a8526 100644 --- a/tests/models/qwen4_exp/test_mtp_hc_prefill.py +++ b/tests/models/qwen4_exp/test_mtp_hc_prefill.py @@ -71,3 +71,198 @@ def test_mtp_prefill_view_borrows_weights_and_state(): assert multi.data_ptr() != hidden.data_ptr() assert torch.equal(owner.state, torch.ones(1)) assert view._target_model is owner + + +@pytest.mark.parametrize("aot", [False, True]) +def test_mtp_singleton_first_compilation_preserves_token_axes( + monkeypatch, tmp_path, aot +): + import json + + from vllm.compilation.decorators import support_torch_compile + from vllm.config import ( + CompilationConfig, + CompilationMode, + ModelConfig, + VllmConfig, + set_current_vllm_config, + ) + from vllm.config.compilation import ( + DynamicShapesConfig, + DynamicShapesType, + PassConfig, + ) + from vllm.forward_context import set_forward_context + + monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "1" if aot else "0") + monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path)) + config_dir = tmp_path / "model" + config_dir.mkdir() + (config_dir / "config.json").write_text( + json.dumps( + { + "model_type": "gpt2", + "architectures": ["GPT2LMHeadModel"], + "n_embd": 8, + "n_head": 2, + "n_layer": 1, + "n_positions": 128, + "vocab_size": 32, + "bos_token_id": 1, + "eos_token_id": 2, + } + ) + ) + + class Layer(nn.Module): + def forward(self, *, hidden_states, positions, **kwargs): + block = hidden_states + positions.to(hidden_states.dtype)[:, None] + return hidden_states, block, torch.zeros_like(hidden_states) + + class Mixer(nn.Module): + def combine_and_mix(self, hidden, block, injection): + multi = hidden + block + return multi, multi[:, :4], None + + @support_torch_compile( + dynamic_arg_dims={**mtp._MTP_TOKEN_AXES, "intermediate_tensors": 0}, + ) + class Predictor(nn.Module): + forward = mtp.Qwen4ExpMultiTokenPredictor.forward + + def __init__(self): + super().__init__() + self.hc_count, self.hidden_size, self.num_mtp_layers = 2, 4, 1 + self.layers = nn.ModuleList([Layer()]) + self.hyper_connection_mixer = Mixer() + for name in ( + "pre_fc_norm_embedding", + "pre_fc_norm_hidden", + "fc_embedding", + "fc_hidden", + ): + setattr(self, name, nn.Identity()) + + cfg = VllmConfig( + model_config=ModelConfig( + model=str(config_dir), + dtype="float32", + max_model_len=128, + skip_tokenizer_init=True, + ), + compilation_config=CompilationConfig( + mode=CompilationMode.VLLM_COMPILE, + backend="inductor", + compile_sizes=[], + cudagraph_mode=CUDAGraphMode.NONE, + dynamic_shapes_config=DynamicShapesConfig(type=DynamicShapesType.UNBACKED), + pass_config=PassConfig( + fuse_norm_quant=False, + fuse_act_quant=False, + fuse_attn_quant=False, + fuse_allreduce_rms=False, + fuse_minimax_qk_norm=False, + enable_qk_norm_rope_fusion=False, + fuse_rope_kvcache=False, + fuse_rope_kvcache_cat_mla=False, + fuse_mla_dual_rms_norm=False, + ), + ), + ) + with torch.no_grad(), set_current_vllm_config(cfg), set_forward_context({}, cfg): + owner = Predictor() + for rows in (1, 6, 16, 2, 1, 64): + ids = torch.arange(rows) + positions = torch.arange(rows) + 7 + hidden = torch.arange(rows * 8).float().reshape(rows, 8) / 128 + embeds = torch.arange(rows * 4).float().reshape(rows, 4) / 64 + reference = owner.forward(ids, positions, hidden, inputs_embeds=embeds) + actual = owner(ids, positions, hidden, inputs_embeds=embeds) + for a, b in zip(actual, reference): + torch.testing.assert_close(a, b, rtol=0, atol=0) + + assert actual[0].shape == (rows, 4) and actual[1].shape == (rows, 8) + if aot: + restored = Predictor() + ids, positions = torch.arange(6), torch.arange(6) + 81 + hidden, embeds = torch.ones(6, 8), torch.ones(6, 4) + actual = restored(ids, positions, hidden, inputs_embeds=embeds) + reference = restored.forward(ids, positions, hidden, inputs_embeds=embeds) + assert restored.was_aot_compile_fn_loaded_from_disk + for a, b in zip(actual, reference): + torch.testing.assert_close(a, b, rtol=0, atol=0) + + +def test_mtp_fallback_compiler_copy_preserves_other_owners(): + from vllm.config import CompilationConfig + from vllm.config.compilation import DynamicShapesType + + root = NS(compilation_config=CompilationConfig(), kernel_config=object()) + root.compilation_config.cache_dir = "original" + root.compilation_config.traced_files = {"original.py"} + assert mtp._make_mtp_prefill_fallback_config(root, qualified=False) is root + copied = mtp._make_mtp_prefill_fallback_config(root, qualified=True) + assert copied.kernel_config is root.kernel_config + assert ( + copied.compilation_config.dynamic_shapes_config.type + == DynamicShapesType.UNBACKED + ) + assert ( + root.compilation_config.dynamic_shapes_config.type == DynamicShapesType.BACKED + ) + assert copied.compilation_config.cache_dir == "" + assert root.compilation_config.cache_dir == "original" + copied.compilation_config.traced_files.add("new.py") + assert root.compilation_config.traced_files == {"original.py"} + + +def test_mtp_constructor_binds_the_explicit_fallback_config(monkeypatch): + from vllm.config import ( + CompilationConfig, + VllmConfig, + get_current_vllm_config, + set_current_vllm_config, + ) + from vllm.config.compilation import DynamicShapesType + + seen = [] + + class Predictor(nn.Module): + def __init__(self, *, vllm_config, **kwargs): + super().__init__() + seen.append((vllm_config, get_current_vllm_config())) + self.layers = nn.ModuleList() + self.make_empty_intermediate_tensors = lambda *args: None + + monkeypatch.setattr(mtp, "Qwen4ExpMultiTokenPredictor", Predictor) + monkeypatch.setattr(mtp, "_hc_mtp_prefill_sequence_qualified", lambda *args: True) + monkeypatch.setattr( + mtp, "get_pp_group", lambda: NS(world_size=1, is_last_rank=True) + ) + monkeypatch.setattr(mtp, "ParallelLMHead", lambda *args, **kwargs: nn.Identity()) + monkeypatch.setattr(mtp, "LogitsProcessor", lambda *args: nn.Identity()) + monkeypatch.setattr(mtp.Qwen4ExpMTP, "set_moe_parameters", lambda *args: None) + for name in ( + "enable_qwen4_exp_low_latency_gemm", + "enable_qwen38_sm70_fp16_gemv", + "enable_qwen38_sm70_fp16_fused_hc", + ): + monkeypatch.setattr(mtp, name, lambda *args: None) + root = VllmConfig(compilation_config=CompilationConfig(mode=0)) + root.model_config = NS( + hf_text_config=NS(vocab_size=32, hidden_size=2560, tie_word_embeddings=False), + dtype=torch.half, + ) + root.cache_config.mamba_cache_mode = "align" + with set_current_vllm_config(root): + mtp.Qwen4ExpMTP(vllm_config=root) + assert get_current_vllm_config() is root + explicit, current = seen[0] + assert current is explicit and current is not root + assert ( + current.compilation_config.dynamic_shapes_config.type + == DynamicShapesType.UNBACKED + ) + assert ( + root.compilation_config.dynamic_shapes_config.type == DynamicShapesType.BACKED + ) diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index b41f89036..80f9b2b23 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -89,6 +89,78 @@ from .sm70_fp16_hc import enable_qwen38_sm70_fp16_fused_hc logger = init_logger(__name__) +_MTP_TOKEN_AXES = { + "input_ids": 0, + "positions": -1, + "hidden_states": 0, + "inputs_embeds": 0, +} + + +def _hc_mtp_prefill_sequence_qualified(vllm_config, config): + parallel = vllm_config.parallel_config + return bool( + vllm_config.kernel_config.layer_execution.hc_mtp_prefill_sequence_parallel + and vllm_config.kernel_config.sm70_mtp_q8_0_experts + and not is_offload_process() + and current_platform.is_device_capability(70) + and vllm_config.model_config.dtype == torch.float16 + and (config.hc_count, config.hidden_size, config.hc_lowrank) == (4, 2560, 320) + and vllm_config.kernel_config.hc_weight_storage == "sharded" + and not vllm_config.kernel_config.sm70_hcx_diagnostics + and vllm_config.lora_config is None + and vllm_config.scheduler_config.max_num_seqs == 1 + and vllm_config.compilation_config.dynamic_shapes_config.type + == DynamicShapesType.BACKED + and ( + parallel.tensor_parallel_size, + parallel.pipeline_parallel_size, + parallel.decode_context_parallel_size, + parallel.prefill_context_parallel_size, + ) + == (4, 1, 1, 1) + ) + + +def _make_mtp_prefill_fallback_config(vllm_config, *, qualified): + if not qualified: + return vllm_config + # The row view handles the large startup batch, so the fallback can first + # compile at M1. BACKED tracing specializes that singleton. Keep this one + # owner's token axes unbacked; decode views retain their original config. + config = copy.copy(vllm_config) + config.compilation_config = copy.copy(vllm_config.compilation_config) + config.compilation_config.dynamic_shapes_config = copy.copy( + vllm_config.compilation_config.dynamic_shapes_config + ) + config.compilation_config.dynamic_shapes_config.type = DynamicShapesType.UNBACKED + config.compilation_config.cache_dir = "" + config.compilation_config.local_cache_dir = "" + config.compilation_config.traced_files = set() + return config + + +def _mtp_token_shape_invariants( + self, + input_ids, + positions, + hidden_states=None, + intermediate_tensors=None, + inputs_embeds=None, + spec_step_idx=0, + hc_sequence_parallel=False, +): + # A fallback can first compile on a singleton after the row-owned prefill + # view handles startup's large batch. Keep that singleton dynamic; guards + # are removed from the loaded AOT callable. + if hidden_states is not None: + torch._check(positions.size(-1) == hidden_states.size(0)) + if input_ids is not None: + torch._check(input_ids.size(0) == hidden_states.size(0)) + if inputs_embeds is not None: + torch._check(inputs_embeds.size(0) == hidden_states.size(0)) + + def _remap_ignored_layers( ignored_layers: list[str], mtp_start_layer_idx: int, @@ -313,7 +385,7 @@ def _make_draft_vllm_config( "intermediate_tensors": 0, "inputs_embeds": 0, "hidden_states": 0, - } + }, ) class Qwen4ExpMultiTokenPredictor(nn.Module): hf_to_vllm_mapper = _QWEN3_5_WEIGHTS_MAPPER | _HC_WEIGHTS_MAPPER @@ -454,6 +526,10 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): spec_step_idx: int = 0, hc_sequence_parallel: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | IntermediateTensors: + if torch.compiler.is_compiling(): + _mtp_token_shape_invariants( + self, input_ids, positions, hidden_states, inputs_embeds=inputs_embeds + ) hc_count = self.hc_count hidden_size = self.hidden_size @@ -679,11 +755,20 @@ class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): and speculative.method == "mtp" and get_pp_group().world_size == 1 ) - self.model = Qwen4ExpMultiTokenPredictor( - vllm_config=vllm_config, - prefix=maybe_prefix(prefix, "mtp"), - share_target_embed_tokens=self.share_target_io_weights, + self._hc_prefill_sequence_qualified = _hc_mtp_prefill_sequence_qualified( + vllm_config, config ) + predictor_config = _make_mtp_prefill_fallback_config( + vllm_config, qualified=self._hc_prefill_sequence_qualified + ) + # The compiler wrapper also binds the current configuration during + # construction. Keep that binding equivalent to the explicit input. + with set_current_vllm_config(predictor_config): + self.model = Qwen4ExpMultiTokenPredictor( + vllm_config=predictor_config, + prefix=maybe_prefix(prefix, "mtp"), + share_target_embed_tokens=self.share_target_io_weights, + ) if get_pp_group().is_last_rank: if self.share_target_io_weights: @@ -710,29 +795,6 @@ class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): enable_qwen38_sm70_fp16_fused_hc(self, config_dtype, vllm_config) object.__setattr__(self, "_sm70_decode_graph_model", None) object.__setattr__(self, "_hc_prefill_sequence_model", None) - parallel = vllm_config.parallel_config - self._hc_prefill_sequence_qualified = bool( - vllm_config.kernel_config.layer_execution.hc_mtp_prefill_sequence_parallel - and vllm_config.kernel_config.sm70_mtp_q8_0_experts - and not is_offload_process() - and current_platform.is_device_capability(70) - and vllm_config.model_config.dtype == torch.float16 - and (config.hc_count, config.hidden_size, config.hc_lowrank) - == (4, 2560, 320) - and vllm_config.kernel_config.hc_weight_storage == "sharded" - and not vllm_config.kernel_config.sm70_hcx_diagnostics - and vllm_config.lora_config is None - and vllm_config.scheduler_config.max_num_seqs == 1 - and vllm_config.compilation_config.dynamic_shapes_config.type - == DynamicShapesType.BACKED - and ( - parallel.tensor_parallel_size, - parallel.pipeline_parallel_size, - parallel.decode_context_parallel_size, - parallel.prefill_context_parallel_size, - ) - == (4, 1, 1, 1) - ) self._sm70_draft_head = None def prepare_loaded_linear_weights(self) -> None: