docs: Add Usage Stats V2 PRD and initial design sketch (VLLM-94)

- Add comprehensive PRD with current state analysis, gaps, competitive research
- Add design sketch with architecture, component design, code sketches
- Document hardware topology, model metadata, anonymized workload requirements
- Update usage_stats.md with links to V2 design docs

Co-authored-by: Simon Mo <simon@inferact.ai>

.buildkite/performance-benchmarks/README.md

@@ -83,7 +83,6 @@ We test the throughput by using `vllm bench serve` with request rate = inf to co
         "server_parameters": {
             "model": "meta-llama/Meta-Llama-3-8B",
             "tensor_parallel_size": 1,
-            "swap_space": 16,
             "disable_log_stats": "",
             "load_format": "dummy"
         },

.buildkite/performance-benchmarks/tests/serving-tests-hpu.json

@@ -10,7 +10,6 @@
         "server_parameters": {
             "model": "meta-llama/Meta-Llama-3.1-8B-Instruct",
             "tensor_parallel_size": 1,
-            "swap_space": 16,
             "disable_log_stats": "",
             "load_format": "dummy",
             "max-model-len": 2048,
@@ -37,7 +36,6 @@
         "server_parameters": {
             "model": "meta-llama/Meta-Llama-3.1-70B-Instruct",
             "tensor_parallel_size": 4,
-            "swap_space": 16,
             "disable_log_stats": "",
             "load_format": "dummy",
             "max-model-len": 2048,
@@ -64,7 +62,6 @@
         "server_parameters": {
             "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
             "tensor_parallel_size": 2,
-            "swap_space": 16,
             "disable_log_stats": "",
             "load_format": "dummy",
             "max-model-len": 2048,
@@ -91,7 +88,6 @@
         "server_parameters": {
             "model": "deepseek-ai/DeepSeek-R1",
             "tensor_parallel_size": 8,
-            "swap_space": 16,
             "disable_log_stats": "",
             "load_format": "dummy",
             "max-model-len": 2048,

.buildkite/performance-benchmarks/tests/serving-tests.json

@@ -5,7 +5,6 @@
         "server_parameters": {
             "model": "meta-llama/Meta-Llama-3.1-8B-Instruct",
             "tensor_parallel_size": 1,
-            "swap_space": 16,
             "disable_log_stats": "",
             "load_format": "dummy"
         },
@@ -23,7 +22,6 @@
         "server_parameters": {
             "model": "meta-llama/Meta-Llama-3.1-70B-Instruct",
             "tensor_parallel_size": 4,
-            "swap_space": 16,
             "disable_log_stats": "",
             "load_format": "dummy"
         },
@@ -41,7 +39,6 @@
         "server_parameters": {
             "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
             "tensor_parallel_size": 2,
-            "swap_space": 16,
             "disable_log_stats": "",
             "load_format": "dummy"
         },
@@ -59,7 +56,6 @@
         "server_parameters": {
             "model": "meta-llama/Meta-Llama-3.1-70B-Instruct", 
             "tensor_parallel_size": 4,
-            "swap_space": 16,
             "speculative_config": {
                 "model": "turboderp/Qwama-0.5B-Instruct",
                 "num_speculative_tokens": 4,

.buildkite/test-amd.yaml

@@ -1486,6 +1486,20 @@ steps:
     - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
     - DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
+- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
+  mirror_hardwares: [amdexperimental, amdproduction]
+  agent_pool: mi325_4
+  # grade: Blocking
+  timeout_in_minutes: 30
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 4
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - tests/v1/kv_connector/nixl_integration/
+  commands:
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+    - CROSS_LAYERS_BLOCKS=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+
 ##### multi gpus test #####
 ##### A100 test #####
 
@@ -2801,7 +2815,7 @@ steps:
   - vllm/v1/attention/selector.py
   - vllm/platforms/cuda.py
   commands:
-    rocm-smi
+    - rocm-smi
     - python3 examples/offline_inference/basic/chat.py
     # Attention
     # num_heads2 broken by https://github.com/flashinfer-ai/flashinfer/issues/1353
@@ -3136,6 +3150,20 @@ steps:
     - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
     - DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
+- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
+  mirror_hardwares: [amdexperimental, amdproduction]
+  agent_pool: mi355_4
+  # grade: Blocking
+  timeout_in_minutes: 30
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 4
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - tests/v1/kv_connector/nixl_integration/
+  commands:
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+    - CROSS_LAYERS_BLOCKS=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+
 ##### multi gpus test #####
 ##### A100 test #####
 
@@ -3283,7 +3311,7 @@ steps:
   commands:
   - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010
 
-- label: Qwen3-30B-A3B-FP8-block Accuracy (B200/MI355)
+- label: Qwen3-30B-A3B-FP8-block Accuracy (B200-MI355)
   mirror_hardwares: [amdexperimental, amdproduction, amdmi355]
   agent_pool: mi355_2
   timeout_in_minutes: 60
@@ -3305,7 +3333,7 @@ steps:
   commands:
   - bash .buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh 0.8 1319 8040
 
-- label: Attention Benchmarks Smoke Test (B200/MI355)
+- label: Attention Benchmarks Smoke Test (B200-MI355)
   device: b200
   mirror_hardwares: [amdexperimental, amdmi355]
   agent_pool: mi355_2

benchmarks/attention_benchmarks/mla_runner.py

@@ -145,7 +145,6 @@ def create_minimal_vllm_config(
     cache_config = CacheConfig(
         block_size=block_size,
         gpu_memory_utilization=0.9,
-        swap_space=0,
         cache_dtype="auto",
         enable_prefix_caching=False,
     )

benchmarks/attention_benchmarks/runner.py

@@ -141,7 +141,6 @@ def _create_vllm_config(
     cache_config = CacheConfig(
         block_size=config.block_size,
         cache_dtype="auto",
-        swap_space=0,
     )
     cache_config.num_gpu_blocks = max_num_blocks
     cache_config.num_cpu_blocks = 0

docs/design/attention_backends.md

@@ -206,7 +206,7 @@ configuration.
 |---------|--------|-----------|-------------|------------|------|--------|-----------|-----|-----------------|--------------|
 | `CUTLASS_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
 | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
-| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
+| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
 | `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
 | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
 | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |

docs/design/metrics.md

@@ -507,10 +507,10 @@ longer relevant in v1:
 - `vllm:num_requests_swapped`
 - `vllm:cpu_cache_usage_perc`
 
-In this mode, when a request is preempted (e.g. to make room in KV
-cache to complete other requests), we swap kv cache blocks out to CPU
-memory. This is also known as "KV cache offloading" and is configured
-with `--swap-space` and `--preemption-mode`.
+In this mode, when a request was preempted (e.g. to make room in KV
+cache to complete other requests), kv cache blocks were swapped out to
+CPU memory. The `--swap-space` flag has been removed as this feature
+is no longer used in V1.
 
 Historically, [vLLM has long supported beam search](https://github.com/vllm-project/vllm/issues/6226). The
 SequenceGroup encapsulated the idea of N Sequences which

docs/design/usage_stats_v2_design.md

@@ -0,0 +1,238 @@
+# Usage Stats V2: Initial Design Sketch
+
+**Status:** Draft  
+**Project:** Usage Stats V2  
+**Linear Issue:** VLLM-94  
+**Depends on:** [Usage Stats V2 PRD](usage_stats_v2_prd.md)
+
+## 1. Design Principles
+
+1. **Additive only:** New fields extend existing schema; no breaking changes
+2. **Fail-safe:** Collection errors are logged and skipped; never raise to caller
+3. **Async-first:** All collection and transmission off the critical path
+4. **Schema versioning:** Include `usage_stats_version` (e.g., `2`) for evolution
+5. **Curlable format:** Flat JSON key-value; `curl -X POST -d @payload.json $URL`
+
+## 2. Architecture Overview
+
+```
+┌─────────────────────────────────────────────────────────────────┐
+│                     vLLM Engine (Worker 0)                        │
+├─────────────────────────────────────────────────────────────────┤
+│  report_usage_stats(vllm_config)                                 │
+│       │                                                           │
+│       ▼                                                           │
+│  ┌─────────────────────────────────────────────────────────────┐ │
+│  │  UsageStatsCollector (new)                                   │ │
+│  │  - collect_hardware()   → GPU topology, cluster               │ │
+│  │  - collect_model()     → HF ID, params, size bucket          │ │
+│  │  - collect_config()    → existing + max_model_len             │ │
+│  │  - collect_runtime()   → workload buckets (Phase 2)          │ │
+│  └─────────────────────────────────────────────────────────────┘ │
+│       │                                                           │
+│       ▼                                                           │
+│  usage_message.report_usage(arch, context, extra_kvs)            │
+│       │                                                           │
+│       ├──► _write_to_file(usage_stats.json)  [append JSONL]       │
+│       └──► _send_to_server(POST stats.vllm.ai)                    │
+└─────────────────────────────────────────────────────────────────┘
+
+┌─────────────────────────────────────────────────────────────────┐
+│  Runtime Aggregator (Phase 2, optional)                          │
+│  - Subscribes to request completion events                       │
+│  - Maintains in-memory bucket counters (prompt, output, batch)    │
+│  - Flushed on heartbeat (every 10 min)                          │
+└─────────────────────────────────────────────────────────────────┘
+```
+
+## 3. Component Design
+
+### 3.1 UsageStatsCollector (New Module)
+
+**Location:** `vllm/usage/usage_collector.py`
+
+**Responsibilities:**
+- Gather hardware, model, and config data
+- Normalize and validate values
+- Return a `dict[str, Any]` of extra key-values to merge with `UsageMessage`
+
+**Interface:**
+```python
+def collect_usage_stats_extras(vllm_config: VllmConfig) -> dict[str, Any]:
+    """Collect additional usage stats. Returns flat KV dict. Never raises."""
+    extras = {}
+    try:
+        extras.update(_collect_hardware_extras())
+        extras.update(_collect_model_extras(vllm_config.model_config))
+        extras.update(_collect_config_extras(vllm_config))
+    except Exception as e:
+        logger.debug("Usage stats collection failed: %s", e)
+    return extras
+```
+
+**Integration point:** `report_usage_stats` in `vllm/v1/utils.py` calls this and merges into `extra_kvs`.
+
+### 3.2 Hardware Collection
+
+**Source:** NVML (pynvml / vllm's pynvml), platform, env vars
+
+| Field | Source | Fallback |
+|-------|--------|----------|
+| gpu_topology_nvlink_links | `nvmlDeviceGetNvLinkState` count | 0 or null |
+| gpu_topology_pcie_gen | `nvmlDeviceGetPcieInfo` → gen | null |
+| gpu_topology_pcie_width | `nvmlDeviceGetPcieInfo` → width | null |
+| gpu_compute_capability | `nvmlDeviceGetCudaComputeCapability` | null |
+| cluster_node_count | `WORLD_SIZE` / `gpu_count` heuristic | 1 |
+| cluster_gpus_per_node | `gpu_count` (single node) or env | null |
+
+**NVML availability:** Use try/import; if NVML unavailable (e.g., CPU, TPU), skip GPU topology fields.
+
+**Code sketch:**
+```python
+def _collect_hardware_extras() -> dict[str, Any]:
+    extras = {}
+    try:
+        import pynvml
+        pynvml.nvmlInit()
+        handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+        # NVLink
+        try:
+            nvlink_count = sum(1 for i in range(pynvml.NVML_NVLINK_MAX_LINKS)
+                              if pynvml.nvmlDeviceGetNvLinkState(handle, i) == 1)
+            extras["gpu_topology_nvlink_links"] = nvlink_count
+        except Exception:
+            extras["gpu_topology_nvlink_links"] = 0
+        # PCIe
+        pcie = pynvml.nvmlDeviceGetPcieInfo(handle)
+        extras["gpu_topology_pcie_gen"] = getattr(pcie, "pcieGen", None)
+        extras["gpu_topology_pcie_width"] = getattr(pcie, "pcieWidth", None)
+        # Compute cap
+        major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
+        extras["gpu_compute_capability"] = f"{major}.{minor}"
+        pynvml.nvmlShutdown()
+    except Exception as e:
+        logger.debug("NVML hardware collection failed: %s", e)
+    return extras
+```
+
+### 3.3 Model Metadata Collection
+
+**Source:** `ModelConfig`, HuggingFace config
+
+| Field | Source |
+|-------|--------|
+| model_hf_id | `model_config.model` (normalize HF path) |
+| model_num_parameters | `model_config.hf_config` or `get_num_params_from_config` |
+| model_size_bucket | Derived from num_parameters (7B, 70B, etc.) |
+| model_hidden_size | `model_config.get_hidden_size()` |
+| model_num_layers | `model_config.get_num_layers(parallel_config)` |
+| model_max_position_embeddings | From `hf_config` or derived |
+
+**Normalization for model_hf_id:**
+- If `model` is HF path (e.g., `meta-llama/Llama-3.1-8B`), use as-is
+- If local path, use `None` or hash of path (avoid leaking paths)
+- Strip query params, revision suffixes for consistency
+
+**Parameter count:** Many HF configs have `num_parameters` or similar; otherwise estimate from `hidden_size * num_layers * 12` (rough transformer formula). Prefer explicit config.
+
+**Code sketch:**
+```python
+def _collect_model_extras(model_config: ModelConfig) -> dict[str, Any]:
+    extras = {}
+    try:
+        model_path = getattr(model_config, "model", None)
+        if model_path and _is_hf_model_id(model_path):
+            extras["model_hf_id"] = _normalize_hf_id(model_path)
+        # num_parameters from config
+        hf_config = getattr(model_config, "hf_config", None)
+        if hf_config:
+            num_params = getattr(hf_config, "num_parameters", None) or ...
+            if num_params:
+                extras["model_num_parameters"] = num_params
+                extras["model_size_bucket"] = _to_size_bucket(num_params)
+        extras["model_hidden_size"] = model_config.get_hidden_size()
+        # num_layers needs parallel_config; may need to pass from vllm_config
+    except Exception as e:
+        logger.debug("Model metadata collection failed: %s", e)
+    return extras
+```
+
+### 3.4 Config Extensions
+
+**Add to existing extra_kvs:**
+- `max_model_len`: `vllm_config.model_config.max_model_len` (after resolution)
+
+### 3.5 Runtime Workload (Phase 2)
+
+**Design:**
+- Add `set_runtime_usage_data` calls at request completion (or equivalent hook)
+- Maintain global counters: `prompt_tokens_bucket`, `output_tokens_bucket`, `batch_size_bucket`
+- Buckets: `[1-64, 65-256, 257-1024, 1025-4096, 4097-16384, 16385+]`
+- On heartbeat, include `prompt_tokens_bucket_counts`, `output_tokens_bucket_counts` as JSON string
+- Reset counters after each report (or sliding window)
+
+**Privacy:** Only bucket indices; no raw values. Aggregation over 10-min window.
+
+**Hook location:** Scheduler or engine output path; must be low-overhead (atomic increment).
+
+### 3.6 Schema Versioning
+
+Add to every payload:
+```json
+"usage_stats_version": 2
+```
+
+Server can route/store by version for backward compatibility.
+
+## 4. Backward Compatibility
+
+| Aspect | Strategy |
+|--------|----------|
+| **Payload** | New fields are additive; server ignores unknown keys |
+| **File format** | Same JSONL append to `usage_stats.json` |
+| **Opt-out** | No change; `VLLM_NO_USAGE_STATS`, `DO_NOT_TRACK`, `do_not_track` file |
+| **Server URL** | Same `VLLM_USAGE_STATS_SERVER`; server must accept new fields |
+| **Curlable** | Payload remains flat KV; `curl -X POST -H "Content-Type: application/json" -d @payload.json $URL` |
+
+## 5. Error Handling and Tolerance
+
+- **Collection:** Each collector (hardware, model, config) wrapped in try/except; failures logged at debug, return partial data
+- **Transmission:** Existing behavior: `requests.exceptions.RequestException` caught, debug log only
+- **File write:** Existing behavior: create dirs, append; failures could raise—consider try/except in `_write_to_file` for extra safety
+- **No blocking:** `report_usage` already uses daemon thread; no change
+
+## 6. Testing Strategy
+
+1. **Unit tests:** `_collect_hardware_extras`, `_collect_model_extras` with mocked NVML, ModelConfig
+2. **Integration:** Run `vllm serve` with usage stats enabled; verify `usage_stats.json` contains new fields
+3. **Opt-out:** Verify `VLLM_NO_USAGE_STATS=1` prevents all collection
+4. **Curlable:** Generate payload, `curl -X POST -d @payload.json https://httpbin.org/post` to validate format
+
+## 7. Dependencies
+
+- **pynvml:** vLLM already has `vllm/third_party/pynvml.py`; use for NVML
+- **No new deps:** All collection uses existing vLLM imports
+
+## 8. Rollout Plan
+
+1. **Phase 1 PR:** Add `UsageStatsCollector`, hardware + model + max_model_len; feature-flagged or always-on for v2
+2. **Server:** Update stats.vllm.ai to accept and store new fields (separate effort)
+3. **Phase 2:** Add runtime aggregator and workload buckets
+4. **Docs:** Update `docs/usage/usage_stats.md` with new fields and example payload
+
+## 9. Open Questions
+
+1. **Model path normalization:** How to handle S3/GCS paths, custom registries?
+2. **Cluster detection:** Reliable multi-node detection without NCCL init?
+3. **Parameter count:** Fallback when `num_parameters` not in config?
+4. **Heartbeat enrichment:** Which counters to add in Phase 1 vs Phase 2?
+
+## 10. Decision Log
+
+| Decision | Rationale |
+|----------|-----------|
+| Flat KV only | Server constraint; simplifies parsing and curlability |
+| NVML for topology | Standard, already in tree; no new deps |
+| Bucket-based workload | Privacy; no raw token counts |
+| Additive schema | Backward compatibility; no migration |
+| Daemon thread | Keeps collection off critical path |

docs/design/usage_stats_v2_prd.md

@@ -0,0 +1,254 @@
+# Usage Stats V2: Product Requirements Document
+
+**Status:** Draft  
+**Project:** Usage Stats V2  
+**Linear Issue:** VLLM-94  
+**Last Updated:** 2025-03-08
+
+## Executive Summary
+
+vLLM's current usage stats capture basic GPU and model configuration but lack the depth needed to drive product decisions around feature deprecation, model support prioritization, and optimization investments. This PRD defines requirements for Usage Stats V2—a comprehensive, privacy-preserving telemetry system that enables data-driven development of vLLM while maintaining backward compatibility, low footprint, and robust queryability.
+
+## 1. Goals and Success Criteria
+
+### Primary Goals
+
+1. **Product Intelligence:** Inform which features to deprecate, which models to prioritize/deprecate, and which optimizations to improve
+2. **Hardware Understanding:** Capture cluster, topology, and networking information to understand deployment patterns
+3. **Model Intelligence:** Beyond architecture class names, capture model metadata (base checkpoints, parameter counts, sizes)
+4. **Workload Characterization:** Understand typical context lengths and input-output distributions (anonymized)
+5. **Future Direction:** Provide data to guide vLLM's roadmap and competitive positioning
+
+### Success Criteria
+
+- **Backward compatible:** Existing opt-out mechanisms and data format remain functional
+- **Low footprint:** Minimal CPU/memory overhead, non-blocking collection
+- **Tolerant:** Failures in collection or transmission do not affect inference
+- **Queryable:** Data easily curlable as key-value pairs; compatible with data warehouse infrastructure
+- **Privacy-preserving:** No sensitive information; workload data sufficiently anonymized
+
+## 2. Current State Analysis
+
+### 2.1 What vLLM Collects Today
+
+**Source:** `vllm/usage/usage_lib.py`, `vllm/v1/utils.py` (report_usage_stats)
+
+| Category | Fields | Notes |
+|----------|--------|-------|
+| **Environment** | provider, num_cpu, cpu_type, cpu_family_model_stepping, total_memory, architecture, platform | Cloud detection via DMI, env vars |
+| **GPU** | gpu_count, gpu_type, gpu_memory_per_device, cuda_runtime | Single device properties; no topology |
+| **Model** | model_architecture | Class name only (e.g., LlamaForCausalLM) |
+| **vLLM Config** | dtype, block_size, gpu_memory_utilization, kv_cache_memory_bytes, quantization, kv_cache_dtype | |
+| **Features** | enable_lora, enable_prefix_caching, enforce_eager, disable_custom_all_reduce | |
+| **Parallelism** | tensor_parallel_size, data_parallel_size, pipeline_parallel_size, enable_expert_parallel, all2all_backend, kv_connector | |
+| **Metadata** | uuid, log_time, source, context (UsageContext) | |
+| **Env Vars** | VLLM_USE_MODELSCOPE, VLLM_USE_FLASHINFER_SAMPLER, etc. | JSON blob |
+
+**Reporting Flow:**
+- One-time report at engine init (from v1 workers: gpu_worker, xpu_worker)
+- Continuous heartbeat every 10 minutes (uuid, log_time, _GLOBAL_RUNTIME_DATA)
+- Writes to `~/.config/vllm/usage_stats.json` (append JSONL)
+- POSTs to `https://stats.vllm.ai`
+
+**Opt-out:** `VLLM_NO_USAGE_STATS`, `DO_NOT_TRACK`, `VLLM_DO_NOT_TRACK`, `~/.config/vllm/do_not_track`
+
+### 2.2 Gaps Identified
+
+| Gap | Impact |
+|-----|--------|
+| **GPU topology** | Cannot distinguish NVLink vs PCIe clusters; affects TP/PP optimization decisions |
+| **Cluster info** | Unknown multi-node vs single-node distribution |
+| **Networking** | No InfiniBand/RoCE vs Ethernet visibility |
+| **Model metadata** | No base model ID, parameter count, or size; hard to prioritize model support |
+| **Context length** | No visibility into typical max_model_len or enabled context |
+| **Workload distribution** | No input/output token distributions; cannot tune batching/scheduling |
+| **Runtime metrics** | Heartbeat has minimal data; no aggregated performance signals |
+
+## 3. Competitive Landscape
+
+### 3.1 Hugging Face Text Generation Inference (TGI)
+
+- **Scope:** Docker-only; startup/shutdown + 15-min heartbeat
+- **Model:** model_type, tokenizer_class, revision
+- **System:** CPU count/type, memory, architecture, platform
+- **GPU:** device name, driver version, memory usage, temperature, utilization, power draw, compute capability, ECC errors
+- **Config:** max_batch_prefill_tokens, max_batch_size, max_input_tokens, max_total_tokens, max_concurrent_requests
+- **Opt-out:** `--usage-stats=no-stack` (omit stack traces), `--usage-stats=off` (full disable)
+
+### 3.2 NVIDIA TensorRT-LLM
+
+- **Metrics:** Prometheus endpoint (`/prometheus/metrics`)
+- **Request-level:** KV cache hit rates, TTFT, TPOT, e2e latency, queue time, finish reasons
+- **Aggregate:** Running/waiting requests, prompt/generation token totals, per-iteration stats
+- **Model:** Model name and version tracking
+
+### 3.3 Takeaways
+
+- TGI: Rich GPU telemetry (temp, power, ECC); batch/config limits; Docker-gated
+- TensorRT-LLM: Production-oriented metrics; Prometheus-native; request-level histograms
+- vLLM opportunity: Combine hardware topology (unique), model metadata (HF-style), and workload distributions (anonymized) for product intelligence
+
+## 4. Requirements
+
+### 4.1 Hardware / GPU (Cluster, Topology, Networking)
+
+| Requirement | Priority | Description |
+|-------------|----------|-------------|
+| H1 | P0 | **GPU topology:** NVLink link count, PCIe gen/width, P2P capability summary |
+| H2 | P0 | **Cluster size:** Number of nodes (if detectable), GPUs per node |
+| H3 | P1 | **Networking:** InfiniBand/RoCE vs Ethernet (via env or driver hints) |
+| H4 | P1 | **Compute capability:** SM version (e.g., 8.0, 9.0) for kernel compatibility |
+| H5 | P2 | **Driver/CUDA:** Driver version, CUDA version (already partially present) |
+
+**Implementation hints:**
+- NVML: `nvmlDeviceGetNvLinkState`, `nvmlDeviceGetPcieInfo`, `nvmlDeviceGetCudaComputeCapability`
+- `nvidia-smi topo -m` (parse or use NVML equivalents)
+- Cluster: `NCCL_*` env vars, `RANK`, `WORLD_SIZE` patterns
+
+### 4.2 Model Metadata
+
+| Requirement | Priority | Description |
+|-------------|----------|-------------|
+| M1 | P0 | **Base model identifier:** HuggingFace model ID or normalized name (e.g., `meta-llama/Llama-3.1-8B`) |
+| M2 | P0 | **Parameter count:** Total parameters (from config or loaded model) |
+| M3 | P1 | **Model size bucket:** e.g., 7B, 70B, 405B for bucketing |
+| M4 | P1 | **Hidden size, num layers:** For architecture variant differentiation |
+| M5 | P2 | **Revision/commit:** Model revision if from HF Hub |
+
+**Privacy:** Use model ID from public registries (HF, etc.); avoid custom paths that could leak org names.
+
+### 4.3 Runtime / Workload (Anonymized)
+
+| Requirement | Priority | Description |
+|-------------|----------|-------------|
+| R1 | P0 | **max_model_len:** Configured context length (already derivable; ensure reported) |
+| R2 | P0 | **Input/output distributions:** Histogram buckets (e.g., prompt_tokens: 1-64, 65-256, 257-1K, 1K-4K, 4K+) |
+| R3 | P1 | **Batch size distribution:** Concurrent request counts over time |
+| R4 | P1 | **Feature usage:** Prefix cache hit rate, speculative decoding usage (if enabled) |
+| R5 | P2 | **Latency percentiles:** TTFT, TPOT (aggregated, no per-request) |
+
+**Anonymization rules:**
+- No raw token counts; only bucket indices or aggregates
+- No request IDs, timestamps finer than minute-level
+- No user-identifiable content
+
+### 4.4 Non-Functional
+
+| Requirement | Description |
+|-------------|-------------|
+| NF1 | **Backward compatibility:** Existing `usage_stats.json` format and server contract; new fields additive |
+| NF2 | **Opt-out:** All existing mechanisms remain; no new tracking without opt-out |
+| NF3 | **Low footprint:** Collection async, <1% CPU; no blocking on network |
+| NF4 | **Tolerance:** Exceptions in collection do not affect inference |
+| NF5 | **Curlable:** Payload as flat key-value JSON; easy to inspect and replay |
+| NF6 | **Warehouse-ready:** Schema stable; suitable for BigQuery/Snowflake/etc. |
+
+## 5. Data Schema (Proposed Additions)
+
+### 5.1 New Fields (Flat KV)
+
+```
+# Hardware
+gpu_topology_nvlink_links: int | null      # NVLink links per GPU (0 if PCIe-only)
+gpu_topology_pcie_gen: int | null         # e.g., 4, 5
+gpu_topology_pcie_width: int | null       # e.g., 16
+gpu_compute_capability: str | null        # e.g., "8.0", "9.0"
+cluster_node_count: int | null            # 1 if single-node
+cluster_gpus_per_node: int | null
+
+# Model
+model_hf_id: str | null                   # e.g., "meta-llama/Llama-3.1-8B"
+model_num_parameters: int | null          # e.g., 8_000_000_000
+model_size_bucket: str | null              # "7B", "70B", "405B"
+model_hidden_size: int | null
+model_num_layers: int | null
+model_max_position_embeddings: int | null
+
+# Runtime (anonymized)
+max_model_len: int                         # configured context length
+prompt_tokens_bucket_counts: str | null   # JSON: {"1-64": N, "65-256": M, ...}
+output_tokens_bucket_counts: str | null   # same
+batch_size_bucket_counts: str | null      # concurrent requests over time
+prefix_cache_hit_rate_avg: float | null   # 0-1, if prefix caching enabled
+```
+
+### 5.2 Heartbeat Enrichment
+
+Continuous heartbeat (every 10 min) could include:
+- `uptime_seconds`
+- `total_requests_served` (counter)
+- `total_prompt_tokens`, `total_output_tokens` (counters)
+- `prefix_cache_hits`, `prefix_cache_queries` (if enabled)
+
+## 6. Privacy and Anonymization
+
+### 6.1 Principles
+
+- **Minimize:** Collect only what is needed for stated goals
+- **Anonymize:** No PII; workload data in coarse buckets only
+- **Transparency:** Document all fields in public docs (as today)
+- **User control:** Opt-out must be easy and complete
+
+### 6.2 Sensitive Data (Never Collect)
+
+- Model paths that could reveal org/user (e.g., `/home/company/secret-model`)
+- Request content, prompts, or outputs
+- IP addresses, hostnames (beyond cloud provider detection)
+- Fine-grained timestamps that could correlate with external logs
+
+### 6.3 Workload Anonymization
+
+- **Prompt/output lengths:** Use fixed buckets (e.g., 1-64, 65-256, 257-1024, 1025-4096, 4097-16384, 16385+)
+- **Aggregation:** Report counts per bucket over a time window (e.g., 10 min)
+- **No linkage:** No request-level data; only aggregates
+
+## 7. Open Source and Product Angles
+
+### 7.1 Open Source
+
+- Public schema and documentation
+- Transparent opt-out
+- Align with DO_NOT_TRACK and similar conventions
+- Publish aggregated reports (e.g., 2024.vllm.ai) to build trust
+
+### 7.2 Product
+
+- Data drives deprecation (low-usage features)
+- Model support prioritization (high-usage architectures)
+- Optimization ROI (where to invest in kernels, memory, etc.)
+- Competitive differentiation (e.g., "vLLM runs on X% of H100 clusters")
+
+### 7.3 Mature PM Considerations
+
+- **Stakeholder trust:** Clear docs, minimal data, strong opt-out
+- **Legal:** GDPR/privacy compliance; avoid consent fatigue
+- **Competitive:** Don't over-collect; avoid perception of surveillance
+- **Technical debt:** Design for schema evolution; version the payload
+
+## 8. Implementation Phases
+
+### Phase 1 (MVP)
+- Add model metadata (M1, M2, M3)
+- Add max_model_len to initial report
+- Add GPU topology (H1, H4) where NVML available
+- No runtime workload yet
+
+### Phase 2
+- Add cluster info (H2)
+- Add workload buckets (R1, R2) with anonymization
+- Enrich heartbeat with aggregate counters
+
+### Phase 3
+- Add networking hints (H3)
+- Add feature usage (R4)
+- Refine schema based on warehouse usage
+
+## 9. Appendix: Current Code References
+
+- **Usage lib:** `vllm/usage/usage_lib.py`
+- **Report call:** `vllm/v1/utils.py::report_usage_stats`
+- **Workers:** `vllm/v1/worker/gpu_worker.py`, `vllm/v1/worker/xpu_worker.py`
+- **Model loader:** `vllm/model_executor/model_loader/utils.py::get_architecture_class_name`
+- **Platform utils:** `vllm/utils/platform_utils.py::cuda_get_device_properties`
+- **NVML bindings:** `vllm/third_party/pynvml.py` (NVLink, PCIe constants)
+- **Config:** `vllm/config/model.py` (ModelConfig: max_model_len, get_hidden_size, get_num_layers)

docs/models/supported_models.md

@@ -469,6 +469,8 @@ th {
 | `Qwen3MoeForCausalLM` | Qwen3MoE | `Qwen/Qwen3-30B-A3B`, etc. | ✅︎ | ✅︎ |
 | `Qwen3NextForCausalLM` | Qwen3NextMoE | `Qwen/Qwen3-Next-80B-A3B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `RWForCausalLM` | Falcon RW | `tiiuae/falcon-40b`, etc. | | ✅︎ |
+| `SarvamMoEForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-30b-a3b`, etc. | ✅︎ | ✅︎ |
+| `SarvamMLAForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-105b-a9b`, etc. | | ✅︎ |
 | `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `SolarForCausalLM` | Solar Pro | `upstage/solar-pro-preview-instruct`, etc. | ✅︎ | ✅︎ |
 | `StableLmForCausalLM` | StableLM | `stabilityai/stablelm-3b-4e1t`, `stabilityai/stablelm-base-alpha-7b-v2`, etc. | | |

docs/serving/integrations/llamaindex.md

@@ -17,7 +17,7 @@ llm = Vllm(
     model="microsoft/Orca-2-7b",
     tensor_parallel_size=4,
     max_new_tokens=100,
-    vllm_kwargs={"swap_space": 1, "gpu_memory_utilization": 0.5},
+    vllm_kwargs={"gpu_memory_utilization": 0.5},
 )
 ```
 

docs/usage/usage_stats.md

@@ -59,3 +59,10 @@ export VLLM_NO_USAGE_STATS=1
 export DO_NOT_TRACK=1
 mkdir -p ~/.config/vllm && touch ~/.config/vllm/do_not_track
 ```
+
+## Usage Stats V2 (Planned)
+
+A comprehensive redesign of usage stats is planned to capture additional hardware topology, model metadata, and anonymized workload information. This will help guide vLLM's product direction, feature deprecation, and optimization priorities. See the design documentation:
+
+- [Usage Stats V2 PRD](../design/usage_stats_v2_prd.md) - Product requirements and competitive analysis
+- [Usage Stats V2 Design](../design/usage_stats_v2_design.md) - Technical design and implementation sketch

tests/conftest.py

@@ -794,7 +794,6 @@ class VllmRunner:
         tensor_parallel_size: int = 1,
         block_size: int = 16 if not torch.xpu.is_available() else 64,
         enable_chunked_prefill: bool | None = False,
-        swap_space: int = 4,
         enforce_eager: bool | None = False,
         # Set this to avoid hanging issue
         default_torch_num_threads: int | None = None,
@@ -831,7 +830,6 @@ class VllmRunner:
                 trust_remote_code=trust_remote_code,
                 dtype=dtype,
                 seed=seed,
-                swap_space=swap_space,
                 enforce_eager=enforce_eager,
                 disable_log_stats=disable_log_stats,
                 tensor_parallel_size=tensor_parallel_size,

tests/distributed/test_torchrun_example.py

@@ -22,7 +22,7 @@ prompts = [
 
 sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
 
-# set different `gpu_memory_utilization` and `swap_space` for different ranks,
+# set different `gpu_memory_utilization` for different ranks,
 # to test if all ranks agree on the same kv cache configuration.
 llm = LLM(
     model="facebook/opt-125m",
@@ -30,7 +30,6 @@ llm = LLM(
     pipeline_parallel_size=int(os.getenv("PP_SIZE", 1)),
     distributed_executor_backend="external_launcher",
     gpu_memory_utilization=random.uniform(0.7, 0.9),
-    swap_space=random.randint(1, 4),
     seed=0,
 )
 

tests/distributed/test_torchrun_example_moe.py

@@ -28,7 +28,7 @@ if dp_size > 1:
 
 sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
 
-# set different `gpu_memory_utilization` and `swap_space` for different ranks,
+# set different `gpu_memory_utilization` for different ranks,
 # to test if all ranks agree on the same kv cache configuration.
 llm = LLM(
     model="microsoft/Phi-mini-MoE-instruct",
@@ -37,7 +37,6 @@ llm = LLM(
     enable_expert_parallel=int(os.getenv("ENABLE_EP", "0")) == 1,
     distributed_executor_backend="external_launcher",
     gpu_memory_utilization=random.uniform(0.7, 0.9),
-    swap_space=random.randint(1, 4),
     seed=0,
 )
 

tests/entrypoints/openai/test_realtime_validation.py

@@ -13,7 +13,7 @@ import websockets
 
 from vllm.assets.audio import AudioAsset
 
-from ...utils import RemoteOpenAIServer
+from ...utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
 from .conftest import add_attention_backend
 
 MISTRAL_FORMAT_ARGS = [
@@ -23,7 +23,7 @@ MISTRAL_FORMAT_ARGS = [
     "mistral",
     "--load_format",
     "mistral",
-]
+] + ROCM_EXTRA_ARGS
 
 MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602"
 
@@ -77,7 +77,9 @@ async def test_multi_chunk_streaming(
 
     add_attention_backend(server_args, rocm_aiter_fa_attention)
 
-    with RemoteOpenAIServer(model_name, server_args) as remote_server:
+    with RemoteOpenAIServer(
+        model_name, server_args, env_dict=ROCM_ENV_OVERRIDES
+    ) as remote_server:
         ws_url = _get_websocket_url(remote_server)
         async with websockets.connect(ws_url) as ws:
             # Receive session.created
@@ -178,7 +180,9 @@ async def test_empty_commit_does_not_crash_engine(
 
     add_attention_backend(server_args, rocm_aiter_fa_attention)
 
-    with RemoteOpenAIServer(model_name, server_args) as remote_server:
+    with RemoteOpenAIServer(
+        model_name, server_args, env_dict=ROCM_ENV_OVERRIDES
+    ) as remote_server:
         ws_url = _get_websocket_url(remote_server)
 
         # --- First connection: empty commit (no audio appended) ----------

tests/entrypoints/openai/test_transcription_validation.py

@@ -6,7 +6,7 @@ import json
 
 import pytest
 
-from ...utils import RemoteOpenAIServer
+from ...utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
 from .conftest import add_attention_backend
 
 MISTRAL_FORMAT_ARGS = [
@@ -19,12 +19,55 @@ MISTRAL_FORMAT_ARGS = [
 ]
 
 
+async def transcribe_and_check(
+    client,
+    model_name: str,
+    file,
+    *,
+    language: str,
+    expected_text: str,
+    expected_seconds: int | None = None,
+    case_sensitive: bool = False,
+):
+    """Run a transcription request and assert the output contains
+    *expected_text* and optionally that usage reports *expected_seconds*.
+
+    Provides detailed failure messages with the actual transcription output.
+    """
+    transcription = await client.audio.transcriptions.create(
+        model=model_name,
+        file=file,
+        language=language,
+        response_format="text",
+        temperature=0.0,
+    )
+    out = json.loads(transcription)
+    out_text = out["text"]
+    out_usage = out["usage"]
+
+    if case_sensitive:
+        assert expected_text in out_text, (
+            f"Expected {expected_text!r} in transcription output, got: {out_text!r}"
+        )
+    else:
+        assert expected_text.lower() in out_text.lower(), (
+            f"Expected {expected_text!r} (case-insensitive) in transcription "
+            f"output, got: {out_text!r}"
+        )
+
+    if expected_seconds is not None:
+        assert out_usage["seconds"] == expected_seconds, (
+            f"Expected {expected_seconds}s of audio, "
+            f"got {out_usage['seconds']}s. Full usage: {out_usage!r}"
+        )
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "model_name", ["mistralai/Voxtral-Mini-3B-2507", "Qwen/Qwen3-ASR-0.6B"]
 )
 async def test_basic_audio(mary_had_lamb, model_name, rocm_aiter_fa_attention):
-    server_args = ["--enforce-eager"]
+    server_args = ["--enforce-eager", *ROCM_EXTRA_ARGS]
 
     if model_name.startswith("mistralai"):
         server_args += MISTRAL_FORMAT_ARGS
@@ -32,20 +75,18 @@ async def test_basic_audio(mary_had_lamb, model_name, rocm_aiter_fa_attention):
     add_attention_backend(server_args, rocm_aiter_fa_attention)
 
     # Based on https://github.com/openai/openai-cookbook/blob/main/examples/Whisper_prompting_guide.ipynb.
-    with RemoteOpenAIServer(model_name, server_args) as remote_server:
+    with RemoteOpenAIServer(
+        model_name, server_args, env_dict=ROCM_ENV_OVERRIDES
+    ) as remote_server:
         client = remote_server.get_async_client()
-        transcription = await client.audio.transcriptions.create(
-            model=model_name,
-            file=mary_had_lamb,
+        await transcribe_and_check(
+            client,
+            model_name,
+            mary_had_lamb,
             language="en",
-            response_format="text",
-            temperature=0.0,
+            expected_text="Mary had a little lamb",
+            expected_seconds=16,
         )
-        out = json.loads(transcription)
-        out_text = out["text"]
-        out_usage = out["usage"]
-        assert "Mary had a little lamb" in out_text
-        assert out_usage["seconds"] == 16, out_usage["seconds"]
 
 
 @pytest.mark.asyncio
@@ -74,20 +115,18 @@ async def test_basic_audio_with_lora(mary_had_lamb, rocm_aiter_fa_attention):
     add_attention_backend(server_args, rocm_aiter_fa_attention)
 
     # Based on https://github.com/openai/openai-cookbook/blob/main/examples/Whisper_prompting_guide.ipynb.
-    with RemoteOpenAIServer(model_name, server_args) as remote_server:
+    with RemoteOpenAIServer(
+        model_name, server_args, env_dict=ROCM_ENV_OVERRIDES
+    ) as remote_server:
         client = remote_server.get_async_client()
-        transcription = await client.audio.transcriptions.create(
-            model=lora_model_name,
-            file=mary_had_lamb,
+        await transcribe_and_check(
+            client,
+            lora_model_name,
+            mary_had_lamb,
             language="en",
-            response_format="text",
-            temperature=0.0,
+            expected_text="mary had a little lamb",
+            expected_seconds=16,
         )
-    out = json.loads(transcription)
-    out_text = out["text"]
-    out_usage = out["usage"]
-    assert "mary had a little lamb" in out_text
-    assert out_usage["seconds"] == 16, out_usage["seconds"]
 
 
 @pytest.mark.asyncio
@@ -97,20 +136,21 @@ async def test_basic_audio_with_lora(mary_had_lamb, rocm_aiter_fa_attention):
 async def test_basic_audio_foscolo(foscolo, rocm_aiter_fa_attention, model_name):
     # Gemma accuracy on some of the audio samples we use is particularly bad,
     # hence we use a different one here. WER is evaluated separately.
-    server_args = ["--enforce-eager"]
+    server_args = ["--enforce-eager", *ROCM_EXTRA_ARGS]
 
     add_attention_backend(server_args, rocm_aiter_fa_attention)
 
     with RemoteOpenAIServer(
-        model_name, server_args, max_wait_seconds=480
+        model_name,
+        server_args,
+        max_wait_seconds=480,
+        env_dict=ROCM_ENV_OVERRIDES,
     ) as remote_server:
         client = remote_server.get_async_client()
-        transcription = await client.audio.transcriptions.create(
-            model=model_name,
-            file=foscolo,
+        await transcribe_and_check(
+            client,
+            model_name,
+            foscolo,
             language="it",
-            response_format="text",
-            temperature=0.0,
+            expected_text="ove il mio corpo fanciulletto giacque",
         )
-        out = json.loads(transcription)["text"]
-        assert "ove il mio corpo fanciulletto giacque" in out

tests/entrypoints/openai/test_video.py

@@ -35,6 +35,8 @@ def server():
         "--trust-remote-code",
         "--limit-mm-per-prompt",
         json.dumps({"video": MAXIMUM_VIDEOS}),
+        "--media-io-kwargs",
+        json.dumps({"video": {"num_frames": 32}}),
     ]
 
     # ROCm: Increase timeouts to handle potential network delays and slower
@@ -127,6 +129,73 @@ async def test_single_chat_session_video(
     assert message.content is not None and len(message.content) >= 0
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+@pytest.mark.parametrize("video_url", [TEST_VIDEO_URLS[0]])
+async def test_request_media_io_kwargs_override_uses_fewer_video_frames(
+    client: openai.AsyncOpenAI, model_name: str, video_url: str
+):
+    messages = dummy_messages_from_video_url(video_url)
+
+    default_resp = await client.chat.completions.create(
+        model=model_name,
+        messages=messages,
+        max_completion_tokens=1,
+        temperature=0.0,
+    )
+    override_resp = await client.chat.completions.create(
+        model=model_name,
+        messages=messages,
+        max_completion_tokens=1,
+        temperature=0.0,
+        extra_body={
+            "media_io_kwargs": {
+                "video": {
+                    "num_frames": 4,
+                }
+            }
+        },
+    )
+
+    assert default_resp.usage is not None
+    assert override_resp.usage is not None
+    assert override_resp.usage.prompt_tokens < default_resp.usage.prompt_tokens
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+@pytest.mark.parametrize("video_url", [TEST_VIDEO_URLS[0]])
+async def test_invalid_num_frames_request_recoverable(
+    client: openai.AsyncOpenAI, model_name: str, video_url: str
+):
+    messages = dummy_messages_from_video_url(video_url)
+
+    with pytest.raises((openai.BadRequestError, openai.APIStatusError)):
+        await client.chat.completions.create(
+            model=model_name,
+            messages=messages,
+            max_completion_tokens=1,
+            temperature=0.0,
+            extra_body={
+                "media_io_kwargs": {
+                    "video": {
+                        "num_frames": "invalid",
+                    }
+                }
+            },
+        )
+
+    # Server should still handle subsequent requests after the failed one.
+    recovery_resp = await client.chat.completions.create(
+        model=model_name,
+        messages=messages,
+        max_completion_tokens=1,
+        temperature=0.0,
+    )
+    recovery_msg = recovery_resp.choices[0].message
+    assert recovery_msg.content is not None and len(recovery_msg.content) >= 0
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
 @pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)

tests/entrypoints/openai/test_vision.py

@@ -12,7 +12,7 @@ from vllm.multimodal.media import MediaWithBytes
 from vllm.multimodal.utils import encode_image_url, fetch_image
 from vllm.platforms import current_platform
 
-from ...utils import RemoteOpenAIServer
+from ...utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
 
 MODEL_NAME = "microsoft/Phi-3.5-vision-instruct"
 MAXIMUM_IMAGES = 2
@@ -48,10 +48,37 @@ def check_output_matches_terms(content: str, term_groups: list[list[str]]) -> bo
     All term groups must be satisfied.
     """
     content_lower = content.lower()
-    for group in term_groups:
-        if not any(term.lower() in content_lower for term in group):
-            return False
-    return True
+    return all(
+        any(term.lower() in content_lower for term in group) for group in term_groups
+    )
+
+
+def assert_non_empty_content(chat_completion, *, context: str = "") -> str:
+    """Assert the first choice has non-empty string content; return it.
+
+    Provides a detailed failure message including the full ChatCompletion
+    response so flaky / model-quality issues are easy to diagnose.
+    """
+    prefix = f"[{context}] " if context else ""
+    choice = chat_completion.choices[0]
+    content = choice.message.content
+
+    assert content is not None, (
+        f"{prefix}Expected non-None content but got None. "
+        f"finish_reason={choice.finish_reason!r}, "
+        f"full message={choice.message!r}, "
+        f"usage={chat_completion.usage!r}"
+    )
+    assert isinstance(content, str), (
+        f"{prefix}Expected str content, got {type(content).__name__}: {content!r}"
+    )
+    assert len(content) > 0, (
+        f"{prefix}Expected non-empty content but got empty string. "
+        f"finish_reason={choice.finish_reason!r}, "
+        f"full message={choice.message!r}, "
+        f"usage={chat_completion.usage!r}"
+    )
+    return content
 
 
 @pytest.fixture(scope="module")
@@ -67,16 +94,22 @@ def server():
         "--trust-remote-code",
         "--limit-mm-per-prompt",
         json.dumps({"image": MAXIMUM_IMAGES}),
+        *ROCM_EXTRA_ARGS,
     ]
 
     # ROCm: Increase timeouts to handle potential network delays and slower
     # video processing when downloading multiple videos from external sources
-    env_overrides = {}
-    if current_platform.is_rocm():
-        env_overrides = {
-            "VLLM_VIDEO_FETCH_TIMEOUT": "120",
-            "VLLM_ENGINE_ITERATION_TIMEOUT_S": "300",
-        }
+    env_overrides = {
+        **ROCM_ENV_OVERRIDES,
+        **(
+            {
+                "VLLM_VIDEO_FETCH_TIMEOUT": "120",
+                "VLLM_ENGINE_ITERATION_TIMEOUT_S": "300",
+            }
+            if current_platform.is_rocm()
+            else {}
+        ),
+    }
 
     with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_overrides) as remote_server:
         yield remote_server
@@ -117,6 +150,51 @@ def dummy_messages_from_image_url(
     ]
 
 
+def describe_image_messages(
+    image_url: str, *, extra_image_fields: dict | None = None
+) -> list[dict]:
+    """Build the system + user messages used by the completions-with-image
+    family of tests. *extra_image_fields* is merged into the top-level
+    image content block (for uuid / bad-key tests)."""
+    image_block: dict = {
+        "type": "image_url",
+        "image_url": {"url": image_url},
+    }
+    if extra_image_fields:
+        image_block.update(extra_image_fields)
+
+    return [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Describe this image."},
+                image_block,
+            ],
+        },
+    ]
+
+
+async def complete_and_check(
+    client: openai.AsyncOpenAI,
+    model_name: str,
+    messages: list[dict],
+    *,
+    context: str,
+    max_completion_tokens: int = 50,
+    temperature: float = 0.0,
+) -> str:
+    """Run a chat completion and assert the output is non-empty.
+    Returns the content string."""
+    chat_completion = await client.chat.completions.create(
+        model=model_name,
+        messages=messages,
+        max_completion_tokens=max_completion_tokens,
+        temperature=temperature,
+    )
+    return assert_non_empty_content(chat_completion, context=context)
+
+
 def get_hf_prompt_tokens(model_name, content, image_url):
     processor = AutoProcessor.from_pretrained(
         model_name, trust_remote_code=True, num_crops=4
@@ -153,7 +231,6 @@ async def test_single_chat_session_image(
     messages = dummy_messages_from_image_url(image_url, content_text)
 
     max_completion_tokens = 10
-    # test single completion
     chat_completion = await client.chat.completions.create(
         model=model_name,
         messages=messages,
@@ -162,32 +239,46 @@ async def test_single_chat_session_image(
         temperature=0.0,
         top_logprobs=5,
     )
-    assert len(chat_completion.choices) == 1
+    assert len(chat_completion.choices) == 1, (
+        f"Expected 1 choice, got {len(chat_completion.choices)}"
+    )
 
     choice = chat_completion.choices[0]
-    assert choice.finish_reason == "length"
+    assert choice.finish_reason == "length", (
+        f"Expected finish_reason='length' (capped at {max_completion_tokens} "
+        f"tokens), got {choice.finish_reason!r}. "
+        f"content={choice.message.content!r}"
+    )
+
     hf_prompt_tokens = get_hf_prompt_tokens(model_name, content_text, image_url)
-    assert chat_completion.usage == openai.types.CompletionUsage(
+    expected_usage = openai.types.CompletionUsage(
         completion_tokens=max_completion_tokens,
         prompt_tokens=hf_prompt_tokens,
         total_tokens=hf_prompt_tokens + max_completion_tokens,
     )
+    assert chat_completion.usage == expected_usage, (
+        f"Usage mismatch: got {chat_completion.usage!r}, expected {expected_usage!r}"
+    )
 
     message = choice.message
-    message = chat_completion.choices[0].message
-    assert message.content is not None and len(message.content) >= 10
-    assert message.role == "assistant"
+    assert message.content is not None and len(message.content) >= 10, (
+        f"Expected content with >=10 chars, got {message.content!r}"
+    )
+    assert message.role == "assistant", (
+        f"Expected role='assistant', got {message.role!r}"
+    )
+
     messages.append({"role": "assistant", "content": message.content})
 
     # test multi-turn dialogue
     messages.append({"role": "user", "content": "express your result in json"})
-    chat_completion = await client.chat.completions.create(
-        model=model_name,
-        messages=messages,
+    await complete_and_check(
+        client,
+        model_name,
+        messages,
+        context=f"multi-turn follow-up for {image_url}",
         max_completion_tokens=10,
     )
-    message = chat_completion.choices[0].message
-    assert message.content is not None and len(message.content) >= 0
 
 
 @pytest.mark.asyncio
@@ -209,7 +300,7 @@ async def test_error_on_invalid_image_url_type(
 
     # image_url should be a dict {"url": "some url"}, not directly a string
     with pytest.raises(openai.BadRequestError):
-        _ = await client.chat.completions.create(
+        await client.chat.completions.create(
             model=model_name,
             messages=messages,
             max_completion_tokens=10,
@@ -235,10 +326,15 @@ async def test_single_chat_session_image_beamsearch(
         top_logprobs=5,
         extra_body=dict(use_beam_search=True),
     )
-    assert len(chat_completion.choices) == 2
-    assert (
-        chat_completion.choices[0].message.content
-        != chat_completion.choices[1].message.content
+    assert len(chat_completion.choices) == 2, (
+        f"Expected 2 beam search choices, got {len(chat_completion.choices)}"
+    )
+
+    content_0 = chat_completion.choices[0].message.content
+    content_1 = chat_completion.choices[1].message.content
+    assert content_0 != content_1, (
+        f"Beam search should produce different outputs for {image_url}, "
+        f"but both returned: {content_0!r}"
     )
 
 
@@ -269,33 +365,46 @@ async def test_single_chat_session_image_base64encoded(
         temperature=0.0,
         top_logprobs=5,
     )
-    assert len(chat_completion.choices) == 1
+    assert len(chat_completion.choices) == 1, (
+        f"Expected 1 choice, got {len(chat_completion.choices)}"
+    )
 
     choice = chat_completion.choices[0]
-    assert choice.finish_reason == "length"
+    assert choice.finish_reason == "length", (
+        f"Expected finish_reason='length', got {choice.finish_reason!r}. "
+        f"content={choice.message.content!r}"
+    )
+
     hf_prompt_tokens = get_hf_prompt_tokens(model_name, content_text, image_url)
-    assert chat_completion.usage == openai.types.CompletionUsage(
+    expected_usage = openai.types.CompletionUsage(
         completion_tokens=max_completion_tokens,
         prompt_tokens=hf_prompt_tokens,
         total_tokens=hf_prompt_tokens + max_completion_tokens,
     )
+    assert chat_completion.usage == expected_usage, (
+        f"Usage mismatch: got {chat_completion.usage!r}, expected {expected_usage!r}"
+    )
 
     message = choice.message
-    message = chat_completion.choices[0].message
-    assert message.content is not None and len(message.content) >= 10
-    assert message.role == "assistant"
+    assert message.content is not None and len(message.content) >= 10, (
+        f"Expected content with >=10 chars, got {message.content!r}"
+    )
+    assert message.role == "assistant", (
+        f"Expected role='assistant', got {message.role!r}"
+    )
+
     messages.append({"role": "assistant", "content": message.content})
 
     # test multi-turn dialogue
     messages.append({"role": "user", "content": "express your result in json"})
-    chat_completion = await client.chat.completions.create(
-        model=model_name,
-        messages=messages,
+    await complete_and_check(
+        client,
+        model_name,
+        messages,
+        context=f"multi-turn base64 follow-up for {raw_image_url}",
         max_completion_tokens=10,
         temperature=0.0,
     )
-    message = chat_completion.choices[0].message
-    assert message.content is not None and len(message.content) >= 0
 
 
 @pytest.mark.asyncio
@@ -321,7 +430,10 @@ async def test_single_chat_session_image_base64encoded_beamsearch(
         temperature=0.0,
         extra_body=dict(use_beam_search=True),
     )
-    assert len(chat_completion.choices) == 2
+    assert len(chat_completion.choices) == 2, (
+        f"Expected 2 beam search choices for image {image_idx} "
+        f"({raw_image_url}), got {len(chat_completion.choices)}"
+    )
 
     # Verify beam search produces two different non-empty outputs
     content_0 = chat_completion.choices[0].message.content
@@ -333,18 +445,28 @@ async def test_single_chat_session_image_base64encoded_beamsearch(
         f"Output 0: {content_0!r}, Output 1: {content_1!r}"
     )
 
-    assert content_0, "First beam search output should not be empty"
-    assert content_1, "Second beam search output should not be empty"
-    assert content_0 != content_1, "Beam search should produce different outputs"
+    assert content_0, (
+        f"First beam output is empty for image {image_idx} ({raw_image_url}). "
+        f"finish_reason={chat_completion.choices[0].finish_reason!r}"
+    )
+    assert content_1, (
+        f"Second beam output is empty for image {image_idx} "
+        f"({raw_image_url}). "
+        f"finish_reason={chat_completion.choices[1].finish_reason!r}"
+    )
+    assert content_0 != content_1, (
+        f"Beam search produced identical outputs for image {image_idx} "
+        f"({raw_image_url}): {content_0!r}"
+    )
 
     # Verify each output contains the required terms for this image
     for i, content in enumerate([content_0, content_1]):
-        if not check_output_matches_terms(content, required_terms):
-            pytest.fail(
-                f"Output {i} '{content}' doesn't contain required terms. "
-                f"Expected all of these term groups (at least one from each): "
-                f"{required_terms}"
-            )
+        assert check_output_matches_terms(content, required_terms), (
+            f"Beam output {i} for image {image_idx} ({raw_image_url}) "
+            f"doesn't match required terms.\n"
+            f"  content: {content!r}\n"
+            f"  required (all groups, >=1 per group): {required_terms}"
+        )
 
 
 @pytest.mark.asyncio
@@ -378,16 +500,29 @@ async def test_chat_streaming_image(
     async for chunk in stream:
         delta = chunk.choices[0].delta
         if delta.role:
-            assert delta.role == "assistant"
+            assert delta.role == "assistant", (
+                f"Expected role='assistant' in stream delta, got {delta.role!r}"
+            )
         if delta.content:
             chunks.append(delta.content)
         if chunk.choices[0].finish_reason is not None:
             finish_reason_count += 1
     # finish reason should only return in last block
-    assert finish_reason_count == 1
-    assert chunk.choices[0].finish_reason == stop_reason
-    assert delta.content
-    assert "".join(chunks) == output
+    assert finish_reason_count == 1, (
+        f"Expected exactly 1 finish_reason across stream chunks, "
+        f"got {finish_reason_count}"
+    )
+    assert chunk.choices[0].finish_reason == stop_reason, (
+        f"Stream finish_reason={chunk.choices[0].finish_reason!r} "
+        f"doesn't match non-stream finish_reason={stop_reason!r}"
+    )
+
+    streamed_text = "".join(chunks)
+    assert streamed_text == output, (
+        f"Streamed output doesn't match non-streamed for {image_url}.\n"
+        f"  streamed:     {streamed_text!r}\n"
+        f"  non-streamed: {output!r}"
+    )
 
 
 @pytest.mark.asyncio
@@ -418,17 +553,19 @@ async def test_multi_image_input(
             max_tokens=5,
             temperature=0.0,
         )
-        completion = completion.choices[0].text
-        assert completion is not None and len(completion) >= 0
+        assert completion.choices[0].text is not None, (
+            "Server failed to produce output after rejecting over-limit "
+            "multi-image request"
+        )
     else:
-        chat_completion = await client.chat.completions.create(
-            model=model_name,
-            messages=messages,
+        await complete_and_check(
+            client,
+            model_name,
+            messages,
+            context=f"multi-image input ({len(image_urls)} images)",
             max_completion_tokens=10,
             temperature=0.0,
         )
-        message = chat_completion.choices[0].message
-        assert message.content is not None and len(message.content) >= 0
 
 
 @pytest.mark.asyncio
@@ -444,30 +581,13 @@ async def test_completions_with_image(
     image_urls: list[str],
 ):
     for image_url in image_urls:
-        chat_completion = await client.chat.completions.create(
-            messages=[
-                {"role": "system", "content": "You are a helpful assistant."},
-                {
-                    "role": "user",
-                    "content": [
-                        {
-                            "type": "text",
-                            "text": "Describe this image.",
-                        },
-                        {
-                            "type": "image_url",
-                            "image_url": {
-                                "url": image_url,
-                            },
-                        },
-                    ],
-                },
-            ],
-            model=model_name,
+        messages = describe_image_messages(image_url)
+        await complete_and_check(
+            client,
+            model_name,
+            messages,
+            context=f"completions_with_image url={image_url}",
         )
-        assert chat_completion.choices[0].message.content is not None
-        assert isinstance(chat_completion.choices[0].message.content, str)
-        assert len(chat_completion.choices[0].message.content) > 0
 
 
 @pytest.mark.asyncio
@@ -483,54 +603,33 @@ async def test_completions_with_image_with_uuid(
     image_urls: list[str],
 ):
     for image_url in image_urls:
-        chat_completion = await client.chat.completions.create(
-            messages=[
-                {"role": "system", "content": "You are a helpful assistant."},
-                {
-                    "role": "user",
-                    "content": [
-                        {
-                            "type": "text",
-                            "text": "Describe this image.",
-                        },
-                        {
-                            "type": "image_url",
-                            "image_url": {
-                                "url": image_url,
-                            },
-                            "uuid": image_url,
-                        },
-                    ],
-                },
-            ],
-            model=model_name,
+        messages = describe_image_messages(
+            image_url,
+            extra_image_fields={"uuid": image_url},
+        )
+        await complete_and_check(
+            client,
+            model_name,
+            messages,
+            context=f"uuid first request url={image_url}",
         )
-        assert chat_completion.choices[0].message.content is not None
-        assert isinstance(chat_completion.choices[0].message.content, str)
-        assert len(chat_completion.choices[0].message.content) > 0
 
-        # Second request, with empty image but the same uuid.
-        chat_completion_with_empty_image = await client.chat.completions.create(
-            messages=[
-                {"role": "system", "content": "You are a helpful assistant."},
-                {
-                    "role": "user",
-                    "content": [
-                        {
-                            "type": "text",
-                            "text": "Describe this image.",
-                        },
-                        {"type": "image_url", "image_url": {}, "uuid": image_url},
-                    ],
-                },
-            ],
-            model=model_name,
+        cached_messages: list[dict] = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "Describe this image."},
+                    {"type": "image_url", "image_url": {}, "uuid": image_url},
+                ],
+            },
+        ]
+        await complete_and_check(
+            client,
+            model_name,
+            cached_messages,
+            context=f"uuid cached (empty image) uuid={image_url}",
         )
-        assert chat_completion_with_empty_image.choices[0].message.content is not None
-        assert isinstance(
-            chat_completion_with_empty_image.choices[0].message.content, str
-        )
-        assert len(chat_completion_with_empty_image.choices[0].message.content) > 0
 
 
 @pytest.mark.asyncio
@@ -540,16 +639,13 @@ async def test_completions_with_empty_image_with_uuid_without_cache_hit(
     model_name: str,
 ):
     with pytest.raises(openai.BadRequestError):
-        _ = await client.chat.completions.create(
+        await client.chat.completions.create(
             messages=[
                 {"role": "system", "content": "You are a helpful assistant."},
                 {
                     "role": "user",
                     "content": [
-                        {
-                            "type": "text",
-                            "text": "Describe this image.",
-                        },
+                        {"type": "text", "text": "Describe this image."},
                         {
                             "type": "image_url",
                             "image_url": {},
@@ -575,29 +671,18 @@ async def test_completions_with_image_with_incorrect_uuid_format(
     image_urls: list[str],
 ):
     for image_url in image_urls:
-        chat_completion = await client.chat.completions.create(
-            messages=[
-                {"role": "system", "content": "You are a helpful assistant."},
-                {
-                    "role": "user",
-                    "content": [
-                        {
-                            "type": "text",
-                            "text": "Describe this image.",
-                        },
-                        {
-                            "type": "image_url",
-                            "image_url": {
-                                "url": image_url,
-                                "incorrect_uuid_key": image_url,
-                            },
-                            "also_incorrect_uuid_key": image_url,
-                        },
-                    ],
-                },
-            ],
-            model=model_name,
+        messages = describe_image_messages(
+            image_url,
+            extra_image_fields={
+                "also_incorrect_uuid_key": image_url,
+            },
+        )
+        # Inject the bad key inside image_url dict too
+        messages[1]["content"][1]["image_url"]["incorrect_uuid_key"] = image_url
+
+        await complete_and_check(
+            client,
+            model_name,
+            messages,
+            context=f"incorrect uuid format url={image_url}",
         )
-        assert chat_completion.choices[0].message.content is not None
-        assert isinstance(chat_completion.choices[0].message.content, str)
-        assert len(chat_completion.choices[0].message.content) > 0

tests/entrypoints/pooling/embed/test_online_vision.py

@@ -127,6 +127,39 @@ def test_chat_image_base64_request(server: RemoteOpenAIServer, model_name: str):
     assert output.usage.prompt_tokens == 767
 
 
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+def test_chat_image_with_media_io_kwargs(server: RemoteOpenAIServer, model_name: str):
+    rgba_image_url = (
+        "https://vllm-public-assets.s3.us-west-2.amazonaws.com"
+        "/vision_model_images/RGBA_comp.png"
+    )
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Represent the user's input."},
+                {"type": "image_url", "image_url": {"url": rgba_image_url}},
+            ],
+        }
+    ]
+
+    response = requests.post(
+        server.url_for("v1/embeddings"),
+        json={
+            "model": model_name,
+            "messages": messages,
+            "media_io_kwargs": {
+                "image": {"rgba_background_color": [0, 0, 0]},
+            },
+        },
+    )
+    response.raise_for_status()
+
+    output = EmbeddingResponse.model_validate(response.json())
+    assert len(output.data) == 1
+    assert len(output.data[0].embedding) == 3072
+
+
 def get_hf_prompt_tokens(model_name, content, image_url):
     processor = AutoProcessor.from_pretrained(
         model_name, trust_remote_code=True, num_crops=4

tests/lora/test_olmoe_tp.py

@@ -3,6 +3,7 @@
 
 
 import shutil
+from collections.abc import Sequence
 
 import pytest
 import torch
@@ -15,7 +16,7 @@ from ..utils import multi_gpu_test
 
 MODEL_PATH = "allenai/OLMoE-1B-7B-0125-Instruct"
 
-PROMPT_TEMPLATE = """I want you to act as a SQL terminal in front of an example database, you need only to return the sql command to me.Below is an instruction that describes a task, Write a response that appropriately completes the request.
+PROMPT_TEMPLATE = """I want you to act as a SQL terminal in front of an example database, you need only to return the sql command to me. Do not return any additional explanation. Below is an instruction that describes a task, Write a response that appropriately completes the request.
 "
 ##Instruction:
 candidate_poll contains tables such as candidate, people. Table candidate has columns such as Candidate_ID, People_ID, Poll_Source, Date, Support_rate, Consider_rate, Oppose_rate, Unsure_rate. Candidate_ID is the primary key.
@@ -39,10 +40,20 @@ EXPECTED_BASE_MODEL_OUTPUT = [
     "SELECT COUNT(Candidate_ID) FROM candidate",
     "SELECT COUNT(Candidate_ID) FROM candidate",
     "SELECT Candidate_ID, COUNT(*) as Total_Candidates\nFROM candidate\nINNER JOIN people ON candidate.People_ID = people.People_ID",  # noqa: E501
-    "SELECT Candidate_ID, Poll_Source FROM candidate WHERE People_ID IN (SELECT People_ID FROM people) ORDER BY COUNT(*) DESC LIMIT 1",  # noqa: E501
+    # There are multiple acceptable responses
+    (
+        "SELECT Candidate_ID, Poll_Source FROM candidate WHERE People_ID IN (SELECT People_ID FROM people) ORDER BY COUNT(*) DESC LIMIT 1",  # noqa: E501
+        "SELECT Candidate_ID, Poll_Source FROM candidate WHERE COUNT(People_ID) = (SELECT COUNT(People_ID) FROM people) ORDER BY Candidate_ID DESC LIMIT 1",  # noqa: E501
+    ),
 ]
 
 
+def _output_matches(generated: str, accepted: str | Sequence[str]) -> bool:
+    if isinstance(accepted, str):
+        accepted = (accepted,)
+    return any(generated.startswith(s) for s in accepted)
+
+
 def generate_and_test(
     llm: vllm.LLM,
     lora_path: str,
@@ -90,9 +101,13 @@ def generate_and_test(
 
         if compare_lower:
             generated_text = generated_text.lower()
-            expected_output = expected_output.lower()
-
-        assert generated_text.startswith(expected_output)
+            if isinstance(expected_output, str):
+                expected_output = (expected_output.lower(),)
+            else:
+                expected_output = tuple(s.lower() for s in expected_output)
+        assert _output_matches(generated_text, expected_output), (
+            f"Output {i}: {generated_text!r} does not match any of {expected_output!r}"
+        )
 
 
 def test_olmoe_lora(olmoe_lora_files):

tests/lora/test_worker.py

@@ -64,7 +64,6 @@ def test_worker_apply_lora(qwen3_lora_files):
         device_config=DeviceConfig("cuda"),
         cache_config=CacheConfig(
             block_size=16,
-            swap_space=0,
             cache_dtype="auto",
         ),
         lora_config=LoRAConfig(

tests/models/quantization/test_gpt_oss.py

@@ -21,6 +21,7 @@ import lm_eval
 import pytest
 from packaging import version
 
+from vllm.platforms.rocm import on_gfx950
 from vllm.utils.torch_utils import cuda_device_count_stateless
 
 MODEL_ACCURACIES = {
@@ -83,11 +84,17 @@ class EvaluationConfig:
 @pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
 @pytest.mark.parametrize("model_name, expected_accuracy", MODEL_ACCURACIES.items())
 def test_gpt_oss_attention_quantization(
-    model_name: str, tp_size: int, expected_accuracy: float
+    model_name: str,
+    tp_size: int,
+    expected_accuracy: float,
+    monkeypatch: pytest.MonkeyPatch,
 ):
     if tp_size > cuda_device_count_stateless():
         pytest.skip("Not enough GPUs to run this test case")
 
+    if "amd/gpt-oss-20b-MoE-Quant-W-MXFP4-A-FP8-KV-FP8" in model_name and on_gfx950():
+        monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
+
     model_args = EvaluationConfig(model_name).get_model_args(tp_size)
 
     extra_run_kwargs = {

tests/models/registry.py

@@ -480,6 +480,18 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
         min_transformers_version="4.56.3",
     ),
     "RWForCausalLM": _HfExamplesInfo("tiiuae/falcon-40b"),
+    "SarvamMoEForCausalLM": _HfExamplesInfo(
+        "sarvamai/sarvam-30b",
+        trust_remote_code=True,
+        max_model_len=4096,
+        is_available_online=True,
+    ),
+    "SarvamMLAForCausalLM": _HfExamplesInfo(
+        "sarvamai/sarvam-105b",
+        trust_remote_code=True,
+        max_model_len=4096,
+        is_available_online=True,
+    ),
     "SeedOssForCausalLM": _HfExamplesInfo(
         "ByteDance-Seed/Seed-OSS-36B-Instruct",
         trust_remote_code=True,

tests/v1/attention/test_mla_backends.py

@@ -327,6 +327,12 @@ class MockSparseMLAAttentionLayer:
         self._k_scale_float = 1.0
         self._v_scale_float = 1.0
 
+        self._decode_concat_quant_fp8_op = _DecodeConcatQuantFP8(
+            static=True,
+            group_shape=GroupShape.PER_TENSOR,
+            compile_native=True,
+        )
+
     def forward_impl(
         self,
         q: torch.Tensor,
@@ -338,6 +344,7 @@ class MockSparseMLAAttentionLayer:
     ) -> torch.Tensor:
         """Forward for sparse MLA - uses forward_mqa for all tokens."""
         kv_cache_dtype = getattr(self.impl, "kv_cache_dtype", "auto")
+        fp8_attention = kv_cache_dtype.startswith("fp8")
 
         # Write to KV cache
         if kv_cache.numel() > 0:
@@ -350,6 +357,9 @@ class MockSparseMLAAttentionLayer:
                 scale=self._k_scale,
             )
 
+        if fp8_attention and kv_cache_dtype != "fp8_ds_mla":
+            kv_cache = kv_cache.view(current_platform.fp8_dtype())
+
         num_tokens = q.shape[0]
 
         # Sparse MLA uses forward_mqa for all tokens
@@ -367,8 +377,14 @@ class MockSparseMLAAttentionLayer:
         # Convert from (N, B, L) to (B, N, L)
         mqa_ql_nope = mqa_ql_nope.transpose(0, 1)
 
-        # Pass as tuple to forward_mqa
-        mqa_q = (mqa_ql_nope, mqa_q_pe)
+        if fp8_attention and self.impl.supports_quant_query_input:
+            assert mqa_ql_nope.shape[0] == mqa_q_pe.shape[0]
+            assert mqa_ql_nope.shape[1] == mqa_q_pe.shape[1]
+            mqa_q = self._decode_concat_quant_fp8_op(
+                mqa_ql_nope, mqa_q_pe, self._q_scale
+            )
+        else:
+            mqa_q = (mqa_ql_nope, mqa_q_pe)
 
         attn_out, _ = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self)
 

tests/v1/attention/test_sparse_mla_backends.py

@@ -191,6 +191,16 @@ def test_sparse_backend_decode_correctness(
     if kv_cache_dtype not in backend_cls.supported_kv_cache_dtypes:
         pytest.skip(f"{backend_cls.get_name()} does not support {kv_cache_dtype}")
 
+    if (
+        backend_cls == FlashMLASparseBackend
+        and kv_cache_dtype.startswith("fp8")
+        and kv_cache_dtype != "fp8_ds_mla"
+    ):
+        pytest.skip(
+            "FlashMLA Sparse Attention backend fp8 only supports "
+            "fp8_ds_mla kv-cache dtype"
+        )
+
     supported_block_sizes = backend_cls.get_supported_kernel_block_sizes()
     if block_size not in supported_block_sizes:
         pytest.skip(
@@ -419,7 +429,7 @@ def test_sparse_backend_decode_correctness(
         num_blocks=vllm_config.cache_config.num_gpu_blocks,
         common_attn_metadata=common_attn_metadata,
         randomize_blocks=False,
-        kv_cache_dtype=kv_cache_dtype if use_fp8_ds_mla_quantization else "auto",
+        kv_cache_dtype=kv_cache_dtype,
         scale=kv_cache_scale,
     )
 

tests/v1/attention/utils.py

@@ -182,7 +182,6 @@ def create_vllm_config(
     cache_config = CacheConfig(
         block_size=block_size,
         cache_dtype="auto",
-        swap_space=0,
     )
     # Set cache blocks for testing
     #   (these may be set during initialization normally)

tests/v1/core/test_scheduler.py

@@ -1776,7 +1776,6 @@ def create_scheduler_with_priority(
     cache_config = CacheConfig(
         block_size=block_size,
         gpu_memory_utilization=0.9,
-        swap_space=0,
         cache_dtype="auto",
         enable_prefix_caching=enable_prefix_caching,
     )
@@ -3726,7 +3725,6 @@ def _create_encoder_decoder_scheduler(
     cache_config = CacheConfig(
         block_size=block_size,
         gpu_memory_utilization=0.9,
-        swap_space=0,
         cache_dtype="auto",
         enable_prefix_caching=False,
     )

tests/v1/core/utils.py

@@ -94,7 +94,6 @@ def create_scheduler(
     cache_config = CacheConfig(
         block_size=block_size,
         gpu_memory_utilization=0.9,
-        swap_space=0,
         cache_dtype="auto",
         enable_prefix_caching=enable_prefix_caching,
     )

tests/v1/e2e/test_async_scheduling.py

@@ -98,7 +98,7 @@ def test_without_spec_decoding(
 
 @single_gpu_only
 @large_gpu_mark(min_gb=16)
-def test_with_spec_decoding(sample_json_schema, monkeypatch: pytest.MonkeyPatch):
+def test_with_eagle3_spec_decoding(sample_json_schema, monkeypatch: pytest.MonkeyPatch):
     """Test consistency and acceptance rates with some different combos of
     preemption, executor, async scheduling, prefill chunking,
     spec decoding model length.
@@ -154,6 +154,42 @@ def test_with_spec_decoding(sample_json_schema, monkeypatch: pytest.MonkeyPatch)
     )
 
 
+def test_with_ngram_gpu_spec_decoding(monkeypatch: pytest.MonkeyPatch):
+    """Test ngram_gpu speculative decoding with different configurations.
+
+    This test specifically validates ngram_gpu behavior with various:
+    - Number of speculative tokens (2-6)
+    - Prompt lookup window sizes (min/max)
+    - Async scheduling enabled (as in production)
+    - Different executors and chunking settings
+    """
+
+    # Variant with larger speculation window
+    ngram_gpu_config = {
+        "method": "ngram_gpu",
+        "num_speculative_tokens": 3,
+        "prompt_lookup_max": 3,
+        "prompt_lookup_min": 2,
+    }
+
+    # Test configurations covering various scenarios
+    # test_preemption, executor, async_scheduling,
+    # spec_config, test_prefill_chunking
+    test_configs = [
+        (False, "mp", False, None, False),
+        (False, "mp", False, ngram_gpu_config, False),
+        (True, "mp", False, ngram_gpu_config, True),
+        (False, "mp", True, ngram_gpu_config, False),
+        (True, "mp", True, ngram_gpu_config, False),
+        (True, "uni", True, ngram_gpu_config, False),
+        (True, "mp", True, ngram_gpu_config, True),
+    ]
+
+    # Use MODEL (Qwen) for ngram_gpu tests as it's lighter weight
+    # and ngram_gpu doesn't require a specific draft model
+    run_tests(monkeypatch, MODEL, test_configs, [{}])
+
+
 @dynamo_config.patch(cache_size_limit=16)
 def run_tests(
     monkeypatch: pytest.MonkeyPatch,
@@ -282,11 +318,12 @@ def run_test(
         else dict(gpu_memory_utilization=0.9)
     )
     spec_mml = (spec_config or {}).get("max_model_len")
+    spec_method = (spec_config or {}).get("method", "none")
     test_config = (
         f"executor={executor}, preemption={test_preemption}, "
         f"async_sched={async_scheduling}, "
         f"chunk_prefill={test_prefill_chunking}, "
-        f"spec_decoding={spec_decoding}, spec_mml={spec_mml}"
+        f"spec_decoding={spec_decoding}, spec_method={spec_method}, spec_mml={spec_mml}"
     )
     print("-" * 80)
     print(f"---- TESTING {test_str}: {test_config}")
@@ -294,7 +331,7 @@ def run_test(
 
     with VllmRunner(
         model,
-        max_model_len=512,
+        max_model_len=4096,
         enable_chunked_prefill=test_prefill_chunking,
         # Force prefill chunking
         max_num_batched_tokens=48 if test_prefill_chunking else None,

tests/v1/e2e/test_spec_decode.py

@@ -183,6 +183,34 @@ def test_ngram_and_suffix_correctness(
     cleanup_dist_env_and_memory()
 
 
+@pytest.mark.parametrize("async_scheduling", [True], ids=["async"])
+@single_gpu_only
+@large_gpu_mark(min_gb=20)
+def test_ngram_gpu_default_with_async_scheduling(
+    async_scheduling: bool,
+):
+    """
+    Test ngram_gpu speculative decoding (k=3) correctness with and without
+    async scheduling, validated via GSM8K accuracy.
+    Uses Qwen/Qwen3-8B (ref GSM8K accuracy: 87%-92%).
+    """
+    qwen3_model = "Qwen/Qwen3-8B"
+    spec_llm = LLM(
+        model=qwen3_model,
+        speculative_config={
+            "method": "ngram_gpu",
+            "prompt_lookup_max": 3,
+            "prompt_lookup_min": 2,
+            "num_speculative_tokens": 2,
+        },
+        max_model_len=4096,
+        async_scheduling=async_scheduling,
+    )
+    evaluate_llm_for_gsm8k(spec_llm, expected_accuracy_threshold=0.8)
+    del spec_llm
+    cleanup_dist_env_and_memory()
+
+
 @single_gpu_only
 @large_gpu_mark(min_gb=20)
 def test_suffix_decoding_acceptance(

tests/v1/engine/test_engine_core.py

@@ -506,7 +506,6 @@ def test_encoder_instance_zero_kv_cache(
     cache_config = CacheConfig(
         block_size=16,
         gpu_memory_utilization=gpu_memory_utilization,
-        swap_space=0,
         cache_dtype="auto",
         enable_prefix_caching=enable_prefix_caching,
     )

tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

@@ -56,24 +56,27 @@ run_tests() {
   echo "✅ All ${label} tests passed!"
 }
 
-# Run tests
+# Set backend
+label="default backend"
+cmdline_args=""
 if [[ -n "${ROCM_ATTN:-}" ]]; then
   echo "ROCM_ATTN is set, running with --attention-backend ROCM_ATTN"
-  run_tests "ROCM_ATTN backend" "--attention-backend ROCM_ATTN"
+  label="ROCM_ATTN backend"
+  cmdline_args=" --attention-backend ROCM_ATTN "
+elif [[ -n "${FLASHINFER:-}" ]]; then
+  echo "FLASHINFER is set, running with --attention-backend FLASHINFER"
+  label="FLASHINFER backend"
+  cmdline_args=" --attention-backend FLASHINFER "
 else
-  run_tests "default backend" ""
-fi
-
-# Check if FLASHINFER is set (non-empty)
-if [[ -n "${FLASHINFER:-}" ]]; then
-  echo "FLASHINFER is set, rerunning with --attention-backend FLASHINFER"
-  run_tests "FLASHINFER backend" "--attention-backend FLASHINFER"
-else
-  echo "FLASHINFER not set, skipping FLASHINFER runs."
+  echo "running with default attention backend"
 fi
 
 # Check if cross-layers is enabled (non-empty)
 if [[ -n "${CROSS_LAYERS_BLOCKS:-}" ]]; then
-  echo "CROSS_LAYERS_BLOCKS is set, rerunning with --enable-cross-layers"
-  run_tests "default backend" "--enable-cross-layers"
+  echo "CROSS_LAYERS_BLOCKS is set, running with --enable-cross-layers"
+  label+=" - CROSS_LAYERS_BLOCKS enabled"
+  cmdline_args+=" --enable-cross-layers "
 fi
+
+# Run tests
+run_tests "${label}" "${cmdline_args}"

tests/v1/kv_connector/unit/test_moriio_connector.py

@@ -206,7 +206,6 @@ def create_vllm_config(
     cache_config = CacheConfig(
         block_size=block_size,
         gpu_memory_utilization=0.9,
-        swap_space=0,
         cache_dtype="auto",
         enable_prefix_caching=True,
     )

tests/v1/kv_connector/unit/utils.py

@@ -118,7 +118,6 @@ def create_vllm_config(
     cache_config = CacheConfig(
         block_size=block_size,
         gpu_memory_utilization=0.9,
-        swap_space=0,
         cache_dtype=cache_dtype,
         enable_prefix_caching=True,
     )

tests/v1/worker/test_gpu_model_runner.py

@@ -96,7 +96,6 @@ def get_vllm_config():
     cache_config = CacheConfig(
         block_size=BLOCK_SIZE,
         gpu_memory_utilization=0.9,
-        swap_space=0,
         cache_dtype="auto",
     )
     parallel_config = ParallelConfig()
@@ -809,7 +808,6 @@ def test_hybrid_attention_mamba_tensor_shapes():
     cache_config = CacheConfig(
         block_size=BLOCK_SIZE,
         gpu_memory_utilization=0.9,
-        swap_space=0,
         cache_dtype="auto",
     )
     parallel_config = ParallelConfig()
@@ -1242,7 +1240,6 @@ def test_cudagraph_sizes_capped_for_mamba_cache():
     cache_config = CacheConfig(
         block_size=BLOCK_SIZE,
         gpu_memory_utilization=0.9,
-        swap_space=0,
         cache_dtype="auto",
     )
     parallel_config = ParallelConfig()

tools/pre_commit/generate_attention_backend_docs.py

@@ -49,6 +49,11 @@ MLA_ATTENTION_FILE = (
 # Backends to skip during doc generation
 SKIP_BACKENDS = {"CUSTOM", "TORCH_SDPA"}
 
+BACKEND_KV_DTYPE_EXCLUDES: dict[str, set[str]] = {
+    # fp8 is an alias for fp8_ds_mla for FlashMLA Sparse
+    "FLASHMLA_SPARSE": {"fp8"},
+}
+
 
 def is_relevant_file(filepath: str) -> bool:
     """Check if a file matches any of the relevant patterns."""
@@ -546,10 +551,19 @@ def analyze_backend(backend_name: str, class_path: str) -> dict[str, Any] | None
             tree, impl_class_name, "can_return_lse_for_decode", False, file_path
         )
 
+    kv_cache_dtypes = parse_kv_cache_dtypes(class_node)
+    if backend_name in BACKEND_KV_DTYPE_EXCLUDES:
+        excluded = BACKEND_KV_DTYPE_EXCLUDES[backend_name]
+        kv_cache_dtypes = ", ".join(
+            d
+            for d in (d.strip() for d in kv_cache_dtypes.split(","))
+            if d not in excluded
+        )
+
     return {
         "name": backend_name,
         "dtypes": parse_supported_dtypes(class_node),
-        "kv_cache_dtypes": parse_kv_cache_dtypes(class_node),
+        "kv_cache_dtypes": kv_cache_dtypes,
         "block_sizes": parse_block_sizes(class_node),
         "head_sizes": parse_head_sizes(class_node),
         "attn_types": parse_attention_types(class_node),

vllm/benchmarks/throughput.py

@@ -38,6 +38,7 @@ from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
 from vllm.inputs import TextPrompt, TokensPrompt
 from vllm.lora.request import LoRARequest
 from vllm.outputs import RequestOutput
+from vllm.platforms import current_platform
 from vllm.sampling_params import BeamSearchParams
 from vllm.tokenizers import TokenizerLike, get_tokenizer
 from vllm.utils.async_utils import merge_async_iterators
@@ -256,17 +257,21 @@ def run_hf(
     max_batch_size: int,
     trust_remote_code: bool,
     disable_detokenize: bool = False,
+    dtype: torch.dtype | None = torch.float16,
+    enable_torch_compile: bool = False,
 ) -> float:
     assert isinstance(tokenizer, PreTrainedTokenizerBase), (
         "the hf backend only supports HF tokenizers"
     )
     llm = AutoModelForCausalLM.from_pretrained(
-        model, dtype=torch.float16, trust_remote_code=trust_remote_code
+        model, dtype=dtype, trust_remote_code=trust_remote_code
     )
     if llm.config.model_type == "llama":
         # To enable padding in the HF backend.
         tokenizer.pad_token = tokenizer.eos_token
-    llm = llm.cuda()
+    llm = llm.to(current_platform.device_type)
+    if enable_torch_compile:
+        llm = torch.compile(llm)
 
     pbar = tqdm(total=len(requests))
     start = time.perf_counter()
@@ -295,7 +300,7 @@ def run_hf(
         # Generate the sequences.
         input_ids = tokenizer(batch, return_tensors="pt", padding=True).input_ids
         llm_outputs = llm.generate(
-            input_ids=input_ids.cuda(),
+            input_ids=input_ids.to(current_platform.device_type),
             do_sample=True,
             num_return_sequences=n,
             temperature=1.0,
@@ -733,6 +738,12 @@ def add_cli_args(parser: argparse.ArgumentParser):
         default=None,
         help="Maximum batch size for HF backend.",
     )
+    parser.add_argument(
+        "--hf-enable-torch-compile",
+        action="store_true",
+        default=False,
+        help="Enable Torch compile for HF backend.",
+    )
     parser.add_argument(
         "--output-json",
         type=str,
@@ -884,6 +895,8 @@ def main(args: argparse.Namespace):
             args.hf_max_batch_size,
             args.trust_remote_code,
             args.disable_detokenize,
+            dtype=args.dtype,
+            enable_torch_compile=args.hf_enable_torch_compile,
         )
     elif args.backend == "vllm-chat":
         elapsed_time, request_outputs = run_vllm_chat(

vllm/compilation/backends.py

@@ -907,6 +907,13 @@ class VllmBackend:
         # Honors opt-outs such as CompilationMode.NONE or VLLM_DISABLE_COMPILE_CACHE.
         disable_cache = not is_compile_cache_enabled(self.inductor_config)
 
+        # TODO(patchy): ngram gpu kernel will cause vllm torch compile cache errors.
+        is_ngram_gpu_enabled = (
+            vllm_config.speculative_config is not None
+            and vllm_config.speculative_config.use_ngram_gpu()
+        )
+        disable_cache = disable_cache or is_ngram_gpu_enabled
+
         if disable_cache:
             logger.info_once("vLLM's torch.compile cache is disabled.", scope="local")
         else:

vllm/compilation/caching.py

@@ -189,13 +189,13 @@ class VllmSerializableFunction(SerializableCallable):  # type: ignore[misc]
         self.shape_env = None
         self.vllm_backend = vllm_backend
         self.sym_tensor_indices = sym_tensor_indices
+        self._fake_mode: Any | None = None
 
         import torch._functorch.config as functorch_config
 
         self.aot_autograd_config = (
             aot_autograd_config or functorch_config.save_config_portable()
         )
-
         sym_input = next(
             (i for i in self.example_inputs if isinstance(i, torch.SymInt)), None
         )
@@ -217,6 +217,7 @@ class VllmSerializableFunction(SerializableCallable):  # type: ignore[misc]
         state.pop("optimized_call")
         state.pop("shape_env")
         state.pop("vllm_backend", None)
+        state.pop("_fake_mode", None)
         for node in state["graph_module"].graph.nodes:
             node.meta.pop("source_fn_stack", None)
             node.meta.pop("nn_module_stack", None)
@@ -351,8 +352,31 @@ class VllmSerializableFunction(SerializableCallable):  # type: ignore[misc]
             return fn.optimized_call(*example_inputs)
 
         fn = cls(**state, optimized_call=optimized_call)
+        fn._fake_mode = fake_mode
         return fn
 
+    def finalize_loading(self, vllm_config: VllmConfig) -> None:
+        """Eagerly initialize the compiled backend and perform all loading.
+
+        Must be called after _verify_source_unchanged has populated
+        compilation_config.traced_files, which is needed for cache dir
+        computation.
+        """
+        if self._fake_mode is None:
+            return  # Already finalized, or mega path (no _fake_mode set)
+
+        from torch._guards import TracingContext, tracing
+
+        from vllm.compilation.backends import VllmBackend
+
+        vllm_backend = VllmBackend(vllm_config, self.prefix, self.is_encoder)
+        with tracing(TracingContext(self._fake_mode)):
+            result = vllm_backend(self.graph_module, list(self.example_inputs))
+            self.optimized_call = result.optimized_call
+            self.vllm_backend = vllm_backend
+
+        self._fake_mode = None
+
     @property
     def co_name(self) -> Literal["VllmSerializableFunction"]:
         """

vllm/compilation/cuda_graph.py

@@ -2,10 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import dataclasses
+import weakref
 from collections import Counter
 from collections.abc import Callable
 from contextlib import ExitStack
-from typing import Any
+from typing import Any, ClassVar
 from unittest.mock import patch
 
 import torch
@@ -162,6 +163,14 @@ class CUDAGraphWrapper:
     guaranteed when VLLM_LOGGING_LEVEL == "DEBUG".
     """
 
+    _all_instances: ClassVar[weakref.WeakSet["CUDAGraphWrapper"]] = weakref.WeakSet()
+
+    @classmethod
+    def clear_all_graphs(cls) -> None:
+        """Clear captured graphs from all CUDAGraphWrapper instances."""
+        for instance in list(cls._all_instances):
+            instance.clear_graphs()
+
     def __init__(
         self,
         runnable: Callable[..., Any],
@@ -192,6 +201,8 @@ class CUDAGraphWrapper:
         # cudagraphs for.
         self.concrete_cudagraph_entries: dict[BatchDescriptor, CUDAGraphEntry] = {}
 
+        CUDAGraphWrapper._all_instances.add(self)
+
     def __getattr__(self, key: str) -> Any:
         # allow accessing the attributes of the runnable.
         if hasattr(self.runnable, key):
@@ -205,6 +216,13 @@ class CUDAGraphWrapper:
         # in case we need to access the original runnable.
         return self.runnable
 
+    @property
+    def cudagraph_wrapper(self) -> "CUDAGraphWrapper":
+        return self
+
+    def clear_graphs(self) -> None:
+        self.concrete_cudagraph_entries.clear()
+
     def __call__(self, *args: Any, **kwargs: Any) -> Any | None:
         forward_context = get_forward_context()
         batch_descriptor = forward_context.batch_descriptor

vllm/compilation/decorators.py

@@ -30,7 +30,7 @@ from vllm.sequence import IntermediateTensors
 from vllm.utils.import_utils import resolve_obj_by_qualname
 from vllm.utils.torch_utils import is_torch_equal_or_newer
 
-from .monitor import start_monitoring_torch_compile
+from .monitor import monitor_profiling_run, monitor_torch_compile
 
 if TYPE_CHECKING:
     # Only added on nightly/2.10 so wrap
@@ -434,17 +434,24 @@ def _support_torch_compile(
             cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}")
             aot_compilation_path = os.path.join(cache_dir, "model")
             try:
-                with (
-                    set_current_vllm_config(self.vllm_config),
-                    open(aot_compilation_path, "rb") as f,
-                ):
-                    start_monitoring_torch_compile(self.vllm_config)
-                    loaded_fn = torch.compiler.load_compiled_function(
-                        f, f_globals=self.forward.__globals__
-                    )
-                _verify_source_unchanged(loaded_fn.source_info(), self.vllm_config)
-                if not self.compilation_config.dynamic_shapes_config.evaluate_guards:
-                    loaded_fn.disable_guard_check()
+                with monitor_torch_compile(self.vllm_config):
+                    with (
+                        set_current_vllm_config(self.vllm_config),
+                        open(aot_compilation_path, "rb") as f,
+                    ):
+                        loaded_fn = torch.compiler.load_compiled_function(
+                            f, f_globals=self.forward.__globals__
+                        )
+                    _verify_source_unchanged(loaded_fn.source_info(), self.vllm_config)
+                    ds_config = self.compilation_config.dynamic_shapes_config
+                    if not ds_config.evaluate_guards:
+                        loaded_fn.disable_guard_check()
+                    # Eagerly load compiled artifacts now that traced_files
+                    # is populated by _verify_source_unchanged.
+                    with maybe_use_cudagraph_partition_wrapper(self.vllm_config):
+                        loaded_fn._artifacts.compiled_fn.finalize_loading(
+                            self.vllm_config
+                        )
                 self.aot_compiled_fn = loaded_fn
                 self.was_aot_compile_fn_loaded_from_disk = True
             except Exception as e:
@@ -465,12 +472,11 @@ def _support_torch_compile(
                 logger.info(
                     "Directly load AOT compilation from path %s", aot_compilation_path
                 )
-                # Apply partition wrapper context for proper CUDA graph capture
-                from .monitor import end_monitoring_torch_compile
-
-                with maybe_use_cudagraph_partition_wrapper(self.vllm_config):
+                with (
+                    monitor_profiling_run(),
+                    maybe_use_cudagraph_partition_wrapper(self.vllm_config),
+                ):
                     output = self.aot_compiled_fn(self, *args, **kwargs)
-                end_monitoring_torch_compile(self.vllm_config)
                 return output
 
         if self.compiled:
@@ -489,8 +495,6 @@ def _support_torch_compile(
             **kwargs,
         )
 
-        # here, it is the starting point of the `torch.compile` process
-        start_monitoring_torch_compile(self.vllm_config)
         original_code_object = self.original_code_object()
         logger.debug("Start compiling function %s", original_code_object)
 
@@ -559,16 +563,26 @@ def _support_torch_compile(
                 # store the path for saving after warmup
                 self._aot_compilation_path = aot_compilation_path
                 self._aot_cache_dir = cache_dir
-                self.aot_compiled_fn = self.aot_compile(*args, **kwargs)
-                # All compilation is done at this point, save the AOT artifact.
-                self.save_aot_compiled_function()
-                output = self.aot_compiled_fn(self, *args, **kwargs)
+                with monitor_torch_compile(self.vllm_config):
+                    self.aot_compiled_fn = self.aot_compile(*args, **kwargs)
+                    # All compilation is done at this point, save the
+                    # AOT artifact.
+                    self.save_aot_compiled_function()
+
+                with monitor_profiling_run():
+                    output = self.aot_compiled_fn(self, *args, **kwargs)
             else:
-                output = TorchCompileWithNoGuardsWrapper.__call__(self, *args, **kwargs)  # type: ignore[arg-type]
+                with monitor_torch_compile(
+                    self.vllm_config,
+                    "torch.compile and initial profiling/warmup "
+                    "run together took %.2f s in total",
+                ):
+                    output = TorchCompileWithNoGuardsWrapper.__call__(
+                        self,  # type: ignore[arg-type]
+                        *args,
+                        **kwargs,
+                    )
 
-        from .monitor import end_monitoring_torch_compile
-
-        end_monitoring_torch_compile(self.vllm_config)
         self.compiled = True
         return output
 

vllm/compilation/monitor.py

@@ -1,46 +1,83 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import contextlib
 import time
+from collections.abc import Generator
 
-from vllm.config import CompilationConfig, CompilationMode, VllmConfig
+from vllm.config import CompilationMode, VllmConfig
 from vllm.logger import init_logger
 
 logger = init_logger(__name__)
 
-context_manager = None
+# Shared global so backends.py can read the start time for Dynamo timing.
 torch_compile_start_time: float = 0.0
 
 
-def start_monitoring_torch_compile(vllm_config: VllmConfig) -> None:
+@contextlib.contextmanager
+def monitor_torch_compile(
+    vllm_config: VllmConfig,
+    message: str = "torch.compile took %.2f s in total",
+) -> Generator[None, None, None]:
+    """Context manager that times torch.compile and manages depyf debugging.
+
+    On normal exit: logs the compile time and exits depyf.
+    On exception: cleans up depyf without logging (compilation failed).
+    """
     global torch_compile_start_time
     torch_compile_start_time = time.perf_counter()
 
-    compilation_config: CompilationConfig = vllm_config.compilation_config
+    compilation_config = vllm_config.compilation_config
+    depyf_cm = None
     path = vllm_config.compile_debug_dump_path()
     if compilation_config.mode == CompilationMode.VLLM_COMPILE and path:
         import depyf
 
         path.mkdir(parents=True, exist_ok=True)
         logger.debug("Dumping depyf output to %s", path)
-        global context_manager
-        context_manager = depyf.prepare_debug(path.as_posix())
-        context_manager.__enter__()
+        depyf_cm = depyf.prepare_debug(path.as_posix())
+        depyf_cm.__enter__()
+
+    try:
+        yield
+    except Exception:
+        raise
+    else:
+        total_compile_time = time.perf_counter() - torch_compile_start_time
+        if compilation_config.mode == CompilationMode.VLLM_COMPILE:
+            logger.info_once(message, total_compile_time, scope="local")
+    finally:
+        if depyf_cm is not None:
+            try:
+                depyf_cm.__exit__(None, None, None)
+            except Exception:
+                logger.warning("Exception during depyf cleanup.", exc_info=True)
 
 
-def end_monitoring_torch_compile(vllm_config: VllmConfig) -> None:
-    compilation_config: CompilationConfig = vllm_config.compilation_config
-    total_compile_time: float = time.perf_counter() - torch_compile_start_time
-    if compilation_config.mode == CompilationMode.VLLM_COMPILE:
-        logger.info_once(
-            "torch.compile and initial profiling run took %.2f s in total",
-            total_compile_time,
-            scope="local",
-        )
-        global context_manager
-        if context_manager is not None:
-            context_manager.__exit__(None, None, None)
-            context_manager = None
+@contextlib.contextmanager
+def monitor_profiling_run() -> Generator[None, None, None]:
+    """Context manager that times the initial profiling run.
+
+    Asserts that no backend compilation occurs during the profiling run
+    (all compilation should have completed before this point).
+    """
+    from vllm.compilation.counter import compilation_counter
+
+    backend_compilations_before = compilation_counter.num_backend_compilations
+    start = time.perf_counter()
+    yield
+    elapsed = time.perf_counter() - start
+    assert (
+        compilation_counter.num_backend_compilations == backend_compilations_before
+    ), (
+        "backend compilation occurred during the initial profiling run; "
+        "all compilation should be complete before the profiling run starts."
+    )
+    logger.info_once(
+        "Initial profiling/warmup run took %.2f s",
+        elapsed,
+        scope="local",
+    )
 
 
 cudagraph_capturing_enabled: bool = True

vllm/config/cache.py

@@ -1,21 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import math
 from dataclasses import field
-from typing import TYPE_CHECKING, Any, Literal
+from typing import Literal
 
 from pydantic import Field, SkipValidation, field_validator
 
 from vllm.config.utils import config
 from vllm.logger import init_logger
-from vllm.utils.mem_constants import GiB_bytes
-from vllm.utils.mem_utils import format_gib, get_cpu_memory
-
-if TYPE_CHECKING:
-    from vllm.config.parallel import ParallelConfig
-else:
-    ParallelConfig = Any
 
 logger = init_logger(__name__)
 
@@ -53,8 +45,6 @@ class CacheConfig:
     not matter if you have another vLLM instance running on the same GPU. For
     example, if you have two vLLM instances running on the same GPU, you can
     set the GPU memory utilization to 0.5 for each instance."""
-    swap_space: float = Field(default=4, ge=0)
-    """Size of the CPU swap space per GPU (in GiB)."""
     cache_dtype: CacheDType = "auto"
     """Data type for kv cache storage. If "auto", will use model data type.
     CUDA 11.8+ supports fp8 (=fp8_e4m3) and fp8_e5m2. ROCm (AMD GPU) supports
@@ -173,7 +163,6 @@ class CacheConfig:
         ignored_factors = {
             # Runtime/derived knobs that don't affect compiled graph shape
             "gpu_memory_utilization",
-            "swap_space",
             "is_attention_free",
             "num_gpu_blocks_override",
             "enable_prefix_caching",
@@ -208,24 +197,3 @@ class CacheConfig:
                 "scaling factor."
             )
         return cache_dtype
-
-    def verify_with_parallel_config(
-        self,
-        parallel_config: ParallelConfig,
-    ) -> None:
-        swap_space_bytes = math.ceil(self.swap_space * GiB_bytes)
-        total_cpu_memory = get_cpu_memory()
-        # FIXME(woosuk): Here, it is assumed that the GPUs in a tensor parallel
-        # group are in the same node. However, the GPUs may span multiple nodes.
-        num_gpus_per_node = parallel_config.tensor_parallel_size
-        cpu_memory_usage = swap_space_bytes * num_gpus_per_node
-
-        msg = (
-            f"{format_gib(cpu_memory_usage)} GiB out of the "
-            f"{format_gib(total_cpu_memory)} GiB total CPU memory "
-            "is allocated for the swap space."
-        )
-        if cpu_memory_usage > 0.7 * total_cpu_memory:
-            raise ValueError("Too large swap space. " + msg)
-        elif cpu_memory_usage > 0.4 * total_cpu_memory:
-            logger.warning("Possibly too large swap space. %s", msg)

vllm/config/speculative.py

@@ -47,6 +47,7 @@ MTPModelTypes = Literal[
     "step3p5_mtp",
 ]
 EagleModelTypes = Literal["eagle", "eagle3", "extract_hidden_states", MTPModelTypes]
+NgramGPUTypes = Literal["ngram_gpu"]
 SpeculativeMethod = Literal[
     "ngram",
     "medusa",
@@ -54,6 +55,7 @@ SpeculativeMethod = Literal[
     "draft_model",
     "suffix",
     EagleModelTypes,
+    NgramGPUTypes,
 ]
 
 
@@ -364,6 +366,8 @@ class SpeculativeConfig:
                     self.quantization = self.target_model_config.quantization
             elif self.method in ("ngram", "[ngram]"):
                 self.model = "ngram"
+            elif self.method == "ngram_gpu":
+                self.model = "ngram_gpu"
             elif self.method == "suffix":
                 self.model = "suffix"
             elif self.method == "extract_hidden_states":
@@ -374,8 +378,9 @@ class SpeculativeConfig:
                 )
 
         if self.method in ("ngram", "[ngram]"):
-            # Unified to "ngram" internally
             self.method = "ngram"
+
+        if self.method in ("ngram", "ngram_gpu"):
             # Set default values if not provided
             if self.prompt_lookup_min is None and self.prompt_lookup_max is None:
                 # TODO(woosuk): Tune these values. They are arbitrarily chosen.
@@ -832,6 +837,9 @@ class SpeculativeConfig:
     def uses_extract_hidden_states(self) -> bool:
         return self.method == "extract_hidden_states"
 
+    def use_ngram_gpu(self) -> bool:
+        return self.method == "ngram_gpu"
+
     def __repr__(self) -> str:
         method = self.method
         model = (

vllm/config/vllm.py

@@ -41,7 +41,7 @@ from .offload import OffloadConfig
 from .parallel import ParallelConfig
 from .profiler import ProfilerConfig
 from .scheduler import SchedulerConfig
-from .speculative import EagleModelTypes, SpeculativeConfig
+from .speculative import EagleModelTypes, NgramGPUTypes, SpeculativeConfig
 from .structured_outputs import StructuredOutputsConfig
 from .utils import SupportsHash, config, replace
 from .weight_transfer import WeightTransferConfig
@@ -674,8 +674,6 @@ class VllmConfig:
 
             self.parallel_config.is_moe_model = self.model_config.is_moe
 
-        self.cache_config.verify_with_parallel_config(self.parallel_config)
-
         if self.lora_config is not None:
             self.lora_config.verify_with_model_config(self.model_config)
 
@@ -698,11 +696,13 @@ class VllmConfig:
             if self.speculative_config is not None:
                 if (
                     self.speculative_config.method not in get_args(EagleModelTypes)
+                    and self.speculative_config.method not in get_args(NgramGPUTypes)
                     and self.speculative_config.method != "draft_model"
                 ):
                     raise ValueError(
                         "Currently, async scheduling is only supported "
-                        "with EAGLE/MTP/Draft Model kind of speculative decoding."
+                        "with EAGLE/MTP/Draft Model/NGram GPU kind of "
+                        "speculative decoding"
                     )
                 if self.speculative_config.disable_padded_drafter_batch:
                     raise ValueError(
@@ -720,6 +720,7 @@ class VllmConfig:
             if (
                 self.speculative_config is not None
                 and self.speculative_config.method not in get_args(EagleModelTypes)
+                and self.speculative_config.method not in get_args(NgramGPUTypes)
             ):
                 logger.warning_once(
                     "Async scheduling not supported with %s-based "

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py

@@ -336,11 +336,21 @@ class LMCacheMPRequestMetadata:
             start_token_idx = start * vllm_block_size
             end_token_idx = end * vllm_block_size
             token_ids = list(tracker.all_token_ids)
+
+            # Compute how many tokens at the start of the retrieve range
+            # overlap with APC-shared blocks. The server must skip writing
+            # to these positions to avoid a cross-stream data race: the
+            # retrieve writes on the LMCache CUDA stream while concurrent
+            # requests may read these APC-shared blocks on the vLLM stream.
+            apc_overlap_blocks = tracker.num_vllm_hit_blocks - start
+            skip_first_n_tokens = apc_overlap_blocks * vllm_block_size
+
             op = LoadStoreOp(
                 token_ids=token_ids,
                 block_ids=block_ids,
                 start=start_token_idx,
                 end=end_token_idx,
+                skip_first_n_tokens=skip_first_n_tokens,
             )
 
             ret = LMCacheMPRequestMetadata(
@@ -700,13 +710,22 @@ class LMCacheMPConnector(KVConnectorBase_V1):
             num_external_tokens (int): the number of tokens that will be
                 loaded from the external KV cache.
         """
-        # NOTE: the `blocks` are NEW BLOCKS allocated for this request.
+        # NOTE: `blocks` comes from kv_cache_manager.get_blocks(request_id),
+        # which returns ALL blocks for the request (not just newly allocated).
+        # This function may be called twice for async-load requests:
+        #   1st call: blocks = initial allocation (APC + fresh)
+        #   2nd call: blocks = all blocks
+        #  (initial + newly allocated for remaining tokens)
+        # We must only append the NEW blocks beyond what's already tracked
+        # to avoid duplication, which would corrupt the store path's block indexing.
         tracker = self._get_request_tracker(request.request_id)
         block_ids = reformat_block_ids(blocks.get_block_ids())
 
-        # No matter we need to retrieve or not, we need to update
-        # the block ids into the tracker
-        tracker.append_block_ids(block_ids)
+        # Only append blocks beyond what's already tracked
+        existing_count = len(tracker.allocated_block_ids)
+        new_block_ids = block_ids[existing_count:]
+        if new_block_ids:
+            tracker.append_block_ids(new_block_ids)
 
         # Update the state of the tracker
         condition = tracker.needs_retrieve()
@@ -721,6 +740,34 @@ class LMCacheMPConnector(KVConnectorBase_V1):
             # Clean up lookup future in scheduler adapter
             self.scheduler_adapter.cleanup_lookup_result(request.request_id)
 
+            # Free locks on chunks that vLLM already computed and won't
+            # retrieve from LMCache.
+            if tracker.num_lmcache_hit_blocks > 0:
+                if not condition:
+                    # No retrieve needed — free ALL locked chunks
+                    free_end = tracker.num_lmcache_hit_blocks * self.vllm_block_size
+                else:
+                    # Note(Roy): Boundary misalignment between vLLM blocks and LMCache
+                    # blocks is handled in free_lookup_locks. It makes sure that if
+                    # the last vLLM computed block ends in the middle of a LMCache
+                    # block, the end LMCache block is not freed (i.e., floor division)
+                    # since it will still be needed by vLLM and such block's lock will
+                    # be freed by vLLM's retrieve.
+                    free_end = tracker.num_vllm_hit_blocks * self.vllm_block_size
+
+                if free_end > 0:
+                    self.scheduler_adapter.free_lookup_locks(
+                        token_ids=list(tracker.all_token_ids),
+                        start=0,
+                        end=free_end,
+                        request_id=request.request_id,
+                    )
+                    logger.debug(
+                        "Free locks of tokens %d-%d since it is cached by vLLM.",
+                        0,
+                        free_end,
+                    )
+
     def build_connector_meta(
         self, scheduler_output: SchedulerOutput
     ) -> KVConnectorMetadata:

vllm/engine/arg_utils.py

@@ -447,7 +447,6 @@ class EngineArgs:
     )
     disable_sliding_window: bool = ModelConfig.disable_sliding_window
     disable_cascade_attn: bool = ModelConfig.disable_cascade_attn
-    swap_space: float = CacheConfig.swap_space
     offload_backend: str = OffloadConfig.offload_backend
     cpu_offload_gb: float = UVAOffloadConfig.cpu_offload_gb
     cpu_offload_params: set[str] = get_field(UVAOffloadConfig, "cpu_offload_params")
@@ -961,7 +960,6 @@ class EngineArgs:
         cache_group.add_argument(
             "--kv-cache-memory-bytes", **cache_kwargs["kv_cache_memory_bytes"]
         )
-        cache_group.add_argument("--swap-space", **cache_kwargs["swap_space"])
         cache_group.add_argument("--kv-cache-dtype", **cache_kwargs["cache_dtype"])
         cache_group.add_argument(
             "--num-gpu-blocks-override", **cache_kwargs["num_gpu_blocks_override"]
@@ -1526,7 +1524,6 @@ class EngineArgs:
             block_size=self.block_size,
             gpu_memory_utilization=self.gpu_memory_utilization,
             kv_cache_memory_bytes=self.kv_cache_memory_bytes,
-            swap_space=self.swap_space,
             cache_dtype=resolved_cache_dtype,  # type: ignore[arg-type]
             is_attention_free=model_config.is_attention_free,
             num_gpu_blocks_override=self.num_gpu_blocks_override,

vllm/entrypoints/chat_utils.py

@@ -462,10 +462,15 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):
     maximum per prompt.
     """
 
-    def __init__(self, model_config: ModelConfig):
+    def __init__(
+        self,
+        model_config: ModelConfig,
+        media_io_kwargs: dict[str, dict[str, Any]] | None = None,
+    ):
         super().__init__()
 
         self._model_config = model_config
+        self._media_io_kwargs = media_io_kwargs
 
         self._items_by_modality = defaultdict[str, list[_T]](list)
         # Track original modality for each vision_chunk item (image or video)
@@ -487,6 +492,14 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):
         model_cls = get_model_cls(self.model_config)
         return cast(type[SupportsMultiModal], model_cls)
 
+    @property
+    def media_io_kwargs(self) -> dict[str, dict[str, Any]] | None:
+        return self._media_io_kwargs or (
+            self._model_config.multimodal_config.media_io_kwargs
+            if self._model_config.multimodal_config
+            else None
+        )
+
     @property
     def allowed_local_media_path(self):
         return self._model_config.allowed_local_media_path
@@ -769,12 +782,10 @@ class MultiModalContentParser(BaseMultiModalContentParser):
         super().__init__()
 
         self._tracker = tracker
-        multimodal_config = self._tracker.model_config.multimodal_config
-        media_io_kwargs = getattr(multimodal_config, "media_io_kwargs", None)
 
         self._connector: MediaConnector = MEDIA_CONNECTOR_REGISTRY.load(
             envs.VLLM_MEDIA_CONNECTOR,
-            media_io_kwargs=media_io_kwargs,
+            media_io_kwargs=tracker.media_io_kwargs,
             allowed_local_media_path=tracker.allowed_local_media_path,
             allowed_media_domains=tracker.allowed_media_domains,
         )
@@ -881,11 +892,9 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser):
         super().__init__()
 
         self._tracker = tracker
-        multimodal_config = self._tracker.model_config.multimodal_config
-        media_io_kwargs = getattr(multimodal_config, "media_io_kwargs", None)
         self._connector: MediaConnector = MEDIA_CONNECTOR_REGISTRY.load(
             envs.VLLM_MEDIA_CONNECTOR,
-            media_io_kwargs=media_io_kwargs,
+            media_io_kwargs=tracker.media_io_kwargs,
             allowed_local_media_path=tracker.allowed_local_media_path,
             allowed_media_domains=tracker.allowed_media_domains,
         )
@@ -1530,13 +1539,14 @@ def parse_chat_messages(
     messages: list[ChatCompletionMessageParam],
     model_config: ModelConfig,
     content_format: ChatTemplateContentFormat,
+    media_io_kwargs: dict[str, dict[str, Any]] | None = None,
 ) -> tuple[
     list[ConversationMessage],
     MultiModalDataDict | None,
     MultiModalUUIDDict | None,
 ]:
     conversation: list[ConversationMessage] = []
-    mm_tracker = MultiModalItemTracker(model_config)
+    mm_tracker = MultiModalItemTracker(model_config, media_io_kwargs=media_io_kwargs)
 
     for msg in messages:
         sub_messages = _parse_chat_message_content(
@@ -1563,13 +1573,16 @@ async def parse_chat_messages_async(
     messages: list[ChatCompletionMessageParam],
     model_config: ModelConfig,
     content_format: ChatTemplateContentFormat,
+    media_io_kwargs: dict[str, dict[str, Any]] | None = None,
 ) -> tuple[
     list[ConversationMessage],
     MultiModalDataDict | None,
     MultiModalUUIDDict | None,
 ]:
     conversation: list[ConversationMessage] = []
-    mm_tracker = AsyncMultiModalItemTracker(model_config)
+    mm_tracker = AsyncMultiModalItemTracker(
+        model_config, media_io_kwargs=media_io_kwargs
+    )
 
     for msg in messages:
         sub_messages = _parse_chat_message_content(

vllm/entrypoints/cli/launch.py

@@ -8,7 +8,7 @@ import uvloop
 from vllm.engine.arg_utils import AsyncEngineArgs
 from vllm.entrypoints.cli.types import CLISubcommand
 from vllm.entrypoints.openai.api_server import (
-    build_and_serve,
+    build_and_serve_renderer,
     setup_server,
 )
 from vllm.entrypoints.openai.cli_args import (
@@ -109,19 +109,17 @@ def cmd_init() -> list[CLISubcommand]:
 async def run_launch_fastapi(args: argparse.Namespace) -> None:
     """Run the online serving layer with FastAPI (no GPU inference)."""
     from vllm.config import VllmConfig
-    from vllm.v1.engine.launch import LaunchEngineClient
 
     # 1. Socket binding
     listen_address, sock = setup_server(args)
 
-    # 2. Create LaunchEngineClient (no GPU)
+    # 2. Build and serve the API server
     engine_args = AsyncEngineArgs.from_cli_args(args)
     model_config = engine_args.create_model_config()
     vllm_config = VllmConfig(model_config=model_config)
-    engine_client = LaunchEngineClient.from_vllm_config(vllm_config)
-
-    # 3. Build app, initialize state, and start serving
-    shutdown_task = await build_and_serve(engine_client, listen_address, sock, args)
+    shutdown_task = await build_and_serve_renderer(
+        vllm_config, listen_address, sock, args
+    )
     try:
         await shutdown_task
     finally:

vllm/entrypoints/llm.py

@@ -164,12 +164,6 @@ class LLM:
             compared with using gpu_memory_utilization. Note that
             kv_cache_memory_bytes (when not-None) ignores
             gpu_memory_utilization
-        swap_space: The size (GiB) of CPU memory per GPU to use as swap space.
-            This can be used for temporarily storing the states of the requests
-            when their `best_of` sampling parameters are larger than 1. If all
-            requests will have `best_of=1`, you can safely set this to 0.
-            Noting that `best_of` is only supported in V0. Otherwise, too small
-            values may cause out-of-memory (OOM) errors.
         cpu_offload_gb: The size (GiB) of CPU memory to use for offloading
             the model weights. This virtually increases the GPU memory space
             you can use to hold the model weights, at the cost of CPU-GPU data
@@ -240,7 +234,6 @@ class LLM:
         chat_template: Path | str | None = None,
         seed: int = 0,
         gpu_memory_utilization: float = 0.9,
-        swap_space: float = 4,
         cpu_offload_gb: float = 0,
         offload_group_size: int = 0,
         offload_num_in_group: int = 1,
@@ -265,6 +258,17 @@ class LLM:
     ) -> None:
         """LLM constructor."""
 
+        if "swap_space" in kwargs:
+            kwargs.pop("swap_space")
+            import warnings
+
+            warnings.warn(
+                "The 'swap_space' parameter is deprecated and ignored. "
+                "It will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
         if "disable_log_stats" not in kwargs:
             kwargs["disable_log_stats"] = True
 
@@ -353,7 +357,6 @@ class LLM:
             seed=seed,
             gpu_memory_utilization=gpu_memory_utilization,
             kv_cache_memory_bytes=kv_cache_memory_bytes,
-            swap_space=swap_space,
             cpu_offload_gb=cpu_offload_gb,
             offload_group_size=offload_group_size,
             offload_num_in_group=offload_num_in_group,

vllm/entrypoints/openai/api_server.py

@@ -22,6 +22,7 @@ from fastapi.middleware.cors import CORSMiddleware
 from starlette.datastructures import State
 
 import vllm.envs as envs
+from vllm.config import VllmConfig
 from vllm.engine.arg_utils import AsyncEngineArgs
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.chat_utils import load_chat_template
@@ -198,7 +199,7 @@ def build_app(
 
     register_sagemaker_api_router(app, supported_tasks)
 
-    if any(task in supported_tasks for task in ("generate", "render")):
+    if "generate" in supported_tasks:
         from vllm.entrypoints.openai.generate.api_router import (
             register_generate_api_routers,
         )
@@ -223,6 +224,13 @@ def build_app(
 
         elastic_ep_attach_router(app)
 
+    if "generate" in supported_tasks or "render" in supported_tasks:
+        from vllm.entrypoints.serve.render.api_router import (
+            attach_router as attach_render_router,
+        )
+
+        attach_render_router(app)
+
     if "transcription" in supported_tasks:
         from vllm.entrypoints.openai.speech_to_text.api_router import (
             attach_router as register_speech_to_text_api_router,
@@ -363,7 +371,7 @@ async def init_app_state(
         trust_request_chat_template=args.trust_request_chat_template,
     )
 
-    if any(task in supported_tasks for task in ("generate", "render")):
+    if "generate" in supported_tasks:
         from vllm.entrypoints.openai.generate.api_router import init_generate_state
 
         await init_generate_state(
@@ -393,6 +401,64 @@ async def init_app_state(
     state.server_load_metrics = 0
 
 
+async def init_render_app_state(
+    vllm_config: VllmConfig,
+    state: State,
+    args: Namespace,
+) -> None:
+    """Initialise FastAPI app state for a CPU-only render server.
+
+    Unlike :func:`init_app_state` this function does not require an
+    :class:`~vllm.engine.protocol.EngineClient`; it bootstraps the
+    preprocessing pipeline (renderer, io_processor, input_processor)
+    directly from the :class:`~vllm.config.VllmConfig`.
+    """
+    from vllm.entrypoints.chat_utils import load_chat_template
+    from vllm.entrypoints.serve.render.serving import OpenAIServingRender
+    from vllm.plugins.io_processors import get_io_processor
+    from vllm.renderers import renderer_from_config
+
+    served_model_names = args.served_model_name or [args.model]
+
+    if args.enable_log_requests:
+        request_logger = RequestLogger(max_log_len=args.max_log_len)
+    else:
+        request_logger = None
+
+    renderer = renderer_from_config(vllm_config)
+    io_processor = get_io_processor(
+        vllm_config, renderer, vllm_config.model_config.io_processor_plugin
+    )
+    resolved_chat_template = load_chat_template(args.chat_template)
+
+    state.openai_serving_render = OpenAIServingRender(
+        model_config=vllm_config.model_config,
+        renderer=renderer,
+        io_processor=io_processor,
+        served_model_names=served_model_names,
+        request_logger=request_logger,
+        chat_template=resolved_chat_template,
+        chat_template_content_format=args.chat_template_content_format,
+        trust_request_chat_template=args.trust_request_chat_template,
+        enable_auto_tools=args.enable_auto_tool_choice,
+        exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
+        tool_parser=args.tool_call_parser,
+        default_chat_template_kwargs=args.default_chat_template_kwargs,
+        log_error_stack=args.log_error_stack,
+    )
+
+    # Expose models endpoint via the render handler.
+    state.openai_serving_models = state.openai_serving_render
+
+    state.vllm_config = vllm_config
+    # Disable stats logging — there is no engine to poll.
+    state.log_stats = False
+    state.engine_client = None
+    state.args = args
+    state.enable_server_load_tracking = False
+    state.server_load_metrics = 0
+
+
 def create_server_socket(addr: tuple[str, int]) -> socket.socket:
     family = socket.AF_INET
     if is_valid_ipv6_address(addr[0]):
@@ -494,7 +560,6 @@ async def build_and_serve(
 
     supported_tasks = await engine_client.get_supported_tasks()
     logger.info("Supported tasks: %s", supported_tasks)
-
     app = build_app(args, supported_tasks)
     await init_app_state(engine_client, app.state, args, supported_tasks)
 
@@ -522,6 +587,51 @@ async def build_and_serve(
     )
 
 
+async def build_and_serve_renderer(
+    vllm_config: VllmConfig,
+    listen_address: str,
+    sock: socket.socket,
+    args: Namespace,
+    **uvicorn_kwargs,
+) -> asyncio.Task:
+    """Build FastAPI app for a CPU-only render server, initialize state, and
+    start serving.
+
+    Returns the shutdown task for the caller to await.
+    """
+
+    # Get uvicorn log config (from file or with endpoint filter)
+    log_config = get_uvicorn_log_config(args)
+    if log_config is not None:
+        uvicorn_kwargs["log_config"] = log_config
+
+    app = build_app(args, ("render",))
+    await init_render_app_state(vllm_config, app.state, args)
+
+    logger.info("Starting vLLM server on %s", listen_address)
+
+    return await serve_http(
+        app,
+        sock=sock,
+        enable_ssl_refresh=args.enable_ssl_refresh,
+        host=args.host,
+        port=args.port,
+        log_level=args.uvicorn_log_level,
+        # NOTE: When the 'disable_uvicorn_access_log' value is True,
+        # no access log will be output.
+        access_log=not args.disable_uvicorn_access_log,
+        timeout_keep_alive=envs.VLLM_HTTP_TIMEOUT_KEEP_ALIVE,
+        ssl_keyfile=args.ssl_keyfile,
+        ssl_certfile=args.ssl_certfile,
+        ssl_ca_certs=args.ssl_ca_certs,
+        ssl_cert_reqs=args.ssl_cert_reqs,
+        ssl_ciphers=args.ssl_ciphers,
+        h11_max_incomplete_event_size=args.h11_max_incomplete_event_size,
+        h11_max_header_count=args.h11_max_header_count,
+        **uvicorn_kwargs,
+    )
+
+
 async def run_server(args, **uvicorn_kwargs) -> None:
     """Run a single-worker API server."""
 

vllm/entrypoints/openai/chat_completion/api_router.py

@@ -71,34 +71,5 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re
     return StreamingResponse(content=generator, media_type="text/event-stream")
 
 
-@router.post(
-    "/v1/chat/completions/render",
-    dependencies=[Depends(validate_json_request)],
-    response_model=list,
-    responses={
-        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
-        HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
-        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
-        HTTPStatus.NOT_IMPLEMENTED.value: {"model": ErrorResponse},
-    },
-)
-async def render_chat_completion(request: ChatCompletionRequest, raw_request: Request):
-    """Render chat completion request and return conversation and engine
-    prompts without generating."""
-    handler = chat(raw_request)
-    if handler is None:
-        base_server = raw_request.app.state.openai_serving_tokenization
-        return base_server.create_error_response(
-            message="The model does not support Chat Completions API"
-        )
-
-    result = await handler.render_chat_request(request)
-
-    if isinstance(result, ErrorResponse):
-        return JSONResponse(content=result.model_dump(), status_code=result.error.code)
-
-    return JSONResponse(content=result)
-
-
 def attach_router(app: FastAPI):
     app.include_router(router)

vllm/entrypoints/openai/chat_completion/protocol.py

@@ -268,6 +268,13 @@ class ChatCompletionRequest(OpenAIBaseModel):
             "Will be accessible by the chat template."
         ),
     )
+    media_io_kwargs: dict[str, dict[str, Any]] | None = Field(
+        default=None,
+        description=(
+            "Additional kwargs to pass to the media IO connectors, "
+            "keyed by modality. Merged with engine-level media_io_kwargs."
+        ),
+    )
     mm_processor_kwargs: dict[str, Any] | None = Field(
         default=None,
         description=("Additional kwargs to pass to the HF processor."),
@@ -366,6 +373,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
                     reasoning_effort=self.reasoning_effort,
                 ),
             ),
+            media_io_kwargs=self.media_io_kwargs,
         )
 
     def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:

vllm/entrypoints/openai/completion/api_router.py

@@ -69,32 +69,5 @@ async def create_completion(request: CompletionRequest, raw_request: Request):
     return StreamingResponse(content=generator, media_type="text/event-stream")
 
 
-@router.post(
-    "/v1/completions/render",
-    dependencies=[Depends(validate_json_request)],
-    response_model=list,
-    responses={
-        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
-        HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
-        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
-    },
-)
-async def render_completion(request: CompletionRequest, raw_request: Request):
-    """render completion request and return engine prompts without generating."""
-    handler = completion(raw_request)
-    if handler is None:
-        base_server = raw_request.app.state.openai_serving_tokenization
-        return base_server.create_error_response(
-            message="The model does not support Completions API"
-        )
-
-    result = await handler.render_completion_request(request)
-
-    if isinstance(result, ErrorResponse):
-        return JSONResponse(content=result.model_dump(), status_code=result.error.code)
-
-    return JSONResponse(content=result)
-
-
 def attach_router(app: FastAPI):
     app.include_router(router)

vllm/entrypoints/openai/engine/serving.py

@@ -900,10 +900,15 @@ class OpenAIServing:
             ),
         )
 
+        mm_config = self.model_config.multimodal_config
+
         tok_params = request.build_tok_params(self.model_config)
         chat_params = request.build_chat_params(
             default_template, default_template_content_format
-        ).with_defaults(default_template_kwargs)
+        ).with_defaults(
+            default_template_kwargs,
+            default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None),
+        )
 
         (conversation,), (engine_prompt,) = await renderer.render_chat_async(
             [messages],

vllm/entrypoints/openai/generate/api_router.py

@@ -111,7 +111,7 @@ async def init_generate_state(
             enable_log_outputs=args.enable_log_outputs,
             enable_log_deltas=args.enable_log_deltas,
         )
-        if any(task in supported_tasks for task in ("generate", "render"))
+        if "generate" in supported_tasks
         else None
     )
     # Warm up chat template processing to avoid first-request latency
@@ -126,7 +126,7 @@ async def init_generate_state(
             enable_prompt_tokens_details=args.enable_prompt_tokens_details,
             enable_force_include_usage=args.enable_force_include_usage,
         )
-        if any(task in supported_tasks for task in ("generate", "render"))
+        if "generate" in supported_tasks
         else None
     )
     state.anthropic_serving_messages = (
@@ -160,3 +160,26 @@ async def init_generate_state(
         if "generate" in supported_tasks
         else None
     )
+
+    # Render endpoints are always backed by OpenAIServingRender so that
+    # /v1/chat/completions/render and /v1/completions/render work on both
+    # generate-mode and render-only servers.
+    from vllm.entrypoints.serve.render.serving import OpenAIServingRender
+
+    state.openai_serving_render = OpenAIServingRender(
+        model_config=engine_client.model_config,
+        renderer=engine_client.renderer,
+        io_processor=engine_client.io_processor,
+        served_model_names=[
+            mp.name for mp in state.openai_serving_models.base_model_paths
+        ],
+        request_logger=request_logger,
+        chat_template=resolved_chat_template,
+        chat_template_content_format=args.chat_template_content_format,
+        trust_request_chat_template=args.trust_request_chat_template,
+        enable_auto_tools=args.enable_auto_tool_choice,
+        exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
+        tool_parser=args.tool_call_parser,
+        default_chat_template_kwargs=args.default_chat_template_kwargs,
+        log_error_stack=args.log_error_stack,
+    )

vllm/entrypoints/openai/responses/protocol.py

@@ -197,6 +197,13 @@ class ResponsesRequest(OpenAIBaseModel):
             "through out the inference process and return in response."
         ),
     )
+    media_io_kwargs: dict[str, dict[str, Any]] | None = Field(
+        default=None,
+        description=(
+            "Additional kwargs to pass to the media IO connectors, "
+            "keyed by modality. Merged with engine-level media_io_kwargs."
+        ),
+    )
     mm_processor_kwargs: dict[str, Any] | None = Field(
         default=None,
         description=("Additional kwargs to pass to the HF processor."),
@@ -276,6 +283,7 @@ class ResponsesRequest(OpenAIBaseModel):
                     reasoning_effort=None if reasoning is None else reasoning.effort,
                 ),
             ),
+            media_io_kwargs=self.media_io_kwargs,
         )
 
     def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:

vllm/entrypoints/pooling/base/io_processor.py

@@ -123,10 +123,15 @@ class PoolingIOProcessor:
             ),
         )
 
+        mm_config = self.model_config.multimodal_config
+
         tok_params = request.build_tok_params(self.model_config)
         chat_params = request.build_chat_params(
             default_template, default_template_content_format
-        ).with_defaults(default_template_kwargs)
+        ).with_defaults(
+            default_template_kwargs,
+            default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None),
+        )
 
         (conversation,), (engine_prompt,) = renderer.render_chat(
             [messages],

vllm/entrypoints/pooling/base/protocol.py

@@ -124,6 +124,13 @@ class ChatRequestMixin(OpenAIBaseModel):
             "Will be accessible by the chat template."
         ),
     )
+    media_io_kwargs: dict[str, dict[str, Any]] | None = Field(
+        default=None,
+        description=(
+            "Additional kwargs to pass to the media IO connectors, "
+            "keyed by modality. Merged with engine-level media_io_kwargs."
+        ),
+    )
     # --8<-- [end:chat-extra-params]
 
     @model_validator(mode="before")
@@ -151,6 +158,7 @@ class ChatRequestMixin(OpenAIBaseModel):
                     continue_final_message=self.continue_final_message,
                 ),
             ),
+            media_io_kwargs=self.media_io_kwargs,
         )
 
 

vllm/entrypoints/serve/instrumentator/health.py

@@ -22,8 +22,12 @@ def engine_client(request: Request) -> EngineClient:
 @router.get("/health", response_class=Response)
 async def health(raw_request: Request) -> Response:
     """Health check."""
+    client = engine_client(raw_request)
+    if client is None:
+        # Render-only servers have no engine; they are always healthy.
+        return Response(status_code=200)
     try:
-        await engine_client(raw_request).check_health()
+        await client.check_health()
         return Response(status_code=200)
     except EngineDeadError:
         return Response(status_code=503)

vllm/entrypoints/serve/render/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/entrypoints/serve/render/api_router.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from http import HTTPStatus
+
+from fastapi import APIRouter, Depends, FastAPI, Request
+from fastapi.responses import JSONResponse
+
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.entrypoints.openai.completion.protocol import CompletionRequest
+from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+from vllm.entrypoints.openai.utils import validate_json_request
+from vllm.entrypoints.serve.render.serving import OpenAIServingRender
+from vllm.entrypoints.utils import create_error_response
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+router = APIRouter()
+
+
+def render(request: Request) -> OpenAIServingRender | None:
+    return getattr(request.app.state, "openai_serving_render", None)
+
+
+@router.post(
+    "/v1/chat/completions/render",
+    dependencies=[Depends(validate_json_request)],
+    response_model=list,
+    responses={
+        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
+        HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
+        HTTPStatus.NOT_IMPLEMENTED.value: {"model": ErrorResponse},
+        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
+    },
+)
+async def render_chat_completion(request: ChatCompletionRequest, raw_request: Request):
+    handler = render(raw_request)
+    if handler is None:
+        error = create_error_response(
+            message="The model does not support Chat Completions Render API",
+            err_type="NotFoundError",
+            status_code=HTTPStatus.NOT_FOUND,
+        )
+        return JSONResponse(
+            status_code=HTTPStatus.NOT_FOUND, content=error.model_dump()
+        )
+
+    result = await handler.render_chat_request(request)
+
+    if isinstance(result, ErrorResponse):
+        return JSONResponse(content=result.model_dump(), status_code=result.error.code)
+
+    return JSONResponse(content=result)
+
+
+@router.post(
+    "/v1/completions/render",
+    dependencies=[Depends(validate_json_request)],
+    response_model=list,
+    responses={
+        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
+        HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
+        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
+    },
+)
+async def render_completion(request: CompletionRequest, raw_request: Request):
+    handler = render(raw_request)
+    if handler is None:
+        error = create_error_response(
+            message="The model does not support Completions Render API",
+            err_type="NotFoundError",
+            status_code=HTTPStatus.NOT_FOUND,
+        )
+        return JSONResponse(
+            status_code=HTTPStatus.NOT_FOUND, content=error.model_dump()
+        )
+
+    result = await handler.render_completion_request(request)
+
+    if isinstance(result, ErrorResponse):
+        return JSONResponse(content=result.model_dump(), status_code=result.error.code)
+
+    return JSONResponse(content=result)
+
+
+def attach_router(app: FastAPI) -> None:
+    app.include_router(router)

vllm/entrypoints/serve/render/serving.py

@@ -0,0 +1,475 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import sys
+import traceback
+from collections.abc import Callable, Sequence
+from http import HTTPStatus
+from typing import Any
+
+import jinja2
+from openai_harmony import Message as OpenAIMessage
+
+from vllm.config import ModelConfig
+from vllm.entrypoints.chat_utils import (
+    ChatTemplateContentFormatOption,
+    ConversationMessage,
+)
+from vllm.entrypoints.logger import RequestLogger
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.entrypoints.openai.completion.protocol import CompletionRequest
+from vllm.entrypoints.openai.engine.protocol import (
+    ErrorInfo,
+    ErrorResponse,
+    ModelCard,
+    ModelList,
+    ModelPermission,
+)
+from vllm.entrypoints.openai.parser.harmony_utils import (
+    get_developer_message,
+    get_system_message,
+    parse_chat_inputs_to_harmony_messages,
+    render_for_completion,
+)
+from vllm.entrypoints.utils import sanitize_message
+from vllm.inputs.data import ProcessorInputs, PromptType, SingletonPrompt, TokensPrompt
+from vllm.logger import init_logger
+from vllm.parser import ParserManager
+from vllm.renderers import BaseRenderer, merge_kwargs
+from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq
+from vllm.tokenizers import TokenizerLike
+from vllm.tool_parsers import ToolParser
+from vllm.utils.mistral import is_mistral_tokenizer
+from vllm.utils.mistral import mt as _mt
+
+logger = init_logger(__name__)
+
+
+class OpenAIServingRender:
+    def __init__(
+        self,
+        model_config: ModelConfig,
+        renderer: BaseRenderer,
+        io_processor: Any,
+        served_model_names: list[str],
+        *,
+        request_logger: RequestLogger | None,
+        chat_template: str | None,
+        chat_template_content_format: ChatTemplateContentFormatOption,
+        trust_request_chat_template: bool = False,
+        enable_auto_tools: bool = False,
+        exclude_tools_when_tool_choice_none: bool = False,
+        tool_parser: str | None = None,
+        default_chat_template_kwargs: dict[str, Any] | None = None,
+        log_error_stack: bool = False,
+    ) -> None:
+        self.model_config = model_config
+        self.renderer = renderer
+        self.io_processor = io_processor
+        self.served_model_names = served_model_names
+        self.request_logger = request_logger
+        self.chat_template = chat_template
+        self.chat_template_content_format: ChatTemplateContentFormatOption = (
+            chat_template_content_format
+        )
+        self.trust_request_chat_template = trust_request_chat_template
+        self.enable_auto_tools = enable_auto_tools
+        self.exclude_tools_when_tool_choice_none = exclude_tools_when_tool_choice_none
+        self.tool_parser: Callable[[TokenizerLike], ToolParser] | None = (
+            ParserManager.get_tool_parser(
+                tool_parser_name=tool_parser,
+                enable_auto_tools=enable_auto_tools,
+                model_name=model_config.model,
+            )
+        )
+        self.default_chat_template_kwargs: dict[str, Any] = (
+            default_chat_template_kwargs or {}
+        )
+        self.log_error_stack = log_error_stack
+        self.use_harmony = model_config.hf_config.model_type == "gpt_oss"
+        self.supports_browsing = False
+        self.supports_code_interpreter = False
+
+    async def render_chat_request(
+        self,
+        request: ChatCompletionRequest,
+    ) -> tuple[list[ConversationMessage], list[ProcessorInputs]] | ErrorResponse:
+        """Copied from OpenAIServingChat.render_chat_request.
+
+        Differences: engine_client.errored check removed (no engine client).
+        """
+        error_check_ret = await self._check_model(request)
+        if error_check_ret is not None:
+            logger.error("Error with model %s", error_check_ret)
+            return error_check_ret
+
+        try:
+            tokenizer = self.renderer.tokenizer
+
+            tool_parser = self.tool_parser
+
+            if is_mistral_tokenizer(tokenizer):
+                # because of issues with pydantic we need to potentially
+                # re-serialize the tool_calls field of the request
+                # for more info: see comment in `maybe_serialize_tool_calls`
+                _mt.maybe_serialize_tool_calls(request)  # type: ignore[arg-type]
+                _mt.truncate_tool_call_ids(request)  # type: ignore[arg-type]
+                _mt.validate_request_params(request)
+
+            # Check if tool parsing is unavailable (common condition)
+            tool_parsing_unavailable = (
+                tool_parser is None
+                and not is_mistral_tokenizer(tokenizer)
+                and not self.use_harmony
+            )
+
+            # Validate tool_choice when tool parsing is required but unavailable
+            if tool_parsing_unavailable and request.tool_choice not in (
+                None,
+                "none",
+            ):
+                if request.tool_choice == "auto" and not self.enable_auto_tools:
+                    # for hf tokenizers, "auto" tools requires
+                    # --enable-auto-tool-choice and --tool-call-parser
+                    return self.create_error_response(
+                        '"auto" tool choice requires '
+                        "--enable-auto-tool-choice and --tool-call-parser to be set"
+                    )
+                elif request.tool_choice != "auto":
+                    # "required" or named tool requires tool parser
+                    return self.create_error_response(
+                        f'tool_choice="{request.tool_choice}" requires '
+                        "--tool-call-parser to be set"
+                    )
+
+            if request.tools is None or (
+                request.tool_choice == "none"
+                and self.exclude_tools_when_tool_choice_none
+            ):
+                tool_dicts = None
+            else:
+                tool_dicts = [tool.model_dump() for tool in request.tools]
+
+            if not self.use_harmony:
+                # Common case.
+                error_check_ret = self._validate_chat_template(
+                    request_chat_template=request.chat_template,
+                    chat_template_kwargs=request.chat_template_kwargs,
+                    trust_request_chat_template=self.trust_request_chat_template,
+                )
+                if error_check_ret is not None:
+                    return error_check_ret
+
+                conversation, engine_prompts = await self._preprocess_chat(
+                    request,
+                    request.messages,
+                    default_template=self.chat_template,
+                    default_template_content_format=self.chat_template_content_format,
+                    default_template_kwargs=self.default_chat_template_kwargs,
+                    tool_dicts=tool_dicts,
+                    tool_parser=tool_parser,
+                )
+            else:
+                # For GPT-OSS.
+                should_include_tools = tool_dicts is not None
+                conversation, engine_prompts = self._make_request_with_harmony(
+                    request, should_include_tools
+                )
+        except (ValueError, TypeError, RuntimeError, jinja2.TemplateError) as e:
+            logger.exception("Error in preprocessing prompt inputs")
+            return self.create_error_response(e)
+
+        return conversation, engine_prompts
+
+    async def render_completion_request(
+        self,
+        request: CompletionRequest,
+    ) -> list[ProcessorInputs] | ErrorResponse:
+        """Copied from OpenAIServingCompletion.render_completion_request.
+
+        Differences: engine_client.errored check removed (no engine client).
+        """
+        error_check_ret = await self._check_model(request)
+        if error_check_ret is not None:
+            return error_check_ret
+
+        # Return error for unsupported features.
+        if request.suffix is not None:
+            return self.create_error_response("suffix is not currently supported")
+
+        if request.echo and request.prompt_embeds is not None:
+            return self.create_error_response("Echo is unsupported with prompt embeds.")
+
+        if request.prompt_logprobs is not None and request.prompt_embeds is not None:
+            return self.create_error_response(
+                "prompt_logprobs is not compatible with prompt embeds."
+            )
+
+        try:
+            engine_prompts = await self._preprocess_completion(
+                request,
+                prompt_input=request.prompt,
+                prompt_embeds=request.prompt_embeds,
+            )
+        except (ValueError, TypeError, RuntimeError, jinja2.TemplateError) as e:
+            logger.exception("Error in preprocessing prompt inputs")
+            return self.create_error_response(e)
+
+        return engine_prompts
+
+    def _make_request_with_harmony(
+        self,
+        request: ChatCompletionRequest,
+        should_include_tools: bool = True,
+    ):
+        """Copied from OpenAIServingChat._make_request_with_harmony."""
+        messages: list[OpenAIMessage] = []
+
+        # because of issues with pydantic we need to potentially
+        # re-serialize the tool_calls field of the request
+        # for more info: see comment in `maybe_serialize_tool_calls`
+        _mt.maybe_serialize_tool_calls(request)  # type: ignore[arg-type]
+
+        # Add system message.
+        # NOTE: In Chat Completion API, browsing is enabled by default
+        # if the model supports it. TODO: Support browsing.
+        assert not self.supports_browsing
+        assert not self.supports_code_interpreter
+        sys_msg = get_system_message(
+            reasoning_effort=request.reasoning_effort,
+            browser_description=None,
+            python_description=None,
+            with_custom_tools=should_include_tools,
+        )
+        messages.append(sys_msg)
+
+        # Add developer message.
+        if request.tools:
+            dev_msg = get_developer_message(
+                tools=request.tools if should_include_tools else None  # type: ignore[arg-type]
+            )
+            messages.append(dev_msg)
+
+        # Add user message.
+        messages.extend(parse_chat_inputs_to_harmony_messages(request.messages))
+
+        # Render prompt token ids.
+        prompt_token_ids = render_for_completion(messages)
+        engine_prompt = TokensPrompt(prompt_token_ids=prompt_token_ids)
+
+        # Add cache_salt if provided in the request
+        if request.cache_salt is not None:
+            engine_prompt["cache_salt"] = request.cache_salt
+
+        return messages, [engine_prompt]
+
+    async def show_available_models(self) -> ModelList:
+        """Returns the models served by this render server."""
+        max_model_len = self.model_config.max_model_len
+        return ModelList(
+            data=[
+                ModelCard(
+                    id=name,
+                    max_model_len=max_model_len,
+                    root=self.model_config.model,
+                    permission=[ModelPermission()],
+                )
+                for name in self.served_model_names
+            ]
+        )
+
+    def create_error_response(
+        self,
+        message: str | Exception,
+        err_type: str = "BadRequestError",
+        status_code: HTTPStatus = HTTPStatus.BAD_REQUEST,
+        param: str | None = None,
+    ) -> ErrorResponse:
+        """Copied from OpenAIServing.create_error_response."""
+        exc: Exception | None = None
+
+        if isinstance(message, Exception):
+            exc = message
+
+            from vllm.exceptions import VLLMValidationError
+
+            if isinstance(exc, VLLMValidationError):
+                err_type = "BadRequestError"
+                status_code = HTTPStatus.BAD_REQUEST
+                param = exc.parameter
+            elif isinstance(exc, (ValueError, TypeError, RuntimeError, OverflowError)):
+                # Common validation errors from user input
+                err_type = "BadRequestError"
+                status_code = HTTPStatus.BAD_REQUEST
+                param = None
+            elif isinstance(exc, NotImplementedError):
+                err_type = "NotImplementedError"
+                status_code = HTTPStatus.NOT_IMPLEMENTED
+                param = None
+            elif exc.__class__.__name__ == "TemplateError":
+                # jinja2.TemplateError (avoid importing jinja2)
+                err_type = "BadRequestError"
+                status_code = HTTPStatus.BAD_REQUEST
+                param = None
+            else:
+                err_type = "InternalServerError"
+                status_code = HTTPStatus.INTERNAL_SERVER_ERROR
+                param = None
+
+            message = str(exc)
+
+        if self.log_error_stack:
+            exc_type, _, _ = sys.exc_info()
+            if exc_type is not None:
+                traceback.print_exc()
+            else:
+                traceback.print_stack()
+
+        return ErrorResponse(
+            error=ErrorInfo(
+                message=sanitize_message(message),
+                type=err_type,
+                code=status_code.value,
+                param=param,
+            )
+        )
+
+    def _is_model_supported(self, model_name: str) -> bool:
+        """Simplified from OpenAIServing._is_model_supported (no LoRA support)."""
+        return model_name in self.served_model_names
+
+    async def _check_model(
+        self,
+        request: Any,
+    ) -> ErrorResponse | None:
+        """Simplified from OpenAIServing._check_model (no LoRA support)."""
+        if self._is_model_supported(request.model):
+            return None
+        return self.create_error_response(
+            message=f"The model `{request.model}` does not exist.",
+            err_type="NotFoundError",
+            status_code=HTTPStatus.NOT_FOUND,
+            param="model",
+        )
+
+    def _validate_chat_template(
+        self,
+        request_chat_template: str | None,
+        chat_template_kwargs: dict[str, Any] | None,
+        trust_request_chat_template: bool,
+    ) -> ErrorResponse | None:
+        """Copied from OpenAIServing._validate_chat_template."""
+        if not trust_request_chat_template and (
+            request_chat_template is not None
+            or (
+                chat_template_kwargs
+                and chat_template_kwargs.get("chat_template") is not None
+            )
+        ):
+            return self.create_error_response(
+                "Chat template is passed with request, but "
+                "--trust-request-chat-template is not set. "
+                "Refused request with untrusted chat template."
+            )
+        return None
+
+    async def _preprocess_completion(
+        self,
+        request: Any,
+        prompt_input: str | list[str] | list[int] | list[list[int]] | None,
+        prompt_embeds: bytes | list[bytes] | None,
+    ) -> list[ProcessorInputs]:
+        """Copied from OpenAIServing._preprocess_completion."""
+        prompts = list[SingletonPrompt | bytes]()
+        if prompt_embeds is not None:  # embeds take higher priority
+            prompts.extend(prompt_to_seq(prompt_embeds))
+        if prompt_input is not None:
+            prompts.extend(prompt_to_seq(prompt_input))
+        return await self._preprocess_cmpl(request, prompts)
+
+    async def _preprocess_cmpl(
+        self,
+        request: Any,
+        prompts: Sequence[PromptType | bytes],
+    ) -> list[ProcessorInputs]:
+        """Copied from OpenAIServing._preprocess_cmpl."""
+        renderer = self.renderer
+        model_config = self.model_config
+
+        parsed_prompts = [
+            (
+                prompt
+                if isinstance(prompt, bytes)
+                else parse_model_prompt(model_config, prompt)
+            )
+            for prompt in prompts
+        ]
+        tok_params = request.build_tok_params(model_config)
+
+        return await renderer.render_cmpl_async(
+            parsed_prompts,
+            tok_params,
+            prompt_extras={
+                k: v
+                for k in ("mm_processor_kwargs", "cache_salt")
+                if (v := getattr(request, k, None)) is not None
+            },
+        )
+
+    async def _preprocess_chat(
+        self,
+        request: Any,
+        messages: list[Any],
+        default_template: str | None,
+        default_template_content_format: ChatTemplateContentFormatOption,
+        default_template_kwargs: dict[str, Any] | None,
+        tool_dicts: list[dict[str, Any]] | None = None,
+        tool_parser: Callable[[TokenizerLike], ToolParser] | None = None,
+    ) -> tuple[list[ConversationMessage], list[ProcessorInputs]]:
+        """Copied from OpenAIServing._preprocess_chat.
+
+        Differences: isinstance check is ChatCompletionRequest-only
+        (ResponsesRequest not supported here); TODO comment dropped accordingly.
+        """
+        renderer = self.renderer
+
+        default_template_kwargs = merge_kwargs(
+            default_template_kwargs,
+            dict(
+                tools=tool_dicts,
+                tokenize=is_mistral_tokenizer(renderer.tokenizer),
+            ),
+        )
+
+        tok_params = request.build_tok_params(self.model_config)
+        chat_params = request.build_chat_params(
+            default_template, default_template_content_format
+        ).with_defaults(default_template_kwargs)
+
+        (conversation,), (engine_prompt,) = await renderer.render_chat_async(
+            [messages],
+            chat_params,
+            tok_params,
+            prompt_extras={
+                k: v
+                for k in ("mm_processor_kwargs", "cache_salt")
+                if (v := getattr(request, k, None)) is not None
+            },
+        )
+
+        # tool parsing is done only if a tool_parser has been set and if
+        # tool_choice is not "none" (if tool_choice is "none" but a tool_parser
+        # is set, we want to prevent parsing a tool_call hallucinated by the LLM
+        if tool_parser is not None:
+            tool_choice = getattr(request, "tool_choice", "none")
+            if tool_choice != "none":
+                if not isinstance(request, ChatCompletionRequest):
+                    msg = (
+                        "Tool usage is only supported "
+                        " for ChatCompletionRequest, but got "
+                        f"{type(request).__name__}"
+                    )
+                    raise NotImplementedError(msg)
+                tokenizer = renderer.get_tokenizer()
+                request = tool_parser(tokenizer).adjust_request(request=request)  # type: ignore[arg-type]
+
+        return conversation, [engine_prompt]

vllm/entrypoints/serve/tokenize/protocol.py

@@ -100,6 +100,13 @@ class TokenizeChatRequest(OpenAIBaseModel):
             "Will be accessible by the chat template."
         ),
     )
+    media_io_kwargs: dict[str, dict[str, Any]] | None = Field(
+        default=None,
+        description=(
+            "Additional kwargs to pass to the media IO connectors, "
+            "keyed by modality. Merged with engine-level media_io_kwargs."
+        ),
+    )
     mm_processor_kwargs: dict[str, Any] | None = Field(
         default=None,
         description="Additional kwargs to pass to the HF processor.",
@@ -134,6 +141,7 @@ class TokenizeChatRequest(OpenAIBaseModel):
                     continue_final_message=self.continue_final_message,
                 ),
             ),
+            media_io_kwargs=self.media_io_kwargs,
         )
 
     def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:

vllm/envs.py

@@ -244,6 +244,7 @@ if TYPE_CHECKING:
     VLLM_CUDA_COMPATIBILITY_PATH: str | None = None
     VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: bool = False
     VLLM_ELASTIC_EP_DRAIN_REQUESTS: bool = False
+    VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: bool = False
 
 
 def get_default_cache_root():
@@ -1628,6 +1629,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_ELASTIC_EP_DRAIN_REQUESTS": lambda: bool(
         int(os.getenv("VLLM_ELASTIC_EP_DRAIN_REQUESTS", "0"))
     ),
+    # If set to 1, enable CUDA graph memory estimation during memory profiling.
+    # This profiles CUDA graph memory usage to provide more accurate KV cache
+    # memory allocation. Disabled by default to preserve existing behavior.
+    "VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS": lambda: bool(
+        int(os.getenv("VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS", "0"))
+    ),
 }
 
 

vllm/model_executor/layers/attention/mla_attention.py

@@ -331,11 +331,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):
             calculate_kv_scales = False
         self.quant_config = quant_config
 
-        # Initialize KV cache quantization attributes
-        self.kv_cache_dtype = kv_cache_dtype
-        self.calculate_kv_scales = calculate_kv_scales
-        _init_kv_cache_quant(self, quant_config, prefix)
-
         dtype = torch.get_default_dtype()
         self.attn_backend = get_attn_backend(
             self.head_size,
@@ -347,6 +342,36 @@ class MLAAttention(nn.Module, AttentionLayerBase):
             num_heads=self.num_heads,
         )
 
+        # FlashMLA Sparse Attention fp8 backend uses "fp8_ds_mla" kv-cache format
+        # Automatically convert fp8 kv-cache format to "fp8_ds_mla"
+        if (
+            self.attn_backend.get_name() == "FLASHMLA_SPARSE"
+            and kv_cache_dtype.startswith("fp8")
+            and kv_cache_dtype != "fp8_ds_mla"
+        ):
+            assert cache_config is not None
+            cache_config.cache_dtype = "fp8_ds_mla"
+            kv_cache_dtype = "fp8_ds_mla"
+            logger.info_once(
+                "Using DeepSeek's fp8_ds_mla KV cache format. To use standard "
+                "fp8 kv-cache format, please set `--attention-backend "
+                "FLASHINFER_MLA_SPARSE`"
+            )
+
+        if (
+            self.attn_backend.get_name() == "FLASHINFER_MLA_SPARSE"
+            and kv_cache_dtype.startswith("fp8")
+        ):
+            logger.info_once(
+                "Using standard fp8 KV cache format. To use DeepSeek's fp8_ds_mla "
+                "KV cache format, please set `--attention-backend FLASHMLA_SPARSE`"
+            )
+
+        # Initialize KV cache quantization attributes
+        self.kv_cache_dtype = kv_cache_dtype
+        self.calculate_kv_scales = calculate_kv_scales
+        _init_kv_cache_quant(self, quant_config, prefix)
+
         if (
             cache_config is not None
             and cache_config.enable_prefix_caching
@@ -905,6 +930,10 @@ def unified_mla_kv_cache_update(
     the data dependency between them to ensure torch.compile preserves ordering.
     """
     forward_context = get_forward_context()
+    if forward_context.attn_metadata is None:
+        # Dummy/profile forwards should not update live KV cache pages.
+        return torch.empty(0, device=kv_c_normed.device, dtype=kv_c_normed.dtype)
+
     attn_layer = forward_context.no_compile_layers[layer_name]
     kv_cache = attn_layer.kv_cache[forward_context.virtual_engine]
 

vllm/model_executor/models/config.py

@@ -31,20 +31,13 @@ class VerifyAndUpdateConfig:
 class DeepseekV32ForCausalLM(VerifyAndUpdateConfig):
     @classmethod
     def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
-        """
-        Updated fp8 cache to custom "fp8_ds_mla" format for DeepSeekV32
-        """
         hf_config = vllm_config.model_config.hf_config
 
         # Mirror the check in vllm/model_executor/models/deepseek_v2.py
         is_v32 = hasattr(hf_config, "index_topk")
         assert is_v32
 
-        # For DeepSeekV3.2, a custom fp8 format is used when fp8 kv-cache is enabled.
         cache_config = vllm_config.cache_config
-        if cache_config.cache_dtype.startswith("fp8"):
-            cache_config.cache_dtype = "fp8_ds_mla"
-            logger.info("Using custom fp8 kv-cache format for DeepSeekV3.2")
         if cache_config.cache_dtype == "bfloat16":
             cache_config.cache_dtype = "auto"
             logger.info("Using bfloat16 kv-cache for DeepSeekV3.2")

vllm/model_executor/models/deepseek_v2.py

@@ -756,7 +756,7 @@ direct_register_custom_op(
 )
 
 
-class DeepSeekV2FusedQkvAProj(MergedColumnParallelLinear):
+class DeepSeekV2FusedQkvAProjLinear(MergedColumnParallelLinear):
     def __init__(
         self,
         input_size: int,
@@ -848,7 +848,7 @@ class DeepseekV2MLAAttention(nn.Module):
         self.max_position_embeddings = max_position_embeddings
 
         if self.q_lora_rank is not None:
-            self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProj(
+            self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear(
                 self.hidden_size,
                 [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
                 quant_config=quant_config,

vllm/model_executor/models/nano_nemotron_vl.py

@@ -17,11 +17,11 @@ from functools import cached_property
 from typing import Annotated, Any, Literal, TypeAlias, TypeVar
 
 import einops
+import numpy as np
 import numpy.typing as npt
 import regex as re
 import torch
 import torch.nn as nn
-import torchvision.transforms as T
 from PIL import Image
 from transformers import BatchFeature, PretrainedConfig, TensorType
 
@@ -214,7 +214,12 @@ NanoNemotronVLVideoInputs: TypeAlias = (
 
 
 def dynamic_preprocess(
-    image, *, image_size=512, max_num_tiles=12, use_thumbnail=True, idx=0
+    image,
+    *,
+    image_size=512,
+    max_num_tiles=12,
+    use_thumbnail=True,
+    idx=0,
 ):
     orig_width, orig_height = image.size
 
@@ -227,35 +232,44 @@ def dynamic_preprocess(
         image_size=image_size,
         use_thumbnail=False,
     )
-    # resize the image
-    resized_img = image.resize((target_width, target_height))
-    processed_images = []
-    for i in range(blocks):
-        box = (
-            (i % (target_width // image_size)) * image_size,
-            (i // (target_width // image_size)) * image_size,
-            ((i % (target_width // image_size)) + 1) * image_size,
-            ((i // (target_width // image_size)) + 1) * image_size,
-        )
-        # split the image
-        split_img = resized_img.crop(box)
-        processed_images.append(split_img)
-    assert len(processed_images) == blocks
-    if use_thumbnail and len(processed_images) != 1:
-        thumbnail_img = image.resize((image_size, image_size))
-        processed_images.append(thumbnail_img)
 
-    processed_images = [
-        img.convert("RGB") if img.mode != "RGB" else img for img in processed_images
-    ]
-    processed_images = [
-        T.Resize((image_size, image_size), interpolation=T.InterpolationMode.BICUBIC)(
-            img
+    image = np.asarray(
+        image.convert("RGB") if image.mode != "RGB" else image, dtype=np.uint8
+    )
+
+    image = torch.from_numpy(image).unsqueeze(0)  # (1, H, W, 3)
+    image = image.permute(0, 3, 1, 2)  # (1, 3, H, W)
+
+    resized_img = torch.nn.functional.interpolate(
+        image,
+        size=(target_height, target_width),
+        mode="bicubic",
+        align_corners=False,
+        antialias=True,
+    )
+    B, C, H, W = resized_img.shape
+    hp, wp = H // image_size, W // image_size
+    patches = (
+        resized_img.reshape(B, C, hp, image_size, wp, image_size)
+        .permute(0, 2, 4, 1, 3, 5)
+        .reshape(B * hp * wp, C, image_size, image_size)
+        / 255.0
+    )
+
+    if use_thumbnail and patches.shape[0] > 1:
+        thumb = (
+            torch.nn.functional.interpolate(
+                image,
+                size=(image_size, image_size),
+                mode="bicubic",
+                align_corners=False,
+                antialias=True,
+            )
+            / 255.0
         )
-        for img in processed_images
-    ]
-    processed_images = [T.ToTensor()(img) for img in processed_images]
-    return processed_images
+        patches = torch.cat([patches, thumb], dim=0)
+
+    return list(patches)
 
 
 def image_to_pixel_values(
@@ -287,22 +301,21 @@ def video_to_pixel_values(
 ) -> torch.Tensor:
     assert max_num_tiles == 1, "Video modality always uses one tile"
 
-    # Convert each frame to a single resized tile tensor consistent
-    # with image path
-    frames_tensors: list[torch.Tensor] = []
-    for frame in video:
-        pil_frame = dynamic_preprocess(
-            Image.fromarray(frame, mode="RGB"),
-            image_size=input_size,
-            max_num_tiles=max_num_tiles,
-            use_thumbnail=use_thumbnail,
-            idx=0,
-        )
-        # dynamic_preprocess returns tensors already; take the single tile
-        assert len(pil_frame) >= 1
-        frames_tensors.append(pil_frame[-1])
+    # (num_frames, H, W, C) -> (num_frames, C, H, W)
+    video_tensor = torch.from_numpy(video).permute(0, 3, 1, 2)
 
-    return torch.stack(frames_tensors)
+    if video_tensor.shape[2] != input_size or video_tensor.shape[3] != input_size:
+        video_tensor = torch.nn.functional.interpolate(
+            video_tensor,
+            size=(input_size, input_size),
+            mode="bicubic",
+            align_corners=False,
+            antialias=True,
+        )
+
+    video_tensor = video_tensor / 255.0
+
+    return video_tensor
 
 
 def input_conditioner(x, norm_mean, norm_std):
@@ -346,12 +359,6 @@ class DynamicResolutionImageTiler:
         self._factor_max = factor_max
         self.norm_mean = torch.tensor(norm_mean).reshape(3, 1, 1)
         self.norm_std = torch.tensor(norm_std).reshape(3, 1, 1)
-        self._transform = T.Compose(
-            [
-                T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img),
-                T.ToTensor(),
-            ]
-        )
         assert downsample_ratio < 1
         reduction_factor = 1 / downsample_ratio
         assert reduction_factor == 2.0
@@ -441,15 +448,25 @@ class DynamicResolutionImageTiler:
         patch_size: tuple[int, int]
 
     def apply_params(self, params: DynamicResolutionParams) -> list[torch.Tensor]:
-        resized_img = params.media.resize(
-            (
-                params.patch_size[0] * self._patch_size,
-                params.patch_size[1] * self._patch_size,
-            )
+        target_size = (
+            params.patch_size[1] * self._patch_size,
+            params.patch_size[0] * self._patch_size,
         )
-        processed_images = [resized_img]
-
-        return [self._transform(img) for img in processed_images]
+        image = np.asarray(
+            params.media.convert("RGB") if params.media.mode != "RGB" else params.media,
+            dtype=np.uint8,
+        )
+        resized_img = (
+            torch.nn.functional.interpolate(
+                torch.from_numpy(image).unsqueeze(0).permute(0, 3, 1, 2),
+                size=target_size,
+                mode="bicubic",
+                align_corners=False,
+                antialias=True,
+            )
+            / 255.0
+        )
+        return list(resized_img)
 
     def process_media(
         self,
@@ -803,6 +820,7 @@ class BaseNanoNemotronVLProcessor(ABC):
             image_repl = self.get_image_repl(feature_size, num_patches)
             parts[i] = parts[i].replace("<image>", image_repl.full)
         text = ["".join(parts)]
+
         return text, image_inputs
 
     def _make_batch_input(self, input_item: Any | list[Any] | None = None):
@@ -922,14 +940,14 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):
             frames_indices_lst = [
                 metadata["frames_indices"] for metadata in video_metadata_lst
             ]
-
+            video_num_patches = torch.tensor(
+                [len(item) for item in pixel_values_lst_video]
+            )
             video_inputs = {
                 "pixel_values_flat_video": input_conditioner(
                     torch.cat(pixel_values_lst_video), self.norm_mean, self.norm_std
                 ),
-                "video_num_patches": torch.tensor(
-                    [len(item) for item in pixel_values_lst_video]
-                ),
+                "video_num_patches": video_num_patches,
                 "frames_indices": frames_indices_lst,
                 "frame_duration_ms": torch.tensor(frame_duration_ms_lst),
             }
@@ -985,6 +1003,7 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):
                     video_repl.full, skip_special_tokens=False
                 )
                 text = [t.replace("<video>", video_repl_text, 1) for t in text]
+
         return text, video_inputs
 
     def _preprocess_audio(

vllm/model_executor/models/registry.py

@@ -191,6 +191,8 @@ _TEXT_GENERATION_MODELS = {
     "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),
     "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),
     "RWForCausalLM": ("falcon", "FalconForCausalLM"),
+    "SarvamMoEForCausalLM": ("sarvam", "SarvamMoEForCausalLM"),
+    "SarvamMLAForCausalLM": ("sarvam", "SarvamMLAForCausalLM"),
     "SeedOssForCausalLM": ("seed_oss", "SeedOssForCausalLM"),
     "Step1ForCausalLM": ("step1", "Step1ForCausalLM"),
     "Step3TextForCausalLM": ("step3_text", "Step3TextForCausalLM"),

vllm/model_executor/models/sarvam.py

@@ -0,0 +1,786 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Copyright 2026 Sarvam AI team. All rights reserved.
+#
+# This code is based on Llama, Deepseek, and Bailing MoE implementations
+# in this library. It has been modified from its original forms to
+# accommodate Sarvam's MoE architectures.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import math
+from collections.abc import Iterable, Iterator
+from itertools import islice
+
+import torch
+from torch import nn
+
+from vllm.config import CacheConfig, ParallelConfig, VllmConfig
+from vllm.distributed import (
+    get_pp_group,
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.fused_moe import SharedFusedMoE
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.sequence import IntermediateTensors
+
+from .bailing_moe import BailingMoeForCausalLM
+from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
+from .utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    is_pp_missing_parameter,
+    make_empty_intermediate_tensors_factory,
+    make_layers,
+    maybe_prefix,
+)
+
+
+def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
+    if scale <= 1:
+        return 1.0
+    return 0.1 * mscale * math.log(scale) + 1.0
+
+
+def _is_gate_expert_bias_name(name: str) -> bool:
+    return name.endswith(".mlp.gate.e_score_correction_bias") or name.endswith(
+        ".gate.e_score_correction_bias"
+    )
+
+
+def _zero_mean_tensor(t: torch.Tensor) -> torch.Tensor:
+    if t.numel() == 0:
+        return t
+    return t - t.mean()
+
+
+def _normalized_weights(
+    weights: Iterable[tuple[str, torch.Tensor]],
+) -> Iterator[tuple[str, torch.Tensor]]:
+    for name, w in weights:
+        if _is_gate_expert_bias_name(name):
+            yield name, _zero_mean_tensor(w)
+        else:
+            yield name, w
+
+
+class SarvamMLAAttention(nn.Module):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        config,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.qk_nope_head_dim = config.qk_nope_head_dim
+        self.qk_rope_head_dim = config.qk_rope_head_dim
+        self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+        self.v_head_dim = config.v_head_dim
+
+        self.q_lora_rank = getattr(config, "q_lora_rank", None)
+        self.kv_lora_rank = config.kv_lora_rank
+
+        self.total_num_heads = config.num_attention_heads
+        tp_size = get_tensor_model_parallel_world_size()
+        assert self.total_num_heads % tp_size == 0
+        self.num_local_heads = self.total_num_heads // tp_size
+
+        self.scaling = self.qk_head_dim**-0.5
+        self.max_position_embeddings = config.max_position_embeddings
+
+        if self.q_lora_rank is not None:
+            self.q_a_proj = ReplicatedLinear(
+                self.hidden_size,
+                self.q_lora_rank,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.q_a_proj",
+            )
+            self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
+            self.q_b_proj = ColumnParallelLinear(
+                self.q_lora_rank,
+                self.total_num_heads * self.qk_head_dim,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.q_b_proj",
+            )
+            self.q_proj = None  # type: ignore
+        else:
+            self.q_proj = ColumnParallelLinear(
+                self.hidden_size,
+                self.total_num_heads * self.qk_head_dim,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.q_proj",
+            )
+            self.q_a_proj = None  # type: ignore
+            self.q_a_layernorm = None  # type: ignore
+            self.q_b_proj = None  # type: ignore
+
+        # KV latent (MQA-style) A-proj
+        self.kv_a_proj_with_mqa = ReplicatedLinear(
+            self.hidden_size,
+            self.kv_lora_rank + self.qk_rope_head_dim,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.kv_a_proj_with_mqa",
+        )
+        self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
+
+        # KV B-proj produces per-head K_nope and V
+        self.kv_b_proj = ColumnParallelLinear(
+            self.kv_lora_rank,
+            self.total_num_heads * (self.qk_nope_head_dim + self.v_head_dim),
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.kv_b_proj",
+        )
+
+        self.o_proj = RowParallelLinear(
+            self.total_num_heads * self.v_head_dim,
+            self.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        self.rotary_emb = get_rope(
+            self.qk_rope_head_dim,
+            # rotary_dim=self.qk_rope_head_dim,
+            max_position=config.max_position_embeddings,
+            rope_parameters=config.rope_parameters,
+            is_neox_style=False,
+        )
+
+        if config.rope_parameters.get("rope_type", None) == "deepseek_yarn":
+            mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False)
+            scaling_factor = config.rope_parameters["factor"]
+            mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
+            self.scaling = self.scaling * mscale * mscale
+
+        mla_modules = MLAModules(
+            kv_a_layernorm=self.kv_a_layernorm,
+            kv_b_proj=self.kv_b_proj,
+            rotary_emb=self.rotary_emb,
+            o_proj=self.o_proj,
+            fused_qkv_a_proj=None,
+            kv_a_proj_with_mqa=self.kv_a_proj_with_mqa,
+            q_a_layernorm=self.q_a_layernorm if self.q_lora_rank is not None else None,
+            q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None,
+            q_proj=self.q_proj if self.q_lora_rank is None else None,
+            indexer=None,
+            indexer_rotary_emb=None,
+            is_sparse=False,
+            topk_indices_buffer=None,
+        )
+
+        self.mla_attn = MultiHeadLatentAttentionWrapper(
+            self.hidden_size,
+            self.num_local_heads,
+            self.scaling,
+            self.qk_nope_head_dim,
+            self.qk_rope_head_dim,
+            self.v_head_dim,
+            self.q_lora_rank,
+            self.kv_lora_rank,
+            mla_modules,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.mla_attn(positions, hidden_states, llama_4_scaling=None)
+
+
+class SarvamMLAMLP(nn.Module):
+    def __init__(
+        self,
+        intermediate_size: int,
+        config,
+        quant_config: QuantizationConfig | None = None,
+        reduce_results: bool = True,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.gate_up_proj = MergedColumnParallelLinear(
+            config.hidden_size,
+            [intermediate_size] * 2,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            intermediate_size,
+            config.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            reduce_results=reduce_results,
+            prefix=f"{prefix}.down_proj",
+        )
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate_up, _ = self.gate_up_proj(x)
+        x = self.act_fn(gate_up)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class SarvamMLAMoE(nn.Module):
+    def __init__(
+        self,
+        config,
+        parallel_config: ParallelConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.config = config
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_rank = get_tensor_model_parallel_rank()
+        self.hidden_size = config.hidden_size
+
+        self.num_experts = config.num_experts
+        self.top_k = config.num_experts_per_tok
+        self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 2.5)
+
+        self.n_group = getattr(config, "n_group", None)
+        self.topk_group = getattr(config, "topk_group", None)
+        self.use_grouped_topk = self.n_group is not None and self.topk_group is not None
+
+        self.norm_expert_prob = getattr(config, "norm_topk_prob", True)
+
+        router_dtype_cfg = getattr(config, "router_dtype", "fp32")
+        if router_dtype_cfg is None:
+            self.router_dtype = None
+        elif router_dtype_cfg == "fp32":
+            self.router_dtype = torch.float32
+        else:
+            self.router_dtype = torch.bfloat16
+
+        self.gate = nn.Linear(
+            self.hidden_size,
+            self.num_experts,
+            bias=False,
+            dtype=self.router_dtype,
+        )
+
+        if getattr(config, "moe_router_enable_expert_bias", True):
+            self.gate.e_score_correction_bias = nn.Parameter(
+                torch.empty(
+                    (self.num_experts,),
+                    dtype=torch.float32,
+                )
+            )
+        else:
+            self.gate.e_score_correction_bias = None
+
+        self.score_function = getattr(config, "score_function", "sigmoid")
+        self.num_shared_experts = getattr(config, "num_shared_experts", 1)
+        if self.num_shared_experts > 0:
+            if hasattr(config, "moe_shared_expert_intermediate_size"):
+                shared_int = config.moe_shared_expert_intermediate_size
+            else:
+                shared_int = config.moe_intermediate_size
+            shared_int *= self.num_shared_experts
+            self.shared_experts = SarvamMLAMLP(
+                intermediate_size=shared_int,
+                config=config,
+                quant_config=quant_config,
+                reduce_results=False,
+                prefix=f"{prefix}.shared_experts",
+            )
+        else:
+            self.shared_experts = None
+
+        self.experts = SharedFusedMoE(
+            shared_experts=self.shared_experts,
+            num_experts=self.num_experts,
+            top_k=self.top_k,
+            hidden_size=self.hidden_size,
+            intermediate_size=config.moe_intermediate_size,
+            reduce_results=False,
+            renormalize=self.norm_expert_prob,
+            quant_config=quant_config,
+            prefix=f"{prefix}.experts",
+            scoring_func=self.score_function,
+            e_score_correction_bias=self.gate.e_score_correction_bias,
+            num_expert_group=self.n_group,
+            topk_group=self.topk_group,
+            use_grouped_topk=self.use_grouped_topk,
+            routed_scaling_factor=self.routed_scaling_factor,
+        )
+
+    def maybe_get_fused_moe(self) -> SharedFusedMoE:
+        return self.experts
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        num_tokens, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        router_logits = self.gate(
+            hidden_states.to(self.router_dtype)
+            if self.router_dtype is not None
+            else hidden_states
+        )
+        router_logits = router_logits.to(hidden_states.dtype)
+        final_hidden = self.experts(
+            hidden_states=hidden_states,
+            router_logits=router_logits,
+        )
+
+        if self.shared_experts is not None:
+            shared_output, expert_output = final_hidden
+        else:
+            shared_output, expert_output = None, final_hidden
+
+        if shared_output is not None:
+            expert_output = expert_output + shared_output
+
+        if self.tp_size > 1:
+            expert_output = self.experts.maybe_all_reduce_tensor_model_parallel(
+                expert_output
+            )
+
+        return expert_output.view(num_tokens, hidden_dim)
+
+
+class SarvamMLABlock(nn.Module):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+        parallel_config = vllm_config.parallel_config
+        layer_idx = int(prefix.split(".")[-1])
+        hidden_size = config.hidden_size
+        dense_intermediate = getattr(config, "intermediate_size", 16384)
+
+        self.input_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps)
+        self.self_attn = SarvamMLAAttention(
+            vllm_config=vllm_config,
+            config=config,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.post_attention_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps)
+        use_moe = hasattr(config, "num_experts") and config.num_experts is not None
+        first_k_dense = getattr(config, "first_k_dense_replace", 1)
+        moe_layer_freq = getattr(config, "moe_layer_freq", 1)
+        if use_moe:
+            is_moe_layer = layer_idx >= first_k_dense and (
+                (layer_idx - first_k_dense) % moe_layer_freq == 0
+            )
+        else:
+            is_moe_layer = False
+
+        if is_moe_layer:
+            self.mlp = SarvamMLAMoE(
+                config=config,
+                parallel_config=parallel_config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.mlp",
+            )
+        else:
+            self.mlp = SarvamMLAMLP(
+                intermediate_size=dense_intermediate,
+                config=config,
+                quant_config=quant_config,
+                reduce_results=True,
+                prefix=f"{prefix}.mlp",
+            )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        positions: torch.Tensor,
+        residual: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+        )
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+
+class SarvamMLAModel(nn.Module):
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+
+        self.config = config
+        self.vocab_size = config.vocab_size
+        self.embed_dim = config.hidden_size
+        self.tie_word_embeddings = getattr(config, "tie_word_embeddings", False)
+        if get_pp_group().is_first_rank or (
+            self.tie_word_embeddings and get_pp_group().is_last_rank
+        ):
+            self.embed_tokens = VocabParallelEmbedding(
+                self.vocab_size,
+                self.embed_dim,
+                quant_config=quant_config,
+                prefix=f"{prefix}.embed_tokens",
+            )
+        else:
+            self.embed_tokens = PPMissingLayer()
+
+        self.embedding_dropout = torch.nn.Dropout(
+            getattr(config, "embedding_dropout", 0.0)
+        )
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers,
+            lambda prefix: SarvamMLABlock(
+                vllm_config=vllm_config,
+                prefix=prefix,
+            ),
+            prefix=f"{prefix}.layers",
+        )
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states", "residual"], config.hidden_size
+        )
+        if get_pp_group().is_last_rank:
+            self.norm = RMSNorm(self.embed_dim, eps=config.rms_norm_eps)
+        else:
+            self.norm = PPMissingLayer()
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        if get_pp_group().is_first_rank:
+            if inputs_embeds is not None:
+                hidden_states = inputs_embeds
+            else:
+                hidden_states = self.embed_input_ids(input_ids)
+            hidden_states = self.embedding_dropout(hidden_states)
+            residual = None
+        else:
+            assert intermediate_tensors is not None
+            hidden_states = intermediate_tensors["hidden_states"]
+            residual = intermediate_tensors["residual"]
+
+        for layer in islice(self.layers, self.start_layer, self.end_layer):
+            hidden_states, residual = layer(
+                hidden_states,
+                positions,
+                residual,
+            )
+        if not get_pp_group().is_last_rank:
+            return IntermediateTensors(
+                {"hidden_states": hidden_states, "residual": residual}
+            )
+        if residual is None:
+            hidden_states = self.norm(hidden_states)
+        else:
+            hidden_states, _ = self.norm(hidden_states, residual)
+        return hidden_states
+
+    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
+        return SharedFusedMoE.make_expert_params_mapping(
+            self,
+            ckpt_gate_proj_name="gate_proj",
+            ckpt_down_proj_name="down_proj",
+            ckpt_up_proj_name="up_proj",
+            num_experts=self.config.num_experts,
+        )
+
+    def load_weights(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> set[str]:
+        """Load weights with stacked gate+up and MoE expert remapping."""
+        weights = _normalized_weights(weights)
+        stacked_params_mapping = [
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+
+        params_dict = dict(self.named_parameters(remove_duplicate=False))
+        loaded_params: set[str] = set()
+        expert_params_mapping = self.get_expert_mapping()
+
+        for name, loaded_weight in weights:
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                if "mlp.experts" in name:
+                    continue
+                new_name = name.replace(weight_name, param_name)
+                if new_name.endswith(".bias") and new_name not in params_dict:
+                    continue
+                if new_name not in params_dict:
+                    continue
+                if is_pp_missing_parameter(new_name, self):
+                    continue
+
+                param = params_dict[new_name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight, shard_id)
+                loaded_params.add(new_name)
+                break
+            else:
+                mapped = False
+                for (
+                    param_name,
+                    weight_name,
+                    expert_id,
+                    shard_id,
+                ) in expert_params_mapping:
+                    if weight_name not in name:
+                        continue
+
+                    new_name = name.replace(weight_name, param_name)
+                    if is_pp_missing_parameter(new_name, self):
+                        continue
+                    if new_name not in params_dict:
+                        continue
+
+                    param = params_dict[new_name]
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(
+                        param,
+                        loaded_weight,
+                        name,
+                        shard_id=shard_id,
+                        expert_id=expert_id,
+                    )
+                    loaded_params.add(new_name)
+                    mapped = True
+                    break
+
+                if mapped:
+                    continue
+
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                if name not in params_dict:
+                    continue
+                if is_pp_missing_parameter(name, self):
+                    continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+                loaded_params.add(name)
+
+        return loaded_params
+
+
+class SarvamMixtureOfExperts(MixtureOfExperts):
+    def extract_moe_parameters(self, example_moe: SarvamMLAMoE | None) -> None:
+        if example_moe is None:
+            raise RuntimeError("No SarvamMLAMoE layer found in model.layers.")
+
+        self.num_logical_experts = example_moe.num_experts
+        self.num_routed_experts = example_moe.num_experts  # routed pool size
+        self.num_shared_experts = getattr(example_moe.config, "num_shared_experts", 1)
+
+        self.num_physical_experts = self.num_logical_experts
+        self.num_local_physical_experts = self.num_logical_experts
+        self.num_redundant_experts = 0
+
+    def update_physical_experts_metadata(
+        self,
+        num_physical_experts: int,
+        num_local_physical_experts: int,
+    ) -> None:
+        self.num_physical_experts = num_physical_experts
+        self.num_local_physical_experts = num_local_physical_experts
+        self.num_redundant_experts = num_physical_experts - self.num_logical_experts
+
+        for moe in self.moe_mlp_layers:
+            moe.n_physical_experts = num_physical_experts
+            moe.n_local_physical_experts = num_local_physical_experts
+            moe.n_redundant_experts = self.num_redundant_experts
+
+            fused = moe.experts
+            if hasattr(fused, "n_local_physical_experts"):
+                fused.n_local_physical_experts = num_local_physical_experts
+            if hasattr(fused, "n_physical_experts"):
+                fused.n_physical_experts = num_physical_experts
+            if hasattr(fused, "n_redundant_experts"):
+                fused.n_redundant_experts = self.num_redundant_experts
+            if hasattr(fused, "update_expert_map"):
+                fused.update_expert_map()
+
+    def set_eplb_state(self, eplb_state) -> None:
+        self.eplb_state = eplb_state
+        for moe in self.moe_layers:
+            if hasattr(moe, "set_eplb_state"):
+                moe.set_eplb_state(eplb_state)
+
+
+class SarvamMLAForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SarvamMixtureOfExperts):
+    packed_modules_mapping = {
+        "q_proj": ["q_proj"],
+        "q_a_proj": ["q_a_proj"],
+        "q_b_proj": ["q_b_proj"],
+        "kv_a_proj_with_mqa": ["kv_a_proj_with_mqa"],
+        "kv_b_proj": ["kv_b_proj"],
+        "gate_up_proj": ["gate_proj", "up_proj"],
+    }
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        self.config = config
+        self.quant_config = quant_config
+
+        self.model = SarvamMLAModel(
+            vllm_config=vllm_config,
+            prefix=maybe_prefix(prefix, "model"),
+        )
+
+        self.tie_word_embeddings = getattr(config, "tie_word_embeddings", False)
+        if get_pp_group().is_last_rank:
+            if self.tie_word_embeddings:
+                self.lm_head = self.model.embed_tokens
+            else:
+                self.lm_head = ParallelLMHead(
+                    config.vocab_size,
+                    config.hidden_size,
+                    quant_config=quant_config,
+                    prefix=maybe_prefix(prefix, "lm_head"),
+                )
+            self.logits_processor = LogitsProcessor(config.vocab_size)
+        else:
+            self.lm_head = PPMissingLayer()
+            self.logits_processor = None  # type: ignore
+
+        self.make_empty_intermediate_tensors = (
+            self.model.make_empty_intermediate_tensors
+        )
+
+        self.expert_weights = []
+        self.num_moe_layers = 0
+
+        self.moe_layers = []
+        self.moe_mlp_layers = []
+
+        example_moe = None
+        for layer in self.model.layers:
+            if isinstance(layer, PPMissingLayer):
+                continue
+            if isinstance(layer.mlp, SarvamMLAMoE):
+                example_moe = layer.mlp
+                self.moe_mlp_layers.append(layer.mlp)
+                self.moe_layers.append(layer.mlp.experts)
+                self.num_moe_layers += 1
+
+        self.extract_moe_parameters(example_moe)
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        return self.model(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        if not get_pp_group().is_last_rank:
+            return None
+        logits = self.logits_processor(self.lm_head, hidden_states)
+        return logits
+
+    def load_weights(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> set[str]:
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=(["lm_head."] if self.tie_word_embeddings else None),
+        )
+        return loader.load_weights(weights)
+
+    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
+        return self.model.get_expert_mapping()
+
+
+class SarvamMoEForCausalLM(BailingMoeForCausalLM):
+    """Same as BailingMoeForCausalLM, but normalizes gate expert_bias pre-load."""
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        return super().load_weights(_normalized_weights(weights))

vllm/multimodal/media/audio.py

@@ -83,11 +83,17 @@ def extract_audio_from_video_bytes(
 
 
 class AudioMediaIO(MediaIO[tuple[npt.NDArray, float]]):
+    """Configuration values can be user-provided either by --media-io-kwargs or
+    by the runtime API field "media_io_kwargs". Ensure proper validation and
+    error handling.
+    """
+
     def __init__(self, **kwargs) -> None:
         super().__init__()
 
         # `kwargs` contains custom arguments from
-        # --media-io-kwargs for this modality.
+        # --media-io-kwargs for this modality, merged with
+        # per-request runtime media_io_kwargs via merge_kwargs().
         # They can be passed to the underlying
         # media loaders (e.g. custom implementations)
         # for flexible control.
@@ -122,6 +128,11 @@ class AudioMediaIO(MediaIO[tuple[npt.NDArray, float]]):
 
 
 class AudioEmbeddingMediaIO(MediaIO[torch.Tensor]):
+    """Configuration values can be user-provided either by --media-io-kwargs or
+    by the runtime API field "media_io_kwargs". Ensure proper validation and
+    error handling.
+    """
+
     def __init__(self) -> None:
         super().__init__()
 

vllm/multimodal/media/base.py

@@ -44,6 +44,28 @@ class MediaWithBytes(Generic[_T]):
 
 
 class MediaIO(ABC, Generic[_T]):
+    """Configuration values can be user-provided either by --media-io-kwargs or
+    by the runtime API field "media_io_kwargs". Ensure proper validation and
+    error handling.
+    """
+
+    @classmethod
+    def merge_kwargs(
+        cls,
+        default_kwargs: dict[str, Any] | None,
+        runtime_kwargs: dict[str, Any] | None,
+    ) -> dict[str, Any]:
+        """Merge config-level kwargs and request-level kwargs.
+
+        By default this performs a shallow merge where runtime kwargs override
+        keys in default kwargs. Subclasses may override to apply modality-
+        specific behavior.
+        """
+        merged = dict(default_kwargs or {})
+        if runtime_kwargs:
+            merged.update(runtime_kwargs)
+        return merged
+
     @abstractmethod
     def load_bytes(self, data: bytes) -> _T:
         raise NotImplementedError

vllm/multimodal/media/connector.py

@@ -32,9 +32,43 @@ atexit.register(global_thread_pool.shutdown)
 
 MEDIA_CONNECTOR_REGISTRY = ExtensionManager()
 
+MODALITY_IO_MAP: dict[str, type[MediaIO]] = {
+    "audio": AudioMediaIO,
+    "image": ImageMediaIO,
+    "video": VideoMediaIO,
+}
+
+
+def merge_media_io_kwargs(
+    defaults: dict[str, dict[str, Any]] | None,
+    overrides: dict[str, dict[str, Any]] | None,
+) -> dict[str, dict[str, Any]] | None:
+    """Merge config-level and per-request media_io_kwargs per modality.
+
+    Each modality key is merged using the corresponding MediaIO subclass's
+    ``merge_kwargs``, which may apply modality-specific logic (e.g.
+    VideoMediaIO clears cross-dependent fps/num_frames fields).
+    """
+    if not defaults and not overrides:
+        return None
+    all_keys = set(defaults or {}) | set(overrides or {})
+    merged = {}
+    for key in all_keys:
+        io_cls = MODALITY_IO_MAP.get(key, MediaIO)
+        merged[key] = io_cls.merge_kwargs(
+            (defaults or {}).get(key),
+            (overrides or {}).get(key),
+        )
+    return merged or None
+
 
 @MEDIA_CONNECTOR_REGISTRY.register("http")
 class MediaConnector:
+    """Configuration values can be user-provided either by --media-io-kwargs or
+    by the runtime API field "media_io_kwargs". Ensure proper validation and
+    error handling.
+    """
+
     def __init__(
         self,
         media_io_kwargs: dict[str, dict[str, Any]] | None = None,

vllm/multimodal/media/image.py

@@ -15,12 +15,18 @@ from .base import MediaIO, MediaWithBytes
 
 
 class ImageMediaIO(MediaIO[Image.Image]):
+    """Configuration values can be user-provided either by --media-io-kwargs or
+    by the runtime API field "media_io_kwargs". Ensure proper validation and
+    error handling.
+    """
+
     def __init__(self, image_mode: str = "RGB", **kwargs) -> None:
         super().__init__()
 
         self.image_mode = image_mode
         # `kwargs` contains custom arguments from
-        # --media-io-kwargs for this modality.
+        # --media-io-kwargs for this modality, merged with
+        # per-request runtime media_io_kwargs via merge_kwargs().
         # They can be passed to the underlying
         # media loaders (e.g. custom implementations)
         # for flexible control.
@@ -88,6 +94,13 @@ class ImageMediaIO(MediaIO[Image.Image]):
 
 
 class ImageEmbeddingMediaIO(MediaIO[torch.Tensor]):
+    """Image embedding MediaIO implementation.
+
+    Configuration values can be user-provided either by --media-io-kwargs or
+    by the runtime API field "media_io_kwargs". Ensure proper validation and
+    error handling.
+    """
+
     def __init__(self) -> None:
         super().__init__()
 

vllm/multimodal/media/video.py

@@ -17,6 +17,28 @@ from .image import ImageMediaIO
 
 
 class VideoMediaIO(MediaIO[tuple[npt.NDArray, dict[str, Any]]]):
+    """Configuration values can be user-provided either by --media-io-kwargs or
+    by the runtime API field "media_io_kwargs". Ensure proper validation and
+    error handling.
+    """
+
+    @classmethod
+    def merge_kwargs(
+        cls,
+        default_kwargs: dict[str, Any] | None,
+        runtime_kwargs: dict[str, Any] | None,
+    ) -> dict[str, Any]:
+        merged = super().merge_kwargs(default_kwargs, runtime_kwargs)
+        # fps and num_frames interact with each other, so if either is
+        # overridden at request time, wipe the other from defaults to
+        # avoid unintuitive cross-field interactions.
+        if runtime_kwargs:
+            if "num_frames" in runtime_kwargs and "fps" not in runtime_kwargs:
+                merged.pop("fps", None)
+            elif "fps" in runtime_kwargs and "num_frames" not in runtime_kwargs:
+                merged.pop("num_frames", None)
+        return merged
+
     def __init__(
         self,
         image_io: ImageMediaIO,
@@ -28,7 +50,8 @@ class VideoMediaIO(MediaIO[tuple[npt.NDArray, dict[str, Any]]]):
         self.image_io = image_io
         self.num_frames = num_frames
         # `kwargs` contains custom arguments from
-        # --media-io-kwargs for this modality.
+        # --media-io-kwargs for this modality, merged with
+        # per-request runtime media_io_kwargs via merge_kwargs().
         # They can be passed to the underlying
         # media loaders (e.g. custom implementations)
         # for flexible control.

vllm/parser/__init__.py

@@ -22,13 +22,6 @@ _PARSERS_TO_REGISTER = {
     ),
 }
 
-# Register lazy parsers
-ParserManager.register_lazy_module(
-    name="minimax_m2",
-    module_path="vllm.parser.minimax_m2_parser",
-    class_name="MiniMaxM2Parser",
-)
-
 
 def register_lazy_parsers():
     for name, (file_name, class_name) in _PARSERS_TO_REGISTER.items():

vllm/renderers/deepseek_v32.py

@@ -49,6 +49,7 @@ class DeepseekV32Renderer(BaseRenderer[DeepseekV32Tokenizer]):
             messages,
             self.model_config,
             content_format="string",
+            media_io_kwargs=params.media_io_kwargs,
         )
 
         prompt_raw = tokenizer.apply_chat_template(
@@ -75,6 +76,7 @@ class DeepseekV32Renderer(BaseRenderer[DeepseekV32Tokenizer]):
             messages,
             self.model_config,
             content_format="string",
+            media_io_kwargs=params.media_io_kwargs,
         )
 
         prompt_raw = tokenizer.apply_chat_template(

vllm/renderers/grok2.py

@@ -49,6 +49,7 @@ class Grok2Renderer(BaseRenderer[Grok2Tokenizer]):
             messages,
             self.model_config,
             content_format="string",
+            media_io_kwargs=params.media_io_kwargs,
         )
 
         prompt_raw = tokenizer.apply_chat_template(
@@ -75,6 +76,7 @@ class Grok2Renderer(BaseRenderer[Grok2Tokenizer]):
             messages,
             self.model_config,
             content_format="string",
+            media_io_kwargs=params.media_io_kwargs,
         )
 
         prompt_raw = tokenizer.apply_chat_template(

vllm/renderers/hf.py

@@ -635,6 +635,7 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
                 tokenizer=tokenizer,
                 model_config=model_config,
             ),
+            media_io_kwargs=params.media_io_kwargs,
         )
 
         prompt_raw = safe_apply_chat_template(
@@ -689,6 +690,7 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
                 tokenizer=tokenizer,
                 model_config=model_config,
             ),
+            media_io_kwargs=params.media_io_kwargs,
         )
 
         prompt_raw = safe_apply_chat_template(

vllm/renderers/mistral.py

@@ -90,6 +90,7 @@ class MistralRenderer(BaseRenderer[MistralTokenizer]):
             messages,
             self.model_config,
             content_format="string",
+            media_io_kwargs=params.media_io_kwargs,
         )
 
         prompt_raw = safe_apply_chat_template(
@@ -116,6 +117,7 @@ class MistralRenderer(BaseRenderer[MistralTokenizer]):
             messages,
             self.model_config,
             content_format="string",
+            media_io_kwargs=params.media_io_kwargs,
         )
 
         prompt_raw = await self._apply_chat_template_async(

vllm/renderers/params.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, TypeVar
 from vllm.exceptions import VLLMValidationError
 from vllm.inputs import EmbedsPrompt, TextPrompt, TokensPrompt
 from vllm.logger import init_logger
+from vllm.multimodal.media.connector import merge_media_io_kwargs
 from vllm.tokenizers import TokenizerLike
 from vllm.utils.import_utils import LazyLoader
 
@@ -52,8 +53,15 @@ class ChatParams:
     chat_template_kwargs: dict[str, Any] = field(default_factory=dict)
     """The kwargs to pass to the chat template."""
 
-    def with_defaults(self, default_chat_template_kwargs: dict[str, Any] | None):
-        if not default_chat_template_kwargs:
+    media_io_kwargs: dict[str, dict[str, Any]] | None = None
+    """Per-modality kwargs for media I/O (loading/decoding images, videos, etc.)."""
+
+    def with_defaults(
+        self,
+        default_chat_template_kwargs: dict[str, Any] | None = None,
+        default_media_io_kwargs: dict[str, dict[str, Any]] | None = None,
+    ):
+        if not default_chat_template_kwargs and not default_media_io_kwargs:
             return self
 
         return ChatParams(
@@ -63,6 +71,10 @@ class ChatParams:
                 default_chat_template_kwargs,
                 self.chat_template_kwargs,
             ),
+            media_io_kwargs=merge_media_io_kwargs(
+                default_media_io_kwargs,
+                self.media_io_kwargs,
+            ),
         )
 
     def get_apply_chat_template_kwargs(self) -> dict[str, Any]:

vllm/renderers/terratorch.py

@@ -43,6 +43,7 @@ class TerratorchRenderer(BaseRenderer):
             messages,
             model_config,
             content_format="string",
+            media_io_kwargs=params.media_io_kwargs,
         )
 
         prompt = parse_dec_only_prompt([1])  # Dummy token IDs
@@ -64,6 +65,7 @@ class TerratorchRenderer(BaseRenderer):
             messages,
             model_config,
             content_format="string",
+            media_io_kwargs=params.media_io_kwargs,
         )
 
         prompt = parse_dec_only_prompt([1])  # Dummy token IDs

vllm/tool_parsers/hermes_tool_parser.py

@@ -385,6 +385,7 @@ class Hermes2ProToolParser(ToolParser):
             prev_arguments = self.prev_tool_call_arr[self.current_tool_id].get(
                 "arguments"
             )
+            assert current_tool_call is not None
             cur_arguments = current_tool_call.get("arguments")
 
             logger.debug("diffing old arguments: %s", prev_arguments)
@@ -489,6 +490,7 @@ class Hermes2ProToolParser(ToolParser):
 
             # handle saving the state for the current tool into
             # the "prev" list for use in diffing for the next iteration
+            assert isinstance(current_tool_call, dict)
             if self.current_tool_id == len(self.prev_tool_call_arr) - 1:
                 self.prev_tool_call_arr[self.current_tool_id] = current_tool_call
             else:

vllm/transformers_utils/config.py

@@ -24,7 +24,10 @@ from transformers.utils import CONFIG_NAME as HF_CONFIG_NAME
 from vllm import envs
 from vllm.logger import init_logger
 from vllm.transformers_utils.repo_utils import is_mistral_model_repo
-from vllm.transformers_utils.utils import parse_safetensors_file_metadata
+from vllm.transformers_utils.utils import (
+    parse_safetensors_file_metadata,
+    without_trust_remote_code,
+)
 
 from .config_parser_base import ConfigParserBase
 from .gguf_utils import (
@@ -140,11 +143,12 @@ class HFConfigParser(ConfigParserBase):
         **kwargs,
     ) -> tuple[dict, PretrainedConfig]:
         kwargs["local_files_only"] = huggingface_hub.constants.HF_HUB_OFFLINE
+        trust_remote_code |= kwargs.get("trust_remote_code", False)
+        kwargs = without_trust_remote_code(kwargs)
         config_dict, _ = PretrainedConfig.get_config_dict(
             model,
             revision=revision,
             code_revision=code_revision,
-            trust_remote_code=trust_remote_code,
             **kwargs,
         )
         # Use custom model class if it's in our registry
@@ -225,7 +229,7 @@ class MistralConfigParser(ConfigParserBase):
                 model,
                 revision=revision,
                 code_revision=code_revision,
-                **kwargs,
+                **without_trust_remote_code(kwargs),
             )
         except OSError:  # Not found
             hf_config_dict = {}
@@ -521,8 +525,7 @@ def maybe_override_with_speculators(
     config_dict, _ = PretrainedConfig.get_config_dict(
         model if gguf_model_repo is None else gguf_model_repo,
         revision=revision,
-        trust_remote_code=trust_remote_code,
-        **kwargs,
+        **without_trust_remote_code(kwargs),
     )
     speculators_config = config_dict.get("speculators_config")
 

vllm/transformers_utils/configs/eagle.py

@@ -5,6 +5,8 @@ import os
 
 from transformers import AutoConfig, DeepseekV2Config, PretrainedConfig
 
+from vllm.transformers_utils.utils import without_trust_remote_code
+
 
 class EAGLEConfig(PretrainedConfig):
     model_type = "eagle"
@@ -79,7 +81,7 @@ class EAGLEConfig(PretrainedConfig):
         **kwargs,
     ) -> "EAGLEConfig":
         config_dict, kwargs = cls.get_config_dict(
-            pretrained_model_name_or_path, **kwargs
+            pretrained_model_name_or_path, **without_trust_remote_code(kwargs)
         )
         return cls.from_dict(config_dict, **kwargs)
 

vllm/transformers_utils/configs/extract_hidden_states.py

@@ -7,6 +7,8 @@ import os
 
 from transformers import PretrainedConfig
 
+from vllm.transformers_utils.utils import without_trust_remote_code
+
 
 class ExtractHiddenStatesConfig(PretrainedConfig):
     model_type = "extract_hidden_states"
@@ -42,7 +44,7 @@ class ExtractHiddenStatesConfig(PretrainedConfig):
         **kwargs,
     ) -> "ExtractHiddenStatesConfig":
         config_dict, kwargs = cls.get_config_dict(
-            pretrained_model_name_or_path, **kwargs
+            pretrained_model_name_or_path, **without_trust_remote_code(kwargs)
         )
         return cls.from_dict(config_dict, **kwargs)
 

vllm/transformers_utils/configs/medusa.py

@@ -5,6 +5,8 @@ import os
 
 from transformers import PretrainedConfig
 
+from vllm.transformers_utils.utils import without_trust_remote_code
+
 
 class MedusaConfig(PretrainedConfig):
     model_type = "medusa"
@@ -42,7 +44,7 @@ class MedusaConfig(PretrainedConfig):
         **kwargs,
     ) -> "MedusaConfig":
         config_dict, kwargs = cls.get_config_dict(
-            pretrained_model_name_or_path, **kwargs
+            pretrained_model_name_or_path, **without_trust_remote_code(kwargs)
         )
         for k in list(config_dict.keys()):
             if "num" in k:

vllm/transformers_utils/configs/speculators/base.py

@@ -11,6 +11,8 @@ from vllm.transformers_utils.configs.speculators.algos import (
 
 __all__ = ["SpeculatorsConfig"]
 
+from vllm.transformers_utils.utils import without_trust_remote_code
+
 
 class SpeculatorsConfig(PretrainedConfig):
     model_type = "speculators"
@@ -22,7 +24,9 @@ class SpeculatorsConfig(PretrainedConfig):
         **kwargs,
     ) -> "SpeculatorsConfig":
         """Load speculators Eagle config and convert to vLLM format."""
-        config_dict, _ = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
+        config_dict, _ = cls.get_config_dict(
+            pretrained_model_name_or_path, **without_trust_remote_code(kwargs)
+        )
 
         vllm_config = cls.extract_transformers_pre_trained_config(config_dict)
         return cls(**vllm_config)

vllm/transformers_utils/utils.py

@@ -27,6 +27,13 @@ def is_cloud_storage(model_or_path: str) -> bool:
     return is_s3(model_or_path) or is_gcs(model_or_path)
 
 
+def without_trust_remote_code(kwargs: dict[str, Any]) -> dict[str, Any]:
+    """Return kwargs without trust_remote_code without modifying original dict."""
+    if "trust_remote_code" not in kwargs:
+        return kwargs
+    return {k: v for k, v in kwargs.items() if k != "trust_remote_code"}
+
+
 def modelscope_list_repo_files(
     repo_id: str,
     revision: str | None = None,

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

@@ -63,6 +63,8 @@ class FlashInferMLASparseBackend(AttentionBackend):
     supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
         "auto",
         "bfloat16",
+        "fp8",
+        "fp8_e4m3",
     ]
 
     @staticmethod
@@ -304,6 +306,11 @@ class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata
         self.bmm1_scale: float | None = None
         self.bmm2_scale: float | None = None
 
+        # fp8 query quantization is required when using fp8 kv_cache,
+        # as the TRTLLM-GEN sparse MLA kernel requires matching dtypes
+        # for query and kv_cache (mixed bf16+fp8 is not supported).
+        self.supports_quant_query_input = True
+
     def forward_mqa(
         self,
         q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

vllm/v1/attention/backends/mla/flashmla_sparse.py

@@ -83,6 +83,7 @@ class FlashMLASparseBackend(AttentionBackend):
         "auto",
         "bfloat16",
         "fp8_ds_mla",
+        "fp8",  # alias for fp8_ds_mla
     ]
 
     @staticmethod
@@ -567,6 +568,12 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
         )
         self.fp8_decode_padded_heads = self._compute_fp8_decode_padded_heads(num_heads)
 
+        if kv_cache_dtype.startswith("fp8"):
+            assert kv_cache_dtype == "fp8_ds_mla", (
+                "FlashMLA Sparse Attention backend fp8 only supports "
+                "fp8_ds_mla kv-cache dtype"
+            )
+
         if kv_cache_dtype == "fp8_ds_mla":
             # Reserve workspace during initialization
             vllm_config = get_current_vllm_config()

vllm/v1/attention/backends/rocm_aiter_fa.py

@@ -1152,11 +1152,10 @@ class AiterFlashAttentionImpl(AttentionImpl):
                 decode_max_query_len = attn_metadata.decode_metadata.max_query_len
 
                 # Use unified_attention for speculative decoding (multi-token)
-                # or when sliding window is enabled
-                if self.sliding_window[0] != -1 or decode_max_query_len > 1:
+                if decode_max_query_len > 1:
                     assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
-                        "Shuffle KV cache layout is not supported with sliding "
-                        "window or speculative decoding (multi-token decode)."
+                        "Shuffle KV cache layout is not supported with "
+                        "speculative decoding (multi-token decode)."
                     )
                     from aiter.ops.triton.unified_attention import (
                         unified_attention,

vllm/v1/cudagraph_dispatcher.py

@@ -334,8 +334,11 @@ class CudagraphDispatcher:
         for mode in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL]:
             descs = list(self.cudagraph_keys[mode])
             if descs:
-                # Sort by num_tokens descending (largest first)
-                descs.sort(key=lambda d: d.num_tokens, reverse=True)
+                # Sort by (num_tokens, num_active_loras) descending
+                descs.sort(
+                    key=lambda d: (d.num_tokens, d.num_active_loras),
+                    reverse=True,
+                )
                 result.append((mode, descs))
 
         return result

vllm/v1/engine/core_client.py

@@ -567,10 +567,7 @@ class MPClient(EngineCoreClient):
                 )
 
                 with launch_core_engines(
-                    vllm_config,
-                    executor_class,
-                    log_stats,
-                    addresses,
+                    vllm_config, executor_class, log_stats, addresses
                 ) as (engine_manager, coordinator, addresses):
                     self.resources.coordinator = coordinator
                     self.resources.engine_manager = engine_manager
@@ -638,10 +635,10 @@ class MPClient(EngineCoreClient):
 
     def shutdown(self, timeout: float | None = None) -> None:
         """Shutdown engine manager under timeout and clean up resources."""
-        self._finalizer.detach()
-        if self.resources.engine_manager is not None:
-            self.resources.engine_manager.shutdown(timeout=timeout)
-        self.resources()
+        if self._finalizer.detach() is not None:
+            if self.resources.engine_manager is not None:
+                self.resources.engine_manager.shutdown(timeout=timeout)
+            self.resources()
 
     def _format_exception(self, e: Exception) -> Exception:
         """If errored, use EngineDeadError so root cause is clear."""
@@ -685,7 +682,7 @@ class MPClient(EngineCoreClient):
             sentinels = [proc.sentinel for proc in engine_processes]
             died = multiprocessing.connection.wait(sentinels)
             _self = self_ref()
-            if not _self or _self.resources.engine_dead:
+            if not _self or not _self._finalizer.alive or _self.resources.engine_dead:
                 return
             _self.resources.engine_dead = True
             proc_name = next(

vllm/v1/engine/launch.py

@@ -1,204 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-LaunchEngineClient: A lightweight EngineClient for GPU-less online serving.
-
-This implements the EngineClient protocol without AsyncLLM or EngineCore,
-enabling preprocessing (tokenization, rendering) and postprocessing
-(detokenization) without GPU inference.
-"""
-
-from collections.abc import AsyncGenerator, Iterable, Mapping
-from typing import Any
-
-from vllm.config import VllmConfig
-from vllm.engine.protocol import EngineClient, StreamingInput
-from vllm.inputs import ProcessorInputs, PromptType
-from vllm.logger import init_logger
-from vllm.lora.request import LoRARequest
-from vllm.outputs import PoolingRequestOutput, RequestOutput
-from vllm.plugins.io_processors import get_io_processor
-from vllm.pooling_params import PoolingParams
-from vllm.renderers import renderer_from_config
-from vllm.sampling_params import SamplingParams
-from vllm.tasks import SupportedTask
-from vllm.v1.engine import EngineCoreRequest, PauseMode
-from vllm.v1.engine.input_processor import InputProcessor
-
-logger = init_logger(__name__)
-
-
-class LaunchEngineClient(EngineClient):
-    """GPU-less EngineClient that only supports preprocessing/postprocessing.
-
-    This is a Null Object at the EngineClient level, bypassing AsyncLLM
-    entirely. It initializes renderer, io_processor, and input_processor
-    for tokenization and rendering, but raises NotImplementedError for
-    any inference-related operations.
-    """
-
-    def __init__(
-        self,
-        vllm_config: VllmConfig,
-    ) -> None:
-        self.vllm_config = vllm_config
-        self.model_config = vllm_config.model_config
-
-        self.renderer = renderer = renderer_from_config(self.vllm_config)
-        self.io_processor = get_io_processor(
-            self.vllm_config,
-            self.renderer,
-            self.model_config.io_processor_plugin,
-        )
-
-        # Convert TokPrompt --> EngineCoreRequest.
-        self.input_processor = InputProcessor(self.vllm_config, renderer)
-
-    @classmethod
-    def from_vllm_config(
-        cls,
-        vllm_config: VllmConfig,
-    ) -> "LaunchEngineClient":
-        """Create a LaunchEngineClient from a VllmConfig without GPU."""
-        return cls(
-            vllm_config=vllm_config,
-        )
-
-    # -- Task support --
-
-    async def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
-        return ("render",)
-
-    # -- Inference (not supported) --
-
-    async def generate(
-        self,
-        prompt: EngineCoreRequest
-        | PromptType
-        | ProcessorInputs
-        | AsyncGenerator[StreamingInput, None],
-        sampling_params: SamplingParams,
-        request_id: str,
-        *,
-        prompt_text: str | None = None,
-        lora_request: LoRARequest | None = None,
-        tokenization_kwargs: dict[str, Any] | None = None,
-        trace_headers: Mapping[str, str] | None = None,
-        priority: int = 0,
-        data_parallel_rank: int | None = None,
-        reasoning_ended: bool | None = None,
-    ) -> AsyncGenerator[RequestOutput, None]:
-        raise NotImplementedError(
-            "LaunchEngineClient does not support inference. "
-            "Use vllm serve for generation requests."
-        )
-        # yield is needed to make this an async generator
-        yield  # type: ignore[misc] # pragma: no cover
-
-    # -- Request management (no-op) --
-
-    async def abort(
-        self, request_id: str | Iterable[str], internal: bool = False
-    ) -> None:
-        pass
-
-    # -- Generation control (no-op) --
-
-    async def pause_generation(
-        self,
-        *,
-        mode: PauseMode = "abort",
-        wait_for_inflight_requests: bool | None = None,
-        clear_cache: bool = True,
-    ) -> None:
-        pass
-
-    async def resume_generation(self) -> None:
-        pass
-
-    async def is_paused(self) -> bool:
-        return False
-
-    def shutdown(self, timeout: float | None = None) -> None:
-        pass
-
-    async def encode(
-        self,
-        prompt: PromptType | ProcessorInputs,
-        pooling_params: PoolingParams,
-        request_id: str,
-        lora_request: LoRARequest | None = None,
-        trace_headers: Mapping[str, str] | None = None,
-        priority: int = 0,
-        tokenization_kwargs: dict[str, Any] | None = None,
-        reasoning_ended: bool | None = None,
-    ) -> AsyncGenerator[PoolingRequestOutput, None]:
-        raise NotImplementedError(
-            "LaunchEngineClient does not support inference. "
-            "Use vllm serve for encoding requests."
-        )
-        yield  # type: ignore[misc] # pragma: no cover
-
-    # -- Observability (no-op / defaults) --
-
-    async def is_tracing_enabled(self) -> bool:
-        return False
-
-    async def do_log_stats(self) -> None:
-        pass
-
-    async def check_health(self) -> None:
-        pass
-
-    async def start_profile(self) -> None:
-        pass
-
-    async def stop_profile(self) -> None:
-        pass
-
-    # -- Cache management (no-op) --
-
-    async def reset_mm_cache(self) -> None:
-        pass
-
-    async def reset_prefix_cache(
-        self, reset_running_requests: bool = False, reset_connector: bool = False
-    ) -> bool:
-        return True
-
-    async def reset_encoder_cache(self) -> None:
-        pass
-
-    # -- Power management (no-op) --
-
-    async def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:
-        pass
-
-    async def wake_up(self, tags: list[str] | None = None) -> None:
-        pass
-
-    async def is_sleeping(self) -> bool:
-        return False
-
-    # -- LoRA (not supported) --
-
-    async def add_lora(self, lora_request: LoRARequest) -> bool:
-        return False
-
-    # -- Status properties --
-
-    @property
-    def is_running(self) -> bool:
-        return True
-
-    @property
-    def is_stopped(self) -> bool:
-        return False
-
-    @property
-    def errored(self) -> bool:
-        return False
-
-    @property
-    def dead_error(self) -> BaseException:
-        return RuntimeError("LaunchEngineClient does not support inference")

vllm/v1/spec_decode/ngram_proposer_gpu.py

@@ -0,0 +1,660 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+GPU-accelerated N-gram proposer using fully async PyTorch tensor operations.
+
+This version uses a fully vectorized approach with unfold and argmax for
+finding the first match across all sequences in parallel.
+"""
+
+import torch
+from torch import nn
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import (
+    CompilationConfig,
+    CompilationMode,
+    CUDAGraphMode,
+    VllmConfig,
+)
+from vllm.forward_context import set_forward_context
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.utils import record_function_or_nullcontext
+from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
+
+
+@support_torch_compile()
+class NgramGPUKernel(nn.Module):
+    """GPU-accelerated N-gram proposer using fully async tensor operations."""
+
+    def __init__(
+        self, vllm_config: VllmConfig, prefix: str = "", device: torch.device = "cuda"
+    ):
+        super().__init__()
+
+        assert vllm_config.speculative_config is not None
+        assert vllm_config.speculative_config.prompt_lookup_min is not None
+        assert vllm_config.speculative_config.prompt_lookup_max is not None
+
+        self.min_n = vllm_config.speculative_config.prompt_lookup_min
+        self.max_n = vllm_config.speculative_config.prompt_lookup_max
+        self.k = vllm_config.speculative_config.num_speculative_tokens
+        self.max_model_len = vllm_config.model_config.max_model_len
+        self.max_num_seqs = vllm_config.scheduler_config.max_num_seqs
+        self.device = device
+
+    def _find_first_and_extract_all_n_parallel(
+        self,
+        token_ids: torch.Tensor,
+        seq_lengths: torch.Tensor,
+        min_ngram_len: int,
+        max_ngram_len: int,
+        num_draft_tokens: int,
+    ) -> torch.Tensor:
+        """
+        Find suffix n-gram matches and extract following tokens.
+        Searches for the earliest prior occurrence of the trailing n-gram,
+        tries multiple lengths, and picks the longest valid match.
+
+        Args:
+            token_ids: Token IDs for each sequence
+            seq_lengths: Actual length of each sequence (excluding padding)
+            min_ngram_len: Minimum n-gram size to search for (e.g., 2)
+            max_ngram_len: Maximum n-gram size to search for (e.g., 5)
+            num_draft_tokens: Number of tokens to extract after match (k)
+
+        Returns:
+            Draft token predictions; -1 means invalid/no match.
+        """
+        batch_size = token_ids.shape[0]
+        max_seq_len = token_ids.shape[1]
+        device = token_ids.device
+        num_ngram_sizes = max_ngram_len - min_ngram_len + 1
+
+        # All n-gram sizes to try.
+        ngram_lengths = torch.arange(min_ngram_len, max_ngram_len + 1, device=device)
+        batch_indices = torch.arange(batch_size, device=device)
+
+        # Earliest match per (sequence, ngram_len); -1 means no match.
+        first_match_positions = torch.full(
+            (batch_size, num_ngram_sizes), -1, dtype=torch.long, device=device
+        )
+
+        for i, ngram_len in enumerate(range(min_ngram_len, max_ngram_len + 1)):
+            # Sliding windows of size ngram_len; unfold is O(1) view.
+            search_windows = token_ids.unfold(1, ngram_len, 1)
+            num_windows = search_windows.shape[1]
+
+            # Trailing suffix (last ngram_len tokens) for each sequence.
+            suffix_starts = seq_lengths - ngram_len
+            suffix_indices = suffix_starts.unsqueeze(1) + torch.arange(
+                ngram_len, device=device
+            )
+            suffix = torch.gather(token_ids, 1, suffix_indices.clamp(min=0))
+
+            # Window matches for each sequence.
+            matches = (search_windows == suffix.unsqueeze(1)).all(dim=-1)
+
+            # Match must leave room for at least one draft token.
+            max_valid_suffix_start = seq_lengths - ngram_len - 1
+            window_positions = torch.arange(num_windows, device=device)
+            valid_mask = window_positions <= max_valid_suffix_start.unsqueeze(1)
+            final_matches = matches & valid_mask
+
+            # Find earliest match (argmax=0 when empty; verify with has_match).
+            first_match_idx = torch.argmax(final_matches.int(), dim=1)
+            has_match = final_matches[batch_indices, first_match_idx]
+
+            # Store valid match positions (window index = position).
+            first_match_positions[:, i] = torch.where(has_match, first_match_idx, -1)
+
+        # Select the longest n-gram with a match.
+        best_ngram_idx = (first_match_positions >= 0).int().flip(dims=[1]).argmax(dim=1)
+        best_ngram_idx = num_ngram_sizes - 1 - best_ngram_idx  # Flip back
+
+        # Match position for the best n-gram.
+        best_match_pos = first_match_positions[batch_indices, best_ngram_idx]
+
+        # Avoid data-dependent branching.
+        has_any_match = best_match_pos >= 0
+
+        # Length of the best matching n-gram.
+        best_ngram_lengths = ngram_lengths[best_ngram_idx]
+
+        # Start position right after the matched suffix.
+        draft_start = torch.where(
+            has_any_match,
+            best_match_pos + best_ngram_lengths,
+            torch.zeros_like(best_match_pos),
+        )
+        tokens_available = seq_lengths - draft_start
+
+        # Gather indices for draft tokens.
+        draft_indices = draft_start.unsqueeze(1) + torch.arange(
+            num_draft_tokens, device=device
+        )
+        draft_indices = draft_indices.clamp(min=0, max=max_seq_len - 1)
+
+        # Extract draft tokens; gather always runs.
+        draft_tokens = torch.gather(token_ids, 1, draft_indices)
+
+        # Mask positions beyond available tokens.
+        position_indices = torch.arange(num_draft_tokens, device=device).unsqueeze(0)
+        valid_positions = position_indices < tokens_available.unsqueeze(1)
+
+        draft_tokens = torch.where(
+            valid_positions,
+            draft_tokens,
+            torch.full_like(draft_tokens, -1),
+        )
+
+        # If no match, mask all positions.
+        draft_tokens = torch.where(
+            has_any_match.unsqueeze(1),
+            draft_tokens,
+            torch.full_like(draft_tokens, -1),
+        )
+
+        return draft_tokens
+
+    def forward(
+        self,
+        num_tokens_no_spec: torch.Tensor,
+        token_ids_gpu: torch.Tensor,
+        combined_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Forward pass for N-gram proposal using GPU tensor operations.
+
+        Args:
+            num_tokens_no_spec: Number of tokens for each sequence [batch_size]
+            token_ids_gpu: Token IDs [batch_size, max_len]
+            combined_mask: Whether each sequence is valid for spec decode [batch_size]
+
+        Returns:
+            draft_tokens: [batch_size, k] on GPU
+            num_valid_draft_tokens: [batch_size] int32 on GPU, count of
+                leading valid (non -1) tokens per request.
+        """
+
+        device = token_ids_gpu.device
+
+        # Infer batch size to preserve dynamic shape.
+        actual_batch_size = token_ids_gpu.shape[0]
+
+        # Allocate in forward so torch.compile can optimize.
+        # NOTE(patchy): Do NOT pre-allocate this as a buffer
+        #               it breaks torch.compile
+        draft_tokens = torch.full(
+            (actual_batch_size, self.k), -1, dtype=torch.int32, device=device
+        )
+
+        results = self._find_first_and_extract_all_n_parallel(
+            token_ids_gpu,
+            num_tokens_no_spec,
+            min_ngram_len=self.min_n,
+            max_ngram_len=self.max_n,
+            num_draft_tokens=self.k,
+        )
+
+        draft_tokens = torch.where(combined_mask.unsqueeze(1), results, -1)
+
+        # Count leading contiguous valid (non -1) tokens per request.
+        is_valid = draft_tokens != -1  # [batch, k]
+        cum_valid = is_valid.int().cumsum(dim=1)  # [batch, k]
+        positions = torch.arange(1, self.k + 1, device=device).unsqueeze(0)
+        num_valid_draft_tokens = (cum_valid == positions).int().sum(dim=1)
+
+        return draft_tokens, num_valid_draft_tokens
+
+    def load_model(self, *args, **kwargs):
+        """No model to load for N-gram proposer."""
+        pass
+
+
+class NgramProposerGPU:
+    def __init__(self, vllm_config: VllmConfig, device: torch.device, runner=None):
+        assert vllm_config.speculative_config is not None
+        assert vllm_config.speculative_config.prompt_lookup_min is not None
+        assert vllm_config.speculative_config.prompt_lookup_max is not None
+
+        compilation_config = CompilationConfig(
+            mode=CompilationMode.VLLM_COMPILE,
+            custom_ops=["none"],
+            splitting_ops=[],
+            compile_sizes=[],
+            inductor_compile_config={
+                "enable_auto_functionalized_v2": False,
+                "max_autotune": True,
+                "aggressive_fusion": True,
+                "triton.autotune_pointwise": True,
+                "coordinate_descent_tuning": True,
+                "use_mixed_mm": False,
+            },
+            cudagraph_mode=CUDAGraphMode.NONE,
+        )
+        model_config = vllm_config.model_config
+        speculative_config = vllm_config.speculative_config
+        scheduler_config = vllm_config.scheduler_config
+
+        self.vllm_config = VllmConfig(
+            compilation_config=compilation_config,
+            model_config=model_config,
+            speculative_config=speculative_config,
+            scheduler_config=scheduler_config,
+        )
+
+        self.min_n = vllm_config.speculative_config.prompt_lookup_min
+        self.max_n = vllm_config.speculative_config.prompt_lookup_max
+        self.k = vllm_config.speculative_config.num_speculative_tokens
+        self.max_model_len = vllm_config.model_config.max_model_len
+        self.max_num_seqs = vllm_config.scheduler_config.max_num_seqs
+        self.device = device
+
+        self.kernel = NgramGPUKernel(
+            vllm_config=self.vllm_config, prefix="ngram_gpu_kernel", device=device
+        )
+        self.kernel.to(device)
+        self.kernel.eval()
+
+        self._dummy_run()
+
+    def _dummy_run(self):
+        token_ids, num_tokens, sampled_flags, valid_mask = self._generate_dummy_data(
+            batch_size=self.max_num_seqs,
+            max_seq_len=self.max_model_len,
+            pattern_len=self.k,
+            device=self.device,
+        )
+
+        combined_mask = sampled_flags & valid_mask & (num_tokens >= self.min_n)
+
+        for _ in range(3):
+            with set_forward_context(None, self.vllm_config):
+                _, _ = self.kernel(num_tokens, token_ids, combined_mask)
+
+    def _generate_dummy_data(
+        self,
+        batch_size: int,
+        max_seq_len: int,
+        pattern_len: int,
+        device: str = "cuda",
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Generate random test data with n-gram repetitions.
+
+        Args:
+            batch_size: Number of sequences in the batch
+            max_seq_len: Maximum sequence length
+            pattern_len: Length of patterns to inject for matching
+            device: Device to place tensors on
+
+        Returns:
+            token_ids: [batch_size, max_seq_len] tensor
+            num_tokens: [batch_size] tensor
+            sampled_flags: [batch_size] bool tensor
+            valid_mask: [batch_size] bool tensor
+        """
+        token_ids = torch.zeros(
+            batch_size,
+            max_seq_len,
+            dtype=torch.int32,
+            device=device,
+        )
+
+        num_tokens = torch.randint(
+            pattern_len, max_seq_len, (batch_size,), dtype=torch.int32, device=device
+        )
+
+        sampled_flags = torch.ones(batch_size, dtype=torch.bool, device=device)
+        valid_mask = torch.ones(batch_size, dtype=torch.bool, device=device)
+
+        return token_ids, num_tokens, sampled_flags, valid_mask
+
+    def propose(
+        self,
+        num_tokens_no_spec: torch.Tensor,  # [batch_size]
+        token_ids_gpu: torch.Tensor,  # [batch_size, max_len]
+        valid_sampled_token_ids_gpu: torch.Tensor,  # [batch_size, num_spec_tokens + 1]
+        valid_sampled_tokens_count: torch.Tensor,  # [batch_size]
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Propose draft tokens using GPU-accelerated n-gram matching.
+
+        Scatter sampled tokens into `token_ids_gpu`, compute temporary
+        updated lengths, then run the kernel.
+
+        Args:
+            num_tokens_no_spec: Number of tokens per sequence (read-only)
+            token_ids_gpu: Token IDs tensor (modified in-place with new tokens)
+            valid_sampled_token_ids_gpu: Newly sampled tokens to scatter
+            valid_sampled_tokens_count: Count of valid tokens per sequence
+
+        Returns:
+            draft_tokens: Proposed draft token IDs [batch_size, k]
+            num_valid_draft_tokens: Count of leading valid draft tokens
+                per request [batch_size]
+        """
+        assert token_ids_gpu.device == self.device
+        assert num_tokens_no_spec.device == self.device
+
+        batch_size = num_tokens_no_spec.shape[0]
+        max_seq_len = token_ids_gpu.shape[1]
+        max_new_tokens = valid_sampled_token_ids_gpu.shape[1]  # num_spec_tokens + 1
+
+        # Scatter newly sampled tokens into token_ids_gpu.
+        offsets = torch.arange(max_new_tokens, device=self.device)
+        write_positions = num_tokens_no_spec.unsqueeze(1) + offsets.unsqueeze(0)
+        valid_write_mask = offsets.unsqueeze(0) < valid_sampled_tokens_count.unsqueeze(
+            1
+        )
+        in_bounds = write_positions < max_seq_len
+        scatter_mask = (
+            valid_write_mask & (valid_sampled_token_ids_gpu != -1) & in_bounds
+        )
+
+        write_positions_long = write_positions.clamp(max=max_seq_len - 1).long()
+        existing_values = token_ids_gpu.gather(1, write_positions_long)
+
+        tokens_cast = valid_sampled_token_ids_gpu.to(token_ids_gpu.dtype)
+        tokens_to_scatter = torch.where(
+            scatter_mask,
+            tokens_cast,
+            existing_values,
+        )
+        token_ids_gpu.scatter_(1, write_positions_long, tokens_to_scatter)
+
+        num_tokens_tmp = num_tokens_no_spec + valid_sampled_tokens_count
+
+        # Compute validity masks.
+        sampled_flags = valid_sampled_tokens_count > 0
+        valid_mask = torch.ones(batch_size, dtype=torch.bool, device=self.device)
+
+        with set_forward_context(None, self.vllm_config):
+            combined_mask = sampled_flags & valid_mask & (num_tokens_tmp >= self.min_n)
+
+            with record_function_or_nullcontext("ngram_proposer_gpu: kernel"):
+                draft_tokens, num_valid_draft_tokens = self.kernel(
+                    num_tokens_tmp,
+                    token_ids_gpu,
+                    combined_mask,
+                )
+
+            return draft_tokens, num_valid_draft_tokens
+
+    def update_token_ids_ngram(
+        self,
+        sampled_token_ids: torch.Tensor | list[list[int]],
+        gpu_input_batch: InputBatch,
+        token_ids_gpu: torch.Tensor,
+        num_tokens_no_spec: torch.Tensor,
+        discard_request_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Prepare speculative decoding inputs on device:
+        compute next token ids and valid counts, honoring discarded requests
+        and rejected tokens, without CPU-GPU sync.
+        """
+        num_reqs = gpu_input_batch.num_reqs
+
+        if isinstance(sampled_token_ids, list):
+            # When disable_padded_drafter_batch=True, sampled_token_ids is
+            # an irregular list[list[int]] where sublists may have different
+            # lengths (including empty lists for discarded requests).
+            # Pad all sublists to the same length with -1 before converting
+            # to tensor.
+            max_len = max(
+                (len(sublist) for sublist in sampled_token_ids),
+                default=0,
+            )
+            # Ensure at least length 1 for tensor creation
+            max_len = max(max_len, 1)
+            padded_list = [
+                sublist + [-1] * (max_len - len(sublist))
+                for sublist in sampled_token_ids
+            ]
+            sampled_token_ids = torch.tensor(
+                padded_list, dtype=torch.int32, device=self.device
+            )
+        assert isinstance(sampled_token_ids, torch.Tensor), (
+            "sampled_token_ids should be a torch.Tensor for ngram_gpu"
+        )
+
+        # Backup last valid token before speculative tokens.
+        backup_indices = (num_tokens_no_spec[:num_reqs] - 1).clamp(min=0).long()
+        backup_next_token_ids = torch.gather(
+            token_ids_gpu[:num_reqs], dim=1, index=backup_indices.unsqueeze(1)
+        ).squeeze(1)
+
+        valid_sampled_token_ids_gpu = sampled_token_ids.clone()
+        # Invalidate sampled tokens for discarded requests.
+        discard_mask_expanded = discard_request_mask[:num_reqs].unsqueeze(1)
+        valid_sampled_token_ids_gpu.masked_fill_(discard_mask_expanded, -1)
+
+        # Mask valid tokens within each request.
+        valid_mask = (valid_sampled_token_ids_gpu != -1) & (
+            valid_sampled_token_ids_gpu < gpu_input_batch.vocab_size
+        )
+
+        # Count valid tokens per request.
+        valid_sampled_tokens_count = valid_mask.sum(dim=1)
+
+        # Rightmost valid index per row.
+        last_valid_indices = valid_sampled_tokens_count - 1
+        last_valid_indices_safe = torch.clamp(last_valid_indices, min=0)
+
+        # Last valid token from each row; undefined if none.
+        selected_tokens = torch.gather(
+            valid_sampled_token_ids_gpu, 1, last_valid_indices_safe.unsqueeze(1)
+        ).squeeze(1)
+
+        # Use last token if valid; otherwise fallback to backup.
+        next_token_ids = torch.where(
+            last_valid_indices != -1,
+            selected_tokens,
+            backup_next_token_ids,
+        )
+
+        return next_token_ids, valid_sampled_tokens_count, valid_sampled_token_ids_gpu
+
+    def load_model(self, *args, **kwargs):
+        self.kernel.load_model(*args, **kwargs)
+
+
+def update_scheduler_for_invalid_drafts(
+    num_valid_draft_tokens_event: torch.cuda.Event,
+    num_valid_draft_tokens_cpu: torch.Tensor,
+    scheduler_output: "SchedulerOutput",
+    req_id_to_index: dict[str, int],
+) -> None:
+    """Trim invalid speculative slots using per-request valid draft counts.
+
+    Args:
+        num_valid_draft_tokens_event: Event for async D2H completion.
+        num_valid_draft_tokens_cpu: CPU buffer of valid draft counts.
+        scheduler_output: Scheduler metadata to update in-place.
+        req_id_to_index: Request-id to batch-index mapping.
+    """
+    req_data = scheduler_output.scheduled_cached_reqs
+    num_valid_draft_tokens_event.synchronize()
+
+    for req_id in req_data.req_ids:
+        req_index = req_id_to_index.get(req_id)
+        if req_index is None:
+            continue
+
+        spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(req_id)
+        if spec_token_ids is None:
+            continue
+
+        scheduled_k = len(spec_token_ids)
+
+        valid_k = int(num_valid_draft_tokens_cpu[req_index].item())
+        valid_k = max(0, min(valid_k, scheduled_k))
+
+        tokens_to_trim = scheduled_k - valid_k
+        scheduler_output.total_num_scheduled_tokens -= tokens_to_trim
+        scheduler_output.num_scheduled_tokens[req_id] -= tokens_to_trim
+
+        if valid_k == 0:
+            scheduler_output.scheduled_spec_decode_tokens.pop(req_id, None)
+        else:
+            scheduler_output.scheduled_spec_decode_tokens[req_id] = spec_token_ids[
+                :valid_k
+            ]
+
+
+def update_ngram_gpu_tensors_incremental(
+    input_batch: InputBatch,
+    token_ids_gpu_tensor: torch.Tensor,
+    num_tokens_no_spec_gpu: torch.Tensor,
+    new_reqs: list[CachedRequestState],
+    device: torch.device,
+    _pinned_idx_buf: torch.Tensor,
+    _pinned_val_buf: torch.Tensor,
+) -> None:
+    """Incrementally update token_ids_gpu_tensor and num_tokens_no_spec_gpu
+    for ngram GPU proposer.
+    """
+    prev_req_id_to_index = input_batch.prev_req_id_to_index
+    curr_req_id_to_index = input_batch.req_id_to_index
+
+    if not curr_req_id_to_index:
+        return
+
+    active_indices = list(curr_req_id_to_index.values())
+    n_active = len(active_indices)
+
+    # Use resident pinned buffers to avoid per-call allocation.
+    active_idx_cpu = _pinned_idx_buf[:n_active]
+    active_idx_cpu.copy_(torch.as_tensor(active_indices, dtype=torch.long))
+
+    active_idx_gpu = active_idx_cpu.to(device=device, non_blocking=True)
+
+    new_req_ids = {req.req_id for req in new_reqs}
+
+    # First run, no previous state.
+    if prev_req_id_to_index is None:
+        for idx in active_indices:
+            num_tokens = input_batch.num_tokens_no_spec[idx]
+            if num_tokens > 0:
+                token_ids_gpu_tensor[idx, :num_tokens].copy_(
+                    input_batch.token_ids_cpu_tensor[idx, :num_tokens],
+                    non_blocking=True,
+                )
+
+        _sync_num_tokens(
+            input_batch,
+            num_tokens_no_spec_gpu,
+            active_idx_cpu,
+            active_idx_gpu,
+            n_active,
+            device,
+            _pinned_val_buf,
+        )
+        return
+
+    # Detect index changes for reorder.
+    reorder_src: list[int] = []
+    reorder_dst: list[int] = []
+
+    for req_id, curr_idx in curr_req_id_to_index.items():
+        if req_id in new_req_ids:
+            continue
+        prev_idx = prev_req_id_to_index.get(req_id)
+        if prev_idx is not None and prev_idx != curr_idx:
+            reorder_src.append(prev_idx)
+            reorder_dst.append(curr_idx)
+
+    if reorder_src:
+        src_tensor = torch.tensor(reorder_src, dtype=torch.long, device=device)
+        dst_tensor = torch.tensor(reorder_dst, dtype=torch.long, device=device)
+
+        temp_token_ids = token_ids_gpu_tensor[src_tensor].clone()
+        temp_num_tokens = num_tokens_no_spec_gpu[src_tensor].clone()
+
+        token_ids_gpu_tensor[dst_tensor] = temp_token_ids
+        num_tokens_no_spec_gpu[dst_tensor] = temp_num_tokens
+
+    # Full copy for new/resumed requests.
+    for req_state in new_reqs:
+        new_req_idx = curr_req_id_to_index.get(req_state.req_id)
+        if new_req_idx is None:
+            continue
+
+        num_tokens = input_batch.num_tokens_no_spec[new_req_idx]
+        if num_tokens > 0:
+            token_ids_gpu_tensor[new_req_idx, :num_tokens].copy_(
+                input_batch.token_ids_cpu_tensor[new_req_idx, :num_tokens],
+                non_blocking=True,
+            )
+
+    # Always batch-sync sequence lengths from CPU for ALL active requests.
+    _sync_num_tokens(
+        input_batch,
+        num_tokens_no_spec_gpu,
+        active_idx_cpu,
+        active_idx_gpu,
+        n_active,
+        device,
+        _pinned_val_buf,
+    )
+
+
+def _sync_num_tokens(
+    input_batch: InputBatch,
+    num_tokens_no_spec_gpu: torch.Tensor,
+    active_idx_cpu: torch.Tensor,
+    active_idx_gpu: torch.Tensor,
+    n_active: int,
+    device: torch.device,
+    _pinned_val_buf: torch.Tensor,
+) -> None:
+    """Batch-sync GPU sequence lengths from CPU source of truth.
+
+    Inputs:
+        input_batch: Batch container with CPU length tensor.
+        num_tokens_no_spec_gpu: Destination GPU length tensor.
+        active_idx_cpu: Active request indices on CPU.
+        active_idx_gpu: Active request indices on GPU.
+        n_active: Number of active requests.
+        device: Target CUDA device.
+        _pinned_val_buf: Resident pinned int32 staging buffer.
+    Outputs:
+        None (updates num_tokens_no_spec_gpu in-place).
+    """
+    src_cpu = input_batch.num_tokens_no_spec_cpu_tensor
+    vals = _pinned_val_buf[:n_active]
+    vals.copy_(src_cpu.index_select(0, active_idx_cpu))
+
+    num_tokens_no_spec_gpu.index_copy_(
+        0,
+        active_idx_gpu,
+        vals.to(device=device, non_blocking=True),
+    )
+
+
+def copy_num_valid_draft_tokens(
+    num_valid_draft_tokens_cpu: torch.Tensor,
+    num_valid_draft_tokens_copy_stream: torch.cuda.Stream,
+    num_valid_draft_tokens_event: torch.cuda.Event,
+    num_valid_draft_tokens: torch.Tensor | None,
+    batch_size: int,
+) -> None:
+    """
+    Async D2H copy of per-request valid draft counts.
+    """
+    if num_valid_draft_tokens is None:
+        return
+
+    num_reqs_to_copy = min(batch_size, num_valid_draft_tokens.shape[0])
+    if num_reqs_to_copy <= 0:
+        return
+
+    default_stream = torch.cuda.current_stream()
+    with torch.cuda.stream(num_valid_draft_tokens_copy_stream):
+        num_valid_draft_tokens_copy_stream.wait_stream(default_stream)
+        num_valid_draft_tokens_cpu[:num_reqs_to_copy].copy_(
+            num_valid_draft_tokens[:num_reqs_to_copy], non_blocking=True
+        )
+        num_valid_draft_tokens_event.record()

vllm/v1/worker/gpu/warmup.py

@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from collections.abc import Callable
+from typing import Any
+
 import numpy as np
 import torch
 
@@ -17,9 +20,14 @@ from vllm.v1.worker.gpu.model_runner import GPUModelRunner
 
 
 @torch.inference_mode()
-def warmup_kernels(model_runner: GPUModelRunner) -> None:
+def warmup_kernels(
+    model_runner: GPUModelRunner,
+    worker_execute_model: Callable[[SchedulerOutput], Any],
+    worker_sample_tokens: Callable[[GrammarOutput | None], Any],
+) -> None:
     """Run two execute_model + sample_tokens iterations to JIT compile
-    triton kernels.
+    triton kernels. We must call the provided worker's execute_model for
+    pipeline parallel coordination.
 
     The first iteration simulates a prefill with requests of 2 prompt
     tokens each. The second iteration simulates a decode step with all
@@ -83,7 +91,7 @@ def warmup_kernels(model_runner: GPUModelRunner) -> None:
 
     # Disable KV connector for warmup run.
     model_runner.kv_connector.set_disabled(True)
-    model_runner.execute_model(prefill_output)
+    worker_execute_model(prefill_output)
 
     if not model_runner.is_pooling_model:
         # Warm up sampler and perform a decode step for non-pooling models.
@@ -101,7 +109,7 @@ def warmup_kernels(model_runner: GPUModelRunner) -> None:
                 structured_output_request_ids=req_ids, grammar_bitmask=grammar_bitmask
             )
 
-        model_runner.sample_tokens(grammar_output)
+        worker_sample_tokens(grammar_output)
 
         # Step 2: Decode all requests with 1 token each.
         cached_req_data = CachedRequestData.make_empty()
@@ -120,12 +128,12 @@ def warmup_kernels(model_runner: GPUModelRunner) -> None:
         decode_output.total_num_scheduled_tokens = num_reqs
         decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups
 
-        model_runner.execute_model(decode_output)
-        model_runner.sample_tokens(None)
+        worker_execute_model(decode_output)
+        worker_sample_tokens(None)
 
     # Clean up - process finish_req_ids.
     cleanup_output = SchedulerOutput.make_empty()
     cleanup_output.finished_req_ids = set(req_ids)
-    model_runner.execute_model(cleanup_output)
+    worker_execute_model(cleanup_output)
     model_runner.kv_connector.set_disabled(False)
     torch.accelerator.synchronize()