vllm_mlx.patches.qwen3_5_mtp

Runtime MTP (Multi-Token Prediction) support for Qwen3.5 models.

View the complete module source at #L1-L512.

API details

Each callable below includes its exact signature, type annotations, inputs, defaults, return contract, documented exceptions, implementation source, and parsed docstring sections when the source provides them.

vllm_mlx.patches.qwen3_5_mtp

Runtime MTP (Multi-Token Prediction) support for Qwen3.5 models.

Qwen3.5 models may include a built-in MTP head that predicts token n+2 from hidden states + token n+1. MTP weights are added to the quantized MLX model via scripts/add_mtp_weights_qwen35.py.

Since mlx_lm's qwen3_5.py does NOT define MTP module/methods, this module provides: - inject_mtp_support(): dynamically creates MTP module, loads weights, and monkey-patches the model class with return_hidden, mtp_forward, and make_mtp_cache - validate_mtp_support(): checks whether a loaded model has working MTP

Supports both Dense (27B) and MoE (122B-A10B, 35B-A3B) architectures.

The actual MTP scheduling logic lives in

• vllm_mlx/scheduler.py (_install_mtp, _mtp_step, _mtp_next)

vllm_mlx.patches.qwen3_5_mtp.logger module-attribute

logger = logging.getLogger(__name__)

vllm_mlx.patches.qwen3_5_mtp._MTP_KEY_PREFIXES module-attribute

_MTP_KEY_PREFIXES = ('mtp.', 'language_model.mtp.')

vllm_mlx.patches.qwen3_5_mtp._QWEN_MTP_RMSNORM_WEIGHT_SUFFIXES module-attribute

_QWEN_MTP_RMSNORM_WEIGHT_SUFFIXES = ('input_layernorm.weight', 'post_attention_layernorm.weight', 'q_norm.weight', 'k_norm.weight', 'pre_fc_norm_hidden.weight', 'pre_fc_norm_embedding.weight', 'norm.weight')

vllm_mlx.patches.qwen3_5_mtp._QWEN_MTP_HIDDEN_STATE_MODES module-attribute

_QWEN_MTP_HIDDEN_STATE_MODES = frozenset({'post_norm', 'pre_norm'})

vllm_mlx.patches.qwen3_5_mtp._strip_mtp_key_prefix

_strip_mtp_key_prefix(key: str) -> str | None

Return an MTP-relative key for supported standalone shard layouts.

Source code in vllm_mlx/patches/qwen3_5_mtp.py

def _strip_mtp_key_prefix(key: str) -> str | None:
    """Return an MTP-relative key for supported standalone shard layouts."""
    for prefix in _MTP_KEY_PREFIXES:
        if key.startswith(prefix):
            return key.removeprefix(prefix)
    return None

vllm_mlx.patches.qwen3_5_mtp._resolve_qwen_mtp_hidden_state_mode

_resolve_qwen_mtp_hidden_state_mode(config: dict) -> str

Resolve the checkpoint's MTP hidden-state contract safely.

Source code in vllm_mlx/patches/qwen3_5_mtp.py

def _resolve_qwen_mtp_hidden_state_mode(config: dict) -> str:
    """Resolve the checkpoint's MTP hidden-state contract safely."""
    text_config = config.get("text_config", config)
    mode = text_config.get(
        "mtp_hidden_state_mode",
        config.get("mtp_hidden_state_mode", "post_norm"),
    )
    if not isinstance(mode, str) or mode not in _QWEN_MTP_HIDDEN_STATE_MODES:
        logger.warning(
            "[MTP inject] Unsupported mtp_hidden_state_mode=%r; using post_norm",
            mode,
        )
        return "post_norm"
    return mode

vllm_mlx.patches.qwen3_5_mtp._select_qwen_mtp_hidden_state

_select_qwen_mtp_hidden_state(mode: str, hidden_states, normed)

Select the representation expected by the checkpoint's MTP head.

Source code in vllm_mlx/patches/qwen3_5_mtp.py

def _select_qwen_mtp_hidden_state(mode: str, hidden_states, normed):
    """Select the representation expected by the checkpoint's MTP head."""
    return hidden_states if mode == "pre_norm" else normed

vllm_mlx.patches.qwen3_5_mtp._is_qwen_mtp_rmsnorm_weight

_is_qwen_mtp_rmsnorm_weight(key: str, weight) -> bool

Return True for MTP RMSNorm weights that use Qwen's offset convention.

Source code in vllm_mlx/patches/qwen3_5_mtp.py

def _is_qwen_mtp_rmsnorm_weight(key: str, weight) -> bool:
    """Return True for MTP RMSNorm weights that use Qwen's offset convention."""
    return weight.ndim == 1 and any(
        key.endswith(suffix) for suffix in _QWEN_MTP_RMSNORM_WEIGHT_SUFFIXES
    )

vllm_mlx.patches.qwen3_5_mtp._apply_qwen_mtp_rmsnorm_offset_fixups

_apply_qwen_mtp_rmsnorm_offset_fixups(mtp_weights: dict) -> int

Apply Qwen raw-offset RMSNorm fixups without double-shifting MLX weights.

Source code in vllm_mlx/patches/qwen3_5_mtp.py

def _apply_qwen_mtp_rmsnorm_offset_fixups(mtp_weights: dict) -> int:
    """Apply Qwen raw-offset RMSNorm fixups without double-shifting MLX weights."""
    norm_fixup_count = 0
    for key, weight in list(mtp_weights.items()):
        if not _is_qwen_mtp_rmsnorm_weight(key, weight):
            continue
        mean_val = weight.mean().item()
        if mean_val < 0.5:
            mtp_weights[key] = weight + 1.0
            norm_fixup_count += 1
    return norm_fixup_count

vllm_mlx.patches.qwen3_5_mtp._fixup_moe_mtp

_fixup_moe_mtp(mtp, inner_model, loaded_keys: set, mx) -> None

Fix missing weights in MoE MTP module.

MoE MTP checkpoints (122B, 35B) only contain: fc, q_proj, o_proj, shared_expert.*, and per-expert weights. Missing: - k_proj, v_proj → zero out (attention becomes no-op) - gate, shared_expert_gate → copy from main model's last full-attn layer - norms → already at identity (weight=1.0), no action needed

Source code in vllm_mlx/patches/qwen3_5_mtp.py

def _fixup_moe_mtp(mtp, inner_model, loaded_keys: set, mx) -> None:
    """Fix missing weights in MoE MTP module.

    MoE MTP checkpoints (122B, 35B) only contain: fc, q_proj, o_proj,
    shared_expert.*, and per-expert weights.  Missing:
    - k_proj, v_proj → zero out (attention becomes no-op)
    - gate, shared_expert_gate → copy from main model's last full-attn layer
    - norms → already at identity (weight=1.0), no action needed
    """
    import mlx.utils

    mtp_layer = mtp.layers[0]

    # Find last full-attention layer in main model for gate weights
    last_fa_layer = None
    for layer in reversed(inner_model.layers):
        if not layer.is_linear:
            last_fa_layer = layer
            break

    if last_fa_layer is None:
        logger.warning("[MTP fixup] No full-attention layer found in main model")
        return

    # Copy expert routing gate if not in checkpoint
    if "layers.0.mlp.gate.weight" not in loaded_keys:
        src = getattr(last_fa_layer.mlp, "gate", None)
        dst = getattr(mtp_layer.mlp, "gate", None)
        if src is not None and dst is not None:
            src_params = mlx.utils.tree_flatten(src.parameters())
            dst.load_weights(src_params)
            mx.eval(dst.parameters())
            logger.info("[MTP fixup] Copied mlp.gate from main model last layer")

    # Copy shared_expert_gate if not in checkpoint
    if "layers.0.mlp.shared_expert_gate.weight" not in loaded_keys:
        src = getattr(last_fa_layer.mlp, "shared_expert_gate", None)
        dst = getattr(mtp_layer.mlp, "shared_expert_gate", None)
        if src is not None and dst is not None:
            src_params = mlx.utils.tree_flatten(src.parameters())
            dst.load_weights(src_params)
            mx.eval(dst.parameters())
            logger.info(
                "[MTP fixup] Copied shared_expert_gate from main model last layer"
            )

    # Zero out k_proj and v_proj → attention becomes no-op
    attn = getattr(mtp_layer, "self_attn", None)
    if attn is None:
        return

    for proj_name in ("k_proj", "v_proj"):
        key = f"layers.0.self_attn.{proj_name}.weight"
        if key not in loaded_keys:
            proj = getattr(attn, proj_name, None)
            if proj is None:
                continue
            # For quantized layers: zero scales+biases → dequantized = 0
            if hasattr(proj, "scales"):
                proj.scales = mx.zeros_like(proj.scales)
                proj.biases = mx.zeros_like(proj.biases)
            else:
                proj.weight = mx.zeros_like(proj.weight)
            mx.eval(proj.parameters())
            logger.info(f"[MTP fixup] Zeroed {proj_name} (not in checkpoint)")

vllm_mlx.patches.qwen3_5_mtp.inject_mtp_support

inject_mtp_support(model: Any, model_path, config: dict) -> bool

Inject MTP module into a loaded Qwen3.5 model.

mlx_lm's qwen3_5.py does not define MTP layers, so we: 1. Create MTP module matching the weight structure 2. Quantize it to match the base model 3. Load MTP weights from model-mtp.safetensors 4. Monkey-patch Model with return_hidden, mtp_forward, make_mtp_cache

Parameters:

• model (Any) –

  A model loaded via mlx_lm (strict=False, MTP weights ignored)

• model_path –

  Path to model directory (contains model-mtp.safetensors)

• config (dict) –

  Parsed config.json dict

Returns:

• bool –

  True if MTP was successfully injected, False otherwise.

Source code in vllm_mlx/patches/qwen3_5_mtp.py

def inject_mtp_support(model: Any, model_path, config: dict) -> bool:
    """Inject MTP module into a loaded Qwen3.5 model.

    mlx_lm's qwen3_5.py does not define MTP layers, so we:
    1. Create MTP module matching the weight structure
    2. Quantize it to match the base model
    3. Load MTP weights from model-mtp.safetensors
    4. Monkey-patch Model with return_hidden, mtp_forward, make_mtp_cache

    Args:
        model: A model loaded via mlx_lm (strict=False, MTP weights ignored)
        model_path: Path to model directory (contains model-mtp.safetensors)
        config: Parsed config.json dict

    Returns:
        True if MTP was successfully injected, False otherwise.
    """
    import mlx.core as mx
    import mlx.nn as nn

    # Navigate nested config: text_config for VLM wrappers
    text_config = config.get("text_config", config)
    hidden_state_mode = _resolve_qwen_mtp_hidden_state_mode(config)
    num_mtp_layers = text_config.get("mtp_num_hidden_layers", 0)
    if num_mtp_layers == 0:
        # Fallback: check flat config for num_nextn_predict_layers
        num_mtp_layers = text_config.get(
            "num_nextn_predict_layers",
            config.get("num_nextn_predict_layers", 0),
        )
    if num_mtp_layers == 0:
        logger.info("[MTP inject] No MTP layers configured, skipping")
        return False

    model_path = Path(model_path)
    # Look for MTP weights in mtp/ subdirectory first (avoids mlx_vlm glob),
    # then fall back to model-mtp.safetensors in model dir.
    mtp_file = model_path / "mtp" / "weights.safetensors"
    if not mtp_file.exists():
        mtp_file = model_path / "model-mtp.safetensors"
    if not mtp_file.exists():
        logger.warning(f"[MTP inject] MTP weights not found in {model_path}")
        return False

    # Get model args — navigate VLM wrapper if needed
    # Model hierarchy: Model → language_model (TextModel) → model (Qwen3_5TextModel)
    text_model = model
    if hasattr(model, "language_model"):
        text_model = model.language_model

    args = text_model.args

    # When loaded via mlx_vlm, args may be a TextConfig object missing fields
    # that mlx_lm's TextModelArgs defines (rope_theta, partial_rotary_factor,
    # rope_scaling, etc.). Build a proper TextModelArgs from the config dict.
    from mlx_lm.models.qwen3_5 import TextModelArgs

    if not isinstance(args, TextModelArgs):
        logger.info("[MTP inject] Building TextModelArgs from config dict")
        args = TextModelArgs.from_dict(text_config)

    # Detect MoE vs Dense from args
    num_experts = getattr(args, "num_experts", 0)
    is_moe = num_experts > 0

    # Import model components
    from mlx_lm.models.base import create_attention_mask, create_ssm_mask
    from mlx_lm.models.cache import KVCache
    from mlx_lm.models.qwen3_5 import DecoderLayer

    logger.info(
        f"[MTP inject] Creating MTP module ({num_mtp_layers} layers, "
        f"{'MoE' if is_moe else 'Dense'}, hidden_state={hidden_state_mode})"
    )

    # MTP decoder uses full attention (not GatedDeltaNet).
    # layer_idx = full_attention_interval - 1 ensures is_linear=False.
    fa_idx = args.full_attention_interval - 1

    class _MTPModule(nn.Module):
        def __init__(self, args, n_layers):
            super().__init__()
            self.pre_fc_norm_hidden = nn.RMSNorm(
                args.hidden_size, eps=args.rms_norm_eps
            )
            self.pre_fc_norm_embedding = nn.RMSNorm(
                args.hidden_size, eps=args.rms_norm_eps
            )
            self.fc = nn.Linear(args.hidden_size * 2, args.hidden_size, bias=False)
            self.layers = [
                DecoderLayer(args, layer_idx=fa_idx) for _ in range(n_layers)
            ]
            self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)

    mtp = _MTPModule(args, num_mtp_layers)

    # --- Load MTP weights in BF16 (no quantization) ---
    # MTP head is extremely sensitive to quantization — even 4-bit destroys
    # prediction quality (0% acceptance).  Keep MTP in full precision.
    # See: https://github.com/vllm-project/vllm/issues/36331
    quant_config = text_config.get("quantization", config.get("quantization", {}))
    bits = quant_config.get("bits", 4) if quant_config else 4
    group_size = quant_config.get("group_size", 64) if quant_config else 64

    logger.info(
        f"[MTP inject] Loading weights from {mtp_file.name} (BF16, no quantization)"
    )
    raw = mx.load(str(mtp_file))
    raw_mtp = {
        clean: value
        for key, value in raw.items()
        if (clean := _strip_mtp_key_prefix(key)) is not None
    }
    del raw

    # Dequantize any quantized weight triplets (weight + scales + biases)
    mtp_weights: dict[str, mx.array] = {}
    processed = set()
    for key in sorted(raw_mtp.keys()):
        if key in processed:
            continue
        if key.endswith(".scales") or key.endswith(".biases"):
            continue

        scales_key = key.replace(".weight", ".scales")
        biases_key = key.replace(".weight", ".biases")

        if scales_key != key and scales_key in raw_mtp and biases_key in raw_mtp:
            # Quantized triplet → dequantize to BF16
            dq = mx.dequantize(
                raw_mtp[key],
                raw_mtp[scales_key],
                raw_mtp[biases_key],
                group_size=group_size,
                bits=bits,
            )
            mtp_weights[key] = dq
            processed.update([key, scales_key, biases_key])
        else:
            # Already FP (norms, fc, shared_expert_gate)
            mtp_weights[key] = raw_mtp[key]
            processed.add(key)
    del raw_mtp

    # --- Convert fused expert format to split format ---
    # Qwen3.6 MTP uses fused expert keys:
    #   layers.X.mlp.experts.gate_up_proj  [n_experts, 2*intermediate, hidden]
    #   layers.X.mlp.experts.down_proj     [n_experts, hidden, intermediate]
    # but mlx_lm's DecoderLayer expects split switch_mlp keys:
    #   layers.X.mlp.switch_mlp.gate_proj.weight  [n_experts, intermediate, hidden]
    #   layers.X.mlp.switch_mlp.up_proj.weight    [n_experts, intermediate, hidden]
    #   layers.X.mlp.switch_mlp.down_proj.weight  [n_experts, hidden, intermediate]
    for key in list(mtp_weights.keys()):
        if ".mlp.experts.gate_up_proj" in key:
            prefix = key.replace(".mlp.experts.gate_up_proj", "")
            w = mtp_weights.pop(key)
            intermediate = w.shape[1] // 2
            gate_key = f"{prefix}.mlp.switch_mlp.gate_proj.weight"
            up_key = f"{prefix}.mlp.switch_mlp.up_proj.weight"
            mtp_weights[gate_key] = w[:, :intermediate, :]
            mtp_weights[up_key] = w[:, intermediate:, :]
            logger.info(
                "[MTP inject] Split fused experts.gate_up_proj -> "
                "switch_mlp.{gate_proj,up_proj}"
            )
        elif ".mlp.experts.down_proj" in key:
            prefix = key.replace(".mlp.experts.down_proj", "")
            w = mtp_weights.pop(key)
            down_key = f"{prefix}.mlp.switch_mlp.down_proj.weight"
            mtp_weights[down_key] = w
            logger.info(
                "[MTP inject] Renamed experts.down_proj -> switch_mlp.down_proj"
            )

    # --- Fixup RMSNorm weights: HuggingFace offset convention ---
    # Qwen3.5/3.6 models store RMSNorm weights as offsets (actual = 1 + stored).
    # The main model's sanitize() handles this, but MTP weights bypass sanitize.
    # Detect raw-offset weights (mean < 0.5) and apply +1.0; skip if already
    # in actual-gamma space (as produced by add_mtp_weights_qwen35.py).
    norm_fixup_count = _apply_qwen_mtp_rmsnorm_offset_fixups(mtp_weights)
    if norm_fixup_count > 0:
        logger.info(
            f"[MTP inject] Applied +1.0 RMSNorm offset to {norm_fixup_count} "
            f"norm weights (raw HF offset detected)"
        )

    mtp.load_weights(list(mtp_weights.items()), strict=False)
    mx.eval(mtp.parameters())

    dq_count = sum(1 for k in mtp_weights if not k.endswith((".scales", ".biases")))
    has_quantized = any(k.endswith(".scales") for k in processed)
    mode = "dequantized from quantized" if has_quantized else "native BF16"
    logger.info(f"[MTP inject] Loaded {dq_count} MTP weight tensors ({mode})")

    # --- Step 4: Fix missing MoE MTP weights ---
    # MoE checkpoints lack: k_proj, v_proj, gate, shared_expert_gate, norms.
    # Norms default to identity (weight=1.0) which is correct.
    # k_proj/v_proj: zero out → attention becomes no-op, MLP does prediction.
    # gate/shared_expert_gate: copy from main model's last full-attention layer.
    if is_moe:
        loaded_key_set = set(mtp_weights.keys())
        _fixup_moe_mtp(mtp, text_model.model, loaded_key_set, mx)

    # --- Attach MTP and monkey-patch model class ---
    text_model.mtp = mtp

    original_class = text_model.__class__

    class _Qwen3_5MTP(original_class):
        """Qwen3.5 with MTP support (injected at runtime)."""

        def __call__(
            self,
            inputs,
            cache=None,
            return_hidden: bool = False,
            input_embeddings=None,
            **kwargs,
        ):
            inner = self.model
            if input_embeddings is not None:
                hidden_states = input_embeddings
            else:
                hidden_states = inner.embed_tokens(inputs)

            if cache is None:
                cache = [None] * len(inner.layers)

            fa_mask = create_attention_mask(hidden_states, cache[inner.fa_idx])
            ssm_mask = create_ssm_mask(hidden_states, cache[inner.ssm_idx])

            for layer, c in zip(inner.layers, cache):
                mask = ssm_mask if layer.is_linear else fa_mask
                hidden_states = layer(hidden_states, mask=mask, cache=c)

            normed = inner.norm(hidden_states)

            if self.args.tie_word_embeddings:
                out = inner.embed_tokens.as_linear(normed)
            else:
                out = self.lm_head(normed)

            if return_hidden:
                return out, _select_qwen_mtp_hidden_state(
                    hidden_state_mode,
                    hidden_states,
                    normed,
                )
            return out

        def mtp_forward(
            self,
            hidden_states,
            next_token_ids,
            cache=None,
            mtp_cache=None,
        ):
            """Run MTP head: predict token n+2 from hidden states + token n+1."""
            input_embeds = self.model.embed_tokens(next_token_ids)
            e = self.mtp.pre_fc_norm_embedding(input_embeds)
            h = self.mtp.pre_fc_norm_hidden(hidden_states)
            x = self.mtp.fc(mx.concatenate([e, h], axis=-1))

            layer = self.mtp.layers[0]
            c = mtp_cache[0] if mtp_cache else None
            mask = create_attention_mask(x, c)
            x = layer(x, mask=mask, cache=c)

            x = self.mtp.norm(x)

            if self.args.tie_word_embeddings:
                return self.model.embed_tokens.as_linear(x)
            return self.lm_head(x)

        def make_mtp_cache(self):
            """Create KV cache for MTP layers."""
            if self.mtp is None:
                return None
            return [KVCache() for _ in self.mtp.layers]

    text_model.__class__ = _Qwen3_5MTP
    logger.info("[MTP inject] Model class patched with MTP support")

    # If we patched the inner language_model, also expose MTP on the outer Model
    if hasattr(model, "language_model") and model.language_model is text_model:
        model.mtp = mtp

    return True

vllm_mlx.patches.qwen3_5_mtp.validate_mtp_support

validate_mtp_support(model: Any) -> bool

Validate that a loaded model has working MTP support.

Checks: 1. model.mtp exists and is not None 2. model.mtp has layers with loaded weights 3. model has return_hidden support in call 4. model has mtp_forward method 5. model has make_mtp_cache method

Parameters:

• model (Any) –

  A model loaded via mlx_lm.load()

Returns:

• bool –

  True if MTP is fully functional, False otherwise.

Source code in vllm_mlx/patches/qwen3_5_mtp.py

def validate_mtp_support(model: Any) -> bool:
    """Validate that a loaded model has working MTP support.

    Checks:
    1. model.mtp exists and is not None
    2. model.mtp has layers with loaded weights
    3. model has return_hidden support in __call__
    4. model has mtp_forward method
    5. model has make_mtp_cache method

    Args:
        model: A model loaded via mlx_lm.load()

    Returns:
        True if MTP is fully functional, False otherwise.
    """
    # Navigate to text model if VLM wrapper
    text_model = model
    if hasattr(model, "language_model"):
        text_model = model.language_model

    mtp = getattr(text_model, "mtp", None)
    if mtp is None:
        args = getattr(text_model, "args", None)
        if args is not None:
            num_mtp = getattr(args, "mtp_num_hidden_layers", 0)
            if num_mtp == 0:
                num_mtp = getattr(args, "num_nextn_predict_layers", 0)
            if num_mtp > 0:
                logger.warning(
                    "[MTP] Model config has MTP layers=%d but model.mtp is None. "
                    "Run scripts/add_mtp_weights_qwen35.py to add weights.",
                    num_mtp,
                )
        return False

    mtp_layers = getattr(mtp, "layers", [])
    if not mtp_layers:
        logger.warning("[MTP] model.mtp exists but has no layers.")
        return False

    import inspect

    call_sig = inspect.signature(type(text_model).__call__)
    if "return_hidden" not in call_sig.parameters:
        logger.warning("[MTP] Model.__call__ does not accept return_hidden parameter.")
        return False

    if not hasattr(text_model, "mtp_forward") or not callable(text_model.mtp_forward):
        logger.warning("[MTP] Model does not have mtp_forward() method.")
        return False

    if not hasattr(text_model, "make_mtp_cache") or not callable(
        text_model.make_mtp_cache
    ):
        logger.warning("[MTP] Model does not have make_mtp_cache() method.")
        return False

    logger.info(
        "[MTP] Qwen3.5 model has working MTP support: %d MTP layer(s)",
        len(mtp_layers),
    )
    return True

Complete contract reference

Expand any definition for its exact inputs, annotations, defaults, return contract, directly raised exceptions, source-grounded behavior, and immutable line link. This section includes private and nested definitions that ordinary API generators omit.

vllm_mlx.patches.qwen3_5_mtp._strip_mtp_key_prefix · function

vllm_mlx.patches.qwen3_5_mtp._strip_mtp_key_prefix(key: str) -> str | None

Return an MTP-relative key for supported standalone shard layouts.

Parameters

Name	Type	Required	Default	Description
key	str	yes	none	Required positional or keyword input.

Returns

• Type: str | None
• Direct return expressions: key.removeprefix(prefix); None

Exceptions and behavior

Function _strip_mtp_key_prefix calls key.startswith, key.removeprefix; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L31-L36.

vllm_mlx.patches.qwen3_5_mtp._resolve_qwen_mtp_hidden_state_mode · function

vllm_mlx.patches.qwen3_5_mtp._resolve_qwen_mtp_hidden_state_mode(config: dict) -> str

Resolve the checkpoint's MTP hidden-state contract safely.

Parameters

Name	Type	Required	Default	Description
config	dict	yes	none	Required positional or keyword input.

Returns

• Type: str
• Direct return expressions: 'post_norm'; mode

Exceptions and behavior

Function _resolve_qwen_mtp_hidden_state_mode calls config.get, text_config.get, isinstance, logger.warning; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L52-L65.

vllm_mlx.patches.qwen3_5_mtp._select_qwen_mtp_hidden_state · function

vllm_mlx.patches.qwen3_5_mtp._select_qwen_mtp_hidden_state(mode: str, hidden_states, normed) -> not annotated

Select the representation expected by the checkpoint's MTP head.

Parameters

Name	Type	Required	Default	Description
mode	str	yes	none	Required positional or keyword input.
hidden_states	not annotated	yes	none	Required positional or keyword input.
normed	not annotated	yes	none	Required positional or keyword input.

Returns

• Type: not annotated
• Direct return expressions: hidden_states if mode == 'pre_norm' else normed

Exceptions and behavior

Function _select_qwen_mtp_hidden_state returns hidden_states if mode == 'pre_norm' else normed. No direct raise statement appears in this definition.

View source #L68-L70.

vllm_mlx.patches.qwen3_5_mtp._is_qwen_mtp_rmsnorm_weight · function

vllm_mlx.patches.qwen3_5_mtp._is_qwen_mtp_rmsnorm_weight(key: str, weight) -> bool

Return True for MTP RMSNorm weights that use Qwen's offset convention.

Parameters

Name	Type	Required	Default	Description
key	str	yes	none	Required positional or keyword input.
weight	not annotated	yes	none	Required positional or keyword input.

Returns

• Type: bool
• Direct return expressions: weight.ndim == 1 and any((key.endswith(suffix) for suffix in _QWEN_MTP_RMSNORM_WEIGHT_SUFFIXES))

Exceptions and behavior

Function _is_qwen_mtp_rmsnorm_weight calls any, key.endswith; returns weight.ndim == 1 and any((key.endswith(suffix) for suffix in _QWEN_MTP_RMSNORM_WEIGHT_SUFFIXES)). No direct raise statement appears in this definition.

View source #L73-L77.

vllm_mlx.patches.qwen3_5_mtp._apply_qwen_mtp_rmsnorm_offset_fixups · function

vllm_mlx.patches.qwen3_5_mtp._apply_qwen_mtp_rmsnorm_offset_fixups(mtp_weights: dict) -> int

Apply Qwen raw-offset RMSNorm fixups without double-shifting MLX weights.

Parameters

Name	Type	Required	Default	Description
mtp_weights	dict	yes	none	Required positional or keyword input.

Returns

• Type: int
• Direct return expressions: norm_fixup_count

Exceptions and behavior

Function _apply_qwen_mtp_rmsnorm_offset_fixups calls list, mtp_weights.items, _is_qwen_mtp_rmsnorm_weight, weight.mean().item; returns norm_fixup_count. No direct raise statement appears in this definition.

View source #L80-L90.

vllm_mlx.patches.qwen3_5_mtp._fixup_moe_mtp · function

vllm_mlx.patches.qwen3_5_mtp._fixup_moe_mtp(mtp, inner_model, loaded_keys: set, mx) -> None

Fix missing weights in MoE MTP module.

Parameters

Name	Type	Required	Default	Description
mtp	not annotated	yes	none	Required positional or keyword input.
inner_model	not annotated	yes	none	Required positional or keyword input.
loaded_keys	set	yes	none	Required positional or keyword input.
mx	not annotated	yes	none	Required positional or keyword input.

Returns

• Type: None
• Direct return expressions: None

Exceptions and behavior

Function _fixup_moe_mtp calls reversed, logger.warning, getattr, mlx.utils.tree_flatten; returns None. No direct raise statement appears in this definition.

View source #L93-L157.

vllm_mlx.patches.qwen3_5_mtp.inject_mtp_support · function

vllm_mlx.patches.qwen3_5_mtp.inject_mtp_support(model: Any, model_path, config: dict) -> bool

Inject MTP module into a loaded Qwen3.5 model.

Parameters

Name	Type	Required	Default	Description
model	Any	yes	none	A model loaded via mlx_lm (strict=False, MTP weights ignored)
model_path	not annotated	yes	none	Path to model directory (contains model-mtp.safetensors)
config	dict	yes	none	Parsed config.json dict

Returns

• Type: bool
• Direct return expressions: False; True

Exceptions and behavior

Function inject_mtp_support calls config.get, _resolve_qwen_mtp_hidden_state_mode, text_config.get, logger.info; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L160-L447.

vllm_mlx.patches.qwen3_5_mtp.inject_mtp_support._MTPModule · nested class

vllm_mlx.patches.qwen3_5_mtp.inject_mtp_support._MTPModule(args, n_layers)

Nested Class inject_mtp_support._MTPModule derives from nn.Module and declares 1 direct member(s).

Parameters

Name	Type	Required	Default	Description
args	not annotated	yes	none	Required positional or keyword input.
n_layers	not annotated	yes	none	Required positional or keyword input.

Returns

• Constructs: vllm_mlx.patches.qwen3_5_mtp.inject_mtp_support._MTPModule

Exceptions and behavior

Nested Class inject_mtp_support._MTPModule derives from nn.Module and declares 1 direct member(s). No direct raise statement appears in this definition.

View source #L239-L252.

vllm_mlx.patches.qwen3_5_mtp.inject_mtp_support._MTPModule.__init__ · nested function

vllm_mlx.patches.qwen3_5_mtp.inject_mtp_support._MTPModule.__init__(args, n_layers) -> not annotated

Nested Function inject_mtp_support._MTPModule.__init__ updates self.pre_fc_norm_hidden, self.pre_fc_norm_embedding, self.fc, self.layers; calls super().__init__, super, nn.RMSNorm, nn.Linear.

Parameters

Name	Type	Required	Default	Description
args	not annotated	yes	none	Required positional or keyword input.
n_layers	not annotated	yes	none	Required positional or keyword input.

Returns

• Type: not annotated

Exceptions and behavior

Nested Function inject_mtp_support._MTPModule.__init__ updates self.pre_fc_norm_hidden, self.pre_fc_norm_embedding, self.fc, self.layers; calls super().__init__, super, nn.RMSNorm, nn.Linear. No direct raise statement appears in this definition.

View source #L240-L252.

vllm_mlx.patches.qwen3_5_mtp.inject_mtp_support._Qwen3_5MTP · nested class

vllm_mlx.patches.qwen3_5_mtp.inject_mtp_support._Qwen3_5MTP()

Qwen3.5 with MTP support (injected at runtime).

Parameters

This callable has no explicit inputs.

Returns

• Constructs: vllm_mlx.patches.qwen3_5_mtp.inject_mtp_support._Qwen3_5MTP

Exceptions and behavior

Nested Class inject_mtp_support._Qwen3_5MTP derives from original_class and declares 3 direct member(s). No direct raise statement appears in this definition.

View source #L368-L438.

vllm_mlx.patches.qwen3_5_mtp.inject_mtp_support._Qwen3_5MTP.__call__ · nested function

vllm_mlx.patches.qwen3_5_mtp.inject_mtp_support._Qwen3_5MTP.__call__(inputs, cache = None, return_hidden: bool = False, input_embeddings = None, **kwargs) -> not annotated

Nested Function inject_mtp_support._Qwen3_5MTP.__call__ calls inner.embed_tokens, len, create_attention_mask, create_ssm_mask; has 2 explicit return paths.

Parameters

Name	Type	Required	Default	Description
inputs	not annotated	yes	none	Required positional or keyword input.
cache	not annotated	no	None	Optional positional or keyword input; defaults to None.
return_hidden	bool	no	False	Optional positional or keyword input; defaults to False.
input_embeddings	not annotated	no	None	Optional positional or keyword input; defaults to None.
**kwargs	not annotated	no	none	Additional variadic keyword inputs accepted by this callable.

Returns

• Type: not annotated
• Direct return expressions: (out, _select_qwen_mtp_hidden_state(hidden_state_mode, hidden_states, normed)); out

Exceptions and behavior

Nested Function inject_mtp_support._Qwen3_5MTP.__call__ calls inner.embed_tokens, len, create_attention_mask, create_ssm_mask; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L371-L408.

vllm_mlx.patches.qwen3_5_mtp.inject_mtp_support._Qwen3_5MTP.mtp_forward · nested function

vllm_mlx.patches.qwen3_5_mtp.inject_mtp_support._Qwen3_5MTP.mtp_forward(hidden_states, next_token_ids, cache = None, mtp_cache = None) -> not annotated

Run MTP head: predict token n+2 from hidden states + token n+1.

Parameters

Name	Type	Required	Default	Description
hidden_states	not annotated	yes	none	Required positional or keyword input.
next_token_ids	not annotated	yes	none	Required positional or keyword input.
cache	not annotated	no	None	Optional positional or keyword input; defaults to None.
mtp_cache	not annotated	no	None	Optional positional or keyword input; defaults to None.

Returns

• Type: not annotated
• Direct return expressions: self.model.embed_tokens.as_linear(x); self.lm_head(x)

Exceptions and behavior

Nested Function inject_mtp_support._Qwen3_5MTP.mtp_forward calls self.model.embed_tokens, self.mtp.pre_fc_norm_embedding, self.mtp.pre_fc_norm_hidden, self.mtp.fc; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L410-L432.

vllm_mlx.patches.qwen3_5_mtp.inject_mtp_support._Qwen3_5MTP.make_mtp_cache · nested function

vllm_mlx.patches.qwen3_5_mtp.inject_mtp_support._Qwen3_5MTP.make_mtp_cache() -> not annotated

Create KV cache for MTP layers.

Parameters

This callable has no explicit inputs.

Returns

• Type: not annotated
• Direct return expressions: None; [KVCache() for _ in self.mtp.layers]

Exceptions and behavior

Nested Function inject_mtp_support._Qwen3_5MTP.make_mtp_cache calls KVCache; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L434-L438.

vllm_mlx.patches.qwen3_5_mtp.validate_mtp_support · function

vllm_mlx.patches.qwen3_5_mtp.validate_mtp_support(model: Any) -> bool

Validate that a loaded model has working MTP support.

Parameters

Name	Type	Required	Default	Description
model	Any	yes	none	A model loaded via mlx_lm.load()

Returns

• Type: bool
• Direct return expressions: False; True

Exceptions and behavior

Function validate_mtp_support calls hasattr, getattr, logger.warning, inspect.signature; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L450-L512.

Complete symbol map

This map also includes private definitions and nested helpers. The signature column exposes every explicit input even when an internal helper has no dedicated parameter prose.

Symbol	Kind	Signature and inputs	What it does	Source
_strip_mtp_key_prefix	function	_strip_mtp_key_prefix(key: str) -> str \| None	Return an MTP-relative key for supported standalone shard layouts.	#L31-L36
_resolve_qwen_mtp_hidden_state_mode	function	_resolve_qwen_mtp_hidden_state_mode(config: dict) -> str	Resolve the checkpoint's MTP hidden-state contract safely.	#L52-L65
_select_qwen_mtp_hidden_state	function	_select_qwen_mtp_hidden_state(mode: str, hidden_states, normed) -> not annotated	Select the representation expected by the checkpoint's MTP head.	#L68-L70
_is_qwen_mtp_rmsnorm_weight	function	_is_qwen_mtp_rmsnorm_weight(key: str, weight) -> bool	Return True for MTP RMSNorm weights that use Qwen's offset convention.	#L73-L77
_apply_qwen_mtp_rmsnorm_offset_fixups	function	_apply_qwen_mtp_rmsnorm_offset_fixups(mtp_weights: dict) -> int	Apply Qwen raw-offset RMSNorm fixups without double-shifting MLX weights.	#L80-L90
_fixup_moe_mtp	function	_fixup_moe_mtp(mtp, inner_model, loaded_keys: set, mx) -> None	Fix missing weights in MoE MTP module.	#L93-L157
inject_mtp_support	function	inject_mtp_support(model: Any, model_path, config: dict) -> bool	Inject MTP module into a loaded Qwen3.5 model.	#L160-L447
inject_mtp_support._MTPModule	nested class	inject_mtp_support._MTPModule(args, n_layers)	Nested Class inject_mtp_support._MTPModule derives from nn.Module and declares 1 direct member(s).	#L239-L252
inject_mtp_support._MTPModule.__init__	nested function	inject_mtp_support._MTPModule.__init__(args, n_layers) -> not annotated	Nested Function inject_mtp_support._MTPModule.__init__ updates self.pre_fc_norm_hidden, self.pre_fc_norm_embedding, self.fc, self.layers; calls super().__init__, super, nn.RMSNorm, nn.Linear.	#L240-L252
inject_mtp_support._Qwen3_5MTP	nested class	inject_mtp_support._Qwen3_5MTP()	Qwen3.5 with MTP support (injected at runtime).	#L368-L438
inject_mtp_support._Qwen3_5MTP.__call__	nested function	inject_mtp_support._Qwen3_5MTP.__call__(inputs, cache = None, return_hidden: bool = False, input_embeddings = None, **kwargs) -> not annotated	Nested Function inject_mtp_support._Qwen3_5MTP.__call__ calls inner.embed_tokens, len, create_attention_mask, create_ssm_mask; has 2 explicit return paths.	#L371-L408
inject_mtp_support._Qwen3_5MTP.mtp_forward	nested function	inject_mtp_support._Qwen3_5MTP.mtp_forward(hidden_states, next_token_ids, cache = None, mtp_cache = None) -> not annotated	Run MTP head: predict token n+2 from hidden states + token n+1.	#L410-L432
inject_mtp_support._Qwen3_5MTP.make_mtp_cache	nested function	inject_mtp_support._Qwen3_5MTP.make_mtp_cache() -> not annotated	Create KV cache for MTP layers.	#L434-L438
validate_mtp_support	function	validate_mtp_support(model: Any) -> bool	Validate that a loaded model has working MTP support.	#L450-L512