[diffusion] feat: support component-scoped quantization overrides (#36035)

docs/docs/sglang-diffusion/api/cli.mdx

@@ -131,6 +131,8 @@ pipeline's registered module name:
 | --- | --- | --- | --- |
 | Replace a component | `--component-paths.<component> {MODEL}` | `--<component>-path {MODEL}` | Load the replacement component's configuration and weights |
 | Replace only its weights | `--component-weights-paths.<component> {WEIGHTS}` | `--<component>-weights-path {WEIGHTS}` | Retain the base component configuration and replace its weights |
+| Quantize an unquantized component online | `--component-quantizations.<component> {METHOD}` | `--<component>-quantization {METHOD}` | Apply a method supported by that component's native loader |
+| Keep selected component layers unquantized | `--component-quantization-ignored-layers.<component> {PATTERN...}` | None | Pass component-local ignored-layer patterns to its online quantizer |
 
 For example, pair a replacement text-encoder configuration with a separate
 single-file checkpoint as follows:
@@ -152,6 +154,7 @@ unquantized during online quantization.
 For a native text encoder:
 
 - `--component-paths.text_encoder {MODEL}` replaces the text-encoder checkpoint; `--text-encoder-path {MODEL}` is its shorter alias
+- `--component-quantizations.text_encoder {METHOD}` applies supported online quantization to an unquantized native encoder; pair it with `--component-quantization-ignored-layers.text_encoder {PATTERN...}` when selected layers must remain unquantized
 - Quantization metadata is auto-detected from that checkpoint. The native loader accepts compatible serialized formats without a model-name allowlist and rejects implementations that do not construct the required quantized layers.
 
 The same contract applies to every weighted component: path routing is generic,

docs/docs/sglang-diffusion/quantization.mdx

@@ -23,6 +23,8 @@ Use these paths:
   (local or Hub)
 - `--quantization`: override the quantization method used by the transformer loader
 - `--quantization-ignored-layers`: transformer layer name patterns to keep unquantized during online quantization (e.g. `attention.to_`)
+- `--component-quantizations.<component>`: explicitly apply supported online quantization to an unquantized component
+- `--component-quantization-ignored-layers.<component>`: component-local layer patterns to keep unquantized during that online quantization
 - `--component-paths.text_encoder`: replace a native text encoder with a checkpoint whose `quantization_config` is auto-detected
 - `--text-encoder-path`: shorter alias for `--component-paths.text_encoder`
 - `--kv-cache-quant`: compress completed causal KV-cache chunks for supported realtime models
@@ -78,13 +80,15 @@ paths:
 | Library-managed Transformers or Diffusers components | Delegates to the upstream `from_pretrained` path and inherits its format support and validation behavior. The local PE model uses this path; compatible formats remain model-specific. |
 | Native auxiliary components that load raw state dicts | Quantized checkpoints are rejected before model construction until that component has a quantized materialization implementation. This includes connectors, duration heads, bridges, diffusion decoders, sound tokenizers, spatial upsamplers, and vocoders. |
 
-`--quantization` is the explicit override for the transformer loader; it is not
+`--quantization` is the explicit override for the primary transformer loader;
-the boundary of component quantization support. Other pre-quantized component
+`--component-quantizations.<component>` expresses the same intent for one
-repositories select their format through their own metadata and the capability
+supported component. Pair the latter with
-of the selected loader. A generic string override without a matching
+`--component-quantization-ignored-layers.<component>` to keep matching layers
-materialization backend would advertise support that the component does not
+unquantized. Pre-quantized component repositories instead select their format
-have, while a quantized weight file without matching config metadata cannot be
+through their own metadata and the capability of the selected loader. A generic
-identified or restored generically.
+string override without a matching materialization backend would advertise
+support that the component does not have, while a quantized weight file without
+matching config metadata cannot be identified or restored generically.
 
 ## Quant Families
 

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

@@ -248,6 +248,7 @@ def _configure_encoder_quantization(
     component_weights_path: str,
     component_name: str,
     explicit_quantization: str | None = None,
+    ignored_layers: list[str] | None = None,
 ) -> None:
     if getattr(model_cls, "manages_checkpoint_quantization", False):
         if explicit_quantization is not None:
@@ -292,7 +293,9 @@ def _configure_encoder_quantization(
             get_quantization_config,
         )
 
-        model_config.quant_config = get_quantization_config(explicit_quantization)()
+        model_config.quant_config = get_quantization_config(explicit_quantization)(
+            ignored_layers=ignored_layers
+        )
         quant_config = model_config.quant_config
     if quant_config is None:
         return
@@ -311,6 +314,7 @@ def _resolve_and_configure_encoder_quantization(
     component_weights_path: str,
     component_name: str,
     explicit_quantization: str | None = None,
+    ignored_layers: list[str] | None = None,
 ) -> type[nn.Module]:
     architectures = getattr(model_config, "architectures", [])
     try:
@@ -351,6 +355,7 @@ def _resolve_and_configure_encoder_quantization(
         component_weights_path,
         component_name,
         explicit_quantization,
+        ignored_layers,
     )
     return model_cls
 
@@ -725,6 +730,7 @@ class TextEncoderLoader(ComponentLoader):
             component_weights_path,
             component_name,
             server_args.component_quantizations.get(component_name),
+            server_args.component_quantization_ignored_layers.get(component_name),
         )
         if issubclass(model_cls, EncoderTensorParallelMixin):
             model_cls.configure_component_paths(

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

@@ -101,7 +101,14 @@ def _server_args_for_transformer_component(
     """Mask global quantized override flags for secondary transformer components."""
     component_weights_path = server_args.component_weights_paths.get(component_name)
     component_quantization = server_args.component_quantizations.get(component_name)
-    if component_weights_path is not None or component_quantization is not None:
+    component_ignored_layers = server_args.component_quantization_ignored_layers.get(
+        component_name
+    )
+    if (
+        component_weights_path is not None
+        or component_quantization is not None
+        or component_ignored_layers 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
@@ -118,6 +125,8 @@ def _server_args_for_transformer_component(
                 component_quantization,
                 component_name,
             )
+        if component_ignored_layers is not None:
+            component_server_args.quantization_ignored_layers = component_ignored_layers
         return component_server_args
 
     if component_name not in ("transformer_2", "unconditional_transformer"):

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

@@ -309,6 +309,10 @@ class ServerArgs(DisaggServerArgsMixin):
     # 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)
+    # Component-local layer name patterns to skip during online quantization.
+    component_quantization_ignored_layers: dict[str, list[str]] = field(
+        default_factory=dict
+    )
     # Optional LTX-2.5 decoder is large enough to load only when requested.
     load_diffusion_decoder: bool = False
 
@@ -1741,6 +1745,31 @@ class ServerArgs(DisaggServerArgsMixin):
                 )
             normalized_quantizations[component] = quantization
         self.component_quantizations = normalized_quantizations
+        normalized_ignored_layers: dict[str, list[str]] = {}
+        for component, layers in self.component_quantization_ignored_layers.items():
+            component = str(component).strip().replace("-", "_")
+            if isinstance(layers, str):
+                layers = [layers]
+            if (
+                not component
+                or not isinstance(layers, (list, tuple))
+                or not all(isinstance(layer, str) and layer.strip() for layer in layers)
+            ):
+                raise ValueError(
+                    "Component quantization ignored layers require a component "
+                    "and non-empty layer patterns"
+                )
+            normalized_ignored_layers[component] = [layer.strip() for layer in layers]
+        missing_quantization = set(normalized_ignored_layers) - set(
+            self.component_quantizations
+        )
+        if missing_quantization:
+            raise ValueError(
+                "Component quantization ignored layers require a matching "
+                "quantization override for: "
+                f"{', '.join(sorted(missing_quantization))}"
+            )
+        self.component_quantization_ignored_layers = normalized_ignored_layers
 
         # Convert string disagg_role to enum (from CLI/config)
         if isinstance(self.disagg_role, str):
@@ -2867,6 +2896,46 @@ class ServerArgs(DisaggServerArgsMixin):
             alias_suffix="-quantization",
         )
 
+    @staticmethod
+    def _extract_component_quantization_ignored_layers(
+        unknown_args: list[str],
+    ) -> tuple[dict[str, list[str]], list[str]]:
+        ignored_layers: dict[str, list[str]] = {}
+        remaining: list[str] = []
+        i = 0
+        prefixes = (
+            "--component-quantization-ignored-layers.",
+            "--component_quantization_ignored_layers.",
+        )
+        while i < len(unknown_args):
+            arg = unknown_args[i]
+            key_part = arg.split("=", 1)[0] if "=" in arg else arg
+            prefix = next(
+                (candidate for candidate in prefixes if key_part.startswith(candidate)),
+                None,
+            )
+            if prefix is None:
+                remaining.append(arg)
+                i += 1
+                continue
+
+            component = key_part[len(prefix) :].replace("-", "_")
+            if "=" in arg:
+                values = [arg.split("=", 1)[1]]
+            else:
+                values = []
+                while i + 1 < len(unknown_args) and not unknown_args[i + 1].startswith(
+                    "-"
+                ):
+                    i += 1
+                    values.append(unknown_args[i])
+            if component and values:
+                ignored_layers[component] = values
+            else:
+                remaining.append(arg)
+            i += 1
+        return ignored_layers, remaining
+
     @staticmethod
     def _extract_component_attention_backends(
         unknown_args: list[str],
@@ -2917,6 +2986,9 @@ class ServerArgs(DisaggServerArgsMixin):
         dynamic_quantizations, remaining = cls._extract_component_quantizations(
             unknown_args
         )
+        dynamic_ignored_layers, remaining = (
+            cls._extract_component_quantization_ignored_layers(remaining)
+        )
         # Extract the more specific weights suffix before the generic path alias.
         dynamic_weights_paths, remaining = cls._extract_component_weights_paths(
             remaining
@@ -2957,6 +3029,13 @@ class ServerArgs(DisaggServerArgsMixin):
             existing.update(dynamic_quantizations)
             provided_args["component_quantizations"] = existing
             explicit_arg_names.add("component_quantizations")
+        if dynamic_ignored_layers:
+            existing = dict(
+                provided_args.get("component_quantization_ignored_layers") or {}
+            )
+            existing.update(dynamic_ignored_layers)
+            provided_args["component_quantization_ignored_layers"] = existing
+            explicit_arg_names.add("component_quantization_ignored_layers")
         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

@@ -511,6 +511,7 @@ class TestIdeogram4(unittest.TestCase):
             nunchaku_config={"enabled": True},
             component_weights_paths={},
             component_quantizations={},
+            component_quantization_ignored_layers={},
         )
         component_args = _server_args_for_transformer_component(
             server_args, "unconditional_transformer"
@@ -532,6 +533,9 @@ class TestIdeogram4(unittest.TestCase):
                 )
             },
             component_quantizations={"unconditional_transformer": "fp8"},
+            component_quantization_ignored_layers={
+                "unconditional_transformer": ["lm_head"]
+            },
         )
 
         component_args = _server_args_for_transformer_component(
@@ -546,6 +550,7 @@ class TestIdeogram4(unittest.TestCase):
         )
         self.assertIsNone(component_args.nunchaku_config)
         self.assertEqual(component_args.quantization, "fp8")
+        self.assertEqual(component_args.quantization_ignored_layers, ["lm_head"])
 
     def test_ideogram_nvfp4_unconditional_transformer_path_uses_sibling_file(self):
         self.assertEqual(

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

@@ -496,6 +496,9 @@ class TestServerArgsPathExpansion(unittest.TestCase):
                 "--image-encoder-weights-path=/custom/image_encoder.safetensors",
                 "--component-quantizations.text_encoder",
                 "kitchen_int8",
+                "--component-quantization-ignored-layers.text_encoder",
+                "model.layers.0",
+                "lm_head",
                 "--transformer-quantization=fp8",
                 "--component-attention-backends.transformer",
                 "fa3",
@@ -543,6 +546,10 @@ class TestServerArgsPathExpansion(unittest.TestCase):
             {"text_encoder": "kitchen_int8", "transformer": "fp8"},
             server_args.component_quantizations,
         )
+        self.assertEqual(
+            {"text_encoder": ["model.layers.0", "lm_head"]},
+            server_args.component_quantization_ignored_layers,
+        )
 
     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

@@ -363,10 +363,12 @@ class TestTextEncoderQuantization(unittest.TestCase):
             "/model/text_encoder",
             "text_encoder",
             explicit_quantization="kitchen_int8",
+            ignored_layers=["lm_head"],
         )
 
         self.assertIsInstance(model_config.quant_config, KitchenInt8Config)
         self.assertFalse(model_config.quant_config.is_checkpoint_int8_serialized)
+        self.assertEqual(model_config.quant_config.ignored_layers, ["lm_head"])
 
     def test_weight_file_metadata_configures_native_encoder(self):
         model_config = SimpleNamespace(quant_config=None)