[diffusion][feature] Add offline SVDQuant W4A4 support

docs/.nav.yml

@@ -54,6 +54,7 @@ nav:
       - BitsAndBytes W4: user_guide/quantization/bitsandbytes.md
       - MXFP8 W8A8: user_guide/quantization/mxfp8.md
       - MXFP4 W4A4: user_guide/quantization/mxfp4.md
+      - SVDQuant W4A4: user_guide/quantization/svdquant.md
       - ModelOpt: user_guide/quantization/modelopt.md
       - AutoRound: user_guide/quantization/autoround.md
       - msModelSlim: user_guide/quantization/msmodelslim.md

docs/user_guide/quantization/svdquant.md

@@ -0,0 +1,50 @@
+# SVDQuant W4A4
+
+## Overview
+
+[SVDQuant](https://arxiv.org/abs/2411.05007) combines four-bit weights
+and activations with a small low-rank branch that corrects part of the
+quantization error. vLLM-Omni consumes an offline-quantized checkpoint;
+it does not calibrate the model while loading.
+
+## Checkpoint contract
+
+Place the following entry in the diffusion transformer's `config.json`:
+
+```json
+{
+  "quantization_config": {
+    "quant_method": "svdquant",
+    "rank": 32,
+    "precision": "nvfp4",
+    "act_unsigned": false,
+    "modules_to_not_convert": []
+  }
+}
+```
+
+For a quantized linear with input size `K`, output size `N`, and correction
+rank `R`, the checkpoint stores:
+
+| Suffix | Shape | dtype |
+| --- | --- | --- |
+| `qweight` | `(N, K / 2)` | `int8` (two packed FP4 values per byte) |
+| `wscales` | `(K / 16, N)` | `float8_e4m3fn` |
+| `proj_down` | `(K, R)` | `bfloat16` |
+| `proj_up` | `(N, R)` | `bfloat16` |
+| `smooth_factor` | `(K,)` | `bfloat16` |
+| `wcscales` | `(N,)` | `bfloat16` |
+| `wtscale` | `(1,)` | `bfloat16` |
+
+`K` must be divisible by 16 on every tensor-parallel rank. Set `wcscales` to
+ones when no per-output correction is needed. Modules listed in
+`modules_to_not_convert` keep their checkpoint precision.
+
+## Runtime support
+
+The compatibility path accepts BF16 inputs and executes an NVFP4 GEMM followed
+by the BF16 rank correction. It supports vLLM's FlashInfer, CUTLASS, and FBGEMM
+NVFP4 tensor layouts; incompatible forced backends fail during model loading.
+SM103 is the currently validated and enabled hardware target. Native fusion of
+the NVFP4 GEMM and rank correction is separate from this checkpoint-loading
+contract.

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),
+        )