From 67853c580482dbc9b69aad5022a0cdd3a1606570 Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 25 Aug 2026 11:26:37 +0800 Subject: [PATCH] [diffusion] feat: dispatch fp8 companions in mixed NVFP4 checkpoints (#36066) --- docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx | 5 +++-- docs/docs/sglang-diffusion/quantization.mdx | 2 +- .../runtime/layers/quantization/modelopt_quant.py | 15 ++++++++++++++- .../runtime/loader/minimax_h3_weights.py | 2 +- .../test/unit/test_transformer_quant.py | 14 +++++++++++++- 5 files changed, 32 insertions(+), 6 deletions(-) diff --git a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx index ac67a5bff..eae27a0d6 100644 --- a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx +++ b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx @@ -319,8 +319,9 @@ 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. +Mixed files may mark selected linears as `int8_tensorwise` or dynamic/static +FP8; SGLang dispatches those layers to their serialized Kitchen INT8 or native +FP8 path automatically. ### Advanced: precomputed AdaLN cache diff --git a/docs/docs/sglang-diffusion/quantization.mdx b/docs/docs/sglang-diffusion/quantization.mdx index 88248d748..bc8fee5b0 100644 --- a/docs/docs/sglang-diffusion/quantization.mdx +++ b/docs/docs/sglang-diffusion/quantization.mdx @@ -162,7 +162,7 @@ backend. --transformer-path for mixed overrides; --transformer-weights-path for raw exports; --model-path for full repos FLUX.1, FLUX.2, Wan2.2, Qwen Image, Qwen Image 2512, Qwen Image Edit, Qwen Image Edit 2511, MiniMax-H3 None - Mixed override repos keep the base model separate; full Qwen Image exports can be loaded directly as --model-path; raw exports such as black-forest-labs/FLUX.2-dev-NVFP4 use the weights-path flow. Comfy markers select their checkpoint layout automatically; omit --quantization. + Mixed override repos keep the base model separate; full Qwen Image exports can be loaded directly as --model-path; raw exports such as black-forest-labs/FLUX.2-dev-NVFP4 use the weights-path flow. Comfy markers select NVFP4 plus INT8 or FP8 companion linears automatically; omit --quantization. gguf 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 79ea6e9fc..38e8e4560 100755 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py +++ b/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) 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 c573fd0d7..d314b8dc4 100644 --- a/python/sglang/multimodal_gen/runtime/loader/minimax_h3_weights.py +++ b/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): " 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 4b226bc6c..4a847ec7a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/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(