From e0570a08596fb9a85d392531bc07e509c38867dc Mon Sep 17 00:00:00 2001 From: Diya Peng <50388704+dpeng123@users.noreply.github.com> Date: Tue, 14 Jul 2026 00:48:48 -0700 Subject: [PATCH] [Quantization] Add BitsAndBytes W4 online quantization for diffusion transformer (#5037) Signed-off-by: diyapeng <50388704+diyapeng@users.noreply.github.com> Co-authored-by: diyapeng <50388704+diyapeng@users.noreply.github.com> --- docs/.nav.yml | 1 + docs/user_guide/quantization/bitsandbytes.md | 120 ++++++++ .../text_to_image/text_to_image.py | 5 +- .../quantization/test_bitsandbytes_config.py | 258 ++++++++++++++++++ .../test_async_omni_diffusion_config.py | 24 ++ vllm_omni/quantization/bitsandbytes_config.py | 176 ++++++++++++ vllm_omni/quantization/factory.py | 8 + 7 files changed, 590 insertions(+), 2 deletions(-) create mode 100644 docs/user_guide/quantization/bitsandbytes.md create mode 100644 tests/diffusion/quantization/test_bitsandbytes_config.py create mode 100644 vllm_omni/quantization/bitsandbytes_config.py diff --git a/docs/.nav.yml b/docs/.nav.yml index f94907a67..6bf161b1b 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -80,6 +80,7 @@ nav: - Quantized KV Cache: user_guide/quantization/quantized_kvcache.md - FP8 W8A8: user_guide/quantization/fp8.md - Int8 W8A8: user_guide/quantization/int8.md + - BitsAndBytes W4: user_guide/quantization/bitsandbytes.md - MXFP8 W8A8: user_guide/quantization/mxfp8.md - MXFP4 W4A4: user_guide/quantization/mxfp4.md - ModelOpt: user_guide/quantization/modelopt.md diff --git a/docs/user_guide/quantization/bitsandbytes.md b/docs/user_guide/quantization/bitsandbytes.md new file mode 100644 index 000000000..37ccd2471 --- /dev/null +++ b/docs/user_guide/quantization/bitsandbytes.md @@ -0,0 +1,120 @@ +# BitsAndBytes W4 Quantization + +## Overview + +BitsAndBytes 4-bit quantization supports weight-only NF4/FP4 diffusion transformer +inference on CUDA GPUs. It quantizes BF16/FP16 weights at load time; activations +stay in BF16/FP16. No pre-quantized checkpoint is required. + +This is an online (dynamic) path: load the normal HuggingFace checkpoint and +quantize during model loading via the `bitsandbytes` CUDA kernels. + +## Hardware Support + +| Device | Support | +|--------|---------| +| NVIDIA CUDA GPU (SM 75+) | ✅ | +| NVIDIA Blackwell GPU (SM 100+) | ✅ | +| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | +| NVIDIA Ampere GPU (SM 80+) | ✅ | +| AMD ROCm | ❌ | +| Intel XPU | ❌ | +| Ascend NPU | ❌ | + +Legend: `✅` supported, `❌` unsupported. Non-CUDA platforms raise in +`get_quant_method()`; CUDA requires compute capability 7.5+ (`SM 75+`). + +Requires the optional `bitsandbytes` package (`pip install bitsandbytes`). + +## Model Type Support + +### Diffusion Model (Qwen-Image, Wan2.2) + +| Model | HF models | CUDA | Mode | Recommendation | +|-------|-----------|:----:|------|----------------| +| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Online W4 weight-only | All heavy linear layers; sensitive embedders stay BF16 | +| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Not validated | Online W4 weight-only | Compare vs BF16 before enabling | +| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Online W4 weight-only | Validate before enabling in docs | + +Other diffusion models may work if their transformer uses supported linear +layers, but they are not validated in this guide. + +### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS) + +| Model | Scope | Status | Notes | +|-------|-------|--------|-------| +| Qwen3-Omni | Thinker language-model stage | Not validated | Prefer checkpoint-supported ModelOpt FP8 or AutoRound paths | +| Qwen3-TTS | TTS language-model stage | Not validated | No BitsAndBytes TTS stage support is documented | + +### Multi-Stage Diffusion Model (BAGEL, GLM-Image) + +| Model | Scope | Status | Notes | +|-------|-------|--------|-------| +| BAGEL | Stage-specific transformer or DiT module | Not validated | Requires explicit stage routing | +| GLM-Image | Stage-specific transformer or DiT module | Not validated | Requires quality comparison with BF16 | + +## Configuration + +Python API: + +```python +from vllm_omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +omni = Omni(model="Tongyi-MAI/Z-Image-Turbo", quantization="bitsandbytes") + +omni_with_skips = Omni( + model="Tongyi-MAI/Z-Image-Turbo", + quantization_config={ + "method": "bitsandbytes", + "ignored_layers": ["to_out"], + }, +) + +outputs = omni.generate( + "A cup of coffee on the table", + OmniDiffusionSamplingParams(num_inference_steps=50), +) +``` + +CLI: + +```bash +python text_to_image.py --model Tongyi-MAI/Z-Image-Turbo --quantization bitsandbytes +python text_to_image.py --model Tongyi-MAI/Z-Image-Turbo --quantization bitsandbytes --ignored-layers "to_out" +vllm serve Tongyi-MAI/Z-Image-Turbo --omni --quantization bitsandbytes +``` + +## Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `method` | str | - | Quantization method (`"bitsandbytes"`) | +| `quant_type` | str | `"nf4"` | 4-bit data type: `"nf4"` (recommended) or `"fp4"` | +| `compress_statistics` | bool | `True` | Double-quantize block scaling statistics for better accuracy | +| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16/FP16 | + +## Validation and Notes + +On Z-Image-Turbo (single GPU, 1024×1024, 50 steps), BitsAndBytes W4 typically +reduces peak VRAM from roughly 24.5 GiB (BF16) to roughly 17 GiB. Compare output +quality against a BF16 baseline before enabling on new models. + +Multi-GPU tensor parallelism (`tensor_parallel_size` > 1) is not validated for +BitsAndBytes in diffusion models; each rank quantizes its own weight shard +independently. + +If quality regresses, use `ignored_layers` to keep sensitive projections in BF16 +(for example `to_out` or `w2`). + +```python +omni = Omni( + model="Tongyi-MAI/Z-Image-Turbo", + quantization="bitsandbytes", + cache_backend="tea_cache", + cache_config={"rel_l1_thresh": 0.2}, +) +``` + +Only add a new model to the supported table after comparing the BitsAndBytes +output against a BF16 baseline and documenting any required `ignored_layers`. diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 374356ef4..61dd29261 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -202,9 +202,10 @@ def parse_args() -> argparse.Namespace: "--quantization", type=str, default=None, - choices=["fp8", "int8"], + choices=["fp8", "int8", "bitsandbytes"], help="Quantization method for the transformer. " - "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8). " + "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), " + "'int8' (Int8 W8A8), 'bitsandbytes' (NF4 4-bit weight-only, CUDA SM 75+). " "Default: None (no quantization, uses BF16).", ) parser.add_argument( diff --git a/tests/diffusion/quantization/test_bitsandbytes_config.py b/tests/diffusion/quantization/test_bitsandbytes_config.py new file mode 100644 index 000000000..8da2cd521 --- /dev/null +++ b/tests/diffusion/quantization/test_bitsandbytes_config.py @@ -0,0 +1,258 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for BitsAndBytes quantization config.""" + +import importlib.util +import sys +import types + +import pytest +import torch +from pytest_mock import MockerFixture +from torch.nn import Module, Parameter +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + +from vllm_omni.platforms import current_omni_platform +from vllm_omni.quantization import build_quant_config +from vllm_omni.quantization.factory import SUPPORTED_QUANTIZATION_METHODS + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + +cuda_available = pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="GPU platform not available.") + +bitsandbytes_available = pytest.mark.skipif( + importlib.util.find_spec("bitsandbytes") is None, + reason="bitsandbytes package not installed", +) + + +def _ensure_bitsandbytes_importable(monkeypatch: pytest.MonkeyPatch) -> None: + """Inject a stub bitsandbytes package when the optional dep is absent. + + ``mocker.patch("bitsandbytes...")`` imports the target module; without a + stub, core-model CI without bitsandbytes raises ModuleNotFoundError. + """ + if importlib.util.find_spec("bitsandbytes") is not None: + return + bnb = types.ModuleType("bitsandbytes") + bnb_functional = types.ModuleType("bitsandbytes.functional") + bnb.functional = bnb_functional + monkeypatch.setitem(sys.modules, "bitsandbytes", bnb) + monkeypatch.setitem(sys.modules, "bitsandbytes.functional", bnb_functional) + + +def test_bitsandbytes_config_creation(): + """Test that BitsAndBytes config can be created.""" + config = build_quant_config("bitsandbytes") + assert config is not None + assert config.get_name() == "bitsandbytes" + + +def test_bitsandbytes_config_with_custom_params(): + """Test BitsAndBytes config with custom parameters.""" + config = build_quant_config( + "bitsandbytes", + quant_type="fp4", + compress_statistics=False, + ignored_layers=["to_out"], + ) + assert config is not None + assert config.quant_type == "fp4" + assert config.compress_statistics is False + assert "to_out" in config.ignored_layers + + +def test_supported_methods(): + """Test that supported methods list includes bitsandbytes.""" + assert "bitsandbytes" in SUPPORTED_QUANTIZATION_METHODS + + +def test_quantization_integration(): + """Test end-to-end quantization flow through OmniDiffusionConfig.""" + from vllm_omni.diffusion.data import OmniDiffusionConfig + + config = OmniDiffusionConfig(model="test", quantization_config="bitsandbytes") + assert config.quantization_config is not None + assert config.quantization_config.get_name() == "bitsandbytes" + + config2 = OmniDiffusionConfig( + model="test", + quantization_config={ + "method": "bitsandbytes", + "quant_type": "nf4", + "compress_statistics": True, + }, + ) + assert config2.quantization_config is not None + assert config2.quantization_config.get_name() == "bitsandbytes" + assert config2.quantization_config.quant_type == "nf4" + assert config2.quantization_config.compress_statistics is True + + +def test_quantization_dict_not_mutated(): + """Test that passing a dict to quantization_config doesn't mutate it.""" + from vllm_omni.diffusion.data import OmniDiffusionConfig + + original_dict = {"method": "bitsandbytes", "quant_type": "nf4"} + dict_copy = original_dict.copy() + + OmniDiffusionConfig(model="test", quantization_config=original_dict) + + assert original_dict == dict_copy + + +def test_get_quant_method(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): + """Test get_quant_method routing for CUDA.""" + from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod + + config = build_quant_config("bitsandbytes") + + def _fake_init(self, quant_config): + pass + + layer = mocker.Mock(spec=LinearBase) + mocker.patch.object(BnBOnlineLinearMethod, "__init__", _fake_init) + + prefix = "test_layer" + + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True) + method = config.get_quant_method(layer, prefix) + assert isinstance(method, BnBOnlineLinearMethod) + + config.ignored_layers = [prefix] + method = config.get_quant_method(layer, prefix) + assert isinstance(method, UnquantizedLinearMethod) + + +class TestBnBOnlineLinearMethod: + @pytest.fixture + def mock_quant_config(self, mocker): + config = mocker.Mock() + config.quant_type = "nf4" + config.compress_statistics = True + return config + + @pytest.fixture + def mock_deps(self, mocker, monkeypatch: pytest.MonkeyPatch): + _ensure_bitsandbytes_importable(monkeypatch) + mock_qweight = torch.ones((64, 32), dtype=torch.uint8) + mock_quant_state = mocker.Mock() + mock_quant = mocker.patch( + "bitsandbytes.functional.quantize_4bit", + return_value=(mock_qweight, mock_quant_state), + create=True, + ) + # x is (2, 16, 32) -> flattened (32, 32); matmul must return (32, 64) + # so apply can reshape back to (2, 16, 64). + mock_matmul = mocker.patch( + "bitsandbytes.matmul_4bit", + return_value=torch.randn(32, 64), + create=True, + ) + return { + "quant": mock_quant, + "matmul": mock_matmul, + "mock_qweight": mock_qweight, + "mock_quant_state": mock_quant_state, + } + + def test_process_weights_after_loading(self, mock_deps, mock_quant_config, mocker): + from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod + + method = BnBOnlineLinearMethod(mock_quant_config) + layer = Module() + layer.weight = Parameter(torch.randn(64, 32)) + mocker.patch.object(torch.Tensor, "cuda", return_value=layer.weight.data) + + method.process_weights_after_loading(layer) + + mock_deps["quant"].assert_called_once() + assert torch.equal(layer.weight.data, mock_deps["mock_qweight"]) + assert layer.quant_state is mock_deps["mock_quant_state"] + assert layer.bnb_shape == (64, 32) + + def test_apply(self, mock_deps, mock_quant_config): + from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod + + method = BnBOnlineLinearMethod(mock_quant_config) + layer = Module() + layer.weight = Parameter(torch.ones((64, 32), dtype=torch.uint8), requires_grad=False) + layer.quant_state = mock_deps["mock_quant_state"] + + x = torch.randn(2, 16, 32, dtype=torch.float16) + bias = torch.randn(64) + output = method.apply(layer, x, bias) + + mock_deps["matmul"].assert_called_once() + assert output.shape == (2, 16, 64) + assert output.dtype == torch.float16 + + +@pytest.fixture +def quant_config(): + from vllm_omni.quantization.bitsandbytes_config import DiffusionBitsAndBytesConfig + + return DiffusionBitsAndBytesConfig( + quant_type="nf4", + compress_statistics=True, + ) + + +@cuda_available +@bitsandbytes_available +class TestCudaBnBSmoke: + """Smoke tests using real bitsandbytes CUDA kernels.""" + + @pytest.fixture + def real_layer(self): + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter( + torch.randn(128, 64, dtype=torch.float16, device="cuda"), + requires_grad=False, + ) + layer.logical_widths = [128] + layer.input_size_per_partition = 64 + layer.output_size_per_partition = 128 + layer.orig_dtype = torch.float16 + return layer + + def test_real_cuda_quantize_4bit_shape_contract(self, quant_config): + import bitsandbytes.functional as bnb_F + + weight = torch.randn(128, 64, dtype=torch.float16, device="cuda") + qweight, quant_state = bnb_F.quantize_4bit( + weight, + quant_type=quant_config.quant_type, + compress_statistics=quant_config.compress_statistics, + ) + + assert qweight.dtype == torch.uint8 + assert quant_state is not None + + def test_real_cuda_online_process_weights_after_loading(self, quant_config, real_layer): + from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod + + method = BnBOnlineLinearMethod(quant_config) + method.process_weights_after_loading(real_layer) + + assert real_layer.weight.dtype == torch.uint8 + assert hasattr(real_layer, "quant_state") + assert real_layer.bnb_shape == (128, 64) + + def test_real_cuda_bnb_apply_forward(self, quant_config): + from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod + + method = BnBOnlineLinearMethod(quant_config) + + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter( + torch.randn(128, 64, dtype=torch.float16, device="cuda"), + requires_grad=False, + ) + method.process_weights_after_loading(layer) + + x = torch.randn(2, 16, 64, dtype=torch.float16, device="cuda") + output = method.apply(layer, x) + + assert output.shape == (2, 16, 128) + assert output.dtype == torch.float16 diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 3d29a4332..ba132dbd2 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -371,3 +371,27 @@ def test_default_stage_config_includes_quantization_config(): stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"quantization_config": quantization_config})[0] assert stage_cfg["engine_args"]["quantization_config"] == quantization_config + + +def test_resolve_stage_configs_injects_quantization_config_into_diffusion_stage(mocker): + fake_diffusion_stage = SimpleNamespace( + stage_type="diffusion", + engine_args=SimpleNamespace(quantization_config=None), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs", + return_value=("dummy.yaml", [fake_diffusion_stage], None), + ) + + engine = AsyncOmniEngine.__new__(AsyncOmniEngine) + engine._strip_single_engine_args = lambda kwargs: kwargs + + _, stage_configs = engine._resolve_stage_configs( + "dummy-model", + { + "stage_configs_path": "dummy.yaml", + "quantization_config": {"method": "bitsandbytes"}, + }, + ) + + assert stage_configs[0].engine_args.quantization_config == {"method": "bitsandbytes"} diff --git a/vllm_omni/quantization/bitsandbytes_config.py b/vllm_omni/quantization/bitsandbytes_config.py new file mode 100644 index 000000000..7d8e7a684 --- /dev/null +++ b/vllm_omni/quantization/bitsandbytes_config.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""BitsAndBytes 4-bit quantization config for diffusion transformers. + +Supports online (dynamic) NF4/FP4 weight-only quantization from BF16/FP16 +checkpoints on CUDA GPUs. +""" + +from typing import TYPE_CHECKING, Any, Optional + +import torch +from torch.nn import Module +from vllm.model_executor.layers.linear import ( + LinearBase, + LinearMethodBase, + UnquantizedLinearMethod, +) +from vllm.model_executor.layers.quantization import QuantizationMethods +from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from vllm.model_executor.layers.quantization.fp8 import _copy_missing_attrs +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + is_layer_skipped, +) +from vllm.model_executor.model_loader.weight_utils import initialize_single_dummy_weight +from vllm.model_executor.parameter import ModelWeightParameter +from vllm.model_executor.utils import replace_parameter + +from vllm_omni.platforms import current_omni_platform +from vllm_omni.quantization.int8_config import LazyWeightMixin + +if TYPE_CHECKING: + from vllm.model_executor.models.utils import WeightsMapper + + +class DiffusionBitsAndBytesConfig(QuantizationConfig): + """BitsAndBytes 4-bit weight-only config for diffusion transformers. + + Supports online (dynamic) quantization from BF16/FP16 checkpoints. + Works on CUDA GPUs with the optional ``bitsandbytes`` package installed. + """ + + def __init__( + self, + quant_type: str = "nf4", + compress_statistics: bool = True, + ignored_layers: list[str] | None = None, + ) -> None: + super().__init__() + + if quant_type not in ("nf4", "fp4"): + raise ValueError(f"Unsupported quant_type {quant_type!r}; expected 'nf4' or 'fp4'") + self.quant_type = quant_type + self.compress_statistics = compress_statistics + self.ignored_layers = ignored_layers or [] + + @classmethod + def get_name(cls) -> QuantizationMethods: + return "bitsandbytes" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.bfloat16, torch.float16] + + @classmethod + def get_min_capability(cls) -> int: + return 75 + + @classmethod + def get_config_filenames(cls) -> list[str]: + return [] + + def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"): + if self.ignored_layers is not None: + self.ignored_layers = hf_to_vllm_mapper.apply_list(self.ignored_layers) + + @classmethod + def from_config(cls, config: dict[str, Any]) -> "DiffusionBitsAndBytesConfig": + quant_type = cls.get_from_keys_or(config, ["quant_type"], "nf4") + compress_statistics = cls.get_from_keys_or(config, ["compress_statistics"], True) + ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None) + + if not ignored_layers: + ignored_layers = cls.get_from_keys_or(config, ["modules_to_not_convert"], None) + return cls( + quant_type=quant_type, + compress_statistics=compress_statistics, + ignored_layers=ignored_layers, + ) + + def get_quant_method( + self, + layer: torch.nn.Module, + prefix: str, + ) -> Optional["QuantizeMethodBase"]: + if isinstance(layer, LinearBase): + if is_layer_skipped( + prefix=prefix, + ignored_layers=self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + ): + return UnquantizedLinearMethod() + if current_omni_platform.is_cuda(): + return BnBOnlineLinearMethod(self) + raise NotImplementedError("BitsAndBytes online quantization is only supported on CUDA.") + return None + + +class BnBOnlineLinearMethod(LazyWeightMixin, LinearMethodBase): + """Online BitsAndBytes 4-bit linear method. + + Loads BF16/FP16 checkpoint weights and quantizes them during loading. + """ + + def __init__(self, quant_config: DiffusionBitsAndBytesConfig): + self.quant_config = quant_config + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + if layer.weight.device == torch.device("meta"): + weight = ModelWeightParameter( + data=torch.empty_like(layer.weight, device=layer._load_device), + input_dim=1, + output_dim=0, + weight_loader=layer.weight.weight_loader, + ) + _copy_missing_attrs(layer.weight, weight) + layer.register_parameter("weight", weight) + initialize_single_dummy_weight(layer.weight) + + import bitsandbytes.functional as bnb_F + + weight = layer.weight.data.contiguous() + if not weight.is_cuda: + weight = weight.cuda() + + original_shape = tuple(weight.shape) + qweight, quant_state = bnb_F.quantize_4bit( + weight, + quant_type=self.quant_config.quant_type, + compress_statistics=self.quant_config.compress_statistics, + ) + + replace_parameter( + layer, + "weight", + torch.nn.Parameter(qweight, requires_grad=False), + ) + layer.quant_state = quant_state + layer.bnb_shape = original_shape + + layer._already_called_process_weights_after_loading = True + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + import bitsandbytes as bnb + + ori_shape = x.shape + ori_dtype = x.dtype + x_2d = x.reshape(-1, ori_shape[-1]) + + out = bnb.matmul_4bit( + x_2d, + layer.weight.t(), + quant_state=layer.quant_state, + bias=bias, + ) + return out.reshape(*ori_shape[:-1], -1).to(ori_dtype) diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py index 2129d8f04..b5d0a8ad7 100644 --- a/vllm_omni/quantization/factory.py +++ b/vllm_omni/quantization/factory.py @@ -88,6 +88,13 @@ def _build_int8(**kw: Any) -> QuantizationConfig: return DiffusionInt8Config(**kw) +def _build_bitsandbytes(**kw: Any) -> QuantizationConfig: + """Lazy import for BitsAndBytes 4-bit diffusion config (CUDA only).""" + from .bitsandbytes_config import DiffusionBitsAndBytesConfig + + return DiffusionBitsAndBytesConfig(**kw) + + def _build_mxfp8(**kw: Any) -> QuantizationConfig: """Lazy import for W8A8 MXFP8 diffusion config (NPU only).""" from .mxfp8_config import DiffusionMXFP8Config @@ -134,6 +141,7 @@ def _build_inc(**kw: Any) -> QuantizationConfig: _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = { "int8": _build_int8, + "bitsandbytes": _build_bitsandbytes, "mxfp8": _build_mxfp8, "mxfp4": _build_mxfp4, "mxfp4_dualscale": _build_mxfp4_dualscale,