[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

@@ -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、长文任务、

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
+    )

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: