anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Quantization] Migrate GGUF diffusion model support to OOT plugin (#4769)

Browse Source 
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
This commit is contained in:
Isotr0py authored and GitHub committed 2026-07-11 18:48:43 +00:00
1 parent ff6f0da271
commit b95905deb1
35 files changed
+272 -1614
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.buildkite/test-ready.yml
+21
View File
@@ -684,3 +684,24 @@ steps:
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Diffusion · GGUF Plugin Test"
soft_fail: true
source_file_dependencies:
- tests/plugin/test_gguf_diffusion.py
- vllm_omni/diffusion/model_loader
commands:
- pip install "vllm-gguf-plugin>=0.0.3"
- timeout 30m pytest -sv tests/plugin/test_gguf_diffusion.py
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
docs/user_guide/quantization/gguf.md
+9
View File
@@ -9,6 +9,15 @@ for the transformer.
GGUF is static quantization: the quantized weights are produced before serving.
!!! note
GGUF support has migrated to OOT [vllm-gguf-plugin](https://github.com/vllm-project/vllm-gguf-plugin). Make sure you have GGUF plugin installed before serving a GGUF model.
Before serving a GGUF model, make sure to install the [vllm-gguf-plugin](https://github.com/vllm-project/vllm-gguf-plugin):
```bash
uv pip install vllm-gguf-plugin>=0.0.3
```
## Hardware Support
| Device | Support |
examples/offline_inference/image_to_video/image_to_video.py
+1 -1
View File
@@ -220,7 +220,7 @@ def parse_args() -> argparse.Namespace:
"--quantization",
type=str,
default=None,
choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"],
choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8"],
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).",
)
examples/offline_inference/text_to_image/text_to_image.py
+3 -16
View File
@@ -202,17 +202,11 @@ def parse_args() -> argparse.Namespace:
"--quantization",
type=str,
default=None,
choices=["fp8", "int8", "gguf"],
choices=["fp8", "int8"],
help="Quantization method for the transformer. "
"Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8), 'gguf' (GGUF quantized weights). "
"Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8). "
"Default: None (no quantization, uses BF16).",
)
parser.add_argument(
"--gguf-model",
type=str,
default=None,
help=("GGUF file path or HF reference for transformer weights. Required when --quantization gguf is set."),
)
parser.add_argument(
"--ignored-layers",
type=str,
@@ -402,14 +396,7 @@ def main():
# ignored_layers is specified so the list flows through OmniDiffusionConfig
quant_kwargs: dict[str, Any] = {}
ignored_layers = [s.strip() for s in args.ignored_layers.split(",") if s.strip()] if args.ignored_layers else None
if args.quantization == "gguf":
if not args.gguf_model:
raise ValueError("--gguf-model is required when --quantization gguf is set.")
quant_kwargs["quantization_config"] = {
"method": "gguf",
"gguf_model": args.gguf_model,
}
elif args.quantization and ignored_layers:
if args.quantization and ignored_layers:
quant_kwargs["quantization_config"] = {
"method": args.quantization,
"ignored_layers": ignored_layers,
examples/offline_inference/text_to_video/text_to_video.md
+1 -1
View File
@@ -231,7 +231,7 @@ python text_to_video.py \
- `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
- `--audio-sample-rate`: audio sample rate for embedded audio (when the
pipeline returns audio; LTX-2.3 outputs 48kHz audio).
- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
- `--quantization`: quantization method (such as `fp8` for FP8).
- `--flow-shift`: scheduler flow_shift parameter.
- `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above).
examples/offline_inference/text_to_video/text_to_video.py
+1 -1
View File
@@ -227,7 +227,7 @@ def parse_args() -> argparse.Namespace:
"--quantization",
type=str,
default=None,
choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"],
choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8"],
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).",
)
tests/config/test_omni_config.py
+11
View File
@@ -751,3 +751,14 @@ def test_diffusion_config_field_classification_covers_current_fields():
"distributed_executor_backend",
} <= omni_config_module._DIFFUSION_SHARED_CONFIG_FIELDS
assert "prompt_file_path" in omni_config_module._DIFFUSION_RUNTIME_CONFIG_FIELDS
def test_diffusion_config_projection_keeps_mapping_quantization_config_serializable():
quantization_config = {
"method": "example_quant",
"weights": "weights.bin",
}
cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(quantization_config=quantization_config)
assert cfg.quantization_config == quantization_config
tests/diffusion/model_loader/gguf_adapters/test_flux2_klein_gguf_adapter.py
-112
View File
@@ -1,112 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the Flux2-Klein GGUF adapter."""
from __future__ import annotations
from types import SimpleNamespace
import pytest
import torch
from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
from vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein import (
Flux2KleinGGUFAdapter,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def _make_od_config(
model_class_name: str = "Flux2Pipeline",
model_type: str = "flux",
):
return SimpleNamespace(
model_class_name=model_class_name,
tf_model_config={"model_type": model_type},
)
def _make_source(prefix: str = "", subfolder: str = "transformer"):
return SimpleNamespace(prefix=prefix, subfolder=subfolder)
def test_flux2_klein_adapter_selected_for_flux_family():
adapter = get_gguf_adapter(
"dummy.gguf",
object(),
_make_source(),
_make_od_config(),
)
assert isinstance(adapter, Flux2KleinGGUFAdapter)
def test_flux2_klein_adapter_matches_flux_model_type():
assert Flux2KleinGGUFAdapter.is_compatible(
_make_od_config(model_class_name="OtherPipeline", model_type="flux-dev"),
object(),
_make_source(),
)
def test_flux2_klein_adapter_renames_core_projection_paths(monkeypatch: pytest.MonkeyPatch):
import vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein as flux_module
monkeypatch.setattr(
flux_module,
"gguf_quant_weights_iterator",
lambda _path: iter(
[
("img_in.qweight", torch.ones((1, 1))),
("time_in.in_layer.qweight_type", torch.tensor(1)),
("double_blocks.0.img_attn.norm.query_norm.weight", torch.full((1, 1), 2.0)),
("final_layer.linear.qweight", torch.full((1, 1), 3.0)),
]
),
)
adapter = Flux2KleinGGUFAdapter(
"dummy.gguf",
object(),
_make_source(),
_make_od_config(),
)
weights = list(adapter.weights_iterator())
names = [name for name, _ in weights]
assert "x_embedder.qweight" in names
assert "time_guidance_embed.timestep_embedder.linear_1.qweight_type" in names
assert "transformer_blocks.0.attn.norm_q.weight" in names
assert "proj_out.qweight" in names
def test_flux2_klein_adapter_swaps_final_adaln_shift_and_scale(monkeypatch: pytest.MonkeyPatch):
import vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein as flux_module
monkeypatch.setattr(
flux_module,
"gguf_quant_weights_iterator",
lambda _path: iter(
[
(
"final_layer.adaLN_modulation.1.scale",
torch.tensor([1.0, 2.0, 3.0, 4.0]),
)
]
),
)
adapter = Flux2KleinGGUFAdapter(
"dummy.gguf",
object(),
_make_source(),
_make_od_config(),
)
weights = list(adapter.weights_iterator())
assert len(weights) == 1
assert weights[0][0] == "norm_out.linear.weight"
assert torch.equal(weights[0][1], torch.tensor([3.0, 4.0, 1.0, 2.0]))
tests/diffusion/model_loader/gguf_adapters/test_qwen_image_gguf_adapter.py
-184
View File
@@ -1,184 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
from types import SimpleNamespace
import pytest
from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
from vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image import (
QwenImageGGUFAdapter,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def _make_od_config(
model_class_name: str = "QwenImagePipeline",
model_type: str = "qwen_image",
):
return SimpleNamespace(
model_class_name=model_class_name,
tf_model_config={"model_type": model_type},
)
def _make_source(prefix: str = "", subfolder: str = "transformer"):
return SimpleNamespace(prefix=prefix, subfolder=subfolder)
class _FakeTransformer:
def __init__(self) -> None:
self._params = {
"img_in.qweight": 1,
"img_in.qweight_type": 1,
"txt_in.qweight": 1,
"txt_in.qweight_type": 1,
"proj_out.qweight": 1,
"proj_out.qweight_type": 1,
"time_text_embed.timestep_embedder.linear_1.qweight": 1,
"time_text_embed.timestep_embedder.linear_1.qweight_type": 1,
"time_text_embed.timestep_embedder.linear_2.qweight": 1,
"time_text_embed.timestep_embedder.linear_2.qweight_type": 1,
"transformer_blocks.0.img_mod.1.qweight": 1,
"transformer_blocks.0.img_mod.1.qweight_type": 1,
"transformer_blocks.0.txt_mod.1.qweight": 1,
"transformer_blocks.0.txt_mod.1.qweight_type": 1,
"transformer_blocks.0.attn.to_qkv.qweight": 1,
"transformer_blocks.0.attn.to_qkv.qweight_type": 1,
"transformer_blocks.0.attn.add_kv_proj.qweight": 1,
"transformer_blocks.0.attn.add_kv_proj.qweight_type": 1,
"transformer_blocks.0.attn.to_out.weight": 1,
"transformer_blocks.0.img_mlp.net.0.proj.qweight": 1,
"transformer_blocks.0.img_mlp.net.0.proj.qweight_type": 1,
}
def named_parameters(self):
return list(self._params.items())
def named_buffers(self):
return []
def get_submodule(self, _prefix: str):
return self
def test_qwen_adapter_selected_for_qwen_image_family():
adapter = get_gguf_adapter(
"dummy.gguf",
_FakeTransformer(),
_make_source(),
_make_od_config(),
)
assert isinstance(adapter, QwenImageGGUFAdapter)
def test_qwen_adapter_matches_multiple_pipeline_variants():
for model_class_name in (
"QwenImagePipeline",
"QwenImageEditPipeline",
"QwenImageEditPlusPipeline",
"QwenImageLayeredPipeline",
):
assert QwenImageGGUFAdapter.is_compatible(
_make_od_config(model_class_name=model_class_name),
_FakeTransformer(),
_make_source(),
)
def test_qwen_adapter_preserves_split_projection_names(monkeypatch: pytest.MonkeyPatch):
import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
monkeypatch.setattr(
qwen_image_module,
"gguf_quant_weights_iterator",
lambda _path: iter(
[
("transformer_blocks.0.attn.to_q.qweight_type", 1),
("transformer_blocks.0.attn.to_q.qweight", 2),
("transformer_blocks.0.attn.to_k.qweight_type", 3),
("transformer_blocks.0.attn.to_k.qweight", 4),
("transformer_blocks.0.attn.to_out.0.qweight_type", 5),
("transformer_blocks.0.attn.to_out.0.qweight", 6),
]
),
)
adapter = QwenImageGGUFAdapter(
"dummy.gguf",
_FakeTransformer(),
_make_source(),
_make_od_config(),
)
weights = list(adapter.weights_iterator())
assert ("transformer_blocks.0.attn.to_q.qweight_type", 1) in weights
assert ("transformer_blocks.0.attn.to_q.qweight", 2) in weights
assert ("transformer_blocks.0.attn.to_k.qweight_type", 3) in weights
assert ("transformer_blocks.0.attn.to_k.qweight", 4) in weights
assert ("transformer_blocks.0.attn.to_out.0.qweight_type", 5) in weights
assert ("transformer_blocks.0.attn.to_out.0.qweight", 6) in weights
def test_qwen_adapter_keeps_top_level_quantized_weights(monkeypatch: pytest.MonkeyPatch):
import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
monkeypatch.setattr(
qwen_image_module,
"gguf_quant_weights_iterator",
lambda _path: iter(
[
("img_in.qweight_type", 1),
("img_in.qweight", 2),
("time_text_embed.timestep_embedder.linear_1.qweight_type", 3),
("time_text_embed.timestep_embedder.linear_1.qweight", 4),
]
),
)
adapter = QwenImageGGUFAdapter(
"dummy.gguf",
_FakeTransformer(),
_make_source(),
_make_od_config(),
)
weights = list(adapter.weights_iterator())
assert ("img_in.qweight_type", 1) in weights
assert ("img_in.qweight", 2) in weights
assert ("time_text_embed.timestep_embedder.linear_1.qweight_type", 3) in weights
assert ("time_text_embed.timestep_embedder.linear_1.qweight", 4) in weights
def test_qwen_adapter_keeps_modulation_quantized_weights(monkeypatch: pytest.MonkeyPatch):
import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
monkeypatch.setattr(
qwen_image_module,
"gguf_quant_weights_iterator",
lambda _path: iter(
[
("transformer_blocks.0.img_mod.1.qweight_type", 1),
("transformer_blocks.0.img_mod.1.qweight", 2),
("transformer_blocks.0.txt_mod.1.qweight_type", 3),
("transformer_blocks.0.txt_mod.1.qweight", 4),
]
),
)
adapter = QwenImageGGUFAdapter(
"dummy.gguf",
_FakeTransformer(),
_make_source(),
_make_od_config(),
)
weights = list(adapter.weights_iterator())
assert ("transformer_blocks.0.img_mod.1.qweight_type", 1) in weights
assert ("transformer_blocks.0.img_mod.1.qweight", 2) in weights
assert ("transformer_blocks.0.txt_mod.1.qweight_type", 3) in weights
assert ("transformer_blocks.0.txt_mod.1.qweight", 4) in weights
tests/diffusion/model_loader/gguf_adapters/test_z_image_gguf_adapter.py
-83
View File
@@ -1,83 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the Z-Image GGUF adapter."""
from __future__ import annotations
from types import SimpleNamespace
import pytest
import torch
from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
from vllm_omni.diffusion.model_loader.gguf_adapters.z_image import (
ZImageGGUFAdapter,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def _make_od_config(
model_class_name: str = "ZImagePipeline",
model_type: str = "z_image",
):
return SimpleNamespace(
model_class_name=model_class_name,
tf_model_config={"model_type": model_type},
)
def _make_source(prefix: str = "", subfolder: str = "transformer"):
return SimpleNamespace(prefix=prefix, subfolder=subfolder)
def test_z_image_adapter_selected_for_z_image_family():
adapter = get_gguf_adapter(
"dummy.gguf",
object(),
_make_source(),
_make_od_config(),
)
assert isinstance(adapter, ZImageGGUFAdapter)
def test_z_image_adapter_matches_model_type_variants():
for model_type in ("z_image", "zimage", "z-image"):
assert ZImageGGUFAdapter.is_compatible(
_make_od_config(model_class_name="OtherPipeline", model_type=model_type),
object(),
_make_source(),
)
def test_z_image_adapter_renames_known_gguf_tensor_paths(monkeypatch: pytest.MonkeyPatch):
import vllm_omni.diffusion.model_loader.gguf_adapters.z_image as z_image_module
monkeypatch.setattr(
z_image_module,
"gguf_quant_weights_iterator",
lambda _path: iter(
[
("model.diffusion_model.final_layer.qweight", torch.ones((1, 1))),
("model.diffusion_model.x_embedder.qweight_type", torch.tensor(1)),
("transformer_blocks.0.attention.out.weight", torch.full((1, 1), 2.0)),
("transformer_blocks.0.attention.qkv.qweight", torch.full((1, 1), 3.0)),
]
),
)
adapter = ZImageGGUFAdapter(
"dummy.gguf",
object(),
_make_source(),
_make_od_config(),
)
weights = list(adapter.weights_iterator())
names = [name for name, _ in weights]
assert "all_final_layer.2-1.qweight" in names
assert "all_x_embedder.2-1.qweight_type" in names
assert "transformer_blocks.0.attention.to_out.0.weight" in names
assert "transformer_blocks.0.attention.to_qkv.qweight" in names
tests/diffusion/model_loader/test_diffusers_loader.py
-23
View File
@@ -16,7 +16,6 @@ from vllm.config.load import LoadConfig
from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
from vllm_omni.diffusion.models.helios import HeliosPipeline
from vllm_omni.diffusion.registry import initialize_model
@@ -111,27 +110,6 @@ def test_empty_source_prefix_keeps_full_model_strict_check():
loader.load_weights(model)
def test_qwen_model_class_selects_qwen_gguf_adapter():
od_config = type(
"Config",
(),
{
"model_class_name": "QwenImagePipeline",
"tf_model_config": {"model_type": "qwen_image"},
},
)()
source = DiffusersPipelineLoader.ComponentSource(
model_or_path="dummy",
subfolder="transformer",
revision=None,
prefix="transformer.",
)
adapter = get_gguf_adapter("dummy.gguf", object(), source, od_config)
assert adapter.__class__.__name__ == "QwenImageGGUFAdapter"
class _ConfigAwareModel(nn.Module):
def __init__(self, *, od_config):
super().__init__()
@@ -186,7 +164,6 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
loader.load_weights = lambda model: None # type: ignore[assignment]
loader._process_weights_after_loading = lambda model, target_device: None # type: ignore[assignment]
loader._is_gguf_quantization = lambda: False # type: ignore[assignment]
monkeypatch.setattr(loader_mod, "resolve_obj_by_qualname", lambda _name: _ConfigAwareModel)
monkeypatch.setattr(loader_mod.torch, "device", lambda _name: _DeviceContext("cpu"))
tests/diffusion/model_loader/test_diffusers_loader_gguf.py
-259
View File
@@ -1,259 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for GGUF-specific DiffusersPipelineLoader behavior."""
from __future__ import annotations
from types import SimpleNamespace
import pytest
import torch
import torch.nn as nn
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_module
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
class _DummyModel(nn.Module):
def __init__(self) -> None:
super().__init__()
self.transformer = nn.Linear(2, 2, bias=True)
self.vae = nn.Linear(2, 2, bias=False)
self.register_buffer("transformer_buffer", torch.ones(1))
self.calls: list[list[str]] = []
self.weights_sources = [
DiffusersPipelineLoader.ComponentSource(
model_or_path="dummy",
subfolder="transformer",
revision=None,
prefix="transformer.",
fall_back_to_pt=True,
),
DiffusersPipelineLoader.ComponentSource(
model_or_path="dummy",
subfolder="vae",
revision=None,
prefix="vae.",
fall_back_to_pt=True,
),
]
def load_weights(self, weights):
loadable = dict(self.named_parameters())
loadable.update(dict(self.named_buffers()))
seen: list[str] = []
loaded: set[str] = set()
for name, tensor in weights:
seen.append(name)
if name in loadable:
target = loadable[name]
target.data.copy_(tensor.to(dtype=target.dtype))
loaded.add(name)
self.calls.append(seen)
return loaded
def _make_loader() -> DiffusersPipelineLoader:
loader = object.__new__(DiffusersPipelineLoader)
loader.load_config = SimpleNamespace(
download_dir="cache-dir",
ignore_patterns=["*.tmp"],
)
loader.od_config = None
loader.quant_config = None
loader.counter_before_loading_weights = 0.0
loader.counter_after_loading_weights = 0.0
return loader
def test_is_gguf_quantization_accepts_dict_config():
loader = _make_loader()
loader.quant_config = {
"method": "gguf",
"gguf_model": "weights.gguf",
}
assert loader._is_gguf_quantization() is True
def test_is_gguf_quantization_rejects_non_gguf_dict_config():
loader = _make_loader()
loader.quant_config = {
"method": "fp8",
"gguf_model": "weights.gguf",
}
assert loader._is_gguf_quantization() is False
def test_is_gguf_quantization_requires_gguf_model_for_dict_config():
loader = _make_loader()
loader.quant_config = {"method": "gguf"}
with pytest.raises(ValueError, match="gguf_model"):
loader._is_gguf_quantization()
def test_is_gguf_quantization_accepts_object_config():
loader = _make_loader()
loader.quant_config = SimpleNamespace(
get_name=lambda: "gguf",
gguf_model="weights.gguf",
)
assert loader._is_gguf_quantization() is True
def test_is_gguf_quantization_uses_fallback_object_without_get_name():
loader = _make_loader()
loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
assert loader._is_gguf_quantization() is True
def test_get_model_loadable_names_collects_parameters_and_buffers():
loader = _make_loader()
model = _DummyModel()
names = loader._get_model_loadable_names(model)
assert "transformer.weight" in names
assert "transformer.bias" in names
assert "transformer_buffer" in names
def test_resolve_gguf_model_path_returns_local_file(tmp_path):
loader = _make_loader()
gguf_file = tmp_path / "model.gguf"
gguf_file.write_bytes(b"gguf")
resolved = loader._resolve_gguf_model_path(str(gguf_file), revision=None)
assert resolved == str(gguf_file)
def test_resolve_gguf_model_path_downloads_explicit_gguf_filename(monkeypatch: pytest.MonkeyPatch):
loader = _make_loader()
monkeypatch.setattr(loader_module.os.path, "isfile", lambda _path: False)
monkeypatch.setattr(
loader_module,
"hf_hub_download",
lambda repo_id, filename, revision, cache_dir: f"{cache_dir}/{repo_id}/{filename}",
)
resolved = loader._resolve_gguf_model_path("owner/repo/model-Q4_0.gguf", revision="main")
assert resolved == "cache-dir/owner/repo/model-Q4_0.gguf"
def test_resolve_gguf_model_path_downloads_by_quant_type(monkeypatch: pytest.MonkeyPatch):
loader = _make_loader()
monkeypatch.setattr(loader_module.os.path, "isfile", lambda _path: False)
monkeypatch.setattr(
loader_module,
"download_gguf",
lambda repo_id, quant_type, cache_dir, revision, ignore_patterns: f"{cache_dir}/{repo_id}/{quant_type}.gguf",
)
resolved = loader._resolve_gguf_model_path("owner/repo:Q4_0", revision="main")
assert resolved == "cache-dir/owner/repo/Q4_0.gguf"
def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.MonkeyPatch):
loader = _make_loader()
loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
loader.od_config = SimpleNamespace(revision="main")
source = DiffusersPipelineLoader.ComponentSource(
model_or_path="dummy",
subfolder="transformer",
revision=None,
prefix="transformer.",
)
class _Adapter:
def weights_iterator(self):
yield "weight", torch.ones((2, 2))
yield "bias", torch.zeros(2)
monkeypatch.setattr(loader, "_resolve_gguf_model_path", lambda gguf_model, revision: "resolved.gguf")
monkeypatch.setattr(loader_module, "get_gguf_adapter", lambda gguf_file, model, source, od_config: _Adapter())
weights = list(loader._get_gguf_weights_iterator(source, object()))
assert weights[0][0] == "transformer.weight"
assert torch.equal(weights[0][1], torch.ones((2, 2)))
assert weights[1][0] == "transformer.bias"
assert torch.equal(weights[1][1], torch.zeros(2))
def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(monkeypatch: pytest.MonkeyPatch):
loader = _make_loader()
model = _DummyModel()
monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
monkeypatch.setattr(
loader,
"_get_gguf_weights_iterator",
lambda source, model: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]),
)
def _hf_iter(source):
if source.subfolder == "transformer":
return iter([("transformer.bias", torch.zeros(2, dtype=torch.float32))])
return iter([("vae.weight", torch.full((2, 2), 2.0, dtype=torch.float32))])
monkeypatch.setattr(loader, "_get_weights_iterator", _hf_iter)
monkeypatch.setattr(
loader,
"_get_expected_parameter_names",
lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
)
loaded = loader._load_weights_with_gguf(model)
assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
assert model.calls[0] == ["transformer.weight"]
assert model.calls[1] == ["transformer.bias"]
assert model.calls[2] == ["vae.weight"]
def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete(monkeypatch: pytest.MonkeyPatch):
loader = _make_loader()
model = _DummyModel()
monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
monkeypatch.setattr(
loader,
"_get_gguf_weights_iterator",
lambda source, model: iter(
[
("transformer.weight", torch.ones((2, 2), dtype=torch.float32)),
("transformer.bias", torch.zeros(2, dtype=torch.float32)),
]
),
)
hf_calls: list[str] = []
def _hf_iter(source):
hf_calls.append(source.subfolder or "")
if source.subfolder == "transformer":
return iter([("transformer.bias", torch.zeros(2, dtype=torch.float32))])
return iter([("vae.weight", torch.full((2, 2), 2.0, dtype=torch.float32))])
monkeypatch.setattr(loader, "_get_weights_iterator", _hf_iter)
monkeypatch.setattr(
loader,
"_get_expected_parameter_names",
lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
)
loaded = loader._load_weights_with_gguf(model)
assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
assert hf_calls == ["vae"]
tests/diffusion/quantization/test_fp8_config.py
+1 -16
View File
@@ -154,21 +154,6 @@ def test_component_config_with_default():
assert config.resolve("vae.encoder") is None
def test_gguf_config():
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig
config = build_quant_config(
{
"method": "gguf",
"gguf_model": "path/to/model.gguf",
}
)
assert isinstance(config, DiffusionGGUFConfig)
assert config.gguf_model == "path/to/model.gguf"
assert config.get_name() == "gguf"
def test_integration_string():
from vllm_omni.diffusion.data import OmniDiffusionConfig
@@ -236,7 +221,7 @@ def test_transformer_config_auto_detects_modelopt_fp8():
def test_supported_methods_includes_vllm():
from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
for method in ["fp8", "gguf", "awq", "gptq", "bitsandbytes", "modelopt"]:
for method in ["fp8", "awq", "gptq", "bitsandbytes", "modelopt"]:
assert method in SUPPORTED_QUANTIZATION_METHODS, f"{method} missing"
tests/diffusion/quantization/test_gguf_config.py
-128
View File
@@ -1,128 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for diffusion GGUF quantization config and linear method."""
from __future__ import annotations
from types import SimpleNamespace
import pytest
import torch
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.gguf_config import (
UNQUANTIZED_TYPES,
DiffusionGGUFConfig,
DiffusionGGUFLinearMethod,
dequant_gemm_gguf,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def test_gguf_config_creation_and_delegation():
config = DiffusionGGUFConfig(
gguf_model="weights.gguf",
unquantized_modules=["proj_out"],
)
assert config.gguf_model == "weights.gguf"
assert config.unquantized_modules == ["proj_out"]
assert config.get_name() == "gguf"
def test_build_quant_config_builds_gguf_config():
config = build_quant_config(
"gguf",
gguf_model="repo/model-Q4_0.gguf",
unquantized_modules=["norm_out"],
)
assert isinstance(config, DiffusionGGUFConfig)
assert config.gguf_model == "repo/model-Q4_0.gguf"
assert config.unquantized_modules == ["norm_out"]
def test_gguf_vllm_config_returns_diffusion_linear_method_for_linear_layers():
linear = object.__new__(LinearBase)
method = DiffusionGGUFConfig(unquantized_modules=[]).get_quant_method(linear, "transformer.img_in")
assert isinstance(method, DiffusionGGUFLinearMethod)
def test_gguf_vllm_config_respects_unquantized_modules():
linear = object.__new__(LinearBase)
method = DiffusionGGUFConfig(unquantized_modules=["proj_out"]).get_quant_method(linear, "transformer.proj_out")
assert isinstance(method, UnquantizedLinearMethod)
def test_gguf_vllm_config_returns_none_for_non_linear_layers():
method = DiffusionGGUFConfig(unquantized_modules=[]).get_quant_method(torch.nn.LayerNorm(4), "norm")
assert method is None
def test_dequant_gemm_gguf_uses_plain_matmul_for_unquantized_types():
qweight_type = next(iter(UNQUANTIZED_TYPES))
x = torch.tensor([[1.0, 2.0]], dtype=torch.float32)
qweight = torch.tensor([[3.0, 4.0], [5.0, 6.0]], dtype=torch.float32)
out = dequant_gemm_gguf(x, qweight, qweight_type)
assert torch.allclose(out, x @ qweight.T)
def test_diffusion_gguf_linear_method_applies_bias_on_unquantized_weight():
qweight_type = next(iter(UNQUANTIZED_TYPES))
method = DiffusionGGUFLinearMethod(DiffusionGGUFConfig(unquantized_modules=[]))
layer = SimpleNamespace(
qweight=torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32),
qweight_type=SimpleNamespace(weight_type=qweight_type),
)
x = torch.tensor([[[1.0, 0.5]]], dtype=torch.float32)
bias = torch.tensor([0.25, -0.5], dtype=torch.float32)
out = method.apply(layer, x, bias)
expected = x @ layer.qweight.T + bias
assert torch.allclose(out, expected)
def test_diffusion_gguf_linear_method_concatenates_sharded_outputs():
qweight_type = next(iter(UNQUANTIZED_TYPES))
method = DiffusionGGUFLinearMethod(DiffusionGGUFConfig(unquantized_modules=[]))
qweight = torch.nn.Parameter(
torch.tensor(
[
[1.0, 0.0],
[0.0, 1.0],
[2.0, 0.0],
[0.0, 2.0],
],
dtype=torch.float32,
),
requires_grad=False,
)
qweight.shard_id = ["left", "right"]
qweight.shard_offset_map = {
"left": (0, 2, 2),
"right": (2, 4, 2),
}
layer = SimpleNamespace(
qweight=qweight,
qweight_type=SimpleNamespace(
shard_weight_type={
"left": qweight_type,
"right": qweight_type,
}
),
)
x = torch.tensor([[1.5, 2.0]], dtype=torch.float32)
out = method.apply(layer, x)
expected_q = x @ qweight[:2].T
expected_k = x @ qweight[2:4].T
assert torch.allclose(out, torch.cat([expected_q, expected_k], dim=-1))
tests/diffusion/quantization/test_trajectory_similarity_tool.py
-1
View File
@@ -99,7 +99,6 @@ def test_candidate_model_can_point_to_offline_checkpoint_without_online_quantiza
candidate_quantization_config_json=None,
candidate_ignored_layers=None,
ignored_layers=None,
candidate_gguf_model=None,
candidate_load_format="default",
)
tests/entrypoints/test_async_omni_diffusion_config.py
+12
View File
@@ -359,3 +359,15 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo
assert not hasattr(stage_configs[0].engine_args, "additional_config")
assert stage_configs[1].engine_args.additional_config == {"torchair_graph_config": {"enabled": True}}
def test_default_stage_config_includes_quantization_config():
"""Ensure structured quantization_config survives default diffusion-stage creation."""
quantization_config = {
"method": "example_quant",
"weights": "weights.bin",
}
stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"quantization_config": quantization_config})[0]
assert stage_cfg["engine_args"]["quantization_config"] == quantization_config
tests/plugin/__init__.py
+1
View File
@@ -0,0 +1 @@
# SPDX-License-Identifier: Apache-2.0
tests/plugin/test_gguf_diffusion.py
+154
View File
@@ -0,0 +1,154 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""End-to-end tests for GGUF quantization on diffusion models.
Validates that GGUF-quantized diffusion models generate valid images and
use less peak GPU memory than BF16 baseline.
Requires vllm-omni to be installed alongside the plugin.
Usage:
pytest tests/diffusion/test_gguf_memory.py -v
"""
from __future__ import annotations
from typing import NamedTuple
import numpy as np
import pytest
import torch
from tests.helpers.env import DeviceMemoryMonitor
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
class DiffusionGGUFTestConfig(NamedTuple):
artifact_prefix: str
gguf_model: str
hf_model: str
min_cosine_similarity: float = 0.95
Z_IMAGE_CONFIG = DiffusionGGUFTestConfig(
artifact_prefix="zimage",
hf_model="Tongyi-MAI/Z-Image-Turbo",
gguf_model="unsloth/Z-Image-Turbo-GGUF:Q4_0",
)
FLUX_CONFIG = DiffusionGGUFTestConfig(
artifact_prefix="flux2_klein",
hf_model="black-forest-labs/FLUX.2-klein-4B",
gguf_model="unsloth/FLUX.2-klein-4B-GGUF:Q8_0",
)
def _image_cosine_similarity(hf_image, gguf_image) -> float:
hf_tensor = torch.as_tensor(np.array(hf_image.convert("RGB"), copy=True), dtype=torch.float32).flatten()
gguf_tensor = torch.as_tensor(np.array(gguf_image.convert("RGB"), copy=True), dtype=torch.float32).flatten()
return torch.nn.functional.cosine_similarity(hf_tensor, gguf_tensor, dim=0).item()
def _generate_single_stage_image(
model: str,
height: int = 256,
width: int = 256,
num_inference_steps: int = 20,
seed: int = 42,
**extra_kwargs,
) -> tuple[list, float]:
"""Generate an image with a single-stage diffusion model.
Returns (images, peak_memory_gib).
"""
omni_kwargs = dict(extra_kwargs)
memory_monitor = DeviceMemoryMonitor(device_index=0)
memory_monitor.start()
try:
with OmniRunner(model, **omni_kwargs) as runner:
generator = torch.Generator(
device=current_omni_platform.device_type,
).manual_seed(seed)
outputs = runner.omni.generate(
"a photo of a cat sitting on a laptop keyboard",
OmniDiffusionSamplingParams(
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=0.0,
generator=generator,
),
)
finally:
memory_monitor.stop()
peak_mem = memory_monitor.peak_used_mb / 1024
first_output = outputs[0]
assert first_output.final_output_type == "image"
if hasattr(first_output, "images") and first_output.images:
images = first_output.images
else:
assert hasattr(first_output, "request_output") and first_output.request_output
request_output = first_output.request_output
if isinstance(request_output, list):
req_out = request_output[0]
else:
req_out = request_output
assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
images = req_out.images
assert len(images) >= 1
assert images[0].width == width
assert images[0].height == height
return images, peak_mem
@hardware_test(res={"cuda": "L4"})
@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.slow
@pytest.mark.parametrize("model", [Z_IMAGE_CONFIG, FLUX_CONFIG], ids=["Z-Image-Turbo", "FLUX.2-klein"])
def test_single_stage_diffusion_gguf(model: DiffusionGGUFTestConfig, monkeypatch: pytest.MonkeyPatch) -> None:
"""Z-Image-Turbo GGUF generates valid images and uses less memory than BF16."""
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
# BF16 baseline
hf_images, mem_bf16 = _generate_single_stage_image(
model=model.hf_model,
)
# GGUF
images, mem_gguf = _generate_single_stage_image(
model=model.hf_model,
diffusion_quantization_config={
"method": "gguf",
"gguf_model": model.gguf_model,
},
)
hf_image_path = f"test_{model.artifact_prefix}_hf.png"
gguf_image_path = f"test_{model.artifact_prefix}_gguf.png"
assert len(hf_images) >= 1
hf_images[0].save(hf_image_path)
assert len(images) >= 1
images[0].save(gguf_image_path)
print(f"Saved HF image: {hf_image_path}")
print(f"Saved GGUF image: {gguf_image_path}")
image_cosine_similarity = _image_cosine_similarity(hf_images[0], images[0])
print(f"{model.artifact_prefix} image cosine similarity: {image_cosine_similarity:.4f}")
assert image_cosine_similarity >= model.min_cosine_similarity, (
f"GGUF image cosine similarity ({image_cosine_similarity:.4f}) should be >= {model.min_cosine_similarity:.4f}"
)
print(f"{model.artifact_prefix} BF16 peak VRAM delta: {mem_bf16:.2f} GiB")
print(f"{model.artifact_prefix} GGUF peak VRAM delta: {mem_gguf:.2f} GiB")
reduction = (mem_bf16 - mem_gguf) / mem_bf16 * 100
print(f"VRAM reduction: {reduction:.1f}%")
assert mem_gguf < mem_bf16, f"GGUF ({mem_gguf:.2f} GiB) should use less VRAM than BF16 ({mem_bf16:.2f} GiB)"
tests/test_version.py
+23
View File
@@ -56,3 +56,26 @@ def test_warning_contains_version_strings():
msg = str(record[0].message)
assert "0.19.0rc2.dev21" in msg
assert "0.18.0" in msg
def test_omni_plugin_loader_invokes_vllm_general_plugins_first(monkeypatch):
import vllm.plugins as vllm_plugins
import vllm_omni.plugins as omni_plugins
calls = []
monkeypatch.setattr(omni_plugins, "omni_plugins_loaded", False)
monkeypatch.setattr(vllm_plugins, "load_general_plugins", lambda: calls.append("vllm"))
monkeypatch.setattr(
omni_plugins,
"load_omni_plugins_by_group",
lambda group: {"omni": lambda: calls.append(group)},
)
omni_plugins.load_omni_general_plugins()
assert calls == ["vllm", omni_plugins.OMNI_DEFAULT_PLUGINS_GROUP]
omni_plugins.load_omni_general_plugins()
assert calls == ["vllm", omni_plugins.OMNI_DEFAULT_PLUGINS_GROUP]
vllm_omni/config/omni_config.py
+1 -1
View File
@@ -551,7 +551,7 @@ class _DiffusionConfigProjection:
elif isinstance(self.quantization_config, str):
self.quantization_config = build_quant_config(self.quantization_config)
elif isinstance(self.quantization_config, Mapping):
self.quantization_config = build_quant_config(dict(self.quantization_config))
self.quantization_config = dict(self.quantization_config)
else:
raise TypeError(
"quantization_config must be str, dict, QuantizationConfig, or None, "
vllm_omni/diffusion/model_loader/diffusers_loader.py
+1 -118
View File
@@ -11,7 +11,6 @@ from pathlib import Path
from typing import cast
import torch
from huggingface_hub import hf_hub_download
from torch import nn
from vllm.config import ModelConfig
from vllm.config.load import LoadConfig
@@ -36,7 +35,6 @@ from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp
from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
get_checkpoint_adapter,
)
from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
from vllm_omni.diffusion.registry import initialize_model
@@ -381,10 +379,7 @@ class DiffusersPipelineLoader:
# DiffusionPipeline.from_pretrained() internally; it does
# not use our native customized pipeline classes.
cast(DiffusersAdapterPipeline, model).load_weights()
elif self._is_gguf_quantization():
self._load_weights_with_gguf(model)
else:
# Quantization does not happen in `load_weights` but after it
self.load_weights(model)
# Process weights after loading for quantization (e.g., FP8 online quantization)
@@ -454,7 +449,7 @@ class DiffusersPipelineLoader:
def _is_expected_quantized_weight(name: str) -> bool:
"""Return True if *name* is a quantization-specific parameter.
Quantization methods (GPTQ, AWQ, FP8, GGUF, Autoround, etc.) create extra
Quantization methods (GPTQ, AWQ, FP8, Autoround, etc.) create extra
parameters that have no counterpart in an unquantized checkpoint.
These are expected to be absent and should not trigger a load error.
"""
@@ -467,8 +462,6 @@ class DiffusersPipelineLoader:
".weight_scale",
".weight_scale_inv",
".input_scale",
# GGUF
".qweight_type",
# INT8 (weight_scale already covered above)
)
return name.endswith(_QUANTIZED_WEIGHT_SUFFIXES)
@@ -500,116 +493,6 @@ class DiffusersPipelineLoader:
if unexpected_missing:
raise ValueError(f"Following weights were not initialized from checkpoint: {unexpected_missing}")
def _get_gguf_model_from_config(self) -> str | None:
"""Extract gguf_model from quant_config (handles both dict and object)."""
if self.quant_config is None:
return None
if isinstance(self.quant_config, dict):
return self.quant_config.get("gguf_model")
return getattr(self.quant_config, "gguf_model", None)
def _is_gguf_quantization(self) -> bool:
"""Check whether or not this pipeline loader is pointing at a GGUF config."""
from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig
maybe_gguf_model = self._get_gguf_model_from_config()
is_gguf = (
isinstance(self.quant_config, DiffusionGGUFConfig)
or (isinstance(self.quant_config, dict) and self.quant_config.get("method") == "gguf")
or (hasattr(self.quant_config, "get_name") and self.quant_config.get_name() == "gguf")
)
if is_gguf:
if not maybe_gguf_model:
raise ValueError("GGUF quantization requires gguf_model")
return True
if isinstance(self.quant_config, dict):
return False
return bool(maybe_gguf_model)
def _is_transformer_source(self, source: "ComponentSource") -> bool:
if source.subfolder == "transformer":
return True
return source.prefix.startswith("transformer.")
def _get_model_loadable_names(self, model: nn.Module) -> set[str]:
# Avoid model.state_dict() here because GGUF uses UninitializedParameter
# which raises during detach(). Collect names directly.
return {name for name, _ in model.named_parameters()} | {name for name, _ in model.named_buffers()}
def _resolve_gguf_model_path(self, gguf_model: str, revision: str | None) -> str:
if os.path.isfile(gguf_model):
return gguf_model
# repo_id/filename.gguf
if "/" in gguf_model and gguf_model.endswith(".gguf"):
repo_id, filename = gguf_model.rsplit("/", 1)
return hf_hub_download(
repo_id=repo_id,
filename=filename,
revision=revision,
cache_dir=self.load_config.download_dir,
)
# repo_id:quant_type
if "/" in gguf_model and ":" in gguf_model:
repo_id, quant_type = gguf_model.rsplit(":", 1)
return download_gguf(
repo_id,
quant_type,
cache_dir=self.load_config.download_dir,
revision=revision,
ignore_patterns=self.load_config.ignore_patterns,
)
raise ValueError(
f"Unrecognized GGUF reference: {gguf_model!r} (expected local file, "
"<repo_id>/<filename>.gguf, or <repo_id>:<quant_type>)"
)
def _get_gguf_weights_iterator(
self,
source: "ComponentSource",
model: nn.Module,
) -> Generator[tuple[str, torch.Tensor], None, None]:
gguf_model = self._get_gguf_model_from_config()
if gguf_model is None:
raise ValueError("GGUF quantization requires quantization_config.gguf_model")
gguf_file = self._resolve_gguf_model_path(gguf_model, self.od_config.revision)
adapter = get_gguf_adapter(gguf_file, model, source, self.od_config)
weights_iter = adapter.weights_iterator()
return ((source.prefix + name, tensor) for (name, tensor) in weights_iter)
def _load_weights_with_gguf(self, model: nn.Module) -> set[str]:
sources = self._get_weight_sources(model)
loaded: set[str] = set()
loadable_names: set[str] | None = None
for source in sources:
if self._is_transformer_source(source):
loaded |= model.load_weights(self._get_gguf_weights_iterator(source, model))
# GGUF checkpoints can be transformer-only or partially quantized.
# Only fall back to HF if this source still has missing loadable weights.
loadable_names = loadable_names or self._get_model_loadable_names(model)
has_missing_for_source = any(
name.startswith(source.prefix) and name not in loaded for name in loadable_names
)
if not has_missing_for_source:
continue
hf_iter = self._get_weights_iterator(source)
hf_iter = (
(name, tensor) for (name, tensor) in hf_iter if name in loadable_names and name not in loaded
)
loaded |= model.load_weights(hf_iter)
else:
loaded |= model.load_weights(self._get_weights_iterator(source))
weights_to_load = self._get_expected_parameter_names(model)
weights_not_loaded = weights_to_load - loaded
if weights_not_loaded:
raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}")
return loaded
def _init_from_load_format(
self,
load_format: str,
vllm_omni/diffusion/model_loader/gguf_adapters/__init__.py
-47
View File
@@ -1,47 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
from .base import GGUFAdapter
from .flux2_klein import Flux2KleinGGUFAdapter
from .qwen_image import QwenImageGGUFAdapter
from .z_image import ZImageGGUFAdapter
if TYPE_CHECKING:
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.model_loader.diffusers_loader import (
DiffusersPipelineLoader,
)
def get_gguf_adapter(
gguf_file: str,
model: torch.nn.Module,
source: DiffusersPipelineLoader.ComponentSource,
od_config: OmniDiffusionConfig,
) -> GGUFAdapter:
adapter_classes = (QwenImageGGUFAdapter, ZImageGGUFAdapter, Flux2KleinGGUFAdapter)
for adapter_cls in adapter_classes:
if adapter_cls.is_compatible(od_config, model, source):
return adapter_cls(gguf_file, model, source, od_config)
model_type = None
if od_config.tf_model_config is not None:
model_type = od_config.tf_model_config.get("model_type")
supported = ", ".join(cls.__name__ for cls in adapter_classes)
raise ValueError(
"No GGUF adapter matched diffusion model "
f"(model_class_name={od_config.model_class_name!r}, model_type={model_type!r}). "
f"Supported adapters: {supported}."
)
__all__ = [
"GGUFAdapter",
"Flux2KleinGGUFAdapter",
"QwenImageGGUFAdapter",
"ZImageGGUFAdapter",
"get_gguf_adapter",
]
vllm_omni/diffusion/model_loader/gguf_adapters/base.py
-109
View File
@@ -1,109 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
from abc import ABC, abstractmethod
from collections.abc import Generator
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
import numpy as np
import torch
try:
import gguf
except ImportError:
# Optional dependency: only required for GGUF-quantized diffusion models.
# Keep this import lazy so a missing ``gguf`` does not break importing the
# diffusion model loader (and thus inspection/loading of every other
# diffusion model that imports through this package).
gguf = None
if TYPE_CHECKING:
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.model_loader.diffusers_loader import (
DiffusersPipelineLoader,
)
@dataclass
class MappedTensor:
name: str
tensor: Any
tensor_type: Any
row_slice: slice | None = None
swap_scale_shift: bool = False
class GGUFAdapter(ABC):
"""Base class for model-specific GGUF adapters."""
def __init__(
self,
gguf_file: str,
model: torch.nn.Module,
source: DiffusersPipelineLoader.ComponentSource,
od_config: OmniDiffusionConfig,
) -> None:
self.gguf_file = gguf_file
self.model = model
self.source = source
self.od_config = od_config
@staticmethod
def is_compatible(
od_config: OmniDiffusionConfig,
model: torch.nn.Module,
source: DiffusersPipelineLoader.ComponentSource,
) -> bool:
return False
@abstractmethod
def weights_iterator(self) -> Generator[tuple[str, torch.Tensor], None, None]:
raise NotImplementedError
# FIXME(Isotr0py): Sync implemnentation with upstream vLLM?
def gguf_quant_weights_iterator(gguf_file: str) -> Generator[tuple[str, torch.Tensor]]:
"""
Iterate over the quant weights in the model gguf files and convert
them to torch tensors.
Be careful of the order of yielding weight types and weights data,
we have to yield all weight types first before yielding any weights.
Otherwise it would cause issue when loading weights with for packed
layer with different quant types.
"""
if gguf is None:
raise ImportError(
"The 'gguf' package is required to load GGUF-quantized diffusion models. "
"Install it with `pip install gguf`."
)
reader = gguf.GGUFReader(gguf_file)
for tensor in reader.tensors:
weight_type = tensor.tensor_type
name = tensor.name
if weight_type.name not in ("F32", "F16"):
weight_type_name = name.replace("weight", "qweight_type")
weight_type = torch.tensor(weight_type)
yield weight_type_name, weight_type
for tensor in reader.tensors:
weight = tensor.data
weight_type = tensor.tensor_type
name = tensor.name
if weight_type.name not in ("F32", "F16"):
name = name.replace("weight", "qweight")
if weight_type.name == "BF16" and tensor.data.dtype == np.uint8:
# BF16 is currently the only "quantization" type that isn't
# actually quantized but is read as a raw byte tensor.
# Reinterpret as `torch.bfloat16` tensor.
weight = weight.view(np.uint16)
if reader.byte_order == "S":
# GGUF endianness != system endianness
weight = weight.byteswap()
param = torch.tensor(weight).view(torch.bfloat16)
else:
param = torch.tensor(weight)
yield name, param
vllm_omni/diffusion/model_loader/gguf_adapters/flux2_klein.py
-98
View File
@@ -1,98 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
from collections.abc import Iterable
import torch
from vllm.model_executor.models.utils import WeightsMapper
from .base import GGUFAdapter, gguf_quant_weights_iterator
FLUX2_TRANSFORMER_KEYS_RENAME_DICT = {
"single_blocks.": "single_transformer_blocks.",
# Image and text input projections
"img_in": "x_embedder",
"txt_in": "context_embedder",
# Timestep and guidance embeddings
"time_in.in_layer": "time_guidance_embed.timestep_embedder.linear_1",
"time_in.out_layer": "time_guidance_embed.timestep_embedder.linear_2",
"guidance_in.in_layer": "time_guidance_embed.guidance_embedder.linear_1",
"guidance_in.out_layer": "time_guidance_embed.guidance_embedder.linear_2",
# Modulation parameters
"double_stream_modulation_img.lin": "double_stream_modulation_img.linear",
"double_stream_modulation_txt.lin": "double_stream_modulation_txt.linear",
"single_stream_modulation.lin": "single_stream_modulation.linear",
# Final output layer
# "final_layer.adaLN_modulation.1": "norm_out.linear", # Handle separately since we need to swap mod params
"final_layer.linear": "proj_out",
}
FLUX2_TRANSFORMER_ADA_LAYER_NORM_KEY_MAP = {
"final_layer.adaLN_modulation.1": "norm_out.linear",
}
FLUX2_TRANSFORMER_DOUBLE_BLOCK_KEY_MAP = {
"double_blocks.": "transformer_blocks.",
# Handle fused QKV projections separately as we need to break into Q, K, V projections
"img_attn.norm.query_norm": "attn.norm_q",
"img_attn.norm.key_norm": "attn.norm_k",
"img_attn.proj": "attn.to_out.0",
"img_mlp.0": "ff.linear_in",
"img_mlp.2": "ff.linear_out",
"txt_attn.norm.query_norm": "attn.norm_added_q",
"txt_attn.norm.key_norm": "attn.norm_added_k",
"txt_attn.proj": "attn.to_add_out",
"txt_mlp.0": "ff_context.linear_in",
"txt_mlp.2": "ff_context.linear_out",
# Additional for fuse qkv
"img_attn.qkv": "attn.to_qkv",
"txt_attn.qkv": "attn.add_kv_proj",
}
FLUX2_TRANSFORMER_SINGLE_BLOCK_KEY_MAP = {
"linear1": "attn.to_qkv_mlp_proj",
"linear2": "attn.to_out",
"norm.query_norm": "attn.norm_q",
"norm.key_norm": "attn.norm_k",
}
class Flux2KleinGGUFAdapter(GGUFAdapter):
"""GGUF adapter for Flux2-Klein models with qkv splitting and adaLN swap."""
@staticmethod
def is_compatible(od_config, model: torch.nn.Module, source) -> bool:
model_class = od_config.model_class_name or ""
if model_class.startswith("Flux2"):
return True
cfg = od_config.tf_model_config
if cfg is not None:
model_type = str(cfg.get("model_type", "")).lower()
if model_type.startswith("flux"):
return True
return False
gguf_to_hf_mapper = WeightsMapper(
# double_stream_modulation
orig_to_new_prefix=FLUX2_TRANSFORMER_KEYS_RENAME_DICT | FLUX2_TRANSFORMER_ADA_LAYER_NORM_KEY_MAP,
orig_to_new_substr=FLUX2_TRANSFORMER_DOUBLE_BLOCK_KEY_MAP | FLUX2_TRANSFORMER_SINGLE_BLOCK_KEY_MAP,
)
def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]:
def custom_weights_adapter(
weights: Iterable[tuple[str, torch.Tensor]],
) -> Iterable[tuple[str, torch.Tensor]]:
for name, weight in weights:
# Handle the special case for adaLN modulation parameters that require swapping shift and scale
if name.endswith(".scale"):
name = name.replace(".scale", ".weight")
if name == "norm_out.linear.weight":
shift, scale = weight.chunk(2, dim=0)
weight = torch.cat([scale, shift], dim=0)
yield name, weight
else:
yield name, weight
weights = gguf_quant_weights_iterator(self.gguf_file)
weights = self.gguf_to_hf_mapper.apply(weights)
yield from custom_weights_adapter(weights)
vllm_omni/diffusion/model_loader/gguf_adapters/qwen_image.py
-27
View File
@@ -1,27 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
from collections.abc import Iterable
import torch
from .base import GGUFAdapter, gguf_quant_weights_iterator
class QwenImageGGUFAdapter(GGUFAdapter):
"""GGUF adapter for the Qwen-Image transformer family."""
@staticmethod
def is_compatible(od_config, model: torch.nn.Module, source) -> bool:
model_class = od_config.model_class_name or ""
if model_class.startswith("QwenImage"):
return True
cfg = od_config.tf_model_config
if cfg is not None:
model_type = str(cfg.get("model_type", "")).lower()
if model_type.startswith("qwen_image"):
return True
return False
def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]:
yield from gguf_quant_weights_iterator(self.gguf_file)
vllm_omni/diffusion/model_loader/gguf_adapters/z_image.py
-43
View File
@@ -1,43 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
from collections.abc import Iterable
import torch
from vllm.model_executor.models.utils import WeightsMapper
from .base import GGUFAdapter, gguf_quant_weights_iterator
Z_IMAGE_KEYS_RENAME_DICT = {
"final_layer.": "all_final_layer.2-1.",
"x_embedder.": "all_x_embedder.2-1.",
".attention.qkv": ".attention.to_qkv",
".attention.k_norm": ".attention.norm_k",
".attention.q_norm": ".attention.norm_q",
".attention.out": ".attention.to_out.0",
"model.diffusion_model.": "",
}
class ZImageGGUFAdapter(GGUFAdapter):
"""GGUF adapter for Z-Image models with QKV/FFN shard support."""
@staticmethod
def is_compatible(od_config, model: torch.nn.Module, source) -> bool:
model_class = od_config.model_class_name or ""
if model_class.startswith("ZImage"):
return True
cfg = od_config.tf_model_config
if cfg is not None:
model_type = str(cfg.get("model_type", "")).lower()
if model_type in {"z_image", "zimage", "z-image"}:
return True
return False
gguf_to_hf_mapper = WeightsMapper(
orig_to_new_substr=Z_IMAGE_KEYS_RENAME_DICT,
)
def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]:
weights = gguf_quant_weights_iterator(self.gguf_file)
yield from self.gguf_to_hf_mapper.apply(weights)
vllm_omni/diffusion/models/flux2/flux2_transformer.py
+1 -1
View File
@@ -1034,7 +1034,7 @@ class Flux2Transformer2DModel(nn.Module):
name = original_name
if name not in params_dict and ".to_out.0." in name:
name = name.replace(".to_out.0.", ".to_out.")
# Some GGUF checkpoints include quantized tensors for modules that
# Some quantized checkpoints include tensors for modules that
# are intentionally left unquantized in this model.
param = params_dict.get(name)
if param is None:
vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py
+1 -1
View File
@@ -1044,7 +1044,7 @@ class Flux2Transformer2DModel(nn.Module):
name = original_name
if name not in params_dict and ".to_out.0." in name:
name = name.replace(".to_out.0.", ".to_out.")
# Some GGUF checkpoints include quantized tensors for modules that
# Some quantized checkpoints include tensors for modules that
# are intentionally left unquantized in this model.
param = params_dict.get(name)
if param is None:
vllm_omni/engine/async_omni_engine.py
+2 -1
View File
@@ -1043,6 +1043,7 @@ class AsyncOmniEngine:
"enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True),
"num_weight_load_threads": kwargs.get("num_weight_load_threads", 4),
"quantization": kwargs.get("quantization", None),
"quantization_config": kwargs.get("quantization_config", None),
"diffusion_kv_cache_dtype": kwargs.get("diffusion_kv_cache_dtype", None),
"diffusion_kv_cache_skip_steps": kwargs.get("diffusion_kv_cache_skip_steps", None),
"diffusion_kv_cache_skip_layers": kwargs.get("diffusion_kv_cache_skip_layers", None),
@@ -1233,7 +1234,7 @@ class AsyncOmniEngine:
kwargs.get("diffusion_attention_config"),
attention_backend=kwargs.get("diffusion_attention_backend"),
)
quantization_config = kwargs.get("diffusion_quantization_config")
quantization_config = kwargs.get("diffusion_quantization_config") or kwargs.get("quantization_config")
if quantization_config is not None:
if (
not hasattr(cfg.engine_args, "quantization_config")
vllm_omni/entrypoints/cli/serve.py
+1 -1
View File
@@ -478,7 +478,7 @@ class OmniServeCommand(CLISubcommand):
default=None,
help=(
"JSON string for diffusion quantization_config. "
'Example: \'{"method":"gguf","gguf_model":"/path/to/model.gguf"}\'.'
'Example: \'{"method":"fp8","activation_scheme":"dynamic"}\'.'
),
)
omni_config_group.add_argument(
vllm_omni/plugins/__init__.py
+4
View File
@@ -68,6 +68,10 @@ def load_omni_general_plugins() -> None:
return
omni_plugins_loaded = True
from vllm.plugins import load_general_plugins
load_general_plugins()
plugins = load_omni_plugins_by_group(group=OMNI_DEFAULT_PLUGINS_GROUP)
# General plugins: we only need to execute the loaded functions.
for func in plugins.values():
vllm_omni/quantization/__init__.py
+3 -2
View File
@@ -12,10 +12,10 @@ Adds per-component quantization for multi-stage models.
"""
from .component_config import ComponentQuantizationConfig
from .factory import SUPPORTED_QUANTIZATION_METHODS, build_quant_config
from .factory import SUPPORTED_QUANTIZATION_METHODS, build_quant_config, register_quantization_override
from .inc_config import OmniINCConfig
# Heavy configs (GGUF, MXFP8) are NOT imported here to avoid pulling in
# Heavy configs are NOT imported here to avoid pulling in
# optional dependencies (pynvml, torch_npu) at module load time.
# Import them directly when needed:
# from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
@@ -25,4 +25,5 @@ __all__ = [
"ComponentQuantizationConfig",
"OmniINCConfig",
"SUPPORTED_QUANTIZATION_METHODS",
"register_quantization_override",
]
vllm_omni/quantization/factory.py
+20 -11
View File
@@ -3,8 +3,8 @@
"""Factory for building quantization configs.
build_quant_config() delegates to vLLM's quantization registry.
The only extension point is _OVERRIDES for methods that need a
different QuantizationConfig subclass in the OMNI context (e.g. GGUF).
Out-of-tree integrations can register Omni-specific builders with
register_quantization_override().
"""
from __future__ import annotations
@@ -81,13 +81,6 @@ from .component_config import ComponentQuantizationConfig # noqa: E402
logger = init_logger(__name__)
def _build_gguf(**kw: Any) -> QuantizationConfig:
"""Lazy import to avoid pulling in CUDA/pynvml at module load time."""
from .gguf_config import DiffusionGGUFConfig
return DiffusionGGUFConfig(**kw)
def _build_int8(**kw: Any) -> QuantizationConfig:
"""Lazy import for Int8 diffusion config (supports CUDA + NPU)."""
from .int8_config import DiffusionInt8Config
@@ -140,7 +133,6 @@ def _build_inc(**kw: Any) -> QuantizationConfig:
_OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = {
"gguf": _build_gguf,
"int8": _build_int8,
"mxfp8": _build_mxfp8,
"mxfp4": _build_mxfp4,
@@ -150,7 +142,24 @@ _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = {
"auto_round": _build_inc,
}
SUPPORTED_QUANTIZATION_METHODS: list[str] = list(dict.fromkeys(QUANTIZATION_METHODS + list(_OVERRIDES.keys())))
def _compute_supported_quantization_methods() -> list[str]:
return list(dict.fromkeys(QUANTIZATION_METHODS + list(_OVERRIDES.keys())))
SUPPORTED_QUANTIZATION_METHODS: list[str] = _compute_supported_quantization_methods()
def _refresh_registered_methods() -> None:
SUPPORTED_QUANTIZATION_METHODS[:] = _compute_supported_quantization_methods()
global _CACHED_ALIAS_MAP
_CACHED_ALIAS_MAP = None
def register_quantization_override(method: str, builder: Callable[..., QuantizationConfig]) -> None:
"""Register an Omni-specific quantization config builder."""
_OVERRIDES[_normalize_method_name(method)] = builder
_refresh_registered_methods()
def _build_reverse_alias_map() -> dict[str, str]:
vllm_omni/quantization/gguf_config.py
-323
View File
@@ -1,323 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""GGUF quantization config for diffusion transformers.
Uses dequant+GEMM instead of the fused kernel path (which expects 2D inputs).
The GGUF quantization module was migrated to an external plugin in upstream
vLLM (commit 6635279d8). This file inlines the base classes that
DiffusionGGUFConfig / DiffusionGGUFLinearMethod depend on, reproduced from
the last upstream version before the migration.
"""
from __future__ import annotations
from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Any
import gguf
import torch
from gguf import GGMLQuantizationType as WeightType
from torch.nn.parameter import Parameter, UninitializedParameter
from vllm import _custom_ops as ops
from vllm.logger import init_logger
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.utils import set_weight_attrs
if TYPE_CHECKING:
from vllm.model_executor.layers.quantization import QuantizationMethods
logger = init_logger(__name__)
# ---------------------------------------------------------------------------
# GGUF base types (inlined from upstream vllm before migration to plugin)
# ---------------------------------------------------------------------------
UNQUANTIZED_TYPES: set[WeightType] = {WeightType.F32, WeightType.F16, WeightType.BF16}
STANDARD_QUANT_TYPES: set[WeightType] = {
WeightType.Q4_0,
WeightType.Q4_1,
WeightType.Q5_0,
WeightType.Q5_1,
WeightType.Q8_0,
WeightType.Q8_1,
}
KQUANT_TYPES: set[WeightType] = {
WeightType.Q2_K,
WeightType.Q3_K,
WeightType.Q4_K,
WeightType.Q5_K,
WeightType.Q6_K,
}
IMATRIX_QUANT_TYPES: set[WeightType] = {
WeightType.IQ1_M,
WeightType.IQ1_S,
WeightType.IQ2_XXS,
WeightType.IQ2_XS,
WeightType.IQ2_S,
WeightType.IQ3_XXS,
WeightType.IQ3_S,
WeightType.IQ4_XS,
WeightType.IQ4_NL,
}
DEQUANT_TYPES: set[WeightType] = STANDARD_QUANT_TYPES | KQUANT_TYPES | IMATRIX_QUANT_TYPES
def is_layer_skipped_gguf(
prefix: str,
unquantized_modules: list[str],
fused_mapping: Mapping[str, list[str]] = MappingProxyType({}),
) -> bool:
# Fused layers like gate_up_proj or qkv_proj will not be fused
# in the safetensors checkpoint. So, we convert the name
# from the fused version to unfused + check to make sure that
# each shard of the fused layer has the same scheme.
proj_name = prefix.split(".")[-1]
if proj_name in fused_mapping:
shard_prefixes = [prefix.replace(proj_name, shard_proj_name) for shard_proj_name in fused_mapping[proj_name]]
is_skipped = None
for shard_prefix in shard_prefixes:
is_shard_skipped = any(shard_prefix in module_name for module_name in unquantized_modules)
if is_skipped is None:
is_skipped = is_shard_skipped
elif is_shard_skipped != is_skipped:
raise ValueError(
f"Detected some but not all shards of {prefix} "
"are quantized. All shards of fused layers "
"to have the same precision."
)
else:
is_skipped = any(module_name in prefix for module_name in unquantized_modules)
assert is_skipped is not None
return is_skipped
class _GGUFUninitializedParameter(UninitializedParameter):
"""UninitializedParameter subclass that stores GGUF weight data in a list.
The actual tensor is materialized later via _create_padded_weight_param.
"""
cls_to_become = Parameter
data_container: list[torch.Tensor]
class GGUFConfig(QuantizationConfig):
"""Config class for GGUF (inlined from upstream vllm)."""
def __init__(self, unquantized_modules: list[str] | None = None) -> None:
super().__init__()
self.unquantized_modules = unquantized_modules or []
def __repr__(self) -> str:
return "GGUFConfig()"
def get_name(self) -> QuantizationMethods:
return "gguf"
def get_supported_act_dtypes(self) -> list[torch.dtype]:
return [torch.half, torch.bfloat16, torch.float32]
@classmethod
def get_min_capability(cls) -> int:
return 60
@classmethod
def get_config_filenames(cls) -> list[str]:
return [] # no extra configs.
@classmethod
def from_config(cls, config: dict[str, Any]) -> GGUFConfig:
return cls()
@classmethod
def override_quantization_method(
cls, hf_quant_cfg: dict[str, Any], user_quant: str | None, hf_config=None
) -> QuantizationMethods | None:
if user_quant == "gguf":
return "gguf"
return None
def get_quant_method(self, layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None:
if isinstance(layer, LinearBase):
if is_layer_skipped_gguf(prefix, self.unquantized_modules, self.packed_modules_mapping):
return UnquantizedLinearMethod()
return GGUFLinearMethod(self)
return None
class GGUFLinearMethod(LinearMethodBase):
"""Linear method for GGUF (inlined from upstream vllm)."""
def __init__(self, quant_config: GGUFConfig):
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,
):
self.params_dtype = params_dtype
output_size_per_partition = sum(output_partition_sizes)
tensor_shape = (output_size_per_partition, input_size_per_partition)
qweight = _GGUFUninitializedParameter(requires_grad=False)
set_weight_attrs(
qweight,
{
"input_dim": 1,
"output_dim": 0,
"tensor_shape": tensor_shape,
"is_gguf_weight": True,
"data_container": [],
"shard_id": [],
"shard_id_map": {},
},
)
set_weight_attrs(qweight, extra_weight_attrs)
layer.register_parameter("qweight", qweight)
qweight_type = Parameter(
torch.empty(len(output_partition_sizes), dtype=torch.uint8),
requires_grad=False,
)
set_weight_attrs(
qweight_type,
{
"is_gguf_weight_type": True,
"weight_type": 0,
"shard_weight_type": {},
"ignore_warning": True,
},
)
set_weight_attrs(qweight_type, extra_weight_attrs)
layer.register_parameter("qweight_type", qweight_type)
def process_weights_after_loading(self, layer: torch.nn.Module):
qweight_type = layer.qweight_type.weight_type
if not (qweight_type in UNQUANTIZED_TYPES or qweight_type in DEQUANT_TYPES):
qweight_type = WeightType(qweight_type)
raise ValueError(f"Unsupported GGUF quantization type {qweight_type} in layer {layer}.")
self._create_padded_weight_param(layer)
def _create_padded_weight_param(self, layer: torch.nn.Module):
"""Create padded weight parameter for GGUF MergedLinear layer."""
qweight = layer.qweight
shard_id_map = qweight.shard_id_map
shard_id = qweight.shard_id
if len(data_container := qweight.data_container) > 1:
dtype = {data.dtype for data in data_container}
assert len(dtype) == 1, ValueError(f"Data container has mixed dtypes: {dtype}")
dtype = next(iter(dtype))
padded_side = max(x.size(1) for x in data_container)
concat_side = sum(x.size(0) for x in data_container)
padded_data = torch.zeros((concat_side, padded_side), dtype=dtype, device=qweight.device)
shard_offset_map = dict[str, tuple[int, int, int]]()
for idx in shard_id:
id_in_container = shard_id_map[idx]
start = sum(x.size(0) for x in data_container[:id_in_container])
end = start + data_container[id_in_container].size(0)
size = data_container[id_in_container].size(1)
padded_data[start:end, :size] = data_container[id_in_container]
shard_offset_map[idx] = (start, end, size)
qweight.data_container.clear()
padded_param = Parameter(padded_data, requires_grad=False)
set_weight_attrs(padded_param, vars(qweight))
set_weight_attrs(padded_param, {"shard_offset_map": shard_offset_map})
layer.register_parameter("qweight", padded_param)
def apply(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
# Base apply: overridden by DiffusionGGUFLinearMethod.
raise NotImplementedError("Use DiffusionGGUFLinearMethod for diffusion GGUF")
# ---------------------------------------------------------------------------
# Diffusion-specific GGUF (dequant+GEMM, replaces fused kernel path)
# ---------------------------------------------------------------------------
def dequant_gemm_gguf(x: torch.Tensor, qweight: torch.Tensor, qweight_type: int) -> torch.Tensor:
if qweight_type in UNQUANTIZED_TYPES:
return x @ qweight.T
block_size, type_size = gguf.GGML_QUANT_SIZES[qweight_type]
shape = (qweight.shape[0], qweight.shape[1] // type_size * block_size)
weight = ops.ggml_dequantize(qweight, qweight_type, *shape, x.dtype)
return x @ weight.T
class DiffusionGGUFLinearMethod(GGUFLinearMethod):
"""GGUF linear method using dequant+GEMM for N-D diffusion tensors."""
def apply(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
shard_id = getattr(layer.qweight, "shard_id", None)
if shard_id:
shard_id = ["q", "k", "v"] if "q" in shard_id else shard_id
qweight = layer.qweight
result = []
for idx in shard_id:
start, end, offset = layer.qweight.shard_offset_map[idx]
qweight_type = layer.qweight_type.shard_weight_type[idx]
result.append(
dequant_gemm_gguf(
x,
qweight[start:end, :offset].contiguous(),
qweight_type,
)
)
out = torch.cat(result, axis=-1)
else:
qweight = layer.qweight
qweight_type = layer.qweight_type.weight_type
out = dequant_gemm_gguf(x, qweight, qweight_type)
if bias is not None:
out.add_(bias)
return out
class DiffusionGGUFConfig(GGUFConfig):
"""GGUF config that carries gguf_model path and uses dequant+GEMM."""
def __init__(
self,
gguf_model: str | None = None,
unquantized_modules: list[str] | None = None,
) -> None:
super().__init__(unquantized_modules=unquantized_modules or [])
self.gguf_model = gguf_model
def get_quant_method(self, layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None:
if isinstance(layer, LinearBase):
if is_layer_skipped_gguf(prefix, self.unquantized_modules, self.packed_modules_mapping):
return UnquantizedLinearMethod()
return DiffusionGGUFLinearMethod(self)
return None
vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py
-6
View File
@@ -280,7 +280,6 @@ def _build_quantization_config(
quantization: str | None,
quantization_config_json: str | None,
ignored_layers: str | None,
gguf_model: str | None,
) -> dict[str, Any] | None:
explicit = _load_json_arg(quantization_config_json)
if explicit is not None:
@@ -292,8 +291,6 @@ def _build_quantization_config(
ignored = _split_csv(ignored_layers)
if ignored:
config["ignored_layers"] = ignored
if gguf_model is not None:
config["gguf_model"] = gguf_model
return config
@@ -305,7 +302,6 @@ def _build_variant_config(args: argparse.Namespace, variant: str) -> VariantConf
quantization=quantization,
quantization_config_json=getattr(args, f"{variant}_quantization_config_json"),
ignored_layers=ignored_layers,
gguf_model=getattr(args, f"{variant}_gguf_model"),
)
return VariantConfig(
label=variant,
@@ -554,8 +550,6 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--reference-ignored-layers", help="Comma-separated reference ignored layer patterns.")
parser.add_argument("--candidate-ignored-layers", help="Comma-separated candidate ignored layer patterns.")
parser.add_argument("--ignored-layers", help="Alias for --candidate-ignored-layers.")
parser.add_argument("--reference-gguf-model", help="Reference GGUF file or HF reference when using gguf.")
parser.add_argument("--candidate-gguf-model", help="Candidate GGUF file or HF reference when using gguf.")
parser.add_argument("--reference-load-format", default="default", help="Reference diffusion load format.")
parser.add_argument("--candidate-load-format", default="default", help="Candidate diffusion load format.")
parser.add_argument("--prompt", default="a cup of coffee on the table")