[diffusion] feat: add per-component quantization overrides (#36084)

python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py

@@ -130,6 +130,9 @@ class ComponentLoader(ABC):
     # components may fall back when that global choice is incompatible; an
     # explicit --component-attention-backends entry remains strict.
     allow_global_attention_backend_fallback = True
+    # Gates only --component-quantizations.<name>. Quantization declared by a
+    # checkpoint is discovered and admitted by the component's normal loader.
+    supports_online_quantization_override = False
 
     _loaders_registered = False
 
@@ -227,6 +230,17 @@ class ComponentLoader(ABC):
         If all of the above methods failed, an error will be thrown
 
         """
+        component_quantization = server_args.component_quantizations.get(component_name)
+        if (
+            component_quantization is not None
+            and not self.supports_online_quantization_override
+        ):
+            raise ValueError(
+                f"{component_name!r} does not support an explicit quantization "
+                "override; "
+                "use a self-describing quantized component checkpoint when supported"
+            )
+
         gpu_mem_before_loading = current_platform.get_available_gpu_memory()
         logger.info(
             "Loading %s from %s. avail mem: %.2f GB",

python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py

@@ -46,6 +46,7 @@ class ImageEncoderLoader(TextEncoderLoader):
             component_model_path,
             component_weights_path,
             component_name,
+            server_args.component_quantizations.get(component_name),
         )
         # real dims are populated now; resolve fold vs replicate
         finalize_encoder_folding(

python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py

@@ -88,6 +88,8 @@ from sglang.srt.environ import envs
 
 logger = init_logger(__name__)
 
+_ONLINE_ENCODER_QUANTIZATIONS = frozenset({"fp8", "kitchen_int8", "mxfp4"})
+
 _TRANSFORMERS_ENCODER_ONLY_CLASSES = {
     "T5EncoderModel": transformers.T5EncoderModel,
     "T5Model": transformers.T5EncoderModel,
@@ -159,8 +161,14 @@ def _configure_encoder_quantization(
     component_model_path: str,
     component_weights_path: str,
     component_name: str,
+    explicit_quantization: str | None = None,
 ) -> None:
     if getattr(model_cls, "manages_checkpoint_quantization", False):
+        if explicit_quantization is not None:
+            raise ComponentCheckpointUnsupportedError(
+                f"{component_name!r} manages its own checkpoint quantization and "
+                "does not support an online quantization override"
+            )
         # Preserve model-owned formats such as Ideogram's bitsandbytes state.
         # Those models parse metadata, construct layers, and attach quant states
         # themselves; running the generic lifecycle as well would process twice.
@@ -182,6 +190,24 @@ def _configure_encoder_quantization(
             f"Cannot configure checkpoint quantization for {component_name!r}: {error}"
         ) from error
     model_config.quant_config = quant_config
+    if explicit_quantization is not None:
+        if quant_config is not None:
+            raise ComponentCheckpointUnsupportedError(
+                f"{component_name!r} already declares checkpoint quantization; "
+                "drop the explicit online quantization override"
+            )
+        if explicit_quantization not in _ONLINE_ENCODER_QUANTIZATIONS:
+            raise ComponentCheckpointUnsupportedError(
+                f"Online quantization {explicit_quantization!r} is not supported "
+                f"for native encoders; choose one of "
+                f"{sorted(_ONLINE_ENCODER_QUANTIZATIONS)}"
+            )
+        from sglang.multimodal_gen.runtime.layers.quantization import (
+            get_quantization_config,
+        )
+
+        model_config.quant_config = get_quantization_config(explicit_quantization)()
+        quant_config = model_config.quant_config
     if quant_config is None:
         return
     if not issubclass(model_cls, EncoderTensorParallelMixin):
@@ -198,6 +224,7 @@ def _resolve_and_configure_encoder_quantization(
     component_model_path: str,
     component_weights_path: str,
     component_name: str,
+    explicit_quantization: str | None = None,
 ) -> type[nn.Module]:
     architectures = getattr(model_config, "architectures", [])
     try:
@@ -218,6 +245,11 @@ def _resolve_and_configure_encoder_quantization(
                 f"Cannot parse checkpoint quantization for {component_name!r}: "
                 f"{quantization_error}"
             ) from quantization_error
+        if explicit_quantization is not None and quant_config is None:
+            raise ComponentCheckpointUnsupportedError(
+                f"Online quantization for {component_name!r} requires an in-tree "
+                f"native encoder; unsupported architectures: {architectures}"
+            ) from resolution_error
         if quant_config is None:
             raise
         raise ComponentCheckpointUnsupportedError(
@@ -232,6 +264,7 @@ def _resolve_and_configure_encoder_quantization(
         component_model_path,
         component_weights_path,
         component_name,
+        explicit_quantization,
     )
     return model_cls
 
@@ -356,6 +389,15 @@ class TextEncoderLoader(ComponentLoader):
 
     component_names = ["text_encoder"]
     expected_library = "transformers"
+    supports_online_quantization_override = True
+
+    def should_raise_customized_load_error(
+        self, server_args: ServerArgs, component_name: str
+    ) -> bool:
+        return (
+            super().should_raise_customized_load_error(server_args, component_name)
+            or component_name in server_args.component_quantizations
+        )
 
     @staticmethod
     def resolve_model_weights_path(
@@ -602,6 +644,7 @@ class TextEncoderLoader(ComponentLoader):
             component_model_path,
             component_weights_path,
             component_name,
+            server_args.component_quantizations.get(component_name),
         )
         if issubclass(model_cls, EncoderTensorParallelMixin):
             model_cls.configure_component_paths(

python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py

@@ -100,8 +100,10 @@ def _server_args_for_transformer_component(
 ) -> ServerArgs:
     """Mask global quantized override flags for secondary transformer components."""
     component_weights_path = server_args.component_weights_paths.get(component_name)
-    if component_weights_path is not None:
+    component_quantization = server_args.component_quantizations.get(component_name)
+    if component_weights_path is not None or component_quantization is not None:
         component_server_args = copy.copy(server_args)
+        if component_weights_path is not None:
             component_server_args.transformer_weights_path = component_weights_path
             component_server_args.nunchaku_config = None
             logger.info(
@@ -109,6 +111,13 @@ def _server_args_for_transformer_component(
                 component_name,
                 component_weights_path,
             )
+        if component_quantization is not None:
+            component_server_args.quantization = component_quantization
+            logger.info(
+                "Using quantization override %s for %s",
+                component_quantization,
+                component_name,
+            )
         return component_server_args
 
     if component_name not in ("transformer_2", "unconditional_transformer"):
@@ -135,6 +144,7 @@ class TransformerLoader(ComponentLoader):
     """Shared loader for (video/audio) DiT transformers."""
 
     allow_global_attention_backend_fallback = False
+    supports_online_quantization_override = True
 
     component_names = [
         "transformer",

python/sglang/multimodal_gen/runtime/server_args/server_args.py

@@ -306,6 +306,9 @@ class ServerArgs(DisaggServerArgsMixin):
     component_paths: dict[str, str] = field(default_factory=dict)
     # Exact weight-file overrides retain the base component configuration.
     component_weights_paths: dict[str, str] = field(default_factory=dict)
+    # Explicit quantization override for one component. Self-describing
+    # checkpoints remain auto-detected and do not need this override.
+    component_quantizations: dict[str, str] = field(default_factory=dict)
     # Optional LTX-2.5 decoder is large enough to load only when requested.
     load_diffusion_decoder: bool = False
 
@@ -1722,6 +1725,22 @@ class ServerArgs(DisaggServerArgsMixin):
             component_weights_paths[component] = path
         self.component_paths = component_paths
         self.component_weights_paths = component_weights_paths
+        normalized_quantizations: dict[str, str] = {}
+        for component, quantization in self.component_quantizations.items():
+            component = str(component).strip().replace("-", "_")
+            quantization = str(quantization).strip().lower()
+            if not component or not quantization:
+                raise ValueError(
+                    "Component quantization entries require a component and method"
+                )
+            previous = normalized_quantizations.get(component)
+            if previous is not None and previous != quantization:
+                raise ValueError(
+                    f"Conflicting quantization overrides for {component!r}: "
+                    f"{previous!r} and {quantization!r}"
+                )
+            normalized_quantizations[component] = quantization
+        self.component_quantizations = normalized_quantizations
 
         # Convert string disagg_role to enum (from CLI/config)
         if isinstance(self.disagg_role, str):
@@ -2833,6 +2852,21 @@ class ServerArgs(DisaggServerArgsMixin):
             alias_suffix="-weights-path",
         )
 
+    @classmethod
+    def _extract_component_quantizations(
+        cls,
+        unknown_args: list[str],
+    ) -> tuple[dict[str, str], list[str]]:
+        """Extract explicit per-component quantization methods."""
+        return cls._extract_dynamic_component_map(
+            unknown_args,
+            option_prefixes=(
+                "--component-quantizations.",
+                "--component_quantizations.",
+            ),
+            alias_suffix="-quantization",
+        )
+
     @staticmethod
     def _extract_component_attention_backends(
         unknown_args: list[str],
@@ -2880,9 +2914,12 @@ class ServerArgs(DisaggServerArgsMixin):
         if unknown_args is None:
             unknown_args = []
 
+        dynamic_quantizations, remaining = cls._extract_component_quantizations(
+            unknown_args
+        )
         # Extract the more specific weights suffix before the generic path alias.
         dynamic_weights_paths, remaining = cls._extract_component_weights_paths(
-            unknown_args
+            remaining
         )
         dynamic_paths, remaining = cls._extract_component_paths(remaining)
         dynamic_attention_backends, remaining = (
@@ -2915,6 +2952,11 @@ class ServerArgs(DisaggServerArgsMixin):
             existing.update(dynamic_weights_paths)
             provided_args["component_weights_paths"] = existing
             explicit_arg_names.add("component_weights_paths")
+        if dynamic_quantizations:
+            existing = dict(provided_args.get("component_quantizations") or {})
+            existing.update(dynamic_quantizations)
+            provided_args["component_quantizations"] = existing
+            explicit_arg_names.add("component_quantizations")
         if dynamic_attention_backends:
             existing = cls._parse_component_attention_backend_map(
                 provided_args.get("component_attention_backends")

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

@@ -510,6 +510,7 @@ class TestIdeogram4(unittest.TestCase):
             transformer_weights_path="/unused/override.safetensors",
             nunchaku_config={"enabled": True},
             component_weights_paths={},
+            component_quantizations={},
         )
         component_args = _server_args_for_transformer_component(
             server_args, "unconditional_transformer"
@@ -530,6 +531,7 @@ class TestIdeogram4(unittest.TestCase):
                     "ideogram4_unconditional_nvfp4_mixed.safetensors"
                 )
             },
+            component_quantizations={"unconditional_transformer": "fp8"},
         )
 
         component_args = _server_args_for_transformer_component(
@@ -543,6 +545,7 @@ class TestIdeogram4(unittest.TestCase):
             "/ckpt/diffusion_models/ideogram4_unconditional_nvfp4_mixed.safetensors",
         )
         self.assertIsNone(component_args.nunchaku_config)
+        self.assertEqual(component_args.quantization, "fp8")
 
     def test_ideogram_nvfp4_unconditional_transformer_path_uses_sibling_file(self):
         self.assertEqual(

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

@@ -28,6 +28,7 @@ class TestImageEncoderQuantizationAdmission(unittest.TestCase):
                 native_only_components=(),
             ),
             component_weights_paths={},
+            component_quantizations={},
             encoder_parallel="replicate",
             resolve_component_attention_backend=lambda _name: (None, None),
         )

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

@@ -494,6 +494,9 @@ class TestServerArgsPathExpansion(unittest.TestCase):
                 "--component-weights-paths.text_encoder",
                 "owner/repo/text_encoder.safetensors",
                 "--image-encoder-weights-path=/custom/image_encoder.safetensors",
+                "--component-quantizations.text_encoder",
+                "kitchen_int8",
+                "--transformer-quantization=fp8",
                 "--component-attention-backends.transformer",
                 "fa3",
             ]
@@ -536,6 +539,10 @@ class TestServerArgsPathExpansion(unittest.TestCase):
             {"transformer": "fa"},
             server_args.component_attention_backends,
         )
+        self.assertEqual(
+            {"text_encoder": "kitchen_int8", "transformer": "fp8"},
+            server_args.component_quantizations,
+        )
 
     def test_serve_cli_defaults_warmup_on(self):
         from sglang.multimodal_gen.runtime.entrypoints.cli.serve import (

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

@@ -334,6 +334,23 @@ class TestTextEncoderQuantization(unittest.TestCase):
         )
         self.assertIs(model_config.quant_config, self.serialized)
 
+    def test_explicit_online_quantization_configures_native_encoder(self):
+        model_config = SimpleNamespace(quant_config=None)
+        self.get_quant_config.return_value = None
+
+        _configure_encoder_quantization(
+            model_config,
+            TextEncoder,
+            {},
+            "/model/text_encoder",
+            "/model/text_encoder",
+            "text_encoder",
+            explicit_quantization="kitchen_int8",
+        )
+
+        self.assertIsInstance(model_config.quant_config, KitchenInt8Config)
+        self.assertFalse(model_config.quant_config.is_checkpoint_int8_serialized)
+
     def test_weight_file_metadata_configures_native_encoder(self):
         model_config = SimpleNamespace(quant_config=None)
         self.get_quant_config.return_value = None

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

@@ -43,6 +43,7 @@ class _FakeServerArgs:
         self.revision = "test-revision"
         self.trust_remote_code = True
         self.layerwise_components = set()
+        self.component_quantizations = {}
 
     def resolve_component_attention_backend(self, _component_name):
         return None, None