[diffusion] feat: dispatch fp8 companions in mixed NVFP4 checkpoints (#36066)

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

@@ -13,6 +13,7 @@ from sglang.multimodal_gen.runtime.layers.linear import (
     LinearMethodBase,
     UnquantizedLinearMethod,
 )
+from sglang.multimodal_gen.runtime.layers.quantization.comfy_fp8 import ComfyFp8Config
 from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
     QuantizationConfig,
     QuantizeMethodBase,
@@ -254,11 +255,12 @@ class ModelOptFp4Config(ModelOptQuantConfig):
         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
+        self._comfy_fp8_config: ComfyFp8Config | 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"}
+        } - {"nvfp4", "int8_tensorwise", "float8_e4m3fn"}
         if unsupported:
             raise ValueError(
                 "NVFP4 checkpoints cannot dispatch companion Comfy formats: "
@@ -272,6 +274,12 @@ class ModelOptFp4Config(ModelOptQuantConfig):
         self._comfy_int8_config = (
             KitchenInt8Config(layer_markers=int8_markers) if int8_markers else None
         )
+        fp8_markers = {
+            prefix: marker
+            for prefix, marker in layer_markers.items()
+            if marker.get("format") == "float8_e4m3fn"
+        }
+        self._comfy_fp8_config = ComfyFp8Config(fp8_markers) if fp8_markers else None
 
     @classmethod
     def get_name(cls) -> str:
@@ -383,6 +391,11 @@ class ModelOptFp4Config(ModelOptQuantConfig):
             and prefix in self._comfy_int8_config.layer_markers
         ):
             return self._comfy_int8_config.get_quant_method(layer, prefix)
+        if (
+            self._comfy_fp8_config is not None
+            and prefix in self._comfy_fp8_config.layer_markers
+        ):
+            return self._comfy_fp8_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

@@ -57,7 +57,7 @@ def resolve_minimax_h3_checkpoint_quantization(
 ) -> QuantizationConfig | None:
     formats = {str(marker.get("format")) for marker in layer_markers.values()}
     if "nvfp4" in formats:
-        unsupported = formats - {"nvfp4", "int8_tensorwise"}
+        unsupported = formats - {"nvfp4", "int8_tensorwise", "float8_e4m3fn"}
         if unsupported:
             raise NotImplementedError(
                 "Unsupported Comfy NVFP4 companion format(s): "

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

@@ -1252,7 +1252,7 @@ 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):
+    def test_minimax_h3_mixed_nvfp4_companions_dispatch_each_layer(self):
         metadata = {
             "_quantization_metadata": json.dumps(
                 {
@@ -1264,6 +1264,7 @@ class TestTransformerQuantHelpers(unittest.TestCase):
                             "convrot": True,
                             "convrot_groupsize": 256,
                         },
+                        "blocks.0.mlp.fc1": {"format": "float8_e4m3fn"},
                     },
                 }
             )
@@ -1282,6 +1283,10 @@ class TestTransformerQuantHelpers(unittest.TestCase):
                         (32, 256), dtype=torch.int8
                     ),
                     "blocks.0.attn.out_proj.weight_scale": torch.ones((32, 1)),
+                    "blocks.0.mlp.fc1.weight": torch.ones(
+                        (32, 64), dtype=torch.float8_e4m3fn
+                    ),
+                    "blocks.0.mlp.fc1.weight_scale": torch.tensor(1.0),
                 },
                 checkpoint.name,
                 metadata=metadata,
@@ -1316,6 +1321,13 @@ class TestTransformerQuantHelpers(unittest.TestCase):
                 ),
                 KitchenInt8LinearMethod,
             )
+        self.assertIsInstance(
+            config.get_quant_method(
+                LinearBase(input_size=64, output_size=32),
+                "blocks.0.mlp.fc1",
+            ),
+            Fp8LinearMethod,
+        )
 
     def test_builder_adds_diffusers_quant_type_for_nvfp4(self):
         updated = _updated_quant_config(