diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py index 84d9171bc..ac3360d81 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py +++ b/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.. 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", diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py index a0e53672d..d83a102ea 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py +++ b/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( diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py index fe5d72a99..0e269acf4 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py +++ b/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( diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py index 46ebed128..ed74ac686 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py @@ -100,15 +100,24 @@ 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) - component_server_args.transformer_weights_path = component_weights_path - component_server_args.nunchaku_config = None - logger.info( - "Using transformer_weights_path override for %s: %s", - component_name, - component_weights_path, - ) + if component_weights_path is not None: + component_server_args.transformer_weights_path = component_weights_path + component_server_args.nunchaku_config = None + logger.info( + "Using transformer_weights_path override for %s: %s", + 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", diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index 29fc41f1c..0fff4346b 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/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") diff --git a/python/sglang/multimodal_gen/test/unit/test_ideogram4.py b/python/sglang/multimodal_gen/test/unit/test_ideogram4.py index 2bfe875e6..031af5f86 100644 --- a/python/sglang/multimodal_gen/test/unit/test_ideogram4.py +++ b/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( diff --git a/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py b/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py index e092a6908..520bc75df 100644 --- a/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py +++ b/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), ) diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index f40b9fa70..de211cfbd 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/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 ( diff --git a/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py b/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py index cc01cd334..27bfd94ba 100644 --- a/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py +++ b/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 diff --git a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py b/python/sglang/multimodal_gen/test/unit/test_vae_loader.py index eadacd8f3..6ec6a11b9 100644 --- a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py +++ b/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