[Diffusion][Quantization] Enable MiniMax-H3 global FP8 with DLO (#5910)

Signed-off-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -110,6 +110,85 @@ def test_empty_source_prefix_keeps_full_model_strict_check():
         loader.load_weights(model)
 
 
+def test_stream_online_quant_weights_offloads_layers_after_processing():
+    from vllm.model_executor.model_loader.reload.layerwise import (
+        get_layerwise_info,
+    )
+
+    events: list[str] = []
+
+    class _OnlineQuantMethod:
+        uses_meta_device = True
+
+    class _TrackedLayer(nn.Linear):
+        def __init__(self, name: str):
+            super().__init__(2, 2, bias=False)
+            self.name = name
+            self.quant_method = _OnlineQuantMethod()
+            get_layerwise_info(self).load_numel_total = self.weight.numel()
+
+        def to(self, *args, **kwargs):
+            events.append(self.name)
+            return super().to(*args, **kwargs)
+
+    class _StreamingModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.first = _TrackedLayer("first")
+            self.second = _TrackedLayer("second")
+
+    model = _StreamingModel()
+    weights = iter(
+        [
+            ("first.weight", torch.zeros((2, 2))),
+            ("second.weight", torch.zeros((2, 2))),
+        ]
+    )
+    streamed = DiffusersPipelineLoader._stream_online_quant_weights_to_cpu(model, weights)
+
+    assert next(streamed)[0] == "first.weight"
+    get_layerwise_info(model.first).reset()
+    assert next(streamed)[0] == "second.weight"
+    assert events == ["first"]
+
+    get_layerwise_info(model.second).reset()
+    with pytest.raises(StopIteration):
+        next(streamed)
+    assert events == ["first", "second"]
+
+
+def test_process_weights_skips_completed_online_quant_layer(monkeypatch):
+    from unittest.mock import Mock
+
+    from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
+    from vllm.model_executor.model_loader.reload import layerwise
+
+    class _TrackedLayer(nn.Linear):
+        def __init__(self):
+            super().__init__(2, 2, bias=False)
+            self.to_calls: list[object] = []
+
+        def to(self, *args, **kwargs):
+            self.to_calls.append(args[0] if args else kwargs.get("device"))
+            return super().to(*args, **kwargs)
+
+    model = nn.Module()
+    model.layer = _TrackedLayer()
+    quant_method = Mock(spec=QuantizeMethodBase)
+    quant_method.uses_meta_device = True
+    model.layer.quant_method = quant_method
+    model.layer._already_called_process_weights_after_loading = True
+    finalize = Mock()
+    monkeypatch.setattr(layerwise, "finalize_layerwise_processing", finalize)
+
+    loader = _make_loader_with_weights([])
+    loader._process_weights_after_loading(model, torch.device("cuda"))
+
+    finalize.assert_called_once_with(model, model_config=None)
+    quant_method.process_weights_after_loading.assert_not_called()
+    assert model.layer.to_calls == []
+
+
 class _ConfigAwareModel(nn.Module):
     def __init__(self, *, od_config):
         super().__init__()

tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py

@@ -7,17 +7,17 @@ from vllm_omni.diffusion.models.flux2.pipeline_flux2 import (
     Flux2Pipeline,
     _resolve_component_quant_config,
 )
+from vllm_omni.quantization import ComponentQuantizationConfig
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def test_resolve_component_quant_config_routes_components():
     text_encoder_config = object()
-    component_config = Mock()
-    component_config.resolve.side_effect = {
-        "text_encoder": text_encoder_config,
-        "transformer": None,
-    }.get
+    component_config = ComponentQuantizationConfig(
+        component_configs={"text_encoder": text_encoder_config},
+        default_config=None,
+    )
 
     assert _resolve_component_quant_config(component_config, "text_encoder") is text_encoder_config
     assert _resolve_component_quant_config(component_config, "transformer") is None

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -5,7 +5,7 @@ import json
 import sys
 from multiprocessing.reduction import ForkingPickler
 from types import SimpleNamespace
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 import numpy as np
 import pytest
@@ -363,8 +363,18 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once(
     if expected_partition == "ref2va":
         assert pipeline._transformer_for_task("ref2va") is pipeline.transformer
     assert [source.model_or_path for source in pipeline.weights_sources] == [
-        str(tmp_path / partition_name) for partition_name in source_partitions
+        *(str(tmp_path / partition_name) for partition_name in source_partitions),
+        str(component_path),
     ]
+    assert [source.subfolder for source in pipeline.weights_sources] == [
+        *("transformer" for _ in source_partitions),
+        "text_encoder",
+    ]
+    expected_prefixes = ["transformer."]
+    if expected_dits == 2:
+        expected_prefixes.append("transformers_ref.")
+    expected_prefixes.append("text_encoder.")
+    assert [source.prefix for source in pipeline.weights_sources] == expected_prefixes
     expected_patterns = (
         pipeline_module.MINIMAX_H3_DOWNLOAD_PATTERNS
         if expected_partition == "combined"
@@ -936,6 +946,66 @@ def test_text_encoder_stub_constructs_without_group_or_weights():
     assert list(encoder.named_parameters()) == []
 
 
+def test_global_quant_config_is_shared_by_dit_and_encoder():
+    from vllm_omni.quantization import build_quant_config, resolve_component_quant_config
+
+    global_config = build_quant_config("fp8")
+
+    assert resolve_component_quant_config(global_config, "transformer") is global_config
+    assert resolve_component_quant_config(global_config, "text_encoder") is global_config
+
+
+def test_minimax_h3_fp8_component_selection():
+    from vllm_omni.quantization import build_quant_config, resolve_component_quant_config
+
+    dit_only = build_quant_config({"transformer": {"method": "fp8"}})
+    encoder_only = build_quant_config({"text_encoder": {"method": "fp8"}})
+
+    assert resolve_component_quant_config(dit_only, "transformer").get_name() == "fp8"
+    assert resolve_component_quant_config(dit_only, "text_encoder") is None
+    assert resolve_component_quant_config(encoder_only, "transformer") is None
+    assert resolve_component_quant_config(encoder_only, "text_encoder").get_name() == "fp8"
+
+
+def test_text_encoder_rejects_serialized_fp8():
+    from vllm_omni.diffusion.models.minimax_h3.encoder import _validate_encoder_quant_config
+
+    config = SimpleNamespace(
+        get_name=lambda: "fp8",
+        is_checkpoint_fp8_serialized=True,
+    )
+    with pytest.raises(ValueError, match="online FP8 only"):
+        _validate_encoder_quant_config(config)
+
+
+def test_text_encoder_linear_delegates_quantization_to_vllm_factory():
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+    import vllm_omni.diffusion.models.minimax_h3.encoder as encoder_module
+
+    group = SimpleNamespace(rank_in_group=0, world_size=1)
+    method = UnquantizedLinearMethod()
+    quant_config = Mock()
+    quant_config.get_quant_method.return_value = method
+    prefix = "text_encoder.text_model.layers.0.mlp.gate_up_proj"
+
+    with (
+        patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0),
+        patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1),
+    ):
+        linear = encoder_module.MiniMaxH3Qwen3VLMergedColumnParallelLinear(
+            group,
+            input_size=4,
+            intermediate_size=8,
+            dtype=torch.bfloat16,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+
+    assert linear.quant_method is method
+    quant_config.get_quant_method.assert_called_once_with(linear, prefix=prefix)
+
+
 def test_no_offload_keeps_text_encoder_resident():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
@@ -978,15 +1048,21 @@ def test_model_offload_uses_hooked_text_encoder_call():
     pipeline.text_encoder.load_to_device.assert_not_called()
 
 
-def test_layerwise_offload_releases_text_encoder():
+@pytest.mark.parametrize(
+    "offload_flag",
+    ["enable_layerwise_offload", "enable_distributed_layerwise_offload"],
+)
+def test_layerwise_offload_releases_text_encoder(offload_flag):
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
     pipeline.od_config = SimpleNamespace(
         enable_cpu_offload=False,
-        enable_layerwise_offload=True,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=False,
     )
+    setattr(pipeline.od_config, offload_flag, True)
     pipeline.text_encoder = Mock()
     expected = torch.ones(2, 3)
     pipeline.text_encoder.encode_ids.return_value = expected
@@ -1988,21 +2064,25 @@ def _one_layer_text_encoder():
     group = _SingleRankEncoderGroup(0)
     layer = nn.Module()
     layer.self_attn = nn.Module()
-    layer.self_attn.qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear(
-        group,
-        hidden_size=_ENCODER_HIDDEN,
-        num_heads=_ENCODER_NUM_HEADS,
-        num_kv_heads=_ENCODER_NUM_KV_HEADS,
-        head_dim=_ENCODER_HEAD_DIM,
-        dtype=torch.float32,
-    )
     layer.mlp = nn.Module()
-    layer.mlp.gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear(
-        group,
-        input_size=_ENCODER_HIDDEN,
-        intermediate_size=_ENCODER_INTERMEDIATE,
-        dtype=torch.float32,
-    )
+    with (
+        patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0),
+        patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1),
+    ):
+        layer.self_attn.qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear(
+            group,
+            hidden_size=_ENCODER_HIDDEN,
+            num_heads=_ENCODER_NUM_HEADS,
+            num_kv_heads=_ENCODER_NUM_KV_HEADS,
+            head_dim=_ENCODER_HEAD_DIM,
+            dtype=torch.float32,
+        )
+        layer.mlp.gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear(
+            group,
+            input_size=_ENCODER_HIDDEN,
+            intermediate_size=_ENCODER_INTERMEDIATE,
+            dtype=torch.float32,
+        )
     layer.input_layernorm = MiniMaxH3Qwen3VLRMSNorm(_ENCODER_HIDDEN)
 
     encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
@@ -2035,6 +2115,7 @@ def _write_text_encoder_checkpoint(path, *, omit=()):
     safetensors.torch.save_file(tensors, str(path / "model.safetensors"))
     index = {"weight_map": dict.fromkeys(tensors, "model.safetensors")}
     (path / "model.safetensors.index.json").write_text(json.dumps(index))
+    return tensors
 
 
 @pytest.mark.parametrize(
@@ -2056,10 +2137,10 @@ def test_encoder_load_reports_missing_fused_source_shards(
     tmp_path, omitted_sources, expected_detail, unreported_target
 ):
     encoder = _one_layer_text_encoder()
-    _write_text_encoder_checkpoint(tmp_path, omit=omitted_sources)
+    weights = _write_text_encoder_checkpoint(tmp_path, omit=omitted_sources)
 
     with pytest.raises(RuntimeError) as excinfo:
-        encoder._load_weights(str(tmp_path))
+        encoder.load_weights(weights.items())
 
     message = str(excinfo.value)
     assert expected_detail in message
@@ -2068,17 +2149,17 @@ def test_encoder_load_reports_missing_fused_source_shards(
 
 def test_encoder_load_rejects_unloaded_plain_param(tmp_path):
     encoder = _one_layer_text_encoder()
-    _write_text_encoder_checkpoint(tmp_path, omit=("input_layernorm.weight",))
+    weights = _write_text_encoder_checkpoint(tmp_path, omit=("input_layernorm.weight",))
 
     with pytest.raises(RuntimeError, match=r"weights not loaded.*input_layernorm"):
-        encoder._load_weights(str(tmp_path))
+        encoder.load_weights(weights.items())
 
 
 def test_encoder_load_places_every_source_in_its_own_rows(tmp_path):
     encoder = _one_layer_text_encoder()
-    _write_text_encoder_checkpoint(tmp_path)
+    weights = _write_text_encoder_checkpoint(tmp_path)
 
-    encoder._load_weights(str(tmp_path))
+    encoder.load_weights(weights.items())
 
     params = dict(encoder.named_parameters())
     q_rows = _ENCODER_NUM_HEADS * _ENCODER_HEAD_DIM

tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py

@@ -211,3 +211,18 @@ def test_pipeline_resolves_transformer_component_quant_config():
     assert transformer_config.ignored_layers == ignored_layers
     assert _resolve_component_quant_config(component_config, "transformer") is transformer_config
     assert _resolve_component_quant_config(transformer_config, "transformer") is transformer_config
+
+
+def test_pipeline_strips_prequantized_text_encoder_config():
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _resolve_minimax_h3_text_encoder_quant_config
+    from vllm_omni.quantization import ComponentQuantizationConfig
+
+    prequantized = Mock()
+    prequantized.get_name.return_value = "modelopt"
+    online_fp8 = Mock()
+    online_fp8.get_name.return_value = "fp8"
+    component_config = ComponentQuantizationConfig({"text_encoder": prequantized})
+
+    assert _resolve_minimax_h3_text_encoder_quant_config(online_fp8) is online_fp8
+    assert _resolve_minimax_h3_text_encoder_quant_config(prequantized) is None
+    assert _resolve_minimax_h3_text_encoder_quant_config(component_config) is None

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -139,6 +139,41 @@ class TestDistributedLayerwiseOffloadHook:
         assert shard.numel() == 4
         assert torch.equal(shard, _make_values(10.0))
 
+    def test_prefetch_preserves_transposed_weight_stride(
+        self,
+        dist_group,
+        patched_offload_runtime,
+    ):
+        """Online-FP8 Cutlass weights must retain their transposed layout."""
+
+        class StridedBlock(nn.Module):
+            def __init__(self, start: float):
+                super().__init__()
+                base = torch.arange(start, start + 12).reshape(3, 4)
+                self.weight = nn.Parameter(base.t(), requires_grad=False)
+
+        current_block = StridedBlock(1.0)
+        next_block = StridedBlock(20.0)
+        expected = next_block.weight.detach().clone()
+        expected_stride = next_block.weight.stride()
+        assert expected_stride == (1, 4)
+
+        hook = DistributedLayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            dp_group=None,
+            dp_size=1,
+            rank=0,
+            copy_stream=DummyStream(),
+            comm_stream=DummyStream(),
+            pin_memory=False,
+        )
+        hook.initialize_hook(current_block)
+        hook.prefetch_layer(slot=0, non_blocking=False)
+
+        assert next_block.weight.stride() == expected_stride
+        assert torch.equal(next_block.weight, expected)
+
     def test_dtensor_wrapper_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
         """DTensor wrapper should be preserved through prefetch/offload cycle."""
         current_block = TinyBlock(_make_values(1.0))
@@ -386,6 +421,24 @@ class TestPinnedResidentLayerGroup:
 
         group.offload()
 
+    def test_load_preserves_noncontiguous_weight_stride(self, patched_offload_runtime):
+        block = nn.Module()
+        weight = torch.arange(6, dtype=torch.float32).reshape(2, 3).t()
+        expected = weight.detach().clone()
+        expected_stride = weight.stride()
+        block.weight = nn.Parameter(weight)
+        group = PinnedResidentLayerGroup(
+            [block],
+            device=torch.device("cpu"),
+            copy_stream=DummyStream(),
+            pin_memory=False,
+        )
+
+        group.load()
+        assert torch.equal(block.weight, expected)
+        assert block.weight.stride() == expected_stride
+        group.offload()
+
     def test_mmap_loader_attrs_survive_to_empty_parameter_replacement(self):
         module = nn.Linear(2, 2, bias=False)
 

tests/diffusion/offloader/test_layerwise_backend.py

@@ -79,11 +79,11 @@ def dist_group():
 
 
 @pytest.fixture
-def patched_offload_runtime(mocker):
-    mocker.patch.object(layerwise_backend_module.current_omni_platform, "Stream", DummyStream)
-    mocker.patch.object(layerwise_backend_module.current_omni_platform, "Event", DummyEvent)
-    mocker.patch.object(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
-    mocker.patch.object(layerwise_backend_module.current_omni_platform, "stream", dummy_stream)
+def patched_offload_runtime(monkeypatch):
+    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Stream", DummyStream)
+    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Event", DummyEvent)
+    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
+    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "stream", dummy_stream)
 
 
 class TinyBlock(nn.Module):
@@ -128,6 +128,33 @@ class TestLayerwiseOffloadHook:
         assert current_block.weight.to_local().shape == torch.Size([4])
         assert not hook.is_materialized
 
+    def test_prefetch_preserves_transposed_weight_stride(self, patched_offload_runtime):
+        """Online-FP8 Cutlass weights must retain their transposed layout."""
+
+        class StridedBlock(nn.Module):
+            def __init__(self, start: float):
+                super().__init__()
+                base = torch.arange(start, start + 12).reshape(3, 4)
+                self.weight = nn.Parameter(base.t(), requires_grad=False)
+
+        current_block = StridedBlock(1.0)
+        next_block = StridedBlock(20.0)
+        expected = next_block.weight.detach().clone()
+        expected_stride = next_block.weight.stride()
+        assert expected_stride == (1, 4)
+
+        hook = LayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            stream=DummyStream(),
+            pin_memory=False,
+        )
+        hook.initialize_hook(current_block)
+        hook.prefetch_layer(non_blocking=False)
+
+        assert next_block.weight.stride() == expected_stride
+        assert torch.equal(next_block.weight, expected)
+
 
 class _DummyBlock(nn.Module):
     def __init__(self):

tests/diffusion/quantization/test_component_routing.py

@@ -16,10 +16,11 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
 )
 from vllm_omni.quantization.component_config import (
     ComponentQuantizationConfig,
+    resolve_component_quant_config,
 )
 from vllm_omni.quantization.inc_config import OmniINCConfig
 
-pytestmark = [pytest.mark.core_model]
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 # ---------------------------------------------------------------------------
@@ -347,6 +348,24 @@ class TestTalkerVisualRouting:
 # ===================================================================
 
 
+class TestResolveComponentQuantConfig:
+    def test_global_config_applies_to_each_quantization_aware_component(self):
+        config = _MockQuantConfig("fp8")
+
+        assert resolve_component_quant_config(config, "transformer") is config
+        assert resolve_component_quant_config(config, "text_encoder") is config
+
+    def test_component_map_is_the_only_scope_override(self):
+        transformer = _MockQuantConfig("fp8")
+        config = ComponentQuantizationConfig(
+            component_configs={"transformer": transformer},
+            default_config=None,
+        )
+
+        assert resolve_component_quant_config(config, "transformer") is transformer
+        assert resolve_component_quant_config(config, "text_encoder") is None
+
+
 class TestComponentResolve:
     def test_longest_prefix_match(self):
         a = _MockQuantConfig("a")