[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

@@ -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

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`.

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(

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

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"}

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)

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,