diff --git a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx index 97b96ac08..278f569c6 100644 --- a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx +++ b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx @@ -311,6 +311,8 @@ 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`. +Mixed exports use the same command: SGLang dispatches each marked layer to its +serialized W4A4 or INT8 ConvRot kernel instead of applying one global method. ### Advanced: precomputed AdaLN cache diff --git a/docs/docs/sglang-diffusion/quantization.mdx b/docs/docs/sglang-diffusion/quantization.mdx index c0e86abd2..136dc9d69 100644 --- a/docs/docs/sglang-diffusion/quantization.mdx +++ b/docs/docs/sglang-diffusion/quantization.mdx @@ -202,11 +202,11 @@ backend. comfy-w4a4-convrot - Safetensors with serialized convrot_w4a4 layer metadata + Safetensors with serialized convrot_w4a4 metadata, optionally mixed with int8_tensorwise layers --transformer-weights-path or --component-paths.text_encoder Native DiTs and encoders with matching parameter mappings; MiniMax-H3 FL2VA / Ref2VA DiTs and Qwen3-VL encoder layouts are recognized comfy-kitchen - Auto-detected; omit --quantization. Packed INT4 weights use the checkpoint's W4A4 ConvRot kernel and linear_dtype. CUDA requires SM75+; TP must preserve the 64-element quantization and ConvRot group boundaries. Offload is supported and FSDP is not. + Auto-detected; omit --quantization. Each layer dispatches to its serialized W4A4 or INT8 ConvRot kernel. CUDA requires SM75+; TP must preserve each format's quantization and ConvRot group boundaries. Offload is supported and FSDP is not. quanto-int8 diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a4_config.py b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a4_config.py index 3dcf43477..cde021f64 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a4_config.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a4_config.py @@ -15,6 +15,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.layers.quantization.kitchen_w4a4 import ( KitchenW4A4LinearMethod, ) @@ -26,7 +29,7 @@ _SUPPORTED_LINEAR_DTYPES = ("int4", "int8") class KitchenW4A4Config(QuantizationConfig): - """Dispatch linears carrying serialized ``convrot_w4a4`` markers.""" + """Dispatch serialized W4A4 linears and their optional INT8 companions.""" def __init__(self, layer_markers: dict[str, dict[str, Any]]) -> None: super().__init__() @@ -46,12 +49,23 @@ class KitchenW4A4Config(QuantizationConfig): self.layer_markers = layer_markers self.checkpoint_uses_native_qkv_layout = True self.selected: list[str] = [] + int8_markers = { + prefix: marker + for prefix, marker in layer_markers.items() + if marker.get("format") == "int8_tensorwise" + } + self._int8_config = ( + KitchenInt8Config(layer_markers=int8_markers) if int8_markers else None + ) for prefix, marker in layer_markers.items(): - if marker.get("format") != "convrot_w4a4": + marker_format = marker.get("format") + if marker_format == "int8_tensorwise": + continue + if marker_format != "convrot_w4a4": raise ValueError( f"Unsupported Comfy W4A4 format for {prefix!r}: " - f"{marker.get('format')!r}" + f"{marker_format!r}" ) self._parse_marker(prefix, marker) @@ -86,6 +100,11 @@ class KitchenW4A4Config(QuantizationConfig): marker = self.layer_markers.get(prefix) if marker is None: return UnquantizedLinearMethod() + if marker.get("format") == "int8_tensorwise": + assert self._int8_config is not None + method = self._int8_config.get_quant_method(layer, prefix) + self.selected.append(prefix) + return method convrot_group_size, linear_dtype = self._parse_marker(prefix, marker) if not self._supports_input_size(layer.input_size, convrot_group_size): @@ -130,6 +149,11 @@ class KitchenW4A4Config(QuantizationConfig): marker = self.layer_markers.get(prefix) if marker is None: return True + if marker.get("format") == "int8_tensorwise": + assert self._int8_config is not None + return self._int8_config.supports_input_partition( + prefix, input_size_per_partition + ) convrot_group_size, _ = self._parse_marker(prefix, marker) return self._supports_input_size(input_size_per_partition, convrot_group_size) diff --git a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py index 7ed801b40..dd91f4a20 100644 --- a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py @@ -288,6 +288,8 @@ def resolve_comfy_checkpoint_quantization( return KitchenW4A8Config(layer_markers) if formats == ["convrot_w4a4"]: return KitchenW4A4Config(layer_markers) + if formats == ["convrot_w4a4", "int8_tensorwise"]: + return KitchenW4A4Config(layer_markers) if formats == ["float8_e4m3fn"]: return ComfyFp8Config(layer_markers) if formats == ["mxfp8"]: 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 747861d66..55801716c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py @@ -519,6 +519,53 @@ class TestTransformerQuantHelpers(unittest.TestCase): self.assertEqual(layer.weight_scale.shape, (3,)) self.assertEqual(layer.weight_scale.dtype, torch.float32) + def test_mixed_w4a4_int8_dispatches_each_serialized_layer(self): + markers = { + "w4a4": { + "format": "convrot_w4a4", + "convrot_groupsize": 256, + "linear_dtype": "int8", + }, + "int8": { + "format": "int8_tensorwise", + "convrot": True, + "convrot_groupsize": 256, + }, + } + with ( + patch( + "sglang.multimodal_gen.runtime.layers.quantization.kitchen_w4a4." + "convrot_w4a4_linear", + new=object(), + ), + patch( + "sglang.multimodal_gen.runtime.layers.quantization.kitchen_int8." + "_load_comfy_kitchen" + ), + ): + config = resolve_minimax_h3_checkpoint_quantization(markers) + w4a4 = ReplicatedLinear( + 256, + 3, + bias=False, + params_dtype=torch.bfloat16, + quant_config=config, + prefix="w4a4", + ) + int8 = ReplicatedLinear( + 256, + 3, + bias=False, + params_dtype=torch.bfloat16, + quant_config=config, + prefix="int8", + ) + + self.assertIsInstance(config, KitchenW4A4Config) + self.assertEqual(w4a4.weight.shape, (3, 128)) + self.assertEqual(int8.weight.shape, (3, 256)) + self.assertEqual(set(config.selected), {"w4a4", "int8"}) + @patch( "sglang.multimodal_gen.runtime.layers.quantization.kitchen_int8." "_load_comfy_kitchen"