From b98646339afa1434fb08033ac90a797fd2537d72 Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Fri, 11 Sep 2026 18:45:27 +0000 Subject: [PATCH 1/9] Fix vLLM 0.28 fakequant cache lifecycle Signed-off-by: Kinjal Patel --- examples/vllm_serve/vllm_ptq_utils.py | 119 ++++++++++++++++++++------ 1 file changed, 94 insertions(+), 25 deletions(-) diff --git a/examples/vllm_serve/vllm_ptq_utils.py b/examples/vllm_serve/vllm_ptq_utils.py index 709d6532fb3..e637b8c45e6 100644 --- a/examples/vllm_serve/vllm_ptq_utils.py +++ b/examples/vllm_serve/vllm_ptq_utils.py @@ -36,6 +36,69 @@ def _create_new_data_cls(data_cls, **kwargs): return data_cls(**filtered_kwargs) +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 + empty_block_ids = tuple([] for _ in kv_cache_groups) + + try: + from vllm.v1.worker.gpu.warmup import _reserved_block_count + except ImportError: + # Older vLLM versions do not expose the V2 warmup allocator and used + # empty block tables for this calibration path. + return [empty_block_ids for _ in sequence_lengths], None + + model_runner = self.model_runner + vllm_config = model_runner.vllm_config + + def block_count(num_tokens: int, kv_cache_spec: Any) -> int: + # 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, + ) + 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. " + "Reduce CALIB_BATCH_SIZE or calibration sequence length." + ) + + 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 calibrate_fun(calib_dataloader: DataLoader, self: Any) -> Callable[[Any], None]: def calibrate_loop(model: Any) -> None: for batch_idx, batch in tqdm(enumerate(calib_dataloader)): @@ -56,7 +119,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 +139,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,6 +161,7 @@ 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, ) try: output = self.execute_model(scheduler_output) @@ -103,33 +169,36 @@ def calibrate_loop(model: Any) -> None: 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. + # 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: - 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.execute_model(cleanup_output) + except Exception: + # Older runners expose cleanup directly instead of accepting + # an empty execute_model step. + try: self.model_runner.finish_requests(cleanup_output) - else: + except Exception: warnings.warn( - "model_runner.finish_requests not found; request state may leak during calibration." + "Failed to clean up request state after calibration batch.", + stacklevel=2, ) - except Exception: - warnings.warn("Failed to clean up request state after calibration batch.") return calibrate_loop From b2e9a7dd78e52f4e4deeb4e1b0f69ffa725756bc Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Fri, 11 Sep 2026 19:22:54 +0000 Subject: [PATCH 2/9] Support vLLM 0.26 calibration block allocation Signed-off-by: Kinjal Patel --- examples/vllm_serve/vllm_ptq_utils.py | 47 ++++++++++++++++++--------- 1 file changed, 32 insertions(+), 15 deletions(-) diff --git a/examples/vllm_serve/vllm_ptq_utils.py b/examples/vllm_serve/vllm_ptq_utils.py index e637b8c45e6..fabaa3a839f 100644 --- a/examples/vllm_serve/vllm_ptq_utils.py +++ b/examples/vllm_serve/vllm_ptq_utils.py @@ -50,26 +50,43 @@ def _allocate_calibration_blocks( kv_cache_groups = kv_cache_config.kv_cache_groups empty_block_ids = tuple([] for _ in kv_cache_groups) + model_runner = self.model_runner + vllm_config = model_runner.vllm_config + try: from vllm.v1.worker.gpu.warmup import _reserved_block_count except ImportError: - # Older vLLM versions do not expose the V2 warmup allocator and used - # empty block tables for this calibration path. - return [empty_block_ids for _ in sequence_lengths], None + try: + from vllm.utils.math_utils import cdiv + from vllm.v1.kv_cache_interface import CrossAttentionSpec, MambaSpec + except ImportError: + # Older vLLM versions used empty block tables for this path. + return [empty_block_ids for _ in sequence_lengths], None + + def block_count(num_tokens: int, kv_cache_spec: Any) -> int: + # 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 - model_runner = self.model_runner - vllm_config = model_runner.vllm_config + else: - def block_count(num_tokens: int, kv_cache_spec: Any) -> int: - # 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, - ) + def block_count(num_tokens: int, kv_cache_spec: Any) -> int: + # 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, + ) next_block_id = 1 # Block 0 is reserved as the null block. block_ids_batch: list[tuple[list[int], ...]] = [] allocated_block_ids: list[int] = [] From 7083c5593760ed53ccd944fe643907bd355a632f Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Fri, 11 Sep 2026 20:45:26 +0000 Subject: [PATCH 3/9] Update vLLM fakequant Docker image Signed-off-by: Kinjal Patel --- examples/vllm_serve/Dockerfile | 3 ++- examples/vllm_serve/README.md | 29 +++++++++++++++++++++++++---- 2 files changed, 27 insertions(+), 5 deletions(-) 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..33ddadf2cc0 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -4,17 +4,26 @@ 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 . +``` + +Alternatively, install vLLM and ModelOpt directly with pip. + ## Calibrate and serve fake quant model in vLLM Step 1: Configure quantization settings. @@ -39,6 +48,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. From 5f9baa64bda3249bf04006b9e075bf33ca49d470 Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Fri, 11 Sep 2026 21:06:30 +0000 Subject: [PATCH 4/9] Address vLLM fakequant review feedback Signed-off-by: Kinjal Patel --- examples/vllm_serve/README.md | 11 +++- examples/vllm_serve/vllm_ptq_utils.py | 53 ++++++++++++------- .../quantization/test_vllm_dynamic_modules.py | 37 +++++++++++++ 3 files changed, 80 insertions(+), 21 deletions(-) diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 33ddadf2cc0..bc2b0fb53f3 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -22,7 +22,16 @@ docker build --build-arg VLLM_VERSION=0.26.0 \ -f examples/vllm_serve/Dockerfile -t vllm-modelopt:v0.26.0 . ``` -Alternatively, install vLLM and ModelOpt directly with pip. +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 diff --git a/examples/vllm_serve/vllm_ptq_utils.py b/examples/vllm_serve/vllm_ptq_utils.py index fabaa3a839f..dde81ee5430 100644 --- a/examples/vllm_serve/vllm_ptq_utils.py +++ b/examples/vllm_serve/vllm_ptq_utils.py @@ -14,7 +14,7 @@ # limitations under the License. import dataclasses -import warnings +import sys from collections.abc import Callable from typing import Any @@ -64,20 +64,19 @@ def _allocate_calibration_blocks( return [empty_block_ids for _ in sequence_lengths], 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" - ): + 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( @@ -87,6 +86,7 @@ def block_count(num_tokens: int, kv_cache_spec: Any) -> int: max_model_len=model_runner.max_model_len, max_encoder_len=0, ) + next_block_id = 1 # Block 0 is reserved as the null block. block_ids_batch: list[tuple[list[int], ...]] = [] allocated_block_ids: list[int] = [] @@ -110,14 +110,33 @@ def block_count(num_tokens: int, kv_cache_spec: Any) -> int: ) 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 - ) + 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: + self.execute_model(cleanup_output) + except Exception as execute_error: + try: + self.model_runner.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"] @@ -204,18 +223,11 @@ def calibrate_loop(model: Any) -> None: structured_output_request_ids={}, grammar_bitmask=None, ) - try: - self.execute_model(cleanup_output) - except Exception: - # Older runners expose cleanup directly instead of accepting - # an empty execute_model step. - try: - self.model_runner.finish_requests(cleanup_output) - except Exception: - warnings.warn( - "Failed to clean up request state after calibration batch.", - stacklevel=2, - ) + _cleanup_calibration_requests( + self, + cleanup_output, + calibration_error=sys.exc_info()[1], + ) return calibrate_loop @@ -257,6 +269,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..38dfc670afa 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py @@ -69,6 +69,43 @@ def _load_example_module(name: str): return module +@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 + + +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 From 37f416fb15f94be3dbefd344e23e29453cb02745 Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Fri, 11 Sep 2026 23:54:11 +0000 Subject: [PATCH 5/9] Clarify vLLM calibration capacity error Signed-off-by: Kinjal Patel --- examples/vllm_serve/vllm_ptq_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/examples/vllm_serve/vllm_ptq_utils.py b/examples/vllm_serve/vllm_ptq_utils.py index dde81ee5430..0df5d006b36 100644 --- a/examples/vllm_serve/vllm_ptq_utils.py +++ b/examples/vllm_serve/vllm_ptq_utils.py @@ -105,8 +105,7 @@ def block_count(num_tokens: int, kv_cache_spec: Any) -> int: 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. " - "Reduce CALIB_BATCH_SIZE or calibration sequence length." + f"{kv_cache_config.num_blocks - 1} non-null blocks are available." ) scheduler_fields = {field.name for field in dataclasses.fields(SchedulerOutput)} From 12a463533fcb397daf15f0856f2463cc5139ecdc Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Fri, 11 Sep 2026 23:58:55 +0000 Subject: [PATCH 6/9] Harden vLLM calibration compatibility fallbacks Signed-off-by: Kinjal Patel --- examples/vllm_serve/vllm_ptq_utils.py | 68 +++++++++++-------- .../quantization/test_vllm_dynamic_modules.py | 20 ++++++ 2 files changed, 58 insertions(+), 30 deletions(-) diff --git a/examples/vllm_serve/vllm_ptq_utils.py b/examples/vllm_serve/vllm_ptq_utils.py index 0df5d006b36..bae00608d10 100644 --- a/examples/vllm_serve/vllm_ptq_utils.py +++ b/examples/vllm_serve/vllm_ptq_utils.py @@ -14,7 +14,7 @@ # limitations under the License. import dataclasses -import sys +import warnings from collections.abc import Callable from typing import Any @@ -48,8 +48,6 @@ def _allocate_calibration_blocks( """ kv_cache_config = self.model_runner.kv_cache_config kv_cache_groups = kv_cache_config.kv_cache_groups - empty_block_ids = tuple([] for _ in kv_cache_groups) - model_runner = self.model_runner vllm_config = model_runner.vllm_config @@ -60,8 +58,12 @@ def _allocate_calibration_blocks( from vllm.utils.math_utils import cdiv from vllm.v1.kv_cache_interface import CrossAttentionSpec, MambaSpec except ImportError: - # Older vLLM versions used empty block tables for this path. - return [empty_block_ids for _ in sequence_lengths], 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 def block_count(num_tokens: int, kv_cache_spec: Any) -> int: """Calculate the vLLM 0.26 warmup block reservation.""" @@ -120,10 +122,17 @@ def _cleanup_calibration_requests( ) -> 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: - self.model_runner.finish_requests(cleanup_output) + finish_requests(cleanup_output) except Exception as finish_error: if calibration_error is not None: finish_error.__cause__ = execute_error @@ -198,35 +207,34 @@ def calibrate_loop(model: Any) -> None: 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: - # 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, - ) - _cleanup_calibration_requests( - self, - cleanup_output, - calibration_error=sys.exc_info()[1], - ) + 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 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 38dfc670afa..ed8a6c86d0f 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py @@ -91,6 +91,26 @@ def test_cleanup_failure_preserves_calibration_error(has_calibration_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") From aa0b228b683df178570dc1cad422b80334aaa43c Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Tue, 15 Sep 2026 18:58:14 +0000 Subject: [PATCH 7/9] Add vLLM calibration block allocation tests Signed-off-by: Kinjal Patel --- examples/vllm_serve/vllm_ptq_utils.py | 49 ++++++++----- .../quantization/test_vllm_dynamic_modules.py | 70 +++++++++++++++++++ 2 files changed, 100 insertions(+), 19 deletions(-) diff --git a/examples/vllm_serve/vllm_ptq_utils.py b/examples/vllm_serve/vllm_ptq_utils.py index bae00608d10..7d41d5e9c88 100644 --- a/examples/vllm_serve/vllm_ptq_utils.py +++ b/examples/vllm_serve/vllm_ptq_utils.py @@ -36,19 +36,10 @@ def _create_new_data_cls(data_cls, **kwargs): return data_cls(**filtered_kwargs) -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 - model_runner = self.model_runner +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: @@ -58,12 +49,7 @@ def _allocate_calibration_blocks( from vllm.utils.math_utils import cdiv from vllm.v1.kv_cache_interface import CrossAttentionSpec, MambaSpec except ImportError: - 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 + return None def block_count(num_tokens: int, kv_cache_spec: Any) -> int: """Calculate the vLLM 0.26 warmup block reservation.""" @@ -89,6 +75,31 @@ def block_count(num_tokens: int, kv_cache_spec: Any) -> int: 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] = [] 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 ed8a6c86d0f..22f6a187826 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py @@ -69,6 +69,76 @@ 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 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.""" From d184d14b897d6bc956e34087a5ee46702e3f4b4f Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Tue, 15 Sep 2026 19:30:37 +0000 Subject: [PATCH 8/9] Test vLLM calibration reservation adapters Signed-off-by: Kinjal Patel --- .../quantization/test_vllm_dynamic_modules.py | 86 +++++++++++++++++++ 1 file changed, 86 insertions(+) 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 22f6a187826..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 @@ -84,6 +85,91 @@ def _calibration_worker(num_blocks: int): ) +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") From 84a342971309a242f5e27b3785725c38b8278f99 Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Thu, 17 Sep 2026 18:58:47 +0000 Subject: [PATCH 9/9] Honor vLLM KV cache zeroing requirements Signed-off-by: Kinjal Patel --- examples/vllm_serve/vllm_ptq_utils.py | 7 +++- .../quantization/test_vllm_dynamic_modules.py | 37 ++++++++++++++++--- 2 files changed, 38 insertions(+), 6 deletions(-) diff --git a/examples/vllm_serve/vllm_ptq_utils.py b/examples/vllm_serve/vllm_ptq_utils.py index 7d41d5e9c88..ca9e276b9b4 100644 --- a/examples/vllm_serve/vllm_ptq_utils.py +++ b/examples/vllm_serve/vllm_ptq_utils.py @@ -122,7 +122,12 @@ def _allocate_calibration_blocks( ) 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 + if "new_block_ids_to_zero" in scheduler_fields: + blocks_to_zero = ( + allocated_block_ids if getattr(kv_cache_config, "needs_kv_cache_zeroing", False) else [] + ) + else: + blocks_to_zero = None return block_ids_batch, blocks_to_zero 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 4ed91e41afd..0f440b93deb 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py @@ -70,16 +70,19 @@ 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"), - ] +def _calibration_worker( + num_blocks: int, + *, + cache_specs=("attention", "mamba"), + needs_kv_cache_zeroing=True, +): + cache_groups = [SimpleNamespace(kv_cache_spec=spec) for spec in cache_specs] return SimpleNamespace( model_runner=SimpleNamespace( kv_cache_config=SimpleNamespace( kv_cache_groups=cache_groups, num_blocks=num_blocks, + needs_kv_cache_zeroing=needs_kv_cache_zeroing, ) ) ) @@ -203,6 +206,30 @@ def test_allocate_calibration_blocks_assigns_non_null_blocks(monkeypatch): assert blocks_to_zero == expected_blocks_to_zero +def test_allocate_calibration_blocks_skips_zeroing_for_attention_only_cache(monkeypatch): + """Attention-only caches have no block zeroer and must receive an empty zeroing list.""" + module = _load_example_module("vllm_ptq_utils") + monkeypatch.setattr( + module, + "_get_calibration_block_count", + Mock(return_value=Mock(return_value=1)), + ) + + block_tables, blocks_to_zero = module._allocate_calibration_blocks( + _calibration_worker( + num_blocks=4, + cache_specs=("attention",), + needs_kv_cache_zeroing=False, + ), + sequence_lengths=[8], + ) + + assert block_tables == [([1],)] + scheduler_fields = {field.name for field in module.dataclasses.fields(module.SchedulerOutput)} + expected_blocks_to_zero = [] 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")