diff --git a/examples/vllm_serve/Dockerfile b/examples/vllm_serve/Dockerfile index 5fe0799c427..a4c26ab1a7b 100644 --- a/examples/vllm_serve/Dockerfile +++ b/examples/vllm_serve/Dockerfile @@ -1,4 +1,5 @@ -FROM vllm/vllm-openai:v0.26.0 +ARG VLLM_VERSION=0.28.0 +FROM vllm/vllm-openai:v${VLLM_VERSION} # Set environment variables ENV PIP_NO_CACHE_DIR=off \ diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index fc4e8a0ebcc..bc2b0fb53f3 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -4,17 +4,35 @@ 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, and 0.26.0. The compact -NVFP4 attention worker documented below requires vLLM 0.15.0 or newer. +The general fakequant example is tested with vLLM 0.9.0, 0.19.1, 0.26.0, and 0.28.0. The +compact NVFP4 attention worker documented below requires vLLM 0.15.0 or newer. ## Prepare environment -Follow the following instruction to build a docker environment, or install vllm with pip. +Use the Dockerfile to build an environment with vLLM 0.28.0: ```bash -docker build -f examples/vllm_serve/Dockerfile -t vllm-modelopt . +docker build -f examples/vllm_serve/Dockerfile -t vllm-modelopt:v0.28.0 . ``` +To build the same environment with another tested vLLM release, override `VLLM_VERSION`: + +```bash +docker build --build-arg VLLM_VERSION=0.26.0 \ + -f examples/vllm_serve/Dockerfile -t vllm-modelopt:v0.26.0 . +``` + +For a direct installation from the ModelOpt repository root, install the tested vLLM +release and the ModelOpt extras used by this example: + +```bash +python3 -m pip install "vllm==0.28.0" +python3 -m pip install -e ".[all,mlflow]" +``` + +See the [ModelOpt installation guide](../../docs/source/getting_started/_installation_for_Linux.rst) +for details about installing partial dependency sets. + ## Calibrate and serve fake quant model in vLLM Step 1: Configure quantization settings. @@ -39,6 +57,18 @@ 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: + +```bash +KV_QUANT_CFG=NVFP4_KV_CFG QUANT_CALIB_SIZE=512 \ + python vllm_serve_fakequant.py -tp 8 \ + --max-model-len 8192 --enforce-eager --host 0.0.0.0 --port 8000 +``` + +Calibration uses dedicated scratch KV-cache blocks, so reducing `--max-num-batched-tokens` +is not required to avoid NaNs. + For vLLM versions that expose `--moe-backend`, this launcher defaults to `--moe-backend triton`. ModelOpt expert fakequant needs a decomposed MoE backend so both expert GEMMs are visible during calibration. diff --git a/examples/vllm_serve/vllm_ptq_utils.py b/examples/vllm_serve/vllm_ptq_utils.py index 709d6532fb3..7d41d5e9c88 100644 --- a/examples/vllm_serve/vllm_ptq_utils.py +++ b/examples/vllm_serve/vllm_ptq_utils.py @@ -36,8 +36,126 @@ def _create_new_data_cls(data_cls, **kwargs): return data_cls(**filtered_kwargs) +def _get_calibration_block_count( + model_runner: Any, +) -> Callable[[int, Any], int] | None: + """Return the block reservation policy supported by the installed vLLM.""" + vllm_config = model_runner.vllm_config + + try: + from vllm.v1.worker.gpu.warmup import _reserved_block_count + except ImportError: + try: + from vllm.utils.math_utils import cdiv + from vllm.v1.kv_cache_interface import CrossAttentionSpec, MambaSpec + except ImportError: + return None + + def block_count(num_tokens: int, kv_cache_spec: Any) -> int: + """Calculate the vLLM 0.26 warmup block reservation.""" + # vLLM 0.26's warmup reservation policy. + if isinstance(kv_cache_spec, CrossAttentionSpec): + num_tokens = 0 + num_blocks = cdiv(num_tokens, kv_cache_spec.block_size) + if isinstance(kv_cache_spec, MambaSpec) and kv_cache_spec.mamba_cache_mode == "align": + num_blocks += kv_cache_spec.num_speculative_blocks + return num_blocks + + else: + + def block_count(num_tokens: int, kv_cache_spec: Any) -> int: + """Calculate the current vLLM warmup block reservation.""" + # Calibration runs before model_state is initialized, so call the + # underlying reservation policy rather than _warmup_block_counter. + return _reserved_block_count( + num_tokens, + kv_cache_spec, + num_lookahead_tokens=vllm_config.num_lookahead_tokens, + max_model_len=model_runner.max_model_len, + max_encoder_len=0, + ) + + return block_count + + +def _allocate_calibration_blocks( + self: Any, sequence_lengths: list[int] +) -> tuple[list[tuple[list[int], ...]], list[int] | None]: + """Allocate scheduler-compatible scratch blocks for calibration requests. + + vLLM 0.28 treats block 0 as the null block. Its GPU runner expects real block + tables for hybrid attention/Mamba models, even for one-shot prefill requests. + Use vLLM's warmup reservation policy so this stays aligned with each cache + group's KVCacheSpec. + """ + kv_cache_config = self.model_runner.kv_cache_config + kv_cache_groups = kv_cache_config.kv_cache_groups + block_count = _get_calibration_block_count(self.model_runner) + + if block_count is None: + warnings.warn( + "vLLM warmup block reservation helpers were not found; falling back to " + "empty block tables. Hybrid attention/Mamba models may produce NaNs.", + stacklevel=2, + ) + return [tuple([] for _ in kv_cache_groups) for _ in sequence_lengths], None + + next_block_id = 1 # Block 0 is reserved as the null block. + block_ids_batch: list[tuple[list[int], ...]] = [] + allocated_block_ids: list[int] = [] + + for sequence_length in sequence_lengths: + request_block_ids = [] + for group in kv_cache_groups: + num_blocks = block_count(sequence_length, group.kv_cache_spec) + block_ids = list(range(next_block_id, next_block_id + num_blocks)) + next_block_id += num_blocks + allocated_block_ids.extend(block_ids) + request_block_ids.append(block_ids) + block_ids_batch.append(tuple(request_block_ids)) + + if next_block_id > kv_cache_config.num_blocks: + raise RuntimeError( + "Calibration batch requires " + f"{next_block_id - 1} KV cache blocks, but only " + f"{kv_cache_config.num_blocks - 1} non-null blocks are available." + ) + + scheduler_fields = {field.name for field in dataclasses.fields(SchedulerOutput)} + blocks_to_zero = allocated_block_ids if "new_block_ids_to_zero" in scheduler_fields else None + return block_ids_batch, blocks_to_zero + + +def _cleanup_calibration_requests( + self: Any, + cleanup_output: SchedulerOutput, + calibration_error: BaseException | None, +) -> None: + """Clean request state without hiding an active calibration error.""" + try: + # Zero-token steps return before forward/sampling, so no sample_tokens call is needed. + self.execute_model(cleanup_output) + except Exception as execute_error: + finish_requests = getattr(self.model_runner, "finish_requests", None) + if finish_requests is None: + if calibration_error is not None: + raise calibration_error from execute_error + raise + + try: + finish_requests(cleanup_output) + except Exception as finish_error: + if calibration_error is not None: + finish_error.__cause__ = execute_error + raise calibration_error from finish_error + raise finish_error from execute_error + + def calibrate_fun(calib_dataloader: DataLoader, self: Any) -> Callable[[Any], None]: + """Create a calibration loop backed by the vLLM worker scheduler.""" + def calibrate_loop(model: Any) -> None: + """Calibrate the model with batches submitted through the scheduler.""" for batch_idx, batch in tqdm(enumerate(calib_dataloader)): input_ids_batch = batch["input_ids"] @@ -56,7 +174,9 @@ def calibrate_loop(model: Any) -> None: input_ids_list_batch = [input_ids_list_batch] num_groups = len(self.model_runner.kv_cache_config.kv_cache_groups) - empty_block_ids = tuple([] for _ in range(num_groups)) + block_ids_batch, new_block_ids_to_zero = _allocate_calibration_blocks( + self, [len(input_ids) for input_ids in input_ids_list_batch] + ) scheduled_new_reqs = [] num_scheduled_tokens = {} @@ -74,7 +194,7 @@ def calibrate_loop(model: Any) -> None: mm_features=[], sampling_params=SamplingParams(max_tokens=1), pooling_params=None, - block_ids=empty_block_ids, + block_ids=block_ids_batch[seq_idx], num_computed_tokens=0, lora_request=None, ) @@ -96,40 +216,36 @@ def calibrate_loop(model: Any) -> None: kv_connector_metadata=None, structured_output_request_ids={}, grammar_bitmask=None, + new_block_ids_to_zero=new_block_ids_to_zero, + ) + # Submit a zero-token scheduler step after the request has been + # registered. This is the vLLM 0.28 cleanup path and removes + # request-scoped attention/Mamba state from the persistent batch. + cleanup_output = _create_new_data_cls( + type(scheduler_output), + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={}, + total_num_scheduled_tokens=0, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[0] * num_groups, + finished_req_ids=set(num_scheduled_tokens), + free_encoder_mm_hashes=[], + kv_connector_metadata=None, + structured_output_request_ids={}, + grammar_bitmask=None, ) try: output = self.execute_model(scheduler_output) if hasattr(self, "sample_tokens"): if output is None: # TODO: make this default when vllm <= 0.11 is outdated self.sample_tokens(None) - finally: - # finish_requests runs before add_requests inside execute_model, so - # req IDs aren't registered yet at that point — call it directly after. - # Wrap in try/except so a cleanup error never masks the original exception. - try: - if hasattr(self.model_runner, "finish_requests"): - cleanup_output = _create_new_data_cls( - type(scheduler_output), - scheduled_new_reqs=[], - scheduled_cached_reqs=scheduler_output.scheduled_cached_reqs, - num_scheduled_tokens={}, - total_num_scheduled_tokens=0, - scheduled_spec_decode_tokens={}, - scheduled_encoder_inputs={}, - num_common_prefix_blocks=scheduler_output.num_common_prefix_blocks, - finished_req_ids=set(num_scheduled_tokens.keys()), - free_encoder_mm_hashes=[], - kv_connector_metadata=None, - structured_output_request_ids={}, - grammar_bitmask=None, - ) - self.model_runner.finish_requests(cleanup_output) - else: - warnings.warn( - "model_runner.finish_requests not found; request state may leak during calibration." - ) - except Exception: - warnings.warn("Failed to clean up request state after calibration batch.") + except BaseException as calibration_error: + _cleanup_calibration_requests(self, cleanup_output, calibration_error) + raise + + _cleanup_calibration_requests(self, cleanup_output, calibration_error=None) return calibrate_loop @@ -171,6 +287,7 @@ def update_kv_cfg_for_mla(model: torch.nn.Module, kv_quant_cfg: list) -> list: def get_quant_config(quant_config: dict[str, Any], model: Any) -> dict[str, Any]: + """Resolve and merge model and KV-cache quantization configuration.""" import copy if quant_config["recipe_path"]: 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 038bbd8e978..4ed91e41afd 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py @@ -28,6 +28,7 @@ from __future__ import annotations +import builtins import gc import importlib.util from pathlib import Path @@ -69,6 +70,218 @@ def _load_example_module(name: str): return module +def _calibration_worker(num_blocks: int): + cache_groups = [ + SimpleNamespace(kv_cache_spec="attention"), + SimpleNamespace(kv_cache_spec="mamba"), + ] + return SimpleNamespace( + model_runner=SimpleNamespace( + kv_cache_config=SimpleNamespace( + kv_cache_groups=cache_groups, + num_blocks=num_blocks, + ) + ) + ) + + +def _patch_vllm_imports(monkeypatch, modules): + real_import = builtins.__import__ + + def fake_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in modules: + imported = modules[name] + if isinstance(imported, BaseException): + raise imported + return imported + return real_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", fake_import) + + +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") + reserved_block_count = Mock(return_value=4) + _patch_vllm_imports( + monkeypatch, + {"vllm.v1.worker.gpu.warmup": SimpleNamespace(_reserved_block_count=reserved_block_count)}, + ) + model_runner = SimpleNamespace( + vllm_config=SimpleNamespace(num_lookahead_tokens=3), + max_model_len=2048, + ) + kv_cache_spec = object() + + block_count = module._get_calibration_block_count(model_runner) + + assert block_count is not None + assert block_count(128, kv_cache_spec) == 4 + reserved_block_count.assert_called_once_with( + 128, + kv_cache_spec, + num_lookahead_tokens=3, + max_model_len=2048, + max_encoder_len=0, + ) + + +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") + + class CrossAttentionSpec: + block_size = 16 + + class MambaSpec: + block_size = 16 + mamba_cache_mode = "align" + num_speculative_blocks = 2 + + cdiv = Mock( + side_effect=lambda numerator, denominator: (numerator + denominator - 1) // denominator + ) + _patch_vllm_imports( + monkeypatch, + { + "vllm.v1.worker.gpu.warmup": ImportError("0.28 helper unavailable"), + "vllm.utils.math_utils": SimpleNamespace(cdiv=cdiv), + "vllm.v1.kv_cache_interface": SimpleNamespace( + CrossAttentionSpec=CrossAttentionSpec, + MambaSpec=MambaSpec, + ), + }, + ) + model_runner = SimpleNamespace( + vllm_config=SimpleNamespace(), + max_model_len=2048, + ) + + block_count = module._get_calibration_block_count(model_runner) + + assert block_count is not None + assert block_count(33, SimpleNamespace(block_size=16)) == 3 + assert block_count(33, CrossAttentionSpec()) == 0 + assert block_count(33, MambaSpec()) == 5 + assert cdiv.call_args_list == [ + ((33, 16),), + ((0, 16),), + ((33, 16),), + ] + + +def test_allocate_calibration_blocks_assigns_non_null_blocks(monkeypatch): + """Scratch block tables must use unique non-null blocks for every request and group.""" + module = _load_example_module("vllm_ptq_utils") + block_count = Mock(side_effect=[1, 2, 2, 1]) + monkeypatch.setattr( + module, + "_get_calibration_block_count", + Mock(return_value=block_count), + ) + + block_tables, blocks_to_zero = module._allocate_calibration_blocks( + _calibration_worker(num_blocks=7), + sequence_lengths=[8, 16], + ) + + assert block_tables == [ + ([1], [2, 3]), + ([4, 5], [6]), + ] + assert block_count.call_args_list == [ + ((8, "attention"),), + ((8, "mamba"),), + ((16, "attention"),), + ((16, "mamba"),), + ] + + scheduler_fields = {field.name for field in module.dataclasses.fields(module.SchedulerOutput)} + expected_blocks_to_zero = ( + [1, 2, 3, 4, 5, 6] if "new_block_ids_to_zero" in scheduler_fields else None + ) + assert blocks_to_zero == expected_blocks_to_zero + + +def test_allocate_calibration_blocks_rejects_insufficient_capacity(monkeypatch): + """Scratch block allocation must account for block 0 being unavailable.""" + module = _load_example_module("vllm_ptq_utils") + monkeypatch.setattr( + module, + "_get_calibration_block_count", + Mock(return_value=Mock(side_effect=[1, 2, 2, 1])), + ) + + with pytest.raises( + RuntimeError, + match=( + r"Calibration batch requires 6 KV cache blocks, " + r"but only 5 non-null blocks are available\." + ), + ): + module._allocate_calibration_blocks( + _calibration_worker(num_blocks=6), + sequence_lengths=[8, 16], + ) + + +@pytest.mark.parametrize("has_calibration_error", [False, True]) +def test_cleanup_failure_preserves_calibration_error(has_calibration_error): + """Cleanup must fail closed without replacing an active calibration error.""" + module = _load_example_module("vllm_ptq_utils") + execute_error = RuntimeError("scheduler cleanup failed") + finish_error = RuntimeError("legacy cleanup failed") + calibration_error = ValueError("calibration failed") if has_calibration_error else None + worker = SimpleNamespace( + execute_model=Mock(side_effect=execute_error), + model_runner=SimpleNamespace(finish_requests=Mock(side_effect=finish_error)), + ) + + expected_error = calibration_error or finish_error + with pytest.raises(type(expected_error)) as raised: + module._cleanup_calibration_requests(worker, object(), calibration_error) + + assert raised.value is expected_error + if calibration_error is not None: + assert calibration_error.__cause__ is finish_error + assert finish_error.__cause__ is execute_error + + +@pytest.mark.parametrize("has_calibration_error", [False, True]) +def test_cleanup_without_legacy_fallback_preserves_primary_error(has_calibration_error): + """Missing legacy cleanup must preserve the most useful primary error.""" + module = _load_example_module("vllm_ptq_utils") + execute_error = RuntimeError("scheduler cleanup failed") + calibration_error = ValueError("calibration failed") if has_calibration_error else None + worker = SimpleNamespace( + execute_model=Mock(side_effect=execute_error), + model_runner=SimpleNamespace(), + ) + + expected_error = calibration_error or execute_error + with pytest.raises(type(expected_error)) as raised: + module._cleanup_calibration_requests(worker, object(), calibration_error) + + assert raised.value is expected_error + if calibration_error is not None: + assert calibration_error.__cause__ is execute_error + + +def test_cleanup_uses_legacy_fallback(): + """A successful legacy cleanup may recover from an unsupported scheduler step.""" + module = _load_example_module("vllm_ptq_utils") + cleanup_output = object() + finish_requests = Mock() + worker = SimpleNamespace( + execute_model=Mock(side_effect=RuntimeError("unsupported scheduler cleanup")), + model_runner=SimpleNamespace(finish_requests=finish_requests), + ) + + module._cleanup_calibration_requests(worker, cleanup_output, calibration_error=None) + + finish_requests.assert_called_once_with(cleanup_output) + + class _NativeAttention(torch.nn.Module): def forward(self, query, key, value, *args, **kwargs): return query, key, value