[diffusion] feat: support loading serialized convrot w4a4 checkpoints (#36039)

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:
 

docs/docs/sglang-diffusion/quantization.mdx

@@ -200,6 +200,14 @@ backend.
       <td><code>comfy-kitchen&gt;=0.2.27</code></td>
       <td>Auto-detected; omit <code>--quantization</code>. 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.</td>
     </tr>
+    <tr>
+      <td><code>comfy-w4a4-convrot</code></td>
+      <td>Safetensors with serialized <code>convrot_w4a4</code> layer metadata</td>
+      <td><code>--transformer-weights-path</code> or <code>--component-paths.text_encoder</code></td>
+      <td>Native DiTs and encoders with matching parameter mappings; MiniMax-H3 FL2VA / Ref2VA DiTs and Qwen3-VL encoder layouts are recognized</td>
+      <td><code>comfy-kitchen</code></td>
+      <td>Auto-detected; omit <code>--quantization</code>. Packed INT4 weights use the checkpoint's W4A4 ConvRot kernel and <code>linear_dtype</code>. CUDA requires SM75+; TP must preserve the 64-element quantization and ConvRot group boundaries. Offload is supported and FSDP is not.</td>
+    </tr>
     <tr>
       <td><code>quanto-int8</code></td>
       <td>One native encoder safetensors file with an embedded Quanto quantization map</td>

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 []

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

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):

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"

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

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 = {

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"