Support vLLM-based Model Quantization with llm_compressor Export

auto_round/algorithms/pipeline.py

@@ -244,15 +244,41 @@ def forward_block_batch(
         return output
 
     def _run_block(self, block, quantizer, input_ids, input_others, device):
-        return quantizer._resolve_block_forward()(
-            block,
-            input_ids,
-            input_others,
-            quantizer.model_context.amp,
-            quantizer.model_context.amp_dtype,
-            device,
-            0,
-        )
+        try:
+            return quantizer._resolve_block_forward()(
+                block,
+                input_ids,
+                input_others,
+                quantizer.model_context.amp,
+                quantizer.model_context.amp_dtype,
+                device,
+                0,
+            )
+        except (AssertionError, RuntimeError) as e:
+            # vLLM attention layers require a ForwardContext that is only set by
+            # the vLLM engine (``llm.generate()``).  Directly calling block.forward
+            # outside the engine is not supported.
+            # Additionally, calibration captures may contain mixed prefill/decode
+            # batches with inconsistent token counts across keys (e.g. positions
+            # captured during a different batch from hidden_states), causing
+            # rotary_embedding to fail with "same number of tokens" RuntimeError.
+            # For RTN-style calibration (iters=0) the reference outputs are not
+            # used in loss-based optimisation; the caller (collect_reference) uses
+            # them only as next-block inputs, which the outer loop overwrites from
+            # ``all_inputs[next_block]`` on every iteration.
+            # Returning ``input_ids`` as a placeholder is therefore safe.
+            _msg = str(e)
+            _is_vllm_err = "Forward context is not set" in _msg or "same number of tokens" in _msg
+            if _is_vllm_err:
+                from auto_round.logger import logger as _ar_logger
+
+                _ar_logger.warning_once(
+                    "vLLM block forward skipped outside engine "
+                    "(ForwardContext unavailable or token-count mismatch in calib data); "
+                    "returning input_ids as placeholder.  This is expected for iters=0 RTN."
+                )
+                return input_ids
+            raise
 
     @torch.no_grad()
     def _collect_outputs(self, block, quantizer, *, source: InputSource, batch_size: int, save: bool = True):

auto_round/algorithms/quantization/base.py

@@ -132,6 +132,11 @@ def _immediate_pack_and_save_module(self, module_name):
             set_module(self.model, module_name, module)
         if self.compress_context.is_immediate_saving:
             module = get_module(self.model, module_name)
+            if module is None:
+                # Module was replaced or unfused under a new name during packing
+                # (e.g. vLLM qkv_proj → q_proj/k_proj/v_proj). Sub-modules are
+                # included in the enclosing block's shard write; skip here.
+                return
             module.to("cpu")
             shard_writer.write(module, module_name, False)
             module.to("meta")

auto_round/algorithms/quantization/rtn/quantizer.py

@@ -77,7 +77,7 @@ def quantize_block(self, ctx) -> dict:
             # stays aligned across experts.
             set_amax_for_all_moe_layers(block, attr_name="act_max")
 
-        for _name, m in block.named_modules():
+        for _name, m in list(block.named_modules()):
             if hasattr(m, "global_name") and check_to_quantized(m):
                 self.quantize_layer(m.global_name)
         return {}
@@ -154,8 +154,12 @@ def quantize_block(self, ctx):
         ):
             # enable moe experts act_max automatic generation for Linear
             set_amax_for_all_moe_layers(block, attr_name="act_max")
-        # Normalize imatrix and quantize layers
-        for name, m in block.named_modules():
+        # Normalize imatrix and quantize layers.
+        # Snapshot named_modules() before iterating: unfuse (e.g. qkv_proj →
+        # q_proj/k_proj/v_proj) inside quantize_layer_outside_block modifies
+        # the block's _modules dict and would raise RuntimeError if we iterated
+        # a live generator.
+        for name, m in list(block.named_modules()):
             if hasattr(m, "imatrix"):
                 m.imatrix /= m.imatrix_cnt
             if hasattr(m, "global_name") and check_to_quantized(m):

auto_round/calibration/hooks.py

@@ -147,7 +147,10 @@ def forward_capture(m, hidden_states=None, *positional_inputs, **kwargs):
                 if positional_inputs:
                     return m.orig_forward(hidden_states=hidden_states, *positional_inputs, **kwargs)
                 else:
-                    return m.orig_forward(hidden_states, **kwargs)
+                    # Use keyword argument to support models where hidden_states is not
+                    # the first parameter (e.g., vLLM's DecoderLayer uses
+                    # ``forward(positions, hidden_states, residual)``).
+                    return m.orig_forward(hidden_states=hidden_states, **kwargs)
             else:
                 # Currently only for Llama-3.2-Vision-Instruct Series
                 return m.orig_forward(*positional_inputs, **kwargs)
@@ -190,7 +193,12 @@ def replace_forward_with_hooks(state) -> None:
     """
 
     def register_hook(n, m, hook_handles):
-        if n in state.to_cached_layers and type(m) not in SUPPORTED_LAYER_TYPES:  # block
+        from auto_round.wrapper import _VLLMLinearBase
+
+        is_linear = isinstance(m, SUPPORTED_LAYER_TYPES) or (
+            _VLLMLinearBase is not None and isinstance(m, _VLLMLinearBase)
+        )
+        if n in state.to_cached_layers and not is_linear:  # block
             m.orig_forward = m.forward
             m.forward = partial(state._get_block_forward_func(n), m)
         elif n in state.to_cached_layers:  # linear / conv1d layer

auto_round/cli/main.py

@@ -14,6 +14,7 @@
 
 import argparse
 import difflib
+import json
 import sys
 
 from auto_round.cli.algorithms import AlgorithmHandler
@@ -84,13 +85,35 @@ def _build_entry_compressor_kwargs(args) -> dict:
 
 
 def _build_entry_model_type_kwargs(args) -> dict:
+    vllm_model_kwargs = None
+    if args.vllm_model_kwargs is not None:
+        try:
+            parsed = json.loads(args.vllm_model_kwargs)
+        except json.JSONDecodeError as err:
+            raise ValueError(f"Invalid --vllm_model_kwargs JSON: {err}") from err
+        if not isinstance(parsed, dict):
+            raise ValueError("--vllm_model_kwargs must be a JSON object.")
+        vllm_model_kwargs = parsed
+
+    if args.enable_vllm_loading and vllm_model_kwargs is not None:
+        tp = vllm_model_kwargs.get("tensor_parallel_size", 1)
+        if tp != 1:
+            raise ValueError(
+                f"--enable_vllm_loading only supports single-GPU quantization "
+                f"(tensor_parallel_size=1), but got tensor_parallel_size={tp}. "
+                "Multi-GPU TP is not supported: the wrapper forward, unfuse logic, "
+                "and act_max collection all assume TP=1."
+            )
+
     return {
         "quant_nontext_module": args.quant_nontext_module,
         "extra_data_dir": args.extra_data_dir,
         "template": args.template,
         "guidance_scale": args.guidance_scale,
         "num_inference_steps": args.num_inference_steps,
         "generator_seed": args.generator_seed,
+        "enable_vllm_loading": args.enable_vllm_loading,
+        "vllm_model_kwargs": vllm_model_kwargs,
     }
 
 

auto_round/cli/parser.py

@@ -144,6 +144,17 @@ def build_quantize_parser(*, prog: str = "auto_round quantize") -> argparse.Argu
     rt.add_argument(
         "--layer_config", default=None, type=str, help="Per-layer quantization overrides encoded as JSON-like text."
     )
+    rt.add_argument(
+        "--enable_vllm_loading",
+        action="store_true",
+        help="Load model via vLLM engine before entering AutoRound compressor flow.",
+    )
+    rt.add_argument(
+        "--vllm_model_kwargs",
+        default=None,
+        type=str,
+        help="JSON string of extra kwargs forwarded to vLLM LLM(...), e.g. '{\"tensor_parallel_size\":1}'.",
+    )
     rt.add_argument(
         "--shared_layers",
         type=str,

auto_round/compressors/base.py

@@ -305,6 +305,8 @@ def __init__(
         model_dtype = kwargs.pop("model_dtype", None)
         trust_remote_code = kwargs.pop("trust_remote_code") if "trust_remote_code" in kwargs else True
         quant_nontext_module = kwargs.pop("quant_nontext_module", False)
+        enable_vllm_loading = kwargs.pop("enable_vllm_loading", False)
+        vllm_model_kwargs = kwargs.pop("vllm_model_kwargs", None)
         device = kwargs.pop("device", None)
         if device is not None:
             logger.warning("`device` is deprecated, please use `device_map` instead")
@@ -388,6 +390,8 @@ def __init__(
             formats=self.formats,
             is_act_quantize=self.quantize_config.is_act_quantize,
             quant_nontext_module=quant_nontext_module,
+            enable_vllm_loading=enable_vllm_loading,
+            vllm_model_kwargs=vllm_model_kwargs,
         )
         # Alternatively, you can use CompressContext.create_context
         self.compress_context = CompressContext(

auto_round/compressors/data_driven.py

@@ -718,7 +718,7 @@ def _quantize_blocks(
         if not self.compress_context.is_immediate_saving:
             self.model = mv_module_from_gpu(self.model)
         for n, m in self.model.named_modules():
-            if hasattr(m, "name"):
+            if "name" in m.__dict__:
                 delattr(m, "name")
 
         del q_input
@@ -875,8 +875,13 @@ def quantize(self) -> tuple[torch.nn.Module, dict[str, Any]]:
         # Dump a summary
         quantized_layers = []
         unquantized_layers = []
+        from auto_round.compressors.utils import _VLLMLinearBase as _VLLMBase
+
         for n, m in self.model_context.model.named_modules():
-            if isinstance(m, tuple(SUPPORTED_LAYER_TYPES)):
+            is_supported = isinstance(m, tuple(SUPPORTED_LAYER_TYPES)) or (
+                _VLLMBase is not None and isinstance(m, _VLLMBase)
+            )
+            if is_supported:
                 if check_to_quantized(m):
                     quantized_layers.append(n)
                 else:
@@ -1098,9 +1103,12 @@ def process_input_others(input_others):
                         input_others[key] = val.to(tmp_dtype)
                     elif isinstance(val, list):
                         input_others[key] = [
-                            to_dtype(v, tmp_dtype)
+                            (
+                                to_dtype(v, tmp_dtype)
+                                if not (isinstance(v, torch.Tensor) and v.dtype in (torch.int32, torch.int64))
+                                else v
+                            )
                             for v in val
-                            if not (isinstance(v, torch.Tensor) and v.dtype in (torch.int32, torch.int64))
                         ]
                 return input_others
 
@@ -1164,7 +1172,25 @@ def process_input_others(input_others):
                 )
                 with ExitStack() as fwd_stack:
                     self.pipeline.enter_block_forward_hooks(ctx, fwd_stack)
-                    input_ids = ctx.collect_reference(fwd_stack)
+                    try:
+                        input_ids = ctx.collect_reference(fwd_stack)
+                    except (AssertionError, RuntimeError, IndexError, TypeError) as _e:
+                        # vLLM blocks cannot be forward-called outside the engine
+                        # (ForwardContext is only set by llm.generate()).
+                        # Token-count mismatches and index errors in calibration
+                        # captures are also expected when prefill/decode batches
+                        # are mixed.  For RTN (iters=0), collect_reference outputs
+                        # are never used for gradient optimisation and the outer
+                        # loop overwrites input_ids from all_inputs[next_block]
+                        # anyway, so returning the current input_ids is safe.
+                        logger.warning_once(
+                            "vLLM collect_reference skipped (%s: %s); "
+                            "using block input as placeholder. "
+                            "This is expected for iters=0 RTN with vLLM loading.",
+                            type(_e).__name__,
+                            str(_e)[:120],
+                        )
+                        # input_ids stays as-is (already the pre-captured input)
 
                 if len(device_manager.device_list) > 1:
                     accelerate.hooks.remove_hook_from_submodules(block)

auto_round/compressors/entry.py

@@ -40,8 +40,14 @@
 }
 _ENTRY_MLLM_KWARGS = {"processor", "image_processor", "template", "extra_data_dir", "quant_nontext_module"}
 _ENTRY_DIFFUSION_KWARGS = {"guidance_scale", "num_inference_steps", "generator_seed"}
+_ENTRY_VLLM_KWARGS = {"enable_vllm_loading", "vllm_model_kwargs"}
 _ENTRY_ALLOWED_KWARGS = (
-    _ENTRY_ROUTE_KWARGS | _ENTRY_COMPRESSOR_KWARGS | _ENTRY_BASE_KWARGS | _ENTRY_MLLM_KWARGS | _ENTRY_DIFFUSION_KWARGS
+    _ENTRY_ROUTE_KWARGS
+    | _ENTRY_COMPRESSOR_KWARGS
+    | _ENTRY_BASE_KWARGS
+    | _ENTRY_MLLM_KWARGS
+    | _ENTRY_DIFFUSION_KWARGS
+    | _ENTRY_VLLM_KWARGS
 )
 
 
@@ -78,6 +84,7 @@ def _split_entry_kwargs(kwargs: dict[str, Any]) -> dict[str, dict[str, Any]]:
         "base": {},
         "mllm": {},
         "diffusion": {},
+        "vllm": {},
     }
     for key, value in kwargs.items():
         if key in _ENTRY_ROUTE_KWARGS:
@@ -90,6 +97,8 @@ def _split_entry_kwargs(kwargs: dict[str, Any]) -> dict[str, dict[str, Any]]:
             buckets["mllm"][key] = value
         elif key in _ENTRY_DIFFUSION_KWARGS:
             buckets["diffusion"][key] = value
+        elif key in _ENTRY_VLLM_KWARGS:
+            buckets["vllm"][key] = value
     return buckets
 
 
@@ -179,6 +188,10 @@ def _get_compressor_class(model_type: str, base_cls: type) -> type:
         from auto_round.compressors.diffusion_mixin import DiffusionMixin
 
         mixin = DiffusionMixin
+    elif model_type == "vllm":
+        from auto_round.compressors.vllm_mixin import VLLMMixin
+
+        mixin = VLLMMixin
     else:
         return base_cls
     combined = type(f"{model_type.capitalize()}{base_cls.__name__}", (mixin, base_cls), {})
@@ -233,19 +246,35 @@ def _resolve_quant_config_for_routing(alg_configs) -> tuple[list, list, Quantiza
     )
 
 
-def _build_model_type_ctor_kwargs(model, base_kwargs, mllm_kwargs, diffusion_kwargs) -> tuple[str, dict[str, Any]]:
+def _build_model_type_ctor_kwargs(
+    model,
+    base_kwargs,
+    mllm_kwargs,
+    diffusion_kwargs,
+    vllm_kwargs,
+) -> tuple[str, dict[str, Any]]:
     from auto_round.utils.model import detect_model_type
 
-    model_type = detect_model_type(model)
+    enable_vllm_loading = bool(vllm_kwargs.get("enable_vllm_loading", False))
+    detected_model_type = detect_model_type(model)
+    if enable_vllm_loading and detected_model_type in {"mllm", "diffusion"}:
+        raise ValueError(
+            "`enable_vllm_loading=True` conflicts with detected model type "
+            f"`{detected_model_type}`. Please disable vllm loading for multimodal/diffusion models."
+        )
+
+    model_type = "vllm" if enable_vllm_loading else detected_model_type
     has_multimodal_assets = mllm_kwargs.get("processor") is not None or mllm_kwargs.get("image_processor") is not None
-    if has_multimodal_assets and model_type != "mllm":
+    if has_multimodal_assets and not enable_vllm_loading and model_type != "mllm":
         model_type = "mllm"
 
     ctor_kwargs = dict(base_kwargs)
     if model_type == "mllm":
         ctor_kwargs.update(mllm_kwargs)
     if model_type == "diffusion":
         ctor_kwargs.update(diffusion_kwargs)
+    if model_type == "vllm":
+        ctor_kwargs.update(vllm_kwargs)
     return model_type, ctor_kwargs
 
 
@@ -349,6 +378,7 @@ def __new__(
         base_kwargs = dict(split_kwargs["base"])
         mllm_kwargs = dict(split_kwargs["mllm"])
         diffusion_kwargs = dict(split_kwargs["diffusion"])
+        vllm_kwargs = dict(split_kwargs["vllm"])
 
         # Resolve string alias(es) to config instance(s) before routing.
         alg_configs = cls._resolve_config(alg_configs)
@@ -410,7 +440,13 @@ def __new__(
             seqlen=seqlen,
             **compressor_kwargs,
         )
-        model_type, ctor_kwargs = _build_model_type_ctor_kwargs(model, base_kwargs, mllm_kwargs, diffusion_kwargs)
+        model_type, ctor_kwargs = _build_model_type_ctor_kwargs(
+            model,
+            base_kwargs,
+            mllm_kwargs,
+            diffusion_kwargs,
+            vllm_kwargs,
+        )
 
         # Preprocessor algorithms (AWQ, …) require a data-driven host so that
         # the per-block preprocessor lifecycle (prepare_block_group ->
@@ -647,11 +683,16 @@ def _build_entry_forward_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
             "num_inference_steps": kwargs.pop("num_inference_steps", 50),
             "generator_seed": kwargs.pop("generator_seed", None),
         }
+        vllm_kwargs = {
+            "enable_vllm_loading": kwargs.pop("enable_vllm_loading", False),
+            "vllm_model_kwargs": kwargs.pop("vllm_model_kwargs", None),
+        }
         return {
             "format": format_name,
             "rotation_config": rotation_config,
             **mllm_kwargs,
             **diffusion_kwargs,
+            **vllm_kwargs,
             **kwargs,
         }
 

auto_round/compressors/vllm/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2026 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.