modeling_mlx_vlm_qwen36_modelopt_hybrid.py

"""MLX-VLM runtime for Qwen3.6 ModelOpt hybrid FP8/NVFP4 checkpoints.

The language model keeps the lossless ModelOpt-to-MLX representation used by
``modeling_mlx_qwen36_modelopt_hybrid.py``.  The vision tower remains in its
original BF16 representation and is delegated to MLX-VLM's native Qwen3.5 MoE
vision implementation.

This module intentionally exports the same public symbols as an MLX-VLM model
package so a model-local loader can select it without changing ``model_type``.
"""

from dataclasses import dataclass, field
from typing import Dict

import mlx.core as mx
import mlx.nn as nn
from mlx.utils import tree_flatten, tree_unflatten

from mlx_vlm.models.qwen3_5_moe import LanguageModel, TextConfig, VisionConfig
from mlx_vlm.models.qwen3_5_moe import Model as BaseModel
from mlx_vlm.models.qwen3_5_moe import ModelConfig as BaseModelConfig
from mlx_vlm.models.qwen3_5_moe import VisionModel
from mlx_vlm.models.switch_layers import SwitchLinear


@dataclass
class ModelConfig(BaseModelConfig):
    mlx_modelopt_quantization: Dict[str, str] = field(default_factory=dict)


class ScaledQuantizedLinear(nn.Module):
    """Weight-quantized dense linear with a ModelOpt tensor scale."""

    def __init__(
        self,
        input_dims: int,
        output_dims: int,
        *,
        group_size: int,
        bits: int,
        mode: str,
        bias: bool = False,
    ):
        super().__init__()
        if input_dims % group_size:
            raise ValueError(
                f"input_dims={input_dims} is not divisible by group_size={group_size}"
            )
        if (input_dims * bits) % 32:
            raise ValueError(
                f"input_dims={input_dims}, bits={bits} cannot be packed into uint32"
            )

        self.group_size = group_size
        self.bits = bits
        self.mode = mode
        self.weight = mx.zeros(
            (output_dims, input_dims * bits // 32), dtype=mx.uint32
        )
        self.scales = mx.zeros(
            (output_dims, input_dims // group_size), dtype=mx.uint8
        )
        self.global_scale = mx.ones((), dtype=mx.float32)
        if bias:
            self.bias = mx.zeros((output_dims,))
        self.freeze()

    @classmethod
    def from_linear(cls, linear: nn.Module, kind: str):
        output_dims, input_dims = linear.weight.shape
        has_bias = linear.get("bias") is not None
        if kind == "scaled_mxfp8":
            params = dict(group_size=32, bits=8, mode="mxfp8")
        elif kind == "scaled_nvfp4":
            params = dict(group_size=16, bits=4, mode="nvfp4")
        else:
            raise ValueError(f"Unsupported dense quantization kind: {kind}")
        return cls(input_dims, output_dims, bias=has_bias, **params)

    def __call__(self, x):
        y = mx.quantized_matmul(
            x,
            self["weight"],
            self["scales"],
            transpose=True,
            group_size=self.group_size,
            bits=self.bits,
            mode=self.mode,
        )
        y = y * self["global_scale"].astype(y.dtype)
        if "bias" in self:
            y = y + self["bias"]
        return y


class ScaledNVFP4SwitchLinear(nn.Module):
    """Expert linear using MLX gather_qmm and per-expert tensor scales."""

    group_size = 16
    bits = 4
    mode = "nvfp4"

    def __init__(
        self,
        input_dims: int,
        output_dims: int,
        num_experts: int,
        *,
        bias: bool = False,
    ):
        super().__init__()
        if input_dims % self.group_size:
            raise ValueError(
                f"input_dims={input_dims} is not divisible by {self.group_size}"
            )
        self.weight = mx.zeros(
            (num_experts, output_dims, input_dims * self.bits // 32),
            dtype=mx.uint32,
        )
        self.scales = mx.zeros(
            (num_experts, output_dims, input_dims // self.group_size),
            dtype=mx.uint8,
        )
        self.global_scales = mx.ones((num_experts,), dtype=mx.float32)
        if bias:
            self.bias = mx.zeros((num_experts, output_dims))
        self.freeze()

    @classmethod
    def from_switch_linear(cls, linear: SwitchLinear):
        num_experts, output_dims, input_dims = linear.weight.shape
        has_bias = linear.get("bias") is not None
        return cls(input_dims, output_dims, num_experts, bias=has_bias)

    @property
    def input_dims(self):
        return self.scales.shape[2] * self.group_size

    @property
    def output_dims(self):
        return self.weight.shape[1]

    @property
    def num_experts(self):
        return self.weight.shape[0]

    def __call__(self, x, indices, sorted_indices=False):
        y = mx.gather_qmm(
            x,
            self["weight"],
            self["scales"],
            rhs_indices=indices,
            transpose=True,
            group_size=self.group_size,
            bits=self.bits,
            mode=self.mode,
            sorted_indices=sorted_indices,
        )
        scale = self["global_scales"][indices].astype(y.dtype)[..., None, None]
        y = y * scale
        if "bias" in self:
            y = y + mx.expand_dims(self["bias"][indices], -2)
        return y


def _replace_quantized_modules(model: nn.Module, quantization: Dict[str, str]):
    leaves = dict(
        tree_flatten(model.leaf_modules(), is_leaf=lambda m: isinstance(m, nn.Module))
    )
    missing = sorted(set(quantization) - set(leaves))
    if missing:
        preview = "\n  ".join(missing[:20])
        raise ValueError(f"Quantized module paths are absent from the model:\n  {preview}")

    for path, kind in quantization.items():
        module = leaves[path]
        if kind in ("scaled_mxfp8", "scaled_nvfp4"):
            if not isinstance(module, nn.Linear):
                raise TypeError(f"{path} is {type(module).__name__}, expected Linear")
            leaves[path] = ScaledQuantizedLinear.from_linear(module, kind)
        elif kind == "scaled_nvfp4_switch":
            if not isinstance(module, SwitchLinear):
                raise TypeError(
                    f"{path} is {type(module).__name__}, expected SwitchLinear"
                )
            leaves[path] = ScaledNVFP4SwitchLinear.from_switch_linear(module)
        else:
            raise ValueError(f"Unknown quantization kind {kind!r} for {path}")

    model.update_modules(tree_unflatten(list(leaves.items())))


class Model(BaseModel):
    def __init__(self, config: ModelConfig):
        super().__init__(config)
        _replace_quantized_modules(self, config.mlx_modelopt_quantization)

    def sanitize(self, weights):
        """Map only raw source keys; converted language keys are already sanitized."""
        sanitized = {}
        for key, value in weights.items():
            if "mtp." in key:
                continue
            if key.startswith("model.language_model.visual"):
                key = key.replace(
                    "model.language_model.visual", "vision_tower", 1
                )
            elif key.startswith("model.language_model"):
                key = key.replace(
                    "model.language_model", "language_model.model", 1
                )
            elif key.startswith("model.visual"):
                key = key.replace("model.visual", "vision_tower", 1)
            elif key.startswith("lm_head"):
                key = key.replace("lm_head", "language_model.lm_head", 1)
            sanitized[key] = value
        return sanitized