diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index 6ce48c555..1a9667433 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -139,7 +139,7 @@ or a native pipeline's registered module name. For a native text encoder: - `--component-paths.text_encoder {MODEL}` replaces the text-encoder checkpoint; `--text-encoder-path {MODEL}` is its shorter alias -- Quantization metadata is auto-detected from that checkpoint. Each native encoder must explicitly support the serialized format; this is not blanket quantization support for every component, and unsupported combinations fail before weight loading. +- 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, while quantized materialization is capability-based. Native auxiliary loaders diff --git a/docs/docs/sglang-diffusion/quantization.mdx b/docs/docs/sglang-diffusion/quantization.mdx index 2a94b8288..b75f1dd96 100644 --- a/docs/docs/sglang-diffusion/quantization.mdx +++ b/docs/docs/sglang-diffusion/quantization.mdx @@ -70,7 +70,7 @@ paths: | Component path | Quantized checkpoint behavior | | --- | --- | | `transformer`, `transformer_2`, `unconditional_transformer`, `audio_dit`, `video_dit` | Uses the SGLang transformer quantization adapters documented below. | -| `text_encoder*`, `image_encoder*` | Requires the native encoder class to declare support for the detected format. MiniMax-H3 FP8 and model-managed integrations such as Ideogram are supported; unknown combinations fail closed. | +| `text_encoder*`, `image_encoder*` | Passes detected metadata to the registered native encoder. Loading continues only when that implementation constructs compatible quantized layers; unknown or incompatible combinations fail closed. | | `vae`, `video_vae`, `audio_vae` | A standard top-level Diffusers `quantization_config` is delegated to `AutoModel.from_pretrained`. Native-only VAEs and nested/compression metadata fail closed. | | 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. | @@ -402,27 +402,6 @@ sglang generate \ Each pattern is matched against the full layer prefix (e.g. `layers.0.attention.to_q`). A layer is skipped and left unquantized if its prefix contains any of the given patterns. -## MiniMax-H3 Text Encoder FP8 - -MiniMax-H3 can load a serialized FP8 checkpoint for the language linear layers -in its native Qwen3-VL text encoder independently of the DiT. Embeddings, -normalization layers, and the Qwen vision tower remain in BF16. - -```bash Command -sglang serve \ - --model-path MiniMaxAI/MiniMax-H3 \ - --model-variant fl2va \ - --component-paths.text_encoder Qwen/Qwen3-VL-32B-Instruct-FP8 \ - --num-gpus 4 \ - --port 30010 -``` - -`--text-encoder-path` is accepted as a shorter alias. No quantization flag is -needed: SGLang detects the checkpoint metadata and only enables formats that -the native encoder explicitly supports. Text-encoder FP8 is approximate, is -not enabled by default, and is rejected by MiniMax-H3's strict -`quality="high"` deployment contract. - ## Transformers Component BnB4 Model components that already have a native Transformers loading path can load 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 3a3d07d6b..80b5ced24 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 @@ -129,27 +129,6 @@ def _configure_encoder_quantization( f"got {model_cls.__name__}" ) - capability = model_cls.checkpoint_quantization_capability - if capability is None: - raise ComponentCheckpointUnsupportedError( - f"{model_cls.__name__} does not support quantized checkpoints for " - f"{component_name!r}: no checkpoint quantization capability is declared" - ) - if capability.backend != "diffusion": - raise ComponentCheckpointUnsupportedError( - f"{model_cls.__name__} declares the {capability.backend!r} checkpoint " - f"quantization backend for {component_name!r}, but the native encoder " - "loader currently supports only the 'diffusion' backend" - ) - - quant_method = quant_config.get_name() - if quant_method not in capability.methods: - raise ComponentCheckpointUnsupportedError( - f"{model_cls.__name__} does not support {component_name!r} checkpoints " - f"quantized with {quant_method!r}; supported methods for the " - f"{capability.backend!r} backend: {sorted(capability.methods)}" - ) - def _resolve_and_configure_encoder_quantization( model_config: EncoderConfig, @@ -244,6 +223,23 @@ def _process_quantized_encoder_weights( return processed_layers +def _require_quantized_encoder_layers( + model: nn.Module, + component_name: str, +) -> None: + if any( + isinstance(module, LinearBase) + and module.quant_method is not None + and not isinstance(module.quant_method, UnquantizedLinearMethod) + for module in model.modules() + ): + return + raise ComponentCheckpointUnsupportedError( + f"The native {type(model).__name__} implementation does not construct " + f"quantized linear layers for {component_name!r}" + ) + + def _checkpoint_bytes(model_path: str) -> int: """On-disk size of a checkpoint, readable before any weight of it is.""" total = 0 @@ -624,6 +620,9 @@ class TextEncoderLoader(ComponentLoader): ) model.bind_encoder_tp_group(encoder_tp_group) + if quant_config is not None: + _require_quantized_encoder_layers(model, component_name) + if component_starts_on_cpu and ( current_platform.is_mps() or _keep_this_checkpoint_mapped(model_path) ): diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/base.py b/python/sglang/multimodal_gen/runtime/models/encoders/base.py index 1053e24bc..f58ee332a 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/base.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/base.py @@ -2,8 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 from abc import ABC, abstractmethod -from dataclasses import dataclass, field -from typing import Literal +from dataclasses import field import torch from torch import nn @@ -155,19 +154,10 @@ def finalize_encoder_folding( config.parallel_folding_mode = None -@dataclass(frozen=True) -class CheckpointQuantizationCapability: - """Quantized-checkpoint contract implemented by a native encoder.""" - - backend: Literal["diffusion", "srt"] - methods: frozenset[str] - - class EncoderTensorParallelMixin: """Keep an encoder on the TP group that was used to build its shards.""" _encoder_tp_group: GroupCoordinator | None = None - checkpoint_quantization_capability: CheckpointQuantizationCapability | None = None # Some encoders own checkpoint quantization end to end because their weight # states or sharding contract cannot use the generic loader lifecycle. manages_checkpoint_quantization = False @@ -192,9 +182,6 @@ class TextEncoder( # Qwen2_5_VLCausalLMOutputWithPast). Off by default so a new encoder is # replicated rather than silently broken; flip it once dp is verified there. supports_dp_encode = False - # Quantized checkpoints are opt-in because an encoder must construct - # quantized linears and load the checkpoint's auxiliary scale parameters. - supported_checkpoint_quantization_methods: frozenset[str] = frozenset() # Some encoders own checkpoint quantization end to end because their weight # states or sharding contract cannot use the generic loader lifecycle. manages_checkpoint_quantization = False diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py index 8f0b71724..2dadd25fa 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py @@ -17,10 +17,7 @@ from sglang.multimodal_gen.configs.models.encoders.minimax_h3_qwen3vl import ( ) from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader -from sglang.multimodal_gen.runtime.models.encoders.base import ( - CheckpointQuantizationCapability, - TextEncoder, -) +from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder from sglang.multimodal_gen.runtime.models.encoders.qwen3vl import Qwen3VLModel MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120 @@ -49,10 +46,6 @@ class MiniMaxH3Qwen3VLEncoder(TextEncoder): layer_names = [*TextEncoder.layer_names, "model.visual.blocks"] supports_dp_encode = True - checkpoint_quantization_capability = CheckpointQuantizationCapability( - backend="diffusion", - methods=frozenset({"fp8"}), - ) @staticmethod def should_materialize_checkpoint_weight(name: str) -> bool: 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 b5746ad7b..cd1d9c6df 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 @@ -52,14 +52,6 @@ class TestImageEncoderQuantizationAdmission(unittest.TestCase): "/model/image_encoder", self.server_args, "image_encoder", "transformers" ) - def test_quantized_clip_checkpoint_is_not_silently_enabled(self): - config = self._component_config("CLIPVisionModelWithProjection", quantized=True) - with self._config_patch(config), self.assertRaisesRegex( - ComponentCheckpointUnsupportedError, - "CLIPVisionModel.*image_encoder.*no checkpoint quantization capability", - ): - self.loader.load_customized("/model/image_encoder", self.server_args) - def test_unknown_quantized_architecture_does_not_fall_back(self): config = self._component_config("UnknownVisionModel", quantized=True) with self._config_patch(config), self.assertRaises( diff --git a/python/sglang/multimodal_gen/test/unit/test_qwen3_encoder.py b/python/sglang/multimodal_gen/test/unit/test_qwen3_encoder.py index 0cc34e9c2..8042b67a2 100644 --- a/python/sglang/multimodal_gen/test/unit/test_qwen3_encoder.py +++ b/python/sglang/multimodal_gen/test/unit/test_qwen3_encoder.py @@ -98,3 +98,31 @@ def test_default_position_ids_batch_shape(): assert torch.equal(layer.position_ids[0], torch.arange(4)) assert torch.equal(layer.position_ids[1], torch.arange(4)) assert layer.attention_lengths == (4, 4) + + +def test_fp8_qkv_scale_uses_the_packed_parameter_loader(): + model = Qwen3ForCausalLM.__new__(Qwen3ForCausalLM) + torch.nn.Module.__init__(model) + layer = torch.nn.Module() + layer.self_attn = torch.nn.Module() + layer.self_attn.qkv_proj = torch.nn.Module() + scale = torch.nn.Parameter(torch.zeros(3, 1), requires_grad=False) + + def load_scale(param, loaded_scale, shard_id): + param.data[{"q": 0, "k": 1, "v": 2}[shard_id]].copy_(loaded_scale) + + scale.weight_loader = load_scale + layer.self_attn.qkv_proj.register_parameter("weight_scale_inv", scale) + model.layers = torch.nn.ModuleList([layer]) + model.config = SimpleNamespace( + arch_config=SimpleNamespace( + stacked_params_mapping=[(".qkv_proj", ".q_proj", "q")] + ) + ) + + loaded = model.load_weights( + [("model.layers.0.self_attn.q_proj.weight_scale_inv", torch.tensor([2.0]))] + ) + + assert loaded == {"layers.0.self_attn.qkv_proj.weight_scale_inv"} + torch.testing.assert_close(scale[:, 0], torch.tensor([2.0, 0.0, 0.0])) 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 1bd7245c5..bd7380ea4 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 @@ -16,12 +16,10 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader TextEncoderLoader, _configure_encoder_quantization, _process_quantized_encoder_weights, + _require_quantized_encoder_layers, _resolve_and_configure_encoder_quantization, ) -from sglang.multimodal_gen.runtime.models.encoders.base import ( - CheckpointQuantizationCapability, - TextEncoder, -) +from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder from sglang.multimodal_gen.runtime.models.encoders.minimax_h3_qwen3vl import ( MiniMaxH3Qwen3VLEncoder, ) @@ -182,25 +180,25 @@ class TestTextEncoderQuantization(unittest.TestCase): self.addCleanup(self.quant_config_patcher.stop) self.serialized = serialized - def test_serialized_fp8_checkpoint_configures_h3_encoder(self): + def test_serialized_checkpoint_configures_native_encoder(self): model_config = SimpleNamespace(quant_config=None) _configure_encoder_quantization( model_config, - MiniMaxH3Qwen3VLEncoder, + TextEncoder, {}, "/model/text_encoder", "text_encoder", ) self.assertIs(model_config.quant_config, self.serialized) - def test_encoder_class_must_opt_in(self): + def test_encoder_must_use_native_loader(self): model_config = SimpleNamespace(quant_config=None) with self.assertRaisesRegex( - ComponentCheckpointUnsupportedError, "does not support" + ComponentCheckpointUnsupportedError, "requires an in-tree native encoder" ): _configure_encoder_quantization( model_config, - TextEncoder, + nn.Module, {}, "/model/text_encoder", "text_encoder", @@ -269,28 +267,6 @@ class TestTextEncoderQuantization(unittest.TestCase): "text_encoder", ) - def test_srt_backend_is_not_admitted_without_an_adapter(self): - model_config = SimpleNamespace(quant_config=None) - capability = CheckpointQuantizationCapability( - backend="srt", - methods=frozenset({"fp8"}), - ) - with mock.patch.object( - MiniMaxH3Qwen3VLEncoder, - "checkpoint_quantization_capability", - capability, - ), self.assertRaisesRegex( - ComponentCheckpointUnsupportedError, - "'srt'.*only the 'diffusion' backend", - ): - _configure_encoder_quantization( - model_config, - MiniMaxH3Qwen3VLEncoder, - {}, - "/model/text_encoder", - "text_encoder", - ) - def test_model_managed_quantization_bypasses_generic_lifecycle(self): model_config = SimpleNamespace(quant_config=None) with mock.patch.object( @@ -341,6 +317,13 @@ class _QuantizedEncoder(nn.Module): class TestQuantizedTextEncoderPostprocess(unittest.TestCase): + def test_rejects_native_encoder_without_quantized_layers(self): + with self.assertRaisesRegex( + ComponentCheckpointUnsupportedError, + "does not construct quantized linear layers", + ): + _require_quantized_encoder_layers(nn.Linear(2, 2), "text_encoder") + def test_processes_quantized_layers_without_moving_the_model(self): quant_method = _RecordingQuantMethod() model = _QuantizedEncoder(quant_method)