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

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