"""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