"""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 import types 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 _quantize_nvfp4_activation(value: torch.Tensor) -> Any: """Apply the checkpoint's exact dynamic NVFP4 activation policy.""" from comfy_kitchen.tensor import QuantizedTensor return QuantizedTensor.from_float(value, "TensorCoreNVFP4Layout") 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: 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 = _quantize_nvfp4_activation(flattened) 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 _runtime_policy( quantization: dict[str, Any], *, format_name: str, ) -> tuple[str, dict[str, Any]]: """Resolve a named runtime profile without changing stored weights.""" profiles = quantization.get("runtime_profiles") if format_name != "native_nvfp4_w4a4" or profiles is None: return "legacy", dict(quantization) if not isinstance(profiles, dict) or not profiles: raise TypeError("runtime_profiles must be a non-empty dictionary") default_profile = str(quantization.get("default_profile", "hybrid_fast")) profile_name = os.environ.get("MAGE_VL_RUNTIME_PROFILE", default_profile) if profile_name not in profiles: raise ValueError( f"unsupported MAGE_VL_RUNTIME_PROFILE: {profile_name}; " f"available={sorted(profiles)}" ) profile = profiles[profile_name] if not isinstance(profile, dict): raise TypeError(f"runtime profile {profile_name} must be a dictionary") effective = dict(quantization) effective.update(profile) return profile_name, effective def _pure_w4a4_projection( module: MageVLNVFP4Linear, quantized_input: Any, ) -> torch.Tensor: output = F.linear( quantized_input, module._weight_quantized_tensor(), None, ) if module.bias_bf16 is not None: output = output + module.bias_bf16 return output def install_nvfp4_shared_gate_up_activation( model: nn.Module, *, threshold: int = 1, ) -> dict[str, Any]: """Reuse one exact dynamic NVFP4 activation for gate and up.""" if threshold <= 0: raise ValueError("shared gate/up activation threshold must be positive") installed = [] for layer in range(36): name = f"language_model.layers.{layer}.mlp" mlp = model.get_submodule(name) if hasattr(mlp, "_nvfp4_shared_gate_up_original_forward"): raise RuntimeError(f"{name}: shared gate/up activation is installed") for role in ("gate_proj", "up_proj"): module = getattr(mlp, role, None) if not isinstance(module, MageVLNVFP4Linear): raise TypeError( f"{name}.{role}: expected MageVLNVFP4Linear, " f"got {type(module).__name__}" ) if module.smallm_backend != "off": raise RuntimeError( f"{name}.{role}: shared activation requires pure W4A4" ) original_forward = mlp.forward object.__setattr__( mlp, "_nvfp4_shared_gate_up_original_forward", original_forward, ) object.__setattr__( mlp, "_nvfp4_shared_gate_up_threshold", int(threshold), ) def shared_forward(self, value: torch.Tensor): input_shape = tuple(value.shape) flattened = value.reshape(-1, input_shape[-1]).contiguous() if flattened.shape[0] > self._nvfp4_shared_gate_up_threshold: return self._nvfp4_shared_gate_up_original_forward(value) quantized_input = _quantize_nvfp4_activation(flattened) gate = _pure_w4a4_projection(self.gate_proj, quantized_input) up = _pure_w4a4_projection(self.up_proj, quantized_input) intermediate = (F.silu(gate) * up).reshape( *input_shape[:-1], self.gate_proj.out_features, ) return self.down_proj(intermediate) object.__setattr__(mlp, "forward", types.MethodType(shared_forward, mlp)) installed.append(name) return { "feature": "shared_exact_dynamic_nvfp4_gate_up_activation", "threshold": int(threshold), "installed_module_count": len(installed), "fallback": "original_qwen_mlp_forward", "stored_weight_payload": "unchanged", } def restore_nvfp4_shared_gate_up_activation(model: nn.Module) -> dict[str, Any]: """Restore MLP callables patched for exact NVFP4 activation reuse.""" restored = [] for layer in range(36): name = f"language_model.layers.{layer}.mlp" mlp = model.get_submodule(name) original = getattr( mlp, "_nvfp4_shared_gate_up_original_forward", None, ) if original is None: continue object.__setattr__(mlp, "forward", original) object.__delattr__(mlp, "_nvfp4_shared_gate_up_original_forward") object.__delattr__(mlp, "_nvfp4_shared_gate_up_threshold") restored.append(name) return { "feature": "shared_exact_dynamic_nvfp4_gate_up_activation", "restored_module_count": len(restored), } class _NVFP4SharedQKVActivationCoordinator: """Reuse one exact dynamic NVFP4 activation across Q/K/V.""" def __init__( self, q_proj: MageVLNVFP4Linear, k_proj: MageVLNVFP4Linear, v_proj: MageVLNVFP4Linear, *, threshold: int, ) -> None: self.projections = {"q": q_proj, "k": k_proj, "v": v_proj} self.original_forwards = { role: module.forward for role, module in self.projections.items() } self.threshold = int(threshold) self.pending: dict[str, Any] | None = None @staticmethod def _signature(value: torch.Tensor) -> tuple[Any, ...]: return ( value.data_ptr(), tuple(value.shape), tuple(value.stride()), value.storage_offset(), value.device, value.dtype, ) def forward(self, role: str, value: torch.Tensor) -> torch.Tensor: input_shape = tuple(value.shape) flattened = value.reshape(-1, input_shape[-1]).contiguous() if flattened.shape[0] > self.threshold: self.pending = None return self.original_forwards[role](value) signature = self._signature(value) if role == "q": quantized_input = _quantize_nvfp4_activation(flattened) outputs = { projection_role: _pure_w4a4_projection( projection, quantized_input, ).reshape(*input_shape[:-1], projection.out_features) for projection_role, projection in self.projections.items() } self.pending = { "signature": signature, "k": outputs["k"], "v": outputs["v"], } return outputs["q"] if self.pending is None or self.pending["signature"] != signature: return self.original_forwards[role](value) output = self.pending[role] if role == "v": self.pending = None return output def install_nvfp4_shared_qkv_activation( model: nn.Module, *, threshold: int = 1, ) -> dict[str, Any]: """Patch pure-W4A4 Q/K/V to share one exact dynamic activation.""" if threshold <= 0: raise ValueError("shared QKV activation threshold must be positive") installed = [] for layer in range(36): name = f"language_model.layers.{layer}.self_attn" attention = model.get_submodule(name) if hasattr(attention, "_nvfp4_shared_qkv_coordinator"): raise RuntimeError(f"{name}: shared QKV activation is installed") projections = {} for role in ("q", "k", "v"): module = getattr(attention, f"{role}_proj", None) if not isinstance(module, MageVLNVFP4Linear): raise TypeError( f"{name}.{role}_proj: expected MageVLNVFP4Linear, " f"got {type(module).__name__}" ) if module.smallm_backend != "off": raise RuntimeError( f"{name}.{role}_proj: shared activation requires pure W4A4" ) projections[role] = module coordinator = _NVFP4SharedQKVActivationCoordinator( projections["q"], projections["k"], projections["v"], threshold=threshold, ) object.__setattr__(attention, "_nvfp4_shared_qkv_coordinator", coordinator) for role, module in projections.items(): def shared_forward( self, value: torch.Tensor, *, _role: str = role, _coordinator: _NVFP4SharedQKVActivationCoordinator = coordinator, ): return _coordinator.forward(_role, value) object.__setattr__( module, "forward", types.MethodType(shared_forward, module), ) installed.append(name) return { "feature": "shared_exact_dynamic_nvfp4_qkv_activation", "threshold": int(threshold), "installed_module_count": len(installed), "fallback": "separate_q_k_v_projection_forwards", "stored_weight_payload": "unchanged", } def restore_nvfp4_shared_qkv_activation(model: nn.Module) -> dict[str, Any]: """Restore Q/K/V callables patched for exact NVFP4 activation reuse.""" restored = [] for layer in range(36): name = f"language_model.layers.{layer}.self_attn" attention = model.get_submodule(name) coordinator = getattr(attention, "_nvfp4_shared_qkv_coordinator", None) if coordinator is None: continue for role in ("q", "k", "v"): module = getattr(attention, f"{role}_proj") object.__setattr__(module, "forward", coordinator.original_forwards[role]) coordinator.pending = None object.__delattr__(attention, "_nvfp4_shared_qkv_coordinator") restored.append(name) return { "feature": "shared_exact_dynamic_nvfp4_qkv_activation", "restored_module_count": len(restored), } 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}") runtime_profile, effective_quantization = _runtime_policy( quantization, format_name=format_name, ) backend, threshold, smallm_roles = _smallm_policy( effective_quantization, format_name=format_name, ) fused_gate_up = False fused_qkv = False fused_gate_up_threshold = 1 fused_qkv_threshold = 1 shared_gate_up = False shared_qkv = False shared_gate_up_threshold = 1 shared_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", ""))) if format_name == "native_nvfp4_w4a4": shared_gate_up = _environment_flag( "MAGE_VL_W4A4_SHARED_GATE_UP_ACTIVATION", bool(effective_quantization.get("shared_gate_up_activation", False)), ) shared_qkv = _environment_flag( "MAGE_VL_W4A4_SHARED_QKV_ACTIVATION", bool(effective_quantization.get("shared_qkv_activation", False)), ) shared_gate_up_threshold = int( os.environ.get( "MAGE_VL_W4A4_SHARED_GATE_UP_THRESHOLD", str(effective_quantization.get("shared_gate_up_threshold", 1)), ) ) shared_qkv_threshold = int( os.environ.get( "MAGE_VL_W4A4_SHARED_QKV_THRESHOLD", str(effective_quantization.get("shared_qkv_threshold", 1)), ) ) if shared_gate_up_threshold <= 0 or shared_qkv_threshold <= 0: raise ValueError("NVFP4 shared-activation thresholds must be positive") if backend != "off" and (shared_gate_up or shared_qkv): raise ValueError( "NVFP4 shared-activation features require MAGE_VL_SMALLM_BACKEND=off" ) 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 = { "runtime_version": int(quantization.get("runtime_version", 1)), "runtime_profile": runtime_profile, "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, "shared_gate_up_activation": shared_gate_up, "shared_gate_up_threshold": shared_gate_up_threshold, "shared_qkv_activation": shared_qkv, "shared_qkv_threshold": shared_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, ) if format_name == "native_nvfp4_w4a4": if shared_gate_up: runtime_manifest["shared_gate_up_install"] = ( install_nvfp4_shared_gate_up_activation( model, threshold=shared_gate_up_threshold, ) ) if shared_qkv: runtime_manifest["shared_qkv_install"] = ( install_nvfp4_shared_qkv_activation( model, threshold=shared_qkv_threshold, ) ) object.__setattr__(model, "_mage_vl_runtime_manifest", runtime_manifest) __all__ = [ "MageVLNVFP4Linear", "MageVLScaledFP8Linear", "apply_mage_vl_quantization", "install_nvfp4_shared_gate_up_activation", "install_nvfp4_shared_qkv_activation", "restore_nvfp4_shared_gate_up_activation", "restore_nvfp4_shared_qkv_activation", ]