Merge pull request #56 from QuixiAI/fix/qwen35-gguf-mtp

Load embedded Qwen GGUF MTP heads without mutating the target

Author: Eric Hartford

perf/optimization_status.md

@@ -34372,6 +34372,43 @@ than PyPI 1.3.0), and the FP8 GEMMs.
   equal the dequantized HF-grouped reference to 2.4e-14 at TP 1/2/4/8.
   Live model parity and throughput remain unqualified (no Metal or MI300X
   on the box); the per-layer gather is new decode work on qwen38-q2kxl-1.
+
+
+## 2026-09-19 - Load embedded Qwen GGUF MTP heads without mutating the target
+
+- Status: retained for review.
+- Baseline: `fix/qwen35-gguf-gdn-layout`; forward-port from the B70 serving integration.
+- Change: Map appended NextN tensors to the draft, preserve dense/MoE dispatch and detach cached GGUF config before draft overrides; retain main vision and GPTQ support.
+- Correctness: Four embedded-MTP configuration/mapping tests plus eight inherited layout tests passed on CPU; hardware model restore and serving remain unqualified.
+- Results: no new throughput measurement or hardware qualification claimed.
+- Decision: submit as a focused PR; preserve current-main behavior outside this fix.
+- Raw artifacts: local test output; no traffic captures or model data committed.
+- Follow-up (2026-10-14), root cause of the duplicate MTP attention
+  registration: build_qwen35_config_from_gguf is @cache'd and an embedded
+  MTP draft parses the same file, so without the parser's deepcopy the
+  target and draft ModelConfig share one config object.
+  SpeculativeConfig.hf_config_override rewrites its model_type and
+  architectures in place, and get_model_architecture reads
+  hf_config.architectures at load time, so the worker built the TARGET as
+  Qwen3_5MTP (17 tensors, one Attention at mtp.layers.0.self_attn.attn);
+  the real draft then failed with "Duplicate layer name". Reproduced on a
+  synthetic 4+1-layer qwen35 GGUF (tiled GDN K=4/V=12, Q4_K ssm_out, one
+  NextN block) on the RTX 5090 by removing the deepcopy. With the deepcopy
+  the registration never repeats on either runner, so
+  `_clear_stale_mtp_attention_context` (which deleted the bogus target's
+  registrations and hid the wrong target model) is removed and
+  qwen3_5_mtp.py is back to main.
+- Follow-up (2026-10-14), embedding/lm_head: the draft name map loaded
+  token_embd (dequantized to bf16, ~2.5 GB on Qwen3.8-27B) and output,
+  which both runners then discard because MTP drafts always take the
+  target's modules (V1 _maybe_share_*, V2 load_eagle_model; Qwen3_5MTP has
+  no has_own_* flags), the same as the HF embedded-MTP path. Both are now
+  unmapped: the draft reads 15 tensors and dequantizes 0 modules. Same
+  seeded V2 completion before and after.
+- Correctness: synthetic GGUF serve smoke on V2 and V1 runners (MTP k=1,
+  prefix caching, seeded temperature 1.0 / top_p 0.95 / top_k 20) reaches
+  generation; random weights, so no quality or acceptance claim. No real
+  Qwen3.8 GGUF or throughput run. Raw (operator box): ~/.local/scratch/results/2026-10-14/qwen35-gguf-mtp/ (builder, runner, logs).
 ## 2026-09-19 - Authenticate registered Hugging Face downloads without losing resume headers
 
 - Status: retained for review.

tests/model_executor/test_qwen35_gguf_mtp.py

@@ -0,0 +1,216 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Focused contracts for the Qwen3.5 GGUF fused GDN projections."""
+
+from types import SimpleNamespace
+from unittest import mock
+
+from vllm.model_executor.model_loader.gguf_adapters.qwen35 import (
+    Qwen35GGUFAdapter,
+)
+from vllm.model_executor.models.qwen3_5 import (
+    _qwen35_mlp_kind,
+)
+from vllm.transformers_utils.configs.qwen3_5 import Qwen3_5TextConfig
+from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeTextConfig
+from vllm.transformers_utils.gguf_config_parser import GGUFConfigParser
+
+
+def test_qwen35_gguf_mtp_maps_only_the_embedded_nextn_block():
+    text_config = SimpleNamespace(num_hidden_layers=64)
+    hf_config = SimpleNamespace(
+        model_type="qwen3_5_mtp",
+        architectures=["Qwen3_5MTP"],
+        get_text_config=lambda: text_config,
+    )
+    adapter = Qwen35GGUFAdapter(hf_config)
+    name_map = adapter.build_name_map(SimpleNamespace(hf_config=hf_config))
+
+    # The draft shares the target's embedding and lm_head; it never loads
+    # (or dequantizes) its own copies.
+    assert "token_embd.weight" not in name_map
+    assert "output.weight" not in name_map
+    assert "output_norm.weight" not in name_map
+    assert name_map["blk.64.nextn.eh_proj.weight"] == "mtp.fc.weight"
+    assert name_map["blk.64.nextn.enorm.weight"] == ("mtp.pre_fc_norm_embedding.weight")
+    assert name_map["blk.64.nextn.hnorm.weight"] == "mtp.pre_fc_norm_hidden.weight"
+    assert name_map["blk.64.nextn.shared_head_norm.weight"] == "mtp.norm.weight"
+    assert name_map["blk.64.attn_q.weight"] == ("mtp.layers.0.self_attn.q_proj.weight")
+    assert name_map["blk.64.ffn_down.weight"] == ("mtp.layers.0.mlp.down_proj.weight")
+    assert not any(name.startswith("blk.63.") for name in name_map)
+
+
+def test_qwen35_mtp_model_type_preserves_dense_and_moe_mlp_dispatch():
+    dense = Qwen3_5TextConfig(num_hidden_layers=1)
+    dense.model_type = "qwen3_5_mtp"
+    dense.architectures = ["Qwen3_5MTP"]
+    assert _qwen35_mlp_kind(dense) == "dense"
+
+    moe = Qwen3_5MoeTextConfig(num_hidden_layers=1)
+    moe.model_type = "qwen3_5_mtp"
+    moe.architectures = ["Qwen3_5MoeMTP"]
+    assert _qwen35_mlp_kind(moe) == "moe"
+
+
+def test_gguf_parser_detaches_cached_config_for_embedded_mtp_override():
+    cached = Qwen3_5TextConfig(num_hidden_layers=64)
+    cached.architectures = ["Qwen3_5ForCausalLM"]
+    parser = GGUFConfigParser()
+
+    with (
+        mock.patch(
+            "vllm.transformers_utils.gguf_config_parser.gguf_architecture",
+            return_value="qwen35",
+        ),
+        mock.patch(
+            "vllm.transformers_utils.gguf_qwen35.build_qwen35_config_from_gguf",
+            return_value=cached,
+        ),
+    ):
+        _, target = parser.parse("model.gguf", trust_remote_code=False)
+        _, draft = parser.parse("model.gguf", trust_remote_code=False)
+
+    draft.model_type = "qwen3_5_mtp"
+    draft.architectures = ["Qwen3_5MTP"]
+
+    assert target is not draft
+    assert target is not cached
+    assert target.model_type == "qwen3_5_text"
+    assert target.architectures == ["Qwen3_5ForCausalLM"]
+
+
+def test_qwen35_gguf_mtp_draft_runs_on_the_target_embedding_and_lm_head():
+    """The GGUF draft never loads token_embd/output, so it must share.
+
+    Qwen3_5MTP declares no has_own_embed_tokens / has_own_lm_head, which is
+    what makes the V2 draft loader hand it the target's modules.
+    """
+    import torch.nn as nn
+
+    from vllm.model_executor.models.qwen3_5_mtp import Qwen3_5MTP
+    from vllm.v1.worker.gpu.spec_decode.eagle import utils as eagle_utils
+
+    draft = Qwen3_5MTP.__new__(Qwen3_5MTP)
+    nn.Module.__init__(draft)
+    draft.model = nn.Module()
+    draft.model.embed_tokens = nn.Embedding(4, 2)
+    draft.model.layers = nn.ModuleList([nn.Module()])
+    draft.lm_head = nn.Linear(2, 4)
+    target = nn.Module()
+    target.model = nn.Module()
+    target.model.embed_tokens = nn.Embedding(4, 2)
+    target.lm_head = nn.Linear(2, 4)
+    config = SimpleNamespace(
+        speculative_config=SimpleNamespace(
+            draft_model_config=object(), kv_cache_dtype=None
+        )
+    )
+
+    with (
+        mock.patch.object(eagle_utils, "get_model", return_value=draft),
+        mock.patch.object(
+            eagle_utils,
+            "get_pp_group",
+            return_value=SimpleNamespace(world_size=1),
+        ),
+    ):
+        loaded = eagle_utils.load_eagle_model(target, config)
+
+    assert loaded.model.embed_tokens is target.model.embed_tokens
+    assert loaded.lm_head is target.lm_head
+
+
+def _mtp_name_map(pp_size: int) -> dict[str, str]:
+    from vllm.config import VllmConfig, set_current_vllm_config
+
+    text_config = SimpleNamespace(num_hidden_layers=64)
+    hf_config = SimpleNamespace(
+        model_type="qwen3_5_mtp",
+        architectures=["Qwen3_5MTP"],
+        get_text_config=lambda: text_config,
+    )
+    vllm_config = VllmConfig()
+    object.__setattr__(vllm_config.parallel_config, "pipeline_parallel_size", pp_size)
+    adapter = Qwen35GGUFAdapter(hf_config)
+    with set_current_vllm_config(vllm_config):
+        return adapter.build_name_map(SimpleNamespace(hf_config=hf_config))
+
+
+def test_qwen35_gguf_mtp_loads_its_own_embedding_only_under_pp():
+    """PP=1: the runners share the target embedding, so token_embd stays
+    unmapped. PP>1: neither runner shares it, so the draft must load it
+    or embed through an uninitialised table. lm_head is shared either way
+    (the target's lives on the draft's last rank)."""
+    single = _mtp_name_map(1)
+    assert "token_embd.weight" not in single
+    assert "output.weight" not in single
+
+    pipelined = _mtp_name_map(2)
+    assert pipelined["token_embd.weight"] == "model.embed_tokens.weight"
+    assert "output.weight" not in pipelined
+    assert pipelined["blk.64.nextn.eh_proj.weight"] == "mtp.fc.weight"
+
+
+def _toy_mtp_draft_and_target():
+    import torch.nn as nn
+
+    from vllm.model_executor.models.qwen3_5_mtp import Qwen3_5MTP
+
+    draft = Qwen3_5MTP.__new__(Qwen3_5MTP)
+    nn.Module.__init__(draft)
+    draft.model = nn.Module()
+    draft.model.embed_tokens = nn.Embedding(4, 2)
+    draft.model.layers = nn.ModuleList([nn.Module()])
+    draft.lm_head = nn.Linear(2, 4)
+    target = nn.Module()
+    target.model = nn.Module()
+    target.model.embed_tokens = nn.Embedding(4, 2)
+    target.lm_head = nn.Linear(2, 4)
+    return draft, target
+
+
+def test_qwen35_gguf_mtp_v1_proposer_shares_target_embedding_and_lm_head():
+    """V1 mirror of the V2 load_eagle_model sharing contract."""
+    from vllm.v1.spec_decode import llm_base_proposer
+
+    draft, target = _toy_mtp_draft_and_target()
+    proposer = llm_base_proposer.SpecDecodeBaseProposer.__new__(
+        llm_base_proposer.SpecDecodeBaseProposer
+    )
+    proposer.model = draft
+    proposer.vllm_config = SimpleNamespace(speculative_config=None)
+    proposer.use_local_argmax_reduction = False
+    with mock.patch.object(
+        llm_base_proposer,
+        "get_pp_group",
+        return_value=SimpleNamespace(world_size=1),
+    ):
+        proposer._maybe_share_embeddings(target)
+        proposer._maybe_share_lm_head(target)
+
+    assert draft.model.embed_tokens is target.model.embed_tokens
+    assert draft.lm_head is target.lm_head
+
+
+def test_qwen35_gguf_mtp_v1_proposer_keeps_draft_embedding_under_pp():
+    """Under PP the V1 runner keeps the draft's own (loaded) embedding."""
+    from vllm.v1.spec_decode import llm_base_proposer
+
+    draft, target = _toy_mtp_draft_and_target()
+    own = draft.model.embed_tokens
+    proposer = llm_base_proposer.SpecDecodeBaseProposer.__new__(
+        llm_base_proposer.SpecDecodeBaseProposer
+    )
+    proposer.model = draft
+    proposer.vllm_config = SimpleNamespace(speculative_config=None)
+    proposer.use_local_argmax_reduction = False
+    with mock.patch.object(
+        llm_base_proposer,
+        "get_pp_group",
+        return_value=SimpleNamespace(world_size=2),
+    ):
+        proposer._maybe_share_embeddings(target)
+        proposer._maybe_share_lm_head(target)
+
+    assert draft.model.embed_tokens is own
+    assert draft.lm_head is target.lm_head

vllm/config/speculative.py

@@ -542,8 +542,16 @@ def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
             hf_config.update(
                 {"n_predict": n_predict, "architectures": ["Exaone4_5_MTP"]}
             )
-        if hf_config.model_type in ("qwen3_5", "qwen3_5_moe"):
-            is_moe = hf_config.model_type == "qwen3_5_moe"
+        if hf_config.model_type in (
+            "qwen3_5",
+            "qwen3_5_moe",
+            "qwen3_5_text",
+            "qwen3_5_moe_text",
+        ):
+            # GGUF text-only checkpoints build the text config directly,
+            # while HF multimodal checkpoints wrap it in qwen3_5. Both carry
+            # the same embedded MTP layout.
+            is_moe = hf_config.model_type in ("qwen3_5_moe", "qwen3_5_moe_text")
             hf_config.model_type = "qwen3_5_mtp"
             n_predict = getattr(hf_config, "mtp_num_hidden_layers", None)
             hf_config.update(

vllm/model_executor/model_loader/gguf_adapters/qwen35.py

@@ -145,6 +145,14 @@ def _vision_linear_modules(depth: int) -> list[str]:
     return names
 
 
+_MTP_EXTRA_RENAMES = {
+    "nextn.eh_proj.weight": "fc.weight",
+    "nextn.enorm.weight": "pre_fc_norm_embedding.weight",
+    "nextn.hnorm.weight": "pre_fc_norm_hidden.weight",
+    "nextn.shared_head_norm.weight": "norm.weight",
+}
+
+
 def _untile_v_head_runs(
     tensor: torch.Tensor, *, axis: int, repeats: int, runs: int, run_width: int
 ) -> torch.Tensor:
@@ -244,18 +252,63 @@ def _untile_packed_v_head_columns(
     )
 
 
+def _pipeline_parallel_size() -> int:
+    from vllm.config import get_current_vllm_config_or_none
+
+    vllm_config = get_current_vllm_config_or_none()
+    if vllm_config is None:
+        return 1
+    return int(vllm_config.parallel_config.pipeline_parallel_size)
+
+
 class Qwen35GGUFAdapter(GGUFWeightsAdapter):
     """The qwen35 hybrid target GGUF (dense Qwen3.5 family)."""
 
     def patch_hf_config(self, model_path: str, hf_config):
         # The config was built from this same GGUF; nothing to patch.
         return hf_config
 
+    @staticmethod
+    def _is_mtp_model(model_config: ModelConfig) -> bool:
+        config = model_config.hf_config
+        architectures = getattr(config, "architectures", None) or []
+        return (
+            getattr(config, "model_type", None) == "qwen3_5_mtp"
+            or "Qwen3_5MTP" in architectures
+            or "Qwen3_5MoeMTP" in architectures
+        )
+
     def build_name_map(self, model_config: ModelConfig) -> dict[str, str]:
         config = model_config.hf_config
         text_config = (
             config.get_text_config() if hasattr(config, "get_text_config") else config
         )
+        if self._is_mtp_model(model_config):
+            # The appended GGUF block is numbered after the target backbone,
+            # while Qwen3_5MTP numbers its compact draft layers from zero.
+            mtp_idx = int(text_config.num_hidden_layers)
+            # Without pipeline parallelism token_embd / output stay unmapped:
+            # both runners (V1 _maybe_share_embeddings/_maybe_share_lm_head,
+            # V2 load_eagle_model) replace the draft's embed_tokens and
+            # lm_head with the target's, so loading them would only
+            # dequantize a second ~2.5 GB embedding table and a second
+            # lm_head that are discarded right after.
+            # With PP > 1 neither runner shares the embedding (the target's
+            # embed_tokens lives on the first rank, the draft on the last),
+            # so the draft must load its own or run on an uninitialised
+            # table. lm_head is still shared under PP: both runners share it
+            # unconditionally and the target's lm_head lives on the same
+            # (last) rank as the draft, so output stays unmapped.
+            name_map: dict[str, str] = {}
+            if _pipeline_parallel_size() > 1:
+                name_map["token_embd.weight"] = "model.embed_tokens.weight"
+            per_layer = {**_COMMON_BLK_RENAMES, **_FULL_ATTN_BLK_RENAMES}
+            for gguf_part, hf_part in per_layer.items():
+                name_map[f"blk.{mtp_idx}.{gguf_part}"] = f"mtp.layers.0.{hf_part}"
+            for gguf_part, hf_part in _MTP_EXTRA_RENAMES.items():
+                name_map[f"blk.{mtp_idx}.{gguf_part}"] = f"mtp.{hf_part}"
+            return name_map
+
         # When the mmproj was found the config is the composite
         # Qwen3_5ForConditionalGeneration; the text half lives under
         # `language_model.` (Muse-Glimmer pattern).