[diffusion] feat: support compact qwen3-vl conditioning for minimax h3 (#36076)

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

@@ -26,6 +26,8 @@ class MiniMaxH3Qwen3VLArchConfig(Qwen3VLArchConfig):
     head_dim: int = 128
     text_len: int = 262144
     hidden_state_skip_layer: int = 0
+    checkpoint_num_hidden_layers: int = MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER
+    conditioning_projection_path: str | None = None
 
 
 @dataclass
@@ -38,6 +40,7 @@ class MiniMaxH3Qwen3VLConfig(Qwen3VLConfig):
         """Select the in-tree extractor after loading the HF architecture."""
 
         arch = self.arch_config
+        arch.checkpoint_num_hidden_layers = int(arch.text_config.num_hidden_layers)
         arch.architectures = ["MiniMaxH3Qwen3VLEncoder"]
         arch.hidden_size = int(arch.text_config.hidden_size)
         arch.intermediate_size = int(arch.text_config.intermediate_size)

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

@@ -600,6 +600,11 @@ class TextEncoderLoader(ComponentLoader):
             component_weights_path,
             component_name,
         )
+        if issubclass(model_cls, EncoderTensorParallelMixin):
+            model_cls.configure_component_paths(
+                encoder_config,
+                server_args.component_paths,
+            )
         encoder_dp_group = get_encoder_data_parallel_group()
         prefer_dp = (
             server_args.batching_max_size > 1

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

@@ -162,6 +162,14 @@ class EncoderTensorParallelMixin:
     # states or sharding contract cannot use the generic loader lifecycle.
     manages_checkpoint_quantization = False
 
+    @classmethod
+    def configure_component_paths(
+        cls,
+        config: EncoderConfig,
+        component_paths: dict[str, str],
+    ) -> None:
+        """Apply optional runtime components before parallel layout is resolved."""
+
     def bind_encoder_tp_group(self, tp_group: GroupCoordinator) -> None:
         self._encoder_tp_group = tp_group
 

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

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Native, TP-foldable Qwen3-VL layer-50 encoder for MiniMax H3."""
+"""Native, TP-foldable Qwen3-VL conditioning encoder for MiniMax H3."""
 
 from __future__ import annotations
 
@@ -9,6 +9,9 @@ from typing import Any
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
+from safetensors import safe_open
+from safetensors.torch import load_file
 
 from sglang.multimodal_gen.configs.models.encoders.base import BaseEncoderOutput
 from sglang.multimodal_gen.configs.models.encoders.minimax_h3_qwen3vl import (
@@ -20,6 +23,10 @@ from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
 from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
 from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder
 from sglang.multimodal_gen.runtime.models.encoders.qwen3vl import Qwen3VLModel
+from sglang.multimodal_gen.runtime.weights.source import (
+    materialize_weight,
+    resolve_weight,
+)
 
 MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120
 _LAYER_WEIGHT_RE = re.compile(r"^model\.language_model\.layers\.(\d+)\.")
@@ -35,17 +42,155 @@ def _map_checkpoint_name(name: str) -> str:
     return _MAP_CHECKPOINT_NAME(name)[0]
 
 
-def _is_unconsumed_checkpoint_weight(name: str) -> bool:
+def _is_unconsumed_checkpoint_weight(name: str, selected_layer: int) -> bool:
-    """Weights intentionally absent from the layer-50 feature extractor."""
+    """Weights intentionally absent from the selected feature extractor."""
 
     if name == "lm_head.weight" or name.startswith("model.language_model.norm."):
         return True
     match = _LAYER_WEIGHT_RE.match(name)
-    return bool(match and int(match.group(1)) >= MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER)
+    return bool(match and int(match.group(1)) >= selected_layer)
+
+
+class _FrozenLinear(nn.Module):
+    def __init__(self, weight: torch.Tensor, bias: torch.Tensor | None) -> None:
+        super().__init__()
+        self.register_buffer("weight", weight)
+        self.register_buffer("bias", bias)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return F.linear(hidden_states, self.weight, self.bias)
+
+
+class MiniMaxH3ConditioningProjection(nn.Module):
+    """Apply the safe ClipProj conditioning format without a custom runtime."""
+
+    _REQUIRED_TENSORS = ("mean_in", "std_in", "mean_out", "std_out")
+
+    @staticmethod
+    def inspect(path: str) -> tuple[int, int, int]:
+        if not path.endswith(".safetensors"):
+            raise ValueError("H3 conditioning projections must use safetensors")
+        with safe_open(path, framework="pt", device="cpu") as handle:
+            metadata = handle.metadata() or {}
+            try:
+                tap = int(metadata["tap"])
+                input_dim = int(handle.get_slice("mean_in").get_shape()[0])
+                output_dim = int(handle.get_slice("mean_out").get_shape()[0])
+            except (IndexError, KeyError, TypeError, ValueError) as error:
+                raise ValueError(
+                    f"Invalid H3 conditioning projection metadata in {path!r}"
+                ) from error
+        return tap, input_dim, output_dim
+
+    def __init__(self, path: str) -> None:
+        super().__init__()
+        self.tap, self.input_dim, self.output_dim = self.inspect(path)
+        tensors = load_file(path, device="cpu")
+        missing = set(self._REQUIRED_TENSORS) - set(tensors)
+        if missing:
+            raise ValueError(
+                f"H3 conditioning projection is missing tensors: {sorted(missing)}"
+            )
+
+        for name in self._REQUIRED_TENSORS:
+            self.register_buffer(name, tensors.pop(name).float())
+        expected_shapes = {
+            "mean_in": (self.input_dim,),
+            "std_in": (self.input_dim,),
+            "mean_out": (self.output_dim,),
+            "std_out": (self.output_dim,),
+        }
+        for name, expected_shape in expected_shapes.items():
+            actual_shape = tuple(self.get_buffer(name).shape)
+            if actual_shape != expected_shape:
+                raise ValueError(
+                    f"H3 conditioning projection {name} has shape "
+                    f"{actual_shape}, expected {expected_shape}"
+                )
+        self.register_buffer(
+            "weight",
+            tensors.pop("W").float() if "W" in tensors else None,
+        )
+        self.register_buffer(
+            "sink_out",
+            tensors.pop("sink_out").float() if "sink_out" in tensors else None,
+        )
+
+        layer_indices = sorted(
+            {
+                int(name.split(".")[1])
+                for name in tensors
+                if re.fullmatch(r"mlp\.\d+\.weight", name)
+            }
+        )
+        layers: list[nn.Module] = []
+        layer_input_dim = self.input_dim
+        for layer_index in layer_indices:
+            weight_name = f"mlp.{layer_index}.weight"
+            bias_name = f"mlp.{layer_index}.bias"
+            weight = tensors.pop(weight_name)
+            bias = tensors.pop(bias_name, None)
+            if weight.ndim != 2 or int(weight.shape[1]) != layer_input_dim:
+                raise ValueError(
+                    f"H3 conditioning projection {weight_name} cannot follow "
+                    f"width {layer_input_dim}: got {tuple(weight.shape)}"
+                )
+            if bias is not None and tuple(bias.shape) != (int(weight.shape[0]),):
+                raise ValueError(
+                    f"H3 conditioning projection {bias_name} has shape "
+                    f"{tuple(bias.shape)}, expected ({int(weight.shape[0])},)"
+                )
+            layers.append(_FrozenLinear(weight, bias))
+            layer_input_dim = int(weight.shape[0])
+        if tensors:
+            raise ValueError(
+                "H3 conditioning projection contains unsupported tensors: "
+                f"{sorted(tensors)}"
+            )
+        if self.weight is None and not layers:
+            raise ValueError("H3 conditioning projection has neither W nor an MLP")
+        if layers and layer_input_dim != self.output_dim:
+            raise ValueError(
+                f"H3 conditioning projection MLP outputs width {layer_input_dim}, "
+                f"expected {self.output_dim}"
+            )
+        if self.weight is not None and tuple(self.weight.shape) != (
+            self.input_dim,
+            self.output_dim,
+        ):
+            raise ValueError(
+                "H3 conditioning projection W has shape "
+                f"{tuple(self.weight.shape)}, expected "
+                f"({self.input_dim}, {self.output_dim})"
+            )
+        self.layers = nn.ModuleList(layers)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if int(hidden_states.shape[-1]) != self.input_dim:
+            raise ValueError(
+                f"H3 conditioning projection expects width {self.input_dim}, "
+                f"got {int(hidden_states.shape[-1])}"
+            )
+        normalized = (hidden_states.float() - self.mean_in) / self.std_in
+        projected = normalized @ self.weight if self.weight is not None else None
+        if self.layers:
+            residual = normalized.to(self.layers[0].weight.dtype)
+            for index, layer in enumerate(self.layers):
+                residual = layer(residual)
+                if index + 1 < len(self.layers):
+                    residual = F.gelu(residual)
+            residual = residual.float()
+            projected = residual if projected is None else projected + residual
+        if projected is None:
+            raise RuntimeError("H3 conditioning projection produced no output")
+        output = projected * self.std_out + self.mean_out
+        if self.sink_out is not None and int(output.shape[-2]) > 0:
+            output[..., 0, :] = self.sink_out
+        return output
 
 
 class MiniMaxH3Qwen3VLEncoder(TextEncoder):
-    """Qwen3-VL-32B multimodal backbone ending at hidden_states[50].
+    """Qwen3-VL multimodal backbone producing MiniMax H3 conditioning.
 
     The component loader builds and loads this module under the encoder-folding
     TP group. A TP=1/SP=8 DiT deployment therefore shards the encoder over all
@@ -54,27 +199,73 @@ class MiniMaxH3Qwen3VLEncoder(TextEncoder):
 
     # The inherited text-layer list covers Qwen's language stack; reference
     # modes also execute the embedded visual tower.
-    layer_names = [*TextEncoder.layer_names, "model.visual.blocks"]
+    layer_names = [
+        *TextEncoder.layer_names,
+        "model.visual.blocks",
+        "conditioning_projection.layers",
+    ]
 
     supports_dp_encode = True
     param_names_mapping = _PARAM_NAMES_MAPPING
 
-    @staticmethod
+    @classmethod
-    def should_materialize_checkpoint_weight(name: str) -> bool:
+    def configure_component_paths(
+        cls,
+        config: MiniMaxH3Qwen3VLConfig,
+        component_paths: dict[str, str],
+    ) -> None:
+        arch = config.arch_config
+        source = component_paths.get("conditioning_projection")
+        if source is None:
+            if (
+                int(arch.hidden_size) != MINIMAX_H3_QWEN3VL_HIDDEN_DIM
+                or int(arch.checkpoint_num_hidden_layers)
+                < MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER
+            ):
+                raise ValueError(
+                    "MiniMax H3 Qwen3-VL encoders smaller than 32B require "
+                    "--component-paths.conditioning_projection"
+                )
+            return
+
+        projection_path = materialize_weight(resolve_weight(source))
+        tap, input_dim, output_dim = MiniMaxH3ConditioningProjection.inspect(
+            projection_path
+        )
+        if input_dim != int(arch.hidden_size):
+            raise ValueError(
+                f"H3 conditioning projection expects encoder width {input_dim}, "
+                f"but the selected text encoder has width {int(arch.hidden_size)}"
+            )
+        if output_dim != MINIMAX_H3_QWEN3VL_HIDDEN_DIM:
+            raise ValueError(
+                f"H3 conditioning projection must output width "
+                f"{MINIMAX_H3_QWEN3VL_HIDDEN_DIM}, got {output_dim}"
+            )
+        if tap <= 0 or tap > int(arch.checkpoint_num_hidden_layers):
+            raise ValueError(
+                f"H3 conditioning projection tap {tap} is outside the selected "
+                f"encoder's {int(arch.checkpoint_num_hidden_layers)} layers"
+            )
+        arch.conditioning_projection_path = projection_path
+        arch.num_hidden_layers = tap
+        arch.text_config.num_hidden_layers = tap
+
+    def should_materialize_checkpoint_weight(self, name: str) -> bool:
         name = _map_checkpoint_name(name)
         return (
             "rotary_emb.inv_freq" not in name
-            and not _is_unconsumed_checkpoint_weight(name)
+            and not _is_unconsumed_checkpoint_weight(name, self.selected_lm_layer)
         )
 
     def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None:
         super().__init__(config)
         arch = config.arch_config
-        selected_layer = MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER
+        selected_layer = int(arch.text_config.num_hidden_layers)
-        if int(arch.text_config.num_hidden_layers) != selected_layer:
+        if selected_layer <= 0 or int(arch.num_hidden_layers) != selected_layer:
             raise ValueError(
-                "MiniMax H3 Qwen3-VL config must be trimmed to "
+                "MiniMax H3 Qwen3-VL language-layer configuration is "
-                f"{selected_layer} language layers before construction"
+                f"inconsistent: {selected_layer} vs {int(arch.num_hidden_layers)}"
             )
         self.model = Qwen3VLModel(
             arch,
@@ -82,12 +273,17 @@ class MiniMaxH3Qwen3VLEncoder(TextEncoder):
             use_tensor_parallel=True,
             prefix="model",
         )
-        # H3 consumes the unnormalized output immediately after layer 49.
+        # H3 and ClipProj consume an unnormalized intermediate residual stream.
         self.model.language_model.norm = nn.Identity()
         self.image_token_id = int(arch.image_token_id)
         self.video_token_id = int(arch.video_token_id)
         self.selected_lm_layer = selected_layer
         self.hidden_dim = MINIMAX_H3_QWEN3VL_HIDDEN_DIM
+        self.conditioning_projection = (
+            MiniMaxH3ConditioningProjection(arch.conditioning_projection_path)
+            if arch.conditioning_projection_path is not None
+            else None
+        )
 
     @property
     def device(self) -> torch.device:
@@ -179,7 +375,10 @@ class MiniMaxH3Qwen3VLEncoder(TextEncoder):
             )
             call_kwargs["video_grid_thw"] = host_video_grid_thw
 
-        hidden = self(**call_kwargs).last_hidden_state[0].to(torch.bfloat16)
+        hidden = self(**call_kwargs).last_hidden_state[0]
+        if self.conditioning_projection is not None:
+            hidden = self.conditioning_projection(hidden)
+        hidden = hidden.to(torch.bfloat16)
         expected_shape = [int(ids.shape[1]), self.hidden_dim]
         if list(hidden.shape) != expected_shape:
             raise ValueError(

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

@@ -180,6 +180,22 @@ class TestServerArgsPathExpansion(unittest.TestCase):
             {"text_encoder": "owner/repo/text_encoder/model.safetensors"},
         )
 
+    def test_supplemental_weight_file_remains_a_component_path(self):
+        args = self._from_dict_without_model_resolution(
+            {
+                "model_path": "/data/my-model",
+                "component_paths": {
+                    "conditioning_projection": "owner/repo/projection.safetensors"
+                },
+            }
+        )
+
+        self.assertEqual(
+            args.component_paths,
+            {"conditioning_projection": "owner/repo/projection.safetensors"},
+        )
+        self.assertEqual(args.component_weights_paths, {})
+
     def test_component_attention_backends_are_normalized(self):
         args = self._from_dict_without_model_resolution(
             {

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

@@ -27,6 +27,7 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader
 )
 from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder
 from sglang.multimodal_gen.runtime.models.encoders.minimax_h3_qwen3vl import (
+    MiniMaxH3ConditioningProjection,
     MiniMaxH3Qwen3VLEncoder,
 )
 
@@ -132,7 +133,9 @@ class TestTextEncoderClassResolution(unittest.TestCase):
 
 class TestMiniMaxH3CheckpointFilter(unittest.TestCase):
     def test_only_known_unconsumed_weights_are_filtered(self):
-        should_load = MiniMaxH3Qwen3VLEncoder.should_materialize_checkpoint_weight
+        encoder = MiniMaxH3Qwen3VLEncoder.__new__(MiniMaxH3Qwen3VLEncoder)
+        encoder.selected_lm_layer = 50
+        should_load = encoder.should_materialize_checkpoint_weight
         expected = {
             "model.language_model.layers.49.self_attn.q_proj.weight": True,
             "model.language_model.layers.50.self_attn.q_proj.weight": False,
@@ -151,10 +154,18 @@ class TestMiniMaxH3CheckpointFilter(unittest.TestCase):
             {name: should_load(name) for name in expected},
             expected,
         )
+        encoder.selected_lm_layer = 24
+        self.assertTrue(
+            should_load("model.language_model.layers.23.mlp.down_proj.weight")
+        )
+        self.assertFalse(
+            should_load("model.language_model.layers.24.mlp.down_proj.weight")
+        )
 
     def test_vision_qkv_checkpoint_name_maps_to_native_projection(self):
         encoder = MiniMaxH3Qwen3VLEncoder.__new__(MiniMaxH3Qwen3VLEncoder)
         torch.nn.Module.__init__(encoder)
+        encoder.selected_lm_layer = 50
         encoder.model = torch.nn.Module()
         encoder.model.visual = torch.nn.Module()
         block = torch.nn.Module()
@@ -175,6 +186,7 @@ class TestMiniMaxH3CheckpointFilter(unittest.TestCase):
     def test_comfy_language_checkpoint_name_maps_to_native_namespace(self):
         encoder = MiniMaxH3Qwen3VLEncoder.__new__(MiniMaxH3Qwen3VLEncoder)
         torch.nn.Module.__init__(encoder)
+        encoder.selected_lm_layer = 50
         encoder.model = torch.nn.Module()
         encoder.model.language_model = torch.nn.Module()
         layer = torch.nn.Module()
@@ -192,6 +204,69 @@ class TestMiniMaxH3CheckpointFilter(unittest.TestCase):
         torch.testing.assert_close(layer.self_attn.q_proj.weight, source)
 
 
+class TestMiniMaxH3ConditioningProjection(unittest.TestCase):
+    @staticmethod
+    def _save_projection(path, tensors):
+        save_file(tensors, path, metadata={"tap": "2"})
+
+    def test_linear_projection_and_attention_sink(self):
+        with tempfile.NamedTemporaryFile(suffix=".safetensors") as checkpoint:
+            tensors = {
+                "W": torch.tensor([[1.0, 2.0], [3.0, 4.0]]),
+                "mean_in": torch.tensor([1.0, -1.0]),
+                "std_in": torch.tensor([2.0, 4.0]),
+                "mean_out": torch.tensor([0.5, -0.5]),
+                "std_out": torch.tensor([2.0, 3.0]),
+                "sink_out": torch.tensor([9.0, 8.0]),
+            }
+            self._save_projection(checkpoint.name, tensors)
+            projection = MiniMaxH3ConditioningProjection(checkpoint.name)
+            hidden = torch.tensor([[5.0, 3.0], [3.0, -1.0]])
+            expected = (hidden - tensors["mean_in"]) / tensors["std_in"]
+            expected = expected @ tensors["W"]
+            expected = expected * tensors["std_out"] + tensors["mean_out"]
+            expected[0] = tensors["sink_out"]
+
+            torch.testing.assert_close(projection(hidden), expected)
+            self.assertEqual(projection.tap, 2)
+
+    def test_mlp_only_projection(self):
+        with tempfile.NamedTemporaryFile(suffix=".safetensors") as checkpoint:
+            tensors = {
+                "mean_in": torch.zeros(2),
+                "std_in": torch.ones(2),
+                "mean_out": torch.tensor([0.5]),
+                "std_out": torch.tensor([2.0]),
+                "mlp.0.weight": torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
+                "mlp.0.bias": torch.tensor([0.25, -0.25]),
+                "mlp.2.weight": torch.tensor([[2.0, -1.0]]),
+                "mlp.2.bias": torch.tensor([0.75]),
+            }
+            self._save_projection(checkpoint.name, tensors)
+            projection = MiniMaxH3ConditioningProjection(checkpoint.name)
+            hidden = torch.tensor([[1.0, 2.0]])
+            residual = torch.nn.functional.linear(
+                hidden, tensors["mlp.0.weight"], tensors["mlp.0.bias"]
+            )
+            residual = torch.nn.functional.gelu(residual)
+            residual = torch.nn.functional.linear(
+                residual, tensors["mlp.2.weight"], tensors["mlp.2.bias"]
+            )
+            expected = residual * tensors["std_out"] + tensors["mean_out"]
+
+            torch.testing.assert_close(projection(hidden), expected)
+
+    def test_small_encoder_requires_matching_projection(self):
+        config = SimpleNamespace(
+            arch_config=SimpleNamespace(
+                hidden_size=2560,
+                checkpoint_num_hidden_layers=36,
+            )
+        )
+        with self.assertRaisesRegex(ValueError, "conditioning_projection"):
+            MiniMaxH3Qwen3VLEncoder.configure_component_paths(config, {})
+
+
 class TestTextEncoderQuantization(unittest.TestCase):
     def setUp(self):
         serialized = Fp8Config(