[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>

docs/design/feature/quantization.md

@@ -76,6 +76,12 @@ equivalent. `auto-round` and `auto_round` use the Omni INC/AutoRound adapter.
 ModelOpt checkpoint metadata is detected from its method, producer metadata, or
 `quant_algo`, then mapped to the corresponding vLLM ModelOpt configuration.
 
+Unless the input is a per-component dictionary, the resolved configuration is
+passed unchanged to every quantization-aware component constructed by the
+pipeline. It affects only layers supported by that component's quantization
+method; it does not rewrite arbitrary `torch.nn` modules. A per-component
+dictionary narrows this global scope through runtime layer-prefix routing.
+
 The factory uses Omni-specific overrides for methods whose diffusion behavior
 or platform dispatch is not provided by the vLLM registry. The current override
 set includes INT8, BitsAndBytes, MXFP8, MXFP4, MXFP4 dual-scale, and INC/
@@ -154,9 +160,11 @@ quantizable vLLM layers. Each method configuration is responsible for:
 Model code must keep precision-sensitive or unsupported modules unquantized.
 The shared `safe_quant_config()` helper is used by model integrations for
 normalization and modulation layers where a generic FP8 configuration would be
-unsafe. Text encoders, audio/vision encoders, VAEs, schedulers, and tokenizers
-are not implicitly quantized just because the main transformer is quantized;
-they require explicit component routing and method support.
+unsafe. Encoders exposed through the quantization factory receive a global
+configuration just like the main transformer, but only their supported vLLM
+quantizable layers are affected. Ordinary `torch.nn` encoder layers, VAEs,
+schedulers, and tokenizers are not rewritten automatically. A component map can
+select a narrower scope, but it cannot add quantization support to a layer.
 
 Weight quantization and quantized KV cache are separate paths. The former is
 resolved through `quantization_config`; runtime attention/KV-cache quantization

docs/user_guide/quantization/fp8.md

@@ -95,7 +95,7 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
 | HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | |
 | HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | |
 | Cosmos3 | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | Yes | Not validated | All layers | None | |
-| MiniMax H3 | `MiniMaxAI/MiniMax-H3` | Yes | Not validated | DiT linears except patch, timestep, and final projections; not compatible with layerwise offload | None | |
+| MiniMax-H3 | `MiniMaxAI/MiniMax-H3` (`FL2VA` / `Ref2VA`) | Yes | Not validated | `quantization="fp8"` quantizes eligible DiT and text-encoder linears; mixed-precision input/output heads stay FP32 | None | ✅︎ |
 
 ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
 
@@ -138,6 +138,39 @@ outputs = omni.generate(
 )
 ```
 
+### Global and per-component scope
+
+A plain, global quantization configuration is passed unchanged to every
+quantization-aware component constructed by a pipeline. This includes an
+eligible encoder when that encoder is implemented with vLLM quantizable
+layers; it does not rewrite arbitrary `torch.nn` modules. A structured
+component map is the only way to narrow that scope. Pipeline integrations that
+do not yet expose an encoder through the quantization factory remain DiT-only.
+
+For a pipeline that exposes both a transformer and a quantization-aware text
+encoder, the scope is:
+
+| Configuration | Transformer | Text encoder | Components without supported quantizable layers |
+|---------------|-------------|--------------|-------------------------------------------------|
+| `quantization="fp8"` | FP8 | FP8 | checkpoint precision |
+| `{"transformer": {"method": "fp8"}}` | FP8 | checkpoint precision | checkpoint precision |
+| `{"text_encoder": {"method": "fp8"}}` | checkpoint precision | FP8 | checkpoint precision |
+
+Use `quantization_config` for component-selective Python configuration. For
+example, quantize only the text decoder:
+
+```python
+omni = Omni(
+    model="<your-model>",
+    quantization_config={"text_encoder": {"method": "fp8"}},
+)
+```
+
+Component keys are runtime prefixes exposed by the pipeline integration; common
+keys include `transformer` and `text_encoder`. Entries may be combined, and
+`ignored_layers` can keep named eligible layers in checkpoint precision. Check
+the model recipe for supported components and their runtime prefixes.
+
 CLI:
 
 ```bash

recipes/MiniMaxAI/MiniMax-H3.md

@@ -357,12 +357,19 @@ No restart is needed: `task=fl2va` routes to `FL2VA/transformer`, while
 
 ### Online FP8 quantization
 
-MiniMax H3 supports load-time FP8 quantization of the DiT. The checkpoint
-remains BF16 on disk; vLLM-Omni quantizes eligible weights while loading and
-uses dynamic activation scaling during inference. By default, attention and
-MLP linears in the token refiner and main DiT blocks, the condition
-projection, and all AdaLN projections use FP8. Patch, timestep, and final
-projections remain FP32; the text encoder and VAEs are unchanged.
+MiniMax H3 supports online FP8 quantization of both the DiT and the Qwen3-VL
+text decoder. The checkpoint remains BF16 on disk; vLLM-Omni creates FP8
+weights at runtime and uses dynamic activation scaling during inference. By
+default, `--quantization fp8` quantizes eligible attention and MLP linears in
+the text decoder, token refiner, and main DiT blocks, as well as the DiT
+condition and AdaLN projections. The Qwen vision tower, embeddings, norms,
+RoPE, both VAEs, and the model's FP32 patch, timestep, and output projections
+keep checkpoint precision.
+
+To select a component, use `--diffusion-quantization-config` with
+`{"transformer":{"method":"fp8"}}` for DiT-only FP8 or
+`{"text_encoder":{"method":"fp8"}}` for text-decoder-only FP8. The two
+entries can be combined. The shorthand below enables both components.
 
 Add this option to an existing H3 server command:
 
@@ -370,6 +377,28 @@ Add this option to an existing H3 server command:
 --quantization fp8
 ```
 
+#### Single 96 GB GPU, no-offload capacity check
+
+Use the FL2VA-only partition for this capacity test. Loading the combined
+service would also load the Ref2VA DiT and would test a different memory
+budget. A no-offload capacity check should contain none of
+`--enable-cpu-offload`, `--enable-layerwise-offload`, or
+`--enable-distributed-layerwise-offload`. VAE tiling changes decode placement
+but does not offload model weights to the CPU.
+
+The run passes the capacity check when the server initializes, the request
+finishes without CUDA OOM or Xid errors, `peak_used_mib` remains below the
+card's reported `memory_total_mib`, and `ffprobe` reports H.264 video plus
+32 kHz stereo AAC audio. Report the measured headroom rather than assuming
+that every nominal 96 GB SKU exposes the same MiB total.
+
+As a capacity proxy only, the same five-second case on one B300 measured a
+92,946 MiB whole-device first-request peak and a 92,146 MiB worker peak. Its
+encode, diffuse, decode, and client wall times were 8.664 s, 134.504 s,
+6.016 s, and 151.583 s. These numbers suggest that a 96 GB card may fit, but
+they are not an RTX PRO 6000 validation: kernels, allocator behavior, and
+reported device capacity differ across GPUs.
+
 Use `ignored_layers` to keep any otherwise eligible linear in BF16. H3
 resolves the `transformer` component before constructing the DiT, so names do
 not start with `transformer.`. Entries are exact runtime linear prefixes; a
@@ -386,10 +415,10 @@ For example, keep the first main block's attention projections in BF16 with:
   '{"transformer":{"method":"fp8","ignored_layers":["blocks.0.attn.qkv_proj","blocks.0.attn.out_proj"]}}'
 ```
 
-The structured option replaces `--quantization fp8`. Online FP8 is currently
-incompatible with H3 layerwise offload because the offload path produces a
-weight stride rejected by the Cutlass FP8 kernel. Use resident FP8 with tensor
-parallelism and VAE tiling instead.
+The structured option replaces `--quantization fp8`. Online FP8 can be used
+with H3 layerwise offload and with distributed layerwise offload's full-weight
+per-rank path (`--dlo-no-use-allgather`). The sharded DLO AllGather path is not
+supported for runtime-created FP8 weights.
 
 ## AMD ROCm (gfx942 / gfx950)
 
@@ -750,6 +779,13 @@ audio spectral cosine 0.9589 (minimum 0.80), and audio RMS ratio 0.9342. The
 resident per-GPU peak was 68.52 GiB for BF16 and 53.51 GiB for FP8, a 22%
 reduction.
 
+For direct human inspection, see the external
+[BF16 versus global-FP8 comparison](https://lishunyang12.github.io/vllm-omni-rankings/scripts/minimax_h3_global_fp8_vs_bf16/),
+which includes matched five-second T2VA, I2VA, and Ref2VA videos. The
+[comparison sources](https://github.com/lishunyang12/vllm-omni-rankings/tree/main/scripts/minimax_h3_global_fp8_vs_bf16)
+also record per-task fidelity metrics and provenance without storing generated
+media in this repository.
+
 ## TeaCache acceleration
 
 TeaCache reuses DiT block residuals across denoising steps when consecutive
@@ -834,7 +870,8 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
   H3 native `tile` mode only.
 - A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use
   pure Ulysses.
-- Online FP8 is not compatible with layerwise offload.
+- Runtime-created FP8 weights do not support the sharded DLO AllGather path;
+  use `--dlo-no-use-allgather` when combining online FP8 with DLO.
 - TeaCache and Cache-DiT cannot be enabled on the same server.
 - Image+audio Ref2VA accepts exactly one image and one audio reference.
 - Video Ref2VA accepts one or more video files, but not an additional standalone

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

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -298,6 +298,60 @@ class DiffusersPipelineLoader:
         for source in sources:
             yield from self._get_weights_iterator(source, model=model)
 
+    @staticmethod
+    def _stream_online_quant_weights_to_cpu(
+        model: nn.Module,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> Generator[tuple[str, torch.Tensor], None, None]:
+        """Offload each online-quantized layer as soon as it is complete.
+
+        Upstream vLLM's online layerwise loader materializes and quantizes a
+        layer synchronously while consuming ``weights``.  Generator execution
+        resumes after that weight has been consumed, which gives us a safe
+        point to move completed layers to CPU before the next layer is loaded.
+        This bounds accelerator residency during CPU-offloaded model startup
+        instead of retaining the entire quantized model until loading ends.
+        """
+        from vllm.model_executor.model_loader.reload.layerwise import (
+            get_layerwise_info,
+        )
+
+        pending = {
+            module
+            for module in model.modules()
+            if getattr(getattr(module, "quant_method", None), "uses_meta_device", False)
+            and get_layerwise_info(module).can_load()
+        }
+        offloaded = 0
+
+        def offload_completed() -> None:
+            nonlocal offloaded
+            for module in tuple(pending):
+                if get_layerwise_info(module).can_load():
+                    continue
+                module.to("cpu")
+                pending.remove(module)
+                offloaded += 1
+
+        for weight in weights:
+            # This runs after the consumer has handled the previous yield.
+            offload_completed()
+            yield weight
+        offload_completed()
+
+        # Quantization workspaces and the old accelerator-side parameter
+        # storages are now reusable cache blocks.  Release them before the
+        # remaining (unquantized) model tensors are copied to CPU; otherwise
+        # the cached quantization footprint and final offload overlap in the
+        # process-level startup peak.
+        if offloaded:
+            torch.accelerator.empty_cache()
+
+        logger.info(
+            "Stream-offloaded %d online-quantized layers to CPU during weight loading",
+            offloaded,
+        )
+
     def _get_weight_sources(self, model: nn.Module) -> tuple["ComponentSource", ...]:
         return tuple(
             cast(
@@ -403,7 +457,10 @@ class DiffusersPipelineLoader:
                     if load_format == "diffusers":
                         cast(DiffusersAdapterPipeline, model).load_weights()
                     else:
-                        self.load_weights(model)
+                        if offload_after_quant:
+                            self.load_weights(model, stream_online_quant_to_cpu=True)
+                        else:
+                            self.load_weights(model)
                     self._process_weights_after_loading(model, target_device)
 
             if offload_after_quant:
@@ -498,6 +555,14 @@ class DiffusersPipelineLoader:
                 continue
 
             if has_online_quant:
+                # finalize_layerwise_processing() and the synchronous online
+                # loader already processed these layers.  Avoid moving their
+                # quantized CPU weights back to the accelerator merely to call
+                # an idempotent no-op; doing so rebuilds a large CUDA allocator
+                # cache and defeats streaming CPU offload's startup-memory bound.
+                if getattr(module, "_already_called_process_weights_after_loading", False):
+                    continue
+
                 # Online quant may leave straggler params on the ``meta`` device.
                 # Move only real (non-meta) params onto the target device for
                 # processing and restore them afterward, mirroring upstream vLLM's
@@ -535,9 +600,17 @@ class DiffusersPipelineLoader:
                 if needs_device_move:
                     module.to(module_device)
 
-    def load_weights(self, model: nn.Module) -> None:
+    def load_weights(
+        self,
+        model: nn.Module,
+        *,
+        stream_online_quant_to_cpu: bool = False,
+    ) -> None:
         weights_to_load = self._get_expected_parameter_names(model)
-        loaded_weights = model.load_weights(self.get_all_weights(model))
+        weights = self.get_all_weights(model)
+        if stream_online_quant_to_cpu:
+            weights = self._stream_online_quant_weights_to_cpu(model, weights)
+        loaded_weights = model.load_weights(weights)
 
         self.counter_after_loading_weights = time.perf_counter()
         logger.info_once(

vllm_omni/diffusion/models/flux2/pipeline_flux2.py

@@ -39,16 +39,11 @@ from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 from vllm_omni.platforms import current_omni_platform
+from vllm_omni.quantization import resolve_component_quant_config as _resolve_component_quant_config
 
 logger = logging.getLogger(__name__)
 
 
-def _resolve_component_quant_config(quant_config, component: str):
-    if hasattr(quant_config, "resolve"):
-        return quant_config.resolve(component)
-    return quant_config
-
-
 class Flux2ImageProcessor(VaeImageProcessor):
     """Image processor to preprocess the reference image for Flux2."""
 

vllm_omni/diffusion/models/minimax_h3/encoder.py

@@ -30,14 +30,16 @@ complete ``[seq, 5120]`` hidden state.
 from __future__ import annotations
 
 import itertools
-import json
 import re
+from collections.abc import Iterable
 from typing import Any
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
 from vllm_omni.diffusion.layers.norm import RMSNorm as DiffusionRMSNorm
 
@@ -47,6 +49,19 @@ MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120
 logger = init_logger(__name__)
 
 
+def _validate_encoder_quant_config(quant_config: QuantizationConfig | None) -> None:
+    if (
+        quant_config is not None
+        and quant_config.get_name() == "fp8"
+        and bool(getattr(quant_config, "is_checkpoint_fp8_serialized", False))
+    ):
+        raise ValueError("MiniMax-H3 text encoder supports online FP8 only, not serialized FP8 checkpoints")
+
+
+def _is_quantized_linear(module: LinearBase) -> bool:
+    return not isinstance(module.quant_method, UnquantizedLinearMethod)
+
+
 def _default_weight_loader(
     param: nn.Parameter,
     loaded_weight: torch.Tensor,
@@ -114,7 +129,7 @@ class MiniMaxH3Qwen3VLVocabParallelEmbedding(nn.Module):
         param.data.copy_(loaded_weight.narrow(0, start_idx, shard_size))
 
 
-class MiniMaxH3Qwen3VLMergedColumnParallelLinear(nn.Module):
+class MiniMaxH3Qwen3VLMergedColumnParallelLinear(LinearBase):
     """Packed gate/up projection sharded along the output dimension."""
 
     def __init__(
@@ -123,23 +138,43 @@ class MiniMaxH3Qwen3VLMergedColumnParallelLinear(nn.Module):
         input_size: int,
         intermediate_size: int,
         dtype: torch.dtype,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
     ) -> None:
-        super().__init__()
         self.group = group
-        self.input_size = input_size
         self.intermediate_size = intermediate_size
         tp_rank, tp_size = _tp_range(group)
         assert intermediate_size % tp_size == 0, (
             f"intermediate_size {intermediate_size} must be divisible by text_encoder_tp_size {tp_size}"
         )
         self.intermediate_size_per_partition = intermediate_size // tp_size
-        self.weight = nn.Parameter(torch.empty(2 * self.intermediate_size_per_partition, input_size, dtype=dtype))
-        self.weight.weight_loader = self.weight_loader  # type: ignore[attr-defined]
+        self.output_dtype = dtype
+        super().__init__(
+            input_size=input_size,
+            output_size=2 * intermediate_size,
+            bias=False,
+            params_dtype=dtype,
+            quant_config=quant_config,
+            prefix=prefix,
+            disable_tp=True,
+        )
+        self.quant_method.create_weights(
+            self,
+            input_size_per_partition=input_size,
+            output_partition_sizes=[
+                self.intermediate_size_per_partition,
+                self.intermediate_size_per_partition,
+            ],
+            input_size=input_size,
+            output_size=2 * intermediate_size,
+            params_dtype=dtype,
+            weight_loader=self.weight_loader,
+        )
         self._tp_rank = tp_rank
         self._tp_size = tp_size
 
     def forward(self, input_: torch.Tensor) -> torch.Tensor:
-        return F.linear(input_, self.weight)
+        return self.quant_method.apply(self, input_)
 
     def weight_loader(
         self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None
@@ -152,7 +187,7 @@ class MiniMaxH3Qwen3VLMergedColumnParallelLinear(nn.Module):
             param.data[0:shard_size].copy_(loaded_weight.narrow(0, start_idx, shard_size))
 
 
-class MiniMaxH3Qwen3VLQKVParallelLinear(nn.Module):
+class MiniMaxH3Qwen3VLQKVParallelLinear(LinearBase):
     """QKV projection with GQA head sharding along the output dimension."""
 
     def __init__(
@@ -163,8 +198,9 @@ class MiniMaxH3Qwen3VLQKVParallelLinear(nn.Module):
         num_kv_heads: int,
         head_dim: int,
         dtype: torch.dtype,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
     ) -> None:
-        super().__init__()
         self.group = group
         self.hidden_size = hidden_size
         self.num_heads = num_heads
@@ -179,15 +215,33 @@ class MiniMaxH3Qwen3VLQKVParallelLinear(nn.Module):
         )
         self.local_num_heads = num_heads // tp_size
         self.local_num_kv_heads = num_kv_heads // tp_size
+        self.output_dtype = dtype
         q_local = self.local_num_heads * head_dim
         kv_local = self.local_num_kv_heads * head_dim
-        self.weight = nn.Parameter(torch.empty(q_local + 2 * kv_local, hidden_size, dtype=dtype))
-        self.weight.weight_loader = self.weight_loader  # type: ignore[attr-defined]
+        output_size = (num_heads + 2 * num_kv_heads) * head_dim
+        super().__init__(
+            input_size=hidden_size,
+            output_size=output_size,
+            bias=False,
+            params_dtype=dtype,
+            quant_config=quant_config,
+            prefix=prefix,
+            disable_tp=True,
+        )
+        self.quant_method.create_weights(
+            self,
+            input_size_per_partition=hidden_size,
+            output_partition_sizes=[q_local, kv_local, kv_local],
+            input_size=hidden_size,
+            output_size=output_size,
+            params_dtype=dtype,
+            weight_loader=self.weight_loader,
+        )
         self._tp_rank = tp_rank
         self._tp_size = tp_size
 
     def forward(self, input_: torch.Tensor) -> torch.Tensor:
-        return F.linear(input_, self.weight)
+        return self.quant_method.apply(self, input_)
 
     def weight_loader(
         self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None
@@ -240,7 +294,7 @@ def _verify_fused_shards_complete(
     raise RuntimeError(f"MiniMax H3 text encoder fused weights are missing source shards: {details}")
 
 
-class MiniMaxH3Qwen3VLRowParallelLinear(nn.Module):
+class MiniMaxH3Qwen3VLRowParallelLinear(LinearBase):
     def __init__(
         self,
         group: Any,
@@ -248,17 +302,33 @@ class MiniMaxH3Qwen3VLRowParallelLinear(nn.Module):
         output_size: int,
         dtype: torch.dtype,
         input_is_parallel: bool = True,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
     ) -> None:
-        super().__init__()
         self.group = group
-        self.input_size = input_size
-        self.output_size = output_size
         self.input_is_parallel = input_is_parallel
+        self.output_dtype = dtype
         tp_rank, tp_size = _tp_range(group)
         assert input_size % tp_size == 0, f"input_size {input_size} must be divisible by text_encoder_tp_size {tp_size}"
         self.input_size_per_partition = input_size // tp_size
-        self.weight = nn.Parameter(torch.empty(output_size, self.input_size_per_partition, dtype=dtype))
-        self.weight.weight_loader = self.weight_loader  # type: ignore[attr-defined]
+        super().__init__(
+            input_size=input_size,
+            output_size=output_size,
+            bias=False,
+            params_dtype=dtype,
+            quant_config=quant_config,
+            prefix=prefix,
+            disable_tp=True,
+        )
+        self.quant_method.create_weights(
+            self,
+            input_size_per_partition=self.input_size_per_partition,
+            output_partition_sizes=[output_size],
+            input_size=input_size,
+            output_size=output_size,
+            params_dtype=dtype,
+            weight_loader=self.weight_loader,
+        )
         self._tp_rank = tp_rank
         self._tp_size = tp_size
 
@@ -268,7 +338,7 @@ class MiniMaxH3Qwen3VLRowParallelLinear(nn.Module):
         else:
             split_input = input_.split(self.input_size_per_partition, dim=-1)
             input_parallel = split_input[self._tp_rank].contiguous()
-        output_parallel = F.linear(input_parallel, self.weight)
+        output_parallel = self.quant_method.apply(self, input_parallel)
         if self._tp_size > 1:
             # Reduce in fp32: the reference path accumulates the full (K=8192)
             # dot product inside a single cuBLAS GEMM before rounding to bf16.
@@ -276,7 +346,7 @@ class MiniMaxH3Qwen3VLRowParallelLinear(nn.Module):
             # and amplify error on this model's large-magnitude activations.
             output_parallel = output_parallel.float()
             self.group.all_reduce(output_parallel)
-            output_parallel = output_parallel.to(self.weight.dtype)
+            output_parallel = output_parallel.to(self.output_dtype)
         return output_parallel
 
     def weight_loader(
@@ -665,7 +735,14 @@ class MiniMaxH3Qwen3VLVisionModel(nn.Module):
 
 
 class MiniMaxH3Qwen3VLTextAttention(nn.Module):
-    def __init__(self, group: Any, config: Any, dtype: torch.dtype) -> None:
+    def __init__(
+        self,
+        group: Any,
+        config: Any,
+        dtype: torch.dtype,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
         super().__init__()
         self.group = group
         self.config = config
@@ -683,6 +760,8 @@ class MiniMaxH3Qwen3VLTextAttention(nn.Module):
             num_kv_heads=self.num_kv_heads,
             head_dim=self.head_dim,
             dtype=dtype,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
         )
         self.o_proj = MiniMaxH3Qwen3VLRowParallelLinear(
             group,
@@ -690,6 +769,8 @@ class MiniMaxH3Qwen3VLTextAttention(nn.Module):
             output_size=self.hidden_size,
             dtype=dtype,
             input_is_parallel=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
         )
         self.q_norm = MiniMaxH3Qwen3VLRMSNorm(self.head_dim, eps=config.rms_norm_eps, dtype=dtype)
         self.k_norm = MiniMaxH3Qwen3VLRMSNorm(self.head_dim, eps=config.rms_norm_eps, dtype=dtype)
@@ -703,14 +784,23 @@ class MiniMaxH3Qwen3VLTextAttention(nn.Module):
         hidden_shape = (*input_shape, -1, self.head_dim)
         q_local = self.qkv_proj.local_num_heads * self.head_dim
         kv_local = self.qkv_proj.local_num_kv_heads * self.head_dim
-        q_weight = self.qkv_proj.weight[0:q_local]
-        k_weight = self.qkv_proj.weight[q_local : q_local + kv_local]
-        v_weight = self.qkv_proj.weight[q_local + kv_local : q_local + 2 * kv_local]
-        # Three separate GEMMs keep the numerics identical to the reference
-        # path (packed QKV changes cuBLAS tiling and shifts bf16 rounding).
-        query_states = self.q_norm(F.linear(hidden_states, q_weight).view(hidden_shape)).transpose(1, 2)
-        key_states = self.k_norm(F.linear(hidden_states, k_weight).view(hidden_shape)).transpose(1, 2)
-        value_states = F.linear(hidden_states, v_weight).view(hidden_shape).transpose(1, 2)
+        if _is_quantized_linear(self.qkv_proj):
+            query, key, value = self.qkv_proj(hidden_states).split(
+                (q_local, kv_local, kv_local),
+                dim=-1,
+            )
+        else:
+            q_weight = self.qkv_proj.weight[0:q_local]
+            k_weight = self.qkv_proj.weight[q_local : q_local + kv_local]
+            v_weight = self.qkv_proj.weight[q_local + kv_local : q_local + 2 * kv_local]
+            # Three separate GEMMs keep the numerics identical to the reference
+            # path (packed QKV changes cuBLAS tiling and shifts bf16 rounding).
+            query = F.linear(hidden_states, q_weight)
+            key = F.linear(hidden_states, k_weight)
+            value = F.linear(hidden_states, v_weight)
+        query_states = self.q_norm(query.view(hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(key.view(hidden_shape)).transpose(1, 2)
+        value_states = value.view(hidden_shape).transpose(1, 2)
 
         cos, sin = position_embeddings
         query_states, key_states = _apply_rotary_pos_emb(query_states, key_states, cos, sin)
@@ -734,7 +824,14 @@ class MiniMaxH3Qwen3VLTextAttention(nn.Module):
 
 
 class MiniMaxH3Qwen3VLTextMLP(nn.Module):
-    def __init__(self, group: Any, config: Any, dtype: torch.dtype) -> None:
+    def __init__(
+        self,
+        group: Any,
+        config: Any,
+        dtype: torch.dtype,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
         super().__init__()
         self.hidden_size = config.hidden_size
         self.intermediate_size = config.intermediate_size
@@ -743,6 +840,8 @@ class MiniMaxH3Qwen3VLTextMLP(nn.Module):
             input_size=self.hidden_size,
             intermediate_size=self.intermediate_size,
             dtype=dtype,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
         )
         self.down_proj = MiniMaxH3Qwen3VLRowParallelLinear(
             group,
@@ -750,22 +849,46 @@ class MiniMaxH3Qwen3VLTextMLP(nn.Module):
             output_size=self.hidden_size,
             dtype=dtype,
             input_is_parallel=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         ip_local = self.gate_up_proj.intermediate_size_per_partition
-        gate = F.linear(x, self.gate_up_proj.weight[0:ip_local])
-        up = F.linear(x, self.gate_up_proj.weight[ip_local : 2 * ip_local])
+        if _is_quantized_linear(self.gate_up_proj):
+            gate, up = self.gate_up_proj(x).split(ip_local, dim=-1)
+        else:
+            gate = F.linear(x, self.gate_up_proj.weight[0:ip_local])
+            up = F.linear(x, self.gate_up_proj.weight[ip_local : 2 * ip_local])
         x = F.silu(gate) * up
         return self.down_proj(x)
 
 
 class MiniMaxH3Qwen3VLTextDecoderLayer(nn.Module):
-    def __init__(self, group: Any, config: Any, dtype: torch.dtype) -> None:
+    def __init__(
+        self,
+        group: Any,
+        config: Any,
+        dtype: torch.dtype,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
         super().__init__()
         self.hidden_size = config.hidden_size
-        self.self_attn = MiniMaxH3Qwen3VLTextAttention(group, config, dtype)
-        self.mlp = MiniMaxH3Qwen3VLTextMLP(group, config, dtype)
+        self.self_attn = MiniMaxH3Qwen3VLTextAttention(
+            group,
+            config,
+            dtype,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.mlp = MiniMaxH3Qwen3VLTextMLP(
+            group,
+            config,
+            dtype,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mlp",
+        )
         self.input_layernorm = MiniMaxH3Qwen3VLRMSNorm(config.hidden_size, eps=config.rms_norm_eps, dtype=dtype)
         self.post_attention_layernorm = MiniMaxH3Qwen3VLRMSNorm(
             config.hidden_size, eps=config.rms_norm_eps, dtype=dtype
@@ -791,7 +914,15 @@ class MiniMaxH3Qwen3VLTextDecoderLayer(nn.Module):
 class MiniMaxH3Qwen3VLTextModel(nn.Module):
     """TP-sharded Qwen3-VL text model returning unnormalized layer-50 states."""
 
-    def __init__(self, group: Any, config: Any, selected_layer: int, dtype: torch.dtype) -> None:
+    def __init__(
+        self,
+        group: Any,
+        config: Any,
+        selected_layer: int,
+        dtype: torch.dtype,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "text_encoder.text_model",
+    ) -> None:
         super().__init__()
         self.group = group
         self.config = config
@@ -804,7 +935,16 @@ class MiniMaxH3Qwen3VLTextModel(nn.Module):
         )
         self.rotary_emb = MiniMaxH3Qwen3VLTextRotaryEmbedding(config)
         self.layers = nn.ModuleList(
-            [MiniMaxH3Qwen3VLTextDecoderLayer(group, config, dtype) for _ in range(self.num_layers)]
+            [
+                MiniMaxH3Qwen3VLTextDecoderLayer(
+                    group,
+                    config,
+                    dtype,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.layers.{layer_idx}",
+                )
+                for layer_idx in range(self.num_layers)
+            ]
         )
 
     def _deepstack_process(
@@ -951,10 +1091,13 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module):
         device: torch.device,
         load_model: bool,
         encoder_group: Any | None = None,
+        quant_config: QuantizationConfig | None = None,
     ) -> None:
         super().__init__()
         self.device_target = device
         self.encoder_group = encoder_group
+        self.quant_config = quant_config
+        _validate_encoder_quant_config(quant_config)
         self.image_token_id = 151655
         self.video_token_id = 151656
         self._tp_size = 1
@@ -975,9 +1118,9 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module):
             config.text_config,
             MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER,
             dtype,
+            quant_config=quant_config,
         )
         self.text_model.to(dtype=dtype)
-        self._load_weights(model_path)
         logger.info(
             "MiniMax H3 Qwen3-VL encoder: %d retained decoder layers, text_encoder_tp_size=%d, vision replicated",
             MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER,
@@ -1032,48 +1175,37 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module):
             return ("text_model." + rest, None)
         return None
 
-    def _load_weights(self, model_path: str) -> None:
-        import os
-
-        import safetensors.torch
-
-        index_path = os.path.join(model_path, "model.safetensors.index.json")
-        with open(index_path, encoding="utf-8") as handle:
-            index = json.load(handle)
-        weight_map: dict[str, str] = index["weight_map"]
-        files = sorted(set(weight_map.values()))
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         params = dict(self.named_parameters())
         loaded: set[str] = set()
-        # Seeded from the modules that own a multi-shard loader, so a fused weight
-        # that receives no source shard at all is caught too, not only a partially
-        # filled one.
+        # Seeded from modules that own a multi-shard loader, so a fused weight
+        # receiving no source shard is caught as well as a partially filled one.
         expected_fused_shards: dict[str, tuple[tuple[str, str | int], ...]] = {
             f"{module_name}.weight": sources
             for module_name, module in self.named_modules()
             if (sources := _fused_source_shards(module)) is not None
         }
         loaded_fused_shards: dict[str, set[str | int]] = {name: set() for name in expected_fused_shards}
-        for shard in files:
-            tensors = safetensors.torch.load_file(os.path.join(model_path, shard), device="cpu")
-            for name, tensor in tensors.items():
-                mapped = self._map_weight_name(name)
-                if mapped is None:
-                    continue
-                param_name, shard_id = mapped
-                param = params.get(param_name)
-                if param is None:
-                    logger.warning("MiniMax H3 text encoder weight %s has no target parameter", name)
-                    continue
-                weight_loader = getattr(param, "weight_loader", _default_weight_loader)
-                weight_loader(param, tensor, shard_id)
-                loaded.add(param_name)
-                if shard_id is not None and param_name in loaded_fused_shards:
-                    loaded_fused_shards[param_name].add(shard_id)
+        for name, tensor in weights:
+            mapped = self._map_weight_name(name)
+            if mapped is None:
+                continue
+            param_name, shard_id = mapped
+            param = params.get(param_name)
+            if param is None:
+                logger.warning("MiniMax H3 text encoder weight %s has no target parameter", name)
+                continue
+            weight_loader = getattr(param, "weight_loader", _default_weight_loader)
+            weight_loader(param, tensor, shard_id)
+            loaded.add(param_name)
+            if shard_id is not None and param_name in loaded_fused_shards:
+                loaded_fused_shards[param_name].add(shard_id)
         _verify_fused_shards_complete(expected_fused_shards, loaded_fused_shards)
         missing = sorted(set(params) - loaded)
         if missing:
             # Listed in full: this aborts startup, so the message is the only diagnostic.
             raise RuntimeError(f"MiniMax H3 text encoder weights not loaded: {len(missing)} params: {missing}")
+        return loaded
 
     def load_to_device(self) -> None:
         if not self.is_loaded:

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -20,6 +20,7 @@ import torch.nn as nn
 from PIL import Image
 from transformers import Qwen2TokenizerFast, Qwen3VLProcessor
 from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
 from vllm_omni.diffusion import envs
 from vllm_omni.diffusion.cache.cachedit import (
@@ -54,6 +55,12 @@ from vllm_omni.model_executor.model_loader.weight_utils import (
     download_weights_from_hf_specific,
 )
 from vllm_omni.platforms import current_omni_platform
+from vllm_omni.quantization import (
+    resolve_component_quant_config as _resolve_component_quant_config,
+)
+from vllm_omni.quantization.component_config import (
+    resolve_encoder_quant_config as _resolve_encoder_quant_config,
+)
 
 from .condition_noise import (
     minimax_h3_audio_cond_noise_aug_rows,
@@ -128,6 +135,13 @@ MINIMAX_H3_TASK_DOWNLOAD_PATTERNS = {
 }
 
 
+def _resolve_minimax_h3_text_encoder_quant_config(
+    quant_config: QuantizationConfig | None,
+) -> QuantizationConfig | None:
+    resolved = _resolve_component_quant_config(quant_config, "text_encoder")
+    return _resolve_encoder_quant_config(resolved)
+
+
 def _minimax_h3_partition_for_task(
     task_type: str | None,
     model: str | None = None,
@@ -175,12 +189,6 @@ def _read_base_schedule(release: Mapping[str, Any]) -> DMD2SigmaSchedule | None:
     return DMD2SigmaSchedule.from_metadata(release)
 
 
-def _resolve_component_quant_config(quant_config, component: str):
-    if hasattr(quant_config, "resolve"):
-        return quant_config.resolve(component)
-    return quant_config
-
-
 def _minimax_h3_post_process(output, output_type: str = "np"):
     """Convert the joint video/audio output without capturing worker state.
 
@@ -688,12 +696,24 @@ class MiniMaxH3Pipeline(
             )
         self.text_encoder_tp_size = text_encoder_tp_size
         self.text_encoder_group = self._build_text_encoder_group(text_encoder_tp_size)
+        load_text_encoder = rank < text_encoder_tp_size
         self.text_encoder = MiniMaxH3Qwen3VLEncoder(
             os.path.join(model_path, "text_encoder"),
             device=self.device,
-            load_model=rank < text_encoder_tp_size,
+            load_model=load_text_encoder,
             encoder_group=self.text_encoder_group,
+            quant_config=_resolve_minimax_h3_text_encoder_quant_config(od_config.quantization_config),
         )
+        if load_text_encoder:
+            self.weights_sources.append(
+                DiffusersPipelineLoader.ComponentSource(
+                    model_or_path=str(model_path),
+                    subfolder="text_encoder",
+                    revision=od_config.revision,
+                    prefix="text_encoder.",
+                    fall_back_to_pt=False,
+                )
+            )
         stage_components = bool(
             od_config.enable_layerwise_offload or getattr(od_config, "enable_distributed_layerwise_offload", False)
         )
@@ -725,7 +745,7 @@ class MiniMaxH3Pipeline(
         def source_prefix(item: tuple[str, torch.Tensor]) -> str:
             name, _ = item
             prefix = name.partition(".")[0] + "."
-            if prefix in {"transformer.", "transformers_ref."}:
+            if prefix in {"transformer.", "transformers_ref.", "text_encoder."}:
                 return prefix
             raise ValueError(f"unexpected MiniMax-H3 weight {name!r}")
 
@@ -735,17 +755,16 @@ class MiniMaxH3Pipeline(
             if prefix in loaded_prefixes:
                 raise ValueError(f"MiniMax-H3 weight source {prefix!r} is not contiguous")
             loaded_prefixes.add(prefix)
-            transformer = getattr(self, prefix.removesuffix("."))
-            loaded = transformer.load_weights((name[len(prefix) :], tensor) for name, tensor in grouped_weights)
-            transformer.post_load_weights()
+            component = getattr(self, prefix.removesuffix("."))
+            loaded = component.load_weights((name[len(prefix) :], tensor) for name, tensor in grouped_weights)
+            if prefix != "text_encoder.":
+                component.post_load_weights()
             loaded_with_prefix.update(prefix + name for name in loaded)
-        # The text encoder and both VAEs load eagerly in ``__init__`` rather
-        # than through ``weights_sources``. Record them for the runner's strict
-        # missing-parameter check. For the text encoder that report is exact
-        # because its loader raises on any unloaded parameter or partially filled
-        # fused parameter; weaken that and gaps here become invisible. The two
-        # VAEs carry no equivalent guarantee.
-        for component_name in ("text_encoder", "video_vae", "audio_vae"):
+        # Both VAEs load eagerly in ``__init__`` rather than through
+        # ``weights_sources``. The text encoder uses the shared component
+        # loader so online quantization and offload processing follow the same
+        # path as the DiT.
+        for component_name in ("video_vae", "audio_vae"):
             component = getattr(self, component_name)
             loaded_with_prefix.update(f"{component_name}.{name}" for name, _ in component.named_parameters())
         return loaded_with_prefix

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -178,7 +178,16 @@ class DistributedLayerwiseOffloadHook(ModelHook):
                         if m["name"] in self.next_block_parameters
                         else self.next_block_buffers[m["name"]]
                     )
-                    repoint.append((target, dtype, m["offset"], m["numel"], m["shape"]))
+                    repoint.append(
+                        (
+                            target,
+                            dtype,
+                            m["offset"],
+                            m["numel"],
+                            m["shape"],
+                            m["stride"],
+                        )
+                    )
             self._cached_repoint.append(repoint)
 
         # Pre-compute AG output sizes (avoid sum() per layer).
@@ -228,9 +237,15 @@ class DistributedLayerwiseOffloadHook(ModelHook):
                     # parameter as an mmap view avoids a private full-model
                     # copy in every worker.
                     local_t = mmap_transform(local_t)
-                weights_with_local.append((name, t, local_t))
+                stride = local_t.stride()
+                storage_numel = (
+                    0
+                    if local_t.numel() == 0
+                    else 1 + sum((size - 1) * axis_stride for size, axis_stride in zip(local_t.shape, stride))
+                )
+                weights_with_local.append((name, t, local_t, storage_numel, stride))
 
-            total_numel = sum(local.numel() for _, _, local in weights_with_local)
+            total_numel = sum(storage_numel for _, _, _, storage_numel, _ in weights_with_local)
 
             # Equal-sized shards (ceil division) for all_gather_into_tensor
             shard_size = (total_numel + dp_size - 1) // dp_size  # ceil
@@ -242,8 +257,13 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             shard = torch.zeros(shard_size, dtype=dtype, device="cpu")
 
             current_offset = 0
-            for name, original_tensor, local_tensor in weights_with_local:
-                numel = local_tensor.numel()
+            for (
+                name,
+                original_tensor,
+                local_tensor,
+                storage_numel,
+                stride,
+            ) in weights_with_local:
                 if dtype not in dtype_metadata:
                     dtype_metadata[dtype] = []
                 # Offsets remain relative to the FULL flattened buffer
@@ -252,27 +272,48 @@ class DistributedLayerwiseOffloadHook(ModelHook):
                     {
                         "name": name,
                         "offset": current_offset,
-                        "numel": numel,
+                        "numel": storage_numel,
                         "shape": local_tensor.shape,
+                        "stride": stride,
                     }
                 )
 
                 # Copy ONLY the portion within [shard_start, shard_end)
                 overlap_start = max(current_offset, shard_start)
-                overlap_end = min(current_offset + numel, shard_end)
+                overlap_end = min(current_offset + storage_numel, shard_end)
                 if overlap_start < overlap_end:
+                    if local_tensor.is_contiguous():
+                        flat_storage = local_tensor.flatten()
+                    else:
+                        # Online FP8 stores Cutlass weights as transposed views
+                        # (e.g. stride=(1, K)). Flattening such a tensor in
+                        # logical order and later rebuilding it with .view()
+                        # changes its layout and makes scaled_mm reject it.
+                        # Pack the physical storage order and preserve the
+                        # original stride for zero-copy reconstruction.
+                        flat_storage = torch.zeros(
+                            storage_numel,
+                            dtype=dtype,
+                            device=local_tensor.device,
+                        )
+                        physical_view = torch.as_strided(
+                            flat_storage,
+                            size=local_tensor.shape,
+                            stride=stride,
+                        )
+                        physical_view.copy_(local_tensor)
                     src_start = overlap_start - current_offset
                     src_end = overlap_end - current_offset
                     dst_start = overlap_start - shard_start
                     dst_end = overlap_end - shard_start
-                    shard[dst_start:dst_end].copy_(local_tensor.flatten()[src_start:src_end])
+                    shard[dst_start:dst_end].copy_(flat_storage[src_start:src_end])
 
                 # Replace original tensor with placeholder (frees CPU storage)
                 set_tensor_storage(
                     original_tensor,
                     make_offload_placeholder(original_tensor),
                 )
-                current_offset += numel
+                current_offset += storage_numel
 
             if pin_memory:
                 shard = shard.pin_memory()
@@ -365,10 +406,14 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             self._shared_slot_group[slot] = self._group_id
 
         # Re-point using cached metadata (avoids per-layer dict lookups).
-        for target, dtype, offset, numel, shape in self._cached_repoint[slot]:
+        for target, dtype, offset, numel, shape, stride in self._cached_repoint[slot]:
             set_tensor_storage(
                 target,
-                gpu_weights[dtype][offset : offset + numel].view(shape),
+                torch.as_strided(
+                    gpu_weights[dtype][offset : offset + numel],
+                    size=shape,
+                    stride=stride,
+                ),
             )
 
     def get_weights(self, slot: int) -> dict[torch.dtype, torch.Tensor] | None:
@@ -577,7 +622,11 @@ class PinnedResidentLayerGroup:
                 for meta in metas:
                     set_tensor_storage(
                         targets[meta["name"]],
-                        gpu_buffer[meta["offset"] : meta["offset"] + meta["numel"]].view(meta["shape"]),
+                        torch.as_strided(
+                            gpu_buffer[meta["offset"] : meta["offset"] + meta["numel"]],
+                            size=meta["shape"],
+                            stride=meta["stride"],
+                        ),
                     )
 
         current_omni_platform.current_stream().wait_event(ready)

vllm_omni/diffusion/offloader/layerwise_backend.py

@@ -139,29 +139,57 @@ class LayerwiseOffloadHook(ModelHook):
             weights_with_local = []
             for name, t in name2weights.items():
                 local_t = t.to_local() if hasattr(t, "to_local") else t
-                weights_with_local.append((name, t, local_t))
-            total_numel = sum(local.numel() for _, _, local in weights_with_local)
+                stride = local_t.stride()
+                storage_numel = (
+                    0
+                    if local_t.numel() == 0
+                    else 1 + sum((size - 1) * axis_stride for size, axis_stride in zip(local_t.shape, stride))
+                )
+                weights_with_local.append((name, t, local_t, storage_numel, stride))
+            total_numel = sum(storage_numel for _, _, _, storage_numel, _ in weights_with_local)
             cpu_tensor = torch.empty(total_numel, dtype=dtype, device="cpu", pin_memory=pin_memory)
 
             current_offset = 0
-            for name, original_tensor, local_tensor in weights_with_local:
-                numel = local_tensor.numel()
-                cpu_tensor[current_offset : current_offset + numel].copy_(local_tensor.flatten())
+            for (
+                name,
+                original_tensor,
+                local_tensor,
+                storage_numel,
+                stride,
+            ) in weights_with_local:
+                if local_tensor.is_contiguous():
+                    flat_storage = local_tensor.flatten()
+                else:
+                    # Cutlass FP8 weights use a transposed physical layout.
+                    # Preserve it across the flattened CPU staging buffer.
+                    flat_storage = torch.zeros(
+                        storage_numel,
+                        dtype=dtype,
+                        device=local_tensor.device,
+                    )
+                    physical_view = torch.as_strided(
+                        flat_storage,
+                        size=local_tensor.shape,
+                        stride=stride,
+                    )
+                    physical_view.copy_(local_tensor)
+                cpu_tensor[current_offset : current_offset + storage_numel].copy_(flat_storage)
                 if dtype not in dtype_metadata:
                     dtype_metadata[dtype] = []
                 dtype_metadata[dtype].append(
                     {
                         "name": name,
                         "offset": current_offset,
-                        "numel": numel,
+                        "numel": storage_numel,
                         "shape": local_tensor.shape,
+                        "stride": stride,
                     }
                 )
 
                 LayerwiseOffloadHook._set_tensor_storage(
                     original_tensor, LayerwiseOffloadHook._make_offload_placeholder(original_tensor)
                 )
-                current_offset += numel
+                current_offset += storage_numel
 
             dtype_cpu_flattened_weights[dtype] = cpu_tensor
 
@@ -210,7 +238,11 @@ class LayerwiseOffloadHook(ModelHook):
 
                 LayerwiseOffloadHook._set_tensor_storage(
                     target_param_or_buf,
-                    gpu_weight[metadata["offset"] : metadata["offset"] + metadata["numel"]].view(metadata["shape"]),
+                    torch.as_strided(
+                        gpu_weight[metadata["offset"] : metadata["offset"] + metadata["numel"]],
+                        size=metadata["shape"],
+                        stride=metadata["stride"],
+                    ),
                 )
 
         self._prefetch_done = evt

vllm_omni/quantization/__init__.py

@@ -11,7 +11,7 @@ Adds per-component quantization for multi-stage models.
     config = build_quant_config({"transformer": {"method": "fp8"}, "vae": None})
 """
 
-from .component_config import ComponentQuantizationConfig
+from .component_config import ComponentQuantizationConfig, resolve_component_quant_config
 from .factory import SUPPORTED_QUANTIZATION_METHODS, build_quant_config, register_quantization_override
 from .inc_config import OmniINCConfig
 
@@ -23,6 +23,7 @@ from .inc_config import OmniINCConfig
 __all__ = [
     "build_quant_config",
     "ComponentQuantizationConfig",
+    "resolve_component_quant_config",
     "OmniINCConfig",
     "SUPPORTED_QUANTIZATION_METHODS",
     "register_quantization_override",

vllm_omni/quantization/component_config.py

@@ -29,6 +29,21 @@ if TYPE_CHECKING:
 PRE_QUANTIZED_METHODS: frozenset[str] = frozenset({"modelopt", "modelopt_fp4", "modelopt_mxfp8", "modelopt_mixed"})
 
 
+def resolve_component_quant_config(
+    quant_config: QuantizationConfig | None,
+    component: str,
+) -> QuantizationConfig | None:
+    """Resolve one pipeline component from a global or component config.
+
+    A plain config is global and therefore applies unchanged to every
+    quantization-aware component. Only ``ComponentQuantizationConfig`` narrows
+    the scope through its explicit prefix map.
+    """
+    if isinstance(quant_config, ComponentQuantizationConfig):
+        return quant_config.resolve(component)
+    return quant_config
+
+
 def resolve_encoder_quant_config(
     quant_config: QuantizationConfig | None,
 ) -> QuantizationConfig | None: