"""Portable Mage-VL FP8 and NVFP4 linear modules. This file is loaded as Hugging Face remote code. The checkpoint config selects one format before the state dictionary is materialized, so the original BF16 language projection weights are never allocated or requested. """ from __future__ import annotations import hashlib import os from functools import lru_cache from pathlib import Path from typing import Any import torch import torch.nn.functional as F from torch import nn LANGUAGE_PROJECTION_ROLES = { "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", } _SMALLM_SOURCE_ROOT: Path | None = None def _configure_smallm_source(model_name_or_path: str) -> None: """Resolve native sources from a local repo or Hugging Face snapshot.""" global _SMALLM_SOURCE_ROOT candidate = Path(model_name_or_path).expanduser() local_root = candidate / "native" / "smallm_gemv" required = ("smallm_gemv.cpp", "smallm_gemv.cu", "smallm_gemv.h") if all((local_root / name).is_file() for name in required): _SMALLM_SOURCE_ROOT = local_root.resolve() return if not model_name_or_path: raise RuntimeError("Mage-VL small-M source repository is unspecified") from transformers.utils.hub import cached_file resolved = [ Path( cached_file( model_name_or_path, f"native/smallm_gemv/{name}", ) ) for name in required ] parents = {path.parent.resolve() for path in resolved} if len(parents) != 1: raise RuntimeError( "small-M native sources resolved to different directories: " f"{sorted(str(value) for value in parents)}" ) _SMALLM_SOURCE_ROOT = parents.pop() def _smallm_source_root() -> Path: if _SMALLM_SOURCE_ROOT is None: raise RuntimeError( "small-M native sources were not configured during model setup" ) return _SMALLM_SOURCE_ROOT @lru_cache(maxsize=1) def _load_smallm_extension() -> Any: from torch.utils.cpp_extension import load source_root = _smallm_source_root() configured_build = os.environ.get("MAGE_VL_SMALLM_BUILD_DIR") build_root = ( Path(configured_build).expanduser().resolve() if configured_build else Path(__file__).resolve().parent / ".native_build" ) build_root.mkdir(parents=True, exist_ok=True) os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0") os.environ.setdefault("MAX_JOBS", "4") source_hash = hashlib.sha256( b"".join( (source_root / name).read_bytes() for name in ( "smallm_gemv.cpp", "smallm_gemv.cu", "smallm_gemv.h", ) ) ).hexdigest()[:12] return load( name=f"mage_vl_smallm_gemv_{source_hash}", sources=[ str(source_root / "smallm_gemv.cpp"), str(source_root / "smallm_gemv.cu"), ], extra_cflags=["-O3"], extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo"], extra_include_paths=[str(source_root)], build_directory=str(build_root), with_cuda=True, verbose=False, is_python_module=True, ) def _smallm_nvfp4_linear( value: torch.Tensor, *, qdata: torch.Tensor, weight_block_scale: torch.Tensor, weight_scale: torch.Tensor, bias: torch.Tensor | None, ) -> torch.Tensor: return _load_smallm_extension().linear( value.contiguous(), qdata, weight_block_scale, weight_scale, bias, ) def _resolve_parent(root: nn.Module, module_name: str) -> tuple[nn.Module, str]: parent_name, separator, leaf = module_name.rpartition(".") if not separator: return root, module_name return root.get_submodule(parent_name), leaf def _empty_like_source( source: nn.Linear, shape: tuple[int, ...], dtype: torch.dtype, ) -> torch.Tensor: return torch.empty(shape, dtype=dtype, device=source.weight.device) class MageVLScaledFP8Linear(nn.Module): """W8A8 prefill with optional resident-weight W8A16 small-M decode.""" def __init__( self, source: nn.Linear, *, role: str, smallm_backend: str, smallm_threshold: int, smallm_roles: set[str], ) -> None: super().__init__() if smallm_backend not in {"off", "w8a16_gemv"}: raise ValueError(f"unsupported FP8 small-M backend: {smallm_backend}") if smallm_threshold <= 0: raise ValueError("FP8 small-M threshold must be positive") self.in_features = int(source.in_features) self.out_features = int(source.out_features) self.role = role self.smallm_backend = ( smallm_backend if role in smallm_roles else "off" ) self.smallm_threshold = int(smallm_threshold) self.register_buffer( "qdata", _empty_like_source( source, (self.out_features, self.in_features), torch.float8_e4m3fn, ), persistent=True, ) self.register_buffer( "weight_scale", _empty_like_source(source, (), torch.float32), persistent=True, ) if source.bias is None: self.bias_bf16 = None else: self.register_buffer( "bias_bf16", _empty_like_source( source, (self.out_features,), torch.bfloat16, ), persistent=True, ) def _weight_quantized_tensor(self) -> Any: from comfy_kitchen.tensor import QuantizedTensor, TensorCoreFP8Layout params = TensorCoreFP8Layout.Params( scale=self.weight_scale, orig_dtype=torch.bfloat16, orig_shape=(self.out_features, self.in_features), ) return QuantizedTensor( self.qdata, "TensorCoreFP8Layout", params, ) def forward(self, value: torch.Tensor) -> torch.Tensor: from comfy_kitchen.tensor import QuantizedTensor input_shape = tuple(value.shape) flattened = value.reshape(-1, input_shape[-1]).contiguous() if ( self.smallm_backend == "w8a16_gemv" and flattened.shape[0] <= self.smallm_threshold ): from .fp8_decode_runtime import smallm_fp8_linear output = smallm_fp8_linear( flattened, qdata=self.qdata, weight_scale=self.weight_scale, bias=self.bias_bf16, ) return output.reshape(*input_shape[:-1], self.out_features) quantized_input = QuantizedTensor.from_float( flattened, "TensorCoreFP8Layout", ) output = F.linear( quantized_input, self._weight_quantized_tensor(), None, ) if self.bias_bf16 is not None: output = output + self.bias_bf16 return output.reshape(*input_shape[:-1], self.out_features) class MageVLNVFP4Linear(nn.Module): """Native W4A4 prefill with optional packed-weight W4A16 small-M decode.""" def __init__( self, source: nn.Linear, *, role: str, smallm_backend: str, smallm_threshold: int, smallm_roles: set[str], ) -> None: super().__init__() self.in_features = int(source.in_features) self.out_features = int(source.out_features) self.role = role self.smallm_backend = ( smallm_backend if role in smallm_roles else "off" ) self.smallm_threshold = int(smallm_threshold) self.register_buffer( "qdata", _empty_like_source( source, (self.out_features, self.in_features // 2), torch.uint8, ), persistent=True, ) self.register_buffer( "weight_scale", _empty_like_source(source, (), torch.float32), persistent=True, ) self.register_buffer( "weight_block_scale", _empty_like_source( source, (self.out_features, self.in_features // 16), torch.float8_e4m3fn, ), persistent=True, ) if source.bias is None: self.bias_bf16 = None else: self.register_buffer( "bias_bf16", _empty_like_source( source, (self.out_features,), torch.bfloat16, ), persistent=True, ) def _weight_quantized_tensor(self) -> Any: from comfy_kitchen.tensor import QuantizedTensor, TensorCoreNVFP4Layout params = TensorCoreNVFP4Layout.Params( scale=self.weight_scale, orig_dtype=torch.bfloat16, orig_shape=(self.out_features, self.in_features), block_scale=self.weight_block_scale, ) return QuantizedTensor( self.qdata, "TensorCoreNVFP4Layout", params, ) def forward(self, value: torch.Tensor) -> torch.Tensor: from comfy_kitchen.tensor import QuantizedTensor input_shape = tuple(value.shape) flattened = value.reshape(-1, input_shape[-1]).contiguous() if ( self.smallm_backend == "w4a16_gemv" and flattened.shape[0] <= self.smallm_threshold ): output = _smallm_nvfp4_linear( flattened, qdata=self.qdata, weight_block_scale=self.weight_block_scale, weight_scale=self.weight_scale, bias=self.bias_bf16, ) return output.reshape(*input_shape[:-1], self.out_features) quantized_input = QuantizedTensor.from_float( flattened, "TensorCoreNVFP4Layout", ) output = F.linear( quantized_input, self._weight_quantized_tensor(), None, ) if self.bias_bf16 is not None: output = output + self.bias_bf16 return output.reshape(*input_shape[:-1], self.out_features) def _smallm_policy( quantization: dict[str, Any], *, format_name: str, ) -> tuple[str, int, set[str]]: backend = os.environ.get( "MAGE_VL_SMALLM_BACKEND", str(quantization.get("smallm_backend", "off")), ) supported_backend = ( "w8a16_gemv" if format_name == "scaled_fp8_w8a8" else "w4a16_gemv" ) if backend not in {"off", supported_backend}: raise ValueError(f"unsupported MAGE_VL_SMALLM_BACKEND: {backend}") threshold = int( os.environ.get( "MAGE_VL_SMALLM_THRESHOLD", str(quantization.get("smallm_threshold", 1)), ) ) if threshold <= 0: raise ValueError("MAGE_VL_SMALLM_THRESHOLD must be positive") configured_roles = quantization.get( "smallm_roles", sorted(LANGUAGE_PROJECTION_ROLES), ) role_text = os.environ.get( "MAGE_VL_SMALLM_ROLES", ",".join(str(value) for value in configured_roles), ) roles = {value.strip() for value in role_text.split(",") if value.strip()} if not roles <= LANGUAGE_PROJECTION_ROLES: raise ValueError( f"invalid small-M roles: {sorted(roles - LANGUAGE_PROJECTION_ROLES)}" ) return backend, threshold, roles def _environment_flag(name: str, default: bool) -> bool: value = os.environ.get(name) if value is None: return bool(default) normalized = value.strip().lower() if normalized in {"1", "true", "yes", "on"}: return True if normalized in {"0", "false", "no", "off"}: return False raise ValueError(f"{name} must be one of 1/0, true/false, yes/no, or on/off") def apply_mage_vl_quantization( model: nn.Module, config: Any, ) -> None: """Replace all 252 Qwen language projections before checkpoint loading.""" quantization = getattr(config, "mage_vl_quantization", None) if not quantization: return if not isinstance(quantization, dict): raise TypeError("mage_vl_quantization must be a dictionary") format_name = quantization.get("format") if format_name not in {"scaled_fp8_w8a8", "native_nvfp4_w4a4"}: raise ValueError(f"unsupported Mage-VL quantization: {format_name}") backend, threshold, smallm_roles = _smallm_policy( quantization, format_name=format_name, ) fused_gate_up = False fused_qkv = False fused_gate_up_threshold = 1 fused_qkv_threshold = 1 if format_name == "scaled_fp8_w8a8": fused_gate_up = _environment_flag( "MAGE_VL_FP8_FUSED_GATE_UP", bool(quantization.get("fused_gate_up", False)), ) fused_qkv = _environment_flag( "MAGE_VL_FP8_FUSED_QKV", bool(quantization.get("fused_qkv", False)), ) fused_gate_up_threshold = int( os.environ.get( "MAGE_VL_FP8_FUSED_GATE_UP_THRESHOLD", str(quantization.get("fused_gate_up_threshold", 1)), ) ) fused_qkv_threshold = int( os.environ.get( "MAGE_VL_FP8_FUSED_QKV_THRESHOLD", str(quantization.get("fused_qkv_threshold", 1)), ) ) if fused_gate_up_threshold <= 0 or fused_qkv_threshold <= 0: raise ValueError("FP8 fusion thresholds must be positive") if backend != "off" or fused_gate_up or fused_qkv: from .fp8_decode_runtime import configure_fp8_decode_sources configure_fp8_decode_sources( str(getattr(config, "_name_or_path", "")) ) if format_name == "native_nvfp4_w4a4" and backend != "off": _configure_smallm_source(str(getattr(config, "_name_or_path", ""))) installed = [] for layer in range(36): for branch, roles in ( ("self_attn", ("q_proj", "k_proj", "v_proj", "o_proj")), ("mlp", ("gate_proj", "up_proj", "down_proj")), ): for role in roles: name = f"language_model.layers.{layer}.{branch}.{role}" parent, leaf = _resolve_parent(model, name) source = getattr(parent, leaf) if not isinstance(source, nn.Linear): raise TypeError( f"{name}: expected nn.Linear, got " f"{type(source).__name__}" ) if format_name == "scaled_fp8_w8a8": replacement = MageVLScaledFP8Linear( source, role=role, smallm_backend=backend, smallm_threshold=threshold, smallm_roles=smallm_roles, ) else: replacement = MageVLNVFP4Linear( source, role=role, smallm_backend=backend, smallm_threshold=threshold, smallm_roles=smallm_roles, ) setattr(parent, leaf, replacement) installed.append(name) if len(installed) != 252: raise RuntimeError( f"expected 252 quantized language projections, got {len(installed)}" ) runtime_manifest = { "format": format_name, "smallm_backend": backend, "smallm_threshold": threshold, "smallm_roles": sorted(smallm_roles), "fused_gate_up": fused_gate_up, "fused_gate_up_threshold": fused_gate_up_threshold, "fused_qkv": fused_qkv, "fused_qkv_threshold": fused_qkv_threshold, } if format_name == "scaled_fp8_w8a8": from .fp8_decode_runtime import ( install_fp8_fused_gate_up, install_fp8_fused_qkv, ) if fused_gate_up: runtime_manifest["gate_up_install"] = install_fp8_fused_gate_up( model, threshold=fused_gate_up_threshold, ) if fused_qkv: runtime_manifest["qkv_install"] = install_fp8_fused_qkv( model, threshold=fused_qkv_threshold, ) object.__setattr__(model, "_mage_vl_runtime_manifest", runtime_manifest) __all__ = [ "MageVLNVFP4Linear", "MageVLScaledFP8Linear", "apply_mage_vl_quantization", ]