Fix vLLM compilation guard for models without marker (#2518)

### What does this PR do?

Type of change: Bug fix

Makes the vLLM `disable_compilation` context manager support inner model
implementations that do not predefine a `do_not_compile` attribute,
including GLM-5.3. The context manager now installs the marker
temporarily and removes it afterward, while preserving and restoring
existing marker values for other vLLM models.

Adds regression coverage for both supported wrapper layouts:
`model.model` and `model.language_model.model`.

### Usage

```python
with disable_compilation(model):
    mtq.quantize(model, quant_cfg, forward_loop=calibrate_loop)
```

No caller changes are required.

### Testing

- Ran `tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py`:
24 passed with vLLM 0.28.
- Ran pre-commit on both changed files: all applicable hooks passed.
- Installed this branch into `vllm/vllm-openai:glm53-flash` on OCI-JHB
and served the GLM-5.3-Flash BF16 checkpoint with
`QUANT_CFG=NVFP4_DEFAULT_CFG`, TP=4, eager mode, and BF16 KV cache.
- GLM passed the previous `do_not_compile` failure point, inserted 1,700
quantizers, enabled 456 weight quantizers, reached a healthy API server,
and returned a relevant manual prompt response.

### Before your PR is "*Ready for review*"

Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).

Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).

- Is this change backward compatible?: ✅
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A — integration compatibility fix; no user-facing API change.
- Did you get Claude approval on this PR?: N/A

### Additional Information

Validated against GLM-5.3-Flash using ModelOpt commit
`869b64fcee0b20be323663449b00e8c52940a289`.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **Bug Fixes**
- Compilation settings are now handled across supported nested model
configurations and restored after calibration, including when errors
occur.
- Calibration inputs correctly exclude padding when an attention mask is
provided and reject empty sequences.
- vLLM warmup reserves the required cache space for supported tail-cache
configurations.
- Serving startup supports an alternate vLLM launcher import path when
the OpenAI entrypoint is unavailable.
- **Compatibility**
- The vLLM serving example now defaults to vLLM 0.30.0 and documents
tested support for Nemotron 3 Nano hybrid attention/Mamba serving on
vLLM 0.26.0 and 0.30.0.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Kinjal Patel <kinjalpravin@nvidia.com>

examples/vllm_serve/Dockerfile

@@ -1,4 +1,4 @@
-ARG VLLM_VERSION=0.28.0
+ARG VLLM_VERSION=0.30.0
 FROM vllm/vllm-openai:v${VLLM_VERSION}
 
 # Set environment variables

examples/vllm_serve/README.md

@@ -4,15 +4,15 @@ This is a simple example to demonstrate calibrating and serving ModelOpt fakequa
 
 Compared with realquant, fakequant is 2-5x slower, but doesn't require dedicated kernel support and facilitates research.
 
-The general fakequant example is tested with vLLM 0.9.0, 0.19.1, 0.26.0, and 0.28.0. The
+The general fakequant example is tested with vLLM 0.9.0, 0.19.1, 0.26.0, 0.28.0, 0.29.0 and 0.30.0. The
 compact NVFP4 attention worker documented below requires vLLM 0.15.0 or newer.
 
 ## Prepare environment
 
-Use the Dockerfile to build an environment with vLLM 0.28.0:
+Use the Dockerfile to build an environment with vLLM 0.30.0:
 
 ```bash
-docker build -f examples/vllm_serve/Dockerfile -t vllm-modelopt:v0.28.0 .
+docker build -f examples/vllm_serve/Dockerfile -t vllm-modelopt:v0.30.0 .
 ```
 
 To build the same environment with another tested vLLM release, override `VLLM_VERSION`:
@@ -26,7 +26,7 @@ For a direct installation from the ModelOpt repository root, install the tested
 release and the ModelOpt extras used by this example:
 
 ```bash
-python3 -m pip install "vllm==0.28.0"
+python3 -m pip install "vllm==0.30.0"
 python3 -m pip install -e ".[all,mlflow]"
 ```
 
@@ -57,8 +57,8 @@ Step 2: Run the following command, with all supported flag as `vllm serve`:
 python vllm_serve_fakequant.py <model_path> -tp 8 --host 0.0.0.0 --port 8000
 ```
 
-Hybrid attention/Mamba models such as Nemotron 3 Nano are supported on vLLM 0.26.0 and
-0.28.0. For example, calibrate and serve with NVFP4 KV-cache fakequant as follows:
+Hybrid attention/Mamba models such as Nemotron 3 Nano are supported on vLLM 0.26.0, 0.28.0, 0.29.0 and
+0.30.0. For example, calibrate and serve with NVFP4 KV-cache fakequant as follows:
 
 ```bash
 KV_QUANT_CFG=NVFP4_KV_CFG QUANT_CALIB_SIZE=512 \

examples/vllm_serve/vllm_ptq_utils.py

@@ -62,9 +62,28 @@ def _get_calibration_block_count(
             return num_blocks
 
     else:
+        try:
+            from vllm.v1.kv_cache_interface import KpoolTailSpec, UniformTypeKVCacheSpecs
+        except ImportError:
+            kpool_tail_spec_types = ()
+            uniform_kv_cache_spec_types = ()
+        else:
+            kpool_tail_spec_types = (KpoolTailSpec,)
+            uniform_kv_cache_spec_types = (UniformTypeKVCacheSpecs,)
 
         def block_count(num_tokens: int, kv_cache_spec: Any) -> int:
             """Calculate the current vLLM warmup block reservation."""
+            # KpoolTailSpec is a one-block circular scratch cache. The generic
+            # vLLM warmup helper sees its SlidingWindowSpec base and reserves
+            # one block per block_size tokens, overflowing the one-block table.
+            unwrapped_spec = (
+                kv_cache_spec.first_spec
+                if isinstance(kv_cache_spec, uniform_kv_cache_spec_types)
+                else kv_cache_spec
+            )
+            if isinstance(unwrapped_spec, kpool_tail_spec_types):
+                return 1
+
             # Calibration runs before model_state is initialized, so call the
             # underlying reservation policy rather than _warmup_block_counter.
             return _reserved_block_count(

examples/vllm_serve/vllm_serve_fakequant.py

@@ -57,10 +57,32 @@ from pathlib import Path
 import uvloop
 import vllm
 from packaging import version
-from vllm.entrypoints.openai.api_server import run_server
-from vllm.entrypoints.openai.cli_args import make_arg_parser
 from vllm_mlflow_utils import MLFLOW_ENV_VARS, add_mlflow_args, resolve_mlflow_args
 
+
+def _is_missing_entrypoint(error: ModuleNotFoundError, entrypoint: str) -> bool:
+    missing_module = error.name
+    return missing_module is not None and (
+        missing_module == entrypoint or entrypoint.startswith(f"{missing_module}.")
+    )
+
+
+run_server_entrypoint = "vllm.entrypoints.launchers.api_server.entry"
+try:
+    from vllm.entrypoints.launchers.api_server.entry import run_server
+except ModuleNotFoundError as error:
+    if not _is_missing_entrypoint(error, run_server_entrypoint):
+        raise
+    from vllm.entrypoints.openai.api_server import run_server
+
+arg_parser_entrypoint = "vllm.entrypoints.cli.serve"
+try:
+    from vllm.entrypoints.cli.serve import make_arg_parser
+except ModuleNotFoundError as error:
+    if not _is_missing_entrypoint(error, arg_parser_entrypoint):
+        raise
+    from vllm.entrypoints.openai.cli_args import make_arg_parser
+
 vllm_version = version.parse(vllm.__version__)
 if vllm_version <= version.parse("0.11.0"):
     from vllm.utils import FlexibleArgumentParser

modelopt/torch/quantization/plugins/vllm.py

@@ -239,30 +239,48 @@ _moe_fakequant_active: contextvars.ContextVar[bool] = contextvars.ContextVar(
 )
 
 
+def _iter_compile_targets(model):
+    """Yield every module in the model tree that exposes vLLM's compile marker."""
+    modules = getattr(model, "modules", None)
+    candidates = modules() if callable(modules) else (model,)
+    yield from (candidate for candidate in candidates if hasattr(candidate, "do_not_compile"))
+
+
 @contextmanager
 def disable_compilation(model):
-    """Disable compilation for a model.
+    """Temporarily disable every vLLM compile wrapper in a model tree.
 
     Args:
-        model: The model to disable compilation for.
+        model: The model whose compile wrappers should be disabled.
     """
-    do_not_compile = True
-    if hasattr(model, "model"):
-        do_not_compile = model.model.do_not_compile
-        model.model.do_not_compile = True
-    elif hasattr(model, "language_model"):
-        do_not_compile = model.language_model.model.do_not_compile
-        model.language_model.model.do_not_compile = True
-    else:
-        raise ValueError("Model does not have a model or language_model attribute")
+    targets = tuple(_iter_compile_targets(model))
+    if not targets:
+        warnings.warn(
+            f"{type(model).__name__} has no modules exposing vLLM's 'do_not_compile' marker, "
+            "so ModelOpt cannot dynamically disable torch.compile during calibration. This "
+            "is harmless when vLLM is already running in eager mode (for example, with "
+            "--enforce-eager or CompilationMode.NONE). Otherwise, calibration may enter a "
+            "compiled path; rerun with --enforce-eager or add vLLM compile-wrapper support "
+            "for this model.",
+            stacklevel=2,
+        )
+        yield
+        return
 
+    states = []
     try:
+        for target in targets:
+            had_do_not_compile = "do_not_compile" in vars(target)
+            previous_do_not_compile = target.do_not_compile
+            states.append((target, had_do_not_compile, previous_do_not_compile))
+            target.do_not_compile = True
         yield
     finally:
-        if hasattr(model, "model"):
-            model.model.do_not_compile = do_not_compile
-        elif hasattr(model, "language_model"):
-            model.language_model.model.do_not_compile = do_not_compile
+        for target, had_do_not_compile, previous_do_not_compile in reversed(states):
+            if had_do_not_compile:
+                target.do_not_compile = previous_do_not_compile
+            else:
+                vars(target).pop("do_not_compile", None)
 
 
 # vLLM Attention stores ``device``/``dtype`` as plain attrs; ``dtype`` may be a string

tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py

@@ -102,6 +102,105 @@ def _patch_vllm_imports(monkeypatch, modules):
     monkeypatch.setattr(builtins, "__import__", fake_import)
 
 
+def _launcher_import_modules():
+    """Build isolated vLLM module stubs for launcher import compatibility tests."""
+    entrypoints = SimpleNamespace(
+        current_run_server=Mock(name="current_run_server"),
+        legacy_run_server=Mock(name="legacy_run_server"),
+        current_arg_parser=Mock(name="current_arg_parser"),
+        legacy_arg_parser=Mock(name="legacy_arg_parser"),
+    )
+    modules = {
+        "uvloop": SimpleNamespace(run=Mock()),
+        "vllm": SimpleNamespace(__version__="0.30.0"),
+        "vllm_mlflow_utils": SimpleNamespace(
+            MLFLOW_ENV_VARS=set(),
+            add_mlflow_args=Mock(),
+            resolve_mlflow_args=Mock(),
+        ),
+        "vllm.entrypoints.launchers.api_server.entry": SimpleNamespace(
+            run_server=entrypoints.current_run_server
+        ),
+        "vllm.entrypoints.openai.api_server": SimpleNamespace(
+            run_server=entrypoints.legacy_run_server
+        ),
+        "vllm.entrypoints.cli.serve": SimpleNamespace(
+            make_arg_parser=entrypoints.current_arg_parser
+        ),
+        "vllm.entrypoints.openai.cli_args": SimpleNamespace(
+            make_arg_parser=entrypoints.legacy_arg_parser
+        ),
+        "vllm.utils.argparse_utils": SimpleNamespace(FlexibleArgumentParser=Mock()),
+        "vllm.executor.ray_distributed_executor": SimpleNamespace(
+            RayDistributedExecutor=SimpleNamespace(ADDITIONAL_ENV_VARS=set())
+        ),
+    }
+    return modules, entrypoints
+
+
+@pytest.mark.parametrize(
+    ("run_server_missing", "arg_parser_missing", "uses_legacy"),
+    [
+        (None, None, False),
+        (
+            "vllm.entrypoints.launchers.api_server.entry",
+            "vllm.entrypoints.cli.serve",
+            True,
+        ),
+        ("vllm.entrypoints", "vllm.entrypoints", True),
+    ],
+    ids=("current", "legacy", "missing-parent"),
+)
+def test_vllm_serve_entrypoint_layouts(
+    monkeypatch, run_server_missing, arg_parser_missing, uses_legacy
+):
+    """Resolve current entrypoints and valid legacy fallbacks."""
+    modules, entrypoints = _launcher_import_modules()
+    if run_server_missing is not None:
+        modules["vllm.entrypoints.launchers.api_server.entry"] = ModuleNotFoundError(
+            name=run_server_missing
+        )
+    if arg_parser_missing is not None:
+        modules["vllm.entrypoints.cli.serve"] = ModuleNotFoundError(name=arg_parser_missing)
+    _patch_vllm_imports(monkeypatch, modules)
+
+    launcher = _load_example_module("vllm_serve_fakequant")
+
+    expected_run_server = (
+        entrypoints.legacy_run_server if uses_legacy else entrypoints.current_run_server
+    )
+    expected_arg_parser = (
+        entrypoints.legacy_arg_parser if uses_legacy else entrypoints.current_arg_parser
+    )
+    assert launcher.run_server is expected_run_server
+    assert launcher.make_arg_parser is expected_arg_parser
+
+
+@pytest.mark.parametrize(
+    ("entrypoint", "fallback"),
+    [
+        (
+            "vllm.entrypoints.launchers.api_server.entry",
+            "vllm.entrypoints.openai.api_server",
+        ),
+        ("vllm.entrypoints.cli.serve", "vllm.entrypoints.openai.cli_args"),
+    ],
+    ids=("run-server", "argument-parser"),
+)
+def test_vllm_serve_entrypoint_dependency_error_propagates(monkeypatch, entrypoint, fallback):
+    """Do not replace a missing entrypoint dependency with a fallback import error."""
+    modules, _ = _launcher_import_modules()
+    dependency_error = ModuleNotFoundError(name="vllm_dependency")
+    modules[entrypoint] = dependency_error
+    modules[fallback] = AssertionError("fallback must not be imported")
+    _patch_vllm_imports(monkeypatch, modules)
+
+    with pytest.raises(ModuleNotFoundError) as raised:
+        _load_example_module("vllm_serve_fakequant")
+
+    assert raised.value is dependency_error
+
+
 def test_get_calibration_block_count_uses_vllm_028_reservation_helper(monkeypatch):
     """The current vLLM adapter must forward every warmup reservation argument."""
     module = _load_example_module("vllm_ptq_utils")
@@ -129,6 +228,42 @@ def test_get_calibration_block_count_uses_vllm_028_reservation_helper(monkeypatc
     )
 
 
+def test_get_calibration_block_count_reserves_one_kpool_tail_block(monkeypatch):
+    """Kpool's circular tail cache owns exactly one physical block per request."""
+    module = _load_example_module("vllm_ptq_utils")
+
+    class KpoolTailSpec:
+        pass
+
+    class UniformTypeKVCacheSpecs:
+        def __init__(self, first_spec):
+            self.first_spec = first_spec
+
+    reserved_block_count = Mock(return_value=32)
+    _patch_vllm_imports(
+        monkeypatch,
+        {
+            "vllm.v1.worker.gpu.warmup": SimpleNamespace(
+                _reserved_block_count=reserved_block_count
+            ),
+            "vllm.v1.kv_cache_interface": SimpleNamespace(
+                KpoolTailSpec=KpoolTailSpec,
+                UniformTypeKVCacheSpecs=UniformTypeKVCacheSpecs,
+            ),
+        },
+    )
+    model_runner = SimpleNamespace(
+        vllm_config=SimpleNamespace(num_lookahead_tokens=0),
+        max_model_len=1024,
+    )
+    block_count = module._get_calibration_block_count(model_runner)
+
+    assert block_count is not None
+    assert block_count(128, KpoolTailSpec()) == 1
+    assert block_count(128, UniformTypeKVCacheSpecs(KpoolTailSpec())) == 1
+    reserved_block_count.assert_not_called()
+
+
 def test_get_calibration_block_count_uses_vllm_026_reservation_policy(monkeypatch):
     """The vLLM 0.26 adapter must preserve its cross-attention and Mamba rules."""
     module = _load_example_module("vllm_ptq_utils")
@@ -448,6 +583,40 @@ def test_quant_vllm_attention_forward_skips_only_in_kernel_qv_quantization():
     assert attention.v_bmm_quantizer.call_count == 2
 
 
+def test_disable_compilation_warns_without_installing_marker():
+    """A non-compile-wrapped model remains unchanged while the no-op risk is visible."""
+    model = torch.nn.Module()
+
+    with pytest.warns(UserWarning, match="rerun with --enforce-eager"), disable_compilation(model):
+        assert not hasattr(model, "do_not_compile")
+
+    assert not hasattr(model, "do_not_compile")
+
+
+def test_disable_compilation_updates_all_markers_and_restores_after_error():
+    """Every language and vision compile wrapper is restored after an exceptional exit."""
+
+    class CompileWrappedModule(torch.nn.Module):
+        do_not_compile = False
+
+    model = CompileWrappedModule()
+    model.do_not_compile = False
+    model.vision_model = CompileWrappedModule()
+    model.vision_model.do_not_compile = True
+    model.language_model = CompileWrappedModule()
+
+    with pytest.raises(RuntimeError, match="quantization failed"), disable_compilation(model):
+        assert model.do_not_compile is True
+        assert model.vision_model.do_not_compile is True
+        assert model.language_model.do_not_compile is True
+        raise RuntimeError("quantization failed")
+
+    assert model.do_not_compile is False
+    assert model.vision_model.do_not_compile is True
+    assert model.language_model.do_not_compile is False
+    assert "do_not_compile" not in vars(model.language_model)
+
+
 def test_attention_kv_defaults_set_only_uncalibrated_dynamic_block16_quantizers():
     calibrated_amax = 7.25
     layer = SimpleNamespace(