[diffusion] feat: support mixed Comfy NVFP4 and INT8 layers (#36061)

python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py

@@ -17,6 +17,9 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config impor
     QuantizationConfig,
     QuantizeMethodBase,
 )
+from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_int8_config import (
+    KitchenInt8Config,
+)
 from sglang.multimodal_gen.runtime.models.parameter import (
     ModelWeightParameter,
     PerTensorScaleParameter,
@@ -250,6 +253,25 @@ class ModelOptFp4Config(ModelOptQuantConfig):
         self.swap_weight_nibbles = swap_weight_nibbles
         self.checkpoint_weight_scale_layout = checkpoint_weight_scale_layout
         self.checkpoint_uses_comfy_quantization = checkpoint_uses_comfy_quantization
+        self._comfy_int8_config: KitchenInt8Config | None = None
+
+    def set_comfy_layer_markers(self, layer_markers: dict[str, dict[str, Any]]) -> None:
+        unsupported = {
+            str(marker.get("format")) for marker in layer_markers.values()
+        } - {"nvfp4", "int8_tensorwise"}
+        if unsupported:
+            raise ValueError(
+                "NVFP4 checkpoints cannot dispatch companion Comfy formats: "
+                + ", ".join(sorted(unsupported))
+            )
+        int8_markers = {
+            prefix: marker
+            for prefix, marker in layer_markers.items()
+            if marker.get("format") == "int8_tensorwise"
+        }
+        self._comfy_int8_config = (
+            KitchenInt8Config(layer_markers=int8_markers) if int8_markers else None
+        )
 
     @classmethod
     def get_name(cls) -> str:
@@ -356,6 +378,11 @@ class ModelOptFp4Config(ModelOptQuantConfig):
         )
 
     def get_quant_method(self, layer: torch.nn.Module, prefix: str):
+        if (
+            self._comfy_int8_config is not None
+            and prefix in self._comfy_int8_config.layer_markers
+        ):
+            return self._comfy_int8_config.get_quant_method(layer, prefix)
         return self._get_quant_method(layer, prefix, Linear=ModelOptFp4LinearMethod)
 
 

python/sglang/multimodal_gen/runtime/loader/minimax_h3_weights.py

@@ -8,6 +8,9 @@ from safetensors import safe_open
 from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
     QuantizationConfig,
 )
+from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
+    ModelOptFp4Config,
+)
 from sglang.multimodal_gen.runtime.utils.quantization_utils import (
     build_nvfp4_config_from_safetensors_list,
     inspect_comfy_quant_markers,
@@ -53,7 +56,13 @@ def resolve_minimax_h3_checkpoint_quantization(
     reverse_param_names_mapping: dict | None = None,
 ) -> QuantizationConfig | None:
     formats = {str(marker.get("format")) for marker in layer_markers.values()}
-    if formats == {"nvfp4"}:
+    if "nvfp4" in formats:
+        unsupported = formats - {"nvfp4", "int8_tensorwise"}
+        if unsupported:
+            raise NotImplementedError(
+                "Unsupported Comfy NVFP4 companion format(s): "
+                + ", ".join(sorted(unsupported))
+            )
         if safetensors_list is None:
             raise ValueError("MiniMax-H3 NVFP4 metadata requires checkpoint files")
         config = build_nvfp4_config_from_safetensors_list(
@@ -61,8 +70,9 @@ def resolve_minimax_h3_checkpoint_quantization(
             param_names_mapping,
             reverse_param_names_mapping,
         )
-        if config is None:
+        if not isinstance(config, ModelOptFp4Config):
             raise ValueError("Could not resolve MiniMax-H3 NVFP4 checkpoint layout")
+        config.set_comfy_layer_markers(layer_markers)
         config.checkpoint_uses_comfy_quantization = True
         config.checkpoint_uses_native_qkv_layout = True
         config.checkpoint_weight_scale_layout = "swizzled"

python/sglang/multimodal_gen/test/unit/test_transformer_quant.py

@@ -73,8 +73,12 @@ from sglang.multimodal_gen.runtime.layers.quantization.fp8 import (
     Fp8Config,
     Fp8LinearMethod,
 )
+from sglang.multimodal_gen.runtime.layers.quantization.kitchen_int8 import (
+    KitchenInt8LinearMethod,
+)
 from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
     ModelOptFp4Config,
+    ModelOptFp4LinearMethod,
     ModelOptFp8Config,
     _prepare_nvfp4_weight_bytes,
 )
@@ -1248,6 +1252,71 @@ class TestTransformerQuantHelpers(unittest.TestCase):
         self.assertEqual(config.checkpoint_weight_scale_layout, "swizzled")
         self.assertTrue(config.swap_weight_nibbles)
 
+    def test_minimax_h3_mixed_nvfp4_int8_dispatches_each_layer(self):
+        metadata = {
+            "_quantization_metadata": json.dumps(
+                {
+                    "format_version": "1.0",
+                    "layers": {
+                        "blocks.0.attn.qkv_proj": {"format": "nvfp4"},
+                        "blocks.0.attn.out_proj": {
+                            "format": "int8_tensorwise",
+                            "convrot": True,
+                            "convrot_groupsize": 256,
+                        },
+                    },
+                }
+            )
+        }
+        with tempfile.NamedTemporaryFile(suffix=".safetensors") as checkpoint:
+            save_file(
+                {
+                    "blocks.0.attn.qkv_proj.weight": torch.zeros(
+                        (32, 8), dtype=torch.uint8
+                    ),
+                    "blocks.0.attn.qkv_proj.weight_scale": torch.ones(
+                        (32, 1), dtype=torch.float8_e4m3fn
+                    ),
+                    "blocks.0.attn.qkv_proj.weight_scale_2": torch.tensor(1.0),
+                    "blocks.0.attn.out_proj.weight": torch.zeros(
+                        (32, 256), dtype=torch.int8
+                    ),
+                    "blocks.0.attn.out_proj.weight_scale": torch.ones((32, 1)),
+                },
+                checkpoint.name,
+                metadata=metadata,
+            )
+            _, markers = inspect_minimax_h3_safetensors([checkpoint.name])
+            config = resolve_minimax_h3_checkpoint_quantization(
+                markers,
+                [checkpoint.name],
+            )
+
+        self.assertIsInstance(config, ModelOptFp4Config)
+        with patch(
+            "sglang.multimodal_gen.runtime.layers.quantization."
+            "modelopt_quant.current_platform.get_device_capability",
+            return_value=DeviceCapability(10, 0),
+        ):
+            self.assertIsInstance(
+                config.get_quant_method(
+                    LinearBase(input_size=16, output_size=32),
+                    "blocks.0.attn.qkv_proj",
+                ),
+                ModelOptFp4LinearMethod,
+            )
+        with patch(
+            "sglang.multimodal_gen.runtime.layers.quantization."
+            "kitchen_int8._load_comfy_kitchen"
+        ):
+            self.assertIsInstance(
+                config.get_quant_method(
+                    LinearBase(input_size=256, output_size=32),
+                    "blocks.0.attn.out_proj",
+                ),
+                KitchenInt8LinearMethod,
+            )
+
     def test_builder_adds_diffusers_quant_type_for_nvfp4(self):
         updated = _updated_quant_config(
             {