[diffusion] feat: support loading serialized fp8 clip image encoders (#36056)

python/sglang/multimodal_gen/configs/models/encoders/base.py

@@ -1,8 +1,10 @@
 # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
 
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
 from dataclasses import dataclass, field
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 import torch
 
@@ -10,6 +12,11 @@ from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig
 from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
 from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
 
+if TYPE_CHECKING:
+    from sglang.srt.layers.quantization.base_config import (
+        QuantizationConfig as SRTQuantizationConfig,
+    )
+
 
 @dataclass
 class EncoderArchConfig(ArchConfig):
@@ -73,7 +80,7 @@ class EncoderConfig(ModelConfig):
     arch_config: ArchConfig = field(default_factory=EncoderArchConfig)
 
     prefix: str = ""
-    quant_config: QuantizationConfig | None = None
+    quant_config: QuantizationConfig | SRTQuantizationConfig | None = None
     lora_config: Any | None = None
 
     # Parallel folding: during the encoding stage the whole DiT replica is idle,

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

@@ -3,7 +3,6 @@ import glob
 import os
 import re
 from collections.abc import Callable, Generator, Iterable
-from itertools import chain
 from typing import cast
 
 import torch
@@ -96,9 +95,14 @@ from sglang.multimodal_gen.runtime.weights.source import (
 from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
 from sglang.srt.environ import envs
 from sglang.srt.layers.linear import LinearBase as SrtLinearBase
+from sglang.srt.layers.quantization.fp8 import Fp8Config as SrtFp8Config
 from sglang.srt.layers.quantization.unquant import (
     UnquantizedLinearMethod as SrtUnquantizedLinearMethod,
 )
+from sglang.srt.model_loader.checkpoint_quantization import (
+    resolve_checkpoint_quant_spec,
+)
+from sglang.srt.model_loader.post_load import stage_module_for_post_load
 
 logger = init_logger(__name__)
 
@@ -129,12 +133,42 @@ def _delegate_standard_bnb4_to_transformers(
         )
 
 
+def _get_srt_encoder_quant_config(
+    component_config: dict,
+    model_cls: type[EncoderTensorParallelMixin],
+) -> SrtFp8Config | None:
+    quant_spec = resolve_checkpoint_quant_spec(component_config)
+    if quant_spec is None:
+        return None
+    if quant_spec.declared_method != "fp8":
+        raise ComponentCheckpointUnsupportedError(
+            "The SRT encoder checkpoint adapter supports only serialized 'fp8', "
+            f"got {quant_spec.declared_method!r}"
+        )
+
+    config = dict(quant_spec.config)
+    config["packed_modules_mapping"] = model_cls.packed_modules_mapping
+    return SrtFp8Config.from_config(config)
+
+
 def _get_encoder_quant_config(
     component_config: dict,
     component_model_path: str,
     component_weights_path: str,
     model_cls: type[nn.Module] | None = None,
 ):
+    if (
+        model_cls is not None
+        and issubclass(model_cls, EncoderTensorParallelMixin)
+        and model_cls.checkpoint_quantization_backend == "srt"
+    ):
+        srt_quant_config = _get_srt_encoder_quant_config(
+            component_config,
+            model_cls,
+        )
+        if srt_quant_config is not None:
+            return srt_quant_config
+
     quant_config = get_quant_config(component_config, component_model_path)
     name_mapper = None
     if model_cls is not None:
@@ -290,28 +324,6 @@ def _resolve_and_configure_encoder_quantization(
     return model_cls
 
 
-def _module_tensor_device(module: nn.Module) -> torch.device | None:
-    """Return the device of a module's own tensors.
-
-    Quantized linear layers are expected to keep their parameters and buffers
-    together.  Failing explicitly is safer than staging only part of a layer.
-    """
-
-    devices = {
-        tensor.device
-        for tensor in chain(
-            module.parameters(recurse=False),
-            module.buffers(recurse=False),
-        )
-    }
-    if len(devices) > 1:
-        raise ValueError(
-            f"Cannot stage {type(module).__name__} with tensors on multiple "
-            f"devices: {sorted(map(str, devices))}"
-        )
-    return next(iter(devices), None)
-
-
 def _process_quantized_encoder_weights(
     model: nn.Module,
     process_device: torch.device | None,
@@ -327,23 +339,12 @@ def _process_quantized_encoder_weights(
             (UnquantizedLinearMethod, SrtUnquantizedLinearMethod),
         ):
             continue
-
-        origin_device = _module_tensor_device(module)
-        should_stage = (
-            process_device is not None
-            and origin_device is not None
-            and origin_device != process_device
-        )
-        if should_stage:
-            module.to(process_device)
-        try:
+        if process_device is None:
             quant_method.process_weights_after_loading(module)
-            processed_layers += 1
-        finally:
-            # Post-load methods may replace parameters or register buffers. Move
-            # the complete layer back so component residency remains authoritative.
-            if should_stage:
-                module.to(origin_device)
+        else:
+            with stage_module_for_post_load(module, process_device):
+                quant_method.process_weights_after_loading(module)
+        processed_layers += 1
     if processed_layers == 0:
         raise ValueError(
             f"The {component_name!r} checkpoint declares quantization, but the "

python/sglang/multimodal_gen/runtime/models/encoders/base.py

@@ -158,6 +158,8 @@ class EncoderTensorParallelMixin:
     """Keep an encoder on the TP group that was used to build its shards."""
 
     _encoder_tp_group: GroupCoordinator | None = None
+    checkpoint_quantization_backend = "diffusion"
+    packed_modules_mapping: dict[str, list[str]] = {}
     # 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/clip.py

@@ -17,12 +17,12 @@ from sglang.multimodal_gen.configs.models.encoders import (
     CLIPTextConfig,
     CLIPVisionConfig,
 )
-from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
 from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
 from sglang.multimodal_gen.runtime.models.encoders.base import ImageEncoder, TextEncoder
 from sglang.multimodal_gen.runtime.models.encoders.vision import (
     resolve_visual_encoder_outputs,
 )
+from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.models.clip import (
     CLIPEncoder,
     CLIPTextEmbeddings,
@@ -337,6 +337,9 @@ class CLIPVisionModel(ImageEncoder):
     config_class = CLIPVisionConfig
     main_input_name = "pixel_values"
     packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
+    # CLIPEncoder and its linears come from SRT, so their serialized checkpoint
+    # format must use the matching SRT quantization implementation.
+    checkpoint_quantization_backend = "srt"
 
     def __init__(self, config: CLIPVisionConfig) -> None:
         super().__init__(config)

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

@@ -11,6 +11,11 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader imp
 from sglang.multimodal_gen.runtime.loader.component_loaders.image_encoder_loader import (
     ImageEncoderLoader,
 )
+from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader import (
+    _configure_encoder_quantization,
+)
+from sglang.multimodal_gen.runtime.models.encoders.clip import CLIPVisionModel
+from sglang.srt.layers.quantization.fp8 import Fp8Config as SRTFp8Config
 
 
 class TestImageEncoderQuantizationAdmission(unittest.TestCase):
@@ -54,6 +59,23 @@ class TestImageEncoderQuantizationAdmission(unittest.TestCase):
             "/model/image_encoder", self.server_args, "image_encoder", "transformers"
         )
 
+    def test_clip_serialized_fp8_uses_srt_quantization(self):
+        encoder_config = CLIPVisionConfig()
+        _configure_encoder_quantization(
+            encoder_config,
+            CLIPVisionModel,
+            self._component_config("CLIPVisionModelWithProjection", quantized=True),
+            "/model/image_encoder",
+            "/model/image_encoder",
+            "image_encoder",
+        )
+
+        self.assertIsInstance(encoder_config.quant_config, SRTFp8Config)
+        self.assertEqual(
+            encoder_config.quant_config.packed_modules_mapping,
+            {"qkv_proj": ["q_proj", "k_proj", "v_proj"]},
+        )
+
     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_text_encoder_loader.py

@@ -37,6 +37,7 @@ from sglang.multimodal_gen.runtime.models.encoders.minimax_h3_qwen3vl import (
     MiniMaxH3Qwen3VLEncoder,
 )
 from sglang.multimodal_gen.runtime.models.encoders.qwen3vl import Qwen3VLTextModel
+from sglang.srt.layers.linear import LinearBase as SrtLinearBase
 
 
 class TestTextEncoderClassResolution(unittest.TestCase):
@@ -640,7 +641,27 @@ class _QuantizedEncoder(nn.Module):
         self.unquantized = nn.Linear(2, 2, bias=False)
 
 
+class _SRTQuantizedLinear(SrtLinearBase):
+    def __init__(self, quant_method):
+        nn.Module.__init__(self)
+        self.weight = nn.Parameter(torch.empty(2, 2), requires_grad=False)
+        self.quant_method = quant_method
+
+
 class TestQuantizedTextEncoderPostprocess(unittest.TestCase):
+    def test_processes_srt_quantized_linear(self):
+        quant_method = _RecordingQuantMethod()
+        model = _SRTQuantizedLinear(quant_method)
+
+        processed = _process_quantized_encoder_weights(
+            model,
+            torch.device("cpu"),
+            "image_encoder",
+        )
+
+        self.assertEqual(processed, 1)
+        self.assertEqual(quant_method.devices, [torch.device("cpu")])
+
     def test_rejects_native_encoder_without_quantized_layers(self):
         with self.assertRaisesRegex(
             ComponentCheckpointUnsupportedError,