diff --git a/examples/vllm_serve/Dockerfile b/examples/vllm_serve/Dockerfile index a4c26ab1a..8bb3f8c05 100644 --- a/examples/vllm_serve/Dockerfile +++ b/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 diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 9c4eec609..3c838b154 100644 --- a/examples/vllm_serve/README.md +++ b/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 -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 \ diff --git a/examples/vllm_serve/vllm_ptq_utils.py b/examples/vllm_serve/vllm_ptq_utils.py index ca9e276b9..bf9b0d85f 100644 --- a/examples/vllm_serve/vllm_ptq_utils.py +++ b/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( diff --git a/examples/vllm_serve/vllm_serve_fakequant.py b/examples/vllm_serve/vllm_serve_fakequant.py index 97134d268..0fe371f9d 100644 --- a/examples/vllm_serve/vllm_serve_fakequant.py +++ b/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 diff --git a/modelopt/torch/quantization/plugins/vllm.py b/modelopt/torch/quantization/plugins/vllm.py index 67191a0f3..e34986860 100644 --- a/modelopt/torch/quantization/plugins/vllm.py +++ b/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 diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py index 0f440b93d..ce0245872 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py +++ b/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(