ci: raise merge-tier entrypoint timeouts (H100 60->90, L4 30->45)

Both entrypoint jobs are killed mid-run by timeout_in_minutes while tests
are still passing, not by any failing test. The budgets had no headroom:
on main the H100 job timed out in 2 of the last 3 gate builds (2648, 2650)
and passed 2651 at 56/60 min; the L4 job passed at 26/30 min only with a
warm image cache. On the vLLM 0.25 wheel, model weight loading is slower
on the CI runners (like-for-like: +10 min total on H100, 1.7B TTS load
11s -> 132s on L4), which turns the coin flip into a guaranteed timeout
(build 2655).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

.buildkite/test-ready.yml

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

.gitignore

@@ -275,8 +275,3 @@ tests/.ci_generated/
 
 # Local Lance precision-matching dump scripts (not part of the repo)
 lance_compare/
-
-# FP8 quantization test sample images: tests/diffusion/quantization/
-# test_quantization_fp8.py saves these into the repo root for manual
-# inspection; they are byproducts and must never be committed.
-test_*_fp8.png

docker/Dockerfile.ci

@@ -14,71 +14,8 @@ RUN apt-get update && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
-
-# Install vLLM (CUDA 13.0 build) from precompiled wheel at the selected commit.
-# Must use direct URL because the wheel has a PEP 440 local version identifier
-# (e.g. +gaeee7ef93.cu130) which pip/uv refuse to install from a PEP 503 package index.
-# The cu130 wheel index lives at /<commit>/cu130/vllm/ but the actual wheel files
-# are stored at the top level /<commit>/<wheel>.
-ENV VLLM_PRECOMPILED_WHEEL_COMMIT=dd10e03f95f94edbea1975c67ace3a35ec9a8a40
-# Keep the base image's CUDA-12 toolkit libs on the loader path. NOTE: this does
-# NOT resolve libnvJitLink.so.13 (that CUDA-13 lib is not in /usr/local/cuda here) —
-# the .so.13 JIT linker is handled separately via ldconfig after the cu130 wheels
-# are installed (see the nvidia-cuda-nvjitlink-cu13 step below).
-ENV LD_LIBRARY_PATH=/usr/local/cuda/targets/x86_64-linux/lib:${LD_LIBRARY_PATH}
-RUN VLLM_WHEEL_URL=$(python3 -c "import urllib.request,re; \
-    from urllib.parse import urljoin, quote; \
-    index_url='https://wheels.vllm.ai/${VLLM_PRECOMPILED_WHEEL_COMMIT}/cu130/vllm/'; \
-    html=urllib.request.urlopen(index_url).read().decode(); \
-    m=re.search(r'href=\"([^\"]*x86_64\.whl)\"',html); \
-    assert m, f'No x86_64 vLLM wheel found at {index_url}'; \
-    print(quote(urljoin(index_url,m.group(1)), safe=':/+%'))") && \
-    echo "Installing vLLM from: ${VLLM_WHEEL_URL}" && \
-    uv pip install --system --force-reinstall "${VLLM_WHEEL_URL}"
-
-
-
 RUN uv pip install --system ".[dev]"
 
-RUN uv pip install --system --upgrade \
-        "flashinfer-python==0.6.13" \
-        "flashinfer-cubin==0.6.13"
-
-RUN uv pip install --system --upgrade \
-    "flashinfer-jit-cache==0.6.13" \
-    --index-url https://flashinfer.ai/whl/cu130
-
-# Force-reinstall torch AFTER flashinfer, because flashinfer --upgrade pulls in
-# a newer torch (2.12.x) whose libtorch_cuda.so lacks
-# at::cuda::getCurrentCUDABlasHandle, causing "undefined symbol" at import time.
-RUN uv pip install --system --force-reinstall \
-    "torch==2.11.0" "torchaudio==2.11.0" "torchvision==0.26.0" \
-    --torch-backend cu130
-
-# CUDA 13 JIT linker: libnvJitLink.so.13 is provided by the nvidia-nvjitlink
-# package that the cu130 torch/vllm wheels above pull into site-packages (the
-# v0.24.0 base image's /usr/local/cuda toolkit is CUDA 12, so the LD_LIBRARY_PATH
-# above never resolves the .so.13). Its lib dir is not on the dynamic-loader path,
-# so bitsandbytes / torch.compile worker subprocesses fail to dlopen it. Locate the
-# already-installed lib and register its dir with ldconfig (system-wide, survives
-# subprocess env resets). Only fall back to the pip wheel (correct name:
-# nvidia-nvjitlink-cu13, NOT nvidia-cuda-nvjitlink-cu13, which does not exist) if
-# the torch deps did not bring it in.
-RUN NVJITLINK_LIB="$(find / -name 'libnvJitLink.so.13*' 2>/dev/null | head -1)" && \
-    if [ -z "$NVJITLINK_LIB" ]; then \
-        uv pip install --system "nvidia-nvjitlink-cu13" && \
-        NVJITLINK_LIB="$(find / -name 'libnvJitLink.so.13*' 2>/dev/null | head -1)"; \
-    fi && \
-    test -n "$NVJITLINK_LIB" && \
-    dirname "$NVJITLINK_LIB" > /etc/ld.so.conf.d/nvjitlink-cu13.conf && \
-    ldconfig && \
-    echo "Registered ${NVJITLINK_LIB} with ldconfig"
-
-# Re-pin numpy LAST. numba requires numpy <= 2.4, but the flashinfer and torch cu130
-# reinstalls above resolve numpy back up to 2.5, which breaks `import numba` at test
-# time ("Numba needs NumPy 2.4 or less"). This must be the final dependency op so
-# nothing can clobber it again.
-RUN uv pip install --system "numpy==2.2.6"
 
 RUN ln -sf /usr/bin/python3 /usr/bin/python
 

docs/user_guide/quantization/gguf.md

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

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

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

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

@@ -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/ar_diffusion/test_kv_cache.py

@@ -9,9 +9,14 @@ KV stack on CPU (block bookkeeping only, no GPU tensors).
 import pytest
 import torch
 from vllm.v1.kv_cache_interface import KVCacheSpecKind, get_kv_cache_spec_kind
-from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
 from vllm.v1.request import RequestStatus
 
+try:
+    from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
+except ModuleNotFoundError:
+    KVCacheSpecRegistry = None
+    from vllm.v1.core.single_type_kv_cache_manager import spec_manager_map
+
 from vllm_omni.experimental.ar_diffusion.kv_cache import (
     ARDiffusionKVCache,
     ARDiffusionKVConfig,
@@ -48,7 +53,10 @@ def test_spec_registration_resolves_to_chunk_window_manager():
     # Without explicit registration the MRO walk would fall back to the parent
     # SlidingWindowManager; assert the subclass manager wins.
     spec = make_spec()
-    assert KVCacheSpecRegistry.get_manager_class(spec) is ChunkWindowManager
+    if KVCacheSpecRegistry is not None:
+        assert KVCacheSpecRegistry.get_manager_class(spec) is ChunkWindowManager
+    else:
+        assert spec_manager_map[type(spec)] is ChunkWindowManager
 
 
 def test_spec_kind_is_sliding_window():

tests/config/test_omni_config.py

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

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

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

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

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

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

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

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

@@ -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/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py

@@ -27,15 +27,8 @@ WIDTH = 1024
 PROMPT = "A brown and white dog is running on the grass."
 MEAN_THRESHOLD = 3e-2
 P99_THRESHOLD = 3e-1
-# NOTE: re-baselined for the vLLM rebase (dev/vllm-align). HunyuanImage-3 is an
-# expert-parallel MoE model; the newer upstream vLLM (65a7b4628409, 585 commits
-# past the v0.24.0 tag that main tracks) shifted MoE kernel numerics, so the
-# structural-similarity metrics against the committed baseline drift to a
-# deterministic SSIM 0.9490 / PSNR 26.21 dB (the mean/p99 pixel-diff checks above
-# still pass — the image renders correctly). Thresholds relaxed to accept this
-# drift as the new bar. On main (older vLLM) the stricter 0.97 / 30.0 hold.
-SSIM_THRESHOLD = 0.94
-PSNR_THRESHOLD = 25.0
+SSIM_THRESHOLD = 0.97
+PSNR_THRESHOLD = 30.0
 
 _REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent
 BASELINE_PATH = _REPO_ROOT / "tests" / "assets" / "hunyuan" / "hunyuan_baseline.png"

tests/engine/test_async_omni_engine_stage_init.py

@@ -874,54 +874,3 @@ def test_build_engine_args_dict_uses_diffusion_attention_config_key():
 
     assert "attention_config" not in engine_args_dict
     assert engine_args_dict["diffusion_attention_config"].default.backend == "FLASH_ATTN"
-
-
-def test_omni_master_server_allocates_globally_unique_route_ports(monkeypatch):
-    """Regression: two stages must never draw the same ZMQ port.
-
-    ``get_open_ports_list`` only dedups within a single call, so per-route
-    allocation used to let a later stage reuse an earlier stage's port. The
-    second engine to ``bind()`` then died with ``zmq.error.ZMQError: Address
-    already in use`` (flaky multi-stage startup, e.g. Qwen3-Omni thinker/talker/
-    code2wav). ``OmniMasterServer`` now dedups every port it hands out.
-    """
-    import itertools
-
-    from vllm_omni.engine import stage_engine_startup as ses
-
-    # A colliding prefix (repeats + the master port 9000) followed by an endless
-    # fresh stream, so the only way to succeed is to redraw the collisions.
-    supply = itertools.chain([9000, 9000, 9001, 9001, 9000], itertools.count(9002))
-
-    def fake_get_open_ports_list(count):
-        return [next(supply) for _ in range(count)]
-
-    monkeypatch.setattr(ses, "get_open_ports_list", fake_get_open_ports_list)
-
-    server = ses.OmniMasterServer(
-        master_address="127.0.0.1",
-        master_port=9000,  # seed: a route must not reuse the registration port
-        stage_ids=[0, 1, 2],
-    )
-
-    ports = []
-    for sid in (0, 1, 2):
-        alloc = server.get_allocation(sid)
-        for addr in (
-            alloc.handshake_bind_address,
-            alloc.input_bind_address,
-            alloc.output_bind_address,
-        ):
-            ports.append(ses._port_from_zmq_address(addr))
-
-    assert len(ports) == len(set(ports)), f"duplicate route ports allocated: {ports}"
-    assert 9000 not in ports, "route reused the master registration port"
-
-
-def test_port_from_zmq_address_parsing():
-    from vllm_omni.engine.stage_engine_startup import _port_from_zmq_address
-
-    assert _port_from_zmq_address("tcp://127.0.0.1:34277") == 34277
-    assert _port_from_zmq_address(None) is None
-    assert _port_from_zmq_address("ipc:///tmp/sock") is None
-    assert _port_from_zmq_address("tcp://host:not-a-port") is None

tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py

@@ -198,37 +198,16 @@ def _assert_realtime_smoke(result: dict) -> None:
     assert result["output_sample_rate"] > 0
 
 
-def _assert_realtime_accuracy(
-    result: dict,
-    threshold: float = 0.8,
-    whisper_model_size: str = "large-v3",
-) -> None:
-    """Assert that whisper transcription of audio output matches model text.
-
-    Args:
-        result: Roundtrip result dict from ``_run_realtime_audio_roundtrip``.
-        threshold: Minimum cosine similarity between whisper text and model text.
-                   Default 0.8 is appropriate for sync mode (exact match expected).
-                   Use 0.65 for async_chunk mode where chunked upload with pacing
-                   delays introduces inherent variability in audio codec quality.
-        whisper_model_size: Whisper model used to transcribe the generated audio
-                   for the accuracy check. Defaults to ``large-v3``: the default
-                   ``small`` model mishears short Chinese TTS clips (observed:
-                   北京→韦京 and a dropped leading sentence, sim=0.443), which
-                   caused spurious sim<threshold failures under async_chunk codec
-                   variability even though audio generation was correct. large-v3
-                   transcribes these clips reliably, so a failure here now points
-                   at the model, not the ASR grader.
-    """
+def _assert_realtime_accuracy(result: dict) -> None:
     final_text = (result["transcription_text"] or "").strip()
     assert final_text, "Expected non-empty transcription (model text stream)"
 
     wav_out = _wav_bytes_from_pcm16(result["output_pcm"], result["output_sample_rate"])
-    whisper_text = convert_audio_bytes_to_text(wav_out, model_size=whisper_model_size).strip()
+    whisper_text = convert_audio_bytes_to_text(wav_out).strip()
     assert whisper_text, "Whisper returned empty string for synthesized output audio"
 
     sim = cosine_similarity_text(whisper_text.lower(), final_text.lower())
-    assert sim > threshold, (
+    assert sim > 0.8, (
         f"Output audio transcript should match model text (sim={sim:.3f}): "
         f"whisper={whisper_text!r}, model_text={final_text!r}"
     )
@@ -255,12 +234,7 @@ class TestQwen3OmniRealtimeWebSocket:
         )
 
         _assert_realtime_smoke(result)
-        # Use threshold=0.65 for async_chunk: chunked upload with send_delay_ms
-        # introduces inherent variability in audio codec quality. Whisper
-        # transcription is nondeterministic and may miss 1-2 words in async
-        # mode even when model text generation is correct.
-        # Sync mode uses the default 0.8 (exact match expected).
-        _assert_realtime_accuracy(result, threshold=0.65)
+        _assert_realtime_accuracy(result)
 
     @pytest.mark.advanced_model
     @pytest.mark.omni

tests/entrypoints/openai_api/test_serving_chat_speaker.py

@@ -87,8 +87,8 @@ def test_create_chat_completion_converts_value_error_to_error_response(mocker: M
     serving_chat.chat_template = None
     serving_chat.chat_template_content_format = "string"
     serving_chat.default_chat_template_kwargs = {}
-    serving_chat.online_renderer = mocker.MagicMock()
-    serving_chat.online_renderer.validate_chat_template.return_value = None
+    serving_chat.openai_serving_render = mocker.MagicMock()
+    serving_chat.openai_serving_render.validate_chat_template.return_value = None
     serving_chat._effective_chat_template_kwargs = mocker.MagicMock(return_value={})
     serving_chat._preprocess_chat = mocker.AsyncMock(
         side_effect=ValueError("Invalid speaker 'uncle_fu'. Supported: ethan, vivian")

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -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/entrypoints/test_stream_finish_reason.py

@@ -144,7 +144,7 @@ def _build_serving_chat():
         engine_client=mock_engine,
         models=models,
         response_role="assistant",
-        online_renderer=mock_render,
+        openai_serving_render=mock_render,
         request_logger=None,
         chat_template=None,
         chat_template_content_format="auto",

tests/model_executor/models/voxcpm2/test_talker_state_eviction.py

@@ -52,10 +52,6 @@ def _make_bare_talker():
     talker._cfm_buffers = None
     talker._last_audio_output_req_ids = []
     talker._batched_fsq_fusion_max_batch = 32
-    # Added in VoxCPM2TalkerForConditionalGeneration.__init__; set to
-    # default (False) to avoid AttributeError when test exercises forward
-    # path that checks this flag.
-    talker._enable_unified_decode_graph = False
     return talker
 
 
@@ -333,27 +329,23 @@ class TestDecodeBatchContract:
 
         for (expected_res, expected_meta), (actual_res, actual_meta) in zip(expected, actual, strict=True):
             torch.testing.assert_close(actual_res, expected_res)
-            torch.testing.assert_close(actual_meta.new_lm_hidden, expected_meta.new_lm_hidden)
+            torch.testing.assert_close(actual_meta["new_lm_hidden"], expected_meta["new_lm_hidden"])
         assert [state.decode_step_count for state in batch_states] == [1, 1]
 
     def test_batch_decode_matches_per_request_state_updates(self) -> None:
-        from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
-            _DecodeResidualMeta,
-        )
-
         seq_talker = self._make_decode_talker()
         batch_talker = self._make_decode_talker()
 
         seq_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
         batch_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
         metas = [
-            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[1.0, 2.0, 3.0, 4.0]])),
-            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[5.0, 6.0, 7.0, 8.0]])),
+            {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])},
+            {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])},
         ]
         batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]])
 
         for state, meta, res_out in zip(seq_states, metas, batch_out, strict=True):
-            seq_talker._finish_decode(state, meta, res_out.unsqueeze(0))
+            seq_talker._finish_decode(state, meta, res_out.unsqueeze(0), torch.device("cpu"))
 
         batch_talker._finish_decode_batch(
             [(state, False, meta) for state, meta in zip(batch_states, metas, strict=True)],
@@ -368,16 +360,12 @@ class TestDecodeBatchContract:
             torch.testing.assert_close(actual.last_audio_patch_gpu, expected.last_audio_patch_gpu)
 
     def test_batch_decode_preserves_per_request_cfm_rng_order(self) -> None:
-        from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
-            _DecodeResidualMeta,
-        )
-
         talker = self._make_decode_talker()
         talker._enable_batched_cfm = False
         states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
         metas = [
-            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[1.0, 2.0, 3.0, 4.0]])),
-            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[5.0, 6.0, 7.0, 8.0]])),
+            {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])},
+            {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])},
         ]
         batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]])
 

tests/plugin/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

tests/plugin/test_gguf_diffusion.py

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

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

@@ -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/core/sched/omni_ar_scheduler.py

@@ -314,6 +314,13 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
         spec_decoding_stats: SpecDecodingStats | None = None
+        kv_connector_stats: KVConnectorStats | None = (
+            kv_connector_output.kv_connector_stats if kv_connector_output else None
+        )
+        if kv_connector_stats and self.connector:
+            kv_stats = self.connector.get_kv_connector_stats()
+            if kv_stats:
+                kv_connector_stats = kv_connector_stats.aggregate(kv_stats)
 
         failed_kv_load_req_ids = None
         if kv_connector_output and kv_connector_output.invalid_block_ids:
@@ -556,20 +563,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         if kv_connector_output:
             self._update_from_kv_xfer_finished(kv_connector_output)
 
-        # Worker-side KV connector stats from the model runner output.
-        kv_connector_stats: KVConnectorStats | None = (
-            kv_connector_output.kv_connector_stats if kv_connector_output else None
-        )
-        if self.connector:
-            # Scheduler-side KV connector stats collected after connector update.
-            scheduler_kv_connector_stats = self.connector.get_kv_connector_stats()
-            if scheduler_kv_connector_stats is not None and not scheduler_kv_connector_stats.is_empty():
-                kv_connector_stats = (
-                    kv_connector_stats.aggregate(scheduler_kv_connector_stats)
-                    if kv_connector_stats is not None
-                    else scheduler_kv_connector_stats
-                )
-
         # collect KV cache events from KV cache manager
         events = self.kv_cache_manager.take_events()
 

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -431,6 +431,14 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
         spec_decoding_stats: SpecDecodingStats | None = None
+        kv_connector_stats: KVConnectorStats | None = (
+            kv_connector_output.kv_connector_stats if kv_connector_output else None
+        )
+        # Merge connector-side stats (align with v0.14.0)
+        if kv_connector_stats and self.connector:
+            kv_stats = self.connector.get_kv_connector_stats()
+            if kv_stats:
+                kv_connector_stats = kv_connector_stats.aggregate(kv_stats)
 
         failed_kv_load_req_ids = None
         if kv_connector_output and getattr(kv_connector_output, "invalid_block_ids", None):
@@ -638,20 +646,6 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         if kv_connector_output:
             self._update_from_kv_xfer_finished(kv_connector_output)
 
-        # Worker-side KV connector stats from the model runner output.
-        kv_connector_stats: KVConnectorStats | None = (
-            kv_connector_output.kv_connector_stats if kv_connector_output else None
-        )
-        if self.connector:
-            # Scheduler-side KV connector stats collected after connector update.
-            scheduler_kv_connector_stats = self.connector.get_kv_connector_stats()
-            if scheduler_kv_connector_stats is not None and not scheduler_kv_connector_stats.is_empty():
-                kv_connector_stats = (
-                    kv_connector_stats.aggregate(scheduler_kv_connector_stats)
-                    if kv_connector_stats is not None
-                    else scheduler_kv_connector_stats
-                )
-
         # Collect and publish KV cache events (align with v0.14.0)
         events = self.kv_cache_manager.take_events()
         if self.connector is not None:

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -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)
@@ -397,77 +392,16 @@ class DiffusersPipelineLoader:
 
         return model.eval()
 
-    @staticmethod
-    def _has_online_quant(model: nn.Module) -> bool:
-        """Whether any layer uses an online-quant method that defers weight
-        materialization onto the ``meta`` device (upstream vLLM
-        ``uses_meta_device=True``, e.g. online FP8)."""
-        for module in model.modules():
-            quant_method = getattr(module, "quant_method", None)
-            if getattr(quant_method, "uses_meta_device", False):
-                return True
-        return False
-
     def _process_weights_after_loading(self, model: nn.Module, target_device: torch.device) -> None:
         """Process weights after loading for quantization methods.
 
         This handles vLLM's quantization methods that need to process weights
         after loading (e.g., FP8 online quantization from BF16/FP16 weights).
         """
-        # Newer upstream vLLM online-quant methods (uses_meta_device=True) create
-        # weights on the ``meta`` device and materialize them just-in-time as each
-        # layer's weights finish loading (via the layerwise online-process loader).
-        # Any "straggler" layers whose weights were not fully materialized during
-        # load (padded / partially-loaded layers) remain on ``meta``. Upstream's
-        # base_loader calls finalize_layerwise_processing() to materialize them;
-        # the diffusion loader must mirror that, otherwise the module.to() below
-        # raises "Cannot copy out of meta tensor; no data!". This whole meta-device
-        # handling is gated on online quant actually being in use, so that the
-        # proven code path for everything else (in particular FSDP/HSDP-sharded
-        # params, whose per-parameter .data cannot be cross-device reassigned) is
-        # left untouched. Import lazily so older vLLM (no meta-device quant) is
-        # unaffected.
-        has_online_quant = self._has_online_quant(model)
-        if has_online_quant:
-            from vllm.model_executor.model_loader.reload.layerwise import (
-                finalize_layerwise_processing,
-            )
-
-            # model_config is only dereferenced by finalize for vLLM Attention /
-            # MLAAttention layers; diffusion DiT models use their own attention and
-            # have none, so passing None is safe here.
-            finalize_layerwise_processing(model, model_config=None)
-
         for _, module in model.named_modules():
             quant_method = getattr(module, "quant_method", None)
-            if quant_method is None or not isinstance(quant_method, QuantizeMethodBase):
-                continue
-
-            if has_online_quant:
-                # Online quant may leave straggler params on the ``meta`` device.
-                # Move only real (non-meta) params onto the target device for
-                # processing and restore them afterward, mirroring upstream vLLM's
-                # device_loading_context — a blanket module.to(target_device) would
-                # raise NotImplementedError on meta params. Online quant initializes
-                # on the accelerator, so params are normally already on the target
-                # device and this loop is a no-op move; the point is to skip meta.
-                original_devices: dict[str, torch.device] = {}
-                for name, param in module.named_parameters():
-                    if param.device.type != "meta" and param.device != target_device:
-                        original_devices[name] = param.device
-                        param.data = param.data.to(target_device)
-
-                quant_method.process_weights_after_loading(module)
-
-                # Restore pre-existing params to their original device; leave any
-                # newly created (e.g. quantized) params on the target device.
-                for name, param in module.named_parameters():
-                    if name in original_devices:
-                        param.data = param.data.to(original_devices[name])
-            else:
-                # No meta params possible here. Preserve the original FSDP/HSDP-aware
-                # whole-module move (module.to()), which correctly handles sharded
-                # DTensor params that per-parameter .data reassignment cannot.
+            if quant_method is not None and isinstance(quant_method, QuantizeMethodBase):
+                # Move module to target device for processing if needed
                 module_device = next(module.parameters(), None)
                 if module_device is not None:
                     module_device = module_device.device
@@ -515,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.
         """
@@ -528,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)
@@ -561,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

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

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

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

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

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

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

@@ -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/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py

@@ -19,7 +19,6 @@ from transformers.utils import cached_file
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.models.gr00t.configs.embodiment_configs import ModalityConfig
-from vllm_omni.diffusion.models.gr00t.configs.gr00t_n1d7 import Gr00tN1d7Config
 from vllm_omni.diffusion.models.gr00t.dataio.embodiment_tags import EmbodimentTag
 from vllm_omni.diffusion.models.gr00t.dataio.state_action.state_action_processor import StateActionProcessor
 from vllm_omni.diffusion.models.gr00t.dataio.utils import parse_modality_configs, to_json_serializable
@@ -553,9 +552,4 @@ class Gr00tN1d7Processor(ProcessorMixin):
         return cls(**processor_kwargs, transformers_loading_kwargs=transformers_loading_kwargs)
 
 
-# transformers >= 5.5 requires the config *class* (not a model_type string) as the
-# first arg to AutoProcessor.register — it does `key.__module__` internally. Passing
-# the string "Gr00tN1d7" raises AttributeError: 'str' object has no attribute
-# '__module__'. Match the AutoModel.register(Gr00tN1d7Config, ...) call in
-# modeling/gr00t_n1d7.py.
-AutoProcessor.register(Gr00tN1d7Config, Gr00tN1d7Processor)
+AutoProcessor.register("Gr00tN1d7", Gr00tN1d7Processor)

vllm_omni/engine/async_omni_engine.py

@@ -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")
@@ -1567,18 +1568,16 @@ class AsyncOmniEngine:
         # and can cause CUDA OOM for subsequent engine instances (especially
         # large models like BAGEL-7B-MoT whose weights alone consume ~134 GiB).
         #
-        # CuMemAllocator.sleep() is NOT idempotent — calling it on already-
-        # slept entries causes CUDA_ERROR_INVALID_VALUE at cumem_allocator
-        # cuMemRelease (double-free of the memory handle).  Use release_pools()
-        # instead, which is the designed cleanup path: it drops MemPool refs
-        # and lets the destructor/free path handle asleep entries correctly
-        # (returns a null handle so the C extension skips unmap/release).
+        # Discard mode (level=2) is correct at shutdown: there is no benefit to
+        # keeping a CPU backup when the engine is being torn down.
         try:
             from vllm.device_allocator.cumem import CuMemAllocator, cumem_available
 
             if cumem_available:
                 allocator = CuMemAllocator.get_instance()
-                allocator.release_pools()
+                # Sleep at level 2 discards all pool memory from the GPU
+                # without creating CPU backups — cheapest and fastest.
+                allocator.sleep()
                 logger.debug("[AsyncOmniEngine] Released CuMem memory pool during shutdown")
         except Exception:
             pass

vllm_omni/engine/output_processor.py

@@ -28,12 +28,6 @@ from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
 
-# Keys whose values may arrive as scalar tensors (e.g. sample rate) but
-# represent metadata, not concatenatable data.  When these appear in
-# ``MultimodalPayload.tensors`` they must be relocated to ``.metadata``
-# before consolidation so the tensor concatenation path does not attempt
-# a dimensioned cat on a 0-d tensor.
-_METADATA_TENSOR_KEYS: frozenset[str] = frozenset({"sr", "sample_rate", "audio_sample_rate"})
 
 # ---------------------------------------------------------------------------
 # Module-level helpers
@@ -218,18 +212,6 @@ class OmniRequestState(RequestState):
             modality = OutputModality.TEXT
         strategy = get_accumulation_strategy(modality)
 
-        # Move known metadata scalar tensors (e.g. sample rate) from
-        # tensors to metadata so they are handled by the REPLACE path
-        # below instead of failing tensor concatenation.
-        for key in _METADATA_TENSOR_KEYS:
-            v = self.mm_accumulated.tensors.pop(key, None)
-            if v is not None:
-                if isinstance(v, list):
-                    # Keep the last value (REPLACE semantics)
-                    self.mm_accumulated.metadata[key] = v[-1]
-                else:
-                    self.mm_accumulated.metadata[key] = v
-
         try:
             for k, v in list(self.mm_accumulated.tensors.items()):
                 if isinstance(v, list) and v and isinstance(v[0], torch.Tensor):

vllm_omni/engine/stage_engine_startup.py

@@ -138,21 +138,6 @@ class StageReplicaResources:
 # ---------------------------------------------------------------------------
 
 
-def _port_from_zmq_address(address: str | None) -> int | None:
-    """Extract the TCP port from a ``tcp://host:port`` ZMQ address, else ``None``.
-
-    Non-TCP transports (``ipc://``, ``inproc://``) and unparsable values return
-    ``None`` so callers can skip seeding them into the port-dedup set.
-    """
-    if not address:
-        return None
-    tail = address.rsplit(":", 1)[-1]
-    try:
-        return int(tail)
-    except ValueError:
-        return None
-
-
 class OmniMasterServer:
     """Registration server for single-stage engine startup."""
 
@@ -173,25 +158,10 @@ class OmniMasterServer:
         self._stage_configs: dict[StageRoute, Any] = {}
         self._stage_coordinator_addresses: dict[StageRoute, StageCoordinatorAddresses] = {}
         self._stage_config_events: dict[StageRoute, threading.Event] = {}
-        # Ports already handed out by *this* server. ``get_open_ports_list`` only
-        # guarantees uniqueness *within* a single call; a per-route call cannot
-        # see ports drawn for other routes, so two stages/replicas can draw the
-        # same ephemeral port and the second engine to ``bind()`` it dies with
-        # ``zmq.error.ZMQError: Address already in use``. Multi-stage models
-        # (e.g. Qwen3-Omni: thinker/talker/code2wav) allocate many routes and hit
-        # this intermittently. We dedup every allocated port against this set,
-        # seeded with the registration port (and coordinator ROUTER port) that
-        # are already bound on the same host.
-        self._allocated_ports: set[int] = set()
-        if master_port:
-            self._allocated_ports.add(int(master_port))
         # Coordinator ROUTER address echoed back in every registration reply
         # so OmniCoordClientForStage knows where to connect from inside the
         # engine subprocess.
         self._coordinator_router_address = coordinator_router_address
-        coord_port = _port_from_zmq_address(coordinator_router_address)
-        if coord_port is not None:
-            self._allocated_ports.add(coord_port)
         # Fires only for *newly assigned* (auto-assigned) replicas, not for
         # head-side pre-allocated slots that already have head-side clients.
         self._on_register = on_register
@@ -265,7 +235,7 @@ class OmniMasterServer:
 
         self._stage_config_events[route] = threading.Event()
         self._stage_coordinator_addresses[route] = StageCoordinatorAddresses()
-        hs_port, inp_port, out_port = self._alloc_unique_ports(3)
+        hs_port, inp_port, out_port = get_open_ports_list(count=3)
         alloc = StageAllocation(
             handshake_bind_address=f"tcp://{self._address}:{hs_port}",
             handshake_connect_address=f"tcp://{self._address}:{hs_port}",
@@ -277,32 +247,6 @@ class OmniMasterServer:
         self._stage_routes[route] = alloc
         return alloc
 
-    def _alloc_unique_ports(self, count: int) -> list[int]:
-        """Return ``count`` open ports unique across every route this server owns.
-
-        ``get_open_ports_list`` dedups only within its own call, so we redraw any
-        port that collides with one already recorded in ``self._allocated_ports``
-        (registration/coordinator ports plus every previously-allocated route
-        port) and register the winners before returning. Callers already hold the
-        relevant allocation context (``__init__`` is single-threaded;
-        registration holds ``self._alloc_lock``).
-        """
-        picked: list[int] = []
-        # Bounded retry budget so a pathological host (near port exhaustion)
-        # fails loudly instead of spinning forever.
-        for _ in range(64):
-            for port in get_open_ports_list(count=count - len(picked)):
-                if port in self._allocated_ports:
-                    continue
-                self._allocated_ports.add(port)
-                picked.append(port)
-                if len(picked) == count:
-                    return picked
-        raise RuntimeError(
-            f"[OmniMasterServer] Could not allocate {count} unique open ports "
-            f"(host={self._address}); {len(self._allocated_ports)} already in use."
-        )
-
     def _next_free_replica_id(self, stage_id: int) -> int:
         """Return the next replica id to assign for an auto-assign registration.
 

vllm_omni/entrypoints/cli/serve.py

@@ -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/entrypoints/openai/api_server.py

@@ -32,7 +32,7 @@ from starlette.routing import Route
 from starlette.types import ASGIApp, Receive, Scope, Send
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
-from vllm.entrypoints.chat_utils import ChatTemplateConfig, load_chat_template
+from vllm.entrypoints.chat_utils import load_chat_template
 from vllm.entrypoints.launcher import serve_http, terminate_if_errored
 from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer
 from vllm.entrypoints.openai.api_server import build_app as build_openai_app
@@ -60,11 +60,12 @@ from vllm.entrypoints.pooling.classify.serving import ServingClassification
 from vllm.entrypoints.pooling.embed.serving import ServingEmbedding as OpenAIServingEmbedding
 from vllm.entrypoints.pooling.pooling.serving import ServingPooling
 from vllm.entrypoints.pooling.scoring.serving import ServingScores
-from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens
+from vllm.entrypoints.serve.disagg.serving import ServingTokens
 
 # vLLM moved `base` from openai.basic.api_router to serve.instrumentator.basic.
 # Keep a fallback for older/newer upstream layouts during rebase windows.
 from vllm.entrypoints.serve.instrumentator.basic import base
+from vllm.entrypoints.serve.render.serving import OpenAIServingRender
 from vllm.entrypoints.serve.tokenize.serving import ServingTokenization
 from vllm.entrypoints.serve.utils.api_utils import (
     load_aware_call,
@@ -84,7 +85,6 @@ from vllm.entrypoints.speech_to_text.translation.serving import (
     OpenAIServingTranslation,
 )
 from vllm.logger import init_logger
-from vllm.renderers.online_renderer import OnlineRenderer
 from vllm.tasks import POOLING_TASKS
 from vllm.tool_parsers import ToolParserManager
 from vllm.utils import random_uuid
@@ -451,7 +451,7 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None:
     # Add process-specific prefix to stdout and stderr.
     decorate_logs("APIServer", skip_if_decorated=True)
 
-    listen_address, sock = setup_openai_server(args, reuse_port=False)
+    listen_address, sock = setup_openai_server(args)
 
     # Unified use of omni_run_server_worker, AsyncOmni automatically handles LLM and Diffusion models
     await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs)
@@ -830,12 +830,6 @@ async def omni_init_app_state(
         if tokenizer is None or getattr(tokenizer, "chat_template", None) is None:
             resolved_chat_template = _load_model_chat_template_json(args.model)
 
-    chat_template_config = ChatTemplateConfig(
-        chat_template=resolved_chat_template,
-        chat_template_content_format=args.chat_template_content_format,
-        trust_request_chat_template=args.trust_request_chat_template,
-    )
-
     if args.tool_server == "demo":
         tool_server: ToolServer | None = DemoToolServer()
         assert isinstance(tool_server, DemoToolServer)
@@ -901,10 +895,11 @@ async def omni_init_app_state(
     await state.openai_serving_models.init_static_loras()
 
     # NOTE: kept aligned with upstream `init_app_state`:
-    # Use OnlineRenderer (replaced OpenAIServingRender which was removed upstream).
-    state.online_renderer = OnlineRenderer(
+    # Use OpenAIServingRender (replaces the old OnlineRenderer + OnlineDerenderer + ServingRender split).
+    state.openai_serving_render = OpenAIServingRender(
         model_config=engine_client.model_config,
         renderer=engine_client.renderer,
+        model_registry=state.openai_serving_models.registry,
         request_logger=request_logger,
         chat_template=resolved_chat_template,
         chat_template_content_format=args.chat_template_content_format,
@@ -921,7 +916,7 @@ async def omni_init_app_state(
         OpenAIServingResponses(
             engine_client,
             state.openai_serving_models,
-            state.online_renderer,
+            openai_serving_render=state.openai_serving_render,
             request_logger=request_logger,
             chat_template=resolved_chat_template,
             chat_template_content_format=args.chat_template_content_format,
@@ -942,7 +937,7 @@ async def omni_init_app_state(
             engine_client,
             state.openai_serving_models,
             args.response_role,
-            online_renderer=state.online_renderer,
+            openai_serving_render=state.openai_serving_render,
             request_logger=request_logger,
             chat_template=resolved_chat_template,
             chat_template_content_format=args.chat_template_content_format,
@@ -969,7 +964,7 @@ async def omni_init_app_state(
         OpenAIServingCompletion(
             engine_client,
             state.openai_serving_models,
-            online_renderer=state.online_renderer,
+            openai_serving_render=state.openai_serving_render,
             request_logger=request_logger,
             return_tokens_as_token_ids=args.return_tokens_as_token_ids,
             enable_prompt_tokens_details=args.enable_prompt_tokens_details,
@@ -984,7 +979,9 @@ async def omni_init_app_state(
             state.openai_serving_models,
             supported_tasks=tuple(supported_tasks),
             request_logger=request_logger,
-            chat_template_config=chat_template_config,
+            chat_template=resolved_chat_template,
+            chat_template_content_format=args.chat_template_content_format,
+            trust_request_chat_template=args.trust_request_chat_template,
         )
         if any(task in POOLING_TASKS for task in supported_tasks)
         else None
@@ -994,7 +991,9 @@ async def omni_init_app_state(
             engine_client,
             state.openai_serving_models,
             request_logger=request_logger,
-            chat_template_config=chat_template_config,
+            chat_template=resolved_chat_template,
+            chat_template_content_format=args.chat_template_content_format,
+            trust_request_chat_template=args.trust_request_chat_template,
         )
         if "embed" in supported_tasks
         else None
@@ -1004,7 +1003,9 @@ async def omni_init_app_state(
             engine_client,
             state.openai_serving_models,
             request_logger=request_logger,
-            chat_template_config=chat_template_config,
+            chat_template=resolved_chat_template,
+            chat_template_content_format=args.chat_template_content_format,
+            trust_request_chat_template=args.trust_request_chat_template,
         )
         if "classify" in supported_tasks
         else None
@@ -1013,9 +1014,8 @@ async def omni_init_app_state(
         ServingScores(
             engine_client,
             state.openai_serving_models,
-            supported_tasks=tuple(supported_tasks),
             request_logger=request_logger,
-            chat_template_config=chat_template_config,
+            score_template=resolved_chat_template,
             log_error_stack=args.log_error_stack,
         )
         if any(t in supported_tasks for t in ("embed", "score", "token_embed"))
@@ -1023,7 +1023,7 @@ async def omni_init_app_state(
     )
     state.serving_tokenization = ServingTokenization(
         state.openai_serving_models,
-        state.online_renderer,
+        state.openai_serving_render,
         request_logger=request_logger,
         chat_template=resolved_chat_template,
         chat_template_content_format=args.chat_template_content_format,
@@ -1055,7 +1055,7 @@ async def omni_init_app_state(
             engine_client,
             state.openai_serving_models,
             args.response_role,
-            online_renderer=state.online_renderer,
+            openai_serving_render=state.openai_serving_render,
             request_logger=request_logger,
             chat_template=resolved_chat_template,
             chat_template_content_format=args.chat_template_content_format,
@@ -1074,7 +1074,7 @@ async def omni_init_app_state(
         ServingTokens(
             engine_client,
             state.openai_serving_models,
-            state.online_renderer,
+            state.openai_serving_render,
             request_logger=request_logger,
             return_tokens_as_token_ids=args.return_tokens_as_token_ids,
             enable_prompt_tokens_details=args.enable_prompt_tokens_details,

vllm_omni/entrypoints/openai/serving_audio_generate.py

@@ -3,7 +3,7 @@ import asyncio
 import torch
 from fastapi import Request
 from fastapi.responses import Response
-from vllm.entrypoints.generate.base.serving import GenerateBaseServing as OpenAIServing
+from vllm.entrypoints.openai.engine.serving import OpenAIServing
 from vllm.logger import init_logger
 from vllm.utils import random_uuid
 

vllm_omni/entrypoints/openai/serving_chat.py

@@ -41,7 +41,6 @@ except ImportError:
     soundfile = None
 
 
-from vllm.entrypoints.generate.base.serving import clamp_prompt_logprobs
 from vllm.entrypoints.launcher import terminate_if_errored
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionNamedToolChoiceParam,
@@ -65,6 +64,7 @@ from vllm.entrypoints.openai.engine.protocol import (
     ToolCall,
     UsageInfo,
 )
+from vllm.entrypoints.openai.engine.serving import clamp_prompt_logprobs
 from vllm.entrypoints.openai.parser.harmony_utils import (
     get_streamable_parser_for_assistant,
 )
@@ -150,21 +150,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
     # Harmony flag (always False for vllm-omni models)
     use_harmony: bool = False
 
-    @property
-    def tool_call_id_type(self) -> str:
-        """Return the tool call ID type, delegating to model config.
-
-        Upstream vLLM removed the stored ``tool_call_id_type`` attribute
-        from ``OpenAIServingChat`` after the ParserManager refactor; the
-        field is now resolved on demand via ``get_tool_call_id_type``.
-        """
-        try:
-            from vllm.entrypoints.chat_utils import get_tool_call_id_type
-
-            return get_tool_call_id_type(self.model_config)
-        except Exception:
-            return "random"
-
     def _should_stream_with_auto_tool_parsing(self, request: ChatCompletionRequest) -> bool:
         """Check if streamed tokens should go through the tool-call parser.
 
@@ -383,7 +368,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 tool_dicts = [tool.model_dump() for tool in request.tools]
 
             if not self.use_harmony:
-                error_check_ret = self.online_renderer.validate_chat_template(
+                error_check_ret = self.openai_serving_render.validate_chat_template(
                     request_chat_template=request.chat_template,
                     chat_template_kwargs=request.chat_template_kwargs,
                     trust_request_chat_template=self.trust_request_chat_template,
@@ -411,7 +396,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 )
             else:
                 should_include_tools = tool_dicts is not None
-                conversation, engine_prompts = self.online_renderer._make_request_with_harmony(
+                conversation, engine_prompts = self.openai_serving_render._make_request_with_harmony(
                     request, should_include_tools
                 )
 

vllm_omni/entrypoints/openai/serving_speech.py

@@ -21,12 +21,12 @@ import torch
 from fastapi import HTTPException, Request, UploadFile
 from fastapi.responses import Response, StreamingResponse
 from transformers.utils.hub import cached_file
-from vllm.entrypoints.generate.base.serving import GenerateBaseServing as OpenAIServing
 from vllm.entrypoints.launcher import terminate_if_errored
 from vllm.entrypoints.openai.engine.protocol import (
     ErrorResponse,
     RequestResponseMetadata,
 )
+from vllm.entrypoints.openai.engine.serving import OpenAIServing
 from vllm.inputs import tokens_input
 from vllm.logger import init_logger
 from vllm.multimodal.media import MediaConnector

vllm_omni/experimental/ar_diffusion/kv_cache/paged.py

@@ -23,7 +23,22 @@ from dataclasses import dataclass
 import torch
 from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager
 from vllm.v1.kv_cache_interface import SlidingWindowSpec
-from vllm.v1.kv_cache_spec_registry import register_kv_cache_spec
+
+try:
+    from vllm.v1.kv_cache_spec_registry import register_kv_cache_spec
+except ModuleNotFoundError:
+
+    def register_kv_cache_spec(*, manager_class, uniform_type_base_spec=None):
+        """Compatibility shim for vLLM versions without kv_cache_spec_registry."""
+
+        def decorator(spec_class):
+            from vllm.v1.core.single_type_kv_cache_manager import spec_manager_map
+
+            spec_manager_map[spec_class] = manager_class
+            return spec_class
+
+        return decorator
+
 
 # ── Slot mapping ────────────────────────────────────────────────────────────
 

vllm_omni/patch.py

@@ -417,70 +417,3 @@ def _patch_fp8_use_quack_fused_bias():
 
 
 _patch_fp8_use_quack_fused_bias()
-
-
-# =============================================================================
-# Patch CuMemAllocator._python_free_callback to fix CUDA double-free on shutdown
-# =============================================================================
-# WHY: CuMemAllocator._python_free_callback guards the asleep-entry double-free
-# skip with ``if data.is_asleep and current_platform.is_rocm():`` — only ROCm
-# gets the safe empty-handle return.  On CUDA, the callback falls through and
-# returns the original handle, causing cuMemRelease on already-freed memory
-# (CUDA_ERROR_INVALID_VALUE) during EngineCore subprocess atexit cleanup.
-#
-# This happens because ``sleep()`` calls ``unmap_and_release()`` on ALL
-# platforms, then sets ``is_asleep = True``.  When the atexit handler
-# (``_shutdown_singleton`` -> ``release_pools()`` -> GC) triggers the free
-# callback, the asleep entries must return an empty chunk list so the C
-# extension skips ``cuMemRelease`` — exactly the same logic already used by
-# the ROCm guard.
-#
-# The fix removes the ``current_platform.is_rocm()`` condition so the guard
-# applies to CUDA (and any future platform that implements cumem).
-#
-# FRAGILITY: Relies on ``_python_free_callback`` being a regular method
-# (not a slot or C extension).  If CuMemAllocator is rewritten in C/Cython,
-# this monkey-patch will silently become a no-op.
-def _patch_cumem_free_callback_cuda() -> None:
-    try:
-        from vllm.device_allocator.cumem import CuMemAllocator
-    except ImportError:
-        _PATCH_LOGGER.debug("[cumem-cuda] CuMemAllocator not available; skipping patch")
-        return
-
-    _original_free_callback = CuMemAllocator._python_free_callback
-
-    if getattr(_original_free_callback, "_omni_cumem_cuda_patched", False):
-        return
-
-    # The upstream bug: `_python_free_callback` only skips the double-free
-    # for ROCm (line ~206).  We wrap the method to extend the guard to all
-    # platforms.
-    def _patched_free_callback(self, ptr: int) -> tuple:
-        data = self.pointer_to_data.pop(ptr)
-        if data.cpu_backup_tensor is not None:
-            data.cpu_backup_tensor = None
-        if data.is_asleep:
-            # sleep() already called unmap_and_release() on this allocation.
-            # Return an empty chunk list so the C extension skips
-            # cuMemRelease, avoiding a double-free.  Same logic as the
-            # existing ROCm guard, but applied to all platforms.
-            device, size, d_mem, _ = data.handle
-            result = (device, size, d_mem, [])
-            _PATCH_LOGGER.debug(
-                "[cumem-cuda] Free callback: asleep entry %s -> empty handle",
-                ptr,
-            )
-            return result
-        # Drain pending kernels before the C extension's cuMemUnmap.
-        torch.accelerator.synchronize(data.handle[0])
-        return data.handle
-
-    _patched_free_callback._omni_cumem_cuda_patched = True
-    CuMemAllocator._python_free_callback = _patched_free_callback
-    _PATCH_LOGGER.info(
-        "[cumem-cuda] CuMemAllocator._python_free_callback patched: asleep guard extended to all platforms."
-    )
-
-
-_patch_cumem_free_callback_cuda()

vllm_omni/plugins/__init__.py

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

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

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

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

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

vllm_omni/worker/gpu_ar_model_runner.py

@@ -159,11 +159,6 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
         async_output_copy_stream = kwargs.pop("async_output_copy_stream")
         vocab_size = kwargs.pop("vocab_size")
         routed_experts = kwargs.pop("routed_experts", None)
-        # Upstream AsyncGPUModelRunnerOutput added check_ep_fault / _has_fault
-        # for EP all2all fault tolerance (PR #43637). Omni doesn't use this
-        # feature but must consume the kwarg to prevent TypeError from stray
-        # kwargs and initialize the attribute so super().get_output() works.
-        kwargs.pop("check_ep_fault", False)
         if kwargs:
             raise TypeError(f"Unexpected OmniAsyncGPUModelRunnerOutput kwargs: {sorted(kwargs)}")
 
@@ -175,7 +170,6 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
         self.vocab_size = vocab_size
         self._logprobs_tensors = logprobs_tensors
         self._routed_experts = routed_experts
-        self._has_fault: torch.Tensor | None = None
 
         default_stream = torch.cuda.current_stream()
         with torch.cuda.stream(async_output_copy_stream):
@@ -226,12 +220,6 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
             if background_exception is not None:
                 raise background_exception
         self._build_model_runner_output_once()
-        # Upstream AsyncGPUModelRunnerOutput.get_output() accesses
-        # self._has_fault for EP all2all fault tolerance (PR #43637).
-        # Ensure the attribute exists even when __init__ was bypassed
-        # (e.g. unit tests using object.__new__).
-        if not hasattr(self, "_has_fault"):
-            self._has_fault = None
         with record_function_or_nullcontext("omni_async_output:get_output/finalize_async_sampled_tokens"):
             return super().get_output()