fix(quant): load offline-quantized compressed-tensors checkpoints directly on CPU under offload

The diffusion loader treated any quant config without the MXFP8/
is_checkpoint_quantized markers as online quantization: it loaded every
component on the accelerator before offloading, so serving the int8
weight-only compressed-tensors artifact OOM'd at model construction
(encoder 17 GiB + VAE + packed DiT skeleton > 24 GiB). compressed-tensors
is always a serialized format with no on-load quantization pass — detect
it (through component wrappers too) and load on the host, round-tripping
each quantized module through the accelerator only for the post-load
Marlin repack, which has no CPU path.

tests/quantization/test_quant_factory.py

@@ -12,6 +12,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
 from vllm_omni.quantization import build_quant_config, resolve_component_quant_config
 from vllm_omni.quantization.component_config import (
     ComponentQuantizationConfig,
+    is_prequantized_config,
     resolve_encoder_quant_config,
 )
 from vllm_omni.quantization.factory import SUPPORTED_QUANTIZATION_METHODS, _build_single
@@ -75,3 +76,18 @@ def test_unknown_method_still_raises() -> None:
 
 def test_supported_list_advertises_hyphenated_registry_names() -> None:
     assert "compressed-tensors" in SUPPORTED_QUANTIZATION_METHODS
+
+
+def test_is_prequantized_config_recognizes_serialized_formats() -> None:
+    # compressed-tensors is always a serialized format: its linear methods
+    # consume packed checkpoint tensors and there is no on-load quantization.
+    assert is_prequantized_config(build_quant_config(_int8_weight_only_disk_config()))
+    component = build_quant_config({"transformer": _int8_weight_only_disk_config()})
+    assert isinstance(component, ComponentQuantizationConfig)
+    assert is_prequantized_config(component)
+    assert not is_prequantized_config(None)
+    # Online (non-serialized) methods stay unrecognized so the loader keeps
+    # the accelerator-side load-quantize-offload path for them.
+    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
+
+    assert not is_prequantized_config(DiffusionFp8Config())

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -63,6 +63,7 @@ from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery, Pipe
 from vllm_omni.diffusion.offloader.offload_plan import get_offload_plan
 from vllm_omni.diffusion.registry import initialize_model
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+from vllm_omni.quantization.component_config import is_prequantized_config
 from vllm_omni.transformers_utils.repo_utils import hf_api
 
 
@@ -695,16 +696,17 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
         }
         if load_format is None:
             load_format = "default"
-        # CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8),
-        # weights are already quantized in the checkpoint — load directly on CPU.
-        # For online quantization, load on device so quantization can run on accelerator,
-        # then move back to CPU afterward.
+        # CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8,
+        # compressed-tensors), weights are already quantized in the checkpoint — load
+        # directly on CPU. For online quantization, load on device so quantization can
+        # run on accelerator, then move back to CPU afterward.
         offload_after_quant = False
+        offline_quant_direct_cpu = False
         if load_device == "cpu" and self.quant_config is not None and device is not None:
             quant_cfg = self.quant_config
             is_offline = getattr(quant_cfg, "data_type", None) == "mx_fp" or getattr(
                 quant_cfg, "is_checkpoint_quantized", False
-            )
+            ) or is_prequantized_config(quant_cfg)
             if not is_offline:
                 load_device = device.type
                 offload_after_quant = True
@@ -713,6 +715,7 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
                     load_device,
                 )
             else:
+                offline_quant_direct_cpu = True
                 logger.info("Offline-quantized model with CPU offload, loading weights directly on CPU")
 
         target_device = torch.device(load_device)
@@ -874,7 +877,14 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
                         else:
                             self.load_weights(model)
                     self._maybe_fuse_distilled_lora(model)
-                    self._process_weights_after_loading(model, target_device)
+                    # Post-load transforms (e.g. the Marlin repack for
+                    # compressed-tensors) only run on the accelerator, so an
+                    # offline-quantized CPU-offload load round-trips each
+                    # quantized module through the device and back to the host.
+                    process_device = (
+                        device if offline_quant_direct_cpu and device is not None else target_device
+                    )
+                    self._process_weights_after_loading(model, process_device)
 
                 # A warm final-layout hit has already completed all
                 # byte-changing work through the restorer.  Shared runtime

vllm_omni/quantization/component_config.py

@@ -33,6 +33,26 @@ PRE_QUANTIZED_METHODS: frozenset[str] = frozenset(
 )
 
 
+def is_prequantized_config(quant_config: QuantizationConfig | None) -> bool:
+    """True when the config can only describe pre-quantized (serialized) weights.
+
+    Such checkpoints need no accelerator-side quantization pass at load, so a
+    CPU-offload load can read weights directly on the host. compressed-tensors
+    is the canonical case: its linear methods consume packed checkpoint tensors
+    (weight_packed/weight_scale) directly and vLLM has no on-load quantization
+    path for the format. A component wrapper is pre-quantized when every routed
+    component (and the default) is itself pre-quantized or unquantized.
+    """
+    if quant_config is None:
+        return False
+    if isinstance(quant_config, ComponentQuantizationConfig):
+        routed = [*quant_config.component_configs.values(), quant_config.default_config]
+        return any(c is not None for c in routed) and all(
+            c is None or is_prequantized_config(c) for c in routed
+        )
+    return quant_config.get_name() == "compressed-tensors"
+
+
 def resolve_component_quant_config(
     quant_config: QuantizationConfig | None,
     component: str,