perf(deepseek): Multi-stream overlap of indexer wk+weights_proj with QKV-A for DeepSeek V3.2 NSA

Signed-off-by: Alexander Matveev <amatveev@redhat.com>

.buildkite/hardware_tests/cpu.yaml

@@ -46,7 +46,7 @@ steps:
   - tests/models/language/pooling/
   commands:
     - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 40m "
       pytest -x -v -s tests/models/language/generation -m cpu_model
       pytest -x -v -s tests/models/language/pooling -m cpu_model"
 
@@ -99,7 +99,7 @@ steps:
     - |
       bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
       pytest -x -v -s tests/models/multimodal/generation --ignore=tests/models/multimodal/generation/test_pixtral.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
-  parallelism: 2
+  parallelism: 3
 
 - label: "Arm CPU Test"
   depends_on: []

.buildkite/image_build/image_build_torch_nightly.sh

@@ -0,0 +1,68 @@
+#!/bin/bash
+set -euo pipefail
+
+# Build a vLLM test image with PyTorch nightly installed.
+# Called by the pipeline generator's "vLLM Against PyTorch Nightly" group.
+
+if [[ $# -lt 5 ]]; then
+  echo "Usage: $0 <registry> <repo> <commit> <branch> <image_tag>"
+  exit 1
+fi
+
+REGISTRY=$1
+REPO=$2
+BUILDKITE_COMMIT=$3
+BRANCH=$4
+IMAGE_TAG=$5
+
+# --- Arguments ---
+echo "--- :mag: Arguments"
+echo "REGISTRY: ${REGISTRY}"
+echo "REPO: ${REPO}"
+echo "BUILDKITE_COMMIT: ${BUILDKITE_COMMIT}"
+echo "BRANCH: ${BRANCH}"
+echo "IMAGE_TAG: ${IMAGE_TAG}"
+
+# --- ECR login ---
+echo "--- :key: ECR login"
+aws ecr-public get-login-password --region us-east-1 \
+  | docker login --username AWS --password-stdin "$REGISTRY"
+aws ecr get-login-password --region us-east-1 \
+  | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
+
+# --- Set up buildx ---
+echo "--- :docker: Setting up buildx"
+docker buildx create --name vllm-builder --driver docker-container --use || true
+docker buildx inspect --bootstrap
+docker buildx ls
+
+# --- Skip if image already exists ---
+echo "--- :mag: Checking if image already exists"
+if docker manifest inspect "$IMAGE_TAG" >/dev/null 2>&1; then
+  echo "Image found: $IMAGE_TAG — skipping build"
+  exit 0
+fi
+echo "Image not found, proceeding with build..."
+
+# --- CUDA 13.0 for nightly builds ---
+# Nightly CI uses CUDA 13.0 while regular CI stays on CUDA 12.9
+NIGHTLY_CUDA_VERSION="13.0.0"
+NIGHTLY_BUILD_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-devel-ubuntu22.04"
+NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu22.04"
+
+echo "--- :docker: Building torch nightly image (CUDA ${NIGHTLY_CUDA_VERSION})"
+docker buildx build --file docker/Dockerfile \
+  --build-arg max_jobs=16 \
+  --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
+  --build-arg USE_SCCACHE=1 \
+  --build-arg PYTORCH_NIGHTLY=1 \
+  --build-arg CUDA_VERSION="${NIGHTLY_CUDA_VERSION}" \
+  --build-arg BUILD_BASE_IMAGE="${NIGHTLY_BUILD_BASE_IMAGE}" \
+  --build-arg FINAL_BASE_IMAGE="${NIGHTLY_FINAL_BASE_IMAGE}" \
+  --build-arg torch_cuda_arch_list="8.0 8.9 9.0 10.0 12.0" \
+  --tag "$IMAGE_TAG" \
+  --push \
+  --target test \
+  --progress plain .
+
+echo "--- :white_check_mark: Torch nightly image build complete: $IMAGE_TAG"

.buildkite/release-pipeline.yaml

@@ -98,8 +98,15 @@ steps:
     commands:
       - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
 
+  - block: "Unblock to build release Docker images"
+    depends_on: ~
+    key: block-build-release-images
+    if: build.env("NIGHTLY") != "1"
+
   - group: "Build release Docker images"
     key: "build-release-images"
+    depends_on: block-build-release-images
+    allow_dependency_failure: true
     steps:
       - label: "Build release image - x86_64 - CUDA 12.9"
         depends_on: ~
@@ -617,6 +624,8 @@ steps:
   - label: ":docker: Build release image - x86_64 - ROCm"
     id: build-rocm-release-image
     depends_on:
+      - step: block-build-release-images
+        allow_failure: true
       - step: build-rocm-base-wheels
         allow_failure: false
     agents:

.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh

@@ -51,6 +51,7 @@ function cpu_tests() {
     set -e
     pytest -x -v -s tests/kernels/test_onednn.py
     pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
+    pytest -x -v -s tests/kernels/core/test_cpu_activation.py
     pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic"
 
   # basic online serving

.buildkite/scripts/hardware_ci/run-cpu-test.sh

@@ -16,5 +16,5 @@ echo "--- :docker: Building Docker image"
 docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test -f docker/Dockerfile.cpu .
 
 # Run the image, setting --shm-size=4g for tensor parallel.
-docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g "$IMAGE_NAME" \
+docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 --shm-size=4g "$IMAGE_NAME" \
         timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"

.buildkite/test-amd.yaml

@@ -769,7 +769,7 @@ steps:
   - tests/kernels/helion/
   - vllm/platforms/rocm.py
   commands:
-  - pip install helion==0.3.3
+  - pip install helion==1.0.0
   - pytest -v -s kernels/helion/
 
 

.buildkite/test_areas/distributed.yaml

@@ -196,7 +196,6 @@ steps:
     - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
     - VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
     - pytest -v -s tests/v1/distributed/test_dbo.py
-    - TP_SIZE=1 DP_SIZE=2 pytest -v -s tests/v1/distributed/test_eagle_dp.py
 
 - label: Distributed Tests (2 GPUs)(B200)
   device: b200

.buildkite/test_areas/kernels.yaml

@@ -155,7 +155,7 @@ steps:
   - vllm/utils/import_utils.py
   - tests/kernels/helion/
   commands:
-    - pip install helion==0.3.3
+    - pip install helion==1.0.0
     - pytest -v -s kernels/helion/
 
  
@@ -200,7 +200,14 @@ steps:
   timeout_in_minutes: 90
   device: h100
   num_devices: 2
-  optional: true
+  source_file_dependencies:
+  - csrc/quantization/cutlass_w8a8/moe/
+  - csrc/moe/
+  - tests/kernels/moe
+  - vllm/model_executor/layers/fused_moe/
+  - vllm/model_executor/layers/quantization/
+  - vllm/distributed/device_communicators/
+  - vllm/config
   commands:
     - pytest -v -s kernels/moe/test_moe_layer.py
 
@@ -209,6 +216,13 @@ steps:
   timeout_in_minutes: 90
   device: b200
   num_devices: 2
-  optional: true
+  source_file_dependencies:
+  - csrc/quantization/cutlass_w8a8/moe/
+  - csrc/moe/
+  - tests/kernels/moe
+  - vllm/model_executor/layers/fused_moe/
+  - vllm/model_executor/layers/quantization/
+  - vllm/distributed/device_communicators/
+  - vllm/config
   commands:
     - pytest -v -s kernels/moe/test_moe_layer.py

.buildkite/test_areas/lm_eval.yaml

@@ -91,6 +91,16 @@ steps:
     - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
 
 
+- label: LM Eval TurboQuant KV Cache
+  timeout_in_minutes: 75
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/turboquant/
+  - vllm/v1/attention/backends/turboquant_attn.py
+  - vllm/v1/attention/ops/triton_turboquant_decode.py
+  - vllm/v1/attention/ops/triton_turboquant_store.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant.txt
+
 - label: GPQA Eval (GPT-OSS) (H100)
   timeout_in_minutes: 120
   device: h100

.buildkite/test_areas/misc.yaml

@@ -224,6 +224,7 @@ steps:
     - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
     - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
     - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
+    - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
   
 - label: Acceptance Length Test (Large Models) # optional
   timeout_in_minutes: 25

.buildkite/test_areas/models_basic.yaml

@@ -1,10 +1,9 @@
 group: Models - Basic
-depends_on: 
+depends_on:
   - image-build
 steps:
 - label: Basic Models Tests (Initialization)
   timeout_in_minutes: 45
-  device: h200_18gb
   torch_nightly: true
   source_file_dependencies:
   - vllm/
@@ -13,10 +12,11 @@ steps:
   commands:
     # Run a subset of model initialization tests
     - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset
+  mirror:
+    torch_nightly: {}
 
 - label: Basic Models Tests (Extra Initialization) %N
   timeout_in_minutes: 45
-  torch_nightly: true
   source_file_dependencies:
   - vllm/model_executor/models/
   - tests/models/test_initialization.py
@@ -27,6 +27,8 @@ steps:
     # test.) Also run if model initialization test file is modified
     - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   parallelism: 2
+  mirror:
+    torch_nightly: {}
 
 - label: Basic Models Tests (Other)
   timeout_in_minutes: 45
@@ -42,10 +44,10 @@ steps:
       device: mi325_1
       depends_on:
       - image-build-amd
-    
+
 
 - label: Basic Models Test (Other CPU) # 5min
-  depends_on: 
+  depends_on:
   - image-build-cpu
   timeout_in_minutes: 10
   source_file_dependencies:
@@ -70,3 +72,18 @@ steps:
     - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
     # Whisper needs spawn method to avoid deadlock
     - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
+
+- label: Transformers Backward Compatibility Models Test
+  working_dir: "/vllm-workspace/"
+  optional: true
+  soft_fail: true
+  commands:
+    - pip install transformers==4.57.5
+    - pytest -v -s tests/models/test_initialization.py
+    - pytest -v -s tests/models/test_transformers.py
+    - pytest -v -s tests/models/multimodal/processing/
+    - pytest -v -s tests/models/multimodal/test_mapping.py
+    - python3 examples/offline_inference/basic/chat.py
+    - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
+    # Whisper needs spawn method to avoid deadlock
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper

.buildkite/test_areas/models_language.yaml

@@ -1,10 +1,9 @@
 group: Models - Language
-depends_on: 
+depends_on:
   - image-build
 steps:
 - label: Language Models Tests (Standard)
   timeout_in_minutes: 25
-  torch_nightly: true
   source_file_dependencies:
   - vllm/
   - tests/models/language
@@ -12,10 +11,11 @@ steps:
     # Test standard language models, excluding a subset of slow tests
     - pip freeze | grep -E 'torch'
     - pytest -v -s models/language -m 'core_model and (not slow_test)'
+  mirror:
+    torch_nightly: {}
 
 - label: Language Models Tests (Extra Standard) %N
   timeout_in_minutes: 45
-  torch_nightly: true
   source_file_dependencies:
   - vllm/model_executor/models/
   - tests/models/language/pooling/test_embedding.py
@@ -27,10 +27,11 @@ steps:
     - pip freeze | grep -E 'torch'
     - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   parallelism: 2
+  mirror:
+    torch_nightly: {}
 
 - label: Language Models Tests (Hybrid) %N
   timeout_in_minutes: 75
-  torch_nightly: true
   source_file_dependencies:
   - vllm/
   - tests/models/language/generation
@@ -42,6 +43,8 @@ steps:
     # 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
   parallelism: 2
+  mirror:
+    torch_nightly: {}
 
 - label: Language Models Test (Extended Generation) # 80min
   timeout_in_minutes: 110
@@ -62,7 +65,7 @@ steps:
       - image-build-amd
       commands:
       - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
-      - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
+      - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
       - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
 
 - label: Language Models Test (PPL)

.buildkite/test_areas/spec_decode.yaml

@@ -42,3 +42,16 @@ steps:
     - tests/v1/e2e/spec_decode/
   commands:
     - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
+
+- label: DFlash Speculators Correctness
+  timeout_in_minutes: 30
+  device: h100
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+    - vllm/v1/spec_decode/
+    - vllm/model_executor/models/qwen3_dflash.py
+    - tests/v1/spec_decode/test_speculators_dflash.py
+  commands:
+    - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
+    - pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test

.github/PULL_REQUEST_TEMPLATE.md

@@ -15,7 +15,6 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
 - [ ] The test plan, such as providing test command.
 - [ ] The test results, such as pasting the results comparison before and after, or e2e results
 - [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model.
-- [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft in the [Google Doc](https://docs.google.com/document/d/1YyVqrgX4gHTtrstbq8oWUImOyPCKSGnJ7xtTpmXzlRs/edit?tab=t.0).
 </details>
 
 **BEFORE SUBMITTING, PLEASE READ <https://docs.vllm.ai/en/latest/contributing>** (anything written below this line will be removed by GitHub Actions)

.github/mergify.yml

@@ -264,6 +264,7 @@ pull_request_rules:
       - files=\.buildkite/ci_config_intel.yaml
       - files=vllm/model_executor/layers/fused_moe/xpu_fused_moe.py
       - files=vllm/model_executor/kernels/linear/mixed_precision/xpu.py
+      - files=vllm/model_executor/kernels/linear/mxfp8/xpu.py
       - files=vllm/model_executor/kernels/linear/scaled_mm/xpu.py
       - files=vllm/distributed/device_communicators/xpu_communicator.py
       - files=vllm/v1/attention/backends/mla/xpu_mla_sparse.py
@@ -271,6 +272,7 @@ pull_request_rules:
       - files=vllm/v1/worker/xpu_worker.py
       - files=vllm/v1/worker/xpu_model_runner.py
       - files=vllm/_xpu_ops.py
+      - files=vllm/kernels/xpu_ops.py
       - files~=^vllm/lora/ops/xpu_ops
       - files=vllm/lora/punica_wrapper/punica_xpu.py
       - files=vllm/platforms/xpu.py
@@ -278,7 +280,6 @@ pull_request_rules:
       - title~=(?i)XPU
       - title~=(?i)Intel
       - title~=(?i)BMG
-      - title~=(?i)Arc
   actions:
     label:
       add:

cmake/cpu_extension.cmake

@@ -360,6 +360,7 @@ set(VLLM_EXT_SRC
 if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
     set(VLLM_EXT_SRC
         "csrc/cpu/shm.cpp"
+        "csrc/cpu/activation_lut_bf16.cpp"
         ${VLLM_EXT_SRC})
 endif()
 

csrc/cpu/activation_lut_bf16.cpp

@@ -0,0 +1,71 @@
+#include "cpu_types.hpp"
+
+#include <array>
+#include <cstdint>
+#include <mutex>
+#include <string>
+
+#include <ATen/ops/empty.h>
+#include <ATen/ops/gelu.h>
+#include <c10/util/BFloat16.h>
+
+constexpr uint32_t ActivationLutSize = 1u << 16;
+
+at::Tensor gelu_reference(const at::Tensor& x) { return at::gelu(x, "none"); }
+
+void maybe_init_activation_lut_bf16(
+    uint16_t* lut, std::once_flag& once,
+    at::Tensor (*activation)(const at::Tensor&)) {
+  std::call_once(once, [&]() {
+    auto lut_input =
+        at::empty({static_cast<int64_t>(ActivationLutSize)},
+                  at::TensorOptions().device(at::kCPU).dtype(at::kFloat));
+    auto* lut_input_ptr = lut_input.data_ptr<float>();
+#pragma omp parallel for
+    for (uint32_t i = 0; i < ActivationLutSize; ++i) {
+      lut_input_ptr[i] = c10::detail::f32_from_bits(static_cast<uint16_t>(i));
+    }
+
+    auto lut_output = activation(lut_input);
+    const auto* lut_output_ptr = lut_output.data_ptr<float>();
+#pragma omp parallel for
+    for (uint32_t i = 0; i < ActivationLutSize; ++i) {
+      lut[i] = c10::detail::round_to_nearest_even(lut_output_ptr[i]);
+    }
+  });
+}
+
+void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
+                         const uint16_t* lut, const char* op_name) {
+  TORCH_CHECK(input.scalar_type() == at::kBFloat16, op_name,
+              ": input must be bfloat16");
+  TORCH_CHECK(out.scalar_type() == at::kBFloat16, op_name,
+              ": out must be bfloat16");
+  TORCH_CHECK(input.is_contiguous(), op_name, ": input must be contiguous");
+  TORCH_CHECK(out.is_contiguous(), op_name, ": out must be contiguous");
+
+  const auto* src =
+      reinterpret_cast<const uint16_t*>(input.data_ptr<at::BFloat16>());
+  auto* dst = reinterpret_cast<uint16_t*>(out.data_ptr<at::BFloat16>());
+  const int64_t n = input.numel();
+
+  CPU_KERNEL_GUARD_IN(activation_lut_bf16_impl)
+#pragma omp parallel for
+  for (int64_t i = 0; i < n; ++i) {
+    dst[i] = lut[src[i]];
+  }
+  CPU_KERNEL_GUARD_OUT(activation_lut_bf16_impl)
+}
+
+void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
+                         const std::string& activation) {
+  if (activation == "gelu") {
+    static std::array<uint16_t, ActivationLutSize> lut{};
+    static std::once_flag once;
+    maybe_init_activation_lut_bf16(lut.data(), once, gelu_reference);
+    activation_lut_bf16(out, input, lut.data(), "gelu_lut");
+    return;
+  }
+
+  TORCH_CHECK(false, "Unsupported activation: ", activation);
+}

csrc/cpu/cpu_attn_impl.hpp

@@ -147,6 +147,9 @@ struct AttentionMetadata {
       case ISA::NEON:
         ss << "NEON, ";
         break;
+      case ISA::VXE:
+        ss << "VXE, ";
+        break;
     }
     ss << "workitem_group_num: " << workitem_group_num
        << ", reduction_item_num: " << reduction_item_num

csrc/cpu/torch_bindings.cpp

@@ -85,6 +85,9 @@ at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros,
                               at::Tensor& w_scales,
                               std::optional<at::Tensor> bias);
 
+void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
+                         const std::string& activation);
+
 torch::Tensor get_scheduler_metadata(
     const int64_t num_req, const int64_t num_heads_q,
     const int64_t num_heads_kv, const int64_t head_dim,
@@ -231,6 +234,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
   ops.impl("gelu_quick", torch::kCPU, &gelu_quick);
 
+#if (defined(__aarch64__) && !defined(__APPLE__))
+
+  ops.def(
+      "activation_lut_bf16(Tensor! out, Tensor input, str activation)"
+      " -> ()");
+  ops.impl("activation_lut_bf16", torch::kCPU, &activation_lut_bf16);
+
+#endif  // (defined(__aarch64__) && !defined(__APPLE__))
+
   // Layernorm
   // Apply Root Mean Square (RMS) Normalization to the input tensor.
   ops.def(

csrc/cpu/utils.hpp

@@ -54,12 +54,34 @@ struct Counter {
 };
 
 inline int64_t get_available_l2_size() {
+#if defined(__s390x__)
+  static int64_t size = []() {
+    uint32_t l2_cache_size = 0;
+    auto caps = at::cpu::get_cpu_capabilities();
+    auto it = caps.find("l2_cache_size");
+    if (it != caps.end()) {
+      l2_cache_size = static_cast<uint32_t>(it->second.toInt());
+    }
+    if (l2_cache_size == 0) {
+      long sys_l2 = sysconf(_SC_LEVEL2_CACHE_SIZE);
+      if (sys_l2 > 0) {
+        l2_cache_size = static_cast<uint32_t>(sys_l2);
+      }
+    }
+    if (l2_cache_size == 0) {
+      l2_cache_size = 256 * 1024;
+    }
+    return static_cast<int64_t>(l2_cache_size) >> 1;  // use 50% of L2 cache
+  }();
+  return size;
+#else
   static int64_t size = []() {
     auto caps = at::cpu::get_cpu_capabilities();
     const uint32_t l2_cache_size = caps.at("l2_cache_size").toInt();
     return l2_cache_size >> 1;  // use 50% of L2 cache
   }();
   return size;
+#endif
 }
 
 template <int32_t alignment_v, typename T>

csrc/ops.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <optional>
+#include <string>
 #include <torch/library.h>
 #include <tuple>
 

csrc/standalone/custom_ar/bench_moe_allgather.py

@@ -1,200 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Benchmark: Lamport all-gather vs NCCL."""
-
-import ctypes
-import os
-import sys
-
-import torch
-import torch.distributed as dist
-
-_cudart = ctypes.CDLL("libcudart.so")
-IPC = 64
-
-
-def _cc(r):
-    if r:
-        raise RuntimeError(f"err {r}")
-
-
-def ipc_buf(sz, rank, ws):
-    p = ctypes.c_void_p()
-    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
-    _cc(_cudart.cudaMemset(p, 0, sz))
-    _cc(_cudart.cudaDeviceSynchronize())
-    h = (ctypes.c_byte * IPC)()
-    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
-    ah = [None] * ws
-    dist.all_gather_object(ah, bytes(h))
-    ptrs = []
-    for i in range(ws):
-        if i == rank:
-            ptrs.append(p.value)
-        else:
-            hh = (ctypes.c_byte * IPC)(*ah[i])
-            pp = ctypes.c_void_p()
-            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
-            ptrs.append(pp.value)
-    return ptrs
-
-
-def gpu_timer(fn, warmup=20, repeats=200):
-    for _ in range(warmup):
-        fn()
-    torch.cuda.synchronize()
-    s = torch.cuda.Event(enable_timing=True)
-    e = torch.cuda.Event(enable_timing=True)
-    s.record()
-    for _ in range(repeats):
-        fn()
-    e.record()
-    torch.cuda.synchronize()
-    return s.elapsed_time(e) / repeats * 1000
-
-
-def gpu_timer_graph(fn, warmup=20, repeats=200):
-    """Time with CUDA graph to exclude CPU overhead."""
-    for _ in range(warmup):
-        fn()
-    torch.cuda.synchronize()
-    g = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(g):
-        fn()
-    for _ in range(5):
-        g.replay()
-    torch.cuda.synchronize()
-    s = torch.cuda.Event(enable_timing=True)
-    e = torch.cuda.Event(enable_timing=True)
-    s.record()
-    for _ in range(repeats):
-        g.replay()
-    e.record()
-    torch.cuda.synchronize()
-    return s.elapsed_time(e) / repeats * 1000
-
-
-def main():
-    dist.init_process_group("nccl")
-    rank = dist.get_rank()
-    ws = dist.get_world_size()
-    torch.cuda.set_device(rank)
-    dev = f"cuda:{rank}"
-
-    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-    if rank == 0:
-        from moe_allgather import _load_lib
-
-        lib = _load_lib()
-    dist.barrier()
-    if rank != 0:
-        from moe_allgather import _load_lib
-
-        lib = _load_lib()
-    dist.barrier()
-    from moe_allgather import MoeAllGather
-
-    max_size = 8 * 1024 * 1024
-    bp = ipc_buf(max_size, rank, ws)
-    dist.barrier()
-
-    class FakeCA:
-        pass
-
-    ca = FakeCA()
-    ca.rank = rank
-    ca.world_size = ws
-    ca.device = torch.device(dev)
-    ca.buffer_ptrs = bp
-    ca.max_size = max_size
-    ag = MoeAllGather(ca)
-    dist.barrier()
-
-    configs = [
-        ("1tok", 1),
-        ("2tok", 2),
-        ("4tok", 4),
-        ("8tok", 8),
-        ("16tok", 16),
-        ("32tok", 32),
-        ("64tok", 64),
-        ("128tok", 128),
-        ("256tok", 256),
-    ]
-    topk = 8
-    hd = 3584
-    sd = 448
-
-    if rank == 0:
-        print(f"world_size={ws}, max_per_rank={ag.max_per_rank} bytes")
-        print(f"{'config':<12} {'lamport_graph':>10} {'nccl_graph':>10} {'speedup':>8}")
-        print("-" * 65)
-
-    for name, N in configs:
-        # Check if data fits in buffer.
-        cursor = 0
-        per_tok = topk * 4 + topk * 4 + hd + sd
-        cursor = N * per_tok
-        cursor = (cursor + 15) & ~15
-        if cursor > ag.max_per_rank:
-            if rank == 0:
-                print(f"{name:<12} {'skip (too large)':>40}")
-            continue
-
-        ids = torch.randint(0, 256, (N, topk), dtype=torch.int32, device=dev)
-        wt = torch.randn(N, topk, dtype=torch.float32, device=dev).abs()
-        hs = torch.randint(0, 255, (N, hd), dtype=torch.uint8, device=dev)
-        sc = torch.randint(0, 255, (N, sd), dtype=torch.uint8, device=dev)
-        inputs = [ids, wt, hs, sc]
-
-        # Custom Lamport kernel.
-        c_outs = [
-            torch.empty(N * ws, *t.shape[1:], dtype=t.dtype, device=dev) for t in inputs
-        ]
-
-        def run_lamport():
-            lib.moe_all_gather(
-                ag._buf_ptrs_ptr,
-                ag._counters_ptr,
-                rank,
-                ws,
-                ag.seg_capacity,
-                ag.rank_stride,
-                inputs,
-                c_outs,
-            )
-
-        # Lamport with CUDA graph.
-        try:
-            lam_g_us = gpu_timer_graph(run_lamport)
-        except Exception as ex:
-            lam_g_us = float("nan")
-            if rank == 0:
-                print(f"  [graph capture failed: {ex}]")
-
-        # NCCL 1×AG (concat into one tensor).
-        cat_inp = torch.cat(
-            [t.reshape(N, -1).contiguous().view(torch.uint8) for t in inputs],
-            dim=1,
-        ).contiguous()
-        cat_out = torch.empty(N * ws, cat_inp.shape[1], dtype=torch.uint8, device=dev)
-
-        def run_nccl():
-            dist.all_gather_into_tensor(cat_out, cat_inp)
-
-        # NCCL with CUDA graph.
-        try:
-            nccl_g_us = gpu_timer_graph(run_nccl)
-        except Exception:
-            nccl_g_us = float("nan")
-
-        if rank == 0:
-            speedup = nccl_g_us / lam_g_us if lam_g_us > 0 else float("nan")
-            print(f"{name:<12} {lam_g_us:>9.1f}µ {nccl_g_us:>9.1f}µ {speedup:>7.2f}x")
-
-    dist.barrier()
-    dist.destroy_process_group()
-
-
-if __name__ == "__main__":
-    main()

csrc/standalone/custom_ar/moe_allgather.cu

@@ -1,301 +0,0 @@
-// Lamport-based MoE all-gather kernel for EP dispatch.
-//
-// Replaces the flag-barrier approach with a Lamport sentinel protocol
-// (inspired by FlashInfer's trtllm_allreduce_fusion).
-//
-// Key advantages over the flag-barrier approach:
-//   - No explicit barriers (sentinels provide per-element synchronization).
-//   - Push model: NVLink writes (fire-and-forget) instead of NVLink reads.
-//   - Triple buffering: no end barrier needed.
-//
-// Gathers the MoE dispatch tensors from all EP ranks:
-//   - topk_ids      [N, topk]     int32
-//   - topk_weights  [N, topk]     float32 / bfloat16
-//   - hidden_states [N, D_h]      uint8 (NVFP4) / bfloat16
-//   - quant_scales  [N, D_s]      (optional)
-//
-// Double-buffer layout in each rank's IPC buffer:
-//   [Segment 0][Segment 1]
-//   Each segment: [Rank 0 slot][Rank 1 slot]...[Rank N-1 slot]
-//   Each rank slot: packed tensors at 16-byte aligned offsets.
-//
-// Sentinel: 0x80000000 (negative-zero in float32).  The writer replaces
-// any data word matching the sentinel with 0 before pushing.  The reader
-// spin-loads (volatile) until no sentinel words remain in the vector.
-
-#include <cuda.h>
-#include <cuda_runtime.h>
-#include <torch/extension.h>
-
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAStream.h>
-
-#define DINLINE __device__ __forceinline__
-
-constexpr uint32_t SENTINEL = 0x80000000u;
-constexpr int kMaxBlocks = 36;
-
-// ---------------------------------------------------------------------------
-// Volatile 128-bit load/store and sentinel helpers
-// ---------------------------------------------------------------------------
-
-static DINLINE int4 ld128v(const void* addr) {
-  int4 v;
-  asm volatile("ld.volatile.global.v4.b32 {%0,%1,%2,%3}, [%4];"
-               : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
-               : "l"(addr));
-  return v;
-}
-
-static DINLINE bool has_sentinel(int4 v) {
-  return reinterpret_cast<uint32_t&>(v.x) == SENTINEL |
-         reinterpret_cast<uint32_t&>(v.y) == SENTINEL |
-         reinterpret_cast<uint32_t&>(v.z) == SENTINEL |
-         reinterpret_cast<uint32_t&>(v.w) == SENTINEL;
-}
-
-static DINLINE int4 remove_sentinel(int4 v) {
-  if (reinterpret_cast<uint32_t&>(v.x) == SENTINEL) v.x = 0;
-  if (reinterpret_cast<uint32_t&>(v.y) == SENTINEL) v.y = 0;
-  if (reinterpret_cast<uint32_t&>(v.z) == SENTINEL) v.z = 0;
-  if (reinterpret_cast<uint32_t&>(v.w) == SENTINEL) v.w = 0;
-  return v;
-}
-
-// ---------------------------------------------------------------------------
-// Lamport all-gather kernel
-// ---------------------------------------------------------------------------
-//
-// Phase 1 — PUSH: each rank writes its packed data to ALL peers' current
-//           segment via regular stores (NVLink push, fire-and-forget).
-// Phase 2 — CLEAR: each rank writes sentinels to the OLDEST segment of
-//           its own buffer, preparing it for reuse.
-// Phase 3 — POLL + SCATTER: each rank volatile-loads from its own current
-//           segment, spinning until sentinels disappear, then scatters
-//           directly to per-tensor output arrays.
-// Phase 4 — ADVANCE: one thread advances the triple-buffer ring counter.
-
-template <int ngpus, int nbufs>
-__global__ void __launch_bounds__(512, 1) moe_allgather_lamport_kernel(
-    int64_t* buf_ptrs,  // [ngpus] IPC buffer base addresses (device)
-    int* counters,      // [0] = unused, [1] = ring (0/1/2), [2] = prev total_sz
-    int rank,
-    int seg_capacity,  // bytes per segment
-    int rank_stride,   // bytes per rank-slot within a segment
-    int total_sz,      // int4 units of actual packed data per rank
-    // inputs (up to 4)
-    const void* inp0, const void* inp1, const void* inp2, const void* inp3,
-    int off0, int sz0, int off1, int sz1, int off2, int sz2, int off3, int sz3,
-    // outputs (up to 4)
-    void* out0, void* out1, void* out2, void* out3) {
-  using V = int4;
-  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  const int stride = gridDim.x * blockDim.x;
-
-  // Read segment index and previous clear size.
-  const int seg = counters[1];            // 0 or 1
-  const int prev_total_sz = counters[2];  // set by previous invocation
-  const int cur_seg = seg;
-  const int old_seg = 1 - seg;
-
-  char* bufs[ngpus];
-#pragma unroll
-  for (int r = 0; r < ngpus; r++)
-    bufs[r] = reinterpret_cast<char*>(buf_ptrs[r]) + cur_seg * seg_capacity;
-
-  // Sentinel vector for clearing.
-  V sent;
-  sent.x = sent.y = sent.z = sent.w = static_cast<int>(SENTINEL);
-
-  // ---- Phase 1: PUSH local data to ALL peers ----
-  // Write to peer_r's buffer at [rank * rank_stride + off_i].
-
-#define PUSH(idx, inp_ptr, off_val, sz_val)                                 \
-  if constexpr (nbufs > (idx)) {                                            \
-    const V* src = reinterpret_cast<const V*>(inp_ptr);                     \
-    for (int i = tid; i < (sz_val); i += stride) {                          \
-      V val = remove_sentinel(src[i]);                                      \
-      _Pragma("unroll") for (int r = 0; r < ngpus; r++) {                   \
-        reinterpret_cast<V*>(bufs[r] + rank * rank_stride + (off_val))[i] = \
-            val;                                                            \
-      }                                                                     \
-    }                                                                       \
-  }
-
-  PUSH(0, inp0, off0, sz0)
-  PUSH(1, inp1, off1, sz1)
-  PUSH(2, inp2, off2, sz2)
-  PUSH(3, inp3, off3, sz3)
-#undef PUSH
-
-  // ---- Phase 2: CLEAR only the previously-written data in oldest segment ----
-  // Only clear what the previous invocation actually wrote (per rank-slot).
-  if (prev_total_sz > 0) {
-    char* clr_base =
-        reinterpret_cast<char*>(buf_ptrs[rank]) + old_seg * seg_capacity;
-#pragma unroll
-    for (int r = 0; r < ngpus; r++) {
-      V* clr = reinterpret_cast<V*>(clr_base + r * rank_stride);
-      for (int i = tid; i < prev_total_sz; i += stride) clr[i] = sent;
-    }
-  }
-
-  // ---- Phase 3: POLL + SCATTER ----
-  // Volatile-load from own buffer; spin until sentinel gone; scatter to output.
-  char* my = bufs[rank];
-
-#define POLL(idx, out_ptr, off_val, sz_val)                                \
-  if constexpr (nbufs > (idx)) {                                           \
-    for (int i = tid; i < (sz_val); i += stride) {                         \
-      _Pragma("unroll") for (int s = 0; s < ngpus; s++) {                  \
-        V val;                                                             \
-        do {                                                               \
-          val = ld128v(                                                    \
-              reinterpret_cast<V*>(my + s * rank_stride + (off_val)) + i); \
-        } while (has_sentinel(val));                                       \
-        reinterpret_cast<V*>(out_ptr)[s * (sz_val) + i] = val;             \
-      }                                                                    \
-    }                                                                      \
-  }
-
-  POLL(0, out0, off0, sz0)
-  POLL(1, out1, off1, sz1)
-  POLL(2, out2, off2, sz2)
-  POLL(3, out3, off3, sz3)
-#undef POLL
-
-  // ---- Phase 4: ADVANCE ring counter + store clear size for next call ----
-  // Stream serialization ensures the next kernel sees these updates.
-  if (blockIdx.x == 0 && threadIdx.x == 0) {
-    counters[1] = 1 - seg;
-    counters[2] = total_sz;
-  }
-}
-
-// ---------------------------------------------------------------------------
-// Sentinel initialization kernel
-// ---------------------------------------------------------------------------
-
-__global__ void lamport_init_kernel(uint32_t* buf, int n) {
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  int stride = gridDim.x * blockDim.x;
-  for (int i = tid; i < n; i += stride) buf[i] = SENTINEL;
-}
-
-// ---------------------------------------------------------------------------
-// Host launcher
-// ---------------------------------------------------------------------------
-
-struct TensorDesc {
-  void* inp;
-  int off;
-  int sz;
-  int64_t nbytes;
-};
-
-static TensorDesc make_desc(torch::Tensor& inp, int64_t& cursor) {
-  TORCH_CHECK(inp.is_contiguous(), "input must be contiguous");
-  int64_t nbytes = inp.numel() * inp.element_size();
-  TORCH_CHECK(nbytes % 16 == 0, "tensor byte size must be multiple of 16, got ",
-              nbytes);
-  cursor = (cursor + 15) & ~15;
-  TensorDesc d;
-  d.inp = inp.data_ptr();
-  d.off = static_cast<int>(cursor);
-  d.sz = static_cast<int>(nbytes / 16);
-  d.nbytes = nbytes;
-  cursor += nbytes;
-  return d;
-}
-
-void lamport_init(int64_t buf_ptr, int64_t nbytes) {
-  auto stream = c10::cuda::getCurrentCUDAStream().stream();
-  int n = static_cast<int>(nbytes / 4);
-  lamport_init_kernel<<<256, 256, 0, stream>>>(
-      reinterpret_cast<uint32_t*>(buf_ptr), n);
-}
-
-void moe_all_gather(int64_t buf_ptrs_ptr, int64_t counters_ptr, int64_t rank,
-                    int64_t world_size, int64_t seg_capacity,
-                    int64_t rank_stride, std::vector<torch::Tensor>& inputs,
-                    std::vector<torch::Tensor>& outputs) {
-  auto stream = c10::cuda::getCurrentCUDAStream().stream();
-  int n = static_cast<int>(inputs.size());
-  TORCH_CHECK(n >= 2 && n <= 4, "2-4 input tensors required");
-  TORCH_CHECK(inputs.size() == outputs.size());
-
-  int64_t cursor = 0;
-  TensorDesc descs[4] = {};
-  for (int i = 0; i < n; i++) descs[i] = make_desc(inputs[i], cursor);
-  TORCH_CHECK(cursor % 16 == 0);
-  int total_sz = static_cast<int>(cursor / 16);
-  TORCH_CHECK(cursor <= rank_stride, "packed data (", cursor,
-              " bytes) exceeds rank_stride (", rank_stride, " bytes)");
-
-  int ws = static_cast<int>(world_size);
-  for (int i = 0; i < n; i++) {
-    TORCH_CHECK(outputs[i].is_contiguous());
-    TORCH_CHECK(outputs[i].numel() == inputs[i].numel() * ws);
-  }
-
-  int r = static_cast<int>(rank);
-  int threads = 512;
-  int blocks =
-      std::max(1, std::min(kMaxBlocks, (total_sz + threads - 1) / threads));
-
-  void *inps[4] = {}, *outs[4] = {};
-  int offs[4] = {}, szs[4] = {};
-  for (int i = 0; i < n; i++) {
-    inps[i] = descs[i].inp;
-    offs[i] = descs[i].off;
-    szs[i] = descs[i].sz;
-    outs[i] = outputs[i].data_ptr();
-  }
-
-  auto* bp = reinterpret_cast<int64_t*>(buf_ptrs_ptr);
-  auto* ct = reinterpret_cast<int*>(counters_ptr);
-  int sc = static_cast<int>(seg_capacity);
-  int rs = static_cast<int>(rank_stride);
-
-#define KL(ng, nb)                                                        \
-  moe_allgather_lamport_kernel<ng, nb><<<blocks, threads, 0, stream>>>(   \
-      bp, ct, r, sc, rs, total_sz, inps[0], inps[1], inps[2], inps[3],    \
-      offs[0], szs[0], offs[1], szs[1], offs[2], szs[2], offs[3], szs[3], \
-      outs[0], outs[1], outs[2], outs[3]);
-
-#define GPU_CASE(ng) \
-  case ng:           \
-    switch (n) {     \
-      case 2:        \
-        KL(ng, 2);   \
-        break;       \
-      case 3:        \
-        KL(ng, 3);   \
-        break;       \
-      case 4:        \
-        KL(ng, 4);   \
-        break;       \
-    }                \
-    break;
-
-  switch (ws) {
-    GPU_CASE(2)
-    GPU_CASE(4)
-    GPU_CASE(6)
-    GPU_CASE(8)
-    default:
-      TORCH_CHECK(false, "world_size must be 2, 4, 6, or 8");
-  }
-#undef GPU_CASE
-#undef KL
-}
-
-// ---------------------------------------------------------------------------
-// Python binding
-// ---------------------------------------------------------------------------
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("moe_all_gather", &moe_all_gather, "Lamport MoE all-gather");
-  m.def("lamport_init", &lamport_init,
-        "Initialize Lamport buffer with sentinels");
-}

csrc/standalone/custom_ar/moe_allgather.py

@@ -1,116 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Lamport-based fused MoE all-gather for EP dispatch.
-
-JIT-compiles the CUDA kernel on first use (cached afterwards).
-Uses a Lamport sentinel protocol (push writes + per-element sync)
-with triple-buffered IPC regions — no explicit barriers.
-
-Usage:
-    ag = MoeAllGather(custom_allreduce)
-    ids_g, wt_g, hs_g, sc_g = ag.gather(topk_ids, topk_weights, hidden, scales)
-"""
-
-from __future__ import annotations
-
-import os
-from pathlib import Path
-
-import torch
-
-_lib = None
-
-
-def _load_lib():
-    global _lib
-    if _lib is not None:
-        return _lib
-
-    from torch.utils.cpp_extension import load
-
-    src = str(Path(__file__).with_name("moe_allgather.cu"))
-    _lib = load(
-        name="moe_allgather_kernel",
-        sources=[src],
-        extra_cuda_cflags=["-O3", "--use_fast_math"],
-        verbose=os.environ.get("MOE_AG_VERBOSE", "") == "1",
-    )
-    return _lib
-
-
-class MoeAllGather:
-    """Lamport-based MoE dispatch all-gather with triple buffering."""
-
-    def __init__(self, ca_comm):
-        self.rank = ca_comm.rank
-        self.world_size = ca_comm.world_size
-        self.device = ca_comm.device
-        self.buffer_ptrs = ca_comm.buffer_ptrs
-        self.max_size = ca_comm.max_size
-
-        ws = self.world_size
-        # Double-buffer layout: 2 segments, each with ws rank-slots.
-        # Safe because kernels in the same stream are serialized, and the
-        # Use the FIRST half of the IPC buffer (second half reserved for
-        # MoeReduceScatter) to avoid overlapping writes.
-        half_size = (self.max_size // 2) & ~15
-
-        # Double-buffer layout within our half: 2 segments, each ws rank-slots.
-        # seg_capacity and rank_stride are 16-byte aligned.
-        self.seg_capacity = (half_size // 2) & ~15
-        self.rank_stride = (self.seg_capacity // ws) & ~15
-        self.max_per_rank = self.rank_stride  # max packed bytes per rank
-
-        # Buffer pointer array on device (no offset — first half).
-        self._buf_ptrs = torch.zeros(
-            8, dtype=torch.int64, device=f"cuda:{self.device.index}"
-        )
-        for i in range(ws):
-            self._buf_ptrs[i] = self.buffer_ptrs[i]
-        self._buf_ptrs_ptr = self._buf_ptrs.data_ptr()
-
-        # Counters on device: [0]=unused, [1]=seg (0/1), [2]=prev_total_sz.
-        self._counters = torch.zeros(
-            3, dtype=torch.int32, device=f"cuda:{self.device.index}"
-        )
-        self._counters_ptr = self._counters.data_ptr()
-
-        # Initialize our half with sentinel values.
-        lib = _load_lib()
-        lib.lamport_init(self.buffer_ptrs[self.rank], half_size)
-        torch.accelerator.synchronize(self.device)
-
-    def gather(
-        self,
-        topk_ids: torch.Tensor,
-        topk_weights: torch.Tensor,
-        hidden_states: torch.Tensor,
-        scales: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
-        lib = _load_lib()
-        ws = self.world_size
-
-        inputs = [topk_ids, topk_weights, hidden_states]
-        if scales is not None:
-            inputs.append(scales)
-
-        outputs = [
-            torch.empty((t.shape[0] * ws, *t.shape[1:]), dtype=t.dtype, device=t.device)
-            for t in inputs
-        ]
-
-        lib.moe_all_gather(
-            self._buf_ptrs_ptr,
-            self._counters_ptr,
-            self.rank,
-            self.world_size,
-            self.seg_capacity,
-            self.rank_stride,
-            inputs,
-            outputs,
-        )
-
-        if scales is not None:
-            return outputs[0], outputs[1], outputs[2], outputs[3]
-        return outputs[0], outputs[1], outputs[2], None

csrc/standalone/custom_ar/moe_reduce_scatter.cu

@@ -1,261 +0,0 @@
-// Lamport-based MoE reduce-scatter kernel for EP combine.
-//
-// JIT-compilable via torch.utils.cpp_extension — no vLLM build required.
-//
-// Reduce-scatters a bf16 tensor [N_total, D] across EP ranks.  Each rank
-// contributes its partial MoE output; the kernel sums all contributions
-// and each rank receives its own slice of the result.
-//
-// Protocol (same as the all-gather variant):
-//   1. PUSH: write own data to all peers' Lamport buffers (NVLink push).
-//   2. CLEAR: write sentinels to old segment of own buffer.
-//   3. POLL + REDUCE: volatile-load all peers' data for own slice,
-//      accumulate in fp32, convert back to bf16, store to output.
-//   4. ADVANCE: toggle double-buffer index.
-//
-// Sentinel: 0x80000000 (two bf16 negative-zeros packed in uint32).
-// For bf16 reduce, replacing -0 with +0 is lossless.
-
-#include <cuda.h>
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-#include <torch/extension.h>
-
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAStream.h>
-
-#define DINLINE __device__ __forceinline__
-
-constexpr uint32_t SENTINEL = 0x80000000u;
-constexpr int kMaxBlocks = 36;
-
-// ---------------------------------------------------------------------------
-// Volatile 128-bit load and sentinel helpers
-// ---------------------------------------------------------------------------
-
-static DINLINE int4 ld128v(const void* addr) {
-  int4 v;
-  asm volatile("ld.volatile.global.v4.b32 {%0,%1,%2,%3}, [%4];"
-               : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
-               : "l"(addr));
-  return v;
-}
-
-static DINLINE bool has_sentinel(int4 v) {
-  return reinterpret_cast<uint32_t&>(v.x) == SENTINEL |
-         reinterpret_cast<uint32_t&>(v.y) == SENTINEL |
-         reinterpret_cast<uint32_t&>(v.z) == SENTINEL |
-         reinterpret_cast<uint32_t&>(v.w) == SENTINEL;
-}
-
-static DINLINE int4 remove_sentinel(int4 v) {
-  if (reinterpret_cast<uint32_t&>(v.x) == SENTINEL) v.x = 0;
-  if (reinterpret_cast<uint32_t&>(v.y) == SENTINEL) v.y = 0;
-  if (reinterpret_cast<uint32_t&>(v.z) == SENTINEL) v.z = 0;
-  if (reinterpret_cast<uint32_t&>(v.w) == SENTINEL) v.w = 0;
-  return v;
-}
-
-// ---------------------------------------------------------------------------
-// bf16 ↔ fp32 helpers for int4 (8 bf16 values = 16 bytes)
-// ---------------------------------------------------------------------------
-
-// Accumulate 8 bf16 values from an int4 into 8 fp32 accumulators.
-static DINLINE void accumulate_bf16(float* acc, int4 v) {
-  const __nv_bfloat16* bp = reinterpret_cast<const __nv_bfloat16*>(&v);
-#pragma unroll
-  for (int k = 0; k < 8; k++) acc[k] += __bfloat162float(bp[k]);
-}
-
-// Convert 8 fp32 accumulators to bf16 and pack into int4.
-static DINLINE int4 fp32_to_bf16_int4(const float* acc) {
-  int4 out;
-  __nv_bfloat16* bp = reinterpret_cast<__nv_bfloat16*>(&out);
-#pragma unroll
-  for (int k = 0; k < 8; k++) bp[k] = __float2bfloat16(acc[k]);
-  return out;
-}
-
-// ---------------------------------------------------------------------------
-// Lamport reduce-scatter kernel
-// ---------------------------------------------------------------------------
-
-template <int ngpus>
-__global__ void __launch_bounds__(512, 1) moe_rs_lamport_kernel(
-    int64_t* buf_ptrs,  // [ngpus] IPC buffer base addresses (device)
-    int* counters,      // [0] = unused, [1] = seg (0/1), [2] = prev total_sz
-    int rank,
-    int seg_capacity,   // bytes per segment
-    int rank_stride,    // bytes per rank-slot within a segment
-    const void* input,  // [N_total, D] bf16 — full input
-    void* output,       // [N_per_rank, D] bf16 — this rank's reduced slice
-    int total_sz,       // int4 units of full input per rank
-    int slice_off,      // int4 offset to this rank's slice within packed data
-    int slice_sz) {     // int4 units of this rank's slice
-  using V = int4;
-  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  const int stride = gridDim.x * blockDim.x;
-
-  // Read segment index and previous clear size.
-  const int seg = counters[1];
-  const int prev_total_sz = counters[2];
-  const int cur_seg = seg;
-  const int old_seg = 1 - seg;
-
-  char* bufs[ngpus];
-#pragma unroll
-  for (int r = 0; r < ngpus; r++)
-    bufs[r] = reinterpret_cast<char*>(buf_ptrs[r]) + cur_seg * seg_capacity;
-
-  V sent;
-  sent.x = sent.y = sent.z = sent.w = static_cast<int>(SENTINEL);
-
-  // ---- Phase 1: PUSH full input to ALL peers ----
-  {
-    const V* src = reinterpret_cast<const V*>(input);
-    for (int i = tid; i < total_sz; i += stride) {
-      V val = remove_sentinel(src[i]);
-#pragma unroll
-      for (int r = 0; r < ngpus; r++)
-        reinterpret_cast<V*>(bufs[r] + rank * rank_stride)[i] = val;
-    }
-  }
-
-  // ---- Phase 2: CLEAR old segment ----
-  if (prev_total_sz > 0) {
-    char* clr_base =
-        reinterpret_cast<char*>(buf_ptrs[rank]) + old_seg * seg_capacity;
-#pragma unroll
-    for (int r = 0; r < ngpus; r++) {
-      V* clr = reinterpret_cast<V*>(clr_base + r * rank_stride);
-      for (int i = tid; i < prev_total_sz; i += stride) clr[i] = sent;
-    }
-  }
-
-  // ---- Phase 3: POLL + REDUCE for own slice ----
-  // Read all ranks' data at [slice_off, slice_off + slice_sz) from own buffer,
-  // sum in fp32, store bf16 result.
-  {
-    char* my = bufs[rank];
-    V* dst = reinterpret_cast<V*>(output);
-
-    for (int i = tid; i < slice_sz; i += stride) {
-      float acc[8] = {0, 0, 0, 0, 0, 0, 0, 0};
-
-#pragma unroll
-      for (int s = 0; s < ngpus; s++) {
-        V val;
-        do {
-          val = ld128v(reinterpret_cast<V*>(my + s * rank_stride) + slice_off +
-                       i);
-        } while (has_sentinel(val));
-        accumulate_bf16(acc, val);
-      }
-
-      dst[i] = fp32_to_bf16_int4(acc);
-    }
-  }
-
-  // ---- Phase 4: ADVANCE ----
-  if (blockIdx.x == 0 && threadIdx.x == 0) {
-    counters[1] = 1 - seg;
-    counters[2] = total_sz;
-  }
-}
-
-// ---------------------------------------------------------------------------
-// Sentinel initialization
-// ---------------------------------------------------------------------------
-
-__global__ void lamport_init_kernel(uint32_t* buf, int n) {
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  int stride = gridDim.x * blockDim.x;
-  for (int i = tid; i < n; i += stride) buf[i] = SENTINEL;
-}
-
-// ---------------------------------------------------------------------------
-// Host launcher
-// ---------------------------------------------------------------------------
-
-void lamport_init(int64_t buf_ptr, int64_t nbytes) {
-  auto stream = c10::cuda::getCurrentCUDAStream().stream();
-  int n = static_cast<int>(nbytes / 4);
-  lamport_init_kernel<<<256, 256, 0, stream>>>(
-      reinterpret_cast<uint32_t*>(buf_ptr), n);
-}
-
-void moe_reduce_scatter(int64_t buf_ptrs_ptr, int64_t counters_ptr,
-                        int64_t rank, int64_t world_size, int64_t seg_capacity,
-                        int64_t rank_stride, torch::Tensor input,
-                        torch::Tensor output) {
-  auto stream = c10::cuda::getCurrentCUDAStream().stream();
-  TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
-  TORCH_CHECK(output.is_contiguous(), "output must be contiguous");
-  TORCH_CHECK(input.scalar_type() == torch::kBFloat16,
-              "input must be bf16, got ", input.scalar_type());
-  TORCH_CHECK(output.scalar_type() == torch::kBFloat16, "output must be bf16");
-
-  int ws = static_cast<int>(world_size);
-  int r = static_cast<int>(rank);
-
-  // Input: [N_total, D], Output: [N_per_rank, D]
-  int64_t N_total = input.size(0);
-  int64_t D = input.size(1);
-  TORCH_CHECK(N_total % ws == 0, "N_total must be divisible by world_size");
-  int64_t N_per_rank = N_total / ws;
-  TORCH_CHECK(output.size(0) == N_per_rank);
-  TORCH_CHECK(output.size(1) == D);
-
-  int64_t input_bytes = input.numel() * input.element_size();
-  TORCH_CHECK(input_bytes % 16 == 0,
-              "input byte size must be multiple of 16, got ", input_bytes);
-  TORCH_CHECK(input_bytes <= rank_stride, "input (", input_bytes,
-              " bytes) exceeds rank_stride (", rank_stride, " bytes)");
-
-  int total_sz = static_cast<int>(input_bytes / 16);
-  int slice_sz = total_sz / ws;
-  int slice_off = r * slice_sz;
-
-  int threads = 512;
-  int blocks =
-      std::max(1, std::min(kMaxBlocks, (total_sz + threads - 1) / threads));
-
-  auto* bp = reinterpret_cast<int64_t*>(buf_ptrs_ptr);
-  auto* ct = reinterpret_cast<int*>(counters_ptr);
-  int sc = static_cast<int>(seg_capacity);
-  int rs = static_cast<int>(rank_stride);
-
-#define LAUNCH(ng)                                                      \
-  moe_rs_lamport_kernel<ng><<<blocks, threads, 0, stream>>>(            \
-      bp, ct, r, sc, rs, input.data_ptr(), output.data_ptr(), total_sz, \
-      slice_off, slice_sz);
-
-  switch (ws) {
-    case 2:
-      LAUNCH(2);
-      break;
-    case 4:
-      LAUNCH(4);
-      break;
-    case 6:
-      LAUNCH(6);
-      break;
-    case 8:
-      LAUNCH(8);
-      break;
-    default:
-      TORCH_CHECK(false, "world_size must be 2, 4, 6, or 8");
-  }
-#undef LAUNCH
-}
-
-// ---------------------------------------------------------------------------
-// Python binding
-// ---------------------------------------------------------------------------
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("moe_reduce_scatter", &moe_reduce_scatter,
-        "Lamport MoE reduce-scatter");
-  m.def("lamport_init", &lamport_init,
-        "Initialize Lamport buffer with sentinels");
-}

csrc/standalone/custom_ar/moe_reduce_scatter.py

@@ -1,102 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Lamport-based MoE reduce-scatter for EP combine.
-
-JIT-compiles the CUDA kernel on first use (cached afterwards).
-Uses the same Lamport sentinel protocol as the all-gather kernel.
-"""
-
-from __future__ import annotations
-
-import os
-from pathlib import Path
-
-import torch
-
-_lib = None
-
-
-def _load_lib():
-    global _lib
-    if _lib is not None:
-        return _lib
-
-    from torch.utils.cpp_extension import load
-
-    src = str(Path(__file__).with_name("moe_reduce_scatter.cu"))
-    _lib = load(
-        name="moe_reduce_scatter_kernel",
-        sources=[src],
-        extra_cuda_cflags=["-O3", "--use_fast_math"],
-        verbose=os.environ.get("MOE_RS_VERBOSE", "") == "1",
-    )
-    return _lib
-
-
-class MoeReduceScatter:
-    """Lamport-based MoE combine reduce-scatter with double buffering."""
-
-    def __init__(self, ca_comm):
-        self.rank = ca_comm.rank
-        self.world_size = ca_comm.world_size
-        self.device = ca_comm.device
-        self.buffer_ptrs = ca_comm.buffer_ptrs
-        self.max_size = ca_comm.max_size
-
-        ws = self.world_size
-        # Use the SECOND half of the IPC buffer (first half reserved for
-        # MoeAllGather) to avoid overlapping writes.
-        half_size = (self.max_size // 2) & ~15
-        self.buffer_offset = half_size
-
-        # Double-buffer layout within our half: 2 segments, each ws rank-slots.
-        self.seg_capacity = (half_size // 2) & ~15
-        self.rank_stride = (self.seg_capacity // ws) & ~15
-        self.max_per_rank = self.rank_stride
-
-        # Buffer pointer array on device — offset to our half.
-        self._buf_ptrs = torch.zeros(
-            8, dtype=torch.int64, device=f"cuda:{self.device.index}"
-        )
-        for i in range(ws):
-            self._buf_ptrs[i] = self.buffer_ptrs[i] + self.buffer_offset
-        self._buf_ptrs_ptr = self._buf_ptrs.data_ptr()
-
-        # Counters: [0]=unused, [1]=seg (0/1), [2]=prev_total_sz.
-        self._counters = torch.zeros(
-            3, dtype=torch.int32, device=f"cuda:{self.device.index}"
-        )
-        self._counters_ptr = self._counters.data_ptr()
-
-        # Initialize our half with sentinels.
-        lib = _load_lib()
-        lib.lamport_init(self.buffer_ptrs[self.rank] + self.buffer_offset, half_size)
-        torch.accelerator.synchronize(self.device)
-
-    def reduce_scatter(
-        self,
-        input: torch.Tensor,
-    ) -> torch.Tensor:
-        """Reduce-scatter input [N_total, D] bf16 → output [N_per_rank, D] bf16."""
-        lib = _load_lib()
-        ws = self.world_size
-
-        assert input.dim() == 2
-        N_total, D = input.shape
-        assert N_total % ws == 0
-        N_per_rank = N_total // ws
-
-        output = torch.empty((N_per_rank, D), dtype=input.dtype, device=input.device)
-
-        lib.moe_reduce_scatter(
-            self._buf_ptrs_ptr,
-            self._counters_ptr,
-            self.rank,
-            self.world_size,
-            self.seg_capacity,
-            self.rank_stride,
-            input,
-            output,
-        )
-        return output

csrc/standalone/custom_ar/moe_rs_fused.cu

@@ -1,308 +0,0 @@
-// Lamport reduce-scatter fused with residual add + RMSNorm.
-//
-// Replaces three separate kernels (RS + residual_add + RMSNorm) with one:
-//   1. PUSH: write MoE output to all peers' Lamport buffers.
-//   2. CLEAR: write sentinels to old segment.
-//   3. POLL+REDUCE+FUSE (per-token):
-//      a. Volatile-load from all peers, sum in fp32.
-//      b. Add residual.
-//      c. Compute RMSNorm (block reduction for variance).
-//      d. Store normed output + updated residual.
-//
-// Saves: one kernel launch (~3-5µs) + one global memory round-trip per layer.
-
-#include <cuda.h>
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-#include <torch/extension.h>
-
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAStream.h>
-
-#define DINLINE __device__ __forceinline__
-
-constexpr uint32_t SENTINEL = 0x80000000u;
-constexpr int kMaxBlocks = 36;
-
-// Each token has D=7168 bf16 values = 896 int4 vectors.
-// With 512 threads: ceil(896/512) = 2 int4 per thread = 16 fp32 values.
-constexpr int kMaxValsPerThread = 16;
-
-static DINLINE int4 ld128v(const void* addr) {
-  int4 v;
-  asm volatile("ld.volatile.global.v4.b32 {%0,%1,%2,%3}, [%4];"
-               : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
-               : "l"(addr));
-  return v;
-}
-
-static DINLINE bool has_sentinel(int4 v) {
-  return reinterpret_cast<uint32_t&>(v.x) == SENTINEL |
-         reinterpret_cast<uint32_t&>(v.y) == SENTINEL |
-         reinterpret_cast<uint32_t&>(v.z) == SENTINEL |
-         reinterpret_cast<uint32_t&>(v.w) == SENTINEL;
-}
-
-static DINLINE int4 remove_sentinel(int4 v) {
-  if (reinterpret_cast<uint32_t&>(v.x) == SENTINEL) v.x = 0;
-  if (reinterpret_cast<uint32_t&>(v.y) == SENTINEL) v.y = 0;
-  if (reinterpret_cast<uint32_t&>(v.z) == SENTINEL) v.z = 0;
-  if (reinterpret_cast<uint32_t&>(v.w) == SENTINEL) v.w = 0;
-  return v;
-}
-
-// bf16 helpers
-static DINLINE void accumulate_bf16(float* acc, int4 v) {
-  const __nv_bfloat16* bp = reinterpret_cast<const __nv_bfloat16*>(&v);
-#pragma unroll
-  for (int k = 0; k < 8; k++) acc[k] += __bfloat162float(bp[k]);
-}
-
-static DINLINE void add_bf16_to_fp32(float* dst, int4 v) {
-  const __nv_bfloat16* bp = reinterpret_cast<const __nv_bfloat16*>(&v);
-#pragma unroll
-  for (int k = 0; k < 8; k++) dst[k] += __bfloat162float(bp[k]);
-}
-
-static DINLINE int4 fp32_to_bf16_int4(const float* vals) {
-  int4 out;
-  __nv_bfloat16* bp = reinterpret_cast<__nv_bfloat16*>(&out);
-#pragma unroll
-  for (int k = 0; k < 8; k++) bp[k] = __float2bfloat16(vals[k]);
-  return out;
-}
-
-// Block-level tree reduction in shared memory.
-static DINLINE float block_reduce_sum(float val, float* smem) {
-  smem[threadIdx.x] = val;
-  __syncthreads();
-  for (int s = blockDim.x / 2; s > 0; s >>= 1) {
-    if (threadIdx.x < s) smem[threadIdx.x] += smem[threadIdx.x + s];
-    __syncthreads();
-  }
-  return smem[0];
-}
-
-// ---------------------------------------------------------------------------
-// Fused reduce-scatter + residual + RMSNorm kernel
-// ---------------------------------------------------------------------------
-
-template <int ngpus>
-__global__ void __launch_bounds__(512, 1) moe_rs_fused_kernel(
-    int64_t* buf_ptrs, int* counters, int rank, int seg_capacity,
-    int rank_stride,
-    const void* input,        // [N_total, D] bf16 — MoE output
-    const void* residual_in,  // [N_per_rank, D] bf16 — skip connection
-    const void* gamma,        // [D] bf16 — RMSNorm weight
-    void* normed_out,         // [N_per_rank, D] bf16 — normed result
-    void* residual_out,       // [N_per_rank, D] bf16 — updated residual
-    int total_sz,             // int4 units of full input
-    int slice_off,            // int4 offset to this rank's slice
-    int slice_sz,             // int4 units of this rank's slice
-    int D_int4,               // int4 units per token (hidden_dim * 2 / 16)
-    int N_per_rank,           // tokens in this rank's slice
-    float eps) {              // RMSNorm epsilon
-  using V = int4;
-  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  const int stride = gridDim.x * blockDim.x;
-
-  const int seg = counters[1];
-  const int prev_total_sz = counters[2];
-  const int cur_seg = seg;
-  const int old_seg = 1 - seg;
-
-  char* bufs[ngpus];
-#pragma unroll
-  for (int r = 0; r < ngpus; r++)
-    bufs[r] = reinterpret_cast<char*>(buf_ptrs[r]) + cur_seg * seg_capacity;
-
-  V sent;
-  sent.x = sent.y = sent.z = sent.w = static_cast<int>(SENTINEL);
-
-  // ---- Phase 1: PUSH full MoE output to ALL peers ----
-  {
-    const V* src = reinterpret_cast<const V*>(input);
-    for (int i = tid; i < total_sz; i += stride) {
-      V val = remove_sentinel(src[i]);
-#pragma unroll
-      for (int r = 0; r < ngpus; r++)
-        reinterpret_cast<V*>(bufs[r] + rank * rank_stride)[i] = val;
-    }
-  }
-
-  // ---- Phase 2: CLEAR old segment ----
-  if (prev_total_sz > 0) {
-    char* clr_base =
-        reinterpret_cast<char*>(buf_ptrs[rank]) + old_seg * seg_capacity;
-#pragma unroll
-    for (int r = 0; r < ngpus; r++) {
-      V* clr = reinterpret_cast<V*>(clr_base + r * rank_stride);
-      for (int i = tid; i < prev_total_sz; i += stride) clr[i] = sent;
-    }
-  }
-
-  // ---- Phase 3: Fused POLL + REDUCE + RESIDUAL + RMSNORM ----
-  // Each block handles one token. Only first N_per_rank blocks participate.
-  if (blockIdx.x < N_per_rank) {
-    const int token = blockIdx.x;
-    const int token_off = slice_off + token * D_int4;  // in the full buffer
-    char* my = bufs[rank];
-
-    extern __shared__ float smem[];
-
-    // Register storage for intermediate fp32 values.
-    float local_vals[kMaxValsPerThread];
-    int n_vals = 0;
-    float partial_sum_sq = 0.0f;
-
-    // Pass 1: poll + reduce + add residual + compute sum_sq
-    for (int pos = threadIdx.x; pos < D_int4; pos += blockDim.x) {
-      float acc[8] = {0, 0, 0, 0, 0, 0, 0, 0};
-
-      // Poll all ranks' data for this position.
-#pragma unroll
-      for (int s = 0; s < ngpus; s++) {
-        V val;
-        do {
-          val = ld128v(reinterpret_cast<V*>(my + s * rank_stride) + token_off +
-                       pos);
-        } while (has_sentinel(val));
-        accumulate_bf16(acc, val);
-      }
-
-      // Add residual.
-      V res = reinterpret_cast<const V*>(residual_in)[token * D_int4 + pos];
-      add_bf16_to_fp32(acc, res);
-
-      // Store in registers and accumulate sum_sq.
-#pragma unroll
-      for (int k = 0; k < 8; k++) {
-        local_vals[n_vals++] = acc[k];
-        partial_sum_sq += acc[k] * acc[k];
-      }
-    }
-
-    // Block-level reduction: total sum of squares.
-    float total_sum_sq = block_reduce_sum(partial_sum_sq, smem);
-    float rms_scale = rsqrtf(total_sum_sq / (D_int4 * 8) + eps);
-
-    // Pass 2: apply RMSNorm, store outputs.
-    n_vals = 0;
-    const V* gamma_v = reinterpret_cast<const V*>(gamma);
-    for (int pos = threadIdx.x; pos < D_int4; pos += blockDim.x) {
-      V gv = gamma_v[pos];
-      const __nv_bfloat16* gp = reinterpret_cast<const __nv_bfloat16*>(&gv);
-
-      // Build normed output and residual output.
-      float normed_fp32[8], res_fp32[8];
-#pragma unroll
-      for (int k = 0; k < 8; k++) {
-        float val = local_vals[n_vals++];
-        res_fp32[k] = val;  // residual_out
-        normed_fp32[k] = val * rms_scale * __bfloat162float(gp[k]);  // normed
-      }
-
-      reinterpret_cast<V*>(normed_out)[token * D_int4 + pos] =
-          fp32_to_bf16_int4(normed_fp32);
-      reinterpret_cast<V*>(residual_out)[token * D_int4 + pos] =
-          fp32_to_bf16_int4(res_fp32);
-    }
-  }
-
-  // ---- Phase 4: ADVANCE ----
-  if (blockIdx.x == 0 && threadIdx.x == 0) {
-    counters[1] = 1 - seg;
-    counters[2] = total_sz;
-  }
-}
-
-// ---------------------------------------------------------------------------
-// Sentinel init + host launcher
-// ---------------------------------------------------------------------------
-
-__global__ void lamport_init_kernel(uint32_t* buf, int n) {
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  int stride = gridDim.x * blockDim.x;
-  for (int i = tid; i < n; i += stride) buf[i] = SENTINEL;
-}
-
-void lamport_init(int64_t buf_ptr, int64_t nbytes) {
-  auto stream = c10::cuda::getCurrentCUDAStream().stream();
-  int n = static_cast<int>(nbytes / 4);
-  lamport_init_kernel<<<256, 256, 0, stream>>>(
-      reinterpret_cast<uint32_t*>(buf_ptr), n);
-}
-
-void moe_rs_fused(int64_t buf_ptrs_ptr, int64_t counters_ptr, int64_t rank,
-                  int64_t world_size, int64_t seg_capacity, int64_t rank_stride,
-                  torch::Tensor input,         // [N_total, D] bf16
-                  torch::Tensor residual_in,   // [N_per_rank, D] bf16
-                  torch::Tensor gamma,         // [D] bf16
-                  torch::Tensor normed_out,    // [N_per_rank, D] bf16
-                  torch::Tensor residual_out,  // [N_per_rank, D] bf16
-                  double eps) {
-  auto stream = c10::cuda::getCurrentCUDAStream().stream();
-  TORCH_CHECK(input.is_contiguous() && residual_in.is_contiguous());
-  TORCH_CHECK(gamma.is_contiguous() && normed_out.is_contiguous());
-  TORCH_CHECK(residual_out.is_contiguous());
-  TORCH_CHECK(input.scalar_type() == torch::kBFloat16);
-
-  int ws = static_cast<int>(world_size);
-  int r = static_cast<int>(rank);
-  int64_t N_total = input.size(0);
-  int64_t D = input.size(1);
-  TORCH_CHECK(N_total % ws == 0);
-  int N_per_rank = static_cast<int>(N_total / ws);
-
-  int64_t input_bytes = input.numel() * input.element_size();
-  TORCH_CHECK(input_bytes % 16 == 0);
-  TORCH_CHECK(input_bytes <= rank_stride);
-
-  int total_sz = static_cast<int>(input_bytes / 16);
-  int D_int4 = static_cast<int>(D * 2 / 16);  // bf16 elements → int4 units
-  int slice_sz = total_sz / ws;
-  int slice_off = r * slice_sz;
-
-  int threads = 512;
-  // Need at least N_per_rank blocks for Phase 3 (one per token).
-  int blocks = std::max(
-      N_per_rank, std::min(kMaxBlocks, (total_sz + threads - 1) / threads));
-
-  auto* bp = reinterpret_cast<int64_t*>(buf_ptrs_ptr);
-  auto* ct = reinterpret_cast<int*>(counters_ptr);
-  int sc = static_cast<int>(seg_capacity);
-  int rs = static_cast<int>(rank_stride);
-  int smem = threads * sizeof(float);
-
-#define LAUNCH(ng)                                                      \
-  moe_rs_fused_kernel<ng><<<blocks, threads, smem, stream>>>(           \
-      bp, ct, r, sc, rs, input.data_ptr(), residual_in.data_ptr(),      \
-      gamma.data_ptr(), normed_out.data_ptr(), residual_out.data_ptr(), \
-      total_sz, slice_off, slice_sz, D_int4, N_per_rank,                \
-      static_cast<float>(eps));
-
-  switch (ws) {
-    case 2:
-      LAUNCH(2);
-      break;
-    case 4:
-      LAUNCH(4);
-      break;
-    case 6:
-      LAUNCH(6);
-      break;
-    case 8:
-      LAUNCH(8);
-      break;
-    default:
-      TORCH_CHECK(false, "world_size must be 2, 4, 6, or 8");
-  }
-#undef LAUNCH
-}
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("moe_rs_fused", &moe_rs_fused,
-        "Fused Lamport reduce-scatter + residual + RMSNorm");
-  m.def("lamport_init", &lamport_init,
-        "Initialize Lamport buffer with sentinels");
-}

csrc/standalone/custom_ar/moe_rs_fused.py

@@ -1,29 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-JIT wrapper for the fused reduce-scatter + residual + RMSNorm kernel.
-"""
-
-from __future__ import annotations
-
-import os
-from pathlib import Path
-
-_lib = None
-
-
-def _load_lib():
-    global _lib
-    if _lib is not None:
-        return _lib
-
-    from torch.utils.cpp_extension import load
-
-    src = str(Path(__file__).with_name("moe_rs_fused.cu"))
-    _lib = load(
-        name="moe_rs_fused_kernel",
-        sources=[src],
-        extra_cuda_cflags=["-O3", "--use_fast_math"],
-        verbose=os.environ.get("MOE_RS_FUSED_VERBOSE", "") == "1",
-    )
-    return _lib

csrc/standalone/custom_ar/test_moe_allgather.py

@@ -1,137 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Test for the Lamport-based MoE all-gather kernel."""
-
-import ctypes
-import os
-import sys
-
-import torch
-import torch.distributed as dist
-
-_cudart = ctypes.CDLL("libcudart.so")
-IPC = 64
-
-
-def _cc(r):
-    if r:
-        raise RuntimeError(f"CUDA err {r}")
-
-
-def ipc_buf(sz, rank, ws):
-    p = ctypes.c_void_p()
-    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
-    _cc(_cudart.cudaMemset(p, 0, sz))
-    _cc(_cudart.cudaDeviceSynchronize())
-    h = (ctypes.c_byte * IPC)()
-    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
-    ah = [None] * ws
-    dist.all_gather_object(ah, bytes(h))
-    ptrs = []
-    for i in range(ws):
-        if i == rank:
-            ptrs.append(p.value)
-        else:
-            hh = (ctypes.c_byte * IPC)(*ah[i])
-            pp = ctypes.c_void_p()
-            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
-            ptrs.append(pp.value)
-    return ptrs
-
-
-def main():
-    dist.init_process_group("nccl")
-    rank = dist.get_rank()
-    ws = dist.get_world_size()
-    torch.cuda.set_device(rank)
-    dev = f"cuda:{rank}"
-
-    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-    if rank == 0:
-        from moe_allgather import MoeAllGather, _load_lib
-
-        _load_lib()
-    dist.barrier()
-    from moe_allgather import MoeAllGather, _load_lib
-
-    _load_lib()
-    dist.barrier()
-
-    max_size = 8 * 1024 * 1024
-    bp = ipc_buf(max_size, rank, ws)
-    dist.barrier()
-
-    # Build a fake ca_comm-like object.
-    class FakeCA:
-        pass
-
-    ca = FakeCA()
-    ca.rank = rank
-    ca.world_size = ws
-    ca.device = torch.device(dev)
-    ca.buffer_ptrs = bp
-    ca.max_size = max_size
-    # meta_ptrs not needed for Lamport approach
-    ag = MoeAllGather(ca)
-    dist.barrier()
-
-    errors = 0
-
-    # Test with various token counts.
-    for N in [1, 4, 16, 64]:
-        topk = 8
-        hd = 3584
-        sd = 448
-        ids = (
-            torch.arange(N * topk, dtype=torch.int32, device=dev) + rank * 1000
-        ).reshape(N, topk)
-        wt = torch.ones(N, topk, dtype=torch.float32, device=dev) * (rank + 1) * 0.1
-        hs = torch.full((N, hd), rank + 1, dtype=torch.uint8, device=dev)
-        sc = torch.full((N, sd), rank + 1, dtype=torch.uint8, device=dev)
-
-        ids_g, wt_g, hs_g, sc_g = ag.gather(ids, wt, hs, sc)
-
-        for src in range(ws):
-            s, e = src * N, (src + 1) * N
-            exp_ids = (
-                torch.arange(N * topk, dtype=torch.int32, device=dev) + src * 1000
-            ).reshape(N, topk)
-            if not torch.equal(ids_g[s:e], exp_ids):
-                print(f"[{rank}] FAIL ids src={src} N={N}")
-                errors += 1
-            exp_wt = torch.full(
-                (N, topk), (src + 1) * 0.1, dtype=torch.float32, device=dev
-            )
-            if not torch.allclose(wt_g[s:e], exp_wt):
-                print(f"[{rank}] FAIL wt src={src} N={N}")
-                errors += 1
-            exp_hs = torch.full((N, hd), src + 1, dtype=torch.uint8, device=dev)
-            if not torch.equal(hs_g[s:e], exp_hs):
-                print(f"[{rank}] FAIL hs src={src} N={N}")
-                errors += 1
-            exp_sc = torch.full((N, sd), src + 1, dtype=torch.uint8, device=dev)
-            if not torch.equal(sc_g[s:e], exp_sc):
-                print(f"[{rank}] FAIL sc src={src} N={N}")
-                errors += 1
-
-        # Without scales.
-        ids_g2, wt_g2, hs_g2, _ = ag.gather(ids, wt, hs)
-        for src in range(ws):
-            s, e = src * N, (src + 1) * N
-            exp_ids = (
-                torch.arange(N * topk, dtype=torch.int32, device=dev) + src * 1000
-            ).reshape(N, topk)
-            if not torch.equal(ids_g2[s:e], exp_ids):
-                print(f"[{rank}] FAIL no-sc ids src={src} N={N}")
-                errors += 1
-
-    dist.barrier()
-    print(
-        f"[rank {rank}] {'PASSED' if errors == 0 else f'FAILED ({errors})'} (ws={ws})"
-    )
-    dist.destroy_process_group()
-    return errors
-
-
-if __name__ == "__main__":
-    sys.exit(main())

csrc/standalone/custom_ar/test_moe_allgather_stress.py

@@ -1,173 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Stress test: random data, check bitwise correctness against NCCL."""
-
-import ctypes
-import os
-import sys
-
-import torch
-import torch.distributed as dist
-
-_cudart = ctypes.CDLL("libcudart.so")
-IPC = 64
-
-
-def _cc(r):
-    if r:
-        raise RuntimeError(f"CUDA err {r}")
-
-
-def ipc_buf(sz, rank, ws):
-    p = ctypes.c_void_p()
-    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
-    _cc(_cudart.cudaMemset(p, 0, sz))
-    _cc(_cudart.cudaDeviceSynchronize())
-    h = (ctypes.c_byte * IPC)()
-    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
-    ah = [None] * ws
-    dist.all_gather_object(ah, bytes(h))
-    ptrs = []
-    for i in range(ws):
-        if i == rank:
-            ptrs.append(p.value)
-        else:
-            hh = (ctypes.c_byte * IPC)(*ah[i])
-            pp = ctypes.c_void_p()
-            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
-            ptrs.append(pp.value)
-    return ptrs
-
-
-def main():
-    dist.init_process_group("nccl")
-    rank = dist.get_rank()
-    ws = dist.get_world_size()
-    torch.cuda.set_device(rank)
-    dev = f"cuda:{rank}"
-
-    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-    if rank == 0:
-        from moe_allgather import _load_lib
-
-        _load_lib()
-    dist.barrier()
-    from moe_allgather import MoeAllGather, _load_lib
-
-    _load_lib()
-    dist.barrier()
-
-    max_size = 8 * 1024 * 1024
-    bp = ipc_buf(max_size, rank, ws)
-    dist.barrier()
-
-    class FakeCA:
-        pass
-
-    ca = FakeCA()
-    ca.rank = rank
-    ca.world_size = ws
-    ca.device = torch.device(dev)
-    ca.buffer_ptrs = bp
-    ca.max_size = max_size
-    ag = MoeAllGather(ca)
-    dist.barrier()
-
-    topk = 8
-    hd = 3584
-    sd = 448
-    errors = 0
-    total_checks = 0
-    sentinel_collisions = 0
-
-    for trial in range(200):
-        # All ranks must use the same N for NCCL reference.
-        N_tensor = torch.randint(1, 65, (1,), device=dev)
-        dist.broadcast(N_tensor, src=0)
-        N = N_tensor.item()
-        # Random data including possible sentinel values
-        ids = torch.randint(0, 256, (N, topk), dtype=torch.int32, device=dev)
-        wt = torch.randn(N, topk, dtype=torch.float32, device=dev)
-        hs = torch.randint(0, 256, (N, hd), dtype=torch.uint8, device=dev)
-        sc = torch.randint(0, 256, (N, sd), dtype=torch.uint8, device=dev)
-
-        # Count sentinel patterns in hidden_states (as uint32 view)
-        hs_u32 = hs.view(torch.int32)
-        sentinel_collisions += (hs_u32 == 0x80000000).sum().item()
-
-        # Custom kernel
-        ids_g, wt_g, hs_g, sc_g = ag.gather(ids, wt, hs, sc)
-
-        # NCCL reference
-        ids_ref = torch.empty(N * ws, topk, dtype=torch.int32, device=dev)
-        wt_ref = torch.empty(N * ws, topk, dtype=torch.float32, device=dev)
-        hs_ref = torch.empty(N * ws, hd, dtype=torch.uint8, device=dev)
-        sc_ref = torch.empty(N * ws, sd, dtype=torch.uint8, device=dev)
-        dist.all_gather_into_tensor(ids_ref, ids)
-        dist.all_gather_into_tensor(wt_ref, wt)
-        dist.all_gather_into_tensor(hs_ref, hs)
-        dist.all_gather_into_tensor(sc_ref, sc)
-
-        # Compare
-        if not torch.equal(ids_g, ids_ref):
-            mismatches = (ids_g != ids_ref).sum().item()
-            if trial < 5 or mismatches > 0:
-                print(f"[{rank}] trial={trial} ids MISMATCH: {mismatches} elements")
-            errors += 1
-        if not torch.equal(wt_g, wt_ref):
-            # Check for -0 vs +0 differences
-            bit_diff = wt_g.view(torch.int32) != wt_ref.view(torch.int32)
-            neg_zero_mask = wt_ref.view(torch.int32) == 0x80000000
-            real_errors = bit_diff & ~neg_zero_mask
-            if real_errors.any():
-                print(
-                    f"[{rank}] trial={trial} wt MISMATCH (non-negzero): {real_errors.sum().item()}"
-                )
-                errors += 1
-        if not torch.equal(hs_g, hs_ref):
-            mismatches = (hs_g != hs_ref).sum().item()
-            # Check if mismatches are due to sentinel collision
-            hs_g_u32 = hs_g.view(torch.int32)
-            hs_ref_u32 = hs_ref.view(torch.int32)
-            diff_mask = hs_g_u32 != hs_ref_u32
-            sentinel_mask = (hs_ref_u32 == 0x80000000) & diff_mask
-            non_sentinel = diff_mask & ~sentinel_mask
-            if non_sentinel.any():
-                print(
-                    f"[{rank}] trial={trial} hs NON-SENTINEL MISMATCH: {non_sentinel.sum().item()}"
-                )
-                errors += 1
-            elif sentinel_mask.any():
-                if trial < 3:
-                    print(
-                        f"[{rank}] trial={trial} hs sentinel collision: "
-                        f"{sentinel_mask.sum().item()} words (expected rare)"
-                    )
-        if not torch.equal(sc_g, sc_ref):
-            mismatches = (sc_g != sc_ref).sum().item()
-            sc_g_u32 = sc_g.view(torch.int32) if sc_g.numel() % 4 == 0 else None
-            if sc_g_u32 is not None:
-                sc_ref_u32 = sc_ref.view(torch.int32)
-                diff_mask = sc_g_u32 != sc_ref_u32
-                sentinel_mask = (sc_ref_u32 == 0x80000000) & diff_mask
-                non_sentinel = diff_mask & ~sentinel_mask
-                if non_sentinel.any():
-                    print(
-                        f"[{rank}] trial={trial} sc NON-SENTINEL MISMATCH: {non_sentinel.sum().item()}"
-                    )
-                    errors += 1
-
-        total_checks += 1
-
-    dist.barrier()
-    print(
-        f"[rank {rank}] {total_checks} trials, {errors} real errors, "
-        f"{sentinel_collisions} sentinel patterns in hs data. "
-        f"{'PASSED' if errors == 0 else 'FAILED'}"
-    )
-    dist.destroy_process_group()
-    return errors
-
-
-if __name__ == "__main__":
-    sys.exit(main())

csrc/standalone/custom_ar/test_moe_reduce_scatter.py

@@ -1,204 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Test + benchmark for Lamport MoE reduce-scatter."""
-
-import ctypes
-import os
-import sys
-
-import torch
-import torch.distributed as dist
-
-_cudart = ctypes.CDLL("libcudart.so")
-IPC = 64
-
-
-def _cc(r):
-    if r:
-        raise RuntimeError(f"CUDA err {r}")
-
-
-def ipc_buf(sz, rank, ws):
-    p = ctypes.c_void_p()
-    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
-    _cc(_cudart.cudaMemset(p, 0, sz))
-    _cc(_cudart.cudaDeviceSynchronize())
-    h = (ctypes.c_byte * IPC)()
-    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
-    ah = [None] * ws
-    dist.all_gather_object(ah, bytes(h))
-    ptrs = []
-    for i in range(ws):
-        if i == rank:
-            ptrs.append(p.value)
-        else:
-            hh = (ctypes.c_byte * IPC)(*ah[i])
-            pp = ctypes.c_void_p()
-            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
-            ptrs.append(pp.value)
-    return ptrs
-
-
-def gpu_timer_graph(fn, warmup=20, repeats=200):
-    for _ in range(warmup):
-        fn()
-    torch.cuda.synchronize()
-    g = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(g):
-        fn()
-    for _ in range(5):
-        g.replay()
-    torch.cuda.synchronize()
-    s = torch.cuda.Event(enable_timing=True)
-    e = torch.cuda.Event(enable_timing=True)
-    s.record()
-    for _ in range(repeats):
-        g.replay()
-    e.record()
-    torch.cuda.synchronize()
-    return s.elapsed_time(e) / repeats * 1000
-
-
-def main():
-    dist.init_process_group("nccl")
-    rank = dist.get_rank()
-    ws = dist.get_world_size()
-    torch.cuda.set_device(rank)
-    dev = f"cuda:{rank}"
-
-    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-    if rank == 0:
-        from moe_reduce_scatter import _load_lib
-
-        _load_lib()
-    dist.barrier()
-    if rank != 0:
-        from moe_reduce_scatter import _load_lib
-
-        _load_lib()
-    dist.barrier()
-    from moe_reduce_scatter import MoeReduceScatter
-
-    max_size = 8 * 1024 * 1024
-    bp = ipc_buf(max_size, rank, ws)
-    dist.barrier()
-
-    class FakeCA:
-        pass
-
-    ca = FakeCA()
-    ca.rank = rank
-    ca.world_size = ws
-    ca.device = torch.device(dev)
-    ca.buffer_ptrs = bp
-    ca.max_size = max_size
-    rs = MoeReduceScatter(ca)
-    dist.barrier()
-
-    D = 7168  # DeepSeek V3 hidden_dim
-    errors = 0
-
-    # ---- Correctness tests ----
-    for N_per_rank in [1, 4, 16]:
-        N_total = N_per_rank * ws
-
-        # Each rank gets a deterministic input.
-        torch.manual_seed(42)
-        # All ranks create the SAME "ground truth" inputs for each rank.
-        all_inputs = [
-            torch.randn(N_total, D, dtype=torch.bfloat16, device=dev) for _ in range(ws)
-        ]
-        # This rank's input is all_inputs[rank].
-        my_input = all_inputs[rank]
-
-        # Custom reduce-scatter.
-        custom_out = rs.reduce_scatter(my_input)
-
-        # NCCL reference: reduce_scatter_tensor.
-        nccl_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
-        dist.reduce_scatter_tensor(nccl_out, my_input)
-
-        # bf16 summation order differs between our kernel and NCCL,
-        # giving ~1-2 ULP differences. Use generous tolerance.
-        max_diff = (custom_out.float() - nccl_out.float()).abs().max().item()
-        if not torch.allclose(custom_out, nccl_out, atol=0.125, rtol=0.01):
-            mismatches = (
-                ((custom_out.float() - nccl_out.float()).abs() > 0.125).sum().item()
-            )
-            print(
-                f"[{rank}] N_per_rank={N_per_rank} MISMATCH: "
-                f"max_diff={max_diff:.6f}, mismatches={mismatches}"
-            )
-            errors += 1
-        else:
-            if rank == 0:
-                print(f"  N_per_rank={N_per_rank}: PASS (max_diff={max_diff:.6f})")
-
-    # ---- Benchmark ----
-    if rank == 0:
-        print(f"\nworld_size={ws}, max_per_rank={rs.max_per_rank} bytes")
-        print(
-            f"{'config':<12} {'lamport':>10} {'lamp_graph':>10} "
-            f"{'nccl':>10} {'nccl_graph':>10} {'speedup':>8}"
-        )
-        print("-" * 65)
-
-    configs = [
-        ("1tok", 1),
-        ("2tok", 2),
-        ("4tok", 4),
-        ("8tok", 8),
-        ("16tok", 16),
-        ("32tok", 32),
-        ("64tok", 64),
-        ("128tok", 128),
-        ("256tok", 256),
-    ]
-
-    for name, N_per_rank in configs:
-        N_total = N_per_rank * ws
-        input_bytes = N_total * D * 2  # bf16
-        if input_bytes > rs.max_per_rank:
-            if rank == 0:
-                print(f"{name:<12} {'skip (too large)':>40}")
-            continue
-
-        inp = torch.randn(N_total, D, dtype=torch.bfloat16, device=dev)
-        c_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
-        n_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
-
-        def run_lamport():
-            rs.reduce_scatter(inp)
-
-        def run_nccl():
-            dist.reduce_scatter_tensor(n_out, inp)
-
-        from bench_moe_allgather import gpu_timer
-
-        lam_us = gpu_timer(run_lamport)
-        try:
-            lam_g_us = gpu_timer_graph(run_lamport)
-        except Exception:
-            lam_g_us = float("nan")
-
-        nccl_us = gpu_timer(run_nccl)
-        try:
-            nccl_g_us = gpu_timer_graph(run_nccl)
-        except Exception:
-            nccl_g_us = float("nan")
-
-        if rank == 0:
-            speedup = nccl_g_us / lam_g_us if lam_g_us > 0 else float("nan")
-            print(
-                f"{name:<12} {lam_us:>9.1f}µ {lam_g_us:>9.1f}µ "
-                f"{nccl_us:>9.1f}µ {nccl_g_us:>9.1f}µ {speedup:>7.2f}x"
-            )
-
-    dist.barrier()
-    print(f"[rank {rank}] {'PASSED' if errors == 0 else f'FAILED ({errors})'}")
-    dist.destroy_process_group()
-    return errors
-
-
-if __name__ == "__main__":
-    sys.exit(main())

csrc/standalone/custom_ar/test_rs_fused.py

@@ -1,270 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Test + benchmark: fused RS + residual + RMSNorm vs separate kernels."""
-
-import ctypes
-import os
-import sys
-
-import torch
-import torch.distributed as dist
-
-_cudart = ctypes.CDLL("libcudart.so")
-IPC = 64
-
-
-def _cc(r):
-    if r:
-        raise RuntimeError(f"CUDA err {r}")
-
-
-def ipc_buf(sz, rank, ws):
-    p = ctypes.c_void_p()
-    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
-    _cc(_cudart.cudaMemset(p, 0, sz))
-    _cc(_cudart.cudaDeviceSynchronize())
-    h = (ctypes.c_byte * IPC)()
-    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
-    ah = [None] * ws
-    dist.all_gather_object(ah, bytes(h))
-    ptrs = []
-    for i in range(ws):
-        if i == rank:
-            ptrs.append(p.value)
-        else:
-            hh = (ctypes.c_byte * IPC)(*ah[i])
-            pp = ctypes.c_void_p()
-            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
-            ptrs.append(pp.value)
-    return ptrs
-
-
-def rms_norm_ref(x, gamma, eps):
-    """Reference RMSNorm in fp32."""
-    xf = x.float()
-    rms = torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)
-    return (xf * rms * gamma.float()).to(x.dtype)
-
-
-def gpu_timer(fn, warmup=20, repeats=200):
-    for _ in range(warmup):
-        fn()
-    torch.cuda.synchronize()
-    s = torch.cuda.Event(enable_timing=True)
-    e = torch.cuda.Event(enable_timing=True)
-    s.record()
-    for _ in range(repeats):
-        fn()
-    e.record()
-    torch.cuda.synchronize()
-    return s.elapsed_time(e) / repeats * 1000
-
-
-def gpu_timer_graph(fn, warmup=20, repeats=200):
-    for _ in range(warmup):
-        fn()
-    torch.cuda.synchronize()
-    g = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(g):
-        fn()
-    for _ in range(5):
-        g.replay()
-    torch.cuda.synchronize()
-    s = torch.cuda.Event(enable_timing=True)
-    e = torch.cuda.Event(enable_timing=True)
-    s.record()
-    for _ in range(repeats):
-        g.replay()
-    e.record()
-    torch.cuda.synchronize()
-    return s.elapsed_time(e) / repeats * 1000
-
-
-def main():
-    dist.init_process_group("nccl")
-    rank = dist.get_rank()
-    ws = dist.get_world_size()
-    torch.cuda.set_device(rank)
-    dev = f"cuda:{rank}"
-
-    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-
-    # Compile fused kernel
-    if rank == 0:
-        from torch.utils.cpp_extension import load
-
-        load(
-            name="moe_rs_fused_kernel",
-            sources=[os.path.join(os.path.dirname(__file__), "moe_rs_fused.cu")],
-            extra_cuda_cflags=["-O3", "--use_fast_math"],
-            verbose=False,
-        )
-    dist.barrier()
-    from torch.utils.cpp_extension import load
-
-    fused_lib = load(
-        name="moe_rs_fused_kernel",
-        sources=[os.path.join(os.path.dirname(__file__), "moe_rs_fused.cu")],
-        extra_cuda_cflags=["-O3", "--use_fast_math"],
-        verbose=False,
-    )
-
-    # Also compile separate RS kernel for comparison
-    from moe_reduce_scatter import MoeReduceScatter
-
-    max_size = 8 * 1024 * 1024
-    bp = ipc_buf(max_size, rank, ws)
-    dist.barrier()
-
-    class FakeCA:
-        pass
-
-    ca = FakeCA()
-    ca.rank = rank
-    ca.world_size = ws
-    ca.device = torch.device(dev)
-    ca.buffer_ptrs = bp
-    ca.max_size = max_size
-
-    rs_separate = MoeReduceScatter(ca)
-
-    # Fused kernel setup (uses same buffer layout as MoeReduceScatter)
-    half_size = (max_size // 2) & ~15
-    buf_offset = half_size
-    fused_seg_cap = (half_size // 2) & ~15
-    fused_rank_stride = (fused_seg_cap // ws) & ~15
-
-    fused_buf_ptrs = torch.zeros(8, dtype=torch.int64, device=dev)
-    for i in range(ws):
-        fused_buf_ptrs[i] = bp[i] + buf_offset
-    fused_counters = torch.zeros(3, dtype=torch.int32, device=dev)
-
-    # Init sentinels for fused kernel's buffer region
-    fused_lib.lamport_init(bp[rank] + buf_offset, half_size)
-    torch.cuda.synchronize()
-    dist.barrier()
-
-    D = 7168
-    eps = 1e-6
-    gamma = torch.randn(D, dtype=torch.bfloat16, device=dev).abs() + 0.5
-    errors = 0
-
-    # ---- Correctness ----
-    for N_per_rank in [1, 4]:
-        N_total = N_per_rank * ws
-        torch.manual_seed(42 + rank)
-        moe_out = torch.randn(N_total, D, dtype=torch.bfloat16, device=dev)
-        residual = torch.randn(N_per_rank, D, dtype=torch.bfloat16, device=dev)
-
-        # Reference: NCCL RS + add + norm
-        rs_ref = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
-        dist.reduce_scatter_tensor(rs_ref, moe_out)
-        ref_residual = residual + rs_ref
-        ref_normed = rms_norm_ref(ref_residual, gamma, eps)
-
-        # Fused kernel
-        normed_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
-        residual_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
-        fused_lib.moe_rs_fused(
-            fused_buf_ptrs.data_ptr(),
-            fused_counters.data_ptr(),
-            rank,
-            ws,
-            fused_seg_cap,
-            fused_rank_stride,
-            moe_out,
-            residual,
-            gamma,
-            normed_out,
-            residual_out,
-            eps,
-        )
-        torch.cuda.synchronize()
-
-        # Compare
-        max_diff_res = (residual_out.float() - ref_residual.float()).abs().max().item()
-        max_diff_norm = (normed_out.float() - ref_normed.float()).abs().max().item()
-
-        ok = max_diff_res < 0.125 and max_diff_norm < 0.125
-        if rank == 0:
-            print(
-                f"  N_per_rank={N_per_rank}: {'PASS' if ok else 'FAIL'} "
-                f"(res_diff={max_diff_res:.4f}, norm_diff={max_diff_norm:.4f})"
-            )
-        if not ok:
-            errors += 1
-
-    # ---- Benchmark ----
-    if rank == 0:
-        print(f"\nBenchmark: D={D}, world_size={ws}")
-        print(
-            f"{'config':<10} {'fused':>10} {'fused_g':>10} "
-            f"{'RS+norm':>10} {'RS+norm_g':>10} {'speedup':>8}"
-        )
-        print("-" * 58)
-
-    for N_per_rank in [1, 2, 4, 8]:
-        N_total = N_per_rank * ws
-        input_bytes = N_total * D * 2
-        if input_bytes > fused_rank_stride:
-            if rank == 0:
-                print(f"{N_per_rank}tok      skip (too large)")
-            continue
-
-        moe_out = torch.randn(N_total, D, dtype=torch.bfloat16, device=dev)
-        residual = torch.randn(N_per_rank, D, dtype=torch.bfloat16, device=dev)
-        normed_out = torch.empty_like(residual)
-        residual_out = torch.empty_like(residual)
-        rs_out = torch.empty_like(residual)
-
-        # Fused
-        def run_fused():
-            fused_lib.moe_rs_fused(
-                fused_buf_ptrs.data_ptr(),
-                fused_counters.data_ptr(),
-                rank,
-                ws,
-                fused_seg_cap,
-                fused_rank_stride,
-                moe_out,
-                residual,
-                gamma,
-                normed_out,
-                residual_out,
-                eps,
-            )
-
-        # Separate: RS + add + norm (our Lamport RS + triton-like ops)
-        def run_separate():
-            rs_separate.reduce_scatter(moe_out)
-            # Simulate add + RMSNorm (in practice this is a fused triton kernel)
-            tmp = residual + rs_out
-            torch.rsqrt(tmp.float().pow(2).mean(-1, keepdim=True) + eps)
-
-        fused_us = gpu_timer(run_fused)
-        try:
-            fused_g = gpu_timer_graph(run_fused)
-        except Exception:
-            fused_g = float("nan")
-
-        sep_us = gpu_timer(run_separate)
-        try:
-            sep_g = gpu_timer_graph(run_separate)
-        except Exception:
-            sep_g = float("nan")
-
-        if rank == 0:
-            speedup = sep_g / fused_g if fused_g > 0 else float("nan")
-            print(
-                f"{N_per_rank}tok     {fused_us:>9.1f}µ {fused_g:>9.1f}µ "
-                f"{sep_us:>9.1f}µ {sep_g:>9.1f}µ {speedup:>7.2f}x"
-            )
-
-    dist.barrier()
-    print(f"[rank {rank}] {'PASSED' if errors == 0 else f'FAILED ({errors})'}")
-    dist.destroy_process_group()
-    return errors
-
-
-if __name__ == "__main__":
-    sys.exit(main())

docker/Dockerfile

@@ -642,7 +642,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     else \
         BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_X86}"; \
     fi; \
-    uv pip install --system accelerate hf_transfer modelscope \
+    uv pip install --system accelerate modelscope \
         "bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}"
 
 # ============================================================
@@ -756,9 +756,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system -e tests/vllm_test_utils
 
 # enable fast downloads from hf (for testing)
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system hf_transfer
-ENV HF_HUB_ENABLE_HF_TRANSFER 1
+ENV HF_XET_HIGH_PERFORMANCE 1
+
+# increase timeout for hf downloads (for testing)
+ENV HF_HUB_DOWNLOAD_TIMEOUT 60
 
 # Copy in the v1 package for testing (it isn't distributed yet)
 COPY vllm/v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1

docker/Dockerfile.cpu

@@ -197,6 +197,12 @@ ADD ./.buildkite/ ./.buildkite/
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install -e tests/vllm_test_utils
 
+# enable fast downloads from hf (for testing)
+ENV HF_XET_HIGH_PERFORMANCE 1
+
+# increase timeout for hf downloads (for testing)
+ENV HF_HUB_DOWNLOAD_TIMEOUT 60
+
 ######################### RELEASE IMAGE #########################
 FROM base AS vllm-openai
 

docker/Dockerfile.nightly_torch

@@ -272,9 +272,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system -e tests/vllm_test_utils
 
 # enable fast downloads from hf (for testing)
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system hf_transfer
-ENV HF_HUB_ENABLE_HF_TRANSFER 1
+ENV HF_XET_HIGH_PERFORMANCE 1
+
+# increase timeout for hf downloads (for testing)
+ENV HF_HUB_DOWNLOAD_TIMEOUT 60
 
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system -r requirements/test/nightly-torch.txt

docker/Dockerfile.rocm

@@ -365,9 +365,10 @@ RUN cd /vllm-workspace \
     && python3 -m pip install pytest-shard
 
 # enable fast downloads from hf (for testing)
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system hf_transfer
-ENV HF_HUB_ENABLE_HF_TRANSFER=1
+ENV HF_XET_HIGH_PERFORMANCE=1
+
+# increase timeout for hf downloads (for testing)
+ENV HF_HUB_DOWNLOAD_TIMEOUT 60
 
 # install audio decode package `torchcodec` from source (required due to 
 # ROCm and torch version mismatch) for tests with datasets package

docker/Dockerfile.s390x

@@ -42,7 +42,7 @@ FROM python-install AS pyarrow
 # Build Apache Arrow
 WORKDIR /tmp
 RUN --mount=type=cache,target=/root/.cache/uv \
-    git clone https://github.com/apache/arrow.git && \
+    git clone https://github.com/apache/arrow.git  -b maint-19.0.1 && \
     cd arrow/cpp && \
     mkdir release && cd release && \
     cmake -DCMAKE_BUILD_TYPE=Release \
@@ -68,19 +68,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install -r requirements-build.txt && \
     python setup.py build_ext --build-type=$ARROW_BUILD_TYPE --bundle-arrow-cpp bdist_wheel
 
-FROM python-install AS numa-build
-# Install numactl (needed for numa.h dependency)
-WORKDIR /tmp
-RUN curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.16.tar.gz && \
-    tar -xvzf v2.0.16.tar.gz && \
-    cd numactl-2.0.16 && \
-    ./autogen.sh && \
-    ./configure && \
-    make
-
-# Set include path
-ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH"
-
 FROM python-install AS rust
 ENV CARGO_HOME=/root/.cargo
 ENV RUSTUP_HOME=/root/.rustup
@@ -91,6 +78,18 @@ RUN curl https://sh.rustup.rs -sSf | sh -s -- -y && \
     rustup default stable && \
     rustup show
 
+FROM python-install AS numa-build
+WORKDIR /tmp
+RUN curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.19.tar.gz && \
+    tar -xvzf v2.0.19.tar.gz && \
+    cd numactl-2.0.19 && \
+    ./autogen.sh && \
+    ./configure && \
+    make
+
+# Set include path
+ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH"
+
 FROM python-install AS torch-vision
 # Install torchvision
 ARG TORCH_VISION_VERSION=v0.26.0
@@ -133,7 +132,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     git clone --recursive https://github.com/numba/llvmlite.git -b v0.44.0 && \
     git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \
     cd llvm-project && mkdir build && cd  build && \
-    uv pip install 'cmake<4' setuptools numpy && \
+    uv pip install 'cmake<4' 'setuptools<70' numpy && \
     export PREFIX=/usr/local && CMAKE_ARGS="${CMAKE_ARGS} -DLLVM_ENABLE_PROJECTS=lld;libunwind;compiler-rt" \
     CFLAGS="$(echo $CFLAGS | sed 's/-fno-plt //g')" \
     CXXFLAGS="$(echo $CXXFLAGS | sed 's/-fno-plt //g')" \
@@ -193,27 +192,22 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     cd opencv-python && \
     python -m build --wheel --installer=uv --outdir /tmp/opencv-python/dist
 
-# Build Outlines Core
-FROM python-install AS outlines-core-builder
+## Todo(r3hankhan123): Remove guidance-builder stage once vLLM upgrades to new version of llguidance that fixes s390x issues. See https://github.com/guidance-ai/llguidance/issues/330
+FROM python-install AS guidance-builder
 WORKDIR /tmp
 ENV CARGO_HOME=/root/.cargo
 ENV RUSTUP_HOME=/root/.rustup
 ENV PATH="$CARGO_HOME/bin:$RUSTUP_HOME/bin:$PATH"
-COPY requirements/common.txt /tmp/requirements/common.txt
-ARG OUTLINES_CORE_VERSION
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,from=rust,source=/root/.cargo,target=/root/.cargo,rw \
     --mount=type=bind,from=rust,source=/root/.rustup,target=/root/.rustup,rw \
-    OUTLINES_CORE_VERSION=${OUTLINES_CORE_VERSION:-$(grep -E '^outlines_core\s*==\s*[0-9.]+' /tmp/requirements/common.txt | grep -Eo '[0-9.]+')} && \
-    if [ -z "${OUTLINES_CORE_VERSION}" ]; then echo "ERROR: Could not determine outlines_core version"; exit 1; fi && \
-    git clone https://github.com/dottxt-ai/outlines-core.git && \
-    cd outlines-core && \
-    git checkout tags/${OUTLINES_CORE_VERSION} && \
-    sed -i "s/version = \"0.0.0\"/version = \"${OUTLINES_CORE_VERSION}\"/" Cargo.toml && \
+    git clone https://github.com/guidance-ai/llguidance.git && \
+    cd llguidance && \
+    git checkout s390x-fix-v2  && \
     uv pip install maturin && \
-    python -m maturin build --release --out dist
+    python -m maturin build --release --out dist --compatibility linux
 
-# Final build stage
+# # Final build stage
 FROM python-install AS vllm-cpu
 ARG PYTHON_VERSION
 ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"
@@ -229,10 +223,12 @@ ENV PKG_CONFIG_PATH="/opt/rh/gcc-toolset-14/root/usr/lib64/pkgconfig:/usr/local/
 ENV PATH="${VIRTUAL_ENV:+${VIRTUAL_ENV}/bin}:/opt/rh/gcc-toolset-14/root/usr/bin:/usr/local/bin:$CARGO_HOME/bin:$RUSTUP_HOME/bin:$PATH"
 ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
+# Force pure Python protobuf to avoid s390x C++ extension crashes
+ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
 COPY . /workspace/vllm
 WORKDIR /workspace/vllm
 
-RUN --mount=type=bind,from=numa-build,src=/tmp/numactl-2.0.16,target=/numactl \
+RUN --mount=type=bind,from=numa-build,src=/tmp/numactl-2.0.19,target=/numactl \
     make -C /numactl install
 
 # Install dependencies, including PyTorch and Apache Arrow
@@ -245,22 +241,22 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,from=numba-builder,source=/tmp/llvmlite/dist,target=/tmp/llvmlite-wheels/ \
     --mount=type=bind,from=numba-builder,source=/tmp/numba/dist,target=/tmp/numba-wheels/ \
     --mount=type=bind,from=opencv-builder,source=/tmp/opencv-python/dist,target=/tmp/opencv-wheels/ \
-    --mount=type=bind,from=outlines-core-builder,source=/tmp/outlines-core/dist,target=/tmp/outlines-core/dist/ \
-     ARROW_WHL_FILE=$(ls /tmp/arrow-wheels/pyarrow-*.whl) && \
+    --mount=type=bind,from=guidance-builder,source=/tmp/llguidance/dist,target=/tmp/guidance-wheels/ \
+     ARROW_WHL_FILE=$(ls /tmp/arrow-wheels/*.whl) && \
      VISION_WHL_FILE=$(ls /tmp/vision-wheels/*.whl) && \
      HF_XET_WHL_FILE=$(ls /tmp/hf-xet-wheels/*.whl) && \
      LLVM_WHL_FILE=$(ls /tmp/llvmlite-wheels/*.whl) && \
      NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \
      OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \
-     OUTLINES_CORE_WHL_FILE=$(ls /tmp/outlines-core/dist/*.whl) && \
-     uv pip install -v \
-        $ARROW_WHL_FILE  \
+     GUIDANCE_WHL_FILE=$(ls /tmp/guidance-wheels/*.whl) && \
+     uv pip install -v \    
+        $ARROW_WHL_FILE \
         $VISION_WHL_FILE \
         $HF_XET_WHL_FILE \
         $LLVM_WHL_FILE \
         $NUMBA_WHL_FILE \
         $OPENCV_WHL_FILE \
-        $OUTLINES_CORE_WHL_FILE \
+        $GUIDANCE_WHL_FILE \
         --index-strategy unsafe-best-match \
         -r requirements/build/cpu.txt \
         -r requirements/cpu.txt
@@ -271,6 +267,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \
     uv pip install "$(echo dist/*.whl)[tensorizer]"
 
+# Remove protobuf C++ extension that crashes on s390x
+RUN rm -rf /opt/vllm/lib64/python${PYTHON_VERSION}/site-packages/google/_upb/*.so \
+           /opt/vllm/lib64/python${PYTHON_VERSION}/site-packages/google/protobuf/pyext/*.so 2>/dev/null || true
+
 # setup non-root user for vllm
 RUN umask 002 && \
     /usr/sbin/useradd --uid 2000 --gid 0 vllm && \

docs/benchmarking/cli.md

@@ -37,6 +37,7 @@ th {
 | 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` |
+| SPEED-Bench | ✅ | ✅ | `curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -` |
 | Custom | ✅ | ✅ | Local file: `data.jsonl` |
 | Custom MM | ✅ | ✅ | Local file: `mm_data.jsonl` |
 
@@ -239,6 +240,69 @@ vllm bench serve \
     --spec-bench-category "summarization"
 ```
 
+#### SPEED-Bench Benchmark with Speculative Decoding
+
+[SPEED-Bench](https://huggingface.co/datasets/nvidia/SPEED-Bench) is a unified and diverse dataset for speculative decoding, supporting acceptance rate and length measurements using the Qualitative split and throughput measurements using the Throughput splits in 5 configuration of input sequence length (1k, 2k, 8k, 16k, 32k).
+
+!!! note
+    This dataset is governed by the [NVIDIA Evaluation Dataset License Agreement](https://huggingface.co/datasets/nvidia/SPEED-Bench/blob/main/License.pdf). For each dataset a user elects to use, the user is responsible for checking if the dataset license is fit for the intended purpose. The `prepare.py` script automatically fetches data from all the source datasets.
+
+First, download the dataset to a folder, using this one liner:
+
+```bash
+curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -
+```
+
+The command supports also the following arguments:
+
+- `--config`: download only a subset of the dataset: `qualitative`, `throughput_1k`, `throughput_2k`, `throughput_8k`, `throughput_16k` and `throughput_32k`. By default, it will download all subsets.
+- `--output_dir`: download to a specified folder. By default, it will download to the current directory.
+
+Start a server with speculative decoding:
+
+```bash
+vllm serve meta-llama/Llama-3.3-70B-Instruct \
+    --speculative-config $'{"method": "eagle3",
+    "num_speculative_tokens": 3,
+    "model": "nvidia/Llama-3.3-70B-Instruct-Eagle3"}'
+```
+
+Run all categories in the Qualitative split:
+
+```bash
+vllm bench serve \
+    --model meta-llama/Llama-3.3-70B-Instruct \
+    --dataset-name speed_bench \
+    --dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench" \
+    --num-prompts -1
+```
+
+Available categories include `[writing, roleplay, reasoning, math, coding, stem, humanities, multilingual, summarization, qa, rag]`.
+
+Run only a specific category like "multilingual":
+
+```bash
+vllm bench serve \
+    --model meta-llama/Llama-3.3-70B-Instruct \
+    --dataset-name speed_bench \
+    --dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench" \
+    --num-prompts -1
+    --speed-bench-category "multilingual"
+```
+
+Run all categories in the Throughput split (2k ISL):
+
+```bash
+vllm bench serve \
+    --model meta-llama/Llama-3.3-70B-Instruct \
+    --dataset-name speed_bench \
+    --speed-bench-dataset-subset throughput_2k
+    --dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench/" \
+    --num-prompts -1
+```
+
+Available categories include `[high_entropy, mixed, low_entropy]`, where high entropy data contains unstructued data such as creative writing while low entropy data contains more structured data such as coding, more details are in the dataset card.
+
 #### Other HuggingFaceDataset Examples
 
 ```bash

docs/design/attention_backends.md

@@ -178,6 +178,7 @@ Priority is **1 = highest** (tried first).
 | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
 | `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
 | `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
+| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | Decoder | Any |
 
 > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
 >

docs/design/cuda_graphs_multimodal.md

@@ -28,6 +28,7 @@ Multiple CUDA Graphs are pre-captured at different **token budget** levels (e.g.
 class BudgetGraphMetadata:
     token_budget: int
     max_batch_size: int
+    max_frames_per_batch: int
     graph: torch.cuda.CUDAGraph
     input_buffer: torch.Tensor       # e.g. pixel_values
     metadata_buffers: dict[str, torch.Tensor]  # e.g. embeddings, seq metadata
@@ -51,6 +52,15 @@ For each graph replay:
 
 When `mm_encoder_tp_mode="data"`, the manager distributes images across TP ranks using load-balanced assignment via `get_load_balance_assignment`, executes locally on each rank, then gathers results back in the original order via `tensor_model_parallel_all_gather`.
 
+### Video inference support (experimental)
+
+Following <https://github.com/vllm-project/vllm/pull/35963> (ViT full CUDA graph support for image inference), <https://github.com/vllm-project/vllm/pull/38061> extends the encoder CUDA graph framework to support video inference for Qwen3-VL. Previously, the CUDA graph capture/replay path only handled image inputs (`pixel_values` + `image_grid_thw`). Video inputs use different keys (`pixel_values_videos` + `video_grid_thw`) and require larger `cu_seqlens` buffers because each video item contributes multiple frames (`T` attention sequences). This PR generalizes the protocol and manager to handle both modalities through a single shared graph manager.
+
+!!! note
+    Video CUDA graphs are automatically disabled when EVS (Efficient Video Sampling) pruning is enabled, since EVS makes the token count data-dependent and incompatible with CUDA graph capture.
+
+    Currently, we only support image-only or video-only inputs when enabling CUDA graph, mixed inputs (image + video) are not supported yet (we will work on it in the near future). Thus, it's recommended to turn off the image modality by `--limit-mm-per-prompt '{"image": 0}'` for video-only inputs.
+
 ## Model integration via `SupportsEncoderCudaGraph`
 
 Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGraph][vllm.model_executor.models.interfaces.SupportsEncoderCudaGraph] protocol. This protocol encapsulates all model-specific logic so that the manager remains model-agnostic. The protocol defines the following methods:
@@ -65,12 +75,17 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
 * `prepare_encoder_cudagraph_replay_buffers(...)` — computes new buffer values from actual batch inputs before replay.
 * `encoder_cudagraph_forward(...)` — forward pass using precomputed buffers (called during capture and replay).
 * `encoder_eager_forward(...)` — fallback eager forward when no graph fits.
-
-Currently supported: **Qwen3-VL** (see `vllm/model_executor/models/qwen3_vl.py`).
+* `get_input_modality(...)` - return the modality of the inputs.
 
 !!! note
     The `SupportsEncoderCudaGraph` protocol is designed to be model-agnostic. New vision encoder models can opt-in by implementing the protocol methods without modifying the manager.
 
+**Supported models:**
+
+| Architecture | Models | CG for Image | CG for Video |
+| ------------ | ------ | ------------ | ------------ |
+| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
+
 !!! note
     Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
 
@@ -80,10 +95,13 @@ Three fields in `CompilationConfig` control encoder CUDA Graphs:
 
 * `cudagraph_mm_encoder` (`bool`, default `False`) — enable CUDA Graph capture for multimodal encoder. When enabled, captures the full encoder forward as a CUDA Graph for each token budget level.
 * `encoder_cudagraph_token_budgets` (`list[int]`, default `[]`) — token budget levels for capture. If empty (default), auto-inferred from model architecture as power-of-2 levels. User-provided values override auto-inference.
-* `encoder_cudagraph_max_images_per_batch` (`int`, default `0`) — maximum number of images per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`.
+* `encoder_cudagraph_max_vision_items_per_batch` (`int`, default `0`) — maximum number of images/videos per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`.
+* `encoder_cudagraph_max_frames_per_batch` (`int`, default `0`) — maximum number of video frames per batch during capture. If 0 (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * 2` (to be optimized).
 
 ## Usage guide
 
+### Image inference
+
 Enable encoder CUDA Graphs via `compilation_config`:
 
 ```bash
@@ -95,7 +113,7 @@ With explicit budgets:
 
 ```bash
 vllm serve Qwen/Qwen3-VL-32B \
-  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_images_per_batch": 8}'
+  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_vision_items_per_batch": 8}'
 ```
 
 Python example:
@@ -107,7 +125,7 @@ compilation_config = {
     "cudagraph_mm_encoder": True,
     # Optional: override auto-inferred budgets
     # "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824],
-    # "encoder_cudagraph_max_images_per_batch": 8,
+    # "encoder_cudagraph_max_vision_items_per_batch": 8,
 }
 
 model = vllm.LLM(
@@ -118,6 +136,44 @@ model = vllm.LLM(
 
 The manager tracks hit/miss statistics and logs them periodically. A "hit" means an image was processed via CUDA Graph replay; a "miss" means eager fallback (image exceeded all budgets).
 
+### Video inference
+
+Enable encoder CUDA Graphs via `compilation_config`:
+
+```bash
+vllm serve Qwen/Qwen3-VL-32B \
+  --limit-mm-per-prompt '{"image": 0}' \
+  --compilation-config '{"cudagraph_mm_encoder": true}'
+```
+
+With explicit budgets:
+
+```bash
+vllm serve Qwen/Qwen3-VL-32B \
+  --limit-mm-per-prompt '{"image": 0}' \
+  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_vision_items_per_batch": 8, "encoder_cudagraph_max_frames_per_batch": 64}'
+```
+
+Python example:
+
+```python
+import vllm
+
+compilation_config = {
+    "cudagraph_mm_encoder": True,
+    # Optional: override auto-inferred budgets
+    # "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824],
+    # "encoder_cudagraph_max_vision_items_per_batch": 8,
+    # "encoder_cudagraph_max_frames_per_batch": 64,
+}
+
+model = vllm.LLM(
+    model="Qwen/Qwen3-VL-32B",
+    limit_mm_per_prompt='{"image": 0}',
+    compilation_config=compilation_config,
+)
+```
+
 ## About the Performance
 
 The following benchmarks were run on Blackwell GPUs (GB200) using `vllm bench mm-processor`. See [#35963](https://github.com/vllm-project/vllm/pull/35963) for full details.
@@ -140,7 +196,7 @@ vllm bench mm-processor \
   --num-prompts 3000 --num-warmups 300 \
   --max-model-len 32768 --seed 42 \
   --mm-encoder-attn-backend FLASH_ATTN \
-  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_images_per_batch": 8}'
+  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_vision_items_per_batch": 8}'
 ```
 
 ### Multi-GPU (4x GB200, TP=4, DP=4)
@@ -165,5 +221,8 @@ vllm bench mm-processor \
   --max-model-len 8192 --seed 42 \
   --mm-encoder-attn-backend FLASHINFER \
   --tensor-parallel-size 4 --mm-encoder-tp-mode data \
-  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_images_per_batch": 8}'
+  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_vision_items_per_batch": 8}'
 ```
+
+!!! note
+    Find more details about benchmarks on GPUs (A100) for video inference at [#38061](https://github.com/vllm-project/vllm/pull/38061).

docs/design/moe_kernel_features.md

@@ -86,7 +86,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
 | cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp4] |
 | cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] |
 | flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
-| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts] |
+| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.OAITritonExperts] |
 | marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
 | trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],</br>[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],</br>[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],</br>[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] |
 | rocm aiter moe | standard | mxfp4,</br>fp8 | G(32),G(128),A,T | silu, gelu,</br>swigluoai | Y | N | `rocm_aiter_fused_experts`,</br>`AiterExperts` |

docs/features/quantization/README.md

@@ -16,6 +16,7 @@ The following are the supported quantization formats for vLLM:
 - [INT8 W8A8](int8.md)
 - [FP8 W8A8](fp8.md)
 - [NVIDIA Model Optimizer](modelopt.md)
+- [Online Quantization](online.md)
 - [AMD Quark](quark.md)
 - [Quantized KV Cache](quantized_kvcache.md)
 - [TorchAO](torchao.md)

docs/features/quantization/online.md

@@ -0,0 +1,94 @@
+# Online Quantization
+
+Online quantization lets you take a BF16/FP16 model and quantize its Linear
+and MoE weights to lower precision (such as FP8) at load time, without needing
+a pre-quantized checkpoint or calibration data. Weights are converted during
+model loading and activations are dynamically scaled during each forward pass.
+
+## Quick Start
+
+Pass a scheme name to the `quantization` parameter:
+
+```python
+from vllm import LLM
+
+# Per-tensor FP8 quantization (one scale per weight tensor)
+llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_tensor")
+
+# Per-block FP8 quantization (128x128 block scaling for weights and 1x128 block scaling for activations)
+llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_block")
+```
+
+Or with the CLI:
+
+```bash
+vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_tensor
+vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_block
+```
+
+## Supported Schemes
+
+| Scheme | Weight recipe | Activation recipe | Notes |
+| ------ | ------------- | ------------------ | ----- |
+| `fp8_per_tensor` | fp8_e4m3 data, fp32 per-tensor scale | fp8_e4m3 data, fp32 per-tensor scale | On some GPUs (Ada, Hopper) linear activations use per-token scaling for better performance |
+| `fp8_per_block` | fp8_e4m3 data, fp32 per-128x128-block scale | fp8_e4m3 data, fp32 per-1x128-block scale | |
+
+Support for additional schemes will be added in future versions of vllm.
+
+## Advanced Configuration
+
+For fine-grained control, use a `quantization_config` dictionary.
+
+### Separate Schemes for Dense and MoE Layers
+
+You can apply different quantization schemes to dense linear layers and MoE expert layers:
+
+```python
+from vllm import LLM
+
+llm = LLM(
+    "ibm-granite/granite-3.0-1b-a400m-base",
+    quantization="fp8_per_tensor",
+    quantization_config={
+        "linear_scheme_override": "fp8_per_block",
+    },
+)
+```
+
+Or,
+
+```python
+from vllm import LLM
+
+llm = LLM(
+    "ibm-granite/granite-3.0-1b-a400m-base",
+    quantization="fp8_per_tensor",
+    quantization_config={
+        "moe_scheme_override": "fp8_per_block",
+    },
+)
+```
+
+### Excluding Layers from Quantization
+
+Use the `ignore` parameter to skip specific layers. It accepts exact layer names and regex patterns (prefixed with `re:`):
+
+```python
+from vllm import LLM
+
+llm = LLM(
+    "ibm-granite/granite-3.0-1b-a400m-base",
+    quantization="fp8_per_tensor",
+    quantization_config={
+        "ignore": [
+            # exact layer name
+            "model.layers.1.self_attn.o_proj",
+            # regex: skip all QKV projections
+            "re:.*[qkv]_proj",
+        ],
+    },
+)
+```
+
+!!! note
+    For fused layers (e.g., `qkv_proj` which fuses `q_proj`, `k_proj`, `v_proj`), the ignore pattern must match the **unfused** shard names (`q_proj`, `k_proj`, `v_proj`), not the fused name.

docs/getting_started/installation/cpu.s390x.inc.md

@@ -3,15 +3,15 @@
 
 vLLM has experimental support for s390x architecture on IBM Z platform. For now, users must build from source to natively run on IBM Z platform.
 
-Currently, the CPU implementation for s390x architecture supports FP32 datatype only.
+Currently, the CPU implementation for s390x architecture supports FP32, BF16 and FP16.
 
 --8<-- [end:installation]
 --8<-- [start:requirements]
 
 - OS: `Linux`
-- SDK: `gcc/g++ >= 12.3.0` or later with Command Line Tools
+- SDK: `gcc/g++ >= 14.0.0` or later with Command Line Tools
 - Instruction Set Architecture (ISA): VXE support is required. Works with Z14 and above.
-- Build install python packages: `pyarrow`, `torch` and `torchvision`
+- Build install python packages: `torchvision`, `llvmlite`, `numba`, `pyarrow (for testing)`, `opencv-headless`
 
 --8<-- [end:requirements]
 --8<-- [start:set-up-using-python]
@@ -24,13 +24,14 @@ Currently, there are no pre-built IBM Z CPU wheels.
 --8<-- [end:pre-built-wheels]
 --8<-- [start:build-wheel-from-source]
 
-Install the following packages from the package manager before building the vLLM. For example on RHEL 9.4:
+Install the following packages from the package manager before building the vLLM. For example on RHEL 9.6:
 
 ```bash
 dnf install -y \
-    which procps findutils tar vim git gcc g++ make patch make cython zlib-devel \
+    which procps findutils tar vim git gcc-toolset-14 gcc-toolset-14-binutils gcc-toolset-14-libatomic-devel zlib-devel \
     libjpeg-turbo-devel libtiff-devel libpng-devel libwebp-devel freetype-devel harfbuzz-devel \
-    openssl-devel openblas openblas-devel wget autoconf automake libtool cmake numactl-devel
+    openssl-devel openblas openblas-devel autoconf automake libtool cmake numpy libsndfile \
+    clang llvm-devel llvm-static clang-devel
 ```
 
 Install rust>=1.80 which is needed for `outlines-core` and `uvloop` python packages installation.
@@ -43,13 +44,13 @@ curl https://sh.rustup.rs -sSf | sh -s -- -y && \
 Execute the following commands to build and install vLLM from source.
 
 !!! tip
-    Please build the following dependencies, `torchvision`, `pyarrow` from source before building vLLM.
+    Please build the following dependencies, `torchvision`, `llvmlite`, `numba`, `llguidance`, `pyarrow`, `opencv-headless` from source before building vLLM.
 
 ```bash
-    sed -i '/^torch/d' requirements/build/cuda.txt    # remove torch from requirements/build/cuda.txt since we use nightly builds
     uv pip install -v \
+        --extra-index-url https://download.pytorch.org/whl/cpu \
         --torch-backend auto \
-        -r requirements/build/cuda.txt \
+        -r requirements/build/cpu.txt \
         -r requirements/cpu.txt \
     VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
         uv pip install dist/*.whl
@@ -57,10 +58,9 @@ Execute the following commands to build and install vLLM from source.
 
 ??? console "pip"
     ```bash
-        sed -i '/^torch/d' requirements/build/cuda.txt    # remove torch from requirements/build/cuda.txt since we use nightly builds
         pip install -v \
-            --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
-            -r requirements/build/cuda.txt \
+            --extra-index-url https://download.pytorch.org/whl/cpu \
+            -r requirements/build/cpu.txt \
             -r requirements/cpu.txt \
         VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
             pip install dist/*.whl

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

@@ -240,7 +240,7 @@ uv pip install vllm==${VLLM_VERSION} \
         # Install dependencies
         pip install --upgrade numba \
             scipy \
-            huggingface-hub[cli,hf_transfer] \
+            huggingface-hub[cli] \
             setuptools_scm
         pip install -r requirements/rocm.txt
 

docs/models/pooling_models/README.md

@@ -59,6 +59,16 @@ please refer to [IO Processor Plugins](../../design/io_processor_plugins.md).
     Within classification tasks, there is a specialized subcategory: Cross-encoder (aka reranker) models. These models
 are a subset of classification models that accept two prompts as input and output num_labels equal to 1.
 
+### Pooling Types
+
+| Pooling Tasks  | Granularity   | Description                                                                                                                                                                                       |
+|----------------|---------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| `CLS` pooling  | Sequence-wise | For BERT‑like (bidirectional self‑attention) models, CLS pooling is used by default. This means the last_hidden_states corresponding to the first token (the [CLS] token) is taken as the output. |
+| `LAST` pooling | Sequence-wise | For GPT‑like (causal self‑attention) models, LAST pooling is used by default. This means the last_hidden_states corresponding to the last token is taken as the output.                           |
+| `MEAN` pooling | Sequence-wise | Many studies have shown that averaging the last_hidden_states over all input tokens performs better on certain downstream tasks. Therefore, more and more models are using MEAN pooling.          |
+| `ALL` pooling  | Token-wise    | Outputs the last_hidden_states for all input tokens.                                                                                                                                              |
+| `STEP` pooling | Token-wise    | Filters and outputs the last_hidden_states corresponding to the token IDs returned by returned_token_ids.                                                                                         |
+
 ### Score Types
 
 The scoring models is designed to compute similarity scores between two input prompts. It supports three model types

docs/models/pooling_models/embed.md

@@ -45,6 +45,7 @@ You can compute pairwise similarity scores to build a similarity matrix using th
 | `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ |
 | `GteModel` | Arctic-Embed-2.0-M | `Snowflake/snowflake-arctic-embed-m-v2.0`. | | |
 | `GteNewModel` | mGTE-TRM (see note) | `Alibaba-NLP/gte-multilingual-base`, etc. | | |
+| `JinaEmbeddingsV5Model`<sup>C</sup> | Qwen3-based with task-specific LoRA adapters | `jinaai/jina-embeddings-v5-text-small` (see note) | ✅︎ | ✅︎ |
 | `LlamaBidirectionalModel`<sup>C</sup> | Llama-based with bidirectional attention | `nvidia/llama-nemotron-embed-1b-v2`, etc. | ✅︎ | ✅︎ |
 | `LlamaModel`<sup>C</sup>, `LlamaForCausalLM`<sup>C</sup>, `MistralModel`<sup>C</sup>, etc. | Llama-based | `intfloat/e5-mistral-7b-instruct`, etc. | ✅︎ | ✅︎ |
 | `ModernBertModel` | ModernBERT-based | `Alibaba-NLP/gte-modernbert-base`, etc. | | |
@@ -73,6 +74,12 @@ You can compute pairwise similarity scores to build a similarity matrix using th
 !!! note
     `jinaai/jina-embeddings-v3` supports multiple tasks through LoRA, while vllm temporarily only supports text-matching tasks by merging LoRA weights.
 
+!!! note
+    `jinaai/jina-embeddings-v5-text-small` ships with four task-specific LoRA adapters
+    (`retrieval`, `text-matching`, `classification`, `clustering`). vLLM merges the
+    selected adapter into the base weights at load time. Choose the task with
+    `--hf-overrides '{"jina_task": "<task>"}'`; the default is `retrieval`.
+
 ### Multimodal Models
 
 !!! note

docs/models/pooling_models/scoring.md

@@ -160,6 +160,8 @@ The following Score API parameters are supported:
 --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params"
 --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params"
 --8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params"
+--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:scoring-common-params"
+--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:score-request-params"
 ```
 
 #### Examples
@@ -370,6 +372,8 @@ The following rerank api parameters are supported:
 --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params"
 --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params"
 --8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params"
+--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:scoring-common-params"
+--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:rerank-request-params"
 ```
 
 #### Examples

docs/models/pooling_models/token_classify.md

@@ -68,7 +68,7 @@ If your model is not in the above list, we will try to automatically convert the
     Forced alignment usage requires `--hf-overrides '{"architectures": ["Qwen3ASRForcedAlignerForTokenClassification"]}'`.
     Please refer to [examples/pooling/token_classify/forced_alignment_offline.py](../../../examples/pooling/token_classify/forced_alignment_offline.py).
 
-### As Reward Models
+### Reward Models
 
 Using token classification models as reward models. For details on reward models, see [Reward Models](reward.md).
 

docs/serving/openai_compatible_server.md

@@ -467,28 +467,11 @@ It consists of two endpoints:
 - `/tokenize` corresponds to calling `tokenizer.encode()`.
 - `/detokenize` corresponds to calling `tokenizer.decode()`.
 
-### Score API
-
-#### Score Template
-
-Some scoring models require a specific prompt format to work correctly. You can specify a custom score template using the `--chat-template` parameter (see [Chat Template](#chat-template)).
-
-Score templates are supported for **cross-encoder** models only. If you are using an **embedding** model for scoring, vLLM does not apply a score template.
-
-Like chat templates, the score template receives a `messages` list. For scoring, each message has a `role` attribute—either `"query"` or `"document"`. For the usual kind of point-wise cross-encoder, you can expect exactly two messages: one query and one document. To access the query and document content, use Jinja's `selectattr` filter:
-
-- **Query**: `{{ (messages | selectattr("role", "eq", "query") | first).content }}`
-- **Document**: `{{ (messages | selectattr("role", "eq", "document") | first).content }}`
-
-This approach is more robust than index-based access (`messages[0]`, `messages[1]`) because it selects messages by their semantic role. It also avoids assumptions about message ordering if additional message types are added to `messages` in the future.
-
-Example template file: [examples/pooling/score/template/nemotron-rerank.jinja](../../examples/pooling/score/template/nemotron-rerank.jinja)
-
 ### Generative Scoring API
 
 The `/generative_scoring` endpoint uses a CausalLM model (e.g., Llama, Qwen, Mistral) to compute the probability of specified token IDs appearing as the next token. Each item (document) is concatenated with the query to form a prompt, and the model predicts how likely each label token is as the next token after that prompt. This lets you score items against a query — for example, asking "Is this the capital of France?" and scoring each city by how likely the model is to answer "Yes".
 
-This endpoint is automatically available when the server is started with a generative model (task `"generate"`). It is separate from the pooling-based [Score API](#score-api), which uses cross-encoder, bi-encoder, or late-interaction models.
+This endpoint is automatically available when the server is started with a generative model (task `"generate"`). It is separate from the pooling-based [Score API](../models/pooling_models/scoring.md#score-api), which uses cross-encoder, bi-encoder, or late-interaction models.
 
 **Requirements:**
 

pyproject.toml

@@ -170,6 +170,9 @@ eles = "eles"
 datas = "datas"
 ser = "ser"
 ure = "ure"
+# Walsh-Hadamard Transform
+wht = "wht"
+WHT = "WHT"
 
 [tool.uv]
 no-build-isolation-package = ["torch"]

requirements/common.txt

@@ -7,7 +7,7 @@ requests >= 2.26.0
 tqdm
 blake3
 py-cpuinfo
-transformers >= 4.56.0, < 5
+transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0
 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
 fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
@@ -19,7 +19,7 @@ pillow  # Required for image processing
 prometheus-fastapi-instrumentator >= 7.0.0
 tiktoken >= 0.6.0  # Required for DBRX tokenizer
 lm-format-enforcer == 0.11.3
-llguidance >= 1.3.0, < 1.4.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "s390x" or platform_machine == "ppc64le"
+llguidance >= 1.3.0, < 1.4.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le"
 outlines_core == 0.2.11
 # required for outlines backend disk cache
 diskcache == 5.6.3
@@ -32,12 +32,14 @@ pyzmq >= 25.0.0
 msgspec
 gguf >= 0.17.0
 mistral_common[image] >= 1.11.0
+av  # required for audio in video IO
 opencv-python-headless >= 4.13.0    # required for video IO
+soundfile  # required for audio IO
 pyyaml
 six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12
 setuptools>=77.0.3,<81.0.0; python_version > '3.11' # Setuptools is used by triton, we need to ensure a modern version is installed for 3.12+ so that it does not try to import distutils, which was removed in 3.12
 einops # Required for Qwen2-VL.
-compressed-tensors == 0.14.0.1 # required for compressed-tensors
+compressed-tensors == 0.15.0.1 # required for compressed-tensors
 depyf==0.20.0 # required for profiling and debugging with compilation config
 cloudpickle # allows pickling lambda functions in model_executor/models/registry.py
 watchfiles # required for http server to monitor the updates of TLS files

requirements/kv_connectors.txt

@@ -1,3 +1,5 @@
 lmcache >= 0.3.9
-nixl[cu13] >= 0.7.1, < 0.10.0 # Required for disaggregated prefill
+nixl[cu13] >= 0.7.1, <= 0.10.1 # Required for disaggregated prefill
+nixl-cu12 >= 0.7.1, <= 0.10.1
+nixl-cu13 >= 0.7.1, <= 0.10.1
 mooncake-transfer-engine >= 0.3.8

requirements/test/cuda.in

@@ -18,10 +18,9 @@ httpx
 librosa # required for audio tests
 vector_quantize_pytorch # required for minicpmo_26 test
 vocos # required for minicpmo_26 test
-peft>=0.15.0 # required for phi-4-mm test
+peft>=0.18.1 # required for phi-4-mm test
 pqdm
 ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests
-resampy # required for audio tests
 sentence-transformers>=5.2.0 # required for embedding tests
 soundfile # required for audio tests
 jiwer # required for audio tests
@@ -39,8 +38,8 @@ opencv-python-headless >= 4.13.0 # required for video test
 datamodel_code_generator # required for minicpm3 test
 lm-eval[api]>=0.4.11 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==4.57.5
-tokenizers==0.22.0
+transformers==5.5.3
+tokenizers==0.22.2
 schemathesis>=3.39.15 # Required for openai schema test.
 # quantization
 bitsandbytes==0.49.2

requirements/test/cuda.txt

@@ -4,7 +4,7 @@ absl-py==2.1.0
     # via
     #   rouge-score
     #   tensorboard
-accelerate==1.0.1
+accelerate==1.13.0
     # via peft
 aenum==3.1.16
     # via lightly
@@ -248,7 +248,6 @@ filelock==3.16.1
     #   huggingface-hub
     #   ray
     #   torch
-    #   transformers
     #   virtualenv
 fiona==1.10.1
     # via torchgeo
@@ -331,7 +330,7 @@ h5py==3.13.0
     # via terratorch
 harfile==0.3.0
     # via schemathesis
-hf-xet==1.1.7
+hf-xet==1.4.3
     # via huggingface-hub
 hiredis==3.0.0
     # via tensorizer
@@ -345,9 +344,10 @@ httpx==0.27.2
     # via
     #   -r requirements/test/cuda.in
     #   diffusers
+    #   huggingface-hub
     #   perceptron
     #   schemathesis
-huggingface-hub==0.36.2
+huggingface-hub==1.10.2
     # via
     #   accelerate
     #   datasets
@@ -555,7 +555,6 @@ numba==0.61.2
     #   -c requirements/cuda.txt
     #   -r requirements/test/cuda.in
     #   librosa
-    #   resampy
 numpy==2.2.6
     # via
     #   -r requirements/test/cuda.in
@@ -596,7 +595,6 @@ numpy==2.2.6
     #   pyogrio
     #   pywavelets
     #   rasterio
-    #   resampy
     #   rioxarray
     #   rouge-score
     #   runai-model-streamer
@@ -756,7 +754,7 @@ pathvalidate==3.2.1
     # via pytablewriter
 patsy==1.0.1
     # via statsmodels
-peft==0.16.0
+peft==0.18.1
     # via -r requirements/test/cuda.in
 perceptron==0.1.4
     # via -r requirements/test/cuda.in
@@ -982,7 +980,7 @@ referencing==0.35.1
     # via
     #   jsonschema
     #   jsonschema-specifications
-regex==2024.9.11
+regex==2026.2.28
     # via
     #   diffusers
     #   nltk
@@ -1002,7 +1000,6 @@ requests==2.32.3
     #   google-api-core
     #   google-cloud-storage
     #   gpt-oss
-    #   huggingface-hub
     #   lightly
     #   lm-eval
     #   mistral-common
@@ -1015,10 +1012,7 @@ requests==2.32.3
     #   starlette-testclient
     #   tacoreader
     #   tiktoken
-    #   transformers
     #   wandb
-resampy==0.4.3
-    # via -r requirements/test/cuda.in
 responses==0.25.3
     # via genai-perf
 rfc3339-validator==0.1.4
@@ -1216,7 +1210,7 @@ timm==1.0.17
     #   segmentation-models-pytorch
     #   terratorch
     #   torchgeo
-tokenizers==0.22.0
+tokenizers==0.22.2
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/cuda.in
@@ -1295,7 +1289,7 @@ tqdm==4.67.3
     #   tacoreader
     #   terratorch
     #   transformers
-transformers==4.57.5
+transformers==5.5.3
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/cuda.in
@@ -1317,7 +1311,9 @@ typepy==1.3.2
 typer==0.15.2
     # via
     #   fastsafetensors
+    #   huggingface-hub
     #   perceptron
+    #   transformers
 types-python-dateutil==2.9.0.20241206
     # via arrow
 typeshed-client==2.8.2

requirements/test/nightly-torch.txt

@@ -29,8 +29,8 @@ opencv-python-headless >= 4.13.0 # required for video test
 datamodel_code_generator # required for minicpm3 test
 lm-eval[api]>=0.4.11 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==4.57.5
-tokenizers==0.22.0
+transformers==5.5.3
+tokenizers==0.22.2
 schemathesis>=3.39.15 # Required for openai schema test.
 # quantization
 bitsandbytes>=0.49.2

requirements/test/rocm.in

@@ -23,7 +23,6 @@ vocos # required for minicpmo_26 test
 peft>=0.15.0 # required for phi-4-mm test
 pqdm
 ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests
-resampy # required for audio tests
 sentence-transformers>=5.2.0 # required for embedding tests
 soundfile # required for audio tests
 jiwer # required for audio tests
@@ -38,8 +37,8 @@ opencv-python-headless>=4.13.0 # required for video test
 datamodel_code_generator # required for minicpm3 test
 lm-eval[api]>=0.4.11 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==4.57.5
-tokenizers==0.22.0
+transformers==5.5.3
+tokenizers==0.22.2
 schemathesis>=3.39.15 # Required for openai schema test
 # quantization
 bitsandbytes==0.49.2
@@ -82,4 +81,3 @@ plotly # required for perf comparison html report
 rapidfuzz
 torchgeo==0.7.0
 multiprocess==0.70.16
-huggingface-hub==0.36.2

requirements/test/rocm.txt

@@ -39,7 +39,7 @@ annotated-doc==0.0.4
     #   typer
 annotated-types==0.7.0
     # via pydantic
-anthropic==0.89.0
+anthropic==0.93.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -76,7 +76,9 @@ attrs==26.1.0
 audioread==3.0.1
     # via librosa
 av==16.1.0
-    # via -r requirements/test/rocm.in
+    # via
+    #   -r requirements/test/../common.txt
+    #   -r requirements/test/rocm.in
 azure-core==1.39.0
     # via
     #   azure-identity
@@ -172,7 +174,7 @@ colorful==0.5.8
     # via ray
 colorlog==6.10.1
     # via optuna
-compressed-tensors==0.14.0.1
+compressed-tensors==0.15.0.1
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -269,9 +271,9 @@ fastapi==0.135.2
     #   model-hosting-container-standards
 fastapi-cli==0.0.24
     # via fastapi
-fastapi-cloud-cli==0.15.1
+fastapi-cloud-cli==0.16.1
     # via fastapi-cli
-fastar==0.9.0
+fastar==0.10.0
     # via fastapi-cloud-cli
 fastparquet==2026.3.0
     # via genai-perf
@@ -290,7 +292,6 @@ filelock==3.25.2
     #   python-discovery
     #   ray
     #   torch
-    #   transformers
     #   virtualenv
 fiona==1.10.1
     # via torchgeo
@@ -384,7 +385,7 @@ h5py==3.16.0
     # via terratorch
 harfile==0.4.0
     # via schemathesis
-hf-xet==1.4.2
+hf-xet==1.4.3
     # via huggingface-hub
 hiredis==3.3.1
     # via tensorizer
@@ -403,6 +404,7 @@ httpx==0.27.2
     #   diffusers
     #   fastapi
     #   fastapi-cloud-cli
+    #   huggingface-hub
     #   mcp
     #   model-hosting-container-standards
     #   openai
@@ -410,9 +412,8 @@ httpx==0.27.2
     #   schemathesis
 httpx-sse==0.4.3
     # via mcp
-huggingface-hub==0.36.2
+huggingface-hub==1.10.2
     # via
-    #   -r requirements/test/rocm.in
     #   accelerate
     #   datasets
     #   diffusers
@@ -484,7 +485,7 @@ jinja2==3.1.6
     #   genai-perf
     #   lm-eval
     #   torch
-jiter==0.13.0
+jiter==0.14.0
     # via
     #   anthropic
     #   openai
@@ -631,7 +632,7 @@ msgpack==1.1.2
     # via
     #   librosa
     #   ray
-msgspec==0.20.0
+msgspec==0.21.0
     # via -r requirements/test/../common.txt
 mteb==2.11.5
     # via -r requirements/test/rocm.in
@@ -663,7 +664,6 @@ numba==0.61.2
     #   -c requirements/rocm.txt
     #   -r requirements/test/rocm.in
     #   librosa
-    #   resampy
 numkong==7.1.1
     # via albucore
 numpy==2.2.6
@@ -709,7 +709,6 @@ numpy==2.2.6
     #   pytrec-eval-terrier
     #   pywavelets
     #   rasterio
-    #   resampy
     #   rioxarray
     #   rouge-score
     #   runai-model-streamer
@@ -742,7 +741,7 @@ omegaconf==2.3.0
     #   lightning
 open-clip-torch==2.32.0
     # via -r requirements/test/rocm.in
-openai==2.30.0
+openai==2.31.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1093,7 +1092,7 @@ python-dotenv==1.2.2
     #   uvicorn
 python-json-logger==4.1.0
     # via -r requirements/test/../common.txt
-python-multipart==0.0.22
+python-multipart==0.0.26
     # via
     #   fastapi
     #   mcp
@@ -1180,7 +1179,6 @@ requests==2.32.5
     #   google-api-core
     #   google-cloud-storage
     #   gpt-oss
-    #   huggingface-hub
     #   lightly
     #   lm-eval
     #   mistral-common
@@ -1194,10 +1192,7 @@ requests==2.32.5
     #   starlette-testclient
     #   tacoreader
     #   tiktoken
-    #   transformers
     #   wandb
-resampy==0.4.3
-    # via -r requirements/test/rocm.in
 responses==0.26.0
     # via genai-perf
 rfc3339-validator==0.1.4
@@ -1338,6 +1333,7 @@ sortedcontainers==2.4.0
     # via hypothesis
 soundfile==0.13.1
     # via
+    #   -r requirements/test/../common.txt
     #   -r requirements/test/rocm.in
     #   genai-perf
     #   librosa
@@ -1428,7 +1424,7 @@ timm==1.0.17
     #   segmentation-models-pytorch
     #   terratorch
     #   torchgeo
-tokenizers==0.22.0
+tokenizers==0.22.2
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1471,7 +1467,7 @@ tqdm==4.67.3
     #   tacoreader
     #   terratorch
     #   transformers
-transformers==4.57.5
+transformers==5.5.3
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1498,7 +1494,9 @@ typer==0.24.1
     #   fastapi-cli
     #   fastapi-cloud-cli
     #   fastsafetensors
+    #   huggingface-hub
     #   perceptron
+    #   transformers
 typeshed-client==2.9.0
     # via jsonargparse
 typing-extensions==4.15.0

requirements/test/xpu.in

@@ -13,7 +13,6 @@ pytest-shard
 absl-py
 accelerate
 arctic-inference
-hf_transfer
 lm_eval[api]
 modelscope
 

requirements/test/xpu.txt

@@ -19,7 +19,9 @@ aiosignal==1.4.0
 albumentations==1.4.6
     # via -r requirements/test/xpu.in
 annotated-doc==0.0.4
-    # via fastapi
+    # via
+    #   fastapi
+    #   typer
 annotated-types==0.7.0
     # via pydantic
 anyio==4.13.0
@@ -64,6 +66,7 @@ click==8.3.1
     #   jiwer
     #   nltk
     #   schemathesis
+    #   typer
     #   uvicorn
 colorama==0.4.6
     # via sacrebleu
@@ -112,7 +115,6 @@ filelock==3.25.2
     #   huggingface-hub
     #   modelscope
     #   torch
-    #   transformers
 frozenlist==1.8.0
     # via
     #   aiohttp
@@ -133,9 +135,7 @@ h11==0.16.0
     #   uvicorn
 harfile==0.4.0
     # via schemathesis
-hf-transfer==0.1.9
-    # via -r requirements/test/xpu.in
-hf-xet==1.4.2
+hf-xet==1.4.3
     # via huggingface-hub
 html2text==2025.4.15
     # via gpt-oss
@@ -144,8 +144,9 @@ httpcore==1.0.9
 httpx==0.28.1
     # via
     #   datasets
+    #   huggingface-hub
     #   schemathesis
-huggingface-hub==0.36.2
+huggingface-hub==1.10.2
     # via
     #   accelerate
     #   datasets
@@ -515,7 +516,6 @@ requests==2.33.1
     #   docker
     #   evaluate
     #   gpt-oss
-    #   huggingface-hub
     #   lm-eval
     #   mistral-common
     #   modelscope
@@ -524,11 +524,11 @@ requests==2.33.1
     #   schemathesis
     #   starlette-testclient
     #   tiktoken
-    #   transformers
 rich==14.3.3
     # via
     #   mteb
     #   schemathesis
+    #   typer
 rouge-score==0.1.2
     # via lm-eval
 rpds-py==0.30.0
@@ -572,6 +572,8 @@ setuptools==80.10.2
     #   modelscope
     #   pytablewriter
     #   torch
+shellingham==1.5.4
+    # via typer
 six==1.17.0
     # via
     #   -c requirements/common.txt
@@ -665,7 +667,7 @@ tqdm==4.67.3
     #   pqdm
     #   sentence-transformers
     #   transformers
-transformers==4.57.6
+transformers==5.5.3
     # via
     #   -c requirements/common.txt
     #   sentence-transformers
@@ -676,6 +678,10 @@ typepy==1.3.4
     #   dataproperty
     #   pytablewriter
     #   tabledata
+typer==0.24.1
+    # via
+    #   huggingface-hub
+    #   transformers
 typing-extensions==4.15.0
     # via
     #   -c requirements/common.txt

setup.py

@@ -693,6 +693,12 @@ class precompiled_wheel_utils:
                 flash_attn_regex = re.compile(
                     r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
                 )
+                # __init__.py and flash_attn_interface.py are source-controlled
+                # in vllm and should not be overwritten (matches cmake exclusions)
+                flash_attn_files_to_skip = {
+                    "vllm/vllm_flash_attn/__init__.py",
+                    "vllm/vllm_flash_attn/flash_attn_interface.py",
+                }
                 triton_kernels_regex = re.compile(
                     r"vllm/third_party/triton_kernels/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
                 )
@@ -705,7 +711,11 @@ class precompiled_wheel_utils:
                     filter(lambda x: x.filename in files_to_copy, wheel.filelist)
                 )
                 file_members += list(
-                    filter(lambda x: flash_attn_regex.match(x.filename), wheel.filelist)
+                    filter(
+                        lambda x: flash_attn_regex.match(x.filename)
+                        and x.filename not in flash_attn_files_to_skip,
+                        wheel.filelist,
+                    )
                 )
                 file_members += list(
                     filter(
@@ -1082,10 +1092,7 @@ setup(
         "instanttensor": ["instanttensor >= 0.1.5"],
         "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"],
         "audio": [
-            "av",
-            "resampy",
             "scipy",
-            "soundfile",
             "mistral_common[audio]",
         ],  # Required for audio processing
         "video": [],  # Kept for backwards compatibility
@@ -1094,7 +1101,7 @@ setup(
         # NOTE: When updating helion version, also update CI files:
         #   - .buildkite/test_areas/kernels.yaml
         #   - .buildkite/test-amd.yaml
-        "helion": ["helion==0.3.3"],
+        "helion": ["helion==1.0.0"],
         # Optional deps for gRPC server (vllm serve --grpc)
         "grpc": ["smg-grpc-servicer[vllm] >= 0.5.0"],
         # Optional deps for OpenTelemetry tracing

tests/basic_correctness/test_cumem.py

@@ -13,6 +13,8 @@ from vllm.utils.mem_constants import GiB_bytes
 
 from ..utils import create_new_process_for_each_test, requires_fp8
 
+DEVICE_TYPE = current_platform.device_type
+
 
 @create_new_process_for_each_test("fork" if not current_platform.is_rocm() else "spawn")
 def test_python_error():
@@ -26,13 +28,13 @@ def test_python_error():
     tensors = []
     with allocator.use_memory_pool():
         # allocate 70% of the total memory
-        x = torch.empty(alloc_bytes, dtype=torch.uint8, device="cuda")
+        x = torch.empty(alloc_bytes, dtype=torch.uint8, device=DEVICE_TYPE)
         tensors.append(x)
     # release the memory
     allocator.sleep()
 
     # allocate more memory than the total memory
-    y = torch.empty(alloc_bytes, dtype=torch.uint8, device="cuda")
+    y = torch.empty(alloc_bytes, dtype=torch.uint8, device=DEVICE_TYPE)
     tensors.append(y)
     with pytest.raises(RuntimeError):
         # when the allocator is woken up, it should raise an error
@@ -44,17 +46,17 @@ def test_python_error():
 def test_basic_cumem():
     # some tensors from default memory pool
     shape = (1024, 1024)
-    x = torch.empty(shape, device="cuda")
+    x = torch.empty(shape, device=DEVICE_TYPE)
     x.zero_()
 
     # some tensors from custom memory pool
     allocator = CuMemAllocator.get_instance()
     with allocator.use_memory_pool():
         # custom memory pool
-        y = torch.empty(shape, device="cuda")
+        y = torch.empty(shape, device=DEVICE_TYPE)
         y.zero_()
         y += 1
-        z = torch.empty(shape, device="cuda")
+        z = torch.empty(shape, device=DEVICE_TYPE)
         z.zero_()
         z += 2
 
@@ -77,16 +79,16 @@ def test_basic_cumem():
 def test_cumem_with_cudagraph():
     allocator = CuMemAllocator.get_instance()
     with allocator.use_memory_pool():
-        weight = torch.eye(1024, device="cuda")
+        weight = torch.eye(1024, device=DEVICE_TYPE)
     with allocator.use_memory_pool(tag="discard"):
-        cache = torch.empty(1024, 1024, device="cuda")
+        cache = torch.empty(1024, 1024, device=DEVICE_TYPE)
 
     def model(x):
         out = x @ weight
         cache[: out.size(0)].copy_(out)
         return out + 1
 
-    x = torch.empty(128, 1024, device="cuda")
+    x = torch.empty(128, 1024, device=DEVICE_TYPE)
 
     # warmup
     model(x)

tests/benchmarks/test_sampling_params.py

@@ -0,0 +1,258 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+import pytest
+
+from vllm.benchmarks.datasets.utils import get_sampling_params
+from vllm.tokenizers import TokenizerLike
+
+
+class _FakeTokenizer(TokenizerLike):
+    """Minimal tokenizer implementing the TokenizerLike protocol
+    for testing get_sampling_params."""
+
+    def __init__(self, vocab_size: int = 1000, num_special_tokens: int = 0) -> None:
+        self._vocab_size = vocab_size
+        self._num_special_tokens = num_special_tokens
+
+    # -- Properties required by TokenizerLike --
+
+    @classmethod
+    def from_pretrained(cls, path_or_repo_id, *a, **kw):  # type: ignore[override]
+        return cls()
+
+    @property
+    def vocab_size(self) -> int:
+        return self._vocab_size
+
+    @property
+    def all_special_tokens(self) -> list[str]:
+        return []
+
+    @property
+    def all_special_ids(self) -> list[int]:
+        return []
+
+    @property
+    def bos_token_id(self) -> int:
+        return 0
+
+    @property
+    def eos_token_id(self) -> int:
+        return 1
+
+    @property
+    def pad_token_id(self) -> int:
+        return 2
+
+    @property
+    def is_fast(self) -> bool:
+        return False
+
+    @property
+    def max_token_id(self) -> int:
+        return self._vocab_size - 1
+
+    @property
+    def max_chars_per_token(self) -> int:
+        return 4
+
+    @property
+    def truncation_side(self) -> str:
+        return "right"
+
+    def num_special_tokens_to_add(self) -> int:
+        return self._num_special_tokens
+
+    def __call__(self, text, text_pair=None, **kw):  # type: ignore[override]
+        raise NotImplementedError
+
+    def get_vocab(self) -> dict[str, int]:
+        return {}
+
+    def get_added_vocab(self) -> dict[str, int]:
+        return {}
+
+    def encode(self, text, **kw) -> list[int]:  # type: ignore[override]
+        raise NotImplementedError
+
+    def apply_chat_template(self, messages, **kw):  # type: ignore[override]
+        raise NotImplementedError
+
+    def convert_tokens_to_ids(self, tokens):  # type: ignore[override]
+        raise NotImplementedError
+
+    def convert_tokens_to_string(self, tokens: list[str]) -> str:
+        raise NotImplementedError
+
+    def decode(self, ids, skip_special_tokens: bool = False) -> str:  # type: ignore[override]
+        raise NotImplementedError
+
+    def convert_ids_to_tokens(  # type: ignore[override]
+        self, ids, skip_special_tokens: bool = False
+    ) -> list[str]:
+        raise NotImplementedError
+
+
+class TestGetSamplingParams:
+    """Tests for ``get_sampling_params`` in ``vllm.benchmarks.datasets.shared``."""
+
+    # -- helpers --
+
+    @staticmethod
+    def _tok(vocab_size: int = 1000, num_special: int = 0) -> _FakeTokenizer:
+        return _FakeTokenizer(vocab_size=vocab_size, num_special_tokens=num_special)
+
+    # -- return shape / dtype --
+
+    def test_returns_three_arrays(self):
+        rng = np.random.default_rng(0)
+        result = get_sampling_params(rng, 5, 0.0, 100, 50, self._tok())
+        assert len(result) == 3
+        for arr in result:
+            assert isinstance(arr, np.ndarray)
+
+    @pytest.mark.parametrize("n", [1, 10, 100])
+    def test_output_length_matches_num_requests(self, n: int):
+        rng = np.random.default_rng(42)
+        input_lens, output_lens, offsets = get_sampling_params(
+            rng, n, 0.0, 64, 32, self._tok()
+        )
+        assert input_lens.shape == (n,)
+        assert output_lens.shape == (n,)
+        assert offsets.shape == (n,)
+
+    # -- fixed lengths (range_ratio = 0) --
+
+    def test_zero_range_ratio_gives_constant_lengths(self):
+        rng = np.random.default_rng(7)
+        input_lens, output_lens, _ = get_sampling_params(
+            rng, 20, 0.0, 128, 64, self._tok()
+        )
+        assert np.all(input_lens == 128)
+        assert np.all(output_lens == 64)
+
+    def test_special_tokens_subtracted_from_input_only(self):
+        rng = np.random.default_rng(7)
+        input_lens, output_lens, _ = get_sampling_params(
+            rng, 10, 0.0, 100, 50, self._tok(num_special=4)
+        )
+        # real_input_len = 100 - 4 = 96, range_ratio 0 → all 96
+        assert np.all(input_lens == 96)
+        # special tokens are not subtracted from output length
+        assert np.all(output_lens == 50)
+
+    # -- range ratios --
+
+    def test_input_range_bounds(self):
+        rng = np.random.default_rng(0)
+        ratio = 0.5
+        base = 200
+        input_lens, _, _ = get_sampling_params(
+            rng, 500, {"input": ratio, "output": 0.0}, base, 50, self._tok()
+        )
+        lo = int(np.floor(base * (1 - ratio)))
+        hi = int(np.ceil(base * (1 + ratio)))
+        assert np.all(input_lens >= lo)
+        assert np.all(input_lens <= hi)
+
+    def test_output_range_bounds(self):
+        rng = np.random.default_rng(0)
+        ratio = 0.3
+        base = 100
+        _, output_lens, _ = get_sampling_params(
+            rng, 500, {"input": 0.0, "output": ratio}, 50, base, self._tok()
+        )
+        lo = max(1, int(np.floor(base * (1 - ratio))))
+        hi = int(np.ceil(base * (1 + ratio)))
+        assert np.all(output_lens >= lo)
+        assert np.all(output_lens <= hi)
+
+    def test_output_low_clamped_to_one(self):
+        """Even with a high ratio that would push output_low to 0,
+        the function clamps it to 1."""
+        rng = np.random.default_rng(0)
+        # output_len=1, ratio=0.99 → floor(1*0.01)=0, should clamp to 1
+        _, output_lens, _ = get_sampling_params(
+            rng, 50, {"input": 0.0, "output": 0.99}, 100, 1, self._tok()
+        )
+        assert np.all(output_lens >= 1)
+
+    # -- offsets bounded by vocab_size --
+
+    @pytest.mark.parametrize("vocab", [100, 32000, 128256])
+    def test_offsets_within_vocab(self, vocab: int):
+        rng = np.random.default_rng(0)
+        _, _, offsets = get_sampling_params(
+            rng, 200, 0.0, 64, 32, self._tok(vocab_size=vocab)
+        )
+        assert np.all(offsets >= 0)
+        assert np.all(offsets < vocab)
+
+    # -- reproducibility --
+
+    def test_same_seed_same_results(self):
+        tok = self._tok()
+        rr = {"input": 0.3, "output": 0.2}
+        a = get_sampling_params(np.random.default_rng(42), 50, rr, 256, 64, tok)
+        b = get_sampling_params(np.random.default_rng(42), 50, rr, 256, 64, tok)
+        for arr_a, arr_b in zip(a, b):
+            np.testing.assert_array_equal(arr_a, arr_b)
+
+    def test_different_seed_different_results(self):
+        tok = self._tok()
+        rr = {"input": 0.3, "output": 0.2}
+        a = get_sampling_params(np.random.default_rng(0), 50, rr, 256, 64, tok)
+        b = get_sampling_params(np.random.default_rng(1), 50, rr, 256, 64, tok)
+        # Extremely unlikely all three arrays match with different seeds
+        assert not all(np.array_equal(arr_a, arr_b) for arr_a, arr_b in zip(a, b))
+
+    # -- validation / error paths --
+
+    @pytest.mark.parametrize("bad_ratio", [-0.1, 1.0, 1.5])
+    def test_invalid_input_range_ratio(self, bad_ratio: float):
+        rng = np.random.default_rng(0)
+        with pytest.raises(ValueError, match="input_range_ratio"):
+            get_sampling_params(
+                rng, 10, {"input": bad_ratio, "output": 0.0}, 100, 50, self._tok()
+            )
+
+    @pytest.mark.parametrize("bad_ratio", [-0.1, 1.0, 1.5])
+    def test_invalid_output_range_ratio(self, bad_ratio: float):
+        rng = np.random.default_rng(0)
+        with pytest.raises(ValueError, match="output_range_ratio"):
+            get_sampling_params(
+                rng, 10, {"input": 0.0, "output": bad_ratio}, 100, 50, self._tok()
+            )
+
+    def test_invalid_dict_missing_keys(self):
+        rng = np.random.default_rng(0)
+        with pytest.raises(ValueError, match="input.*output"):
+            get_sampling_params(rng, 10, {"input": 0.1}, 100, 50, self._tok())
+
+    def test_input_len_zero_with_special_tokens(self):
+        """input_len < num_special_tokens → real_input_len = 0, which is fine
+        (range [0, 0])."""
+        rng = np.random.default_rng(0)
+        input_lens, _, _ = get_sampling_params(
+            rng, 5, 0.0, 5, 50, self._tok(num_special=10)
+        )
+        # real_input_len = max(0, 5 - 10) = 0
+        assert np.all(input_lens == 0)
+
+    # -- edge cases --
+
+    def test_single_request(self):
+        rng = np.random.default_rng(0)
+        i, o, off = get_sampling_params(rng, 1, 0.0, 100, 50, self._tok())
+        assert i.shape == (1,)
+        assert o.shape == (1,)
+        assert off.shape == (1,)
+
+    def test_large_num_requests(self):
+        rng = np.random.default_rng(0)
+        i, o, off = get_sampling_params(rng, 10_000, 0.5, 512, 128, self._tok())
+        assert i.shape == (10_000,)
+        assert o.shape == (10_000,)
+        assert off.shape == (10_000,)

tests/benchmarks/test_txt_slices_dataset.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import json
+from pathlib import Path
+
+import pytest
+from transformers import AutoTokenizer, PreTrainedTokenizerBase
+
+from vllm.benchmarks.datasets import CustomDataset
+from vllm.benchmarks.datasets.create_txt_slices_dataset import create_txt_slices_jsonl
+
+
+@pytest.fixture(scope="session")
+def hf_tokenizer() -> PreTrainedTokenizerBase:
+    # Use a small, commonly available tokenizer
+    return AutoTokenizer.from_pretrained("gpt2")
+
+
+text_content = """
+Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor
+incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud
+exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
+Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat
+nulla pariatur. Excepteur sint occaecat cupidatat non proident,
+sunt in culpa qui officia deserunt mollit anim id est laborum.
+"""
+
+
+@pytest.mark.benchmark
+def test_create_txt_slices_jsonl(
+    hf_tokenizer: PreTrainedTokenizerBase, tmp_path: Path
+) -> None:
+    """Test that create_txt_slices_jsonl produces valid JSONL for CustomDataset."""
+    txt_path = tmp_path / "input.txt"
+    jsonl_path = tmp_path / "input.txt.jsonl"
+
+    txt_path.write_text(text_content)
+
+    create_txt_slices_jsonl(
+        input_path=str(txt_path),
+        output_path=str(jsonl_path),
+        tokenizer_name="gpt2",
+        num_prompts=10,
+        input_len=10,
+        output_len=10,
+    )
+
+    # Verify the JSONL file is valid and has the expected structure
+    records = [json.loads(line) for line in jsonl_path.read_text().splitlines()]
+
+    assert len(records) == 10
+    for record in records:
+        assert "prompt" in record
+        assert "output_tokens" in record
+        assert isinstance(record["prompt"], str)
+        assert record["output_tokens"] == 10
+
+    # Verify the JSONL file can be loaded by CustomDataset
+    dataset = CustomDataset(dataset_path=str(jsonl_path))
+    samples = dataset.sample(
+        tokenizer=hf_tokenizer,
+        num_requests=10,
+        output_len=10,
+        skip_chat_template=True,
+    )
+
+    assert len(samples) == 10
+    assert all(sample.expected_output_len == 10 for sample in samples)

tests/compile/passes/distributed/test_async_tp.py

@@ -31,6 +31,7 @@ from vllm.platforms import current_platform
 from vllm.utils.system_utils import update_environment_variables
 from vllm.utils.torch_utils import set_random_seed
 
+DEVICE_TYPE = current_platform.device_type
 FP8_DTYPE = current_platform.fp8_dtype()
 
 prompts = [
@@ -299,7 +300,7 @@ def async_tp_pass_on_test_model(
 ):
     set_random_seed(0)
 
-    device = torch.device(f"cuda:{local_rank}")
+    device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
     torch.accelerator.set_device_index(device)
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
@@ -324,7 +325,7 @@ def async_tp_pass_on_test_model(
             fuse_gemm_comms=True,
         ),
     )
-    vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))
+    vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))
 
     # this is a fake model name to construct the model config
     # in the vllm_config, it's not really used.

tests/compile/passes/distributed/test_fusion_all_reduce.py

@@ -37,6 +37,8 @@ from vllm.platforms import current_platform
 from vllm.utils.system_utils import update_environment_variables
 from vllm.utils.torch_utils import set_random_seed
 
+DEVICE_TYPE = current_platform.device_type
+
 
 class TestAllReduceRMSNormModel(torch.nn.Module):
     def __init__(
@@ -268,7 +270,7 @@ def all_reduce_fusion_pass_on_test_model(
 ):
     set_random_seed(0)
 
-    device = torch.device(f"cuda:{local_rank}")
+    device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
     torch.accelerator.set_device_index(device)
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
@@ -300,7 +302,7 @@ def all_reduce_fusion_pass_on_test_model(
     vllm_config.compilation_config.pass_config = PassConfig(
         fuse_allreduce_rms=True, eliminate_noops=True
     )
-    vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))
+    vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))
     vllm_config.parallel_config.rank = local_rank  # Setup rank for debug path
 
     # this is a fake model name to construct the model config

tests/compile/passes/distributed/test_sequence_parallelism.py

@@ -35,6 +35,8 @@ from vllm.platforms import current_platform
 from vllm.utils.system_utils import update_environment_variables
 from vllm.utils.torch_utils import set_random_seed
 
+DEVICE_TYPE = current_platform.device_type
+
 pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")
 
 FP8_DTYPE = current_platform.fp8_dtype()
@@ -228,7 +230,7 @@ def sequence_parallelism_pass_on_test_model(
 ):
     set_random_seed(0)
 
-    device = torch.device(f"cuda:{local_rank}")
+    device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
     torch.accelerator.set_device_index(device)
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
@@ -258,7 +260,7 @@ def sequence_parallelism_pass_on_test_model(
             eliminate_noops=True,
         ),
     )  # NoOp needed for fusion
-    device_config = DeviceConfig(device=torch.device("cuda"))
+    device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))
 
     # this is a fake model name to construct the model config
     # in the vllm_config, it's not really used.

tests/compile/passes/test_fusion_attn.py

@@ -41,6 +41,7 @@ from vllm.v1.attention.backend import AttentionMetadata
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.kv_cache_interface import AttentionSpec, get_kv_quant_mode
 
+DEVICE_TYPE = current_platform.device_type
 FP8_DTYPE = current_platform.fp8_dtype()
 FP4_DTYPE = torch.uint8
 
@@ -300,7 +301,7 @@ def test_attention_quant_pattern(
 
     custom_ops_list = custom_ops.split(",") if custom_ops else []
 
-    device = torch.device("cuda:0")
+    device = torch.device(f"{DEVICE_TYPE}:0")
     torch.set_default_dtype(dtype)
     torch.manual_seed(42)
 

tests/compile/passes/test_mla_attn_quant_fusion.py

@@ -45,6 +45,7 @@ from vllm.v1.kv_cache_interface import MLAAttentionSpec
 
 FP8_DTYPE = current_platform.fp8_dtype()
 FP4_DTYPE = torch.uint8
+DEVICE_TYPE = current_platform.device_type
 
 
 class MLAAttentionQuantPatternModel(torch.nn.Module):
@@ -356,7 +357,7 @@ def test_mla_attention_quant_pattern(
 
     custom_ops_list = custom_ops.split(",") if custom_ops else []
 
-    device = torch.device("cuda:0")
+    device = torch.device(f"{DEVICE_TYPE}:0")
     torch.set_default_dtype(dtype)
     torch.manual_seed(42)
 

tests/compile/passes/test_noop_elimination.py

@@ -8,6 +8,9 @@ import vllm
 from tests.compile.backend import TestBackend
 from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
 from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConfig
+from vllm.platforms import current_platform
+
+DEVICE_TYPE = current_platform.device_type
 
 
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
@@ -17,7 +20,7 @@ from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConf
 )
 @pytest.mark.parametrize("hidden_size", [64, 4096])
 def test_noop_elimination(dtype, num_tokens, hidden_size, buffer_size):
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     torch.set_default_dtype(dtype)
     torch.manual_seed(1)
 
@@ -88,7 +91,7 @@ def test_non_noop_slice_preserved():
     Regression test for a bug where end=-1 was treated like an inferred
     dimension (reshape semantics) leading to incorrect elimination.
     """
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     x = torch.randn(16, 16)
 
     class SliceModel(torch.nn.Module):

tests/compile/passes/test_scatter_split_replace.py

@@ -13,6 +13,9 @@ from vllm.compilation.passes.utility.scatter_split_replace import (
 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
+from vllm.platforms import current_platform
+
+DEVICE_TYPE = current_platform.device_type
 
 
 class ScatterSplitReplacementModel(nn.Module):
@@ -61,7 +64,7 @@ class ScatterSplitReplacementModel(nn.Module):
 
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 def test_scatter_split_replace(dtype):
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     torch.set_default_dtype(dtype)
     torch.manual_seed(0)
 

tests/compile/passes/test_split_coalescing.py

@@ -8,6 +8,9 @@ import vllm
 from tests.compile.backend import TestBackend
 from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
 from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConfig
+from vllm.platforms import current_platform
+
+DEVICE_TYPE = current_platform.device_type
 
 
 class SplitCoalescingModel(torch.nn.Module):
@@ -28,7 +31,7 @@ class SplitCoalescingModel(torch.nn.Module):
 
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 def test_split_coalescing(dtype):
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     torch.set_default_dtype(dtype)
     torch.manual_seed(0)
 

tests/compile/test_config.py

@@ -31,6 +31,8 @@ from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
 # This import automatically registers `torch.ops.silly.attention`
 from . import silly_attention  # noqa: F401
 
+DEVICE_TYPE = current_platform.device_type
+
 
 def test_version():
     # Test the version comparison logic using the private function
@@ -456,7 +458,7 @@ def test_cached_compilation_config(default_vllm_config):
     from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
 
     dtype = torch.bfloat16
-    device = torch.device("cuda:0")
+    device = torch.device(f"{DEVICE_TYPE}:0")
     batch_size, num_qo_heads, head_size = 8, 16, 128
 
     # access and cache default compilation config
@@ -478,7 +480,7 @@ def test_cached_compilation_config(default_vllm_config):
         query_quant = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
         query_quant = torch.compile(query_quant)
 
-        _q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
+        _q_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE)
         query = torch.randn(
             batch_size, num_qo_heads * head_size, dtype=dtype, device=device
         )

tests/compile/test_dynamic_shapes_compilation.py

@@ -222,3 +222,47 @@ def test_model_specialization_with_evaluate_guards(
         torch.randn(1, 10).cuda(),
         is_01_specialization=True,
     )
+
+
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
+def test_piecewise_backend_empty_sym_shape_indices():
+    """Test that PiecewiseBackend handles empty sym_shape_indices correctly.
+
+    When all inputs have static shapes (no torch.SymInt), sym_shape_indices
+    will be empty. The fix in PiecewiseBackend.__call__ handles this case
+    by using the first compiled range_entry.
+    """
+    gc.collect()
+    torch.accelerator.empty_cache()
+    torch.accelerator.synchronize()
+
+    # Use small max_model_len and max_num_batched_tokens to encourage
+    # static shape compilation with empty sym_shape_indices
+    llm = LLM(
+        model="Qwen/Qwen3-0.6B",
+        max_model_len=512,
+        max_num_batched_tokens=1,
+        compilation_config={
+            "mode": CompilationMode.VLLM_COMPILE,
+            "dynamic_shapes_config": {
+                "type": DynamicShapesType.BACKED.value,
+            },
+        },
+    )
+
+    sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
+
+    # Generate with static shape inputs
+    output = llm.generate("Hello, my name is", sampling_params=sampling_params)
+    result = output[0].outputs[0].text
+    assert len(result) > 0, "Should generate non-empty output"
+
+    # Generate again to verify compilation works with empty sym_shape_indices
+    output = llm.generate("The capital of France is", sampling_params=sampling_params)
+    result = output[0].outputs[0].text
+    assert len(result) > 0, "Should generate non-empty output on second run"
+
+    del llm
+    gc.collect()
+    torch.accelerator.empty_cache()
+    torch.accelerator.synchronize()

tests/compile/test_graph_partition.py

@@ -15,10 +15,13 @@ from vllm.compilation.backends import (
     split_graph,
 )
 from vllm.compilation.passes.fx_utils import find_op_nodes
+from vllm.platforms import current_platform
 
 # This import automatically registers `torch.ops.silly.attention`
 from . import silly_attention  # noqa: F401
 
+DEVICE_TYPE = current_platform.device_type
+
 
 def test_getitem_moved_to_producer_subgraph():
     """
@@ -151,7 +154,7 @@ def test_consecutive_ops_in_split():
         final_result = torch.sigmoid(attn_inout)
         return final_result
 
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
 
     # Create the traced FX graph for the model
     x = torch.randn(8, 4)
@@ -329,7 +332,7 @@ def test_builtin_empty_only_partition_is_merged():
         "Expected two builtin empty_like nodes in merged non-splitting subgraph"
     )
 
-    x = torch.randn(2, 3, device="cuda")
+    x = torch.randn(2, 3, device=DEVICE_TYPE)
     output_original = gm(x)
     output_split = split_gm(x)
     assert torch.allclose(output_original, output_split), "Output mismatch after split"

tests/compile/test_rotary_embedding_compile.py

@@ -16,6 +16,8 @@ from vllm.config.compilation import CompilationMode, CUDAGraphMode
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.platforms import current_platform
 
+DEVICE_TYPE = current_platform.device_type
+
 
 @support_torch_compile
 class RotaryEmbeddingCompileModule(torch.nn.Module):
@@ -45,7 +47,7 @@ def test_rotary_embedding_torch_compile_with_custom_op(monkeypatch):
     monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1")
     monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "0")
 
-    device = "cuda"
+    device = DEVICE_TYPE
     positions = torch.arange(16, device=device)
     query = torch.randn(16, 32, device=device, dtype=torch.bfloat16)
     key = torch.randn(16, 32, device=device, dtype=torch.bfloat16)

tests/compile/test_structured_logging.py

@@ -17,8 +17,10 @@ from vllm.config.compilation import (
 )
 from vllm.config.scheduler import SchedulerConfig
 from vllm.forward_context import set_forward_context
+from vllm.platforms import current_platform
 
 MLP_SIZE = 64
+DEVICE_TYPE = current_platform.device_type
 
 
 @support_torch_compile
@@ -71,7 +73,7 @@ class TraceStructuredCapture:
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 def test_vllm_structured_logging_artifacts(use_fresh_inductor_cache):
     """Test that all expected vLLM artifacts are logged during compilation."""
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
 
     capture = TraceStructuredCapture()
 

tests/config/base_model_arch_groundtruth.json

@@ -356,6 +356,23 @@
         "is_multimodal_model": false,
         "dtype": "torch.float32"
     },
+    "stepfun-ai/Step-3.5-Flash": {
+        "architectures": [
+            "Step3p5ForCausalLM"
+        ],
+        "model_type": "step3p5",
+        "text_model_type": "step3p5",
+        "hidden_size": 4096,
+        "total_num_hidden_layers": 45,
+        "total_num_attention_heads": 64,
+        "head_size": 128,
+        "vocab_size": 128896,
+        "total_num_kv_heads": 8,
+        "num_experts": 288,
+        "is_deepseek_mla": false,
+        "is_multimodal_model": false,
+        "dtype": "torch.bfloat16"
+    },
     "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16": {
         "architectures": [
             "NemotronHForCausalLM"

tests/config/test_model_arch_config.py

@@ -16,6 +16,7 @@ BASE_TRUST_REMOTE_CODE_MODELS = {
     "nvidia/Llama-3_3-Nemotron-Super-49B-v1",
     "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",
     "XiaomiMiMo/MiMo-7B-RL",
+    "stepfun-ai/Step-3.5-Flash",
     # Excluded: Not available online right now
     # "FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1",
     "meituan-longcat/LongCat-Flash-Chat",

tests/conftest.py

@@ -364,6 +364,7 @@ class HfRunner:
         model_name: str,
         dtype: str = "auto",
         *,
+        revision: str | None = None,
         model_kwargs: dict[str, Any] | None = None,
         trust_remote_code: bool = True,
         is_sentence_transformer: bool = False,
@@ -383,6 +384,7 @@ class HfRunner:
             self._init(
                 model_name=model_name,
                 dtype=dtype,
+                revision=revision,
                 model_kwargs=model_kwargs,
                 trust_remote_code=trust_remote_code,
                 is_sentence_transformer=is_sentence_transformer,
@@ -396,6 +398,7 @@ class HfRunner:
         model_name: str,
         dtype: str = "auto",
         *,
+        revision: str | None = None,
         model_kwargs: dict[str, Any] | None = None,
         trust_remote_code: bool = True,
         is_sentence_transformer: bool = False,
@@ -410,6 +413,15 @@ class HfRunner:
             model_name,
             trust_remote_code=trust_remote_code,
         )
+        # HF runner should use the HF config so that it's consistent with the HF model
+        if self.config.__module__.startswith("vllm.transformers_utils.configs"):
+            from transformers.models.auto.configuration_auto import CONFIG_MAPPING
+
+            del CONFIG_MAPPING._extra_content[self.config.model_type]
+            self.config = AutoConfig.from_pretrained(
+                model_name,
+                trust_remote_code=trust_remote_code,
+            )
         self.device = self.get_default_device()
         self.dtype = dtype = _get_and_verify_dtype(
             self.model_name,
@@ -428,6 +440,7 @@ class HfRunner:
 
             self.model = SentenceTransformer(
                 model_name,
+                revision=revision,
                 device=self.device,
                 model_kwargs=model_kwargs,
                 trust_remote_code=trust_remote_code,
@@ -438,6 +451,7 @@ class HfRunner:
 
             self.model = CrossEncoder(
                 model_name,
+                revision=revision,
                 device=self.device,
                 automodel_args=model_kwargs,
                 trust_remote_code=trust_remote_code,
@@ -447,6 +461,7 @@ class HfRunner:
                 nn.Module,
                 auto_cls.from_pretrained(
                     model_name,
+                    revision=revision,
                     trust_remote_code=trust_remote_code,
                     **model_kwargs,
                 ),

tests/entrypoints/llm/test_generate.py

@@ -91,6 +91,12 @@ def test_multiple_priority(llm: LLM):
         outputs = llm.generate(PROMPTS, sampling_params=None, priority=[])
 
 
+def test_single_prompt_priority(llm: LLM):
+    # Single string prompts should be normalized to one request.
+    outputs = llm.generate(PROMPTS[0], sampling_params=None, priority=[0])
+    assert len(outputs) == 1
+
+
 def test_max_model_len():
     max_model_len = 20
     llm = LLM(

tests/entrypoints/openai/responses/test_function_call.py

@@ -249,40 +249,74 @@ async def test_function_calling_with_streaming_expected_arguments(
                 "additionalProperties": False,
             },
             "strict": True,
-        }
+        },
+        {
+            "type": "function",
+            "name": "get_time",
+            "description": "Get current local time for provided location.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "location": {"type": "string"},
+                },
+                "required": ["location"],
+                "additionalProperties": False,
+            },
+            "strict": True,
+        },
     ]
 
     stream_response = await client.responses.create(
         model=model_name,
-        input="Can you tell me what the current weather is in Berlin?",
+        input=(
+            "Use tools only. Call get_weather for Berlin and get_time for Tokyo. "
+            "Do not answer directly."
+        ),
         tools=tools,
         stream=True,
     )
 
-    tool_call_item = None
-    completed_event = None
+    tool_call_items = {}
+    arguments_done_events = {}
+    completed_events = {}
     async for event in stream_response:
         if (
             event.type == "response.output_item.added"
             and event.item.type == "function_call"
         ):
-            tool_call_item = event.item
-        elif event.type == "response.function_call_arguments.delta" and tool_call_item:
+            tool_call_items[event.output_index] = event.item
+        elif event.type == "response.function_call_arguments.delta":
+            tool_call_item = tool_call_items[event.output_index]
             tool_call_item.arguments += event.delta
+        elif event.type == "response.function_call_arguments.done":
+            arguments_done_events[event.output_index] = event
         elif (
             event.type == "response.output_item.done"
             and event.item.type == "function_call"
         ):
-            completed_event = event
-    assert tool_call_item is not None
-    assert tool_call_item.type == "function_call"
-    assert tool_call_item.name == "get_weather"
-    assert completed_event is not None
-    assert tool_call_item.arguments == completed_event.item.arguments
-    assert tool_call_item.name == completed_event.item.name
-    args = json.loads(tool_call_item.arguments)
-    assert "location" in args
-    assert args["location"] is not None
+            completed_events[event.output_index] = event
+    assert len(tool_call_items) >= 2
+    assert len(arguments_done_events) >= 2
+    assert len(completed_events) >= 2
+
+    tool_calls_by_name = {
+        event.item.name: (
+            tool_call_items[output_index],
+            arguments_done_events[output_index],
+            event.item,
+        )
+        for output_index, event in completed_events.items()
+    }
+    assert {"get_weather", "get_time"}.issubset(tool_calls_by_name)
+    for added_item, arguments_done_event, completed_item in tool_calls_by_name.values():
+        assert added_item.type == "function_call"
+        assert added_item.arguments == arguments_done_event.arguments
+        assert added_item.arguments == completed_item.arguments
+        assert added_item.name == arguments_done_event.name
+        assert added_item.name == completed_item.name
+        args = json.loads(added_item.arguments)
+        assert "location" in args
+        assert args["location"] is not None
 
 
 @pytest.mark.asyncio

tests/entrypoints/openai/responses/test_serving_responses.py

@@ -27,7 +27,9 @@ from openai.types.responses.tool import (
 import vllm.envs as envs
 from vllm.entrypoints.mcp.tool_server import ToolServer
 from vllm.entrypoints.openai.engine.protocol import (
+    DeltaFunctionCall,
     DeltaMessage,
+    DeltaToolCall,
     ErrorResponse,
     RequestResponseMetadata,
 )
@@ -928,3 +930,197 @@ class TestStreamingReasoningToContentTransition:
         ]
         assert len(item_done_events) == 1
         assert isinstance(item_done_events[0].item, ResponseReasoningItem)
+
+
+class TestAutoToolStreaming:
+    @staticmethod
+    async def _collect_events(delta_sequence: list[DeltaMessage]):
+        serving = _make_serving_instance_with_reasoning()
+        _mock_parser_with_reasoning(serving, delta_sequence)
+
+        contexts = [
+            _make_simple_context_with_output("chunk", [i])
+            for i in range(len(delta_sequence))
+        ]
+
+        async def result_generator():
+            for ctx in contexts:
+                yield ctx
+
+        request = ResponsesRequest(
+            input="hi",
+            tools=[
+                {
+                    "type": "function",
+                    "name": "get_weather",
+                    "description": "Get weather.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {"location": {"type": "string"}},
+                        "required": ["location"],
+                        "additionalProperties": False,
+                    },
+                }
+            ],
+            tool_choice="auto",
+            stream=True,
+        )
+        sampling_params = SamplingParams(max_tokens=64)
+        metadata = RequestResponseMetadata(request_id="req")
+        _identity_increment._counter = 0  # type: ignore
+
+        events = []
+        async for event in serving._process_simple_streaming_events(
+            request=request,
+            sampling_params=sampling_params,
+            result_generator=result_generator(),
+            context=SimpleContext(),
+            model_name="test-model",
+            tokenizer=MagicMock(),
+            request_metadata=metadata,
+            created_time=0,
+            _increment_sequence_number_and_return=_identity_increment,
+        ):
+            events.append(event)
+        return events
+
+    @pytest.mark.skip_global_cleanup
+    @pytest.mark.asyncio
+    async def test_auto_multi_tool_streaming_opens_one_item_per_tool(self, monkeypatch):
+        monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False)
+
+        delta_sequence = [
+            DeltaMessage(
+                tool_calls=[
+                    DeltaToolCall(
+                        id="call_vienna",
+                        type="function",
+                        index=0,
+                        function=DeltaFunctionCall(
+                            name="get_weather",
+                            arguments="",
+                        ),
+                    )
+                ]
+            ),
+            DeltaMessage(
+                tool_calls=[
+                    DeltaToolCall(
+                        index=0,
+                        function=DeltaFunctionCall(
+                            arguments='{"location":"Vienna"}',
+                        ),
+                    )
+                ]
+            ),
+            DeltaMessage(
+                tool_calls=[
+                    DeltaToolCall(
+                        id="call_berlin",
+                        type="function",
+                        index=1,
+                        function=DeltaFunctionCall(
+                            name="get_weather",
+                            arguments='{"location":"Berlin"}',
+                        ),
+                    )
+                ]
+            ),
+        ]
+        events = await self._collect_events(delta_sequence)
+
+        function_items = [
+            event
+            for event in events
+            if event.type == "response.output_item.added"
+            and getattr(event.item, "type", None) == "function_call"
+        ]
+        assert len(function_items) == 2
+        assert [event.item.name for event in function_items] == [
+            "get_weather",
+            "get_weather",
+        ]
+        assert [event.output_index for event in function_items] == [0, 1]
+
+        argument_deltas = [
+            event.delta
+            for event in events
+            if event.type == "response.function_call_arguments.delta"
+        ]
+        assert argument_deltas == [
+            '{"location":"Vienna"}',
+            '{"location":"Berlin"}',
+        ]
+
+        argument_done = [
+            event
+            for event in events
+            if event.type == "response.function_call_arguments.done"
+        ]
+        assert [event.arguments for event in argument_done] == [
+            '{"location":"Vienna"}',
+            '{"location":"Berlin"}',
+        ]
+        assert [event.output_index for event in argument_done] == [0, 1]
+
+        function_done = [
+            event
+            for event in events
+            if event.type == "response.output_item.done"
+            and getattr(event.item, "type", None) == "function_call"
+        ]
+        assert [event.item.arguments for event in function_done] == [
+            '{"location":"Vienna"}',
+            '{"location":"Berlin"}',
+        ]
+        assert [event.output_index for event in function_done] == [0, 1]
+
+    @pytest.mark.skip_global_cleanup
+    @pytest.mark.asyncio
+    async def test_auto_tool_choice_first_delta_tool_call_does_not_duplicate_item(
+        self, monkeypatch
+    ):
+        monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False)
+
+        delta_sequence = [
+            DeltaMessage(
+                tool_calls=[
+                    DeltaToolCall(
+                        id="call_test",
+                        type="function",
+                        index=0,
+                        function=DeltaFunctionCall(
+                            name="get_weather",
+                            arguments="",
+                        ),
+                    )
+                ]
+            ),
+            DeltaMessage(
+                tool_calls=[
+                    DeltaToolCall(
+                        index=0,
+                        function=DeltaFunctionCall(
+                            arguments='{"location":"Berlin"}',
+                        ),
+                    )
+                ]
+            ),
+        ]
+        events = await self._collect_events(delta_sequence)
+
+        function_items = [
+            event
+            for event in events
+            if event.type == "response.output_item.added"
+            and getattr(event.item, "type", None) == "function_call"
+        ]
+        assert len(function_items) == 1
+        assert function_items[0].item.name == "get_weather"
+
+        argument_deltas = [
+            event.delta
+            for event in events
+            if event.type == "response.function_call_arguments.delta"
+        ]
+        assert "".join(argument_deltas) == '{"location":"Berlin"}'

tests/entrypoints/openai/test_tool_calls_serialization.py

@@ -0,0 +1,150 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for tool_calls Iterable → list materialisation.
+
+Regression tests for https://github.com/vllm-project/vllm/issues/34792.
+
+Setting VLLM_LOGGING_LEVEL=debug caused tool calling to break for Mistral
+models because:
+  1. The OpenAI Python SDK types tool_calls as Iterable[...] in
+     ChatCompletionAssistantMessageParam.
+  2. Pydantic v2, when validating from Python objects (not from raw JSON),
+     wraps Iterable fields in a one-shot lazy iterator.
+  3. Debug logging called model_dump_json() which consumed that iterator.
+  4. The Mistral tokenizer then saw empty tool_calls and raised
+     "ValueError: Unexpected tool call id ...".
+"""
+
+import pytest
+
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+
+
+def _make_tool_call(tc_id: str, name: str, args: str) -> dict:
+    return {
+        "id": tc_id,
+        "type": "function",
+        "function": {"name": name, "arguments": args},
+    }
+
+
+def _make_request(messages: list) -> ChatCompletionRequest:
+    return ChatCompletionRequest(
+        model="test-model",
+        messages=messages,
+    )
+
+
+def test_tool_calls_list_preserved_after_model_dump():
+    """tool_calls in assistant messages must be readable after model_dump_json.
+
+    When the request is built from Python dicts (as in the Anthropic → OpenAI
+    conversion path), Pydantic v2 previously wrapped the Iterable tool_calls
+    in a one-shot iterator.  model_dump_json() consumed it, leaving subsequent
+    readers (e.g. the Mistral tokenizer) with an empty sequence.
+    """
+    tool_call = _make_tool_call("call_abc123", "get_weather", '{"city": "Paris"}')
+    messages = [
+        {"role": "user", "content": "What is the weather in Paris?"},
+        {"role": "assistant", "content": None, "tool_calls": [tool_call]},
+        {
+            "role": "tool",
+            "tool_call_id": "call_abc123",
+            "content": '{"temperature": 20}',
+        },
+    ]
+
+    req = _make_request(messages)
+
+    # Simulate debug logging: serialize the model (this was the trigger)
+    _ = req.model_dump_json()
+
+    # The assistant message must still have accessible tool_calls afterwards
+    assistant_msg = req.messages[1]
+    assert isinstance(assistant_msg, dict)
+    tool_calls = assistant_msg.get("tool_calls")
+    assert tool_calls is not None, "tool_calls must not be None after model_dump_json"
+    assert isinstance(tool_calls, list), "tool_calls must be a list"
+    assert len(tool_calls) > 0, "tool_calls must not be empty after model_dump_json"
+
+
+def test_tool_calls_from_generator_are_materialised():
+    """tool_calls passed as a generator must be converted to list on validation."""
+    tool_call = _make_tool_call("call_gen1", "search", '{"query": "vllm"}')
+
+    def tool_calls_gen():
+        yield tool_call
+
+    messages = [
+        {"role": "user", "content": "Search for vllm"},
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": tool_calls_gen(),  # one-shot generator
+        },
+    ]
+
+    req = _make_request(messages)
+    assistant_msg = req.messages[1]
+    assert isinstance(assistant_msg, dict)
+
+    # Iterate twice — must not raise or return empty on second pass
+    tool_calls_first = list(assistant_msg.get("tool_calls", []))
+    tool_calls_second = list(assistant_msg.get("tool_calls", []))
+
+    assert len(tool_calls_first) == 1, "First read must return the tool call"
+    assert len(tool_calls_second) == 1, "Second read must also return the tool call"
+
+
+def test_tool_calls_list_passthrough():
+    """tool_calls already provided as a list must remain a list."""
+    tool_call = _make_tool_call("call_list1", "calculate", '{"expr": "2+2"}')
+    messages = [
+        {"role": "user", "content": "Calculate 2+2"},
+        {"role": "assistant", "content": None, "tool_calls": [tool_call]},
+    ]
+
+    req = _make_request(messages)
+    assistant_msg = req.messages[1]
+    assert isinstance(assistant_msg, dict)
+    assert isinstance(assistant_msg.get("tool_calls"), list)
+
+
+def test_messages_without_tool_calls_unaffected():
+    """Messages without tool_calls must be handled correctly."""
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "Hello!"},
+        {"role": "assistant", "content": "Hi there!"},
+    ]
+
+    req = _make_request(messages)
+    # None of the messages should have tool_calls injected
+    for msg in req.messages:
+        assert isinstance(msg, dict)
+        assert msg.get("tool_calls") is None or msg.get("tool_calls") == []
+
+
+@pytest.mark.parametrize("num_tool_calls", [1, 3])
+def test_multiple_tool_calls_materialised(num_tool_calls: int):
+    """Multiple tool calls in a single message are all preserved."""
+    tool_calls = [
+        _make_tool_call(f"call_{i}", f"func_{i}", f'{{"arg": {i}}}')
+        for i in range(num_tool_calls)
+    ]
+    messages = [
+        {"role": "user", "content": "Do things"},
+        {"role": "assistant", "content": None, "tool_calls": iter(tool_calls)},
+    ]
+
+    req = _make_request(messages)
+    assistant_msg = req.messages[1]
+    assert isinstance(assistant_msg, dict)
+
+    result_tool_calls = assistant_msg.get("tool_calls")
+    assert isinstance(result_tool_calls, list)
+    assert len(result_tool_calls) == num_tool_calls
+
+    # Verify after model_dump_json too
+    _ = req.model_dump_json()
+    assert len(assistant_msg.get("tool_calls", [])) == num_tool_calls

tests/entrypoints/pooling/embed/test_io_processor.py

@@ -4,6 +4,7 @@
 
 import pytest
 
+from vllm import PoolingParams
 from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
 from vllm.entrypoints.pooling.embed.protocol import (
     CohereEmbedContent,
@@ -218,6 +219,7 @@ class TestPreProcessCohereOnline:
     def _make_context(**request_kwargs) -> PoolingServeContext[CohereEmbedRequest]:
         return PoolingServeContext(
             request=CohereEmbedRequest(model="test", **request_kwargs),
+            pooling_params=PoolingParams(),
             model_name="test",
             request_id="embd-test",
         )
@@ -233,13 +235,13 @@ class TestPreProcessCohereOnline:
         ctx = self._make_context(texts=["hello"])
         calls: list[tuple[str, object]] = []
 
-        def preprocess_completion(request, prompt_input, prompt_embeds):
+        def preprocess_cmpl_online(request, prompt_input, prompt_embeds):
             calls.append(("completion", prompt_input))
             return ["completion"]
 
         handler._get_task_instruction_prefix = lambda _input_type: None
         handler._has_chat_template = lambda: False
-        handler._preprocess_completion_online = preprocess_completion
+        handler._preprocess_cmpl_online = preprocess_cmpl_online
         handler._batch_render_chat = lambda *_args, **_kwargs: (
             pytest.fail("text-only request should not require chat rendering")
         )
@@ -254,7 +256,7 @@ class TestPreProcessCohereOnline:
         ctx = self._make_context(texts=["hello"], input_type="query")
         calls: list[tuple[str, object]] = []
 
-        def preprocess_completion(request, prompt_input, prompt_embeds):
+        def preprocess_cmpl(request, prompt_input, prompt_embeds):
             calls.append(("completion", prompt_input))
             return ["fallback"]
 
@@ -263,7 +265,7 @@ class TestPreProcessCohereOnline:
         handler._batch_render_chat = lambda *_args, **_kwargs: (
             pytest.fail("chat rendering should be skipped without a template")
         )
-        handler._preprocess_completion_online = preprocess_completion
+        handler._preprocess_cmpl_online = preprocess_cmpl
 
         handler._pre_process_cohere_online(ctx)
 
@@ -297,7 +299,7 @@ class TestPreProcessCohereOnline:
         handler._get_task_instruction_prefix = lambda _input_type: "query: "
         handler._has_chat_template = lambda: True
         handler._batch_render_chat = batch_render_chat
-        handler._preprocess_completion_online = lambda *_args, **_kwargs: (
+        handler._preprocess_cmpl_online = lambda *_args, **_kwargs: (
             pytest.fail("completion path should be skipped when a template exists")
         )
 

tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-BF16.yaml

@@ -8,4 +8,5 @@ server_args: >-
   --max-model-len 4096
   --tensor-parallel-size 8
   --enable-expert-parallel
+  --mamba-backend flashinfer
   --speculative-config '{"method":"mtp","num_speculative_tokens":5}'

tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4.yaml

@@ -8,4 +8,5 @@ server_args: >-
   --max-model-len 4096
   --tensor-parallel-size 2
   --enable-expert-parallel
+  --mamba-backend flashinfer
   --speculative-config '{"method":"mtp","num_speculative_tokens":5}'

tests/evals/gsm8k/configs/Qwen3-4B-TQ-k3v4nc.yaml

@@ -0,0 +1,5 @@
+model_name: "Qwen/Qwen3-4B"
+accuracy_threshold: 0.78
+num_questions: 1319
+num_fewshot: 5
+server_args: "--kv-cache-dtype turboquant_k3v4_nc --enforce-eager --max-model-len 4096"

tests/evals/gsm8k/configs/Qwen3-4B-TQ-k8v4.yaml

@@ -0,0 +1,5 @@
+model_name: "Qwen/Qwen3-4B"
+accuracy_threshold: 0.80
+num_questions: 1319
+num_fewshot: 5
+server_args: "--kv-cache-dtype turboquant_k8v4 --enforce-eager --max-model-len 4096"

tests/evals/gsm8k/configs/Qwen3-4B-TQ-t3nc.yaml

@@ -0,0 +1,5 @@
+model_name: "Qwen/Qwen3-4B"
+accuracy_threshold: 0.75
+num_questions: 1319
+num_fewshot: 5
+server_args: "--kv-cache-dtype turboquant_3bit_nc --enforce-eager --max-model-len 4096"

tests/evals/gsm8k/configs/Qwen3-4B-TQ-t4nc.yaml

@@ -0,0 +1,5 @@
+model_name: "Qwen/Qwen3-4B"
+accuracy_threshold: 0.80
+num_questions: 1319
+num_fewshot: 5
+server_args: "--kv-cache-dtype turboquant_4bit_nc --enforce-eager --max-model-len 4096"

tests/evals/gsm8k/configs/models-turboquant.txt

@@ -0,0 +1,4 @@
+Qwen3-4B-TQ-k8v4.yaml
+Qwen3-4B-TQ-t4nc.yaml
+Qwen3-4B-TQ-k3v4nc.yaml
+Qwen3-4B-TQ-t3nc.yaml

tests/kernels/core/test_cpu_activation.py

@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from tests.kernels.allclose_default import get_default_atol, get_default_rtol
+from tests.kernels.utils import opcheck
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.utils.torch_utils import set_random_seed
+
+if not current_platform.is_cpu():
+    pytest.skip("skipping CPU-only tests", allow_module_level=True)
+
+from vllm.model_executor.layers.activation import (
+    GELU,
+    FastGELU,
+    GeluAndMul,
+    NewGELU,
+    QuickGELU,
+    SiluAndMul,
+)
+
+DTYPES = [torch.bfloat16, torch.float32]
+NUM_TOKENS = [7, 83]
+D = [512, 2048]
+SEEDS = [0]
+
+
+@pytest.mark.parametrize(
+    ("activation_cls", "fn"),
+    [
+        (SiluAndMul, torch.ops._C.silu_and_mul),
+        (GeluAndMul, torch.ops._C.gelu_and_mul),
+        (GeluAndMul, torch.ops._C.gelu_tanh_and_mul),
+    ],
+)
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
+@pytest.mark.parametrize("d", D)
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("seed", SEEDS)
+@torch.inference_mode()
+def test_cpu_act_and_mul(
+    default_vllm_config,
+    activation_cls: type[torch.nn.Module],
+    fn: object,
+    num_tokens: int,
+    d: int,
+    dtype: torch.dtype,
+    seed: int,
+) -> None:
+    set_random_seed(seed)
+    x = torch.randn(num_tokens, 2 * d, dtype=dtype)
+
+    layer = activation_cls()
+    out = layer(x)
+    ref_out = layer.forward_native(x)
+
+    torch.testing.assert_close(
+        out, ref_out, atol=get_default_atol(out), rtol=get_default_rtol(out)
+    )
+
+    output_shape = x.shape[:-1] + (x.shape[-1] // 2,)
+    raw_out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
+    opcheck(fn, (raw_out, x))
+
+
+@pytest.mark.parametrize(
+    ("activation_cls", "fn", "op_args"),
+    [
+        (NewGELU, torch.ops._C.gelu_new, ()),
+        (FastGELU, torch.ops._C.gelu_fast, ()),
+        (QuickGELU, torch.ops._C.gelu_quick, ()),
+        pytest.param(
+            GELU,
+            getattr(torch.ops._C, "activation_lut_bf16", None),
+            ("gelu",),
+            marks=pytest.mark.skipif(
+                current_platform.get_cpu_architecture() != CpuArchEnum.ARM,
+                reason="activation_lut_bf16 is only built on Arm CPU",
+            ),
+        ),
+    ],
+)
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
+@pytest.mark.parametrize("d", D)
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("seed", SEEDS)
+@torch.inference_mode()
+def test_cpu_unary_activation(
+    default_vllm_config,
+    activation_cls: type[torch.nn.Module],
+    fn: object,
+    op_args: tuple[str, ...],
+    num_tokens: int,
+    d: int,
+    dtype: torch.dtype,
+    seed: int,
+) -> None:
+    set_random_seed(seed)
+    x = torch.randn(num_tokens, d, dtype=dtype)
+    layer = activation_cls()
+    out = layer(x)
+    ref_out = layer.forward_native(x)
+    torch.testing.assert_close(
+        out, ref_out, atol=get_default_atol(out), rtol=get_default_rtol(out)
+    )
+    # gelu with activation_lut_bf16 only makes sense for BF16
+    if not (activation_cls is GELU and dtype != torch.bfloat16):
+        raw_out = torch.empty_like(x)
+        opcheck(fn, (raw_out, x, *op_args))

tests/kernels/helion/test_register.py

@@ -36,9 +36,11 @@ from vllm.kernels.helion.register import (
 )
 
 if _HOP_AVAILABLE:
+    from helion._compat import supports_torch_compile_fusion
     from helion._compiler._dynamo.higher_order_ops import (
         helion_kernel_wrapper_mutation,
     )
+    from torch._inductor.utils import run_and_get_code
 
 
 def _add_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
@@ -1003,3 +1005,49 @@ class TestTorchCompileHOP:
             "Compiled execution result doesn't match eager execution. "
             f"Max difference: {torch.max(torch.abs(compiled_result - eager_result))}"
         )
+
+    @pytest.mark.skipif(
+        not (_HOP_AVAILABLE and supports_torch_compile_fusion()),
+        reason="Requires PyTorch with Helion inductor fusion support",
+    )
+    def test_inductor_backend_compiles_helion_hop(self):
+        """Test torch.compile with inductor backend and Helion fusion enabled."""
+
+        configs = {"default": helion.Config(block_sizes=[4, 4])}
+
+        with dummy_kernel_registry(configs=configs) as register:
+            add_helion_kernel = register(
+                op_name="test_inductor_add_kernel",
+                config_picker=lambda args, keys: "default",
+                helion_settings=helion.Settings(
+                    torch_compile_fusion=True, static_shapes=False
+                ),
+            )(_add_kernel)
+
+        def f(x, y):
+            x = x * 2.0
+            y = y + 1.0
+            out = add_helion_kernel(x, y)
+            return out.relu()
+
+        torch._dynamo.reset()
+        compiled_f = torch.compile(f, backend="inductor", fullgraph=True)
+
+        x = torch.randn(4, 4, device="cuda")
+        y = torch.randn(4, 4, device="cuda")
+
+        compiled_result, source_codes = run_and_get_code(compiled_f, x, y)
+        eager_result = f(x, y)
+
+        assert torch.allclose(compiled_result, eager_result, atol=1e-5, rtol=1e-5), (
+            "Inductor-compiled result doesn't match eager execution. "
+            f"Max difference: {torch.max(torch.abs(compiled_result - eager_result))}"
+        )
+
+        # With fusion enabled, prologue/epilogue ops should be fused into
+        # a single triton kernel rather than generating separate kernels.
+        kernel_count = sum(code.count("@triton.jit") for code in source_codes)
+        assert kernel_count == 1, (
+            f"Expected 1 fused triton kernel, got {kernel_count}. "
+            "Prologue/epilogue ops were not fused into the Helion kernel."
+        )

tests/kernels/mamba/test_ssu_dispatch.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.config.mamba import MambaBackendEnum, MambaConfig
+from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
+    FlashInferSSUBackend,
+    TritonSSUBackend,
+    get_mamba_ssu_backend,
+    initialize_mamba_ssu_backend,
+    selective_state_update,
+)
+from vllm.utils.torch_utils import set_random_seed
+
+try:
+    import flashinfer.mamba  # noqa: F401
+
+    HAS_FLASHINFER = True
+except ImportError:
+    HAS_FLASHINFER = False
+
+
+def test_default_backend_is_triton():
+    initialize_mamba_ssu_backend(MambaConfig())
+    backend = get_mamba_ssu_backend()
+    assert isinstance(backend, TritonSSUBackend)
+    assert backend.name == "triton"
+
+
+def test_explicit_triton_backend():
+    initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
+    backend = get_mamba_ssu_backend()
+    assert isinstance(backend, TritonSSUBackend)
+
+
+@pytest.mark.skipif(not HAS_FLASHINFER, reason="flashinfer not installed")
+def test_flashinfer_backend_init():
+    initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.FLASHINFER))
+    backend = get_mamba_ssu_backend()
+    assert isinstance(backend, FlashInferSSUBackend)
+    assert backend.name == "flashinfer"
+
+
+def test_uninitialized_backend_raises():
+    import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
+
+    old = mod._mamba_ssu_backend
+    mod._mamba_ssu_backend = None
+    with pytest.raises(RuntimeError, match="not been initialized"):
+        get_mamba_ssu_backend()
+    mod._mamba_ssu_backend = old
+
+
+@pytest.mark.skipif(HAS_FLASHINFER, reason="flashinfer is installed")
+def test_flashinfer_import_error():
+    with pytest.raises(ImportError, match="FlashInfer is required"):
+        FlashInferSSUBackend(MambaConfig())
+
+
+def test_triton_basic_call():
+    set_random_seed(0)
+    initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
+    device = "cuda"
+    batch_size = 2
+    dim = 64
+    dstate = 16
+
+    state = torch.randn(batch_size, dim, dstate, device=device)
+    x = torch.randn(batch_size, dim, device=device)
+    out = torch.empty_like(x)
+    dt = torch.randn(batch_size, dim, device=device)
+    dt_bias = torch.rand(dim, device=device) - 4.0
+    A = -torch.rand(dim, dstate, device=device)
+    B = torch.randn(batch_size, dstate, device=device)
+    C = torch.randn(batch_size, dstate, device=device)
+    D = torch.randn(dim, device=device)
+
+    selective_state_update(
+        state,
+        x,
+        dt,
+        A,
+        B,
+        C,
+        D=D,
+        dt_bias=dt_bias,
+        dt_softplus=True,
+        out=out,
+    )
+    assert not torch.isnan(out).any()

tests/kernels/moe/modular_kernel_tools/common.py

@@ -46,6 +46,7 @@ from vllm.utils.import_utils import (
     has_deep_gemm,
     has_mori,
 )
+from vllm.utils.math_utils import next_power_of_2
 
 from .mk_objects import (
     TestMoEQuantConfig,
@@ -604,13 +605,6 @@ def make_modular_kernel(
     vllm_config: VllmConfig,
     quant_config: FusedMoEQuantConfig,
 ) -> mk.FusedMoEKernel:
-    def next_power_of_2(x):
-        import math
-
-        if x == 0:
-            return 1
-        return 2 ** math.ceil(math.log2(x))
-
     # make moe config
     moe_parallel_config: FusedMoEParallelConfig = FusedMoEParallelConfig.make(
         tp_size_=get_tensor_model_parallel_world_size(),

tests/kernels/moe/modular_kernel_tools/parallel_utils.py

@@ -126,7 +126,7 @@ def parallel_launch_with_config(
     world_size: int,
     worker: Callable[Concatenate[ProcessGroupInfo, VllmConfig, Any, P], None],
     vllm_config: VllmConfig,
-    env_dict: dict[Any, Any],
+    env_dict: dict[Any, Any] | None,
     *args: P.args,
     **kwargs: P.kwargs,
 ) -> None:

tests/kernels/moe/test_deepep_deepgemm_moe.py

@@ -29,6 +29,7 @@ from vllm.utils.deep_gemm import (
     is_deep_gemm_supported,
 )
 from vllm.utils.import_utils import has_deep_ep, has_deep_gemm
+from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.worker.workspace import init_workspace_manager
 
@@ -84,14 +85,6 @@ def with_dp_metadata(M: int, world_size: int):
         yield
 
 
-def next_power_of_2(x):
-    import math
-
-    if x == 0:
-        return 1
-    return 2 ** math.ceil(math.log2(x))
-
-
 def make_block_quant_fp8_weights(
     e: int,
     n: int,

tests/kernels/moe/test_flashinfer.py

@@ -32,6 +32,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
 from vllm.model_executor.layers.quantization.utils.fp8_utils import input_to_float8
 from vllm.model_executor.models.llama4 import Llama4MoE
 from vllm.platforms import current_platform
+from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
 
 try:
@@ -174,6 +175,7 @@ class TestData:
             routing_method=layer.routing_method_type,
             activation=activation,
             device=w13_quantized.device,
+            max_num_tokens=next_power_of_2(m),
         )
 
         return TestData(
@@ -348,6 +350,7 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(
             in_dtype=torch.bfloat16,
             is_act_and_mul=activation.is_gated,
             routing_method=RoutingMethodType.TopK,
+            max_num_tokens=next_power_of_2(m),
         )
 
         kernel = mk.FusedMoEKernel(

tests/kernels/moe/test_flashinfer_moe.py

@@ -29,6 +29,7 @@ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
 from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
 from vllm.platforms import current_platform
 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
+from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
 
 if not has_flashinfer_cutlass_fused_moe() or not current_platform.has_device_capability(
@@ -105,6 +106,7 @@ def test_flashinfer_fp4_moe_no_graph(
             in_dtype=dtype,
             is_act_and_mul=is_gated_act,
             routing_method=RoutingMethodType.TopK,
+            max_num_tokens=next_power_of_2(m),
         )
 
         flashinfer_experts = FusedMoEKernel(