Merge branch 'main' into wentao-optimize-pooling-by-ragged-tensor

Signed-off-by: yewentao256 <zhyanwentao@126.com>

.buildkite/image_build/image_build_torch_nightly.sh

@@ -46,7 +46,7 @@ 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.2"
+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"
 

.buildkite/release-pipeline.yaml

@@ -37,7 +37,7 @@ steps:
         agents:
           queue: arm64_cpu_queue_release
         commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04  --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04  --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
           - "mkdir artifacts"
           - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
           - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
@@ -76,7 +76,7 @@ steps:
         agents:
           queue: cpu_queue_release
         commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
           - "mkdir artifacts"
           - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
           - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
@@ -121,7 +121,7 @@ steps:
           queue: cpu_queue_release
         commands:
           - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
           - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
           # re-tag to default image tag and push, just in case arm64 build fails
           - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
@@ -134,7 +134,7 @@ steps:
           queue: arm64_cpu_queue_release
         commands:
           - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
           - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
 
       - label: "Build release image - x86_64 - CUDA 12.9"
@@ -167,7 +167,7 @@ steps:
           queue: cpu_queue_release
         commands:
           - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
           - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
           - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
           - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
@@ -179,7 +179,7 @@ steps:
           queue: arm64_cpu_queue_release
         commands:
           - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
           - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
 
       - label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"

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

@@ -25,100 +25,22 @@ export PYTHONPATH=".."
 ###############################################################################
 
 cleanup_docker() {
-  # Share the same lock with image pull to avoid cleanup/pull races on one node.
-  local docker_lock="/tmp/docker-pull.lock"
-  exec 9>"$docker_lock"
-  flock 9
-
   docker_root=$(docker info -f '{{.DockerRootDir}}')
   if [ -z "$docker_root" ]; then
     echo "Failed to determine Docker root directory." >&2
-    flock -u 9
-    return 1
+    exit 1
   fi
   echo "Docker root directory: $docker_root"
 
   disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
   threshold=70
   if [ "$disk_usage" -gt "$threshold" ]; then
-    echo "Disk usage is above $threshold%. Running aggressive CI image cleanup..."
-    cleanup_old_ci_images "${REGISTRY}/${REPO}" "${image_name}" "${DOCKER_IMAGE_CLEANUP_HOURS:-72}" 1
+    echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
+    docker image prune -f
+    docker volume prune -f && docker system prune --force --filter "until=72h" --all
+    echo "Docker images and volumes cleanup completed."
   else
-    echo "Disk usage is below $threshold%. Checking old CI images anyway."
-    cleanup_old_ci_images "${REGISTRY}/${REPO}" "${image_name}" "${DOCKER_IMAGE_CLEANUP_HOURS:-72}" 0
-  fi
-  echo "Old CI image cleanup completed."
-
-  flock -u 9
-}
-
-cleanup_old_ci_images() {
-  local repo_prefix="$1"
-  local current_image_ref="$2"
-  local ttl_hours="$3"
-  local aggressive_cleanup="$4"
-
-  if [[ -z "$repo_prefix" || "$repo_prefix" == "/" ]]; then
-    echo "Skip old-image cleanup: invalid repo prefix '${repo_prefix}'"
-    return 0
-  fi
-
-  if ! [[ "$ttl_hours" =~ ^[0-9]+$ ]]; then
-    echo "Invalid DOCKER_IMAGE_CLEANUP_HOURS='${ttl_hours}', fallback to 72"
-    ttl_hours=72
-  fi
-
-  local now_epoch cutoff_epoch
-  now_epoch=$(date +%s)
-  cutoff_epoch=$((now_epoch - ttl_hours * 3600))
-
-  local -a used_image_ids
-  mapfile -t used_image_ids < <(docker ps -aq | xargs -r docker inspect --format '{{.Image}}' | sort -u)
-
-  local removed_count=0
-  local examined_count=0
-  declare -A seen_ids=()
-
-  while read -r image_ref image_id; do
-    [[ -z "$image_ref" || -z "$image_id" ]] && continue
-    ((examined_count++))
-
-    # Keep the image this job is going to use.
-    if [[ "$image_ref" == "$current_image_ref" ]]; then
-      continue
-    fi
-
-    # Avoid duplicate deletes when multiple tags point to same image id.
-    if [[ -n "${seen_ids[$image_id]:-}" ]]; then
-      continue
-    fi
-    seen_ids[$image_id]=1
-
-    # Never delete images that are used by any container on this node.
-    if printf '%s\n' "${used_image_ids[@]}" | grep -qx "$image_id"; then
-      continue
-    fi
-
-    local created created_epoch
-    created=$(docker image inspect -f '{{.Created}}' "$image_id" 2>/dev/null || true)
-    [[ -z "$created" ]] && continue
-    created_epoch=$(date -d "$created" +%s 2>/dev/null || true)
-    [[ -z "$created_epoch" ]] && continue
-
-    if (( created_epoch < cutoff_epoch )) || [[ "$aggressive_cleanup" == "1" ]]; then
-      if docker image rm -f "$image_id" >/dev/null 2>&1; then
-        ((removed_count++))
-      fi
-    fi
-  done < <(docker image ls --no-trunc "$repo_prefix" --format '{{.Repository}}:{{.Tag}} {{.ID}}')
-
-  # Also trim old dangling layers; this is safe and does not remove referenced images.
-  docker image prune -f --filter "until=${ttl_hours}h" >/dev/null 2>&1 || true
-
-  if [[ "$aggressive_cleanup" == "1" ]]; then
-    echo "Examined ${examined_count} images under ${repo_prefix}, removed ${removed_count} unused images under disk pressure."
-  else
-    echo "Examined ${examined_count} images under ${repo_prefix}, removed ${removed_count} old images (>${ttl_hours}h)."
+    echo "Disk usage is below $threshold%. No cleanup needed."
   fi
 }
 
@@ -343,6 +265,8 @@ fi
 
 remove_docker_container() {
   docker rm -f "${container_name}" || true
+  docker image rm -f "${image_name}" || true
+  docker system prune -f || true
 }
 trap remove_docker_container EXIT
 

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

@@ -12,7 +12,9 @@ docker build -t "${image_name}" -f docker/Dockerfile.xpu .
 
 # Setup cleanup
 remove_docker_container() {
-  docker rm -f "${container_name}" || true
+  docker rm -f "${container_name}" || true;
+  docker image rm -f "${image_name}" || true;
+  docker system prune -f || true;
 }
 trap remove_docker_container EXIT
 

.buildkite/test_areas/models_multimodal.yaml

@@ -28,7 +28,6 @@ steps:
     - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
     - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
-    - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
   mirror:
     amd:
       device: mi325_1

CMakeLists.txt

@@ -34,7 +34,7 @@ install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS)
 # Supported python versions.  These versions will be searched in order, the
 # first match will be selected.  These should be kept in sync with setup.py.
 #
-set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13" "3.14")
+set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13")
 
 # Supported AMD GPU architectures.
 set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")
@@ -310,9 +310,7 @@ set(VLLM_EXT_SRC
   "csrc/torch_bindings.cpp")
 
 if(VLLM_GPU_LANG STREQUAL "CUDA")
-  list(APPEND VLLM_EXT_SRC
-    "csrc/minimax_reduce_rms_kernel.cu"
-    "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
+  list(APPEND VLLM_EXT_SRC "csrc/minimax_reduce_rms_kernel.cu")
 
   SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
 
@@ -1053,8 +1051,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   list(APPEND VLLM_MOE_EXT_SRC
     "csrc/moe/moe_wna16.cu"
     "csrc/moe/grouped_topk_kernels.cu"
-    "csrc/moe/router_gemm.cu"
-    "csrc/moe/topk_softplus_sqrt_kernels.cu")
+    "csrc/moe/router_gemm.cu")
 endif()
 
 if(VLLM_GPU_LANG STREQUAL "CUDA")

README.md

@@ -14,7 +14,7 @@ Easy, fast, and cheap LLM serving for everyone
 | <a href="https://docs.vllm.ai"><b>Documentation</b></a> | <a href="https://blog.vllm.ai/"><b>Blog</b></a> | <a href="https://arxiv.org/abs/2309.06180"><b>Paper</b></a> | <a href="https://x.com/vllm_project"><b>Twitter/X</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> |
 </p>
 
-🔥 We have built a vLLM website to help you get started with vLLM. Please visit [vllm.ai](https://vllm.ai) to learn more.
+🔥 We have built a vllm website to help you get started with vllm. Please visit [vllm.ai](https://vllm.ai) to learn more.
 For events, please visit [vllm.ai/events](https://vllm.ai/events) to join us.
 
 ---
@@ -50,7 +50,7 @@ vLLM is flexible and easy to use with:
 - Efficient multi-LoRA support for dense and MoE layers
 - Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
 
-vLLM seamlessly supports 200+ model architectures on Hugging Face, including:
+vLLM seamlessly supports 200+ model architectures on HuggingFace, including:
 
 - Decoder-only LLMs (e.g., Llama, Qwen, Gemma)
 - Mixture-of-Expert LLMs (e.g., Mixtral, DeepSeek-V3, Qwen-MoE, GPT-OSS)

benchmarks/attention_benchmarks/mla_runner.py

@@ -404,7 +404,6 @@ def _build_attention_metadata(
         query_start_loc=q_start_gpu,
         query_start_loc_cpu=q_start_cpu,
         seq_lens=seq_lens_gpu,
-        seq_lens_cpu_upper_bound=seq_lens_cpu,
         _seq_lens_cpu=seq_lens_cpu,
         _num_computed_tokens_cpu=num_computed_tokens_cpu,
         slot_mapping=slot_mapping,

benchmarks/kernels/benchmark_cutlass_moe_fp8.py

@@ -16,7 +16,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
     maybe_make_prepare_finalize,
 )
 from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
-from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import CutlassExpertsFp8
+from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
 from vllm.platforms import current_platform
 from vllm.utils.argparse_utils import FlexibleArgumentParser

benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py

@@ -22,7 +22,7 @@ from vllm.model_executor.layers.fused_moe.config import (
     fp8_w8a8_moe_quant_config,
     nvfp4_moe_quant_config,
 )
-from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
+from vllm.model_executor.layers.fused_moe.cutlass_moe import (
     CutlassExpertsFp4,
 )
 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk

benchmarks/kernels/benchmark_grouped_gemm_cutlass.py

@@ -13,7 +13,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
     maybe_make_prepare_finalize,
 )
 from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
-from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import CutlassExpertsFp8
+from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
 from vllm.model_executor.layers.fused_moe.fused_moe import (
     fused_experts,
     fused_topk,

cmake/external_projects/deepgemm.cmake

@@ -20,7 +20,7 @@ else()
   FetchContent_Declare(
     deepgemm
     GIT_REPOSITORY https://github.com/deepseek-ai/DeepGEMM.git
-    GIT_TAG 891d57b4db1071624b5c8fa0d1e51cb317fa709f
+    GIT_TAG 477618cd51baffca09c4b0b87e97c03fe827ef03
     GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
     GIT_PROGRESS TRUE
     CONFIGURE_COMMAND ""
@@ -120,11 +120,6 @@ if(DEEPGEMM_ARCHS)
     COMPONENT _deep_gemm_C
     FILES_MATCHING PATTERN "*.py")
 
-  install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/mega/"
-    DESTINATION vllm/third_party/deep_gemm/mega
-    COMPONENT _deep_gemm_C
-    FILES_MATCHING PATTERN "*.py")
-
   # Generate envs.py (normally generated by DeepGEMM's setup.py build step)
   file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/deep_gemm_envs.py"
     "# Pre-installed environment variables\npersistent_envs = dict()\n")

cmake/external_projects/flashmla.cmake

@@ -19,7 +19,7 @@ else()
   FetchContent_Declare(
         flashmla
         GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
-        GIT_TAG a6ec2ba7bd0a7dff98b3f4d3e6b52b159c48d78b
+        GIT_TAG 692917b1cda61b93ac9ee2d846ec54e75afe87b1
         GIT_PROGRESS TRUE
         CONFIGURE_COMMAND ""
         BUILD_COMMAND ""

csrc/activation_kernels.cu

@@ -11,74 +11,29 @@
 namespace vllm {
 
 template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
-          bool act_first, bool HAS_CLAMP>
+          bool act_first>
 __device__ __forceinline__ scalar_t compute(const scalar_t& x,
-                                            const scalar_t& y,
-                                            const float limit) {
-  if constexpr (act_first) {
-    scalar_t gate = x;
-    scalar_t up = y;
-    if constexpr (HAS_CLAMP) {
-      gate = (scalar_t)fminf((float)gate, limit);
-      up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
-    }
-    return ACT_FN(gate) * up;
-  } else {
-    scalar_t gate = x;
-    scalar_t up = y;
-    if constexpr (HAS_CLAMP) {
-      gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
-      up = (scalar_t)fminf((float)up, limit);
-    }
-    return gate * ACT_FN(up);
-  }
+                                            const scalar_t& y) {
+  return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
 }
 
 template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
-          bool act_first, bool HAS_CLAMP>
+          bool act_first>
 __device__ __forceinline__ packed_t packed_compute(const packed_t& x,
-                                                   const packed_t& y,
-                                                   const float limit) {
-  if constexpr (act_first) {
-    packed_t gate = x;
-    packed_t up = y;
-    if constexpr (HAS_CLAMP) {
-      float2 g = cast_to_float2(gate);
-      float2 u = cast_to_float2(up);
-      g.x = fminf(g.x, limit);
-      g.y = fminf(g.y, limit);
-      u.x = fmaxf(fminf(u.x, limit), -limit);
-      u.y = fmaxf(fminf(u.y, limit), -limit);
-      gate = cast_to_packed<packed_t>(g);
-      up = cast_to_packed<packed_t>(u);
-    }
-    return packed_mul(PACKED_ACT_FN(gate), up);
-  } else {
-    packed_t gate = x;
-    packed_t up = y;
-    if constexpr (HAS_CLAMP) {
-      float2 g = cast_to_float2(gate);
-      float2 u = cast_to_float2(up);
-      g.x = fmaxf(fminf(g.x, limit), -limit);
-      g.y = fmaxf(fminf(g.y, limit), -limit);
-      u.x = fminf(u.x, limit);
-      u.y = fminf(u.y, limit);
-      gate = cast_to_packed<packed_t>(g);
-      up = cast_to_packed<packed_t>(u);
-    }
-    return packed_mul(gate, PACKED_ACT_FN(up));
-  }
+                                                   const packed_t& y) {
+  return act_first ? packed_mul(PACKED_ACT_FN(x), y)
+                   : packed_mul(x, PACKED_ACT_FN(y));
 }
 
 // Activation and gating kernel template.
 template <typename scalar_t, typename packed_t,
           scalar_t (*ACT_FN)(const scalar_t&),
           packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
-          bool use_vec, bool HAS_CLAMP, bool use_256b = false>
+          bool use_vec, bool use_256b = false>
 __global__ void act_and_mul_kernel(
     scalar_t* __restrict__ out,          // [..., d]
     const scalar_t* __restrict__ input,  // [..., 2, d]
-    const int d, const float limit) {
+    const int d) {
   const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
   const scalar_t* y_ptr = x_ptr + d;
   scalar_t* out_ptr = out + blockIdx.x * d;
@@ -103,9 +58,8 @@ __global__ void act_and_mul_kernel(
       }
 #pragma unroll
       for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
-        x.elts[j] =
-            packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
-                x.elts[j], y.elts[j], limit);
+        x.elts[j] = packed_compute<packed_t, PACKED_ACT_FN, act_first>(
+            x.elts[j], y.elts[j]);
       }
       if constexpr (use_256b) {
         st256(x, &out_vec[i]);
@@ -118,8 +72,7 @@ __global__ void act_and_mul_kernel(
     for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
       const scalar_t x = VLLM_LDG(&x_ptr[idx]);
       const scalar_t y = VLLM_LDG(&y_ptr[idx]);
-      out_ptr[idx] =
-          compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
+      out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
     }
   }
 }
@@ -198,11 +151,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
 
 // Launch activation and gating kernel.
 // Use ACT_FIRST (bool) indicating whether to apply the activation function
-// first. HAS_CLAMP (bool) enables pre-activation clamping: gate input is
-// clamped (max only) and up input is clamped (both sides) before the
-// activation function is applied.
-#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST,        \
-                                      HAS_CLAMP, LIMIT)                        \
+// first.
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST)        \
   auto dtype = input.scalar_type();                                            \
   int d = input.size(-1) / 2;                                                  \
   int64_t num_tokens = input.numel() / input.size(-1);                         \
@@ -227,8 +177,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
             scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
             KERNEL<scalar_t>,                                                  \
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
-            ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>(     \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
+            ACT_FIRST, true, true><<<grid, block, 0, stream>>>(                \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);          \
       });                                                                      \
     } else {                                                                   \
       VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {          \
@@ -236,8 +186,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
             scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
             KERNEL<scalar_t>,                                                  \
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
-            ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>(    \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
+            ACT_FIRST, true, false><<<grid, block, 0, stream>>>(               \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);          \
       });                                                                      \
     }                                                                          \
   } else {                                                                     \
@@ -247,8 +197,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
           scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,        \
           KERNEL<scalar_t>,                                                    \
           PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
-          ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>(            \
-          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);     \
+          ACT_FIRST, false><<<grid, block, 0, stream>>>(                       \
+          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);            \
     });                                                                        \
   }
 
@@ -256,14 +206,7 @@ void silu_and_mul(torch::Tensor& out,    // [..., d]
                   torch::Tensor& input)  // [..., 2 * d]
 {
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true, false, 0.0f);
-}
-
-void silu_and_mul_clamp(torch::Tensor& out,    // [..., d]
-                        torch::Tensor& input,  // [..., 2 * d]
-                        double limit) {
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true, true, (float)limit);
+                                true);
 }
 
 void mul_and_silu(torch::Tensor& out,    // [..., d]
@@ -272,21 +215,21 @@ void mul_and_silu(torch::Tensor& out,    // [..., d]
   // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
   // applies the silu to the latter half of the input.
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                false, false, 0.0f);
+                                false);
 }
 
 void gelu_and_mul(torch::Tensor& out,    // [..., d]
                   torch::Tensor& input)  // [..., 2 * d]
 {
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
-                                true, false, 0.0f);
+                                true);
 }
 
 void gelu_tanh_and_mul(torch::Tensor& out,    // [..., d]
                        torch::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(
-      vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
+                                vllm::packed_gelu_tanh_kernel, true);
 }
 
 namespace vllm {

csrc/cache_kernels.cu

@@ -599,11 +599,6 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
   const int head_idx = (blockIdx.y * blockDim.x + threadIdx.x) * VEC_SIZE;
   // Find batch index within a block
   __shared__ int batch_idx[BLOCK_Y_SIZE];
-  if (threadIdx.x == 0) {
-    batch_idx[threadIdx.y] = -1;
-  }
-  __syncthreads();
-
   for (int iter = 0; iter < cuda_utils::ceil_div(batch_size, int(blockDim.x));
        iter++) {
     int tid = iter * blockDim.x + threadIdx.x;
@@ -616,18 +611,16 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
     }
   }
 
-  __syncthreads();
+#ifndef USE_ROCM
+  __syncwarp();
+#endif
 
-  // num_tokens may be an allocation upper bound when Python avoids a D2H sync.
-  // Only tokens covered by the exact device-side cu_seq_lens are valid to
-  // gather.
-  const int batch = batch_idx[threadIdx.y];
-  if (head_idx >= head_dim || token_idx >= num_tokens || batch < 0) {
+  if (head_idx >= head_dim || token_idx >= num_tokens) {
     return;
   }
-  const int inbatch_seq_idx = token_idx - cu_seq_lens[batch];
-  const int block_idx =
-      block_table[batch * num_blocks + inbatch_seq_idx / cache_block_size];
+  const int inbatch_seq_idx = token_idx - cu_seq_lens[batch_idx[threadIdx.y]];
+  const int block_idx = block_table[batch_idx[threadIdx.y] * num_blocks +
+                                    inbatch_seq_idx / cache_block_size];
   const int64_t src_block_offset = block_idx * block_stride;
   const int64_t cache_inblock_offset =
       (inbatch_seq_idx % cache_block_size) * head_dim + head_idx;
@@ -1497,9 +1490,6 @@ void concat_mla_q(torch::Tensor& ql_nope,  // [num_tokens, num_heads, nope_dim]
   TORCH_CHECK(ql_nope.stride(2) == 1, "ql_nope must have stride 1 in dim 2");
   TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2");
   TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2");
-  TORCH_CHECK(ql_nope.scalar_type() == at::ScalarType::Half ||
-                  ql_nope.scalar_type() == at::ScalarType::BFloat16,
-              "ql_nope must be float16 or bfloat16 dtype");
 
   if (num_tokens == 0) return;
 
@@ -1511,7 +1501,7 @@ void concat_mla_q(torch::Tensor& ql_nope,  // [num_tokens, num_heads, nope_dim]
   const at::cuda::OptionalCUDAGuard device_guard(device_of(ql_nope));
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
-  VLLM_DISPATCH_HALF_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
+  VLLM_DISPATCH_FLOATING_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
     vllm::ConcatMLAQKernel<scalar_t, 512><<<grid_size, block_size, 0, stream>>>(
         q_out.data_ptr<scalar_t>(), ql_nope.data_ptr<scalar_t>(),
         q_pe.data_ptr<scalar_t>(), num_tokens, num_heads, q_out.stride(0),

csrc/cpu/pos_encoding.cpp

@@ -178,12 +178,7 @@ void rotary_embedding_gptj_impl(
 
 void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
                       std::optional<torch::Tensor> key, int64_t head_size,
-                      torch::Tensor& cos_sin_cache, bool is_neox,
-                      int64_t rope_dim_offset, bool inverse) {
-  TORCH_CHECK(rope_dim_offset == 0,
-              "rope_dim_offset != 0 is not supported on CPU");
-  TORCH_CHECK(!inverse, "inverse rotary embedding is not supported on CPU");
-
+                      torch::Tensor& cos_sin_cache, bool is_neox) {
   int num_tokens = positions.numel();
   int rot_dim = cos_sin_cache.size(1);
   int num_heads = query.size(-1) / head_size;

csrc/cpu/torch_bindings.cpp

@@ -263,8 +263,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def(
       "rotary_embedding(Tensor positions, Tensor! query,"
       "                 Tensor!? key, int head_size,"
-      "                 Tensor cos_sin_cache, bool is_neox, int "
-      "rope_dim_offset=0, bool inverse=False) -> ()");
+      "                 Tensor cos_sin_cache, bool is_neox) -> ()");
   ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);
 
   // Quantization

csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu

@@ -1,477 +0,0 @@
-/*
- * SPDX-License-Identifier: Apache-2.0
- * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
- *
- * Horizontally-fused DeepseekV4-MLA kernel:
- *   - Q side:  per-head RMSNorm (no weight) + GPT-J RoPE on last ROPE_DIM
- *   - KV side: GPT-J RoPE on last ROPE_DIM + UE8M0 FP8 quant on NoPE + paged
- *              cache insert
- *
- * Structured after `applyMLARopeAndAssignQKVKernelGeneration` in
- * TensorRT-LLM's mlaKernels.cu: one kernel, one grid, with head-slot
- * dispatch choosing Q vs KV work per warp.  The per-warp RMSNorm/RoPE
- * skeleton is adapted from vllm-deepseek_v4's existing
- * `fusedQKNormRopeKernel` (csrc/fused_qknorm_rope_kernel.cu).
- *
- * Assumptions (hard-coded for DeepseekV4 attention):
- *   HEAD_DIM  = 512
- *   ROPE_DIM  = 64   (RoPE applied to dims [NOPE_DIM, HEAD_DIM))
- *   NOPE_DIM  = 448
- *   QUANT_BLOCK = 64 (UE8M0 FP8 quant block)
- *   FP8_MAX   = 448.0f
- *   is_neox=false (GPT-J interleaved pairs)
- *   cos_sin_cache layout [max_pos, rope_dim] = cos || sin (cos first, sin
- *     second along last dim; each half is rope_dim/2 = 32 values)
- *
- * Cache layout per paged-cache block (block_size tokens):
- *   [0,            bs*576):          token data, 448 fp8 + 128 bf16 each
- *   [bs*576,       bs*576 + bs*8):   UE8M0 scales, 7 real + 1 pad per token
- */
-
-#include <cmath>
-#include <cuda_fp8.h>
-#include <cuda_runtime.h>
-#include <type_traits>
-
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <torch/cuda.h>
-
-#include "cuda_compat.h"
-#include "dispatch_utils.h"
-#include "type_convert.cuh"
-
-#ifndef FINAL_MASK
-  #define FINAL_MASK 0xffffffffu
-#endif
-
-namespace vllm {
-namespace deepseek_v4_fused_ops {
-
-namespace {
-inline int getSMVersion() {
-  auto* props = at::cuda::getCurrentDeviceProperties();
-  return props->major * 10 + props->minor;
-}
-}  // namespace
-
-// ────────────────────────────────────────────────────────────────────────────
-// Constants
-// ────────────────────────────────────────────────────────────────────────────
-constexpr int kHeadDim = 512;
-constexpr int kRopeDim = 64;
-constexpr int kNopeDim = kHeadDim - kRopeDim;  // 448
-constexpr int kQuantBlock = 64;
-constexpr int kNumQuantBlocks = kNopeDim / kQuantBlock;   // 7
-constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1;  // 8 (7 real + 1 pad)
-constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2;  // 448 + 128 = 576
-constexpr float kFp8Max = 448.0f;
-
-// Per-warp layout:  32 lanes × 16 elems/lane = 512 elems = HEAD_DIM.
-constexpr int kNumLanes = 32;
-constexpr int kElemsPerLane = kHeadDim / kNumLanes;  // 16
-
-// ────────────────────────────────────────────────────────────────────────────
-// Small inline helpers
-// ────────────────────────────────────────────────────────────────────────────
-__device__ __forceinline__ float warp4MaxAbs(float val) {
-  // Reduce absolute max across 4 consecutive lanes (lane id & 3 group).
-  float peer = __shfl_xor_sync(FINAL_MASK, val, 1);
-  val = fmaxf(val, peer);
-  peer = __shfl_xor_sync(FINAL_MASK, val, 2);
-  val = fmaxf(val, peer);
-  return val;
-}
-
-template <typename T>
-__device__ __forceinline__ float warpSum(float val) {
-#pragma unroll
-  for (int mask = 16; mask > 0; mask >>= 1) {
-    val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
-  }
-  return val;
-}
-
-// ────────────────────────────────────────────────────────────────────────────
-// Kernel
-// ────────────────────────────────────────────────────────────────────────────
-//
-// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) /
-// warps_per_block) Block: blockDim.x = 256 threads (8 warps per block) Each
-// warp handles one (token, head_slot) pair. head_slot < num_heads_q          →
-// Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q         → KV
-// branch (RoPE + UE8M0 quant + insert)
-//
-// With DP padding, q/kv/position_ids can have more rows than slot_mapping.
-// The Q branch covers all `num_tokens_full` rows (downstream attention uses
-// them).  The KV branch only inserts the first `num_tokens_insert` tokens
-// (= slot_mapping length) into the paged cache.
-//
-template <typename scalar_t_in>
-__global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
-    scalar_t_in* __restrict__ q_inout,         // [N, H, 512] bf16, in place
-    scalar_t_in const* __restrict__ kv_in,     // [N, 512] bf16
-    uint8_t* __restrict__ k_cache,             // [num_blocks, block_stride]
-    int64_t const* __restrict__ slot_mapping,  // [num_tokens_insert] i64
-    int64_t const* __restrict__ position_ids,  // [N] i64
-    float const* __restrict__ cos_sin_cache,   // [max_pos, 64] fp32
-    float const eps,
-    int const num_tokens_full,    // = q.size(0) = kv.size(0)
-    int const num_tokens_insert,  // = slot_mapping.size(0), ≤ num_tokens_full
-    int const num_heads_q,        // H
-    int const cache_block_size,   // tokens per paged-cache block
-    int const kv_block_stride) {  // bytes per paged-cache block
-#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
-  // BF16 _typeConvert specialization is unavailable on pre-Ampere.  The
-  // DeepseekV4 kernel only runs with bf16 inputs in practice, so compile a
-  // no-op stub for sm_70/sm_75 to keep multi-arch builds happy.
-  if constexpr (std::is_same_v<scalar_t_in, c10::BFloat16>) {
-    return;
-  } else {
-#endif
-    using Converter = vllm::_typeConvert<scalar_t_in>;
-
-    int const warpsPerBlock = blockDim.x / 32;
-    int const warpId = threadIdx.x / 32;
-    int const laneId = threadIdx.x % 32;
-    int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
-
-    int const total_slots_per_token = num_heads_q + 1;
-    int const tokenIdx = globalWarpIdx / total_slots_per_token;
-    int const slotIdx = globalWarpIdx % total_slots_per_token;
-    if (tokenIdx >= num_tokens_full) return;
-
-    bool const isKV = (slotIdx == num_heads_q);
-    // KV branch: skip DP-padded tokens (no slot reserved for them).
-    if (isKV && tokenIdx >= num_tokens_insert) return;
-
-    // PDL: wait for predecessor kernel (upstream q/kv producer) to signal
-    // before touching any global memory.  No-op when PDL is not enabled on
-    // the launch.  The CUDA runtime wrapper emits the griddepcontrol.wait
-    // PTX with the required memory clobber internally.
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-    cudaGridDependencySynchronize();
-#endif
-
-    // Dim range this lane owns within the 512-wide head.
-    int const dim_base = laneId * kElemsPerLane;  // in [0, 512) step 16
-
-    // ── Load 16 bf16 → 16 fp32 registers (one 16-byte + one 16-byte LDG) ────
-    float elements[kElemsPerLane];
-    float sumOfSquares = 0.0f;
-
-    scalar_t_in const* src_ptr;
-    if (isKV) {
-      src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
-    } else {
-      int64_t const q_row_offset =
-          (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
-          dim_base;
-      src_ptr = q_inout + q_row_offset;
-    }
-
-    // Two 16-byte loads per thread (8 bf16 each).  Use uint4 as the vector
-    // type and bitcast to scalar_t_in packed pairs for conversion.
-    uint4 v0 = *reinterpret_cast<uint4 const*>(src_ptr);
-    uint4 v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
-
-    {
-      typename Converter::packed_hip_type const* p0 =
-          reinterpret_cast<typename Converter::packed_hip_type const*>(&v0);
-      typename Converter::packed_hip_type const* p1 =
-          reinterpret_cast<typename Converter::packed_hip_type const*>(&v1);
-// Each packed_hip_type holds 2 bf16 → 4 packed = 8 elems per uint4.
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        float2 f2 = Converter::convert(p0[i]);
-        elements[2 * i] = f2.x;
-        elements[2 * i + 1] = f2.y;
-      }
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        float2 f2 = Converter::convert(p1[i]);
-        elements[8 + 2 * i] = f2.x;
-        elements[8 + 2 * i + 1] = f2.y;
-      }
-    }
-
-    // ── Q branch: RMSNorm with no weight (has_weight=False) ─────────────────
-    // Variance + rsqrt + multiply all in fp32, no intermediate bf16 round.
-    // The downstream bf16 round only happens at the final store.
-    if (!isKV) {
-#pragma unroll
-      for (int i = 0; i < kElemsPerLane; i++) {
-        sumOfSquares += elements[i] * elements[i];
-      }
-      sumOfSquares = warpSum<float>(sumOfSquares);
-      float const rms_rcp =
-          rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
-#pragma unroll
-      for (int i = 0; i < kElemsPerLane; i++) {
-        elements[i] = elements[i] * rms_rcp;
-      }
-    }
-
-    // ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ─────────────────────────────
-    // All math in fp32.  cos_sin_cache is loaded as fp32 (its native storage).
-    bool const is_rope_lane = dim_base >= kNopeDim;
-    if (is_rope_lane) {
-      int64_t const pos = position_ids[tokenIdx];
-      constexpr int kHalfRope = kRopeDim / 2;  // 32
-      float const* cos_ptr = cos_sin_cache + pos * kRopeDim;
-      float const* sin_ptr = cos_ptr + kHalfRope;
-
-      int const rope_local_base = dim_base - kNopeDim;  // in [0, 64) step 16
-#pragma unroll
-      for (int p = 0; p < kElemsPerLane / 2; p++) {
-        int const pair_dim = rope_local_base + 2 * p;
-        int const half_idx = pair_dim / 2;
-        float const cos_v = VLLM_LDG(cos_ptr + half_idx);
-        float const sin_v = VLLM_LDG(sin_ptr + half_idx);
-        float const x_even = elements[2 * p];
-        float const x_odd = elements[2 * p + 1];
-        elements[2 * p] = x_even * cos_v - x_odd * sin_v;
-        elements[2 * p + 1] = x_even * sin_v + x_odd * cos_v;
-      }
-    }
-
-    // ═══════════════════════════════════════════════════════════════════════
-    // Q branch: cast to bf16 and store back in place.
-    // ═══════════════════════════════════════════════════════════════════════
-    if (!isKV) {
-      uint4 out0, out1;
-      typename Converter::packed_hip_type* po0 =
-          reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
-      typename Converter::packed_hip_type* po1 =
-          reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        po0[i] = Converter::convert(
-            make_float2(elements[2 * i], elements[2 * i + 1]));
-      }
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        po1[i] = Converter::convert(
-            make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
-      }
-      scalar_t_in* dst =
-          q_inout +
-          (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
-          dim_base;
-      *reinterpret_cast<uint4*>(dst) = out0;
-      *reinterpret_cast<uint4*>(dst + 8) = out1;
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-      cudaTriggerProgrammaticLaunchCompletion();
-#endif
-      return;
-    }
-
-    // ═══════════════════════════════════════════════════════════════════════
-    // KV branch.
-    // ═══════════════════════════════════════════════════════════════════════
-    int64_t const slot_id = slot_mapping[tokenIdx];
-    if (slot_id < 0) {
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-      cudaTriggerProgrammaticLaunchCompletion();
-#endif
-      return;
-    }
-
-    int64_t const block_idx = slot_id / cache_block_size;
-    int64_t const pos_in_block = slot_id % cache_block_size;
-    uint8_t* block_base =
-        k_cache + block_idx * static_cast<int64_t>(kv_block_stride);
-    uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes;
-    uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim;
-    uint8_t* token_scale_ptr =
-        block_base + static_cast<int64_t>(cache_block_size) * kTokenDataBytes +
-        pos_in_block * kScaleBytesPerToken;
-
-    // Round K to bf16 first, matching the unfused reference path where K is
-    // materialized as bf16 before K quantization.  absmax, clamp, and FP8
-    // quant below all run on these bf16-rounded values.
-#pragma unroll
-    for (int i = 0; i < kElemsPerLane; i++) {
-      elements[i] = Converter::convert(Converter::convert(elements[i]));
-    }
-
-    // Per-quant-block absmax must be computed by ALL 32 lanes (warp-collective
-    // shuffle requires full participation).  RoPE lanes contribute garbage,
-    // but their values are gated out below via `!is_rope_lane`.
-    float local_absmax = 0.0f;
-#pragma unroll
-    for (int i = 0; i < kElemsPerLane; i++) {
-      local_absmax = fmaxf(local_absmax, fabsf(elements[i]));
-    }
-    float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f);
-    float const exponent = ceilf(log2f(absmax / kFp8Max));
-    float const inv_scale = exp2f(-exponent);
-
-    if (!is_rope_lane) {
-      // ── NoPE lane: UE8M0 FP8 quant ───────────────────────────────────────
-      uint8_t out_bytes[kElemsPerLane];
-#pragma unroll
-      for (int i = 0; i < kElemsPerLane; i++) {
-        float scaled = elements[i] * inv_scale;
-        scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
-        __nv_fp8_storage_t s =
-            __nv_cvt_float_to_fp8(scaled, __NV_SATFINITE, __NV_E4M3);
-        out_bytes[i] = static_cast<uint8_t>(s);
-      }
-      // One 16-byte STG per lane.
-      *reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
-          *reinterpret_cast<uint4 const*>(out_bytes);
-
-      // Lane (4k) of each 4-lane group writes the scale byte for block k<7.
-      if ((laneId & 3) == 0) {
-        int const q_block_idx = laneId >> 2;  // 0..6 for NoPE lanes
-        float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f);
-        token_scale_ptr[q_block_idx] = static_cast<uint8_t>(encoded);
-      }
-      // Lane 0 also writes the padding byte at index 7.
-      if (laneId == 0) {
-        token_scale_ptr[kNumQuantBlocks] = 0;  // pad
-      }
-    } else {
-      // ── RoPE lane: cast back to bf16 and store to cache bf16 tail ────────
-      uint4 out0, out1;
-      typename Converter::packed_hip_type* po0 =
-          reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
-      typename Converter::packed_hip_type* po1 =
-          reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        po0[i] = Converter::convert(
-            make_float2(elements[2 * i], elements[2 * i + 1]));
-      }
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        po1[i] = Converter::convert(
-            make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
-      }
-      int const rope_local_base = dim_base - kNopeDim;  // in [0, 64)
-      scalar_t_in* bf16_dst =
-          reinterpret_cast<scalar_t_in*>(token_bf16_ptr) + rope_local_base;
-      *reinterpret_cast<uint4*>(bf16_dst) = out0;
-      *reinterpret_cast<uint4*>(bf16_dst + 8) = out1;
-    }
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-    cudaTriggerProgrammaticLaunchCompletion();
-#endif
-#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
-  }
-#endif
-}
-
-// ────────────────────────────────────────────────────────────────────────────
-// Launch wrapper
-// ────────────────────────────────────────────────────────────────────────────
-template <typename scalar_t_in>
-void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
-    scalar_t_in* q_inout, scalar_t_in const* kv_in, uint8_t* k_cache,
-    int64_t const* slot_mapping, int64_t const* position_ids,
-    float const* cos_sin_cache, float const eps, int const num_tokens_full,
-    int const num_tokens_insert, int const num_heads_q,
-    int const cache_block_size, int const kv_block_stride,
-    cudaStream_t stream) {
-  constexpr int kBlockSize = 256;
-  constexpr int kWarpsPerBlock = kBlockSize / 32;
-  int64_t const total_warps =
-      static_cast<int64_t>(num_tokens_full) * (num_heads_q + 1);
-  int const grid =
-      static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
-
-  // PDL: enable programmatic stream serialization whenever the hardware
-  // supports it (SM90+).  On pre-Hopper GPUs the attribute is unavailable,
-  // so leave numAttrs = 0 and launch as a regular kernel.
-  static int const sm_version = getSMVersion();
-  // Host-side guard: the device kernel body is compiled as a no-op for
-  // bf16 on pre-Ampere (sm_70/sm_75) because _typeConvert<BFloat16> is
-  // unavailable there.  Refuse the launch loudly instead of silently
-  // skipping the work.
-  TORCH_CHECK(
-      sm_version >= 80,
-      "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert requires sm_80+ "
-      "(Ampere or newer); got sm_",
-      sm_version);
-  cudaLaunchConfig_t config;
-  config.gridDim = dim3(grid);
-  config.blockDim = dim3(kBlockSize);
-  config.dynamicSmemBytes = 0;
-  config.stream = stream;
-  cudaLaunchAttribute attrs[1];
-  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-  attrs[0].val.programmaticStreamSerializationAllowed = 1;
-  config.attrs = attrs;
-  config.numAttrs = (sm_version >= 90) ? 1 : 0;
-
-  cudaLaunchKernelEx(
-      &config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>,
-      q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
-      num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
-      kv_block_stride);
-}
-
-}  // namespace deepseek_v4_fused_ops
-}  // namespace vllm
-
-// ────────────────────────────────────────────────────────────────────────────
-// Torch op wrapper
-// ────────────────────────────────────────────────────────────────────────────
-void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::Tensor& q,                    // [N, H, 512] bf16, in place
-    torch::Tensor const& kv,             // [N, 512] bf16 (read-only)
-    torch::Tensor& k_cache,              // [num_blocks, block_bytes] uint8
-    torch::Tensor const& slot_mapping,   // [N] int64
-    torch::Tensor const& position_ids,   // [N] int64
-    torch::Tensor const& cos_sin_cache,  // [max_pos, rope_dim] bf16
-    double eps, int64_t cache_block_size) {
-  TORCH_CHECK(q.is_cuda() && q.is_contiguous(), "q must be contiguous CUDA");
-  TORCH_CHECK(kv.is_cuda() && kv.is_contiguous(), "kv must be contiguous CUDA");
-  TORCH_CHECK(k_cache.is_cuda(), "k_cache must be CUDA");
-  TORCH_CHECK(slot_mapping.is_cuda() && slot_mapping.dtype() == torch::kInt64,
-              "slot_mapping must be int64 CUDA");
-  TORCH_CHECK(position_ids.is_cuda() && position_ids.dtype() == torch::kInt64,
-              "position_ids must be int64 CUDA");
-  TORCH_CHECK(cos_sin_cache.is_cuda(), "cos_sin_cache must be CUDA");
-  TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
-  TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
-  TORCH_CHECK(q.dtype() == kv.dtype(), "q and kv dtype must match");
-  TORCH_CHECK(k_cache.dtype() == torch::kUInt8, "k_cache must be uint8");
-  TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
-              "cos_sin_cache shape [max_pos, 64]");
-  TORCH_CHECK(cos_sin_cache.dtype() == torch::kFloat32,
-              "cos_sin_cache must be float32");
-
-  // With DP padding, slot_mapping can be shorter than q/kv/positions.
-  // Q-norm+RoPE runs on all q.size(0) rows (downstream attention uses them);
-  // KV quant+insert runs only on the first slot_mapping.size(0) rows.
-  int const num_tokens_full = static_cast<int>(q.size(0));
-  int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
-  TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
-                  static_cast<int>(position_ids.size(0)) == num_tokens_full,
-              "q/kv/position_ids row counts must match");
-  TORCH_CHECK(num_tokens_insert <= num_tokens_full,
-              "slot_mapping must not exceed q row count");
-  int const num_heads_q = static_cast<int>(q.size(1));
-  int const cache_block_size_i = static_cast<int>(cache_block_size);
-  int const kv_block_stride = static_cast<int>(k_cache.stride(0));
-
-  at::cuda::OptionalCUDAGuard device_guard(device_of(q));
-  auto stream = at::cuda::getCurrentCUDAStream();
-
-  VLLM_DISPATCH_HALF_TYPES(
-      q.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
-        using qkv_scalar_t = scalar_t;
-        vllm::deepseek_v4_fused_ops::
-            launchFusedDeepseekV4QNormRopeKVRopeQuantInsert<qkv_scalar_t>(
-                reinterpret_cast<qkv_scalar_t*>(q.data_ptr()),
-                reinterpret_cast<qkv_scalar_t const*>(kv.data_ptr()),
-                reinterpret_cast<uint8_t*>(k_cache.data_ptr()),
-                reinterpret_cast<int64_t const*>(slot_mapping.data_ptr()),
-                reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
-                cos_sin_cache.data_ptr<float>(), static_cast<float>(eps),
-                num_tokens_full, num_tokens_insert, num_heads_q,
-                cache_block_size_i, kv_block_stride, stream);
-      });
-}

csrc/layernorm_kernels.cu

@@ -77,8 +77,7 @@ __global__ void rms_norm_kernel(
 #pragma unroll
     for (int j = 0; j < VEC_SIZE; j++) {
       float x = static_cast<float>(src1.val[j]);
-      float w = static_cast<float>(src2.val[j]);
-      dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
+      dst.val[j] = ((scalar_t)(x * s_variance)) * src2.val[j];
     }
     v_out[i] = dst;
   }
@@ -135,17 +134,10 @@ fused_add_rms_norm_kernel(
   for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
     int id = blockIdx.x * vec_hidden_size + idx;
     int64_t strided_id = blockIdx.x * vec_input_stride + idx;
-    _f16Vec<scalar_t, width> res = residual_v[id];
-    _f16Vec<scalar_t, width> w = weight_v[idx];
-    _f16Vec<scalar_t, width> out;
-    using Converter = _typeConvert<scalar_t>;
-#pragma unroll
-    for (int j = 0; j < width; ++j) {
-      float x = Converter::convert(res.data[j]);
-      float wf = Converter::convert(w.data[j]);
-      out.data[j] = Converter::convert(x * s_variance * wf);
-    }
-    input_v[strided_id] = out;
+    _f16Vec<scalar_t, width> temp = residual_v[id];
+    temp *= s_variance;
+    temp *= weight_v[idx];
+    input_v[strided_id] = temp;
   }
 }
 
@@ -182,8 +174,8 @@ fused_add_rms_norm_kernel(
 
   for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
     float x = (float)residual[blockIdx.x * hidden_size + idx];
-    float w = (float)weight[idx];
-    input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
+    input[blockIdx.x * input_stride + idx] =
+        ((scalar_t)(x * s_variance)) * weight[idx];
   }
 }
 

csrc/layernorm_quant_kernels.cu

@@ -65,16 +65,9 @@ __global__ void rms_norm_static_fp8_quant_kernel(
 #pragma unroll
     for (int j = 0; j < VEC_SIZE; j++) {
       float x = static_cast<float>(src1.val[j]);
-      float w = static_cast<float>(src2.val[j]);
-      // Round normalized result through scalar_t to match the precision of the
-      // unfused composite (rms_norm writes scalar_t, then
-      // static_scaled_fp8_quant re-loads it as float before FP8 conversion).
-      // Without this round, the fused path is strictly more accurate and
-      // disagrees with the composite at exact E4M3 quantization tie boundaries.
-      scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
+      float const out_norm = ((scalar_t)(x * s_variance)) * src2.val[j];
       out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
-          scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
-                                                scale_inv);
+          scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
     }
   }
 }
@@ -134,21 +127,13 @@ fused_add_rms_norm_static_fp8_quant_kernel(
 
   for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
     int id = blockIdx.x * vec_hidden_size + idx;
-    _f16Vec<scalar_t, width> res = residual_v[id];
-    _f16Vec<scalar_t, width> w = weight_v[idx];
-    using Converter = _typeConvert<scalar_t>;
-    using HipT = typename Converter::hip_type;
+    _f16Vec<scalar_t, width> temp = residual_v[id];
+    temp *= s_variance;
+    temp *= weight_v[idx];
 #pragma unroll
     for (int i = 0; i < width; ++i) {
-      float x = Converter::convert(res.data[i]);
-      float wf = Converter::convert(w.data[i]);
-      // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
-      // to match the unfused composite path at FP8 boundaries. We use the
-      // backend's hip_type for the intermediate since c10::Half/BFloat16 has
-      // ambiguous conversions on CUDA and no implicit conversion on ROCm.
-      HipT out_norm_h = Converter::convert(x * s_variance * wf);
-      out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
-          Converter::convert(out_norm_h), scale_inv);
+      out[id * width + i] =
+          scaled_fp8_conversion<true, fp8_type>(float(temp.data[i]), scale_inv);
     }
   }
 }
@@ -191,12 +176,9 @@ fused_add_rms_norm_static_fp8_quant_kernel(
 
   for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
     float x = (float)residual[blockIdx.x * hidden_size + idx];
-    float w = (float)weight[idx];
-    // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
-    // to match the unfused composite path at FP8 boundaries.
-    scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
-    out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
-        static_cast<float>(out_norm), scale_inv);
+    float const out_norm = ((scalar_t)(x * s_variance)) * weight[idx];
+    out[blockIdx.x * hidden_size + idx] =
+        scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
   }
 }
 

csrc/moe/moe_ops.h

@@ -12,15 +12,6 @@ void topk_sigmoid(torch::Tensor& topk_weights, torch::Tensor& topk_indices,
                   torch::Tensor& gating_output, bool renormalize,
                   std::optional<torch::Tensor> bias);
 
-void topk_softplus_sqrt(torch::Tensor& topk_weights,
-                        torch::Tensor& topk_indices,
-                        torch::Tensor& token_expert_indices,
-                        torch::Tensor& gating_output, bool renormalize,
-                        double routed_scaling_factor,
-                        const c10::optional<torch::Tensor>& correction_bias,
-                        const c10::optional<torch::Tensor>& input_ids,
-                        const c10::optional<torch::Tensor>& tid2eid);
-
 void moe_sum(torch::Tensor& input, torch::Tensor& output);
 
 void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,

csrc/moe/topk_softplus_sqrt_kernels.cu

@@ -1,715 +0,0 @@
-/*
- * Adapted from
- * https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu
- * Copyright (c) 2024, The vLLM team.
- * SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION &
- * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-#include <type_traits>
-#include <torch/all.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-#include "../cuda_compat.h"
-#include "../cub_helpers.h"
-#ifndef USE_ROCM
-  #include <cuda_bf16.h>
-  #include <cuda_fp16.h>
-#else
-  #include <hip/hip_bf16.h>
-  #include <hip/hip_fp16.h>
-typedef __hip_bfloat16 __nv_bfloat16;
-typedef __hip_bfloat162 __nv_bfloat162;
-#endif
-
-#define MAX(a, b) ((a) > (b) ? (a) : (b))
-#define MIN(a, b) ((a) < (b) ? (a) : (b))
-
-namespace vllm {
-namespace moe {
-
-/// Aligned array type
-template <typename T,
-          /// Number of elements in the array
-          int N,
-          /// Alignment requirement in bytes
-          int Alignment = sizeof(T) * N>
-struct alignas(Alignment) AlignedArray {
-  T data[N];
-};
-
-template <typename T>
-__device__ __forceinline__ float toFloat(T value) {
-  if constexpr (std::is_same_v<T, float>) {
-    return value;
-  } else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
-    return __bfloat162float(value);
-  } else if constexpr (std::is_same_v<T, __half>) {
-    return __half2float(value);
-  }
-}
-
-#define FINAL_MASK 0xffffffff
-template <typename T>
-__inline__ __device__ T warpReduceSum(T val) {
-#pragma unroll
-  for (int mask = 16; mask > 0; mask >>= 1)
-    val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
-  return val;
-}
-
-// ====================== TopK softplus_sqrt things
-// ===============================
-
-/*
-  A Top-K gating softplus_sqrt written to exploit when the number of experts in
-  the MoE layers are a small power of 2. This allows us to cleanly share the
-  rows among the threads in a single warp and eliminate communication between
-  warps (so no need to use shared mem).
-
-  It fuses the sigmoid, max and argmax into a single kernel.
-
-  Limitations:
-  1) This implementation is optimized for when the number of experts is a small
-  power of 2. Additionally it also supports when number of experts is multiple
-  of 64 which is still faster than the computing sigmoid and topK separately
-  (only tested on CUDA yet). 2) This implementation assumes k is small, but will
-  work for any k.
-*/
-
-template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG,
-          int WARP_SIZE_PARAM, bool USE_HASH, typename IndType,
-          typename InputType = float>
-__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
-    void topkGatingSoftplusSqrt(
-        const InputType* input, const bool* finished, float* output,
-        const int num_rows, IndType* indices, int* source_rows, const int k,
-        const int start_expert, const int end_expert, const bool renormalize,
-        double routed_scaling_factor, const float* correction_bias,
-        const IndType* input_ids, const IndType* tid2eid) {
-  static_assert(std::is_same_v<InputType, float> ||
-                    std::is_same_v<InputType, __nv_bfloat16> ||
-                    std::is_same_v<InputType, __half>,
-                "InputType must be float, __nv_bfloat16, or __half");
-
-  // We begin by enforcing compile time assertions and setting up compile time
-  // constants.
-  static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG),
-                "BYTES_PER_LDG must be power of 2");
-  static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
-
-  // Number of bytes each thread pulls in per load
-  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
-  static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
-  static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
-  static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
-
-  if constexpr (std::is_same_v<InputType, __nv_bfloat16> ||
-                std::is_same_v<InputType, __half>) {
-    static_assert(ELTS_PER_LDG == 1 || ELTS_PER_LDG % 2 == 0,
-                  "ELTS_PER_LDG must be 1 or even for 16-bit conversion");
-  }
-
-  // Restrictions based on previous section.
-  static_assert(
-      VPT % ELTS_PER_LDG == 0,
-      "The elements per thread must be a multiple of the elements per ldg");
-  static_assert(WARP_SIZE_PARAM % THREADS_PER_ROW == 0,
-                "The threads per row must cleanly divide the threads per warp");
-  static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW),
-                "THREADS_PER_ROW must be power of 2");
-  static_assert(THREADS_PER_ROW <= WARP_SIZE_PARAM,
-                "THREADS_PER_ROW can be at most warp size");
-
-  // We have NUM_EXPERTS elements per row. We specialize for small #experts
-  static constexpr int ELTS_PER_WARP = WARP_SIZE_PARAM * VPT;
-  static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
-  static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
-
-  // Restrictions for previous section.
-  static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0,
-                "The elts per row must cleanly divide the total elt per warp");
-
-  // ===================== From this point, we finally start computing run-time
-  // variables. ========================
-
-  // Compute CTA and warp rows. We pack multiple rows into a single warp, and a
-  // block contains WARPS_PER_CTA warps. This, each block processes a chunk of
-  // rows. We start by computing the start row for each block.
-  const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
-
-  // Now, using the base row per thread block, we compute the base row per warp.
-  const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
-
-  // The threads in a warp are split into sub-groups that will work on a row.
-  // We compute row offset for each thread sub-group
-  const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
-  const int thread_row = warp_base_row + thread_row_in_warp;
-
-  // Threads with indices out of bounds should early exit here.
-  if (thread_row >= num_rows) {
-    return;
-  }
-  const bool row_is_active = finished ? !finished[thread_row] : true;
-
-  // We finally start setting up the read pointers for each thread. First, each
-  // thread jumps to the start of the row it will read.
-  const InputType* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
-
-  // Now, we compute the group each thread belong to in order to determine the
-  // first column to start loads.
-  const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
-  const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
-  const InputType* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
-
-  // Finally, we pull in the data from global mem
-  float row_chunk[VPT];
-
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.wait;");
-#endif
-
-  // NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
-  // to float
-  if constexpr (std::is_same_v<InputType, float>) {
-    using VecType = AlignedArray<float, ELTS_PER_LDG>;
-    VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk);
-    const VecType* vec_thread_read_ptr =
-        reinterpret_cast<const VecType*>(thread_read_ptr);
-#pragma unroll
-    for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
-      row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
-    }
-  } else if constexpr (std::is_same_v<InputType, __nv_bfloat16>) {
-    if constexpr (ELTS_PER_LDG >= 2) {
-      using VecType = AlignedArray<__nv_bfloat16, ELTS_PER_LDG>;
-      float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
-      const VecType* vec_thread_read_ptr =
-          reinterpret_cast<const VecType*>(thread_read_ptr);
-#pragma unroll
-      for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
-        VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
-        int base_idx_f2 = ii * ELTS_PER_LDG / 2;
-#pragma unroll
-        for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
-          row_chunk_f2[base_idx_f2 + jj] = __bfloat1622float2(
-              *reinterpret_cast<const __nv_bfloat162*>(vec.data + jj * 2));
-        }
-      }
-    } else {  // ELTS_PER_LDG == 1
-#pragma unroll
-      for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
-        const __nv_bfloat16* scalar_ptr =
-            thread_read_ptr + ii * THREADS_PER_ROW;
-        row_chunk[ii] = __bfloat162float(*scalar_ptr);
-      }
-    }
-  } else if constexpr (std::is_same_v<InputType, __half>) {
-    if constexpr (ELTS_PER_LDG >= 2) {
-      using VecType = AlignedArray<__half, ELTS_PER_LDG>;
-      float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
-      const VecType* vec_thread_read_ptr =
-          reinterpret_cast<const VecType*>(thread_read_ptr);
-#pragma unroll
-      for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
-        VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
-        int base_idx_f2 = ii * ELTS_PER_LDG / 2;
-#pragma unroll
-        for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
-          row_chunk_f2[base_idx_f2 + jj] = __half22float2(
-              *reinterpret_cast<const __half2*>(vec.data + jj * 2));
-        }
-      }
-    } else {  // ELTS_PER_LDG == 1
-#pragma unroll
-      for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
-        const __half* scalar_ptr = thread_read_ptr + ii * THREADS_PER_ROW;
-        row_chunk[ii] = __half2float(*scalar_ptr);
-      }
-    }
-  }
-  constexpr float threshold = 20.0f;
-  constexpr float beta = 1.0f;
-
-  // Hash MoE path: indices are predetermined from lookup table
-  if constexpr (USE_HASH) {
-    const IndType token_id = input_ids[thread_row];
-    const IndType* expert_indices_for_token = tid2eid + token_id * k;
-#pragma unroll
-    for (int ii = 0; ii < VPT; ++ii) {
-      float val = row_chunk[ii];
-      float val_b = val * beta;
-      val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
-      row_chunk[ii] = sqrtf(val);
-    }
-    float selected_sum = 0.f;
-#pragma unroll
-    for (int k_idx = 0; k_idx < k; ++k_idx) {
-      const int expert = expert_indices_for_token[k_idx];
-      const int idx = k * thread_row + k_idx;
-      for (int ii = 0; ii < VPT; ++ii) {
-        const int group_id = ii / ELTS_PER_LDG;
-        const int local_id = ii % ELTS_PER_LDG;
-        const int expert_idx = first_elt_read_by_thread +
-                               group_id * THREADS_PER_ROW * ELTS_PER_LDG +
-                               local_id;
-        if (expert == expert_idx) {
-          indices[idx] = expert;
-          selected_sum += row_chunk[ii];
-          break;
-        }
-      }
-    }
-    // Compute per-thread scale (using warp reduction when renormalizing).
-    if (renormalize) {
-      selected_sum = warpReduceSum(selected_sum);
-    }
-    float scale = static_cast<float>(routed_scaling_factor);
-    if (renormalize) {
-      const float denom = selected_sum > 0.f ? selected_sum : 1.f;
-      scale /= denom;
-    }
-
-#pragma unroll
-    for (int k_idx = 0; k_idx < k; ++k_idx) {
-      const int expert = expert_indices_for_token[k_idx];
-      const int idx = k * thread_row + k_idx;
-      for (int ii = 0; ii < VPT; ++ii) {
-        const int group_id = ii / ELTS_PER_LDG;
-        const int local_id = ii % ELTS_PER_LDG;
-        const int expert_idx = first_elt_read_by_thread +
-                               group_id * THREADS_PER_ROW * ELTS_PER_LDG +
-                               local_id;
-        if (expert == expert_idx) {
-          output[idx] = row_chunk[ii] * scale;
-          break;
-        }
-      }
-    }
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-    asm volatile("griddepcontrol.launch_dependents;");
-#endif
-    return;
-  }
-
-#pragma unroll
-  for (int ii = 0; ii < VPT; ++ii) {
-    float val = row_chunk[ii];
-    float val_b = val * beta;
-    // Compute softplus: log(1 + exp(val)) with numerical stability
-    // When val > threshold, softplus(x) ≈ x to avoid exp overflow
-    val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
-    val = sqrtf(val);
-    if (correction_bias) {
-      const int group_id = ii / ELTS_PER_LDG;
-      const int local_id = ii % ELTS_PER_LDG;
-      const int expert_idx = first_elt_read_by_thread +
-                             group_id * THREADS_PER_ROW * ELTS_PER_LDG +
-                             local_id;
-      val = val + correction_bias[expert_idx];
-    }
-    row_chunk[ii] = val;
-  }
-
-  // Original TopK path: find top-k experts by score
-  // Now, sigmoid_res contains the sigmoid of the row chunk. Now, I want to find
-  // the topk elements in each row, along with the max index.
-  int start_col = first_elt_read_by_thread;
-  static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
-
-  float selected_sum = 0.f;
-  for (int k_idx = 0; k_idx < k; ++k_idx) {
-    // First, each thread does the local argmax
-    float max_val = row_chunk[0];
-    int expert = start_col;
-#pragma unroll
-    for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
-         ++ldg, col += COLS_PER_GROUP_LDG) {
-#pragma unroll
-      for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
-        float val = row_chunk[ldg * ELTS_PER_LDG + ii];
-
-        // No check on the experts here since columns with the smallest index
-        // are processed first and only updated if > (not >=)
-        if (val > max_val) {
-          max_val = val;
-          expert = col + ii;
-        }
-      }
-    }
-
-// Now, we perform the argmax reduce. We use the butterfly pattern so threads
-// reach consensus about the max. This will be useful for K > 1 so that the
-// threads can agree on "who" had the max value. That thread can then blank out
-// their max with -inf and the warp can run more iterations...
-#pragma unroll
-    for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
-      float other_max =
-          VLLM_SHFL_XOR_SYNC_WIDTH(max_val, mask, THREADS_PER_ROW);
-      int other_expert =
-          VLLM_SHFL_XOR_SYNC_WIDTH(expert, mask, THREADS_PER_ROW);
-
-      // We want lower indices to "win" in every thread so we break ties this
-      // way
-      if (other_max > max_val ||
-          (other_max == max_val && other_expert < expert)) {
-        max_val = other_max;
-        expert = other_expert;
-      }
-    }
-
-    // Write the max for this k iteration to global memory.
-    if (thread_group_idx == 0) {
-      // Add a guard to ignore experts not included by this node
-      const bool node_uses_expert =
-          expert >= start_expert && expert < end_expert;
-      const bool should_process_row = row_is_active && node_uses_expert;
-
-      // The lead thread from each sub-group will write out the final results to
-      // global memory. (This will be a single) thread per row of the
-      // input/output matrices.
-      const int idx = k * thread_row + k_idx;
-      if (correction_bias != nullptr) {
-        max_val -= correction_bias[expert];
-      }
-      output[idx] = max_val;
-      indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
-      source_rows[idx] = k_idx * num_rows + thread_row;
-      if (renormalize) {
-        selected_sum += max_val;
-      }
-    }
-
-    // Finally, we clear the value in the thread with the current max if there
-    // is another iteration to run.
-    if (k_idx + 1 < k) {
-      const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
-      const int thread_to_clear_in_group =
-          (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
-
-      // Only the thread in the group which produced the max will reset the
-      // "winning" value to -inf.
-      if (thread_group_idx == thread_to_clear_in_group) {
-        const int offset_for_expert = expert % ELTS_PER_LDG;
-        // Safe to set to any negative value since row_chunk values must be
-        // between 0 and 1.
-        row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] =
-            -10000.f;
-      }
-    }
-  }
-
-  // Apply renormalization and routed scaling factor to final weights.
-  if (thread_group_idx == 0) {
-    float scale = static_cast<float>(routed_scaling_factor);
-    if (renormalize) {
-      const float denom = selected_sum > 0.f ? selected_sum : 1.f;
-      scale /= denom;
-    }
-    for (int k_idx = 0; k_idx < k; ++k_idx) {
-      const int idx = k * thread_row + k_idx;
-      output[idx] = output[idx] * scale;
-    }
-  }
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.launch_dependents;");
-#endif
-}
-
-namespace detail {
-// Constructs some constants needed to partition the work across threads at
-// compile time.
-template <int EXPERTS, int BYTES_PER_LDG, int WARP_SIZE_PARAM,
-          typename InputType>
-struct TopkConstants {
-  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
-  static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0 ||
-                    EXPERTS % (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0,
-                "");
-  static constexpr int VECs_PER_THREAD =
-      MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM));
-  static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
-  static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
-  static const int ROWS_PER_WARP = WARP_SIZE_PARAM / THREADS_PER_ROW;
-};
-}  // namespace detail
-
-#define DISPATCH_HASH(use_hash, USE_HASH, ...)                                 \
-  if (use_hash) {                                                              \
-    const bool USE_HASH = true;                                                \
-    static_assert(USE_HASH == true, "USE_HASH must be compile-time constant"); \
-    __VA_ARGS__                                                                \
-  } else {                                                                     \
-    const bool USE_HASH = false;                                               \
-    static_assert(USE_HASH == false,                                           \
-                  "USE_HASH must be compile-time constant");                   \
-    __VA_ARGS__                                                                \
-  }
-
-template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM,
-          int MAX_BYTES_PER_LDG, typename IndType, typename InputType>
-void topkGatingSoftplusSqrtLauncherHelper(
-    const InputType* input, const bool* finished, float* output,
-    IndType* indices, int* source_row, const int num_rows, const int k,
-    const int start_expert, const int end_expert, const bool renormalize,
-    double routed_scaling_factor, const float* correction_bias,
-    const bool use_hash, const IndType* input_ids, const IndType* tid2eid,
-    cudaStream_t stream) {
-  static constexpr int BYTES_PER_LDG =
-      MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
-  using Constants =
-      detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
-  static constexpr int VPT = Constants::VPT;
-  static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
-  const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
-  const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
-  dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
-  DISPATCH_HASH(use_hash, USE_HASH, {
-    auto* kernel =
-        &topkGatingSoftplusSqrt<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG,
-                                WARP_SIZE_PARAM, USE_HASH, IndType, InputType>;
-#ifndef USE_ROCM
-    cudaLaunchConfig_t config = {};
-    config.gridDim = num_blocks;
-    config.blockDim = block_dim;
-    config.dynamicSmemBytes = 0;
-    config.stream = stream;
-    cudaLaunchAttribute attrs[1];
-    attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-    attrs[0].val.programmaticStreamSerializationAllowed = 1;
-    config.numAttrs = 1;
-    config.attrs = attrs;
-    cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows,
-                       indices, source_row, k, start_expert, end_expert,
-                       renormalize, routed_scaling_factor, correction_bias,
-                       input_ids, tid2eid);
-#else
-    kernel<<<num_blocks, block_dim, 0, stream>>>(
-        input, finished, output, num_rows, indices, source_row, k, start_expert,
-        end_expert, renormalize, routed_scaling_factor, correction_bias,
-        input_ids, tid2eid);
-#endif
-  })
-}
-
-#ifndef USE_ROCM
-  #define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES)           \
-    static_assert(WARP_SIZE == 32,                                             \
-                  "Unsupported warp size. Only 32 is supported for CUDA");     \
-    topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, WARP_SIZE, \
-                                         MAX_BYTES>(                           \
-        gating_output, nullptr, topk_weights, topk_indices,                    \
-        token_expert_indices, num_tokens, topk, 0, num_experts, renormalize,   \
-        routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid,  \
-        stream);
-#else
-  #define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES)           \
-    if (WARP_SIZE == 64) {                                                     \
-      topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 64,      \
-                                           MAX_BYTES>(                         \
-          gating_output, nullptr, topk_weights, topk_indices,                  \
-          token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
-          routed_scaling_factor, correction_bias, use_hash, input_ids,         \
-          tid2eid, stream);                                                    \
-    } else if (WARP_SIZE == 32) {                                              \
-      topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32,      \
-                                           MAX_BYTES>(                         \
-          gating_output, nullptr, topk_weights, topk_indices,                  \
-          token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
-          routed_scaling_factor, correction_bias, use_hash, input_ids,         \
-          tid2eid, stream);                                                    \
-    } else {                                                                   \
-      assert(false &&                                                          \
-             "Unsupported warp size. Only 32 and 64 are supported for ROCm");  \
-    }
-#endif
-
-template <typename IndType, typename InputType>
-void topkGatingSoftplusSqrtKernelLauncher(
-    const InputType* gating_output, float* topk_weights, IndType* topk_indices,
-    int* token_expert_indices, const int num_tokens, const int num_experts,
-    const int topk, const bool renormalize, double routed_scaling_factor,
-    const float* correction_bias, const bool use_hash, const IndType* input_ids,
-    const IndType* tid2eid, cudaStream_t stream) {
-  static constexpr int WARPS_PER_TB = 4;
-  static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
-#ifndef USE_ROCM
-  // for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
-  // elements can be loaded by a warp
-  static constexpr int BYTES_PER_LDG_MULTIPLE_64 =
-      (std::is_same_v<InputType, __nv_bfloat16> ||
-       std::is_same_v<InputType, __half>)
-          ? 4
-          : 8;
-#endif
-  switch (num_experts) {
-    case 1:
-      LAUNCH_SOFTPLUS_SQRT(1, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-    case 2:
-      LAUNCH_SOFTPLUS_SQRT(2, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-    case 4:
-      LAUNCH_SOFTPLUS_SQRT(4, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-    case 8:
-      LAUNCH_SOFTPLUS_SQRT(8, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-    case 16:
-      LAUNCH_SOFTPLUS_SQRT(16, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-    case 32:
-      LAUNCH_SOFTPLUS_SQRT(32, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-    case 64:
-      LAUNCH_SOFTPLUS_SQRT(64, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-    case 128:
-      LAUNCH_SOFTPLUS_SQRT(128, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-    case 256:
-      LAUNCH_SOFTPLUS_SQRT(256, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-    case 512:
-      LAUNCH_SOFTPLUS_SQRT(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-      // (CUDA only) support multiples of 64 when num_experts is not power of 2.
-      // ROCm uses WARP_SIZE 64 so 8 bytes loading won't fit for some of
-      // num_experts, alternatively we can test 4 bytes loading and enable it in
-      // future.
-#ifndef USE_ROCM
-    case 192:
-      LAUNCH_SOFTPLUS_SQRT(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
-      break;
-    case 320:
-      LAUNCH_SOFTPLUS_SQRT(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
-      break;
-    case 384:
-      LAUNCH_SOFTPLUS_SQRT(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
-      break;
-    case 448:
-      LAUNCH_SOFTPLUS_SQRT(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
-      break;
-    case 576:
-      LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
-      break;
-#endif
-    default: {
-      TORCH_CHECK(false, "Unsupported expert number: ", num_experts);
-    }
-  }
-}
-
-}  // namespace moe
-}  // namespace vllm
-
-template <typename ComputeType>
-void dispatch_topk_softplus_sqrt_launch(
-    const ComputeType* gating_output, torch::Tensor& topk_weights,
-    torch::Tensor& topk_indices, torch::Tensor& token_expert_indices,
-    int num_tokens, int num_experts, int topk, bool renormalize,
-    double routed_scaling_factor,
-    const c10::optional<torch::Tensor>& correction_bias,
-    const c10::optional<torch::Tensor>& input_ids,
-    const c10::optional<torch::Tensor>& tid2eid, cudaStream_t stream) {
-  const float* bias_ptr = nullptr;
-  if (correction_bias.has_value()) {
-    bias_ptr = correction_bias.value().data_ptr<float>();
-  }
-  bool use_hash = false;
-  if (tid2eid.has_value()) {
-    TORCH_CHECK(input_ids.has_value(), "input_ids is required for hash MoE");
-    use_hash = true;
-  }
-  if (topk_indices.scalar_type() == at::ScalarType::Int) {
-    const int* input_ids_ptr = nullptr;
-    const int* tid2eid_ptr = nullptr;
-    if (tid2eid.has_value()) {
-      input_ids_ptr = input_ids.value().data_ptr<int>();
-      tid2eid_ptr = tid2eid.value().data_ptr<int>();
-    }
-
-    vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int, ComputeType>(
-        gating_output, topk_weights.data_ptr<float>(),
-        topk_indices.data_ptr<int>(), token_expert_indices.data_ptr<int>(),
-        num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
-        bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
-  } else if (topk_indices.scalar_type() == at::ScalarType::UInt32) {
-    const uint32_t* input_ids_ptr = nullptr;
-    const uint32_t* tid2eid_ptr = nullptr;
-    if (tid2eid.has_value()) {
-      input_ids_ptr = input_ids.value().data_ptr<uint32_t>();
-      tid2eid_ptr = tid2eid.value().data_ptr<uint32_t>();
-    }
-    vllm::moe::topkGatingSoftplusSqrtKernelLauncher<uint32_t, ComputeType>(
-        gating_output, topk_weights.data_ptr<float>(),
-        topk_indices.data_ptr<uint32_t>(), token_expert_indices.data_ptr<int>(),
-        num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
-        bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
-  } else {
-    TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long);
-
-    const int64_t* input_ids_ptr = nullptr;
-    const int64_t* tid2eid_ptr = nullptr;
-    if (tid2eid.has_value()) {
-      input_ids_ptr = input_ids.value().data_ptr<int64_t>();
-      tid2eid_ptr = tid2eid.value().data_ptr<int64_t>();
-    }
-
-    vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int64_t, ComputeType>(
-        gating_output, topk_weights.data_ptr<float>(),
-        topk_indices.data_ptr<int64_t>(), token_expert_indices.data_ptr<int>(),
-        num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
-        bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
-  }
-}
-
-void topk_softplus_sqrt(
-    torch::Tensor& topk_weights,          // [num_tokens, topk]
-    torch::Tensor& topk_indices,          // [num_tokens, topk]
-    torch::Tensor& token_expert_indices,  // [num_tokens, topk]
-    torch::Tensor& gating_output,         // [num_tokens, num_experts]
-    bool renormalize, double routed_scaling_factor,
-    const c10::optional<torch::Tensor>& correction_bias,
-    const c10::optional<torch::Tensor>& input_ids,
-    const c10::optional<torch::Tensor>& tid2eid) {
-  const int num_experts = gating_output.size(-1);
-  const auto num_tokens = gating_output.numel() / num_experts;
-  const int topk = topk_weights.size(-1);
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-
-  if (gating_output.scalar_type() == at::ScalarType::Float) {
-    dispatch_topk_softplus_sqrt_launch<float>(
-        gating_output.data_ptr<float>(), topk_weights, topk_indices,
-        token_expert_indices, num_tokens, num_experts, topk, renormalize,
-        routed_scaling_factor, correction_bias, input_ids, tid2eid, stream);
-  } else if (gating_output.scalar_type() == at::ScalarType::Half) {
-    dispatch_topk_softplus_sqrt_launch<__half>(
-        reinterpret_cast<const __half*>(gating_output.data_ptr<at::Half>()),
-        topk_weights, topk_indices, token_expert_indices, num_tokens,
-        num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
-        input_ids, tid2eid, stream);
-  } else if (gating_output.scalar_type() == at::ScalarType::BFloat16) {
-    dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>(
-        reinterpret_cast<const __nv_bfloat16*>(
-            gating_output.data_ptr<at::BFloat16>()),
-        topk_weights, topk_indices, token_expert_indices, num_tokens,
-        num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
-        input_ids, tid2eid, stream);
-  } else {
-    TORCH_CHECK(false, "Unsupported gating_output data type: ",
-                gating_output.scalar_type());
-  }
-}

csrc/moe/torch_bindings.cpp

@@ -16,14 +16,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
       "bias) -> ()");
   m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);
 
-#ifndef USE_ROCM
-  m.def(
-      "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
-      "token_expert_indices, Tensor gating_output, bool renormalize, float "
-      "routed_scaling_factor, Tensor? "
-      "bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
-  m.impl("topk_softplus_sqrt", torch::kCUDA, &topk_softplus_sqrt);
-#endif
   // Calculate the result of moe by summing up the partial results
   // from all selected experts.
   m.def("moe_sum(Tensor input, Tensor! output) -> ()");

csrc/ops.h

@@ -100,11 +100,6 @@ void fused_qk_norm_rope(torch::Tensor& qkv, int64_t num_heads_q,
                         bool is_neox, torch::Tensor& position_ids,
                         int64_t forced_token_heads_per_warp);
 
-void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::Tensor& q, torch::Tensor const& kv, torch::Tensor& k_cache,
-    torch::Tensor const& slot_mapping, torch::Tensor const& position_ids,
-    torch::Tensor const& cos_sin_cache, double eps, int64_t cache_block_size);
-
 void apply_repetition_penalties_(torch::Tensor& logits,
                                  const torch::Tensor& prompt_mask,
                                  const torch::Tensor& output_mask,
@@ -158,13 +153,10 @@ void silu_and_mul_per_block_quant(torch::Tensor& out,
 
 void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
                       std::optional<torch::Tensor> key, int64_t head_size,
-                      torch::Tensor& cos_sin_cache, bool is_neox,
-                      int64_t rope_dim_offset, bool inverse);
+                      torch::Tensor& cos_sin_cache, bool is_neox);
 
 void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
 
-void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
-
 void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
                         torch::Tensor& scale);
 

csrc/persistent_topk.cuh

@@ -18,6 +18,7 @@ namespace persistent {
 // Constants
 // ============================================================================
 
+constexpr int TopK = 2048;
 constexpr int kThreadsPerBlock = 1024;
 constexpr int RADIX = 256;
 
@@ -127,12 +128,11 @@ struct RadixRowState {
 
 struct PersistentTopKParams {
   const float* __restrict__ input;  // [num_rows, stride]
-  int32_t* __restrict__ output;     // [num_rows, top_k]
+  int32_t* __restrict__ output;     // [num_rows, TopK]
   int32_t* __restrict__ lengths;    // [num_rows]
   RadixRowState* row_states;        // large path: per-group state
   uint32_t num_rows;
   uint32_t stride;
-  uint32_t top_k;           // actual k value for output stride
   uint32_t chunk_size;      // large path: elements per CTA
   uint32_t ctas_per_group;  // 1=medium, >1=large
   uint32_t max_seq_len;     // max seq_len across all rows (for early CTA exit)
@@ -154,7 +154,6 @@ __device__ __forceinline__ uint32_t decode_bin(float x) {
   return key >> 5;
 }
 
-template <int TopK>
 __device__ __noinline__ void histogram_2048_topk(
     const float* __restrict__ logits, int32_t* __restrict__ output_indices,
     int32_t seq_len) {
@@ -419,7 +418,6 @@ __device__ __noinline__ void histogram_2048_topk(
 // by: DarkSharpness
 // which at the same time is an optimized topk kernel copied from tilelang
 // kernel
-template <int TopK>
 __device__ __noinline__ void histogram_256_topk(
     const float* __restrict__ logits, int* __restrict__ output_indices,
     int logits_offset, int seq_len) {
@@ -651,7 +649,7 @@ __device__ __forceinline__ void wait_ge(int* ptr, int target_val,
 // Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215
 // ============================================================================
 
-template <int TopK, uint32_t VEC_SIZE>
+template <uint32_t VEC_SIZE>
 __device__ void radix_topk(const float* __restrict__ row_input,
                            int32_t* __restrict__ row_output, uint32_t seq_len,
                            uint32_t my_chunk_start, uint32_t chunk_size,
@@ -859,7 +857,7 @@ __device__ void radix_topk(const float* __restrict__ row_input,
 // see filtered_topk.cuh)
 // ============================================================================
 
-template <int TopK = 2048, uint32_t VEC_SIZE = 1>
+template <uint32_t VEC_SIZE = 1>
 __global__ void __launch_bounds__(kThreadsPerBlock, 2)
     persistent_topk_kernel(PersistentTopKParams params) {
   const uint32_t tx = threadIdx.x;
@@ -917,7 +915,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
     if (row_idx >= params.num_rows) break;
 
     const uint32_t seq_len = params.lengths[row_idx];
-    int32_t* row_output = params.output + row_idx * params.top_k;
+    int32_t* row_output = params.output + row_idx * TopK;
     const float* row_input = params.input + row_idx * params.stride;
 
     if (seq_len <= RADIX_THRESHOLD) {
@@ -929,19 +927,19 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
             row_output[i] = (i < seq_len) ? static_cast<int32_t>(i) : -1;
           }
         } else if (seq_len <= static_cast<uint32_t>(HIST2048_THRESHOLD)) {
-          histogram_2048_topk<TopK>(row_input, row_output, seq_len);
+          histogram_2048_topk(row_input, row_output, seq_len);
         } else {
-          histogram_256_topk<TopK>(row_input, row_output, 0, seq_len);
+          histogram_256_topk(row_input, row_output, 0, seq_len);
         }
       }
       continue;
     }
 
     const uint32_t my_chunk_start = cta_in_group * chunk_size;
-    radix_topk<TopK, VEC_SIZE>(
-        row_input, row_output, seq_len, my_chunk_start, chunk_size,
-        local_histogram, suffix_sum, shared_scalars, shared_ordered, state,
-        cta_in_group, ctas_per_group, barrier_phase, iter, tx);
+    radix_topk<VEC_SIZE>(row_input, row_output, seq_len, my_chunk_start,
+                         chunk_size, local_histogram, suffix_sum,
+                         shared_scalars, shared_ordered, state, cta_in_group,
+                         ctas_per_group, barrier_phase, iter, tx);
   }
 }
 
@@ -1013,6 +1011,7 @@ struct FilteredTopKTraits<float> {
   }
 };
 
+constexpr uint32_t FILTERED_TOPK_MAX_K = 2048;
 constexpr uint32_t FILTERED_TOPK_BLOCK_THREADS = 1024;
 constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE =
     16 * 1024;  // 16K indices per buffer
@@ -1026,7 +1025,7 @@ constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC =
  * \tparam IdType Index type (int32_t)
  * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
  */
-template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048>
+template <typename DType, typename IdType, int VEC_SIZE>
 __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
     FilteredTopKUnifiedKernel(const DType* __restrict__ input,
                               IdType* __restrict__ output,
@@ -1060,7 +1059,7 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
   alignas(128) __shared__ int s_counter;
   alignas(128) __shared__ int s_threshold_bin_id;
   alignas(128) __shared__ int s_num_input[2];
-  alignas(128) __shared__ int s_indices[MAX_K];
+  alignas(128) __shared__ int s_indices[FILTERED_TOPK_MAX_K];
 
   auto& s_histogram = s_histogram_buf[0];
 
@@ -1281,7 +1280,7 @@ constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) {
   return static_cast<int>(g);
 }
 
-template <typename DType, typename IdType, uint32_t MAX_K = 2048>
+template <typename DType, typename IdType>
 cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
                                         IdType* lengths, uint32_t num_rows,
                                         uint32_t top_k_val, uint32_t max_len,
@@ -1298,7 +1297,7 @@ cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
 
 #define DISPATCH_VEC_SIZE(VS)                                               \
   if (vec_size == VS) {                                                     \
-    auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS, MAX_K>;      \
+    auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS>;             \
     FLASHINFER_CUDA_CALL(cudaFuncSetAttribute(                              \
         kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));   \
     FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \

csrc/pos_encoding_kernels.cu

@@ -9,29 +9,28 @@ namespace vllm {
 
 template <typename scalar_t, bool IS_NEOX>
 inline __device__ void apply_token_rotary_embedding(
-    scalar_t* __restrict__ arr, const float* __restrict__ cos_ptr,
-    const float* __restrict__ sin_ptr, int rot_offset, int embed_dim,
-    const bool inverse) {
+    scalar_t* __restrict__ arr, const scalar_t* __restrict__ cos_ptr,
+    const scalar_t* __restrict__ sin_ptr, int rot_offset, int embed_dim) {
   int x_index, y_index;
-  float cos_f, sin_f;
+  scalar_t cos, sin;
   if (IS_NEOX) {
+    // GPT-NeoX style rotary embedding.
     x_index = rot_offset;
     y_index = embed_dim + rot_offset;
-    cos_f = VLLM_LDG(cos_ptr + x_index);
-    sin_f = VLLM_LDG(sin_ptr + x_index);
+    cos = VLLM_LDG(cos_ptr + x_index);
+    sin = VLLM_LDG(sin_ptr + x_index);
   } else {
+    // GPT-J style rotary embedding.
     x_index = 2 * rot_offset;
     y_index = 2 * rot_offset + 1;
-    cos_f = VLLM_LDG(cos_ptr + x_index / 2);
-    sin_f = VLLM_LDG(sin_ptr + x_index / 2);
+    cos = VLLM_LDG(cos_ptr + x_index / 2);
+    sin = VLLM_LDG(sin_ptr + x_index / 2);
   }
-  if (inverse) {
-    sin_f = -sin_f;
-  }
-  const float x_f = static_cast<float>(arr[x_index]);
-  const float y_f = static_cast<float>(arr[y_index]);
-  arr[x_index] = static_cast<scalar_t>(x_f * cos_f - y_f * sin_f);
-  arr[y_index] = static_cast<scalar_t>(y_f * cos_f + x_f * sin_f);
+
+  const scalar_t x = arr[x_index];
+  const scalar_t y = arr[y_index];
+  arr[x_index] = x * cos - y * sin;
+  arr[y_index] = y * cos + x * sin;
 }
 
 template <typename scalar_t, bool IS_NEOX>
@@ -43,23 +42,22 @@ inline __device__ void apply_rotary_embedding(
                                    // [batch_size, seq_len, num_kv_heads,
                                    // head_size] or [num_tokens, num_kv_heads,
                                    // head_size]
-    const float* cache_ptr, const int head_size, const int num_heads,
+    const scalar_t* cache_ptr, const int head_size, const int num_heads,
     const int num_kv_heads, const int rot_dim, const int token_idx,
     const int64_t query_stride, const int64_t key_stride,
-    const int64_t head_stride, const int64_t rope_dim_offset,
-    const bool inverse) {
+    const int64_t head_stride) {
   const int embed_dim = rot_dim / 2;
-  const float* cos_ptr = cache_ptr;
-  const float* sin_ptr = cache_ptr + embed_dim;
+  const scalar_t* cos_ptr = cache_ptr;
+  const scalar_t* sin_ptr = cache_ptr + embed_dim;
 
   const int nq = num_heads * embed_dim;
   for (int i = threadIdx.x; i < nq; i += blockDim.x) {
     const int head_idx = i / embed_dim;
     const int64_t token_head =
-        token_idx * query_stride + head_idx * head_stride + rope_dim_offset;
+        token_idx * query_stride + head_idx * head_stride;
     const int rot_offset = i % embed_dim;
     apply_token_rotary_embedding<scalar_t, IS_NEOX>(
-        query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
+        query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
   }
 
   if (key != nullptr) {
@@ -67,10 +65,10 @@ inline __device__ void apply_rotary_embedding(
     for (int i = threadIdx.x; i < nk; i += blockDim.x) {
       const int head_idx = i / embed_dim;
       const int64_t token_head =
-          token_idx * key_stride + head_idx * head_stride + rope_dim_offset;
+          token_idx * key_stride + head_idx * head_stride;
       const int rot_offset = i % embed_dim;
       apply_token_rotary_embedding<scalar_t, IS_NEOX>(
-          key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
+          key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
     }
   }
 }
@@ -86,18 +84,19 @@ __global__ void rotary_embedding_kernel(
                                  // [batch_size, seq_len, num_kv_heads,
                                  // head_size] or [num_tokens, num_kv_heads,
                                  // head_size]
-    const float* __restrict__ cos_sin_cache,  // [max_position, rot_dim] fp32
+    const scalar_t* __restrict__ cos_sin_cache,  // [max_position, 2, rot_dim //
+                                                 // 2]
     const int rot_dim, const int64_t query_stride, const int64_t key_stride,
     const int64_t head_stride, const int num_heads, const int num_kv_heads,
-    const int head_size, const int64_t rope_dim_offset, const bool inverse) {
+    const int head_size) {
+  // Each thread block is responsible for one token.
   const int token_idx = blockIdx.x;
   int64_t pos = positions[token_idx];
-  const float* cache_ptr = cos_sin_cache + pos * rot_dim;
+  const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim;
 
   apply_rotary_embedding<scalar_t, IS_NEOX>(
       query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim,
-      token_idx, query_stride, key_stride, head_stride, rope_dim_offset,
-      inverse);
+      token_idx, query_stride, key_stride, head_stride);
 }
 
 }  // namespace vllm
@@ -116,7 +115,7 @@ void rotary_embedding(
     // [num_tokens, num_heads, head_size]
     int64_t head_size,
     torch::Tensor& cos_sin_cache,  // [max_position, rot_dim]
-    bool is_neox, int64_t rope_dim_offset, bool inverse) {
+    bool is_neox) {
   // num_tokens = batch_size * seq_len
   int64_t num_tokens = positions.numel();
   int positions_ndim = positions.dim();
@@ -155,8 +154,6 @@ void rotary_embedding(
   int seq_dim_idx = positions_ndim - 1;
   int64_t query_stride = query.stride(seq_dim_idx);
   int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0;
-
-  TORCH_CHECK((rot_dim + rope_dim_offset) <= head_size);
   // Determine head stride: for [*, heads, head_size] use stride of last dim;
   // for flat [*, heads*head_size], heads blocks are contiguous of size
   // head_size
@@ -168,23 +165,20 @@ void rotary_embedding(
   dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
   const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  auto cache_f32 = cos_sin_cache.to(torch::kFloat32);
   VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "rotary_embedding", [&] {
     if (is_neox) {
       vllm::rotary_embedding_kernel<scalar_t, true><<<grid, block, 0, stream>>>(
           positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
           key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
-          cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
-          head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
-          inverse);
+          cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride, key_stride,
+          head_stride, num_heads, num_kv_heads, head_size);
     } else {
       vllm::rotary_embedding_kernel<scalar_t, false>
           <<<grid, block, 0, stream>>>(
               positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
               key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
-              cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
-              head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
-              inverse);
+              cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride,
+              key_stride, head_stride, num_heads, num_kv_heads, head_size);
     }
   });
 }

csrc/sampler.cu

@@ -258,13 +258,7 @@ __device__ bool processHistogramStep(
   auto processBins = [&](float logit, int idx) {
     if (isPartialMatch<patternShift>(logit, logitPattern)) {
       uint32_t binIdx = extractBinIdx<step>(logit);
-      // Only write elements with binIdx < thresholdBinIdx when:
-      // 1. This is step 0 and the threshold bin is small enough (no step 1)
-      // 2. This is step >= 1 (where pattern matching filters correctly)
-      // This prevents duplicates when step 0 and step 1 both run.
-      bool shouldWriteDirectly =
-          (step == 0 && smemFinalBinSize[0] <= kNumFinalItems) || (step >= 1);
-      if (binIdx < thresholdBinIdx && shouldWriteDirectly) {
+      if (binIdx < thresholdBinIdx) {
         // The element is part of the top-k selection
         int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1);
 

csrc/topk.cu

@@ -10,17 +10,33 @@
   #include "persistent_topk.cuh"
 #endif
 
-namespace {
-
+void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
+                     torch::Tensor& output, torch::Tensor& workspace, int64_t k,
+                     int64_t max_seq_len) {
 #ifndef USE_ROCM
-template <int TopK>
-void launch_persistent_topk(const torch::Tensor& logits,
-                            const torch::Tensor& lengths, torch::Tensor& output,
-                            torch::Tensor& workspace, int64_t max_seq_len) {
-  namespace P = vllm::persistent;
+  TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
+  TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
+  TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
+  TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
+  TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
+  TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
+  TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
+  TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
+              "lengths must be 1D or 2D");
+  TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
+  TORCH_CHECK(output.dim() == 2, "output must be 2D");
 
   const int64_t num_rows = logits.size(0);
   const int64_t stride = logits.size(1);
+
+  TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
+  TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
+              "output size mismatch");
+  namespace P = vllm::persistent;
+
+  TORCH_CHECK(k == P::TopK, "k must be 2048");
+  TORCH_CHECK(k <= stride, "k out of range");
+
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
   static int num_sms = 0;
@@ -34,17 +50,18 @@ void launch_persistent_topk(const torch::Tensor& logits,
   }
 
   if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
-    cudaError_t status =
-        vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
-            logits.data_ptr<float>(), output.data_ptr<int32_t>(),
-            lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
-            static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride), stream);
+    cudaError_t status = vllm::FilteredTopKRaggedTransform<float, int32_t>(
+        logits.data_ptr<float>(), output.data_ptr<int32_t>(),
+        lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
+        static_cast<uint32_t>(k), static_cast<uint32_t>(stride), stream);
     TORCH_CHECK(status == cudaSuccess,
                 "FilteredTopK failed: ", cudaGetErrorString(status));
   } else {
     TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
     TORCH_CHECK(workspace.dtype() == torch::kUInt8, "workspace must be uint8");
 
+    // Smem cap: smaller smem → more CTAs/group → more per-row parallelism for
+    // large path. Empirically tuned.
     int effective_max_smem;
     if (num_rows <= 4) {
       effective_max_smem =
@@ -82,73 +99,18 @@ void launch_persistent_topk(const torch::Tensor& logits,
     size_t smem_size = P::kFixedSmemLarge + chunk_size * sizeof(uint32_t);
     if (smem_size < P::kSmemMedium) smem_size = P::kSmemMedium;
 
-    // Query occupancy for the instantiation that will actually launch;
-    // overestimating it deadlocks the cooperative barrier.
     int occupancy = 1;
-    cudaError_t occ_err = cudaSuccess;
-    if (vec_size == 4) {
-      occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
-          smem_size);
-    } else if (vec_size == 2) {
-      occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &occupancy, P::persistent_topk_kernel<TopK, 2>, P::kThreadsPerBlock,
-          smem_size);
-    } else {
-      occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &occupancy, P::persistent_topk_kernel<TopK, 1>, P::kThreadsPerBlock,
-          smem_size);
-    }
-    TORCH_CHECK(occ_err == cudaSuccess,
-                "persistent_topk occupancy query failed: ",
-                cudaGetErrorString(occ_err));
+    cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+        &occupancy, P::persistent_topk_kernel<4>, P::kThreadsPerBlock,
+        smem_size);
     if (occupancy < 1) occupancy = 1;
 
-    // The cooperative spin-wait barrier only runs when at least one row hits
-    // the radix path (seq_len > RADIX_THRESHOLD). Below that, non-CTA-0 CTAs
-    // early-exit, so oversubscription can't deadlock and headroom is wasted.
-    const bool needs_cooperative =
-        static_cast<uint32_t>(max_seq_len) > P::RADIX_THRESHOLD;
-
-    const uint32_t hw_resident_cap =
-        static_cast<uint32_t>(num_sms) * static_cast<uint32_t>(occupancy);
-    uint32_t max_resident_ctas = hw_resident_cap;
-    if (needs_cooperative) {
-      // Reserve one CTA per SM when occupancy allows; fall back to a single
-      // CTA when occupancy == 1 (the most deadlock-prone case — any straggler
-      // kernel that takes the only slot on one SM hangs the barrier). Never
-      // drop below one full group's worth.
-      uint32_t headroom = (occupancy > 1) ? static_cast<uint32_t>(num_sms) : 1u;
-      if (max_resident_ctas >= headroom + ctas_per_group) {
-        max_resident_ctas -= headroom;
-      }
-    }
+    uint32_t max_resident_ctas = static_cast<uint32_t>(num_sms) * occupancy;
     uint32_t num_groups = std::min(max_resident_ctas / ctas_per_group,
                                    static_cast<uint32_t>(num_rows));
     if (num_groups == 0) num_groups = 1;
     uint32_t total_ctas = num_groups * ctas_per_group;
 
-    // If the cooperative launch wouldn't fit, fall back to FilteredTopK
-    // instead of deadlocking. Only relevant when needs_cooperative.
-    if (needs_cooperative && total_ctas > hw_resident_cap) {
-      TORCH_CHECK(max_smem_per_block >= 128 * 1024,
-                  "persistent_topk would oversubscribe and the FilteredTopK "
-                  "fallback requires >=128KB smem per block (have ",
-                  max_smem_per_block, "). total_ctas=", total_ctas,
-                  " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
-                  ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
-                  ", smem=", smem_size, ").");
-      cudaError_t status =
-          vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
-              logits.data_ptr<float>(), output.data_ptr<int32_t>(),
-              lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
-              static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride),
-              stream);
-      TORCH_CHECK(status == cudaSuccess,
-                  "FilteredTopK fallback failed: ", cudaGetErrorString(status));
-      return;
-    }
-
     size_t state_bytes = num_groups * sizeof(P::RadixRowState);
     TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
                 "workspace too small, need ", state_bytes, " bytes");
@@ -159,16 +121,15 @@ void launch_persistent_topk(const torch::Tensor& logits,
     params.lengths = lengths.data_ptr<int32_t>();
     params.num_rows = static_cast<uint32_t>(num_rows);
     params.stride = static_cast<uint32_t>(stride);
-    params.top_k = static_cast<uint32_t>(TopK);
     params.chunk_size = chunk_size;
     params.row_states =
         reinterpret_cast<P::RadixRowState*>(workspace.data_ptr<uint8_t>());
     params.ctas_per_group = ctas_per_group;
     params.max_seq_len = static_cast<uint32_t>(max_seq_len);
 
-  #define LAUNCH_PERSISTENT(TOPK_VAL, VS)                                     \
+  #define LAUNCH_PERSISTENT(VS)                                               \
     do {                                                                      \
-      auto kernel = &P::persistent_topk_kernel<TOPK_VAL, VS>;                 \
+      auto kernel = &P::persistent_topk_kernel<VS>;                           \
       cudaError_t err = cudaFuncSetAttribute(                                 \
           kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);    \
       TORCH_CHECK(err == cudaSuccess,                                         \
@@ -177,11 +138,11 @@ void launch_persistent_topk(const torch::Tensor& logits,
     } while (0)
 
     if (vec_size == 4) {
-      LAUNCH_PERSISTENT(TopK, 4);
+      LAUNCH_PERSISTENT(4);
     } else if (vec_size == 2) {
-      LAUNCH_PERSISTENT(TopK, 2);
+      LAUNCH_PERSISTENT(2);
     } else {
-      LAUNCH_PERSISTENT(TopK, 1);
+      LAUNCH_PERSISTENT(1);
     }
   #undef LAUNCH_PERSISTENT
   }
@@ -189,46 +150,6 @@ void launch_persistent_topk(const torch::Tensor& logits,
   cudaError_t err = cudaGetLastError();
   TORCH_CHECK(err == cudaSuccess,
               "persistent_topk failed: ", cudaGetErrorString(err));
-}
-#endif
-
-}  // anonymous namespace
-
-void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
-                     torch::Tensor& output, torch::Tensor& workspace, int64_t k,
-                     int64_t max_seq_len) {
-#ifndef USE_ROCM
-  TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
-  TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
-  TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
-  TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
-  TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
-  TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
-  TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
-  TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
-              "lengths must be 1D or 2D");
-  TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
-  TORCH_CHECK(output.dim() == 2, "output must be 2D");
-
-  const int64_t num_rows = logits.size(0);
-  const int64_t stride = logits.size(1);
-
-  TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
-  TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
-              "output size mismatch");
-  TORCH_CHECK(k == 512 || k == 1024 || k == 2048,
-              "persistent_topk supports k=512, k=1024, or k=2048, got k=", k);
-
-  if (k == 512) {
-    launch_persistent_topk<512>(logits, lengths, output, workspace,
-                                max_seq_len);
-  } else if (k == 1024) {
-    launch_persistent_topk<1024>(logits, lengths, output, workspace,
-                                 max_seq_len);
-  } else {
-    launch_persistent_topk<2048>(logits, lengths, output, workspace,
-                                 max_seq_len);
-  }
 #else
   TORCH_CHECK(false, "persistent_topk is not supported on ROCm");
 #endif

csrc/torch_bindings.cpp

@@ -106,12 +106,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
   ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);
 
-  // SwiGLU activation with input clamping.
-  ops.def(
-      "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
-      "-> ()");
-  ops.impl("silu_and_mul_with_clamp", torch::kCUDA, &silu_and_mul_clamp);
-
   ops.def(
       "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
   ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
@@ -183,19 +177,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "int forced_token_heads_per_warp=-1) -> ()");
   ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);
 
-#ifndef USE_ROCM
-  // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
-  // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
-  // kernel launch.
-  ops.def(
-      "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
-      "Tensor! q, Tensor kv, Tensor! k_cache, "
-      "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
-      "float eps, int cache_block_size) -> ()");
-  ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
-           &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
-#endif
-
   // Apply repetition penalties to logits in-place
   ops.def(
       "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -259,8 +240,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def(
       "rotary_embedding(Tensor positions, Tensor! query,"
       "                 Tensor!? key, int head_size,"
-      "                 Tensor cos_sin_cache, bool is_neox, int "
-      "rope_dim_offset=0, bool inverse=False) -> ()");
+      "                 Tensor cos_sin_cache, bool is_neox) -> ()");
   ops.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);
 
   // Quantization ops

docker/Dockerfile

@@ -22,7 +22,7 @@
 #   docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json
 # =============================================================================
 
-ARG CUDA_VERSION=13.0.2
+ARG CUDA_VERSION=13.0.0
 ARG PYTHON_VERSION=3.12
 ARG UBUNTU_VERSION=22.04
 
@@ -478,6 +478,9 @@ FROM ${FINAL_BASE_IMAGE} AS vllm-base
 
 ARG CUDA_VERSION
 ARG PYTHON_VERSION
+ARG DEADSNAKES_MIRROR_URL
+ARG DEADSNAKES_GPGKEY_URL
+ARG GET_PIP_URL
 
 ENV DEBIAN_FRONTEND=noninteractive
 WORKDIR /vllm-workspace
@@ -487,35 +490,43 @@ WORKDIR /vllm-workspace
 RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
     echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment
 
-# Install Python (via uv / python-build-standalone) and system dependencies.
-# This replaces the deadsnakes PPA, removing the build-time dependency on
-# Launchpad and matching how the build-stage (`base`) installs Python.
-# python-build-standalone bundles dev headers, the venv module, and
-# python3-config, so the python3.X-dev / python3.X-venv apt packages
-# are not needed.
+# Install Python and system dependencies
 RUN apt-get update -y \
     && apt-get install -y --no-install-recommends \
+        software-properties-common \
         curl \
         sudo \
         ffmpeg \
         libsm6 \
         libxext6 \
         libgl1 \
+    && if [ ! -z ${DEADSNAKES_MIRROR_URL} ] ; then \
+        if [ ! -z "${DEADSNAKES_GPGKEY_URL}" ] ; then \
+            mkdir -p -m 0755 /etc/apt/keyrings ; \
+            curl -L ${DEADSNAKES_GPGKEY_URL} | gpg --dearmor > /etc/apt/keyrings/deadsnakes.gpg ; \
+            sudo chmod 644 /etc/apt/keyrings/deadsnakes.gpg ; \
+            echo "deb [signed-by=/etc/apt/keyrings/deadsnakes.gpg] ${DEADSNAKES_MIRROR_URL} $(lsb_release -cs) main" > /etc/apt/sources.list.d/deadsnakes.list ; \
+        fi ; \
+    else \
+        for i in 1 2 3; do \
+            add-apt-repository -y ppa:deadsnakes/ppa && break || \
+            { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
+        done ; \
+    fi \
+    && apt-get update -y \
+    && apt-get install -y --no-install-recommends \
+        python${PYTHON_VERSION} \
+        python${PYTHON_VERSION}-dev \
+        python${PYTHON_VERSION}-venv \
         libibverbs-dev \
     && rm -rf /var/lib/apt/lists/* \
-    && curl -LsSf https://astral.sh/uv/install.sh | sh \
-    && $HOME/.local/bin/uv venv /opt/venv --python ${PYTHON_VERSION} \
-    && rm -f /usr/bin/python3 /usr/bin/python3-config /usr/bin/pip \
-    && ln -s /opt/venv/bin/python3 /usr/bin/python3 \
-    && ln -s /opt/venv/bin/python${PYTHON_VERSION} /usr/bin/python${PYTHON_VERSION} \
-    && ln -s /opt/venv/bin/python3-config /usr/bin/python3-config \
-    && ln -s /opt/venv/bin/pip /usr/bin/pip \
+    && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
+    && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
+    && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
+    && rm -f /usr/lib/python${PYTHON_VERSION}/EXTERNALLY-MANAGED \
+    && curl -sS ${GET_PIP_URL} | python${PYTHON_VERSION} \
     && python3 --version && python3 -m pip --version
 
-# Activate virtual environment and add uv to PATH
-ENV PATH="/opt/venv/bin:/root/.local/bin:$PATH"
-ENV VIRTUAL_ENV="/opt/venv"
-
 # Install CUDA development tools for runtime JIT compilation
 # (FlashInfer, DeepGEMM, EP kernels all require compilation at runtime)
 RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
@@ -529,9 +540,7 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
         libcurand-dev-${CUDA_VERSION_DASH} \
         libcublas-${CUDA_VERSION_DASH} \
         # Required by fastsafetensors (fixes #20384)
-        libnuma-dev \
-        # numactl CLI for NUMA binding at runtime
-        numactl && \
+        libnuma-dev && \
     # Fixes nccl_allocator requiring nccl.h at runtime
     # https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22
     # NCCL packages don't use the cuda-MAJOR-MINOR naming convention,
@@ -540,6 +549,9 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
     apt-get install -y --no-install-recommends --allow-change-held-packages libnccl-dev=${NCCL_VER} libnccl2=${NCCL_VER} && \
     rm -rf /var/lib/apt/lists/*
 
+# Install uv for faster pip installs
+RUN python3 -m pip install uv
+
 # Environment for uv
 ENV UV_HTTP_TIMEOUT=500
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
@@ -729,7 +741,7 @@ ENV HF_XET_HIGH_PERFORMANCE 1
 ENV HF_HUB_DOWNLOAD_TIMEOUT 60
 
 # Copy in the v1 package for testing (it isn't distributed yet)
-COPY vllm/v1 /opt/venv/lib/python${PYTHON_VERSION}/site-packages/vllm/v1
+COPY vllm/v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
 
 # Source code is used in the `python_only_compile.sh` test
 # We hide it inside `src/` so that this source code

docker/Dockerfile.nightly_torch

@@ -77,7 +77,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system $pkgs --index-url https://download.pytorch.org/whl/nightly/cu128
 
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system numba==0.65.0
+    uv pip install --system numba==0.61.2
 
 RUN --mount=type=cache,target=/root/.cache/uv \
 uv pip install --system -r requirements/common.txt

docker/Dockerfile.rocm

@@ -124,10 +124,10 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
 
 # RIXL/UCX build stages
 FROM base AS build_rixl
-ARG RIXL_BRANCH="bf4a7214"
+ARG RIXL_BRANCH="f33a5599"
 ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
-ARG UCX_BRANCH="7009d7a1"
-ARG UCX_REPO="https://github.com/openucx/ucx.git"
+ARG UCX_BRANCH="da3fac2a"
+ARG UCX_REPO="https://github.com/ROCm/ucx.git"
 ENV ROCM_PATH=/opt/rocm
 ENV UCX_HOME=/usr/local/ucx
 ENV RIXL_HOME=/usr/local/rixl
@@ -165,7 +165,7 @@ RUN cd /usr/local/src && \
         --disable-doxygen-doc \
         --enable-optimizations \
         --enable-devel-headers \
-        --with-rocm=${ROCM_PATH} \
+        --with-rocm=/opt/rocm \
         --with-verbs \
         --with-dm \
         --enable-mt && \
@@ -186,12 +186,7 @@ RUN git clone ${RIXL_REPO} /opt/rixl && \
     ninja install
 
 # Generate RIXL wheel
-# Exclude libcore and libpull from auditwheel: transitive dependencies
-# that are not shipped in the wheel and vary across base images.
-RUN cd /opt/rixl && \
-    sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
-        contrib/build-wheel.sh && \
-    mkdir -p /app/install && \
+RUN cd /opt/rixl && mkdir -p /app/install && \
     ./contrib/build-wheel.sh \
         --output-dir /app/install \
         --rocm-dir ${ROCM_PATH} \
@@ -436,10 +431,6 @@ COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-pac
 ENV MIOPEN_DEBUG_CONV_DIRECT=0
 ENV MIOPEN_DEBUG_CONV_GEMM=0
 
-# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
-# See: https://github.com/ROCm/rocm-libraries/issues/6266
-ENV HSA_ENABLE_IPC_MODE_LEGACY=1
-
 # Source code is used in the `python_only_compile.sh` test
 # We hide it inside `src/` so that this source code
 # will not be imported by other tests

docker/versions.json

@@ -2,7 +2,7 @@
   "_comment": "Auto-generated from Dockerfile ARGs. Do not edit manually. Run: python tools/generate_versions_json.py",
   "variable": {
     "CUDA_VERSION": {
-      "default": "13.0.2"
+      "default": "13.0.0"
     },
     "PYTHON_VERSION": {
       "default": "3.12"
@@ -11,10 +11,10 @@
       "default": "22.04"
     },
     "BUILD_BASE_IMAGE": {
-      "default": "nvidia/cuda:13.0.2-devel-ubuntu22.04"
+      "default": "nvidia/cuda:13.0.0-devel-ubuntu22.04"
     },
     "FINAL_BASE_IMAGE": {
-      "default": "nvidia/cuda:13.0.2-base-ubuntu22.04"
+      "default": "nvidia/cuda:13.0.0-base-ubuntu22.04"
     },
     "GET_PIP_URL": {
       "default": "https://bootstrap.pypa.io/get-pip.py"

docs/benchmarking/cli.md

@@ -108,38 +108,6 @@ P99 ITL (ms):                            8.39
 ==================================================
 ```
 
-#### Results Visualization
-
-The `--plot-timeline` and `--plot-dataset-stats` can be used to generate respectively the requests completion timeline and dataset prompt and output tokens statistics, which can be useful for debugging purpose or for deeper analysis.
-
-```bash
-vllm bench serve \
-    --backend vllm \
-    --model meta-llama/Llama-3.1-8B-Instruct \
-    --endpoint /v1/completions \
-    --dataset-name sharegpt \
-    --dataset-path <your data path>/ShareGPT_V3_unfiltered_cleaned_split.json \
-    --num-prompts 100 \
-    --plot-timeline \
-    --timeline-itl-thresholds 2,5 \
-    --plot-dataset-stats \
-    --save-result
-```
-
-##### Interactive Timeline
-
-The generated timeline is an interactive visualization in the form of an HTML file that can be rendered in most browsers. To customize the ITL color thresholds, one can use `--timeline-itl-thresholds` flag (default: 25ms, 50ms)
-
-Example output:
-
-<iframe src="../../assets/contributing/vllm_bench_serve_timeline.html" width="100%" height="600" frameborder="0"></iframe>
-
-##### Dataset statistics
-
-The generated figure shows the input prompt and output tokens distribution.
-
-Example output: ![Dataset Statistics](../assets/contributing/vllm_bench_serve_dataset_stats.png)
-
 #### Custom Dataset
 
 If the dataset you want to benchmark is not supported yet in vLLM, even then you can benchmark on it using `CustomDataset`. Your data needs to be in `.jsonl` format and needs to have "prompt" field per entry, e.g., data.jsonl

docs/design/attention_backends.md

@@ -172,7 +172,7 @@ Priority is **1 = highest** (tried first).
 | `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
 | `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
-| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ❌ | ✅ | All | ≥10.0 |
+| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
 | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
 | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
 | `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
@@ -213,7 +213,7 @@ configuration.
 | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
 | `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
 | `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
-| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
+| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
 | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
 | `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |

docs/design/moe_kernel_features.md

@@ -36,7 +36,7 @@ th {
 | deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] |
 | deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] |
 | flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
-| flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
+| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
 
 !!! info "Table key"
     1. All types: mxfp4, nvfp4, int4, int8, fp8
@@ -83,8 +83,8 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
 | triton | standard | all<sup>1</sup> | G,A,T | silu, gelu,</br>swigluoai,</br>silu_no_mul,</br>gelu_no_mul | Y | Y | [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts],</br>[`TritonExperts`][vllm.model_executor.layers.fused_moe.fused_moe.TritonExperts] |
 | triton (batched) | batched | all<sup>1</sup> | G,A,T | silu, gelu | <sup>6</sup> | Y | [`BatchedTritonExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedTritonExperts] |
 | deep gemm | standard,</br>batched | fp8 | G(128),A,T | silu, gelu | <sup>6</sup> | Y | </br>[`DeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe.DeepGemmExperts],</br>[`BatchedDeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe.BatchedDeepGemmExperts] |
-| cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp4] |
-| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassBatchedExpertsFp8] |
+| 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.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] |

docs/features/context_extension.md

@@ -1,70 +0,0 @@
-# Context Extension
-
-!!! note
-    The `--rope-scaling` parameter used in older versions of vLLM is no longer supported. Please use the `--hf-overrides` method with `rope_parameters` instead.
-This directory contains examples for extending the context length of models using vLLM.
-
-## Offline Inference Example
-
-The [`context_extension.py`](../../examples/offline_inference/context_extension) script demonstrates how to extend the context length of a Qwen model using the YARN method (rope_parameters) and run a simple chat example.
-
-### Usage
-
-```bash
-python examples/offline_inference/context_extension.py
-```
-
-## OpenAI Online Method
-
-You can also use vLLM's OpenAI-compatible API to serve models with extended context length.
-
-### Usage
-
-Run the vLLM server with the following command to extend the context length using YARN:
-
-```bash
-vllm serve Qwen/Qwen3-0.6B \
-  --hf-overrides '{"rope_parameters": {"factor": 4.0, "original_max_position_embeddings": 32768, "rope_theta": 1000000, "rope_type": "yarn"}}' \
-  --max-model-len 131072
-```
-
-### Client Example
-
-After starting the server, you can use the OpenAI Python client to interact with it:
-
-```python
-from openai import OpenAI
-
-client = OpenAI(
-    base_url="http://localhost:8000/v1",
-    api_key="token-abc123"  # Dummy API key, required by the client
-)
-
-response = client.chat.completions.create(
-    model="Qwen/Qwen3-0.6B",
-    messages=[
-        {"role": "system", "content": "You are a helpful assistant"},
-        {"role": "user", "content": "Hello"}
-    ],
-    max_tokens=128,
-    temperature=0.8,
-    top_p=0.95
-)
-
-print(response.choices[0].message.content)
-```
-
-### Key Parameters
-
-The available parameters depend on the `rope_type` you choose. For detailed information about all supported RoPE types and their specific parameters, please refer to the [Hugging Face Transformers RoPE documentation](https://huggingface.co/docs/transformers/main/en/internal/rope_utils#transformers.RopeParameters).
-
-Common parameters include:
-
-- `rope_type`: The type of RoPE implementation (e.g., "yarn", "linear", "dynamic")
-- `factor`: The factor by which to extend the context length
-- `original_max_position_embeddings`: The original maximum position embeddings of the model
-
-The following parameters are specific to vLLM:
-
-- `max_model_len`: The new maximum sequence length after extension (original * factor).
-  Used for KV cache pre‑allocation and request limit at serving time.

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

@@ -375,8 +375,8 @@ For (G)B300, we recommend using CUDA 13, as shown in the following command.
 
     ```bash
     DOCKER_BUILDKIT=1 docker build \
-    --build-arg CUDA_VERSION=13.0.2 \
-    --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
+    --build-arg CUDA_VERSION=13.0.1 \
+    --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 \
     --build-arg max_jobs=256 \
     --build-arg nvcc_threads=2 \
     --build-arg RUN_WHEEL_CHECK=false \

docs/models/pooling_models/README.md

@@ -78,7 +78,7 @@ The scoring models is designed to compute similarity scores between two input pr
 |-----------------------|---------------|----------------------------------------------|--------------------|--------------------------|
 | `classify` (see note) | Sequence-wise | reranker score for each sequence             | `cross-encoder`    | linear classifier        |
 | `embed`               | Sequence-wise | vector representations for each sequence     | `bi-encoder`       | cosine similarity        |
-| `token_classify`      | Token-wise    | probability vector of classes for each token | N/A                | N/A                      |
+| `token_classify`      | Token-wise    | probability vector of classes for each token | nan                | nan                      |
 | `token_embed`         | Token-wise    | vector representations for each token        | `late-interaction` | late interaction(MaxSim) |
 
 !!! note
@@ -86,15 +86,14 @@ The scoring models is designed to compute similarity scores between two input pr
 
 ### Pooling Usages
 
-| Pooling Usages              | Description                                                                                                                                               |
-|-----------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------|
-| Classification Usages       | Predicting which predefined category, class, or label best corresponds to a given input.                                                                  |
-| Embedding Usages            | Converts unstructured data (text, images, audio, etc.) into structured numerical vectors (embeddings).                                                    |
-| Token Classification Usages | Token-wise classification                                                                                                                                 |
-| Token Embedding Usages      | Token-wise embedding                                                                                                                                      |
-| Reward Usages               | Evaluates the quality of outputs generated by a language model, acting as a proxy for human preferences.                                                  |
-| Scoring Usages              | Computes similarity scores between two inputs. It supports three model types (aka `score_type`): `cross-encoder`, `late-interaction`, and `bi-encoder`.   |
-| Plugins Usages              | Allow users to customize input and output processors. For more information, please refer to [IO Processor Plugins](../../design/io_processor_plugins.md). |
+| Pooling Usages              | Description                                                                                                                                             |
+|-----------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------|
+| Classification Usages       | Predicting which predefined category, class, or label best corresponds to a given input.                                                                |
+| Embedding Usages            | Converts unstructured data (text, images, audio, etc.) into structured numerical vectors (embeddings).                                                  |
+| Token Classification Usages | Token-wise classification                                                                                                                               |
+| Token Embedding Usages      | Token-wise embedding                                                                                                                                    |
+| Scoring Usages              | Computes similarity scores between two inputs. It supports three model types (aka `score_type`): `cross-encoder`, `late-interaction`, and `bi-encoder`. |
+| Reward Usages               | Evaluates the quality of outputs generated by a language model, acting as a proxy for human preferences.                                                |
 
 We also have some special models that support multiple pooling tasks, or have specific usage scenarios, or support special inputs and outputs.
 
@@ -102,9 +101,9 @@ For more detailed information, please refer to the link below.
 
 - [Classification Usages](classify.md)
 - [Embedding Usages](embed.md)
+- [Reward Usages](reward.md)
 - [Token Classification Usages](token_classify.md)
 - [Token Embedding Usages](token_embed.md)
-- [Reward Usages](reward.md)
 - [Scoring Usages](scoring.md)
 - [Specific Model Examples](specific_models.md)
 
@@ -114,17 +113,15 @@ Each pooling model in vLLM supports one or more of these tasks according to
 [Pooler.get_supported_tasks][vllm.model_executor.layers.pooler.Pooler.get_supported_tasks],
 enabling the corresponding APIs.
 
-### Offline APIs corresponding to pooling usages
+### Offline APIs corresponding to pooling tasks
 
-| Pooling Usages              | Dedicated API       | Pooling task for `LLM.encode` API | Score Types                | scoring function         |
-|-----------------------------|---------------------|-----------------------------------|----------------------------|--------------------------|
-| Classification Usages       | `LLM.classify(...)` | `classify`                        | `cross-encoder` (see note) | linear classifier        |
-| Embedding Usages            | `LLM.embed(...)`    | `embed`                           | `bi-encoder`               | cosine similarity        |
-| Token Classification Usages | N/A                 | `token_classify`                  | N/A                        | N/A                      |
-| Token Embedding Usages      | N/A                 | `token_embed`                     | `late-interaction`         | late interaction(MaxSim) |
-| Reward Usages               | N/A                 | `classify` & `token_classify`     | N/A                        | N/A                      |
-| Scoring Usages              | `LLM.score(...)`    | N/A                               | N/A                        | N/A                      |
-| Plugins Usages              | N/A                 | `plugin`                          | N/A                        | N/A                      |
+| Task             | APIs                                                                                  |
+|------------------|---------------------------------------------------------------------------------------|
+| `embed`          | `LLM.embed(...)`, `LLM.encode(..., pooling_task="embed")`, `LLM.score(...)`(see note) |
+| `classify`       | `LLM.classify(...)`, `LLM.encode(..., pooling_task="classify")`, `LLM.score(...)`     |
+| `token_classify` | `LLM.reward(...)`, `LLM.encode(..., pooling_task="token_classify")`                   |
+| `token_embed`    | `LLM.encode(..., pooling_task="token_embed")`, `LLM.score(...)`                       |
+| `plugin`         | `LLM.encode(..., pooling_task="plugin")`                                              |
 
 !!! note
     Only when a classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
@@ -150,7 +147,7 @@ It is primarily designed for [score models](scoring.md).
 
 The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM.
 
-Please use one of the more specific methods or set the task directly when using `LLM.encode`, refer to the [table above](#offline-apis-corresponding-to-pooling-usages).
+Please use one of the more specific methods or set the task directly when using `LLM.encode`, refer to the [table above](#offline-apis-corresponding-to-pooling-tasks).
 
 ### Examples
 
@@ -186,12 +183,9 @@ Our Pooling API (`/pooling`) is similar to `LLM.encode`, being applicable to all
 
 The input format is the same as [Embeddings API](embed.md#openai-compatible-embeddings-api), but the output data can contain an arbitrary nested list, not just a 1-D list of floats.
 
-Please use one of the more specific APIs or set the task directly when using the Pooling API, refer to the [table above](#offline-apis-corresponding-to-pooling-usages).
+Please use one of the more specific APIs or set the task directly when using the Pooling API, refer to the [table above](#offline-apis-corresponding-to-pooling-tasks).
 
-Code examples:
-
-- [Online example](../../../examples/pooling/reward/token_reward_online.py)
-- [Offline example](../../../examples/pooling/reward/token_reward_offline.py)
+Code example: [examples/pooling/pooling/pooling_online.py](../../../examples/pooling/pooling/pooling_online.py)
 
 ### Examples
 
@@ -292,10 +286,10 @@ Pooling models now support token-wise task.
 
 ### Score task
 
-`score` task is deprecated and will be removed in v0.20. Please use `classify` instead. Only when a
-classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
+`score` task have has been removed in v0.21, use `classify` instead. Only when a classification model outputs num_labels
+equal to 1 can it be used as a scoring model and have its scoring API enabled.
 
 ### Pooling multitask support
 
-Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task is not what you want,
+Pooling multitask support has been removed in v0.21. When the default pooling task is not what you want,
 you need to manually specify it via `PoolerConfig(task=<task>)` offline or `--pooler-config.task <task>` online.

docs/models/pooling_models/reward.md

@@ -134,13 +134,3 @@ print(f"Data: {data!r}")
 ## Online Serving
 
 Please refer to the [pooling API](README.md#pooling-api). Pooling task corresponding to reward model types refer to the [table above](#summary).
-
-## More examples
-
-More examples can be found here: [examples/pooling/reward](../../../examples/pooling/reward)
-
-## Deprecated Features
-
-### `LLM.reward`
-
-`llm.reward` api is deprecated and will be removed in v0.23. Please use `LLM.encode` with `pooling_task="classify"` or `pooling_task="token_classify"` instead.

docs/models/supported_models.md

@@ -384,7 +384,6 @@ th {
 | `DeepseekForCausalLM` | DeepSeek | `deepseek-ai/deepseek-llm-67b-base`, `deepseek-ai/deepseek-llm-7b-chat`, etc. | ✅︎ | ✅︎ |
 | `DeepseekV2ForCausalLM` | DeepSeek-V2 | `deepseek-ai/DeepSeek-V2`, `deepseek-ai/DeepSeek-V2-Chat`, etc. | ✅︎ | ✅︎ |
 | `DeepseekV3ForCausalLM` | DeepSeek-V3 | `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`, etc. | ✅︎ | ✅︎ |
-| `DeepseekV4ForCausalLM` | DeepSeek-V4 | `deepseek-ai/DeepSeek-V4-Flash`, `deepseek-ai/DeepSeek-V4-Pro`, etc. | | |
 | `Dots1ForCausalLM` | dots.llm1 | `rednote-hilab/dots.llm1.base`, `rednote-hilab/dots.llm1.inst`, etc. | | ✅︎ |
 | `DotsOCRForCausalLM` | dots_ocr | `rednote-hilab/dots.ocr` | ✅︎ | ✅︎ |
 | `Ernie4_5ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | ✅︎ | ✅︎ |
@@ -644,10 +643,10 @@ Some models are supported only via the [Transformers modeling backend](#transfor
 !!! note
     `Gemma3nForConditionalGeneration` is only supported on V1 due to shared KV caching and it depends on `timm>=1.0.17` to make use of its
     MobileNet-v5 vision backbone.
-
+  
     Performance is not yet fully optimized mainly due to:
-
-    - Both audio and vision MM encoders use `transformers.AutoModel` implementation.
+  
+    - Both audio and vision MM encoders use `transformers.AutoModel` implementation.  
     - There's no PLE caching or out-of-memory swapping support, as described in [Google's blog](https://developers.googleblog.com/en/introducing-gemma-3n/). These features might be too model-specific for vLLM, and swapping in particular may be better suited for constrained setups.
 
 !!! note

examples/basic/offline_inference/reward.py

@@ -1,15 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Example offline usage of token reward models.
-
-The key distinction between sequence classification and token classification
-lies in their output granularity: sequence classification produces a single
-result for an entire input sequence, whereas token classification yields a
-result for each individual token within the sequence.
-"""
-
 from argparse import Namespace
 
 from vllm import LLM, EngineArgs
@@ -45,14 +36,14 @@ def main(args: Namespace):
     llm = LLM(**vars(args))
 
     # Generate rewards. The output is a list of PoolingRequestOutput.
-    outputs = llm.encode(prompts, pooling_task="token_classify")
+    outputs = llm.reward(prompts)
 
     # Print the outputs.
     print("\nGenerated Outputs:\n" + "-" * 60)
     for prompt, output in zip(prompts, outputs):
         rewards = output.outputs.data
         print(f"Prompt: {prompt!r}")
-        print_embeddings(rewards.tolist(), prefix="Reward")
+        print_embeddings(rewards, prefix="Reward")
         print("-" * 60)
 
 

examples/pooling/pooling/pooling_online.py

@@ -1,17 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """
-Example online usage of token reward models.
+Example online usage of Pooling API.
 
 Run `vllm serve <model> --runner pooling`
 to start up the server in vLLM. e.g.
 
 vllm serve internlm/internlm2-1_8b-reward --trust-remote-code
-
-The key distinction between sequence classification and token classification
-lies in their output granularity: sequence classification produces a single
-result for an entire input sequence, whereas token classification yields a
-result for each individual token within the sequence.
 """
 
 import argparse

examples/pooling/reward/sequence_reward_offline.py

@@ -1,62 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-Example offline usage of sequence reward models.
-
-The key distinction between sequence classification and token classification
-lies in their output granularity: sequence classification produces a single
-result for an entire input sequence, whereas token classification yields a
-result for each individual token within the sequence.
-"""
-
-from argparse import Namespace
-
-from vllm import LLM, EngineArgs
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-from vllm.utils.print_utils import print_embeddings
-
-
-def parse_args():
-    parser = FlexibleArgumentParser()
-    parser = EngineArgs.add_cli_args(parser)
-    # Set example specific arguments
-    parser.set_defaults(
-        model="Skywork/Skywork-Reward-V2-Qwen3-0.6B",
-        runner="pooling",
-        enforce_eager=True,
-        max_model_len=1024,
-        trust_remote_code=True,
-    )
-    return parser.parse_args()
-
-
-def main(args: Namespace):
-    # Sample prompts.
-    prompts = [
-        "Hello, my name is",
-        "The president of the United States is",
-        "The capital of France is",
-        "The future of AI is",
-    ]
-
-    # Create an LLM.
-    # You should pass runner="pooling" for reward models
-    llm = LLM(**vars(args))
-
-    # Generate rewards. The output is a list of PoolingRequestOutput.
-    # Use pooling_task="classify" for sequence reward models.
-    outputs = llm.encode(prompts, pooling_task="classify")
-
-    # Print the outputs.
-    print("\nGenerated Outputs:\n" + "-" * 60)
-    for prompt, output in zip(prompts, outputs):
-        rewards = output.outputs.data
-        print(f"Prompt: {prompt!r}")
-        print_embeddings(rewards.tolist(), prefix="Reward")
-        print("-" * 60)
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    main(args)

examples/pooling/reward/sequence_reward_online.py

@@ -1,71 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Example online usage of sequence reward models.
-
-Run `vllm serve <model> --runner pooling`
-to start up the server in vLLM. e.g.
-
-vllm serve Skywork/Skywork-Reward-V2-Qwen3-0.6B
-
-The key distinction between sequence classification and token classification
-lies in their output granularity: sequence classification produces a single
-result for an entire input sequence, whereas token classification yields a
-result for each individual token within the sequence.
-"""
-
-import argparse
-import pprint
-
-import requests
-
-
-def post_http_request(prompt: dict, api_url: str) -> requests.Response:
-    headers = {"User-Agent": "Test Client"}
-    response = requests.post(api_url, headers=headers, json=prompt)
-    return response
-
-
-def parse_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--host", type=str, default="localhost")
-    parser.add_argument("--port", type=int, default=8000)
-
-    return parser.parse_args()
-
-
-def main(args):
-    base_url = f"http://{args.host}:{args.port}"
-    models_url = base_url + "/v1/models"
-    pooing_url = base_url + "/pooling"
-
-    response = requests.get(models_url)
-    model = response.json()["data"][0]["id"]
-
-    # Input like Completions API
-    prompt = {"model": model, "input": "vLLM is great!"}
-    pooling_response = post_http_request(prompt=prompt, api_url=pooing_url)
-    print("-" * 50)
-    print("Pooling Response:")
-    pprint.pprint(pooling_response.json())
-    print("-" * 50)
-
-    # Input like Chat API
-    prompt = {
-        "model": model,
-        "messages": [
-            {
-                "role": "user",
-                "content": [{"type": "text", "text": "vLLM is great!"}],
-            }
-        ],
-    }
-    pooling_response = post_http_request(prompt=prompt, api_url=pooing_url)
-    print("Pooling Response:")
-    pprint.pprint(pooling_response.json())
-    print("-" * 50)
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    main(args)

examples/pooling/token_embed/jina_embeddings_v4_offline.py

@@ -4,68 +4,74 @@
 import torch
 
 from vllm import LLM
+from vllm.config import PoolerConfig
 from vllm.inputs import TextPrompt
 from vllm.multimodal.utils import fetch_image
 
-# Initialize model
-model = LLM(
-    model="jinaai/jina-embeddings-v4-vllm-text-matching",
-    runner="pooling",
-    max_model_len=1024,
-    gpu_memory_utilization=0.8,
-)
 
-# Create text prompts
-text1 = "Ein wunderschöner Sonnenuntergang am Strand"
-text1_prompt = TextPrompt(prompt=f"Query: {text1}")
+def main():
+    # Initialize model
+    model = LLM(
+        model="jinaai/jina-embeddings-v4-vllm-text-matching",
+        pooler_config=PoolerConfig(task="token_embed"),
+        runner="pooling",
+        max_model_len=1024,
+        gpu_memory_utilization=0.8,
+    )
 
-text2 = "浜辺に沈む美しい夕日"
-text2_prompt = TextPrompt(prompt=f"Query: {text2}")
+    # Create text prompts
+    text1 = "Ein wunderschöner Sonnenuntergang am Strand"
+    text1_prompt = TextPrompt(prompt=f"Query: {text1}")
 
-# Create image prompt
-image = fetch_image(
-    "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg"  # noqa: E501
-)
-image_prompt = TextPrompt(
-    prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n",  # noqa: E501
-    multi_modal_data={"image": image},
-)
+    text2 = "浜辺に沈む美しい夕日"
+    text2_prompt = TextPrompt(prompt=f"Query: {text2}")
 
-# Encode all prompts
-prompts = [text1_prompt, text2_prompt, image_prompt]
-outputs = model.encode(prompts, pooling_task="token_embed")
+    # Create image prompt
+    image = fetch_image(
+        "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg"  # noqa: E501
+    )
+    image_prompt = TextPrompt(
+        prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n",  # noqa: E501
+        multi_modal_data={"image": image},
+    )
+
+    # Encode all prompts
+    prompts = [text1_prompt, text2_prompt, image_prompt]
+    outputs = model.encode(prompts, pooling_task="token_embed")
+
+    def get_embeddings(outputs):
+        VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
+
+        embeddings = []
+        for output in outputs:
+            if VISION_START_TOKEN_ID in output.prompt_token_ids:
+                # Gather only vision tokens
+                img_start_pos = torch.where(
+                    torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
+                )[0][0]
+                img_end_pos = torch.where(
+                    torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
+                )[0][0]
+                embeddings_tensor = output.outputs.data.detach().clone()[
+                    img_start_pos : img_end_pos + 1
+                ]
+            else:
+                # Use all tokens for text-only prompts
+                embeddings_tensor = output.outputs.data.detach().clone()
+
+            # Pool and normalize embeddings
+            pooled_output = (
+                embeddings_tensor.sum(dim=0, dtype=torch.float32)
+                / embeddings_tensor.shape[0]
+            )
+            embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
+        return embeddings
+
+    embeddings = get_embeddings(outputs)
+
+    for embedding in embeddings:
+        print(embedding.shape)
 
 
-def get_embeddings(outputs):
-    VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
-
-    embeddings = []
-    for output in outputs:
-        if VISION_START_TOKEN_ID in output.prompt_token_ids:
-            # Gather only vision tokens
-            img_start_pos = torch.where(
-                torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
-            )[0][0]
-            img_end_pos = torch.where(
-                torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
-            )[0][0]
-            embeddings_tensor = output.outputs.data.detach().clone()[
-                img_start_pos : img_end_pos + 1
-            ]
-        else:
-            # Use all tokens for text-only prompts
-            embeddings_tensor = output.outputs.data.detach().clone()
-
-        # Pool and normalize embeddings
-        pooled_output = (
-            embeddings_tensor.sum(dim=0, dtype=torch.float32)
-            / embeddings_tensor.shape[0]
-        )
-        embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
-    return embeddings
-
-
-embeddings = get_embeddings(outputs)
-
-for embedding in embeddings:
-    print(embedding.shape)
+if __name__ == "__main__":
+    main()

examples/pooling/token_embed/multi_vector_retrieval_offline.py

@@ -4,6 +4,7 @@
 from argparse import Namespace
 
 from vllm import LLM, EngineArgs
+from vllm.config import PoolerConfig
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 
@@ -13,6 +14,7 @@ def parse_args():
     # Set example specific arguments
     parser.set_defaults(
         model="BAAI/bge-m3",
+        pooler_config=PoolerConfig(task="token_embed"),
         runner="pooling",
         enforce_eager=True,
     )
@@ -32,15 +34,6 @@ def main(args: Namespace):
     # You should pass runner="pooling" for embedding models
     llm = LLM(**vars(args))
 
-    # Generate embedding. The output is a list of EmbeddingRequestOutputs.
-    outputs = llm.embed(prompts)
-
-    # Print the outputs.
-    print("\nGenerated Outputs:\n" + "-" * 60)
-    for prompt, output in zip(prompts, outputs):
-        embeds = output.outputs.embedding
-        print(len(embeds))
-
     # Generate embedding for each token. The output is a list of PoolingRequestOutput.
     outputs = llm.encode(prompts, pooling_task="token_embed")
 
@@ -50,6 +43,20 @@ def main(args: Namespace):
         multi_vector = output.outputs.data
         print(multi_vector.shape)
 
+    query = "What is the capital of France?"
+    documents = [
+        "The capital of Brazil is Brasilia.",
+        "The capital of France is Paris.",
+    ]
+    # Generate scores.
+    outputs = llm.score(query, documents)
+    # Print the outputs.
+    print("\nGenerated Outputs:\n" + "-" * 60)
+    for document, output in zip(documents, outputs):
+        score = output.outputs.score
+        print(f"Pair: {[query, document]!r} \nScore: {score}")
+        print("-" * 60)
+
 
 if __name__ == "__main__":
     args = parse_args()

examples/pooling/token_embed/multi_vector_retrieval_online.py

@@ -7,10 +7,11 @@ Example online usage of Pooling API for multi vector retrieval.
 Run `vllm serve <model> --runner pooling`
 to start up the server in vLLM. e.g.
 
-vllm serve BAAI/bge-m3
+vllm serve BAAI/bge-m3 --pooler-config.task token_embed
 """
 
 import argparse
+import pprint
 
 import requests
 import torch
@@ -32,7 +33,8 @@ def parse_args():
 
 
 def main(args):
-    api_url = f"http://{args.host}:{args.port}/pooling"
+    pooling_url = f"http://{args.host}:{args.port}/pooling"
+    score_url = f"http://{args.host}:{args.port}/score"
     model_name = args.model
 
     prompts = [
@@ -43,11 +45,23 @@ def main(args):
     ]
     prompt = {"model": model_name, "input": prompts}
 
-    pooling_response = post_http_request(prompt=prompt, api_url=api_url)
+    pooling_response = post_http_request(prompt=prompt, api_url=pooling_url)
     for output in pooling_response.json()["data"]:
         multi_vector = torch.tensor(output["data"])
         print(multi_vector.shape)
 
+    queries = "What is the capital of France?"
+    documents = [
+        "The capital of Brazil is Brasilia.",
+        "The capital of France is Paris.",
+    ]
+    prompt = {"model": model_name, "queries": queries, "documents": documents}
+    score_response = post_http_request(prompt=prompt, api_url=score_url)
+    print("\nPrompt when queries is string and documents is a list:")
+    pprint.pprint(prompt)
+    print("\nScore Response:")
+    pprint.pprint(score_response.json())
+
 
 if __name__ == "__main__":
     args = parse_args()

pyproject.toml

@@ -24,14 +24,13 @@ classifiers = [
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
-    "Programming Language :: Python :: 3.14",
     "Intended Audience :: Developers",
     "Intended Audience :: Information Technology",
     "Intended Audience :: Science/Research",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
     "Topic :: Scientific/Engineering :: Information Analysis",
 ]
-requires-python = ">=3.10,<3.15"
+requires-python = ">=3.10,<3.14"
 dynamic = [ "version", "dependencies", "optional-dependencies"]
 
 [project.urls]
@@ -124,8 +123,7 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer
     "benchmarks/sonnet.txt", "tests/lora/data/*", "build/*",
     "examples/pooling/token_embed/*", "tests/models/language/pooling/*",
     "vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*", "tests/entrypoints/openai/speech_to_text/test_transcription_validation.py",
-    "docs/governance/process.md", "docs/assets/contributing/vllm_bench_serve_timeline.html", 
-    "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*"]
+    "docs/governance/process.md", "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*"]
 ignore-hidden = false
 
 [tool.typos.default]

requirements/common.txt

@@ -20,7 +20,7 @@ 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 == "ppc64le"
-outlines_core == 0.2.14
+outlines_core == 0.2.11
 # required for outlines backend disk cache
 diskcache == 5.6.3
 lark == 1.2.2

requirements/cpu.txt

@@ -4,7 +4,7 @@
 
 setuptools==77.0.3 # this version can reuse CMake build dir
 
-numba == 0.65.0; platform_machine != "s390x" # Required for N-gram speculative decoding
+numba == 0.61.2; platform_machine != "s390x" # Required for N-gram speculative decoding
 
 # Dependencies for CPUs
 torch==2.11.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64"

requirements/cuda.txt

@@ -1,7 +1,7 @@
 # Common dependencies
 -r common.txt
 
-numba == 0.65.0 # Required for N-gram speculative decoding
+numba == 0.61.2 # Required for N-gram speculative decoding
 
 # Dependencies for NVIDIA GPUs
 torch==2.11.0
@@ -11,8 +11,6 @@ torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytor
 # FlashInfer should be updated together with the Dockerfile
 flashinfer-python==0.6.8.post1
 flashinfer-cubin==0.6.8.post1
-apache-tvm-ffi==0.1.9
-tilelang==0.1.9
 # Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to
 # breaking changes in 1.19.0
 nvidia-cudnn-frontend>=1.13.0,<1.19.0

requirements/rocm.txt

@@ -5,7 +5,7 @@
 grpcio==1.78.0
 grpcio-reflection==1.78.0
 
-numba == 0.65.0 # Required for N-gram speculative decoding
+numba == 0.61.2 # Required for N-gram speculative decoding
 
 # Dependencies for AMD GPUs
 datasets
@@ -20,4 +20,4 @@ conch-triton-kernels==1.2.1
 timm>=1.0.17
 # amd-quark: required for Quark quantization on ROCm 
 # To be consistent with test_quark.py
-amd-quark>=0.8.99
+amd-quark>=0.8.99

requirements/test/cuda.in

@@ -54,7 +54,7 @@ grpcio==1.78.0
 grpcio-reflection==1.78.0
 
 arctic-inference == 0.1.1 # Required for suffix decoding test
-numba == 0.65.0 # Required for N-gram speculative decoding
+numba == 0.61.2 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs,azure]==0.15.7
 fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage

requirements/test/cuda.txt

@@ -479,7 +479,7 @@ lightning-utilities==0.14.3
     #   lightning
     #   pytorch-lightning
     #   torchmetrics
-llvmlite==0.47.0
+llvmlite==0.44.0
     # via numba
 lm-eval==0.4.11
     # via -r requirements/test/cuda.in
@@ -550,7 +550,7 @@ nltk==3.9.1
     # via rouge-score
 num2words==0.5.14
     # via -r requirements/test/cuda.in
-numba==0.65.0
+numba==0.61.2
     # via
     #   -c requirements/cuda.txt
     #   -r requirements/test/cuda.in

requirements/test/nightly-torch.txt

@@ -40,7 +40,7 @@ buildkite-test-collector==0.1.9
 genai_perf>=0.0.8
 tritonclient>=2.51.0
 
-numba == 0.65.0 # Required for N-gram speculative decoding
+numba == 0.61.2 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs,azure]==0.15.7
 fastsafetensors>=0.2.2

requirements/test/rocm.in

@@ -52,7 +52,7 @@ grpcio==1.78.0
 grpcio-reflection==1.78.0
 
 arctic-inference==0.1.1 # Required for suffix decoding test
-numba==0.65.0 # Required for N-gram speculative decoding
+numba==0.61.2 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs,azure]==0.15.7
 fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.2.2 # PyPI only ships CUDA wheels

requirements/test/rocm.txt

@@ -559,7 +559,7 @@ llguidance==1.3.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
-llvmlite==0.47.0
+llvmlite==0.44.0
     # via numba
 lm-eval==0.4.11
     # via -r requirements/test/rocm.in
@@ -655,7 +655,7 @@ nltk==3.9.3
     # via rouge-score
 num2words==0.5.14
     # via -r requirements/test/rocm.in
-numba==0.65.0
+numba==0.61.2
     # via
     #   -c requirements/rocm.txt
     #   -r requirements/test/rocm.in
@@ -811,7 +811,7 @@ orjson==3.11.7
     # via
     #   genai-perf
     #   kaleido
-outlines-core==0.2.14
+outlines-core==0.2.11
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt

requirements/test/xpu.txt

@@ -244,7 +244,7 @@ lazy-loader==0.5
     #   scikit-image
 librosa==0.10.2.post1
     # via -r requirements/test/xpu.in
-llvmlite==0.47.0
+llvmlite==0.44.0
     # via numba
 lm-eval==0.4.11
     # via -r requirements/test/xpu.in
@@ -304,7 +304,7 @@ nltk==3.9.4
     # via rouge-score
 num2words==0.5.14
     # via -r requirements/test/xpu.in
-numba==0.65.0
+numba==0.61.2
     # via
     #   -c requirements/xpu.txt
     #   librosa

requirements/xpu.txt

@@ -9,7 +9,7 @@ setuptools>=77.0.3,<81.0.0
 wheel
 jinja2>=3.1.6
 datasets # for benchmark scripts
-numba == 0.65.0 # Required for N-gram speculative decoding
+numba == 0.61.2 # Required for N-gram speculative decoding
 --extra-index-url=https://download.pytorch.org/whl/xpu
 torch==2.11.0+xpu
 torchaudio

setup.py

@@ -927,9 +927,7 @@ def get_vllm_version() -> str:
     elif _is_tpu():
         version += f"{sep}tpu"
     elif _is_cpu():
-        # Check the local VLLM_TARGET_DEVICE (may be set by auto-detect above),
-        # not envs.VLLM_TARGET_DEVICE, so CPU-only hosts still get `+cpu`.
-        if VLLM_TARGET_DEVICE == "cpu":
+        if envs.VLLM_TARGET_DEVICE == "cpu":
             version += f"{sep}cpu"
     elif _is_xpu():
         version += f"{sep}xpu"

tests/compile/correctness_e2e/test_sequence_parallel.py

@@ -261,8 +261,6 @@ def _compare_sp(
         },
         "use_inductor_graph_partition": use_inductor_graph_partition,
     }
-    if not use_inductor_graph_partition:
-        compilation_config["splitting_ops"] = []
 
     tp_sp_args = [
         *common_args,

tests/compile/fusions_e2e/conftest.py

@@ -116,11 +116,6 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
         model_kwargs["attention_config"] = {"backend": attn_backend.backend.name}
         model_kwargs["tensor_parallel_size"] = tp_size
 
-        # Cap warmup memory: tests use small max_model_len (1024) but the
-        # engine default max_num_batched_tokens is 16384. Warming up large
-        # models (e.g. Llama-4-Scout-FP8) at 16384 tokens may trigger OOM.
-        model_kwargs.setdefault("max_num_batched_tokens", 8192)
-
         # Sparse MLA models (DSv3.2) hit an over-strict inductor assertion in
         # decompose_auto_functionalized when +rotary_embedding is forced into
         # the compile graph. Disable qk_norm+rope fusion (which auto-enables

tests/compile/h100/test_startup.py

@@ -34,10 +34,7 @@ def _run_vllm(vllm_runner):
             mode=CompilationMode.VLLM_COMPILE,
             cudagraph_mode=CUDAGraphMode.NONE,
         ),
-        # Phi-tiny-MoE uses SWA, whose admission cap is `cdiv(L, block_size) + 1`
-        # at default block_size=16 — i.e. 17 blocks for max_model_len=256. Use
-        # 32 for headroom.
-        num_gpu_blocks_override=32,
+        num_gpu_blocks_override=8,
     ):
         pass
 
@@ -193,7 +190,7 @@ def _run_model(vllm_runner, spec: ModelStartupSpec):
             cudagraph_mode=CUDAGraphMode.NONE,
             pass_config=PassConfig(fuse_allreduce_rms=False),
         ),
-        num_gpu_blocks_override=16,
+        num_gpu_blocks_override=8,
     ):
         pass
 

tests/compile/passes/distributed/test_async_tp.py

@@ -19,7 +19,6 @@ from vllm.config import (
     VllmConfig,
     set_current_vllm_config,
 )
-from vllm.config.utils import Range
 from vllm.distributed import (
     tensor_model_parallel_all_gather,
     tensor_model_parallel_reduce_scatter,
@@ -289,22 +288,6 @@ def test_async_tp_pass_replace(
     run_torch_spawn(async_tp_pass_on_test_model, num_processes)
 
 
-def test_async_tp_pass_requires_full_graph_compilation():
-    vllm_config = VllmConfig()
-    vllm_config.compilation_config.use_inductor_graph_partition = False
-    vllm_config.compilation_config.splitting_ops = [
-        "vllm::unified_attention_with_output"
-    ]
-
-    async_tp_pass = object.__new__(AsyncTPPass)
-    async_tp_pass.compilation_config = vllm_config.compilation_config
-
-    with pytest.raises(
-        AssertionError, match="AsyncTPPass requires full-graph compilation"
-    ):
-        async_tp_pass.is_applicable_for_range(Range(start=8, end=8))
-
-
 def async_tp_pass_on_test_model(
     local_rank: int,
     world_size: int,

tests/compile/passes/distributed/test_sequence_parallelism.py

@@ -22,7 +22,6 @@ from vllm.config import (
     get_current_vllm_config,
     set_current_vllm_config,
 )
-from vllm.config.utils import Range
 from vllm.distributed import tensor_model_parallel_all_reduce
 from vllm.distributed.parallel_state import (
     init_distributed_environment,
@@ -217,24 +216,6 @@ def test_sequence_parallelism_pass(
     run_torch_spawn(sequence_parallelism_pass_on_test_model, num_processes)
 
 
-def test_sequence_parallelism_pass_requires_full_graph_compilation():
-    vllm_config = VllmConfig()
-    vllm_config.compilation_config.use_inductor_graph_partition = False
-    vllm_config.compilation_config.splitting_ops = [
-        "vllm::unified_attention_with_output"
-    ]
-
-    sequence_parallelism_pass = object.__new__(SequenceParallelismPass)
-    sequence_parallelism_pass.compilation_config = vllm_config.compilation_config
-    sequence_parallelism_pass.min_token_num = 1
-
-    with pytest.raises(
-        AssertionError,
-        match="SequenceParallelismPass requires full-graph compilation",
-    ):
-        sequence_parallelism_pass.is_applicable_for_range(Range(start=8, end=8))
-
-
 def sequence_parallelism_pass_on_test_model(
     local_rank: int,
     world_size: int,

tests/compile/passes/test_functionalization.py

@@ -117,9 +117,9 @@ class TestFusedAddRMSNorm(torch.nn.Module):
         else:
             return norm_output, residual_output
 
-    def example_inputs(self, batch_size=8, seq_len=16):
-        hidden_states = torch.randn((batch_size * seq_len, self.hidden_size))
-        residual = torch.randn((batch_size * seq_len, self.intermediate_size))
+    def example_inputs(self, batch_size=8, hidden_size=16, seq_len=16):
+        hidden_states = torch.randn((batch_size * seq_len, hidden_size))
+        residual = torch.randn((batch_size * seq_len, hidden_size))
         return (hidden_states, residual)
 
     def ops_in_model(self, do_fusion):

tests/compile/passes/test_silu_mul_quant_fusion.py

@@ -10,7 +10,7 @@ import vllm.envs as envs
 from tests.compile.backend import TestBackend
 from tests.kernels.quantization.nvfp4_utils import quant_nvfp4_tensor
 from tests.utils import TestFP8Layer
-from vllm._aiter_ops import IS_AITER_FOUND, rocm_aiter_ops
+from vllm._aiter_ops import IS_AITER_FOUND
 from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
 from vllm.compilation.passes.fusion.act_quant_fusion import (
     FUSED_OPS,
@@ -157,13 +157,12 @@ class TestSiluMulGroupFp8QuantModel(torch.nn.Module):
             activation_quant_key=self.act_quant_key,
             input_dtype=dtype,
         )
+        self.w = torch.rand(hidden_size, hidden_size).to(dtype=FP8_DTYPE).t()
 
-        if not current_platform.is_fp8_fnuz():
-            kernel = self.w8a8_block_fp8_linear.kernel
-            orig_quant = kernel.quant_fp8
-            kernel.quant_fp8 = lambda *a, use_triton=False, **kw: orig_quant(
-                *a, use_triton=True, **kw
-            )
+        scale_hidden_size = (hidden_size + 128 - 1) // 128
+        self.wscale = torch.rand(
+            (scale_hidden_size, scale_hidden_size), dtype=torch.float32
+        )
 
         self.enable_silu_mul_custom_op = self.silu_and_mul.enabled()
 
@@ -175,9 +174,6 @@ class TestSiluMulGroupFp8QuantModel(torch.nn.Module):
     def ops_in_model_before(self):
         return [
             SILU_MUL_OP if self.enable_silu_mul_custom_op else torch.ops.aten.mul,
-            rocm_aiter_ops.get_group_quant_op()
-            if current_platform.is_fp8_fnuz()
-            else torch.ops.vllm.triton_per_token_group_quant_fp8.default,
         ]
 
     def ops_in_model_after(self):
@@ -328,6 +324,7 @@ def test_fusion_silu_and_mul_quant(
     with set_current_vllm_config(config), monkeypatch.context() as m:
         fusion_passes = [ActivationQuantFusionPass(config)]
         if IS_AITER_FOUND and model_class is TestSiluMulGroupFp8QuantModel:
+            from vllm._aiter_ops import rocm_aiter_ops
             from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
                 RocmAiterSiluMulFp8GroupQuantFusionPass,
             )
@@ -355,16 +352,10 @@ def test_fusion_silu_and_mul_quant(
             atol, rtol = 1e-3, 1e-3
         elif isinstance(model, TestSiluMulNvfp4QuantModel):
             atol, rtol = 1e-1, 1e-1
-        elif isinstance(model, TestSiluMulGroupFp8QuantModel):
+        elif isinstance(
+            model, (TestSiluMulGroupFp8QuantModel, TestSiluMulBlockQuantModel)
+        ):
             atol, rtol = 5e-2, 5e-2
-        elif isinstance(model, TestSiluMulBlockQuantModel):
-            if current_platform.is_rocm():
-                atol, rtol = 1e-3, 1e-3
-            else:
-                # CUDA fused kernel computes silu*mul in fp32 while the reference
-                # goes through bf16/fp16 storage, so group maxima (and thus scales)
-                # can shift by one FP8-e4m3 code (~1/8 relative step).
-                atol, rtol = 5e-2, 5e-2
 
         torch.testing.assert_close(
             result[0].to(dtype=dtype), result2[0].to(dtype=dtype), atol=atol, rtol=rtol

tests/compile/test_config.py

@@ -205,22 +205,6 @@ def test_enforce_eager(vllm_runner, monkeypatch):
         pass
 
 
-@pytest.mark.forked
-def test_torch_compile_disable(vllm_runner, monkeypatch):
-    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
-    monkeypatch.setenv("TORCH_COMPILE_DISABLE", "1")
-    monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
-
-    with (
-        compilation_counter.expect(num_graphs_seen=0, stock_torch_compile_count=0),
-        vllm_runner(
-            "facebook/opt-125m",
-            gpu_memory_utilization=0.4,
-        ) as _,
-    ):
-        pass
-
-
 def test_splitting_ops_dynamic():
     # Default config
     config = VllmConfig()
@@ -405,12 +389,9 @@ def test_should_split():
         (None, 0, 1, False, 2048, CUDAGraphMode.NONE, 0),
         # truncated to nearest multiple of 8 or 16
         (None, 257, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256),
-        # max_num_batched_tokens <= max_cudagraph_capture_size should always be
-        # captured even if not landing on a 16-stride step
-        (None, 2048, 1, False, 257, CUDAGraphMode.FULL_AND_PIECEWISE, 257),
         # max from list
         ([1, 2, 4, 15], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 15),
-        # SP forces full-graph compilation, sizes are filtered by TP
+        # filtered out 15 due to SP
         ([1, 2, 4, 15], None, 2, True, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 4),
         # limited by the max_tokens
         ([1, 2, 4, 15], None, 1, False, 8, CUDAGraphMode.FULL_AND_PIECEWISE, 4),
@@ -468,123 +449,6 @@ def test_cudagraph_sizes_post_init(
         )
 
 
-@pytest.mark.skipif(
-    not current_platform.support_static_graph_mode(),
-    reason="Skip if not cudagraph mode supported",
-)
-@pytest.mark.parametrize(
-    (
-        "cudagraph_mode",
-        "use_inductor_graph_partition",
-        "expected_enable_sp",
-        "expected_cudagraph_mode",
-        "expected_piecewise_compile",
-        "expected_capture_sizes",
-        "expected_max_size",
-    ),
-    [
-        (CUDAGraphMode.PIECEWISE, False, True, CUDAGraphMode.FULL, False, [2, 4], 4),
-        (
-            CUDAGraphMode.FULL_DECODE_ONLY,
-            False,
-            True,
-            CUDAGraphMode.FULL_DECODE_ONLY,
-            False,
-            [2, 4],
-            4,
-        ),
-        (
-            CUDAGraphMode.FULL_AND_PIECEWISE,
-            False,
-            True,
-            CUDAGraphMode.FULL,
-            False,
-            [2, 4],
-            4,
-        ),
-        (
-            CUDAGraphMode.FULL_AND_PIECEWISE,
-            True,
-            True,
-            CUDAGraphMode.FULL_AND_PIECEWISE,
-            True,
-            [2, 4],
-            4,
-        ),
-    ],
-)
-def test_sequence_parallelism_requires_full_graph_compilation(
-    cudagraph_mode: CUDAGraphMode,
-    use_inductor_graph_partition: bool,
-    expected_enable_sp: bool,
-    expected_cudagraph_mode: CUDAGraphMode,
-    expected_piecewise_compile: bool,
-    expected_capture_sizes: list[int],
-    expected_max_size: int,
-):
-    with patch.object(current_platform, "device_count", return_value=2):
-        vllm_config = VllmConfig(
-            parallel_config=ParallelConfig(tensor_parallel_size=2),
-            scheduler_config=SchedulerConfig(
-                max_num_seqs=128,
-                max_num_batched_tokens=2048,
-                max_model_len=2048,
-                is_encoder_decoder=False,
-            ),
-        )
-        vllm_config.model_config = MagicMock(
-            dtype=torch.float16,
-            enforce_eager=False,
-            is_moe=False,
-            disable_cascade_attn=False,
-            get_hidden_size=MagicMock(return_value=4096),
-        )
-        vllm_config.compilation_config = CompilationConfig(
-            mode=CompilationMode.VLLM_COMPILE,
-            cudagraph_capture_sizes=[1, 2, 4, 15],
-            max_cudagraph_capture_size=None,
-            compile_sizes=["cudagraph_capture_sizes"],
-            use_inductor_graph_partition=use_inductor_graph_partition,
-            pass_config=PassConfig(
-                enable_sp=True,
-                fuse_gemm_comms=True,
-                fuse_norm_quant=True,
-                fuse_act_quant=True,
-                eliminate_noops=True,
-                sp_min_token_num=512,
-            ),
-            cudagraph_mode=cudagraph_mode,
-        )
-        vllm_config.compilation_config.set_splitting_ops_for_v1(
-            all2all_backend=vllm_config.parallel_config.all2all_backend,
-            data_parallel_size=1,
-        )
-        vllm_config._set_compile_ranges()
-        vllm_config._set_cudagraph_sizes()
-
-    assert (
-        vllm_config.compilation_config.use_inductor_graph_partition
-        == use_inductor_graph_partition
-    )
-    assert (
-        bool(vllm_config.compilation_config.splitting_ops) == expected_piecewise_compile
-    )
-    assert vllm_config.compilation_config.pass_config.enable_sp == expected_enable_sp
-    assert (
-        vllm_config.compilation_config.pass_config.fuse_gemm_comms == expected_enable_sp
-    )
-    assert vllm_config.compilation_config.cudagraph_mode == expected_cudagraph_mode
-    assert (
-        vllm_config.compilation_config.cudagraph_capture_sizes == expected_capture_sizes
-    )
-    assert (
-        vllm_config.compilation_config.max_cudagraph_capture_size == expected_max_size
-    )
-    assert (
-        511 in vllm_config.compilation_config.compile_ranges_endpoints
-    ) == expected_enable_sp
-
-
 def test_cached_compilation_config(default_vllm_config):
     import torch
     from torch._inductor.utils import run_and_get_code

tests/conftest.py

@@ -1183,7 +1183,7 @@ class VllmRunner:
         return [req_output.outputs.data for req_output in req_outputs]
 
     def reward(self, prompts: list[str]) -> list[list[float]]:
-        req_outputs = self.llm.encode(prompts, pooling_task="token_classify")
+        req_outputs = self.llm.reward(prompts)
         return [req_output.outputs.data for req_output in req_outputs]
 
     def score(

tests/distributed/test_eplb_execute.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import asyncio
 import random
 
 import pytest
@@ -360,14 +361,16 @@ def _test_async_transfer_layer_without_mtp_worker(
         communicator.set_stream(cuda_stream)
 
         for layer_idx in range(num_layers):
-            transfer_metadata = transfer_layer(
-                old_layer_indices=old_indices_cpu[layer_idx],
-                new_layer_indices=new_indices_cpu[layer_idx],
-                expert_weights=expert_weights[layer_idx],
-                expert_weights_buffer=expert_buffer,
-                ep_group=ep_group,
-                communicator=communicator,
-                cuda_stream=cuda_stream,
+            transfer_metadata = asyncio.run(
+                transfer_layer(
+                    old_layer_indices=old_indices_cpu[layer_idx],
+                    new_layer_indices=new_indices_cpu[layer_idx],
+                    expert_weights=expert_weights[layer_idx],
+                    expert_weights_buffer=expert_buffer,
+                    ep_group=ep_group,
+                    communicator=communicator,
+                    cuda_stream=cuda_stream,
+                )
             )
             cuda_stream.synchronize()
             move_from_buffer(

tests/entrypoints/openai/chat_completion/test_chat_error.py

@@ -164,58 +164,6 @@ async def test_chat_error_non_stream():
         await serving_chat.create_chat_completion(request)
 
 
-@pytest.mark.asyncio
-async def test_openai_chat_keeps_mm_cache_for_engine_execution():
-    mock_engine = MagicMock(spec=AsyncLLM)
-    mock_engine.errored = False
-    mock_engine.model_config = MockModelConfig()
-    mock_engine.input_processor = MagicMock()
-    mock_engine.renderer = _build_renderer(mock_engine.model_config)
-
-    serving_chat = _build_serving_chat(mock_engine)
-
-    request = ChatCompletionRequest(
-        model=MODEL_NAME,
-        messages=[{"role": "user", "content": "Test prompt"}],
-    )
-
-    result = await serving_chat.render_chat_request(request)
-
-    assert isinstance(result, tuple)
-    assert (
-        serving_chat.openai_serving_render.preprocess_chat.call_args.kwargs[
-            "skip_mm_cache"
-        ]
-        is False
-    )
-
-
-@pytest.mark.asyncio
-async def test_renderer_only_chat_request_skips_mm_cache():
-    mock_engine = MagicMock(spec=AsyncLLM)
-    mock_engine.errored = False
-    mock_engine.model_config = MockModelConfig()
-    mock_engine.input_processor = MagicMock()
-    mock_engine.renderer = _build_renderer(mock_engine.model_config)
-
-    serving_chat = _build_serving_chat(mock_engine)
-
-    request = ChatCompletionRequest(
-        model=MODEL_NAME,
-        messages=[{"role": "user", "content": "Test prompt"}],
-    )
-
-    result = await serving_chat.openai_serving_render.render_chat_request(request)
-
-    assert result.token_ids == [1, 2, 3]
-    assert (
-        serving_chat.openai_serving_render.preprocess_chat.call_args.kwargs[
-            "skip_mm_cache"
-        ]
-        is True
-    )
-
-
 @pytest.mark.asyncio
 async def test_chat_error_stream():
     """test finish_reason='error' returns 500 InternalServerError (streaming)"""

tests/entrypoints/openai/completion/test_completion_error.py

@@ -3,7 +3,7 @@
 
 from dataclasses import dataclass, field
 from typing import Any
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import MagicMock
 
 import pytest
 
@@ -148,66 +148,6 @@ async def test_completion_error_non_stream():
         await serving_completion.create_completion(request)
 
 
-@pytest.mark.asyncio
-async def test_openai_completion_keeps_mm_cache_for_engine_execution():
-    mock_engine = MagicMock(spec=AsyncLLM)
-    mock_engine.errored = False
-    mock_engine.model_config = MockModelConfig()
-    mock_engine.input_processor = MagicMock()
-    mock_engine.renderer = _build_renderer(mock_engine.model_config)
-
-    serving_completion = _build_serving_completion(mock_engine)
-    serving_completion.openai_serving_render.preprocess_completion = AsyncMock(
-        return_value=[{"prompt_token_ids": [1, 2, 3]}]
-    )
-
-    request = CompletionRequest(
-        model=MODEL_NAME,
-        prompt="Test prompt",
-    )
-
-    result = await serving_completion.render_completion_request(request)
-
-    assert isinstance(result, list)
-    assert (
-        serving_completion.openai_serving_render.preprocess_completion.call_args.kwargs[
-            "skip_mm_cache"
-        ]
-        is False
-    )
-
-
-@pytest.mark.asyncio
-async def test_renderer_only_completion_request_skips_mm_cache():
-    mock_engine = MagicMock(spec=AsyncLLM)
-    mock_engine.errored = False
-    mock_engine.model_config = MockModelConfig()
-    mock_engine.input_processor = MagicMock()
-    mock_engine.renderer = _build_renderer(mock_engine.model_config)
-
-    serving_completion = _build_serving_completion(mock_engine)
-    serving_completion.openai_serving_render.preprocess_completion = AsyncMock(
-        return_value=[{"prompt_token_ids": [1, 2, 3]}]
-    )
-
-    request = CompletionRequest(
-        model=MODEL_NAME,
-        prompt="Test prompt",
-    )
-
-    result = await serving_completion.openai_serving_render.render_completion_request(
-        request
-    )
-
-    assert isinstance(result, list)
-    assert (
-        serving_completion.openai_serving_render.preprocess_completion.call_args.kwargs[
-            "skip_mm_cache"
-        ]
-        is True
-    )
-
-
 @pytest.mark.asyncio
 async def test_completion_error_stream():
     """test finish_reason='error' returns 500 InternalServerError (streaming)"""

tests/entrypoints/openai/completion/test_shutdown.py

@@ -311,7 +311,7 @@ async def test_abort_timeout_exits_quickly(wait_for_engine_idle: float):
             pytest.fail("Process did not exit after SIGTERM with abort timeout")
 
         exit_time = time.time() - start_time
-        assert exit_time < 2.1, f"Default shutdown took too long: {exit_time:.1f}s"
+        assert exit_time < 2, f"Default shutdown took too long: {exit_time:.1f}s"
         assert proc.returncode in (0, -15, None), f"Unexpected: {proc.returncode}"
 
         await _assert_children_cleaned_up(child_pids)

tests/entrypoints/pooling/classify/test_offline.py

@@ -1,13 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import logging
 import weakref
 
 import pytest
 import torch
 
 from tests.models.utils import softmax
-from vllm import LLM, ClassificationRequestOutput, PoolingParams, PoolingRequestOutput
+from vllm import LLM, ClassificationRequestOutput, PoolingParams
 from vllm.distributed import cleanup_dist_env_and_memory
 from vllm.tasks import PoolingTask
 
@@ -66,18 +65,6 @@ def test_list_prompts(llm: LLM):
         assert len(outputs[i].outputs.probs) == num_labels
 
 
-@pytest.mark.skip_global_cleanup
-def test_token_classify(llm: LLM, caplog_vllm):
-    with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
-        outputs = llm.encode(prompt, pooling_task="token_classify", use_tqdm=False)
-        assert "deprecated" in caplog_vllm.text
-
-    assert len(outputs) == 1
-    assert isinstance(outputs[0], PoolingRequestOutput)
-    assert outputs[0].prompt_token_ids == prompt_token_ids
-    assert outputs[0].outputs.data.shape == (len(prompt_token_ids), num_labels)
-
-
 @pytest.mark.skip_global_cleanup
 def test_pooling_params(llm: LLM):
     def get_outputs(use_activation):
@@ -110,10 +97,12 @@ def test_score_api(llm: LLM):
         llm.score("ping", "pong", use_tqdm=False)
 
 
-@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
 def test_unsupported_tasks(llm: LLM, task: PoolingTask):
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "token_classify":
+        err_msg = "Try switching the model's pooling_task via.+"
     else:
         err_msg = "Embedding API is not supported by this model.+"
     with pytest.raises(ValueError, match=err_msg):

tests/entrypoints/pooling/classify/test_online.py

@@ -436,26 +436,7 @@ async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
-async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
-    task = "token_classify"
-    response = requests.post(
-        server.url_for("pooling"),
-        json={
-            "model": model_name,
-            "input": input_text,
-            "encoding_format": "float",
-            "task": task,
-        },
-    )
-    poolings = PoolingResponse.model_validate(response.json())
-    assert len(poolings.data) == 1
-    assert len(poolings.data[0].data) == 8
-    assert len(poolings.data[0].data[0]) == 2
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
 async def test_pooling_not_supported(
     server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -469,8 +450,11 @@ async def test_pooling_not_supported(
         },
     )
     assert response.json()["error"]["type"] == "BadRequestError"
+
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "token_classify":
+        err_msg = "Try switching the model's pooling_task via"
     else:
         err_msg = f"Unsupported task: {task!r}"
     assert response.json()["error"]["message"].startswith(err_msg)

tests/entrypoints/pooling/embed/test_offline.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import logging
 import weakref
 
 import pytest
@@ -38,11 +37,11 @@ def llm():
         seed=0,
         attention_config=attention_config,
     )
+    assert embedding_size == llm.model_config.embedding_size
 
     yield weakref.proxy(llm)
 
     del llm
-
     cleanup_dist_env_and_memory()
 
 
@@ -74,16 +73,6 @@ def test_list_prompts(llm: LLM):
         assert len(outputs[i].outputs.embedding) == embedding_size
 
 
-@pytest.mark.skip_global_cleanup
-def test_token_embed(llm: LLM, caplog_vllm):
-    with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
-        outputs = llm.encode(prompt, pooling_task="token_embed", use_tqdm=False)
-        assert "deprecated" in caplog_vllm.text
-
-    multi_vector = outputs[0].outputs.data
-    assert multi_vector.shape == (11, 384)
-
-
 @pytest.mark.skip_global_cleanup
 def test_pooling_params(llm: LLM):
     def get_outputs(normalize):
@@ -107,10 +96,14 @@ def test_pooling_params(llm: LLM):
     )
 
 
-@pytest.mark.parametrize("task", ["token_classify", "classify", "plugin"])
+@pytest.mark.parametrize(
+    "task", ["token_classify", "classify", "token_embed", "plugin"]
+)
 def test_unsupported_tasks(llm: LLM, task: PoolingTask):
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "token_embed":
+        err_msg = "Try switching the model's pooling_task via.+"
     else:
         err_msg = "Classification API is not supported by this model.+"
     with pytest.raises(ValueError, match=err_msg):

tests/entrypoints/pooling/embed/test_online.py

@@ -732,28 +732,9 @@ async def test_pooling_embed(server: RemoteOpenAIServer, model_name: str):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
-async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
-    task = "token_embed"
-    response = requests.post(
-        server.url_for("pooling"),
-        json={
-            "model": model_name,
-            "input": input_text,
-            "encoding_format": "float",
-            "task": task,
-        },
-    )
-
-    poolings = PoolingResponse.model_validate(response.json())
-
-    assert len(poolings.data) == 1
-    assert len(poolings.data[0].data) == len(input_tokens)
-    assert len(poolings.data[0].data[0]) == 384
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
+@pytest.mark.parametrize(
+    "task", ["classify", "token_classify", "token_embed", "plugin"]
+)
 async def test_pooling_not_supported(
     server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -769,6 +750,8 @@ async def test_pooling_not_supported(
     assert response.json()["error"]["type"] == "BadRequestError"
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "token_embed":
+        err_msg = "Try switching the model's pooling_task via"
     else:
         err_msg = f"Unsupported task: {task!r}"
     assert response.json()["error"]["message"].startswith(err_msg)

tests/entrypoints/pooling/scoring/test_cross_encoder_online.py

@@ -452,25 +452,6 @@ async def test_pooling_classify(server: RemoteOpenAIServer):
     assert len(poolings.data[0].data) == 1
 
 
-@pytest.mark.asyncio
-async def test_pooling_token_classify(server: RemoteOpenAIServer):
-    response = requests.post(
-        server.url_for("pooling"),
-        json={
-            "model": MODEL_NAME,
-            "task": "token_classify",
-            "input": input_text,
-            "encoding_format": "float",
-        },
-    )
-
-    poolings = PoolingResponse.model_validate(response.json())
-
-    assert len(poolings.data) == 1
-    assert len(poolings.data[0].data) == len(input_tokens)
-    assert len(poolings.data[0].data[0]) == 1
-
-
 @pytest.mark.asyncio
 async def test_rerank_max_tokens_per_doc(
     server: RemoteOpenAIServer,
@@ -544,7 +525,7 @@ async def test_rerank_max_tokens_per_doc_validation(
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
 async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
     response = requests.post(
         server.url_for("pooling"),
@@ -558,6 +539,8 @@ async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
     assert response.json()["error"]["type"] == "BadRequestError"
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "token_classify":
+        err_msg = "Try switching the model's pooling_task via"
     else:
         err_msg = f"Unsupported task: {task!r}"
     assert response.json()["error"]["message"].startswith(err_msg)

tests/entrypoints/pooling/token_classify/test_offline.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import logging
 import weakref
 
 import pytest
@@ -60,22 +59,19 @@ def test_token_ids_prompts(llm: LLM):
 
 @pytest.mark.skip_global_cleanup
 def test_score_api(llm: LLM):
-    err_msg = "Scoring API is only enabled for num_labels == 1."
+    err_msg = "This model does not support the Scoring API."
     with pytest.raises(ValueError, match=err_msg):
         llm.score("ping", "pong", use_tqdm=False)
 
 
 @pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
 def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
-    if task == "classify":
-        with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
-            llm.encode(prompt, pooling_task=task, use_tqdm=False)
-        assert "deprecated" in caplog_vllm.text
+    if task == "plugin":
+        err_msg = "No IOProcessor plugin installed."
+    elif task == "classify":
+        err_msg = "Try switching the model's pooling_task via.+"
     else:
-        if task == "plugin":
-            err_msg = "No IOProcessor plugin installed."
-        else:
-            err_msg = "Embedding API is not supported by this model.+"
+        err_msg = "Embedding API is not supported by this model.+"
 
-        with pytest.raises(ValueError, match=err_msg):
-            llm.encode(prompt, pooling_task=task, use_tqdm=False)
+    with pytest.raises(ValueError, match=err_msg):
+        llm.encode(prompt, pooling_task=task, use_tqdm=False)

tests/entrypoints/pooling/token_classify/test_online.py

@@ -50,7 +50,7 @@ async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: st
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
+@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
 async def test_pooling_not_supported(
     server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -63,9 +63,12 @@ async def test_pooling_not_supported(
             "task": task,
         },
     )
+    assert response.json()["error"]["type"] == "BadRequestError"
 
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "classify":
+        err_msg = "Try switching the model's pooling_task via"
     else:
         err_msg = f"Unsupported task: {task!r}"
     assert response.json()["error"]["message"].startswith(err_msg)

tests/entrypoints/pooling/token_embed/test_offline.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import logging
 import weakref
 
 import pytest
@@ -64,15 +63,12 @@ def test_token_ids_prompts(llm: LLM):
 
 @pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
 def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
-    if task == "embed":
-        with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
-            llm.encode(prompt, pooling_task=task, use_tqdm=False)
-        assert "deprecated" in caplog_vllm.text
+    if task == "plugin":
+        err_msg = "No IOProcessor plugin installed."
+    elif task == "embed":
+        err_msg = "Try switching the model's pooling_task via.+"
     else:
-        if task == "plugin":
-            err_msg = "No IOProcessor plugin installed."
-        else:
-            err_msg = "Classification API is not supported by this model.+"
+        err_msg = "Classification API is not supported by this model.+"
 
-        with pytest.raises(ValueError, match=err_msg):
-            llm.encode(prompt, pooling_task=task, use_tqdm=False)
+    with pytest.raises(ValueError, match=err_msg):
+        llm.encode(prompt, pooling_task=task, use_tqdm=False)

tests/entrypoints/pooling/token_embed/test_online.py

@@ -73,7 +73,7 @@ async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
 async def test_pooling_not_supported(
     server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -86,9 +86,12 @@ async def test_pooling_not_supported(
             "task": task,
         },
     )
+    assert response.json()["error"]["type"] == "BadRequestError"
 
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "embed":
+        err_msg = "Try switching the model's pooling_task via"
     else:
         err_msg = f"Unsupported task: {task!r}"
     assert response.json()["error"]["message"].startswith(err_msg)

tests/entrypoints/serve/disagg/test_generate_stream.py

@@ -12,10 +12,7 @@ from vllm.config.multimodal import MultiModalConfig
 from vllm.entrypoints.openai.engine.protocol import StreamOptions
 from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
-from vllm.entrypoints.serve.disagg.protocol import (
-    GenerateRequest,
-    GenerateResponse,
-)
+from vllm.entrypoints.serve.disagg.protocol import GenerateRequest
 from vllm.entrypoints.serve.disagg.serving import ServingTokens
 from vllm.entrypoints.serve.render.serving import OpenAIServingRender
 from vllm.logprobs import Logprob
@@ -167,36 +164,6 @@ def _parse_sse_chunks(chunks: list[str]) -> list[Any]:
     return parsed
 
 
-@pytest.mark.asyncio
-async def test_serve_tokens_skips_mm_cache_for_remote_engine_execution():
-    engine = _mock_engine()
-
-    async def mock_generate(*args, **kwargs):
-        yield _make_request_output(
-            "req-1", token_ids=[10], finish_reason="stop", finished=True
-        )
-
-    engine.generate = MagicMock(side_effect=mock_generate)
-    serving = _build_serving_tokens(engine)
-
-    request = GenerateRequest(
-        token_ids=[1, 2, 3],
-        sampling_params=SamplingParams(max_tokens=1),
-        model=MODEL_NAME,
-        stream=False,
-    )
-
-    response = await serving.serve_tokens(request)
-
-    assert isinstance(response, GenerateResponse)
-    assert (
-        serving.openai_serving_render.preprocess_completion.call_args.kwargs[
-            "skip_mm_cache"
-        ]
-        is True
-    )
-
-
 @pytest.mark.asyncio
 async def test_stream_basic():
     """Streaming returns SSE chunks with correct token_ids and ends with [DONE]."""

tests/entrypoints/serve/tokenize/test_serving_tokenization.py

@@ -1,140 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from dataclasses import dataclass, field
-from typing import Any
-from unittest.mock import AsyncMock, MagicMock
-
-import pytest
-
-from vllm.config.multimodal import MultiModalConfig
-from vllm.entrypoints.openai.models.protocol import BaseModelPath
-from vllm.entrypoints.openai.models.serving import OpenAIServingModels
-from vllm.entrypoints.serve.render.serving import OpenAIServingRender
-from vllm.entrypoints.serve.tokenize.protocol import (
-    TokenizeChatRequest,
-    TokenizeCompletionRequest,
-)
-from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
-from vllm.v1.engine.async_llm import AsyncLLM
-
-MODEL_NAME = "openai-community/gpt2"
-BASE_MODEL_PATHS = [
-    BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME),
-]
-
-
-@dataclass
-class MockHFConfig:
-    model_type: str = "any"
-
-
-@dataclass
-class MockModelConfig:
-    task = "generate"
-    runner_type = "generate"
-    model = MODEL_NAME
-    tokenizer = MODEL_NAME
-    trust_remote_code = False
-    tokenizer_mode = "auto"
-    max_model_len = 100
-    tokenizer_revision = None
-    multimodal_config = MultiModalConfig()
-    hf_config = MockHFConfig()
-    hf_text_config = MockHFConfig()
-    logits_processors: list[str] | None = None
-    diff_sampling_param: dict | None = None
-    allowed_local_media_path: str = ""
-    allowed_media_domains: list[str] | None = None
-    encoder_config = None
-    generation_config: str = "auto"
-    media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
-    skip_tokenizer_init = False
-    is_encoder_decoder: bool = False
-    is_multimodal_model: bool = False
-    renderer_num_workers: int = 1
-
-    def get_diff_sampling_param(self):
-        return self.diff_sampling_param or {}
-
-
-def _build_serving_tokenization(engine: AsyncLLM) -> OpenAIServingTokenization:
-    models = OpenAIServingModels(
-        engine_client=engine,
-        base_model_paths=BASE_MODEL_PATHS,
-    )
-    serving_render = OpenAIServingRender(
-        model_config=engine.model_config,
-        renderer=engine.renderer,
-        model_registry=models.registry,
-        request_logger=None,
-        chat_template=None,
-        chat_template_content_format="auto",
-    )
-    return OpenAIServingTokenization(
-        engine,
-        models,
-        openai_serving_render=serving_render,
-        request_logger=None,
-        chat_template=None,
-        chat_template_content_format="auto",
-    )
-
-
-@pytest.mark.asyncio
-async def test_tokenize_chat_skips_mm_cache_for_renderer_only_path():
-    mock_engine = MagicMock(spec=AsyncLLM)
-    mock_engine.errored = False
-    mock_engine.model_config = MockModelConfig()
-    mock_engine.input_processor = MagicMock()
-    mock_engine.renderer = MagicMock()
-
-    serving = _build_serving_tokenization(mock_engine)
-    serving.openai_serving_render.preprocess_chat = AsyncMock(
-        return_value=(
-            [{"role": "user", "content": "Test"}],
-            [{"prompt_token_ids": [1, 2, 3]}],
-        )
-    )
-
-    request = TokenizeChatRequest(
-        model=MODEL_NAME,
-        messages=[{"role": "user", "content": "Test prompt"}],
-    )
-
-    response = await serving.create_tokenize(request, MagicMock(headers={}))
-
-    assert response.tokens == [1, 2, 3]
-    assert (
-        serving.openai_serving_render.preprocess_chat.call_args.kwargs["skip_mm_cache"]
-        is True
-    )
-
-
-@pytest.mark.asyncio
-async def test_tokenize_completion_skips_mm_cache_for_renderer_only_path():
-    mock_engine = MagicMock(spec=AsyncLLM)
-    mock_engine.errored = False
-    mock_engine.model_config = MockModelConfig()
-    mock_engine.input_processor = MagicMock()
-    mock_engine.renderer = MagicMock()
-
-    serving = _build_serving_tokenization(mock_engine)
-    serving.openai_serving_render.preprocess_completion = AsyncMock(
-        return_value=[{"prompt_token_ids": [1, 2, 3]}]
-    )
-
-    request = TokenizeCompletionRequest(
-        model=MODEL_NAME,
-        prompt="Test prompt",
-    )
-
-    response = await serving.create_tokenize(request, MagicMock(headers={}))
-
-    assert response.tokens == [1, 2, 3]
-    assert (
-        serving.openai_serving_render.preprocess_completion.call_args.kwargs[
-            "skip_mm_cache"
-        ]
-        is True
-    )

tests/kernels/attention/test_deepgemm_attention.py

@@ -9,8 +9,8 @@ from vllm.platforms import current_platform
 from vllm.utils.deep_gemm import (
     _ceil_to_ue8m0,
     calc_diff,
-    fp8_fp4_mqa_logits,
-    fp8_fp4_paged_mqa_logits,
+    fp8_mqa_logits,
+    fp8_paged_mqa_logits,
     get_num_sms,
     get_paged_mqa_logits_metadata,
 )
@@ -127,8 +127,8 @@ def test_deepgemm_fp8_mqa_logits(clean_logits: bool):
 
                 q_fp8 = q.to(torch.float8_e4m3fn)
                 kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0,), False)
-                logits = fp8_fp4_mqa_logits(
-                    (q_fp8, None), kv_fp8, weights, ks, ke, clean_logits=clean_logits
+                logits = fp8_mqa_logits(
+                    q_fp8, kv_fp8, weights, ks, ke, clean_logits=clean_logits
                 )
 
                 ref_logits = _ref_fp8_mqa_logits(
@@ -150,7 +150,7 @@ def test_deepgemm_fp8_mqa_logits(clean_logits: bool):
                 assert diff < 1e-3, f"{diff=}"
 
 
-def _ref_fp8_fp4_paged_mqa_logits(
+def _ref_fp8_paged_mqa_logits(
     q: torch.Tensor,
     kv_cache: torch.Tensor,
     weights: torch.Tensor,
@@ -205,10 +205,8 @@ def _ref_fp8_fp4_paged_mqa_logits(
 @pytest.mark.skipif(
     not current_platform.has_device_capability(90), reason="SM90 and SM100 only"
 )
-def test_deepgemm_fp8_fp4_paged_mqa_logits():
-    # NOTE: clean_logits=True is incompatible with the 2D context_lens
-    # required by csrc/apis/attention.hpp; only the False path is exercised.
-    clean_logits = False
+@pytest.mark.parametrize("clean_logits", [True, False])
+def test_deepgemm_fp8_paged_mqa_logits(clean_logits: bool):
     torch.manual_seed(0)
     random.seed(0)
 
@@ -260,29 +258,21 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits():
                 q_fp8 = q.to(torch.float8_e4m3fn)
                 kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache)
 
-                # deep_gemm paged MQA logits requires 2D context_lens of
-                # shape (B, next_n) (csrc/apis/attention.hpp:332-335);
-                # see indexer.py:607-608. For each batch/next_n token, the
-                # effective context length is context_lens[b] - next_n + j + 1.
-                next_n_arange = torch.arange(next_n, device="cuda", dtype=torch.int32)
-                context_lens_2d = (
-                    context_lens.unsqueeze(-1) - next_n + 1 + next_n_arange
-                ).contiguous()
                 schedule_metadata = get_paged_mqa_logits_metadata(
-                    context_lens_2d, blocksize, get_num_sms()
+                    context_lens, blocksize, get_num_sms()
                 )
-                logits = fp8_fp4_paged_mqa_logits(
-                    (q_fp8, None),
+                logits = fp8_paged_mqa_logits(
+                    q_fp8,
                     kv_cache_fp8,
                     weights,
-                    context_lens_2d,
+                    context_lens,
                     block_tables,
                     schedule_metadata,
                     max_model_len,
                     clean_logits=clean_logits,
                 )
 
-                ref_logits = _ref_fp8_fp4_paged_mqa_logits(
+                ref_logits = _ref_fp8_paged_mqa_logits(
                     q,
                     kv_cache,
                     weights,

tests/kernels/core/test_activation.py

@@ -16,7 +16,6 @@ from vllm.model_executor.layers.activation import (
     NewGELU,
     QuickGELU,
     SiluAndMul,
-    SiluAndMulWithClamp,
     SwigluOAIAndMul,
     SwigluStepAndMul,
     swiglustep_and_mul_triton,
@@ -117,85 +116,6 @@ def test_act_and_mul(
         opcheck(fn, (out, x))
 
 
-SWIGLU_LIMITS = [3.0, 7.0, 15.0]
-
-
-@pytest.mark.parametrize("swiglu_limit", SWIGLU_LIMITS)
-@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
-@pytest.mark.parametrize("d", D)
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize("seed", SEEDS)
-@pytest.mark.parametrize("device", CUDA_DEVICES)
-@torch.inference_mode()
-def test_silu_and_mul_with_clamp(
-    default_vllm_config,
-    swiglu_limit: float,
-    num_tokens: int,
-    d: int,
-    dtype: torch.dtype,
-    seed: int,
-    device: str,
-) -> None:
-    """SiluAndMulWithClamp: cuda kernel must match native reference."""
-    set_random_seed(seed)
-    torch.set_default_device(device)
-    # Use large values to ensure clamping is exercised.
-    x = torch.randn(num_tokens, 2 * d, dtype=dtype) * swiglu_limit * 2
-
-    layer = SiluAndMulWithClamp(swiglu_limit, compile_native=False)
-    out = layer(x)
-    ref_out = layer.forward_native(x)
-
-    rtol = {
-        torch.float16: 2e-3,
-        torch.bfloat16: 2e-2,
-        torch.float: 1.3e-6,
-    }
-    torch.testing.assert_close(
-        out, ref_out, atol=get_default_atol(out), rtol=rtol[out.dtype]
-    )
-
-    # Verify clamping is actually being applied: the clamped output should
-    # differ from the unclamped SiluAndMul output when inputs are large.
-    unclamped_out = SiluAndMul.forward_native(x)
-    assert not torch.equal(ref_out.float(), unclamped_out.float()), (
-        "Input was not large enough to exercise the clamp; increase scale"
-    )
-
-    # Verify gate clamping semantics with a controlled scalar case.
-    # gate=large_val is clamped to limit first, then silu(limit) * 1.0.
-    x_gate = torch.tensor(
-        [[swiglu_limit * 20.0, 1.0]], dtype=torch.float32, device=device
-    )
-    out_gate = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_gate)
-    expected_gate = torch.nn.functional.silu(
-        torch.tensor(swiglu_limit, dtype=torch.float32)
-    ).item()
-    torch.testing.assert_close(
-        out_gate,
-        torch.tensor([[expected_gate]], dtype=torch.float32, device=device),
-        atol=1e-3,
-        rtol=1e-3,
-    )
-
-    # Verify up clamping semantics: up >> limit gets clamped to limit.
-    x_up = torch.tensor(
-        [[1.0, swiglu_limit * 20.0]], dtype=torch.float32, device=device
-    )
-    out_up = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_up)
-    silu_1 = torch.nn.functional.silu(torch.tensor(1.0)).item()
-    torch.testing.assert_close(
-        out_up,
-        torch.tensor([[silu_1 * swiglu_limit]], dtype=torch.float32, device=device),
-        atol=1e-3,
-        rtol=1e-3,
-    )
-
-    # opcheck
-    out_buf = torch.empty(x.shape[:-1] + (d,), dtype=dtype, device=device)
-    opcheck(torch.ops._C.silu_and_mul_with_clamp, (out_buf, x, swiglu_limit))
-
-
 @pytest.mark.parametrize(
     "activation",
     [

tests/kernels/core/test_fused_q_kv_rmsnorm.py

@@ -1,81 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Correctness + large-token-count launch tests for fused_q_kv_rmsnorm.
-
-Before the grid-dim fix the kernel used grid ``(2, num_tokens)``, which hit
-CUDA's 65535 grid-y cap for ``num_tokens >= 65536`` and failed with
-``Triton Error [CUDA]: invalid argument`` at every large chunked-prefill
-profile run. These tests pin the new grid layout.
-"""
-
-from __future__ import annotations
-
-import pytest
-import torch
-
-from vllm.platforms import current_platform
-from vllm.v1.attention.ops.deepseek_v4_ops import fused_q_kv_rmsnorm
-
-pytestmark = pytest.mark.skipif(
-    not current_platform.is_cuda_alike(),
-    reason="fused_q_kv_rmsnorm requires a CUDA/ROCm device",
-)
-
-
-def _ref_rmsnorm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
-    x_f32 = x.to(torch.float32)
-    variance = x_f32.pow(2).mean(dim=-1, keepdim=True)
-    y = x_f32 * torch.rsqrt(variance + eps) * w.to(torch.float32)
-    return y.to(x.dtype)
-
-
-@pytest.mark.parametrize("num_tokens", [1, 17, 1024, 8192])
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
-def test_fused_q_kv_rmsnorm_correctness(num_tokens: int, dtype: torch.dtype):
-    torch.manual_seed(0)
-    device = "cuda"
-    q_size, kv_size = 192, 576
-    qr = torch.randn(num_tokens, q_size, dtype=dtype, device=device)
-    kv = torch.randn(num_tokens, kv_size, dtype=dtype, device=device)
-    qw = torch.randn(q_size, dtype=dtype, device=device)
-    kvw = torch.randn(kv_size, dtype=dtype, device=device)
-    eps = 1e-6
-
-    qr_out, kv_out = fused_q_kv_rmsnorm(qr, kv, qw, kvw, eps)
-
-    qr_ref = _ref_rmsnorm(qr, qw, eps)
-    kv_ref = _ref_rmsnorm(kv, kvw, eps)
-
-    tol = dict(rtol=1e-2, atol=1e-2)
-    torch.testing.assert_close(qr_out, qr_ref, **tol)
-    torch.testing.assert_close(kv_out, kv_ref, **tol)
-
-
-@pytest.mark.parametrize("num_tokens", [65535, 65536, 131072])
-def test_fused_q_kv_rmsnorm_launches_past_grid_y_cap(num_tokens: int):
-    """Regression guard: grid used to be (2, num_tokens), hitting CUDA's
-    65535 grid-y cap at num_tokens >= 65536. The new grid (num_tokens, 2)
-    lifts that bound to 2**31-1."""
-    device = "cuda"
-    dtype = torch.bfloat16
-    q_size, kv_size = 192, 576
-    qr = torch.randn(num_tokens, q_size, dtype=dtype, device=device)
-    kv = torch.randn(num_tokens, kv_size, dtype=dtype, device=device)
-    qw = torch.randn(q_size, dtype=dtype, device=device)
-    kvw = torch.randn(kv_size, dtype=dtype, device=device)
-
-    qr_out, kv_out = fused_q_kv_rmsnorm(qr, kv, qw, kvw, 1e-6)
-    # spot-check a couple of rows against the torch reference
-    for row in (0, num_tokens // 2, num_tokens - 1):
-        torch.testing.assert_close(
-            qr_out[row],
-            _ref_rmsnorm(qr[row : row + 1], qw, 1e-6)[0],
-            rtol=1e-2,
-            atol=1e-2,
-        )
-        torch.testing.assert_close(
-            kv_out[row],
-            _ref_rmsnorm(kv[row : row + 1], kvw, 1e-6)[0],
-            rtol=1e-2,
-            atol=1e-2,
-        )

tests/kernels/moe/modular_kernel_tools/mk_objects.py

@@ -223,7 +223,7 @@ if has_deep_ep() and not current_platform.has_device_capability(100):
     )
 
 if has_mori():
-    from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import (
+    from vllm.model_executor.layers.fused_moe.mori_prepare_finalize import (
         MoriPrepareAndFinalize,
     )
 
@@ -367,9 +367,7 @@ else:
     CutlassExpertsFp8 = None
 
 if cutlass_fp4_supported():
-    from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
-        CutlassExpertsFp4,
-    )
+    from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp4
 
     register_experts(
         CutlassExpertsFp4,

tests/kernels/moe/test_batched_deepgemm.py

@@ -10,12 +10,10 @@ from vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe import (
     BatchedDeepGemmExperts,
 )
 from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
+    BatchedPrepareAndFinalize,
     BatchedTritonExperts,
 )
 from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
-from vllm.model_executor.layers.fused_moe.prepare_finalize.batched import (
-    BatchedPrepareAndFinalize,
-)
 from vllm.utils.deep_gemm import calc_diff, is_deep_gemm_supported
 
 from .test_deepgemm import make_block_quant_fp8_weights

tests/kernels/moe/test_cutlass_moe.py

@@ -21,7 +21,7 @@ from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEQuantConfig,
     fp8_w8a8_moe_quant_config,
 )
-from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
+from vllm.model_executor.layers.fused_moe.cutlass_moe import (
     CutlassExpertsFp8,
     run_cutlass_moe_fp8,
 )

tests/kernels/moe/test_deepgemm.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """
-Unit-test DeepGEMM FP8 and FP4 kernels (no DeepEP).
+Unit-test DeepGEMM FP8 kernels (no DeepEP).
 Compare DeepGEMM path against the Triton fallback inside vLLM's fused_experts.
 """
 
@@ -21,8 +21,6 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
     maybe_make_prepare_finalize,
 )
 from vllm.model_executor.layers.fused_moe.config import (
-    FusedMoEQuantConfig,
-    FusedMoEQuantDesc,
     fp8_w8a8_moe_quant_config,
 )
 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
@@ -206,195 +204,3 @@ def test_deepgemm_vs_triton(m, n, k, topk, num_experts, monkeypatch, workspace_i
             f"DeepGEMM path was not executed during the test. "
             f"Call counter: {call_counter['cnt']}"
         )
-
-
-# ---------------------------------------------------------------------------
-# FP4 weight tests (DeepGEMM m_grouped_fp8_fp4_gemm_nt_contiguous)
-# ---------------------------------------------------------------------------
-
-
-def make_mxfp4_weights(
-    e: int,
-    n: int,
-    k: int,
-):
-    """
-    Generate (w1, w2) expert weights in MXFP4 packed format with float32 scales,
-    plus BF16 reference weights for validation.
-
-      w1 shape: (E, 2N, K//2) uint8    — packed FP4
-      w2 shape: (E, K, N//2)  uint8    — packed FP4
-      w1_s shape: (E, 2N, K//32) float32  — per-row block-32 scales
-      w2_s shape: (E, K, N//32)  float32  — per-row block-32 scales
-      w1_bf16: (E, 2N, K)   — original BF16 for reference
-      w2_bf16: (E, K, N)    — original BF16 for reference
-    """
-    from deep_gemm.utils.math import per_token_cast_to_fp4
-
-    dtype = torch.bfloat16
-    gran_k = 32  # MXFP4 block size
-
-    # bf16 reference weights — scale by 1/sqrt(dim) for numerical stability
-    w1_bf16 = torch.randn(e, 2 * n, k, device="cuda", dtype=dtype) * (k**-0.5)
-    w2_bf16 = torch.randn(e, k, n, device="cuda", dtype=dtype) * (n**-0.5)
-
-    # Quantize per-expert to FP4
-    w1 = torch.empty(e, 2 * n, k // 2, device="cuda", dtype=torch.uint8)
-    w2 = torch.empty(e, k, n // 2, device="cuda", dtype=torch.uint8)
-    w1_s = torch.empty(
-        e, 2 * n, math.ceil(k / gran_k), device="cuda", dtype=torch.float32
-    )
-    w2_s = torch.empty(e, k, math.ceil(n / gran_k), device="cuda", dtype=torch.float32)
-
-    for i in range(e):
-        w1[i], w1_s[i] = per_token_cast_to_fp4(
-            w1_bf16[i].float(), use_ue8m0=True, gran_k=gran_k
-        )
-        w2[i], w2_s[i] = per_token_cast_to_fp4(
-            w2_bf16[i].float(), use_ue8m0=True, gran_k=gran_k
-        )
-
-    return w1, w2, w1_s, w2_s, w1_bf16, w2_bf16
-
-
-def _bf16_moe_reference(x, w1, w2, topk_weights, topk_ids):
-    """BF16 token-loop MoE reference for correctness testing."""
-    import torch.nn.functional as F
-
-    num_tokens, hidden_size = x.shape
-    intermediate = w1.shape[1] // 2
-    top_k = topk_ids.shape[1]
-
-    output = torch.zeros(num_tokens, hidden_size, dtype=torch.float32, device=x.device)
-    for t in range(num_tokens):
-        for kk in range(top_k):
-            e = topk_ids[t, kk].item()
-            w = topk_weights[t, kk].item()
-            fc1 = x[t : t + 1].float() @ w1[e].float().T
-            linear = fc1[:, :intermediate]
-            gate = fc1[:, intermediate:]
-            act = F.silu(gate) * linear
-            fc2 = act @ w2[e].float().T
-            output[t] += w * fc2[0]
-    return output.to(torch.bfloat16)
-
-
-def run_single_fp4_case(m, n, k, topk, num_experts):
-    """
-    Run one (M,N,K) configuration with FP4 weights on DeepGEMM and assert
-    DeepGEMM FP4 == BF16 reference within tolerance.
-    """
-    tokens_bf16 = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * (k**-0.5)
-
-    # FP4 expert weight tensors + BF16 originals for reference
-    w1, w2, w1_s, w2_s, w1_bf16, w2_bf16 = make_mxfp4_weights(num_experts, n, k)
-
-    router_logits = torch.randn(m, num_experts, device="cuda", dtype=torch.float32)
-    topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1)
-    topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1)
-
-    from vllm.model_executor.layers.quantization.utils.quant_utils import (
-        GroupShape,
-    )
-    from vllm.platforms import current_platform
-
-    _fp8_dtype = current_platform.fp8_dtype()
-    _block_shape = GroupShape(128, 128)
-    quant_config = FusedMoEQuantConfig(
-        _a1=FusedMoEQuantDesc(_fp8_dtype, _block_shape, None, None, None, None),
-        _a2=FusedMoEQuantDesc(_fp8_dtype, _block_shape, None, None, None, None),
-        _w1=FusedMoEQuantDesc("mxfp4", None, w1_s, None, None, None),
-        _w2=FusedMoEQuantDesc("mxfp4", None, w2_s, None, None, None),
-    )
-    moe_config = make_dummy_moe_config()
-
-    from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
-        DeepGemmFP4Experts,
-    )
-
-    deep_gemm_fp4_experts = mk.FusedMoEKernel(
-        prepare_finalize=maybe_make_prepare_finalize(
-            moe=moe_config,
-            quant_config=quant_config,
-            allow_new_interface=True,
-            use_monolithic=False,
-        ),
-        fused_experts=DeepGemmFP4Experts(
-            moe_config=moe_config,
-            quant_config=quant_config,
-        ),
-        inplace=False,
-    )
-
-    # DeepGEMM FP4 path
-    out_deepgemm_fp4 = deep_gemm_fp4_experts.apply(
-        hidden_states=tokens_bf16,
-        w1=w1,
-        w2=w2,
-        topk_weights=topk_weights,
-        topk_ids=topk_ids,
-        global_num_experts=num_experts,
-        activation=MoEActivation.SILU,
-        apply_router_weight_on_input=False,
-        expert_map=None,
-    )
-
-    # BF16 reference using the same original weights
-    out_ref = _bf16_moe_reference(tokens_bf16, w1_bf16, w2_bf16, topk_weights, topk_ids)
-
-    # FP4 vs BF16 reference: quantization error from FP4 weights + FP8 activations
-    diff = calc_diff(out_deepgemm_fp4, out_ref)
-    assert diff < 0.05, f"FP4 diff exceeded 5%: {diff}"
-
-
-# DeepSeek V4 dims: H=4096, I=2048, so N=2*I=4096, K=H=4096.
-# FP4 quantization with block_k=32 needs large K for good accuracy.
-FP4_MNKs = [
-    (128, 4096, 4096),  # DeepSeek V4 shape
-    (256, 2048, 2048),  # Half-size variant
-]
-
-FP4_TOPKS = [2]
-FP4_NUM_EXPERTS = [8]
-
-
-@pytest.mark.parametrize(("m", "n", "k"), FP4_MNKs)
-@pytest.mark.parametrize("topk", FP4_TOPKS)
-@pytest.mark.parametrize("num_experts", FP4_NUM_EXPERTS)
-@pytest.mark.skipif(not is_deep_gemm_supported(), reason="Requires deep_gemm kernels")
-def test_deepgemm_fp4_vs_triton(
-    m, n, k, topk, num_experts, monkeypatch, workspace_init
-):
-    pytest.importorskip("deep_gemm.utils.math")
-    with monkeypatch.context() as mp:
-        mp.setenv("VLLM_USE_DEEP_GEMM", "1")
-
-        _DeepGemmFP4Experts = importlib.import_module(
-            "vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe"
-        ).DeepGemmFP4Experts
-
-        call_counter = {"cnt": 0}
-
-        orig_fn = _DeepGemmFP4Experts.apply
-
-        def _spy_apply(*args, **kwargs):
-            call_counter["cnt"] += 1
-            return orig_fn(*args, **kwargs)
-
-        monkeypatch.setattr(_DeepGemmFP4Experts, "apply", _spy_apply)
-        if topk > num_experts:
-            pytest.skip(f"topk={topk} > num_experts={num_experts}")
-
-        run_single_fp4_case(
-            m=m,
-            n=n,
-            k=k,
-            topk=topk,
-            num_experts=num_experts,
-        )
-
-        # ensure that the DeepGEMM FP4 path was indeed taken.
-        assert call_counter["cnt"] == 1, (
-            f"DeepGEMM FP4 path was not executed during the test. "
-            f"Call counter: {call_counter['cnt']}"
-        )

tests/kernels/moe/test_moe_layer.py

@@ -1632,7 +1632,6 @@ def _parallel_worker(
                 if all2all_manager is not None:
                     all2all_manager.destroy()
             total = total + 1
-            torch.distributed.barrier()
 
     skipped = total - (passed + failed)
 

tests/kernels/moe/test_nvfp4_moe.py

@@ -19,7 +19,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
     maybe_make_prepare_finalize,
 )
 from vllm.model_executor.layers.fused_moe.config import nvfp4_moe_quant_config
-from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
+from vllm.model_executor.layers.fused_moe.cutlass_moe import (
     CutlassExpertsFp4,
 )
 from vllm.model_executor.layers.fused_moe.prepare_finalize import (