[Misc] fix qwen3.5 config (#34604)

.buildkite/hardware_tests/amd.yaml

@@ -10,7 +10,7 @@ steps:
       docker build
       --build-arg max_jobs=16
       --build-arg REMOTE_VLLM=1
-      --build-arg ARG_PYTORCH_ROCM_ARCH='gfx942;gfx950'
+      --build-arg ARG_PYTORCH_ROCM_ARCH='gfx90a;gfx942'
       --build-arg VLLM_BRANCH=$BUILDKITE_COMMIT
       --tag "rocm/vllm-ci:${BUILDKITE_COMMIT}"
       -f docker/Dockerfile.rocm

.buildkite/image_build/image_build.sh

@@ -8,7 +8,7 @@ clean_docker_tag() {
 }
 
 print_usage_and_exit() {
-    echo "Usage: $0 <registry> <repo> <commit> <branch> <image_tag> [<image_tag_latest>]"
+    echo "Usage: $0 <registry> <repo> <commit> <branch> <vllm_use_precompiled> <vllm_merge_base_commit> <cache_from> <cache_to>"
     exit 1
 }
 
@@ -142,16 +142,11 @@ resolve_parent_commit() {
 
 print_bake_config() {
     echo "--- :page_facing_up: Resolved bake configuration"
-    # Write to a temp directory to avoid polluting the repo root (which is the
-    # Docker build context). Files left in the repo root get COPY'd into the
-    # image and can cause duplicate artifact uploads from downstream steps.
-    local bake_tmp
-    bake_tmp="$(mktemp -d)"
-    BAKE_CONFIG_FILE="${bake_tmp}/bake-config-build-${BUILDKITE_BUILD_NUMBER:-local}.json"
+    BAKE_CONFIG_FILE="bake-config-build-${BUILDKITE_BUILD_NUMBER:-local}.json"
     docker buildx bake -f "${VLLM_BAKE_FILE_PATH}" -f "${CI_HCL_PATH}" --print "${TARGET}" | tee "${BAKE_CONFIG_FILE}" || true
     echo "Saved bake config to ${BAKE_CONFIG_FILE}"
     echo "--- :arrow_down: Uploading bake config to Buildkite"
-    (cd "$(dirname "${BAKE_CONFIG_FILE}")" && buildkite-agent artifact upload "$(basename "${BAKE_CONFIG_FILE}")")
+    buildkite-agent artifact upload "${BAKE_CONFIG_FILE}"
 }
 
 #################################
@@ -159,7 +154,7 @@ print_bake_config() {
 #################################
 print_instance_info
 
-if [[ $# -lt 5 ]]; then
+if [[ $# -lt 7 ]]; then
     print_usage_and_exit
 fi
 
@@ -168,8 +163,10 @@ REGISTRY=$1
 REPO=$2
 BUILDKITE_COMMIT=$3
 BRANCH=$4
-IMAGE_TAG=$5
-IMAGE_TAG_LATEST=${6:-} # only used for main branch, optional
+VLLM_USE_PRECOMPILED=$5
+VLLM_MERGE_BASE_COMMIT=$6
+IMAGE_TAG=$7
+IMAGE_TAG_LATEST=${8:-} # only used for main branch, optional
 
 # build config
 TARGET="test-ci"
@@ -196,6 +193,8 @@ export CACHE_FROM
 export CACHE_FROM_BASE_BRANCH
 export CACHE_FROM_MAIN
 export CACHE_TO
+export VLLM_USE_PRECOMPILED
+export VLLM_MERGE_BASE_COMMIT
 
 # print args
 echo "--- :mag: Arguments"
@@ -203,6 +202,8 @@ echo "REGISTRY: ${REGISTRY}"
 echo "REPO: ${REPO}"
 echo "BUILDKITE_COMMIT: ${BUILDKITE_COMMIT}"
 echo "BRANCH: ${BRANCH}"
+echo "VLLM_USE_PRECOMPILED: ${VLLM_USE_PRECOMPILED}"
+echo "VLLM_MERGE_BASE_COMMIT: ${VLLM_MERGE_BASE_COMMIT}"
 echo "IMAGE_TAG: ${IMAGE_TAG}"
 echo "IMAGE_TAG_LATEST: ${IMAGE_TAG_LATEST}"
 

.buildkite/image_build/image_build.yaml

@@ -5,7 +5,8 @@ steps:
     depends_on: []
     timeout_in_minutes: 600
     commands:
-    - if [[ "$BUILDKITE_BRANCH" == "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $IMAGE_TAG $IMAGE_TAG_LATEST; else .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $IMAGE_TAG; fi
+    - if [[ "$BUILDKITE_BRANCH" != "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $VLLM_USE_PRECOMPILED $VLLM_MERGE_BASE_COMMIT $IMAGE_TAG; fi
+    - if [[ "$BUILDKITE_BRANCH" == "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $VLLM_USE_PRECOMPILED $VLLM_MERGE_BASE_COMMIT $IMAGE_TAG $IMAGE_TAG_LATEST; fi
     retry:
       automatic:
         - exit_status: -1  # Agent was lost

.buildkite/image_build/image_build_cpu.sh

@@ -11,10 +11,10 @@ REPO=$2
 BUILDKITE_COMMIT=$3
 
 # authenticate with AWS ECR
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
+aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin $REGISTRY
 
 # skip build if image already exists
-if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu) ]]; then
+if [[ -z $(docker manifest inspect $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu) ]]; then
   echo "Image not found, proceeding with build..."
 else
   echo "Image found"
@@ -24,13 +24,13 @@ fi
 # build
 docker build --file docker/Dockerfile.cpu \
   --build-arg max_jobs=16 \
-  --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
+  --build-arg buildkite_commit=$BUILDKITE_COMMIT \
   --build-arg VLLM_CPU_AVX512BF16=true \
   --build-arg VLLM_CPU_AVX512VNNI=true \
   --build-arg VLLM_CPU_AMXBF16=true \
-  --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu \
+  --tag $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu \
   --target vllm-test \
   --progress plain .
 
 # push
-docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu
+docker push $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu

.buildkite/image_build/image_build_cpu_arm64.sh

@@ -11,10 +11,10 @@ REPO=$2
 BUILDKITE_COMMIT=$3
 
 # authenticate with AWS ECR
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
+aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin $REGISTRY
 
 # skip build if image already exists
-if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu) ]]; then
+if [[ -z $(docker manifest inspect $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu) ]]; then
   echo "Image not found, proceeding with build..."
 else
   echo "Image found"
@@ -24,10 +24,10 @@ fi
 # build
 docker build --file docker/Dockerfile.cpu \
   --build-arg max_jobs=16 \
-  --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
-  --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu \
+  --build-arg buildkite_commit=$BUILDKITE_COMMIT \
+  --tag $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu \
   --target vllm-test \
   --progress plain .
 
 # push
-docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu
+docker push $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu

.buildkite/image_build/image_build_hpu.sh

@@ -11,10 +11,10 @@ REPO=$2
 BUILDKITE_COMMIT=$3
 
 # authenticate with AWS ECR
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
+aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin $REGISTRY
 
 # skip build if image already exists
-if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-hpu) ]]; then
+if [[ -z $(docker manifest inspect $REGISTRY/$REPO:$BUILDKITE_COMMIT-hpu) ]]; then
   echo "Image not found, proceeding with build..."
 else
   echo "Image found"
@@ -25,10 +25,10 @@ fi
 docker build \
   --file tests/pytorch_ci_hud_benchmark/Dockerfile.hpu \
   --build-arg max_jobs=16 \
-  --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
-  --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-hpu \
+  --build-arg buildkite_commit=$BUILDKITE_COMMIT \
+  --tag $REGISTRY/$REPO:$BUILDKITE_COMMIT-hpu \
   --progress plain \
   https://github.com/vllm-project/vllm-gaudi.git
 
 # push
-docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-hpu
+docker push $REGISTRY/$REPO:$BUILDKITE_COMMIT-hpu

.buildkite/lm-eval-harness/run-lm-eval-chartqa-vllm-vlm-baseline.sh

@@ -2,7 +2,7 @@
 # We can use this script to compute baseline accuracy on chartqa for vllm.
 #
 # Make sure you have lm-eval-harness installed:
-#   pip install "lm-eval[api]>=0.4.11"
+#   pip install "lm-eval[api]>=0.4.9.2"
 
 usage() {
     echo``
@@ -41,4 +41,4 @@ lm_eval --model vllm-vlm \
   --tasks chartqa \
   --batch_size auto \
   --apply_chat_template \
-  --limit "$LIMIT"
+  --limit $LIMIT

.buildkite/lm-eval-harness/run-lm-eval-gsm-hf-baseline.sh

@@ -2,7 +2,7 @@
 # We can use this script to compute baseline accuracy on GSM for transformers.
 #
 # Make sure you have lm-eval-harness installed:
-#   pip install "lm-eval[api]>=0.4.11"
+#   pip install "lm-eval[api]>=0.4.9.2"
 
 usage() {
     echo``

.buildkite/lm-eval-harness/run-lm-eval-gsm-vllm-baseline.sh

@@ -3,7 +3,7 @@
 # We use this for fp8, which HF does not support.
 #
 # Make sure you have lm-eval-harness installed:
-#   pip install "lm-eval[api]>=0.4.11"
+#   pip install "lm-eval[api]>=0.4.9.2"
 
 usage() {
     echo``

.buildkite/lm-eval-harness/run-lm-eval-mmlupro-vllm-baseline.sh

@@ -3,7 +3,7 @@
 # We use this for fp8, which HF does not support.
 #
 # Make sure you have lm-eval-harness installed:
-#   pip install "lm-eval[api]>=0.4.11"
+#   pip install "lm-eval[api]>=0.4.9.2"
 
 usage() {
     echo``
@@ -20,11 +20,14 @@ usage() {
     echo
 }
 
-while getopts "m:l:f:t:" OPT; do
+while getopts "m:b:l:f:t:" OPT; do
   case ${OPT} in
     m )
         MODEL="$OPTARG"
         ;;
+    b )
+        BATCH_SIZE="$OPTARG"
+        ;;
     l )
         LIMIT="$OPTARG"
         ;;

.buildkite/performance-benchmarks/scripts/run-performance-benchmarks.sh

@@ -15,11 +15,11 @@ DTYPE_FILTER="${DTYPE_FILTER:-}"
 check_gpus() {
   if command -v nvidia-smi; then
     # check the number of GPUs and GPU type.
-    declare -g gpu_count=$(nvidia-smi --list-gpus | grep -c . || true)
+    declare -g gpu_count=$(nvidia-smi --list-gpus | wc -l)
   elif command -v amd-smi; then
-    declare -g gpu_count=$(amd-smi list | grep -c 'GPU' || true)
+    declare -g gpu_count=$(amd-smi list | grep 'GPU' | wc -l)
   elif command -v hl-smi; then
-    declare -g gpu_count=$(hl-smi --list | grep -ci "Module ID" || true)
+    declare -g gpu_count=$(hl-smi --list | grep -i "Module ID" | wc -l)
   fi
 
   if [[ $gpu_count -gt 0 ]]; then
@@ -47,7 +47,7 @@ check_cpus() {
   declare -g numa_count=$(lscpu | grep "NUMA node(s):" | awk '{print $3}')
   if [[ $numa_count -gt 0 ]]; then
     echo "NUMA found."
-    echo "$numa_count"
+    echo $numa_count
   else
     echo "Need at least 1 NUMA to run benchmarking."
     exit 1
@@ -434,7 +434,7 @@ run_serving_tests() {
 
       # iterate over different max_concurrency
       for max_concurrency in $max_concurrency_list; do
-        new_test_name="${test_name}_qps_${qps}_concurrency_${max_concurrency}"
+        new_test_name=$test_name"_qps_"$qps"_concurrency_"$max_concurrency
         echo " new test name $new_test_name"
         # pass the tensor parallel size, the compilation mode, and the optimization
         # level to the client so that they can be used on the benchmark dashboard
@@ -471,7 +471,7 @@ run_serving_tests() {
 
     # clean up
     if [[ "${DRY_RUN:-0}" != "1" ]]; then
-      kill -9 "$server_pid"
+      kill -9 $server_pid
       kill_gpu_processes
     fi
   done

.buildkite/scripts/annotate-rocm-release.sh

@@ -25,7 +25,7 @@ S3_REGION="${AWS_DEFAULT_REGION:-us-west-2}"
 S3_URL="http://${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com"
 
 # Format ROCm version for path (e.g., "7.1" -> "rocm710")
-ROCM_VERSION_PATH="rocm$(echo "${ROCM_VERSION}" | tr -d '.')"
+ROCM_VERSION_PATH="rocm$(echo ${ROCM_VERSION} | tr -d '.')"
 ROCM_PATH="rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}"
 buildkite-agent annotate --style 'success' --context 'rocm-release-workflow' << EOF
 ## ROCm Wheel and Docker Image Releases

.buildkite/scripts/cache-rocm-base-wheels.sh

@@ -83,7 +83,7 @@ case "${1:-}" in
             exit 1
         fi
 
-        WHEEL_COUNT=$(find artifacts/rocm-base-wheels -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l)
+        WHEEL_COUNT=$(ls artifacts/rocm-base-wheels/*.whl 2>/dev/null | wc -l)
         if [[ "$WHEEL_COUNT" -eq 0 ]]; then
             echo "ERROR: No wheels found in artifacts/rocm-base-wheels/" >&2
             exit 1
@@ -110,9 +110,9 @@ case "${1:-}" in
 
         echo ""
         echo "Downloaded wheels:"
-        find artifacts/rocm-base-wheels -maxdepth 1 -name '*.whl' -exec ls -lh {} \;
+        ls -lh artifacts/rocm-base-wheels/
 
-        WHEEL_COUNT=$(find artifacts/rocm-base-wheels -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l)
+        WHEEL_COUNT=$(ls artifacts/rocm-base-wheels/*.whl 2>/dev/null | wc -l)
         echo ""
         echo "Total: $WHEEL_COUNT wheels"
         echo "========================================"

.buildkite/scripts/cherry-pick-from-milestone.sh

@@ -134,7 +134,7 @@ log_info "Fetching merged PRs from milestone '${MILESTONE}'..."
 
 # Store PR data in a temp file
 PR_DATA=$(mktemp)
-trap 'rm -f "$PR_DATA"' EXIT
+trap "rm -f $PR_DATA" EXIT
 
 if ! gh pr list --state merged --search "milestone:${MILESTONE}" \
     --limit 1000 \

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

@@ -1,57 +1,25 @@
 #!/bin/bash
 
-# This script runs tests inside the corresponding ROCm docker container.
-# It handles both single-node and multi-node test configurations.
-#
-# Multi-node detection: Instead of matching on fragile group names, we detect
-# multi-node jobs structurally by looking for the bracket command syntax
-# "[node0_cmds] && [node1_cmds]" or via the NUM_NODES environment variable.
-#
-###############################################################################
-# QUOTING / COMMAND PASSING
-#
-# Passing commands as positional arguments ($*) is fragile when the command
-# string itself contains double quotes, e.g.:
-#
-#   bash run-amd-test.sh "export FLAGS="value" && pytest -m "not slow""
-#
-# The outer shell resolves the nested quotes *before* this script runs, so
-# the script receives mangled input it cannot fully recover.
-#
-# Preferred: pass commands via the VLLM_TEST_COMMANDS environment variable:
-#
-#   export VLLM_TEST_COMMANDS='export FLAGS="value" && pytest -m "not slow"'
-#   bash run-amd-test.sh
-#
-# Single-quoted assignment preserves all inner double quotes verbatim.
-# The $* path is kept for backward compatibility but callers should migrate.
-###############################################################################
+# This script runs test inside the corresponding ROCm docker container.
 set -o pipefail
 
 # Export Python path
 export PYTHONPATH=".."
 
-###############################################################################
-# Helper Functions
-###############################################################################
+# Print ROCm version
+echo "--- Confirming Clean Initial State"
+while true; do
+        sleep 3
+        if grep -q clean /opt/amdgpu/etc/gpu_state; then
+                echo "GPUs state is \"clean\""
+                break
+        fi
+done
 
-wait_for_clean_gpus() {
-  local timeout=${1:-300}
-  local start=$SECONDS
-  echo "--- Waiting for clean GPU state (timeout: ${timeout}s)"
-  while true; do
-    if grep -q clean /opt/amdgpu/etc/gpu_state; then
-      echo "GPUs state is \"clean\""
-      return
-    fi
-    if (( SECONDS - start >= timeout )); then
-      echo "Error: GPUs did not reach clean state within ${timeout}s" >&2
-      exit 1
-    fi
-    sleep 3
-  done
-}
+echo "--- ROCm info"
+rocminfo
 
+# cleanup older docker images
 cleanup_docker() {
   # Get Docker's root directory
   docker_root=$(docker info -f '{{.DockerRootDir}}')
@@ -60,12 +28,15 @@ cleanup_docker() {
     exit 1
   fi
   echo "Docker root directory: $docker_root"
-
+  # Check disk usage of the filesystem where Docker's root directory is located
   disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
+  # Define the threshold
   threshold=70
   if [ "$disk_usage" -gt "$threshold" ]; then
     echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
+    # Remove dangling images (those that are not tagged and not used by any container)
     docker image prune -f
+    # Remove unused volumes / force the system prune for old images as well.
     docker volume prune -f && docker system prune --force --filter "until=72h" --all
     echo "Docker images and volumes cleanup completed."
   else
@@ -74,410 +45,193 @@ cleanup_docker() {
 }
 
 cleanup_network() {
-  local max_nodes=${NUM_NODES:-2}
-  for node in $(seq 0 $((max_nodes - 1))); do
-    if docker ps -a -q -f name="node${node}" | grep -q .; then
-      docker stop "node${node}" || true
+  for node in $(seq 0 $((NUM_NODES-1))); do
+    if docker pr -a -q -f name="node${node}" | grep -q .; then
+      docker stop "node${node}"
     fi
   done
-  if docker network ls | grep -q docker-net; then
-    docker network rm docker-net || true
+  if docker network ls | grep docker-net; then
+    docker network rm docker-net
   fi
 }
 
-is_multi_node() {
-  local cmds="$1"
-  # Primary signal: NUM_NODES environment variable set by the pipeline
-  if [[ "${NUM_NODES:-1}" -gt 1 ]]; then
-    return 0
-  fi
-  # Fallback: detect the bracket syntax structurally
-  # Pattern: [...] && [...] (per-node command arrays)
-  if [[ "$cmds" =~ \[.*\].*\&\&.*\[.*\] ]]; then
-    return 0
-  fi
-  return 1
-}
-
-###############################################################################
-# Pytest marker/keyword re-quoting
-#
-# When commands are passed through Buildkite -> shell -> $* -> bash -c,
-# quotes around multi-word pytest -m/-k expressions get stripped:
-#   pytest -v -s -m 'not cpu_test' v1/core
-# becomes:
-#   pytest -v -s -m not cpu_test v1/core
-#
-# pytest then interprets "cpu_test" as a file path, not part of the marker.
-#
-# This function detects unquoted expressions after -m/-k and re-quotes them
-# by collecting tokens until a recognizable boundary is reached:
-#   - test path (contains '/')
-#   - test file (ends with '.py')
-#   - another pytest flag (--xxx or -x single-char flags)
-#   - command separator (&& || ; |)
-#   - environment variable assignment (FOO=bar)
-#
-# Single-word markers (e.g. -m cpu_test, -m hybrid_model) pass through
-# unquoted since they have no spaces and work fine.
-#
-# Already-quoted expressions (containing literal single quotes) are passed
-# through untouched to avoid double-quoting values injected by
-# apply_rocm_test_overrides.
-#
-# NOTE: This ONLY fixes -m/-k flags. It cannot recover arbitrary inner
-# double-quotes stripped by the calling shell (see header comment).
-# Use VLLM_TEST_COMMANDS to avoid the problem entirely.
-###############################################################################
-re_quote_pytest_markers() {
-  local input="$1"
-  local output=""
-  local collecting=false
-  local marker_buf=""
-
-  # Flatten newlines for consistent tokenization
-  local flat="${input//$'\n'/ }"
-
-  # Disable globbing to prevent *.py etc. from expanding during read -ra
-  local restore_glob
-  restore_glob="$(shopt -p -o noglob 2>/dev/null || true)"
-  set -o noglob
-  local -a words
-  read -ra words <<< "$flat"
-  eval "$restore_glob"
-
-  for word in "${words[@]}"; do
-    if $collecting; then
-      # If the token we're about to collect already contains a literal
-      # single quote, the expression was already quoted upstream.
-      # Flush and stop collecting.
-      if [[ "$word" == *"'"* ]]; then
-        if [[ -n "$marker_buf" ]]; then
-          # Should not normally happen (partial buf + quote), flush raw
-          output+="${marker_buf} "
-          marker_buf=""
-        fi
-        output+="${word} "
-        collecting=false
-        continue
-      fi
-
-      local is_boundary=false
-      case "$word" in
-        # Command separators
-        "&&"|"||"|";"|"|")
-          is_boundary=true ;;
-        # Long flags (--ignore, --shard-id, etc.)
-        --*)
-          is_boundary=true ;;
-        # Short flags (-v, -s, -x, etc.) but NOT negative marker tokens
-        # like "not" which don't start with "-". Also skip -k/-m which
-        # would start a new marker (handled below).
-        -[a-zA-Z])
-          is_boundary=true ;;
-        # Test path (contains /)
-        */*)
-          is_boundary=true ;;
-        # Test file (ends with .py, possibly with ::method)
-        *.py|*.py::*)
-          is_boundary=true ;;
-        # Environment variable assignment preceding a command (FOO=bar)
-        *=*)
-          # Only treat as boundary if it looks like VAR=value, not
-          # pytest filter expressions like num_gpus=2 inside markers
-          if [[ "$word" =~ ^[A-Z_][A-Z0-9_]*= ]]; then
-            is_boundary=true
-          fi
-          ;;
-      esac
-
-      if $is_boundary; then
-        # Flush the collected marker expression
-        if [[ "$marker_buf" == *" "* || "$marker_buf" == *"("* ]]; then
-          output+="'${marker_buf}' "
-        else
-          output+="${marker_buf} "
-        fi
-        collecting=false
-        marker_buf=""
-        # Check if this boundary word itself starts a new -m/-k
-        if [[ "$word" == "-m" || "$word" == "-k" ]]; then
-          output+="${word} "
-          collecting=true
-        else
-          output+="${word} "
-        fi
-      else
-        # Accumulate into marker buffer
-        if [[ -n "$marker_buf" ]]; then
-          marker_buf+=" ${word}"
-        else
-          marker_buf="${word}"
-        fi
-      fi
-    elif [[ "$word" == "-m" || "$word" == "-k" ]]; then
-      output+="${word} "
-      collecting=true
-      marker_buf=""
-    else
-      output+="${word} "
-    fi
-  done
-
-  # Flush any trailing marker expression (marker at end of command)
-  if $collecting && [[ -n "$marker_buf" ]]; then
-    if [[ "$marker_buf" == *" "* || "$marker_buf" == *"("* ]]; then
-      output+="'${marker_buf}'"
-    else
-      output+="${marker_buf}"
-    fi
-  fi
-
-  echo "${output% }"
-}
-
-###############################################################################
-# ROCm-specific pytest command rewrites
-#
-# These apply ignore flags and environment overrides for tests that are not
-# yet supported or behave differently on ROCm hardware. Kept as a single
-# function so new exclusions are easy to add in one place.
-###############################################################################
-
-apply_rocm_test_overrides() {
-  local cmds="$1"
-
-  # --- Model registry filter ---
-  if [[ $cmds == *"pytest -v -s models/test_registry.py"* ]]; then
-    cmds=${cmds//"pytest -v -s models/test_registry.py"/"pytest -v -s models/test_registry.py -k 'not BambaForCausalLM and not GritLM and not Mamba2ForCausalLM and not Zamba2ForCausalLM'"}
-  fi
-
-  # --- LoRA: disable custom paged attention ---
-  if [[ $cmds == *"pytest -v -s lora"* ]]; then
-    cmds=${cmds//"pytest -v -s lora"/"VLLM_ROCM_CUSTOM_PAGED_ATTN=0 pytest -v -s lora"}
-  fi
-
-  # --- Kernel ignores ---
-  if [[ $cmds == *" kernels/core"* ]]; then
-    cmds="${cmds} \
-    --ignore=kernels/core/test_fused_quant_layernorm.py \
-    --ignore=kernels/core/test_permute_cols.py"
-  fi
-
-  if [[ $cmds == *" kernels/attention"* ]]; then
-    cmds="${cmds} \
-    --ignore=kernels/attention/test_attention_selector.py \
-    --ignore=kernels/attention/test_encoder_decoder_attn.py \
-    --ignore=kernels/attention/test_flash_attn.py \
-    --ignore=kernels/attention/test_flashinfer.py \
-    --ignore=kernels/attention/test_prefix_prefill.py \
-    --ignore=kernels/attention/test_cascade_flash_attn.py \
-    --ignore=kernels/attention/test_mha_attn.py \
-    --ignore=kernels/attention/test_lightning_attn.py \
-    --ignore=kernels/attention/test_attention.py"
-  fi
-
-  if [[ $cmds == *" kernels/quantization"* ]]; then
-    cmds="${cmds} \
-    --ignore=kernels/quantization/test_int8_quant.py \
-    --ignore=kernels/quantization/test_machete_mm.py \
-    --ignore=kernels/quantization/test_block_fp8.py \
-    --ignore=kernels/quantization/test_block_int8.py \
-    --ignore=kernels/quantization/test_marlin_gemm.py \
-    --ignore=kernels/quantization/test_cutlass_scaled_mm.py \
-    --ignore=kernels/quantization/test_int8_kernel.py"
-  fi
-
-  if [[ $cmds == *" kernels/mamba"* ]]; then
-    cmds="${cmds} \
-    --ignore=kernels/mamba/test_mamba_mixer2.py \
-    --ignore=kernels/mamba/test_causal_conv1d.py \
-    --ignore=kernels/mamba/test_mamba_ssm_ssd.py"
-  fi
-
-  if [[ $cmds == *" kernels/moe"* ]]; then
-    cmds="${cmds} \
-    --ignore=kernels/moe/test_moe.py \
-    --ignore=kernels/moe/test_cutlass_moe.py \
-    --ignore=kernels/moe/test_triton_moe_ptpc_fp8.py"
-  fi
-
-  # --- Entrypoint ignores ---
-  if [[ $cmds == *" entrypoints/openai "* ]]; then
-    cmds=${cmds//" entrypoints/openai "/" entrypoints/openai \
-    --ignore=entrypoints/openai/test_audio.py \
-    --ignore=entrypoints/openai/test_shutdown.py \
-    --ignore=entrypoints/openai/test_completion.py \
-    --ignore=entrypoints/openai/test_models.py \
-    --ignore=entrypoints/openai/test_lora_adapters.py \
-    --ignore=entrypoints/openai/test_return_tokens_as_ids.py \
-    --ignore=entrypoints/openai/test_root_path.py \
-    --ignore=entrypoints/openai/test_tokenization.py \
-    --ignore=entrypoints/openai/test_prompt_validation.py "}
-  fi
-
-  if [[ $cmds == *" entrypoints/llm "* ]]; then
-    cmds=${cmds//" entrypoints/llm "/" entrypoints/llm \
-    --ignore=entrypoints/llm/test_chat.py \
-    --ignore=entrypoints/llm/test_accuracy.py \
-    --ignore=entrypoints/llm/test_init.py \
-    --ignore=entrypoints/llm/test_prompt_validation.py "}
-  fi
-
-  # Clean up escaped newlines from --ignore appends
-  cmds=$(echo "$cmds" | sed 's/ \\ / /g')
-
-  echo "$cmds"
-}
-
-###############################################################################
-# Main
-###############################################################################
-
-# --- GPU initialization ---
-echo "--- Confirming Clean Initial State"
-wait_for_clean_gpus
-
-echo "--- ROCm info"
-rocminfo
-
-# --- Docker housekeeping ---
+# Call the cleanup docker function
 cleanup_docker
 
 echo "--- Resetting GPUs"
-echo "reset" > /opt/amdgpu/etc/gpu_state
-wait_for_clean_gpus
 
-# --- Pull test image ---
+echo "reset" > /opt/amdgpu/etc/gpu_state
+
+while true; do
+        sleep 3
+        if grep -q clean /opt/amdgpu/etc/gpu_state; then
+                echo "GPUs state is \"clean\""
+                break
+        fi
+done
+
 echo "--- Pulling container"
 image_name="rocm/vllm-ci:${BUILDKITE_COMMIT}"
 container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
 docker pull "${image_name}"
 
 remove_docker_container() {
-  docker rm -f "${container_name}" || docker image rm -f "${image_name}" || true
+   docker rm -f "${container_name}" || docker image rm -f "${image_name}" || true
 }
 trap remove_docker_container EXIT
 
-# --- Prepare commands ---
 echo "--- Running container"
 
 HF_CACHE="$(realpath ~)/huggingface"
 mkdir -p "${HF_CACHE}"
 HF_MOUNT="/root/.cache/huggingface"
 
-# ---- Command source selection ----
-# Prefer VLLM_TEST_COMMANDS (preserves all inner quoting intact).
-# Fall back to $* for backward compatibility, but warn that inner
-# double-quotes will have been stripped by the calling shell.
-if [[ -n "${VLLM_TEST_COMMANDS:-}" ]]; then
-  commands="${VLLM_TEST_COMMANDS}"
-  echo "Commands sourced from VLLM_TEST_COMMANDS (quoting preserved)"
-else
-  commands="$*"
-  if [[ -z "$commands" ]]; then
-    echo "Error: No test commands provided." >&2
-    echo "Usage:" >&2
-    echo "  Preferred:  VLLM_TEST_COMMANDS='...' bash $0" >&2
-    echo "  Legacy:     bash $0 \"commands here\"" >&2
-    exit 1
-  fi
-  echo "Commands sourced from positional args (legacy mode)"
-  echo "WARNING: Inner double-quotes in the command string may have been"
-  echo "  stripped by the calling shell. If you see syntax errors, switch to:"
-  echo "  export VLLM_TEST_COMMANDS='your commands here'"
-  echo "  bash $0"
-fi
-
+commands=$@
 echo "Raw commands: $commands"
 
-# Fix quoting before ROCm overrides (so overrides see correct structure)
-commands=$(re_quote_pytest_markers "$commands")
-echo "After re-quoting: $commands"
+commands=${commands//"pytest -v -s basic_correctness/test_basic_correctness.py"/"pytest -v -s basic_correctness/test_basic_correctness.py"}
 
-commands=$(apply_rocm_test_overrides "$commands")
+if [[ $commands == *"pytest -v -s models/test_registry.py"* ]]; then
+  commands=${commands//"pytest -v -s models/test_registry.py"/"pytest -v -s models/test_registry.py -k 'not BambaForCausalLM and not GritLM and not Mamba2ForCausalLM and not Zamba2ForCausalLM'"}
+fi
+
+commands=${commands//"pytest -v -s compile/test_basic_correctness.py"/"pytest -v -s compile/test_basic_correctness.py"}
+
+if [[ $commands == *"pytest -v -s lora"* ]]; then
+  commands=${commands//"pytest -v -s lora"/"VLLM_ROCM_CUSTOM_PAGED_ATTN=0 pytest -v -s lora"}
+fi
+
+#ignore certain kernels tests
+if [[ $commands == *" kernels/core"* ]]; then
+  commands="${commands} \
+  --ignore=kernels/core/test_fused_quant_layernorm.py \
+  --ignore=kernels/core/test_permute_cols.py"
+fi
+
+if [[ $commands == *" kernels/attention"* ]]; then
+  commands="${commands} \
+  --ignore=kernels/attention/test_attention_selector.py \
+  --ignore=kernels/attention/test_encoder_decoder_attn.py \
+  --ignore=kernels/attention/test_flash_attn.py \
+  --ignore=kernels/attention/test_flashinfer.py \
+  --ignore=kernels/attention/test_prefix_prefill.py \
+  --ignore=kernels/attention/test_cascade_flash_attn.py \
+  --ignore=kernels/attention/test_mha_attn.py \
+  --ignore=kernels/attention/test_lightning_attn.py \
+  --ignore=kernels/attention/test_attention.py"
+fi
+
+if [[ $commands == *" kernels/quantization"* ]]; then
+  commands="${commands} \
+  --ignore=kernels/quantization/test_int8_quant.py \
+  --ignore=kernels/quantization/test_machete_mm.py \
+  --ignore=kernels/quantization/test_block_fp8.py \
+  --ignore=kernels/quantization/test_block_int8.py \
+  --ignore=kernels/quantization/test_marlin_gemm.py \
+  --ignore=kernels/quantization/test_cutlass_scaled_mm.py \
+  --ignore=kernels/quantization/test_int8_kernel.py"
+fi
+
+if [[ $commands == *" kernels/mamba"* ]]; then
+  commands="${commands} \
+  --ignore=kernels/mamba/test_mamba_mixer2.py \
+  --ignore=kernels/mamba/test_causal_conv1d.py \
+  --ignore=kernels/mamba/test_mamba_ssm_ssd.py"
+fi
+
+if [[ $commands == *" kernels/moe"* ]]; then
+  commands="${commands} \
+  --ignore=kernels/moe/test_moe.py \
+  --ignore=kernels/moe/test_cutlass_moe.py \
+  --ignore=kernels/moe/test_triton_moe_ptpc_fp8.py"
+fi
+
+#ignore certain Entrypoints/openai tests
+if [[ $commands == *" entrypoints/openai "* ]]; then
+  commands=${commands//" entrypoints/openai "/" entrypoints/openai \
+  --ignore=entrypoints/openai/test_audio.py \
+  --ignore=entrypoints/openai/test_shutdown.py \
+  --ignore=entrypoints/openai/test_completion.py \
+  --ignore=entrypoints/openai/test_models.py \
+  --ignore=entrypoints/openai/test_lora_adapters.py \
+  --ignore=entrypoints/openai/test_return_tokens_as_ids.py \
+  --ignore=entrypoints/openai/test_root_path.py \
+  --ignore=entrypoints/openai/test_tokenization.py \
+  --ignore=entrypoints/openai/test_prompt_validation.py "}
+fi
+
+#ignore certain Entrypoints/llm tests
+if [[ $commands == *" entrypoints/llm "* ]]; then
+  commands=${commands//" entrypoints/llm "/" entrypoints/llm \
+  --ignore=entrypoints/llm/test_chat.py \
+  --ignore=entrypoints/llm/test_accuracy.py \
+  --ignore=entrypoints/llm/test_init.py \
+  --ignore=entrypoints/llm/test_prompt_validation.py "}
+fi
+
+commands=$(echo "$commands" | sed 's/ \\ / /g')
 echo "Final commands: $commands"
 
+# --ignore=entrypoints/openai/test_encoder_decoder.py \
+# --ignore=entrypoints/openai/test_embedding.py \
+# --ignore=entrypoints/openai/test_oot_registration.py
+# --ignore=entrypoints/openai/test_accuracy.py \
+# --ignore=entrypoints/openai/test_models.py <= Fails on MI250 but passes on MI300 as of 2025-03-13
+
+
 MYPYTHONPATH=".."
 
-# Verify GPU access
+# Test that we're launching on the machine that has
+# proper access to GPUs
 render_gid=$(getent group render | cut -d: -f3)
 if [[ -z "$render_gid" ]]; then
   echo "Error: 'render' group not found. This is required for GPU access." >&2
   exit 1
 fi
 
-# --- RDMA device passthrough (conditional) ---
-# If the host has RDMA devices, pass them through so tests like
-# test_moriio_connector can access ibverbs. On hosts without RDMA
-# hardware the tests will gracefully skip via _rdma_available().
-RDMA_FLAGS=""
-if [ -d /dev/infiniband ]; then
-  echo "RDMA devices detected on host, enabling passthrough"
-  RDMA_FLAGS="--device /dev/infiniband --cap-add=IPC_LOCK"
-else
-  echo "No RDMA devices found on host, RDMA tests will be skipped"
-fi
+if [[ $commands == *"VLLM_TEST_GROUP_NAME=mi325_4-2-node-tests-4-gpus-in-total"* ]]; then
 
-# --- Route: multi-node vs single-node ---
-if is_multi_node "$commands"; then
-  echo "--- Multi-node job detected"
   export DCKR_VER=$(docker --version | sed 's/Docker version \(.*\), build .*/\1/')
 
-  # Parse the bracket syntax:  prefix ; [node0_cmds] && [node1_cmds]
-  #   BASH_REMATCH[1] = prefix (everything before first bracket)
-  #   BASH_REMATCH[2] = comma-separated node0 commands
-  #   BASH_REMATCH[3] = comma-separated node1 commands
-  if [[ "$commands" =~ ^(.*)\[(.*)"] && ["(.*)\]$ ]]; then
-    prefix=$(echo "${BASH_REMATCH[1]}" | sed 's/;//g')
-    echo "PREFIX: ${prefix}"
-
-    export composite_command="(command rocm-smi || true)"
-    saved_IFS=$IFS
-    IFS=','
-    read -ra node0 <<< "${BASH_REMATCH[2]}"
-    read -ra node1 <<< "${BASH_REMATCH[3]}"
-    IFS=$saved_IFS
-
-    if [[ ${#node0[@]} -ne ${#node1[@]} ]]; then
-      echo "Warning: node0 has ${#node0[@]} commands, node1 has ${#node1[@]}. They will be paired by index."
-    fi
-
-    for i in "${!node0[@]}"; do
-      command_node_0=$(echo "${node0[i]}" | sed 's/\"//g')
-      command_node_1=$(echo "${node1[i]}" | sed 's/\"//g')
-
-      step_cmd="./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 ${image_name} '${command_node_0}' '${command_node_1}'"
-      echo "COMMANDS: ${step_cmd}"
-      composite_command="${composite_command} && ${step_cmd}"
-    done
-
-    /bin/bash -c "${composite_command}"
-    cleanup_network
+  if [[ "$commands" =~ ^(.*)"["(.*)"] && ["(.*)"]"$ ]]; then
+      prefix=$( echo "${BASH_REMATCH[1]}" | sed 's/;//g')
+      echo "PREFIX: ${prefix}"
+      export composite_command="(command rocm-smi || true)"
+      myIFS=$IFS
+      IFS=','
+      read -ra node0 <<< ${BASH_REMATCH[2]}
+      read -ra node1 <<< ${BASH_REMATCH[3]}
+      IFS=$myIFS
+      for i in "${!node0[@]}";do 
+        command_node_0=$(echo ${node0[i]} | sed 's/\"//g')
+        command_node_1=$(echo ${node1[i]} | sed 's/\"//g')
+        
+        export commands="./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 ${image_name} '${command_node_0}' '${command_node_1}'"
+        echo "COMMANDS: ${commands}"
+        composite_command=$(echo "${composite_command} && ${commands}")
+      done
+      /bin/bash -c "${composite_command}"
+      cleanup_network
   else
-    echo "Multi-node job detected but failed to parse bracket command syntax."
-    echo "Expected format: prefix ; [node0_cmd1, node0_cmd2] && [node1_cmd1, node1_cmd2]"
-    echo "Got: $commands"
-    cleanup_network
-    exit 111
+      echo "Failed to parse node commands! Exiting."
+      cleanup_network
+      exit 111
   fi
 else
-  echo "--- Single-node job"
   echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
   docker run \
-    --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
-    $RDMA_FLAGS \
-    --network=host \
-    --shm-size=16gb \
-    --group-add "$render_gid" \
-    --rm \
-    -e HF_TOKEN \
-    -e AWS_ACCESS_KEY_ID \
-    -e AWS_SECRET_ACCESS_KEY \
-    -v "${HF_CACHE}:${HF_MOUNT}" \
-    -e "HF_HOME=${HF_MOUNT}" \
-    -e "PYTHONPATH=${MYPYTHONPATH}" \
-    --name "${container_name}" \
-    "${image_name}" \
-    /bin/bash -c "${commands}"
+          --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
+          --network=host \
+          --shm-size=16gb \
+          --group-add "$render_gid" \
+          --rm \
+          -e HF_TOKEN \
+          -e AWS_ACCESS_KEY_ID \
+          -e AWS_SECRET_ACCESS_KEY \
+          -v "${HF_CACHE}:${HF_MOUNT}" \
+          -e "HF_HOME=${HF_MOUNT}" \
+          -e "PYTHONPATH=${MYPYTHONPATH}" \
+          --name "${container_name}" \
+          "${image_name}" \
+          /bin/bash -c "${commands}"
 fi

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

@@ -27,7 +27,7 @@ function cpu_tests() {
   podman exec -it "$container_id" bash -c "
     export TORCH_COMPILE_DISABLE=1
     set -xve
-    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m" >> "$HOME"/test_basic.log
+    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m" >> $HOME/test_basic.log
 
   # Run basic model test
   podman exec -it "$container_id" bash -c "
@@ -43,7 +43,7 @@ function cpu_tests() {
     pytest -v -s tests/models/language/generation/test_common.py::test_models[False-False-5-32-google/gemma-1.1-2b-it]
     pytest -v -s tests/models/language/pooling/test_classification.py::test_models[float-jason9693/Qwen2.5-1.5B-apeach]
     # TODO: Below test case tests/models/language/pooling/test_embedding.py::test_models[True-ssmits/Qwen2-7B-Instruct-embed-base] fails on ppc64le. Disabling it for time being.
-    # pytest -v -s tests/models/language/pooling/test_embedding.py -m cpu_model" >> "$HOME"/test_rest.log
+    # pytest -v -s tests/models/language/pooling/test_embedding.py -m cpu_model" >> $HOME/test_rest.log
 }
 
 # All of CPU tests are expected to be finished less than 40 mins.

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

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

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

@@ -7,7 +7,7 @@ set -exuo pipefail
 # Try building the docker image
 image_name="hpu/upstream-vllm-ci:${BUILDKITE_COMMIT}"
 container_name="hpu-upstream-vllm-ci-${BUILDKITE_COMMIT}-container"
-cat <<EOF | docker build -t "${image_name}" -f - .
+cat <<EOF | docker build -t ${image_name} -f - .
 FROM gaudi-base-image:latest
 
 COPY ./ /workspace/vllm
@@ -39,12 +39,12 @@ EOF
 # functions, while other platforms only need one remove_docker_container
 # function.
 EXITCODE=1
-remove_docker_containers() { docker rm -f "${container_name}" || true; }
+remove_docker_containers() { docker rm -f ${container_name} || true; }
 trap 'remove_docker_containers; exit $EXITCODE;' EXIT
 remove_docker_containers
 
 echo "Running HPU plugin v1 test"
-docker run --rm --runtime=habana --name="${container_name}" --network=host \
+docker run --rm --runtime=habana --name=${container_name} --network=host \
   -e HABANA_VISIBLE_DEVICES=all \
   -e VLLM_SKIP_WARMUP=true \
   -e PT_HPU_ENABLE_LAZY_COLLECTIVES=true \

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

@@ -41,7 +41,6 @@ get_config() {
         echo "Error: file '${TEST_RUN_CONFIG_FILE}' does not exist in the warehouse" >&2
         exit 1
     fi
-    # shellcheck source=/dev/null
     source "${TEST_RUN_CONFIG_FILE}"
     echo "Base docker image name that get from configuration: ${BASE_IMAGE_NAME}"
     return 0
@@ -49,8 +48,9 @@ get_config() {
 
 # get test running configuration.
 fetch_vllm_test_cfg
+get_config
 # Check if the function call was successful. If not, exit the script.
-if ! get_config; then
+if [ $? -ne 0 ]; then
   exit 1
 fi
 
@@ -62,14 +62,14 @@ agent_idx=$(echo "${BUILDKITE_AGENT_NAME}" | awk -F'-' '{print $(NF-1)}')
 echo "agent_idx: ${agent_idx}"
 builder_name="cachebuilder${agent_idx}"
 builder_cache_dir="/mnt/docker-cache${agent_idx}"
-mkdir -p "${builder_cache_dir}"
+mkdir -p ${builder_cache_dir}
 
 # Try building the docker image
 cat <<EOF | DOCKER_BUILDKIT=1 docker build \
-    --add-host cache-service-vllm.nginx-pypi-cache.svc.cluster.local:"${PYPI_CACHE_HOST}" \
-    --builder "${builder_name}" --cache-from type=local,src="${builder_cache_dir}" \
-                           --cache-to type=local,dest="${builder_cache_dir}",mode=max \
-    --progress=plain --load -t "${image_name}" -f - .
+    --add-host cache-service-vllm.nginx-pypi-cache.svc.cluster.local:${PYPI_CACHE_HOST} \
+    --builder ${builder_name} --cache-from type=local,src=${builder_cache_dir} \
+                           --cache-to type=local,dest=${builder_cache_dir},mode=max \
+    --progress=plain --load -t ${image_name} -f - .
 FROM ${BASE_IMAGE_NAME}
 
 # Define environments
@@ -116,7 +116,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
     source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
     source /usr/local/Ascend/nnal/atb/set_env.sh && \
-    export LD_LIBRARY_PATH=\$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/$(uname -i)-linux/devlib && \
+    export LD_LIBRARY_PATH=\$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
     python3 -m pip install -v -e /workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
@@ -139,7 +139,7 @@ trap remove_docker_container EXIT
 # Generate corresponding --device args based on BUILDKITE_AGENT_NAME
 # Ascend NPU BUILDKITE_AGENT_NAME format is {hostname}-{agent_idx}-{npu_card_num}cards, and agent_idx starts from 1.
 #   e.g. atlas-a2-001-1-2cards means this is the 1-th agent on atlas-a2-001 host, and it has 2 NPU cards.
-#   returns one argument per line: --device, /dev/davinciX, ...
+#   returns --device /dev/davinci0 --device /dev/davinci1
 parse_and_gen_devices() {
     local input="$1"
     local index cards_num
@@ -151,24 +151,29 @@ parse_and_gen_devices() {
         return 1
     fi
 
+    local devices=""
     local i=0
     while (( i < cards_num )); do
         local dev_idx=$(((index - 1)*cards_num + i ))
-        printf '%s\n' "--device"
-        printf '%s\n' "/dev/davinci${dev_idx}"
+        devices="$devices --device /dev/davinci${dev_idx}"
         ((i++))
     done
+
+    # trim leading space
+    devices="${devices#"${devices%%[![:space:]]*}"}"
+    # Output devices: assigned to the caller variable
+    printf '%s' "$devices"
 }
 
-mapfile -t device_args < <(parse_and_gen_devices "${BUILDKITE_AGENT_NAME}") || exit 1
+devices=$(parse_and_gen_devices "${BUILDKITE_AGENT_NAME}") || exit 1
 
 # Run the image and execute the Out-Of-Tree (OOT) platform interface test case on Ascend NPU hardware.
 # This test checks whether the OOT platform interface is functioning properly in conjunction with
 # the hardware plugin vllm-ascend.
 model_cache_dir=/mnt/modelscope${agent_idx}
-mkdir -p "${model_cache_dir}"
+mkdir -p ${model_cache_dir}
 docker run \
-    "${device_args[@]}" \
+    ${devices} \
     --device /dev/davinci_manager \
     --device /dev/devmm_svm \
     --device /dev/hisi_hdc \
@@ -177,7 +182,7 @@ docker run \
     -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
     -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
     -v /etc/ascend_install.info:/etc/ascend_install.info \
-    -v "${model_cache_dir}":/root/.cache/modelscope \
+    -v ${model_cache_dir}:/root/.cache/modelscope \
     --entrypoint="" \
     --name "${container_name}" \
     "${image_name}" \

.buildkite/scripts/hardware_ci/run-tpu-v1-test-part2.sh

@@ -61,7 +61,7 @@ echo "Results will be stored in: $RESULTS_DIR"
 echo "--- Installing Python dependencies ---"
 python3 -m pip install --progress-bar off git+https://github.com/thuml/depyf.git \
     && python3 -m pip install --progress-bar off pytest pytest-asyncio tpu-info \
-    && python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.11" \
+    && python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.9.2" \
     && python3 -m pip install --progress-bar off hf-transfer tblib==3.1.0
 echo "--- Python dependencies installed ---"
 

.buildkite/scripts/hardware_ci/run-tpu-v1-test.sh

@@ -61,7 +61,7 @@ echo "Results will be stored in: $RESULTS_DIR"
 echo "--- Installing Python dependencies ---"
 python3 -m pip install --progress-bar off git+https://github.com/thuml/depyf.git \
     && python3 -m pip install --progress-bar off pytest pytest-asyncio tpu-info \
-    && python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.11" \
+    && python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.9.2" \
     && python3 -m pip install --progress-bar off hf-transfer tblib==3.1.0
 echo "--- Python dependencies installed ---"
 

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

@@ -8,7 +8,7 @@ image_name="xpu/vllm-ci:${BUILDKITE_COMMIT}"
 container_name="xpu_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
 
 # Try building the docker image
-docker build -t "${image_name}" -f docker/Dockerfile.xpu .
+docker build -t ${image_name} -f docker/Dockerfile.xpu .
 
 # Setup cleanup
 remove_docker_container() {

.buildkite/scripts/push-nightly-builds.sh

@@ -21,16 +21,16 @@ echo "Pushing original tag $ORIG_TAG_NAME$ORIG_TAG_SUFFIX to new nightly tag nam
 
 # pull original arch-dependent images from AWS ECR Public
 aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:"$ORIG_TAG_NAME"-x86_64"$ORIG_TAG_SUFFIX"
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:"$ORIG_TAG_NAME"-aarch64"$ORIG_TAG_SUFFIX"
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:$ORIG_TAG_NAME-x86_64$ORIG_TAG_SUFFIX
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:$ORIG_TAG_NAME-aarch64$ORIG_TAG_SUFFIX
 # tag arch-dependent images
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:"$ORIG_TAG_NAME"-x86_64"$ORIG_TAG_SUFFIX" vllm/vllm-openai:"$TAG_NAME"-x86_64
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:"$ORIG_TAG_NAME"-aarch64"$ORIG_TAG_SUFFIX" vllm/vllm-openai:"$TAG_NAME"-aarch64
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$ORIG_TAG_NAME-x86_64$ORIG_TAG_SUFFIX vllm/vllm-openai:$TAG_NAME-x86_64
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$ORIG_TAG_NAME-aarch64$ORIG_TAG_SUFFIX vllm/vllm-openai:$TAG_NAME-aarch64
 # push arch-dependent images to DockerHub
-docker push vllm/vllm-openai:"$TAG_NAME"-x86_64
-docker push vllm/vllm-openai:"$TAG_NAME"-aarch64
+docker push vllm/vllm-openai:$TAG_NAME-x86_64
+docker push vllm/vllm-openai:$TAG_NAME-aarch64
 # push arch-independent manifest to DockerHub
-docker manifest create vllm/vllm-openai:"$TAG_NAME" vllm/vllm-openai:"$TAG_NAME"-x86_64 vllm/vllm-openai:"$TAG_NAME"-aarch64 --amend
-docker manifest create vllm/vllm-openai:"$TAG_NAME"-"$BUILDKITE_COMMIT" vllm/vllm-openai:"$TAG_NAME"-x86_64 vllm/vllm-openai:"$TAG_NAME"-aarch64 --amend
-docker manifest push vllm/vllm-openai:"$TAG_NAME"
-docker manifest push vllm/vllm-openai:"$TAG_NAME"-"$BUILDKITE_COMMIT"
+docker manifest create vllm/vllm-openai:$TAG_NAME vllm/vllm-openai:$TAG_NAME-x86_64 vllm/vllm-openai:$TAG_NAME-aarch64 --amend
+docker manifest create vllm/vllm-openai:$TAG_NAME-$BUILDKITE_COMMIT vllm/vllm-openai:$TAG_NAME-x86_64 vllm/vllm-openai:$TAG_NAME-aarch64 --amend
+docker manifest push vllm/vllm-openai:$TAG_NAME
+docker manifest push vllm/vllm-openai:$TAG_NAME-$BUILDKITE_COMMIT

.buildkite/scripts/run-prime-rl-test.sh

@@ -0,0 +1,64 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Setup script for Prime-RL integration tests
+# This script prepares the environment for running Prime-RL tests with nightly vLLM
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
+PRIME_RL_REPO="https://github.com/PrimeIntellect-ai/prime-rl.git"
+PRIME_RL_DIR="${REPO_ROOT}/prime-rl"
+
+if command -v rocm-smi &> /dev/null || command -v rocminfo &> /dev/null; then
+    echo "AMD GPU detected. Prime-RL currently only supports NVIDIA. Skipping..."
+    exit 0
+fi
+
+echo "Setting up Prime-RL integration test environment..."
+
+# Clean up any existing Prime-RL directory
+if [ -d "${PRIME_RL_DIR}" ]; then
+    echo "Removing existing Prime-RL directory..."
+    rm -rf "${PRIME_RL_DIR}"
+fi
+
+# Install UV if not available
+if ! command -v uv &> /dev/null; then
+    echo "Installing UV package manager..."
+    curl -LsSf https://astral.sh/uv/install.sh | sh
+    source $HOME/.local/bin/env
+fi
+
+# Clone Prime-RL repository at specific branch for reproducible tests
+PRIME_RL_BRANCH="integ-vllm-main"
+echo "Cloning Prime-RL repository at branch: ${PRIME_RL_BRANCH}..."
+git clone --branch "${PRIME_RL_BRANCH}" --single-branch "${PRIME_RL_REPO}" "${PRIME_RL_DIR}"
+cd "${PRIME_RL_DIR}"
+
+echo "Setting up UV project environment..."
+export UV_PROJECT_ENVIRONMENT=/usr/local
+ln -s /usr/bin/python3 /usr/local/bin/python
+
+# Remove vllm pin from pyproject.toml
+echo "Removing vllm pin from pyproject.toml..."
+sed -i '/vllm==/d' pyproject.toml
+
+# Sync Prime-RL dependencies
+echo "Installing Prime-RL dependencies..."
+uv sync --inexact && uv sync --inexact --all-extras
+
+# Verify installation
+echo "Verifying installations..."
+uv run python -c "import vllm; print(f'vLLM version: {vllm.__version__}')"
+uv run python -c "import prime_rl; print('Prime-RL imported successfully')"
+
+echo "Prime-RL integration test environment setup complete!"
+
+echo "Running Prime-RL integration tests..."
+export WANDB_MODE=offline # this makes this test not require a WANDB_API_KEY
+uv run pytest -vs tests/integration/test_rl.py -m gpu
+
+echo "Prime-RL integration tests completed!"

.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh

@@ -51,14 +51,14 @@ for BACK in "${BACKENDS[@]}"; do
     --enable-eplb \
     --trust-remote-code \
     --max-model-len 2048 \
-    --all2all-backend "$BACK" \
-    --port "$PORT" &
+    --all2all-backend $BACK \
+    --port $PORT &
   SERVER_PID=$!
-  wait_for_server "$PORT"
+  wait_for_server $PORT
 
   TAG=$(echo "$MODEL" | tr '/: \\n' '_____')
   OUT="${OUT_DIR}/${TAG}_${BACK}.json"
-  python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}"
+  python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port $PORT --num-questions ${NUM_Q} --save-results ${OUT}
   python3 - <<PY
 import json; acc=json.load(open('${OUT}'))['accuracy']
 print(f"${MODEL} ${BACK}: accuracy {acc:.3f}")

.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh

@@ -1,57 +0,0 @@
-#!/usr/bin/env bash
-set -euxo pipefail
-
-# Nightly e2e test for prefetch offloading with a MoE model.
-# Runs DeepSeek-V2-Lite with prefetch offloading of MoE expert weights
-# and validates GSM8K accuracy matches baseline (no offloading).
-#
-# args: [THRESHOLD] [NUM_QUESTIONS] [START_PORT]
-THRESHOLD=${1:-0.25}
-NUM_Q=${2:-1319}
-PORT=${3:-8030}
-OUT_DIR=${OUT_DIR:-/tmp/vllm-scheduled}
-mkdir -p "${OUT_DIR}"
-
-wait_for_server() {
-  local port=$1
-  timeout 600 bash -c '
-    until curl -sf "http://127.0.0.1:'"$port"'/health" > /dev/null; do
-      sleep 1
-    done'
-}
-
-MODEL="deepseek-ai/DeepSeek-V2-Lite"
-
-cleanup() {
-  if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
-    kill "${SERVER_PID}" 2>/dev/null || true
-    for _ in {1..20}; do
-      kill -0 "${SERVER_PID}" 2>/dev/null || break
-      sleep 0.5
-    done
-    kill -9 "${SERVER_PID}" 2>/dev/null || true
-  fi
-}
-trap cleanup EXIT
-
-vllm serve "$MODEL" \
-  --max-model-len 2048 \
-  --offload-group-size 8 \
-  --offload-num-in-group 2 \
-  --offload-prefetch-step 1 \
-  --offload-params w13_weight w2_weight \
-  --port "$PORT" &
-SERVER_PID=$!
-wait_for_server "$PORT"
-
-TAG=$(echo "$MODEL" | tr '/: \\n' '_____')
-OUT="${OUT_DIR}/${TAG}_prefetch_offload.json"
-python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}"
-python3 - <<PY
-import json; acc=json.load(open('${OUT}'))['accuracy']
-print(f"${MODEL} prefetch_offload: accuracy {acc:.3f}")
-assert acc >= ${THRESHOLD}, f"${MODEL} prefetch_offload accuracy {acc}"
-PY
-
-cleanup
-SERVER_PID=

.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh

@@ -47,20 +47,20 @@ for BACK in "${BACKENDS[@]}"; do
   vllm serve "$MODEL" \
     --enforce-eager \
     --enable-eplb \
-    --all2all-backend "$BACK" \
+    --all2all-backend $BACK \
     --eplb-config '{"window_size":10, "step_interval":100, "num_redundant_experts":0, "log_balancedness":true}' \
-    --tensor-parallel-size "${TENSOR_PARALLEL_SIZE}" \
-    --data-parallel-size "${DATA_PARALLEL_SIZE}" \
+    --tensor-parallel-size ${TENSOR_PARALLEL_SIZE} \
+    --data-parallel-size ${DATA_PARALLEL_SIZE} \
     --enable-expert-parallel \
     --trust-remote-code \
     --max-model-len 2048 \
-    --port "$PORT" &
+    --port $PORT &
   SERVER_PID=$!
-  wait_for_server "$PORT"
+  wait_for_server $PORT
 
   TAG=$(echo "$MODEL" | tr '/: \\n' '_____')
   OUT="${OUT_DIR}/${TAG}_${BACK}.json"
-  python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}"
+  python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port $PORT --num-questions ${NUM_Q} --save-results ${OUT}
   python3 - <<PY
 import json; acc=json.load(open('${OUT}'))['accuracy']
 print(f"${MODEL} ${BACK}: accuracy {acc:.3f}")

.buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh

@@ -51,20 +51,20 @@ for BACK in "${BACKENDS[@]}"; do
     --tensor-parallel-size 4 \
     --enable-expert-parallel \
     --enable-eplb \
-    --all2all-backend "$BACK" \
+    --all2all-backend $BACK \
     --eplb-config '{"window_size":200,"step_interval":600,"use_async":true}' \
     --speculative-config '{"method":"qwen3_next_mtp","num_speculative_tokens":1}' \
     --trust-remote-code \
     --max-model-len 2048 \
     --gpu-memory-utilization 0.9 \
     "${PLATFORM_ARGS[@]}" \
-    --port "$PORT" &
+    --port $PORT &
   SERVER_PID=$!
-  wait_for_server "$PORT"
+  wait_for_server $PORT
 
   TAG=$(echo "$MODEL" | tr '/: \\n' '_____')
   OUT="${OUT_DIR}/${TAG}_${BACK}.json"
-  python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}"
+  python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port $PORT --num-questions ${NUM_Q} --save-results ${OUT}
   python3 - <<PY
 import json; acc=json.load(open('${OUT}'))['accuracy']
 print(f"${MODEL} ${BACK}: accuracy {acc:.3f}")

.buildkite/scripts/tpu/docker_run_bm.sh

@@ -9,11 +9,10 @@ ENV_FILE=$1
 
 # For testing on local vm, use `set -a` to export all variables
 source /etc/environment
-# shellcheck source=/dev/null
-source "$ENV_FILE"
+source $ENV_FILE
 
 remove_docker_container() { 
-    docker rm -f "$CONTAINER_NAME" || true;
+    docker rm -f $CONTAINER_NAME || true;
 }
 
 trap remove_docker_container EXIT
@@ -42,13 +41,13 @@ echo
 echo "starting docker...$CONTAINER_NAME"
 echo    
 docker run \
- -v "$DOWNLOAD_DIR":"$DOWNLOAD_DIR" \
- --env-file "$ENV_FILE" \
+ -v $DOWNLOAD_DIR:$DOWNLOAD_DIR \
+ --env-file $ENV_FILE \
  -e HF_TOKEN="$HF_TOKEN" \
- -e TARGET_COMMIT="$BUILDKITE_COMMIT" \
- -e MODEL="$MODEL" \
+ -e TARGET_COMMIT=$BUILDKITE_COMMIT \
+ -e MODEL=$MODEL \
  -e WORKSPACE=/workspace \
- --name "$CONTAINER_NAME" \
+ --name $CONTAINER_NAME \
  -d \
  --privileged \
  --network host \

.buildkite/scripts/tpu/run_bm.sh

@@ -42,21 +42,21 @@ echo "lanching vllm..."
 echo "logging to $VLLM_LOG"
 echo
 
-vllm serve "$MODEL" \
+vllm serve $MODEL \
  --seed 42 \
- --max-num-seqs "$MAX_NUM_SEQS" \
- --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \
- --tensor-parallel-size "$TENSOR_PARALLEL_SIZE" \
+ --max-num-seqs $MAX_NUM_SEQS \
+ --max-num-batched-tokens $MAX_NUM_BATCHED_TOKENS \
+ --tensor-parallel-size $TENSOR_PARALLEL_SIZE \
  --no-enable-prefix-caching \
- --download_dir "$DOWNLOAD_DIR" \
- --max-model-len "$MAX_MODEL_LEN" > "$VLLM_LOG" 2>&1 &
+ --download_dir $DOWNLOAD_DIR \
+ --max-model-len $MAX_MODEL_LEN > "$VLLM_LOG" 2>&1 &
 
 
 echo "wait for 20 minutes.."
 echo
 # sleep 1200
 # wait for 10 minutes...
-for _ in {1..120}; do
+for i in {1..120}; do
     # TODO: detect other type of errors.
     if grep -Fq "raise RuntimeError" "$VLLM_LOG"; then
         echo "Detected RuntimeError, exiting."
@@ -78,11 +78,11 @@ echo "logging to $BM_LOG"
 echo
 vllm bench serve \
     --backend vllm \
-    --model "$MODEL"  \
+    --model $MODEL  \
     --dataset-name sonnet \
     --dataset-path benchmarks/sonnet_4x.txt \
-    --sonnet-input-len "$INPUT_LEN" \
-    --sonnet-output-len "$OUTPUT_LEN" \
+    --sonnet-input-len $INPUT_LEN \
+    --sonnet-output-len $OUTPUT_LEN \
     --ignore-eos > "$BM_LOG"
 
 echo "completed..."

.buildkite/scripts/upload-nightly-wheels.sh

@@ -76,15 +76,16 @@ mkdir -p "$INDICES_OUTPUT_DIR"
 # this indices have relative paths that could work as long as it is next to the wheel directory in s3
 # i.e., the wheels are always in s3://vllm-wheels/<commit>/
 # and indices can be placed in /<commit>/, or /nightly/, or /<version>/
-alias_args=()
-if [[ -n "$DEFAULT_VARIANT_ALIAS" ]]; then
-    alias_args=(--alias-to-default "$DEFAULT_VARIANT_ALIAS")
+if [[ ! -z "$DEFAULT_VARIANT_ALIAS" ]]; then
+    alias_arg="--alias-to-default $DEFAULT_VARIANT_ALIAS"
+else
+    alias_arg=""
 fi
 
 # HACK: we do not need regex module here, but it is required by pre-commit hook
 # To avoid any external dependency, we simply replace it back to the stdlib re module
 sed -i 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py
-$PYTHON .buildkite/scripts/generate-nightly-index.py --version "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "commit $BUILDKITE_COMMIT" "${alias_args[@]}"
+$PYTHON .buildkite/scripts/generate-nightly-index.py --version "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "commit $BUILDKITE_COMMIT" $alias_arg
 
 # copy indices to /<commit>/ unconditionally
 echo "Uploading indices to $S3_COMMIT_PREFIX"
@@ -99,9 +100,9 @@ fi
 # re-generate and copy to /<pure_version>/ only if it does not have "dev" in the version
 if [[ "$version" != *"dev"* ]]; then
     echo "Re-generating indices for /$pure_version/"
-    rm -rf "${INDICES_OUTPUT_DIR:?}/*"
+    rm -rf "$INDICES_OUTPUT_DIR/*"
     mkdir -p "$INDICES_OUTPUT_DIR"
     # wheel-dir is overridden to be the commit directory, so that the indices point to the correct wheel path
-    $PYTHON .buildkite/scripts/generate-nightly-index.py --version "$pure_version" --wheel-dir "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "version $pure_version" "${alias_args[@]}"
+    $PYTHON .buildkite/scripts/generate-nightly-index.py --version "$pure_version" --wheel-dir "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "version $pure_version" $alias_arg
     aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/$pure_version/"
 fi

.buildkite/scripts/upload-release-wheels-pypi.sh

@@ -7,7 +7,7 @@ SUBPATH=$BUILDKITE_COMMIT
 S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
 
 RELEASE_VERSION=$(buildkite-agent meta-data get release-version)
-GIT_VERSION=$(git describe --exact-match --tags "$BUILDKITE_COMMIT" 2>/dev/null)
+GIT_VERSION=$(git describe --exact-match --tags $BUILDKITE_COMMIT 2>/dev/null)
 
 echo "Release version from Buildkite: $RELEASE_VERSION"
 
@@ -55,7 +55,7 @@ mkdir -p $DIST_DIR
 aws s3 cp --recursive --exclude "*" --include "vllm-${PURE_VERSION}*.whl" --exclude "*dev*" --exclude "*rc[0-9]*" "$S3_COMMIT_PREFIX" $DIST_DIR
 echo "Wheels copied to local directory"
 # generate source tarball
-git archive --format=tar.gz --output="$DIST_DIR/vllm-${PURE_VERSION}.tar.gz" "$BUILDKITE_COMMIT"
+git archive --format=tar.gz --output="$DIST_DIR/vllm-${PURE_VERSION}.tar.gz" $BUILDKITE_COMMIT
 ls -la $DIST_DIR
 
 # upload wheels to PyPI (only default variant, i.e. files without '+' in the name)
@@ -65,6 +65,6 @@ if [[ -z "$PYPI_WHEEL_FILES" ]]; then
   exit 1
 fi
 
-python3 -m twine check "$PYPI_WHEEL_FILES"
-python3 -m twine upload --non-interactive --verbose "$PYPI_WHEEL_FILES"
+python3 -m twine check $PYPI_WHEEL_FILES
+python3 -m twine upload --non-interactive --verbose $PYPI_WHEEL_FILES
 echo "Wheels uploaded to PyPI"

.buildkite/scripts/upload-rocm-wheels.sh

@@ -55,7 +55,7 @@ mkdir -p all-rocm-wheels
 cp artifacts/rocm-base-wheels/*.whl all-rocm-wheels/ 2>/dev/null || true
 cp artifacts/rocm-vllm-wheel/*.whl all-rocm-wheels/ 2>/dev/null || true
 
-WHEEL_COUNT=$(find all-rocm-wheels -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l)
+WHEEL_COUNT=$(ls all-rocm-wheels/*.whl 2>/dev/null | wc -l)
 echo "Total wheels to upload: $WHEEL_COUNT"
 
 if [ "$WHEEL_COUNT" -eq 0 ]; then
@@ -115,7 +115,7 @@ if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]] |
 fi
 
 # Extract version from vLLM wheel and update version-specific index
-VLLM_WHEEL=$(find all-rocm-wheels -maxdepth 1 -name 'vllm*.whl' 2>/dev/null | head -1)
+VLLM_WHEEL=$(ls all-rocm-wheels/vllm*.whl 2>/dev/null | head -1)
 if [ -n "$VLLM_WHEEL" ]; then
     VERSION=$(unzip -p "$VLLM_WHEEL" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
     echo "Version in wheel: $VERSION"

.buildkite/test_areas/distributed.yaml

@@ -103,8 +103,8 @@ steps:
   - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
   # NEW rlhf examples
   - cd new_weight_syncing
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_nccl.py
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_ipc.py
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py
 
 - label: Distributed Tests (8 GPUs)(H100)
   timeout_in_minutes: 10
@@ -146,7 +146,6 @@ steps:
   num_devices: 2
   commands:
     - pytest -v -s tests/distributed/test_context_parallel.py
-    - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py
     - VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
     - pytest -v -s tests/v1/distributed/test_dbo.py
 
@@ -166,7 +165,6 @@ steps:
   num_devices: 2
   num_nodes: 2
   no_plugin: true
-  optional: true # TODO: revert once infra issue solved
   source_file_dependencies:
   - vllm/distributed/
   - vllm/engine/
@@ -199,18 +197,7 @@ steps:
     - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
     - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
-- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 4
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
-
-- label: Pipeline + Context Parallelism (4 GPUs)
+- label: Pipeline + Context Parallelism (4 GPUs))
   timeout_in_minutes: 60
   working_dir: "/vllm-workspace/tests"
   num_devices: 4

.buildkite/test_areas/e2e_integration.yaml

@@ -29,11 +29,15 @@ steps:
   commands:
   - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1
 
-- label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100)
-  timeout_in_minutes: 60
-  device: h100
+- label: Prime-RL Integration (2 GPUs)
+  timeout_in_minutes: 30
   optional: true
-  num_devices: 1
+  soft_fail: true
+  num_devices: 2
   working_dir: "/vllm-workspace"
+  source_file_dependencies:
+  - vllm/
+  - .buildkite/scripts/run-prime-rl-test.sh
   commands:
-  - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh 0.25 200 8030
+    - nvidia-smi
+    - bash .buildkite/scripts/run-prime-rl-test.sh

.buildkite/test_areas/engine.yaml

@@ -28,11 +28,3 @@ steps:
     - pytest -v -s v1/engine/test_preprocess_error_handling.py
     # Run the rest of v1/engine tests
     - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd
-      commands:
-      - pytest -v -s v1/e2e
-      - pytest -v -s v1/engine

.buildkite/test_areas/entrypoints.yaml

@@ -65,11 +65,6 @@ steps:
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s entrypoints/pooling
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd
 
 - label: Entrypoints Integration (Responses API)
   timeout_in_minutes: 50

.buildkite/test_areas/expert_parallelism.yaml

@@ -20,19 +20,4 @@ steps:
   - tests/distributed/test_eplb_execute.py
   commands:
   - pytest -v -s distributed/test_eplb_execute.py
-  - pytest -v -s distributed/test_eplb_spec_decode.py
-
-- label: Elastic EP Scaling Test
-  timeout_in_minutes: 20
-  device: b200
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 4
-  source_file_dependencies:
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/executor/
-  - vllm/compilation/
-  - tests/distributed/
-  commands:
-  - pytest -v -s distributed/test_elastic_ep.py
+  - pytest -v -s distributed/test_eplb_spec_decode.py

.buildkite/test_areas/kernels.yaml

@@ -70,7 +70,7 @@ steps:
   - tests/kernels/moe/test_batched_deepgemm.py
   - tests/kernels/attention/test_deepgemm_attention.py
   commands:
-    - pytest -v -s kernels/quantization/test_block_fp8.py
+    - pytest -v -s kernels/quantization/test_block_fp8.py -k deep_gemm
     - pytest -v -s kernels/moe/test_deepgemm.py
     - pytest -v -s kernels/moe/test_batched_deepgemm.py
     - pytest -v -s kernels/attention/test_deepgemm_attention.py
@@ -115,7 +115,6 @@ steps:
     - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
     - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
     - pytest -v -s tests/kernels/moe/test_flashinfer.py
-    - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
     - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
     # e2e
     - pytest -v -s tests/models/quantization/test_nvfp4.py
@@ -155,7 +154,9 @@ steps:
   commands:
     - pytest -v -s kernels/moe/test_deepep_deepgemm_moe.py
     - pytest -v -s kernels/moe/test_deepep_moe.py
-
+    - pytest -v -s kernels/moe/test_pplx_cutlass_moe.py
+    # - pytest -v -s kernels/moe/test_pplx_moe.py - failing on main
+  
 - label: Kernels Fp4 MoE Test (B200)
   timeout_in_minutes: 60
   device: b200

.buildkite/test_areas/lm_eval.yaml

@@ -73,29 +73,3 @@ steps:
   num_devices: 2
   commands:
     - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
-
-- label: GPQA Eval (GPT-OSS) (H100)
-  timeout_in_minutes: 120
-  device: h100
-  optional: true
-  num_devices: 2
-  source_file_dependencies:
-  - csrc/
-  - vllm/model_executor/layers/quantization
-  - tests/evals/gpt_oss/
-  commands:
-    - uv pip install --system 'gpt-oss[eval]==0.0.5'
-    - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-h100.txt
-
-- label: GPQA Eval (GPT-OSS) (B200)
-  timeout_in_minutes: 120
-  device: b200
-  optional: true
-  num_devices: 2
-  source_file_dependencies:
-  - csrc/
-  - vllm/model_executor/layers/quantization
-  - tests/evals/gpt_oss/
-  commands:
-    - uv pip install --system 'gpt-oss[eval]==0.0.5'
-    - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt

.buildkite/test_areas/misc.yaml

@@ -9,7 +9,6 @@ steps:
     - tests/v1
   commands:
     - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     # split the test to avoid interference
     - pytest -v -s -m 'not cpu_test' v1/core
     - pytest -v -s v1/executor
@@ -17,7 +16,6 @@ steps:
     - pytest -v -s v1/sample
     - pytest -v -s v1/logits_processors
     - pytest -v -s v1/worker
-    # TODO: create another `optional` test group for slow tests
     - pytest -v -s -m 'not slow_test' v1/spec_decode
     - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
     - pytest -v -s -m 'not cpu_test' v1/metrics
@@ -27,11 +25,6 @@ steps:
     # Integration test for streaming correctness (requires special branch).
     - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
     - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd
 
 - label: V1 Others (CPU)
   depends_on:
@@ -115,11 +108,9 @@ steps:
   timeout_in_minutes: 50
   source_file_dependencies:
   - vllm/
-  - tests/detokenizer
   - tests/multimodal
   - tests/utils_
   commands:
-  - pytest -v -s detokenizer
   - pytest -v -s -m 'not cpu_test' multimodal
   - pytest -v -s utils_
 
@@ -132,7 +123,6 @@ steps:
   - tests/test_inputs.py
   - tests/test_outputs.py
   - tests/test_pooling_params.py
-  - tests/test_ray_env.py
   - tests/multimodal
   - tests/renderers
   - tests/standalone_tests/lazy_imports.py
@@ -146,7 +136,6 @@ steps:
   - pytest -v -s test_inputs.py
   - pytest -v -s test_outputs.py
   - pytest -v -s test_pooling_params.py
-  - pytest -v -s test_ray_env.py
   - pytest -v -s -m 'cpu_test' multimodal
   - pytest -v -s renderers
   - pytest -v -s tokenizers_
@@ -154,6 +143,20 @@ steps:
   - pytest -v -s transformers_utils
   - pytest -v -s config
 
+- label: GPT-OSS Eval (B200)
+  timeout_in_minutes: 60
+  working_dir: "/vllm-workspace/"
+  device: b200
+  optional: true
+  source_file_dependencies:
+  - tests/evals/gpt_oss
+  - vllm/model_executor/models/gpt_oss.py
+  - vllm/model_executor/layers/quantization/mxfp4.py
+  - vllm/v1/attention/backends/flashinfer.py
+  commands:
+    - uv pip install --system 'gpt-oss[eval]==0.0.5'
+    - pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py --model openai/gpt-oss-20b --metric 0.58
+
 - label: Batch Invariance (H100)
   timeout_in_minutes: 25
   device: h100

.buildkite/test_areas/models_language.yaml

@@ -55,15 +55,6 @@ steps:
     - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
     - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
     - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd
-      commands:
-      - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
-      - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
-      - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
 
 - label: Language Models Test (PPL)
   timeout_in_minutes: 110
@@ -82,11 +73,6 @@ steps:
   - tests/models/language/pooling
   commands:
     - pytest -v -s models/language/pooling -m 'not core_model'
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd
 
 - label: Language Models Test (MTEB)
   timeout_in_minutes: 110

.buildkite/test_areas/plugins.yaml

@@ -19,10 +19,6 @@ steps:
   - pip install -e ./plugins/prithvi_io_processor_plugin
   - pytest -v -s plugins_tests/test_io_processor_plugins.py
   - pip uninstall prithvi_io_processor_plugin -y
-  # test bge_m3_sparse io_processor plugin
-  - pip install -e ./plugins/bge_m3_sparse_plugin
-  - pytest -v -s plugins_tests/test_bge_m3_sparse_io_processor_plugins.py
-  - pip uninstall bge_m3_sparse_plugin -y
   # end io_processor plugins test
   # begin stat_logger plugins test
   - pip install -e ./plugins/vllm_add_dummy_stat_logger

.buildkite/test_areas/samplers.yaml

@@ -18,4 +18,4 @@ steps:
       depends_on:
       - image-build-amd
       commands:
-      - pytest -v -s samplers
+      - pytest -v -s -m 'not skip_v1' samplers

.github/.bc-linter.yml

@@ -0,0 +1,24 @@
+# doc: https://github.com/pytorch/test-infra/blob/main/tools/stronghold/docs/bc_linter_config.md
+version: 1
+paths:
+# We temporarily disable globally, and will only enable with `annotations.include`
+# include:
+#   - "vllm/v1/attetion/*.py"
+#   - "vllm/v1/core/*.py"
+exclude:
+  - "**/*.py"
+
+scan:
+  functions: true        # check free functions and methods
+  classes: true          # check classes/dataclasses
+  public_only: true      # ignore names starting with "_" at any level
+
+annotations:
+  include:               # decorators that force‑include a symbol
+    - name: "bc_linter_include"  # matched by simple name or dotted suffix
+      propagate_to_members: false # for classes, include methods/inner classes
+  exclude:               # decorators that force‑exclude a symbol
+    - name: "bc_linter_skip"     # matched by simple name or dotted suffix
+      propagate_to_members: true  # for classes, exclude methods/inner classes
+
+excluded_violations: []  # e.g. ["ParameterRenamed", "FieldTypeChanged"]

.github/CODEOWNERS

@@ -2,17 +2,17 @@
 # for more info about CODEOWNERS file
 
 # This lists cover the "core" components of vLLM that require careful review
-/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng
+/vllm/compilation @zou3519 @youkaichao @ProExpertProg
 /vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery
 /vllm/lora @jeejeelee
-/vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni
+/vllm/model_executor/layers/attention @LucasWilkinson
 /vllm/model_executor/layers/fused_moe @mgoin @pavanimajety
 /vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety
 /vllm/model_executor/layers/mamba @tdoublep
 /vllm/model_executor/model_loader @22quinn
 /vllm/model_executor/layers/batch_invariant.py @yewentao256 
 /vllm/multimodal @DarkLight1337 @ywang96 @NickLucche @tjtanaa
-/vllm/vllm_flash_attn @LucasWilkinson @MatthewBonanni
+/vllm/vllm_flash_attn @LucasWilkinson
 CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 
 # Any change to the VllmConfig changes can have a large user-facing impact,
@@ -43,25 +43,22 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /vllm/tool_parsers @aarnphm @chaunceyjiang
 
 # vLLM V1
-/vllm/v1/attention @LucasWilkinson @MatthewBonanni
+/vllm/v1/attention @LucasWilkinson
 /vllm/v1/attention/backend.py @WoosukKwon @zhuohan123 @youkaichao @alexm-redhat @njhill
 /vllm/v1/attention/backends/mla @pavanimajety
 /vllm/v1/attention/backends/flashinfer.py @mgoin @pavanimajety
 /vllm/v1/attention/backends/triton_attn.py @tdoublep
 /vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery
 /vllm/v1/sample @22quinn @houseroad @njhill
-/vllm/v1/spec_decode @benchislett @luccafong @MatthewBonanni
+/vllm/v1/spec_decode @benchislett @luccafong
 /vllm/v1/structured_output @mgoin @russellb @aarnphm @benchislett
 /vllm/v1/kv_cache_interface.py @heheda12345
 /vllm/v1/kv_offload @ApostaC @orozery
-/vllm/v1/engine @njhill
-/vllm/v1/executor @njhill
-/vllm/v1/worker @njhill
-/vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery @NickLucche
+/vllm/v1/worker/gpu/kv_connector.py @orozery
+/vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery
 
 # Model runner V2
-/vllm/v1/worker/gpu @WoosukKwon @njhill
-/vllm/v1/worker/gpu/kv_connector.py @orozery
+/vllm/v1/worker/gpu @WoosukKwon
 
 # Test ownership
 /.buildkite/lm-eval-harness @mgoin 

.github/mergify.yml

@@ -259,7 +259,8 @@ pull_request_rules:
       - files=benchmarks/run_structured_output_benchmark.sh
       - files=docs/features/structured_outputs.md
       - files=examples/offline_inference/structured_outputs.py
-      - files=examples/online_serving/structured_outputs/structured_outputs.py
+      - files=examples/online_serving/openai_chat_completion_structured_outputs.py
+      - files=examples/online_serving/openai_chat_completion_structured_outputs_with_reasoning.py
       - files~=^tests/v1/structured_output/
       - files=tests/v1/entrypoints/llm/test_struct_output_generate.py
       - files~=^vllm/v1/structured_output/

.github/workflows/bc-lint.yml

@@ -0,0 +1,29 @@
+name: BC Lint
+
+on:
+  pull_request:
+    types:
+      - opened
+      - synchronize
+      - reopened
+      - labeled
+      - unlabeled
+
+jobs:
+  bc_lint:
+    if: github.repository_owner == 'vllm-project'
+    runs-on: ubuntu-latest
+    steps:
+      - name: Run BC Lint Action
+        uses: pytorch/test-infra/.github/actions/bc-lint@main
+        with:
+          repo: ${{ github.event.pull_request.head.repo.full_name }}
+          base_sha: ${{ github.event.pull_request.base.sha }}
+          head_sha: ${{ github.event.pull_request.head.sha }}
+          suppression: ${{ contains(github.event.pull_request.labels.*.name, 'suppress-bc-linter') }}
+          docs_link: 'https://github.com/pytorch/test-infra/wiki/BC-Linter'
+          config_dir: .github
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: true

.gitignore

@@ -3,8 +3,6 @@
 
 # vllm-flash-attn built from source
 vllm/vllm_flash_attn/*
-!vllm/vllm_flash_attn/__init__.py
-!vllm/vllm_flash_attn/flash_attn_interface.py
 
 # OpenAI triton kernels copied from source
 vllm/third_party/triton_kernels/*

CMakeLists.txt

@@ -725,7 +725,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   # CUTLASS MoE kernels
 
   # The MoE kernel cutlass_moe_mm requires CUDA 12.3 or later (and ONLY works
-  # on Hopper). get_cutlass_(batched_)moe_mm_data should only be compiled
+  # on Hopper). get_cutlass_(pplx_)moe_mm_data should only be compiled
   # if it's possible to compile MoE kernels that use its output.
   cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a" "${CUDA_ARCHS}")
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS)
@@ -771,24 +771,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     endif()
   endif()
 
-  # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
-  endif()
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS)
-    set(DSV3_FUSED_A_GEMM_SRC "csrc/dsv3_fused_a_gemm.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${DSV3_FUSED_A_GEMM_SRC}"
-      CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}")
-    list(APPEND VLLM_EXT_SRC ${DSV3_FUSED_A_GEMM_SRC})
-    message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}")
-  else()
-    message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found "
-                   "in CUDA target architectures.")
-  endif()
-
   # moe_data.cu is used by all CUTLASS MoE kernels.
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
     cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
@@ -971,8 +953,7 @@ set(VLLM_MOE_EXT_SRC
 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/grouped_topk_kernels.cu")
 endif()
 
 if(VLLM_GPU_LANG STREQUAL "CUDA")
@@ -1101,27 +1082,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     message(STATUS "Not building Marlin MOE kernels as no compatible archs found"
                    " in CUDA target architectures")
   endif()
-
-  # DeepSeek V3 router GEMM kernel - requires SM90+
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
-  endif()
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_ROUTER_GEMM_ARCHS)
-    set(DSV3_ROUTER_GEMM_SRC
-      "csrc/moe/dsv3_router_gemm_entry.cu"
-      "csrc/moe/dsv3_router_gemm_float_out.cu"
-      "csrc/moe/dsv3_router_gemm_bf16_out.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${DSV3_ROUTER_GEMM_SRC}"
-      CUDA_ARCHS "${DSV3_ROUTER_GEMM_ARCHS}")
-    list(APPEND VLLM_MOE_EXT_SRC "${DSV3_ROUTER_GEMM_SRC}")
-    message(STATUS "Building DSV3 router GEMM kernel for archs: ${DSV3_ROUTER_GEMM_ARCHS}")
-  else()
-    message(STATUS "Not building DSV3 router GEMM kernel as no compatible archs found"
-                   " (requires SM90+ and CUDA >= 12.0)")
-  endif()
 endif()
 
 message(STATUS "Enabling moe extension.")

benchmarks/attention_benchmarks/__init__.py

@@ -15,6 +15,7 @@ from .common import (
     BenchmarkConfig,
     BenchmarkResult,
     MockLayer,
+    MockModelConfig,
     ResultsFormatter,
     get_attention_scale,
     is_mla_backend,
@@ -35,6 +36,7 @@ __all__ = [
     "ResultsFormatter",
     # Mock objects
     "MockLayer",
+    "MockModelConfig",
     # Utilities
     "setup_mla_dims",
     "get_attention_scale",

benchmarks/attention_benchmarks/common.py

@@ -10,6 +10,7 @@ from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import Any
 
+import numpy as np
 import torch
 from batch_spec import get_batch_type, parse_batch_spec
 from rich.console import Console
@@ -61,7 +62,10 @@ class MockHfConfig:
 # Import AttentionLayerBase at module level to avoid circular dependencies
 try:
     from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+
+    _HAS_ATTENTION_LAYER_BASE = True
 except ImportError:
+    _HAS_ATTENTION_LAYER_BASE = False
     AttentionLayerBase = object  # Fallback
 
 
@@ -163,6 +167,95 @@ class MockLayer(AttentionLayerBase):
         return self._kv_cache_spec
 
 
+class MockModelConfig:
+    """Mock model configuration."""
+
+    def __init__(
+        self,
+        num_q_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        dtype: torch.dtype = torch.float16,
+        max_model_len: int = 32768,
+    ):
+        self._n_q = num_q_heads
+        self._n_kv = num_kv_heads
+        self._d = head_dim
+        self.dtype = dtype
+        self.max_model_len = max_model_len
+
+    def get_num_attention_heads(self, _=None) -> int:
+        return self._n_q
+
+    def get_num_kv_heads(self, _=None) -> int:
+        return self._n_kv
+
+    def get_head_size(self) -> int:
+        return self._d
+
+    def get_num_layers(self) -> int:
+        """Mock method for layer count queries."""
+        return 1
+
+    def get_sliding_window_for_layer(self, _layer_idx: int):
+        """Mock method for sliding window queries."""
+        return None
+
+    def get_logits_soft_cap_for_layer(self, _layer_idx: int):
+        """Mock method for logits soft cap queries."""
+        return None
+
+    def get_sm_scale_for_layer(self, _layer_idx: int) -> float:
+        """Mock method for SM scale queries."""
+        return 1.0 / (self.get_head_size() ** 0.5)
+
+
+class MockParallelConfig:
+    """Mock parallel configuration."""
+
+    pass
+
+
+class MockCompilationConfig:
+    """Mock compilation configuration."""
+
+    def __init__(self):
+        self.full_cuda_graph = False
+        self.static_forward_context = {}
+
+
+class MockVLLMConfig:
+    """Mock VLLM configuration."""
+
+    def __init__(self):
+        self.compilation_config = MockCompilationConfig()
+
+
+class MockRunner:
+    """Mock GPU runner for metadata builders."""
+
+    def __init__(
+        self,
+        seq_lens: np.ndarray,
+        query_start_locs: np.ndarray,
+        device: torch.device,
+        num_q_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        dtype: torch.dtype,
+    ):
+        self.model_config = MockModelConfig(num_q_heads, num_kv_heads, head_dim, dtype)
+        self.parallel_config = MockParallelConfig()
+        self.vllm_config = MockVLLMConfig()
+        self.seq_lens_np = seq_lens
+        self.query_start_loc_np = query_start_locs
+        self.device = device
+        self.attention_chunk_size = None
+        self.num_query_heads = num_q_heads
+        self.num_kv_heads = num_kv_heads
+        self.dtype = dtype
+
+
 @dataclass
 class ParameterSweep:
     """Configuration for sweeping a backend parameter."""

benchmarks/auto_tune/auto_tune.sh

@@ -46,10 +46,10 @@ echo "VLLM_LOGGING_LEVEL=$VLLM_LOGGING_LEVEL"
 echo "RESULT_FILE=$RESULT"
 echo "====================== AUTO TUNEPARAMETERS ===================="
 
-rm -rf "$LOG_FOLDER"
-rm -rf "$PROFILE_PATH"
-mkdir -p "$LOG_FOLDER"
-mkdir -p "$PROFILE_PATH"
+rm -rf $LOG_FOLDER
+rm -rf $PROFILE_PATH
+mkdir -p $LOG_FOLDER
+mkdir -p $PROFILE_PATH
 
 cd "$BASE/vllm"
 
@@ -114,7 +114,7 @@ start_server() {
 
     # wait for 10 minutes...
     server_started=0
-    for _ in {1..60}; do
+    for i in {1..60}; do
         # This line checks whether the server is still alive or not,
         # since that we should always have permission to send signal to the server process.
         kill -0 $server_pid 2> /dev/null || break
@@ -145,12 +145,12 @@ run_benchmark() {
     local vllm_log="$LOG_FOLDER/vllm_log_${max_num_seqs}_${max_num_batched_tokens}.txt"
     echo "vllm_log: $vllm_log"
     echo
-    rm -f "$vllm_log"
+    rm -f $vllm_log
     pkill -if "vllm serve" || true
 
     echo "starting server..."
     # Call start_server without a profile_dir to avoid profiling overhead
-    start_server "$gpu_memory_utilization" "$max_num_seqs" "$max_num_batched_tokens" "$vllm_log" ""
+    start_server $gpu_memory_utilization $max_num_seqs $max_num_batched_tokens $vllm_log ""
     result=$?
     if [[ "$result" -eq 1 ]]; then
         echo "server failed to start. gpu_memory_utilization:$gpu_memory_utilization, max_num_seqs:$max_num_seqs, max_num_batched_tokens: $max_num_batched_tokens"
@@ -168,15 +168,15 @@ run_benchmark() {
     # --profile flag is removed from this call
     vllm bench serve \
         --backend vllm \
-        --model "$MODEL"  \
+        --model $MODEL  \
         --dataset-name random \
         --random-input-len $adjusted_input_len \
-        --random-output-len "$OUTPUT_LEN" \
+        --random-output-len $OUTPUT_LEN \
         --ignore-eos \
         --disable-tqdm \
         --request-rate inf \
         --percentile-metrics ttft,tpot,itl,e2el \
-        --goodput e2el:"$MAX_LATENCY_ALLOWED_MS" \
+        --goodput e2el:$MAX_LATENCY_ALLOWED_MS \
         --num-prompts 1000 \
         --random-prefix-len $prefix_len \
         --host "$HOSTNAME" \
@@ -195,20 +195,20 @@ run_benchmark() {
         request_rate=$((${throughput%.*} + 1))
         while ((request_rate > 0)); do
             # clear prefix cache
-            curl -X POST http://"${HOSTNAME}":8004/reset_prefix_cache
+            curl -X POST http://${HOSTNAME}:8004/reset_prefix_cache
             sleep 5
             bm_log="$LOG_FOLDER/bm_log_${max_num_seqs}_${max_num_batched_tokens}_requestrate_${request_rate}.txt"
             vllm bench serve \
                 --backend vllm \
-                --model "$MODEL"  \
+                --model $MODEL  \
                 --dataset-name random \
                 --random-input-len $adjusted_input_len \
-                --random-output-len "$OUTPUT_LEN" \
+                --random-output-len $OUTPUT_LEN \
                 --ignore-eos \
                 --disable-tqdm \
                 --request-rate $request_rate \
                 --percentile-metrics ttft,tpot,itl,e2el \
-                --goodput e2el:"$MAX_LATENCY_ALLOWED_MS" \
+                --goodput e2el:$MAX_LATENCY_ALLOWED_MS \
                 --num-prompts 100 \
                 --random-prefix-len $prefix_len \
                 --host "$HOSTNAME" \
@@ -255,7 +255,7 @@ gpu_memory_utilization=0.98
 find_gpu_memory_utilization=0
 while (( $(echo "$gpu_memory_utilization >= 0.9" | bc -l) )); do
     # Pass empty string for profile_dir argument
-    start_server "$gpu_memory_utilization" "${num_seqs_list[-1]}" "${num_batched_tokens_list[-1]}" "$LOG_FOLDER/vllm_log_gpu_memory_utilization_$gpu_memory_utilization.log" ""
+    start_server $gpu_memory_utilization "${num_seqs_list[-1]}" "${num_batched_tokens_list[-1]}" "$LOG_FOLDER/vllm_log_gpu_memory_utilization_$gpu_memory_utilization.log" ""
     result=$?
     if [[ "$result" -eq 0 ]]; then
         find_gpu_memory_utilization=1
@@ -274,7 +274,7 @@ fi
 
 for num_seqs in "${num_seqs_list[@]}"; do
     for num_batched_tokens in "${num_batched_tokens_list[@]}"; do
-        run_benchmark "$num_seqs" "$num_batched_tokens" "$gpu_memory_utilization"
+        run_benchmark $num_seqs $num_batched_tokens $gpu_memory_utilization
     done
 done
 echo "finish permutations"
@@ -285,7 +285,7 @@ echo "finish permutations"
 if (( $(echo "$best_throughput > 0" | bc -l) )); then
     echo
     echo "Benchmark tuning finished. Now running profiling on the best configuration found..."
-    echo "Best config: max_num_seqs: $best_max_num_seqs, max_num_batched_tokens: $best_num_batched_tokens, throughput: $best_throughput, goodput: $best_goodput"
+    echo "Best config: max_num_seqs: $best_max_num_seqs, max_num_batched_tokens: $best_num_batched_tokens, throughput: $best_throughput"
     echo
 
     vllm_log="$LOG_FOLDER/vllm_log_BEST_PROFILE.txt"
@@ -293,7 +293,7 @@ if (( $(echo "$best_throughput > 0" | bc -l) )); then
 
     # Start server with the best params and profiling ENABLED
     echo "Starting server for profiling..."
-    start_server "$gpu_memory_utilization" "$best_max_num_seqs" "$best_num_batched_tokens" "$vllm_log" "$PROFILE_PATH"
+    start_server $gpu_memory_utilization $best_max_num_seqs $best_num_batched_tokens "$vllm_log" "$PROFILE_PATH"
 
     # Run benchmark with the best params and the --profile flag
     echo "Running benchmark with profiling..."
@@ -301,15 +301,15 @@ if (( $(echo "$best_throughput > 0" | bc -l) )); then
     adjusted_input_len=$(( INPUT_LEN - prefix_len ))
     vllm bench serve \
         --backend vllm \
-        --model "$MODEL" \
+        --model $MODEL \
         --dataset-name random \
         --random-input-len $adjusted_input_len \
-        --random-output-len "$OUTPUT_LEN" \
+        --random-output-len $OUTPUT_LEN \
         --ignore-eos \
         --disable-tqdm \
-        --request-rate "$best_request_rate" \
+        --request-rate $best_request_rate \
         --percentile-metrics ttft,tpot,itl,e2el \
-        --goodput e2el:"$MAX_LATENCY_ALLOWED_MS" \
+        --goodput e2el:$MAX_LATENCY_ALLOWED_MS \
         --num-prompts 100 \
         --random-prefix-len $prefix_len \
         --host "$HOSTNAME" \

benchmarks/auto_tune/batch_auto_tune.sh

@@ -64,7 +64,7 @@ for i in $(seq 0 $(($num_runs - 1))); do
   else
     STATUS="FAILURE"
     ((FAILURE_COUNT++))
-    FAILED_RUNS+=("Run #$((i+1)): $(echo "$run_object" | jq -c .)")
+    FAILED_RUNS+=("Run #$((i+1)): $(echo $run_object | jq -c .)")
   fi
 
   RUN_OUTPUT=$(<"$RUN_OUTPUT_FILE")

benchmarks/backend_request_func.py

@@ -649,3 +649,9 @@ ASYNC_REQUEST_FUNCS = {
     "sglang": async_request_openai_completions,
     "llama.cpp": async_request_openai_completions,
 }
+
+OPENAI_COMPATIBLE_BACKENDS = [
+    k
+    for k, v in ASYNC_REQUEST_FUNCS.items()
+    if v in (async_request_openai_completions, async_request_openai_chat_completions)
+]

benchmarks/benchmark_topk_topp.py

@@ -1,471 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark comparing Triton vs PyTorch sort-based top-k/top-p implementations.
-
-Compares:
-- apply_top_k_top_p_triton (Triton binary search)
-- apply_top_k_top_p (PyTorch sort-based)
-
-Scenarios:
-- top_k only (whole batch, partial batch)
-- top_p only (whole batch, partial batch)
-- mix of top_k and top_p
-"""
-
-import argparse
-import gc
-from dataclasses import dataclass
-
-import torch
-
-from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch
-from vllm.v1.sample.ops.topk_topp_triton import (
-    apply_top_k_top_p_triton,
-    reset_buffer_cache,
-)
-
-
-@dataclass
-class BenchmarkConfig:
-    """Configuration for a benchmark run."""
-
-    name: str
-    batch_size: int
-    vocab_size: int
-    # k and p can be tensors or None
-    k_values: torch.Tensor | None  # [batch_size] or None
-    p_values: torch.Tensor | None  # [batch_size] or None
-    description: str
-    ops_pct: float = 0.0  # Percentage of ops relative to batch size
-
-
-def calculate_ops_pct(
-    k_values: torch.Tensor | None,
-    p_values: torch.Tensor | None,
-    vocab_size: int,
-    batch_size: int,
-) -> float:
-    """
-    Calculate the percentage of active top-k and top-p operations.
-
-    Returns percentage where 100% = batch_size ops.
-    E.g., if all rows have both top-k and top-p active, returns 200%.
-    """
-    active_ops = 0
-
-    if k_values is not None:
-        # Count rows where k < vocab_size (active top-k filtering)
-        active_ops += (k_values < vocab_size).sum().item()
-
-    if p_values is not None:
-        # Count rows where p < 1.0 (active top-p filtering)
-        active_ops += (p_values < 1.0).sum().item()
-
-    return (active_ops / batch_size) * 100 if batch_size > 0 else 0.0
-
-
-def create_logits(
-    batch_size: int, vocab_size: int, device: str = "cuda"
-) -> torch.Tensor:
-    """Create random logits mimicking a realistic LLM distribution.
-
-    Uses a Zipf-like probability distribution (rank^-1.1) converted to logits
-    via log, then randomly permuted per row. This produces a peaked distribution
-    where a small number of tokens capture most probability mass, similar to
-    real model outputs.
-    """
-    # Create Zipf-like probabilities: p(rank) ~ rank^(-alpha)
-    ranks = torch.arange(1, vocab_size + 1, dtype=torch.float32, device=device)
-    probs = ranks.pow(-1.1)
-    probs = probs / probs.sum()
-
-    # Convert to logits (log-probabilities, unnormalized is fine)
-    base_logits = probs.log()
-
-    # Broadcast to batch and randomly permute each row
-    logits = base_logits.unsqueeze(0).expand(batch_size, -1).clone()
-    for i in range(batch_size):
-        logits[i] = logits[i, torch.randperm(vocab_size, device=device)]
-
-    return logits
-
-
-def measure_memory() -> tuple[int, int]:
-    """Return (allocated, reserved) memory in bytes."""
-    torch.cuda.synchronize()
-    return torch.cuda.memory_allocated(), torch.cuda.max_memory_allocated()
-
-
-def reset_memory_stats():
-    """Reset peak memory statistics."""
-    reset_buffer_cache()
-    torch.cuda.reset_peak_memory_stats()
-    torch.cuda.empty_cache()
-    gc.collect()
-
-
-def benchmark_function(
-    func,
-    logits: torch.Tensor,
-    k: torch.Tensor | None,
-    p: torch.Tensor | None,
-    warmup_iters: int = 5,
-    benchmark_iters: int = 20,
-) -> tuple[float, int]:
-    """
-    Benchmark a function and return (avg_time_ms, peak_memory_bytes).
-
-    Returns average time in milliseconds and peak memory usage.
-    """
-    # Warmup
-    for _ in range(warmup_iters):
-        logits_copy = logits.clone()
-        func(logits_copy, k, p)
-    torch.cuda.synchronize()
-
-    # Reset memory stats before benchmark
-    reset_memory_stats()
-
-    # Benchmark
-    start_events = [
-        torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)
-    ]
-    end_events = [torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)]
-
-    for i in range(benchmark_iters):
-        logits_copy = logits.clone()
-        start_events[i].record()
-        func(logits_copy, k, p)
-        end_events[i].record()
-
-    torch.cuda.synchronize()
-
-    # Calculate timing
-    times = [
-        start_events[i].elapsed_time(end_events[i]) for i in range(benchmark_iters)
-    ]
-    avg_time = sum(times) / len(times)
-
-    # Get peak memory
-    _, peak_memory = measure_memory()
-
-    return avg_time, peak_memory
-
-
-def create_benchmark_configs(
-    batch_sizes: list[int],
-    vocab_sizes: list[int],
-    device: str = "cuda",
-) -> list[BenchmarkConfig]:
-    """Create all benchmark configurations."""
-    configs = []
-
-    for vocab_size in vocab_sizes:
-        for batch_size in batch_sizes:
-            # 1. Top-k only - whole batch (all rows have k < vocab_size)
-            k_all = torch.full((batch_size,), 50, dtype=torch.int32, device=device)
-            configs.append(
-                BenchmarkConfig(
-                    name=f"topk_whole_b{batch_size}_v{vocab_size // 1000}k",
-                    batch_size=batch_size,
-                    vocab_size=vocab_size,
-                    k_values=k_all,
-                    p_values=None,
-                    description=f"Top-k only (whole batch, k=50), "
-                    f"batch={batch_size}, vocab={vocab_size}",
-                    ops_pct=calculate_ops_pct(k_all, None, vocab_size, batch_size),
-                )
-            )
-
-            # 2. Top-k only - partial batch (half have k=50, half have k=vocab_size)
-            k_partial = torch.full((batch_size,), 50, dtype=torch.int32, device=device)
-            k_partial[batch_size // 2 :] = vocab_size  # No filtering for second half
-            configs.append(
-                BenchmarkConfig(
-                    name=f"topk_partial_b{batch_size}_v{vocab_size // 1000}k",
-                    batch_size=batch_size,
-                    vocab_size=vocab_size,
-                    k_values=k_partial,
-                    p_values=None,
-                    description=f"Top-k only (partial batch, 50% k=50, 50% k=vocab), "
-                    f"batch={batch_size}, vocab={vocab_size}",
-                    ops_pct=calculate_ops_pct(k_partial, None, vocab_size, batch_size),
-                )
-            )
-
-            # 3. Top-p only - whole batch (all rows have p < 1.0)
-            p_all = torch.full((batch_size,), 0.9, dtype=torch.float32, device=device)
-            configs.append(
-                BenchmarkConfig(
-                    name=f"topp_whole_b{batch_size}_v{vocab_size // 1000}k",
-                    batch_size=batch_size,
-                    vocab_size=vocab_size,
-                    k_values=None,
-                    p_values=p_all,
-                    description=f"Top-p only (whole batch, p=0.9), "
-                    f"batch={batch_size}, vocab={vocab_size}",
-                    ops_pct=calculate_ops_pct(None, p_all, vocab_size, batch_size),
-                )
-            )
-
-            # 4. Top-p only - partial batch (half have p=0.9, half have p=1.0)
-            p_partial = torch.full(
-                (batch_size,), 0.9, dtype=torch.float32, device=device
-            )
-            p_partial[batch_size // 2 :] = 1.0  # No filtering for second half
-            configs.append(
-                BenchmarkConfig(
-                    name=f"topp_partial_b{batch_size}_v{vocab_size // 1000}k",
-                    batch_size=batch_size,
-                    vocab_size=vocab_size,
-                    k_values=None,
-                    p_values=p_partial,
-                    description=f"Top-p only (partial batch, 50% p=0.9, 50% p=1.0), "
-                    f"batch={batch_size}, vocab={vocab_size}",
-                    ops_pct=calculate_ops_pct(None, p_partial, vocab_size, batch_size),
-                )
-            )
-
-            # 5. Mix of top-k and top-p (both applied to whole batch)
-            k_mix = torch.full((batch_size,), 100, dtype=torch.int32, device=device)
-            p_mix = torch.full((batch_size,), 0.9, dtype=torch.float32, device=device)
-            configs.append(
-                BenchmarkConfig(
-                    name=f"topk_topp_whole_b{batch_size}_v{vocab_size // 1000}k",
-                    batch_size=batch_size,
-                    vocab_size=vocab_size,
-                    k_values=k_mix,
-                    p_values=p_mix,
-                    description=f"Top-k + Top-p (whole batch, k=100, p=0.9), "
-                    f"batch={batch_size}, vocab={vocab_size}",
-                    ops_pct=calculate_ops_pct(k_mix, p_mix, vocab_size, batch_size),
-                )
-            )
-
-            # 6. Mix with partial application (some rows k only, some p only, some both)
-            k_mixed = torch.full(
-                (batch_size,), vocab_size, dtype=torch.int32, device=device
-            )
-            p_mixed = torch.full((batch_size,), 1.0, dtype=torch.float32, device=device)
-            # First third: k only
-            third = batch_size // 3
-            k_mixed[:third] = 50
-            # Second third: p only
-            p_mixed[third : 2 * third] = 0.5
-            # Last third: both k and p
-            k_mixed[2 * third :] = 100
-            p_mixed[2 * third :] = 0.9
-            configs.append(
-                BenchmarkConfig(
-                    name=f"mixed_partial_b{batch_size}_v{vocab_size // 1000}k",
-                    batch_size=batch_size,
-                    vocab_size=vocab_size,
-                    k_values=k_mixed,
-                    p_values=p_mixed,
-                    description=f"Mixed partial (1/3 k=50, 1/3 p=0.9, 1/3 both), "
-                    f"batch={batch_size}, vocab={vocab_size}",
-                    ops_pct=calculate_ops_pct(k_mixed, p_mixed, vocab_size, batch_size),
-                )
-            )
-
-    return configs
-
-
-def format_memory(bytes_val: int) -> str:
-    """Format memory in human-readable form."""
-    if bytes_val >= 1024**3:
-        return f"{bytes_val / (1024**3):.2f} GB"
-    elif bytes_val >= 1024**2:
-        return f"{bytes_val / (1024**2):.2f} MB"
-    elif bytes_val >= 1024:
-        return f"{bytes_val / 1024:.2f} KB"
-    return f"{bytes_val} B"
-
-
-def run_benchmark(
-    configs: list[BenchmarkConfig],
-    warmup_iters: int = 5,
-    benchmark_iters: int = 20,
-    verbose: bool = True,
-):
-    """Run all benchmarks and print results."""
-    results = []
-
-    print("=" * 100)
-    print("Top-k/Top-p Benchmark: Triton vs PyTorch Sort-based")
-    print("=" * 100)
-    print()
-
-    for config in configs:
-        if verbose:
-            print(f"Running: {config.description}")
-
-        # Create fresh logits for this config
-        logits = create_logits(config.batch_size, config.vocab_size)
-
-        # Benchmark Triton
-        reset_memory_stats()
-        triton_time, triton_mem = benchmark_function(
-            apply_top_k_top_p_triton,
-            logits,
-            config.k_values,
-            config.p_values,
-            warmup_iters,
-            benchmark_iters,
-        )
-
-        # Benchmark PyTorch
-        reset_memory_stats()
-        pytorch_time, pytorch_mem = benchmark_function(
-            apply_top_k_top_p_pytorch,
-            logits,
-            config.k_values,
-            config.p_values,
-            warmup_iters,
-            benchmark_iters,
-        )
-
-        speedup = pytorch_time / triton_time if triton_time > 0 else float("inf")
-        mem_ratio = pytorch_mem / triton_mem if triton_mem > 0 else float("inf")
-
-        result = {
-            "config": config,
-            "triton_time_ms": triton_time,
-            "pytorch_time_ms": pytorch_time,
-            "triton_mem": triton_mem,
-            "pytorch_mem": pytorch_mem,
-            "speedup": speedup,
-            "mem_ratio": mem_ratio,
-        }
-        results.append(result)
-
-        if verbose:
-            print(f"  Triton:  {triton_time:.3f} ms, {format_memory(triton_mem)}")
-            print(f"  PyTorch: {pytorch_time:.3f} ms, {format_memory(pytorch_mem)}")
-            print(f"  Speedup: {speedup:.2f}x, Memory ratio: {mem_ratio:.2f}x")
-            print()
-
-        # Clean up
-        del logits
-        reset_memory_stats()
-
-    return results
-
-
-def print_summary_table(results: list[dict]):
-    """Print a summary table of results."""
-    print()
-    print("=" * 130)
-    print("SUMMARY TABLE")
-    print("=" * 130)
-    print()
-
-    # Header
-    header = (
-        f"{'Scenario':<40} {'Batch':>6} {'Vocab':>7} {'Ops%':>6} "
-        f"{'Triton (ms)':>12} {'PyTorch (ms)':>13} {'Speedup':>8} "
-        f"{'Tri Mem':>10} {'Pyt Mem':>10}"
-    )
-    print(header)
-    print("-" * 130)
-
-    # Group by scenario type
-    current_vocab = None
-    for result in results:
-        config = result["config"]
-
-        # Add separator between vocab sizes
-        if current_vocab != config.vocab_size:
-            if current_vocab is not None:
-                print("-" * 130)
-            current_vocab = config.vocab_size
-
-        scenario = config.name.split("_b")[0]  # Extract scenario name
-        print(
-            f"{scenario:<40} {config.batch_size:>6} {config.vocab_size:>7} "
-            f"{config.ops_pct:>5.0f}% "
-            f"{result['triton_time_ms']:>12.3f} {result['pytorch_time_ms']:>13.3f} "
-            f"{result['speedup']:>7.2f}x "
-            f"{format_memory(result['triton_mem']):>10} "
-            f"{format_memory(result['pytorch_mem']):>10}"
-        )
-
-    print("=" * 130)
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Benchmark Triton vs PyTorch sort-based top-k/top-p implementations"
-    )
-    parser.add_argument(
-        "--batch-sizes",
-        type=int,
-        nargs="+",
-        default=[1, 4, 16, 64, 128, 512, 1024, 2048],
-        help="Batch sizes to test (default: 1 4 16 64)",
-    )
-    parser.add_argument(
-        "--vocab-sizes",
-        type=int,
-        nargs="+",
-        default=[32768, 131072],  # 32k, 128k
-        help="Vocabulary sizes to test (default: 32768 131072)",
-    )
-    parser.add_argument(
-        "--warmup-iters",
-        type=int,
-        default=5,
-        help="Number of warmup iterations (default: 5)",
-    )
-    parser.add_argument(
-        "--benchmark-iters",
-        type=int,
-        default=20,
-        help="Number of benchmark iterations (default: 20)",
-    )
-    parser.add_argument(
-        "--quiet",
-        action="store_true",
-        help="Only print summary table",
-    )
-
-    args = parser.parse_args()
-
-    # Print configuration
-    print(f"Batch sizes: {args.batch_sizes}")
-    print(f"Vocab sizes: {args.vocab_sizes}")
-    print(f"Warmup iterations: {args.warmup_iters}")
-    print(f"Benchmark iterations: {args.benchmark_iters}")
-    print()
-
-    # Check CUDA
-    if not torch.cuda.is_available():
-        print("ERROR: CUDA is not available. This benchmark requires a GPU.")
-        return
-
-    device_name = torch.cuda.get_device_name(0)
-    print(f"GPU: {device_name}")
-    print()
-
-    # Create configs
-    configs = create_benchmark_configs(
-        args.batch_sizes,
-        args.vocab_sizes,
-    )
-
-    # Run benchmarks
-    results = run_benchmark(
-        configs,
-        warmup_iters=args.warmup_iters,
-        benchmark_iters=args.benchmark_iters,
-        verbose=not args.quiet,
-    )
-
-    # Print summary
-    print_summary_table(results)
-
-
-if __name__ == "__main__":
-    main()

benchmarks/benchmark_utils.py

@@ -1,7 +1,78 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import argparse
+import json
+import math
+import os
 import time
 from types import TracebackType
+from typing import Any
+
+
+def convert_to_pytorch_benchmark_format(
+    args: argparse.Namespace, metrics: dict[str, list], extra_info: dict[str, Any]
+) -> list:
+    """
+    Save the benchmark results in the format used by PyTorch OSS benchmark with
+    on metric per record
+    https://github.com/pytorch/pytorch/wiki/How-to-integrate-with-PyTorch-OSS-benchmark-database
+    """
+    records = []
+    if not os.environ.get("SAVE_TO_PYTORCH_BENCHMARK_FORMAT", False):
+        return records
+
+    for name, benchmark_values in metrics.items():
+        record = {
+            "benchmark": {
+                "name": "vLLM benchmark",
+                "extra_info": {
+                    "args": vars(args),
+                },
+            },
+            "model": {
+                "name": args.model,
+            },
+            "metric": {
+                "name": name,
+                "benchmark_values": benchmark_values,
+                "extra_info": extra_info,
+            },
+        }
+
+        tp = record["benchmark"]["extra_info"]["args"].get("tensor_parallel_size")
+        # Save tensor_parallel_size parameter if it's part of the metadata
+        if not tp and "tensor_parallel_size" in extra_info:
+            record["benchmark"]["extra_info"]["args"]["tensor_parallel_size"] = (
+                extra_info["tensor_parallel_size"]
+            )
+
+        records.append(record)
+
+    return records
+
+
+class InfEncoder(json.JSONEncoder):
+    def clear_inf(self, o: Any):
+        if isinstance(o, dict):
+            return {k: self.clear_inf(v) for k, v in o.items()}
+        elif isinstance(o, list):
+            return [self.clear_inf(v) for v in o]
+        elif isinstance(o, float) and math.isinf(o):
+            return "inf"
+        return o
+
+    def iterencode(self, o: Any, *args, **kwargs) -> Any:
+        return super().iterencode(self.clear_inf(o), *args, **kwargs)
+
+
+def write_to_json(filename: str, records: list) -> None:
+    with open(filename, "w") as f:
+        json.dump(
+            records,
+            f,
+            cls=InfEncoder,
+            default=lambda o: f"<{type(o).__name__} object is not JSON serializable>",
+        )
 
 
 # Collect time and generate time metrics

benchmarks/cutlass_benchmarks/utils.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 # Cutlass bench utils
+from collections.abc import Iterable
 
 import torch
 
@@ -85,3 +86,15 @@ def make_rand_sparse_tensors(
 
     # Compressed B, Metadata, Original A, B
     return b_compressed, e, a, b
+
+
+def make_n_rand_sparse_tensors(
+    num_tensors: int, dtype: torch.dtype, m: int, n: int, k: int
+) -> tuple[Iterable[torch.Tensor], Iterable[torch.Tensor]]:
+    ABs = []
+    for _ in range(num_tensors):
+        b_comp, e, a, b = make_rand_sparse_tensors(dtype, m, n, k)
+        if b_comp is not None:
+            ABs.append(make_rand_sparse_tensors(dtype, m, n, k))
+    BComps, Es, As, Bs = zip(*ABs)
+    return list(BComps), list(Es), list(As), list(Bs)

benchmarks/disagg_benchmarks/rate_limiter.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+import time
+
+
+class RateLimiter:
+    """Token bucket rate limiter implementation"""
+
+    def __init__(self, rate_limit):
+        self.rate_limit = rate_limit  # Requests per second
+        self.num_available_tokens = rate_limit  # Available tokens
+        self.last_refill = time.monotonic()  # Last token refill time
+        self.lock = asyncio.Lock()  # Synchronization lock
+
+    async def acquire(self):
+        """Acquire a token from the rate limiter"""
+        while True:
+            async with self.lock:
+                current_time = time.monotonic()
+                elapsed = current_time - self.last_refill
+
+                # Refill num_available_tokens if more than 1 second has passed
+                if elapsed > 1.0:
+                    self.num_available_tokens = self.rate_limit
+                    self.last_refill = current_time
+
+                # Check if num_available_tokens are available
+                if self.num_available_tokens > 0:
+                    self.num_available_tokens -= 1
+                    return True
+
+                # Calculate wait time if no num_available_tokens available
+                wait_time = 1.0 - elapsed
+            await asyncio.sleep(wait_time)
+
+    async def __aenter__(self):
+        """Enter async context manager - acquire token"""
+        await self.acquire()
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        """Exit async context manager - no cleanup needed"""
+        pass

benchmarks/disagg_benchmarks/request_queue.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+from collections import deque
+
+
+class RequestQueue:
+    """Request queue manager with concurrency control"""
+
+    def __init__(self, max_concurrent, max_queue_size):
+        # Maximum concurrent requests
+        self.max_concurrent = max_concurrent
+        self.max_queue_size = max_queue_size  # Maximum queue size
+        # Concurrency control
+        self.semaphore = asyncio.Semaphore(max_concurrent)
+        self.queue = deque()  # Request queue
+        self.queue_size = 0  # Current queue size
+        self.lock = asyncio.Lock()  # Sync queue Lock
+
+    async def enqueue(self, task):
+        """Add a request task to the queue"""
+        async with self.lock:
+            if self.queue_size >= self.max_queue_size:
+                return False
+
+            self.queue.append(task)
+            self.queue_size += 1
+            return True
+
+    async def process(self):
+        """Process queued requests using semaphore for concurrency control"""
+        while True:
+            if self.queue:
+                async with self.semaphore, self.lock:
+                    task = self.queue.popleft()
+                    self.queue_size -= 1
+                    await task
+            await asyncio.sleep(0.01)  # Yield control to event loop

benchmarks/fused_kernels/layernorm_rms_benchmarks.py

@@ -13,7 +13,6 @@ from torch.utils.benchmark import Measurement as TMeasurement
 from tqdm import tqdm
 
 import vllm._custom_ops as ops
-from vllm.benchmarks.lib.utils import default_vllm_config
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
     per_token_group_quant_fp8,
@@ -292,7 +291,6 @@ def print_timers(timers: Iterable[TMeasurement]):
     compare.print()
 
 
-@default_vllm_config()
 def main():
     torch.set_default_device("cuda")
     bench_params = get_bench_params()

benchmarks/kernels/bench_block_fp8_gemm.py

@@ -8,7 +8,6 @@ os.environ["VLLM_USE_DEEP_GEMM"] = "0"
 
 import torch
 
-from vllm.benchmarks.lib.utils import default_vllm_config
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
     W8A8BlockFp8LinearOp,
 )
@@ -41,7 +40,6 @@ DEEPSEEK_V3_SHAPES = [
 ]
 
 
-@default_vllm_config()
 def build_w8a8_block_fp8_runner(M, N, K, block_size, device, use_cutlass):
     """Build runner function for w8a8 block fp8 matmul."""
     factor_for_scale = 1e-2

benchmarks/kernels/bench_per_token_quant_fp8.py

@@ -7,7 +7,6 @@ from unittest.mock import patch
 import pandas as pd
 import torch
 
-from vllm.benchmarks.lib.utils import default_vllm_config
 from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
 from vllm.triton_utils import triton
@@ -85,7 +84,6 @@ def calculate_diff(
 configs = []
 
 
-@default_vllm_config()
 def benchmark_quantization(
     batch_size,
     hidden_size,

benchmarks/kernels/benchmark_activation.py

@@ -7,7 +7,6 @@ import itertools
 import torch
 
 import vllm.model_executor.layers.activation  # noqa F401
-from vllm.benchmarks.lib.utils import default_vllm_config
 from vllm.model_executor.custom_op import op_registry
 from vllm.triton_utils import triton
 from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -19,7 +18,6 @@ intermediate_size = [3072, 9728, 12288]
 configs = list(itertools.product(batch_size_range, seq_len_range, intermediate_size))
 
 
-@default_vllm_config()
 def benchmark_activation(
     batch_size: int,
     seq_len: int,

benchmarks/kernels/benchmark_device_communicators.py

@@ -30,9 +30,6 @@ import torch.distributed as dist
 from torch.distributed import ProcessGroup
 
 from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
-from vllm.distributed.device_communicators.flashinfer_all_reduce import (
-    FlashInferAllReduce,
-)
 from vllm.distributed.device_communicators.pynccl import (
     PyNcclCommunicator,
     register_nccl_symmetric_ops,
@@ -47,7 +44,7 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser
 logger = init_logger(__name__)
 
 # Default sequence lengths to benchmark
-DEFAULT_SEQUENCE_LENGTHS = [16, 64, 128, 512, 1024, 2048, 4096, 8192]
+DEFAULT_SEQUENCE_LENGTHS = [128, 512, 1024, 2048, 4096, 8192]
 
 # Fixed hidden size and dtype for all benchmarks
 HIDDEN_SIZE = 8192
@@ -84,7 +81,6 @@ class CommunicatorBenchmark:
         self.symm_mem_comm = None
         self.symm_mem_comm_multimem = None
         self.symm_mem_comm_two_shot = None
-        self.fi_ar_comm = None
 
         self._init_communicators()
 
@@ -165,22 +161,6 @@ class CommunicatorBenchmark:
             )
             self.symm_mem_comm_two_shot = None
 
-        try:
-            self.fi_ar_comm = FlashInferAllReduce(
-                group=self.cpu_group,
-                device=self.device,
-            )
-            if not self.fi_ar_comm.disabled:
-                logger.info("Rank %s: FlashInferAllReduce initialized", self.rank)
-            else:
-                logger.info("Rank %s: FlashInferAllReduce disabled", self.rank)
-                self.fi_ar_comm = None
-        except Exception as e:
-            logger.warning(
-                "Rank %s: Failed to initialize FlashInferAllReduce: %s", self.rank, e
-            )
-            self.fi_ar_comm = None
-
     def benchmark_allreduce(
         self, sequence_length: int, num_warmup: int, num_trials: int
     ) -> dict[str, float]:
@@ -200,8 +180,7 @@ class CommunicatorBenchmark:
                     lambda t, c=comm: c.custom_all_reduce(t),
                     lambda t, c=comm: c.should_custom_ar(t),
                     comm.capture(),
-                    {"VLLM_CUSTOM_ALLREDUCE_ALGO": "1stage"},
-                    None,  # no destroy function
+                    "1stage",  # env variable value
                 )
             )
             # CustomAllreduce two-shot
@@ -211,8 +190,7 @@ class CommunicatorBenchmark:
                     lambda t, c=comm: c.custom_all_reduce(t),
                     lambda t, c=comm: c.should_custom_ar(t),
                     comm.capture(),
-                    {"VLLM_CUSTOM_ALLREDUCE_ALGO": "2stage"},
-                    None,  # no destroy function
+                    "2stage",  # env variable value
                 )
             )
 
@@ -224,8 +202,7 @@ class CommunicatorBenchmark:
                     lambda t, c=comm: c.all_reduce(t),
                     lambda t: True,  # Always available if initialized
                     nullcontext(),
-                    {},  # no env variable needed
-                    None,  # no destroy function
+                    None,  # no env variable needed
                 )
             )
             communicators.append(
@@ -234,8 +211,7 @@ class CommunicatorBenchmark:
                     lambda t: torch.ops.vllm.all_reduce_symmetric_with_copy(t),
                     lambda t: True,  # Always available if initialized
                     nullcontext(),
-                    {},  # no env variable needed
-                    None,  # no destroy function
+                    None,  # no env variable needed
                 )
             )
 
@@ -247,8 +223,7 @@ class CommunicatorBenchmark:
                     lambda t, c=comm: c.all_reduce(t),
                     lambda t, c=comm: c.should_use_symm_mem(t),
                     nullcontext(),
-                    {},  # no env variable needed
-                    None,  # no destroy function
+                    None,  # no env variable needed
                 )
             )
 
@@ -260,67 +235,29 @@ class CommunicatorBenchmark:
                     lambda t, c=comm: c.all_reduce(t),
                     lambda t, c=comm: c.should_use_symm_mem(t),
                     nullcontext(),
-                    {},  # no env variable needed
-                    None,  # no destroy function needed
-                )
-            )
-
-        if self.fi_ar_comm is not None:
-            comm = self.fi_ar_comm
-            communicators.append(
-                (
-                    "flashinfer_trtllm",
-                    lambda t, c=comm: c.all_reduce(t),
-                    lambda t, c=comm: c.should_use_fi_ar(t),
-                    nullcontext(),
-                    {"VLLM_FLASHINFER_ALLREDUCE_BACKEND": "trtllm"},
-                    lambda c=comm: c.destroy(),
-                )
-            )
-            communicators.append(
-                (
-                    "flashinfer_mnnvl",
-                    lambda t, c=comm: c.all_reduce(t),
-                    lambda t, c=comm: c.should_use_fi_ar(t),
-                    nullcontext(),
-                    {"VLLM_FLASHINFER_ALLREDUCE_BACKEND": "mnnvl"},
-                    lambda c=comm: c.destroy(),
+                    None,  # no env variable needed
                 )
             )
 
         # Benchmark each communicator
-        for (
-            name,
-            allreduce_fn,
-            should_use_fn,
-            context,
-            env_dict,
-            destroy_fn,
-        ) in communicators:
-            # Save original values and apply new environment variables
-            saved_env = {key: os.environ.get(key) for key in env_dict}
-            for key, value in env_dict.items():
-                os.environ[key] = value
-            try:
-                latency = self.benchmark_allreduce_single(
-                    sequence_length,
-                    allreduce_fn,
-                    should_use_fn,
-                    context,
-                    num_warmup,
-                    num_trials,
-                )
-                if latency is not None:
-                    results[name] = latency
-            finally:
-                if destroy_fn is not None:
-                    destroy_fn()
-                # Restore environment variables to their original state
-                for key, original_value in saved_env.items():
-                    if original_value is None:
-                        os.environ.pop(key, None)
-                    else:
-                        os.environ[key] = original_value
+        for name, allreduce_fn, should_use_fn, context, env_var in communicators:
+            # Set environment variable if needed
+            if env_var is not None:
+                os.environ["VLLM_CUSTOM_ALLREDUCE_ALGO"] = env_var
+            else:
+                # Clear the environment variable to avoid interference
+                os.environ.pop("VLLM_CUSTOM_ALLREDUCE_ALGO", None)
+
+            latency = self.benchmark_allreduce_single(
+                sequence_length,
+                allreduce_fn,
+                should_use_fn,
+                context,
+                num_warmup,
+                num_trials,
+            )
+            if latency is not None:
+                results[name] = latency
 
         return results
 

benchmarks/kernels/benchmark_fused_collective.py

@@ -5,11 +5,8 @@
 Benchmark for FlashInfer fused collective operations vs standard operations.
 
 This benchmark compares:
-1. FlashInfer's allreduce_fusion with trtllm backend
-   (fused allreduce + rmsnorm + optional FP8/FP4 quant)
-2. FlashInfer's allreduce_fusion with mnnvl backend
-   (fused allreduce + rmsnorm only, no quantization support)
-3. Standard tensor_model_parallel_all_reduce + separate rmsnorm/quant operations
+1. FlashInfer's allreduce_fusion (fused allreduce + rmsnorm + optional quant)
+2. Standard tensor_model_parallel_all_reduce + separate rmsnorm/quant operations
 
 Usage with torchrun:
     torchrun --nproc_per_node=2 benchmark_fused_collective.py
@@ -51,12 +48,8 @@ SCALED_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant
 logger = init_logger(__name__)
 
 # Try to import FlashInfer
-TorchDistBackend = None
 try:
     import flashinfer.comm as flashinfer_comm  # type: ignore
-    from flashinfer.comm.mnnvl import (  # type: ignore
-        TorchDistBackend,
-    )
 
     if not (
         hasattr(flashinfer_comm, "allreduce_fusion")
@@ -81,15 +74,11 @@ _FI_MAX_SIZES = {
     8: 64 * MiB,  # 64MB
 }
 
-# Global workspace tensors for FlashInfer (keyed by backend name)
-_FI_WORKSPACES: dict = {}
-
-# Backends to benchmark
-FLASHINFER_BACKENDS = ["trtllm", "mnnvl"]
+# Global workspace tensor for FlashInfer
+_FI_WORKSPACE = None
 
 
 def setup_flashinfer_workspace(
-    backend: str,
     world_size: int,
     rank: int,
     hidden_dim: int,
@@ -97,54 +86,41 @@ def setup_flashinfer_workspace(
     dtype: torch.dtype,
 ):
     """Setup FlashInfer workspace for fused allreduce operations."""
-    global FI_WORKSPACES
+    global _FI_WORKSPACE
 
     if flashinfer_comm is None:
-        return None
+        return None, None
 
     if world_size not in _FI_MAX_SIZES:
         logger.warning("FlashInfer not supported for world size %s", world_size)
-        return None
+        return None, None
 
     try:
-        kwargs = {}
-        if TorchDistBackend is not None:
-            kwargs["comm_backend"] = TorchDistBackend(group=dist.group.WORLD)
-
         workspace = flashinfer_comm.create_allreduce_fusion_workspace(
-            backend=backend,
+            backend="trtllm",
             world_size=world_size,
             rank=rank,
             max_token_num=max_token_num,
             hidden_dim=hidden_dim,
             dtype=dtype,
-            **kwargs,
         )
 
-        _FI_WORKSPACES[backend] = workspace
+        _FI_WORKSPACE = workspace
         return workspace
     except Exception as e:
-        logger.error(
-            "Failed to setup FlashInfer workspace (backend=%s): %s", backend, e
-        )
+        logger.error("Failed to setup FlashInfer workspace: %s", e)
         return None
 
 
-def cleanup_flashinfer_workspaces():
-    """Cleanup all FlashInfer workspaces."""
-    if flashinfer_comm is None:
+def cleanup_flashinfer_workspace(workspace):
+    """Cleanup FlashInfer workspace."""
+    if flashinfer_comm is None or workspace is None:
         return
 
-    for backend, workspace in _FI_WORKSPACES.items():
-        try:
-            workspace.destroy()
-        except Exception as e:
-            logger.error(
-                "Failed to cleanup FlashInfer workspace (backend=%s): %s",
-                backend,
-                e,
-            )
-    _FI_WORKSPACES.clear()
+    try:
+        workspace.destroy()
+    except Exception as e:
+        logger.error("Failed to cleanup FlashInfer workspace: %s", e)
 
 
 class FlashInferFusedAllReduceParams:
@@ -158,7 +134,7 @@ class FlashInferFusedAllReduceParams:
         self.fp32_acc = True
         self.max_token_num = max_token_num
 
-    def get_flashinfer_fused_allreduce_kwargs(self):
+    def get_trtllm_fused_allreduce_kwargs(self):
         return {
             "launch_with_pdl": self.launch_with_pdl,
             "fp32_acc": self.fp32_acc,
@@ -171,12 +147,11 @@ def flashinfer_fused_allreduce_rmsnorm(
     rms_gamma: torch.Tensor,
     rms_eps: float,
     allreduce_params: "FlashInferFusedAllReduceParams",
-    workspace: object,
     use_oneshot: bool,
     norm_out: torch.Tensor | None = None,
 ):
     """FlashInfer fused allreduce + rmsnorm operation."""
-    if flashinfer_comm is None or workspace is None:
+    if flashinfer_comm is None or _FI_WORKSPACE is None:
         raise RuntimeError("FlashInfer not available or workspace not initialized")
 
     if norm_out is None:
@@ -185,13 +160,9 @@ def flashinfer_fused_allreduce_rmsnorm(
     else:
         residual_out = input_tensor
 
-    layout_code = None
-    if workspace.backend == "trtllm":
-        layout_code = flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4
-
     flashinfer_comm.allreduce_fusion(
         input=input_tensor,
-        workspace=workspace,
+        workspace=_FI_WORKSPACE,
         pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
         residual_in=residual,
         residual_out=residual_out,
@@ -200,10 +171,10 @@ def flashinfer_fused_allreduce_rmsnorm(
         rms_eps=rms_eps,
         quant_out=None,
         scale_out=None,
-        layout_code=layout_code,
+        layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
         scale_factor=None,
         use_oneshot=use_oneshot,
-        **allreduce_params.get_flashinfer_fused_allreduce_kwargs(),
+        **allreduce_params.get_trtllm_fused_allreduce_kwargs(),
     )
 
 
@@ -214,16 +185,12 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
     rms_eps: float,
     scale_factor: torch.Tensor,
     allreduce_params: FlashInferFusedAllReduceParams,
-    workspace: object,
     use_oneshot: bool = True,
     norm_out: torch.Tensor | None = None,
     quant_out: torch.Tensor | None = None,
 ):
-    """FlashInfer fused allreduce + rmsnorm + FP8 quantization.
-
-    Note: Only supported by the trtllm backend.
-    """
-    if flashinfer_comm is None or workspace is None:
+    """FlashInfer fused allreduce + rmsnorm + FP8 quantization."""
+    if flashinfer_comm is None or _FI_WORKSPACE is None:
         raise RuntimeError("FlashInfer not available or workspace not initialized")
 
     if norm_out is None:
@@ -234,7 +201,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
 
     flashinfer_comm.allreduce_fusion(
         input=input_tensor,
-        workspace=workspace,
+        workspace=_FI_WORKSPACE,
         pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant,
         residual_in=residual,
         residual_out=residual_out,
@@ -246,7 +213,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
         layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
         scale_factor=scale_factor,
         use_oneshot=use_oneshot,
-        **allreduce_params.get_flashinfer_fused_allreduce_kwargs(),
+        **allreduce_params.get_trtllm_fused_allreduce_kwargs(),
     )
 
 
@@ -257,17 +224,13 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
     rms_eps: float,
     input_global_scale: torch.Tensor,
     allreduce_params: FlashInferFusedAllReduceParams,
-    workspace: object,
     quant_out: torch.Tensor,
     use_oneshot: bool,
     output_scale: torch.Tensor,
     norm_out: torch.Tensor | None = None,
 ):
-    """FlashInfer fused allreduce + rmsnorm + FP4 quantization.
-
-    Note: Only supported by the trtllm backend.
-    """
-    if flashinfer_comm is None or workspace is None:
+    """FlashInfer fused allreduce + rmsnorm + FP4 quantization."""
+    if flashinfer_comm is None or _FI_WORKSPACE is None:
         raise RuntimeError("FlashInfer not available or workspace not initialized")
 
     if norm_out is None:
@@ -278,7 +241,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
 
     flashinfer_comm.allreduce_fusion(
         input=input_tensor,
-        workspace=workspace,
+        workspace=_FI_WORKSPACE,
         pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant,
         residual_in=residual,
         residual_out=residual_out,
@@ -290,7 +253,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
         layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
         scale_factor=input_global_scale,
         use_oneshot=use_oneshot,
-        **allreduce_params.get_flashinfer_fused_allreduce_kwargs(),
+        **allreduce_params.get_trtllm_fused_allreduce_kwargs(),
     )
 
 
@@ -423,16 +386,13 @@ def run_benchmarks(
     dtype: torch.dtype,
     use_residual: bool,
     allreduce_params: FlashInferFusedAllReduceParams | None,
-    workspaces: dict,
     quant_modes: set[str],
     no_oneshot: bool,
 ):
     """Run all benchmarks for given configuration.
 
     Args:
-        allreduce_params: Shared parameters for FlashInfer fused allreduce.
-        workspaces: Dict mapping backend name ("trtllm", "mnnvl") to workspace.
-        quant_modes: Set of quantization modes: "none", "fp8", "fp4".
+        quant_mode: "none", "fp8_only", "fp4_only", or "all"
     """
     (
         input_tensor,
@@ -448,18 +408,18 @@ def run_benchmarks(
 
     rms_eps = 1e-6
     results = {}
+    vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
     use_oneshot_options = [False] if no_oneshot else [True, False]
 
+    # Create RMSNorm and QuantFP8 layers once for native benchmarks
+
     if "none" in quant_modes:
         # Standard AllReduce + RMSNorm
-        # Re-create VllmFusedAllreduce per config so CustomOp binds the
-        # correct forward method (native vs custom kernel).
         for custom_op in ["-rms_norm", "+rms_norm"]:
             with set_current_vllm_config(
                 VllmConfig(compilation_config=CompilationConfig(custom_ops=[custom_op]))
             ):
                 try:
-                    vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
                     suffix = (
                         "_custom_rms_norm" if "+" in custom_op else "_native_rms_norm"
                     )
@@ -478,7 +438,6 @@ def run_benchmarks(
             VllmConfig(compilation_config=CompilationConfig(custom_ops=["-rms_norm"]))
         ):
             try:
-                vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
                 standard_allreduce_rmsnorm_native_compiled = torch.compile(
                     vllm_fused_allreduce.allreduce_rmsnorm,
                     fullgraph=True,
@@ -494,11 +453,10 @@ def run_benchmarks(
                 logger.error("Standard AllReduce+RMSNorm Native Compiled failed: %s", e)
                 results["standard_allreduce_rmsnorm_native_compiled"] = float("inf")
 
-        # FlashInfer Fused AllReduce + RMSNorm (all backends)
-        for backend, workspace in workspaces.items():
+        # FlashInfer Fused AllReduce + RMSNorm Oneshot/Twoshot
+        if flashinfer_comm is not None and allreduce_params is not None:
             for use_oneshot in use_oneshot_options:
                 suffix = "_oneshot" if use_oneshot else "_twoshot"
-                key = f"flashinfer_{backend}_fused_allreduce_rmsnorm{suffix}"
                 try:
                     time_ms = benchmark_operation(
                         flashinfer_fused_allreduce_rmsnorm,
@@ -508,17 +466,14 @@ def run_benchmarks(
                         rms_gamma=rms_gamma,
                         rms_eps=rms_eps,
                         allreduce_params=allreduce_params,
-                        workspace=workspace,
                         use_oneshot=use_oneshot,
                     )
-                    results[key] = time_ms
+                    results[f"flashinfer_fused_allreduce_rmsnorm{suffix}"] = time_ms
                 except Exception as e:
-                    logger.error(
-                        "FlashInfer (%s) Fused AllReduce+RMSNorm failed: %s",
-                        backend,
-                        e,
+                    logger.error("FlashInfer Fused AllReduce+RMSNorm failed: %s", e)
+                    results[f"flashinfer_fused_allreduce_rmsnorm{suffix}"] = float(
+                        "inf"
                     )
-                    results[key] = float("inf")
 
     if "fp8" in quant_modes:
         # Standard AllReduce + RMSNorm + FP8 Quant
@@ -527,7 +482,7 @@ def run_benchmarks(
                 "_custom_rms_norm" if "+" in rms_norm_custom_op else "_native_rms_norm"
             )
             for quant_fp8_custom_op in ["-quant_fp8", "+quant_fp8"]:
-                op_suffix = suffix + (
+                suffix += (
                     "_custom_quant_fp8"
                     if "+" in quant_fp8_custom_op
                     else "_native_quant_fp8"
@@ -540,17 +495,16 @@ def run_benchmarks(
                     )
                 ):
                     try:
-                        vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
                         time_ms = benchmark_operation(
                             vllm_fused_allreduce.allreduce_rmsnorm_fp8_quant,
                             input_tensor,
                             residual=residual,
                             scale_factor=scale_fp8,
                         )
-                        results[f"standard_allreduce{op_suffix}"] = time_ms
+                        results[f"standard_allreduce{suffix}"] = time_ms
                     except Exception as e:
                         logger.error("Standard AllReduce+RMSNorm+FP8 failed: %s", e)
-                        results[f"standard_allreduce{op_suffix}"] = float("inf")
+                        results[f"standard_allreduce{suffix}"] = float("inf")
 
         # Standard AllReduce + RMSNorm + FP8 Quant Native Compiled
         with set_current_vllm_config(
@@ -561,7 +515,6 @@ def run_benchmarks(
             )
         ):
             try:
-                vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
                 standard_allreduce_rmsnorm_fp8_quant_native_compiled = torch.compile(
                     vllm_fused_allreduce.allreduce_rmsnorm_fp8_quant,
                     fullgraph=True,
@@ -584,12 +537,10 @@ def run_benchmarks(
                     "inf"
                 )
 
-        # FlashInfer Fused AllReduce + RMSNorm + FP8 Quant (trtllm only)
-        if "trtllm" in workspaces:
-            trtllm_ws = workspaces["trtllm"]
+        # FlashInfer Fused AllReduce + RMSNorm + FP8 Quant Oneshot
+        if flashinfer_comm is not None and allreduce_params is not None:
             for use_oneshot in use_oneshot_options:
                 suffix = "_oneshot" if use_oneshot else "_twoshot"
-                key = f"flashinfer_trtllm_fused_allreduce_rmsnorm_fp8_quant{suffix}"
                 try:
                     time_ms = benchmark_operation(
                         flashinfer_fused_allreduce_rmsnorm_fp8_quant,
@@ -601,16 +552,19 @@ def run_benchmarks(
                         scale_factor=scale_fp8,
                         quant_out=quant_out_fp8,
                         allreduce_params=allreduce_params,
-                        workspace=trtllm_ws,
                         use_oneshot=use_oneshot,
                     )
-                    results[key] = time_ms
+                    results[f"flashinfer_fused_allreduce_rmsnorm_fp8_quant{suffix}"] = (
+                        time_ms
+                    )
                 except Exception as e:
                     logger.error(
-                        "FlashInfer (trtllm) Fused AllReduce+RMSNorm+FP8 failed: %s",
+                        "FlashInfer Fused AllReduce+RMSNorm+FP8 Oneshot failed: %s",
                         e,
                     )
-                    results[key] = float("inf")
+                    results[f"flashinfer_fused_allreduce_rmsnorm_fp8_quant{suffix}"] = (
+                        float("inf")
+                    )
 
     if "fp4" in quant_modes and current_platform.has_device_capability(100):
         # Standard AllReduce + RMSNorm + FP4 Quant
@@ -626,7 +580,6 @@ def run_benchmarks(
                 )
             ):
                 try:
-                    vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
                     time_ms = benchmark_operation(
                         vllm_fused_allreduce.allreduce_rmsnorm_fp4_quant,
                         input_tensor,
@@ -645,7 +598,6 @@ def run_benchmarks(
             VllmConfig(compilation_config=CompilationConfig(custom_ops=["-rms_norm"]))
         ):
             try:
-                vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
                 standard_allreduce_rmsnorm_fp4_quant_native_compiled = torch.compile(
                     vllm_fused_allreduce.allreduce_rmsnorm_fp4_quant,
                     fullgraph=True,
@@ -670,12 +622,10 @@ def run_benchmarks(
                     "inf"
                 )
 
-        # FlashInfer Fused AllReduce + RMSNorm + FP4 Quant (trtllm only)
-        if "trtllm" in workspaces:
-            trtllm_ws = workspaces["trtllm"]
+        # FlashInfer Fused AllReduce + RMSNorm + FP4 Quant Oneshot
+        if flashinfer_comm is not None and allreduce_params is not None:
             for use_oneshot in use_oneshot_options:
                 suffix = "_oneshot" if use_oneshot else "_twoshot"
-                key = f"flashinfer_trtllm_fused_allreduce_rmsnorm_fp4_quant{suffix}"
                 try:
                     time_ms = benchmark_operation(
                         flashinfer_fused_allreduce_rmsnorm_fp4_quant,
@@ -686,18 +636,49 @@ def run_benchmarks(
                         rms_eps=rms_eps,
                         input_global_scale=scale_fp4,
                         allreduce_params=allreduce_params,
-                        workspace=trtllm_ws,
                         quant_out=fp4_quant_out,
                         output_scale=fp4_output_scale,
                         use_oneshot=use_oneshot,
                     )
-                    results[key] = time_ms
+                    results[f"flashinfer_fused_allreduce_rmsnorm_fp4_quant{suffix}"] = (
+                        time_ms
+                    )
                 except Exception as e:
                     logger.error(
-                        "FlashInfer (trtllm) Fused AllReduce+RMSNorm+FP4 failed: %s",
+                        "FlashInfer Fused AllReduce+RMSNorm+FP4 Oneshot failed: %s",
                         e,
                     )
-                    results[key] = float("inf")
+                    results[f"flashinfer_fused_allreduce_rmsnorm_fp4_quant{suffix}"] = (
+                        float("inf")
+                    )
+
+        # FlashInfer Fused AllReduce + RMSNorm + FP4 Quant Two-shot
+        if flashinfer_comm is not None and allreduce_params is not None:
+            try:
+                time_ms = benchmark_operation(
+                    flashinfer_fused_allreduce_rmsnorm_fp4_quant,
+                    input_tensor,
+                    residual=residual,
+                    norm_out=norm_out,
+                    rms_gamma=rms_gamma,
+                    rms_eps=rms_eps,
+                    input_global_scale=scale_fp4,
+                    allreduce_params=allreduce_params,
+                    quant_out=fp4_quant_out,
+                    output_scale=fp4_output_scale,
+                    use_oneshot=False,
+                )
+                results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_twoshot"] = (
+                    time_ms
+                )
+            except Exception as e:
+                logger.error(
+                    "FlashInfer Fused AllReduce+RMSNorm+FP4 Two-shot failed: %s",
+                    e,
+                )
+                results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_twoshot"] = float(
+                    "inf"
+                )
 
     return results
 
@@ -1035,7 +1016,8 @@ def main():
 
     configs = list(itertools.product(args.num_tokens, dtypes, residual_options))
 
-    # Setup FlashInfer workspaces for all backends
+    # Setup FlashInfer workspace if available
+    workspace = None
     allreduce_params = None
 
     if flashinfer_comm is not None:
@@ -1050,17 +1032,15 @@ def main():
             args.hidden_dim * max_element_size
         )
 
-        for backend in FLASHINFER_BACKENDS:
-            setup_flashinfer_workspace(
-                backend=backend,
-                world_size=world_size,
-                rank=rank,
-                hidden_dim=args.hidden_dim,
-                max_token_num=max_num_token,
-                dtype=workspace_dtype,
-            )
+        workspace = setup_flashinfer_workspace(
+            world_size,
+            rank,
+            args.hidden_dim,
+            max_num_token,
+            dtype=workspace_dtype,
+        )
 
-        if _FI_WORKSPACES:
+        if workspace is not None:
             allreduce_params = FlashInferFusedAllReduceParams(
                 max_token_num=max_num_token,
             )
@@ -1086,7 +1066,6 @@ def main():
                 dtype,
                 use_residual,
                 allreduce_params,
-                workspaces=_FI_WORKSPACES,
                 quant_modes=quant_modes,
                 no_oneshot=args.no_oneshot,
             )
@@ -1125,13 +1104,11 @@ def main():
 
     finally:
         # Cleanup
-        cleanup_flashinfer_workspaces()
+        if workspace is not None:
+            cleanup_flashinfer_workspace(workspace)
 
         dist.barrier()
 
 
 if __name__ == "__main__":
-    from vllm.config import VllmConfig, set_current_vllm_config
-
-    with set_current_vllm_config(VllmConfig()):
-        main()
+    main()

benchmarks/kernels/benchmark_layernorm.py

@@ -5,14 +5,12 @@ import time
 
 import torch
 
-from vllm.benchmarks.lib.utils import default_vllm_config
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
 
 
 @torch.inference_mode()
-@default_vllm_config()
 def main(
     num_tokens: int,
     hidden_size: int,

benchmarks/kernels/benchmark_moe_defaults.py

@@ -1,278 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark comparing old vs new default fused MoE configs.
-
-Runs the triton fused_moe kernel with three configurations for each scenario:
-  1. Tuned config (from JSON file, if available) — the target to match
-  2. Old default (the hardcoded defaults before this change)
-  3. New default (the improved defaults)
-
-Usage:
-    python benchmarks/kernels/benchmark_moe_defaults.py
-
-Produces a table showing kernel time (us) and speedup of new vs old defaults.
-"""
-
-import torch
-
-from vllm.model_executor.layers.fused_moe import fused_topk, override_config
-from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
-from vllm.model_executor.layers.fused_moe.fused_moe import (
-    fused_experts,
-    get_default_config,
-    get_moe_configs,
-)
-from vllm.platforms import current_platform
-from vllm.triton_utils import triton
-from vllm.utils.torch_utils import set_random_seed
-
-FP8_DTYPE = current_platform.fp8_dtype()
-
-
-def old_default_config(M, E, N, K, topk, dtype=None, block_shape=None):
-    """The original defaults before https://github.com/vllm-project/vllm/pull/34846,
-    for comparison."""
-    if dtype == "fp8_w8a8" and block_shape is not None:
-        return {
-            "BLOCK_SIZE_M": 64,
-            "BLOCK_SIZE_N": block_shape[0],
-            "BLOCK_SIZE_K": block_shape[1],
-            "GROUP_SIZE_M": 32,
-            "SPLIT_K": 1,
-            "num_warps": 4,
-            "num_stages": 3 if not current_platform.is_rocm() else 2,
-        }
-    elif M <= E:
-        return {
-            "BLOCK_SIZE_M": 16,
-            "BLOCK_SIZE_N": 32,
-            "BLOCK_SIZE_K": 64,
-            "GROUP_SIZE_M": 1,
-            "SPLIT_K": 1,
-        }
-    else:
-        return {
-            "BLOCK_SIZE_M": 64,
-            "BLOCK_SIZE_N": 64,
-            "BLOCK_SIZE_K": 32,
-            "GROUP_SIZE_M": 8,
-            "SPLIT_K": 1,
-        }
-
-
-def benchmark_config(
-    config,
-    M,
-    E,
-    N,
-    K,
-    topk,
-    dtype,
-    use_fp8=False,
-    block_shape=None,
-    num_iters=100,
-):
-    """Time a single kernel config. Returns kernel time in microseconds."""
-    init_dtype = torch.float16 if use_fp8 else dtype
-
-    a = torch.randn(M, K, device="cuda", dtype=init_dtype) / 10
-    w1 = torch.randn(E, 2 * N, K, device="cuda", dtype=init_dtype) / 10
-    w2 = torch.randn(E, K, N, device="cuda", dtype=init_dtype) / 10
-
-    w1_scale = None
-    w2_scale = None
-    a1_scale = None
-    a2_scale = None
-    if use_fp8:
-        if block_shape is not None:
-            bsn, bsk = block_shape
-            n_tiles_w1 = triton.cdiv(2 * N, bsn)
-            k_tiles_w1 = triton.cdiv(K, bsk)
-            n_tiles_w2 = triton.cdiv(K, bsn)
-            k_tiles_w2 = triton.cdiv(N, bsk)
-            w1_scale = torch.rand(
-                E, n_tiles_w1, k_tiles_w1, device="cuda", dtype=torch.float32
-            )
-            w2_scale = torch.rand(
-                E, n_tiles_w2, k_tiles_w2, device="cuda", dtype=torch.float32
-            )
-        else:
-            w1_scale = torch.rand(E, device="cuda", dtype=torch.float32)
-            w2_scale = torch.rand(E, device="cuda", dtype=torch.float32)
-        a1_scale = torch.rand(1, device="cuda", dtype=torch.float32)
-        a2_scale = torch.rand(1, device="cuda", dtype=torch.float32)
-        # Only weights are stored in fp8; activations stay in bf16/fp16
-        # and get dynamically quantized inside the kernel.
-        w1 = w1.to(FP8_DTYPE)
-        w2 = w2.to(FP8_DTYPE)
-
-    quant_config = FusedMoEQuantConfig.make(
-        quant_dtype=torch.float8_e4m3fn if use_fp8 else None,
-        w1_scale=w1_scale,
-        w2_scale=w2_scale,
-        a1_scale=a1_scale,
-        a2_scale=a2_scale,
-        block_shape=block_shape,
-    )
-
-    gating = torch.randn(M, E, device="cuda", dtype=torch.float32)
-
-    # Warmup
-    for _ in range(20):
-        with override_config(config):
-            topk_weights, topk_ids, _ = fused_topk(a, gating, topk, renormalize=True)
-            fused_experts(
-                a,
-                w1,
-                w2,
-                topk_weights,
-                topk_ids,
-                quant_config=quant_config,
-            )
-    torch.cuda.synchronize()
-
-    # Benchmark
-    start = torch.cuda.Event(enable_timing=True)
-    end = torch.cuda.Event(enable_timing=True)
-    start.record()
-    for _ in range(num_iters):
-        with override_config(config):
-            topk_weights, topk_ids, _ = fused_topk(a, gating, topk, renormalize=True)
-            fused_experts(
-                a,
-                w1,
-                w2,
-                topk_weights,
-                topk_ids,
-                quant_config=quant_config,
-            )
-    end.record()
-    torch.cuda.synchronize()
-    return start.elapsed_time(end) / num_iters * 1000  # ms -> us
-
-
-# Model configurations: (name, E, N, K, topk, dtype_str, use_fp8, block_shape)
-# N = moe_intermediate_size // tp_size (the value used in config file lookup)
-MODELS = [
-    # --- Few experts ---
-    ("Mixtral bf16", 8, 7168, 4096, 2, None, False, None),
-    ("Mixtral fp8", 8, 7168, 4096, 2, "fp8_w8a8", True, None),
-    # --- Many experts: real model shapes at tp=1 ---
-    # Qwen2-MoE-57B: E=60, topk=4, N=1408, K=2048
-    ("Qwen2-MoE bf16", 60, 1408, 2048, 4, None, False, None),
-    # DeepSeek-V2: E=64, topk=6, N=1407, K=4096
-    # (use 1408 to avoid odd alignment; real model is 1407)
-    ("DeepSeek-V2 bf16", 64, 1408, 4096, 6, None, False, None),
-    # OLMoE-7B: E=64, topk=8, N=2048, K=2048
-    ("OLMoE bf16", 64, 2048, 2048, 8, None, False, None),
-    # GLM-4-100B-A10B: E=128, topk=8, N=1408, K=4096
-    ("GLM-4-MoE bf16", 128, 1408, 4096, 8, None, False, None),
-    # Qwen3-30B-A3B: E=128, topk=8, N=768, K=2048
-    ("Qwen3-MoE bf16", 128, 768, 2048, 8, None, False, None),
-    # DeepSeek-V3 / MiMo-V2-Flash: E=256, topk=8, N=2048, K=7168
-    ("DeepSeek-V3 bf16", 256, 2048, 7168, 8, None, False, None),
-    # Qwen3.5-70B-A22B (Qwen3-Next): E=512, topk=10, N=512, K=2048
-    ("Qwen3-Next bf16", 512, 512, 2048, 10, None, False, None),
-    # E=128 N=1856 bf16
-    ("E128 N1856 bf16", 128, 1856, 4096, 8, None, False, None),
-    # E=256 N=512 bf16 (DS-V3 tp=4)
-    ("DS-V3 tp4 bf16", 256, 512, 7168, 8, None, False, None),
-    # E=512 N=512 bf16 (Qwen3-Next tp=1)
-    ("Qwen3-Next bf16", 512, 512, 2048, 10, None, False, None),
-    # E=512 N=256 bf16 (Qwen3-Next tp=2)
-    ("Qwen3-Next tp2", 512, 256, 2048, 10, None, False, None),
-    # --- FP8 block quant (many experts) ---
-    # DS-V3 tp=4: E=256, N=512, fp8 block
-    ("DS-V3 tp4 fp8blk", 256, 512, 7168, 8, "fp8_w8a8", True, [128, 128]),
-    # DS-V3 tp=8: E=256, N=256, fp8 block
-    ("DS-V3 tp8 fp8blk", 256, 256, 7168, 8, "fp8_w8a8", True, [128, 128]),
-    # Qwen3-Next tp=2 fp8 block
-    ("Qwen3-Next tp2 fp8blk", 512, 256, 2048, 10, "fp8_w8a8", True, [128, 128]),
-]
-
-BATCH_SIZES = [1, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096]
-
-
-def main():
-    set_random_seed(0)
-    torch.set_default_device("cuda")
-    dtype = torch.bfloat16
-
-    for name, E, N, K, topk, dtype_str, use_fp8, block_shape in MODELS:
-        print(f"\n{'=' * 90}")
-        print(f"  {name}  (E={E}, N={N}, K={K}, topk={topk})")
-        print(f"{'=' * 90}")
-
-        # Try to load tuned config
-        block_n = block_shape[0] if block_shape else None
-        block_k = block_shape[1] if block_shape else None
-        tuned = get_moe_configs(E, N, dtype_str, block_n, block_k)
-        has_tuned = tuned is not None
-        print(f"  Tuned config available: {has_tuned}")
-
-        hdr = (
-            f"{'Batch':>6} | {'Tuned (us)':>11} | {'Old (us)':>11} | "
-            f"{'New (us)':>11} | {'New/Old':>8} | {'New/Tuned':>10}"
-        )
-        print(f"  {hdr}")
-        print(f"  {'-' * len(hdr)}")
-
-        for M in BATCH_SIZES:
-            old_cfg = old_default_config(M, E, N, K, topk, dtype_str, block_shape)
-            new_cfg = get_default_config(M, E, N, K, topk, dtype_str, block_shape)
-
-            if has_tuned:
-                tuned_cfg = tuned[min(tuned.keys(), key=lambda x: abs(x - M))]
-                t_tuned = benchmark_config(
-                    tuned_cfg,
-                    M,
-                    E,
-                    N,
-                    K,
-                    topk,
-                    dtype,
-                    use_fp8=use_fp8,
-                    block_shape=block_shape,
-                )
-            else:
-                t_tuned = None
-
-            t_old = benchmark_config(
-                old_cfg,
-                M,
-                E,
-                N,
-                K,
-                topk,
-                dtype,
-                use_fp8=use_fp8,
-                block_shape=block_shape,
-            )
-            t_new = benchmark_config(
-                new_cfg,
-                M,
-                E,
-                N,
-                K,
-                topk,
-                dtype,
-                use_fp8=use_fp8,
-                block_shape=block_shape,
-            )
-
-            ratio_new_old = t_new / t_old
-            tuned_str = f"{t_tuned:11.2f}" if t_tuned else f"{'N/A':>11}"
-            ratio_tuned = f"{t_new / t_tuned:10.2f}x" if t_tuned else f"{'N/A':>10}"
-            # flag regressions where new default is >5% slower than old
-            marker = " <--" if ratio_new_old > 1.05 else ""
-
-            print(
-                f"  {M:>6} | {tuned_str} | {t_old:11.2f} | {t_new:11.2f} "
-                f"| {ratio_new_old:7.2f}x | {ratio_tuned}{marker}"
-            )
-
-
-if __name__ == "__main__":
-    main()

benchmarks/kernels/benchmark_mrope.py

@@ -36,7 +36,6 @@ from typing import Any
 import numpy as np
 import torch
 
-from vllm.benchmarks.lib.utils import default_vllm_config
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.transformers_utils.config import get_config
 from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -79,7 +78,6 @@ def calculate_stats(times: list[float]) -> dict[str, float]:
     }
 
 
-@default_vllm_config()
 def benchmark_mrope(
     model_name: str,
     num_tokens: int,

benchmarks/kernels/benchmark_rope.py

@@ -5,7 +5,6 @@ import itertools
 
 import torch
 
-from vllm.benchmarks.lib.utils import default_vllm_config
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.triton_utils import triton
 from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -30,7 +29,6 @@ def get_benchmark(head_size, rotary_dim, is_neox_style, device):
             args={},
         )
     )
-    @default_vllm_config()
     def benchmark(batch_size, seq_len, num_heads, provider):
         dtype = torch.bfloat16
         max_position = 8192

benchmarks/run_structured_output_benchmark.sh

@@ -71,7 +71,7 @@ while [[ $# -gt 0 ]]; do
       usage
       ;;
     *)
-      printf "Unknown argument: %s\n" "$1"
+      echo "Unknown argument: $1\n"
       usage
       ;;
   esac
@@ -84,17 +84,15 @@ mkdir -p "$OUTPUT_DIR"
 QPS_VALUES=(25 20 15 10 5 1)
 
 # Common parameters
-COMMON_PARAMS=(
-  --backend "$BACKEND"
-  --model "$MODEL"
-  --dataset "$DATASET"
-  --structured-output-ratio "$STRUCTURED_OUTPUT_RATIO"
-  --save-results
-  --result-dir "$OUTPUT_DIR"
-  --output-len "$MAX_NEW_TOKENS"
-  --port "$PORT"
-  --tokenizer-mode "$TOKENIZER_MODE"
-)
+COMMON_PARAMS="--backend $BACKEND \
+               --model $MODEL \
+               --dataset $DATASET \
+               --structured-output-ratio $STRUCTURED_OUTPUT_RATIO \
+               --save-results \
+               --result-dir $OUTPUT_DIR \
+               --output-len $MAX_NEW_TOKENS \
+               --port $PORT \
+               --tokenizer-mode $TOKENIZER_MODE"
 
 echo "Starting structured output benchmark with model: $MODEL"
 echo "Backend: $BACKEND"
@@ -111,17 +109,17 @@ for qps in "${QPS_VALUES[@]}"; do
   GIT_BRANCH=$(git rev-parse --abbrev-ref HEAD 2>/dev/null || echo "unknown")
 
   # Construct filename for this run
-  FILENAME="${BACKEND}_${qps}qps_$(basename "$MODEL")_${DATASET}_${GIT_HASH}_${GIT_BRANCH}.json"
+  FILENAME="${BACKEND}_${qps}qps_$(basename $MODEL)_${DATASET}_${GIT_HASH}.json"
 
   NUM_PROMPTS=$(echo "$TOTAL_SECONDS * $qps" | bc)
   NUM_PROMPTS=${NUM_PROMPTS%.*}  # Remove fractional part
   echo "Running benchmark with $NUM_PROMPTS prompts"
 
   # Run the benchmark
-  python "$SCRIPT_DIR/benchmark_serving_structured_output.py" "${COMMON_PARAMS[@]}" \
-    --request-rate "$qps" \
+  python "$SCRIPT_DIR/benchmark_serving_structured_output.py" $COMMON_PARAMS \
+    --request-rate $qps \
     --result-filename "$FILENAME" \
-    --num-prompts "$NUM_PROMPTS"
+    --num-prompts $NUM_PROMPTS
 
   echo "Completed benchmark with QPS: $qps"
   echo "----------------------------------------"

cmake/cpu_extension.cmake

@@ -13,16 +13,27 @@ endif()
 #
 # Define environment variables for special configurations
 #
-set(ENABLE_X86_ISA $ENV{VLLM_CPU_X86})
-set(ENABLE_ARM_BF16 $ENV{VLLM_CPU_ARM_BF16})
+set(ENABLE_AVX2 $ENV{VLLM_CPU_AVX2})
+set(ENABLE_AVX512 $ENV{VLLM_CPU_AVX512})
+set(ENABLE_AVX512BF16 $ENV{VLLM_CPU_AVX512BF16})
+set(ENABLE_AVX512VNNI $ENV{VLLM_CPU_AVX512VNNI})
+set(ENABLE_AMXBF16 $ENV{VLLM_CPU_AMXBF16})
 
 include_directories("${CMAKE_SOURCE_DIR}/csrc")
 
+
 set (ENABLE_NUMA TRUE)
 
 #
 # Check the compile flags
 #
+
+if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64")
+    list(APPEND CXX_COMPILE_FLAGS
+        "-mf16c"
+    )
+endif()
+
 if(MACOSX_FOUND)
     list(APPEND CXX_COMPILE_FLAGS
         "-DVLLM_CPU_EXTENSION")
@@ -66,6 +77,18 @@ function(check_sysctl TARGET OUT)
     endif()
 endfunction()
 
+
+function (is_avx512_disabled OUT)
+    set(DISABLE_AVX512 $ENV{VLLM_CPU_DISABLE_AVX512})
+    if(DISABLE_AVX512 AND DISABLE_AVX512 STREQUAL "true")
+        set(${OUT} ON PARENT_SCOPE)
+    else()
+        set(${OUT} OFF PARENT_SCOPE)
+    endif()
+endfunction()
+
+is_avx512_disabled(AVX512_DISABLED)
+
 if (MACOSX_FOUND AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
     message(STATUS "Apple Silicon Detected")
     set(APPLE_SILICON_FOUND TRUE)
@@ -73,6 +96,8 @@ if (MACOSX_FOUND AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
     check_sysctl(hw.optional.neon ASIMD_FOUND)
     check_sysctl(hw.optional.arm.FEAT_BF16 ARM_BF16_FOUND)
 else()
+    find_isa(${CPUINFO} "avx2" AVX2_FOUND)
+    find_isa(${CPUINFO} "avx512f" AVX512_FOUND)
     find_isa(${CPUINFO} "Power11" POWER11_FOUND)
     find_isa(${CPUINFO} "POWER10" POWER10_FOUND)
     find_isa(${CPUINFO} "POWER9" POWER9_FOUND)
@@ -82,32 +107,73 @@ else()
     find_isa(${CPUINFO} "v" RVV_FOUND) # Check for RISC-V RVV support
 
     # Support cross-compilation by allowing override via environment variables
-    if (ENABLE_ARM_BF16)
-        set(ARM_BF16_FOUND ON)
-        message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable")
+    if (ENABLE_AVX2)
+        set(AVX2_FOUND ON)
+        message(STATUS "AVX2 support enabled via VLLM_CPU_AVX2 environment variable")
+    endif()
+    if (ENABLE_AVX512)
+        set(AVX512_FOUND ON)
+        message(STATUS "AVX512 support enabled via VLLM_CPU_AVX512 environment variable")
     endif()
 endif()
 
-if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64" OR ENABLE_X86_ISA)
-    set(ENABLE_X86_ISA ON)
-    if (NOT (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
-            CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12.3))
-        message(FATAL_ERROR "X86 backend requires gcc/g++ >= 12.3")
-    endif()
-    list(APPEND CXX_COMPILE_FLAGS "-mf16c")
-    list(APPEND CXX_COMPILE_FLAGS_AVX512 ${CXX_COMPILE_FLAGS})
-    list(APPEND CXX_COMPILE_FLAGS_AVX2 ${CXX_COMPILE_FLAGS})
-    list(APPEND CXX_COMPILE_FLAGS_AVX512
+if (AVX512_FOUND AND NOT AVX512_DISABLED)
+    list(APPEND CXX_COMPILE_FLAGS
         "-mavx512f"
         "-mavx512vl"
         "-mavx512bw"
-        "-mavx512dq"
-        "-mavx512bf16"
-        "-mavx512vnni"
-        "-mamx-bf16"
-        "-mamx-tile")
-    list(APPEND CXX_COMPILE_FLAGS_AVX2
-        "-mavx2")
+        "-mavx512dq")
+
+    find_isa(${CPUINFO} "avx512_bf16" AVX512BF16_FOUND)
+    if (AVX512BF16_FOUND OR ENABLE_AVX512BF16)
+        if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
+            CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12.3)
+            list(APPEND CXX_COMPILE_FLAGS "-mavx512bf16")
+            set(ENABLE_AVX512BF16 ON)
+        else()
+            set(ENABLE_AVX512BF16 OFF)
+            message(WARNING "Disable AVX512-BF16 ISA support, requires gcc/g++ >= 12.3")
+        endif()
+    else()
+        set(ENABLE_AVX512BF16 OFF)
+        message(WARNING "Disable AVX512-BF16 ISA support, no avx512_bf16 found in local CPU flags." " If cross-compilation is required, please set env VLLM_CPU_AVX512BF16=1.")
+    endif()
+
+    find_isa(${CPUINFO} "avx512_vnni" AVX512VNNI_FOUND)
+    if (AVX512VNNI_FOUND OR ENABLE_AVX512VNNI)
+        if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
+            CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12.3)
+            list(APPEND CXX_COMPILE_FLAGS "-mavx512vnni")
+            set(ENABLE_AVX512VNNI ON)
+        else()
+            set(ENABLE_AVX512VNNI OFF)
+            message(WARNING "Disable AVX512-VNNI ISA support, requires gcc/g++ >= 12.3")
+        endif()
+    else()
+        set(ENABLE_AVX512VNNI OFF)
+        message(WARNING "Disable AVX512-VNNI ISA support, no avx512_vnni found in local CPU flags." " If cross-compilation is required, please set env VLLM_CPU_AVX512VNNI=1.")
+    endif()
+
+    find_isa(${CPUINFO} "amx_bf16" AMXBF16_FOUND)
+    if (AMXBF16_FOUND OR ENABLE_AMXBF16)
+        if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
+            CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12.3)
+            list(APPEND CXX_COMPILE_FLAGS "-mamx-bf16" "-mamx-tile")
+            set(ENABLE_AMXBF16 ON)
+            add_compile_definitions(-DCPU_CAPABILITY_AMXBF16)
+        else()
+            set(ENABLE_AMXBF16 OFF)
+            message(WARNING "Disable AMX_BF16 ISA support, requires gcc/g++ >= 12.3")
+        endif()
+    else()
+        set(ENABLE_AMXBF16 OFF)
+        message(WARNING "Disable AMX_BF16 ISA support, no amx_bf16 found in local CPU flags." " If cross-compilation is required, please set env VLLM_CPU_AMXBF16=1.")
+    endif()
+    
+elseif (AVX2_FOUND)
+    list(APPEND CXX_COMPILE_FLAGS "-mavx2")
+    message(WARNING "vLLM CPU backend using AVX2 ISA")
+    
 elseif (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
     message(STATUS "PowerPC detected")
     if (POWER9_FOUND)
@@ -148,12 +214,12 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
         list(APPEND CXX_COMPILE_FLAGS "-march=rv64gc")
     endif()
 else()
-    message(FATAL_ERROR "vLLM CPU backend requires X86, Power9+ ISA, S390X ISA, ARMv8 or RISC-V support.")
+    message(FATAL_ERROR "vLLM CPU backend requires AVX512, AVX2, Power9+ ISA, S390X ISA, ARMv8 or RISC-V support.")
 endif()
 
 
-# Build oneDNN for GEMM kernels
-if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
+# Build oneDNN for GEMM kernels (only for x86-AVX512 /ARM platforms)
+if ((AVX512_FOUND AND NOT AVX512_DISABLED) OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
     # Fetch and build Arm Compute Library (ACL) as oneDNN's backend for AArch64
     # TODO [fadara01]: remove this once ACL can be fetched and built automatically as a dependency of oneDNN
     set(ONEDNN_AARCH64_USE_ACL OFF CACHE BOOL "")
@@ -258,21 +324,13 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
     set(ONEDNN_ENABLE_WORKLOAD "INFERENCE")
     set(ONEDNN_ENABLE_PRIMITIVE "MATMUL;REORDER")
     set(ONEDNN_BUILD_GRAPH "OFF")
-    set(ONEDNN_ENABLE_JIT_PROFILING "ON")
+    set(ONEDNN_ENABLE_JIT_PROFILING "OFF")
     set(ONEDNN_ENABLE_ITT_TASKS "OFF")
-    set(ONEDNN_ENABLE_MAX_CPU_ISA "ON")
-    set(ONEDNN_ENABLE_CPU_ISA_HINTS "ON")
-    set(ONEDNN_VERBOSE "ON")
+    set(ONEDNN_ENABLE_MAX_CPU_ISA "OFF")
+    set(ONEDNN_ENABLE_CPU_ISA_HINTS "OFF")
+    set(ONEDNN_VERBOSE "OFF")
     set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
 
-    # TODO: Refactor this
-    if (ENABLE_X86_ISA)
-        # Note: only enable oneDNN for AVX512
-        list(APPEND DNNL_COMPILE_FLAGS ${CXX_COMPILE_FLAGS_AVX512})
-    else()
-        list(APPEND DNNL_COMPILE_FLAGS ${CXX_COMPILE_FLAGS})
-    endif()
-
     set(VLLM_BUILD_TYPE ${CMAKE_BUILD_TYPE})
     set(CMAKE_BUILD_TYPE "Release") # remove oneDNN debug symbols to reduce size
     FetchContent_MakeAvailable(oneDNN)
@@ -285,20 +343,14 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
         PRIVATE ${oneDNN_SOURCE_DIR}/src
     )
     target_link_libraries(dnnl_ext dnnl torch)
-    target_compile_options(dnnl_ext PRIVATE ${DNNL_COMPILE_FLAGS} -fPIC)
+    target_compile_options(dnnl_ext PRIVATE ${CXX_COMPILE_FLAGS} -fPIC)
     list(APPEND LIBS dnnl_ext)
     set(USE_ONEDNN ON)
 else()
     set(USE_ONEDNN OFF)
 endif()
 
-# TODO: Refactor this
-if (ENABLE_X86_ISA)
-    message(STATUS "CPU extension (AVX512) compile flags: ${CXX_COMPILE_FLAGS_AVX512}")
-    message(STATUS "CPU extension (AVX2) compile flags: ${CXX_COMPILE_FLAGS_AVX2}")
-else()
-    message(STATUS "CPU extension compile flags: ${CXX_COMPILE_FLAGS}")
-endif()
+message(STATUS "CPU extension compile flags: ${CXX_COMPILE_FLAGS}")
 
 if(ENABLE_NUMA)
     list(APPEND LIBS numa)
@@ -333,6 +385,25 @@ set(VLLM_EXT_SRC
     "csrc/cpu/cpu_attn.cpp"
     "csrc/cpu/torch_bindings.cpp")
 
+if (AVX512_FOUND AND NOT AVX512_DISABLED)
+    set(VLLM_EXT_SRC
+        "csrc/cpu/shm.cpp"
+        "csrc/cpu/cpu_wna16.cpp"
+        "csrc/cpu/cpu_fused_moe.cpp"
+        ${VLLM_EXT_SRC})
+    if (ENABLE_AVX512BF16 AND ENABLE_AVX512VNNI)
+        set(VLLM_EXT_SRC
+            "csrc/cpu/sgl-kernels/gemm.cpp"
+            "csrc/cpu/sgl-kernels/gemm_int8.cpp"
+            "csrc/cpu/sgl-kernels/gemm_fp8.cpp"
+            "csrc/cpu/sgl-kernels/moe.cpp"
+            "csrc/cpu/sgl-kernels/moe_int8.cpp"
+            "csrc/cpu/sgl-kernels/moe_fp8.cpp"
+            ${VLLM_EXT_SRC})
+        add_compile_definitions(-DCPU_CAPABILITY_AVX512)
+    endif()
+endif()
+
 if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
     set(VLLM_EXT_SRC
         "csrc/cpu/shm.cpp"
@@ -345,83 +416,21 @@ if(USE_ONEDNN)
         ${VLLM_EXT_SRC})
 endif()
 
-if (ENABLE_X86_ISA)
-    set(VLLM_EXT_SRC_AVX512
-        "csrc/cpu/sgl-kernels/gemm.cpp"
-        "csrc/cpu/sgl-kernels/gemm_int8.cpp"
-        "csrc/cpu/sgl-kernels/gemm_fp8.cpp"
-        "csrc/cpu/sgl-kernels/moe.cpp"
-        "csrc/cpu/sgl-kernels/moe_int8.cpp"
-        "csrc/cpu/sgl-kernels/moe_fp8.cpp"
-        "csrc/cpu/shm.cpp"
-        "csrc/cpu/cpu_wna16.cpp"
-        "csrc/cpu/cpu_fused_moe.cpp"
-        "csrc/cpu/utils.cpp"
-        "csrc/cpu/cpu_attn.cpp"
-        "csrc/cpu/dnnl_kernels.cpp"
-        "csrc/cpu/torch_bindings.cpp"
-        # TODO: Remove these files
-        "csrc/cpu/activation.cpp"
-        "csrc/cpu/layernorm.cpp"
-        "csrc/cpu/mla_decode.cpp"
-        "csrc/cpu/pos_encoding.cpp"
-        "csrc/moe/dynamic_4bit_int_moe_cpu.cpp") 
+message(STATUS "CPU extension source files: ${VLLM_EXT_SRC}")
 
-    set(VLLM_EXT_SRC_AVX2 
-        "csrc/cpu/utils.cpp"
-        "csrc/cpu/cpu_attn.cpp"
-        "csrc/cpu/torch_bindings.cpp"
-        # TODO: Remove these files
-        "csrc/cpu/activation.cpp"
-        "csrc/cpu/layernorm.cpp"
-        "csrc/cpu/mla_decode.cpp"
-        "csrc/cpu/pos_encoding.cpp"
-        "csrc/moe/dynamic_4bit_int_moe_cpu.cpp") 
+#
+# Define extension targets
+#
 
-    message(STATUS "CPU extension (AVX512) source files: ${VLLM_EXT_SRC_AVX512}")
-    message(STATUS "CPU extension (AVX2) source files: ${VLLM_EXT_SRC_AVX2}")
-
-    define_extension_target(
-        _C
-        DESTINATION vllm
-        LANGUAGE CXX
-        SOURCES ${VLLM_EXT_SRC_AVX512}
-        LIBRARIES ${LIBS}
-        COMPILE_FLAGS ${CXX_COMPILE_FLAGS_AVX512}
-        USE_SABI 3
-        WITH_SOABI
-    )
-
-    # For SGL kernels
-    target_compile_definitions(_C PRIVATE "-DCPU_CAPABILITY_AVX512")
-    # For AMX kernels
-    target_compile_definitions(_C PRIVATE "-DCPU_CAPABILITY_AMXBF16")
-
-    define_extension_target(
-        _C_AVX2
-        DESTINATION vllm
-        LANGUAGE CXX
-        SOURCES ${VLLM_EXT_SRC_AVX2}
-        LIBRARIES ${LIBS}
-        COMPILE_FLAGS ${CXX_COMPILE_FLAGS_AVX2}
-        USE_SABI 3
-        WITH_SOABI
-    )
-else()
-    message(STATUS "CPU extension source files: ${VLLM_EXT_SRC}")
-    #
-    # Define extension targets
-    #
-    define_extension_target(
-        _C
-        DESTINATION vllm
-        LANGUAGE CXX
-        SOURCES ${VLLM_EXT_SRC}
-        LIBRARIES ${LIBS}
-        COMPILE_FLAGS ${CXX_COMPILE_FLAGS}
-        USE_SABI 3
-        WITH_SOABI
-    )
-endif()
+define_extension_target(
+    _C
+    DESTINATION vllm
+    LANGUAGE CXX
+    SOURCES ${VLLM_EXT_SRC}
+    LIBRARIES ${LIBS}
+    COMPILE_FLAGS ${CXX_COMPILE_FLAGS}
+    USE_SABI 3
+    WITH_SOABI
+)
 
 message(STATUS "Enabling C extension.")

cmake/external_projects/flashmla.cmake

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

cmake/external_projects/vllm_flash_attn.cmake

@@ -17,8 +17,7 @@ endif()
 # They should be identical but if they aren't, this is a massive footgun.
 #
 # The vllm-flash-attn install rules are nested under vllm to make sure the library gets installed in the correct place.
-# To only install vllm-flash-attn, use --component _vllm_fa2_C (for FA2), --component _vllm_fa3_C (for FA3),
-# or --component _vllm_fa4_cutedsl_C (for FA4 CuteDSL Python files).
+# To only install vllm-flash-attn, use --component _vllm_fa2_C (for FA2) or --component _vllm_fa3_C (for FA3).
 # If no component is specified, vllm-flash-attn is still installed.
 
 # If VLLM_FLASH_ATTN_SRC_DIR is set, vllm-flash-attn is installed from that directory instead of downloading.
@@ -39,7 +38,7 @@ else()
   FetchContent_Declare(
           vllm-flash-attn
           GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG 140c00c0241bb60cc6e44e7c1be9998d4b20d8d2
+          GIT_TAG 5824e6e2008271063c3229ab3e7032bd74abbbc6
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
@@ -47,62 +46,38 @@ else()
 endif()
 
 
-# Install rules for FA components need the install prefix nested under vllm/
-# These run at install time, before the FA library's own install rules
-foreach(_FA_COMPONENT _vllm_fa2_C _vllm_fa3_C)
-  install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY FALSE)" COMPONENT ${_FA_COMPONENT})
-  install(CODE "set(OLD_CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}\")" COMPONENT ${_FA_COMPONENT})
-  install(CODE "set(CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}/vllm/\")" COMPONENT ${_FA_COMPONENT})
-endforeach()
+# Ensure the vllm/vllm_flash_attn directory exists before installation
+install(CODE "file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn\")" ALL_COMPONENTS)
+
+# Make sure vllm-flash-attn install rules are nested under vllm/
+# This is here to support installing all components under the same prefix with cmake --install.
+# setup.py installs every component separately but uses the same prefix for all.
+# ALL_COMPONENTS is used to avoid duplication for FA2 and FA3,
+# and these statements don't hurt when installing neither component.
+install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY FALSE)" ALL_COMPONENTS)
+install(CODE "set(OLD_CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}\")" ALL_COMPONENTS)
+install(CODE "set(CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}/vllm/\")" ALL_COMPONENTS)
 
 # Fetch the vllm-flash-attn library
 FetchContent_MakeAvailable(vllm-flash-attn)
 message(STATUS "vllm-flash-attn is available at ${vllm-flash-attn_SOURCE_DIR}")
 
-# Restore the install prefix after FA's install rules
-foreach(_FA_COMPONENT _vllm_fa2_C _vllm_fa3_C)
-  install(CODE "set(CMAKE_INSTALL_PREFIX \"\${OLD_CMAKE_INSTALL_PREFIX}\")" COMPONENT ${_FA_COMPONENT})
-  install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" COMPONENT ${_FA_COMPONENT})
-endforeach()
+# Restore the install prefix
+install(CODE "set(CMAKE_INSTALL_PREFIX \"\${OLD_CMAKE_INSTALL_PREFIX}\")" ALL_COMPONENTS)
+install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS)
 
-# Install shared Python files for both FA2 and FA3 components
-foreach(_FA_COMPONENT _vllm_fa2_C _vllm_fa3_C)
-  # Ensure the vllm/vllm_flash_attn directory exists before installation
-  install(CODE "file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn\")"
-    COMPONENT ${_FA_COMPONENT})
+# Copy over the vllm-flash-attn python files (duplicated for fa2 and fa3, in
+# case only one is built, in the case both are built redundant work is done)
+install(
+  DIRECTORY ${vllm-flash-attn_SOURCE_DIR}/vllm_flash_attn/
+  DESTINATION vllm/vllm_flash_attn
+  COMPONENT _vllm_fa2_C
+  FILES_MATCHING PATTERN "*.py"
+)
 
-  # Copy vllm_flash_attn python files (except __init__.py and flash_attn_interface.py
-  # which are source-controlled in vllm)
-  install(
-    DIRECTORY ${vllm-flash-attn_SOURCE_DIR}/vllm_flash_attn/
-    DESTINATION vllm/vllm_flash_attn
-    COMPONENT ${_FA_COMPONENT}
-    FILES_MATCHING PATTERN "*.py"
-    PATTERN "__init__.py" EXCLUDE
-    PATTERN "flash_attn_interface.py" EXCLUDE
-  )
-
-endforeach()
-
-#
-# FA4 CuteDSL component
-# This is a Python-only component that copies the flash_attn/cute directory
-# and transforms imports to match our package structure.
-#
-add_custom_target(_vllm_fa4_cutedsl_C)
-
-# Copy flash_attn/cute directory (needed for FA4) and transform imports
-# The cute directory uses flash_attn.cute imports internally, which we replace
-# with vllm.vllm_flash_attn.cute to match our package structure.
-install(CODE "
-  file(GLOB_RECURSE CUTE_PY_FILES \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute/*.py\")
-  foreach(SRC_FILE \${CUTE_PY_FILES})
-    file(RELATIVE_PATH REL_PATH \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute\" \${SRC_FILE})
-    set(DST_FILE \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn/cute/\${REL_PATH}\")
-    get_filename_component(DST_DIR \${DST_FILE} DIRECTORY)
-    file(MAKE_DIRECTORY \${DST_DIR})
-    file(READ \${SRC_FILE} FILE_CONTENTS)
-    string(REPLACE \"flash_attn.cute\" \"vllm.vllm_flash_attn.cute\" FILE_CONTENTS \"\${FILE_CONTENTS}\")
-    file(WRITE \${DST_FILE} \"\${FILE_CONTENTS}\")
-  endforeach()
-" COMPONENT _vllm_fa4_cutedsl_C)
+install(
+  DIRECTORY ${vllm-flash-attn_SOURCE_DIR}/vllm_flash_attn/
+  DESTINATION vllm/vllm_flash_attn
+  COMPONENT _vllm_fa3_C
+  FILES_MATCHING PATTERN "*.py"
+)

csrc/activation_kernels.cu

@@ -5,11 +5,115 @@
 #include <cmath>
 
 #include "cuda_compat.h"
-#include "cuda_vec_utils.cuh"
 #include "dispatch_utils.h"
 
 namespace vllm {
 
+struct alignas(32) u32x8_t {
+  uint32_t u0, u1, u2, u3, u4, u5, u6, u7;
+};
+
+__device__ __forceinline__ void ld256(u32x8_t& val, const u32x8_t* ptr) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
+  asm volatile("ld.global.nc.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];\n"
+               : "=r"(val.u0), "=r"(val.u1), "=r"(val.u2), "=r"(val.u3),
+                 "=r"(val.u4), "=r"(val.u5), "=r"(val.u6), "=r"(val.u7)
+               : "l"(ptr));
+#else
+  const uint4* uint_ptr = reinterpret_cast<const uint4*>(ptr);
+  uint4 top_half = __ldg(&uint_ptr[0]);
+  uint4 bottom_half = __ldg(&uint_ptr[1]);
+  val.u0 = top_half.x;
+  val.u1 = top_half.y;
+  val.u2 = top_half.z;
+  val.u3 = top_half.w;
+  val.u4 = bottom_half.x;
+  val.u5 = bottom_half.y;
+  val.u6 = bottom_half.z;
+  val.u7 = bottom_half.w;
+#endif
+}
+
+__device__ __forceinline__ void st256(u32x8_t& val, u32x8_t* ptr) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
+  asm volatile("st.global.v8.u32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};\n"
+               :
+               : "l"(ptr), "r"(val.u0), "r"(val.u1), "r"(val.u2), "r"(val.u3),
+                 "r"(val.u4), "r"(val.u5), "r"(val.u6), "r"(val.u7)
+               : "memory");
+#else
+  uint4* uint_ptr = reinterpret_cast<uint4*>(ptr);
+  uint_ptr[0] = make_uint4(val.u0, val.u1, val.u2, val.u3);
+  uint_ptr[1] = make_uint4(val.u4, val.u5, val.u6, val.u7);
+#endif
+}
+
+template <bool support_256>
+struct VecTraits;
+
+template <>
+struct VecTraits<true> {
+  static constexpr int ARCH_MAX_VEC_SIZE = 32;
+  using vec_t = u32x8_t;
+};
+
+template <>
+struct VecTraits<false> {
+  static constexpr int ARCH_MAX_VEC_SIZE = 16;
+  using vec_t = int4;
+};
+
+template <typename T>
+struct PackedTraits;
+
+template <>
+struct PackedTraits<c10::BFloat16> {
+  using packed_t = __nv_bfloat162;
+};
+
+template <>
+struct PackedTraits<c10::Half> {
+  using packed_t = __half2;
+};
+
+template <>
+struct PackedTraits<float> {
+  using packed_t = float2;
+};
+
+template <typename packed_t>
+__device__ __forceinline__ float2 cast_to_float2(const packed_t& val) {
+  if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
+    return __bfloat1622float2(val);
+  } else if constexpr (std::is_same_v<packed_t, __half2>) {
+    return __half22float2(val);
+  } else if constexpr (std::is_same_v<packed_t, float2>) {
+    return float2(val);
+  }
+}
+
+template <typename packed_t>
+__device__ __forceinline__ packed_t cast_to_packed(const float2& val) {
+  if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
+    return __float22bfloat162_rn(val);
+  } else if constexpr (std::is_same_v<packed_t, __half2>) {
+    return __float22half2_rn(val);
+  } else if constexpr (std::is_same_v<packed_t, float2>) {
+    return float2(val);
+  }
+}
+
+template <typename packed_t>
+__device__ __forceinline__ packed_t packed_mul(const packed_t& x,
+                                               const packed_t& y) {
+  if constexpr (std::is_same_v<packed_t, __nv_bfloat162> ||
+                std::is_same_v<packed_t, __half2>) {
+    return __hmul2(x, y);
+  } else if constexpr (std::is_same_v<packed_t, float2>) {
+    return make_float2(x.x * y.x, x.y * y.y);
+  }
+}
+
 template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
           bool act_first>
 __device__ __forceinline__ scalar_t compute(const scalar_t& x,
@@ -25,6 +129,16 @@ __device__ __forceinline__ packed_t packed_compute(const packed_t& x,
                    : packed_mul(x, PACKED_ACT_FN(y));
 }
 
+// Check if all pointers are 16-byte aligned for int4 vectorized access
+__host__ __device__ __forceinline__ bool is_16byte_aligned(const void* ptr) {
+  return (reinterpret_cast<uintptr_t>(ptr) & 15) == 0;
+}
+
+// Check if all pointers are 16-byte aligned for longlong4_32a vectorized access
+__host__ __device__ __forceinline__ bool is_32byte_aligned(const void* ptr) {
+  return (reinterpret_cast<uintptr_t>(ptr) & 31) == 0;
+}
+
 // Activation and gating kernel template.
 template <typename scalar_t, typename packed_t,
           scalar_t (*ACT_FN)(const scalar_t&),
@@ -39,32 +153,36 @@ __global__ void act_and_mul_kernel(
   scalar_t* out_ptr = out + blockIdx.x * d;
 
   if constexpr (use_vec) {
-    using cuda_t = typename CUDATypeConverter<scalar_t>::Type;
-    using pvec_t = PackedVec<cuda_t, use_256b>;
+    // Fast path: 128-bit/256-bit vectorized loop
+    using vec_t = typename VecTraits<use_256b>::vec_t;
+    constexpr int ARCH_MAX_VEC_SIZE = VecTraits<use_256b>::ARCH_MAX_VEC_SIZE;
+    constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(packed_t);
 
-    const pvec_t* x_vec = reinterpret_cast<const pvec_t*>(x_ptr);
-    const pvec_t* y_vec = reinterpret_cast<const pvec_t*>(y_ptr);
-    pvec_t* out_vec = reinterpret_cast<pvec_t*>(out_ptr);
-    const int num_vecs = d / 2 / pvec_t::NUM_ELTS;
+    const vec_t* x_vec = reinterpret_cast<const vec_t*>(x_ptr);
+    const vec_t* y_vec = reinterpret_cast<const vec_t*>(y_ptr);
+    vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
+    const int num_vecs = d / 2 / VEC_SIZE;
 
     for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
-      pvec_t x, y;
+      vec_t x, y;
       if constexpr (use_256b) {
         ld256(x, &x_vec[i]);
         ld256(y, &y_vec[i]);
       } else {
-        ld128(x, &x_vec[i]);
-        ld128(y, &y_vec[i]);
+        x = VLLM_LDG(&x_vec[i]);
+        y = VLLM_LDG(&y_vec[i]);
       }
+      auto* xp = reinterpret_cast<packed_t*>(&x);
+      auto* yp = reinterpret_cast<packed_t*>(&y);
 #pragma unroll
-      for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
-        x.elts[j] = packed_compute<packed_t, PACKED_ACT_FN, act_first>(
-            x.elts[j], y.elts[j]);
+      for (int j = 0; j < VEC_SIZE; j++) {
+        xp[j] =
+            packed_compute<packed_t, PACKED_ACT_FN, act_first>(xp[j], yp[j]);
       }
       if constexpr (use_256b) {
         st256(x, &out_vec[i]);
       } else {
-        st128(x, &out_vec[i]);
+        out_vec[i] = x;
       }
     }
   } else {
@@ -152,54 +270,51 @@ 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.
-#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);                         \
-  if (num_tokens == 0) {                                                       \
-    return;                                                                    \
-  }                                                                            \
-  dim3 grid(num_tokens);                                                       \
-  int cc_major = at::cuda::getCurrentDeviceProperties()->major;                \
-  int support_vec =                                                            \
-      (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128)            \
-          ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE                           \
-          : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE;                         \
-  int vec_size = support_vec / at::elementSize(dtype);                         \
-  const bool use_vec = (d % vec_size == 0);                                    \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
-  if (use_vec) {                                                               \
-    dim3 block(std::min(d / vec_size, 1024));                                  \
-    if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) {         \
-      VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {          \
-        vllm::act_and_mul_kernel<                                              \
-            scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
-            KERNEL<scalar_t>,                                                  \
-            PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
-            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", [&] {          \
-        vllm::act_and_mul_kernel<                                              \
-            scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
-            KERNEL<scalar_t>,                                                  \
-            PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
-            ACT_FIRST, true, false><<<grid, block, 0, stream>>>(               \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);          \
-      });                                                                      \
-    }                                                                          \
-  } else {                                                                     \
-    dim3 block(std::min(d, 1024));                                             \
-    VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {            \
-      vllm::act_and_mul_kernel<                                                \
-          scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,        \
-          KERNEL<scalar_t>,                                                    \
-          PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
-          ACT_FIRST, false><<<grid, block, 0, stream>>>(                       \
-          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);            \
-    });                                                                        \
+#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);                      \
+  if (num_tokens == 0) {                                                    \
+    return;                                                                 \
+  }                                                                         \
+  dim3 grid(num_tokens);                                                    \
+  int cc_major = at::cuda::getCurrentDeviceProperties()->major;             \
+  int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16;         \
+  int vec_size = support_vec / at::elementSize(dtype);                      \
+  const bool use_vec = (d % vec_size == 0);                                 \
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));         \
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();             \
+  if (use_vec) {                                                            \
+    dim3 block(std::min(d / vec_size, 1024));                               \
+    if (cc_major >= 10 && num_tokens > 128) {                               \
+      VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {       \
+        vllm::act_and_mul_kernel<                                           \
+            scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,      \
+            KERNEL<scalar_t>,                                               \
+            PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
+            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", [&] {       \
+        vllm::act_and_mul_kernel<                                           \
+            scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,      \
+            KERNEL<scalar_t>,                                               \
+            PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
+            ACT_FIRST, true, false><<<grid, block, 0, stream>>>(            \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);       \
+      });                                                                   \
+    }                                                                       \
+  } else {                                                                  \
+    dim3 block(std::min(d, 1024));                                          \
+    VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {         \
+      vllm::act_and_mul_kernel<                                             \
+          scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,        \
+          KERNEL<scalar_t>,                                                 \
+          PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>,   \
+          ACT_FIRST, false><<<grid, block, 0, stream>>>(                    \
+          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);         \
+    });                                                                     \
   }
 
 void silu_and_mul(torch::Tensor& out,    // [..., d]
@@ -261,31 +376,35 @@ __global__ void act_and_mul_kernel_with_param(
   scalar_t* out_ptr = out + blockIdx.x * d;
 
   if constexpr (use_vec) {
-    using cuda_t = typename CUDATypeConverter<scalar_t>::Type;
-    using pvec_t = PackedVec<cuda_t, use_256b>;
+    // Fast path: 128-bit/256-bit vectorized loop
+    using vec_t = typename VecTraits<use_256b>::vec_t;
+    constexpr int ARCH_MAX_VEC_SIZE = VecTraits<use_256b>::ARCH_MAX_VEC_SIZE;
+    constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(packed_t);
 
-    const pvec_t* x_vec = reinterpret_cast<const pvec_t*>(x_ptr);
-    const pvec_t* y_vec = reinterpret_cast<const pvec_t*>(y_ptr);
-    pvec_t* out_vec = reinterpret_cast<pvec_t*>(out_ptr);
-    const int num_vecs = d / 2 / pvec_t::NUM_ELTS;
+    const vec_t* x_vec = reinterpret_cast<const vec_t*>(x_ptr);
+    const vec_t* y_vec = reinterpret_cast<const vec_t*>(y_ptr);
+    vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
+    const int num_vecs = d / 2 / VEC_SIZE;
 
     for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
-      pvec_t x, y;
+      vec_t x, y;
       if constexpr (use_256b) {
         ld256(x, &x_vec[i]);
         ld256(y, &y_vec[i]);
       } else {
-        ld128(x, &x_vec[i]);
-        ld128(y, &y_vec[i]);
+        x = VLLM_LDG(&x_vec[i]);
+        y = VLLM_LDG(&y_vec[i]);
       }
+      auto* xp = reinterpret_cast<packed_t*>(&x);
+      auto* yp = reinterpret_cast<packed_t*>(&y);
 #pragma unroll
-      for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
-        x.elts[j] = packed_mul(PACKED_ACT_FN(x.elts[j], param), y.elts[j]);
+      for (int j = 0; j < VEC_SIZE; j++) {
+        xp[j] = packed_mul(PACKED_ACT_FN(xp[j], param), yp[j]);
       }
       if constexpr (use_256b) {
         st256(x, &out_vec[i]);
       } else {
-        st128(x, &out_vec[i]);
+        out_vec[i] = x;
       }
     }
   } else {
@@ -378,24 +497,21 @@ __global__ void swigluoai_and_mul_kernel(
   }                                                                            \
   dim3 grid(num_tokens);                                                       \
   int cc_major = at::cuda::getCurrentDeviceProperties()->major;                \
-  int support_vec =                                                            \
-      (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128)            \
-          ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE                           \
-          : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE;                         \
+  int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16;            \
   int vec_size = support_vec / at::elementSize(dtype);                         \
   const bool use_vec = (d % vec_size == 0);                                    \
   const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
   if (use_vec) {                                                               \
     dim3 block(std::min(d / vec_size, 1024));                                  \
-    if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) {         \
+    if (cc_major >= 10 && num_tokens > 128) {                                  \
       VLLM_DISPATCH_FLOATING_TYPES(                                            \
           dtype, "act_and_mul_kernel_with_param", [&] {                        \
             vllm::act_and_mul_kernel_with_param<                               \
-                scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,  \
+                scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,     \
                 KERNEL<scalar_t>,                                              \
                 PACKED_KERNEL<                                                 \
-                    typename vllm::PackedTypeConverter<scalar_t>::Type>,       \
+                    typename vllm::PackedTraits<scalar_t>::packed_t>,          \
                 true, true><<<grid, block, 0, stream>>>(                       \
                 out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d,       \
                 PARAM);                                                        \
@@ -404,10 +520,10 @@ __global__ void swigluoai_and_mul_kernel(
       VLLM_DISPATCH_FLOATING_TYPES(                                            \
           dtype, "act_and_mul_kernel_with_param", [&] {                        \
             vllm::act_and_mul_kernel_with_param<                               \
-                scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,  \
+                scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,     \
                 KERNEL<scalar_t>,                                              \
                 PACKED_KERNEL<                                                 \
-                    typename vllm::PackedTypeConverter<scalar_t>::Type>,       \
+                    typename vllm::PackedTraits<scalar_t>::packed_t>,          \
                 true, false><<<grid, block, 0, stream>>>(                      \
                 out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d,       \
                 PARAM);                                                        \
@@ -417,9 +533,9 @@ __global__ void swigluoai_and_mul_kernel(
     dim3 block(std::min(d, 1024));                                             \
     VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel_with_param", [&] { \
       vllm::act_and_mul_kernel_with_param<                                     \
-          scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,        \
+          scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,           \
           KERNEL<scalar_t>,                                                    \
-          PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
+          PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>,      \
           false><<<grid, block, 0, stream>>>(                                  \
           out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, PARAM);     \
     });                                                                        \
@@ -511,17 +627,14 @@ __global__ void activation_kernel(
   }                                                                      \
   dim3 grid(num_tokens);                                                 \
   int cc_major = at::cuda::getCurrentDeviceProperties()->major;          \
-  int support_vec =                                                      \
-      (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128)      \
-          ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE                     \
-          : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE;                   \
+  int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16;      \
   int vec_size = support_vec / at::elementSize(dtype);                   \
   const bool use_vec = (d % vec_size == 0);                              \
   const at::cuda::OptionalCUDAGuard device_guard(device_of(input));      \
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();          \
   if (use_vec) {                                                         \
     dim3 block(std::min(d / vec_size, 1024));                            \
-    if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) {   \
+    if (cc_major >= 10 && num_tokens > 128) {                            \
       VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {     \
         vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, true>  \
             <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \

csrc/cache_kernels.cu

@@ -1305,8 +1305,7 @@ void indexer_k_quant_and_cache(
   const at::cuda::OptionalCUDAGuard device_guard(device_of(k));
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
-  static const std::string kv_cache_dtype = "fp8_e4m3";
-  DISPATCH_BY_KV_CACHE_DTYPE(k.dtype(), kv_cache_dtype,
+  DISPATCH_BY_KV_CACHE_DTYPE(k.dtype(), "fp8_e4m3",
                              CALL_INDEXER_K_QUANT_AND_CACHE);
 }
 

csrc/cpu/cpu_attn.cpp

@@ -16,8 +16,6 @@ torch::Tensor get_scheduler_metadata(
     isa = cpu_attention::ISA::VEC16;
   } else if (isa_hint == "neon") {
     isa = cpu_attention::ISA::NEON;
-  } else if (isa_hint == "vxe") {
-    isa = cpu_attention::ISA::VXE;
   } else {
     TORCH_CHECK(false, "Unsupported CPU attention ISA hint: " + isa_hint);
   }
@@ -102,8 +100,6 @@ void cpu_attn_reshape_and_cache(
       return cpu_attention::ISA::VEC16;
     } else if (isa == "neon") {
       return cpu_attention::ISA::NEON;
-    } else if (isa == "vxe") {
-      return cpu_attention::ISA::VXE;
     } else {
       TORCH_CHECK(false, "Invalid ISA type: " + isa);
     }

csrc/cpu/cpu_attn_impl.hpp

@@ -12,7 +12,7 @@
 #include "cpu/utils.hpp"
 
 namespace cpu_attention {
-enum class ISA { AMX, VEC, VEC16, NEON, VXE };
+enum class ISA { AMX, VEC, VEC16, NEON };
 
 template <ISA isa, typename scalar_t, int64_t head_dim>
 class AttentionImpl {};

csrc/cpu/cpu_attn_vxe.hpp

@@ -1,386 +0,0 @@
-#ifndef CPU_ATTN_VXE_HPP
-#define CPU_ATTN_VXE_HPP
-
-#include "cpu_attn_impl.hpp"
-#include <vecintrin.h>
-#include <type_traits>
-
-namespace cpu_attention {
-
-namespace {
-
-// s390x Vector = 16 bytes (128 bits)
-#define BLOCK_SIZE_ALIGNMENT 32
-#define HEAD_SIZE_ALIGNMENT 32
-#define MAX_Q_HEAD_NUM_PER_ITER 16
-
-template <typename kv_cache_t>
-FORCE_INLINE void load_row8_B_as_f32(const kv_cache_t* p, __vector float& b0,
-                                     __vector float& b1);
-
-// [1] Float Specialization
-template <>
-FORCE_INLINE void load_row8_B_as_f32<float>(const float* p, __vector float& b0,
-                                            __vector float& b1) {
-  // Explicitly cast to long long for offset, and float* for pointer
-  b0 = vec_xl((long long)0, const_cast<float*>(p));
-  b1 = vec_xl((long long)0, const_cast<float*>(p + 4));
-}
-
-// [2] BFloat16 Specialization (Big Endian Fix)
-template <>
-FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p,
-                                                    __vector float& b0,
-                                                    __vector float& b1) {
-  // 1. Load 8 BF16s (16 bytes) into one vector
-  // Explicit cast to unsigned short* for vec_xl to return vector unsigned short
-  __vector unsigned short raw = vec_xl((long long)0, (unsigned short*)p);
-
-  // 2. Prepare Zero vector
-  __vector unsigned short zeros = vec_splat_u16(0);
-
-  // 3. Merge High/Low to expand BF16 -> Float32
-  // On Big Endian, a float is [BF16_bits | 16_zero_bits]
-  b0 = (__vector float)vec_mergeh(raw, zeros);
-  b1 = (__vector float)vec_mergel(raw, zeros);
-}
-
-template <>
-FORCE_INLINE void load_row8_B_as_f32<c10::Half>(const c10::Half* p,
-                                                __vector float& b0,
-                                                __vector float& b1) {
-  alignas(16) float tmp[8];
-
-  // Manual unroll / conversion
-  tmp[0] = static_cast<float>(p[0]);
-  tmp[1] = static_cast<float>(p[1]);
-  tmp[2] = static_cast<float>(p[2]);
-  tmp[3] = static_cast<float>(p[3]);
-  tmp[4] = static_cast<float>(p[4]);
-  tmp[5] = static_cast<float>(p[5]);
-  tmp[6] = static_cast<float>(p[6]);
-  tmp[7] = static_cast<float>(p[7]);
-
-  // Explicit arguments for intrinsic: (long long offset, float* ptr)
-  b0 = vec_xl((long long)0, (float*)tmp);
-  b1 = vec_xl((long long)0, (float*)(tmp + 4));
-}
-
-template <int32_t M, typename kv_cache_t>
-FORCE_INLINE void gemm_micro_s390x_Mx8_Ku4(
-    const float* __restrict A,       // [M x K]
-    const kv_cache_t* __restrict B,  // [K x 8]
-    float* __restrict C,             // [M x 8]
-    int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) {
-  static_assert(1 <= M && M <= 8, "M must be in [1,8]");
-
-// Helper macros to unroll codegen for M rows
-#define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
-#define IF_M(i) if constexpr (M > (i))
-
-  // 1. Define A pointers
-#define DECL_A(i) const float* a##i = A + (i) * lda;
-  ROWS_APPLY(DECL_A)
-#undef DECL_A
-
-  // 2. Define Accumulators (2 vectors covers 8 columns)
-#define DECL_ACC(i) __vector float acc##i##_0, acc##i##_1;
-  ROWS_APPLY(DECL_ACC)
-#undef DECL_ACC
-
-  // 3. Initialize Accumulators (Load C or Zero)
-#define INIT_ACC(i)                                                    \
-  IF_M(i) {                                                            \
-    if (accumulate) {                                                  \
-      acc##i##_0 =                                                     \
-          vec_xl((long long)0, const_cast<float*>(C + (i) * ldc + 0)); \
-      acc##i##_1 =                                                     \
-          vec_xl((long long)0, const_cast<float*>(C + (i) * ldc + 4)); \
-    } else {                                                           \
-      acc##i##_0 = vec_splats(0.0f);                                   \
-      acc##i##_1 = vec_splats(0.0f);                                   \
-    }                                                                  \
-  }
-  ROWS_APPLY(INIT_ACC)
-#undef INIT_ACC
-
-  int32_t k = 0;
-
-  for (; k + 3 < K; k += 4) {
-    // Load 4 values of A for each Row M: A[k...k+3]
-#define LOAD_A4(i)        \
-  __vector float a##i##v; \
-  IF_M(i) a##i##v = vec_xl((long long)0, const_cast<float*>(a##i + k));
-    ROWS_APPLY(LOAD_A4)
-#undef LOAD_A4
-
-    // Helper: FMA for specific lane L of A
-    // s390x: vec_madd(b, vec_splat(a, lane), acc)
-#define FMAS_LANE(i, aiv, L)                        \
-  IF_M(i) {                                         \
-    __vector float a_broad = vec_splat(aiv, L);     \
-    acc##i##_0 = vec_madd(b0, a_broad, acc##i##_0); \
-    acc##i##_1 = vec_madd(b1, a_broad, acc##i##_1); \
-  }
-
-    // Unroll K=0..3
-    {
-      __vector float b0, b1;
-      load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 0) * ldb, b0, b1);
-#define STEP_K0(i) FMAS_LANE(i, a##i##v, 0)
-      ROWS_APPLY(STEP_K0)
-#undef STEP_K0
-    }
-    {
-      __vector float b0, b1;
-      load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 1) * ldb, b0, b1);
-#define STEP_K1(i) FMAS_LANE(i, a##i##v, 1)
-      ROWS_APPLY(STEP_K1)
-#undef STEP_K1
-    }
-    {
-      __vector float b0, b1;
-      load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 2) * ldb, b0, b1);
-#define STEP_K2(i) FMAS_LANE(i, a##i##v, 2)
-      ROWS_APPLY(STEP_K2)
-#undef STEP_K2
-    }
-
-    {
-      __vector float b0, b1;
-      load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 3) * ldb, b0, b1);
-#define STEP_K3(i) FMAS_LANE(i, a##i##v, 3)
-      ROWS_APPLY(STEP_K3)
-#undef STEP_K3
-    }
-#undef FMAS_LANE
-  }
-
-  for (; k < K; ++k) {
-    __vector float b0, b1;
-    load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb, b0, b1);
-#define TAIL_ROW(i)                              \
-  IF_M(i) {                                      \
-    __vector float ai = vec_splats(*(a##i + k)); \
-    acc##i##_0 = vec_madd(b0, ai, acc##i##_0);   \
-    acc##i##_1 = vec_madd(b1, ai, acc##i##_1);   \
-  }
-    ROWS_APPLY(TAIL_ROW)
-#undef TAIL_ROW
-  }
-
-#define STORE_ROW(i)                           \
-  IF_M(i) {                                    \
-    vec_xst(acc##i##_0, 0, C + (i) * ldc + 0); \
-    vec_xst(acc##i##_1, 0, C + (i) * ldc + 4); \
-  }
-  ROWS_APPLY(STORE_ROW)
-#undef STORE_ROW
-
-#undef ROWS_APPLY
-#undef IF_M
-}
-
-template <int32_t N, typename kv_cache_t>
-FORCE_INLINE void gemm_macro_s390x_Mx8_Ku4(const float* __restrict A,
-                                           const kv_cache_t* __restrict B,
-                                           float* __restrict C, int32_t M,
-                                           int32_t K, int64_t lda, int64_t ldb,
-                                           int64_t ldc, bool accumulate) {
-  static_assert(N % 8 == 0, "N must be a multiple of 8");
-  for (int32_t m = 0; m < M;) {
-    int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1;
-    const float* Ab = A + m * lda;
-    float* Cb = C + m * ldc;
-
-    for (int32_t n = 0; n < N; n += 8) {
-      const kv_cache_t* Bn = B + n;
-      float* Cn = Cb + n;
-      switch (mb) {
-        case 8:
-          gemm_micro_s390x_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
-                                                  accumulate);
-          break;
-        case 4:
-          gemm_micro_s390x_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
-                                                  accumulate);
-          break;
-        case 2:
-          gemm_micro_s390x_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
-                                                  accumulate);
-          break;
-        default:
-          gemm_micro_s390x_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
-                                                  accumulate);
-          break;
-      }
-    }
-    m += mb;
-  }
-}
-
-template <typename kv_cache_t>
-class TileGemmS390X {
- public:
-  template <AttentionGemmPhase phase, int32_t k_size>
-  FORCE_INLINE static void gemm(const int32_t m_size,
-                                float* __restrict__ a_tile,
-                                kv_cache_t* __restrict__ b_tile,
-                                float* __restrict__ c_tile, const int64_t lda,
-                                const int64_t ldb, const int64_t ldc,
-                                const int32_t block_size,
-                                const int32_t dynamic_k_size,
-                                const bool accum_c) {
-    if constexpr (phase == AttentionGemmPhase::QK) {
-      gemm_macro_s390x_Mx8_Ku4<BLOCK_SIZE_ALIGNMENT, kv_cache_t>(
-          a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c);
-    } else {
-      gemm_macro_s390x_Mx8_Ku4<HEAD_SIZE_ALIGNMENT, kv_cache_t>(
-          a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc,
-          accum_c);
-    }
-  }
-};
-
-}  // namespace
-
-template <typename scalar_t, int64_t head_dim>
-class AttentionImpl<ISA::VXE, scalar_t, head_dim> {
- public:
-  using query_t = scalar_t;
-  using q_buffer_t = float;
-  using kv_cache_t = scalar_t;
-  using logits_buffer_t = float;
-  using partial_output_buffer_t = float;
-  using prob_buffer_t = float;
-
-  constexpr static int64_t BlockSizeAlignment = BLOCK_SIZE_ALIGNMENT;
-  constexpr static int64_t HeadDimAlignment = HEAD_SIZE_ALIGNMENT;
-  constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER;
-  constexpr static int64_t HeadDim = head_dim;
-  constexpr static ISA ISAType = ISA::VXE;
-  constexpr static bool scale_on_logits =
-      false;  // Scale is applied to Q during copy
-
- public:
-  AttentionImpl() {}
-
-  template <template <typename tile_gemm_t> typename attention>
-  FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
-    attention<TileGemmS390X<kv_cache_t>> attention_iteration;
-    attention_iteration(CPU_ATTENTION_PARAMS);
-  }
-
-  // Strides for Memory Layout
-  constexpr static int64_t k_cache_token_group_stride(
-      const int32_t block_size) {
-    return BlockSizeAlignment;  // [head_dim, block_size] layout
-  }
-
-  constexpr static int64_t v_cache_token_group_stride(
-      const int32_t block_size) {
-    return head_dim * BlockSizeAlignment;
-  }
-
-  constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
-    return HeadDimAlignment;
-  }
-
-  static void copy_q_heads_tile(scalar_t* __restrict__ src,
-                                float* __restrict__ q_buffer,
-                                const int32_t q_num,
-                                const int32_t q_heads_per_kv,
-                                const int64_t q_num_stride,
-                                const int64_t q_head_stride, float scale) {
-    __vector float scale_vec = vec_splats(scale);
-    constexpr bool is_bf16 = std::is_same<scalar_t, c10::BFloat16>::value;
-
-    // Process 8 elements at a time (32 bytes of float output)
-    for (int32_t i = 0; i < q_num; ++i) {
-      for (int32_t h = 0; h < q_heads_per_kv; ++h) {
-        scalar_t* curr_src = src + i * q_num_stride + h * q_head_stride;
-        float* curr_dst =
-            q_buffer + i * q_heads_per_kv * head_dim + h * head_dim;
-
-        int32_t d = 0;
-        for (; d <= head_dim - 8; d += 8) {
-          if constexpr (is_bf16) {
-            __vector float v0, v1;
-            // Reuse our Big-Endian-Safe loader
-            load_row8_B_as_f32<scalar_t>(curr_src + d, v0, v1);
-
-            v0 = vec_mul(v0, scale_vec);
-            v1 = vec_mul(v1, scale_vec);
-
-            vec_xst(v0, 0, curr_dst + d);
-            vec_xst(v1, 0, curr_dst + d + 4);
-          } else {
-            __vector float v0 = vec_xl((long long)0, (float*)curr_src + d);
-            __vector float v1 = vec_xl((long long)0, (float*)curr_src + d + 4);
-
-            v0 = vec_mul(v0, scale_vec);
-            v1 = vec_mul(v1, scale_vec);
-
-            vec_xst(v0, 0, curr_dst + d);
-            vec_xst(v1, 0, curr_dst + d + 4);
-          }
-        }
-
-        for (; d < head_dim; ++d) {
-          float val = static_cast<float>(curr_src[d]);
-          curr_dst[d] = val * scale;
-        }
-      }
-    }
-  }
-
-  static void reshape_and_cache(
-      const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
-      scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
-      const int64_t* __restrict__ slot_mapping, const int64_t token_num,
-      const int64_t key_token_num_stride, const int64_t value_token_num_stride,
-      const int64_t head_num, const int64_t key_head_num_stride,
-      const int64_t value_head_num_stride, const int64_t num_blocks,
-      const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
-      const int64_t block_size, const int64_t block_size_stride) {
-#pragma omp parallel for collapse(2)
-    for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
-      for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
-        const int64_t pos = slot_mapping[token_idx];
-        if (pos < 0) continue;
-
-        const int64_t block_idx = pos / block_size;
-        const int64_t block_offset = pos % block_size;
-
-        {
-          const scalar_t* key_src = key + token_idx * key_token_num_stride +
-                                    head_idx * key_head_num_stride;
-          scalar_t* key_dst = key_cache + block_idx * num_blocks_stride +
-                              head_idx * cache_head_num_stride + block_offset;
-
-          for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
-            key_dst[j] = key_src[i];
-          }
-        }
-
-        {
-          const scalar_t* val_src = value + token_idx * value_token_num_stride +
-                                    head_idx * value_head_num_stride;
-          scalar_t* val_dst = value_cache + block_idx * num_blocks_stride +
-                              head_idx * cache_head_num_stride +
-                              block_offset * head_dim;
-
-          std::memcpy(val_dst, val_src, sizeof(scalar_t) * head_dim);
-        }
-      }
-    }
-  }
-};
-
-}  // namespace cpu_attention
-
-#undef BLOCK_SIZE_ALIGNMENT
-#undef HEAD_SIZE_ALIGNMENT
-#undef MAX_Q_HEAD_NUM_PER_ITER
-
-#endif

csrc/cpu/generate_cpu_attn_dispatch.py

@@ -19,11 +19,10 @@ ISA_TYPES = {
     "VEC": 1,
     "VEC16": 2,
     "NEON": 3,
-    "VXE": 4,
 }
 
 # ISAs supported for head_dims divisible by 32
-ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE"]
+ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16"]
 
 # ISAs supported for head_dims divisible by 16 only
 ISA_FOR_16 = ["VEC16"]
@@ -119,10 +118,6 @@ def generate_header_file() -> str:
   #include "cpu_attn_neon.hpp"
 #endif
 
-#ifdef __s390x__
-  #include "cpu_attn_vxe.hpp"
-#endif
-
 """
 
     header += generate_helper_function()
@@ -168,25 +163,6 @@ def generate_header_file() -> str:
     } \\
   }()
 
-"""
-
-    # s390x with VXE
-    header += """#elif defined(__s390x__)
-#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
-  [&] { \\
-    int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
-    switch (encoded_params) { \\
-"""
-    header += generate_cases_for_isa_group(["VXE", "VEC", "VEC16"])
-    header += """
-      default: { \\
-        TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
-                    std::to_string(HEAD_DIM) + " isa=" + \\
-                    std::to_string(static_cast<int>(ISA_TYPE))); \\
-      } \\
-    } \\
-  }()
-
 """
 
     # Fallback: VEC and VEC16 only
@@ -206,7 +182,7 @@ def generate_header_file() -> str:
     } \\
   }()
 
-#endif  /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ */
+#endif  /* CPU_CAPABILITY_AMXBF16 / __aarch64__ */
 
 #endif  // CPU_ATTN_DISPATCH_GENERATED_H
 """

csrc/cpu/torch_bindings.cpp

@@ -4,10 +4,6 @@
 
 #include <torch/library.h>
 
-// Note: overwrite the external defination for sharing same name between
-// libraries use different ISAs.
-#define TORCH_EXTENSION_NAME _C
-
 std::string init_cpu_threads_env(const std::string& cpu_ids);
 
 void release_dnnl_matmul_handler(int64_t handler);
@@ -328,12 +324,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "str act, str isa) -> ()");
   ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
 #endif
-  ops.def("init_cpu_threads_env(str cpu_ids) -> str", &init_cpu_threads_env);
-  ops.def(
+}
+
+TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _utils), utils) {
+  // CPU utils
+  utils.def("init_cpu_threads_env(str cpu_ids) -> str", &init_cpu_threads_env);
+}
+
+TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cpu), cpu_ops) {
+  cpu_ops.def(
       "mla_decode_kvcache("
       "   Tensor! out, Tensor query, Tensor kv_cache,"
       "   float scale, Tensor block_tables, Tensor seq_lens) -> ()");
-  ops.impl("mla_decode_kvcache", torch::kCPU, &mla_decode_kvcache);
+  cpu_ops.impl("mla_decode_kvcache", torch::kCPU, &mla_decode_kvcache);
 }
 
 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

csrc/cuda_vec_utils.cuh

@@ -1,334 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-#pragma once
-
-#include <c10/util/BFloat16.h>
-#include <c10/util/Half.h>
-#include <cassert>
-
-#ifdef USE_ROCM
-  #include <hip/hip_runtime.h>
-#else
-  #include <cuda_bf16.h>
-  #include <cuda_fp16.h>
-  #include <cuda_runtime.h>
-#endif
-
-// Device-side: SM100+ architecture with CUDA 12.9+ toolkit, which
-// together enable 256-bit (v8.u32) PTX load/store instructions.
-// Use for PTX instruction selection with architecture fallback paths.
-#if !defined(USE_ROCM) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \
-    defined(CUDA_VERSION) && CUDA_VERSION >= 12090
-  #define VLLM_256B_PTX_ENABLED 1
-#else
-  #define VLLM_256B_PTX_ENABLED 0
-#endif
-
-namespace vllm {
-
-// ============================================================
-// Types and traits
-// ============================================================
-
-// 256-bit (32-byte) aligned vector type: 8 x uint32_t
-struct alignas(32) u32x8_t {
-  uint32_t d[8];
-};
-
-// VecTraits — select between 128-bit (int4) and 256-bit
-// (u32x8_t) vector types at compile time.
-template <bool support_256>
-struct VecTraits;
-
-template <>
-struct VecTraits<true> {
-  static constexpr int ARCH_MAX_VEC_SIZE = 32;
-  using vec_t = u32x8_t;
-};
-
-template <>
-struct VecTraits<false> {
-  static constexpr int ARCH_MAX_VEC_SIZE = 16;
-  using vec_t = int4;
-};
-
-// PackedTypeConverter — map between CUDA scalar and packed types
-//   half  <-> half2,  __nv_bfloat16 <-> __nv_bfloat162, etc.
-template <typename T>
-struct PackedTypeConverter {
-  static_assert(sizeof(T) == 0,
-                "PackedTypeConverter is not specialized for this type.");
-};
-
-template <>
-struct PackedTypeConverter<half2> {
-  using Type = half;
-};
-
-template <>
-struct PackedTypeConverter<half> {
-  using Type = half2;
-};
-
-template <>
-struct PackedTypeConverter<__nv_bfloat162> {
-  using Type = __nv_bfloat16;
-};
-
-template <>
-struct PackedTypeConverter<__nv_bfloat16> {
-  using Type = __nv_bfloat162;
-};
-
-template <>
-struct PackedTypeConverter<float> {
-  using Type = float2;
-};
-
-template <>
-struct PackedTypeConverter<float2> {
-  using Type = float;
-};
-
-template <>
-struct PackedTypeConverter<c10::Half> {
-  using Type = half2;
-};
-
-template <>
-struct PackedTypeConverter<c10::BFloat16> {
-  using Type = __nv_bfloat162;
-};
-
-// CUDATypeConverter — map PyTorch scalar types to CUDA scalar
-//   c10::Half -> half,  c10::BFloat16 -> __nv_bfloat16
-template <typename T>
-struct CUDATypeConverter {
-  using Type = T;
-};
-
-template <>
-struct CUDATypeConverter<c10::Half> {
-  using Type = half;
-};
-
-template <>
-struct CUDATypeConverter<c10::BFloat16> {
-  using Type = __nv_bfloat16;
-};
-
-// PackedVec — typed vector container for packed element access.
-//   Derives alignment and element count from VecTraits.
-//   Type is the CUDA scalar type (e.g. half, __nv_bfloat16).
-template <class Type, bool use_256b>
-struct alignas(VecTraits<use_256b>::ARCH_MAX_VEC_SIZE) PackedVec {
-  static constexpr int NUM_ELTS =
-      VecTraits<use_256b>::ARCH_MAX_VEC_SIZE /
-      sizeof(typename PackedTypeConverter<Type>::Type);
-  typename PackedTypeConverter<Type>::Type elts[NUM_ELTS];
-};
-
-// ============================================================
-// Load / store primitives
-// ============================================================
-
-// 256-bit load / store — SM100+ only (PTX v8 instructions).
-__device__ __forceinline__ void ld256(u32x8_t& val, const u32x8_t* ptr) {
-#if VLLM_256B_PTX_ENABLED
-  asm volatile("ld.global.nc.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];\n"
-               : "=r"(val.d[0]), "=r"(val.d[1]), "=r"(val.d[2]), "=r"(val.d[3]),
-                 "=r"(val.d[4]), "=r"(val.d[5]), "=r"(val.d[6]), "=r"(val.d[7])
-               : "l"(ptr));
-#else
-  assert(false && "ld256 requires SM100+ with CUDA 12.9+");
-#endif
-}
-
-__device__ __forceinline__ void st256(u32x8_t& val, u32x8_t* ptr) {
-#if VLLM_256B_PTX_ENABLED
-  asm volatile("st.global.v8.u32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};\n"
-               :
-               : "l"(ptr), "r"(val.d[0]), "r"(val.d[1]), "r"(val.d[2]),
-                 "r"(val.d[3]), "r"(val.d[4]), "r"(val.d[5]), "r"(val.d[6]),
-                 "r"(val.d[7])
-               : "memory");
-#else
-  assert(false && "st256 requires SM100+ with CUDA 12.9+");
-#endif
-}
-
-// Generic ld256 / st256 for any 32-byte aligned type (e.g. PackedVec).
-// Non-template overloads above are preferred for u32x8_t.
-template <typename T>
-__device__ __forceinline__ void ld256(T& val, const T* ptr) {
-  static_assert(sizeof(T) == 32, "ld256 requires a 32-byte type");
-  ld256(reinterpret_cast<u32x8_t&>(val), reinterpret_cast<const u32x8_t*>(ptr));
-}
-
-template <typename T>
-__device__ __forceinline__ void st256(T& val, T* ptr) {
-  static_assert(sizeof(T) == 32, "st256 requires a 32-byte type");
-  st256(reinterpret_cast<u32x8_t&>(val), reinterpret_cast<u32x8_t*>(ptr));
-}
-
-// 128-bit load / store via __ldg (read-only cache hint).
-template <typename T>
-__device__ __forceinline__ void ld128(T& val, const T* ptr) {
-  static_assert(sizeof(T) == 16, "ld128 requires a 16-byte type");
-  *reinterpret_cast<int4*>(&val) = __ldg(reinterpret_cast<const int4*>(ptr));
-}
-
-template <typename T>
-__device__ __forceinline__ void st128(T& val, T* ptr) {
-  static_assert(sizeof(T) == 16, "st128 requires a 16-byte type");
-  *reinterpret_cast<int4*>(ptr) = *reinterpret_cast<int4*>(&val);
-}
-
-// 256-bit cache-streaming (.cs) load / store  — SM100+ only.
-__forceinline__ __device__ u32x8_t ld256_cs(const u32x8_t* addr) {
-#if VLLM_256B_PTX_ENABLED
-  u32x8_t val;
-  asm volatile("ld.global.cs.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];"
-               : "=r"(val.d[0]), "=r"(val.d[1]), "=r"(val.d[2]), "=r"(val.d[3]),
-                 "=r"(val.d[4]), "=r"(val.d[5]), "=r"(val.d[6]), "=r"(val.d[7])
-               : "l"(addr));
-  return val;
-#else
-  assert(false && "ld256_cs requires SM100+ with CUDA 12.9+");
-  return {};
-#endif
-}
-
-__forceinline__ __device__ void st256_cs(u32x8_t* addr, u32x8_t val) {
-#if VLLM_256B_PTX_ENABLED
-  asm volatile(
-      "st.global.cs.v8.u32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" ::"l"(addr),
-      "r"(val.d[0]), "r"(val.d[1]), "r"(val.d[2]), "r"(val.d[3]), "r"(val.d[4]),
-      "r"(val.d[5]), "r"(val.d[6]), "r"(val.d[7]));
-#else
-  assert(false && "st256_cs requires SM100+ with CUDA 12.9+");
-#endif
-}
-
-// 32-bit cache-streaming (.cs) load / store  — SM100+ only.
-__forceinline__ __device__ int ld32_cs(const int* addr) {
-#if VLLM_256B_PTX_ENABLED
-  int val;
-  asm volatile("ld.global.cs.b32 %0, [%1];" : "=r"(val) : "l"(addr));
-  return val;
-#else
-  assert(false && "ld32_cs requires SM100+ with CUDA 12.9+");
-  return 0;
-#endif
-}
-
-__forceinline__ __device__ void st32_cs(int* addr, int val) {
-#if VLLM_256B_PTX_ENABLED
-  asm volatile("st.global.cs.b32 [%0], %1;" ::"l"(addr), "r"(val));
-#else
-  assert(false && "st32_cs requires SM100+ with CUDA 12.9+");
-#endif
-}
-
-// Predicated 256-bit / 128-bit cache-global (.cg) loads.
-// Returns zero if pred is false.  SM100+ only.
-__device__ __forceinline__ void ld256_cg_or_zero(u32x8_t& val, const void* ptr,
-                                                 bool pred) {
-#if VLLM_256B_PTX_ENABLED
-  asm volatile(
-      "{\n"
-      "  .reg .pred pr;\n"
-      "  setp.ne.u32 pr, %8, 0;\n"
-      "  mov.u32 %0, 0;\n"
-      "  mov.u32 %1, 0;\n"
-      "  mov.u32 %2, 0;\n"
-      "  mov.u32 %3, 0;\n"
-      "  mov.u32 %4, 0;\n"
-      "  mov.u32 %5, 0;\n"
-      "  mov.u32 %6, 0;\n"
-      "  mov.u32 %7, 0;\n"
-      "  @pr ld.global.cg.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%9];\n"
-      "}\n"
-      : "=r"(val.d[0]), "=r"(val.d[1]), "=r"(val.d[2]), "=r"(val.d[3]),
-        "=r"(val.d[4]), "=r"(val.d[5]), "=r"(val.d[6]), "=r"(val.d[7])
-      : "r"((int)pred), "l"(ptr));
-#else
-  assert(false && "ld256_cg_or_zero requires SM100+ with CUDA 12.9+");
-#endif
-}
-
-__device__ __forceinline__ void ld128_cg_or_zero(uint4& val, const void* ptr,
-                                                 bool pred) {
-#if VLLM_256B_PTX_ENABLED
-  uint32_t r0, r1, r2, r3;
-
-  asm volatile(
-      "{\n"
-      "  .reg .pred pr;\n"
-      "  setp.ne.u32 pr, %4, 0;\n"
-      "  mov.u32 %0, 0;\n"
-      "  mov.u32 %1, 0;\n"
-      "  mov.u32 %2, 0;\n"
-      "  mov.u32 %3, 0;\n"
-      "  @pr ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%5];\n"
-      "}\n"
-      : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
-      : "r"((int)pred), "l"(ptr));
-
-  val = uint4{r0, r1, r2, r3};
-#else
-  assert(false && "ld128_cg_or_zero requires SM100+ with CUDA 12.9+");
-#endif
-}
-
-// ============================================================
-// Alignment helpers
-// ============================================================
-
-__host__ __device__ __forceinline__ bool is_16byte_aligned(const void* ptr) {
-  return (reinterpret_cast<uintptr_t>(ptr) & 15) == 0;
-}
-
-__host__ __device__ __forceinline__ bool is_32byte_aligned(const void* ptr) {
-  return (reinterpret_cast<uintptr_t>(ptr) & 31) == 0;
-}
-
-// ============================================================
-// Packed type conversion and arithmetic
-// ============================================================
-
-template <typename packed_t>
-__device__ __forceinline__ float2 cast_to_float2(const packed_t& val) {
-  if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
-    return __bfloat1622float2(val);
-  } else if constexpr (std::is_same_v<packed_t, __half2>) {
-    return __half22float2(val);
-  } else if constexpr (std::is_same_v<packed_t, float2>) {
-    return float2(val);
-  }
-}
-
-template <typename packed_t>
-__device__ __forceinline__ packed_t cast_to_packed(const float2& val) {
-  if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
-    return __float22bfloat162_rn(val);
-  } else if constexpr (std::is_same_v<packed_t, __half2>) {
-    return __float22half2_rn(val);
-  } else if constexpr (std::is_same_v<packed_t, float2>) {
-    return float2(val);
-  }
-}
-
-template <typename packed_t>
-__device__ __forceinline__ packed_t packed_mul(const packed_t& x,
-                                               const packed_t& y) {
-  if constexpr (std::is_same_v<packed_t, __nv_bfloat162> ||
-                std::is_same_v<packed_t, __half2>) {
-    return __hmul2(x, y);
-  } else if constexpr (std::is_same_v<packed_t, float2>) {
-    return make_float2(x.x * y.x, x.y * y.y);
-  }
-}
-
-}  // namespace vllm

csrc/dsv3_fused_a_gemm.cu

@@ -1,751 +0,0 @@
-/*
- * Adapted from
- * https://github.com/sgl-project/sglang/blob/main/sgl-kernel/csrc/gemm/dsv3_fused_a_gemm.cu
- * which was adapted from
- * https://github.com/NVIDIA/TensorRT-LLM/blob/619709fc33bd5dc268f19d6a741fe7ed51c0f8f5/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3FusedAGemm.cu
- *
- * Copyright (c) 2019-2024, NVIDIA CORPORATION.  All rights reserved.
- * Copyright (c) 2021, NAVER Corp.  Authored by CLOVA.
- *
- * 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 <ATen/ATen.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-#include <torch/all.h>
-
-#include "core/registration.h"
-
-#include <cstdlib>
-#include <mutex>
-
-namespace {
-
-inline int getSMVersion() {
-  auto* props = at::cuda::getCurrentDeviceProperties();
-  return props->major * 10 + props->minor;
-}
-
-inline bool getEnvEnablePDL() {
-  static std::once_flag flag;
-  static bool enablePDL = false;
-  std::call_once(flag, [&]() {
-    if (getSMVersion() >= 90) {
-      char const* env = std::getenv("TRTLLM_ENABLE_PDL");
-      enablePDL = env && env[0] == '1' && env[1] == '\0';
-    }
-  });
-  return enablePDL;
-}
-
-}  // namespace
-
-using bf16_t = __nv_bfloat16;
-
-__device__ void hmma_16_8_16_f32acc_bf16ab(float (&d_reg)[4],
-                                           const bf16_t (&a_reg)[8],
-                                           const bf16_t (&b_reg)[4],
-                                           float const (&c_reg)[4]) {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
-  uint32_t a0 = *reinterpret_cast<uint32_t const*>(a_reg + 0);
-  uint32_t a1 = *reinterpret_cast<uint32_t const*>(a_reg + 2);
-  uint32_t a2 = *reinterpret_cast<uint32_t const*>(a_reg + 4);
-  uint32_t a3 = *reinterpret_cast<uint32_t const*>(a_reg + 6);
-  uint32_t b0 = *reinterpret_cast<uint32_t const*>(b_reg + 0);
-  uint32_t b1 = *reinterpret_cast<uint32_t const*>(b_reg + 2);
-  asm volatile(
-      "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
-      "{%0,  %1,  %2,  %3},"
-      "{%4,  %5,  %6,  %7},"
-      "{%8,  %9},"
-      "{%10, %11, %12, %13};\n"
-      : "=f"(d_reg[0]), "=f"(d_reg[1]), "=f"(d_reg[2]), "=f"(d_reg[3])
-      : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "f"(d_reg[0]),
-        "f"(d_reg[1]), "f"(d_reg[2]), "f"(d_reg[3]));
-#endif
-}
-
-extern "C" {
-__device__ uint32_t __nvvm_get_smem_pointer(void*);
-}
-
-__device__ void ldgsts_128(void const* gPtr, void* sPtr, uint32_t pred) {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
-  if (pred) {
-    uint32_t smemPtrAsUint32 = __nvvm_get_smem_pointer(sPtr);
-    asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2;\n" ::"r"(
-                     smemPtrAsUint32),
-                 "l"(gPtr), "n"(16));
-  }
-#endif
-}
-
-__device__ void ldsm_x4(void* smem_ptr, uint32_t* reg_ptr) {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
-  asm volatile(
-      "ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n"
-      : "=r"(reg_ptr[0]), "=r"(reg_ptr[1]), "=r"(reg_ptr[2]), "=r"(reg_ptr[3])
-      : "r"(__nvvm_get_smem_pointer(smem_ptr)));
-#endif
-}
-
-template <class Type>
-__device__ int apply_swizzle_343_on_elem_row_col(int row_idx_, int col_idx_) {
-  uint32_t row_idx = *reinterpret_cast<uint32_t*>(&row_idx_);
-  uint32_t col_idx = *reinterpret_cast<uint32_t*>(&col_idx_);
-  row_idx = row_idx % 8;
-  row_idx = row_idx * (16 / sizeof(Type));
-  col_idx = col_idx ^ row_idx;
-  return *reinterpret_cast<int*>(&col_idx);
-}
-
-__device__ void initialize_barrier(
-    uint64_t* smem_barrier,  // 64 bits user-manged barrier in smem
-    int thread_count =
-        1)  // Thread count expected to arrive/wait on this barrier
-{
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
-  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
-  asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(smem_int_ptr),
-               "r"(thread_count));
-#endif
-}
-
-// Barrier wait
-__device__ void wait_barrier(
-    uint64_t* smem_barrier,  // 64 bits user-manged barrier in smem
-    int phase_bit)           // Current phase bit the barrier waiting to flip
-{
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
-  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
-  asm volatile(
-      "{\n"
-      ".reg .pred                P1;\n"
-      "LAB_WAIT:\n"
-      "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
-      "@P1                       bra DONE;\n"
-      "bra                   LAB_WAIT;\n"
-      "DONE:\n"
-      "}\n" ::"r"(smem_int_ptr),
-      "r"(phase_bit));
-#endif
-}
-
-__device__ bool try_wait_barrier(uint64_t* smem_ptr, int phase_bit) {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
-  uint32_t wait_complete;
-  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_ptr);
-  asm volatile(
-      "{\n\t"
-      ".reg .pred P1; \n\t"
-      "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2; \n\t"
-      "selp.b32 %0, 1, 0, P1; \n\t"
-      "}"
-      : "=r"(wait_complete)
-      : "r"(smem_int_ptr), "r"(phase_bit));
-  return static_cast<bool>(wait_complete);
-#endif
-  return false;
-}
-
-// Barrier arrive
-__device__ void arrive_barrier(
-    uint64_t* smem_barrier)  // 64 bits user-manged barrier in smem
-{
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
-  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
-  asm volatile(
-      "{\n"
-      ".reg .b64 state; \n"
-      "mbarrier.arrive.shared::cta.b64   state, [%0];\n"
-      "}\n" ::"r"(smem_int_ptr));
-#endif
-}
-
-__device__ void ldgsts_arrive(uint64_t* smem_barrier) {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
-  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
-  asm volatile("cp.async.mbarrier.arrive.noinc.shared.b64 [%0];"
-               :
-               : "r"(smem_int_ptr));
-#endif
-}
-
-template <int gemm_k, int tile_m, int tile_k, int stage_cnt>
-struct GmemLoaderA {
-  static constexpr int elem_bytes = 2;
-  static constexpr int vec_bytes = 16;
-  static constexpr int vec_elems = vec_bytes / elem_bytes;
-  static constexpr int thread_cnt = 64;
-  static_assert((tile_m * tile_k) % (vec_elems * thread_cnt) == 0);
-  static constexpr int a_inst_cnt_per_iter =
-      (tile_m * tile_k) / (vec_elems * thread_cnt);
-  static_assert(gemm_k % tile_k == 0);
-  static constexpr int k_iter_cnt = gemm_k / tile_k;
-
-  // Extra params to keep the order of k reduction...
-  static constexpr int mma_warp_cnt = 4;
-  static constexpr int per_mma_warp_k = tile_k / mma_warp_cnt;
-  static constexpr int k_each_chunk = gemm_k / mma_warp_cnt;
-
- private:
-  __device__ int k_project(int tile_k_idx) {
-    return (tile_k_idx / per_mma_warp_k * k_each_chunk) +
-           (tile_k_idx % per_mma_warp_k);
-  }
-
- public:
-  __device__ GmemLoaderA(bf16_t const* gmem_a_local_, bf16_t* smem_a_,
-                         uint64_t* smem_barrier_)
-      : gmem_a(gmem_a_local_),
-        smem_a(smem_a_),
-        smem_barrier(smem_barrier_),
-        local_tid(threadIdx.x % thread_cnt) {}
-
-  __device__ void prepare() {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
-  // swizzle, that's what we want.
-  #pragma unroll
-    for (int i = 0; i < a_inst_cnt_per_iter; i++) {
-      int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
-      int m_idx = linear_idx / tile_k;
-      int k_idx = linear_idx % tile_k;
-      k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
-      a_smem_offsets[i] = m_idx * tile_k + k_idx;
-    }
-#endif
-  }
-
-  __device__ void issue_mainloop() {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
-  #pragma unroll 1
-    for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
-      if (need_wait) {
-        wait_barrier(smem_barrier + 1 + stage_idx * 2, phase_bit);
-      }
-      int next_stage_idx = stage_idx + 1;
-      int next_phase_bit =
-          next_stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
-      next_stage_idx = next_stage_idx == stage_cnt ? 0 : next_stage_idx;
-      if (loop_idx != k_iter_cnt - 1) {
-        need_wait = !try_wait_barrier(smem_barrier + 1 + next_stage_idx * 2,
-                                      next_phase_bit);
-      }
-
-  #pragma unroll
-      for (int i = 0; i < a_inst_cnt_per_iter; i++) {
-        int smem_offset = a_smem_offsets[i];
-        bf16_t* smem_ptr_this_iter =
-            smem_a + stage_idx * tile_m * tile_k + smem_offset;
-        int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
-        int m_idx = linear_idx / tile_k;
-        int k_idx = linear_idx % tile_k;
-        int gmem_offset = m_idx * gemm_k + k_project(k_idx);
-        bf16_t const* gmem_ptr_this_iter = gmem_a + gmem_offset;
-        ldgsts_128(gmem_ptr_this_iter, smem_ptr_this_iter, true);
-      }
-      ldgsts_arrive(smem_barrier + stage_idx * 2);
-
-      stage_idx = next_stage_idx;
-      phase_bit = next_phase_bit;
-      gmem_a += per_mma_warp_k;
-    }
-#endif
-  }
-
-  bf16_t const* gmem_a;
-  bf16_t* smem_a;
-  uint64_t* smem_barrier;
-  int local_tid;
-  int stage_idx = 0;
-  int phase_bit = 1;
-  bool need_wait = true;
-
-  // per smem_stage, store with swizzle information
-  int a_smem_offsets[a_inst_cnt_per_iter];
-};
-
-template <int gemm_k, int tile_n, int tile_k, int stage_cnt>
-struct GmemLoaderB {
-  static constexpr int elem_bytes = 2;
-  static constexpr int vec_bytes = 16;
-  static constexpr int vec_elems = vec_bytes / elem_bytes;
-  static constexpr int thread_cnt = 64;
-  static_assert((tile_n * tile_k) % (vec_elems * thread_cnt) == 0);
-  static constexpr int b_inst_cnt_per_iter =
-      (tile_n * tile_k) / (vec_elems * thread_cnt);
-  static_assert(gemm_k % tile_k == 0);
-  static constexpr int k_iter_cnt = gemm_k / tile_k;
-
-  // Extra params to keep the order of k reduction...
-  static constexpr int mma_warp_cnt = 4;
-  static constexpr int per_mma_warp_k = tile_k / mma_warp_cnt;
-  static constexpr int k_each_chunk = gemm_k / mma_warp_cnt;
-
- private:
-  __device__ int k_project(int tile_k_idx) {
-    return (tile_k_idx / per_mma_warp_k * k_each_chunk) +
-           (tile_k_idx % per_mma_warp_k);
-  }
-
- public:
-  __device__ GmemLoaderB(bf16_t const* gmem_b_local_, bf16_t* smem_b_,
-                         uint64_t* smem_barrier_, int gemm_n_)
-      : gmem_b(gmem_b_local_),
-        smem_b(smem_b_),
-        smem_barrier(smem_barrier_),
-        gemm_n(gemm_n_),
-        local_tid(threadIdx.x % thread_cnt) {}
-
-  __device__ void prepare() {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
-  // swizzle, that's what we want.
-  #pragma unroll
-    for (int i = 0; i < b_inst_cnt_per_iter; i++) {
-      int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
-      int n_idx = linear_idx / tile_k;
-      int k_idx = linear_idx % tile_k;
-      k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(n_idx, k_idx);
-      b_smem_offsets[i] = n_idx * tile_k + k_idx;
-      preds[i] = n_idx < gemm_n;
-    }
-#endif
-  }
-
-  __device__ void issue_mainloop() {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
-    asm volatile("griddepcontrol.wait;");
-  #pragma unroll 1
-    for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
-      if (need_wait) {
-        wait_barrier(smem_barrier + 1 + stage_idx * 2, phase_bit);
-      }
-      int next_stage_idx = stage_idx + 1;
-      int next_phase_bit =
-          next_stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
-      next_stage_idx = next_stage_idx == stage_cnt ? 0 : next_stage_idx;
-      if (loop_idx != k_iter_cnt - 1) {
-        need_wait = !try_wait_barrier(smem_barrier + 1 + next_stage_idx * 2,
-                                      next_phase_bit);
-      }
-  #pragma unroll
-      for (int i = 0; i < b_inst_cnt_per_iter; i++) {
-        int smem_offset = b_smem_offsets[i];
-        bf16_t* smem_ptr_this_iter =
-            smem_b + stage_idx * tile_n * tile_k + smem_offset;
-        int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
-        int n_idx = linear_idx / tile_k;
-        int k_idx = linear_idx % tile_k;
-        int gmem_offset = n_idx * gemm_k + k_project(k_idx);
-        bf16_t const* gmem_ptr_this_iter = gmem_b + gmem_offset;
-        ldgsts_128(gmem_ptr_this_iter, smem_ptr_this_iter, preds[i]);
-      }
-      ldgsts_arrive(smem_barrier + stage_idx * 2);
-
-      stage_idx = next_stage_idx;
-      phase_bit = next_phase_bit;
-      gmem_b += per_mma_warp_k;
-    }
-#endif
-  }
-
-  bf16_t const* gmem_b;
-  bf16_t* smem_b;
-  uint64_t* smem_barrier;
-  int gemm_n;
-  int local_tid;
-  int stage_idx = 0;
-  int phase_bit = 1;
-  bool need_wait = true;
-
-  // per smem_stage, store with swizzle information
-  int b_smem_offsets[b_inst_cnt_per_iter];
-  uint32_t preds[b_inst_cnt_per_iter];
-};
-
-template <int gemm_m, int gemm_k, int tile_m, int tile_n, int tile_k,
-          int stage_cnt>
-struct MmaComputer {
-  static constexpr int elem_bytes = 2;
-  static constexpr int thread_cnt = 128;
-  static_assert(gemm_k % tile_k == 0);
-  static_assert(tile_k % (thread_cnt / 32) == 0);
-  static constexpr int per_warp_tile_k = tile_k / (thread_cnt / 32);
-  static constexpr int k_iter_cnt = gemm_k / tile_k;
-  static constexpr int k_phase_cnt = per_warp_tile_k / 16;
-  static constexpr int m_iter_cnt = (tile_m + 15) / 16;
-  static constexpr int n_iter_cnt =
-      (tile_n + 7) /
-      8;  // Possible to have non-1 n_iter_cnt for ab_swap m16 case.
-  static_assert(m_iter_cnt == 1);
-  static_assert(n_iter_cnt == 1 || n_iter_cnt == 2);
-
-  __device__ MmaComputer(bf16_t* gmem_c_local_, bf16_t* smem_a_,
-                         bf16_t* smem_b_, uint64_t* smem_barrier_,
-                         int warp_idx_, int gemm_n_)
-      : gmem_c(gmem_c_local_),
-        smem_a(smem_a_),
-        smem_b(smem_b_),
-        smem_barrier(smem_barrier_),
-        warp_idx(warp_idx_ - (thread_cnt / 32)),
-        gemm_n(gemm_n_) {}
-
- private:
-  __device__ constexpr int internal_b_atom_func(int tid) {
-    if constexpr (tile_n < 8) {
-      return (tid % tile_n) + ((tid % 8) / tile_n * 0) + tid / 8 * 8 * tile_n;
-    } else {
-      return (tid % 8) + ((tid % 32) / 8 * (tile_n * 8));
-    }
-  }
-
- public:
-  __device__ void prepare() {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
-  #pragma unroll
-    for (int i = 0; i < k_phase_cnt; i++) {
-      int linear_idx = (lane_idx % 16) + (lane_idx / 16) * 128 + i * 256;
-      int m_idx = linear_idx % tile_m;
-      int k_idx = linear_idx / tile_m + warp_k_offset_in_tile_k;
-      k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
-      a_smem_offsets[0][i] = m_idx * tile_k + k_idx;
-    }
-  #pragma unroll
-    for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
-  #pragma unroll
-      for (int i = 0; i < k_phase_cnt; i += 2) {  // Special i+=2 for B.
-        int linear_idx =
-            internal_b_atom_func(lane_idx) + i * tile_n * 16 + n_iter_idx * 8;
-        int n_idx = linear_idx % tile_n;
-        int k_idx = linear_idx / tile_n + warp_k_offset_in_tile_k;
-        k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(n_idx, k_idx);
-        b_smem_offsets[n_iter_idx][i] = n_idx * tile_k + k_idx;
-      }
-    }
-#endif
-  }
-
-  __device__ void issue_mainloop() {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
-  #pragma unroll 1
-    for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
-      wait_barrier(smem_barrier + 0 + stage_idx * 2, phase_bit);
-
-  #pragma unroll
-      for (int i = 0; i < k_phase_cnt; i++) {
-        int smem_offset = a_smem_offsets[0][i];
-        bf16_t* smem_ptr_this_iter =
-            smem_a + stage_idx * tile_m * tile_k + smem_offset;
-        ldsm_x4(smem_ptr_this_iter, reinterpret_cast<uint32_t*>(a_reg[0][i]));
-      }
-
-  #pragma unroll
-      for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
-  #pragma unroll
-        for (int i = 0; i < k_phase_cnt; i += 2) {
-          int smem_offset = b_smem_offsets[n_iter_idx][i];
-          bf16_t* smem_ptr_this_iter =
-              smem_b + stage_idx * tile_n * tile_k + smem_offset;
-          ldsm_x4(smem_ptr_this_iter,
-                  reinterpret_cast<uint32_t*>(b_reg[n_iter_idx][i]));
-        }
-      }
-
-  #pragma unroll
-      for (int k_iter_idx = 0; k_iter_idx < k_phase_cnt; k_iter_idx++) {
-  #pragma unroll
-        for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
-          hmma_16_8_16_f32acc_bf16ab(
-              acc_reg[0][n_iter_idx], a_reg[0][k_iter_idx],
-              b_reg[n_iter_idx][k_iter_idx], acc_reg[0][n_iter_idx]);
-        }
-      }
-      ::arrive_barrier(smem_barrier + 1 + stage_idx * 2);
-      stage_idx += 1;
-      phase_bit = stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
-      stage_idx = stage_idx == stage_cnt ? 0 : stage_idx;
-    }
-#endif
-  }
-
-  __device__ void epi() {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
-    asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt));
-    // reorganize the acc_reg
-    constexpr int thread_m = 2;
-    constexpr int thread_n = 2 * n_iter_cnt;
-    constexpr int cta_mma_n = n_iter_cnt * 8;
-    float acc_reg_reorg[thread_m][thread_n];
-
-    for (int i = 0; i < thread_m; i++) {
-      for (int j = 0; j < thread_n; j++) {
-        acc_reg_reorg[i][j] = acc_reg[0][j / 2][(j % 2) + (i * 2)];
-      }
-    }
-
-    // 4 x cosize(smem_c_layout)
-    float* smem_c = reinterpret_cast<float*>(smem_a);
-    // coord -> index
-    auto smem_c_index_func = [&](int m_idx, int n_idx) {
-      int group_rows = 32 / cta_mma_n;
-      int group_cnt = 2;
-      return (m_idx % group_rows * cta_mma_n) +
-             (m_idx / group_rows * (32 + group_cnt)) + n_idx;
-    };
-    constexpr int cosize_smem_c = ((tile_m * cta_mma_n) / 32) * (32 + 2);
-
-  // This should be optimized to STS.64 but can not be STS.128 due to the bank
-  // index.
-  #pragma unroll
-    for (int m_idx_thread = 0; m_idx_thread < thread_m; m_idx_thread++) {
-  #pragma unroll
-      for (int n_idx_thread = 0; n_idx_thread < thread_n; n_idx_thread++) {
-        int m_idx = (lane_idx / 4) + m_idx_thread * 8;
-        int n_idx =
-            ((lane_idx % 4) * 2) + (n_idx_thread % 2) + (n_idx_thread / 2) * 8;
-        smem_c[cosize_smem_c * warp_idx + smem_c_index_func(m_idx, n_idx)] =
-            acc_reg_reorg[m_idx_thread][n_idx_thread];
-      }
-    }
-    asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt));
-
-    if (warp_idx == 0) {
-      constexpr int final_acc_reg_cnt = (tile_m * tile_n + 31) / 32;
-      float acc_final[final_acc_reg_cnt]{};
-
-  #pragma unroll
-      for (int reg_idx = 0; reg_idx < final_acc_reg_cnt; reg_idx++) {
-        int linear_idx = reg_idx * 32 + lane_idx;
-        int m_idx = linear_idx % tile_m;
-        int n_idx = linear_idx / tile_m;
-        acc_final[reg_idx] +=
-            smem_c[smem_c_index_func(m_idx, n_idx) + 0 * cosize_smem_c] +
-            smem_c[smem_c_index_func(m_idx, n_idx) + 1 * cosize_smem_c] +
-            smem_c[smem_c_index_func(m_idx, n_idx) + 2 * cosize_smem_c] +
-            smem_c[smem_c_index_func(m_idx, n_idx) + 3 * cosize_smem_c];
-      }
-
-  #pragma unroll
-      for (int reg_idx = 0; reg_idx < final_acc_reg_cnt; reg_idx++) {
-        int linear_idx = reg_idx * 32 + lane_idx;
-        int m_idx = linear_idx % tile_m;
-        int n_idx = linear_idx / tile_m;
-        if (m_idx < tile_m && n_idx < gemm_n) {
-          gmem_c[n_idx * gemm_m + m_idx] = acc_final[reg_idx];
-        }
-      }
-    }
-#endif
-  }
-
-  bf16_t* gmem_c;
-  bf16_t* smem_a;
-  bf16_t* smem_b;
-  uint64_t* smem_barrier;
-  int warp_idx;
-  int gemm_n;
-  int stage_idx = 0;
-  int phase_bit = 0;
-  int lane_idx = threadIdx.x % 32;
-  int warp_k_offset_in_tile_k = warp_idx * per_warp_tile_k;
-
-  int a_smem_offsets[m_iter_cnt][k_phase_cnt];
-  int b_smem_offsets[n_iter_cnt][k_phase_cnt];
-
-  bf16_t a_reg[m_iter_cnt][k_phase_cnt][8];
-  bf16_t b_reg[n_iter_cnt][k_phase_cnt][4];
-  float acc_reg[m_iter_cnt][n_iter_cnt][4]{};
-};
-
-// AB swapped, kernel is k-major, k-major, m-major
-template <int batch_size, int gemm_m, int gemm_k, int tile_m, int tile_n,
-          int tile_k, int stage_cnt>
-__global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
-    bf16_t* output, bf16_t const* mat_a, bf16_t const* mat_b, int gemm_n) {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
-  constexpr int load_thread_cnt = 128;
-  constexpr int compute_thread_cnt = 128;
-  constexpr int thread_cnt = load_thread_cnt + compute_thread_cnt;
-  (void)thread_cnt;
-  static_assert(gemm_m % 16 == 0);
-  static_assert(gemm_k % tile_k == 0);
-  static_assert(gemm_m % tile_m == 0);
-  static_assert(
-      tile_k == 128 || tile_k == 256 || tile_k == 512 ||
-      tile_k == 1024);  // tile_k must be larger than 64 since 4 warp splitK.
-  static_assert(tile_m == 16);
-  constexpr int g2s_vec_bytes = 16;
-  constexpr int a_elem_bytes = 2;
-  constexpr int b_elem_bytes = 2;
-  static_assert((tile_m * a_elem_bytes + tile_n * b_elem_bytes) * tile_k *
-                    stage_cnt <=
-                225 * 1024);
-  static_assert((tile_m * tile_k * a_elem_bytes) %
-                    (load_thread_cnt * g2s_vec_bytes) ==
-                0);
-  static_assert((tile_n * tile_k * b_elem_bytes) %
-                    (load_thread_cnt * g2s_vec_bytes) ==
-                0);
-
-  extern __shared__ char smem[];
-  uint64_t* smem_barrier = reinterpret_cast<uint64_t*>(
-      smem);  // producer,consumer; producer,consumer; ...
-  bf16_t* smem_a = reinterpret_cast<bf16_t*>(smem + (stage_cnt * 8 * 2 + 1024) /
-                                                        1024 * 1024);
-  bf16_t* smem_b = smem_a + tile_m * tile_k * stage_cnt;
-
-  int cta_m_idx = tile_m * blockIdx.x;
-  int cta_n_idx = tile_n * blockIdx.y;
-  bf16_t const* gmem_a_local = mat_a + cta_m_idx * gemm_k;
-  bf16_t const* gmem_b_local = mat_b + cta_n_idx * gemm_k;
-  bf16_t* gmem_c_local = output + cta_n_idx * gemm_m + cta_m_idx;
-
-  int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
-
-  if (warp_idx == 4) {
-    for (int i = 0; i < stage_cnt; i++) {
-      initialize_barrier(smem_barrier + i * 2 + 0,
-                         load_thread_cnt);  // producer
-      initialize_barrier(smem_barrier + i * 2 + 1,
-                         compute_thread_cnt);  // consumer
-    }
-  }
-  __syncthreads();
-
-  if (warp_idx < 2) {
-    GmemLoaderA<gemm_k, tile_m, tile_k, stage_cnt> a_loader(
-        gmem_a_local, smem_a, smem_barrier);
-    a_loader.prepare();
-    a_loader.issue_mainloop();
-  } else if (warp_idx < 4) {
-    GmemLoaderB<gemm_k, tile_n, tile_k, stage_cnt> b_loader(
-        gmem_b_local, smem_b, smem_barrier, gemm_n);
-    b_loader.prepare();
-    b_loader.issue_mainloop();
-  } else {
-    MmaComputer<gemm_m, gemm_k, tile_m, tile_n, tile_k, stage_cnt> mma_computer(
-        gmem_c_local, smem_a, smem_b, smem_barrier, warp_idx, gemm_n);
-    mma_computer.prepare();
-    mma_computer.issue_mainloop();
-    mma_computer.epi();
-  }
-  asm volatile("griddepcontrol.launch_dependents;");
-#endif
-}
-
-template <typename T, int kHdIn, int kHdOut, int kTileN>
-void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
-                      cudaStream_t const stream) {
-  constexpr int gemm_m = kHdOut;  // 2112
-  int const gemm_n = num_tokens;  // 1-16
-  constexpr int gemm_k = kHdIn;   // 7168
-  constexpr int batch_size = 1;
-  std::swap(mat_a, mat_b);
-  constexpr int tile_m = 16;
-  constexpr int tile_n = kTileN;                        // 8 or 16
-  constexpr int tile_k = std::max(256, 1024 / tile_n);  // 256
-  constexpr int max_stage_cnt =
-      1024 * 192 / ((tile_m + tile_n) * tile_k * sizeof(bf16_t));
-  constexpr int k_iter_cnt = gemm_k / tile_k;
-  constexpr int stage_cnt =
-      k_iter_cnt > max_stage_cnt ? max_stage_cnt : k_iter_cnt;
-  int cta_m_cnt = gemm_m / tile_m;
-  int cta_n_cnt = (gemm_n + tile_n - 1) / tile_n;
-  constexpr int barrier_bytes = (stage_cnt * 16 + 1023) / 1024 * 1024;
-  constexpr int smem_bytes =
-      ((tile_m * 2 + tile_n * 2) * tile_k * stage_cnt + barrier_bytes + 1023) /
-      1024 * 1024;
-
-  dim3 grid(cta_m_cnt, cta_n_cnt, 1);
-  dim3 block_size(256);
-  cudaLaunchConfig_t config;
-  config.gridDim = grid;
-  config.blockDim = block_size;
-  config.dynamicSmemBytes = smem_bytes;
-  config.stream = stream;
-  cudaLaunchAttribute attrs[1];
-  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-  attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
-  config.numAttrs = 1;
-  config.attrs = attrs;
-  if (smem_bytes >= (48 * 1024)) {
-    cudaFuncSetAttribute(fused_a_gemm_kernel<batch_size, gemm_m, gemm_k, tile_m,
-                                             tile_n, tile_k, stage_cnt>,
-                         cudaFuncAttributeMaxDynamicSharedMemorySize,
-                         smem_bytes);
-  }
-  cudaLaunchKernelEx(&config,
-                     fused_a_gemm_kernel<batch_size, gemm_m, gemm_k, tile_m,
-                                         tile_n, tile_k, stage_cnt>,
-                     output, mat_a, mat_b, gemm_n);
-}
-
-template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 8>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
-    cudaStream_t);
-
-template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
-    cudaStream_t);
-
-void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
-                       torch::Tensor const& mat_b) {
-  TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2);
-  int const num_tokens = mat_a.size(0);
-  int const hd_in = mat_a.size(1);
-  int const hd_out = mat_b.size(1);
-
-  constexpr int kHdIn = 7168;
-  constexpr int kHdOut = 2112;
-  TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
-              "required 1 <= mat_a.shape[0] <= 16")
-  TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168")
-  TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112")
-  TORCH_CHECK(output.size(0) == num_tokens,
-              "required output.shape[0] == mat_a.shape[0]")
-  TORCH_CHECK(output.size(1) == hd_out,
-              "required output.shape[1] == mat_b.shape[1]")
-
-  TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor");
-  TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor");
-  TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor");
-
-  TORCH_CHECK(mat_a.scalar_type() == torch::kBFloat16 &&
-                  mat_b.scalar_type() == torch::kBFloat16,
-              "Only BFloat16 input dtype is supported")
-  TORCH_CHECK(output.scalar_type() == torch::kBFloat16,
-              "Only BFloat16 output dtype is supported")
-
-  TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
-
-  auto stream = at::cuda::getCurrentCUDAStream(mat_a.get_device());
-  if (num_tokens <= 8) {
-    invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>(
-        reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
-        reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-        reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
-        stream);
-  } else {
-    invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 16>(
-        reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
-        reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-        reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
-        stream);
-  }
-}
-
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-  m.impl("dsv3_fused_a_gemm", &dsv3_fused_a_gemm);
-}

csrc/mamba/mamba_ssm/selective_scan.h

@@ -15,9 +15,9 @@
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 struct SSMParamsBase {
-    using index_t = size_t;
+    using index_t = uint32_t;
 
-    int batch, dim, seqlen, dstate, n_groups;
+    int batch, dim, seqlen, dstate, n_groups, n_chunks;
     int dim_ngroups_ratio;
     bool is_variable_B;
     bool is_variable_C;
@@ -72,8 +72,6 @@ struct SSMParamsBase {
     void *__restrict__ block_idx_first_scheduled_token_ptr;  // (batch,) - first block to write
     void *__restrict__ block_idx_last_scheduled_token_ptr;   // (batch,) - last block to write
     void *__restrict__ initial_state_idx_ptr;  // (batch,) - index of the initial state to use
-    void *__restrict__ cu_chunk_seqlen_ptr;      // (nchunks+1,) - cumulative chunk token offsets
-    void *__restrict__ last_chunk_indices_ptr;   // (batch,) - index of last chunk per sequence
 };
 
 

csrc/mamba/mamba_ssm/selective_scan_fwd.cu

@@ -81,6 +81,7 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
     constexpr bool kIsVariableC = Ktraits::kIsVariableC;
     constexpr bool kHasZ = Ktraits::kHasZ;
     constexpr bool kVarlen = Ktraits::kVarlen;
+    constexpr int kNThreads = Ktraits::kNThreads;
     constexpr int kNItems = Ktraits::kNItems;
     constexpr int kNRows = Ktraits::kNRows;
     constexpr bool kDirectIO = Ktraits::kDirectIO;
@@ -160,8 +161,17 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
         }
     }
 
+
+    // for (int state_idx = threadIdx.x; state_idx < params.dstate; state_idx += blockDim.x) {
+    //     smem_a[state_idx] = A[state_idx * params.A_dstate_stride];
+    //     smem_bc[state_idx] = B[state_idx * params.B_dstate_stride] * C[state_idx * params.C_dstate_stride];
+    // }
+
+    constexpr int kChunkSize = kNThreads * kNItems;
+
     // Use block_size for chunking when APC is enabled, otherwise use 2048 for backwards compatibility
-    const int block_size = params.cache_enabled ? params.block_size : 2048;
+    const int iteration_chunk_size = params.cache_enabled ? params.block_size : 2048;
+    const int n_chunks = (seqlen + iteration_chunk_size - 1) / iteration_chunk_size;
 
     const int* batch_cache_indices = cache_indices != nullptr ?
                                      cache_indices + batch_id * params.cache_indices_stride : nullptr;
@@ -171,44 +181,10 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
                                           reinterpret_cast<const int*>(params.block_idx_last_scheduled_token_ptr) : nullptr;
     const int* initial_state_idx = params.initial_state_idx_ptr != nullptr ?
                                    reinterpret_cast<const int*>(params.initial_state_idx_ptr) : nullptr;
-    const int* cu_chunk_seqlen = params.cu_chunk_seqlen_ptr != nullptr ?
-                                 reinterpret_cast<const int*>(params.cu_chunk_seqlen_ptr) : nullptr;
-    const int* last_chunk_indices = params.last_chunk_indices_ptr != nullptr ?
-                                    reinterpret_cast<const int*>(params.last_chunk_indices_ptr) : nullptr;
 
     const size_t load_cache_slot = params.cache_enabled && batch_cache_indices != nullptr ? batch_cache_indices[initial_state_idx[batch_id]] : cache_index;
 
-    const int block_idx_first = (params.cache_enabled && block_idx_first_scheduled != nullptr) ?
-                                 block_idx_first_scheduled[batch_id] : 0;
-
-    // Determine chunk boundaries from pre-computed metadata (APC mode)
-    // or fall back to simple block_size chunking.
-    int first_chunk_idx, n_chunks;
-    int current_position;
-
-    if (cu_chunk_seqlen != nullptr && last_chunk_indices != nullptr) {
-        const int last_chunk_idx = last_chunk_indices[batch_id];
-        first_chunk_idx = (batch_id == 0) ? 0 : last_chunk_indices[batch_id - 1] + 1;
-        n_chunks = last_chunk_idx - first_chunk_idx + 1;
-        // Derive current_position: if the first chunk is partial (fills remainder
-        // of a started block), offset into the block accordingly.
-        const int first_chunk_tokens = cu_chunk_seqlen[first_chunk_idx + 1] - cu_chunk_seqlen[first_chunk_idx];
-        const int chunk_start_offset = (n_chunks > 1 && first_chunk_tokens < block_size)
-                                        ? (block_size - first_chunk_tokens) : 0;
-        current_position = block_idx_first * block_size + chunk_start_offset;
-    } else {
-        first_chunk_idx = 0;
-        n_chunks = (seqlen + block_size - 1) / block_size;
-        current_position = 0;
-    }
-
-    int tokens_processed = 0;
-
     for (int chunk = 0; chunk < n_chunks; ++chunk) {
-        const int chunk_tokens = (cu_chunk_seqlen != nullptr)
-            ? cu_chunk_seqlen[first_chunk_idx + chunk + 1] - cu_chunk_seqlen[first_chunk_idx + chunk]
-            : min(block_size, seqlen - tokens_processed);
-        if (chunk_tokens <= 0) break;
         input_t u_vals[kNRows][kNItems], delta_vals_load[kNRows][kNItems];
 
         __syncthreads();
@@ -217,12 +193,12 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
             if constexpr (!kDirectIO) {
                 if (r > 0) { __syncthreads(); }
             }
-            load_input<Ktraits>(u + r * params.u_d_stride, u_vals[r], smem_load, chunk_tokens);
+            load_input<Ktraits>(u + r * params.u_d_stride, u_vals[r], smem_load, seqlen - chunk * kChunkSize);
             if constexpr (!kDirectIO) { __syncthreads(); }
-            load_input<Ktraits>(delta + r * params.delta_d_stride, delta_vals_load[r], smem_load, chunk_tokens);
+            load_input<Ktraits>(delta + r * params.delta_d_stride, delta_vals_load[r], smem_load, seqlen - chunk * kChunkSize);
         }
-        u += chunk_tokens;
-        delta += chunk_tokens;
+        u += kChunkSize;
+        delta += kChunkSize;
     
         float delta_vals[kNRows][kNItems], delta_u_vals[kNRows][kNItems], out_vals[kNRows][kNItems];
         #pragma unroll
@@ -256,7 +232,7 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
             weight_t B_vals[kNItems], C_vals[kNItems];
             if constexpr (kIsVariableB) {
                 load_weight<Ktraits>(Bvar + state_idx * params.B_dstate_stride, B_vals,
-                    smem_load_weight, chunk_tokens);
+                    smem_load_weight, (seqlen - chunk * kChunkSize) * (1));
                 if constexpr (!kIsVariableC) {
                     #pragma unroll
                     for (int r = 0; r < kNRows; ++r) {
@@ -267,7 +243,7 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
             if constexpr (kIsVariableC) {
                 auto &smem_load_weight_C = !kIsVariableB ? smem_load_weight : smem_load_weight1;
                 load_weight<Ktraits>(Cvar + state_idx * params.C_dstate_stride, C_vals,
-                    smem_load_weight_C, chunk_tokens);
+                    smem_load_weight_C, (seqlen - chunk * kChunkSize) * (1));
                 if constexpr (!kIsVariableB) {
                     #pragma unroll
                     for (int r = 0; r < kNRows; ++r) {
@@ -290,8 +266,10 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
                 for (int i = 0; i < kNItems; ++i) {
                     thread_data[i] = make_float2(exp2f(delta_vals[r][i] * A_val[r]),
                                                  !kIsVariableB ? delta_u_vals[r][i] : B_vals[i] * delta_u_vals[r][i]);
-                    if (threadIdx.x * kNItems + i >= chunk_tokens) {
-                        thread_data[i] = make_float2(1.f, 0.f);
+                    if (seqlen % (kNItems * kNThreads) != 0) {  // So that the last state is correct
+                        if (threadIdx.x * kNItems + i >= seqlen - chunk * kChunkSize) {
+                            thread_data[i] = make_float2(1.f, 0.f);
+                        }
                     }
                 }
                 // Initialize running total
@@ -323,14 +301,14 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
                 if (threadIdx.x == 0) {
                     smem_running_prefix[state_idx + r * MAX_DSTATE] = prefix_op.running_prefix;
 
-                    // Store state at the end of each aligned chunk when cache is enabled
+                    // Store state at the end of each chunk when cache is enabled
                     if (params.cache_enabled && batch_cache_indices != nullptr) {
+
                         size_t cache_slot;
                         if (chunk == n_chunks - 1) {
                             cache_slot = batch_cache_indices[block_idx_last_scheduled[batch_id]];
                         } else {
-                            const int block_idx_completed = (current_position + chunk_tokens - 1) / block_size;
-                            cache_slot = batch_cache_indices[block_idx_completed];
+                            cache_slot = batch_cache_indices[block_idx_first_scheduled[batch_id] + chunk];
                         }
 
                         size_t state_offset = cache_slot * params.ssm_states_batch_stride +
@@ -353,41 +331,38 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
             }
         }
         input_t *out = reinterpret_cast<input_t *>(params.out_ptr) + sequence_start_index * params.out_batch_stride
-            + dim_id * kNRows * params.out_d_stride + tokens_processed;
+            + dim_id * kNRows * params.out_d_stride + chunk * kChunkSize;
         __syncthreads();
         #pragma unroll
         for (int r = 0; r < kNRows; ++r) {
             if constexpr (!kDirectIO) {
                 if (r > 0) { __syncthreads(); }
             }
-            store_output<Ktraits>(out + r * params.out_d_stride, out_vals[r], smem_store, chunk_tokens);
+            store_output<Ktraits>(out + r * params.out_d_stride, out_vals[r], smem_store, seqlen - chunk * kChunkSize);
         }
 
         if constexpr (kHasZ) {
             input_t *z = reinterpret_cast<input_t *>(params.z_ptr) + sequence_start_index * params.z_batch_stride
-                + dim_id * kNRows * params.z_d_stride + tokens_processed;
+                + dim_id * kNRows * params.z_d_stride + chunk * kChunkSize;
             input_t *out_z = reinterpret_cast<input_t *>(params.out_z_ptr) + sequence_start_index * params.out_z_batch_stride
-                + dim_id * kNRows * params.out_z_d_stride + tokens_processed;
+                + dim_id * kNRows * params.out_z_d_stride + chunk * kChunkSize;
             #pragma unroll
             for (int r = 0; r < kNRows; ++r) {
                 input_t z_vals[kNItems];
                 __syncthreads();
-                load_input<Ktraits>(z + r * params.z_d_stride, z_vals, smem_load, chunk_tokens);
+                load_input<Ktraits>(z + r * params.z_d_stride, z_vals, smem_load, seqlen - chunk * kChunkSize);
                 #pragma unroll
                 for (int i = 0; i < kNItems; ++i) {
                     float z_val = z_vals[i];
                     out_vals[r][i] *= z_val / (1 + expf(-z_val));
                 }
                 __syncthreads();
-                store_output<Ktraits>(out_z + r * params.out_z_d_stride, out_vals[r], smem_store, chunk_tokens);
+                store_output<Ktraits>(out_z + r * params.out_z_d_stride, out_vals[r], smem_store, seqlen - chunk * kChunkSize);
             }
         }
 
-        Bvar += chunk_tokens;
-        Cvar += chunk_tokens;
-
-        tokens_processed += chunk_tokens;
-        current_position += chunk_tokens;
+        Bvar += kChunkSize * 1;
+        Cvar += kChunkSize * 1;
     }
 }
 
@@ -531,9 +506,7 @@ void set_ssm_params_fwd(SSMParamsBase &params,
                         int64_t block_size,
                         const std::optional<torch::Tensor> &block_idx_first_scheduled_token,
                         const std::optional<torch::Tensor> &block_idx_last_scheduled_token,
-                        const std::optional<torch::Tensor> &initial_state_idx,
-                        const std::optional<torch::Tensor> &cu_chunk_seqlen,
-                        const std::optional<torch::Tensor> &last_chunk_indices) {
+                        const std::optional<torch::Tensor> &initial_state_idx) {
 
     // Reset the parameters
     memset(&params, 0, sizeof(params));
@@ -575,8 +548,6 @@ void set_ssm_params_fwd(SSMParamsBase &params,
     params.block_idx_first_scheduled_token_ptr = block_idx_first_scheduled_token.has_value() ? block_idx_first_scheduled_token.value().data_ptr() : nullptr;
     params.block_idx_last_scheduled_token_ptr = block_idx_last_scheduled_token.has_value() ? block_idx_last_scheduled_token.value().data_ptr() : nullptr;
     params.initial_state_idx_ptr = initial_state_idx.has_value() ? initial_state_idx.value().data_ptr() : nullptr;
-    params.cu_chunk_seqlen_ptr = cu_chunk_seqlen.has_value() ? cu_chunk_seqlen.value().data_ptr() : nullptr;
-    params.last_chunk_indices_ptr = last_chunk_indices.has_value() ? last_chunk_indices.value().data_ptr() : nullptr;
 
     // All stride are in elements, not bytes.
     params.A_d_stride = A.stride(0);
@@ -662,9 +633,7 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
                   int64_t block_size,
                   const std::optional<torch::Tensor> &block_idx_first_scheduled_token,
                   const std::optional<torch::Tensor> &block_idx_last_scheduled_token,
-                  const std::optional<torch::Tensor> &initial_state_idx,
-                  const std::optional<torch::Tensor> &cu_chunk_seqlen,
-                  const std::optional<torch::Tensor> &last_chunk_indices) {
+                  const std::optional<torch::Tensor> &initial_state_idx) {
     auto input_type = u.scalar_type();
     auto weight_type = A.scalar_type();
     TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
@@ -809,9 +778,7 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
                        block_size,
                        block_idx_first_scheduled_token,
                        block_idx_last_scheduled_token,
-                       initial_state_idx,
-                       cu_chunk_seqlen,
-                       last_chunk_indices
+                       initial_state_idx
                        );
 
     

csrc/moe/dsv3_router_gemm_bf16_out.cu

@@ -1,291 +0,0 @@
-/*
- * Adapted from SGLang's sgl-kernel implementation, which was adapted from
- * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu
- * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp
- *
- * Copyright (c) 2019-2023, NVIDIA CORPORATION.  All rights reserved.
- *
- * 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 <ATen/ATen.h>
-#include <ATen/cuda/CUDAContext.h>
-
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-
-#include "dsv3_router_gemm_utils.h"
-
-// Custom FMA implementation using PTX assembly instructions
-__device__ __forceinline__ void fma(float2& d, float2 const& a, float2 const& b,
-                                    float2 const& c) {
-  asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n"
-               : "=l"(reinterpret_cast<uint64_t&>(d))
-               : "l"(reinterpret_cast<uint64_t const&>(a)),
-                 "l"(reinterpret_cast<uint64_t const&>(b)),
-                 "l"(reinterpret_cast<uint64_t const&>(c)));
-}
-
-// Convert 8 bfloat16 values from a uint4 to float array - optimized conversion
-template <int VPT>
-__device__ __forceinline__ void bf16_uint4_to_float8(uint4 const& vec,
-                                                     float* dst) {
-  __nv_bfloat16* bf16_ptr =
-      reinterpret_cast<__nv_bfloat16*>(const_cast<uint4*>(&vec));
-
-#pragma unroll
-  for (int i = 0; i < VPT; i++) {
-    dst[i] = __bfloat162float(bf16_ptr[i]);
-  }
-}
-
-template <typename T, int kBlockSize, int VPT, int kNumTokens, int kNumExperts,
-          int kHiddenDim>
-__global__ __launch_bounds__(128, 1) void router_gemm_kernel_bf16_output(
-    __nv_bfloat16* out, T const* mat_a, T const* mat_b) {
-  // Each block handles one expert column
-  int const n_idx = blockIdx.x;
-  int const tid = threadIdx.x;
-  constexpr int kWarpSize = 32;
-  constexpr int kNumWarps = kBlockSize / kWarpSize;
-  // Constants for this kernel
-  constexpr int k_elems_per_k_iteration = VPT * kBlockSize;
-  constexpr int k_iterations =
-      kHiddenDim / k_elems_per_k_iteration;  // Total K iterations
-
-  // Initialize accumulators for all M rows
-  float acc[kNumTokens] = {};
-
-  // Shared memory for warp-level reduction
-  __shared__ float sm_reduction[kNumTokens][kNumWarps];  // kNumWarps
-
-  // B matrix is in column-major order, so we can directly load a column for the
-  // n_idx expert
-  T const* b_col = mat_b + n_idx * kHiddenDim;
-
-  // Pre-compute k_base values for each iteration to help compiler optimize
-  int k_bases[k_iterations];
-#pragma unroll
-  for (int ki = 0; ki < k_iterations; ki++) {
-    k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT;
-  }
-
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.wait;");
-#endif
-
-  // Process the GEMM in chunks
-  for (int ki = 0; ki < k_iterations; ki++) {
-    int const k_base = k_bases[ki];
-
-    // Load B matrix values using vector load (8 bf16 values)
-    uint4 b_vec = *reinterpret_cast<uint4 const*>(b_col + k_base);
-
-    // Convert B values to float
-    float b_float[VPT];
-    bf16_uint4_to_float8<VPT>(b_vec, b_float);
-
-// Process each token
-#pragma unroll
-    for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
-      // Load both rows of A matrix using vector loads
-      uint4 a_vec = *reinterpret_cast<uint4 const*>(
-          mat_a + (m_idx * kHiddenDim) + k_base);
-
-      // Convert A values to float
-      float a_float[VPT];
-      bf16_uint4_to_float8<VPT>(a_vec, a_float);
-
-// Process elements in this chunk
-#pragma unroll
-      for (int k = 0; k < VPT; k++) {
-        float a = a_float[k];
-        float b = b_float[k];
-        acc[m_idx] += a * b;
-      }
-    }
-  }
-
-  // Perform warp-level reduction
-  int const warpSize = 32;
-  int const warpId = tid / warpSize;
-  int const laneId = tid % warpSize;
-
-  // Register for warp-level reduction results
-  float warp_result[kNumTokens];
-
-#pragma unroll
-  for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
-    warp_result[m_idx] = acc[m_idx];
-  }
-
-// Perform warp-level reduction using optimized butterfly pattern
-#pragma unroll
-  for (int m = 0; m < kNumTokens; m++) {
-    float sum = warp_result[m];
-
-    // Butterfly reduction pattern
-    sum += __shfl_xor_sync(0xffffffff, sum, 16);
-    sum += __shfl_xor_sync(0xffffffff, sum, 8);
-    sum += __shfl_xor_sync(0xffffffff, sum, 4);
-    sum += __shfl_xor_sync(0xffffffff, sum, 2);
-    sum += __shfl_xor_sync(0xffffffff, sum, 1);
-
-    // Only the first thread in each warp stores to shared memory
-    if (laneId == 0) {
-      sm_reduction[m][warpId] = sum;
-    }
-  }
-
-  __syncthreads();
-
-  // Final reduction across warps (only first thread)
-  if (tid == 0) {
-#pragma unroll
-    for (int m = 0; m < kNumTokens; m++) {
-      float final_sum = 0.0f;
-
-// Sum across the kNumWarps
-#pragma unroll
-      for (int w = 0; w < kNumWarps; w++) {
-        final_sum += sm_reduction[m][w];
-      }
-
-      // Write final result
-      out[m * kNumExperts + n_idx] = __float2bfloat16(final_sum);
-    }
-  }
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.launch_dependents;");
-#endif
-}
-
-template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
-void invokeRouterGemmBf16Output(__nv_bfloat16* output, T const* mat_a,
-                                T const* mat_b, cudaStream_t stream) {
-  constexpr int VPT = 16 / sizeof(T);
-  constexpr int kBlockSize = 128;
-  cudaLaunchConfig_t config;
-  config.gridDim = kNumExperts;
-  config.blockDim = kBlockSize;
-  config.dynamicSmemBytes = 0;
-  config.stream = stream;
-  cudaLaunchAttribute attrs[1];
-  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-  attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
-  config.numAttrs = 1;
-  config.attrs = attrs;
-  cudaLaunchKernelEx(
-      &config,
-      router_gemm_kernel_bf16_output<T, kBlockSize, VPT, kNumTokens,
-                                     kNumExperts, kHiddenDim>,
-      output, mat_a, mat_b);
-}
-
-// Template instantiations for DEFAULT_NUM_EXPERTS experts
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 1, 256, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 2, 256, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 3, 256, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 4, 256, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 5, 256, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 6, 256, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 7, 256, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 8, 256, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 9, 256, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 10, 256, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 11, 256, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 12, 256, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 13, 256, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 14, 256, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 256, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 256, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-// Template instantiations for KIMI_K2_NUM_EXPERTS experts
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 1, 384, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 2, 384, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 3, 384, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 4, 384, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 5, 384, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 6, 384, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 7, 384, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 8, 384, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 9, 384, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 10, 384, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 11, 384, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 12, 384, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 13, 384, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 14, 384, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 384, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 384, 7168>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);

csrc/moe/dsv3_router_gemm_entry.cu

@@ -1,169 +0,0 @@
-/*
- * Adapted from SGLang's sgl-kernel implementation, which was adapted from
- * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu
- * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp
- *
- * Copyright (c) 2019-2023, NVIDIA CORPORATION.  All rights reserved.
- *
- * 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 <ATen/ATen.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <torch/all.h>
-
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-
-#include "core/registration.h"
-#include "dsv3_router_gemm_utils.h"
-
-static constexpr int DEFAULT_NUM_EXPERTS = 256;
-static constexpr int KIMI_K2_NUM_EXPERTS = 384;
-static constexpr int DEFAULT_HIDDEN_DIM = 7168;
-
-template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
-void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b,
-                                 cudaStream_t stream);
-
-template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
-void invokeRouterGemmBf16Output(__nv_bfloat16* output, T const* mat_a,
-                                T const* mat_b, cudaStream_t stream);
-
-template <int kBegin, int kEnd, int kNumExperts, int kHiddenDim>
-struct LoopUnroller {
-  static void unroll_float_output(int num_tokens, float* output,
-                                  __nv_bfloat16 const* input,
-                                  __nv_bfloat16 const* weights,
-                                  cudaStream_t stream) {
-    if (num_tokens == kBegin) {
-      invokeRouterGemmFloatOutput<__nv_bfloat16, kBegin, kNumExperts,
-                                  kHiddenDim>(output, input, weights, stream);
-    } else {
-      LoopUnroller<kBegin + 1, kEnd, kNumExperts,
-                   kHiddenDim>::unroll_float_output(num_tokens, output, input,
-                                                    weights, stream);
-    }
-  }
-
-  static void unroll_bf16_output(int num_tokens, __nv_bfloat16* output,
-                                 __nv_bfloat16 const* input,
-                                 __nv_bfloat16 const* weights,
-                                 cudaStream_t stream) {
-    if (num_tokens == kBegin) {
-      invokeRouterGemmBf16Output<__nv_bfloat16, kBegin, kNumExperts,
-                                 kHiddenDim>(output, input, weights, stream);
-    } else {
-      LoopUnroller<kBegin + 1, kEnd, kNumExperts,
-                   kHiddenDim>::unroll_bf16_output(num_tokens, output, input,
-                                                   weights, stream);
-    }
-  }
-};
-
-template <int kEnd, int kNumExperts, int kHiddenDim>
-struct LoopUnroller<kEnd, kEnd, kNumExperts, kHiddenDim> {
-  static void unroll_float_output(int num_tokens, float* output,
-                                  __nv_bfloat16 const* input,
-                                  __nv_bfloat16 const* weights,
-                                  cudaStream_t stream) {
-    if (num_tokens == kEnd) {
-      invokeRouterGemmFloatOutput<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>(
-          output, input, weights, stream);
-    } else {
-      throw std::invalid_argument("Invalid num_tokens, only supports 1 to 16");
-    }
-  }
-
-  static void unroll_bf16_output(int num_tokens, __nv_bfloat16* output,
-                                 __nv_bfloat16 const* input,
-                                 __nv_bfloat16 const* weights,
-                                 cudaStream_t stream) {
-    if (num_tokens == kEnd) {
-      invokeRouterGemmBf16Output<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>(
-          output, input, weights, stream);
-    } else {
-      throw std::invalid_argument("Invalid num_tokens, only supports 1 to 16");
-    }
-  }
-};
-
-void dsv3_router_gemm(at::Tensor& output,       // [num_tokens, num_experts]
-                      const at::Tensor& mat_a,  // [num_tokens, hidden_dim]
-                      const at::Tensor& mat_b   // [num_experts, hidden_dim]
-) {
-  TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2);
-
-  const int num_tokens = mat_a.size(0);
-  const int num_experts = mat_b.size(0);
-  const int hidden_dim = mat_a.size(1);
-
-  TORCH_CHECK(mat_a.size(1) == mat_b.size(1),
-              "mat_a and mat_b must have the same hidden_dim");
-  TORCH_CHECK(hidden_dim == DEFAULT_HIDDEN_DIM,
-              "Expected hidden_dim=", DEFAULT_HIDDEN_DIM,
-              ", but got hidden_dim=", hidden_dim);
-  TORCH_CHECK(
-      num_experts == DEFAULT_NUM_EXPERTS || num_experts == KIMI_K2_NUM_EXPERTS,
-      "Expected num_experts=", DEFAULT_NUM_EXPERTS,
-      " or num_experts=", KIMI_K2_NUM_EXPERTS,
-      ", but got num_experts=", num_experts);
-  TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
-              "currently num_tokens must be less than or equal to 16 for "
-              "router_gemm");
-  TORCH_CHECK(mat_a.dtype() == at::kBFloat16, "mat_a must be bf16");
-  TORCH_CHECK(mat_b.dtype() == at::kBFloat16, "mat_b must be bf16");
-  TORCH_CHECK(output.dtype() == at::kFloat || output.dtype() == at::kBFloat16,
-              "output must be float32 or bf16");
-
-  auto const sm = getSMVersion();
-  TORCH_CHECK(sm >= 90 && sm <= 103, "required SM_103 >= CUDA ARCH >= SM_90");
-
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-
-  if (output.dtype() == at::kFloat) {
-    if (num_experts == DEFAULT_NUM_EXPERTS) {
-      LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
-          unroll_float_output(
-              num_tokens, reinterpret_cast<float*>(output.mutable_data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
-    } else if (num_experts == KIMI_K2_NUM_EXPERTS) {
-      LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
-          unroll_float_output(
-              num_tokens, reinterpret_cast<float*>(output.mutable_data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
-    }
-  } else if (output.dtype() == at::kBFloat16) {
-    if (num_experts == DEFAULT_NUM_EXPERTS) {
-      LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
-          unroll_bf16_output(
-              num_tokens,
-              reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
-    } else if (num_experts == KIMI_K2_NUM_EXPERTS) {
-      LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
-          unroll_bf16_output(
-              num_tokens,
-              reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
-    }
-  }
-}
-
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-  m.impl("dsv3_router_gemm", &dsv3_router_gemm);
-}

csrc/moe/dsv3_router_gemm_float_out.cu

@@ -1,291 +0,0 @@
-/*
- * Adapted from SGLang's sgl-kernel implementation, which was adapted from
- * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu
- * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp
- *
- * Copyright (c) 2019-2023, NVIDIA CORPORATION.  All rights reserved.
- *
- * 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 <ATen/ATen.h>
-#include <ATen/cuda/CUDAContext.h>
-
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-
-#include "dsv3_router_gemm_utils.h"
-
-// Custom FMA implementation using PTX assembly instructions
-__device__ __forceinline__ void fma(float2& d, float2 const& a, float2 const& b,
-                                    float2 const& c) {
-  asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n"
-               : "=l"(reinterpret_cast<uint64_t&>(d))
-               : "l"(reinterpret_cast<uint64_t const&>(a)),
-                 "l"(reinterpret_cast<uint64_t const&>(b)),
-                 "l"(reinterpret_cast<uint64_t const&>(c)));
-}
-
-// Convert 8 bfloat16 values from a uint4 to float array - optimized conversion
-template <int VPT>
-__device__ __forceinline__ void bf16_uint4_to_float8(uint4 const& vec,
-                                                     float* dst) {
-  __nv_bfloat16* bf16_ptr =
-      reinterpret_cast<__nv_bfloat16*>(const_cast<uint4*>(&vec));
-
-#pragma unroll
-  for (int i = 0; i < VPT; i++) {
-    dst[i] = __bfloat162float(bf16_ptr[i]);
-  }
-}
-
-template <typename T, int kBlockSize, int VPT, int kNumTokens, int kNumExperts,
-          int kHiddenDim>
-__global__ __launch_bounds__(128, 1) void router_gemm_kernel_float_output(
-    float* out, T const* mat_a, T const* mat_b) {
-  // Each block handles one expert column
-  int const n_idx = blockIdx.x;
-  int const tid = threadIdx.x;
-  constexpr int kWarpSize = 32;
-  constexpr int kNumWarps = kBlockSize / kWarpSize;
-  // Constants for this kernel
-  constexpr int k_elems_per_k_iteration = VPT * kBlockSize;
-  constexpr int k_iterations =
-      kHiddenDim / k_elems_per_k_iteration;  // Total K iterations
-
-  // Initialize accumulators for all M rows
-  float acc[kNumTokens] = {};
-
-  // Shared memory for warp-level reduction
-  __shared__ float sm_reduction[kNumTokens][kNumWarps];  // kNumWarps
-
-  // B matrix is in column-major order, so we can directly load a column for the
-  // n_idx expert
-  T const* b_col = mat_b + n_idx * kHiddenDim;
-
-  // Pre-compute k_base values for each iteration to help compiler optimize
-  int k_bases[k_iterations];
-#pragma unroll
-  for (int ki = 0; ki < k_iterations; ki++) {
-    k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT;
-  }
-
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.wait;");
-#endif
-
-  // Process the GEMM in chunks
-  for (int ki = 0; ki < k_iterations; ki++) {
-    int const k_base = k_bases[ki];
-
-    // Load B matrix values using vector load (8 bf16 values)
-    uint4 b_vec = *reinterpret_cast<uint4 const*>(b_col + k_base);
-
-    // Convert B values to float
-    float b_float[VPT];
-    bf16_uint4_to_float8<VPT>(b_vec, b_float);
-
-// Process each token
-#pragma unroll
-    for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
-      // Load both rows of A matrix using vector loads
-      uint4 a_vec = *reinterpret_cast<uint4 const*>(
-          mat_a + (m_idx * kHiddenDim) + k_base);
-
-      // Convert A values to float
-      float a_float[VPT];
-      bf16_uint4_to_float8<VPT>(a_vec, a_float);
-
-// Process elements in this chunk
-#pragma unroll
-      for (int k = 0; k < VPT; k++) {
-        float a = a_float[k];
-        float b = b_float[k];
-        acc[m_idx] += a * b;
-      }
-    }
-  }
-
-  // Perform warp-level reduction
-  int const warpSize = 32;
-  int const warpId = tid / warpSize;
-  int const laneId = tid % warpSize;
-
-  // Register for warp-level reduction results
-  float warp_result[kNumTokens];
-
-#pragma unroll
-  for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
-    warp_result[m_idx] = acc[m_idx];
-  }
-
-// Perform warp-level reduction using optimized butterfly pattern
-#pragma unroll
-  for (int m = 0; m < kNumTokens; m++) {
-    float sum = warp_result[m];
-
-    // Butterfly reduction pattern
-    sum += __shfl_xor_sync(0xffffffff, sum, 16);
-    sum += __shfl_xor_sync(0xffffffff, sum, 8);
-    sum += __shfl_xor_sync(0xffffffff, sum, 4);
-    sum += __shfl_xor_sync(0xffffffff, sum, 2);
-    sum += __shfl_xor_sync(0xffffffff, sum, 1);
-
-    // Only the first thread in each warp stores to shared memory
-    if (laneId == 0) {
-      sm_reduction[m][warpId] = sum;
-    }
-  }
-
-  __syncthreads();
-
-  // Final reduction across warps (only first thread)
-  if (tid == 0) {
-#pragma unroll
-    for (int m = 0; m < kNumTokens; m++) {
-      float final_sum = 0.0f;
-
-// Sum across the kNumWarps
-#pragma unroll
-      for (int w = 0; w < kNumWarps; w++) {
-        final_sum += sm_reduction[m][w];
-      }
-
-      // Write final result
-      out[m * kNumExperts + n_idx] = final_sum;
-    }
-  }
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.launch_dependents;");
-#endif
-}
-
-template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
-void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b,
-                                 cudaStream_t stream) {
-  constexpr int VPT = 16 / sizeof(T);
-  constexpr int kBlockSize = 128;
-  cudaLaunchConfig_t config;
-  config.gridDim = kNumExperts;
-  config.blockDim = kBlockSize;
-  config.dynamicSmemBytes = 0;
-  config.stream = stream;
-  cudaLaunchAttribute attrs[1];
-  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-  attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
-  config.numAttrs = 1;
-  config.attrs = attrs;
-  cudaLaunchKernelEx(
-      &config,
-      router_gemm_kernel_float_output<T, kBlockSize, VPT, kNumTokens,
-                                      kNumExperts, kHiddenDim>,
-      output, mat_a, mat_b);
-}
-
-// Template instantiations for DEFAULT_NUM_EXPERTS experts
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 1, 256, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 2, 256, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 3, 256, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 4, 256, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 5, 256, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 6, 256, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 7, 256, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 8, 256, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 9, 256, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 10, 256, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 11, 256, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 12, 256, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 13, 256, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 14, 256, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 256, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 256, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-// Template instantiations for KIMI_K2_NUM_EXPERTS experts
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 1, 384, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 2, 384, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 3, 384, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 4, 384, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 5, 384, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 6, 384, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 7, 384, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 8, 384, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 9, 384, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 10, 384, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 11, 384, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 12, 384, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 13, 384, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 14, 384, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 384, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
-
-template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 384, 7168>(
-    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);

csrc/moe/dsv3_router_gemm_utils.h

@@ -1,43 +0,0 @@
-/*
- * Adapted from SGLang's sgl-kernel implementation, which was adapted from
- * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu
- * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp
- *
- * Copyright (c) 2019-2023, NVIDIA CORPORATION.  All rights reserved.
- *
- * 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.
- */
-
-#pragma once
-
-#include <ATen/cuda/CUDAContext.h>
-
-#include <cstdlib>
-#include <mutex>
-
-inline int getSMVersion() {
-  auto* props = at::cuda::getCurrentDeviceProperties();
-  return props->major * 10 + props->minor;
-}
-
-inline bool getEnvEnablePDL() {
-  static std::once_flag flag;
-  static bool enablePDL = false;
-  std::call_once(flag, [&]() {
-    if (getSMVersion() >= 90) {
-      const char* env = std::getenv("TRTLLM_ENABLE_PDL");
-      enablePDL = env && env[0] == '1' && env[1] == '\0';
-    }
-  });
-  return enablePDL;
-}

csrc/moe/grouped_topk_kernels.cu

@@ -1,6 +1,6 @@
 /*
  * Adapted from
- * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/v0.21.0/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu
  * Copyright (c) 2025, The vLLM team.
  * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION &
  * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
@@ -17,10 +17,8 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-#include "moeTopKFuncs.cuh"
 #include <c10/cuda/CUDAStream.h>
 #include <torch/all.h>
-#include <cmath>
 #include <cuda_fp16.h>
 #include <cuda_bf16.h>
 #include <cuda/std/limits>
@@ -32,17 +30,7 @@ namespace vllm {
 namespace moe {
 
 constexpr unsigned FULL_WARP_MASK = 0xffffffff;
-static constexpr int WARP_SIZE = 32;
-static constexpr int NumNemotronExperts = 512;
-static constexpr int NumKimiK2Experts = 384;
-static constexpr int NumDeepseekExperts = 256;
-static constexpr int MaxSupportedExpertCount =
-    std::max({NumNemotronExperts, NumKimiK2Experts, NumDeepseekExperts});
-static constexpr int MaxNumExpertsUnit = 128;
-static constexpr int NumTopGroupScores = 2;
-static constexpr int DefaultMaxNumTopExperts = 8;
-static constexpr int MaxSupportedTopExperts = 22;
-static constexpr int MaxNumTopGroups = 4;
+constexpr int32_t WARP_SIZE = 32;
 
 namespace warp_topk {
 
@@ -669,335 +657,76 @@ __global__ void grouped_topk_fused_kernel(
 #endif
 }
 
-template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
-          int MaxNumExperts, bool UseGroups,
-          int MaxNumTopExperts = DefaultMaxNumTopExperts>
-__global__ void grouped_topk_fused_small_expert_count_kernel(
-    T* scores, float* topkValues, IdxT* topkIndices, BiasT const* routingBias,
-    int64_t const numTokens, int64_t const numGroup, int64_t const topkGroup,
-    int64_t const topk, int64_t const numExperts,
-    int64_t const numExpertsPerGroup, bool const renormalize,
-    double const routedScalingFactor) {
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  cudaGridDependencySynchronize();
-#endif
-  // declare shared memory structure
-  // number of experts is bounded by number of threads
-  __shared__ float __attribute((aligned(128))) smemScoreSigmoid[MaxNumExperts];
-  __shared__ float __attribute((aligned(128))) smemScoreBias[MaxNumExperts];
-  // number of expert groups is bounded by number of warps
-  int constexpr NumWarps = MaxNumExperts / WARP_SIZE;
-  __shared__ float __attribute((aligned(128))) smemGroupScores[NumWarps];
-
-  // needed for warp reduce
-  auto block = cg::this_thread_block();
-  auto warp = cg::tiled_partition<WARP_SIZE>(block);
-
-  // for the final reduction of weight norm, only some lanes need to participate
-  int32_t laneIdx = threadIdx.x % WARP_SIZE;
-  int32_t warpIdx = __shfl_sync(0xffffffff, threadIdx.x / WARP_SIZE, 0);
-
-  if constexpr (UseGroups) {
-    if (warpIdx >= numGroup) {
-      return;
-    }
-  }
-  // note that for invalid scores, we simply use a negative value:
-  // they work well even with the compacted format used in topK, and
-  // sigmoid / bias activated scores cannot be negative
-  const float invalidScoreFloat = float{-INFINITY};
-
-  // load bias already; each warp represents one expert group
-  auto threadExpert = threadIdx.x;
-  bool expertSelected = threadExpert < numExperts;
-  if constexpr (UseGroups) {
-    threadExpert = warpIdx * numExpertsPerGroup + laneIdx;
-    expertSelected = laneIdx < numExpertsPerGroup;
-  }
-
-  auto scoreIdx = int64_t{blockIdx.x} * int64_t{numExperts} + threadExpert;
-  auto biasVal = expertSelected ? static_cast<float>(routingBias[threadExpert])
-                                : invalidScoreFloat;
-  topkValues += blockIdx.x * topk;
-  topkIndices += blockIdx.x * topk;
-
-  // get our assigned thread score; each warp represents one expert group
-  float score =
-      expertSelected ? static_cast<float>(scores[scoreIdx]) : invalidScoreFloat;
-  auto scoreSigmoid = apply_scoring<SF>(score);
-  // write the sigmoid score to shared for later use
-  if (expertSelected) {
-    smemScoreSigmoid[threadExpert] = scoreSigmoid;
-  }
-
-  // get the score with bias
-  // note that with invalid values, because sigmoid is < 1 and bias is -1,
-  // we must get a negative value, which is smaller than any valid value
-  auto scoreBias = float{scoreSigmoid + float{biasVal}};
-
-  if (expertSelected) {
-    smemScoreBias[threadExpert] = scoreBias;
-  }
-
-  // registers for top group score reduction
-  float topExpGroupScores[NumTopGroupScores];
-  [[maybe_unused]] int32_t topExpGroupIdx[NumTopGroupScores];
-  float topGroups[MaxNumTopGroups];  // bound of numGroup
-  int32_t topGroupIdx[MaxNumTopGroups];
-  float expertScoreGroup[MaxNumTopGroups];
-  int32_t expertIdxGroup[MaxNumTopGroups];
-  float topScores[MaxNumTopExperts];  // bound of topk
-  int32_t topExperts[MaxNumTopExperts];
-
-  if constexpr (UseGroups) {
-    reduce_topk::reduceTopK(warp, topExpGroupScores, topExpGroupIdx, scoreBias,
-                            threadExpert,
-                            /* minValue */ invalidScoreFloat);
-
-    // get the final group score and write it to shared
-    if (warp.thread_rank() == 0) {
-      auto groupScore = topExpGroupScores[0] + topExpGroupScores[1];
-      smemGroupScores[warpIdx] = groupScore;
-    }
-  }
-
-  // make group scores available to all warps
-  __syncthreads();
-
-  if constexpr (UseGroups) {
-    if (warpIdx == 0) {
-      // a single warp performs the selection of top groups, and goes on to
-      // select the final experts
-      float groupScore =
-          laneIdx < numGroup ? smemGroupScores[laneIdx] : invalidScoreFloat;
-
-      reduce_topk::reduceTopK(warp, topGroups, topGroupIdx, groupScore, laneIdx,
-                              /* minValue */ invalidScoreFloat);
-      // final expert selection: get relevant indexes and scores from shared
-#pragma unroll
-      for (int ii = 0; ii < MaxNumTopGroups; ++ii) {  // bound of numGroup
-        auto groupIdx = topGroupIdx[ii];
-        expertIdxGroup[ii] = groupIdx * numExpertsPerGroup + laneIdx;
-
-        expertScoreGroup[ii] = (ii < topkGroup) && expertSelected
-                                   ? smemScoreBias[expertIdxGroup[ii]]
-                                   : invalidScoreFloat;
-      }
-
-      reduce_topk::reduceTopK(warp, topScores, topExperts, expertScoreGroup,
-                              expertIdxGroup, /* minValue */ invalidScoreFloat,
-                              topk);
-    }
-  } else if constexpr (MaxNumExperts > MaxNumExpertsUnit) {
-    // without groups, and the expert number is larger than MaxNumExpertsUnit,
-    // we need to use multiple warps to calculate the intermediate topk results
-
-    int constexpr NumExpertWarps = (MaxNumExperts - 1) / MaxNumExpertsUnit + 1;
-    int constexpr NumInterTopK = NumExpertWarps * MaxNumTopExperts;
-    __shared__ float
-        __attribute((aligned(128))) smemInterTopScores[NumInterTopK];
-    __shared__ int32_t
-        __attribute((aligned(128))) smemInterTopExperts[NumInterTopK];
-    if (warpIdx < NumExpertWarps) {
-      int offset = warpIdx * WARP_SIZE * MaxNumTopGroups;
-#pragma unroll
-      for (int ii = 0; ii < MaxNumTopGroups; ++ii) {
-        auto expertIdx = ii * WARP_SIZE + laneIdx;
-        expertIdxGroup[ii] = offset + expertIdx;
-        expertScoreGroup[ii] = offset + expertIdx < numExperts
-                                   ? smemScoreBias[offset + expertIdx]
-                                   : invalidScoreFloat;
-      }
-      reduce_topk::reduceTopK(warp, topScores, topExperts, expertScoreGroup,
-                              expertIdxGroup,
-                              /* minValue */ invalidScoreFloat, topk);
-
-      if (laneIdx < topk) {
-        smemInterTopScores[warpIdx * MaxNumTopExperts + laneIdx] =
-            topScores[laneIdx];
-        smemInterTopExperts[warpIdx * MaxNumTopExperts + laneIdx] =
-            topExperts[laneIdx];
-      } else if (laneIdx >= topk && laneIdx < MaxNumTopExperts) {
-        smemInterTopScores[warpIdx * MaxNumTopExperts + laneIdx] =
-            invalidScoreFloat;
-        smemInterTopExperts[warpIdx * MaxNumTopExperts + laneIdx] =
-            MaxNumExperts - 1;
-      }
-    }
-    __syncthreads();
-    if (warpIdx == 0) {
-      int constexpr NumInterTopKPerThread = (NumInterTopK - 1) / WARP_SIZE + 1;
-      float intermediateScore[NumInterTopKPerThread];
-      int32_t intermediateExpert[NumInterTopKPerThread];
-      for (int i = laneIdx; i < NumInterTopKPerThread * WARP_SIZE;
-           i += WARP_SIZE) {
-        int ii = i / WARP_SIZE;
-        if (i < NumInterTopK) {
-          intermediateScore[ii] = smemInterTopScores[i];
-          intermediateExpert[ii] = smemInterTopExperts[i];
-        } else {
-          intermediateScore[ii] = invalidScoreFloat;
-          intermediateExpert[ii] = MaxNumExperts - 1;
-        }
-      }
-      reduce_topk::reduceTopK(warp, topScores, topExperts, intermediateScore,
-                              intermediateExpert,
-                              /* minValue */ invalidScoreFloat, topk);
-    }
-  } else {
-    // without groups, and the expert number is smaller than MaxNumExpertsUnit
-    // each thread just takes `MaxNumTopGroups` experts
-    if (warpIdx == 0) {
-#pragma unroll
-      for (int ii = 0; ii < MaxNumTopGroups; ++ii) {
-        auto expertIdx = ii * WARP_SIZE + laneIdx;
-        expertIdxGroup[ii] = expertIdx;
-        expertScoreGroup[ii] = expertIdx < numExperts ? smemScoreBias[expertIdx]
-                                                      : invalidScoreFloat;
-      }
-      reduce_topk::reduceTopK(warp, topScores, topExperts, expertScoreGroup,
-                              expertIdxGroup,
-                              /* minValue */ invalidScoreFloat, topk);
-    }
-  }
-
-  if (warpIdx == 0) {
-    // determine our lane's expert index and write to output
-    int32_t expertIdx =
-        laneIdx < topk ? topExperts[laneIdx] : MaxNumExperts - 1;
-    float scoreNorm = laneIdx < topk ? smemScoreSigmoid[expertIdx] : 0.F;
-    float finalScore = static_cast<float>(scoreNorm * routedScalingFactor);
-    // norm the value
-    if (renormalize) {
-      auto redNorm = cg::reduce(warp, scoreNorm, cg::plus<float>{});
-      finalScore /= (redNorm + 1e-20);
-    }
-    // store the topk scores and experts to output
-    if (laneIdx < topk) {
-      topkValues[laneIdx] = finalScore;
-      topkIndices[laneIdx] = expertIdx;
-    }
-  }
-
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  cudaTriggerProgrammaticLaunchCompletion();
-#endif
-}
-
-template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
+template <typename T, typename BiasT, typename IdxT>
 void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
                    BiasT const* bias, int64_t const num_tokens,
                    int64_t const num_experts, int64_t const n_group,
                    int64_t const topk_group, int64_t const topk,
                    bool const renormalize, double const routed_scaling_factor,
-                   bool enable_pdl = false, cudaStream_t const stream = 0) {
+                   int const scoring_func, bool enable_pdl = false,
+                   cudaStream_t const stream = 0) {
   cudaLaunchConfig_t config;
+  // One block per token; one warp per group.
+  config.gridDim = static_cast<uint32_t>(num_tokens);
+  config.blockDim = static_cast<uint32_t>(n_group) * WARP_SIZE;
+  // Dynamic shared memory: WarpSelect staging + per-group topk buffers.
+  int32_t const num_warps = static_cast<int32_t>(n_group);
+  size_t const val_bytes =
+      static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(T);
+  size_t const val_bytes_aligned =
+      warp_topk::round_up_to_multiple_of<256>(val_bytes);
+  size_t const idx_bytes =
+      static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(int32_t);
+  size_t const internal_bytes = val_bytes_aligned + idx_bytes;
+  size_t const extra_bytes = 16 + static_cast<size_t>(n_group) * sizeof(T);
+  config.dynamicSmemBytes = internal_bytes + extra_bytes;
   config.stream = stream;
   cudaLaunchAttribute attrs[1];
   attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
   attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl;
   config.numAttrs = 1;
   config.attrs = attrs;
-
-  // Check if we can use the optimized
-  // grouped_topk_fused_small_expert_count_kernel
-  bool const is_single_group =
-      (n_group == 1) && (topk_group == 1) &&
-      (num_experts <= MaxSupportedExpertCount) &&
-      (topk <= DefaultMaxNumTopExperts || topk == MaxSupportedTopExperts);
-
-  int64_t const experts_per_group = num_experts / n_group;
-  bool const is_multi_group =
-      (n_group > 1) && (num_experts <= NumDeepseekExperts) &&
-      (experts_per_group <= WARP_SIZE) &&
-      (experts_per_group * topk_group <= MaxNumExpertsUnit) &&
-      (topk <= DefaultMaxNumTopExperts) && (topk_group <= MaxNumTopGroups);
-
-  if (is_single_group || is_multi_group) {
-    auto* kernel_instance =
-        &grouped_topk_fused_small_expert_count_kernel<T, BiasT, IdxT, SF,
-                                                      NumDeepseekExperts, true>;
-    int num_threads = NumDeepseekExperts;
-    if (is_single_group) {
-      // Special case for Nemotron, which selects top 22 from 512 experts, and 1
-      // group only.
-      if (num_experts == NumNemotronExperts && n_group == 1 &&
-          topk == MaxSupportedTopExperts) {
-        kernel_instance = &grouped_topk_fused_small_expert_count_kernel<
-            T, BiasT, IdxT, SF, NumNemotronExperts, false,
-            MaxSupportedTopExperts>;
-        num_threads = NumNemotronExperts;
-      } else if (num_experts > NumKimiK2Experts &&
-                 num_experts <= MaxSupportedExpertCount) {
-        kernel_instance = &grouped_topk_fused_small_expert_count_kernel<
-            T, BiasT, IdxT, SF, MaxSupportedExpertCount, false>;
-        num_threads = MaxSupportedExpertCount;
-      } else if (num_experts > MaxNumExpertsUnit &&
-                 num_experts <= NumKimiK2Experts) {
-        kernel_instance = &grouped_topk_fused_small_expert_count_kernel<
-            T, BiasT, IdxT, SF, NumKimiK2Experts, false>;
-        num_threads = NumKimiK2Experts;
-      } else {
-        kernel_instance = &grouped_topk_fused_small_expert_count_kernel<
-            T, BiasT, IdxT, SF, MaxNumExpertsUnit, false>;
-        num_threads = MaxNumExpertsUnit;
-      }
+  auto const sf = static_cast<ScoringFunc>(scoring_func);
+  switch (sf) {
+    case SCORING_NONE: {
+      auto* kernel_instance =
+          &grouped_topk_fused_kernel<T, BiasT, IdxT, SCORING_NONE>;
+      cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values,
+                         topk_indices, bias, num_tokens, num_experts, n_group,
+                         topk_group, topk, renormalize, routed_scaling_factor);
+      return;
     }
-    config.gridDim = num_tokens;
-    config.blockDim = num_threads;
-    config.dynamicSmemBytes = 0;
-    cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values,
-                       topk_indices, bias, num_tokens, n_group, topk_group,
-                       topk, num_experts, num_experts / n_group, renormalize,
-                       routed_scaling_factor);
-  } else {
-    auto* kernel_instance = &grouped_topk_fused_kernel<T, BiasT, IdxT, SF>;
-    // One block per token; one warp per group.
-    config.gridDim = static_cast<uint32_t>(num_tokens);
-    config.blockDim = static_cast<uint32_t>(n_group) * WARP_SIZE;
-    // Dynamic shared memory: WarpSelect staging + per-group topk buffers.
-    int32_t const num_warps = static_cast<int32_t>(n_group);
-    size_t const val_bytes =
-        static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(T);
-    size_t const val_bytes_aligned =
-        warp_topk::round_up_to_multiple_of<256>(val_bytes);
-    size_t const idx_bytes =
-        static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(int32_t);
-    size_t const internal_bytes = val_bytes_aligned + idx_bytes;
-    size_t const extra_bytes = 16 + static_cast<size_t>(n_group) * sizeof(T);
-    config.dynamicSmemBytes = internal_bytes + extra_bytes;
-    cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values,
-                       topk_indices, bias, num_tokens, num_experts, n_group,
-                       topk_group, topk, renormalize, routed_scaling_factor);
+    case SCORING_SIGMOID: {
+      auto* kernel_instance =
+          &grouped_topk_fused_kernel<T, BiasT, IdxT, SCORING_SIGMOID>;
+      cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values,
+                         topk_indices, bias, num_tokens, num_experts, n_group,
+                         topk_group, topk, renormalize, routed_scaling_factor);
+      return;
+    }
+    default:
+      // should be guarded by higher level checks.
+      TORCH_CHECK(false, "Unsupported scoring_func in invokeNoAuxTc");
   }
 }
 
-#define INSTANTIATE_NOAUX_TC(T, BiasT, IdxT, SF)                             \
-  template void invokeNoAuxTc<T, BiasT, IdxT, SF>(                           \
+#define INSTANTIATE_NOAUX_TC(T, BiasT, IdxT)                                 \
+  template void invokeNoAuxTc<T, BiasT, IdxT>(                               \
       T * scores, float* topk_values, IdxT* topk_indices, BiasT const* bias, \
       int64_t const num_tokens, int64_t const num_experts,                   \
       int64_t const n_group, int64_t const topk_group, int64_t const topk,   \
       bool const renormalize, double const routed_scaling_factor,            \
-      bool enable_pdl, cudaStream_t const stream);
+      int const scoring_func, bool enable_pdl, cudaStream_t const stream);
 
-INSTANTIATE_NOAUX_TC(float, float, int32_t, SCORING_SIGMOID);
-INSTANTIATE_NOAUX_TC(float, half, int32_t, SCORING_SIGMOID);
-INSTANTIATE_NOAUX_TC(float, __nv_bfloat16, int32_t, SCORING_SIGMOID);
-INSTANTIATE_NOAUX_TC(half, float, int32_t, SCORING_SIGMOID);
-INSTANTIATE_NOAUX_TC(half, half, int32_t, SCORING_SIGMOID);
-INSTANTIATE_NOAUX_TC(half, __nv_bfloat16, int32_t, SCORING_SIGMOID);
-INSTANTIATE_NOAUX_TC(__nv_bfloat16, float, int32_t, SCORING_SIGMOID);
-INSTANTIATE_NOAUX_TC(__nv_bfloat16, half, int32_t, SCORING_SIGMOID);
-INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t, SCORING_SIGMOID);
-INSTANTIATE_NOAUX_TC(float, float, int32_t, SCORING_NONE);
-INSTANTIATE_NOAUX_TC(float, half, int32_t, SCORING_NONE);
-INSTANTIATE_NOAUX_TC(float, __nv_bfloat16, int32_t, SCORING_NONE);
-INSTANTIATE_NOAUX_TC(half, float, int32_t, SCORING_NONE);
-INSTANTIATE_NOAUX_TC(half, half, int32_t, SCORING_NONE);
-INSTANTIATE_NOAUX_TC(half, __nv_bfloat16, int32_t, SCORING_NONE);
-INSTANTIATE_NOAUX_TC(__nv_bfloat16, float, int32_t, SCORING_NONE);
-INSTANTIATE_NOAUX_TC(__nv_bfloat16, half, int32_t, SCORING_NONE);
-INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t, SCORING_NONE);
+INSTANTIATE_NOAUX_TC(float, float, int32_t);
+INSTANTIATE_NOAUX_TC(float, half, int32_t);
+INSTANTIATE_NOAUX_TC(float, __nv_bfloat16, int32_t);
+INSTANTIATE_NOAUX_TC(half, float, int32_t);
+INSTANTIATE_NOAUX_TC(half, half, int32_t);
+INSTANTIATE_NOAUX_TC(half, __nv_bfloat16, int32_t);
+INSTANTIATE_NOAUX_TC(__nv_bfloat16, float, int32_t);
+INSTANTIATE_NOAUX_TC(__nv_bfloat16, half, int32_t);
+INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t);
 }  // end namespace moe
 }  // namespace vllm
 
@@ -1033,53 +762,46 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
       {num_tokens, topk}, torch::dtype(torch::kInt32).device(torch::kCUDA));
 
   auto stream = c10::cuda::getCurrentCUDAStream(scores.get_device());
-  auto const sf = static_cast<vllm::moe::ScoringFunc>(scoring_func);
 
-#define LAUNCH_KERNEL_SF(T, BiasT, IdxT)                                      \
-  do {                                                                        \
-    switch (sf) {                                                             \
-      case vllm::moe::SCORING_NONE:                                           \
-        vllm::moe::invokeNoAuxTc<T, BiasT, IdxT, vllm::moe::SCORING_NONE>(    \
-            reinterpret_cast<T*>(scores.mutable_data_ptr()),                  \
-            reinterpret_cast<float*>(topk_values.mutable_data_ptr()),         \
-            reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()),         \
-            reinterpret_cast<BiasT const*>(bias.data_ptr()), num_tokens,      \
-            num_experts, n_group, topk_group, topk, renormalize,              \
-            routed_scaling_factor, false, stream);                            \
-        break;                                                                \
-      case vllm::moe::SCORING_SIGMOID:                                        \
-        vllm::moe::invokeNoAuxTc<T, BiasT, IdxT, vllm::moe::SCORING_SIGMOID>( \
-            reinterpret_cast<T*>(scores.mutable_data_ptr()),                  \
-            reinterpret_cast<float*>(topk_values.mutable_data_ptr()),         \
-            reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()),         \
-            reinterpret_cast<BiasT const*>(bias.data_ptr()), num_tokens,      \
-            num_experts, n_group, topk_group, topk, renormalize,              \
-            routed_scaling_factor, false, stream);                            \
-        break;                                                                \
-      default:                                                                \
-        throw std::invalid_argument("Unsupported scoring_func");              \
-        break;                                                                \
-    }                                                                         \
-  } while (0)
-
-#define LAUNCH_KERNEL(T, IdxT)                                         \
-  do {                                                                 \
-    switch (bias_type) {                                               \
-      case torch::kFloat16:                                            \
-        LAUNCH_KERNEL_SF(T, half, IdxT);                               \
-        break;                                                         \
-      case torch::kFloat32:                                            \
-        LAUNCH_KERNEL_SF(T, float, IdxT);                              \
-        break;                                                         \
-      case torch::kBFloat16:                                           \
-        LAUNCH_KERNEL_SF(T, __nv_bfloat16, IdxT);                      \
-        break;                                                         \
-      default:                                                         \
-        throw std::invalid_argument(                                   \
-            "Invalid bias dtype, only supports float16, float32, and " \
-            "bfloat16");                                               \
-        break;                                                         \
-    }                                                                  \
+#define LAUNCH_KERNEL(T, IdxT)                                               \
+  do {                                                                       \
+    switch (bias_type) {                                                     \
+      case torch::kFloat16:                                                  \
+        vllm::moe::invokeNoAuxTc<T, half, IdxT>(                             \
+            reinterpret_cast<T*>(scores.mutable_data_ptr()),                 \
+            reinterpret_cast<float*>(topk_values.mutable_data_ptr()),        \
+            reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()),        \
+            reinterpret_cast<half const*>(bias.data_ptr()), num_tokens,      \
+            num_experts, n_group, topk_group, topk, renormalize,             \
+            routed_scaling_factor, static_cast<int>(scoring_func), false,    \
+            stream);                                                         \
+        break;                                                               \
+      case torch::kFloat32:                                                  \
+        vllm::moe::invokeNoAuxTc<T, float, IdxT>(                            \
+            reinterpret_cast<T*>(scores.mutable_data_ptr()),                 \
+            reinterpret_cast<float*>(topk_values.mutable_data_ptr()),        \
+            reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()),        \
+            reinterpret_cast<float const*>(bias.data_ptr()), num_tokens,     \
+            num_experts, n_group, topk_group, topk, renormalize,             \
+            routed_scaling_factor, static_cast<int>(scoring_func), false,    \
+            stream);                                                         \
+        break;                                                               \
+      case torch::kBFloat16:                                                 \
+        vllm::moe::invokeNoAuxTc<T, __nv_bfloat16, IdxT>(                    \
+            reinterpret_cast<T*>(scores.mutable_data_ptr()),                 \
+            reinterpret_cast<float*>(topk_values.mutable_data_ptr()),        \
+            reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()),        \
+            reinterpret_cast<__nv_bfloat16 const*>(bias.data_ptr()),         \
+            num_tokens, num_experts, n_group, topk_group, topk, renormalize, \
+            routed_scaling_factor, static_cast<int>(scoring_func), false,    \
+            stream);                                                         \
+        break;                                                               \
+      default:                                                               \
+        throw std::invalid_argument(                                         \
+            "Invalid bias dtype, only supports float16, float32, and "       \
+            "bfloat16");                                                     \
+        break;                                                               \
+    }                                                                        \
   } while (0)
 
   switch (data_type) {
@@ -1102,6 +824,5 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
       break;
   }
 #undef LAUNCH_KERNEL
-#undef LAUNCH_KERNEL_SF
   return {topk_values, topk_indices};
 }

csrc/moe/moeTopKFuncs.cuh

@@ -1,257 +0,0 @@
-/*
- * Adapted from
- * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh
- * Copyright (c) 2026, The vLLM team.
- * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION. 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.
- */
-#pragma once
-
-#include <cooperative_groups.h>
-#include <cooperative_groups/reduce.h>
-#include <cub/cub.cuh>
-
-namespace vllm {
-namespace moe {
-namespace reduce_topk {
-namespace cg = cooperative_groups;
-static constexpr int kWARP_SIZE = 32;
-
-template <typename T_>
-struct TopKRedType {
-  using T = T_;
-  static_assert(
-      std::is_same_v<T, float> || std::is_same_v<T, half> ||
-          std::is_same_v<T, __nv_bfloat16> || std::is_same_v<T, int>,
-      "Top K reduction only implemented for int, float, float16 and bfloat16");
-
-  using TypeCmp = std::conditional_t<sizeof(T) == 4, uint64_t, uint32_t>;
-  using IdxT = std::conditional_t<sizeof(T) == 4, int32_t, int16_t>;
-
-  static constexpr int kMoveBits = (sizeof(T) == 4) ? 32 : 16;
-  static constexpr int kMaxIdx = 65535;
-  TypeCmp compValIdx;
-
-  static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) {
-    auto valueBits = cub::Traits<T>::TwiddleIn(
-        reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(val));
-    TypeCmp compactTmp = valueBits;
-    compactTmp = (compactTmp << kMoveBits) | (0xFFFF & (kMaxIdx - idx));
-    // Use 65535 minus idx to give higher priority to elements with smaller
-    // indices.
-    return compactTmp;
-  }
-
-  static __host__ __device__ void unpack(T& value, int32_t& index,
-                                         TypeCmp cmp) {
-    // Since “65535-idx” is always smaller than 65536 and positive, we can
-    // directly use it as the lower 16 bits
-    index = kMaxIdx - static_cast<int32_t>((cmp & 0xFFFF));
-
-    auto compactTmp = cmp >> kMoveBits;
-    auto valueBits = cub::Traits<T>::TwiddleOut(
-        reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(compactTmp));
-    value = reinterpret_cast<T&>(valueBits);
-  }
-
-  __host__ __device__ TopKRedType() = default;
-
-  __host__ __device__ TopKRedType(T val, int32_t idx)
-      : compValIdx(makeCmpVal(val, idx)) {}
-
-  __host__ __device__ operator TypeCmp() const noexcept { return compValIdx; }
-
-  __device__ inline TypeCmp reduce(
-      cg::thread_block_tile<kWARP_SIZE> const& warp) {
-    return cg::reduce(warp, compValIdx, cg::greater<TypeCmp>{});
-  }
-};
-
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-template <int K_, bool Enable_>
-struct TopKIdx {
-  // by default, empty
-};
-
-template <int K_>
-struct TopKIdx<K_, true> {
-  static constexpr int K = K_;
-  int32_t val[K];
-};
-
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-#define TOPK_SWAP(I, J)                                         \
-  {                                                             \
-    auto pairMin = min(topK[I].compValIdx, topK[J].compValIdx); \
-    auto pairMax = max(topK[I].compValIdx, topK[J].compValIdx); \
-    topK[I].compValIdx = pairMax;                               \
-    topK[J].compValIdx = pairMin;                               \
-  }
-
-template <int N, typename RedType>
-struct Sort;
-
-template <typename RedType>
-struct Sort<1, RedType> {
-  static __device__ void run(RedType* topK) {}
-};
-
-template <typename RedType>
-struct Sort<2, RedType> {
-  static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); }
-};
-
-template <typename RedType>
-struct Sort<3, RedType> {
-  static __device__ void run(RedType* topK) {
-    TOPK_SWAP(0, 1);
-    TOPK_SWAP(1, 2);
-    TOPK_SWAP(0, 1);
-  }
-};
-
-template <typename RedType>
-struct Sort<4, RedType> {
-  static __device__ void run(RedType* topK) {
-    TOPK_SWAP(0, 2);
-    TOPK_SWAP(1, 3);
-    TOPK_SWAP(0, 1);
-    TOPK_SWAP(2, 3);
-    TOPK_SWAP(1, 2);
-  }
-};
-
-template <int K, typename Type>
-__forceinline__ __device__ void reduceTopK(
-    cg::thread_block_tile<kWARP_SIZE> const& warp, Type (&out)[K],
-    int32_t (&outIdx)[K], Type value, int32_t idx, Type const minValue,
-    int actualK = K) {
-  static_assert(K > 0, "Top K must have K > 0");
-  static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
-  using RedType = TopKRedType<Type>;
-  RedType topK{value, idx};
-  typename RedType::TypeCmp packedMax{};
-#pragma unroll
-  for (int kk = 0; kk < actualK; ++kk) {
-    topK =
-        kk > 0 && packedMax == topK.compValIdx ? RedType{minValue, idx} : topK;
-    // get the next largest value
-    packedMax = topK.reduce(warp);
-    RedType::unpack(out[kk], outIdx[kk], packedMax);
-  }
-};
-
-template <int K, typename Type, int N, bool IsSorted = false>
-__device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> const& warp,
-                               Type (&out)[K], int32_t (&outIdx)[K],
-                               Type (&value)[N], int32_t (&idx)[N],
-                               Type minValue, int actualK = K) {
-  static_assert(K > 0, "Top K must have K > 0");
-  static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
-  static_assert(N > 0, "Top K must have N > 0");
-  static_assert(N < 5,
-                "Only support candidates number less than or equal to 128");
-  using RedType = TopKRedType<Type>;
-  RedType topK[N];
-#pragma unroll
-  for (int nn = 0; nn < N; ++nn) {
-    topK[nn] = RedType{value[nn], idx[nn]};
-  }
-
-  if constexpr (!IsSorted) {
-    Sort<N, RedType>::run(topK);
-  }
-  typename RedType::TypeCmp packedMax{};
-#pragma unroll
-  for (int kk = 0; kk < actualK; ++kk) {
-    bool update = kk > 0 && packedMax == topK[0].compValIdx;
-#pragma unroll
-    for (int nn = 0; nn < N; ++nn) {
-      topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]}
-                 : update              ? topK[nn + 1]
-                                       : topK[nn];
-    }
-    // get the next largest value
-    packedMax = topK[0].reduce(warp);
-    RedType::unpack(out[kk], outIdx[kk], packedMax);
-  }
-};
-
-template <int K, typename Type, int N>
-__forceinline__ __device__ void reduceTopK(
-    cg::thread_block_tile<kWARP_SIZE> const& warp, Type (&out)[K],
-    int32_t (&outIdx)[K], Type (&value)[N], int32_t (&idx)[N],
-    Type const minValue, int actualK = K) {
-  static_assert(K > 0, "Top K must have K > 0");
-  static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
-  static_assert(N > 0, "Top K must have N > 0");
-  static_assert(
-      N <= 16,
-      "Only support candidates number less than or equal to 16*32=512");
-  static_assert(N <= 4 || N % 4 == 0,
-                "Only support candidates number is a multiple of 4*32=128 or "
-                "less than or equal to 4");
-  using RedType = TopKRedType<Type>;
-
-  if constexpr (N <= 4) {
-    reduceTopKFunc<K, Type, N>(warp, out, outIdx, value, idx, minValue,
-                               actualK);
-  } else {
-    constexpr int numLoops = N / 4;
-    constexpr int numResults = (numLoops * K - 1) / kWARP_SIZE + 1;
-
-    Type topKBufferValue[numResults];
-    int32_t topKBufferIdx[numResults];
-    int32_t laneIdx = threadIdx.x % kWARP_SIZE;
-
-    for (int ii = 0; ii < numResults; ++ii) {
-      topKBufferValue[ii] = minValue;
-      topKBufferIdx[ii] = ii * kWARP_SIZE - 1;
-    }
-    for (int loop = 0; loop < numLoops; ++loop) {
-      int start = loop * 4;
-      Type topKValue[K];
-      int32_t topKIdx[K];
-      Type inValue[4];
-      int32_t inIdx[4];
-      for (int i = 0; i < 4; ++i) {
-        inValue[i] = value[start + i];
-        inIdx[i] = idx[start + i];
-      }
-      reduceTopKFunc<K, Type, 4>(warp, topKValue, topKIdx, inValue, inIdx,
-                                 minValue, actualK);
-      int inOffset = laneIdx % K;
-      if (laneIdx >= loop * K && laneIdx < (loop + 1) * K) {
-        topKBufferValue[0] = topKValue[inOffset];
-        topKBufferIdx[0] = topKIdx[inOffset];
-      }
-      if (loop == numLoops - 1 && (laneIdx < (numLoops * K - kWARP_SIZE))) {
-        topKBufferValue[1] = topKValue[inOffset];
-        topKBufferIdx[1] = topKIdx[inOffset];
-      }
-    }
-
-    reduceTopKFunc<K, Type, numResults>(warp, out, outIdx, topKBufferValue,
-                                        topKBufferIdx, minValue, actualK);
-  }
-};
-
-#undef TOPK_SWAP
-
-}  // namespace reduce_topk
-}  // namespace moe
-}  // namespace vllm

csrc/moe/moe_align_sum_kernels.cu

@@ -172,7 +172,7 @@ __device__ void _moe_align_block_size(
     }
   }
 
-  // Fill remaining expert_ids with -1
+  // Fill remaining expert_ids with 0
   const size_t fill_start_idx =
       cumsum[cumsum_offset + num_experts] / block_size + threadIdx.x;
   for (size_t i = fill_start_idx; i < max_num_m_blocks; i += blockDim.x) {
@@ -265,7 +265,7 @@ __device__ void _moe_align_block_size_small_batch_expert(
     }
   }
 
-  // Fill remaining expert_ids with -1
+  // Fill remaining expert_ids with 0
   const size_t fill_start_idx = cumsum[num_experts] / block_size + tid;
   for (size_t i = fill_start_idx; i < max_num_m_blocks; i += stride) {
     expert_ids[expert_ids_offset + i] = inactive_expert_id;
@@ -332,7 +332,7 @@ __global__ void moe_align_block_size_kernel(
       topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map,
       num_experts, padded_num_experts, experts_per_warp, block_size, numel,
       cumsum, max_num_tokens_padded, CEILDIV(max_num_tokens_padded, block_size),
-      0, -1, topk_num, nullptr, has_expert_map);
+      0, 0, topk_num, nullptr, has_expert_map);
 }
 
 template <typename scalar_t>
@@ -373,7 +373,7 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
   _moe_align_block_size_small_batch_expert<scalar_t, fill_threads>(
       topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map,
       num_experts, block_size, numel, max_num_tokens_padded,
-      CEILDIV(max_num_tokens_padded, block_size), -1, 0, topk_num, nullptr,
+      CEILDIV(max_num_tokens_padded, block_size), 0, 0, topk_num, nullptr,
       has_expert_map);
 }