From bfeae4e79a8dc4600e006f1a5fbc85321a01c1a3 Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 24 Aug 2026 19:13:35 +0800 Subject: [PATCH] [diffusion] feat: support loading serialized convrot w4a4 checkpoints (#36039) --- .../cookbook/diffusion/MiniMax/MiniMax-H3.mdx | 16 ++ docs/docs/sglang-diffusion/quantization.mdx | 8 + .../configs/kitchen_w4a4_config.py | 137 ++++++++++++++++++ .../layers/quantization/kitchen_w4a4.py | 95 ++++++++++++ .../component_loaders/text_encoder_loader.py | 8 +- .../runtime/loader/transformer_load_utils.py | 8 + .../runtime/utils/quantization_utils.py | 35 +++++ .../test/unit/test_text_encoder_loader.py | 42 ++++++ .../test/unit/test_transformer_quant.py | 59 ++++++++ 9 files changed, 407 insertions(+), 1 deletion(-) create mode 100644 python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a4_config.py create mode 100644 python/sglang/multimodal_gen/runtime/layers/quantization/kitchen_w4a4.py diff --git a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx index 4a2bf9edb..97b96ac08 100644 --- a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx +++ b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx @@ -306,6 +306,12 @@ and packed INT4 tensors automatically. Do not add `--quantization`. TP remains subject to each row-parallel shard preserving the checkpoint's ConvRot group boundary, and FSDP is rejected. +W4A4 ConvRot files are also detected from their layer metadata. Install a +current `comfy-kitchen`, then pass a full or pruned FL2VA / Ref2VA file such as +`Merserk/MiniMax-H3-INT4-ConvRot/minimax_h3_fl2va_pruned_int4_convrot.safetensors` +to `--transformer-weights-path`. The packed weights stay INT4 and the runtime +honors each layer's activation mode; omit `--quantization`. + ### Advanced: precomputed AdaLN cache The [model card](https://huggingface.co/MiniMaxAI/MiniMax-H3) notes that about @@ -969,6 +975,16 @@ Install `comfy-kitchen>=0.2.27` and omit `--quantization`. SGLang automatically loads its W4A8 language linears and tensorwise INT8 embedding; the unmarked vision tower remains BF16. +W4A4 Qwen3-VL files use the same overlay, for example: + +```bash Overlay +--component-paths.text_encoder \ + Merserk/MiniMax-H3-INT4-ConvRot/qwen3vl_32b_minimax_h3_int4_convrot.safetensors +``` + +This checkpoint keeps its unmarked embedding and vision tower in their source +precision; no explicit component quantization option is needed. + The same component option accepts a self-describing Quanto qint8 file without an additional quantization flag: diff --git a/docs/docs/sglang-diffusion/quantization.mdx b/docs/docs/sglang-diffusion/quantization.mdx index a19f0d2b1..c0e86abd2 100644 --- a/docs/docs/sglang-diffusion/quantization.mdx +++ b/docs/docs/sglang-diffusion/quantization.mdx @@ -200,6 +200,14 @@ backend. comfy-kitchen>=0.2.27 Auto-detected; omit --quantization. Requires SM80+ and validates packed weights, group/channel scales, and optional codebooks before model construction. Mixed encoder files may keep their embedding tensorwise INT8. TP must preserve ConvRot group boundaries; offload is supported and FSDP is not. + + comfy-w4a4-convrot + Safetensors with serialized convrot_w4a4 layer metadata + --transformer-weights-path or --component-paths.text_encoder + Native DiTs and encoders with matching parameter mappings; MiniMax-H3 FL2VA / Ref2VA DiTs and Qwen3-VL encoder layouts are recognized + comfy-kitchen + Auto-detected; omit --quantization. Packed INT4 weights use the checkpoint's W4A4 ConvRot kernel and linear_dtype. CUDA requires SM75+; TP must preserve the 64-element quantization and ConvRot group boundaries. Offload is supported and FSDP is not. + quanto-int8 One native encoder safetensors file with an embedded Quanto quantization map diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a4_config.py b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a4_config.py new file mode 100644 index 000000000..3dcf43477 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a4_config.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Config for serialized Comfy Kitchen ConvRot W4A4 weights.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from sglang.multimodal_gen.runtime.layers.linear import ( + LinearBase, + UnquantizedLinearMethod, +) +from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.multimodal_gen.runtime.layers.quantization.kitchen_w4a4 import ( + KitchenW4A4LinearMethod, +) +from sglang.multimodal_gen.runtime.platforms import current_platform + +_QUANT_GROUP_SIZE = 64 +_SUPPORTED_CONVROT_GROUP_SIZES = (16, 64, 256) +_SUPPORTED_LINEAR_DTYPES = ("int4", "int8") + + +class KitchenW4A4Config(QuantizationConfig): + """Dispatch linears carrying serialized ``convrot_w4a4`` markers.""" + + def __init__(self, layer_markers: dict[str, dict[str, Any]]) -> None: + super().__init__() + if current_platform.is_mps(): + raise ValueError("Serialized W4A4 checkpoints are not supported on MPS") + if current_platform.is_cuda(): + capability = current_platform.get_device_capability() + if ( + capability is not None + and capability.to_int() < self.get_min_capability() + ): + raise ValueError( + "Serialized W4A4 checkpoints require CUDA compute capability " + f">= {self.get_min_capability() / 10:.1f}; got " + f"{capability.to_int() / 10:.1f}" + ) + self.layer_markers = layer_markers + self.checkpoint_uses_native_qkv_layout = True + self.selected: list[str] = [] + + for prefix, marker in layer_markers.items(): + if marker.get("format") != "convrot_w4a4": + raise ValueError( + f"Unsupported Comfy W4A4 format for {prefix!r}: " + f"{marker.get('format')!r}" + ) + self._parse_marker(prefix, marker) + + @classmethod + def get_name(cls) -> str: + return "kitchen_w4a4" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.bfloat16, torch.float16] + + @classmethod + def get_min_capability(cls) -> int: + return 75 + + @classmethod + def get_config_filenames(cls) -> list[str]: + return [] + + @classmethod + def from_config(cls, config: dict[str, Any]) -> KitchenW4A4Config: + raise ValueError( + "kitchen_w4a4 is inferred from per-layer checkpoint metadata; " + "it is not an online quantization method" + ) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> QuantizeMethodBase | None: + if not isinstance(layer, LinearBase): + return None + marker = self.layer_markers.get(prefix) + if marker is None: + return UnquantizedLinearMethod() + + convrot_group_size, linear_dtype = self._parse_marker(prefix, marker) + if not self._supports_input_size(layer.input_size, convrot_group_size): + raise ValueError( + f"Serialized W4A4 layer {prefix!r} has input size " + f"{layer.input_size}, incompatible with quant_group_size=" + f"{_QUANT_GROUP_SIZE} and convrot_groupsize={convrot_group_size}" + ) + self.selected.append(prefix) + return KitchenW4A4LinearMethod( + convrot_group_size=convrot_group_size, + linear_dtype=linear_dtype, + ) + + @staticmethod + def _parse_marker(prefix: str, marker: dict[str, Any]) -> tuple[int, str]: + convrot_group_size = int(marker.get("convrot_groupsize", 256)) + if convrot_group_size not in _SUPPORTED_CONVROT_GROUP_SIZES: + raise ValueError( + f"Serialized W4A4 layer {prefix!r} has unsupported " + f"convrot_groupsize={convrot_group_size}; expected one of " + f"{_SUPPORTED_CONVROT_GROUP_SIZES}" + ) + linear_dtype = str(marker.get("linear_dtype", "int4")) + if linear_dtype not in _SUPPORTED_LINEAR_DTYPES: + raise ValueError( + f"Serialized W4A4 layer {prefix!r} has unsupported " + f"linear_dtype={linear_dtype!r}; expected one of " + f"{_SUPPORTED_LINEAR_DTYPES}" + ) + return convrot_group_size, linear_dtype + + @staticmethod + def _supports_input_size(input_size: int, convrot_group_size: int) -> bool: + return ( + input_size % _QUANT_GROUP_SIZE == 0 and input_size % convrot_group_size == 0 + ) + + def supports_input_partition( + self, prefix: str, input_size_per_partition: int + ) -> bool: + marker = self.layer_markers.get(prefix) + if marker is None: + return True + convrot_group_size, _ = self._parse_marker(prefix, marker) + return self._supports_input_size(input_size_per_partition, convrot_group_size) + + def get_scaled_act_names(self) -> list[str]: + return [] diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/kitchen_w4a4.py b/python/sglang/multimodal_gen/runtime/layers/quantization/kitchen_w4a4.py new file mode 100644 index 000000000..10e49badc --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/kitchen_w4a4.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Serialized ConvRot W4A4 linear backed by Comfy Kitchen.""" + +from __future__ import annotations + +import torch +from torch.nn.parameter import Parameter + +from sglang.multimodal_gen.runtime.layers.linear import LinearMethodBase +from sglang.multimodal_gen.runtime.utils.weight_attrs import set_weight_attrs + +try: + from comfy_kitchen import convrot_w4a4_linear +except ImportError: # pragma: no cover - optional dependency + convrot_w4a4_linear = None + +_QUANT_GROUP_SIZE = 64 + + +class KitchenW4A4LinearMethod(LinearMethodBase): + """Load packed INT4 weights and execute the ConvRot W4A4 kernel.""" + + def __init__(self, *, convrot_group_size: int, linear_dtype: str) -> None: + if convrot_w4a4_linear is None: + raise ImportError( + "W4A4 checkpoints require a current comfy-kitchen build " + "(`pip install -U comfy-kitchen`)." + ) + self.convrot_group_size = convrot_group_size + self.linear_dtype = linear_dtype + + 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, + ) -> None: + del input_size, output_size, params_dtype + if input_size_per_partition % self.convrot_group_size: + raise ValueError( + "W4A4 needs input_size_per_partition " + f"({input_size_per_partition}) divisible by ConvRot group size " + f"{self.convrot_group_size}" + ) + if input_size_per_partition % _QUANT_GROUP_SIZE: + raise ValueError( + "W4A4 needs input_size_per_partition " + f"({input_size_per_partition}) divisible by quantization group size " + f"{_QUANT_GROUP_SIZE}" + ) + + output_size_per_partition = sum(output_partition_sizes) + weight = Parameter( + torch.empty( + output_size_per_partition, + input_size_per_partition // 2, + dtype=torch.int8, + ), + requires_grad=False, + ) + set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) + set_weight_attrs(weight, extra_weight_attrs) + layer.register_parameter("weight", weight) + + weight_scale = Parameter( + torch.empty(output_size_per_partition, dtype=torch.float32), + requires_grad=False, + ) + set_weight_attrs(weight_scale, {"output_dim": 0}) + set_weight_attrs(weight_scale, extra_weight_attrs) + layer.register_parameter("weight_scale", weight_scale) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + assert convrot_w4a4_linear is not None + return convrot_w4a4_linear( + x.contiguous(), + layer.weight, + layer.weight_scale, + bias=bias, + convrot_groupsize=self.convrot_group_size, + quant_group_size=_QUANT_GROUP_SIZE, + linear_dtype=self.linear_dtype, + ) + + +__all__ = ["KitchenW4A4LinearMethod"] diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py index a0165af11..952bfcc21 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py @@ -34,6 +34,9 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config impor from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_int8_config import ( KitchenInt8Config, ) +from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_w4a4_config import ( + KitchenW4A4Config, +) from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_w4a8_config import ( KitchenW4A8Config, ) @@ -368,7 +371,10 @@ def _require_quantized_encoder_layers( f"The native {type(model).__name__} implementation does not construct " f"quantized linear layers for {component_name!r}" ) - if isinstance(quant_config, (ComfyFp8Config, KitchenInt8Config, KitchenW4A8Config)): + if isinstance( + quant_config, + (ComfyFp8Config, KitchenInt8Config, KitchenW4A4Config, KitchenW4A8Config), + ): expected = set(quant_config.layer_markers) selected = set(quant_config.selected) elif isinstance(quant_config, QuantoInt8Config): diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py index aec5354ff..f98a89a26 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py @@ -21,6 +21,9 @@ from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_int8_config import ( KitchenInt8Config, ) +from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_w4a4_config import ( + KitchenW4A4Config, +) from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_w4a8_config import ( KitchenW4A8Config, ) @@ -180,11 +183,16 @@ class TransformerQuantLoadSpec: def is_serialized_kitchen_w4a8(self) -> bool: return isinstance(self.quant_config, KitchenW4A8Config) + @property + def is_serialized_kitchen_w4a4(self) -> bool: + return isinstance(self.quant_config, KitchenW4A4Config) + @property def uses_comfy_layer_markers(self) -> bool: return ( self.is_comfy_fp8 or self.is_serialized_kitchen_int8 + or self.is_serialized_kitchen_w4a4 or self.is_serialized_kitchen_w4a8 or ( _get_quant_config_name(self.quant_config) == "mxfp8" diff --git a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py index 862416c6e..7ed801b40 100644 --- a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py @@ -16,6 +16,9 @@ from sglang.multimodal_gen.runtime.layers.quantization.comfy_fp8 import ComfyFp8 from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_int8_config import ( KitchenInt8Config, ) +from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_w4a4_config import ( + KitchenW4A4Config, +) from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_w4a8_config import ( KitchenW4A8Config, ) @@ -140,6 +143,7 @@ def inspect_comfy_quant_markers( "float8_e4m3fn", "int8_tensorwise", "asym_w4a8_int8", + "convrot_w4a4", ): continue missing = required - checkpoint_meta.keys() @@ -209,6 +213,35 @@ def inspect_comfy_quant_markers( f"Comfy W4A8 layer {prefix!r} has an incompatible correction tensor" ) continue + if marker_format == "convrot_w4a4": + weight_dtype, weight_shape = checkpoint_meta[f"{prefix}.weight"] + scale_dtype, scale_shape = checkpoint_meta[f"{prefix}.weight_scale"] + if weight_dtype != "I8" or scale_dtype != "F32": + raise ValueError( + f"Comfy W4A4 layer {prefix!r} needs I8 packed weights and " + f"F32 scales, got {weight_dtype} and {scale_dtype}" + ) + if len(weight_shape) != 2 or scale_shape != (weight_shape[0],): + raise ValueError( + f"Comfy W4A4 layer {prefix!r} has incompatible weight/scale " + f"shapes: {weight_shape} and {scale_shape}" + ) + logical_input_size = weight_shape[1] * 2 + convrot_group_size = int(marker.get("convrot_groupsize", 256)) + if convrot_group_size not in (16, 64, 256): + raise ValueError( + f"Comfy W4A4 layer {prefix!r} has unsupported " + f"convrot_groupsize={convrot_group_size}" + ) + if logical_input_size % 64 or logical_input_size % convrot_group_size: + raise ValueError( + f"Comfy W4A4 layer {prefix!r} has input size " + f"{logical_input_size}, incompatible with quant_group_size=64 " + f"and convrot_groupsize={convrot_group_size}" + ) + continue + if marker_format != "int8_tensorwise": + continue weight_dtype, weight_shape = checkpoint_meta[f"{prefix}.weight"] scale_dtype, scale_shape = checkpoint_meta[f"{prefix}.weight_scale"] if weight_dtype == "I8" and scale_dtype == "F32" and scale_shape == (): @@ -253,6 +286,8 @@ def resolve_comfy_checkpoint_quantization( return KitchenW4A8Config(layer_markers) if formats == ["asym_w4a8_int8", "int8_tensorwise"]: return KitchenW4A8Config(layer_markers) + if formats == ["convrot_w4a4"]: + return KitchenW4A4Config(layer_markers) if formats == ["float8_e4m3fn"]: return ComfyFp8Config(layer_markers) if formats == ["mxfp8"]: diff --git a/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py b/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py index e7024c0de..18ce38b3c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py +++ b/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py @@ -13,6 +13,9 @@ from sglang.multimodal_gen.runtime.layers.linear import LinearBase from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_int8_config import ( KitchenInt8Config, ) +from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_w4a4_config import ( + KitchenW4A4Config, +) from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_w4a8_config import ( KitchenW4A8Config, ) @@ -414,6 +417,45 @@ class TestTextEncoderQuantization(unittest.TestCase): {"model.visual.blocks.0.attn.qkv_proj"}, ) + def test_comfy_w4a4_weight_file_configures_native_encoder(self): + self.get_quant_config.return_value = None + marker = json.dumps( + {"format": "convrot_w4a4", "convrot_groupsize": 256} + ).encode() + with tempfile.NamedTemporaryFile(suffix=".safetensors") as checkpoint: + save_file( + { + "model.layers.0.self_attn.q_proj.weight": torch.ones( + (2, 128), dtype=torch.int8 + ), + "model.layers.0.self_attn.q_proj.weight_scale": torch.ones(2), + "model.layers.0.self_attn.q_proj.comfy_quant": torch.tensor( + list(marker), dtype=torch.uint8 + ), + }, + checkpoint.name, + ) + model_config = SimpleNamespace(quant_config=None) + with mock.patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "text_encoder_loader.get_quant_config_from_safetensors_metadata", + return_value=None, + ): + _configure_encoder_quantization( + model_config, + MiniMaxH3Qwen3VLEncoder, + {}, + "/model/text_encoder", + checkpoint.name, + "text_encoder", + ) + + self.assertIsInstance(model_config.quant_config, KitchenW4A4Config) + self.assertEqual( + set(model_config.quant_config.layer_markers), + {"model.language_model.layers.0.self_attn.q_proj"}, + ) + def test_mixed_w4a8_weight_file_maps_embedding_and_linear_markers(self): self.get_quant_config.return_value = None layers = { diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py index 8218e9501..747861d66 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py @@ -60,6 +60,9 @@ from sglang.multimodal_gen.runtime.layers.quantization.comfy_fp8 import ( from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_int8_config import ( KitchenInt8Config, ) +from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_w4a4_config import ( + KitchenW4A4Config, +) from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_w4a8_config import ( KitchenW4A8Config, ) @@ -460,6 +463,62 @@ class TestTransformerQuantHelpers(unittest.TestCase): self.assertEqual(layer.weight_codebook.shape, (16,)) self.assertIsNone(layer.weight_correction) + def test_minimax_h3_w4a4_marker_resolves_packed_kitchen(self): + marker = json.dumps( + { + "format": "convrot_w4a4", + "convrot_groupsize": 256, + "linear_dtype": "int8", + } + ).encode() + with tempfile.NamedTemporaryFile(suffix=".safetensors") as checkpoint: + save_file( + { + "blocks.0.mlp.fc1.weight": torch.ones((2, 128), dtype=torch.int8), + "blocks.0.mlp.fc1.weight_scale": torch.ones(2), + "blocks.0.mlp.fc1.comfy_quant": torch.tensor( + list(marker), dtype=torch.uint8 + ), + }, + checkpoint.name, + ) + + _, markers = inspect_minimax_h3_safetensors([checkpoint.name]) + + config = resolve_minimax_h3_checkpoint_quantization(markers) + self.assertIsInstance(config, KitchenW4A4Config) + self.assertTrue(config.supports_input_partition("blocks.0.mlp.fc1", 256)) + self.assertFalse(config.supports_input_partition("blocks.0.mlp.fc1", 128)) + self.assertFalse(_needs_device_weight_postprocess(config)) + + @patch( + "sglang.multimodal_gen.runtime.layers.quantization.kitchen_w4a4." + "convrot_w4a4_linear", + new=object(), + ) + def test_serialized_w4a4_constructs_packed_weight_and_row_scale(self): + config = KitchenW4A4Config( + { + "proj": { + "format": "convrot_w4a4", + "convrot_groupsize": 256, + } + } + ) + layer = ReplicatedLinear( + 256, + 3, + bias=False, + params_dtype=torch.bfloat16, + quant_config=config, + prefix="proj", + ) + + self.assertEqual(layer.weight.shape, (3, 128)) + self.assertEqual(layer.weight.dtype, torch.int8) + self.assertEqual(layer.weight_scale.shape, (3,)) + self.assertEqual(layer.weight_scale.dtype, torch.float32) + @patch( "sglang.multimodal_gen.runtime.layers.quantization.kitchen_int8." "_load_comfy_kitchen"