Fix vLLM fakequant calibration for hybrid attention models (#2414)

### What does this PR do? Type of change: Bug fix Fix fakequant calibration for hybrid attention/Mamba models, including NVIDIA Nemotron-3-Nano, on vLLM 0.26 and 0.28. The manual calibration scheduler path previously submitted requests with empty KV-cache block tables. Hybrid models require scheduler-compatible cache state during prefill; on current vLLM releases the empty tables caused the Mamba state to use the reserved null block and calibration activations became NaN. Request cleanup also no longer matched the vLLM 0.28 execution lifecycle, which could leave request-scoped state in the persistent batch. This PR: - Allocates non-null scratch blocks for every KV-cache group using the vLLM warmup reservation policy. - Supports both the vLLM 0.28 reservation helper and the equivalent vLLM 0.26 calculation. - Passes newly allocated blocks through `new_block_ids_to_zero` when that scheduler field is available. - Validates that the calibration batch fits in the configured cache and reports how to reduce calibration demand if it does not. - Cleans up calibration requests through a zero-token scheduler step on current vLLM, with a direct cleanup fallback for older runners. - Updates the example Dockerfile to default to vLLM 0.28.0 while retaining vLLM 0.26.0 through `VLLM_VERSION`. - Documents the validated Nemotron-3-Nano NVFP4 KV-cache workflow and clarifies that reducing `--max-num-batched-tokens` is not required. ### Usage Build the default vLLM 0.28.0 image: ```bash docker build -f examples/vllm_serve/Dockerfile \ -t vllm-modelopt:v0.28.0 . ``` Build with vLLM 0.26.0: ```bash docker build --build-arg VLLM_VERSION=0.26.0 \ -f examples/vllm_serve/Dockerfile \ -t vllm-modelopt:v0.26.0 . ``` Calibrate and serve Nemotron-3-Nano with NVFP4 KV-cache fakequant: ```bash KV_QUANT_CFG=NVFP4_KV_CFG QUANT_CALIB_SIZE=512 \ python examples/vllm_serve/vllm_serve_fakequant.py \ <nemotron3_nano_model_path> \ --trust-remote-code --enforce-eager -tp 8 \ --max-model-len 8192 --host 0.0.0.0 --port 8000 ``` ### Testing Validated on omniml-a0 with `NVIDIA-Nemotron-3-Nano-30B-A3B-BF16`, tensor parallel size 8, `NVFP4_KV_CFG`, `QUANT_CALIB_SIZE=512`, and `--max-model-len 8192`. No `--max-num-batched-tokens` override was used. - vLLM 0.28.0: - All 512 calibration samples completed. - No NaNs or cache-cleanup warnings were observed. - The server started and `/health` passed. - An OpenAI-compatible completion request returned coherent generated text. - vLLM 0.26.0: - Repeated the same 512-sample TP8 calibration with the official `vllm/vllm-openai:v0.26.0` image. - No NaNs were observed. - The server started, passed `/health`, and returned coherent generated text. - Docker: - Built and verified the updated vLLM 0.28.0 image. - Focused tests: - `tests/examples/vllm_serve/test_vllm_mlflow_utils.py`: 32 passed. - Cleanup failure, missing legacy API, and legacy fallback tests: 5 passed on both vLLM 0.26.0 and 0.28.0. - Repository hooks: - Targeted pre-commit hooks for every changed Python, Markdown, and Docker file: passed. - `git diff --check`: passed. ### Before your PR is "*Ready for review*" - 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?: ✅ — added focused coverage for fail-closed cleanup, exception chaining, and the legacy cleanup fallback; the full regression was also validated end to end. - Did you update Changelog?: N/A - Did you get Claude approval on this PR?: N/A ### Additional Information The change is quantization-format agnostic. It corrects the calibration scheduler and cache lifecycle rather than special-casing `NVFP4_KV_CFG` or using an NVFP4 cast path.  ## Summary by CodeRabbit - **New Features** - Added support for configuring the vLLM version through `VLLM_VERSION`, with vLLM 0.28.0 as the default. - Added calibration and serving guidance for hybrid attention/Mamba models, including Nemotron 3 Nano with NVFP4 KV-cache fake quantization. - **Bug Fixes** - Improved calibration block handling across supported vLLM versions. - Improved calibration cleanup to preserve original errors and provide reliable fallback behavior when standard cleanup is unavailable. - **Documentation** - Documented tested versions, direct installation commands, ModelOpt setup, serving options, and guidance to avoid NaNs during batched serving.  --------- Signed-off-by: Kinjal Patel <kinjalpravin@nvidia.com>
2026-10-02 03:14:52 +08:00 · 2026-09-17 14:12:34 -07:00
parent b16356776b
commit 02b58eb146
4 changed files with 428 additions and 35 deletions
@@ -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 \
@@ -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 <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:
+
+```bash
+KV_QUANT_CFG=NVFP4_KV_CFG QUANT_CALIB_SIZE=512 \
+  python vllm_serve_fakequant.py <nemotron3_nano_model_path> -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.
@@ -36,8 +36,131 @@ 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)}
+    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
+
+
+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 +179,9 @@ def calibrate_fun(calib_dataloader: DataLoader, self: Any) -> Callable[[Any], No
                    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 +199,7 @@ def calibrate_fun(calib_dataloader: DataLoader, self: Any) -> Callable[[Any], No
                    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 +221,36 @@ def calibrate_fun(calib_dataloader: DataLoader, self: Any) -> Callable[[Any], No
                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 +292,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"]: