From ddea7b9156d600c04728e0dfa644a1d1e8ce3b6d Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 25 Aug 2026 09:19:47 +0800 Subject: [PATCH] [diffusion] feat: support mixed Comfy NVFP4 and INT8 layers (#36061) --- .../cookbook/diffusion/MiniMax/MiniMax-H3.mdx | 2 + .../layers/quantization/modelopt_quant.py | 27 ++++++++ .../runtime/loader/minimax_h3_weights.py | 14 +++- .../test/unit/test_transformer_quant.py | 69 +++++++++++++++++++ 4 files changed, 110 insertions(+), 2 deletions(-) diff --git a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx index 6f0683f85..ac67a5bff 100644 --- a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx +++ b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx @@ -319,6 +319,8 @@ compute capability 10.0 or newer. Pass a pruned FL2VA / Ref2VA file such as `Abiray/Minimax-H3-nvfp4-INT4-INT8-Convrot/MiniMax_H3_FL2VA_pruned_nvfp4.safetensors` to `--transformer-weights-path` and omit `--quantization`. SGLang infers the packed group size and Comfy scale layout from the checkpoint; FSDP is rejected. +Mixed files may mark selected linears as `int8_tensorwise`; SGLang dispatches +those layers to the serialized Kitchen INT8 ConvRot path automatically. ### Advanced: precomputed AdaLN cache diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py index b10c5aa12..79ea6e9fc 100755 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/loader/minimax_h3_weights.py b/python/sglang/multimodal_gen/runtime/loader/minimax_h3_weights.py index ad9059c05..c573fd0d7 100644 --- a/python/sglang/multimodal_gen/runtime/loader/minimax_h3_weights.py +++ b/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" 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 ebf124ae3..4b226bc6c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/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( {