vllm_mlx.patches.qwen3_next_mtp

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

View the complete module source at #L1-L261.

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_next_mtp

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

Qwen3-Next 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.py.

Since mlx_lm's qwen3_next.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

The actual MTP scheduling logic lives in

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

vllm_mlx.patches.qwen3_next_mtp.logger module-attribute

logger = logging.getLogger(__name__)

vllm_mlx.patches.qwen3_next_mtp.inject_mtp_support

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

Inject MTP module into a loaded Qwen3-Next model.

mlx_lm's qwen3_next.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_next_mtp.py

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

    mlx_lm's qwen3_next.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

    num_mtp_layers = config.get("num_nextn_predict_layers", 0)
    if num_mtp_layers == 0:
        logger.info("[MTP inject] num_nextn_predict_layers=0, skipping")
        return False

    model_path = Path(model_path)
    mtp_file = model_path / "model-mtp.safetensors"
    if not mtp_file.exists():
        logger.warning(f"[MTP inject] model-mtp.safetensors not found in {model_path}")
        return False

    args = model.args

    # 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_next import Qwen3NextDecoderLayer

    # --- Step 1: Create MTP module ---
    logger.info(f"[MTP inject] Creating MTP module ({num_mtp_layers} layers)")

    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)
            # MTP decoder uses full attention (not linear/delta-net)
            fa_idx = args.full_attention_interval - 1
            self.layers = [
                Qwen3NextDecoderLayer(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)

    # --- Step 2: Quantize MTP module to match base model ---
    quant_config = config.get("quantization", {})
    if quant_config:
        bits = quant_config.get("bits", 6)
        group_size = quant_config.get("group_size", 64)

        def _mtp_quant_pred(path, module):
            # Only quantize Linear layers
            if not isinstance(module, nn.Linear):
                return False
            # fc kept as FP (concat projection)
            if path == "fc":
                return False
            # Gate layers at 8-bit (matching base model)
            if path.endswith("mlp.gate") or path.endswith("shared_expert_gate"):
                return {"group_size": 64, "bits": 8}
            return True

        nn.quantize(
            mtp, group_size=group_size, bits=bits, class_predicate=_mtp_quant_pred
        )
        logger.info(f"[MTP inject] Quantized MTP: {bits}-bit, group_size={group_size}")

    # --- Step 3: Load MTP weights ---
    logger.info(f"[MTP inject] Loading weights from {mtp_file.name}")
    raw = mx.load(str(mtp_file))
    mtp_weights = {
        k.removeprefix("mtp."): v for k, v in raw.items() if k.startswith("mtp.")
    }
    mtp.load_weights(list(mtp_weights.items()), strict=False)
    mx.eval(mtp.parameters())
    logger.info(f"[MTP inject] Loaded {len(mtp_weights)} MTP weight tensors")

    # --- Step 4: Attach MTP and monkey-patch model class ---
    model.mtp = mtp

    original_class = model.__class__

    class _Qwen3NextMTP(original_class):
        """Qwen3-Next with MTP support (injected at runtime)."""

        def __call__(
            self,
            inputs,
            cache=None,
            return_hidden: bool = False,
        ):
            inner = self.model
            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, hidden_states  # pre-norm hidden states
            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)
            h = self.mtp.pre_fc_norm_hidden(hidden_states)
            e = self.mtp.pre_fc_norm_embedding(input_embeds)
            x = self.mtp.fc(mx.concatenate([h, e], 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]

    model.__class__ = _Qwen3NextMTP
    logger.info("[MTP inject] Model class patched with MTP support")
    return True

vllm_mlx.patches.qwen3_next_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 (MTP module instantiated) 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_next_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 (MTP module instantiated)
    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.
    """
    # Check 1: MTP module exists
    mtp = getattr(model, "mtp", None)
    if mtp is None:
        num_mtp = 0
        # Try model.args (Qwen3-Next) and model.language_model.args (Qwen3.5)
        args = getattr(model, "args", None)
        if args is None:
            lm = getattr(model, "language_model", None)
            if lm is not None:
                args = getattr(lm, "args", None)
        if args is not None:
            num_mtp = getattr(args, "num_nextn_predict_layers", 0)
        if num_mtp > 0:
            logger.warning(
                "[MTP] Model config has num_nextn_predict_layers=%d but "
                "model.mtp is None. MTP weights may not be in the model files. "
                "Run scripts/add_mtp_weights.py to add them.",
                num_mtp,
            )
        else:
            logger.info(
                "[MTP] Model does not have MTP config " "(num_nextn_predict_layers=0)."
            )
        return False

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

    # Check 3: return_hidden support
    import inspect

    call_sig = inspect.signature(type(model).__call__)
    if "return_hidden" not in call_sig.parameters:
        logger.warning(
            "[MTP] Model.__call__ does not accept return_hidden parameter. "
            "The mlx_lm model implementation may be outdated."
        )
        return False

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

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

    # All checks passed
    args = getattr(model, "args", None)
    num_layers = getattr(args, "num_nextn_predict_layers", 0)
    logger.info(
        "[MTP] Model has working MTP support: "
        "%s MTP layer(s), %d predictor decoder layer(s)",
        num_layers,
        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_next_mtp.inject_mtp_support · function

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

Inject MTP module into a loaded Qwen3-Next 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, logger.info, Path, mtp_file.exists; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L27-L181.

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

vllm_mlx.patches.qwen3_next_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_next_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 #L68-L83.

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

vllm_mlx.patches.qwen3_next_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 #L69-L83.

vllm_mlx.patches.qwen3_next_mtp.inject_mtp_support._mtp_quant_pred · nested function

vllm_mlx.patches.qwen3_next_mtp.inject_mtp_support._mtp_quant_pred(path, module) -> not annotated

Nested Function inject_mtp_support._mtp_quant_pred calls isinstance, path.endswith; has 3 explicit return paths.

Parameters

Name	Type	Required	Default	Description
path	not annotated	yes	none	Required positional or keyword input.
module	not annotated	yes	none	Required positional or keyword input.

Returns

• Type: not annotated
• Direct return expressions: False; {'group_size': 64, 'bits': 8}; True

Exceptions and behavior

Nested Function inject_mtp_support._mtp_quant_pred calls isinstance, path.endswith; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L93-L103.

vllm_mlx.patches.qwen3_next_mtp.inject_mtp_support._Qwen3NextMTP · nested class

vllm_mlx.patches.qwen3_next_mtp.inject_mtp_support._Qwen3NextMTP()

Qwen3-Next with MTP support (injected at runtime).

Parameters

This callable has no explicit inputs.

Returns

• Constructs: vllm_mlx.patches.qwen3_next_mtp.inject_mtp_support._Qwen3NextMTP

Exceptions and behavior

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

View source #L125-L177.

vllm_mlx.patches.qwen3_next_mtp.inject_mtp_support._Qwen3NextMTP.__call__ · nested function

vllm_mlx.patches.qwen3_next_mtp.inject_mtp_support._Qwen3NextMTP.__call__(inputs, cache = None, return_hidden: bool = False) -> not annotated

Nested Function inject_mtp_support._Qwen3NextMTP.__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.

Returns

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

Exceptions and behavior

Nested Function inject_mtp_support._Qwen3NextMTP.__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 #L128-L150.

vllm_mlx.patches.qwen3_next_mtp.inject_mtp_support._Qwen3NextMTP.mtp_forward · nested function

vllm_mlx.patches.qwen3_next_mtp.inject_mtp_support._Qwen3NextMTP.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._Qwen3NextMTP.mtp_forward calls self.model.embed_tokens, self.mtp.pre_fc_norm_hidden, self.mtp.pre_fc_norm_embedding, self.mtp.fc; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L152-L171.

vllm_mlx.patches.qwen3_next_mtp.inject_mtp_support._Qwen3NextMTP.make_mtp_cache · nested function

vllm_mlx.patches.qwen3_next_mtp.inject_mtp_support._Qwen3NextMTP.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._Qwen3NextMTP.make_mtp_cache calls KVCache; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L173-L177.

vllm_mlx.patches.qwen3_next_mtp.validate_mtp_support · function

vllm_mlx.patches.qwen3_next_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 getattr, logger.warning, logger.info, inspect.signature; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L184-L261.

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
inject_mtp_support	function	inject_mtp_support(model: Any, model_path, config: dict) -> bool	Inject MTP module into a loaded Qwen3-Next model.	#L27-L181
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).	#L68-L83
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.	#L69-L83
inject_mtp_support._mtp_quant_pred	nested function	inject_mtp_support._mtp_quant_pred(path, module) -> not annotated	Nested Function inject_mtp_support._mtp_quant_pred calls isinstance, path.endswith; has 3 explicit return paths.	#L93-L103
inject_mtp_support._Qwen3NextMTP	nested class	inject_mtp_support._Qwen3NextMTP()	Qwen3-Next with MTP support (injected at runtime).	#L125-L177
inject_mtp_support._Qwen3NextMTP.__call__	nested function	inject_mtp_support._Qwen3NextMTP.__call__(inputs, cache = None, return_hidden: bool = False) -> not annotated	Nested Function inject_mtp_support._Qwen3NextMTP.__call__ calls inner.embed_tokens, len, create_attention_mask, create_ssm_mask; has 2 explicit return paths.	#L128-L150
inject_mtp_support._Qwen3NextMTP.mtp_forward	nested function	inject_mtp_support._Qwen3NextMTP.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.	#L152-L171
inject_mtp_support._Qwen3NextMTP.make_mtp_cache	nested function	inject_mtp_support._Qwen3NextMTP.make_mtp_cache() -> not annotated	Create KV cache for MTP layers.	#L173-L177
validate_mtp_support	function	validate_mtp_support(model: Any) -> bool	Validate that a loaded model has working MTP support.	#L184-L261