[diffusion] feat: admit compatible quantized native encoders (#35962)

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)
             ):

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

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:

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(

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]))

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)