[Bugfix][Ascend NPU][Diffusion] Load over-wide unquantized fallback weights straight into host memory under DLOFix/mxfp online offload after quant (#7009)

Signed-off-by: brandneway <gyuan4892@gmail.com>

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -353,9 +353,14 @@ modes can resolve the same `None` value to different schedules.
   Cache-DiT on any selected component. Prompt-embedding cache is also
   incompatible when the text encoder uses `allgather`. Rank-local cache hits
   could otherwise make ranks enter different weight collectives.
-- Per-tensor online FP8 linears use the ordinary loader and can run with either
-  DiT transfer path. With DiT `allgather`, every rank temporarily materializes
-  the complete FP8 model in host memory before DLO retains only its shard.
+- Per-tensor online FP8, INT8, and MXFP8 linears use the ordinary loader and
+  can run with either DiT transfer path. With DiT `allgather`, every rank
+  temporarily materializes the complete quantized model in host memory before
+  DLO retains only its shard. Layers too wide for the NPU quantization kernel
+  (`npu_quant_matmul` rejects an output dimension past 65535) stay
+  unquantized and load straight into host memory under offload-after-quant —
+  their runtime layout is a plain contiguous bf16 weight, the same layout DLO
+  shards on the ordinary path, so they are allowed under `allgather` too.
   Other online quantization methods require rank-local transfer for the
   affected component until their runtime layouts are validated.
 - Resident leading layers require DiT `rank-local` transfer and a model

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -1241,6 +1241,169 @@ def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch):
         loader.load_model(load_device="cpu")
 
 
+def test_dlo_allgather_allows_unquantized_host_fallback():
+    """Host-loaded unquantized fallback layers are plain contiguous bf16 —
+    the same runtime layout DLO already shards on the ordinary path — so the
+    allowlist must not reject them."""
+    from vllm_omni.quantization.int8_config import UnquantizedHostLinearMethod
+
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = object.__new__(UnquantizedHostLinearMethod)
+
+    assert DiffusersPipelineLoader._unsupported_dlo_allgather_online_quant_methods(model) == ()
+
+
+def test_dlo_load_model_keeps_host_fallback_on_cpu_through_post_load_sweep(monkeypatch, mocker):
+    """Through load_model(): DLO + online quant must build the model inside
+    load_unquantizable_fallback_on_cpu(), and the over-wide fallback weight
+    must survive the post-load sweep (the _process_weights_after_loading pass
+    plus model.to("cpu")) as the same host tensor it was loaded into — never
+    bounced to the accelerator."""
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.quantization import int8_config
+    from vllm_omni.quantization.int8_config import (
+        NPU_QUANT_MATMUL_MAX_OUT_FEATURES,
+        DiffusionInt8Config,
+        NPUInt8OnlineLinearMethod,
+        UnquantizedHostLinearMethod,
+    )
+
+    # Same stand-in as TestHostFallbackLoading: the eager fallback delegates to
+    # UnquantizedLinearMethod, which reads the TP group.
+    mock_group = mocker.Mock()
+    mock_group.rank_in_group = 0
+    mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1)
+    mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0)
+    mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group)
+
+    od_config = _make_dlo_online_quant_config()
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+
+    out_features = NPU_QUANT_MATMUL_MAX_OUT_FEATURES + 1
+    built_with_ctx: list[bool] = []
+
+    def copy_loader(param, loaded_weight, *args, **kwargs):
+        param.data.copy_(loaded_weight)
+
+    def build_model(*_args, **_kwargs):
+        built_with_ctx.append(int8_config._LOAD_UNQUANTIZABLE_FALLBACK_ON_CPU.get())
+        quant_config = DiffusionInt8Config(is_checkpoint_int8_serialized=False, activation_scheme="dynamic")
+        method = NPUInt8OnlineLinearMethod(quant_config)
+        layer = nn.Module()
+        layer.quant_method = method
+        method.create_weights(
+            layer,
+            input_size_per_partition=8,
+            output_partition_sizes=[out_features],
+            input_size=8,
+            output_size=out_features,
+            params_dtype=torch.bfloat16,
+            weight_loader=copy_loader,
+        )
+        model = nn.Module()
+        model.transformer = layer
+        return model
+
+    def load_weights(_model, **_kwargs):
+        layer = _model.transformer
+        loaded = torch.arange(out_features * 8, dtype=torch.float32).reshape(out_features, 8).to(torch.bfloat16)
+        layer.weight.weight_loader(layer.weight, loaded)
+
+    loader._init_from_load_format = build_model  # type: ignore[method-assign]
+    loader.load_weights = load_weights  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(
+            None,
+            "online quantization requires the ordinary loader",
+        ),
+    )
+    # Diffusion DiT models have no vLLM Attention layers, so the upstream
+    # layerwise finalize is a no-op here; stub it to keep the real
+    # _process_weights_after_loading sweep CPU-only.
+    monkeypatch.setattr(
+        "vllm.model_executor.model_loader.reload.layerwise.finalize_layerwise_processing",
+        lambda *_args, **_kwargs: None,
+    )
+
+    model = loader.load_model(load_device="cpu", device=torch.device("cpu"))
+
+    # Construction ran inside load_unquantizable_fallback_on_cpu(), so the
+    # over-wide layer took the host-loading fallback...
+    assert built_with_ctx == [True]
+    assert type(model.transformer.quant_method) is UnquantizedHostLinearMethod
+    # ...and the real post-load sweep skipped it via the fully-loaded flag and
+    # model.to("cpu") left the loaded host tensor untouched.
+    assert model.transformer._already_called_process_weights_after_loading
+    assert model.transformer.weight.device.type == "cpu"
+    expected = torch.arange(out_features * 8, dtype=torch.float32).reshape(out_features, 8).to(torch.bfloat16)
+    assert torch.equal(model.transformer.weight, expected)
+
+
+@pytest.mark.parametrize(("offload_after_quant", "ctx_expected"), [(True, True), (False, False)])
+def test_hsdp_enters_host_fallback_context_only_when_offloading_after_quant(mocker, offload_after_quant, ctx_expected):
+    """The HSDP path must apply the same host-fallback bound as the ordinary
+    path: with a quant config it initializes on the accelerator, so the
+    load_unquantizable_fallback_on_cpu() context is what keeps over-wide
+    fallback weights off the device before sharding."""
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+    from vllm_omni.quantization import int8_config
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=2,
+        ),
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    loader.quant_config = object()
+
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    ctx_seen: list[bool] = []
+
+    def build_model(*_args, **_kwargs):
+        ctx_seen.append(int8_config._LOAD_UNQUANTIZABLE_FALLBACK_ON_CPU.get())
+        return model
+
+    loader._init_from_load_format = build_model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: None  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: None  # type: ignore[method-assign]
+    mocker.patch.object(
+        loader_mod.ModuleDiscovery,
+        "discover",
+        return_value=PipelineModules(
+            dits=[model.transformer],
+            dit_names=["transformer"],
+            vaes=[],
+            encoders=[],
+            encoder_names=[],
+            resident_modules=[],
+            resident_names=[],
+        ),
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.quantization.hsdp_fp8.prepare_fp8_layers_for_fsdp",
+        side_effect=lambda _model: None,
+    )
+    mocker.patch.object(
+        loader_mod,
+        "apply_hsdp_to_model",
+        side_effect=lambda *_args, **_kwargs: None,
+    )
+
+    loader._load_model_with_hsdp(torch.device("cpu"), offload_after_quant=offload_after_quant)
+
+    assert ctx_seen == [ctx_expected]
+
+
 def test_dlo_allgather_online_mxfp8_uses_ordinary_loader(monkeypatch):
     mxfp8_config = pytest.importorskip("vllm_omni.quantization.mxfp8_config")
     NPUMxfp8OnlineLinearMethod = mxfp8_config.NPUMxfp8OnlineLinearMethod

tests/diffusion/quantization/test_int8_config.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for Int8 quantization config."""
 
 import pytest
@@ -12,7 +12,7 @@ from vllm_omni.platforms import current_omni_platform
 from vllm_omni.quantization import build_quant_config
 from vllm_omni.quantization.factory import SUPPORTED_QUANTIZATION_METHODS
 
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 npu_available = pytest.mark.skipif(not current_omni_platform.is_npu(), reason="NPU platform not available.")
 
@@ -248,6 +248,124 @@ class TestNPUQuantMatmulShapeFallback:
         assert layer.quant_method is method
 
 
+class TestHostFallbackLoading:
+    """Over-wide fallback layers load straight into host memory when the loader
+    will offload the model back to CPU after online quantization (DLO).
+
+    Without the context the eager path is preserved: the fallback weight is a
+    real tensor created at construction time.
+    """
+
+    @pytest.fixture(autouse=True)
+    def _mock_tp(self, mocker):
+        # Same stand-in as TestNPUQuantMatmulShapeFallback: the eager fallback
+        # delegates to UnquantizedLinearMethod, which reads the TP group.
+        mock_group = mocker.Mock()
+        mock_group.rank_in_group = 0
+        mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1)
+        mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0)
+        mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group)
+
+    @pytest.fixture
+    def quant_config(self):
+        from vllm_omni.quantization.int8_config import DiffusionInt8Config
+
+        return DiffusionInt8Config(is_checkpoint_int8_serialized=False, activation_scheme="dynamic")
+
+    def _make_over_wide_layer(self, quant_config, weight_loader):
+        from vllm_omni.quantization.int8_config import (
+            NPU_QUANT_MATMUL_MAX_OUT_FEATURES,
+            NPUInt8OnlineLinearMethod,
+        )
+
+        out_features = NPU_QUANT_MATMUL_MAX_OUT_FEATURES + 1
+        method = NPUInt8OnlineLinearMethod(quant_config)
+        layer = Module()
+        layer.quant_method = method
+        method.create_weights(
+            layer,
+            input_size_per_partition=8,
+            output_partition_sizes=[out_features],
+            input_size=8,
+            output_size=out_features,
+            params_dtype=torch.bfloat16,
+            weight_loader=weight_loader,
+        )
+        return layer, out_features
+
+    @staticmethod
+    def _loaded_values(out_features):
+        return torch.arange(out_features * 8, dtype=torch.float32).reshape(out_features, 8).to(torch.bfloat16)
+
+    def test_fallback_without_context_keeps_eager_weights(self, quant_config):
+        layer, _ = self._make_over_wide_layer(quant_config, lambda param, loaded_weight, *a, **k: None)
+
+        assert type(layer.quant_method) is UnquantizedLinearMethod
+        assert layer.weight.device.type != "meta"
+
+    def test_fallback_under_context_defers_weight_to_meta(self, quant_config):
+        from vllm_omni.quantization.int8_config import (
+            UnquantizedHostLinearMethod,
+            load_unquantizable_fallback_on_cpu,
+        )
+
+        with load_unquantizable_fallback_on_cpu():
+            layer, _ = self._make_over_wide_layer(quant_config, lambda param, loaded_weight, *a, **k: None)
+
+        assert isinstance(layer.quant_method, UnquantizedHostLinearMethod)
+        assert layer.quant_method.uses_meta_device
+        assert layer.weight.device.type == "meta"
+
+    def test_fallback_weight_loads_straight_to_cpu(self, quant_config):
+        from vllm_omni.quantization.int8_config import load_unquantizable_fallback_on_cpu
+
+        def copy_loader(param, loaded_weight, *args, **kwargs):
+            param.data.copy_(loaded_weight)
+
+        with load_unquantizable_fallback_on_cpu():
+            layer, out_features = self._make_over_wide_layer(quant_config, copy_loader)
+
+        loaded = self._loaded_values(out_features)
+        layer.weight.weight_loader(layer.weight, loaded)
+
+        assert layer.weight.device.type == "cpu"
+        assert torch.equal(layer.weight, loaded)
+        # Marks the post-load sweep to skip the layer instead of bouncing the
+        # host tensor to the accelerator and back for a no-op process call.
+        assert layer._already_called_process_weights_after_loading
+
+    def test_chunked_load_marks_completion_only_when_full(self, quant_config):
+        from vllm_omni.quantization.int8_config import load_unquantizable_fallback_on_cpu
+
+        state = {"offset": 0}
+
+        def half_loader(param, loaded_weight, *args, **kwargs):
+            start = state["offset"]
+            param.data[start : start + loaded_weight.shape[0]].copy_(loaded_weight)
+            state["offset"] += loaded_weight.shape[0]
+
+        with load_unquantizable_fallback_on_cpu():
+            layer, out_features = self._make_over_wide_layer(quant_config, half_loader)
+
+        loaded = self._loaded_values(out_features)
+        halves = loaded.chunk(2, dim=0)
+        layer.weight.weight_loader(layer.weight, halves[0])
+        assert not getattr(layer, "_already_called_process_weights_after_loading", False)
+
+        layer.weight.weight_loader(layer.weight, halves[1])
+        assert layer._already_called_process_weights_after_loading
+        assert torch.equal(layer.weight, loaded)
+
+    def test_context_does_not_leak(self, quant_config):
+        from vllm_omni.quantization.int8_config import load_unquantizable_fallback_on_cpu
+
+        with load_unquantizable_fallback_on_cpu():
+            pass
+        layer, _ = self._make_over_wide_layer(quant_config, lambda param, loaded_weight, *a, **k: None)
+
+        assert type(layer.quant_method) is UnquantizedLinearMethod
+
+
 class TestOffloadAfterQuant:
     def test_only_methods_advertising_the_capability_are_asked(self):
         from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader

tests/diffusion/quantization/test_mxfp8_config.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for MXFP8 quantization config and linear method dispatch.
 
 Coverage:
@@ -9,11 +9,13 @@ Coverage:
 - Weight / scale shape-transform arithmetic from process_weights_after_loading (CPU)
 - build_quant_config integration
 - MXFP8_QUANT_CONFIG structure as the auto-detection contract
+- Offload-after-quant capability of the shared online-method mixin (CPU)
 """
 
 import pytest
 import torch
 from pytest_mock import MockerFixture
+from torch.nn import Module
 from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
 
 from vllm_omni.platforms import current_omni_platform
@@ -335,3 +337,99 @@ def test_supported_methods_include_mxfp8():
     from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
 
     assert "mxfp8" in SUPPORTED_QUANTIZATION_METHODS
+
+
+# ---------------------------------------------------------------------------
+# Offload-after-quant capability (CPU, no NPU ops)
+# ---------------------------------------------------------------------------
+
+
+class TestOffloadAfterQuant:
+    """Online MXFP methods advertise ``supports_offload_after_quant`` so the
+    diffusion loader can stream each layer back to host memory right after it
+    is quantized, instead of accumulating the whole model on the accelerator
+    until the final ``model.to("cpu")``.
+
+    The capability lives on ``_LazyWeightMixin``, which is shared by all
+    online MXFP8/MXFP4 methods, so pinning it here covers mxfp4 too.
+    """
+
+    @pytest.fixture(autouse=True)
+    def _mock_tp(self, mocker: MockerFixture):
+        # create_weights resolves the TP group while computing per-partition
+        # shapes; these tests never initialize torch.distributed.
+        mock_group = mocker.Mock()
+        mock_group.rank_in_group = 0
+        mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1)
+        mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0)
+        mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group)
+
+    def _make_method(self):
+        from vllm_omni.quantization.mxfp8_config import (
+            DiffusionMXFP8Config,
+            NPUMxfp8OnlineLinearMethod,
+        )
+
+        return NPUMxfp8OnlineLinearMethod(DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=False))
+
+    def test_only_methods_advertising_the_capability_are_asked(self):
+        from vllm_omni.diffusion.model_loader.diffusers_loader import (
+            DiffusersPipelineLoader,
+        )
+
+        class MetaDeviceOnlyMethod:
+            """Stands in for upstream online FP8: lazy, but no offload hook."""
+
+            uses_meta_device = True
+
+        model = torch.nn.Sequential(Module(), Module())
+        lazy_mxfp8 = self._make_method()
+        model[0].quant_method = lazy_mxfp8
+        model[1].quant_method = MetaDeviceOnlyMethod()
+
+        marked = DiffusersPipelineLoader._request_offload_after_quant(model)
+
+        assert marked == 1
+        assert lazy_mxfp8._offload_after_quant
+
+    def _load_layer_in_one_chunk(self, method, layer):
+        method.create_weights(
+            layer,
+            input_size_per_partition=8,
+            output_partition_sizes=[4],
+            input_size=8,
+            output_size=4,
+            params_dtype=torch.bfloat16,
+            weight_loader=lambda param, loaded_weight, *args, **kwargs: param.copy_(loaded_weight),
+        )
+        full_weight = torch.randn(4, 8, dtype=torch.bfloat16)
+        layer.weight.weight_loader(layer.weight, full_weight)
+
+    def test_finished_layer_returns_to_cpu_when_enabled(self, mocker: MockerFixture):
+        method = self._make_method()
+        method.enable_offload_after_quant()
+        mocker.patch.object(method, "process_weights_after_loading")
+
+        layer = Module()
+        moved_to = []
+        layer.to = lambda device, *args, **kwargs: moved_to.append(device)  # type: ignore[method-assign]
+
+        self._load_layer_in_one_chunk(method, layer)
+
+        method.process_weights_after_loading.assert_called_once_with(layer)
+        assert layer._already_called_process_weights_after_loading
+        assert moved_to == ["cpu"]
+
+    def test_finished_layer_stays_put_by_default(self, mocker: MockerFixture):
+        method = self._make_method()
+        mocker.patch.object(method, "process_weights_after_loading")
+
+        layer = Module()
+        moved_to = []
+        layer.to = lambda device, *args, **kwargs: moved_to.append(device)  # type: ignore[method-assign]
+
+        self._load_layer_in_one_chunk(method, layer)
+
+        method.process_weights_after_loading.assert_called_once_with(layer)
+        assert layer._already_called_process_weights_after_loading
+        assert moved_to == []

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -679,10 +679,21 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
         with set_default_torch_dtype(self.od_config.dtype):
             if self.parallel_config.use_hsdp:
                 model = self._load_model_with_hsdp(
-                    target_device=device, load_format=load_format, custom_pipeline_name=custom_pipeline_name
+                    target_device=device,
+                    load_format=load_format,
+                    custom_pipeline_name=custom_pipeline_name,
+                    offload_after_quant=offload_after_quant,
                 )
             else:
-                model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False)
+                # The model is headed back to host memory right after online
+                # quantization, so over-wide NPU-unquantizable fallback weights
+                # load straight into host memory instead of round-tripping
+                # through the accelerator (~24 GiB startup peak on MiniMax H3).
+                from vllm_omni.quantization.int8_config import load_unquantizable_fallback_on_cpu
+
+                fallback_ctx = load_unquantizable_fallback_on_cpu() if offload_after_quant else contextlib.nullcontext()
+                with fallback_ctx:
+                    model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False)
 
                 resolved_offload = resolve_offload(self.od_config)
                 distributed_offload = resolved_offload.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
@@ -721,8 +732,9 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
                     if unsupported_methods:
                         raise ValueError(
                             "DLO+AllGather supports online quantization only for "
-                            "per-tensor FP8, INT8, and MXFP8 linears; unsupported "
-                            f"online methods: {', '.join(sorted(unsupported_methods))}. "
+                            "per-tensor FP8, INT8, and MXFP8 linears "
+                            "(host-loaded unquantized fallback layers are also allowed); "
+                            f"unsupported online methods: {', '.join(sorted(unsupported_methods))}. "
                             "Use rank-local transfer for the affected component or "
                             "disable online quantization."
                         )
@@ -930,6 +942,11 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
         scale, packing, or aliasing layouts (e.g. dual-scale fp4 pairs,
         swizzled or NZ hardware formats) and remain fail-closed until
         validated.
+
+        ``UnquantizedHostLinearMethod`` is also allowed: it backs layers too
+        wide for npu_quant_matmul, loads their weights straight into host
+        memory, and its runtime layout is a plain contiguous bf16 weight —
+        identical to the ordinary unquantized path DLO already shards.
         """
         from vllm.model_executor.layers.quantization.online.fp8 import (
             Fp8PerTensorOnlineLinearMethod,
@@ -938,6 +955,7 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
         from vllm_omni.quantization.int8_config import (
             Int8OnlineLinearMethod,
             NPUInt8OnlineLinearMethod,
+            UnquantizedHostLinearMethod,
         )
 
         try:
@@ -959,6 +977,7 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
             Fp8PerTensorOnlineLinearMethod,
             Int8OnlineLinearMethod,
             NPUInt8OnlineLinearMethod,
+            UnquantizedHostLinearMethod,
             *mxfp8_online_methods,
         )
 
@@ -1208,6 +1227,7 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
         target_device: torch.device,
         load_format: str = "default",
         custom_pipeline_name: str | type[nn.Module] | None = None,
+        offload_after_quant: bool = False,
     ) -> nn.Module:
         """Load model with HSDP sharding for inference.
 
@@ -1231,40 +1251,54 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
         # mapping (QKV fusion, etc.).
         if load_format == "diffusers":
             raise ValueError("HSDP is not supported with the diffusers adapter load format")
-        model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True)
-        world_size = 1
-        rank = 0
-        if torch.distributed.is_initialized():
-            world_size = torch.distributed.get_world_size()
-            rank = torch.distributed.get_rank()
+        # Same host-fallback bound as the ordinary path: with a quant config the
+        # HSDP path initializes on the accelerator (hsdp_defer_to_cpu=False), so
+        # over-wide NPU-unquantizable fallback weights would otherwise round-trip
+        # through the device before apply_hsdp_to_model shards them. Loading them
+        # straight into host memory keeps the load-time device peak at the
+        # quantizable layers alone; sharding then distributes the fallback like
+        # any other parameter. Broadcast loading excludes online quantization
+        # already, so the context below only ever matters on the ordinary
+        # per-rank branch -- but spanning both is harmless.
+        from vllm_omni.quantization.int8_config import load_unquantizable_fallback_on_cpu
 
-        has_online_quant = self._has_online_quant(model) or (
-            self.quant_config is not None and not getattr(self.quant_config, "is_checkpoint_quantized", False)
-        )
-        enable_broadcast = bool(getattr(self.od_config, "enable_broadcast_weight_load", False)) and world_size > 1
+        fallback_ctx = load_unquantizable_fallback_on_cpu() if offload_after_quant else contextlib.nullcontext()
+        with fallback_ctx:
+            model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True)
+            world_size = 1
+            rank = 0
+            if torch.distributed.is_initialized():
+                world_size = torch.distributed.get_world_size()
+                rank = torch.distributed.get_rank()
 
-        if enable_broadcast and has_online_quant:
-            logger.info(
-                "Worker %d: Online quantization detected; falling back to ordinary per-rank weight loading for HSDP",
-                rank,
+            has_online_quant = self._has_online_quant(model) or (
+                self.quant_config is not None and not getattr(self.quant_config, "is_checkpoint_quantized", False)
             )
-            enable_broadcast = False
+            enable_broadcast = bool(getattr(self.od_config, "enable_broadcast_weight_load", False)) and world_size > 1
 
-        if enable_broadcast:
-            if rank == 0:
+            if enable_broadcast and has_online_quant:
+                logger.info(
+                    "Worker %d: Online quantization detected; falling back to ordinary "
+                    "per-rank weight loading for HSDP",
+                    rank,
+                )
+                enable_broadcast = False
+
+            if enable_broadcast:
+                if rank == 0:
+                    self.load_weights(model)
+                    self._maybe_fuse_distilled_lora(model)
+                self._broadcast_model_weights(model, target_device=target_device, src_rank=0)
+                if (
+                    rank != 0
+                    and getattr(self.od_config, "lora_backend", None) in (LoRABackend.DISTILL, "distill")
+                    and getattr(self.od_config, "lora_path", None)
+                    and hasattr(model, "load_lora_weights")
+                ):
+                    setattr(model, "lora_is_fused", True)
+            else:
                 self.load_weights(model)
                 self._maybe_fuse_distilled_lora(model)
-            self._broadcast_model_weights(model, target_device=target_device, src_rank=0)
-            if (
-                rank != 0
-                and getattr(self.od_config, "lora_backend", None) in (LoRABackend.DISTILL, "distill")
-                and getattr(self.od_config, "lora_path", None)
-                and hasattr(model, "load_lora_weights")
-            ):
-                setattr(model, "lora_is_fused", True)
-        else:
-            self.load_weights(model)
-            self._maybe_fuse_distilled_lora(model)
 
         # Quantization methods must finish while parameters are ordinary local
         # tensors. Some post-load transforms use operations (for example,

vllm_omni/quantization/int8_config.py

@@ -7,6 +7,8 @@ on CUDA and NPU platforms.
 """
 
 from collections.abc import Callable
+from contextlib import contextmanager
+from contextvars import ContextVar
 from typing import TYPE_CHECKING, Any, Optional
 
 import torch
@@ -62,6 +64,32 @@ NPU_QUANT_MATMUL_MAX_OUT_FEATURES = 65535
 
 logger = init_logger(__name__)
 
+# Set by the diffusion loader while it constructs a model whose weights will be
+# offloaded back to host memory after online quantization (DLO). Over-wide
+# layers that npu_quant_matmul cannot run then load straight into host memory
+# instead of being built on the accelerator and moved off at the end of
+# loading — on MiniMax H3 the 50 fallback adaln layers are ~24 GiB of bf16,
+# which is the dominant startup memory peak.
+_LOAD_UNQUANTIZABLE_FALLBACK_ON_CPU: ContextVar[bool] = ContextVar(
+    "int8_load_unquantizable_fallback_on_cpu", default=False
+)
+
+
+@contextmanager
+def load_unquantizable_fallback_on_cpu():
+    """Create npu_quant_matmul-incompatible fallback weights on meta and load
+    them straight into host memory.
+
+    Only valid when the whole model returns to the host after loading (the
+    loader's offload-after-quant path); otherwise the fallback weights would
+    stay on CPU while the rest of the model runs on the accelerator.
+    """
+    token = _LOAD_UNQUANTIZABLE_FALLBACK_ON_CPU.set(True)
+    try:
+        yield
+    finally:
+        _LOAD_UNQUANTIZABLE_FALLBACK_ON_CPU.reset(token)
+
 
 def _fell_back_to_unquantized_npu(
     layer: torch.nn.Module,
@@ -90,7 +118,15 @@ def _fell_back_to_unquantized_npu(
         output_size_per_partition,
         NPU_QUANT_MATMUL_MAX_OUT_FEATURES,
     )
-    fallback = UnquantizedLinearMethod()
+    if _LOAD_UNQUANTIZABLE_FALLBACK_ON_CPU.get():
+        logger.info_once(
+            "Loading over-wide unquantized fallback weights straight into host memory "
+            "(offload-after-quant is active); they only reach the accelerator via "
+            "the offload backend's runtime prefetch."
+        )
+        fallback: LinearMethodBase = UnquantizedHostLinearMethod()
+    else:
+        fallback = UnquantizedLinearMethod()
     layer.quant_method = fallback
     fallback.create_weights(
         layer,
@@ -297,6 +333,14 @@ class LazyWeightMixin:
         """
         self._offload_after_quant = True
 
+    def _lazy_load_device(self) -> torch.device:
+        """Device the meta weight materializes on at first load.
+
+        Subclasses that never want the accelerator involved (e.g. the
+        over-wide host fallback) override this to return host memory.
+        """
+        return torch.get_default_device()
+
     def create_weights(
         self,
         layer: torch.nn.Module,
@@ -347,7 +391,7 @@ class LazyWeightMixin:
             # process_weights_after_loading
             target_loaded_numel = layer.weight.numel()
             if layer._loaded_numel == target_loaded_numel:
-                self.process_weights_after_loading(layer)
+                self.process_weights_after_loading(layer)  # type: ignore[attr-defined]
 
                 # Prevent the usual `process_weights_after_loading` call from doing
                 # anything
@@ -378,10 +422,46 @@ class LazyWeightMixin:
             weight_loader=patched_weight_loader,
         )
         # stash the correct device for `patched_weight_loader`
-        layer._load_device = torch.get_default_device()
+        layer._load_device = self._lazy_load_device()
         layer.register_parameter("weight", weight)
 
 
+class UnquantizedHostLinearMethod(LazyWeightMixin, UnquantizedLinearMethod):
+    """Unquantized linear method whose weight loads straight into host memory.
+
+    Layers wider than ``NPU_QUANT_MATMUL_MAX_OUT_FEATURES`` cannot be quantized,
+    so under the ordinary path their bf16 weights materialize on the accelerator
+    at construction time and stay there until the whole model is moved off after
+    loading. When the loader has entered ``load_unquantizable_fallback_on_cpu()``
+    (offload-after-quant, i.e. DLO), that round trip is pure startup peak — the
+    weights end up pinned on the host either way and only visit the accelerator
+    through the offload backend's runtime prefetch.
+
+    This method reuses ``LazyWeightMixin``'s meta + just-in-time materialization
+    wholesale (the weight really is deferred on meta until loading, like the
+    online quant methods, so the loader's meta-aware bookkeeping treats the
+    layer correctly) and only overrides the materialization target to host
+    memory: checkpoint chunks are CPU tensors, so loading stays CPU->CPU with
+    no accelerator round trip.
+    """
+
+    # Offload-after-quant marking is deliberately not advertised even though
+    # the mixin provides the hook: the weight never visits the accelerator,
+    # so there is nothing to return to host.
+    supports_offload_after_quant: bool = False
+
+    def _lazy_load_device(self) -> torch.device:
+        return torch.device("cpu")
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        """Unquantized weights need no post-load transform.
+
+        The mixin calls this when the weight is fully loaded, right before it
+        flags the layer — and the flag is what keeps the loader's post-load
+        sweep from bouncing the host tensor to the accelerator and back.
+        """
+
+
 class Int8LinearMethod(BaseInt8LinearMethod):
     """
     Linear method for Int8

vllm_omni/quantization/mxfp8_config.py

@@ -167,6 +167,21 @@ class _LazyWeightMixin:
 
     uses_meta_device: bool = True
 
+    # This mixin knows when a layer's weight is final, so it can hand the layer
+    # back to the host right there. Loaders that intend to offload the whole
+    # model after loading opt in per layer via ``enable_offload_after_quant``.
+    supports_offload_after_quant: bool = True
+    _offload_after_quant: bool = False
+
+    def enable_offload_after_quant(self) -> None:
+        """Return each layer to host memory as soon as it has been quantized.
+
+        Caps the load-time device footprint at one layer instead of the whole
+        model. A quant method instance belongs to a single layer, so this is not
+        a global switch.
+        """
+        self._offload_after_quant = True
+
     def create_weights(
         self,
         layer: torch.nn.Module,
@@ -202,12 +217,17 @@ class _LazyWeightMixin:
             param = layer.weight
             counter = CopyNumelCounter()
             with counter:
-                res = weight_loader(param, loaded_weight, *args, **kwargs)
+                res = weight_loader(param, loaded_weight, *args, **kwargs)  # type: ignore[misc]
             layer._loaded_numel += counter.copied_numel
 
             if layer._loaded_numel == layer.weight.numel():
-                self.process_weights_after_loading(layer)
+                self.process_weights_after_loading(layer)  # type: ignore[attr-defined]
                 layer._already_called_process_weights_after_loading = True
+
+                # This layer's weight is final, so nothing needs it on the
+                # accelerator until inference.
+                if self._offload_after_quant:
+                    layer.to("cpu")
             return res
 
         weight = ModelWeightParameter(