[diffusion][feature] Add offline SVDQuant W4A4 support (#6162)

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

tests/diffusion/quantization/test_svdquant_config.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+from vllm.model_executor.layers.linear import (
+    LinearBase,
+    UnquantizedLinearMethod,
+)
+
+from vllm_omni.quantization import svdquant_config
+from vllm_omni.quantization.factory import (
+    SUPPORTED_QUANTIZATION_METHODS,
+    build_quant_config,
+)
+from vllm_omni.quantization.svdquant_config import (
+    DiffusionSVDQuantConfig,
+    DiffusionSVDQuantLinearMethod,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_config_from_checkpoint_dict() -> None:
+    config = DiffusionSVDQuantConfig.from_config(
+        {
+            "rank": 32,
+            "precision": "nvfp4",
+            "act_unsigned": False,
+            "modules_to_not_convert": ["blocks.0.adaln_proj.linear"],
+        }
+    )
+
+    assert config.rank == 32
+    assert config.precision == "nvfp4"
+    assert config.modules_to_not_convert == ["blocks.0.adaln_proj.linear"]
+    assert config.get_supported_act_dtypes() == [torch.bfloat16]
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "message"),
+    [
+        ({"rank": 0}, "rank must be positive"),
+        ({"precision": "int4"}, "NVFP4 checkpoints only"),
+        ({"act_unsigned": True}, "unsigned activations"),
+    ],
+)
+def test_config_rejects_unsupported_phase1_options(
+    kwargs: dict[str, object],
+    message: str,
+) -> None:
+    with pytest.raises(ValueError, match=message):
+        DiffusionSVDQuantConfig(**kwargs)
+
+
+def test_factory_builds_svdquant_with_extra_checkpoint_fields() -> None:
+    config = build_quant_config(
+        {
+            "quant_method": "svdquant",
+            "rank": 16,
+            "precision": "nvfp4",
+            "w4a16_modules": ["adaln_proj.linear"],
+            "w4a16_group_size": 64,
+        }
+    )
+
+    assert isinstance(config, DiffusionSVDQuantConfig)
+    assert config.rank == 16
+    assert "svdquant" in SUPPORTED_QUANTIZATION_METHODS
+
+
+def test_skipped_and_non_linear_modules_remain_unquantized() -> None:
+    config = DiffusionSVDQuantConfig(modules_to_not_convert=["blocks.0.adaln_proj.linear"])
+    linear = Mock(spec=LinearBase)
+
+    assert isinstance(
+        config.get_quant_method(linear, "blocks.0.adaln_proj.linear"),
+        UnquantizedLinearMethod,
+    )
+    assert config.get_quant_method(torch.nn.ReLU(), "blocks.0.act") is None
+
+
+def test_linear_method_creates_canonical_partitioned_weights(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(svdquant_config, "_assert_supported", lambda: None)
+    method = DiffusionSVDQuantLinearMethod(DiffusionSVDQuantConfig(rank=32))
+    layer = torch.nn.Module()
+
+    method.create_weights(
+        layer,
+        input_size_per_partition=128,
+        output_partition_sizes=[64, 64],
+        input_size=128,
+        output_size=128,
+        params_dtype=torch.bfloat16,
+    )
+
+    assert layer.qweight.shape == (128, 64)
+    assert layer.qweight.dtype == torch.int8
+    assert layer.wscales.shape == (8, 128)
+    assert layer.wscales.dtype == torch.float8_e4m3fn
+    assert layer.proj_down.shape == (128, 32)
+    assert layer.proj_up.shape == (128, 32)
+    assert layer.smooth_factor.shape == (128,)
+    assert layer.wcscales.shape == (128,)
+    assert layer.wtscale.shape == (1,)
+
+    assert layer.qweight.input_dim == 1
+    assert layer.qweight.output_dim == 0
+    assert layer.wscales.input_dim == 0
+    assert layer.wscales.output_dim == 1
+    assert layer.proj_down.input_dim == 0
+    assert layer.proj_up.output_dim == 0
+    assert layer.wcscales.output_dim == 0
+
+
+def test_linear_method_rejects_misaligned_input_partition(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(svdquant_config, "_assert_supported", lambda: None)
+    method = DiffusionSVDQuantLinearMethod(DiffusionSVDQuantConfig())
+
+    with pytest.raises(ValueError, match="divisible by the block size 16"):
+        method.create_weights(
+            torch.nn.Module(),
+            input_size_per_partition=24,
+            output_partition_sizes=[32],
+            input_size=24,
+            output_size=32,
+            params_dtype=torch.bfloat16,
+        )
+
+
+def test_active_linear_uses_svdquant_method(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(svdquant_config, "_assert_supported", lambda: None)
+    config = DiffusionSVDQuantConfig()
+    linear = Mock(spec=LinearBase)
+
+    method = config.get_quant_method(linear, "blocks.0.attn.out_proj")
+
+    assert isinstance(method, DiffusionSVDQuantLinearMethod)

tests/diffusion/quantization/test_svdquant_linear.py

@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.quantization import svdquant_config as svdquant
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _FakeNvfp4Kernel:
+    def __init__(self, base_weight: torch.Tensor | None = None) -> None:
+        self.base_weight = base_weight
+        self.processed = False
+        self.last_input: torch.Tensor | None = None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        del layer
+        self.processed = True
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        del layer
+        assert bias is None
+        assert self.base_weight is not None
+        self.last_input = x.clone()
+        return torch.mm(x, self.base_weight)
+
+
+def _register_parameter(
+    layer: torch.nn.Module,
+    name: str,
+    value: torch.Tensor,
+) -> None:
+    layer.register_parameter(
+        name,
+        torch.nn.Parameter(value, requires_grad=False),
+    )
+
+
+def test_supports_only_validated_datacenter_blackwell() -> None:
+    assert not svdquant._supports_capability(SimpleNamespace(major=10, minor=0))
+    assert svdquant._supports_capability(SimpleNamespace(major=10, minor=3))
+    assert not svdquant._supports_capability(SimpleNamespace(major=11, minor=0))
+    assert not svdquant._supports_capability(SimpleNamespace(major=12, minor=0))
+    assert not svdquant._supports_capability(None)
+
+
+def test_rejects_nvfp4_backend_with_incompatible_layer_contract(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    svdquant._nvfp4_kernel.cache_clear()
+    monkeypatch.setattr(
+        svdquant,
+        "init_nvfp4_linear_kernel",
+        lambda: _FakeNvfp4Kernel(),
+    )
+
+    with pytest.raises(RuntimeError, match="incompatible with the SVDQuant checkpoint"):
+        svdquant._nvfp4_kernel()
+
+    svdquant._nvfp4_kernel.cache_clear()
+
+
+def test_prepare_weights_uses_existing_nvfp4_layout_contract(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    kernel = _FakeNvfp4Kernel()
+    monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
+    monkeypatch.setattr(svdquant, "_nvfp4_kernel", lambda: kernel)
+    method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
+    layer = torch.nn.Module()
+    qweight = torch.arange(16, dtype=torch.int8).reshape(4, 4)
+    wscales = torch.arange(8, dtype=torch.float32).to(torch.float8_e4m3fn).reshape(2, 4)
+    _register_parameter(layer, "qweight", qweight)
+    _register_parameter(layer, "wscales", wscales)
+    _register_parameter(layer, "wtscale", torch.tensor([2.0], dtype=torch.bfloat16))
+    _register_parameter(
+        layer,
+        "wcscales",
+        torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.bfloat16),
+    )
+
+    method.process_weights_after_loading(layer)
+
+    assert kernel.processed
+    assert not hasattr(layer, "qweight")
+    assert not hasattr(layer, "wscales")
+    assert layer.weight.dtype == torch.uint8
+    assert torch.equal(layer.weight.view(torch.int8), qweight)
+    assert torch.equal(layer.weight_scale, wscales.transpose(0, 1))
+    torch.testing.assert_close(layer.alpha, torch.tensor([2.0]))
+    assert layer.input_global_scale_inv.dtype == torch.float32
+    assert layer.output_channel_scale.shape == (4,)
+
+
+def test_prepare_weights_drops_identity_output_scale(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    kernel = _FakeNvfp4Kernel()
+    monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
+    monkeypatch.setattr(svdquant, "_nvfp4_kernel", lambda: kernel)
+    method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
+    layer = torch.nn.Module()
+    _register_parameter(layer, "qweight", torch.zeros(4, 4, dtype=torch.int8))
+    _register_parameter(
+        layer,
+        "wscales",
+        torch.ones(2, 4, dtype=torch.float8_e4m3fn),
+    )
+    _register_parameter(layer, "wtscale", torch.ones(1, dtype=torch.bfloat16))
+    _register_parameter(layer, "wcscales", torch.ones(4, dtype=torch.bfloat16))
+
+    method.process_weights_after_loading(layer)
+
+    assert layer.output_channel_scale is None
+
+
+def test_compatibility_path_applies_qkv_scale_and_rank_correction(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    base_weight = torch.tensor(
+        [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
+        dtype=torch.bfloat16,
+    )
+    kernel = _FakeNvfp4Kernel(base_weight)
+    monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
+    monkeypatch.setattr(svdquant, "_nvfp4_kernel", lambda: kernel)
+    method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
+
+    layer = torch.nn.Module()
+    layer.smooth_factor = torch.tensor([2.0, 4.0], dtype=torch.bfloat16)
+    layer.proj_down = torch.tensor([[1.0], [2.0]], dtype=torch.bfloat16)
+    layer.proj_up = torch.tensor([[1.0], [2.0], [3.0]], dtype=torch.bfloat16)
+    layer.output_channel_scale = torch.tensor([1.0, 2.0, 4.0], dtype=torch.bfloat16)
+    layer.output_size_per_partition = 3
+    x = torch.tensor([[2.0, 4.0], [4.0, 8.0]], dtype=torch.bfloat16)
+    bias = torch.tensor([1.0, 2.0, 3.0], dtype=torch.bfloat16)
+
+    actual = method.apply(layer, x, bias)
+
+    smoothed = x / layer.smooth_factor
+    expected = torch.mm(smoothed, base_weight)
+    expected = expected * layer.output_channel_scale
+    expected = expected + torch.mm(torch.mm(x, layer.proj_down), layer.proj_up.T)
+    expected = expected + bias
+    torch.testing.assert_close(actual, expected)
+    torch.testing.assert_close(kernel.last_input, smoothed)
+
+
+def test_compatibility_path_rejects_fp16_activations(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
+    method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
+    layer = torch.nn.Module()
+    with pytest.raises(ValueError, match="requires BF16 activations"):
+        method.apply(
+            layer,
+            torch.zeros(1, 16, dtype=torch.float16),
+        )

tests/diffusion/quantization/test_svdquant_tp_loading.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+
+from vllm_omni.quantization import svdquant_config as svdquant
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _linear_without_init(linear_type: type[torch.nn.Module], tp_rank: int):
+    layer = object.__new__(linear_type)
+    torch.nn.Module.__init__(layer)
+    layer.tp_rank = tp_rank
+    return layer
+
+
+def _values(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor:
+    values = torch.arange(torch.Size(shape).numel(), dtype=torch.int64).remainder(127)
+    return values.reshape(shape).to(dtype)
+
+
+def _method(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
+    return svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig(rank=8))
+
+
+def test_column_parallel_loader_shards_svdquant_outputs(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    layer = _linear_without_init(ColumnParallelLinear, tp_rank=1)
+    method = _method(monkeypatch)
+    method.create_weights(
+        layer,
+        input_size_per_partition=128,
+        output_partition_sizes=[32],
+        input_size=128,
+        output_size=64,
+        params_dtype=torch.bfloat16,
+        weight_loader=layer.weight_loader,
+    )
+
+    loaded_tensors = {
+        "qweight": _values((64, 64), torch.int8),
+        "wscales": _values((8, 64), torch.float8_e4m3fn),
+        "proj_up": _values((64, 8), torch.bfloat16),
+        "wcscales": _values((64,), torch.bfloat16),
+    }
+    for name, loaded in loaded_tensors.items():
+        parameter = getattr(layer, name)
+        parameter.weight_loader(parameter, loaded)
+        expected = loaded.narrow(parameter.output_dim, 32, 32)
+        torch.testing.assert_close(parameter, expected)
+
+
+def test_row_parallel_loader_shards_svdquant_inputs(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    layer = _linear_without_init(RowParallelLinear, tp_rank=1)
+    method = _method(monkeypatch)
+    method.create_weights(
+        layer,
+        input_size_per_partition=64,
+        output_partition_sizes=[64],
+        input_size=128,
+        output_size=64,
+        params_dtype=torch.bfloat16,
+        weight_loader=layer.weight_loader,
+    )
+
+    loaded_tensors = {
+        "qweight": _values((64, 64), torch.int8),
+        "wscales": _values((8, 64), torch.float8_e4m3fn),
+        "proj_down": _values((128, 8), torch.bfloat16),
+        "smooth_factor": _values((128,), torch.bfloat16),
+    }
+    for name, loaded in loaded_tensors.items():
+        parameter = getattr(layer, name)
+        parameter.weight_loader(parameter, loaded)
+        shard_size = parameter.shape[parameter.input_dim]
+        expected = loaded.narrow(parameter.input_dim, shard_size, shard_size)
+        torch.testing.assert_close(parameter, expected)
+
+
+def test_qkv_parallel_loader_shards_fused_svdquant_outputs(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    layer = _linear_without_init(QKVParallelLinear, tp_rank=1)
+    layer.total_num_heads = 4
+    layer.total_num_kv_heads = 2
+    layer.num_heads = 2
+    layer.num_kv_heads = 1
+    layer.num_kv_head_replicas = 1
+    layer.head_size = 8
+    layer.v_head_size = 8
+    method = _method(monkeypatch)
+    method.create_weights(
+        layer,
+        input_size_per_partition=128,
+        output_partition_sizes=[16, 8, 8],
+        input_size=128,
+        output_size=64,
+        params_dtype=torch.bfloat16,
+        weight_loader=layer.weight_loader,
+    )
+
+    loaded_tensors = {
+        "qweight": _values((64, 64), torch.int8),
+        "wscales": _values((8, 64), torch.float8_e4m3fn),
+        "proj_up": _values((64, 8), torch.bfloat16),
+        "wcscales": _values((64,), torch.bfloat16),
+    }
+    expected_global_outputs = [slice(16, 32), slice(40, 48), slice(56, 64)]
+    for name, loaded in loaded_tensors.items():
+        parameter = getattr(layer, name)
+        parameter.weight_loader(parameter, loaded)
+        expected = torch.cat(
+            [
+                loaded.narrow(parameter.output_dim, part.start, part.stop - part.start)
+                for part in expected_global_outputs
+            ],
+            dim=parameter.output_dim,
+        )
+        torch.testing.assert_close(parameter, expected)