add git

Signed-off-by: khluu <khluu000@gmail.com>

.buildkite/test-amd.yaml

@@ -132,7 +132,7 @@ steps:
   - tests/entrypoints/
   commands:
   - pytest -v -s entrypoints/openai/tool_parsers
-  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/openai --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/instrumentator --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling
+  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/openai --ignore=entrypoints/rpc --ignore=entrypoints/instrumentator --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling
 
 - label: Entrypoints Integration Test (LLM) # 30min
   timeout_in_minutes: 40
@@ -179,14 +179,14 @@ steps:
   torch_nightly: true
   source_file_dependencies:
   - vllm/
-  - tests/entrypoints/sleep
   - tests/entrypoints/rpc
+  - tests/entrypoints/instrumentator
   - tests/tool_use
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/sleep
+  - pytest -v -s entrypoints/instrumentator
+  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
   - pytest -v -s tool_use
-  - PYTHONPATH=/vllm-workspace  pytest -v -s entrypoints/rpc
 
 - label: Entrypoints Integration Test (Pooling)
   timeout_in_minutes: 50
@@ -1334,7 +1334,7 @@ steps:
   - pytest -v -s ./compile/test_wrapper.py
   - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
   - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
-  - pytest -v -s distributed/test_sequence_parallel.py
+  - pytest -v -s compile/correctness_e2e/test_sequence_parallel.py
   - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
   - pytest -v -s v1/worker/test_worker_memory_snapshot.py
 

.buildkite/test-pipeline.yaml

@@ -118,7 +118,7 @@ steps:
   - tests/entrypoints/
   commands:
   - pytest -v -s entrypoints/openai/tool_parsers
-  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling
+  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling
 
 - label: Entrypoints Integration Test (LLM) # 30min
   timeout_in_minutes: 40
@@ -148,7 +148,7 @@ steps:
   - tests/entrypoints/test_chat_utils
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/  --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/instrumentator --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/  --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
   - pytest -v -s entrypoints/test_chat_utils.py
 
 - label: Entrypoints Integration Test (API Server 2)
@@ -159,13 +159,13 @@ steps:
   torch_nightly: true
   source_file_dependencies:
   - vllm/
-  - tests/entrypoints/sleep
   - tests/entrypoints/rpc
+  - tests/entrypoints/instrumentator
   - tests/tool_use
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/sleep
-  - PYTHONPATH=/vllm-workspace  pytest -v -s entrypoints/rpc
+  - pytest -v -s entrypoints/instrumentator
+  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
   - pytest -v -s tool_use
 
 - label: Entrypoints Integration Test (Pooling)
@@ -862,7 +862,7 @@ steps:
   commands:
     # Install fast path packages for testing against transformers
     # Note: also needed to run plamo2 model in vLLM
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.2.5'
+    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
     - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
     # Shard hybrid language model tests
     - pytest -v -s models/language/generation \
@@ -881,7 +881,7 @@ steps:
   commands:
     # Install fast path packages for testing against transformers
     # Note: also needed to run plamo2 model in vLLM
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.2.5'
+    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
     - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
     - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
 

.buildkite/test_areas/benchmarks.yaml

@@ -17,3 +17,15 @@ steps:
   - tests/benchmarks/
   commands:
   - pytest -v -s benchmarks/
+
+- label: Attention Benchmarks Smoke Test (B200)
+  device: b200
+  num_gpus: 2
+  optional: true
+  working_dir: "/vllm-workspace/"
+  timeout_in_minutes: 10
+  source_file_dependencies:
+  - benchmarks/attention_benchmarks/
+  - vllm/v1/attention/
+  commands:
+  - python3 benchmarks/attention_benchmarks/benchmark.py --backends flash flashinfer --batch-specs "8q1s1k" --repeats 1 --warmup-iters 1

.buildkite/test_areas/compile.yaml

@@ -121,13 +121,10 @@ steps:
   optional: true
   commands:
     - nvidia-smi
-    # Run all models and attn backends but only Inductor partition and native custom ops
-    # -k "inductor_partition and not +rms_norm and not +quant_fp8"
+    # Run all models but only FLASHINFER, Inductor partition and native custom ops
     # Qwen requires +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
-    # -k "inductor_partition and not +rms_norm and +quant_fp8 and qwen3"
-    # Run just llama3 (fp8 & fp4) for all config combinations
-    # -k "llama-3"
-    - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and not +quant_fp8" -k "inductor_partition and not +rms_norm and +quant_fp8 and qwen3" -k "llama-3"
+    # Run just llama3 (fp8 & fp4) for all config combinations (only inductor partition)
+    - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and (FLASHINFER and not +rms_norm and (not +quant_fp8 or +quant_fp8 and qwen3) or llama-3)"
 
 - label: Fusion E2E TP2 Quick (H100)
   timeout_in_minutes: 20
@@ -162,7 +159,7 @@ steps:
     - tests/compile/fusions_e2e/
   commands:
     - nvidia-smi
-    # Run just llama3 (fp4 & fp8 & bf16) for all config combinations
+    # Run just llama3 (fp8 & bf16) for all config combinations
     - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "llama-3"
 
 - label: Fusion E2E TP2 AsyncTP Config Sweep (H100)
@@ -197,7 +194,8 @@ steps:
     - tests/compile/fusions_e2e/
   commands:
     - nvidia-smi
-    # Run all models and attn backends but only Inductor partition and native custom ops
+    # Run all models but only FLASHINFER, Inductor partition and native custom ops
+    # include qwen with +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
     # for ar-rms-quant-fp4, also sweep llama3
-    - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "inductor_partition and not +rms_norm and not +quant_fp8" -k "Llama-3.1-8B-Instruct-FP4"
-    - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and not +quant_fp8"
+    - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "(FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and qwen3)) or Llama-3.1-8B-Instruct-FP4"
+    - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and qwen3)"

.buildkite/test_areas/entrypoints.yaml

@@ -42,15 +42,13 @@ steps:
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
-  - tests/tool_use
-  - tests/entrypoints/sleep
-  - tests/entrypoints/instrumentator
   - tests/entrypoints/rpc
+  - tests/entrypoints/instrumentator
+  - tests/tool_use
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
   - pytest -v -s entrypoints/instrumentator
-  - pytest -v -s entrypoints/sleep
+  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
   - pytest -v -s tool_use
 
 - label: Entrypoints Integration (Pooling)

.buildkite/test_areas/models_language.yaml

@@ -40,7 +40,7 @@ steps:
   commands:
     # Install fast path packages for testing against transformers
     # Note: also needed to run plamo2 model in vLLM
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.2.5'
+    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
     - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
     # Shard hybrid language model tests
     - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
@@ -56,7 +56,7 @@ steps:
   commands:
     # Install fast path packages for testing against transformers
     # Note: also needed to run plamo2 model in vLLM
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.2.5'
+    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
     - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
     - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
 

.github/workflows/cleanup_pr_body.yml

@@ -19,6 +19,7 @@ jobs:
         uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
         with:
           python-version: '3.12'
+          cache: 'pip'
 
       - name: Install Python dependencies
         run: |

.pre-commit-config.yaml

@@ -143,6 +143,11 @@ repos:
     name: Check attention backend documentation is up to date
     entry: python tools/pre_commit/generate_attention_backend_docs.py --check
     language: python
+  - id: check-boolean-context-manager
+    name: Check for boolean ops in with-statements
+    entry: python tools/pre_commit/check_boolean_context_manager.py
+    language: python
+    types: [python]
   # Keep `suggestion` last
   - id: suggestion
     name: Suggestion

.readthedocs.yaml

@@ -9,13 +9,14 @@ build:
     python: "3.12"
   jobs:
     post_checkout:
-      - git fetch --unshallow || true
+      - git fetch origin main --unshallow --no-tags --filter=blob:none || true
+    pre_create_environment:
+      - pip install uv
+    create_environment:
+      - uv venv $READTHEDOCS_VIRTUALENV_PATH
+    install:
+      - uv pip install --python $READTHEDOCS_VIRTUALENV_PATH/bin/python --no-cache-dir -r requirements/docs.txt 
 
 mkdocs:
   configuration: mkdocs.yaml
   fail_on_warning: true
-
-# Optionally declare the Python requirements required to build your docs
-python:
-  install:
-    - requirements: requirements/docs.txt

CMakeLists.txt

@@ -293,6 +293,7 @@ set(VLLM_EXT_SRC
   "csrc/fused_qknorm_rope_kernel.cu"
   "csrc/layernorm_quant_kernels.cu"
   "csrc/sampler.cu"
+  "csrc/topk.cu"
   "csrc/cuda_view.cu"
   "csrc/quantization/gptq/q_gemm.cu"
   "csrc/quantization/w8a8/int8/scaled_quant.cu"

benchmarks/attention_benchmarks/batch_spec.py

@@ -229,3 +229,40 @@ def get_batch_stats(requests: list[BatchRequest]) -> dict:
             sum(r.kv_len for r in requests) / len(requests) if requests else 0
         ),
     }
+
+
+def get_batch_type(batch_spec: str, spec_decode_threshold: int = 8) -> str:
+    """
+    Classify a batch spec into a type string.
+
+    Args:
+        batch_spec: Batch specification string (e.g., "q2k", "8q1s1k", "2q2k_8q1s1k")
+        spec_decode_threshold: Max q_len to be considered spec-decode vs extend
+
+    Returns:
+        Type string: "prefill", "decode", "spec-decode", "extend", or "mixed (types...)"
+    """
+    requests = parse_batch_spec(batch_spec)
+
+    # Classify each request
+    types_present = set()
+    for req in requests:
+        if req.is_decode:
+            types_present.add("decode")
+        elif req.is_prefill:
+            types_present.add("prefill")
+        elif req.is_extend:
+            # Distinguish spec-decode (small q_len) from extend (chunked prefill)
+            if req.q_len <= spec_decode_threshold:
+                types_present.add("spec-decode")
+            else:
+                types_present.add("extend")
+
+    if len(types_present) == 1:
+        return types_present.pop()
+    elif len(types_present) > 1:
+        # Sort for consistent output
+        sorted_types = sorted(types_present)
+        return f"mixed ({'+'.join(sorted_types)})"
+    else:
+        return "unknown"

benchmarks/attention_benchmarks/common.py

@@ -12,6 +12,7 @@ from typing import Any
 
 import numpy as np
 import torch
+from batch_spec import get_batch_type, parse_batch_spec
 from rich.console import Console
 from rich.table import Table
 
@@ -316,12 +317,14 @@ class ResultsFormatter:
             backends: List of backend names being compared
             compare_to_fastest: Show percentage comparison to fastest
         """
-        # Group by batch spec
+        # Group by batch spec, preserving first-occurrence order
         by_spec = {}
+        specs_order = []
         for r in results:
             spec = r.config.batch_spec
             if spec not in by_spec:
                 by_spec[spec] = {}
+                specs_order.append(spec)
             by_spec[spec][r.config.backend] = r
 
         # Create shortened backend names for display
@@ -337,6 +340,8 @@ class ResultsFormatter:
 
         table = Table(title="Attention Benchmark Results")
         table.add_column("Batch\nSpec", no_wrap=True)
+        table.add_column("Type", no_wrap=True)
+        table.add_column("Batch\nSize", justify="right", no_wrap=True)
 
         multi = len(backends) > 1
         for backend in backends:
@@ -350,12 +355,14 @@ class ResultsFormatter:
                 table.add_column(col_rel, justify="right", no_wrap=False)
 
         # Add rows
-        for spec in sorted(by_spec.keys()):
+        for spec in specs_order:
             spec_results = by_spec[spec]
             times = {b: r.mean_time for b, r in spec_results.items() if r.success}
             best_time = min(times.values()) if times else 0.0
 
-            row = [spec]
+            batch_type = get_batch_type(spec)
+            batch_size = len(parse_batch_spec(spec))
+            row = [spec, batch_type, str(batch_size)]
             for backend in backends:
                 if backend in spec_results:
                     r = spec_results[backend]

benchmarks/attention_benchmarks/configs/standard_attention.yaml

@@ -25,10 +25,18 @@ batch_specs:
   - "4q1k_16q1s2k"     # 4 prefill + 16 decode
   - "2q4k_32q1s1k"     # 2 large prefill + 32 decode
 
-  # Context extension
-  - "q1ks2k"          # 1k query, 2k sequence (chunked prefill)
+  # Speculative decode (q <= 8)
+  - "16q2s1k"         # 16 requests, 2 spec tokens, 1k KV cache
+  - "16q4s1k"         # 16 requests, 4 spec tokens, 1k KV cache
+  - "16q8s1k"         # 16 requests, 8 spec tokens, 1k KV cache
+  - "32q4s2k"         # 32 requests, 4 spec tokens, 2k KV cache
+  - "8q8s4k"          # 8 requests, 8 spec tokens, 4k KV cache
+
+  # Context extension (chunked prefill)
+  - "q1ks2k"          # 1k query, 2k sequence
   - "2q1ks4k"         # 2 requests: 1k query, 4k sequence
 
+# Available backends: flash, triton, flashinfer
 backends:
   - flash
   - triton

benchmarks/attention_benchmarks/runner.py

@@ -8,7 +8,9 @@ This module provides helpers for running standard attention backends
 (FlashAttention, Triton, FlashInfer) with real vLLM integration.
 """
 
+import logging
 import types
+from contextlib import contextmanager
 
 import numpy as np
 import torch
@@ -24,8 +26,13 @@ from vllm.config import (
     ParallelConfig,
     SchedulerConfig,
     VllmConfig,
+    set_current_vllm_config,
+)
+from vllm.v1.attention.backends.utils import (
+    CommonAttentionMetadata,
+    get_kv_cache_layout,
+    set_kv_cache_layout,
 )
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
 from vllm.v1.kv_cache_interface import FullAttentionSpec
 
 # ============================================================================
@@ -37,22 +44,14 @@ _BACKEND_CONFIG = {
     "flash": {
         "module": "vllm.v1.attention.backends.flash_attn",
         "backend_class": "FlashAttentionBackend",
-        "dtype": torch.float16,
-        "cache_layout": "standard",
-        # ^ [2, num_blocks, block_size, num_kv_heads, head_dim]
     },
     "triton": {
         "module": "vllm.v1.attention.backends.triton_attn",
         "backend_class": "TritonAttentionBackend",
-        "dtype": torch.float32,
-        "cache_layout": "standard",
     },
     "flashinfer": {
         "module": "vllm.v1.attention.backends.flashinfer",
         "backend_class": "FlashInferBackend",
-        "dtype": torch.float16,
-        "cache_layout": "flashinfer",
-        # ^ [num_blocks, 2, block_size, num_kv_heads, head_dim]
     },
 }
 
@@ -66,6 +65,18 @@ def _get_backend_config(backend: str) -> dict:
     return _BACKEND_CONFIG[backend]
 
 
+@contextmanager
+def log_warnings_and_errors_only():
+    """Temporarily set vLLM logger to WARNING level."""
+    logger = logging.getLogger("vllm")
+    old_level = logger.level
+    logger.setLevel(logging.WARNING)
+    try:
+        yield
+    finally:
+        logger.setLevel(old_level)
+
+
 # ============================================================================
 # Metadata Building Helpers
 # ============================================================================
@@ -88,11 +99,7 @@ def _build_common_attn_metadata(
     query_start_loc_cpu = query_start_loc.cpu()
 
     seq_lens = torch.tensor(kv_lens, dtype=torch.int32, device=device)
-    seq_lens_cpu = seq_lens.cpu()
-    max_seq_len = int(seq_lens_cpu.max())
-
-    context_lens = [kv - q for kv, q in zip(kv_lens, q_lens)]
-    num_computed_tokens_cpu = torch.tensor(context_lens, dtype=torch.int32)
+    max_seq_len = int(seq_lens.max().item())
 
     max_blocks = (max(kv_lens) + block_size - 1) // block_size
     num_blocks = batch_size * max_blocks
@@ -107,8 +114,6 @@ def _build_common_attn_metadata(
         query_start_loc=query_start_loc,
         query_start_loc_cpu=query_start_loc_cpu,
         seq_lens=seq_lens,
-        seq_lens_cpu=seq_lens_cpu,
-        num_computed_tokens_cpu=num_computed_tokens_cpu,
         num_reqs=batch_size,
         num_actual_tokens=total_tokens,
         max_query_len=max_query_len,
@@ -121,7 +126,6 @@ def _build_common_attn_metadata(
 
 def _create_vllm_config(
     config: BenchmarkConfig,
-    dtype: torch.dtype,
     max_num_blocks: int,
 ) -> VllmConfig:
     """Create a VllmConfig for benchmarking with mock model methods."""
@@ -129,7 +133,7 @@ def _create_vllm_config(
         model="meta-llama/Meta-Llama-3-8B",
         tokenizer="meta-llama/Meta-Llama-3-8B",
         trust_remote_code=False,
-        dtype=dtype,
+        dtype="auto",  # Use model's native dtype
         seed=0,
         max_model_len=1024,
     )
@@ -198,6 +202,7 @@ def _create_backend_impl(
     backend_cfg: dict,
     config: BenchmarkConfig,
     device: torch.device,
+    dtype: torch.dtype,
 ):
     """Create backend implementation instance."""
     import importlib
@@ -206,7 +211,6 @@ def _create_backend_impl(
     backend_class = getattr(backend_module, backend_cfg["backend_class"])
 
     scale = get_attention_scale(config.head_dim)
-    dtype = backend_cfg["dtype"]
 
     impl = backend_class.get_impl_cls()(
         num_heads=config.num_q_heads,
@@ -227,7 +231,7 @@ def _create_backend_impl(
 
     layer = MockLayer(device, kv_cache_spec=kv_cache_spec)
 
-    return backend_class, impl, layer, dtype
+    return backend_class, impl, layer
 
 
 def _create_metadata_builder(
@@ -235,11 +239,44 @@ def _create_metadata_builder(
     kv_cache_spec: FullAttentionSpec,
     vllm_config: VllmConfig,
     device: torch.device,
+    backend_name: str = "",
 ):
     """Create metadata builder instance."""
-    return backend_class.get_builder_cls()(
+    layer_names = ["layer_0"]
+    builder_cls = backend_class.get_builder_cls()
+
+    # Flashinfer needs get_per_layer_parameters mocked since we don't have
+    # real model layers registered
+    if backend_name == "flashinfer":
+        import unittest.mock
+
+        from vllm.v1.attention.backends.utils import PerLayerParameters
+
+        def mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls):
+            head_size = vllm_config.model_config.get_head_size()
+            return {
+                layer_name: PerLayerParameters(
+                    window_left=-1,  # No sliding window
+                    logits_soft_cap=0.0,  # No soft cap
+                    sm_scale=1.0 / (head_size**0.5),  # Standard scale
+                )
+                for layer_name in layer_names
+            }
+
+        with unittest.mock.patch(
+            "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters",
+            mock_get_per_layer_parameters,
+        ):
+            return builder_cls(
+                kv_cache_spec=kv_cache_spec,
+                layer_names=layer_names,
+                vllm_config=vllm_config,
+                device=device,
+            )
+
+    return builder_cls(
         kv_cache_spec=kv_cache_spec,
-        layer_names=["layer_0"],
+        layer_names=layer_names,
         vllm_config=vllm_config,
         device=device,
     )
@@ -281,39 +318,44 @@ def _create_input_tensors(
 def _create_kv_cache(
     config: BenchmarkConfig,
     max_num_blocks: int,
-    cache_layout: str,
+    backend_class,
     device: torch.device,
     dtype: torch.dtype,
 ) -> list:
-    """Create KV cache tensors for all layers."""
-    if cache_layout == "flashinfer":
-        # FlashInfer layout: [num_blocks, 2, block_size, num_kv_heads, head_dim]
-        cache_list = [
-            torch.zeros(
-                max_num_blocks,
-                2,
-                config.block_size,
-                config.num_kv_heads,
-                config.head_dim,
-                device=device,
-                dtype=dtype,
-            )
-            for _ in range(config.num_layers)
-        ]
-    else:
-        # Standard layout: [2, num_blocks, block_size, num_kv_heads, head_dim]
-        cache_list = [
-            torch.zeros(
-                2,
-                max_num_blocks,
-                config.block_size,
-                config.num_kv_heads,
-                config.head_dim,
-                device=device,
-                dtype=dtype,
-            )
-            for _ in range(config.num_layers)
-        ]
+    """Create KV cache tensors for all layers using the backend's methods.
+
+    Uses the backend's get_kv_cache_shape() and get_kv_cache_stride_order()
+    to create the cache with the correct shape and memory layout.
+    """
+    # Get the logical shape from the backend
+    cache_shape = backend_class.get_kv_cache_shape(
+        num_blocks=max_num_blocks,
+        block_size=config.block_size,
+        num_kv_heads=config.num_kv_heads,
+        head_size=config.head_dim,
+    )
+
+    # Get the stride order for custom memory layout
+    try:
+        stride_order = backend_class.get_kv_cache_stride_order()
+        assert len(stride_order) == len(cache_shape)
+    except (AttributeError, NotImplementedError):
+        stride_order = tuple(range(len(cache_shape)))
+
+    # Permute shape to physical layout order
+    physical_shape = tuple(cache_shape[i] for i in stride_order)
+
+    # Compute inverse permutation to get back to logical view
+    inv_order = [stride_order.index(i) for i in range(len(stride_order))]
+
+    cache_list = []
+    for _ in range(config.num_layers):
+        # Allocate in physical layout order (contiguous in memory)
+        cache = torch.zeros(*physical_shape, device=device, dtype=dtype)
+        # Permute to logical view
+        cache = cache.permute(*inv_order)
+        cache_list.append(cache)
+
     return cache_list
 
 
@@ -418,53 +460,72 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
     kv_lens = [r.kv_len for r in requests]
     total_q = sum(q_lens)
     max_kv = max(kv_lens)
+    batch_size = len(q_lens)
 
-    max_num_blocks = (max_kv + config.block_size - 1) // config.block_size
+    # Calculate total blocks needed: batch_size * max_blocks_per_request
+    max_blocks_per_request = (max_kv + config.block_size - 1) // config.block_size
+    max_num_blocks = batch_size * max_blocks_per_request
 
-    backend_class, impl, layer, dtype = _create_backend_impl(
-        backend_cfg, config, device
-    )
+    # Suppress vLLM logs during setup to reduce spam
+    with log_warnings_and_errors_only():
+        # Create vllm_config first - uses model's native dtype via "auto"
+        vllm_config = _create_vllm_config(config, max_num_blocks)
+        dtype = vllm_config.model_config.dtype
 
-    common_metadata = _build_common_attn_metadata(
-        q_lens, kv_lens, config.block_size, device
-    )
+        # Wrap everything in set_current_vllm_config context
+        # This is required for backends like flashinfer that need global config
+        with set_current_vllm_config(vllm_config):
+            backend_class, impl, layer = _create_backend_impl(
+                backend_cfg, config, device, dtype
+            )
 
-    kv_cache_spec = FullAttentionSpec(
-        block_size=config.block_size,
-        num_kv_heads=config.num_kv_heads,
-        head_size=config.head_dim,
-        dtype=dtype,
-    )
+            # Set KV cache layout if the backend requires a specific one
+            # (e.g., FlashInfer requires HND on SM100/Blackwell for TRTLLM attention)
+            required_layout = backend_class.get_required_kv_cache_layout()
+            if required_layout is not None:
+                set_kv_cache_layout(required_layout)
+                get_kv_cache_layout.cache_clear()
 
-    vllm_config = _create_vllm_config(config, dtype, max_num_blocks)
+            common_metadata = _build_common_attn_metadata(
+                q_lens, kv_lens, config.block_size, device
+            )
 
-    builder = _create_metadata_builder(
-        backend_class, kv_cache_spec, vllm_config, device
-    )
+            kv_cache_spec = FullAttentionSpec(
+                block_size=config.block_size,
+                num_kv_heads=config.num_kv_heads,
+                head_size=config.head_dim,
+                dtype=dtype,
+            )
 
-    attn_metadata = builder.build(
-        common_prefix_len=0,
-        common_attn_metadata=common_metadata,
-    )
+            builder = _create_metadata_builder(
+                backend_class, kv_cache_spec, vllm_config, device, config.backend
+            )
 
-    q_list, k_list, v_list = _create_input_tensors(config, total_q, device, dtype)
+            attn_metadata = builder.build(
+                common_prefix_len=0,
+                common_attn_metadata=common_metadata,
+            )
 
-    cache_list = _create_kv_cache(
-        config, max_num_blocks, backend_cfg["cache_layout"], device, dtype
-    )
+            q_list, k_list, v_list = _create_input_tensors(
+                config, total_q, device, dtype
+            )
 
-    times, mem_stats = _run_single_benchmark(
-        config,
-        impl,
-        layer,
-        q_list,
-        k_list,
-        v_list,
-        cache_list,
-        attn_metadata,
-        device,
-        dtype,
-    )
+            cache_list = _create_kv_cache(
+                config, max_num_blocks, backend_class, device, dtype
+            )
+
+            times, mem_stats = _run_single_benchmark(
+                config,
+                impl,
+                layer,
+                q_list,
+                k_list,
+                v_list,
+                cache_list,
+                attn_metadata,
+                device,
+                dtype,
+            )
 
     mean_time = np.mean(times)
     throughput = total_q / mean_time if mean_time > 0 else 0

benchmarks/kernels/benchmark_fused_collective.py

@@ -5,7 +5,7 @@
 Benchmark for FlashInfer fused collective operations vs standard operations.
 
 This benchmark compares:
-1. FlashInfer's trtllm_allreduce_fusion (fused allreduce + rmsnorm + optional quant)
+1. FlashInfer's allreduce_fusion (fused allreduce + rmsnorm + optional quant)
 2. Standard tensor_model_parallel_all_reduce + separate rmsnorm/quant operations
 
 Usage with torchrun:
@@ -24,7 +24,6 @@ import torch.distributed as dist  # type: ignore
 
 from vllm.config.vllm import CompilationConfig, VllmConfig, set_current_vllm_config
 from vllm.distributed import (
-    get_tp_group,
     tensor_model_parallel_all_reduce,
 )
 from vllm.distributed.parallel_state import (
@@ -52,11 +51,12 @@ logger = init_logger(__name__)
 try:
     import flashinfer.comm as flashinfer_comm  # type: ignore
 
-    if not hasattr(flashinfer_comm, "trtllm_allreduce_fusion"):
+    if not (
+        hasattr(flashinfer_comm, "allreduce_fusion")
+        and hasattr(flashinfer_comm, "create_allreduce_fusion_workspace")
+    ):
         flashinfer_comm = None
-        logger.warning(
-            "FlashInfer comm module found but missing trtllm_allreduce_fusion"
-        )
+        logger.warning("FlashInfer comm module found but missing allreduce_fusion API")
 except ImportError:
     flashinfer_comm = None
     logger.warning("FlashInfer not found, only benchmarking standard operations")
@@ -75,7 +75,7 @@ _FI_MAX_SIZES = {
 }
 
 # Global workspace tensor for FlashInfer
-_FI_WORKSPACE_TENSOR = None
+_FI_WORKSPACE = None
 
 
 def setup_flashinfer_workspace(
@@ -83,10 +83,10 @@ def setup_flashinfer_workspace(
     rank: int,
     hidden_dim: int,
     max_token_num: int,
-    use_fp32_lamport: bool = False,
+    dtype: torch.dtype,
 ):
     """Setup FlashInfer workspace for fused allreduce operations."""
-    global _FI_WORKSPACE_TENSOR
+    global _FI_WORKSPACE
 
     if flashinfer_comm is None:
         return None, None
@@ -96,33 +96,29 @@ def setup_flashinfer_workspace(
         return None, None
 
     try:
-        # Create IPC workspace
-        ipc_handles, workspace_tensor = (
-            flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce_fusion(
-                tp_rank=rank,
-                tp_size=world_size,
-                max_token_num=max_token_num,
-                hidden_dim=hidden_dim,
-                group=get_tp_group().device_group,
-                use_fp32_lamport=use_fp32_lamport,
-            )
+        workspace = flashinfer_comm.create_allreduce_fusion_workspace(
+            backend="trtllm",
+            world_size=world_size,
+            rank=rank,
+            max_token_num=max_token_num,
+            hidden_dim=hidden_dim,
+            dtype=dtype,
         )
 
-        _FI_WORKSPACE_TENSOR = workspace_tensor
-        return ipc_handles, workspace_tensor
+        _FI_WORKSPACE = workspace
+        return workspace
     except Exception as e:
         logger.error("Failed to setup FlashInfer workspace: %s", e)
-        return None, None
+        return None
 
 
-def cleanup_flashinfer_workspace(ipc_handles):
+def cleanup_flashinfer_workspace(workspace):
     """Cleanup FlashInfer workspace."""
-    if flashinfer_comm is None or ipc_handles is None:
+    if flashinfer_comm is None or workspace is None:
         return
 
     try:
-        group = get_tp_group().device_group
-        flashinfer_comm.trtllm_destroy_ipc_workspace_for_all_reduce(ipc_handles, group)
+        workspace.destroy()
     except Exception as e:
         logger.error("Failed to cleanup FlashInfer workspace: %s", e)
 
@@ -132,25 +128,15 @@ class FlashInferFusedAllReduceParams:
 
     def __init__(
         self,
-        rank: int,
-        world_size: int,
-        use_fp32_lamport: bool = False,
         max_token_num: int = 1024,
     ):
-        self.rank = rank
-        self.world_size = world_size
-        self.use_fp32_lamport = use_fp32_lamport
-        self.trigger_completion_at_end = True
         self.launch_with_pdl = True
         self.fp32_acc = True
         self.max_token_num = max_token_num
 
     def get_trtllm_fused_allreduce_kwargs(self):
         return {
-            "world_rank": self.rank,
-            "world_size": self.world_size,
             "launch_with_pdl": self.launch_with_pdl,
-            "trigger_completion_at_end": self.trigger_completion_at_end,
             "fp32_acc": self.fp32_acc,
         }
 
@@ -165,7 +151,7 @@ def flashinfer_fused_allreduce_rmsnorm(
     norm_out: torch.Tensor | None = None,
 ):
     """FlashInfer fused allreduce + rmsnorm operation."""
-    if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None:
+    if flashinfer_comm is None or _FI_WORKSPACE is None:
         raise RuntimeError("FlashInfer not available or workspace not initialized")
 
     if norm_out is None:
@@ -174,18 +160,15 @@ def flashinfer_fused_allreduce_rmsnorm(
     else:
         residual_out = input_tensor
 
-    flashinfer_comm.trtllm_allreduce_fusion(
-        allreduce_in=input_tensor,
-        token_num=input_tensor.shape[0],
+    flashinfer_comm.allreduce_fusion(
+        input=input_tensor,
+        workspace=_FI_WORKSPACE,
+        pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
         residual_in=residual,
         residual_out=residual_out,
         norm_out=norm_out,
         rms_gamma=rms_gamma,
         rms_eps=rms_eps,
-        hidden_dim=input_tensor.shape[-1],
-        workspace_ptrs=_FI_WORKSPACE_TENSOR,
-        pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
-        allreduce_out=None,
         quant_out=None,
         scale_out=None,
         layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
@@ -207,7 +190,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
     quant_out: torch.Tensor | None = None,
 ):
     """FlashInfer fused allreduce + rmsnorm + FP8 quantization."""
-    if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None:
+    if flashinfer_comm is None or _FI_WORKSPACE is None:
         raise RuntimeError("FlashInfer not available or workspace not initialized")
 
     if norm_out is None:
@@ -216,18 +199,15 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
     else:
         residual_out = input_tensor
 
-    flashinfer_comm.trtllm_allreduce_fusion(
-        allreduce_in=input_tensor,
-        token_num=input_tensor.shape[0],
+    flashinfer_comm.allreduce_fusion(
+        input=input_tensor,
+        workspace=_FI_WORKSPACE,
+        pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant,
         residual_in=residual,
         residual_out=residual_out,
         norm_out=norm_out,
         rms_gamma=rms_gamma,
         rms_eps=rms_eps,
-        hidden_dim=input_tensor.shape[-1],
-        workspace_ptrs=_FI_WORKSPACE_TENSOR,
-        pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant,
-        allreduce_out=None,
         quant_out=quant_out,
         scale_out=None,
         layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
@@ -250,7 +230,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
     norm_out: torch.Tensor | None = None,
 ):
     """FlashInfer fused allreduce + rmsnorm + FP4 quantization."""
-    if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None:
+    if flashinfer_comm is None or _FI_WORKSPACE is None:
         raise RuntimeError("FlashInfer not available or workspace not initialized")
 
     if norm_out is None:
@@ -259,18 +239,15 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
     else:
         residual_out = input_tensor
 
-    flashinfer_comm.trtllm_allreduce_fusion(
-        allreduce_in=input_tensor,
-        token_num=input_tensor.shape[0],
+    flashinfer_comm.allreduce_fusion(
+        input=input_tensor,
+        workspace=_FI_WORKSPACE,
+        pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant,
         residual_in=residual,
         residual_out=residual_out,
         norm_out=norm_out,
         rms_gamma=rms_gamma,
         rms_eps=rms_eps,
-        hidden_dim=input_tensor.shape[-1],
-        workspace_ptrs=_FI_WORKSPACE_TENSOR,
-        pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant,
-        allreduce_out=None,
         quant_out=quant_out,
         scale_out=output_scale,
         layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
@@ -1040,23 +1017,31 @@ def main():
     configs = list(itertools.product(args.num_tokens, dtypes, residual_options))
 
     # Setup FlashInfer workspace if available
-    ipc_handles = None
+    workspace = None
     allreduce_params = None
 
     if flashinfer_comm is not None:
         # Use the largest hidden dimension for workspace setup
+        max_element_size = max(torch.finfo(dt).bits // 8 for dt in dtypes)
+        workspace_dtype = (
+            torch.float32
+            if max_element_size == 4
+            else (torch.bfloat16 if torch.bfloat16 in dtypes else torch.float16)
+        )
         max_num_token = _FI_MAX_SIZES.get(world_size) // (
-            args.hidden_dim * world_size * 2
+            args.hidden_dim * max_element_size
         )
 
-        ipc_handles, workspace_tensor = setup_flashinfer_workspace(
-            world_size, rank, args.hidden_dim, max_num_token
+        workspace = setup_flashinfer_workspace(
+            world_size,
+            rank,
+            args.hidden_dim,
+            max_num_token,
+            dtype=workspace_dtype,
         )
 
-        if workspace_tensor is not None:
+        if workspace is not None:
             allreduce_params = FlashInferFusedAllReduceParams(
-                rank=rank,
-                world_size=world_size,
                 max_token_num=max_num_token,
             )
 
@@ -1119,8 +1104,8 @@ def main():
 
     finally:
         # Cleanup
-        if ipc_handles is not None:
-            cleanup_flashinfer_workspace(ipc_handles)
+        if workspace is not None:
+            cleanup_flashinfer_workspace(workspace)
 
         dist.barrier()
 

benchmarks/kernels/benchmark_moe.py

@@ -686,6 +686,7 @@ def get_model_params(config):
         "DeepseekV2ForCausalLM",
         "DeepseekV3ForCausalLM",
         "DeepseekV32ForCausalLM",
+        "GlmMoeDsaForCausalLM",
         "Glm4MoeForCausalLM",
         "Glm4MoeLiteForCausalLM",
         "NemotronHForCausalLM",

csrc/cpu/dnnl_helper.cpp

@@ -237,12 +237,20 @@ W8A8MatMulPrimitiveHandler::W8A8MatMulPrimitiveHandler(const Args& args)
   };
   dnnl::memory::desc original_b_md({b_k_size_, b_n_size_}, b_type_,
                                    {b_k_stride_, b_n_stride_});
+#ifdef __aarch64__
+  // dummy M size for prepacking weights
+  // Prepacking weights improves performance and avoid runtime reorders
+  constexpr dnnl_dim_t kProbeM = 128;
+#else
+  constexpr dnnl_dim_t kProbeM = DNNL_RUNTIME_DIM_VAL;
+#endif
+
   prepack_weight(args.b_ptr, original_b_md,
                  create_primitive_desc(
-                     MSizeCacheKey{.a_m_size = DNNL_RUNTIME_DIM_VAL,
+                     MSizeCacheKey{.a_m_size = kProbeM,
                                    .use_bias = false,
                                    .bias_type = dnnl::memory::data_type::undef},
-                     true)
+                     /*first_time=*/true)
                      .weights_desc());
   init_runtime_memory_cache(args);
 }

csrc/cpu/mla_decode.cpp

@@ -38,16 +38,7 @@ struct KernelVecType<c10::BFloat16> {
   using qk_vec_type = vec_op::BF16Vec32;
   using v_load_vec_type = vec_op::BF16Vec16;
 };
-
-#elif defined(__s390x__)
-template <>
-struct KernelVecType<c10::BFloat16> {
-  using qk_load_vec_type = vec_op::BF16Vec16;
-  using qk_vec_type = vec_op::FP32Vec16;
-  using v_load_vec_type = vec_op::BF16Vec16;
-};
-
-#elif defined(__aarch64__)
+#else
 template <>
 struct KernelVecType<c10::BFloat16> {
   using qk_load_vec_type = vec_op::BF16Vec16;

csrc/ops.h

@@ -114,6 +114,10 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,
                           int64_t numRows, int64_t stride0, int64_t stride1,
                           int64_t topK);
 
+void large_context_topk(const torch::Tensor& score, torch::Tensor& indices,
+                        const torch::Tensor& lengths,
+                        std::optional<torch::Tensor> row_starts_opt);
+
 void rms_norm_static_fp8_quant(torch::Tensor& out, torch::Tensor& input,
                                torch::Tensor& weight, torch::Tensor& scale,
                                double epsilon);

csrc/rocm/skinny_gemms.cu

@@ -1568,7 +1568,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
         {
   #endif
           unsigned int kOff = k + (thrd * A_CHUNK);
-          unsigned int kOffcp = min__(K - A_CHUNK, k_str + kOff);
+          unsigned int kOffcp =
+              k_str + kOff;  // min__(K - A_CHUNK, k_str + kOff);
           for (unsigned int n = 0; n < N; n += CHUNKK * sprdN) {
             __builtin_amdgcn_global_load_lds(
                 (int*)(&A[min__(

csrc/sampler.cu

@@ -725,4 +725,4 @@ void top_k_per_row_prefill(const torch::Tensor& logits,
                      static_cast<int>(stride0), static_cast<int>(stride1),
                      static_cast<int>(topK), kSortingAlgorithmThreshold);
   }
-}
+}

csrc/topk.cu

@@ -0,0 +1,373 @@
+// Portions of this file are adapted from SGLang PR:
+// https://github.com/sgl-project/sglang/pull/11194
+// and
+// https://github.com/sgl-project/sglang/pull/17747
+
+#include "cuda_compat.h"
+#include "dispatch_utils.h"
+
+#include <torch/cuda.h>
+#include <c10/cuda/CUDAGuard.h>
+
+#ifndef USE_ROCM
+  #include <cub/cub.cuh>
+#else
+  #include <hipcub/hipcub.hpp>
+#endif
+
+namespace vllm {
+
+constexpr int TopK = 2048;              // DeepSeek V3 sparse attention top-k
+constexpr int kThreadsPerBlock = 1024;  // Threads per block
+
+// Shared memory budget
+#if defined(USE_ROCM)
+constexpr size_t kSmem = 48 * 1024;  // ROCm default: 48KB
+#else
+// Reduced from 128KB to 32KB to improve occupancy.
+// Each radix pass needs at most ~TopK candidates in the threshold bin,
+// so 4K entries per round (2 rounds = 8K entries = 32KB) is sufficient.
+constexpr size_t kSmem = 8 * 1024 * sizeof(uint32_t);  // 32KB (bytes)
+#endif
+
+struct FastTopKParams {
+  const float* __restrict__ input;         // [batch, seq_len] Logits
+  const int32_t* __restrict__ row_starts;  // [batch] Offset into each row
+                                           // (optional)
+  int32_t* __restrict__ indices;           // [batch, TopK] Output top-k indices
+  int32_t* __restrict__ lengths;           // [batch] Sequence lengths per row
+  int64_t input_stride;                    // Stride between rows
+};
+
+__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t {
+  uint32_t bits = __float_as_uint(x);
+  return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
+}
+
+__device__ __forceinline__ auto convert_to_uint8(float x) -> uint8_t {
+  __half h = __float2half_rn(x);
+  uint16_t bits = __half_as_ushort(h);
+  uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits)
+                                 : static_cast<uint16_t>(bits | 0x8000);
+  return static_cast<uint8_t>(key >> 8);
+}
+
+__device__ void naive_topk_cuda(const float* __restrict__ logits,
+                                int32_t* __restrict__ output_indices,
+                                int32_t seq_len) {
+  const int thread_id = threadIdx.x;
+  for (int i = thread_id; i < TopK; i += kThreadsPerBlock) {
+    output_indices[i] = (i < seq_len) ? i : -1;
+  }
+}
+
+// Adapted from:
+// https://github.com/sgl-project/sglang/blob/v0.5.8/sgl-kernel/csrc/elementwise/topk.cu#L87
+// by: DarkSharpness
+// which at the same time is an optimized topk kernel copied from tilelang
+// kernel
+__device__ void fast_topk_cuda_tl(
+    const float* __restrict__ logits,  // Input logits [seq_len]
+    int* __restrict__ output_indices,  // Output top-k indices [TopK]
+    int logits_offset,                 // Starting offset in logits array
+    int seq_len)                       // Number of valid logits to process
+{
+  constexpr int RADIX = 256;
+  constexpr int MAX_BUFFERED_ITEMS = kSmem / (2 * sizeof(int));
+
+  alignas(128) __shared__ int shared_histogram[2][RADIX + 128];
+  alignas(128) __shared__ int shared_output_count;
+  alignas(128) __shared__ int shared_threshold_bin;
+  alignas(128) __shared__ int shared_buffered_count[2];
+
+  extern __shared__ int buffered_indices[][MAX_BUFFERED_ITEMS];
+
+  const int thread_id = threadIdx.x;
+  int remaining_k = TopK;
+
+  // Pass 0: Build coarse 8-bit histogram using FP16 high bits
+  if (thread_id < RADIX + 1) {
+    shared_histogram[0][thread_id] = 0;
+  }
+  __syncthreads();
+
+  for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) {
+    const auto bin = convert_to_uint8(logits[idx + logits_offset]);
+    ::atomicAdd(&shared_histogram[0][bin], 1);
+  }
+  __syncthreads();
+
+  // Helper: Compute cumulative sum (suffix sum) over histogram using ping-pong
+  // buffers
+  auto compute_cumulative_sum = [&]() {
+    static_assert(1 << 8 == RADIX,
+                  "Radix must be 256 for 8 unrolled iterations");
+#pragma unroll 8
+    for (int i = 0; i < 8; ++i) {
+      if (C10_LIKELY(thread_id < RADIX)) {
+        const int stride = 1 << i;
+        const int src_buffer = i & 1;
+        const int dst_buffer = src_buffer ^ 1;
+
+        int value = shared_histogram[src_buffer][thread_id];
+        if (thread_id < RADIX - stride) {
+          value += shared_histogram[src_buffer][thread_id + stride];
+        }
+        shared_histogram[dst_buffer][thread_id] = value;
+      }
+      __syncthreads();
+    }
+  };
+
+  compute_cumulative_sum();
+
+  // Find threshold bin where cumsum crosses remaining_k
+  if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k &&
+      shared_histogram[0][thread_id + 1] <= remaining_k) {
+    shared_threshold_bin = thread_id;
+    shared_buffered_count[0] = 0;
+    shared_output_count = 0;
+  }
+  __syncthreads();
+
+  const int threshold_bin = shared_threshold_bin;
+  remaining_k -= shared_histogram[0][threshold_bin + 1];
+
+  // Early exit if threshold bin perfectly matches remaining_k
+  if (remaining_k == 0) {
+    for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) {
+      const int bin = convert_to_uint8(logits[idx + logits_offset]);
+      if (bin > threshold_bin) {
+        const int output_pos = ::atomicAdd(&shared_output_count, 1);
+        output_indices[output_pos] = idx;
+      }
+    }
+    __syncthreads();
+    return;
+  }
+
+  // Prepare for refinement passes: Process threshold bin
+  __syncthreads();
+  if (thread_id < RADIX + 1) {
+    shared_histogram[0][thread_id] = 0;
+  }
+  __syncthreads();
+
+  // Scan all elements and:
+  // 1. Write indices > threshold_bin to output
+  // 2. Buffer indices == threshold_bin for refinement
+  // 3. Build histogram for next refinement pass (fused optimization)
+  for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) {
+    const float logit_value = logits[idx + logits_offset];
+    const int bin = convert_to_uint8(logit_value);
+
+    if (bin > threshold_bin) {
+      // in top-k, write to output
+      const int output_pos = ::atomicAdd(&shared_output_count, 1);
+      output_indices[output_pos] = idx;
+    } else if (bin == threshold_bin) {
+      // Candidate for top-k, needs refinement
+      const int buffer_pos = ::atomicAdd(&shared_buffered_count[0], 1);
+      if (C10_LIKELY(buffer_pos < MAX_BUFFERED_ITEMS)) {
+        buffered_indices[0][buffer_pos] = idx;
+        // Fused: Build histogram for next pass
+        const uint32_t fp32_bits = convert_to_uint32_v2(logit_value);
+        const int next_bin = (fp32_bits >> 24) & 0xFF;
+        ::atomicAdd(&shared_histogram[0][next_bin], 1);
+      }
+    }
+  }
+  __syncthreads();
+
+  // ============================================================================
+  // Passes 1-4: Refine using 8-bit passes over FP32 bits
+  // ============================================================================
+  // FP32 bits [31:0] split into 4 bytes processed MSB-first:
+  // Pass 1: bits [31:24], Pass 2: bits [23:16], Pass 3: bits [15:8], Pass 4:
+  // bits [7:0]
+#pragma unroll 4
+  for (int pass = 0; pass < 4; ++pass) {
+    __shared__ int shared_final_k;  // For final pass: remaining slots to fill
+    const int src_buffer = pass % 2;
+    const int dst_buffer = src_buffer ^ 1;
+
+    // Clamp buffered count to prevent overflow
+    const int raw_buffered = shared_buffered_count[src_buffer];
+    const int num_buffered =
+        (raw_buffered < MAX_BUFFERED_ITEMS) ? raw_buffered : MAX_BUFFERED_ITEMS;
+
+    compute_cumulative_sum();
+
+    // Find threshold bin for this pass
+    if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k &&
+        shared_histogram[0][thread_id + 1] <= remaining_k) {
+      shared_threshold_bin = thread_id;
+      shared_buffered_count[dst_buffer] = 0;
+      shared_final_k = remaining_k - shared_histogram[0][thread_id + 1];
+    }
+    __syncthreads();
+
+    const int threshold_bin = shared_threshold_bin;
+    remaining_k -= shared_histogram[0][threshold_bin + 1];
+
+    // Bit offset for this pass: 24, 16, 8, 0
+    const int bit_offset = 24 - pass * 8;
+
+    // Early exit if threshold bin perfectly matches
+    if (remaining_k == 0) {
+      for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) {
+        const int idx = buffered_indices[src_buffer][i];
+        const uint32_t fp32_bits =
+            convert_to_uint32_v2(logits[idx + logits_offset]);
+        const int bin = (fp32_bits >> bit_offset) & 0xFF;
+        if (bin > threshold_bin) {
+          const int output_pos = ::atomicAdd(&shared_output_count, 1);
+          output_indices[output_pos] = idx;
+        }
+      }
+      __syncthreads();
+      break;
+    }
+
+    // Continue refinement
+    __syncthreads();
+    if (thread_id < RADIX + 1) {
+      shared_histogram[0][thread_id] = 0;
+    }
+    __syncthreads();
+
+    for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) {
+      const int idx = buffered_indices[src_buffer][i];
+      const float logit_value = logits[idx + logits_offset];
+      const uint32_t fp32_bits = convert_to_uint32_v2(logit_value);
+      const int bin = (fp32_bits >> bit_offset) & 0xFF;
+
+      if (bin > threshold_bin) {
+        // Definitely in top-k
+        const int output_pos = ::atomicAdd(&shared_output_count, 1);
+        output_indices[output_pos] = idx;
+      } else if (bin == threshold_bin) {
+        if (pass == 3) {
+          // Final pass (bits [7:0]): No more refinement possible
+          // Fill remaining slots in reverse order to maintain descending order
+          const int slot = ::atomicAdd(&shared_final_k, -1);
+          if (slot > 0) {
+            output_indices[TopK - slot] = idx;
+          }
+        } else {
+          // Buffer for next pass and build next histogram
+          const int buffer_pos =
+              ::atomicAdd(&shared_buffered_count[dst_buffer], 1);
+          if (C10_LIKELY(buffer_pos < MAX_BUFFERED_ITEMS)) {
+            buffered_indices[dst_buffer][buffer_pos] = idx;
+            // Fused: Build histogram for next pass
+            const int next_bit_offset = bit_offset - 8;
+            const int next_bin = (fp32_bits >> next_bit_offset) & 0xFF;
+            ::atomicAdd(&shared_histogram[0][next_bin], 1);
+          }
+        }
+      }
+    }
+    __syncthreads();
+  }
+}
+
+__global__ __launch_bounds__(kThreadsPerBlock) void topk_kernel(
+    const FastTopKParams params) {
+  const auto& [input, row_starts, indices, lengths, input_stride] = params;
+  const uint64_t batch_idx = blockIdx.x;
+  const int logits_offset = row_starts == nullptr ? 0 : row_starts[batch_idx];
+  const int seq_len = lengths[batch_idx];
+  int* output_indices = indices + batch_idx * TopK;
+  const float* logits = input + batch_idx * input_stride;
+
+  if (seq_len <= TopK) {
+    // Shortcut: All elements are in top-k
+    return naive_topk_cuda(logits, output_indices, seq_len);
+  } else {
+    return fast_topk_cuda_tl(logits, output_indices, logits_offset, seq_len);
+  }
+}
+
+FastTopKParams get_params(
+    const at::Tensor& score, const at::Tensor& lengths,
+    std::optional<at::Tensor> row_starts_opt = std::nullopt,
+    std::optional<at::Tensor> indices_opt = std::nullopt) {
+  const int64_t batch_size = score.size(0);
+
+  TORCH_CHECK(score.dim() == 2 && score.stride(1) == 1,
+              "score must be 2D with contiguous rows");
+  TORCH_CHECK(lengths.dim() == 1 && lengths.is_contiguous() &&
+                  lengths.size(0) == batch_size,
+              "lengths must be 1D contiguous with size matching batch");
+
+  const int32_t* row_starts_ptr = nullptr;
+  if (row_starts_opt.has_value()) {
+    const auto& row_starts = *row_starts_opt;
+    TORCH_CHECK(row_starts.dim() == 1 && row_starts.size(0) == batch_size,
+                "row_starts must be 1D with size matching batch");
+    row_starts_ptr = row_starts.data_ptr<int32_t>();
+  }
+
+  int32_t* indices_ptr = nullptr;
+  if (indices_opt.has_value()) {
+    const auto& indices = *indices_opt;
+    TORCH_CHECK(indices.dim() == 2 && indices.is_contiguous() &&
+                    indices.size(0) == batch_size && indices.size(1) == TopK,
+                "indices must be 2D contiguous [batch, TopK]");
+    indices_ptr = indices.data_ptr<int32_t>();
+  }
+
+  return FastTopKParams{
+      .input = score.data_ptr<float>(),
+      .row_starts = row_starts_ptr,
+      .indices = indices_ptr,
+      .lengths = lengths.data_ptr<int32_t>(),
+      .input_stride = score.stride(0),
+  };
+}
+
+template <auto* kernel_func, size_t smem_bytes>
+void setup_kernel_smem_once() {
+  static const cudaError_t result = []() -> cudaError_t {
+#ifdef USE_ROCM
+    auto func_ptr = reinterpret_cast<const void*>(kernel_func);
+#else
+    auto func_ptr = kernel_func;
+#endif
+    return cudaFuncSetAttribute(
+        func_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
+  }();
+
+  TORCH_CHECK(
+      result == cudaSuccess,
+      "Failed to set kernel shared memory limit: ", cudaGetErrorString(result));
+}
+
+}  // namespace vllm
+
+void large_context_topk(
+    const torch::Tensor& logits, torch::Tensor& indices,
+    const torch::Tensor& seq_lens,
+    c10::optional<torch::Tensor> row_starts = c10::nullopt) {
+  TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor");
+  TORCH_CHECK(indices.is_cuda(), "indices must be a CUDA tensor");
+  TORCH_CHECK(seq_lens.is_cuda(), "seq_lens must be a CUDA tensor");
+  if (row_starts.has_value()) {
+    TORCH_CHECK(row_starts->is_cuda(), "row_starts must be a CUDA tensor");
+  }
+
+  const auto params = vllm::get_params(logits, seq_lens, row_starts, indices);
+  const int64_t batch_size = logits.size(0);
+
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const dim3 grid(static_cast<uint32_t>(batch_size));
+  const dim3 block(vllm::kThreadsPerBlock);
+
+  vllm::setup_kernel_smem_once<vllm::topk_kernel, vllm::kSmem>();
+  vllm::topk_kernel<<<grid, block, vllm::kSmem, stream>>>(params);
+
+  const cudaError_t result = cudaGetLastError();
+  TORCH_CHECK(result == cudaSuccess,
+              "large_context_topk kernel failed: ", cudaGetErrorString(result));
+}

csrc/torch_bindings.cpp

@@ -190,6 +190,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "int numRows, int stride0, int stride1, int topK) -> ()");
   ops.impl("top_k_per_row_decode", torch::kCUDA, &top_k_per_row_decode);
 
+  ops.def(
+      "large_context_topk(Tensor score, Tensor indices, Tensor lengths, "
+      "Tensor? "
+      "row_starts_opt) -> ()");
+  ops.impl("large_context_topk", torch::kCUDA, &large_context_topk);
+
   // Layernorm-quant
   // Apply Root Mean Square (RMS) Normalization to the input tensor.
   ops.def(

docker/Dockerfile

@@ -506,6 +506,7 @@ RUN apt-get update -y \
         curl \
         sudo \
         python3-pip \
+        git \
         ffmpeg \
         libsm6 \
         libxext6 \

docker/Dockerfile.cpu

@@ -134,6 +134,7 @@ WORKDIR /vllm-workspace
 # Copy test requirements
 COPY requirements/test.in requirements/cpu-test.in
 
+# TODO: Update to 2.9.0 when there is a new build for intel_extension_for_pytorch for that version
 RUN \
     sed -i '/mamba_ssm/d' requirements/cpu-test.in && \
     remove_packages_not_supported_on_aarch64() { \

docker/Dockerfile.rocm

@@ -1,7 +1,7 @@
 # default base image
 ARG REMOTE_VLLM="0"
 ARG COMMON_WORKDIR=/app
-ARG BASE_IMAGE=rocm/vllm-dev:base_custom_releases_rocm_v0.16.0_20260211
+ARG BASE_IMAGE=rocm/vllm-dev:base
 
 # Sccache configuration (only used in release pipeline)
 ARG USE_SCCACHE

docker/Dockerfile.rocm_base

@@ -1,5 +1,5 @@
-ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2-complete
-ARG TRITON_BRANCH="57c693b6"
+ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.0-complete
+ARG TRITON_BRANCH="f332c492"
 ARG TRITON_REPO="https://github.com/ROCm/triton.git"
 ARG PYTORCH_BRANCH="89075173"
 ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git"
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
 ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
 ARG FA_BRANCH="0e60e394"
 ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
-ARG AITER_BRANCH="v0.1.10.post2"
+ARG AITER_BRANCH="6af8b687"
 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
 ARG MORI_BRANCH="2d02c6a9"
 ARG MORI_REPO="https://github.com/ROCm/mori.git"
@@ -142,7 +142,6 @@ ARG PYTORCH_VISION_REPO
 ARG PYTORCH_AUDIO_REPO
 ARG USE_SCCACHE
 
-RUN apt-get update && apt-get install -y pkg-config liblzma-dev
 RUN git clone ${PYTORCH_REPO} pytorch
 RUN cd pytorch && git checkout ${PYTORCH_BRANCH} \
     && pip install -r requirements.txt && git submodule update --init --recursive \
@@ -240,7 +239,7 @@ RUN pip install pyyaml && cd aiter \
            export HIP_CLANG_PATH=/opt/sccache-wrappers \
            && sccache --show-stats; \
        fi \
-    && GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
+    && PREBUILD_KERNELS=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
     && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
     && ls /app/aiter/dist/*.whl
 RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install

docs/benchmarking/cli.md

@@ -30,6 +30,7 @@ th {
 | HuggingFace-Other | ✅ | ✅ | `lmms-lab/LLaVA-OneVision-Data`, `Aeala/ShareGPT_Vicuna_unfiltered` |
 | HuggingFace-MTBench | ✅ | ✅ | `philschmid/mt-bench` |
 | HuggingFace-Blazedit | ✅ | ✅ | `vdaita/edit_5k_char`, `vdaita/edit_10k_char` |
+| HuggingFace-ASR | ✅ | ✅ | `openslr/librispeech_asr`, `facebook/voxpopuli`,  `LIUM/tedlium`, `edinburghcstr/ami`,        `speechcolab/gigaspeech`,        `kensho/spgispeech` |
 | Spec Bench | ✅ | ✅ | `wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl` |
 | Custom | ✅ | ✅ | Local file: `data.jsonl` |
 | Custom MM | ✅ | ✅ | Local file: `mm_data.jsonl` |
@@ -299,6 +300,22 @@ vllm bench serve \
     --blazedit-max-distance 0.99
 ```
 
+`openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech`
+
+```bash
+vllm bench serve \
+    --model openai/whisper-large-v3-turbo \
+    --backend openai-audio \
+    --dataset-name hf \
+    --dataset-path facebook/voxpopuli --hf-subset en --hf-split test --no-stream --trust-remote-code \
+    --num-prompts 99999999 \
+    --no-oversample \
+    --endpoint /v1/audio/transcriptions \
+    --ready-check-timeout-sec 600 \
+    --save-result \
+    --max-concurrency 512
+```
+
 #### Running With Sampling Parameters
 
 When using OpenAI-compatible backends such as `vllm`, optional sampling

docs/design/attention_backends.md

@@ -152,6 +152,7 @@ Priority is **1 = highest** (tried first).
 | **Sink** | Attention sink support (for StreamingLLM) |
 | **Sparse** | Sparse attention support (MLA only) |
 | **MM Prefix** | Multimodal prefix full attention support |
+| **DCP** | Decode Context Parallelism support (`--decode-context-parallel-size`) |
 | **Attention Types** | Supported attention patterns (Decoder, Encoder, Enc-Dec) |
 | **Compute Cap.** | Required CUDA compute capability (N/A for non-CUDA backends) |
 
@@ -159,20 +160,20 @@ Priority is **1 = highest** (tried first).
 
 ## Standard Attention (MHA, MQA, GQA) Backends
 
-| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | Attention Types | Compute Cap. |
-|---------|---------|--------|-----------|-------------|------------|------|-----------|-----------------|--------------|
-| `CPU_ATTN` |  | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256 | ❌ | ❌ | All | N/A |
-| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | Decoder | 7.x-9.x |
-| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | Decoder | 10.x |
-| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | All | ≥8.0 |
-| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | All | 9.x |
-| `FLASH_ATTN_DIFFKV` |  | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | Decoder | Any |
-| `FLEX_ATTENTION` |  | fp16, bf16, fp32 | `auto`, `bfloat16` | Any | Any | ❌ | ✅ | Decoder, Encoder Only | Any |
-| `ROCM_AITER_FA` |  | fp16, bf16 | `auto` | 16, 32 | 64, 128, 256 | ❌ | ❌ | Decoder | N/A |
-| `ROCM_AITER_UNIFIED_ATTN` |  | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | Decoder | N/A |
-| `ROCM_ATTN` |  | fp16, bf16, fp32 | `auto` | 16, 32, 544 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | Decoder | N/A |
-| `TREE_ATTN` |  | fp16, bf16 | `auto` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | Decoder | Any |
-| `TRITON_ATTN` |  | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | All | Any |
+| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. |
+|---------|---------|--------|-----------|-------------|------------|------|-----------|-----|-----------------|--------------|
+| `CPU_ATTN` |  | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | All | N/A |
+| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
+| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
+| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
+| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
+| `FLASH_ATTN_DIFFKV` |  | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
+| `FLEX_ATTENTION` |  | fp16, bf16, fp32 | `auto`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
+| `ROCM_AITER_FA` |  | fp16, bf16 | `auto` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
+| `ROCM_AITER_UNIFIED_ATTN` |  | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | Decoder | N/A |
+| `ROCM_ATTN` |  | fp16, bf16, fp32 | `auto` | 16, 32, 544 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
+| `TREE_ATTN` |  | fp16, bf16 | `auto` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
+| `TRITON_ATTN` |  | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
 
 > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
 >
@@ -199,14 +200,14 @@ configuration.
 
 ### Decode Backends
 
-| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | Attention Types | Compute Cap. |
-|---------|--------|-----------|-------------|------------|------|--------|-----------|-----------------|--------------|
-| `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 |
-| `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 |
-| `ROCM_AITER_MLA` | fp16, bf16 | `auto` | 1 | Any | ❌ | ❌ | ❌ | Decoder | N/A |
-| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto` | Any | 576 | ❌ | ❌ | ❌ | Decoder | N/A |
-| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | Decoder | N/A |
-| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | Decoder | Any |
+| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
+|---------|--------|-----------|-------------|------------|------|--------|-----------|-----|-----------------|--------------|
+| `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 |
+| `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 |
+| `ROCM_AITER_MLA` | fp16, bf16 | `auto` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
+| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto` | Any | 576 | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
+| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
+| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |

docs/design/moe_kernel_features.md

@@ -32,7 +32,7 @@ th {
 
 | Backend | Output act. format | Quant. types | Quant. format | Async | Apply Weight On Input | Subclass |
 |---------|--------------------|--------------|---------------|-------|-----------------------|-----------|
-| naive | standard | all<sup>1</sup> | G,A,T | N | <sup>6</sup> | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE.forward_impl] |
+| naive | standard | all<sup>1</sup> | G,A,T | N | <sup>6</sup> | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE |
 | pplx | batched | fp8,int8 | G,A,T | Y | Y | [`PplxPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.pplx_prepare_finalize.PplxPrepareAndFinalize] |
 | deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |
 | deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |

docs/features/multimodal_inputs.md

@@ -521,7 +521,7 @@ First, launch the OpenAI-compatible server:
 
 ```bash
 vllm serve microsoft/Phi-3.5-vision-instruct --runner generate \
-  --trust-remote-code --max-model-len 4096 --limit-mm-per-prompt '{"image":2}'
+  --trust-remote-code --max-model-len 4096 --limit-mm-per-prompt.image 2
 ```
 
 Then, you can use the OpenAI client as follows:

docs/getting_started/installation/gpu.xpu.inc.md

@@ -6,11 +6,10 @@ vLLM initially supports basic model inference and serving on Intel GPU platform.
 # --8<-- [start:requirements]
 
 - Supported Hardware: Intel Data Center GPU, Intel ARC GPU
-- OneAPI requirements: oneAPI 2025.3
-- Dependency: [vllm-xpu-kernels](https://github.com/vllm-project/vllm-xpu-kernels): a package provide all necessary vllm custom kernel when running vLLM on Intel GPU platform, 
+- OneAPI requirements: oneAPI 2025.1
 - Python: 3.12
 !!! warning
-    The provided vllm-xpu-kernels whl is Python3.12 specific so this version is a MUST.
+    The provided IPEX whl is Python3.12 specific so this version is a MUST.
 
 # --8<-- [end:requirements]
 # --8<-- [start:set-up-using-python]
@@ -25,7 +24,7 @@ Currently, there are no pre-built XPU wheels.
 # --8<-- [end:pre-built-wheels]
 # --8<-- [start:build-wheel-from-source]
 
-- First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers) and [Intel OneAPI](https://www.intel.com/content/www/us/en/developer/tools/oneapi/base-toolkit.html) 2025.3 or later.
+- First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers) and [Intel OneAPI](https://www.intel.com/content/www/us/en/developer/tools/oneapi/base-toolkit.html) 2025.1 or later.
 - Second, install Python packages for vLLM XPU backend building:
 
 ```bash
@@ -38,7 +37,7 @@ pip install -v -r requirements/xpu.txt
 - Then, build and install vLLM XPU backend:
 
 ```bash
-VLLM_TARGET_DEVICE=xpu pip install --no-build-isolation -e . -v
+VLLM_TARGET_DEVICE=xpu python setup.py install
 ```
 
 # --8<-- [end:build-wheel-from-source]

docs/models/supported_models.md

@@ -658,7 +658,7 @@ On the other hand, modalities separated by `/` are mutually exclusive.
 See [this page](../features/multimodal_inputs.md) on how to pass multi-modal inputs to the model.
 
 !!! tip
-    For hybrid-only models such as Llama-4, Step3 and Mistral-3, a text-only mode can be enabled by setting all supported multimodal modalities to 0 (e.g, `--limit-mm-per-prompt '{"image":0}`) so that their multimodal modules will not be loaded to free up more GPU memory for KV cache.
+    For hybrid-only models such as Llama-4, Step3, Mistral-3 and Qwen-3.5, a text-only mode can be enabled by setting all supported multimodal modalities to 0 (`--language-model-only`) so that their multimodal modules will not be loaded to free up more GPU memory for KV cache.
 
 !!! note
     vLLM currently supports adding LoRA adapters to the language backbone for most multimodal models. Additionally, vLLM now experimentally supports adding LoRA to the tower and connector modules for some multimodal models. See [this page](../features/lora.md).
@@ -738,6 +738,8 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `Qwen2VLForConditionalGeneration` | QVQ, Qwen2-VL | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/QVQ-72B-Preview`, `Qwen/Qwen2-VL-7B-Instruct`, `Qwen/Qwen2-VL-72B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `Qwen2_5_VLForConditionalGeneration` | Qwen2.5-VL | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/Qwen2.5-VL-3B-Instruct`, `Qwen/Qwen2.5-VL-72B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `Qwen2_5OmniThinkerForConditionalGeneration` | Qwen2.5-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `Qwen/Qwen2.5-Omni-3B`, `Qwen/Qwen2.5-Omni-7B` | ✅︎ | ✅︎ |
+| `Qwen3_5ForConditionalGeneration` | Qwen3.5 | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/Qwen3.5-9B-Instruct`, etc. | ✅︎ | ✅︎ |
+| `Qwen3_5MoeForConditionalGeneration` | Qwen3.5-MOE | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/Qwen3.5-35B-A3B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `Qwen3VLForConditionalGeneration` | Qwen3-VL | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/Qwen3-VL-4B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `Qwen3VLMoeForConditionalGeneration` | Qwen3-VL-MOE | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/Qwen3-VL-30B-A3B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `Qwen3OmniMoeThinkerForConditionalGeneration` | Qwen3-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `Qwen/Qwen3-Omni-30B-A3B-Instruct`, `Qwen/Qwen3-Omni-30B-A3B-Thinking` | ✅︎ | ✅︎ |

examples/offline_inference/kv_load_failure_recovery/README.md

@@ -28,3 +28,4 @@ It demonstrates vLLM's ability to recover from KV load failures in both synchron
 
 ```bash
 ./run.sh
+```

examples/offline_inference/mistral-small.py

@@ -18,11 +18,11 @@ from vllm.assets.image import ImageAsset
 # # Mistral format
 # vllm serve mistralai/Mistral-Small-3.1-24B-Instruct-2503 \
 #   --tokenizer-mode mistral --config-format mistral --load-format mistral \
-#   --limit-mm-per-prompt '{"image":4}' --max-model-len 16384
+#   --limit-mm-per-prompt.image 4 --max-model-len 16384
 #
 # # HF format
 # vllm serve mistralai/Mistral-Small-3.1-24B-Instruct-2503 \
-#   --limit-mm-per-prompt '{"image":4}' --max-model-len 16384
+#   --limit-mm-per-prompt.image 4 --max-model-len 16384
 # ```
 #
 # - Client:

examples/offline_inference/run_one_batch.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from vllm import LLM, EngineArgs
+from vllm.config import ProfilerConfig
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+DEFAULT_MAX_TOKENS = 16
+
+
+def create_parser() -> FlexibleArgumentParser:
+    parser = FlexibleArgumentParser()
+    EngineArgs.add_cli_args(parser)
+    parser.set_defaults(model="meta-llama/Llama-3.2-1B-Instruct")
+
+    batch_group = parser.add_argument_group("Batch parameters")
+    batch_group.add_argument("--batch-size", type=int, default=1)
+    batch_group.add_argument("--prompt-size", type=int, default=128)
+    batch_group.add_argument("--prompt-prefix", type=str, default="Hello, my name is")
+
+    profile_group = parser.add_argument_group("Profiling parameters")
+    profile_group.add_argument(
+        "--profile",
+        choices=["none", "prefill", "decode", "both"],
+        default="none",
+    )
+    profile_group.add_argument(
+        "--profile-dir",
+        type=str,
+        default="",
+        help="Required when --profile is not 'none'.",
+    )
+
+    return parser
+
+
+def _build_prompt(prefix: str, prompt_size: int) -> str:
+    if prompt_size <= 0:
+        return ""
+    if not prefix:
+        prefix = " "
+    if len(prefix) >= prompt_size:
+        return prefix[:prompt_size]
+    repeat_count = (prompt_size + len(prefix) - 1) // len(prefix)
+    return (prefix * repeat_count)[:prompt_size]
+
+
+def _build_profiler_config(
+    profile: str, profile_dir: str, max_tokens: int
+) -> ProfilerConfig | None:
+    if profile == "none":
+        return None
+    if not profile_dir:
+        raise ValueError("--profile-dir must be set when profiling is enabled.")
+    if profile == "prefill":
+        delay_iterations = 0
+        max_iterations = 1
+    elif profile == "decode":
+        delay_iterations = 1
+        max_iterations = max(1, max_tokens)
+    else:
+        delay_iterations = 0
+        max_iterations = 0
+
+    return ProfilerConfig(
+        profiler="torch",
+        torch_profiler_dir=profile_dir,
+        delay_iterations=delay_iterations,
+        max_iterations=max_iterations,
+    )
+
+
+def main(args: dict) -> None:
+    max_tokens = DEFAULT_MAX_TOKENS
+    batch_size = args.pop("batch_size")
+    prompt_size = args.pop("prompt_size")
+    prompt_prefix = args.pop("prompt_prefix")
+    profile = args.pop("profile")
+    profile_dir = args.pop("profile_dir")
+
+    profiler_config = _build_profiler_config(profile, profile_dir, max_tokens)
+    if profiler_config is not None:
+        args["profiler_config"] = profiler_config
+
+    llm = LLM(**args)
+
+    sampling_params = llm.get_default_sampling_params()
+    sampling_params.max_tokens = max_tokens
+    sampling_params.min_tokens = max_tokens
+    sampling_params.ignore_eos = True
+
+    prompt = _build_prompt(prompt_prefix, prompt_size)
+    prompts = [prompt] * batch_size
+
+    if profile != "none":
+        llm.start_profile()
+    outputs = llm.generate(prompts, sampling_params)
+    if profile != "none":
+        llm.stop_profile()
+
+    print("-" * 50)
+    for output in outputs:
+        generated_text = output.outputs[0].text
+        print(f"Prompt: {output.prompt!r}\nGenerated text: {generated_text!r}")
+        print("-" * 50)
+
+
+if __name__ == "__main__":
+    parser = create_parser()
+    main(vars(parser.parse_args()))

examples/offline_inference/spec_decode.py

@@ -5,14 +5,9 @@ from transformers import AutoTokenizer
 
 from vllm import LLM, SamplingParams
 from vllm.benchmarks.datasets import add_dataset_parser, get_samples
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.v1.metrics.reader import Counter, Vector
 
-try:
-    from vllm.utils.argparse_utils import FlexibleArgumentParser
-except ImportError:
-    from argparse import ArgumentParser as FlexibleArgumentParser
-
-
 QUESTION = "What is the content of each image?"
 IMAGE_URLS = [
     "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/duck.jpg",

examples/online_serving/openai_chat_completion_client_for_multimodal.py

@@ -10,7 +10,7 @@ vllm serve llava-hf/llava-1.5-7b-hf
 
 (multi-image inference with Phi-3.5-vision-instruct)
 vllm serve microsoft/Phi-3.5-vision-instruct --runner generate \
-    --trust-remote-code --max-model-len 4096 --limit-mm-per-prompt '{"image":2}'
+    --trust-remote-code --max-model-len 4096 --limit-mm-per-prompt.image 2
 
 (audio inference with Ultravox)
 vllm serve fixie-ai/ultravox-v0_5-llama-3_2-1b \

examples/pooling/classify/vision_classification_online.py

@@ -7,7 +7,7 @@ NOTE:
     vllm serve muziyongshixin/Qwen2.5-VL-7B-for-VideoCls \
          --runner pooling \
          --max-model-len 5000 \
-         --limit-mm-per-prompt '{"video": 1}' \
+         --limit-mm-per-prompt.video 1 \
          --hf-overrides '{"text_config": {"architectures": ["Qwen2_5_VLForSequenceClassification"]}}'
 """
 

requirements/build.txt

@@ -9,5 +9,5 @@ wheel
 jinja2>=3.1.6
 regex
 build
-protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.*
+protobuf
 grpcio-tools==1.78.0 # Required for grpc entrypoints

requirements/common.txt

@@ -7,9 +7,9 @@ requests >= 2.26.0
 tqdm
 blake3
 py-cpuinfo
-transformers >= 4.56.0, < 5
+transformers @ git+https://github.com/huggingface/transformers.git@main
 tokenizers >= 0.21.1  # Required for fast incremental detokenization.
-protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994
+protobuf # Required by LlamaTokenizer, gRPC.
 fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
 aiohttp >= 3.13.3
 openai >= 1.99.1  # For Responses API with reasoning content

requirements/xpu.txt

@@ -15,4 +15,4 @@ torch==2.10.0+xpu
 torchaudio
 torchvision
 
-vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.2/vllm_xpu_kernels-0.1.2-cp312-cp312-linux_x86_64.whl
+vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.1/vllm_xpu_kernels-0.1.1-cp312-cp312-linux_x86_64.whl

tests/benchmarks/test_serve_cli.py

@@ -1,15 +1,76 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import subprocess
+import tempfile
+import time
+from pathlib import Path
 
 import pytest
+import requests
+import urllib3
 
 from ..utils import RemoteOpenAIServer
 
 MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
 
 
-@pytest.fixture(scope="module")
+def generate_self_signed_cert(cert_dir: Path) -> tuple[Path, Path]:
+    """Generate a self-signed certificate for testing."""
+    cert_file = cert_dir / "cert.pem"
+    key_file = cert_dir / "key.pem"
+
+    # Generate self-signed certificate using openssl
+    subprocess.run(
+        [
+            "openssl",
+            "req",
+            "-x509",
+            "-newkey",
+            "rsa:2048",
+            "-keyout",
+            str(key_file),
+            "-out",
+            str(cert_file),
+            "-days",
+            "1",
+            "-nodes",
+            "-subj",
+            "/CN=localhost",
+        ],
+        check=True,
+        capture_output=True,
+    )
+    return cert_file, key_file
+
+
+class RemoteOpenAIServerSSL(RemoteOpenAIServer):
+    """RemoteOpenAIServer subclass that supports SSL with self-signed certs."""
+
+    @property
+    def url_root(self) -> str:
+        return f"https://{self.host}:{self.port}"
+
+    def _wait_for_server(self, *, url: str, timeout: float):
+        """Override to use HTTPS with SSL verification disabled."""
+        # Suppress InsecureRequestWarning for self-signed certs
+        urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+        start = time.time()
+        while True:
+            try:
+                if requests.get(url, verify=False).status_code == 200:
+                    break
+            except Exception:
+                result = self._poll()
+                if result is not None and result != 0:
+                    raise RuntimeError("Server exited unexpectedly.") from None
+
+                time.sleep(0.5)
+                if time.time() - start > timeout:
+                    raise RuntimeError("Server failed to start in time.") from None
+
+
+@pytest.fixture(scope="function")
 def server():
     args = ["--max-model-len", "1024", "--enforce-eager", "--load-format", "dummy"]
 
@@ -17,6 +78,27 @@ def server():
         yield remote_server
 
 
+@pytest.fixture(scope="function")
+def ssl_server():
+    """Start a vLLM server with SSL enabled using a self-signed certificate."""
+    with tempfile.TemporaryDirectory() as cert_dir:
+        cert_file, key_file = generate_self_signed_cert(Path(cert_dir))
+        args = [
+            "--max-model-len",
+            "1024",
+            "--enforce-eager",
+            "--load-format",
+            "dummy",
+            "--ssl-certfile",
+            str(cert_file),
+            "--ssl-keyfile",
+            str(key_file),
+        ]
+
+        with RemoteOpenAIServerSSL(MODEL_NAME, args) as remote_server:
+            yield remote_server
+
+
 @pytest.mark.benchmark
 def test_bench_serve(server):
     # Test default model detection and input/output len
@@ -42,6 +124,31 @@ def test_bench_serve(server):
     assert result.returncode == 0, f"Benchmark failed: {result.stderr}"
 
 
+@pytest.mark.benchmark
+def test_bench_serve_insecure(ssl_server):
+    """Test --insecure flag with an HTTPS server using a self-signed certificate."""
+    base_url = f"https://{ssl_server.host}:{ssl_server.port}"
+    command = [
+        "vllm",
+        "bench",
+        "serve",
+        "--base-url",
+        base_url,
+        "--input-len",
+        "32",
+        "--output-len",
+        "4",
+        "--num-prompts",
+        "5",
+        "--insecure",
+    ]
+    result = subprocess.run(command, capture_output=True, text=True)
+    print(result.stdout)
+    print(result.stderr)
+
+    assert result.returncode == 0, f"Benchmark failed: {result.stderr}"
+
+
 @pytest.mark.benchmark
 def test_bench_serve_chat(server):
     command = [

tests/compile/passes/distributed/test_fusion_all_reduce.py

@@ -202,9 +202,10 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
 @pytest.mark.skipif(
     not find_spec("flashinfer")
-    or not has_module_attribute("flashinfer.comm", "trtllm_allreduce_fusion"),
+    or not has_module_attribute("flashinfer.comm", "allreduce_fusion")
+    or not has_module_attribute("flashinfer.comm", "create_allreduce_fusion_workspace"),
     reason="flashinfer is not found or flashinfer "
-    "is not compiled with trtllm_allreduce_fusion",
+    "is not compiled with allreduce_fusion",
 )
 def test_all_reduce_fusion_pass_replace(
     test_model: torch.nn.Module,

tests/compile/passes/test_functionalization.py

@@ -1,11 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import copy
-
 import pytest
 import torch
 
+import vllm.envs as envs
 from tests.compile.backend import TestBackend
 from tests.utils import TestFP8Layer
 from vllm.compilation.passes.fusion.act_quant_fusion import (
@@ -32,7 +31,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
 )
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.platforms import current_platform
-from vllm.utils.torch_utils import direct_register_custom_op
 
 TEST_FP8 = current_platform.supports_fp8()
 FP8_DTYPE = current_platform.fp8_dtype()
@@ -200,82 +198,23 @@ class TestRotaryEmbeddingSliceScatter(torch.nn.Module):
         return [torch.ops.aten.slice_scatter.default]
 
 
-class TestFunctionWithMutatedArgsAndReturn(torch.nn.Module):
-    OP_REGISTERED = False
-
-    def __init__(self):
-        super().__init__()
-        self.register_test_custom_op()
-
-    @classmethod
-    def register_test_custom_op(cls):
-        if not cls.OP_REGISTERED:
-
-            def function_with_mutated_args_and_return_impl(
-                x: torch.Tensor,
-            ) -> torch.Tensor:
-                ret = x + 1
-                x.add_(2)
-                return ret
-
-            def function_with_mutated_args_and_return_fake(
-                x: torch.Tensor,
-            ) -> torch.Tensor:
-                return torch.empty_like(x)
-
-            direct_register_custom_op(
-                op_name="function_with_mutated_args_and_return",
-                op_func=function_with_mutated_args_and_return_impl,
-                mutates_args=["x"],
-                fake_impl=function_with_mutated_args_and_return_fake,
-            )
-
-            cls.OP_REGISTERED = True
-
-    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        # Clone x to avoid mutating the original tensor
-        ret = torch.ops.vllm.function_with_mutated_args_and_return(x)
-        return x, ret
-
-    def example_inputs(self, num_tokens=32):
-        hidden_states = torch.randn(num_tokens)
-        return (hidden_states,)
-
-    def ops_in_model(self, do_fusion):
-        return [torch.ops.vllm.function_with_mutated_args_and_return.default]
-
-    def ops_not_in_model(self):
-        return []
-
-
-MODELS_AND_DO_FUSION = {
-    TestSiluMul: [True, False],
-    TestFusedAddRMSNorm: [True, False],
-    TestRotaryEmbedding: [False],
-    TestRotaryEmbeddingSliceScatter: [False],
-    TestFunctionWithMutatedArgsAndReturn: [False],
-}
+MODELS = [
+    TestSiluMul,
+    TestFusedAddRMSNorm,
+    TestRotaryEmbedding,
+    TestRotaryEmbeddingSliceScatter,
+]
 
 
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.parametrize(
-    "model_class, do_fusion",
-    [
-        (model_class, do_fusion)
-        for model_class, fusions in MODELS_AND_DO_FUSION.items()
-        for do_fusion in fusions
-    ],
-)
-@pytest.mark.skipif(
-    not current_platform.is_cuda_alike(),
-    reason="Only test on cuda and rocm platform",
-)
+@pytest.mark.parametrize("model_class", MODELS)
+@pytest.mark.parametrize("do_fusion", [True, False])
+@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE != "cuda", reason="Only test on CUDA")
 def test_fix_functionalization(
     model_class: torch.nn.Module, do_fusion: bool, dtype: torch.dtype
 ):
     torch.set_default_device("cuda")
     torch.set_default_dtype(dtype)
-    torch.manual_seed(0)
 
     vllm_config = VllmConfig(
         model_config=ModelConfig(dtype=dtype),
@@ -307,14 +246,8 @@ def test_fix_functionalization(
         backend_no_func = TestBackend(*passes)
 
         model = model_class()
-        inputs_func = model.example_inputs()
-        inputs_no_func = copy.deepcopy(inputs_func)
-        model_func = model_class()
-        model_no_func = copy.deepcopy(model_func)
-        model_func = torch.compile(model_func, backend=backend_func)
-        model_no_func = torch.compile(model_no_func, backend=backend_no_func)
-        model_func(*inputs_func)
-        model_no_func(*inputs_no_func)
+        torch.compile(model, backend=backend_func)(*model.example_inputs())
+        torch.compile(model, backend=backend_no_func)(*model.example_inputs())
 
         # check if the functionalization pass is applied
         for op in model.ops_in_model(do_fusion):
@@ -332,8 +265,3 @@ def test_fix_functionalization(
                     found[op] = True
         assert all(found[op] for op in model.ops_in_model(do_fusion))
         assert all(not found.get(op) for op in model.ops_not_in_model())
-
-        # TODO (Rohan138): compare the outputs from model_func and model_no_func
-        # currently runs into errors while comparing `TestFusedAddRMSNorm`
-        # Linked issue: https://github.com/vllm-project/vllm/issues/34996
-        # torch.testing.assert_close(outputs_func, outputs_no_func)

tests/compile/passes/test_rope_kvcache_fusion.py

@@ -1,325 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-import torch
-
-import vllm.config
-from tests.compile.backend import TestBackend
-from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
-from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
-from vllm.compilation.passes.fusion.matcher_utils import ROTARY_OP
-from vllm.compilation.passes.fusion.rope_kvcache_fusion import RopeKVCacheFusionPass
-from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
-from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
-from vllm.compilation.passes.utility.scatter_split_replace import (
-    ScatterSplitReplacementPass,
-)
-from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
-from vllm.config import (
-    CacheConfig,
-    CompilationConfig,
-    CompilationMode,
-    ModelConfig,
-    PassConfig,
-    VllmConfig,
-)
-from vllm.forward_context import get_forward_context, set_forward_context
-from vllm.model_executor.layers.attention import Attention
-from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
-from vllm.platforms import current_platform
-from vllm.v1.attention.backend import (
-    AttentionBackend,
-    CommonAttentionMetadata,
-)
-from vllm.v1.attention.backends.registry import AttentionBackendEnum
-from vllm.v1.kv_cache_interface import AttentionSpec
-
-INDEX_SELECT_OP = torch.ops.aten.index.Tensor
-VLLM_UNIFIED_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_kv_cache_update
-FP8_DTYPE = current_platform.fp8_dtype()
-
-
-class QKRoPEKVCacheTestModel(torch.nn.Module):
-    def __init__(
-        self,
-        vllm_config: VllmConfig,
-        attn_backend: AttentionBackendEnum,
-        num_heads: int,
-        num_kv_heads: int,
-        head_size: int,
-        is_neox: bool,
-        dtype: torch.dtype,
-        device: torch.device,
-        prefix: str = "model.layers.0.self_attn.attn",
-    ):
-        super().__init__()
-        self.num_heads = num_heads
-        self.num_kv_heads = num_kv_heads
-        self.head_size = head_size
-        self.block_size = vllm_config.cache_config.block_size
-        self.q_size = num_heads * head_size
-        self.kv_size = num_kv_heads * head_size
-        self.is_neox = is_neox
-        self.dtype = dtype
-        self.device = device
-        self.layer_name = prefix
-
-        self.rotary_emb = RotaryEmbedding(
-            head_size,
-            rotary_dim=head_size,
-            max_position_embeddings=4096,
-            base=10000,
-            is_neox_style=is_neox,
-            dtype=self.dtype,
-        )
-
-        # Whether to check for the RoPE custom op or component index_select
-        self.enable_rope_custom_op = self.rotary_emb.enabled()
-
-        # Register layer metadata for the fusion pass via Attention.
-        self.attn = Attention(
-            num_heads=num_heads,
-            head_size=head_size,
-            scale=1.0 / head_size**0.5,
-            num_kv_heads=num_kv_heads,
-            cache_config=vllm_config.cache_config,
-            quant_config=vllm_config.quant_config,
-            prefix=prefix,
-            attn_backend=attn_backend.get_class(),
-        )
-        self.attn_backend: type[AttentionBackend] = self.attn.get_attn_backend()
-        assert not self.attn_backend.forward_includes_kv_cache_update, (
-            f"Attention backend {self.attn_backend} does not support fuse_rope_kvcache."
-        )
-        self.attn._k_scale = self.attn._k_scale.to(device)
-        self.attn._v_scale = self.attn._v_scale.to(device)
-
-        kv_cache_dtype_str = vllm_config.cache_config.cache_dtype
-        self.kv_cache_dtype = (
-            FP8_DTYPE if kv_cache_dtype_str.startswith("fp8") else self.dtype
-        )
-
-        # Initialize attn MetadataBuilder
-        self.builder = self.attn.attn_backend.get_builder_cls()(
-            kv_cache_spec=AttentionSpec(
-                block_size=self.block_size,
-                num_kv_heads=self.num_kv_heads,
-                head_size=head_size,
-                dtype=self.kv_cache_dtype,
-            ),
-            layer_names=[self.attn.layer_name],
-            vllm_config=vllm_config,
-            device=device,
-        )
-
-    def build_attn_metadata(self, batch_size: int) -> CommonAttentionMetadata:
-        """Initialize attention metadata."""
-        # Create common attn metadata
-        batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
-        common_attn_metadata = create_common_attn_metadata(
-            batch_spec, self.block_size, self.device, arange_block_indices=True
-        )
-
-        max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
-        num_blocks = batch_size * max_blocks
-
-        # Fetch the attention backend and kv cache shape and stride order
-        attn_backend = self.attn.attn_backend
-        kv_cache_shape = attn_backend.get_kv_cache_shape(
-            num_blocks, self.block_size, self.num_kv_heads, self.head_size
-        )
-        try:
-            kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
-        except (AttributeError, NotImplementedError):
-            kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
-
-        kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
-        inv_order = [
-            kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
-        ]
-
-        # Create dummy KV cache
-        raw_tensor = torch.zeros(
-            2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
-            dtype=self.kv_cache_dtype,
-            device=self.device,
-        )
-        raw_tensor = raw_tensor.view(kv_cache_shape)
-        kv_cache = raw_tensor.permute(*inv_order)
-
-        self.attn.kv_cache = [kv_cache]
-
-        # Build attn metadata
-        attn_metadata = self.builder.build(
-            common_prefix_len=0, common_attn_metadata=common_attn_metadata
-        )
-
-        return attn_metadata
-
-    def forward(
-        self, qkv: torch.Tensor, positions: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        # Create copy so inplace ops do not modify the original tensors
-        qkv = qkv.clone()
-        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-        q, k = self.rotary_emb(positions, q, k)
-
-        # Instead of a full forward pass, match only the KV cache update op here
-        q = q.view(-1, self.num_heads, self.head_size)
-        k = k.view(-1, self.num_kv_heads, self.head_size)
-        v = v.view(-1, self.num_kv_heads, self.head_size)
-        kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(
-            k, v, self.layer_name
-        )
-        return q, k, v, kv_cache_dummy_dep
-
-    def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
-        ops = []
-        if self.enable_rope_custom_op:
-            ops.append(ROTARY_OP)
-        else:
-            ops.append(INDEX_SELECT_OP)
-        ops.append(torch.ops.vllm.unified_kv_cache_update.default)
-        return ops
-
-    def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
-        return [torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default]
-
-
-@pytest.mark.parametrize(
-    "attn_backend",
-    [
-        AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
-        AttentionBackendEnum.TRITON_ATTN,
-        AttentionBackendEnum.ROCM_ATTN,
-    ],
-)
-@pytest.mark.parametrize("enable_rope_custom_op", [True])  # [True, False])
-@pytest.mark.parametrize("num_heads", [64])
-@pytest.mark.parametrize("num_kv_heads", [8])
-@pytest.mark.parametrize("head_size", [64])
-@pytest.mark.parametrize("block_size", [16])
-@pytest.mark.parametrize("is_neox", [True, False])
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
-@pytest.mark.skipif(
-    not is_aiter_found_and_supported(),
-    reason="Only test on ROCm with AITER installed and supported",
-)
-def test_rope_kvcache_fusion(
-    attn_backend: AttentionBackendEnum,
-    enable_rope_custom_op: bool,
-    num_heads: int,
-    num_kv_heads: int,
-    head_size: int,
-    block_size: int,
-    is_neox: bool,
-    dtype: torch.dtype,
-    kv_cache_dtype: str,
-    monkeypatch: pytest.MonkeyPatch,
-):
-    torch.set_default_device("cuda")
-    torch.set_default_dtype(dtype)
-    torch.manual_seed(0)
-
-    custom_ops: list[str] = []
-    if enable_rope_custom_op:
-        custom_ops.append("+rotary_embedding")
-
-    vllm_config = VllmConfig(
-        model_config=ModelConfig(dtype=dtype),
-        cache_config=CacheConfig(
-            block_size=block_size,
-            cache_dtype=kv_cache_dtype,
-        ),
-        compilation_config=CompilationConfig(
-            mode=CompilationMode.VLLM_COMPILE,
-            custom_ops=custom_ops,
-            pass_config=PassConfig(
-                fuse_rope_kvcache=True,
-                eliminate_noops=True,
-            ),
-        ),
-    )
-
-    with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
-        m.setenv("VLLM_ROCM_USE_AITER", "1")
-        rocm_aiter_ops.refresh_env_variables()
-
-        model = QKRoPEKVCacheTestModel(
-            vllm_config=vllm_config,
-            attn_backend=attn_backend,
-            num_heads=num_heads,
-            num_kv_heads=num_kv_heads,
-            head_size=head_size,
-            is_neox=is_neox,
-            dtype=dtype,
-            device=torch.get_default_device(),
-        )
-
-        fusion_pass = RopeKVCacheFusionPass(vllm_config)
-        passes = [
-            NoOpEliminationPass(vllm_config),
-            SplitCoalescingPass(vllm_config),
-            ScatterSplitReplacementPass(vllm_config),
-            fusion_pass,
-            PostCleanupPass(vllm_config),
-        ]
-        backend = TestBackend(*passes)
-
-        T = 5
-
-        qkv = torch.randn(
-            T, num_heads * head_size + 2 * num_kv_heads * head_size, dtype=dtype
-        )
-        pos = torch.arange(T, dtype=torch.long)
-
-        qkv_unfused = qkv.clone()
-        pos_unfused = pos.clone()
-
-        with set_forward_context(None, vllm_config):
-            forward_context = get_forward_context()
-            attn_metadata = model.build_attn_metadata(T)
-            forward_context.slot_mapping = {
-                model.layer_name: attn_metadata.slot_mapping
-            }
-            q_unfused, k_unfused, v_unfused, dummy = model(qkv_unfused, pos_unfused)
-            attn_layer = forward_context.no_compile_layers[model.layer_name]
-            kv_cache_unfused = attn_layer.kv_cache[forward_context.virtual_engine]
-        del dummy
-
-        torch._dynamo.mark_dynamic(qkv, 0)
-        torch._dynamo.mark_dynamic(pos, 0)
-        with set_forward_context(None, vllm_config):
-            model_fused = torch.compile(model, backend=backend)
-            forward_context = get_forward_context()
-            attn_metadata = model_fused.build_attn_metadata(T)
-            forward_context.slot_mapping = {
-                model.layer_name: attn_metadata.slot_mapping
-            }
-            q_fused, k_fused, v_fused, dummy = model_fused(qkv, pos)
-            attn_layer = forward_context.no_compile_layers[model.layer_name]
-            kv_cache_fused = attn_layer.kv_cache[forward_context.virtual_engine]
-        del dummy
-
-        assert fusion_pass.matched_count == 1
-
-        backend.check_before_ops(model.ops_in_model_before())
-        backend.check_after_ops(model.ops_in_model_after())
-
-        if dtype == torch.float16:
-            ATOL, RTOL = (2e-3, 2e-3)
-        else:
-            ATOL, RTOL = (1e-2, 1e-2)
-
-        torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL)
-        torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL)
-        torch.testing.assert_close(v_unfused, v_fused, atol=ATOL, rtol=RTOL)
-        # Cannot compare fp8_* directly here, cast to model dtype instead
-        torch.testing.assert_close(
-            kv_cache_unfused.view(dtype),
-            kv_cache_fused.view(dtype),
-            atol=ATOL,
-            rtol=RTOL,
-        )

tests/compile/passes/test_scatter_split_replace.py

@@ -1,107 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-import torch
-import torch.nn as nn
-
-import vllm
-from tests.compile.backend import TestBackend
-from vllm.compilation.passes.utility.scatter_split_replace import (
-    ScatterSplitReplacementPass,
-)
-from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
-from vllm.config import CompilationConfig, CompilationMode, VllmConfig
-from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
-
-
-class ScatterSplitReplacementModel(nn.Module):
-    """Model with a rope+getitem+slice_scatter+split_with_sizes sequence."""
-
-    def __init__(
-        self,
-        num_heads: int,
-        num_kv_heads: int,
-        head_size: int,
-        dtype: torch.dtype,
-    ):
-        super().__init__()
-        self.q_size = num_heads * head_size
-        self.kv_size = num_kv_heads * head_size
-
-        self.rotary_emb = RotaryEmbedding(
-            head_size,
-            rotary_dim=head_size,
-            max_position_embeddings=4096,
-            base=10000,
-            is_neox_style=True,
-            dtype=dtype,
-        )
-
-    def forward(self, qkv: torch.Tensor, positions: torch.Tensor):
-        # Create copy so inplace ops do not modify the original tensors
-        qkv = qkv.clone()
-        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-        q, k = self.rotary_emb(positions, q, k)
-        q = q + 1
-        k = k + 2
-        v = v + 3
-        return q, k, v
-
-    def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
-        return [
-            torch.ops.aten.slice_scatter.default,
-            torch.ops.aten.split_with_sizes.default,
-            torch.ops.aten.getitem.default,
-        ]
-
-    def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
-        return [torch.ops.aten.getitem.default]
-
-
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-def test_scatter_split_replace(dtype):
-    torch.set_default_device("cuda")
-    torch.set_default_dtype(dtype)
-    torch.manual_seed(0)
-
-    num_heads = 8
-    num_kv_heads = 4
-    head_size = 64
-
-    vllm_config = VllmConfig(
-        compilation_config=CompilationConfig(
-            mode=CompilationMode.VLLM_COMPILE,
-            custom_ops=["+rotary_embedding"],
-        ),
-    )
-    with vllm.config.set_current_vllm_config(vllm_config):
-        # ScatterSplitReplacementPass requires SplitCoalescingPass to be run before it
-        coalesce_pass = SplitCoalescingPass(vllm_config)
-        replace_pass = ScatterSplitReplacementPass(vllm_config)
-        passes = [coalesce_pass, replace_pass]
-        backend = TestBackend(*passes)
-
-        model = ScatterSplitReplacementModel(num_heads, num_kv_heads, head_size, dtype)
-
-        T = 5
-        qkv = torch.randn(
-            T, num_heads * head_size + 2 * num_kv_heads * head_size, dtype=dtype
-        )
-        pos = torch.arange(T, dtype=torch.long)
-
-        qkv_eager = qkv.clone()
-        pos_eager = pos.clone()
-        result_eager = model(qkv_eager, pos_eager)
-
-        torch._dynamo.mark_dynamic(qkv, 0)
-        torch._dynamo.mark_dynamic(pos, 0)
-
-        model_compiled = torch.compile(model, backend=backend)
-        result_compiled = model_compiled(qkv, pos)
-
-        for eager, compiled in zip(result_eager, result_compiled):
-            torch.testing.assert_close(eager, compiled)
-
-        assert backend.op_count(torch.ops.aten.slice_scatter.default) == 0
-        assert backend.op_count(torch.ops.aten.split_with_sizes.default) == 1

tests/distributed/test_eplb_execute.py

@@ -295,12 +295,11 @@ def _test_async_transfer_layer_without_mtp_worker(
     for layer_idx in range(num_layers):
         is_unchanged, is_received_locally, recv_metadata = asyncio.run(
             transfer_layer(
-                old_global_expert_indices=old_indices_cpu,
-                new_global_expert_indices=new_indices_cpu,
-                expert_weights=expert_weights,
+                old_layer_indices=old_indices_cpu[layer_idx],
+                new_layer_indices=new_indices_cpu[layer_idx],
+                expert_weights=expert_weights[layer_idx],
                 expert_weights_buffer=expert_buffer,
                 ep_group=ep_group,
-                layer=layer_idx,
                 cuda_stream=cuda_stream,
             )
         )

tests/entrypoints/openai/responses/test_parsable_context.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import importlib
+import importlib.util
 import json
 
 import pytest
@@ -179,12 +179,12 @@ async def test_mcp_tool_call(client: OpenAI, model_name: str):
     assert response.output[2].type == "reasoning"
     # make sure the correct math is in the final output
     assert response.output[3].type == "message"
-    assert "56088" in response.output[3].content[0].text
+    assert any(s in response.output[3].content[0].text for s in ("56088", "56,088"))
 
     # test raw input_messages / output_messages
     assert len(response.input_messages) == 1
     assert len(response.output_messages) == 3
-    assert "56088" in response.output_messages[2]["message"]
+    assert any(s in response.output_messages[2]["message"] for s in ("56088", "56,088"))
 
 
 @pytest.mark.asyncio

tests/entrypoints/test_context.py

@@ -8,6 +8,7 @@ from openai_harmony import Author, Message, Role, StreamState, TextContent
 
 from vllm.entrypoints.openai.responses.context import (
     HarmonyContext,
+    SimpleContext,
     StreamingHarmonyContext,
     TurnMetrics,
 )
@@ -597,3 +598,248 @@ def test_turn_metrics_copy_and_reset():
     assert copied_metrics.output_tokens == 20
     assert copied_metrics.cached_input_tokens == 5
     assert copied_metrics.tool_output_tokens == 3
+
+
+# ==================== SimpleContext Tests ====================
+
+
+def create_simple_context_output(
+    text="",
+    token_ids=None,
+    prompt="Test prompt",
+    prompt_token_ids=None,
+    num_cached_tokens=0,
+    logprobs=None,
+    finished=True,
+):
+    """Helper to create a RequestOutput with customizable text for
+    SimpleContext tests."""
+    if token_ids is None:
+        token_ids = []
+    return RequestOutput(
+        request_id="test-id",
+        prompt=prompt,
+        prompt_token_ids=prompt_token_ids,
+        prompt_logprobs=None,
+        outputs=[
+            CompletionOutput(
+                index=0,
+                text=text,
+                token_ids=token_ids,
+                cumulative_logprob=0.0,
+                logprobs=logprobs,
+                finish_reason=None,
+                stop_reason=None,
+            )
+        ],
+        finished=finished,
+        num_cached_tokens=num_cached_tokens,
+    )
+
+
+def test_simple_context_output_messages_empty():
+    """output_messages should be empty before any output is appended."""
+    context = SimpleContext()
+    assert context.output_messages == []
+
+
+def test_simple_context_output_messages_single_call():
+    """Non-streaming: single append_output produces a single output message."""
+    context = SimpleContext()
+    output = create_simple_context_output(
+        text="Hello world",
+        token_ids=[10, 20, 30],
+        prompt_token_ids=[1, 2, 3],
+    )
+    context.append_output(output)
+
+    messages = context.output_messages
+    assert len(messages) == 1
+    assert messages[0].message == "Hello world"
+    assert messages[0].tokens == [10, 20, 30]
+    assert messages[0].type == "raw_message_tokens"
+
+
+def test_simple_context_output_messages_streaming_consolidation():
+    """Streaming: multiple append_output calls consolidate into one message."""
+    context = SimpleContext()
+
+    # Simulate 3 streaming deltas
+    context.append_output(
+        create_simple_context_output(
+            text="Hello",
+            token_ids=[10],
+            prompt_token_ids=[1, 2, 3],
+        )
+    )
+    context.append_output(
+        create_simple_context_output(
+            text=" world",
+            token_ids=[20],
+            prompt_token_ids=[1, 2, 3],
+        )
+    )
+    context.append_output(
+        create_simple_context_output(
+            text="!",
+            token_ids=[30],
+            prompt_token_ids=[1, 2, 3],
+        )
+    )
+
+    messages = context.output_messages
+    assert len(messages) == 1
+    assert messages[0].message == "Hello world!"
+    assert messages[0].tokens == [10, 20, 30]
+
+
+def test_simple_context_output_messages_many_deltas():
+    """Streaming with many small deltas still produces a single message."""
+    context = SimpleContext()
+
+    words = ["The", " quick", " brown", " fox", " jumps"]
+    for i, word in enumerate(words):
+        context.append_output(
+            create_simple_context_output(
+                text=word,
+                token_ids=[100 + i],
+                prompt_token_ids=[1, 2],
+            )
+        )
+
+    messages = context.output_messages
+    assert len(messages) == 1
+    assert messages[0].message == "The quick brown fox jumps"
+    assert messages[0].tokens == [100, 101, 102, 103, 104]
+
+
+def test_simple_context_input_messages():
+    """input_messages is populated on the first append_output call."""
+    context = SimpleContext()
+    assert context.input_messages == []
+
+    context.append_output(
+        create_simple_context_output(
+            text="Hi",
+            token_ids=[10],
+            prompt="My prompt text",
+            prompt_token_ids=[1, 2, 3],
+        )
+    )
+
+    assert len(context.input_messages) == 1
+    assert context.input_messages[0].message == "My prompt text"
+    assert context.input_messages[0].tokens == [1, 2, 3]
+
+    # Second call should not add another input message
+    context.append_output(
+        create_simple_context_output(
+            text=" there",
+            token_ids=[20],
+            prompt="My prompt text",
+            prompt_token_ids=[1, 2, 3],
+        )
+    )
+
+    assert len(context.input_messages) == 1
+
+
+def test_simple_context_token_counting():
+    """Token counting accumulates across streaming deltas."""
+    context = SimpleContext()
+
+    context.append_output(
+        create_simple_context_output(
+            text="a",
+            token_ids=[10, 11],
+            prompt_token_ids=[1, 2, 3, 4, 5],
+            num_cached_tokens=2,
+        )
+    )
+    context.append_output(
+        create_simple_context_output(
+            text="b",
+            token_ids=[12],
+            prompt_token_ids=[1, 2, 3, 4, 5],
+            num_cached_tokens=2,
+        )
+    )
+
+    assert context.num_prompt_tokens == 5
+    assert context.num_output_tokens == 3  # 2 + 1
+    assert context.num_cached_tokens == 2
+
+
+def test_simple_context_final_output():
+    """final_output reconstructs accumulated text and token_ids."""
+    context = SimpleContext()
+
+    context.append_output(
+        create_simple_context_output(
+            text="foo",
+            token_ids=[1, 2],
+            prompt_token_ids=[10],
+        )
+    )
+    context.append_output(
+        create_simple_context_output(
+            text="bar",
+            token_ids=[3],
+            prompt_token_ids=[10],
+        )
+    )
+
+    final = context.final_output
+    assert final is not None
+    assert final.outputs[0].text == "foobar"
+    assert final.outputs[0].token_ids == (1, 2, 3)
+
+
+def test_simple_context_output_messages_empty_text_with_tokens():
+    """output_messages should be returned when tokens exist even if text is
+    empty (e.g. special tokens)."""
+    context = SimpleContext()
+    context.append_output(
+        create_simple_context_output(
+            text="",
+            token_ids=[99],
+            prompt_token_ids=[1],
+        )
+    )
+
+    messages = context.output_messages
+    assert len(messages) == 1
+    assert messages[0].message == ""
+    assert messages[0].tokens == [99]
+
+
+def test_simple_context_output_messages_no_mutation():
+    """Each call to output_messages returns a fresh list; callers can't
+    corrupt internal state."""
+    context = SimpleContext()
+    context.append_output(
+        create_simple_context_output(
+            text="hello",
+            token_ids=[1],
+            prompt_token_ids=[10],
+        )
+    )
+
+    msgs1 = context.output_messages
+    msgs2 = context.output_messages
+    assert msgs1 is not msgs2
+    assert msgs1[0].message == msgs2[0].message
+
+    # Appending more output updates the property
+    context.append_output(
+        create_simple_context_output(
+            text=" world",
+            token_ids=[2],
+            prompt_token_ids=[10],
+        )
+    )
+
+    msgs3 = context.output_messages
+    assert len(msgs3) == 1
+    assert msgs3[0].message == "hello world"
+    assert msgs3[0].tokens == [1, 2]

tests/kernels/moe/modular_kernel_tools/common.py

@@ -585,6 +585,7 @@ def make_modular_kernel(
         tp_size_=get_tensor_model_parallel_world_size(),
         pcp_size_=get_pcp_group().world_size,
         dp_size_=get_dp_group().world_size,
+        sp_size_=1,
         vllm_parallel_config=vllm_config.parallel_config,
     )
 
@@ -594,6 +595,7 @@ def make_modular_kernel(
         hidden_dim=config.K,
         intermediate_size_per_partition=config.N,
         num_local_experts=config.num_local_experts,
+        num_logical_experts=config.E,
         moe_parallel_config=moe_parallel_config,
         in_dtype=config.dtype,
         max_num_tokens=next_power_of_2(config.M),

tests/kernels/moe/utils.py

@@ -52,6 +52,7 @@ def make_dummy_moe_config(
         hidden_dim=hidden_dim,
         intermediate_size_per_partition=intermediate_size_per_partition,
         num_local_experts=num_experts,
+        num_logical_experts=num_experts,
         moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
         activation="silu",
         in_dtype=in_dtype,

tests/kernels/test_top_k_per_row.py

@@ -275,3 +275,114 @@ def test_top_k_per_row_decode_large_vocab_size(clean_logits: bool) -> None:
     _run_top_k_per_row_decode_test(
         top_k, batch_size, next_n, vocab_size, clean_logits, data_generation
     )
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
+@pytest.mark.parametrize("clean_logits", [True, False])
+@torch.inference_mode()
+def test_deepseek_hybrid_topk(clean_logits: bool) -> None:
+    torch.set_default_device("cuda:0")
+
+    top_k = 2048
+
+    # Test case 1: Short sequences (< 8192)
+    batch_size_short = 4
+    next_n = 1
+    num_rows_short = batch_size_short * next_n
+
+    # Create sequences with max length < 8192
+    seq_lens_short = torch.randint(
+        4000, 8000, (batch_size_short,), dtype=torch.int32, device="cuda"
+    )
+
+    row_starts_short = torch.zeros(num_rows_short, dtype=torch.int32, device="cuda")
+    row_indices_short = torch.arange(num_rows_short, device="cuda") // next_n
+    next_n_offset_short = torch.arange(num_rows_short, device="cuda") % next_n
+    row_ends_short = (
+        seq_lens_short[row_indices_short] - next_n + next_n_offset_short + 1
+    )
+
+    logits_short = create_random_logits(
+        row_starts_short, row_ends_short, torch.float32, 42, clean_logits, "random"
+    )
+
+    indices_vllm = torch.empty(
+        (num_rows_short, top_k), dtype=torch.int32, device="cuda"
+    )
+
+    # Use vllm's kernel for short sequences
+    torch.ops._C.top_k_per_row_decode(
+        logits_short,
+        next_n,
+        seq_lens_short,
+        indices_vllm,
+        num_rows_short,
+        logits_short.stride(0),
+        logits_short.stride(1),
+        top_k,
+    )
+
+    # Test case 2: Long sequences (>= 8192) - should use large_context_topk kernel
+    batch_size_long = 4
+    num_rows_long = batch_size_long * next_n
+
+    # Create sequences with max length >= 8192
+    seq_lens_long = torch.randint(
+        8192, 16384, (batch_size_long,), dtype=torch.int32, device="cuda"
+    )
+
+    row_starts_long = torch.zeros(num_rows_long, dtype=torch.int32, device="cuda")
+    row_indices_long = torch.arange(num_rows_long, device="cuda") // next_n
+    next_n_offset_long = torch.arange(num_rows_long, device="cuda") % next_n
+    row_ends_long = seq_lens_long[row_indices_long] - next_n + next_n_offset_long + 1
+
+    logits_long = create_random_logits(
+        row_starts_long, row_ends_long, torch.float32, 43, clean_logits, "random"
+    )
+
+    indices = torch.empty((num_rows_long, top_k), dtype=torch.int32, device="cuda")
+
+    # Use large_context_topk kernel for long sequences
+    if next_n == 1:
+        lengths = seq_lens_long
+    else:
+        offsets = torch.arange(next_n, device=logits_long.device, dtype=torch.int32)
+        lengths = (seq_lens_long.unsqueeze(1) - next_n + 1 + offsets).flatten()
+
+    torch.ops._C.large_context_topk(
+        logits_long,
+        indices,
+        lengths,
+        None,
+    )
+
+    torch_indices_short = torch.empty(
+        (num_rows_short, top_k), dtype=torch.int32, device="cuda"
+    )
+    for i in range(num_rows_short):
+        row_end = int(row_ends_short[i])
+        k_i = min(top_k, row_end)
+        idx = logits_short[i, :row_end].topk(k_i, dim=-1)[1]
+        torch_indices_short[i, :k_i] = idx
+
+    assert compare_top_k_results(
+        logits_short,
+        indices_vllm,
+        torch_indices_short,
+        row_starts_short,
+        row_ends_short,
+        top_k,
+    ), "top_k_per_row_decode kernel (short sequences) doesn't match torch.topk"
+
+    torch_indices_long = torch.empty(
+        (num_rows_long, top_k), dtype=torch.int32, device="cuda"
+    )
+    for i in range(num_rows_long):
+        row_end = int(row_ends_long[i])
+        k_i = min(top_k, row_end)
+        idx = logits_long[i, :row_end].topk(k_i, dim=-1)[1]
+        torch_indices_long[i, :k_i] = idx
+
+    assert compare_top_k_results(
+        logits_long, indices, torch_indices_long, row_starts_long, row_ends_long, top_k
+    ), "large_context_topk kernel (long sequences) doesn't match torch.topk"

tests/models/language/generation/test_hybrid.py

@@ -7,6 +7,7 @@ import pytest
 
 from tests.models.registry import HF_EXAMPLE_MODELS
 from tests.utils import multi_gpu_test
+from vllm import LLM
 from vllm.engine.arg_utils import EngineArgs
 from vllm.platforms import current_platform
 from vllm.sampling_params import SamplingParams
@@ -769,3 +770,30 @@ def test_apc_multiple_prompts_partial_cached_outputs(
             name_0="vllm_no_cache",
             name_1=f"vllm_cache_it_{r_idx + 1}",
         )
+
+
+# we have to use a real large model to get reasonable results
+# the model can't be a hybrid model as we need block_size 16
+@pytest.mark.parametrize("model", ["tiiuae/falcon-mamba-7b"])
+def test_apc_common_prefix_same_batch(
+    model: str,
+    monkeypatch,
+) -> None:
+    # Required to put the two requests in the same batch
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    llm = LLM(
+        model=model,
+        enforce_eager=True,
+        block_size=16,
+        mamba_block_size=16,
+        enable_prefix_caching=True,
+        seed=42,
+    )
+    prompts = [
+        "hello what is one plus one what is one plus one what is one plus one the answer is",  # noqa: E501
+        "hello what is one plus one what is one plus one what is one plus one the answer is",  # noqa: E501
+    ]
+    sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=20)
+    outputs = llm.generate(prompts, sampling_params)
+    for output in outputs:
+        assert "two" in output.outputs[0].text

tests/models/multimodal/generation/test_common.py

@@ -377,7 +377,7 @@ VLM_TEST_SETTINGS = {
         use_tokenizer_eos=True,
         vllm_output_post_proc=model_utils.fuyu_vllm_to_hf_output,
         num_logprobs=10,
-        image_size_factors=[(), (0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
+        image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
         marks=[large_gpu_mark(min_gb=32)],
     ),
     "gemma3": VLMTestInfo(
@@ -437,7 +437,7 @@ VLM_TEST_SETTINGS = {
         max_num_seqs=2,
         get_stop_token_ids=lambda tok: [151329, 151336, 151338],
         num_logprobs=10,
-        image_size_factors=[(), (0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
+        image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
         auto_cls=AutoModelForImageTextToText,
         marks=[large_gpu_mark(min_gb=32)],
     ),
@@ -468,7 +468,7 @@ VLM_TEST_SETTINGS = {
         max_num_seqs=2,
         get_stop_token_ids=lambda tok: [151329, 151336, 151338],
         num_logprobs=10,
-        image_size_factors=[(), (0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
+        image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
         auto_cls=AutoModelForImageTextToText,
         marks=[large_gpu_mark(min_gb=32)],
     ),

tests/models/registry.py

@@ -275,6 +275,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
         "zai-org/GLM-4.7-Flash",
         min_transformers_version="5.0.0",
     ),
+    "GlmMoeDsaForCausalLM": _HfExamplesInfo(
+        "zai-org/GLM-5", min_transformers_version="5.0.1", is_available_online=False
+    ),
     "GPT2LMHeadModel": _HfExamplesInfo("openai-community/gpt2", {"alias": "gpt2"}),
     "GPTBigCodeForCausalLM": _HfExamplesInfo(
         "bigcode/starcoder",
@@ -967,6 +970,26 @@ _MULTIMODAL_EXAMPLE_MODELS = {
         max_model_len=4096,
         min_transformers_version="4.57",
     ),
+    "Qwen3_5ForConditionalGeneration": _HfExamplesInfo(
+        "Qwen/Qwen3.5-9B-Instruct",
+        max_model_len=4096,
+        min_transformers_version="5.1.0",
+    ),
+    "Qwen3_5MoeForConditionalGeneration": _HfExamplesInfo(
+        "Qwen/Qwen3.5-35B-A3B-Instruct",
+        max_model_len=4096,
+        min_transformers_version="5.1.0",
+    ),
+    "Qwen3_5MTP": _HfExamplesInfo(
+        "Qwen/Qwen3.5-9B-Instruct",
+        speculative_model="Qwen/Qwen3.5-9B-Instruct",
+        min_transformers_version="5.1.0",
+    ),
+    "Qwen3_5MoeMTP": _HfExamplesInfo(
+        "Qwen/Qwen3.5-35B-A3B-Instruct",
+        speculative_model="Qwen/Qwen3.5-35B-A3B-Instruct",
+        min_transformers_version="5.1.0",
+    ),
     "Qwen3OmniMoeForConditionalGeneration": _HfExamplesInfo(
         "Qwen/Qwen3-Omni-30B-A3B-Instruct",
         max_model_len=4096,

tests/models/test_initialization.py

@@ -97,7 +97,7 @@ def can_initialize(
             "pickle error when loading `transformers.models.auto.CONFIG_MAPPING`"
         )
 
-    if model_arch == "DeepseekV32ForCausalLM":
+    if model_arch in ["DeepseekV32ForCausalLM", "GlmMoeDsaForCausalLM"]:
         from vllm.platforms import current_platform
 
         capability = current_platform.get_device_capability()

tests/quantization/test_cpu_wna16.py

@@ -17,7 +17,7 @@ DTYPE = ["bfloat16"]
 
 @pytest.mark.parametrize("model", MODELS)
 @pytest.mark.parametrize("dtype", DTYPE)
-def test_cpu_quant(vllm_runner, model, dtype):
+def test_ipex_quant(vllm_runner, model, dtype):
     with vllm_runner(model, dtype=dtype) as llm:
         output = llm.generate_greedy(["The capital of France is"], max_tokens=32)
     assert output

tests/quantization/test_ipex_quant.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Test model set-up and inference for quantized HF models supported
+on the CPU/GPU backend using IPEX (including AWQ/GPTQ).
+
+Validating the configuration and printing results for manual checking.
+
+Run `pytest tests/quantization/test_ipex_quant.py`.
+"""
+
+import pytest
+
+from vllm.platforms import current_platform
+
+MODELS = [
+    "AMead10/Llama-3.2-1B-Instruct-AWQ",
+    "shuyuej/Llama-3.2-1B-Instruct-GPTQ",  # with g_idx
+]
+DTYPE = ["bfloat16"]
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cpu() and not current_platform.is_xpu(),
+    reason="only supports Intel CPU/XPU backend.",
+)
+@pytest.mark.parametrize("model", MODELS)
+@pytest.mark.parametrize("dtype", DTYPE)
+def test_ipex_quant(vllm_runner, model, dtype):
+    with vllm_runner(model, dtype=dtype, enforce_eager=True) as llm:
+        output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
+    assert output
+    print(output)

tests/tokenizers_/test_basic.py

@@ -4,7 +4,6 @@ from typing import _get_protocol_attrs  # type: ignore
 
 import pytest
 from transformers import (
-    PreTrainedTokenizer,
     PreTrainedTokenizerBase,
     PreTrainedTokenizerFast,
 )
@@ -25,10 +24,6 @@ def _assert_tokenizer_like(tokenizer: object):
 
 
 def test_tokenizer_like_protocol():
-    tokenizer = get_tokenizer("gpt2", use_fast=False)
-    assert isinstance(tokenizer, PreTrainedTokenizer)
-    _assert_tokenizer_like(tokenizer)
-
     tokenizer = get_tokenizer("gpt2", use_fast=True)
     assert isinstance(tokenizer, PreTrainedTokenizerFast)
     _assert_tokenizer_like(tokenizer)

tests/v1/attention/test_attention_backends.py

@@ -179,7 +179,7 @@ def create_and_prepopulate_kv_cache(
         block_table[i, :num_blocks_for_seq] = inv_perm[start:end]
         start_block_idx += num_blocks_for_seq
 
-    # Create a realistic slot mapping that corresponds to the block table
+        # Create a realistic slot mapping that corresponds to the block table
     for i in range(batch_size):
         token_offsets = torch.arange(int(query_lens[i])) + int(context_lens[i])
         block_indices = token_offsets // block_size

tests/v1/attention/test_mla_backends.py

@@ -27,7 +27,7 @@ from vllm.v1.attention.backend import CommonAttentionMetadata
 from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported
-from vllm.v1.kv_cache_interface import FullAttentionSpec
+from vllm.v1.kv_cache_interface import MLAAttentionSpec
 
 BACKENDS_TO_TEST = [
     AttentionBackendEnum.CUTLASS_MLA,
@@ -512,7 +512,7 @@ class MockMLAAttentionLayer(AttentionLayerBase):
 
 def run_attention_backend(
     backend: AttentionBackendEnum,
-    kv_cache_spec: FullAttentionSpec,
+    kv_cache_spec: MLAAttentionSpec,
     layer_names: list[str],
     vllm_config,
     device: torch.device,
@@ -989,7 +989,7 @@ def test_backend_correctness(
         kv_cache = kv_cache_per_block_size[block_size]
 
         # Create kv_cache_spec with the correct block_size for this backend
-        backend_kv_cache_spec = FullAttentionSpec(
+        backend_kv_cache_spec = MLAAttentionSpec(
             block_size=block_size,
             num_kv_heads=vllm_config.model_config.get_num_kv_heads(
                 vllm_config.parallel_config
@@ -997,6 +997,7 @@ def test_backend_correctness(
             head_size=vllm_config.model_config.get_head_size(),
             dtype=vllm_config.model_config.dtype,
             sliding_window=vllm_config.model_config.get_sliding_window(),
+            cache_dtype_str=vllm_config.cache_config.cache_dtype,
         )
 
         backend_output = run_attention_backend(

tests/v1/core/test_kv_cache_utils.py

@@ -1046,6 +1046,99 @@ def test_get_kv_cache_configs_multiple_workers():
         )
 
 
+@pytest.mark.parametrize(
+    "asymmetric_memory",
+    [False, True],
+    ids=["symmetric", "asymmetric"],
+)
+def test_get_kv_cache_configs_pp_sharding(asymmetric_memory):
+    model_config = ModelConfig(max_model_len=512)
+    vllm_config = VllmConfig(model_config=model_config)
+
+    ref_kv_cache_spec = new_kv_cache_spec()
+    pp_kv_cache_specs = [
+        {"layer1": ref_kv_cache_spec},
+        {"layer2": ref_kv_cache_spec},
+    ]
+
+    expected_num_blocks = model_config.max_model_len // ref_kv_cache_spec.block_size + 1
+    avail_memory = ref_kv_cache_spec.page_size_bytes * expected_num_blocks
+
+    # With per-worker validation, each worker only needs memory for its own
+    # layers. Worker 2 having more memory shouldn't affect worker 1's config.
+    available_memory = (
+        [avail_memory, avail_memory * 2] if asymmetric_memory else [avail_memory] * 2
+    )
+
+    kv_cache_configs = get_kv_cache_configs(
+        vllm_config,
+        pp_kv_cache_specs,
+        available_memory,
+    )
+
+    assert kv_cache_configs == [
+        KVCacheConfig(
+            num_blocks=expected_num_blocks,
+            kv_cache_tensors=[
+                KVCacheTensor(
+                    size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks,
+                    shared_by=["layer1"],
+                ),
+            ],
+            kv_cache_groups=[KVCacheGroupSpec(["layer1"], ref_kv_cache_spec)],
+        ),
+        KVCacheConfig(
+            num_blocks=expected_num_blocks,
+            kv_cache_tensors=[
+                KVCacheTensor(
+                    size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks,
+                    shared_by=["layer2"],
+                ),
+            ],
+            kv_cache_groups=[KVCacheGroupSpec(["layer2"], ref_kv_cache_spec)],
+        ),
+    ]
+
+
+def test_project_kv_cache_groups_to_worker():
+    spec_a = new_kv_cache_spec()
+    spec_b = new_kv_cache_spec(num_kv_heads=4)
+
+    global_groups = [
+        KVCacheGroupSpec(["layer1", "layer2", "layer3"], spec_a),
+    ]
+    worker_spec = {"layer1": spec_a, "layer2": spec_a}
+    projected = kv_cache_utils._project_kv_cache_groups_to_worker(
+        global_groups, worker_spec
+    )
+    assert len(projected) == 1
+    assert projected[0].layer_names == ["layer1", "layer2"]
+    assert projected[0].kv_cache_spec is spec_a
+
+    projected = kv_cache_utils._project_kv_cache_groups_to_worker(
+        global_groups, {"layer4": spec_a}
+    )
+    assert len(projected) == 1
+    assert projected[0].layer_names == []
+    assert projected[0].kv_cache_spec is spec_a
+
+    uniform_spec = UniformTypeKVCacheSpecs(
+        block_size=16,
+        kv_cache_specs={"layer1": spec_a, "layer2": spec_b, "layer3": spec_a},
+    )
+    global_groups_uniform = [
+        KVCacheGroupSpec(["layer1", "layer2", "layer3"], uniform_spec),
+    ]
+    projected = kv_cache_utils._project_kv_cache_groups_to_worker(
+        global_groups_uniform, {"layer1": spec_a, "layer3": spec_a}
+    )
+    assert len(projected) == 1
+    assert projected[0].layer_names == ["layer1", "layer3"]
+    proj_spec = projected[0].kv_cache_spec
+    assert isinstance(proj_spec, UniformTypeKVCacheSpecs)
+    assert set(proj_spec.kv_cache_specs.keys()) == {"layer1", "layer3"}
+
+
 def test_merge_kv_cache_spec():
     same_layer_specs = [
         new_kv_cache_spec(num_kv_heads=32),

tests/v1/core/test_prefix_caching.py

@@ -857,6 +857,8 @@ def test_prefill_hybrid_model_combinations(spec_types: list[str]):
     # Should have blocks for all groups
     assert len(blocks.get_block_ids()) == num_groups
 
+    manager.new_step_starts()
+
     # Second request: should hit cached blocks for common prefix
     req1 = make_request("1", common_token_ids + [4] * 5, block_size, hash_fn)
     computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

tests/v1/core/test_scheduler.py

@@ -3675,3 +3675,72 @@ def test_abort_request_finished_recving():
     # verify request is deleted
     assert request.request_id not in scheduler.requests
     assert not scheduler.finished_recving_kv_req_ids
+
+
+def test_eagle3_mm_encoder_cache_with_shift():
+    """Test EAGLE3 encoder scheduling accounts for shift_computed_tokens.
+
+    Regression test for issue #32469: When EAGLE3 is enabled with
+    disable_chunked_mm_input=True, ensure encoder inputs are scheduled
+    when tokens overlap the MM range, properly accounting for
+    shift_computed_tokens in the boundary calculation.
+
+    Without the fix, the scheduler would fail to schedule encoder inputs
+    at the boundary, causing "Encoder cache miss" errors.
+    """
+    scheduler = create_scheduler(
+        model="llava-hf/llava-1.5-7b-hf",
+        max_num_batched_tokens=1024,
+        disable_chunked_mm_input=True,
+        max_model_len=2048,
+        num_speculative_tokens=4,  # This enables EAGLE with shift=1
+    )
+
+    mm_start_pos = 100
+    mm_length = 576
+
+    mm_positions = [
+        [PlaceholderRange(offset=mm_start_pos, length=mm_length)],
+    ]
+
+    requests = create_requests(
+        num_requests=1,
+        num_tokens=mm_start_pos + mm_length + 100,
+        mm_positions=mm_positions,
+    )
+
+    # Start with some tokens already computed to simulate decoding
+    request = requests[0]
+    request.num_computed_tokens = 0
+
+    scheduler.add_request(request)
+    output = scheduler.schedule()
+
+    assert output is not None
+    shift_computed_tokens = 1
+    req_id = request.request_id
+
+    assert req_id in output.num_scheduled_tokens
+    num_scheduled = output.num_scheduled_tokens[req_id]
+
+    mm_feature = request.mm_features[0]
+    start_pos = mm_feature.mm_position.offset
+    tokens_end = request.num_computed_tokens + num_scheduled
+    scheduled_end_with_shift = tokens_end + shift_computed_tokens
+
+    # Assert that we scheduled into the MM range (test setup verification)
+    assert scheduled_end_with_shift > start_pos, (
+        f"Test setup error: expected to schedule into MM range. "
+        f"scheduled_end_with_shift={scheduled_end_with_shift}, "
+        f"start_pos={start_pos}"
+    )
+
+    # The key assertion: when scheduled tokens overlap MM range
+    # (accounting for EAGLE's shift), encoder MUST be scheduled.
+    # Without the fix, this would fail at the boundary case.
+    assert req_id in output.scheduled_encoder_inputs, (
+        f"Encoder input missing: scheduled {num_scheduled} tokens "
+        f"(computed={request.num_computed_tokens}, end={tokens_end}, "
+        f"shifted_end={scheduled_end_with_shift}) overlapping MM at "
+        f"{start_pos}. The fix must schedule encoder inputs."
+    )

tools/pre_commit/check_boolean_context_manager.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Lint: detect `with a() and b():` (boolean op in with-statement context).
+
+Using `and`/`or` to combine context managers is almost always a bug:
+
+    with ctx_a() and ctx_b():   # BUG: only ctx_b is entered
+    with ctx_a() or  ctx_b():   # BUG: only ctx_a is entered
+
+The correct way to combine context managers is:
+
+    with ctx_a(), ctx_b():          # comma-separated
+    with (ctx_a(), ctx_b()):        # parenthesized (Python 3.10+)
+    with contextlib.ExitStack() ... # ExitStack
+"""
+
+import ast
+import sys
+
+
+def check_file(filepath: str) -> list[str]:
+    try:
+        with open(filepath, encoding="utf-8") as f:
+            source = f.read()
+    except (OSError, UnicodeDecodeError):
+        return []
+
+    try:
+        tree = ast.parse(source, filename=filepath)
+    except SyntaxError:
+        return []
+
+    violations = []
+    for node in ast.walk(tree):
+        if isinstance(node, (ast.With, ast.AsyncWith)):
+            for item in node.items:
+                if isinstance(item.context_expr, ast.BoolOp):
+                    op = "and" if isinstance(item.context_expr.op, ast.And) else "or"
+                    violations.append(
+                        f"{filepath}:{item.context_expr.lineno}: "
+                        f"boolean `{op}` used to combine context managers "
+                        f"in `with` statement — use a comma instead"
+                    )
+    return violations
+
+
+def main() -> int:
+    if len(sys.argv) < 2:
+        print("Usage: check_boolean_context_manager.py <file> ...", file=sys.stderr)
+        return 1
+
+    all_violations = []
+    for filepath in sys.argv[1:]:
+        all_violations.extend(check_file(filepath))
+
+    if all_violations:
+        print(
+            "❌ Boolean operator used to combine context managers in `with` "
+            "statement.\n"
+            "   `with a() and b():` only enters `b()` as a context manager.\n"
+            "   Use `with a(), b():` or `with (a(), b()):` instead.\n"
+        )
+        for v in all_violations:
+            print(f"  {v}")
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

vllm/_aiter_ops.py

@@ -87,10 +87,6 @@ def _rocm_aiter_fused_moe_impl(
     a2_scale: torch.Tensor | None = None,
     num_local_tokens: torch.Tensor | None = None,
     output_dtype: torch.dtype | None = None,
-    hidden_pad: int = 0,
-    intermediate_pad: int = 0,
-    bias1: torch.Tensor | None = None,
-    bias2: torch.Tensor | None = None,
 ) -> torch.Tensor:
     from aiter import ActivationType, QuantType
     from aiter.fused_moe import fused_moe
@@ -114,10 +110,6 @@ def _rocm_aiter_fused_moe_impl(
         a2_scale,
         num_local_tokens=num_local_tokens,
         dtype=output_dtype,
-        hidden_pad=hidden_pad,
-        intermediate_pad=intermediate_pad,
-        bias1=bias1,
-        bias2=bias2,
     )
 
 
@@ -315,28 +307,6 @@ def _rocm_aiter_grouped_topk_fake(
     pass
 
 
-def _rocm_aiter_fused_topk_impl(
-    x: torch.Tensor,
-    router_logits: torch.Tensor,
-    top_k: int,
-    gate_up: bool,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    from aiter.fused_moe import fused_topk
-
-    # fused_topk returns (topk_weights, topk_indices)
-    return fused_topk(x, router_logits, top_k, gate_up)
-
-
-def _rocm_aiter_fused_topk_fake(
-    x: torch.Tensor,
-    router_logits: torch.Tensor,
-    top_k: int,
-    gate_up: bool,
-) -> None:
-    # tuple[torch.Tensor, torch.Tensor]:
-    pass
-
-
 # Cache whether aiter supports FP8 MLA parameters
 _AITER_MLA_SUPPORTS_FP8: bool | None = None
 
@@ -971,70 +941,6 @@ class rocm_aiter_ops:
         cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
         cls._TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM
 
-    @staticmethod
-    def get_aiter_activation_type(activation_str: str):
-        """
-        Given an activation type as a string, returns the corresponding aiter ActivationType enum.
-        Supported activation types: "no", "none", "silu", "gelu", "swiglu".
-        Returns None if the mapping fails.
-
-        Args:
-            activation_str (str): Activation type as string.
-
-        Returns:
-            Aiter ActivationType enum value, or None if not found.
-        """
-        # Import only locally, since aiter may not always be available.
-        try:
-            from aiter import ActivationType
-        except ImportError:
-            return None
-
-        if not isinstance(activation_str, str):
-            return None
-
-        name = activation_str.strip().lower()
-        mapping = {
-            "none": ActivationType.No,
-            "no": ActivationType.No,
-            "silu": ActivationType.Silu,
-            "gelu": ActivationType.Gelu,
-            "swiglu": ActivationType.Swiglu,
-        }
-        return mapping.get(name)
-
-    @staticmethod
-    def get_aiter_quant_type(quant_type_str: str):
-        """
-        Given a quantization type as a string, returns the corresponding aiter QuantType enum.
-        Supported quantization types: "no", "per_tensor", "per_token", "per_1x32", "per_1x128", "per_128x128".
-        Returns None if the mapping fails.
-
-        Args:
-            quant_type_str (str): Quantization type as string.
-
-        Returns:
-            Aiter QuantType enum value, or None if not found.
-        """
-        try:
-            from aiter import QuantType
-        except ImportError:
-            return None
-
-        if not isinstance(quant_type_str, str):
-            return None
-
-        name = quant_type_str.strip().lower()
-        mapping = {
-            "no": QuantType.No,
-            "per_tensor": QuantType.per_Tensor,
-            "per_token": QuantType.per_Token,
-            "per_1x32": QuantType.per_1x32,
-            "per_1x128": QuantType.per_1x128,
-            "per_128x128": QuantType.per_128x128,
-        }
-        return mapping.get(name)
-
     @classmethod
     @if_aiter_supported
     def is_enabled(cls) -> bool:
@@ -1164,14 +1070,6 @@ class rocm_aiter_ops:
                 dispatch_key=current_platform.dispatch_key,
             )
 
-            direct_register_custom_op(
-                op_name="rocm_aiter_fused_topk",
-                op_func=_rocm_aiter_fused_topk_impl,
-                mutates_args=[],
-                fake_impl=_rocm_aiter_fused_topk_fake,
-                dispatch_key=current_platform.dispatch_key,
-            )
-
             direct_register_custom_op(
                 op_name="rocm_aiter_mla_decode_fwd",
                 op_func=_rocm_aiter_mla_decode_fwd_impl,
@@ -1393,10 +1291,6 @@ class rocm_aiter_ops:
         a2_scale: torch.Tensor | None = None,
         num_local_tokens: torch.Tensor | None = None,
         output_dtype: torch.dtype | None = None,
-        hidden_pad: int = 0,
-        intermediate_pad: int = 0,
-        bias1: torch.Tensor | None = None,
-        bias2: torch.Tensor | None = None,
     ) -> torch.Tensor:
         return torch.ops.vllm.rocm_aiter_fused_moe(
             hidden_states,
@@ -1414,10 +1308,6 @@ class rocm_aiter_ops:
             a2_scale,
             num_local_tokens,
             output_dtype,
-            hidden_pad,
-            intermediate_pad,
-            bias1,
-            bias2,
         )
 
     @staticmethod
@@ -1522,15 +1412,6 @@ class rocm_aiter_ops:
             routed_scaling_factor,
         )
 
-    @staticmethod
-    def fused_topk(
-        x: torch.Tensor,
-        router_logits: torch.Tensor,
-        top_k: int,
-        gate_up: bool,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        return torch.ops.vllm.rocm_aiter_fused_topk(x, router_logits, top_k, gate_up)
-
     @staticmethod
     def mla_decode_fwd(
         q: torch.Tensor,
@@ -1637,45 +1518,6 @@ class rocm_aiter_ops:
         query = query.view(query_shape)
         key = key.view(key_shape)
 
-    @staticmethod
-    def triton_rope_and_cache(
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        positions: torch.Tensor,
-        cos_sin_cache: torch.Tensor,
-        is_neox: bool,
-        key_cache: torch.Tensor,
-        value_cache: torch.Tensor,
-        layer_slot_mapping: torch.Tensor,
-        k_scale: torch.Tensor,
-        v_scale: torch.Tensor,
-        flash_layout: bool,
-        apply_scale: bool,
-    ):
-        from aiter.ops.triton.fused_kv_cache import fused_qk_rope_reshape_and_cache
-
-        cos, sin = cos_sin_cache.chunk(2, dim=-1)
-        fused_qk_rope_reshape_and_cache(
-            query,
-            key,
-            value,
-            key_cache,
-            value_cache,
-            layer_slot_mapping,
-            positions,
-            cos,
-            sin,
-            k_scale,
-            v_scale,
-            is_neox,
-            flash_layout=flash_layout,
-            apply_scale=apply_scale,
-            q_out=query,
-            k_out=key,
-            output_zeros=False,
-        )
-
     @staticmethod
     def batched_gemm_a16wfp4(
         X: torch.Tensor,
@@ -1787,47 +1629,6 @@ class rocm_aiter_ops:
 
         return shuffle_weight(tensor, layout=layout)
 
-    @staticmethod
-    def shuffle_weight_a16w4(
-        tensor: "torch.Tensor",
-        nLane: int,
-        gate_up: bool,
-    ) -> "torch.Tensor":
-        """
-        Shuffles the weight tensor into (A16W4) layout for AITER kernels.
-
-        Args:
-            tensor: The input weight tensor to be shuffled.
-            layout: The block layout to use, defaults to (16, 4).
-
-        Returns:
-            torch.Tensor: The shuffled tensor.
-        """
-        from aiter.ops.shuffle import shuffle_weight_a16w4
-
-        return shuffle_weight_a16w4(tensor, nLane, gate_up)
-
-    @staticmethod
-    def shuffle_scale_a16w4(
-        tensor: "torch.Tensor",
-        num_experts: int,
-        gate_up: bool,
-    ) -> "torch.Tensor":
-        """
-        Shuffles the scale tensor into (A16W4) layout for AITER kernels.
-
-        Args:
-            tensor: The input scale tensor to be shuffled.
-            num_experts: Number of experts, needed for reshaping logic.
-            gate_up: Whether the scale is for w13 (True) or w2 (False).
-
-        Returns:
-            torch.Tensor: The shuffled scale tensor.
-        """
-        from aiter.ops.shuffle import shuffle_scale_a16w4
-
-        return shuffle_scale_a16w4(tensor, num_experts, gate_up)
-
     @staticmethod
     def shuffle_weights(
         *tensors: torch.Tensor, layout: tuple[int, int] = (16, 16)

vllm/_ipex_ops.py

@@ -53,7 +53,7 @@ if hasattr(torch.ops._xpu_C, "int4_gemm_w4a16"):
         return torch.empty((M, N), dtype=input.dtype, device=input.device)
 
 
-class xpu_ops:
+class ipex_ops:
     @staticmethod
     def flash_attn_varlen_func(
         q: torch.Tensor,
@@ -73,7 +73,7 @@ class xpu_ops:
         cu_seqlens_k: torch.Tensor | None = None,
         # passed in qwen vl
         dropout_p: float = 0.0,
-        # The following parameters are not used in xpu kernel currently,
+        # The following parameters are not used in ipex kernel currently,
         # we keep API compatible to CUDA's.
         scheduler_metadata=None,
         fa_version: int = 2,
@@ -153,6 +153,6 @@ class xpu_ops:
         sm_margin=0,  # Can be tuned if some SMs are used for communication
     ) -> None:
         logger.warning_once(
-            "get_scheduler_metadata is not implemented for xpu_ops, returning None."
+            "get_scheduler_metadata is not implemented for ipex_ops, returning None."
         )
         return None

vllm/benchmarks/datasets.py

@@ -39,6 +39,7 @@ from vllm.lora.utils import get_adapter_absolute_path
 from vllm.multimodal import MultiModalDataDict
 from vllm.multimodal.image import convert_image_mode
 from vllm.tokenizers import TokenizerLike
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.utils.import_utils import PlaceholderModule
 
 try:
@@ -57,11 +58,6 @@ try:
 except ImportError:
     librosa = PlaceholderModule("librosa")
 
-try:
-    from vllm.utils.argparse_utils import FlexibleArgumentParser
-except ImportError:
-    from argparse import ArgumentParser as FlexibleArgumentParser
-
 logger = logging.getLogger(__name__)
 
 # -----------------------------------------------------------------------------
@@ -1443,6 +1439,20 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
         help="Maximum distance for blazedit dataset. Min: 0, Max: 1.0",
     )
 
+    asr_group = parser.add_argument_group("asr dataset options")
+    asr_group.add_argument(
+        "--asr-max-audio-len-sec",
+        type=float,
+        default=float("inf"),
+        help="Maximum audio length in seconds for ASR dataset.",
+    )
+    asr_group.add_argument(
+        "--asr-min-audio-len-sec",
+        type=float,
+        default=0.0,
+        help="Minimum audio length in seconds for ASR dataset.",
+    )
+
     random_group = parser.add_argument_group("random dataset options")
     add_random_dataset_base_args(random_group)
 
@@ -1744,27 +1754,27 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
             or args.hf_name in VisionArenaDataset.SUPPORTED_DATASET_PATHS
         ):
             dataset_class = VisionArenaDataset
-            args.hf_split = "train"
+            args.hf_split = args.hf_split if args.hf_split else "train"
             args.hf_subset = None
         elif (
             args.dataset_path in MMVUDataset.SUPPORTED_DATASET_PATHS
             or args.hf_name in MMVUDataset.SUPPORTED_DATASET_PATHS
         ):
             dataset_class = MMVUDataset
-            args.hf_split = "validation"
+            args.hf_split = args.hf_split if args.hf_split else "validation"
             args.hf_subset = None
         elif (
             args.dataset_path in InstructCoderDataset.SUPPORTED_DATASET_PATHS
             or args.hf_name in InstructCoderDataset.SUPPORTED_DATASET_PATHS
         ):
             dataset_class = InstructCoderDataset
-            args.hf_split = "train"
+            args.hf_split = args.hf_split if args.hf_split else "train"
         elif (
             args.dataset_path in MTBenchDataset.SUPPORTED_DATASET_PATHS
             or args.hf_name in MTBenchDataset.SUPPORTED_DATASET_PATHS
         ):
             dataset_class = MTBenchDataset
-            args.hf_split = "train"
+            args.hf_split = args.hf_split if args.hf_split else "train"
         elif (
             args.dataset_path in MultiModalConversationDataset.SUPPORTED_DATASET_PATHS
             or args.hf_name in MultiModalConversationDataset.SUPPORTED_DATASET_PATHS
@@ -1780,22 +1790,26 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
             or args.hf_name in AIMODataset.SUPPORTED_DATASET_PATHS
         ):
             dataset_class = AIMODataset
-            args.hf_split = "train"
+            args.hf_split = args.hf_split if args.hf_split else "train"
         elif (
             args.dataset_path in NextEditPredictionDataset.SUPPORTED_DATASET_PATHS  # noqa: E501
             or args.hf_name in NextEditPredictionDataset.SUPPORTED_DATASET_PATHS
         ):
             dataset_class = NextEditPredictionDataset
-            args.hf_split = "train"
+            args.hf_split = args.hf_split if args.hf_split else "train"
         elif (
             args.dataset_path in ASRDataset.SUPPORTED_DATASET_PATHS
             or args.hf_name in ASRDataset.SUPPORTED_DATASET_PATHS
         ):
             dataset_class = ASRDataset
-            args.hf_split = "train"
+            args.hf_split = args.hf_split if args.hf_split else "train"
+            hf_kwargs = {
+                "asr_min_audio_len_sec": args.asr_min_audio_len_sec,
+                "asr_max_audio_len_sec": args.asr_max_audio_len_sec,
+            }
         elif args.dataset_path in BlazeditDataset.SUPPORTED_DATASET_PATHS:
             dataset_class = BlazeditDataset
-            args.hf_split = "train"
+            args.hf_split = args.hf_split if args.hf_split else "train"
             hf_kwargs = {
                 "min_distance": args.blazedit_min_distance,
                 "max_distance": args.blazedit_max_distance,
@@ -1805,13 +1819,13 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
             or args.hf_name in MLPerfDataset.SUPPORTED_DATASET_PATHS
         ):
             dataset_class = MLPerfDataset
-            args.hf_split = "train"
+            args.hf_split = args.hf_split if args.hf_split else "train"
         elif (
             args.dataset_path in MMStarDataset.SUPPORTED_DATASET_PATHS
             or args.hf_name in MMStarDataset.SUPPORTED_DATASET_PATHS
         ):
             dataset_class = MMStarDataset
-            args.hf_split = "val"
+            args.hf_split = args.hf_split if args.hf_split else "val"
             args.hf_subset = None
         else:
             supported_datasets = set(
@@ -1847,6 +1861,7 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
             no_stream=args.no_stream,
             hf_name=args.hf_name,
             disable_shuffle=args.disable_shuffle,
+            trust_remote_code=args.trust_remote_code,
         ).sample(
             num_requests=args.num_prompts,
             tokenizer=tokenizer,
@@ -2057,32 +2072,38 @@ class CustomDataset(BenchmarkDataset):
                 break
             prompt = item["prompt"]
 
-            new_output_len = output_len
-            if output_len is None or output_len == -1:
-                # check that the request has an 'output_tokens' field
-                if "output_tokens" not in item:
-                    raise ValueError(
-                        "If no output length is provided the "
-                        "custom dataset must contain an 'output_tokens' field."
+            if tokenizer is None:
+                new_output_len = 1
+            else:
+                new_output_len = output_len
+                if output_len is None or output_len == -1:
+                    # check that the request has an 'output_tokens' field
+                    if "output_tokens" not in item:
+                        raise ValueError(
+                            "If no output length is provided the "
+                            "custom dataset must contain an 'output_tokens' field."
+                        )
+                    # Use number of output tokens from the request data
+                    try:
+                        new_output_len = int(item["output_tokens"])
+                    except (ValueError, TypeError) as e:
+                        raise ValueError(
+                            f"Invalid value for 'output_tokens' in custom dataset: "
+                            f"'{item['output_tokens']}'. Must be an integer."
+                        ) from e
+
+            if tokenizer is None:
+                prompt_len = 1
+            else:
+                # apply template
+                if not skip_chat_template:
+                    prompt = tokenizer.apply_chat_template(
+                        [{"role": "user", "content": prompt}],
+                        add_generation_prompt=True,
+                        tokenize=False,
                     )
-                # Use number of output tokens from the request data
-                try:
-                    new_output_len = int(item["output_tokens"])
-                except (ValueError, TypeError) as e:
-                    raise ValueError(
-                        f"Invalid value for 'output_tokens' in custom dataset: "
-                        f"'{item['output_tokens']}'. Must be an integer."
-                    ) from e
 
-            # apply template
-            if not skip_chat_template:
-                prompt = tokenizer.apply_chat_template(
-                    [{"role": "user", "content": prompt}],
-                    add_generation_prompt=True,
-                    tokenize=False,
-                )
-
-            prompt_len = len(tokenizer(prompt).input_ids)
+                prompt_len = len(tokenizer(prompt).input_ids)
             sampled_requests.append(
                 SampleRequest(
                     prompt=prompt,
@@ -2405,6 +2426,7 @@ class HuggingFaceDataset(BenchmarkDataset):
         no_stream: bool = False,
         dataset_subset: str | None = None,
         hf_name: str | None = None,
+        trust_remote_code: bool = False,
         **kwargs,
     ) -> None:
         super().__init__(dataset_path=dataset_path, **kwargs)
@@ -2413,6 +2435,7 @@ class HuggingFaceDataset(BenchmarkDataset):
         self.dataset_subset = dataset_subset
         self.load_stream = not no_stream
         self.hf_name = hf_name or dataset_path
+        self.trust_remote_code = trust_remote_code
         self.load_data()
 
     def load_data(self) -> None:
@@ -2422,6 +2445,7 @@ class HuggingFaceDataset(BenchmarkDataset):
             name=self.dataset_subset,
             split=self.dataset_split,
             streaming=self.load_stream,
+            trust_remote_code=self.trust_remote_code,
         )
         if not getattr(self, "disable_shuffle", False):
             self.data = self.data.shuffle(seed=self.random_seed)
@@ -3071,13 +3095,9 @@ class ASRDataset(HuggingFaceDataset):
         "kensho/spgispeech",
     }
 
-    DEFAULT_OUTPUT_LEN = 128
+    DEFAULT_OUTPUT_LEN = 1024
     IS_MULTIMODAL = True
 
-    # TODO Whisper-specific. Abstract interface when more models are supported.
-    TRANSCRIPTION_PREAMBLE = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>"
-    skip_long_audios: bool = True
-
     def sample(
         self,
         tokenizer: TokenizerLike,
@@ -3088,22 +3108,28 @@ class ASRDataset(HuggingFaceDataset):
         **kwargs,
     ) -> list:
         output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
-        prompt = ASRDataset.TRANSCRIPTION_PREAMBLE
+        if "openai" in tokenizer.name_or_path:
+            prompt = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>"
+        else:
+            prompt = ""
         prompt_len = len(tokenizer(prompt).input_ids)
         sampled_requests = []
         ind = 0
         skipped = 0
+        asr_min_audio_len_sec = kwargs.get("asr_min_audio_len_sec")
+        asr_max_audio_len_sec = kwargs.get("asr_max_audio_len_sec")
+        durations = []
         for item in self.data:
             if len(sampled_requests) >= num_requests:
                 break
             audio = item["audio"]
             y, sr = audio["array"], audio["sampling_rate"]
             duration_s = librosa.get_duration(y=y, sr=sr)
-            # Whisper max supported duration
-            if self.skip_long_audios and duration_s > 30:
+            if duration_s < asr_min_audio_len_sec or duration_s > asr_max_audio_len_sec:
                 skipped += 1
                 continue
 
+            durations.append(duration_s)
             mm_content = {"audio": (y, sr)}
             sampled_requests.append(
                 SampleRequest(
@@ -3122,6 +3148,20 @@ class ASRDataset(HuggingFaceDataset):
                 " what Whisper supports.",
                 skipped,
             )
+
+        logger.info("Number of audio samples: %d", len(durations))
+        avg_duration = sum(durations) / len(durations) if durations else 0
+        min_duration = min(durations) if durations else 0
+        max_duration = max(durations) if durations else 0
+        median_duration = np.median(durations) if durations else 0
+        logger.info(
+            "Audio duration statistics (s): avg=%.2f, min=%.2f, max=%.2f, median=%.2f",
+            avg_duration,
+            min_duration,
+            max_duration,
+            median_duration,
+        )
+
         self.maybe_oversample_requests(
             sampled_requests, num_requests, request_id_prefix, no_oversample
         )

vllm/benchmarks/lib/endpoint_request_func.py

@@ -93,6 +93,7 @@ class RequestFuncOutput:
     prompt_len: int = 0
     error: str = ""
     start_time: float = 0.0
+    input_audio_duration: float = 0.0  # in seconds
 
 
 class RequestFunc(Protocol):
@@ -422,6 +423,8 @@ async def async_request_openai_audio(
 
         output = RequestFuncOutput()
         output.prompt_len = request_func_input.prompt_len
+        output.input_audio_duration = soundfile.info(f).duration
+        f.seek(0)
 
         generated_text = ""
         ttft = 0.0
@@ -442,7 +445,9 @@ async def async_request_openai_audio(
 
                         messages = handler.add_chunk(chunk_bytes)
                         for message in messages:
-                            chunk = message.decode("utf-8").removeprefix("data: ")
+                            if type(message) is bytes:
+                                message = message.decode("utf-8")
+                            chunk = message.removeprefix("data: ")
                             if chunk != "[DONE]":
                                 timestamp = time.perf_counter()
                                 data = json.loads(chunk)
@@ -741,6 +746,37 @@ async def async_request_infinity_embeddings_clip(
     )
 
 
+async def async_request_vllm_pooling(
+    request_func_input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput:
+    api_url = request_func_input.api_url
+    _validate_api_url(api_url, "vLLM Pooling API", "pooling")
+
+    payload = {
+        "model": request_func_input.model_name
+        if request_func_input.model_name
+        else request_func_input.model,
+        "truncate_prompt_tokens": -1,
+    }
+
+    payload = payload | request_func_input.prompt
+
+    _update_payload_common(payload, request_func_input)
+
+    headers = _get_headers("application/json")
+    _update_headers_common(headers, request_func_input)
+
+    return await _run_pooling_request(
+        session,
+        api_url,
+        payload=payload,
+        headers=headers,
+        pbar=pbar,
+    )
+
+
 # TODO: Add more request functions for different API protocols.
 ASYNC_REQUEST_FUNCS: dict[str, RequestFunc] = {
     "vllm": async_request_openai_completions,
@@ -755,6 +791,7 @@ ASYNC_REQUEST_FUNCS: dict[str, RequestFunc] = {
     "infinity-embeddings": async_request_infinity_embeddings,
     "infinity-embeddings-clip": async_request_infinity_embeddings_clip,
     # (Infinity embedding server does not support vlm2vec)
+    "vllm-pooling": async_request_vllm_pooling,
     "vllm-rerank": async_request_vllm_rerank,
 }
 

vllm/benchmarks/serve.py

@@ -26,6 +26,7 @@ import json
 import os
 import random
 import shutil
+import ssl
 import time
 import uuid
 import warnings
@@ -60,11 +61,14 @@ TERM_PLOTLIB_AVAILABLE = (importlib.util.find_spec("termplotlib") is not None) a
 
 
 async def get_first_model_from_server(
-    base_url: str, headers: dict | None = None
+    base_url: str,
+    headers: dict | None = None,
+    ssl_context: ssl.SSLContext | bool | None = None,
 ) -> tuple[str, str]:
     """Fetch the first model from the server's /v1/models endpoint."""
     models_url = f"{base_url}/v1/models"
-    async with aiohttp.ClientSession() as session:
+    connector = aiohttp.TCPConnector(ssl=ssl_context)
+    async with aiohttp.ClientSession(connector=connector) as session:
         try:
             async with session.get(models_url, headers=headers) as response:
                 response.raise_for_status()
@@ -193,6 +197,7 @@ class BenchmarkMetrics:
     # Max output tokens per second and concurrent requests at that peak
     max_output_tokens_per_s: float
     max_concurrent_requests: int
+    rtfx: float = 0.0  # Inverse Real-Time Factor for ASR benchmarks
 
 
 @dataclass
@@ -412,21 +417,25 @@ def calculate_metrics(
     all_tpots: list[float] = []
     ttfts: list[float] = []
     e2els: list[float] = []
+    input_audio_duration = 0.0
     for i in range(len(outputs)):
         if outputs[i].success:
             output_len = outputs[i].output_tokens
 
             if not output_len:
-                # We use the tokenizer to count the number of output tokens
-                # for some serving backends instead of looking at
-                # len(outputs[i].itl) since multiple output tokens may be
-                # bundled together
-                # Note : this may inflate the output token count slightly
-                output_len = len(
-                    tokenizer(
-                        outputs[i].generated_text, add_special_tokens=False
-                    ).input_ids
-                )
+                if tokenizer is None:
+                    output_len = 1
+                else:
+                    # We use the tokenizer to count the number of output tokens
+                    # for some serving backends instead of looking at
+                    # len(outputs[i].itl) since multiple output tokens may be
+                    # bundled together
+                    # Note : this may inflate the output token count slightly
+                    output_len = len(
+                        tokenizer(
+                            outputs[i].generated_text, add_special_tokens=False
+                        ).input_ids
+                    )
             actual_output_lens.append(output_len)
             total_input += input_requests[i].prompt_len
             tpot = 0
@@ -439,6 +448,7 @@ def calculate_metrics(
             itls += outputs[i].itl
             ttfts.append(outputs[i].ttft)
             e2els.append(outputs[i].latency)
+            input_audio_duration += outputs[i].input_audio_duration
             completed += 1
         else:
             actual_output_lens.append(0)
@@ -583,6 +593,7 @@ def calculate_metrics(
         ],
         max_output_tokens_per_s=max_output_tokens_per_s,
         max_concurrent_requests=max_concurrent_requests,
+        rtfx=input_audio_duration / dur_s,
     )
 
     return metrics, actual_output_lens
@@ -615,6 +626,7 @@ async def benchmark(
     ramp_up_start_rps: int | None = None,
     ramp_up_end_rps: int | None = None,
     ready_check_timeout_sec: int = 600,
+    ssl_context: ssl.SSLContext | bool | None = None,
 ):
     try:
         request_func = ASYNC_REQUEST_FUNCS[endpoint_type]
@@ -622,6 +634,8 @@ async def benchmark(
         raise ValueError(f"Unknown backend: {endpoint_type}") from None
 
     # Reuses connections across requests to reduce TLS handshake overhead.
+    # Use ssl_context if provided, otherwise default to True for https URLs
+    ssl_setting = ssl_context if ssl_context is not None else ("https://" in api_url)
     connector = aiohttp.TCPConnector(
         limit=max_concurrency or 0,
         limit_per_host=max_concurrency or 0,
@@ -630,7 +644,7 @@ async def benchmark(
         keepalive_timeout=60,
         enable_cleanup_closed=True,
         force_close=False,
-        ssl=("https://" in api_url),
+        ssl=ssl_setting,
     )
 
     session = aiohttp.ClientSession(
@@ -908,7 +922,7 @@ async def benchmark(
         print("{:<40} {:<10.2f}".format("Request rate configured (RPS):", request_rate))
     print("{:<40} {:<10.2f}".format("Benchmark duration (s):", benchmark_duration))
     print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input))
-    if isinstance(metrics, BenchmarkMetrics):
+    if isinstance(metrics, BenchmarkMetrics) and tokenizer:
         print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output))
     print(
         "{:<40} {:<10.2f}".format(
@@ -922,26 +936,35 @@ async def benchmark(
             )
         )
     if isinstance(metrics, BenchmarkMetrics):
-        print(
-            "{:<40} {:<10.2f}".format(
-                "Output token throughput (tok/s):", metrics.output_throughput
+        if tokenizer:
+            print(
+                "{:<40} {:<10.2f}".format(
+                    "Output token throughput (tok/s):", metrics.output_throughput
+                )
             )
-        )
-        print(
-            "{:<40} {:<10.2f}".format(
-                "Peak output token throughput (tok/s):", metrics.max_output_tokens_per_s
+            print(
+                "{:<40} {:<10.2f}".format(
+                    "Peak output token throughput (tok/s):",
+                    metrics.max_output_tokens_per_s,
+                )
             )
-        )
         print(
             "{:<40} {:<10.2f}".format(
                 "Peak concurrent requests:", metrics.max_concurrent_requests
             )
         )
-    print(
-        "{:<40} {:<10.2f}".format(
-            "Total token throughput (tok/s):", metrics.total_token_throughput
+        if metrics.rtfx > 0.0:
+            print(
+                "{:<40} {:<10.2f}".format(
+                    "RTFx (Inverse Real-Time Factor):", metrics.rtfx
+                )
+            )
+    if tokenizer:
+        print(
+            "{:<40} {:<10.2f}".format(
+                "Total token throughput (tok/s):", metrics.total_token_throughput
+            )
         )
-    )
 
     if isinstance(metrics, BenchmarkMetrics):
         result = {
@@ -963,6 +986,7 @@ async def benchmark(
             "errors": [output.error for output in outputs],
             "max_output_tokens_per_s": metrics.max_output_tokens_per_s,
             "max_concurrent_requests": metrics.max_concurrent_requests,
+            "rtfx": metrics.rtfx,
         }
     else:
         result = {
@@ -1029,7 +1053,7 @@ async def benchmark(
             print("{:<40} {:<10.2f}".format(f"P{p_word} {metric_name} (ms):", value))
             result[f"p{p_word}_{metric_attribute_name}_ms"] = value
 
-    if task_type == TaskType.GENERATION:
+    if task_type == TaskType.GENERATION and tokenizer:
         process_one_metric("ttft", "TTFT", "Time to First Token")
         process_one_metric("tpot", "TPOT", "Time per Output Token (excl. 1st token)")
         process_one_metric("itl", "ITL", "Inter-token Latency")
@@ -1501,6 +1525,20 @@ def add_cli_args(parser: argparse.ArgumentParser):
         type=json.loads,
         default=None,
     )
+    parser.add_argument(
+        "--skip-tokenizer-init",
+        action="store_true",
+        default=False,
+        help="Skip initialization of tokenizer and detokenizer",
+    )
+
+    parser.add_argument(
+        "--insecure",
+        action="store_true",
+        default=False,
+        help="Disable SSL certificate verification. Use this option when "
+        "connecting to servers with self-signed certificates.",
+    )
 
 
 def main(args: argparse.Namespace) -> dict[str, Any]:
@@ -1553,23 +1591,38 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:
             else:
                 raise ValueError("Invalid header format. Please use KEY=VALUE format.")
 
+    # SSL context configuration
+    ssl_context: ssl.SSLContext | bool | None = None
+    if args.insecure:
+        # Disable SSL certificate verification
+        ssl_context = False
+    elif "https://" in base_url:
+        # Use default SSL context for HTTPS
+        ssl_context = True
+
     # Fetch model from server if not specified
     if args.model is None:
         print("Model not specified, fetching first model from server...")
-        model_name, model_id = await get_first_model_from_server(base_url, headers)
+        model_name, model_id = await get_first_model_from_server(
+            base_url, headers, ssl_context
+        )
         print(f"First model name: {model_name}, first model id: {model_id}")
     else:
         model_name = args.served_model_name
         model_id = args.model
 
-    tokenizer_id = args.tokenizer if args.tokenizer is not None else model_id
-    tokenizer_mode = args.tokenizer_mode
-
-    tokenizer = get_tokenizer(
-        tokenizer_id,
-        tokenizer_mode=tokenizer_mode,
-        trust_remote_code=args.trust_remote_code,
-    )
+    if args.skip_tokenizer_init:
+        tokenizer_id = None
+        tokenizer_mode = None
+        tokenizer = None
+    else:
+        tokenizer_id = args.tokenizer if args.tokenizer is not None else model_id
+        tokenizer_mode = args.tokenizer_mode
+        tokenizer = get_tokenizer(
+            tokenizer_id,
+            tokenizer_mode=tokenizer_mode,
+            trust_remote_code=args.trust_remote_code,
+        )
 
     if args.dataset_name is None:
         raise ValueError(
@@ -1680,6 +1733,7 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:
         ramp_up_start_rps=args.ramp_up_start_rps,
         ramp_up_end_rps=args.ramp_up_end_rps,
         ready_check_timeout_sec=args.ready_check_timeout_sec,
+        ssl_context=ssl_context,
     )
 
     # Save config and results to json

vllm/compilation/compiler_interface.py

@@ -257,7 +257,20 @@ class InductorStandaloneAdaptor(CompilerInterface):
         if use_aot:
             compile_kwargs["aot"] = True  # type: ignore[assignment]
 
-        compiled_graph = standalone_compile(graph, example_inputs, **compile_kwargs)
+        # Inductor's pre-grad passes don't do anything for vLLM.
+        # The pre-grad passes get run even on cache-hit and negatively impact
+        # vllm cold compile times by O(1s)
+        # Can remove this after the following issue gets fixed
+        # https://github.com/pytorch/pytorch/issues/174502
+        if envs.VLLM_ENABLE_PREGRAD_PASSES:
+            ctx: Any = contextlib.nullcontext()
+        else:
+            ctx = patch(
+                "torch._inductor.compile_fx._recursive_pre_grad_passes",
+                lambda gm, _: gm,
+            )
+        with ctx:
+            compiled_graph = standalone_compile(graph, example_inputs, **compile_kwargs)
 
         if use_aot:
             from torch._inductor.standalone_compile import AOTCompiledArtifact

vllm/compilation/passes/fusion/allreduce_rms_fusion.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import contextlib
 from importlib.util import find_spec
 from types import ModuleType
 
@@ -36,7 +37,9 @@ if find_spec("flashinfer"):
     try:
         import flashinfer.comm as _flashinfer_comm
 
-        if hasattr(_flashinfer_comm, "trtllm_allreduce_fusion"):
+        if hasattr(_flashinfer_comm, "allreduce_fusion") and hasattr(
+            _flashinfer_comm, "create_allreduce_fusion_workspace"
+        ):
             flashinfer_comm = _flashinfer_comm
     except ImportError:
         pass
@@ -79,7 +82,7 @@ _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB: dict[int, dict[int, float]] = {
 
 
 if flashinfer_comm is not None:
-    _FI_WORKSPACE_TENSOR = None
+    _FI_WORKSPACE = None
     MiB = 1024 * 1024
 
     def call_trtllm_fused_allreduce_norm(
@@ -87,10 +90,8 @@ if flashinfer_comm is not None:
         residual: torch.Tensor,
         rms_gamma: torch.Tensor,
         rms_eps: float,
-        world_rank: int,
         world_size: int,
         launch_with_pdl: bool,
-        trigger_completion_at_end: bool,
         fp32_acc: bool,
         max_token_num: int,
         pattern_code: int,
@@ -121,7 +122,7 @@ if flashinfer_comm is not None:
             max_one_shot_size is None or current_tensor_size <= max_one_shot_size * MiB
         )
 
-        assert _FI_WORKSPACE_TENSOR is not None, (
+        assert _FI_WORKSPACE is not None, (
             "Flashinfer must be enabled when using flashinfer"
         )
         if norm_out is None:
@@ -134,24 +135,18 @@ if flashinfer_comm is not None:
             residual_out = allreduce_in
         # For the sizes that are smaller than the max size,
         # we only use flashinfer one shot allreduce
-        flashinfer_comm.trtllm_allreduce_fusion(
-            allreduce_in=allreduce_in,
-            token_num=allreduce_in.shape[0],
+        flashinfer_comm.allreduce_fusion(
+            input=allreduce_in,
+            workspace=_FI_WORKSPACE,
+            pattern=pattern_code,
             residual_in=residual,
             residual_out=residual_out,
             norm_out=norm_out,
             rms_gamma=rms_gamma,
             rms_eps=rms_eps,
-            world_rank=world_rank,
-            world_size=world_size,
-            hidden_dim=allreduce_in.shape[-1],
-            workspace_ptrs=_FI_WORKSPACE_TENSOR,
             launch_with_pdl=launch_with_pdl,
             use_oneshot=use_oneshot,
-            trigger_completion_at_end=trigger_completion_at_end,
             fp32_acc=fp32_acc,
-            pattern_code=pattern_code,
-            allreduce_out=None,
             quant_out=quant_out,
             scale_out=scale_out,
             # in vllm we only support swizzled layout
@@ -164,10 +159,8 @@ if flashinfer_comm is not None:
         residual: torch.Tensor,
         rms_gamma: torch.Tensor,
         rms_eps: float,
-        world_rank: int,
         world_size: int,
         launch_with_pdl: bool,
-        trigger_completion_at_end: bool,
         fp32_acc: bool,
         max_token_num: int,
         pattern_code: int,
@@ -200,25 +193,18 @@ class FlashInferFusedAllReduceParams:
 
     def __init__(
         self,
-        rank: int,
         world_size: int,
-        use_fp32_lamport: bool = False,
         max_token_num: int = 1024,
     ) -> None:
-        self.rank = rank
         self.world_size = world_size
-        self.use_fp32_lamport = use_fp32_lamport
-        self.trigger_completion_at_end = True
         self.launch_with_pdl = True
         self.fp32_acc = True
         self.max_token_num = max_token_num
 
     def get_trtllm_fused_allreduce_kwargs(self) -> dict[str, bool | int]:
         return {
-            "world_rank": self.rank,
             "world_size": self.world_size,
             "launch_with_pdl": self.launch_with_pdl,
-            "trigger_completion_at_end": self.trigger_completion_at_end,
             "fp32_acc": self.fp32_acc,
             "max_token_num": self.max_token_num,
         }
@@ -712,7 +698,6 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
         self.hidden_dim = config.model_config.get_hidden_size()
         self.group = get_tp_group().device_group
         rank = get_tensor_model_parallel_rank()
-        use_fp32_lamport = self.model_dtype == torch.float32
         if flashinfer_comm is None:
             logger.warning(
                 "Flashinfer is not installed or comm module not found, "
@@ -730,7 +715,7 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
                 self.tp_size,
             )
             return
-        element_size = 4 if use_fp32_lamport else 2
+        element_size = torch.tensor([], dtype=self.model_dtype).element_size()
         self.max_token_num = max_size // (self.hidden_dim * element_size)
         # take the min to save workspace size and we'll never use more
         # than max_num_batched_tokens anyways
@@ -744,23 +729,19 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
             scope="global",
         )
 
-        self.ipc_handles, workspace_tensor = (
-            flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce_fusion(
-                tp_rank=rank,
-                tp_size=self.tp_size,
-                max_token_num=self.max_token_num,
-                hidden_dim=self.hidden_dim,
-                group=self.group,
-                use_fp32_lamport=use_fp32_lamport,
-            )
+        self.workspace = flashinfer_comm.create_allreduce_fusion_workspace(
+            backend="trtllm",
+            world_size=self.tp_size,
+            rank=rank,
+            max_token_num=self.max_token_num,
+            hidden_dim=self.hidden_dim,
+            dtype=self.model_dtype,
         )
 
-        global _FI_WORKSPACE_TENSOR
-        _FI_WORKSPACE_TENSOR = workspace_tensor
+        global _FI_WORKSPACE
+        _FI_WORKSPACE = self.workspace
         self.allreduce_params = FlashInferFusedAllReduceParams(
-            rank=rank,
             world_size=self.tp_size,
-            use_fp32_lamport=use_fp32_lamport,
             max_token_num=self.max_token_num,
         )
 
@@ -832,7 +813,6 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
     def __del__(self) -> None:
         if getattr(self, "disabled", True):
             return
-        if flashinfer_comm is not None:
-            flashinfer_comm.trtllm_destroy_ipc_workspace_for_all_reduce(
-                self.ipc_handles, self.group
-            )
+        if getattr(self, "workspace", None) is not None:
+            with contextlib.suppress(Exception):
+                self.workspace.destroy()

vllm/compilation/passes/fusion/rope_kvcache_fusion.py

@@ -1,230 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import torch
-import torch._inductor.pattern_matcher as pm
-from torch import fx
-from torch._higher_order_ops import auto_functionalized
-from torch._inductor.fx_passes.post_grad import view_to_reshape
-from torch._inductor.pattern_matcher import PatternMatcherPass
-
-from vllm.config import VllmConfig, get_layers_from_vllm_config
-from vllm.config.utils import Range
-from vllm.logger import init_logger
-from vllm.model_executor.layers.attention.attention import (
-    Attention,
-    get_attention_context,
-)
-from vllm.utils.torch_utils import direct_register_custom_op
-
-from ..inductor_pass import enable_fake_mode
-from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
-from .matcher_utils import (
-    MatcherRotaryEmbedding,
-)
-from .rms_quant_fusion import (
-    empty_bf16,
-    empty_i64,
-)
-
-logger = init_logger(__name__)
-
-
-def fused_rope_and_unified_kv_cache_update_impl(
-    query: torch.Tensor,
-    key: torch.Tensor,
-    value: torch.Tensor,
-    positions: torch.Tensor,
-    cos_sin_cache: torch.Tensor,
-    is_neox: bool,
-    layer_name: str = "",
-) -> torch.Tensor:
-    """
-    This impl fetches the KV cache and slot mapping from the forward context,
-    then calls the layer impl's `AttentionImpl.do_rope_and_kv_cache_update` method.
-    It also returns a dummy tensor, similar to `Attention.unified_kv_cache_update`,
-    that is passed to unified_attention to signal a side effect and
-    the data dependency between them to ensure torch.compile preserves ordering.
-    """
-    _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
-    if layer_slot_mapping is not None:
-        attn_layer.impl.do_rope_and_kv_cache_update(
-            attn_layer,
-            query,
-            key,
-            value,
-            positions,
-            cos_sin_cache,
-            is_neox,
-            kv_cache,
-            layer_slot_mapping,
-        )
-
-    return torch.empty(0, device=kv_cache.device, dtype=kv_cache.dtype)
-
-
-def fused_rope_and_unified_kv_cache_update_fake(
-    query: torch.Tensor,
-    key: torch.Tensor,
-    value: torch.Tensor,
-    positions: torch.Tensor,
-    cos_sin_cache: torch.Tensor,
-    is_neox: bool,
-    layer_name: str = "",
-) -> torch.Tensor:
-    return torch.empty(0, device=query.device, dtype=query.dtype)
-
-
-direct_register_custom_op(
-    op_name="fused_rope_and_unified_kv_cache_update",
-    op_func=fused_rope_and_unified_kv_cache_update_impl,
-    mutates_args=["query", "key"],
-    fake_impl=fused_rope_and_unified_kv_cache_update_fake,
-)
-
-
-class RopeReshapeKVCachePattern:
-    """
-    This pattern matches the following unfused inplace ops:
-      q, k = rotary_embedding(positions, q, k, head_size, cos_sin_cache, is_neox)
-      kv_cache_dummy = unified_kv_cache_update(k, v, layer_name)
-
-    and replaces it with the fused inplace op:
-      kv_cache_dummy = fused_rope_and_unified_kv_cache_update(
-        q, k, v, positions, cos_sin_cache, is_neox, layer_name
-      )
-    """
-
-    FUSED_OP = torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default
-
-    def __init__(
-        self,
-        layer: Attention,
-        is_neox: bool,
-    ) -> None:
-        self.layer_name = layer.layer_name
-        self.num_heads = layer.num_heads
-        self.num_kv_heads = layer.num_kv_heads
-        self.head_size = layer.head_size
-        self.head_size_v = layer.head_size_v
-        self.is_neox = is_neox
-
-        self.q_size = self.num_heads * self.head_size
-        self.k_size = self.num_kv_heads * self.head_size
-        self.v_size = self.num_kv_heads * self.head_size_v
-
-        self.rope_matcher = MatcherRotaryEmbedding(
-            is_neox=self.is_neox,
-            head_size=self.head_size,
-            num_heads=self.num_heads,
-            num_kv_heads=self.num_kv_heads,
-        )
-
-    def get_inputs(self) -> list[torch.Tensor]:
-        # Sample inputs to help pattern tracing
-        T = 5
-        L = 4096
-        qkv = empty_bf16(T, self.q_size + self.k_size + self.v_size)
-        positions = empty_i64(T)
-        cos_sin_cache = empty_bf16(L, self.head_size)
-        return [
-            qkv,
-            positions,
-            cos_sin_cache,
-        ]
-
-    def register(self, pm_pass: PatternMatcherPass) -> None:
-        def pattern(
-            qkv: torch.Tensor,
-            positions: torch.Tensor,
-            cos_sin_cache: torch.Tensor,
-        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-            q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
-            q, k = self.rope_matcher(positions, q, k, cos_sin_cache)
-            q = q.view(-1, self.num_heads, self.head_size)
-            k = k.view(-1, self.num_kv_heads, self.head_size)
-            v = v.view(-1, self.num_kv_heads, self.head_size_v)
-            dummy = torch.ops.vllm.unified_kv_cache_update(k, v, self.layer_name)
-            return dummy, q, k, v
-
-        def replacement(
-            qkv: torch.Tensor,
-            positions: torch.Tensor,
-            cos_sin_cache: torch.Tensor,
-        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-            q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
-            q = q.view(-1, self.num_heads, self.head_size)
-            k = k.view(-1, self.num_kv_heads, self.head_size)
-            v = v.view(-1, self.num_kv_heads, self.head_size_v)
-            results = auto_functionalized(
-                self.FUSED_OP,
-                query=q,
-                key=k,
-                value=v,
-                positions=positions,
-                cos_sin_cache=cos_sin_cache,
-                is_neox=self.is_neox,
-                layer_name=self.layer_name,
-            )
-            return results[0], results[1], results[2], v
-
-        # NOTE: use view_to_reshape to unify view/reshape to simplify
-        # pattern and increase matching opportunities
-        def fwd_and_view_to_reshape(*args, **kwargs) -> fx.GraphModule:
-            gm = pm.fwd_only(*args, **kwargs)
-            view_to_reshape(gm)
-            return gm
-
-        pm.register_replacement(
-            pattern, replacement, self.get_inputs(), fwd_and_view_to_reshape, pm_pass
-        )
-
-
-class RopeKVCacheFusionPass(VllmPatternMatcherPass):
-    """
-    This pass fuses the rotary embedding and KV cache update operations
-    into a single fused kernel if available.
-
-    It uses the pattern matcher and matches each layer manually, as strings
-    cannot be wildcarded. This also lets us check support on attention layers
-    upon registration instead of during pattern matching.
-
-    This fusion eliminates the need for separate kernel launches and
-    intermediate memory operations between the RoPE and cache update steps.
-    """
-
-    @enable_fake_mode
-    def __init__(self, config: VllmConfig) -> None:
-        super().__init__(config)
-
-        self.patterns: PatternMatcherPass = PatternMatcherPass(
-            pass_name="rope_kv_cache_fusion_pass"
-        )
-
-        cc = config.compilation_config
-        self.max_token_num = cc.pass_config.rope_kvcache_fusion_max_token_num
-
-        attn_layers = get_layers_from_vllm_config(config, Attention)
-        for _, layer in attn_layers.items():
-            if layer.impl.fused_rope_kvcache_supported():
-                for is_neox in [True, False]:
-                    RopeReshapeKVCachePattern(
-                        layer=layer,
-                        is_neox=is_neox,
-                    ).register(self.patterns)
-
-        self.dump_patterns(config, self.patterns)
-
-    @VllmInductorPass.time_and_log
-    def __call__(self, graph: fx.Graph) -> None:
-        self.matched_count = self.patterns.apply(graph)
-        logger.debug("Replaced %s patterns", self.matched_count)
-
-    def is_applicable_for_range(self, compile_range: Range) -> bool:
-        # This pass works best for the small-batch decode setting.
-        # For large-batch e.g. prefill, it is better to use two separate kernels
-        # since they are compute bound and the fused kernels require further tuning.
-        return compile_range.end <= self.max_token_num
-
-    def uuid(self) -> str:
-        return VllmInductorPass.hash_source(self, RopeReshapeKVCachePattern)

vllm/compilation/passes/pass_manager.py

@@ -28,9 +28,7 @@ if current_platform.is_cuda_alike():
     from .fusion.attn_quant_fusion import AttnFusionPass
     from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass
     from .fusion.rms_quant_fusion import RMSNormQuantFusionPass
-    from .fusion.rope_kvcache_fusion import RopeKVCacheFusionPass
     from .fusion.sequence_parallelism import SequenceParallelismPass
-    from .utility.scatter_split_replace import ScatterSplitReplacementPass
     from .utility.split_coalescing import SplitCoalescingPass
 
 if current_platform.is_cuda():
@@ -138,11 +136,6 @@ class PostGradPassManager(CustomGraphPass):  # type: ignore[misc]
             if self.pass_config.fuse_act_padding and rocm_aiter_ops.is_enabled():
                 self.passes += [RocmAiterTritonAddRMSNormPadFusionPass(config)]
 
-            if self.pass_config.fuse_rope_kvcache:
-                self.passes += [SplitCoalescingPass(config)]
-                self.passes += [ScatterSplitReplacementPass(config)]
-                self.passes += [RopeKVCacheFusionPass(config)]
-
             if self.pass_config.fuse_attn_quant:
                 self.passes += [AttnFusionPass(config)]
 

vllm/compilation/passes/utility/fix_functionalization.py

@@ -162,24 +162,6 @@ class FixFunctionalizationPass(VllmInductorPass):
                     "position_ids",
                 )
                 self.defunctionalize(graph, node, mutated_args=mutated_args, args=args)
-            elif (
-                hasattr(torch.ops.vllm, "fused_rope_and_unified_kv_cache_update")
-                and at_target
-                == torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default
-            ):
-                mutated_args = {
-                    1: "query",
-                    2: "key",
-                }
-                self.defunctionalize(graph, node, mutated_args=mutated_args)
-            # only used for test_functionalization::TestFunctionWithMutatedArgsAndReturn
-            elif (
-                hasattr(torch.ops.vllm, "function_with_mutated_args_and_return")
-                and at_target
-                == torch.ops.vllm.function_with_mutated_args_and_return.default
-            ):
-                mutated_args = {1: "x"}
-                self.defunctionalize(graph, node, mutated_args=mutated_args)
             else:
                 continue  # skip the count
 
@@ -226,20 +208,13 @@ class FixFunctionalizationPass(VllmInductorPass):
         self, node: torch.fx.Node, mutated_args: dict[int, torch.fx.Node | str]
     ) -> None:
         """
-        Replace mutated getitem users of the auto-functionalized node with the
+        Replace all getitem users of the auto-functionalized node with the
         mutated arguments.
         :param node: The auto-functionalized node
         :param mutated_args: The mutated arguments, indexed by getitem index.
         If the value of an arg is a string, `node.kwargs[arg]` is used.
         """
         for idx, user in self.getitem_users(node).items():
-            # Some functionalized nodes may return both a result at getitem[0]
-            # as well as mutated args at getitem[1:...]
-            if idx == 0:
-                assert idx not in mutated_args, (
-                    f"result at getitem[0] should not be in mutated_args for {node}"
-                )
-                continue
             arg = mutated_args[idx]
             arg = node.kwargs[arg] if isinstance(arg, str) else arg
             user.replace_all_uses_with(arg)
@@ -282,20 +257,10 @@ class FixFunctionalizationPass(VllmInductorPass):
         with graph.inserting_before(node):
             function = node.args[0]
             if args is None:
-                fn_node = graph.call_function(function, kwargs=node.kwargs)
+                graph.call_function(function, kwargs=node.kwargs)
             else:
                 # Args passed as strings refer to items in node.kwargs
                 args = tuple(
                     node.kwargs[arg] if isinstance(arg, str) else arg for arg in args
                 )
-                fn_node = graph.call_function(function, args=args)
-
-        # If the function returns a value as well as mutating args inplace,
-        # the functionalized node will have a getitem[0] user that holds this value
-        # Replace getitem[0] user of the auto-functionalized node
-        # with the new defunctionalized node directly if it exists
-        users = self.getitem_users(node)
-        if 0 in users:
-            user = users[0]
-            user.replace_all_uses_with(fn_node)
-            self._remove(user)
+                graph.call_function(function, args=args)

vllm/compilation/passes/utility/scatter_split_replace.py

@@ -1,134 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Replace ``slice_scatter`` and ``split_with_sizes`` nodes with a single
-assignment if there are no users for the inplace tensor written to by
-the slice_scatter call.
-
-The inplace rotary_embedding custom op takes in mutable query and key inputs
-that are split+getitem outputs of a single qkv tensor.
-When functionalized, we fetch the rotated query and key from the functionalized op
-using `getitem` calls. However, we also write to the qkv tensor inplace using a
-`slice_scatter`, then split the inplace tensor to get the output tensors again.
-Instead, if the inplace tensor has no subsequent users, we can just replace the
-`slice_scatter` and `split_with_sizes` nodes with the `getitem` calls.
-
-This is already done in fix_functionalization::FixFunctionalizationPass, but
-writing a custom pass for it before defunctionalization allows matching against the
-qkv split+rotary_embedding subpattern as part of e.g. the RoPE+KVCache fusion pass.
-"""
-
-import operator
-
-import torch
-from torch import fx
-from torch._higher_order_ops.auto_functionalize import auto_functionalized
-
-from vllm.logger import init_logger
-
-from ..fx_utils import is_func
-from ..vllm_inductor_pass import VllmInductorPass
-
-logger = init_logger(__name__)
-
-
-class ScatterSplitReplacementPass(VllmInductorPass):
-    """Replace getitem+slice_scatter+split nodes with a single getitem when
-    the inplace subtensor written to by the slice_scatter has no other users.
-
-    Here's an example graph with q_size = 512, kv_size = 64:
-    split_with_sizes_1 = torch.ops.aten.split_with_sizes.default(qkv, (512, 64, 64), -1)
-    at = auto_functionalized(torch.ops._C.rotary_embedding.default(positions, q, k))
-    q = operator.getitem(at, 1)
-    k = operator.getitem(at, 2)
-    torch.ops.aten.slice_scatter.default(qkv, q, [0, 512], -1)
-    torch.ops.aten.slice_scatter.default(qkv, k, [512, 512 + 64], -1)
-    split_with_sizes_2 = torch.ops.aten.split_with_sizes.default(qkv, (512, 64, 64), -1)
-    q = operator.getitem(split_with_sizes_2, 0)
-    k = operator.getitem(split_with_sizes_2, 1)
-    v = operator.getitem(split_with_sizes_2, 2)
-
-    After this pass, this sequence of nodes is replaced with:
-    split_with_sizes_1 = torch.ops.aten.split_with_sizes.default(qkv, (512, 64, 64), -1)
-    at = auto_functionalized(torch.ops._C.rotary_embedding.default(positions, q, k))
-    q = operator.getitem(at, 1)
-    k = operator.getitem(at, 2)
-    v = operator.getitem(split_with_sizes_1, 2)
-    """
-
-    @VllmInductorPass.time_and_log
-    def __call__(self, graph: fx.Graph) -> None:
-        count = 0
-
-        for node in graph.nodes:
-            if not is_func(node, auto_functionalized):
-                continue
-
-            kwargs = node.kwargs
-            at_target = node.args[0]
-
-            if at_target == torch.ops._C.rotary_embedding.default:
-                query = kwargs["query"]
-                key = kwargs["key"]
-                getitem_nodes = {}
-                for user in node.users:
-                    if is_func(user, operator.getitem):
-                        getitem_nodes[user.args[1]] = user
-
-                if (
-                    is_func(query, operator.getitem)
-                    and is_func(key, operator.getitem)
-                    and query.args[0] == key.args[0]
-                    and is_func(query.args[0], torch.ops.aten.split_with_sizes.default)
-                    and all(
-                        is_func(user, torch.ops.aten.slice_scatter.default)
-                        for getitem_node in getitem_nodes.values()
-                        for user in getitem_node.users
-                    )
-                ):
-                    # Pattern where query and key are slices of a qkv tensor.
-                    # While functionalized, results at [1] and [2] are scattered
-                    # back into qkv, then split again to get query and key.
-                    # If the inplace tensor has no other users, we can replace
-                    # the slice_scatter+split nodes with the original results.
-                    for user in getitem_nodes[1].users:
-                        slice_scatter_1_node = user
-                    if not is_func(
-                        slice_scatter_1_node, torch.ops.aten.slice_scatter.default
-                    ):
-                        continue
-
-                    for user in getitem_nodes[2].users:
-                        slice_scatter_2_node = user
-                    if not is_func(
-                        slice_scatter_2_node, torch.ops.aten.slice_scatter.default
-                    ):
-                        continue
-
-                    for user in slice_scatter_2_node.users:
-                        split_node = user
-                    if not is_func(split_node, torch.ops.aten.split_with_sizes.default):
-                        continue
-
-                    split_getitem_users = {}
-                    for user in split_node.users:
-                        if is_func(user, operator.getitem):
-                            split_getitem_users[user.args[1]] = user
-
-                    # Replace query node
-                    split_getitem_users[0].replace_all_uses_with(getitem_nodes[1])
-                    graph.erase_node(split_getitem_users[0])
-                    # Replace key node
-                    split_getitem_users[1].replace_all_uses_with(getitem_nodes[2])
-                    graph.erase_node(split_getitem_users[1])
-                    # Redirect value node to original qkv tensor
-                    split_getitem_users[2].replace_input_with(split_node, query.args[0])
-
-                    # Erase unused nodes
-                    graph.erase_node(split_node)
-                    graph.erase_node(slice_scatter_2_node)
-                    graph.erase_node(slice_scatter_1_node)
-
-                    count += 1
-
-        logger.debug("Eliminated %d slice_scatter+split nodes", count)

vllm/config/attention.py

@@ -43,6 +43,9 @@ class AttentionConfig:
     disable_flashinfer_q_quantization: bool = False
     """If set, when using fp8 kv, do not quantize Q to fp8."""
 
+    use_prefill_query_quantization: bool = False
+    """If set, quantize query for attention in prefill."""
+
     def compute_hash(self) -> str:
         """
         Provide a hash that uniquely identifies all the configs

vllm/config/compilation.py

@@ -127,13 +127,6 @@ class PassConfig:
     # ROCm/AITER specific fusions
     fuse_act_padding: bool = Field(default=None)
     """Fuse the custom RMSNorm + padding ops."""
-    fuse_rope_kvcache: bool = Field(default=None)
-    """Fuse the QK rope + KV cache ops."""
-
-    rope_kvcache_fusion_max_token_num: int = 256
-    """The threshold for ROCm AITER RoPE+KVCache fusion e.g. for small batch decode.
-    Larger batch sizes e.g. during prefill will use the unfused kernels.
-    """
 
     fi_allreduce_fusion_max_size_mb: float | None = None
     """The threshold of the communicated tensor sizes under which
@@ -206,7 +199,6 @@ class PassConfig:
         "fuse_gemm_comms",
         "fuse_allreduce_rms",
         "fuse_act_padding",
-        "fuse_rope_kvcache",
         mode="wrap",
     )
     @classmethod
@@ -252,12 +244,6 @@ class PassConfig:
                 "The fusion will be disabled."
             )
             self.fuse_act_padding = False
-        if self.fuse_rope_kvcache and not current_platform.is_rocm():
-            logger.warning_once(
-                "KV cache fusion currently only enabled on ROCm. "
-                "The fusion will be disabled."
-            )
-            self.fuse_rope_kvcache = False
 
 
 class DynamicShapesType(str, enum.Enum):
@@ -839,19 +825,6 @@ class CompilationConfig:
             # TODO(zhuhaoran): support rope native forward match and remove this.
             # Linked issue: https://github.com/vllm-project/vllm/issues/28042
             self.custom_ops.append("+rotary_embedding")
-        if self.pass_config.fuse_rope_kvcache:
-            from vllm._aiter_ops import rocm_aiter_ops
-
-            if rocm_aiter_ops.is_triton_rotary_embed_enabled():
-                logger.warning(
-                    "Cannot use VLLM_ROCM_USE_AITER_TRITON_ROPE with "
-                    "fuse_rope_kvcache. Disabling fuse_rope_kvcache."
-                )
-                self.pass_config.fuse_rope_kvcache = False
-            else:
-                # TODO(Rohan138): support rope native forward match and remove this.
-                # Linked issue: https://github.com/vllm-project/vllm/issues/28042
-                self.custom_ops.append("+rotary_embedding")
 
         if (
             is_torch_equal_or_newer("2.9.0.dev")

vllm/config/ec_transfer.py

@@ -7,8 +7,8 @@ from typing import Any, Literal, get_args
 
 from vllm.config.utils import config
 
-ECProducer = Literal["ec_producer"]
-ECConsumer = Literal["ec_consumer"]
+ECProducer = Literal["ec_producer", "ec_both"]
+ECConsumer = Literal["ec_consumer", "ec_both"]
 ECRole = Literal[ECProducer, ECConsumer]
 
 
@@ -33,7 +33,7 @@ class ECTransferConfig:
 
     ec_role: ECRole | None = None
     """Whether this vLLM instance produces, consumes EC cache, or both. Choices
-    are 'ec_producer', 'ec_consumer'."""
+    are 'ec_producer', 'ec_consumer', 'ec_both'."""
 
     ec_rank: int | None = None
     """The rank of this vLLM instance in the EC cache transfer. Typical value:

vllm/config/model.py

@@ -297,6 +297,7 @@ class ModelConfig:
     multimodal_config: MultiModalConfig | None = None
     """Configuration for multimodal model. If `None`, this will be inferred
     from the architecture of `self.model`."""
+    language_model_only: InitVar[bool] = False
     limit_mm_per_prompt: InitVar[dict[str, int | dict[str, int]] | None] = None
     enable_mm_embeds: InitVar[bool | None] = None
     media_io_kwargs: InitVar[dict[str, dict[str, Any]] | None] = None
@@ -411,6 +412,7 @@ class ModelConfig:
     def __post_init__(
         self,
         # Multimodal config init vars
+        language_model_only: bool,
         limit_mm_per_prompt: dict[str, int | dict[str, int]] | None,
         enable_mm_embeds: bool | None,
         media_io_kwargs: dict[str, dict[str, Any]] | None,
@@ -576,6 +578,7 @@ class ModelConfig:
                 mm_encoder_tp_mode = "weights"
 
             mm_config_kwargs = dict(
+                language_model_only=language_model_only,
                 limit_per_prompt=limit_mm_per_prompt,
                 enable_mm_embeds=enable_mm_embeds,
                 media_io_kwargs=media_io_kwargs,
@@ -1116,6 +1119,9 @@ class ModelConfig:
     @cached_property
     def is_mm_prefix_lm(self) -> bool:
         """Whether to use bidirectional attention for mm positions."""
+        if hasattr(self.hf_config, "is_mm_prefix_lm"):
+            return bool(self.hf_config.is_mm_prefix_lm)
+        # fallback to list of known models
         MM_PREFIX_LM_MODELS = (
             "gemma3",
             "molmo2",
@@ -1218,8 +1224,8 @@ class ModelConfig:
             if attn_type_list:
                 return sum(t == 1 for t in attn_type_list[start:end])
 
-            # Hybrid model Qwen3Next
-            layer_types_value = getattr(self.hf_config, "layer_types", None)
+            # Hybrid model Qwen3Next Qwen3.5 Series
+            layer_types_value = getattr(self.hf_text_config, "layer_types", None)
             if layer_types_value is not None:
                 if block_type == "attention":
                     return sum(

vllm/config/multimodal.py

@@ -54,20 +54,24 @@ DummyOptions: TypeAlias = (
 class MultiModalConfig:
     """Controls the behavior of multimodal models."""
 
+    language_model_only: bool = False
+    """If True, disables all multimodal inputs by setting all modality limits to 0.
+    Equivalent to setting `--limit-mm-per-prompt` to 0 for every modality."""
     limit_per_prompt: dict[str, DummyOptions] = Field(default_factory=dict)
-    """The maximum number of input items and options allowed per 
-        prompt for each modality.
+    """The maximum number of input items and options allowed per
+    prompt for each modality.
+
     Defaults to 999 for each modality.
 
     Legacy format (count only):
         {"image": 16, "video": 2}
 
     Configurable format (with options):
-        {"video": {"count": 1, "num_frames": 32, "width": 512, "height": 512}, 
+        {"video": {"count": 1, "num_frames": 32, "width": 512, "height": 512},
         "image": {"count": 5, "width": 512, "height": 512}}
 
     Mixed format (combining both):
-        {"image": 16, "video": {"count": 1, "num_frames": 32, "width": 512, 
+        {"image": 16, "video": {"count": 1, "num_frames": 32, "width": 512,
         "height": 512}}
     """
     enable_mm_embeds: bool = False
@@ -215,6 +219,7 @@ class MultiModalConfig:
         the final hidden states.
         """
         factors: list[Any] = [
+            self.language_model_only,
             self.mm_encoder_attn_backend.name
             if self.mm_encoder_attn_backend is not None
             else None,
@@ -228,6 +233,9 @@ class MultiModalConfig:
         Get the maximum number of input items allowed per prompt
         for the given modality (backward compatible).
         """
+        if self.language_model_only:
+            return 0
+
         limit_data = self.limit_per_prompt.get(modality)
 
         if limit_data is None:

vllm/config/speculative.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Literal, get_args
 from pydantic import Field, SkipValidation, model_validator
 from typing_extensions import Self
 
+from vllm.config import LoadConfig
 from vllm.config.model import ModelConfig
 from vllm.config.parallel import ParallelConfig
 from vllm.config.utils import config
@@ -37,6 +38,7 @@ MTPModelTypes = Literal[
     "ernie_mtp",
     "exaone_moe_mtp",
     "qwen3_next_mtp",
+    "qwen3_5_mtp",
     "longcat_flash_mtp",
     "mtp",
     "pangu_ultra_moe_mtp",
@@ -159,6 +161,10 @@ class SpeculativeConfig:
     tokens with estimated probability (based on frequency counts) greater than
     or equal to this value."""
 
+    draft_load_config: LoadConfig | None = None
+    """Load config for the draft model. If not specified, will use the load
+    config from the target model."""
+
     def compute_hash(self) -> str:
         """
         WARNING: Whenever a new field is added to this config,
@@ -181,7 +187,7 @@ class SpeculativeConfig:
     @staticmethod
     def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
         initial_architecture = hf_config.architectures[0]
-        if hf_config.model_type in ("deepseek_v3", "deepseek_v32"):
+        if hf_config.model_type in ("deepseek_v3", "deepseek_v32", "glm_moe_dsa"):
             hf_config.model_type = "deepseek_mtp"
         if hf_config.model_type == "deepseek_mtp":
             n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
@@ -263,6 +269,16 @@ class SpeculativeConfig:
                 {"n_predict": n_predict, "architectures": ["ExaoneMoeMTP"]}
             )
 
+        if hf_config.model_type in ("qwen3_5", "qwen3_5_moe"):
+            is_moe = hf_config.model_type == "qwen3_5_moe"
+            hf_config.model_type = "qwen3_5_mtp"
+            n_predict = getattr(hf_config, "mtp_num_hidden_layers", None)
+            hf_config.update(
+                {
+                    "n_predict": n_predict,
+                    "architectures": ["Qwen3_5MoeMTP" if is_moe else "Qwen3_5MTP"],
+                }
+            )
         if hf_config.model_type == "longcat_flash":
             hf_config.model_type = "longcat_flash_mtp"
             n_predict = getattr(hf_config, "num_nextn_predict_layers", 1)

vllm/config/vllm.py

@@ -1365,20 +1365,6 @@ class VllmConfig:
                         "allreduce-rms fusion will be enabled for all num_tokens."
                     )
 
-        if compilation_config.pass_config.fuse_rope_kvcache:
-            max_token_num = (
-                compilation_config.pass_config.rope_kvcache_fusion_max_token_num
-            )
-            if max_token_num is not None:
-                if compile_range_end is not None and max_token_num < compile_range_end:
-                    computed_compile_ranges_split_points.append(max_token_num)
-                else:
-                    logger.debug(
-                        "Max num batched tokens below rope+kvcache fusion threshold, "
-                        "rope+kvcache fusion enabled for num_tokens <= %d.",
-                        compile_range_end,
-                    )
-
         if compilation_config.compile_ranges_split_points is not None:
             for x in compilation_config.compile_ranges_split_points:
                 assert isinstance(x, int)

vllm/distributed/device_communicators/shm_broadcast.py

@@ -488,6 +488,12 @@ class MessageQueue:
                 for i in range(1, self.buffer.n_reader + 1):
                     # set read flag to 0, meaning it is not read yet
                     metadata_buffer[i] = 0
+                # Memory fence here ensures the order of the buffer and flag
+                # writes. This guarantees that when `metadata_buffer[0] = 1` is
+                # visible to readers, `buf` can be completely ready. Without
+                # this, some CPU architectures with weak ordering may incur
+                # memory inconsistency.
+                memory_fence()
                 # mark the block as written
                 metadata_buffer[0] = 1
                 # Memory fence ensures the write is visible to readers on other cores

vllm/distributed/ec_transfer/ec_connector/base.py

@@ -63,6 +63,7 @@ class ECConnectorBase(ABC):
         self._role = role
         if vllm_config.ec_transfer_config is not None:
             self._is_producer = vllm_config.ec_transfer_config.is_ec_producer
+            self._is_consumer = vllm_config.ec_transfer_config.is_ec_consumer
         else:
             raise ValueError("ec_transfer_config must be set for ECConnectorBase")
 
@@ -74,6 +75,10 @@ class ECConnectorBase(ABC):
     def is_producer(self) -> bool:
         return self._is_producer
 
+    @property
+    def is_consumer(self) -> bool:
+        return self._is_consumer
+
     # ==============================
     # Worker-side methods
     # ==============================

vllm/distributed/eplb/async_worker.py

@@ -11,13 +11,13 @@ from typing import TYPE_CHECKING
 import torch
 from torch.distributed import ProcessGroup
 
-from vllm.distributed.parallel_state import get_ep_group
+from vllm.distributed.parallel_state import get_eplb_group
 from vllm.logger import init_logger
 
 from .rebalance_execute import transfer_layer
 
 if TYPE_CHECKING:
-    from .eplb_state import EplbState
+    from .eplb_state import EplbModelState, EplbState
 
 logger = init_logger(__name__)
 
@@ -27,8 +27,8 @@ def start_async_worker(
     rank_mapping: dict[int, int] | None = None,
     is_profile: bool = False,
 ) -> threading.Thread:
-    ep_group = get_ep_group().device_group
-    rank = ep_group.rank()
+    eplb_group = get_eplb_group().device_group
+    rank = eplb_group.rank()
     device_index = state.cuda_device_index
     assert state.is_async
 
@@ -42,7 +42,7 @@ def start_async_worker(
             loop.run_until_complete(
                 transfer_run_periodically(
                     state=state,
-                    ep_group=ep_group,
+                    eplb_group=eplb_group,
                     cuda_stream=cuda_stream,
                     is_profile=is_profile,
                     rank_mapping=rank_mapping,
@@ -58,9 +58,53 @@ def start_async_worker(
     return thread
 
 
+def run_rebalance_experts(
+    model_state: "EplbModelState",
+    eplb_state: "EplbState",
+    physical_to_logical_map_cpu: torch.Tensor,
+) -> None:
+    assert model_state.eplb_stats is not None
+    eplb_stats = model_state.eplb_stats
+
+    # Wait for the main thread's all-reduce and clone to complete before
+    # accessing the global_expert_load_window tensor.
+    assert model_state.window_ready_event is not None
+    model_state.window_ready_event.wait()
+    model_state.window_ready_event = None
+
+    # Move the global expert load window to CPU for computation.
+    global_expert_load_window = eplb_stats.global_expert_load_window.cpu()
+    # Compute new expert mappings for the model
+    (
+        new_physical_to_logical_map,
+        new_logical_to_physical_map,
+        new_logical_replica_count,
+    ) = eplb_state.policy.rebalance_experts(
+        global_expert_load_window,
+        eplb_stats.num_replicas,
+        eplb_stats.num_groups,
+        eplb_stats.num_nodes,
+        eplb_stats.num_gpus,
+        physical_to_logical_map_cpu,
+    )
+    assert new_physical_to_logical_map.device == torch.device("cpu")
+
+    model_state.new_physical_to_logical_map = new_physical_to_logical_map
+
+    max_slots = model_state.logical_to_physical_map.shape[-1]
+    padded_logical = torch.nn.functional.pad(
+        new_logical_to_physical_map,
+        (0, max(0, max_slots - new_logical_to_physical_map.shape[-1])),
+        value=-1,
+    ).to(model_state.logical_to_physical_map.device)
+    new_replica = new_logical_replica_count.to(model_state.logical_replica_count.device)
+    model_state.new_logical_to_physical_map = padded_logical
+    model_state.new_logical_replica_count = new_replica
+
+
 async def transfer_run_periodically(
     state: "EplbState",
-    ep_group: ProcessGroup,
+    eplb_group: ProcessGroup,
     cuda_stream: torch.cuda.Stream,
     is_profile: bool = False,
     rank_mapping: dict[int, int] | None = None,
@@ -71,23 +115,51 @@ async def transfer_run_periodically(
 
         assert state.is_async
         for model_state in state.model_states.values():
+            rebalancing_algorithm_executed = False
+            physical_to_logical_map_cpu = None
             current_num_layers = model_state.model.num_moe_layers
             while (
                 model_state.rebalanced
                 and model_state.layer_to_transfer < current_num_layers
             ):
-                if (
-                    not model_state.ep_buffer_ready
-                    and model_state.rebalanced
-                    and model_state.new_physical_to_logical_map is not None
-                ):
-                    await asyncio.to_thread(model_state.buffer_lock.acquire)
+                if not model_state.ep_buffer_ready and model_state.rebalanced:
+                    # Polling the lock directly in the async thread avoids
+                    # the thread switch overhead of asyncio.to_thread.
+                    # This is typically faster than offloading to a worker thread.
+                    while not model_state.buffer_lock.acquire(blocking=False):
+                        await asyncio.sleep(0)
                     try:
                         if model_state.layer_to_transfer >= current_num_layers:
                             break
+                        if (
+                            not rebalancing_algorithm_executed
+                            or model_state.new_physical_to_logical_map is None
+                        ):
+                            # Move the physical_to_logical_map to CPU
+                            # for rebalancing and transfer_layer.
+                            physical_to_logical_map_cpu = (
+                                model_state.physical_to_logical_map.cpu()
+                            )
+                            run_rebalance_experts(
+                                model_state, state, physical_to_logical_map_cpu
+                            )
+                            rebalancing_algorithm_executed = True
+                            logger.info(
+                                "Async worker computed new indices for model %s",
+                                model_state.model_name,
+                            )
+
+                        assert model_state.new_physical_to_logical_map is not None
+                        assert physical_to_logical_map_cpu is not None
+
+                        layer_idx = model_state.layer_to_transfer
+                        old_layer_indices = physical_to_logical_map_cpu[layer_idx]
+                        new_layer_indices = model_state.new_physical_to_logical_map[
+                            layer_idx
+                        ]
 
                         # Wait for the main thread to finish consuming the buffer
-                        # before overwriting it
+                        # before initiating an EPLB transfer on another layer.
                         if model_state.buffer_consumed_event is not None:
                             cuda_stream.wait_event(model_state.buffer_consumed_event)
                             model_state.buffer_consumed_event = None
@@ -97,13 +169,12 @@ async def transfer_run_periodically(
                             model_state.is_received_locally,
                             model_state.recv_metadata,
                         ) = await transfer_layer(
-                            old_global_expert_indices=model_state.physical_to_logical_map,
-                            new_global_expert_indices=model_state.new_physical_to_logical_map,
-                            expert_weights=model_state.model.expert_weights,
+                            old_layer_indices=old_layer_indices,
+                            new_layer_indices=new_layer_indices,
+                            expert_weights=model_state.model.expert_weights[layer_idx],
                             expert_weights_buffer=model_state.expert_buffer,
-                            ep_group=ep_group,
+                            ep_group=eplb_group,
                             is_profile=is_profile,
-                            layer=model_state.layer_to_transfer,
                             cuda_stream=cuda_stream,
                             rank_mapping=rank_mapping,
                         )

vllm/distributed/eplb/eplb_state.py

@@ -55,6 +55,35 @@ from .rebalance_execute import (
 logger = init_logger(__name__)
 
 
+@dataclass
+class EplbStats:
+    """
+    Model stats used in EPLB rebalancing algorithm.
+    """
+
+    global_expert_load_window: torch.Tensor
+    """
+    Experts load window.
+    Shape: (window_size, num_moe_layers, num_physical_experts)
+    """
+    num_replicas: int
+    """
+    Number of physical experts.
+    """
+    num_groups: int
+    """
+    Number of expert groups.
+    """
+    num_nodes: int
+    """
+    Number of nodes.
+    """
+    num_gpus: int
+    """
+    Number of GPUs.
+    """
+
+
 @dataclass
 class EplbModelState:
     """EPLB metrics."""
@@ -156,6 +185,11 @@ class EplbModelState:
     CUDA event recorded after the main thread finishes consuming the buffer.
     The async worker waits on this before writing to the buffer again.
     """
+    window_ready_event: torch.cuda.Event | None
+    """
+    CUDA event recorded after all-reduce and clone on the main thread.
+    The async worker waits on this before accessing global_expert_load_window.
+    """
     ep_buffer_ready: int
     """
     The flag indicates whether the expert buffer is ready for transfer.
@@ -173,6 +207,10 @@ class EplbModelState:
     """
     Whether the async EPLB needs to poll peers for buffer readiness.
     """
+    eplb_stats: EplbStats | None
+    """
+    EPLB stats for the model.
+    """
     is_unchanged: np.ndarray
     """
     intermediate variable between `move_to_buffer` and `move_to_workspace`.
@@ -508,10 +546,12 @@ class EplbState:
             buffer_lock=threading.Lock(),
             buffer_ready_event=None,
             buffer_consumed_event=None,
+            window_ready_event=None,
             ep_buffer_ready=0,
             layer_to_transfer=0,
             rebalanced=False,
             pending_global_ready_check=False,
+            eplb_stats=None,
             is_unchanged=np.array([]),
             is_received_locally=np.array([]),
             recv_metadata=RecvMetadata(
@@ -642,20 +682,6 @@ class EplbState:
                         ep_group=ep_group,
                         is_profile=is_profile,
                     )
-                    if (
-                        eplb_model_state.layer_to_transfer
-                        >= eplb_model_state.model.num_moe_layers
-                    ):
-                        self.post_eplb(eplb_model_state, is_profile)
-                        eplb_model_state.rebalanced = False
-                        eplb_model_state.layer_to_transfer = 0
-                        eplb_model_state.pending_global_ready_check = False
-                        logger.info(
-                            "finish async transfer for model %s rank %d layer %d",
-                            eplb_model_state.model_name,
-                            ep_group.rank(),
-                            eplb_model_state.model.num_moe_layers,
-                        )
 
         if self.expert_rearrangement_step >= self.expert_rearrangement_step_interval:
             if self.is_async and any(
@@ -802,21 +828,21 @@ class EplbState:
         for eplb_model_state, global_expert_load_window in zip(
             self.model_states.values(), global_expert_load_windows
         ):
-            # Get new expert mappings for the model
-            (
-                new_physical_to_logical_map,
-                new_logical_to_physical_map,
-                new_logical_replica_count,
-            ) = self.policy.rebalance_experts(
-                global_expert_load_window,
-                num_replicas,
-                num_groups,
-                num_nodes,
-                num_gpus,
-                eplb_model_state.physical_to_logical_map,
-            )
-
             if not self.is_async or is_profile:
+                # Get new expert mappings for the model
+                (
+                    new_physical_to_logical_map,
+                    new_logical_to_physical_map,
+                    new_logical_replica_count,
+                ) = self.policy.rebalance_experts(
+                    global_expert_load_window,
+                    num_replicas,
+                    num_groups,
+                    num_nodes,
+                    num_gpus,
+                    eplb_model_state.physical_to_logical_map,
+                )
+
                 # Update expert weights
                 rearrange_expert_weights_inplace(
                     eplb_model_state.physical_to_logical_map,
@@ -873,27 +899,25 @@ class EplbState:
                         gpu_elapsed,
                     )
             else:
-                max_slots = eplb_model_state.logical_to_physical_map.shape[-1]
-                padded_logical = torch.nn.functional.pad(
-                    new_logical_to_physical_map,
-                    (0, max(0, max_slots - new_logical_to_physical_map.shape[-1])),
-                    value=-1,
-                ).to(eplb_model_state.logical_to_physical_map.device)
-                new_replica = new_logical_replica_count.to(
-                    eplb_model_state.logical_replica_count.device
+                eplb_model_state.eplb_stats = EplbStats(
+                    # We copy the tensor to snapshot the global_expert_load_window
+                    # on the main thread so that async worker can access it safely
+                    # while the main thread is running.
+                    global_expert_load_window=global_expert_load_window.clone(),
+                    num_replicas=num_replicas,
+                    num_groups=num_groups,
+                    num_nodes=num_nodes,
+                    num_gpus=num_gpus,
                 )
-
-                # Move map to cpu in advance
-                eplb_model_state.new_physical_to_logical_map = (
-                    new_physical_to_logical_map.cpu()
-                )
-                eplb_model_state.new_logical_to_physical_map = padded_logical
-                eplb_model_state.new_logical_replica_count = new_replica
+                # Record event after clone to signal async worker
+                # that load stats data is ready
+                sync_event = torch.cuda.Event()
+                sync_event.record()
+                eplb_model_state.window_ready_event = sync_event
 
                 eplb_model_state.rebalanced = True
                 eplb_model_state.layer_to_transfer = 0
                 eplb_model_state.pending_global_ready_check = True
-
         # Signal async thread to start transferring layers
         if self.is_async and (not is_profile):
             self.rearrange_event.set()
@@ -925,11 +949,13 @@ class EplbState:
 
         target_device = model_state.physical_to_logical_map.device
         new_physical = model_state.new_physical_to_logical_map
+        # If the number of physical experts has changed, then the new map needs to
+        # be copied synchronously to avoid a race condition with the async worker
         if model_state.physical_to_logical_map.shape[1] != new_physical.shape[1]:
             model_state.physical_to_logical_map = new_physical.to(target_device)
         else:
             model_state.physical_to_logical_map[layer].copy_(
-                new_physical[layer].to(target_device)
+                new_physical[layer].to(target_device, non_blocking=True)
             )
 
         logical_device = model_state.logical_to_physical_map.device
@@ -1004,11 +1030,9 @@ class EplbState:
                 model_state.layer_to_transfer
             ]
             expert_weights_buffer = model_state.expert_buffer
-            new_indices = (
-                model_state.new_physical_to_logical_map[model_state.layer_to_transfer]
-                .cpu()
-                .numpy()
-            )
+            new_indices = model_state.new_physical_to_logical_map[
+                model_state.layer_to_transfer
+            ].numpy()
             move_from_buffer(
                 expert_weights=expert_weights,
                 expert_weights_buffers=expert_weights_buffer,
@@ -1019,7 +1043,7 @@ class EplbState:
                 ep_rank=ep_group.rank(),
             )
             # Record event after consuming buffer to signal async thread
-            # that it's safe to overwrite the buffer
+            # that it's safe to overwrite the intermediate buffer
             consumed_event = torch.cuda.Event()
             consumed_event.record()
             model_state.buffer_consumed_event = consumed_event
@@ -1034,6 +1058,18 @@ class EplbState:
                 model_state.model_name,
                 transferred_layer,
             )
+            if model_state.layer_to_transfer >= model_state.model.num_moe_layers:
+                self.post_eplb(model_state, is_profile)
+                model_state.rebalanced = False
+                model_state.layer_to_transfer = 0
+                model_state.pending_global_ready_check = False
+                logger.info(
+                    "finish async transfer for model %s rank %d layer %d",
+                    model_state.model_name,
+                    ep_group.rank(),
+                    model_state.model.num_moe_layers,
+                )
+
         finally:
             try:
                 model_state.buffer_lock.release()
@@ -1048,9 +1084,7 @@ class EplbState:
         assert model_state.new_physical_to_logical_map is not None
         assert model_state.new_logical_to_physical_map is not None
         assert model_state.new_logical_replica_count is not None
-        if not is_profile:
-            for layer_idx in range(model_state.physical_to_logical_map.shape[0]):
-                self._update_layer_mapping_from_new(model_state, layer_idx)
+
         model_state.new_physical_to_logical_map = None
         model_state.new_logical_to_physical_map = None
         model_state.new_logical_replica_count = None

vllm/distributed/eplb/eplb_utils.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Utility functions for EPLB (Expert Parallel Load Balancing)."""
+
+import os
+
+from vllm.config import ParallelConfig
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def override_envs_for_eplb(parallel_config: ParallelConfig) -> None:
+    """
+    Override environment variables for EPLB when specific conditions are met.
+
+    Args:
+        parallel_config: The parallel configuration object.
+    """
+    is_data_parallel = parallel_config.data_parallel_size > 1
+    is_eplb_enabled = parallel_config.enable_eplb
+    async_eplb = parallel_config.eplb_config.use_async
+    is_deepep_ll = parallel_config.all2all_backend == "deepep_low_latency"
+
+    # Override NCCL_MAX_CTAS to avoid hangs when using async EPLB with the
+    # DeepEP low-latency backend.
+    #
+    # The hang happens when two ranks interleave kernel launches differently
+    # between NCCL collectives (used by async EPLB weight exchange) and DeepEP
+    # low-latency (LL) kernels. DeepEP LL uses a cooperative launch and tries
+    # to reserve a large fraction of the GPU's SMs; if those SMs are currently
+    # occupied by NCCL, the DeepEP LL launch blocks until enough SMs are
+    # freed.
+    #
+    # If rank A enters DeepEP LL in main thread while rank B is still executing
+    # NCCL in async thread, rank A can block waiting for SMs, while rank B can
+    # block inside NCCL waiting for rank A to participate in the collective.
+    # This circular wait causes a deadlock.
+    # Limiting NCCL occupancy via NCCL_MAX_CTAS leaves space for the DeepEP
+    # cooperative kernel to launch and complete, breaking the deadlock.
+    # See: https://github.com/deepseek-ai/DeepEP/issues/496
+    if is_data_parallel and is_eplb_enabled and is_deepep_ll and async_eplb:
+        current_value_str = os.getenv("NCCL_MAX_CTAS")
+
+        if current_value_str and current_value_str.isdigit():
+            return
+
+        override_value = 8
+        os.environ["NCCL_MAX_CTAS"] = str(override_value)
+        logger.info_once(
+            f"EPLB: Setting NCCL_MAX_CTAS={override_value} "
+            "for expert parallel with EPLB and deepep_low_latency backend",
+            scope="global",
+        )

vllm/distributed/eplb/rebalance_execute.py

@@ -434,13 +434,12 @@ def move_from_buffer(
 
 
 async def transfer_layer(
-    old_global_expert_indices: torch.Tensor,
-    new_global_expert_indices: torch.Tensor,
-    expert_weights: Sequence[Sequence[torch.Tensor]],
+    old_layer_indices: torch.Tensor,
+    new_layer_indices: torch.Tensor,
+    expert_weights: Sequence[torch.Tensor],
     expert_weights_buffer: Sequence[torch.Tensor],
     ep_group: ProcessGroup,
     is_profile: bool = False,
-    layer: int = 0,
     cuda_stream: torch.cuda.Stream | None = None,
     rank_mapping: dict[int, int] | None = None,
 ) -> MoveToBufferResult:
@@ -451,56 +450,64 @@ async def transfer_layer(
     while keys are physical.
 
     Args:
-        old_global_expert_indices: Shape (num_moe_layers, num_physical_experts).
-        new_global_expert_indices: Shape (num_moe_layers, num_physical_experts).
-        expert_weights: A sequence of shape (num_moe_layers)(weight_count)
-            of tensors of shape (num_local_physical_experts, hidden_size_i).
-            For example, a linear layer may have up and down projection,
-            so weight_count = 2. Each weight's hidden size can be different.
+        old_layer_indices: Shape (num_physical_experts,).
+        new_layer_indices: Shape (num_physical_experts,).
+        expert_weights: Iterable of weight tensors for this layer, each with shape
+            (num_local_physical_experts, hidden_size_i).
+            For example, a linear layer may have up and down projection.
+        expert_weights_buffer: Intermediate buffers (one per weight tensor).
         ep_group: The device process group for expert parallelism.
         is_profile (bool): If `True`, do not perform any actual weight copy.
             This is used during profile run, where we only perform dummy
             communications to reserve enough memory for the buffers.
+        cuda_stream: CUDA stream for async copies (can be None for sync mode).
+        rank_mapping: Optional rank mapping for elastic expert parallelism.
 
     Returns:
-        is_unchanged (np.ndarray): (1, num_local_experts), True where expert
+        is_unchanged (np.ndarray): (num_local_experts,), True where expert
             is left unchanged.
-        is_received_locally (np.ndarray): (1, num_local_experts), True where expert
+        is_received_locally (np.ndarray): (num_local_experts,), True where expert
             can be received locally.
         RecvMetadata: Metadata needed for completing remote weight transfers.
     """
     ep_size = ep_group.size()
     if rank_mapping is not None:
+        # Add a layer dimension for compatibility with mapping functions
+        old_layer_indices_2d = old_layer_indices.unsqueeze(0)
+        new_layer_indices_2d = new_layer_indices.unsqueeze(0)
+
         if len(rank_mapping) == ep_group.size():
             # scale down
-            new_global_expert_indices = _map_new_expert_indices_with_rank_mapping(
-                new_global_expert_indices,
+            new_layer_indices_2d = _map_new_expert_indices_with_rank_mapping(
+                new_layer_indices_2d,
                 rank_mapping,
             )
         else:
             # scale up
-            old_global_expert_indices = _map_old_expert_indices_with_rank_mapping(
-                old_global_expert_indices,
+            old_layer_indices_2d = _map_old_expert_indices_with_rank_mapping(
+                old_layer_indices_2d,
                 rank_mapping,
                 ep_group.size(),
             )
 
-    assert old_global_expert_indices.shape[1] == new_global_expert_indices.shape[1]
-    num_moe_layers, num_physical_experts = old_global_expert_indices.shape
-    assert len(expert_weights) == num_moe_layers
+        # Remove the layer dimension
+        old_layer_indices = old_layer_indices_2d.squeeze(0)
+        new_layer_indices = new_layer_indices_2d.squeeze(0)
+
+    assert old_layer_indices.shape == new_layer_indices.shape
+    num_physical_experts = old_layer_indices.shape[0]
     assert len(expert_weights[0]) >= 1
-    num_local_physical_experts = expert_weights[0][0].shape[0]
-    assert new_global_expert_indices.shape == (num_moe_layers, num_physical_experts)
+    num_local_physical_experts = expert_weights[0].shape[0]
     assert num_physical_experts == ep_size * num_local_physical_experts
 
-    old_global_expert_indices_np = old_global_expert_indices.cpu().numpy()
-    new_global_expert_indices_np = new_global_expert_indices.cpu().numpy()
+    old_layer_indices_np = old_layer_indices.cpu().numpy()
+    new_layer_indices_np = new_layer_indices.cpu().numpy()
 
     is_unchanged, is_received_locally, recv_metadata = move_to_buffer(
         num_local_experts=num_local_physical_experts,
-        old_indices=old_global_expert_indices_np[layer],
-        new_indices=new_global_expert_indices_np[layer],
-        expert_weights=expert_weights[layer],
+        old_indices=old_layer_indices_np,
+        new_indices=new_layer_indices_np,
+        expert_weights=expert_weights,
         expert_weights_buffers=expert_weights_buffer,
         cuda_stream=cuda_stream,
         ep_group=ep_group,

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py

@@ -20,16 +20,42 @@ from lmcache.v1.multiprocess.protocol import RequestType, get_response_class
 logger = init_logger(__name__)
 
 
-def wrap_kv_caches(kv_caches: dict[str, KVCache]) -> KVCache:
+def wrap_kv_caches(kv_caches: dict[str, torch.Tensor]) -> KVCache:
     logger.info("KV caches keys are %s", list(kv_caches.keys()))
     return [CudaIPCWrapper(tensor) for tensor in kv_caches.values()]
 
 
+def striding_block_hashes(
+    block_hashes: list[bytes], blocks_in_chunk: int
+) -> Iterable[bytes]:
+    """Extract chunk-level hashes from block hashes by striding.
+
+    In hash-based vLLM, each vLLM block has its own hash.  LMCache chunks
+    span ``blocks_in_chunk`` consecutive blocks.  The representative hash
+    for a chunk is the hash of the **last** block in that chunk (because
+    each block hash already encodes its prefix).  So we start at index
+    ``blocks_in_chunk - 1`` and stride by ``blocks_in_chunk``.
+    """
+    return islice(block_hashes, blocks_in_chunk - 1, None, blocks_in_chunk)
+
+
 def send_lmcache_request(
     mq_client: MessageQueueClient,
     request_type: RequestType,
     payloads: list[Any],
 ) -> MessagingFuture[Any]:
+    """
+    Helper function to send the request to the LMCache multiprocess server
+
+    Args:
+        mq_client: The LMCache multiprocess mode message queue client
+        request_type: The request type
+        payloads: The request payloads
+
+    Returns:
+        A messaging future for the request
+    """
+
     future = mq_client.submit_request(
         request_type, payloads, get_response_class(request_type)
     )
@@ -39,40 +65,44 @@ def send_lmcache_request(
 def get_lmcache_chunk_size(
     mq_client: MessageQueueClient,
 ) -> int:
+    """
+    Helper function to get the LMCache chunk size from the server
+
+    Args:
+        mq_client: The LMCache multiprocess mode message queue client
+
+    Returns:
+        An integer representing the LMCache chunk size
+    """
     future = send_lmcache_request(mq_client, RequestType.GET_CHUNK_SIZE, [])
     chunk_size = future.result()
     return chunk_size
 
 
-def striding_block_hashes(
-    block_hashes: list[bytes],
-    blocks_in_chunk,
-) -> Iterable[bytes]:
-    """Striding the block hashes to get the block hashes for each chunk.
-    For example, if blocks_in_chunk is 16, then we will get the block hashes
-    for the 16th, 32nd, 48th, ... blocks.
-    """
-    return islice(block_hashes, blocks_in_chunk - 1, None, blocks_in_chunk)
-
-
 @dataclass
 class LoadStoreOp:
-    block_hashes: list[bytes]
     block_ids: list[int]
+    """Block ids for the load/store operation"""
+
+    token_ids: list[int] | None = None
+    """Token IDs for the load/store operation (token mode)"""
+
+    block_hashes: list[bytes] | None = None
+    """Block hashes for the load/store operation (hash mode)"""
+
+    start: int = 0
+    """Start token index (token mode only)"""
+
+    end: int = 0
+    """End token index (token mode only)"""
 
     def __len__(self) -> int:
-        return len(self.block_hashes)
-
-    def __post_init__(self):
-        assert len(self.block_hashes) == len(self.block_ids), (
-            "The number of block hashes should be equal to the number of block ids "
-            f"But got {len(self.block_hashes)} and {len(self.block_ids)}"
-        )
+        return len(self.block_ids)
 
 
 StoreResult = bool
 RetrieveResult = list[bool]
-LookupResult = list[bool]
+LookupResult = int
 
 
 class LMCacheMPSchedulerAdapter:
@@ -95,10 +125,6 @@ class LMCacheMPSchedulerAdapter:
             kv_rank: The kv rank used for LMCache keys
             vllm_block_size: The block size used in vLLM
         """
-        logger.warning(
-            "Importing LMCacheMPSchedulerAdapter is deprecated. "
-            "Please update your LMCache to the latest version."
-        )
         self.mq_client = MessageQueueClient(server_url, context)
 
         # Request futures
@@ -116,22 +142,89 @@ class LMCacheMPSchedulerAdapter:
         self.blocks_in_chunk = self.chunk_size // vllm_block_size
 
     @_lmcache_nvtx_annotate
-    def maybe_submit_lookup_request(self, request_id: str, block_hashes: list[bytes]):
+    def maybe_submit_lookup_request(
+        self,
+        request_id: str,
+        block_hashes: list[bytes] | None = None,
+        token_ids: list[int] | None = None,
+    ) -> None:
+        """
+        Submit a new lookup request to LMCache if there is no ongoing request.
+
+        Supports both token-based and hash-based vLLM:
+        - token_ids: token IDs (token-based vLLM) -> single token-mode key
+        - block_hashes: block hashes (hash-based vLLM) -> strided hash-mode keys
+
+        Exactly one of block_hashes or token_ids must be provided.
+
+        Args:
+            request_id: The ID of the lookup request. The same ID indicates it's
+                from the same request
+            block_hashes: Block hashes to lookup from LMCache (hash mode)
+            token_ids: Token IDs to lookup from LMCache (token mode)
+
+        Returns:
+            None
+
+        Notes:
+            This function will have a side-effect: submitting a look up request to
+            LMCache, which will essentially 'lock' the KV cache chunks in the LMCache
+            for later retrieve operations.
+            In the meantime, this function will record the lookup request, and the
+            status of the look up request can be checked by `check_lookup_result`.
+        """
         if request_id in self.lookup_futures:
             # Skip if there is already a lookup request
             return
 
-        s = striding_block_hashes(block_hashes, self.blocks_in_chunk)
-        keys = [self._create_key(block_hash) for block_hash in s]
+        assert (block_hashes is None) != (token_ids is None), (
+            "Exactly one of block_hashes or token_ids must be provided"
+        )
+
+        if block_hashes is not None:
+            # Hash mode: stride block hashes -> N hash-mode keys
+            chunk_hashes = list(
+                striding_block_hashes(block_hashes, self.blocks_in_chunk)
+            )
+            keys = [
+                self._create_hash_key(ch, request_id=request_id) for ch in chunk_hashes
+            ]
+        else:
+            # Token mode: truncate to chunk-aligned length
+            assert token_ids is not None
+            aligned_end = (len(token_ids) // self.chunk_size) * self.chunk_size
+            if aligned_end == 0:
+                return
+            keys = [
+                self._create_key(
+                    token_ids,
+                    start=0,
+                    end=aligned_end,
+                    request_id=request_id,
+                ).no_worker_id_version()
+            ]
+
         future = send_lmcache_request(
             self.mq_client,
             RequestType.LOOKUP,
-            [keys, True],
+            [keys],
         )
         self.lookup_futures[request_id] = future
 
     @_lmcache_nvtx_annotate
     def check_lookup_result(self, request_id: str) -> int | None:
+        """
+        Check the result of a previously submitted lookup request.
+
+        Args:
+            request_id: The ID of the lookup request submitted in
+                `maybe_submit_lookup_request`
+
+        Returns:
+            An integer representing the total number of tokens matched
+            in LMCache (prefix matching), or
+            None if the lookup request is not finished yet.
+        """
         assert request_id in self.lookup_futures, (
             f"Lookup request for request_id={request_id} has not been submitted"
         )
@@ -141,7 +234,7 @@ class LMCacheMPSchedulerAdapter:
             return None
 
         result = future.result()
-        num_chunks = sum(result)
+        num_chunks = result
         return num_chunks * self.chunk_size
 
     def num_blocks_per_chunk(self) -> int:
@@ -159,14 +252,47 @@ class LMCacheMPSchedulerAdapter:
         """
         self.lookup_futures.pop(request_id, None)
 
+    def end_session(self, request_id: str) -> None:
+        """
+        Notify LMCache server to remove the session for a finished request.
+        Args:
+            request_id: The ID of the finished request.
+        """
+        send_lmcache_request(
+            self.mq_client,
+            RequestType.END_SESSION,
+            [request_id],
+        )
+
     # Helper functions
-    def _create_key(self, block_hash: bytes) -> IPCCacheEngineKey:
-        """Convert a block hash to an IPC cache engine key"""
+    def _create_key(
+        self,
+        token_ids: list[int],
+        start: int = 0,
+        end: int = 0,
+        request_id: str | None = None,
+    ) -> IPCCacheEngineKey:
+        """Convert token IDs to an IPC cache engine key"""
         return IPCCacheEngineKey(
             model_name=self.model_name,
             world_size=self.world_size,
             worker_id=self.worker_id,
-            chunk_hash=block_hash,
+            token_ids=tuple(token_ids),
+            start=start,
+            end=end,
+            request_id=request_id,
+        )
+
+    def _create_hash_key(
+        self, chunk_hash: bytes, request_id: str | None = None
+    ) -> IPCCacheEngineKey:
+        """Create a hash-mode IPC cache engine key"""
+        return IPCCacheEngineKey(
+            model_name=self.model_name,
+            world_size=self.world_size,
+            worker_id=None,
+            chunk_hash=chunk_hash,
+            request_id=request_id,
         )
 
 
@@ -180,10 +306,6 @@ class LMCacheMPWorkerAdapter:
         kv_rank: int,
         vllm_block_size: int,
     ):
-        logger.warning(
-            "Importing LMCacheMPWorkerAdapter is deprecated. "
-            "Please update your LMCache to the latest version."
-        )
         self.mq_client = MessageQueueClient(server_url, context)
 
         # Instance id for GPU worker
@@ -201,7 +323,10 @@ class LMCacheMPWorkerAdapter:
             str, tuple[MessagingFuture[RetrieveResult], list[str]]
         ] = {}
 
+        # The store requests that have finished execution in LMCache
         self.finished_stores: set[str] = set()
+        # The finished request ids that are passed via vLLM and also
+        # have corresponding store requests submitted to LMCache before
         self.previously_finished: set[str] = set()
 
         self.model_name = model_name
@@ -215,7 +340,14 @@ class LMCacheMPWorkerAdapter:
         )
         self.blocks_in_chunk = chunk_size // vllm_block_size
 
-    def register_kv_caches(self, kv_caches: dict[str, KVCache]):
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+        """
+        Register the kv caches with LMCache server
+
+        Args:
+            kv_caches: A dict of kv caches to register. The keys are the
+                layer names and the values are the corresponding tensors.
+        """
         # Register kv cache and send the request
         self.kv_caches = kv_caches
         logger.info("Registering kv caches")
@@ -230,7 +362,29 @@ class LMCacheMPWorkerAdapter:
     def submit_store_request(
         self, request_id: str, op: LoadStoreOp, event: torch.cuda.Event
     ):
-        keys = self._block_hashes_to_keys(op.block_hashes)
+        """
+        Submit a KV cache store request to LMCache
+
+        Args:
+            request_id: The ID of the request
+            op: The LoadStoreOp describing the store operation.
+            event: The CUDA event that is recorded after the current
+                model inference step
+        """
+        if op.block_hashes is not None:
+            # Hash mode
+            chunk_hashes = list(
+                striding_block_hashes(op.block_hashes, self.blocks_in_chunk)
+            )
+            keys = [
+                self._create_hash_key(ch, request_id=request_id) for ch in chunk_hashes
+            ]
+        else:
+            # Token mode
+            assert op.token_ids is not None
+            keys = [
+                self._create_key(op.token_ids, op.start, op.end, request_id=request_id)
+            ]
         future = send_lmcache_request(
             self.mq_client,
             RequestType.STORE,
@@ -242,7 +396,29 @@ class LMCacheMPWorkerAdapter:
     def submit_retrieve_request(
         self, request_id: str, op: LoadStoreOp, event: torch.cuda.Event
     ):
-        keys = self._block_hashes_to_keys(op.block_hashes)
+        """
+        Submit a KV cache retrieve request to LMCache
+
+        Args:
+            request_id: The ID of the request
+            op: The LoadStoreOp describing the retrieve operation.
+            event: The CUDA event that is recorded after the current
+                model inference step
+        """
+        if op.block_hashes is not None:
+            # Hash mode
+            chunk_hashes = list(
+                striding_block_hashes(op.block_hashes, self.blocks_in_chunk)
+            )
+            keys = [
+                self._create_hash_key(ch, request_id=request_id) for ch in chunk_hashes
+            ]
+        else:
+            # Token mode
+            assert op.token_ids is not None
+            keys = [
+                self._create_key(op.token_ids, op.start, op.end, request_id=request_id)
+            ]
         future = send_lmcache_request(
             self.mq_client,
             RequestType.RETRIEVE,
@@ -257,17 +433,47 @@ class LMCacheMPWorkerAdapter:
         ops: list[LoadStoreOp],
         event: torch.cuda.Event,
     ):
-        keys = []
-        block_ids = []
-        for op in ops:
-            keys.extend(self._block_hashes_to_keys(op.block_hashes))
+        """
+        Submit a batched store request to LMCache
+
+        Args:
+            request_ids: The IDs of the requests
+            ops: The LoadStoreOps describing the store operations. Should have
+                the same length as request_ids
+            event: The CUDA event that is recorded after the current
+                model inference step
+        """
+        all_keys: list[IPCCacheEngineKey] = []
+        block_ids: list[int] = []
+        for request_id, op in zip(request_ids, ops, strict=False):
+            if op.block_hashes is not None:
+                chunk_hashes = list(
+                    striding_block_hashes(op.block_hashes, self.blocks_in_chunk)
+                )
+                keys = [
+                    self._create_hash_key(ch, request_id=request_id)
+                    for ch in chunk_hashes
+                ]
+                all_keys.extend(keys)
+            else:
+                assert op.token_ids is not None
+                all_keys.append(
+                    self._create_key(
+                        op.token_ids, op.start, op.end, request_id=request_id
+                    )
+                )
             block_ids.extend(op.block_ids)
         future = send_lmcache_request(
             self.mq_client,
             RequestType.STORE,
-            [keys, self.instance_id, block_ids, event.ipc_handle()],
+            [
+                all_keys,
+                self.instance_id,
+                block_ids,
+                event.ipc_handle(),
+            ],
         ).to_cuda_future()
-        self.store_futures[request_ids[0]] = (future, request_ids[1:])
+        self.store_futures[request_ids[0]] = (future, list(request_ids[1:]))
 
     @_lmcache_nvtx_annotate
     def batched_submit_retrieve_requests(
@@ -276,34 +482,83 @@ class LMCacheMPWorkerAdapter:
         ops: list[LoadStoreOp],
         event: torch.cuda.Event,
     ):
-        keys = []
-        block_ids = []
+        """
+        Submit a batched retrieve request to LMCache
 
-        for op in ops:
-            keys.extend(self._block_hashes_to_keys(op.block_hashes))
+        Args:
+            request_ids: The IDs of the requests
+            ops: The LoadStoreOps describing the retrieve operations. Should have
+                the same length as request_ids
+            event: The CUDA event that is recorded after the current
+                model inference step
+        """
+        all_keys: list[IPCCacheEngineKey] = []
+        block_ids: list[int] = []
+        for request_id, op in zip(request_ids, ops, strict=False):
+            if op.block_hashes is not None:
+                chunk_hashes = list(
+                    striding_block_hashes(op.block_hashes, self.blocks_in_chunk)
+                )
+                keys = [
+                    self._create_hash_key(ch, request_id=request_id)
+                    for ch in chunk_hashes
+                ]
+                all_keys.extend(keys)
+            else:
+                assert op.token_ids is not None
+                all_keys.append(
+                    self._create_key(
+                        op.token_ids, op.start, op.end, request_id=request_id
+                    )
+                )
             block_ids.extend(op.block_ids)
         future = send_lmcache_request(
             self.mq_client,
             RequestType.RETRIEVE,
-            [keys, self.instance_id, block_ids, event.ipc_handle()],
+            [
+                all_keys,
+                self.instance_id,
+                block_ids,
+                event.ipc_handle(),
+            ],
         ).to_cuda_future()
-        self.retrieve_futures[request_ids[0]] = (future, request_ids[1:])
+        self.retrieve_futures[request_ids[0]] = (future, list(request_ids[1:]))
 
     @_lmcache_nvtx_annotate
     def get_finished(
-        self, finished_req_ids: set[str]
+        self, finished_req_ids_from_engine: set[str]
     ) -> tuple[set[str] | None, set[str] | None]:
+        """
+        Check and get the finished store and retrieve requests.
+
+        Args:
+            finished_req_ids_from_engine: the set of request ids that are
+                reported as finished from the vLLM engine side.
+
+        Returns:
+            A tuple of two sets:
+            - The first set contains the finished store request ids. The returned
+                store request ids MUST be seen before in the
+                `finished_req_ids_from_engine`.
+            - The second set contains the finished retrieve request ids.
+
+        Notes:
+            When enabling async scheduling in vLLM, the same request ID may appear
+            multiple times in `finished_req_ids_from_engine`. The adapter should
+            take care of deduplicating the request IDs and only return the request
+            IDs that have not been returned before.
+        """
         finished_stores = set()
         finished_retrieves = set()
-        for request_id, (future, other_reqs) in self.store_futures.items():
-            if not future.query():
+        for request_id, (s_future, other_reqs) in self.store_futures.items():
+            if not s_future.query():
                 continue
 
-            result = future.result()
+            s_result = s_future.result()
             finished_stores.add(request_id)
             finished_stores.update(other_reqs)
 
-            if not result:
+            if not s_result:
                 # TODO: add error handling here
                 logger.error(
                     "Something went wrong when processing the "
@@ -311,21 +566,21 @@ class LMCacheMPWorkerAdapter:
                     request_id,
                 )
 
-        for request_id, (future, other_reqs) in self.retrieve_futures.items():
-            if not future.query():
+        for request_id, (r_future, other_reqs) in self.retrieve_futures.items():
+            if not r_future.query():
                 continue
 
-            result = future.result()
+            r_result = r_future.result()
             finished_retrieves.add(request_id)
             finished_retrieves.update(other_reqs)
 
-            if not all(result):
+            if not all(r_result):
                 # TODO: add error handing here
                 logger.error(
                     "Something went wrong when processing the "
                     "retrieve request for request_id=%s, result=%s",
                     request_id,
-                    result,
+                    r_result,
                 )
 
         # Remove the finished requests from the tracking dicts
@@ -338,7 +593,7 @@ class LMCacheMPWorkerAdapter:
         self.finished_stores.update(finished_stores)
 
         ret_stores = set()
-        for req_id in finished_req_ids:
+        for req_id in finished_req_ids_from_engine:
             if req_id in self.finished_stores or req_id in self.store_futures:
                 self.previously_finished.add(req_id)
             else:
@@ -357,7 +612,9 @@ class LMCacheMPWorkerAdapter:
         return self.blocks_in_chunk
 
     def shutdown(self):
-        # Unregister kv cache
+        """
+        Shutdown the LMCache MP worker adapter
+        """
         logger.info("Unregistering kv caches")
         send_lmcache_request(
             self.mq_client, RequestType.UNREGISTER_KV_CACHE, [self.instance_id]
@@ -378,18 +635,32 @@ class LMCacheMPWorkerAdapter:
 
         return safe_finished_s
 
-    def _create_key(self, block_hash: bytes) -> IPCCacheEngineKey:
-        """Convert a block hash to an IPC cache engine key"""
+    def _create_key(
+        self,
+        token_ids: list[int],
+        start: int = 0,
+        end: int = 0,
+        request_id: str | None = None,
+    ) -> IPCCacheEngineKey:
+        """Convert token IDs to an IPC cache engine key"""
         return IPCCacheEngineKey(
             model_name=self.model_name,
             world_size=self.world_size,
             worker_id=self.worker_id,
-            chunk_hash=block_hash,
+            token_ids=tuple(token_ids),
+            start=start,
+            end=end,
+            request_id=request_id,
         )
 
-    def _block_hashes_to_keys(
-        self, block_hashes: list[bytes]
-    ) -> list[IPCCacheEngineKey]:
-        """Convert block hashes to IPC cache engine keys"""
-        s = striding_block_hashes(block_hashes, self.blocks_in_chunk)
-        return [self._create_key(block_hash) for block_hash in s]
+    def _create_hash_key(
+        self, chunk_hash: bytes, request_id: str | None = None
+    ) -> IPCCacheEngineKey:
+        """Create a hash-mode IPC cache engine key"""
+        return IPCCacheEngineKey(
+            model_name=self.model_name,
+            world_size=self.world_size,
+            worker_id=self.worker_id,
+            chunk_hash=chunk_hash,
+            request_id=request_id,
+        )

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

@@ -3,7 +3,7 @@
 import enum
 from collections.abc import Iterable
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Literal
 
 import torch
 import zmq
@@ -130,12 +130,6 @@ def create_worker_adapter(
     )
 
 
-def convert_block_hashes_to_bytes(
-    block_hashes: list["BlockHash"],
-) -> list[bytes]:
-    return cast(list[bytes], block_hashes)
-
-
 class LMCacheMPRequestState(enum.Enum):
     """
     State machine:
@@ -266,6 +260,7 @@ class LMCacheMPRequestMetadata:
         Args:
             tracker: The request tracker to generate the metadata from.
             blocks_in_chunk: the number of blocks in a LMCache data chunk
+            vllm_block_size: the block size used in vLLM
         """
         # Store the blocks that has block hashes
         # NOTE: the invariant here is that `num_stored_blocks` should
@@ -282,15 +277,21 @@ class LMCacheMPRequestMetadata:
         if num_chunks >= 1:
             start = tracker.num_stored_blocks
             end = start + num_chunks * blocks_in_chunk
-            block_hashes = convert_block_hashes_to_bytes(
-                tracker.block_hashes[start:end]
-            )
             block_ids = tracker.allocated_block_ids[start:end]
+            start_token_idx = start * vllm_block_size
+            end_token_idx = end * vllm_block_size
+            token_ids = list(tracker.all_token_ids)
+            op = LoadStoreOp(
+                token_ids=token_ids,
+                block_ids=block_ids,
+                start=start_token_idx,
+                end=end_token_idx,
+            )
 
             ret = LMCacheMPRequestMetadata(
                 request_id=tracker.request_id,
                 direction="STORE",
-                op=LoadStoreOp(block_hashes=block_hashes, block_ids=block_ids),
+                op=op,
             )
 
             # Update the request tracker
@@ -303,6 +304,7 @@ class LMCacheMPRequestMetadata:
     def GetRetrieveMetadata(
         tracker: LMCacheMPRequestTracker,
         blocks_in_chunk: int,
+        vllm_block_size: int,
     ) -> "LMCacheMPRequestMetadata | None":
         """
         Generate the retrieve metadata for the current request tracker.
@@ -310,6 +312,7 @@ class LMCacheMPRequestMetadata:
         Args:
             tracker: The request tracker to generate the metadata from.
             blocks_in_chunk: the number of blocks in a LMCache data chunk
+            vllm_block_size: the block size used in vLLM
         """
         if not tracker.is_ready_for_retrieving():
             return None
@@ -330,15 +333,21 @@ class LMCacheMPRequestMetadata:
             "number of LMCache hit blocks. "
         )
         if end > start:
-            block_hashes = convert_block_hashes_to_bytes(
-                tracker.block_hashes[start:end]
-            )
             block_ids = tracker.allocated_block_ids[start:end]
+            start_token_idx = start * vllm_block_size
+            end_token_idx = end * vllm_block_size
+            token_ids = list(tracker.all_token_ids)
+            op = LoadStoreOp(
+                token_ids=token_ids,
+                block_ids=block_ids,
+                start=start_token_idx,
+                end=end_token_idx,
+            )
 
             ret = LMCacheMPRequestMetadata(
                 request_id=tracker.request_id,
                 direction="RETRIEVE",
-                op=LoadStoreOp(block_hashes=block_hashes, block_ids=block_ids),
+                op=op,
             )
             return ret
 
@@ -643,7 +652,8 @@ class LMCacheMPConnector(KVConnectorBase_V1):
             return 0, False
 
         self.scheduler_adapter.maybe_submit_lookup_request(
-            request.request_id, convert_block_hashes_to_bytes(request.block_hashes)
+            request.request_id,
+            token_ids=list(request.all_token_ids),
         )
 
         ret = self.scheduler_adapter.check_lookup_result(request.request_id)
@@ -766,6 +776,9 @@ class LMCacheMPConnector(KVConnectorBase_V1):
         """
         # Clean up request tracker to prevent memory leak
         self._cleanup_request_tracker(request.request_id)
+        # Notify LMCache to end the session for this request
+        self.scheduler_adapter.end_session(request.request_id)
+
         return True, None
 
     def take_events(self) -> Iterable["KVCacheEvent"]:
@@ -846,7 +859,9 @@ class LMCacheMPConnector(KVConnectorBase_V1):
             if request_tracker.state != LMCacheMPRequestState.WAITING_FOR_LOAD:
                 continue
             r_metadata = LMCacheMPRequestMetadata.GetRetrieveMetadata(
-                request_tracker, blocks_per_chunk
+                request_tracker,
+                blocks_per_chunk,
+                vllm_block_size=self.vllm_block_size,
             )
             if r_metadata is not None:
                 metadata.add_request_metadata(r_metadata)

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

@@ -926,6 +926,17 @@ class NixlConnectorWorker:
         else:
             self.use_host_buffer = self.kv_buffer_device == "cpu"
 
+        # reserve different cores for start_load_kv() from model_forward()
+        if self.device_type == "cpu":
+            numa_core_list = current_platform.discover_numa_topology()
+            # setup one last core in each numa for kv transfer.
+            rsv_cores_for_kv = [
+                max(each_numa_core_list) for each_numa_core_list in numa_core_list
+            ]
+
+            if rsv_cores_for_kv:
+                os.sched_setaffinity(0, rsv_cores_for_kv)
+
         # support for oot platform which can't register nixl memory
         # type based on kv_buffer_device
         nixl_memory_type = current_platform.get_nixl_memory_type()

vllm/distributed/parallel_state.py

@@ -1143,6 +1143,18 @@ def get_ep_group() -> GroupCoordinator:
     return _EP
 
 
+_EPLB: GroupCoordinator | None = None
+
+
+def get_eplb_group() -> GroupCoordinator:
+    assert _EPLB is not None, (
+        "EPLB group is not initialized. "
+        "EPLB group is only created for MoE models when EPLB is enabled. "
+        "Ensure parallel_config.enable_eplb is True."
+    )
+    return _EPLB
+
+
 _PCP: GroupCoordinator | None = None
 
 
@@ -1440,12 +1452,29 @@ def initialize_model_parallel(
         _EP = init_model_parallel_group(
             group_ranks, get_world_group().local_rank, backend, group_name="ep"
         )
+
+        # Create EPLB group with the same ranks as EP if EPLB is enabled.
+        # This is a separate process group to isolate EPLB communications
+        # from MoE forward pass collectives and prevent deadlocks when
+        # using torch.distributed in execution with torch.distributed in EPLB.
+        global _EPLB
+        assert _EPLB is None, "EPLB group is already initialized"
+        if (
+            config is not None
+            and config.parallel_config is not None
+            and config.parallel_config.enable_eplb
+        ):
+            # Reuse the same group_ranks from EP
+            _EPLB = init_model_parallel_group(
+                group_ranks, get_world_group().local_rank, backend, group_name="eplb"
+            )
     # If no EP group needed, _EP remains None
+    # If no EPLB group needed, _EPLB remains None
 
     logger.info_once(
         "rank %s in world size %s is assigned as "
         "DP rank %s, PP rank %s, PCP rank %s, "
-        "TP rank %s, EP rank %s",
+        "TP rank %s, EP rank %s, EPLB rank %s",
         rank,
         world_size,
         _DP.rank_in_group,
@@ -1453,6 +1482,7 @@ def initialize_model_parallel(
         _PCP.rank_in_group,
         _TP.rank_in_group,
         _EP.rank_in_group if _EP is not None else "N/A",
+        _EPLB.rank_in_group if _EPLB is not None else "N/A",
     )
 
 
@@ -1514,6 +1544,8 @@ def prepare_communication_buffer_for_model(model: torch.nn.Module):
         _DP.prepare_communication_buffer_for_model(model)
     if _EP is not None:
         _EP.prepare_communication_buffer_for_model(model)
+    if _EPLB is not None:
+        _EPLB.prepare_communication_buffer_for_model(model)
 
 
 def model_parallel_is_initialized():
@@ -1608,6 +1640,11 @@ def destroy_model_parallel():
         _EP.destroy()
     _EP = None
 
+    global _EPLB
+    if _EPLB:
+        _EPLB.destroy()
+    _EPLB = None
+
 
 def destroy_distributed_environment():
     global _WORLD, _NODE_COUNT