diff --git a/docs/.nav.yml b/docs/.nav.yml index 7b4052e9c..1de85696a 100644 --- a/docs/.nav.yml +++ b/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 diff --git a/docs/user_guide/quantization/svdquant.md b/docs/user_guide/quantization/svdquant.md new file mode 100644 index 000000000..4fbd8040d --- /dev/null +++ b/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. diff --git a/tests/diffusion/quantization/test_svdquant_config.py b/tests/diffusion/quantization/test_svdquant_config.py new file mode 100644 index 000000000..adbb88ed0 --- /dev/null +++ b/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) diff --git a/tests/diffusion/quantization/test_svdquant_linear.py b/tests/diffusion/quantization/test_svdquant_linear.py new file mode 100644 index 000000000..bdca88d9d --- /dev/null +++ b/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), + ) diff --git a/tests/diffusion/quantization/test_svdquant_tp_loading.py b/tests/diffusion/quantization/test_svdquant_tp_loading.py new file mode 100644 index 000000000..14138dc1c --- /dev/null +++ b/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) diff --git a/tests/model_executor/models/test_encoder_quant_config.py b/tests/model_executor/models/test_encoder_quant_config.py index 28a1efa71..694e92a02 100644 --- a/tests/model_executor/models/test_encoder_quant_config.py +++ b/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 diff --git a/vllm_omni/quantization/component_config.py b/vllm_omni/quantization/component_config.py index 4d74f7134..bf6f3b6b2 100644 --- a/vllm_omni/quantization/component_config.py +++ b/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( diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py index 245793871..f3b6e8cd2 100644 --- a/vllm_omni/quantization/factory.py +++ b/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, diff --git a/vllm_omni/quantization/svdquant_config.py b/vllm_omni/quantization/svdquant_config.py new file mode 100644 index 000000000..80478361b --- /dev/null +++ b/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", +]