vllm_mlx.text_model_from_vlm

Construct an mlx_lm TextModel from mlx_vlm-loaded model weights.

View the complete module source at #L1-L272.

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.text_model_from_vlm

Construct an mlx_lm TextModel from mlx_vlm-loaded model weights.

When mlx_vlm loads a model, it strips MTP weights in sanitize(). This module builds a parallel mlx_lm TextModel that: 1. Shares backbone + lm_head weights with the vlm model (zero-copy) 2. Loads MTP weights from safetensors on disk 3. Provides full mlx_lm API: return_hidden, n_confirmed, mtp_forward, make_mtp_cache

vllm_mlx.text_model_from_vlm.logger module-attribute

logger = logging.getLogger(__name__)

vllm_mlx.text_model_from_vlm._TEXT_MODEL_FAMILIES module-attribute

_TEXT_MODEL_FAMILIES: tuple[tuple[str, str, tuple[str, str]], ...] = (('gemma4', 'mlx_lm.models.gemma4_text', ('Model', 'ModelArgs')),)

vllm_mlx.text_model_from_vlm._DEFAULT_TEXT_MODEL module-attribute

_DEFAULT_TEXT_MODEL = ('mlx_lm.models.qwen3_5', ('TextModel', 'TextModelArgs'))

vllm_mlx.text_model_from_vlm._import_text_model_classes

_import_text_model_classes(model_type: str)

Return (Model, ModelArgs) for a text config's model_type.

Source code in vllm_mlx/text_model_from_vlm.py

def _import_text_model_classes(model_type: str):
    """Return ``(Model, ModelArgs)`` for a text config's ``model_type``."""
    import importlib

    module_name, (model_attr, args_attr) = _DEFAULT_TEXT_MODEL
    matched = False
    for prefix, family_module, (family_model, family_args) in sorted(
        _TEXT_MODEL_FAMILIES, key=lambda f: len(f[0]), reverse=True
    ):
        if model_type.startswith(prefix):
            module_name, model_attr, args_attr = (
                family_module,
                family_model,
                family_args,
            )
            matched = True
            break

    if not matched:
        logger.debug(
            "No text-model family matches model_type=%r; falling back to %s",
            model_type,
            module_name,
        )
    module = importlib.import_module(module_name)
    return getattr(module, model_attr), getattr(module, args_attr)

vllm_mlx.text_model_from_vlm.build_text_model

build_text_model(vlm_model: Any, model_path: str | Path) -> Any | None

Build an mlx_lm TextModel from a vlm-loaded model's weights.

Parameters:

• vlm_model (Any) –

  The mlx_vlm-loaded model (has .language_model attribute)

• model_path (str | Path) –

  Path to the model directory (contains config.json + safetensors)

Returns:

• Any | None –

  mlx_lm TextModel with MTP support, or None on failure.

Source code in vllm_mlx/text_model_from_vlm.py

def build_text_model(vlm_model: Any, model_path: str | Path) -> Any | None:
    """Build an mlx_lm TextModel from a vlm-loaded model's weights.

    Args:
        vlm_model: The mlx_vlm-loaded model (has .language_model attribute)
        model_path: Path to the model directory (contains config.json + safetensors)

    Returns:
        mlx_lm TextModel with MTP support, or None on failure.
    """
    if vlm_model is None:
        return None

    model_path = Path(model_path) if model_path else None
    if model_path is None or not (model_path / "config.json").exists():
        return None

    model_type = ""
    text_model_cls = None
    try:
        config = json.loads((model_path / "config.json").read_text())
        text_config = config.get("text_config", config)
        model_type = text_config.get("model_type") or config.get("model_type", "")
        TextModel, TextModelArgs = _import_text_model_classes(model_type)
        text_model_cls = f"{TextModel.__module__}.{TextModel.__qualname__}"
        logger.debug(
            "Building TextModel for model_type=%r using %s", model_type, text_model_cls
        )

        # Build args with proper __post_init__ (handles partial_rotary_factor,
        # rope_scaling, head_dim derivation)
        args = TextModelArgs.from_dict(text_config)
        text_model = TextModel(args)

        # Collect all weights first: backbone from vlm + MTP from safetensors
        vlm_lm = vlm_model.language_model
        vlm_weights = mlx.utils.tree_flatten(vlm_lm.parameters())
        mtp_weights = _load_mtp_weights(model_path)

        all_weight_names = set(name for name, _ in vlm_weights)
        all_weight_names.update(name for name, _ in mtp_weights)

        # Quantize the TextModel skeleton to match source weights.
        # Use a predicate that only quantizes layers that have .scales in source.
        # This prevents quantizing layers like mtp.fc which are BF16.
        quantization = text_config.get("quantization", config.get("quantization", None))
        if quantization is not None:
            # Per-layer quantization overrides (e.g. 8-bit MoE gates over a 4-bit
            # body) may be keyed with the source checkpoint's "language_model."
            # wrapper prefix, which the extracted TextModel doesn't carry. Index
            # them by suffix so the prefix-stripped module path resolves to the
            # correct bits/group_size; without this the override is ignored and the
            # layer is quantized at the global width, producing a quantized_matmul
            # shape mismatch at decode.
            per_layer_overrides = {
                k: v for k, v in quantization.items() if isinstance(v, dict)
            }

            def _class_predicate(path, module):
                if not hasattr(module, "to_quantized"):
                    return False
                if f"{path}.scales" not in all_weight_names:
                    return False
                if path in quantization:
                    return quantization[path]
                for key, override in per_layer_overrides.items():
                    if key.endswith("." + path):
                        return override
                return True

            nn.quantize(
                text_model,
                group_size=quantization.get("group_size", 64),
                bits=quantization.get("bits", 8),
                class_predicate=_class_predicate,
            )

        # Transfer backbone + lm_head weights from vlm language_model (zero-copy).
        # strict=False because TextModel has MTP params that vlm doesn't have yet.
        text_model.load_weights(vlm_weights, strict=False)

        logger.info(
            "Transferred %d weight arrays from vlm language_model", len(vlm_weights)
        )

        # Load MTP weights from safetensors
        if mtp_weights:
            text_model.load_weights(mtp_weights, strict=False)
            logger.info("Loaded %d MTP weights from safetensors", len(mtp_weights))
        else:
            logger.warning("No MTP weights found in %s", model_path.name)

        # Inject MTP if TextModel doesn't have native MTP support.
        # mlx_lm's qwen3_5.TextModel strips MTP weights in sanitize(),
        # so we inject MTP module + methods at runtime.
        if not hasattr(text_model, "mtp") or text_model.mtp is None:
            num_mtp = text_config.get("mtp_num_hidden_layers", 0)
            if num_mtp == 0:
                num_mtp = text_config.get("num_nextn_predict_layers", 0)
            if num_mtp > 0:
                from .patches.qwen3_5_mtp import inject_mtp_support

                inject_mtp_support(text_model, model_path, config)

        if hasattr(text_model, "mtp") and text_model.mtp is not None:
            mx.eval(text_model.mtp.parameters())
            num_mtp = text_config.get(
                "mtp_num_hidden_layers",
                text_config.get("num_nextn_predict_layers", 0),
            )
            logger.info("TextModel built with MTP support (%d layers)", num_mtp)
        else:
            logger.info("TextModel built without MTP")

        # Put the derived TextModel in eval mode. mlx_lm.load / mlx_vlm.load both
        # eval() their models (this is also what LM Studio's mlx-engine does), but
        # this freshly-constructed TextModel defaults to training=True. Hybrid
        # layers (Qwen3.5/3.6 gated-delta) select their compute path with
        # `use_kernel = not self.training`, so in training mode prefill/decode fall
        # to the slow Python recurrence instead of the Metal kernel.
        text_model.train(False)

        # Realize every array the model holds before it leaves the build
        # thread — including underscore-private module attributes such as
        # RoPE._freqs, which parameters() excludes. MLX lazy graphs are tagged
        # to the stream of the thread that recorded them; a lazy array
        # surviving into generation dies with "There is no Stream(gpu, N) in
        # current thread" the moment a worker on another thread evaluates it
        # (Gemma 4: the scaled-RoPE _freqs of the first full_attention layer).
        if hasattr(text_model, "modules"):
            mx.eval(
                [
                    v
                    for module in text_model.modules()
                    for v in module.values()
                    if isinstance(v, mx.array)
                ]
            )

        return text_model

    except ImportError as e:
        logger.error("Cannot import mlx_lm TextModel (need PR #990): %s", e)
        return None
    except Exception as e:
        # Name the model_type and the class that was picked. Without them this
        # is a bare TypeError from inside someone else's constructor, and the
        # engine carries on with _text_model=None — a route quietly losing its
        # backend, which reads like a warning rather than the failure it is.
        logger.error(
            "Failed to build TextModel from vlm (model_type=%r, class=%s): %s",
            model_type,
            text_model_cls or "<not selected>",
            e,
            exc_info=True,
        )
        return None

vllm_mlx.text_model_from_vlm._load_mtp_weights

_load_mtp_weights(model_path: Path) -> list[tuple[str, array]]

Load MTP weights from safetensors, stripping the language_model. prefix.

mlx_vlm's sanitize() strips mtp.* keys during model loading, but the weights are still on disk in the safetensors files.

Source code in vllm_mlx/text_model_from_vlm.py

def _load_mtp_weights(model_path: Path) -> list[tuple[str, mx.array]]:
    """Load MTP weights from safetensors, stripping the language_model. prefix.

    mlx_vlm's sanitize() strips mtp.* keys during model loading,
    but the weights are still on disk in the safetensors files.
    """
    index_file = model_path / "model.safetensors.index.json"
    if not index_file.exists():
        return []

    index = json.loads(index_file.read_text())
    weight_map = index.get("weight_map", {})

    # Find MTP keys and their shard files
    mtp_keys: dict[str, tuple[str, str]] = {}
    for key, shard in weight_map.items():
        if ".mtp." in key:
            # Strip "language_model." prefix to match mlx_lm namespace
            clean = (
                key.replace("language_model.", "", 1)
                if key.startswith("language_model.")
                else key
            )
            mtp_keys[key] = (clean, shard)

    if not mtp_keys:
        return []

    # Group by shard to minimize I/O
    shards: dict[str, list[tuple[str, str]]] = {}
    for orig, (clean, shard) in mtp_keys.items():
        shards.setdefault(shard, []).append((orig, clean))

    weights = []
    for shard_file, key_pairs in shards.items():
        shard_path = model_path / shard_file
        if not shard_path.exists():
            logger.warning("MTP shard not found: %s", shard_file)
            continue
        shard_data = mx.load(str(shard_path))
        for orig, clean in key_pairs:
            if orig in shard_data:
                weights.append((clean, shard_data[orig]))

    return weights

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.text_model_from_vlm._import_text_model_classes · function

vllm_mlx.text_model_from_vlm._import_text_model_classes(model_type: str) -> not annotated

Return (Model, ModelArgs) for a text config's model_type.

Parameters

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

Returns

• Type: not annotated
• Direct return expressions: (getattr(module, model_attr), getattr(module, args_attr))

Exceptions and behavior

Function _import_text_model_classes calls sorted, model_type.startswith, logger.debug, importlib.import_module; returns (getattr(module, model_attr), getattr(module, args_attr)). No direct raise statement appears in this definition.

View source #L41-L66.

vllm_mlx.text_model_from_vlm.build_text_model · function

vllm_mlx.text_model_from_vlm.build_text_model(vlm_model: Any, model_path: str | Path) -> Any | None

Build an mlx_lm TextModel from a vlm-loaded model's weights.

Parameters

Name	Type	Required	Default	Description
vlm_model	Any	yes	none	The mlx_vlm-loaded model (has .language_model attribute)
model_path	str \| Path	yes	none	Path to the model directory (contains config.json + safetensors)

Returns

• Type: Any | None
• Direct return expressions: None; text_model

Exceptions and behavior

Function build_text_model calls Path, (model_path / 'config.json').exists, json.loads, (model_path / 'config.json').read_text; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L69-L225.

vllm_mlx.text_model_from_vlm.build_text_model._class_predicate · nested function

vllm_mlx.text_model_from_vlm.build_text_model._class_predicate(path, module) -> not annotated

Nested Function build_text_model._class_predicate calls hasattr, per_layer_overrides.items, key.endswith; has 4 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; quantization[path]; override; True

Exceptions and behavior

Nested Function build_text_model._class_predicate calls hasattr, per_layer_overrides.items, key.endswith; has 4 explicit return paths. No direct raise statement appears in this definition.

View source #L127-L137.

vllm_mlx.text_model_from_vlm._load_mtp_weights · function

vllm_mlx.text_model_from_vlm._load_mtp_weights(model_path: Path) -> list[tuple[str, mx.array]]

Load MTP weights from safetensors, stripping the language_model.

Parameters

Name	Type	Required	Default	Description
model_path	Path	yes	none	Required positional or keyword input.

Returns

• Type: list[tuple[str, mx.array]]
• Direct return expressions: []; weights

Exceptions and behavior

Function _load_mtp_weights calls index_file.exists, json.loads, index_file.read_text, index.get; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L228-L272.

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
_import_text_model_classes	function	_import_text_model_classes(model_type: str) -> not annotated	Return (Model, ModelArgs) for a text config's model_type.	#L41-L66
build_text_model	function	build_text_model(vlm_model: Any, model_path: str \| Path) -> Any \| None	Build an mlx_lm TextModel from a vlm-loaded model's weights.	#L69-L225
build_text_model._class_predicate	nested function	build_text_model._class_predicate(path, module) -> not annotated	Nested Function build_text_model._class_predicate calls hasattr, per_layer_overrides.items, key.endswith; has 4 explicit return paths.	#L127-L137
_load_mtp_weights	function	_load_mtp_weights(model_path: Path) -> list[tuple[str, mx.array]]	Load MTP weights from safetensors, stripping the language_model.	#L228-L272