# Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/quantization/modelopt_quant.py from __future__ import annotations import logging import re from functools import lru_cache from typing import Any, Dict, List, Optional import torch from sglang.multimodal_gen.runtime.layers.linear import ( LinearMethodBase, UnquantizedLinearMethod, ) from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, QuantizeMethodBase, ) from sglang.multimodal_gen.runtime.models.parameter import ( ModelWeightParameter, PerTensorScaleParameter, ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.weight_attrs import set_weight_attrs from sglang.srt.layers.quantization.fp8_utils import ( apply_fp8_linear, cutlass_fp8_supported, ) from sglang.srt.layers.quantization.modelopt_quant import ( pad_nvfp4_activation_for_cutlass, pad_nvfp4_weight, slice_nvfp4_output, ) from sglang.srt.layers.quantization.utils import ( convert_to_channelwise, is_layer_skipped, requantize_with_max_scale, ) from sglang.srt.layers.utils.common import copy_or_rebind_param from sglang.srt.utils.common import is_flashinfer_available, round_up logger = logging.getLogger(__name__) if is_flashinfer_available(): import flashinfer else: flashinfer = None @lru_cache(maxsize=1) def _get_fp4_quantize_op(): return current_platform.get_modelopt_fp4_quantize_op() @lru_cache(maxsize=1) def _get_fp4_gemm_op(): return current_platform.get_modelopt_fp4_gemm_op() def _prepare_nvfp4_weight_bytes( weight: torch.Tensor, *, swap_weight_nibbles: bool ) -> torch.Tensor: """Normalize serialized NVFP4 bytes before padding for the runtime kernel.""" if not swap_weight_nibbles: return weight.contiguous() return ((weight >> 4) | (weight << 4)).contiguous() def _swizzled_nvfp4_scales_to_linear(scales: torch.Tensor) -> torch.Tensor: """Convert FlashInfer/CUTLASS-swizzled FP4 scales back to row-major layout.""" scale_ndim = scales.ndim if scale_ndim == 2: scales = scales.unsqueeze(0) assert scales.ndim == 3 B, M, K = scales.shape M_padded = round_up(M, 128) K_padded = round_up(K, 4) if M != M_padded or K != K_padded: padded = torch.zeros( (B, M_padded, K_padded), dtype=scales.dtype, device=scales.device ) padded[:B, :M, :K] = scales scales = padded linear = scales.reshape(B, M_padded // 128, K_padded // 4, 32, 4, 4) linear = linear.permute(0, 1, 4, 3, 2, 5).contiguous() linear = linear.reshape(B, M_padded, K_padded)[:, :M, :K] return linear.squeeze(0) if scale_ndim == 2 else linear def _require_flashinfer(): if flashinfer is None: raise RuntimeError( "flashinfer is required for the diffusion NVFP4 FlashInfer path." ) return flashinfer class ModelOptQuantConfig(QuantizationConfig): def __init__( self, exclude_modules: Optional[List[str]], packed_modules_mapping: Optional[Dict[str, List[str]]], ): super().__init__() self.packed_modules_mapping = packed_modules_mapping or {} self.exclude_modules = exclude_modules or [] def _get_quant_method( self, layer: torch.nn.Module, prefix: str, *, Linear: type[LinearMethodBase], ) -> Optional[QuantizeMethodBase]: from sglang.multimodal_gen.runtime.layers.linear import LinearBase if isinstance(layer, LinearBase): if self.is_layer_excluded(prefix) or ( self.packed_modules_mapping and is_layer_skipped(prefix, [], self.packed_modules_mapping) ): return UnquantizedLinearMethod() return Linear(self) return None @classmethod def get_config_filenames(cls) -> List[str]: return ["hf_quant_config.json"] def get_scaled_act_names(self) -> List[str]: return [] @classmethod def override_quantization_method(cls, hf_quant_config, user_quant) -> Optional[str]: if hf_quant_config is None: return None quant_algo = ( hf_quant_config.get("quant_algo") or hf_quant_config.get("quantization", {}).get("quant_algo") or "" ).upper() if user_quant in {"modelopt", "modelopt_fp8"} and "FP8" in quant_algo: return "modelopt_fp8" if user_quant in {"modelopt", "modelopt_fp4"} and ( "NVFP4" in quant_algo or "FP4" in quant_algo ): return "modelopt_fp4" return None def is_layer_excluded(self, prefix: str) -> bool: for pattern in self.exclude_modules: regex_str = re.escape(pattern).replace(r"\*", r".*") if re.fullmatch(regex_str, prefix): return True return False class ModelOptFp8Config(ModelOptQuantConfig): """Config class for ModelOpt FP8 diffusion checkpoints.""" def __init__( self, is_checkpoint_fp8_serialized: bool = False, exclude_modules: Optional[List[str]] = None, packed_modules_mapping: Optional[Dict[str, List[str]]] = None, ) -> None: super().__init__(exclude_modules, packed_modules_mapping) self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized if is_checkpoint_fp8_serialized: logger.warning( "Detected ModelOpt FP8 checkpoint. The format is experimental and subject to change." ) @classmethod def get_name(cls) -> str: return "modelopt_fp8" @classmethod def get_supported_act_dtypes(cls) -> List[torch.dtype]: return [torch.bfloat16, torch.half] @classmethod def get_min_capability(cls) -> int: return 89 @classmethod def from_config( cls, config: Dict[str, Any], ignore_remap: Optional[Dict[str, str]] = None, ) -> ModelOptFp8Config: quant_method = config.get("quant_algo") exclude_modules = config.get("ignore") if quant_method is None: try: quantization_section = cls.get_from_keys(config, ["quantization"]) quant_method = quantization_section.get("quant_algo") exclude_modules = quantization_section.get("exclude_modules") except ValueError as exc: raise ValueError( "Cannot find 'quant_algo' in the model's quantization config." ) from exc if quant_method is None or "FP8" not in quant_method: raise ValueError( "ModelOptFp8Config only supports static FP8 quantization in SGLang diffusion." ) if ignore_remap and exclude_modules: exclude_modules = [ignore_remap.get(p, p) for p in exclude_modules] return cls( is_checkpoint_fp8_serialized=True, exclude_modules=exclude_modules, packed_modules_mapping=config.get("packed_modules_mapping"), ) def get_quant_method(self, layer: torch.nn.Module, prefix: str): return self._get_quant_method(layer, prefix, Linear=ModelOptFp8LinearMethod) class ModelOptFp4Config(ModelOptQuantConfig): """Config class for NVFP4.""" def __init__( self, is_checkpoint_nvfp4_serialized: bool = False, group_size: int = None, exclude_modules: List[str] = None, packed_modules_mapping: Optional[Dict[str, List[str]]] = None, checkpoint_uses_packed_qkv: bool = False, swap_weight_nibbles: bool = False, checkpoint_weight_scale_layout: str = "linear", ) -> None: super().__init__(exclude_modules, packed_modules_mapping) self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized if is_checkpoint_nvfp4_serialized: logger.warning( "Detected nvfp4 checkpoint. Please note that the " "format is experimental and subject to change." ) self.group_size = group_size self.checkpoint_uses_packed_qkv = checkpoint_uses_packed_qkv self.swap_weight_nibbles = swap_weight_nibbles self.checkpoint_weight_scale_layout = checkpoint_weight_scale_layout @classmethod def get_name(cls) -> str: return "modelopt_fp4" @classmethod def get_supported_act_dtypes(cls) -> List[torch.dtype]: return [torch.bfloat16, torch.half, torch.float8_e4m3fn] @classmethod def get_min_capability(cls) -> int: return 100 @staticmethod def common_group_size(cfg: dict) -> int: """Return the unique group_size across the config; raise if missing/mismatched.""" sizes = set() def _add_group_size_from_dict(config: dict): group_size = config.get("group_size") if isinstance(group_size, int): sizes.add(group_size) # Top-level and 'quantization' block _add_group_size_from_dict(cfg) quantization = cfg.get("quantization") if isinstance(quantization, dict): _add_group_size_from_dict(quantization) # config_groups: accept group-level or nested dicts (e.g., weights/input_activations) for config_groups in (cfg.get("config_groups") or {}).values(): if isinstance(config_groups, dict): _add_group_size_from_dict(config_groups) for config_group in config_groups.values(): if isinstance(config_group, dict): _add_group_size_from_dict(config_group) if not sizes: raise ValueError("No group_size found in config.") if len(sizes) > 1: raise ValueError(f"Inconsistent group_size values: {sorted(sizes)}") return next(iter(sizes)) @classmethod def from_config(cls, config: Dict[str, Any]) -> ModelOptFp4Config: group_size = None exclude_modules = [] swap_weight_nibbles = False # Flat format (config.json quantization_config) quant_method = config.get("quant_algo") if quant_method is not None: group_size = config.get("group_size") if group_size is None: config_groups = config.get("config_groups", {}) if config_groups: first_group = next(iter(config_groups.values()), {}) group_size = first_group.get("weights", {}).get("group_size") exclude_modules = config.get("ignore", []) swap_weight_nibbles = config.get( "swap_weight_nibbles", config.get("checkpoint_uses_packed_qkv", False), ) else: # Nested format (hf_quant_config.json) try: quant_config = cls.get_from_keys(config, ["quantization"]) quant_method = quant_config["quant_algo"] group_size = ModelOptFp4Config.common_group_size(config) exclude_modules = quant_config.get("exclude_modules", []) swap_weight_nibbles = quant_config.get( "swap_weight_nibbles", config.get( "swap_weight_nibbles", config.get("checkpoint_uses_packed_qkv", False), ), ) except (ValueError, KeyError): raise ValueError("Cannot find 'quant_algo' in quantization config.") if quant_method not in ["NVFP4"]: raise ValueError( f"Only NVFP4 quantization is supported for diffusion, got '{quant_method}'." ) if group_size is None or exclude_modules is None: raise ValueError( "NVFP4 quantization requires group_size and exclude_modules " "in the quantization config" ) return cls( is_checkpoint_nvfp4_serialized=True, group_size=group_size, exclude_modules=exclude_modules, packed_modules_mapping=config.get("packed_modules_mapping"), checkpoint_uses_packed_qkv=config.get("checkpoint_uses_packed_qkv", False), swap_weight_nibbles=swap_weight_nibbles, checkpoint_weight_scale_layout=config.get( "checkpoint_weight_scale_layout", "linear" ), ) def get_quant_method(self, layer: torch.nn.Module, prefix: str): return self._get_quant_method(layer, prefix, Linear=ModelOptFp4LinearMethod) class ModelOptFp8LinearMethod(LinearMethodBase): """Linear method for ModelOpt static FP8 checkpoints.""" def __init__(self, quant_config: ModelOptFp8Config): self.quant_config = quant_config self.cutlass_fp8_supported = cutlass_fp8_supported() def create_weights( self, layer: torch.nn.Module, input_size_per_partition: int, output_partition_sizes: List[int], input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs, ): del input_size, output_size output_size_per_partition = sum(output_partition_sizes) weight_loader = extra_weight_attrs.get("weight_loader") layer.logical_widths = output_partition_sizes layer.input_size_per_partition = input_size_per_partition layer.output_size_per_partition = output_size_per_partition weight_dtype = ( torch.float8_e4m3fn if self.quant_config.is_checkpoint_fp8_serialized else params_dtype ) layer.register_parameter( "weight", ModelWeightParameter( data=torch.empty( output_size_per_partition, input_size_per_partition, dtype=weight_dtype, ), input_dim=1, output_dim=0, weight_loader=weight_loader, ), ) if self.quant_config.is_checkpoint_fp8_serialized: for scale_name in ["weight_scale", "input_scale"]: layer.register_parameter( scale_name, PerTensorScaleParameter( data=torch.full( (len(output_partition_sizes),), torch.finfo(torch.float32).min, dtype=torch.float32, ), weight_loader=weight_loader, ), ) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: max_w_scale, quantized_weight = requantize_with_max_scale( layer.weight, layer.weight_scale, layer.logical_widths ) # Preserve the parameter subclass metadata while rebinding to the # transposed FP8 view expected by the runtime. layer.weight.data = quantized_weight.t().detach() layer.weight.requires_grad_(False) if self.cutlass_fp8_supported: max_w_scale = convert_to_channelwise(max_w_scale, layer.logical_widths) copy_or_rebind_param(layer, "weight_scale", max_w_scale) copy_or_rebind_param(layer, "input_scale", layer.input_scale.max()) def apply( self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: return apply_fp8_linear( input=x, weight=layer.weight, weight_scale=layer.weight_scale, input_scale=layer.input_scale, bias=bias, cutlass_fp8_supported=self.cutlass_fp8_supported, ) class ModelOptFp4LinearMethod(LinearMethodBase): """NVFP4 linear method using the selected FP4 GEMM backend.""" def __init__(self, quant_config: ModelOptFp4Config): self.quant_config = quant_config def create_weights( self, layer: torch.nn.Module, input_size_per_partition: int, output_partition_sizes: List[int], input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs, ): del input_size, output_size if not self.quant_config.is_checkpoint_nvfp4_serialized: raise ValueError( "NVFP4 quantization was selected, " " dynamic quantization is not supported." ) if input_size_per_partition % 16 != 0: raise ValueError( f"Unsupported model when input features size is {input_size_per_partition}, not multiple of 16, for NVFP4 quantization." ) output_size_per_partition = sum(output_partition_sizes) weight_loader = extra_weight_attrs.get("weight_loader") layer.logical_widths = output_partition_sizes layer.input_size_per_partition = input_size_per_partition layer.output_size_per_partition = output_size_per_partition weight_dtype = ( torch.float8_e4m3fn if self.quant_config.is_checkpoint_nvfp4_serialized else params_dtype ) weight = ModelWeightParameter( data=torch.empty( output_size_per_partition, input_size_per_partition // 2, dtype=torch.uint8, ), input_dim=1, output_dim=0, weight_loader=weight_loader, ) layer.register_parameter("weight", weight) input_scale = PerTensorScaleParameter( data=torch.empty(len(output_partition_sizes), dtype=torch.float32), weight_loader=weight_loader, ) set_weight_attrs(input_scale, {"missing_param_init": "ones"}) layer.register_parameter("input_scale", input_scale) weight_scale_2 = PerTensorScaleParameter( data=torch.empty(len(output_partition_sizes), dtype=torch.float32), weight_loader=weight_loader, ) set_weight_attrs(weight_scale_2, {"missing_param_init": "ones"}) layer.register_parameter("weight_scale_2", weight_scale_2) weight_scale = ModelWeightParameter( data=torch.empty( output_size_per_partition, input_size_per_partition // self.quant_config.group_size, dtype=weight_dtype, ), input_dim=1, output_dim=0, weight_loader=weight_loader, ) set_weight_attrs(weight_scale, {"missing_param_init": "ones"}) layer.register_parameter("weight_scale", weight_scale) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: input_scale_2 = layer.input_scale.max().to(torch.float32) weight_scale_2 = layer.weight_scale_2.max().to(torch.float32) copy_or_rebind_param( layer, "alpha", (input_scale_2 * weight_scale_2).to(torch.float32) ) copy_or_rebind_param( layer, "input_scale_inv", (1 / input_scale_2).to(torch.float32) ) layer.output_size_per_partition = layer.weight.shape[0] w = layer.weight.data w_swapped = _prepare_nvfp4_weight_bytes( w, swap_weight_nibbles=getattr( self.quant_config, "swap_weight_nibbles", False ), ) scales = layer.weight_scale if ( getattr(self.quant_config, "checkpoint_weight_scale_layout", "linear") == "swizzled" ): scales = _swizzled_nvfp4_scales_to_linear(scales) _, flashinfer_backend = _get_fp4_gemm_op() if flashinfer_backend == "trtllm": flashinfer_ops = _require_flashinfer() weight, _ = pad_nvfp4_weight(w_swapped, n_alignment=128, k_alignment=0) if scales.shape[0] != weight.shape[0]: pad_n = weight.shape[0] - scales.shape[0] scales = torch.nn.functional.pad(scales, (0, 0, 0, pad_n)) scale_k = scales.shape[1] weights_padding_cols = 0 if scale_k % 4 != 0: padded_scale_k = round_up(scale_k, 4) pad_scale_k = padded_scale_k - scale_k scales = torch.nn.functional.pad(scales, (0, pad_scale_k, 0, 0)) pad_weight_k = pad_scale_k * 8 weight = torch.nn.functional.pad(weight, (0, pad_weight_k, 0, 0)) weights_padding_cols = pad_weight_k epilogue_tile_m = 128 shuffled_scale_shape = scales.shape if not weight.is_cuda: weight = weight.cuda() if scales.device != weight.device: scales = scales.to(device=weight.device) weight = flashinfer_ops.shuffle_matrix_a( weight.view(torch.uint8), epilogue_tile_m ) scales = ( flashinfer_ops.shuffle_matrix_sf_a( scales.view(torch.uint8), epilogue_tile_m ) .reshape(shuffled_scale_shape) .view(torch.float8_e4m3fn) ) layer.weights_padding_cols = weights_padding_cols copy_or_rebind_param(layer, "weight", weight) copy_or_rebind_param(layer, "weight_scale_interleaved", scales) return weight, weights_padding_cols = pad_nvfp4_weight(w_swapped) layer.weights_padding_cols = weights_padding_cols copy_or_rebind_param(layer, "weight", weight) scale_ndim = scales.ndim if scale_ndim == 2: scales = scales.unsqueeze(0) assert scales.ndim == 3 B, M, K = scales.shape M_padded = round_up(M, 128) K_padded = round_up(K, 4) padded_scales = torch.zeros((B, M_padded, K_padded), dtype=scales.dtype) padded_scales[:B, :M, :K] = scales _, flashinfer_backend = _get_fp4_gemm_op() uses_flux1_scale_layout = not getattr( self.quant_config, "checkpoint_uses_packed_qkv", False ) and getattr(layer, "prefix", "").startswith( ("transformer_blocks.", "single_transformer_blocks.") ) if flashinfer_backend is None or uses_flux1_scale_layout: # CUTLASS and FLUX.1 CUDNN paths need the TMA scale layout. padded_scales = padded_scales.reshape( B, M_padded // 128, 4, 32, K_padded // 4, 4 ) padded_scales = padded_scales.permute(0, 1, 4, 3, 2, 5) padded_scales = padded_scales.contiguous().cuda() padded_scales = ( padded_scales.reshape(M_padded, K_padded) if scale_ndim == 2 else padded_scales.reshape(B, M_padded, K_padded) ) copy_or_rebind_param(layer, "weight_scale_interleaved", padded_scales) def apply( self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: output_dtype = x.dtype input_shape = x.shape x = x.view(-1, input_shape[-1]) output_size = layer.output_size_per_partition output_shape = list(input_shape[:-1]) + [output_size] fp4_quantize = _get_fp4_quantize_op() if fp4_quantize is None: raise RuntimeError( "No FP4 quantization kernel available. Install flashinfer or sgl_kernel." ) x_fp4, x_scale_interleaved = fp4_quantize(x, layer.input_scale_inv) weights_padding_cols = getattr(layer, "weights_padding_cols", 0) x_fp4 = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_cols) w = layer.weight w_scale_interleaved = layer.weight_scale_interleaved if x_scale_interleaved.dtype == torch.uint8: x_scale_interleaved = x_scale_interleaved.view(torch.float8_e4m3fn) if w_scale_interleaved.dtype == torch.uint8: w_scale_interleaved = w_scale_interleaved.view(torch.float8_e4m3fn) fp4_gemm, flashinfer_backend = _get_fp4_gemm_op() if flashinfer_backend is not None: out = fp4_gemm( x_fp4, w.T, x_scale_interleaved, w_scale_interleaved.T, layer.alpha, output_dtype, backend=flashinfer_backend, ) elif fp4_gemm is not None: out = fp4_gemm( x_fp4, w, x_scale_interleaved, w_scale_interleaved, layer.alpha, output_dtype, ) else: raise RuntimeError( "No FP4 GEMM kernel available. Install flashinfer or sgl_kernel." ) out = slice_nvfp4_output(out, output_size) if bias is not None: out = out + bias return out.view(*output_shape)