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

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

docs/.nav.yml

@@ -53,6 +53,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),
+        )

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)

tests/model_executor/models/test_encoder_quant_config.py

@@ -1,12 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Regression test for #2686: pre-quantized methods must not apply
 quant config to vision / audio encoders.
 
-For modelopt FP8/FP4/MXFP8 checkpoints the Thinker LM is the only
-quantized component.  Vision and audio encoder weights are BF16 with no
-FP8 scale tensors — passing quant_config to them causes FP8 kernels to
-run on BF16 weights, producing garbage embeddings.
+For serialized ModelOpt and SVDQuant checkpoints the target model component is
+the only quantized component. Vision and audio encoder weights are BF16 with no
+serialized scale or correction tensors, so the same config must not be applied
+to them.
 
 **Import note (pytest + ``tests/model_executor`` only):** :mod:`vllm_omni.quantization`
 :file:`__init__` pulls ``factory``/``inc`` and reproduces vLLM CustomOp double
@@ -58,8 +58,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 @pytest.mark.parametrize("method", sorted(PRE_QUANTIZED_METHODS))
 def test_pre_quantized_returns_none(method: str) -> None:
     """visual_quant_config and audio_quant_config must be None for
-    pre-quantized methods (modelopt, modelopt_fp4, modelopt_mxfp8,
-    modelopt_mixed)."""
+    pre-quantized methods (ModelOpt variants and SVDQuant)."""
     mock_config = MagicMock()
     mock_config.get_name.return_value = method
 
@@ -101,5 +100,5 @@ def test_component_config_passed_through() -> None:
 
 def test_pre_quantized_methods_contains_expected() -> None:
     """Guard against accidental removal of a known pre-quantized method."""
-    expected = {"modelopt", "modelopt_fp4", "modelopt_mxfp8", "modelopt_mixed"}
+    expected = {"modelopt", "modelopt_fp4", "modelopt_mxfp8", "modelopt_mixed", "svdquant"}
     assert PRE_QUANTIZED_METHODS == expected

vllm_omni/quantization/component_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
 """Per-component quantization routing for multi-stage models.
 
 Routes get_quant_method() to different configs based on longest-prefix match:
@@ -26,10 +26,11 @@ if TYPE_CHECKING:
     )
 
 
-# Pre-quantized checkpoints (modelopt FP8/FP4/MXFP8/mixed) only quantize the
-# Thinker LM.  Vision and audio encoder weights remain in BF16 with no
-# corresponding scale tensors in the checkpoint.
-PRE_QUANTIZED_METHODS: frozenset[str] = frozenset({"modelopt", "modelopt_fp4", "modelopt_mxfp8", "modelopt_mixed"})
+# These pre-quantized formats require serialized scale or correction tensors
+# that the vision and audio encoder checkpoints do not provide.
+PRE_QUANTIZED_METHODS: frozenset[str] = frozenset(
+    {"modelopt", "modelopt_fp4", "modelopt_mxfp8", "modelopt_mixed", "svdquant"}
+)
 
 
 def resolve_component_quant_config(

vllm_omni/quantization/factory.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
 """Factory for building quantization configs.
 
 build_quant_config() delegates to vLLM's quantization registry.
@@ -55,12 +55,12 @@ def _register_humming_stubs() -> None:
         registry[name] = mod
 
     # wire parent references
-    registry["humming"].config = registry["humming.config"]
-    registry["humming"].dtypes = registry["humming.dtypes"]
-    registry["humming"].layer = registry["humming.layer"]
-    registry["humming"].schema = registry["humming.schema"]
-    registry["humming"].utils = registry["humming.utils"]
-    registry["humming.utils"].weight = registry["humming.utils.weight"]
+    setattr(registry["humming"], "config", registry["humming.config"])
+    setattr(registry["humming"], "dtypes", registry["humming.dtypes"])
+    setattr(registry["humming"], "layer", registry["humming.layer"])
+    setattr(registry["humming"], "schema", registry["humming.schema"])
+    setattr(registry["humming"], "utils", registry["humming.utils"])
+    setattr(registry["humming.utils"], "weight", registry["humming.utils.weight"])
 
     for name, mod in registry.items():
         sys.modules[name] = mod
@@ -125,6 +125,13 @@ def _build_mxfp4_dualscale(**kw: Any) -> QuantizationConfig:
     return DiffusionMXFP4DualScaleMixedConfig(**kw)
 
 
+def _build_svdquant(**kw: Any) -> QuantizationConfig:
+    """Build the serialized SVDQuant diffusion checkpoint loader."""
+    from .svdquant_config import DiffusionSVDQuantConfig
+
+    return DiffusionSVDQuantConfig.from_config(kw)
+
+
 def _build_inc(**kw: Any) -> QuantizationConfig:
     """Lazy import for INC/AutoRound config with checkpoint kwarg normalization."""
     from .inc_config import OmniINCConfig
@@ -145,6 +152,7 @@ _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = {
     "mxfp8": _build_mxfp8,
     "mxfp4": _build_mxfp4,
     "mxfp4_dualscale": _build_mxfp4_dualscale,
+    "svdquant": _build_svdquant,
     "inc": _build_inc,
     "auto-round": _build_inc,
     "auto_round": _build_inc,

vllm_omni/quantization/svdquant_config.py

@@ -0,0 +1,409 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Serialized SVDQuant NVFP4 support for diffusion transformers.
+
+The checkpoint stores NVFP4 weights plus a rank-R correction for each
+quantized linear. The four-bit GEMM uses vLLM's existing NVFP4 kernel
+registry, while the rank correction uses ordinary BF16 matrix multiplication.
+Native SVDQuant fusion is a separate optimization and is not required to load
+or run the checkpoint.
+"""
+
+from __future__ import annotations
+
+import functools
+from typing import TYPE_CHECKING, Any
+
+import torch
+from torch.nn import Parameter
+from vllm.logger import init_logger
+from vllm.model_executor.kernels.linear import (
+    CutlassNvFp4LinearKernel,
+    FbgemmNvFp4LinearKernel,
+    FlashInferB12xNvFp4LinearKernel,
+    FlashInferCudnnNvFp4LinearKernel,
+    FlashInferCuteDslNvFp4LinearKernel,
+    FlashInferCutlassNvFp4LinearKernel,
+    FlashInferTrtllmNvFp4LinearKernel,
+    NvFp4LinearKernel,
+    init_nvfp4_linear_kernel,
+)
+from vllm.model_executor.layers.linear import (
+    LinearBase,
+    LinearMethodBase,
+    UnquantizedLinearMethod,
+)
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    is_layer_skipped,
+)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.utils import set_weight_attrs
+from vllm.platforms import current_platform
+from vllm.platforms.interface import DeviceCapability
+
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization import QuantizationMethods
+
+logger = init_logger(__name__)
+
+_SUPPORTED_CAPABILITIES = {(10, 3)}
+_COMPATIBLE_NVFP4_KERNELS = (
+    CutlassNvFp4LinearKernel,
+    FbgemmNvFp4LinearKernel,
+    FlashInferB12xNvFp4LinearKernel,
+    FlashInferCudnnNvFp4LinearKernel,
+    FlashInferCuteDslNvFp4LinearKernel,
+    FlashInferCutlassNvFp4LinearKernel,
+    FlashInferTrtllmNvFp4LinearKernel,
+)
+
+
+def _supports_capability(capability: DeviceCapability | None) -> bool:
+    return (
+        capability is not None
+        and (
+            capability.major,
+            capability.minor,
+        )
+        in _SUPPORTED_CAPABILITIES
+    )
+
+
+def _assert_supported() -> None:
+    if not current_platform.is_cuda():
+        raise RuntimeError("SVDQuant NVFP4 requires a CUDA device")
+    capability = current_platform.get_device_capability()
+    if not _supports_capability(capability):
+        device = current_platform.device_name
+        sm = capability.to_int() if capability is not None else "unknown"
+        raise RuntimeError(f"SVDQuant NVFP4 is validated on SM103 only; got {device!r} (SM{sm})")
+
+
+@functools.cache
+def _nvfp4_kernel() -> NvFp4LinearKernel:
+    kernel = init_nvfp4_linear_kernel()
+    if not isinstance(kernel, _COMPATIBLE_NVFP4_KERNELS):
+        raise RuntimeError(
+            "The selected vLLM NVFP4 backend is incompatible with the "
+            f"SVDQuant checkpoint layout: {type(kernel).__name__}. Use one of "
+            "the FlashInfer, CUTLASS, or FBGEMM NVFP4 backends."
+        )
+    return kernel
+
+
+class DiffusionSVDQuantConfig(QuantizationConfig):
+    """Configuration for serialized NVFP4 W4A4 plus low-rank correction."""
+
+    def __init__(
+        self,
+        rank: int = 32,
+        precision: str = "nvfp4",
+        act_unsigned: bool = False,
+        modules_to_not_convert: list[str] | None = None,
+    ) -> None:
+        super().__init__()
+        if rank <= 0:
+            raise ValueError(f"SVDQuant rank must be positive, got {rank}")
+        if precision != "nvfp4":
+            raise ValueError(
+                f"Phase 1 SVDQuant supports serialized NVFP4 checkpoints only; got precision={precision!r}"
+            )
+        if act_unsigned:
+            raise ValueError("Phase 1 SVDQuant does not support unsigned activations")
+        self.rank = rank
+        self.precision = precision
+        self.modules_to_not_convert = modules_to_not_convert or []
+
+    def __repr__(self) -> str:
+        return f"DiffusionSVDQuantConfig(rank={self.rank}, precision={self.precision!r})"
+
+    @classmethod
+    def get_name(cls) -> QuantizationMethods:
+        return "svdquant"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+        return [torch.bfloat16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 103
+
+    @classmethod
+    def get_config_filenames(cls) -> list[str]:
+        return ["quantization_config.json"]
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> DiffusionSVDQuantConfig:
+        return cls(
+            rank=config.get("rank", 32),
+            precision=config.get("precision", "nvfp4"),
+            act_unsigned=config.get("act_unsigned", False),
+            modules_to_not_convert=config.get("modules_to_not_convert"),
+        )
+
+    def get_quant_method(
+        self,
+        layer: torch.nn.Module,
+        prefix: str,
+    ) -> QuantizeMethodBase | None:
+        if not isinstance(layer, LinearBase):
+            return None
+        if is_layer_skipped(
+            prefix,
+            self.modules_to_not_convert,
+            self.packed_modules_mapping,
+            skip_with_substr=True,
+        ):
+            return UnquantizedLinearMethod()
+        return DiffusionSVDQuantLinearMethod(self)
+
+
+class DiffusionSVDQuantLinearMethod(LinearMethodBase):
+    """Load and execute a serialized SVDQuant linear layer."""
+
+    def __init__(self, quant_config: DiffusionSVDQuantConfig) -> None:
+        _assert_supported()
+        self.quant_config = quant_config
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: list[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs: Any,
+    ) -> None:
+        del params_dtype
+        weight_loader = extra_weight_attrs.pop(
+            "weight_loader",
+            default_weight_loader,
+        )
+        if input_size_per_partition % 16 != 0:
+            raise ValueError(
+                "SVDQuant NVFP4 requires each input partition to be divisible "
+                f"by the block size 16; got {input_size_per_partition}"
+            )
+        output_size_per_partition = sum(output_partition_sizes)
+        rank = self.quant_config.rank
+
+        qweight = Parameter(
+            torch.empty(
+                output_size_per_partition,
+                input_size_per_partition // 2,
+                dtype=torch.int8,
+            ),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            qweight,
+            {
+                "input_dim": 1,
+                "output_dim": 0,
+                "weight_loader": weight_loader,
+            },
+        )
+
+        wscales = Parameter(
+            torch.empty(
+                input_size_per_partition // 16,
+                output_size_per_partition,
+                dtype=torch.float8_e4m3fn,
+            ),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            wscales,
+            {
+                "input_dim": 0,
+                "output_dim": 1,
+                "weight_loader": weight_loader,
+            },
+        )
+
+        proj_down = Parameter(
+            torch.empty(
+                input_size_per_partition,
+                rank,
+                dtype=torch.bfloat16,
+            ),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            proj_down,
+            {
+                "input_dim": 0,
+                "weight_loader": weight_loader,
+            },
+        )
+
+        proj_up = Parameter(
+            torch.empty(
+                output_size_per_partition,
+                rank,
+                dtype=torch.bfloat16,
+            ),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            proj_up,
+            {
+                "output_dim": 0,
+                "weight_loader": weight_loader,
+            },
+        )
+
+        smooth_factor = Parameter(
+            torch.empty(
+                input_size_per_partition,
+                dtype=torch.bfloat16,
+            ),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            smooth_factor,
+            {
+                "input_dim": 0,
+                "weight_loader": weight_loader,
+            },
+        )
+
+        wcscales = Parameter(
+            torch.ones(
+                output_size_per_partition,
+                dtype=torch.bfloat16,
+            ),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            wcscales,
+            {
+                "output_dim": 0,
+                "weight_loader": weight_loader,
+            },
+        )
+
+        wtscale = Parameter(
+            torch.ones(1, dtype=torch.bfloat16),
+            requires_grad=False,
+        )
+        set_weight_attrs(wtscale, {"weight_loader": default_weight_loader})
+
+        layer.register_parameter("qweight", qweight)
+        layer.register_parameter("wscales", wscales)
+        layer.register_parameter("proj_down", proj_down)
+        layer.register_parameter("proj_up", proj_up)
+        layer.register_parameter("smooth_factor", smooth_factor)
+        layer.register_parameter("wcscales", wcscales)
+        layer.register_parameter("wtscale", wtscale)
+
+        del input_size, output_size
+        layer.output_size_per_partition = output_size_per_partition
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        """Adapt the canonical row-major checkpoint to vLLM's NVFP4 ABI."""
+        qweight = layer.qweight
+        wscales = layer.wscales
+        del layer.qweight
+        del layer.wscales
+        layer.register_parameter(
+            "weight",
+            Parameter(
+                qweight.detach().view(torch.uint8),
+                requires_grad=False,
+            ),
+        )
+        layer.register_parameter(
+            "weight_scale",
+            Parameter(
+                wscales.detach().transpose(0, 1).contiguous(),
+                requires_grad=False,
+            ),
+        )
+
+        layer.register_parameter(
+            "input_global_scale_inv",
+            Parameter(
+                torch.ones(
+                    1,
+                    dtype=torch.float32,
+                    device=layer.weight.device,
+                ),
+                requires_grad=False,
+            ),
+        )
+
+        wtscale = layer.wtscale.detach().to(dtype=torch.float32)
+        del layer.wtscale
+        layer.register_parameter(
+            "alpha",
+            Parameter(wtscale, requires_grad=False),
+        )
+
+        channel_scale = layer.wcscales.detach()
+        del layer.wcscales
+        if torch.all(channel_scale == 1).item():
+            layer.output_channel_scale = None
+        else:
+            layer.register_parameter(
+                "output_channel_scale",
+                Parameter(channel_scale, requires_grad=False),
+            )
+
+        _nvfp4_kernel().process_weights_after_loading(layer)
+        logger.info_once(
+            "SVDQuant NVFP4 is using vLLM's compatibility path; the rank correction is not fused into the GEMM."
+        )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Compute the base NVFP4 GEMM plus the BF16 rank correction."""
+        if x.dtype != torch.bfloat16:
+            raise ValueError(f"SVDQuant NVFP4 requires BF16 activations; got {x.dtype}")
+
+        original_shape = x.shape
+        x_2d = x.reshape(-1, original_shape[-1]).contiguous()
+
+        # The residual branch consumes the original activation. Only the
+        # four-bit base GEMM consumes the smoothed activation.
+        smoothed = x_2d / layer.smooth_factor
+        out = _nvfp4_kernel().apply_weights(
+            layer=layer,
+            x=smoothed,
+            bias=None,
+        )
+
+        channel_scale = getattr(layer, "output_channel_scale", None)
+        if channel_scale is not None:
+            # Fused QKV can store independent Q/K/V outer scales. Apply them
+            # explicitly until a vector-alpha GEMM epilogue is available.
+            out.mul_(channel_scale)
+
+        correction_input = torch.mm(x_2d, layer.proj_down)
+        out = torch.addmm(
+            out,
+            correction_input,
+            layer.proj_up.transpose(0, 1),
+        )
+        if bias is not None:
+            out.add_(bias)
+        return out.reshape(
+            *original_shape[:-1],
+            layer.output_size_per_partition,
+        )
+
+
+__all__ = [
+    "DiffusionSVDQuantConfig",
+    "DiffusionSVDQuantLinearMethod",
+]