Fix vLLM fakequant calibration for hybrid attention models

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 \

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 <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.

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"]: