Bump actions/setup-python from 6.1.0 to 6.2.0

Bumps [actions/setup-python](https://github.com/actions/setup-python) from 6.1.0 to 6.2.0.
- [Release notes](https://github.com/actions/setup-python/releases)
- [Commits](https://github.com/actions/setup-python/compare/83679a892e2d95755f2dac6acb0bfd1e9ac5d548...a309ff8b426b58ec0e2a45f0f869d46889d02405)

---
updated-dependencies:
- dependency-name: actions/setup-python
  dependency-version: 6.2.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>

.buildkite/ci_config_rocm.yaml

@@ -1,23 +0,0 @@
-name: vllm_rocm_ci
-job_dirs:
-  - ".buildkite/hardware_tests"
-run_all_patterns:
-  - "docker/Dockerfile.rocm"
-  - "docker/Dockerfile.rocm_base"
-  - "docker/ci-rocm.hcl"
-  - "docker/docker-bake-rocm.hcl"
-  - ".buildkite/hardware_tests/amd.yaml"
-  - ".buildkite/scripts/ci-bake-rocm.sh"
-  - ".buildkite/scripts/hardware_ci/run-amd-test.py"
-  - ".buildkite/scripts/hardware_ci/run-amd-test.sh"
-  - "CMakeLists.txt"
-  - "requirements/common.txt"
-  - "requirements/rocm.txt"
-  - "requirements/build/rocm.txt"
-  - "requirements/test/rocm.txt"
-  - "setup.py"
-  - "csrc/"
-  - "cmake/"
-run_all_exclude_patterns:
-  - "csrc/cpu/"
-  - "cmake/cpu_extension.cmake"

.buildkite/hardware_tests/amd.yaml

@@ -1,73 +1,42 @@
-group: Hardware - AMD Build
+group: Hardware - AMD Build 
 steps:
-  # Ensure ci_base is up-to-date before building the test image.
-  # Compares a content hash of ci_base-affecting files against the remote
-  # image label. If hashes match the build is skipped (< 30 s); if they
-  # differ ci_base is rebuilt and pushed automatically.
-  - label: "AMD: :docker: ensure ci_base"
-    key: ensure-ci-base-amd
+  - label: "AMD: :docker: build image"
+    key: image-build-amd
     depends_on: []
     device: amd_cpu
     no_plugin: true
     commands:
-      - bash .buildkite/scripts/ci-bake-rocm.sh ci-base-rocm-ci-with-deps
-    env:
-      DOCKER_BUILDKIT: "1"
-      VLLM_BAKE_FILE: "docker/docker-bake-rocm.hcl"
-      PYTORCH_ROCM_ARCH: "gfx90a;gfx942;gfx950"
-      REMOTE_VLLM: "1"
-      VLLM_BRANCH: "$BUILDKITE_COMMIT"
-    retry:
-      automatic:
-        - exit_status: -1  # Agent was lost
-          limit: 1
-        - exit_status: -10  # Agent was lost
-          limit: 1
-
-  - label: "AMD: :docker: build test image and artifacts"
-    key: image-build-amd
-    depends_on:
-      - ensure-ci-base-amd
-    device: amd_cpu
-    no_plugin: true
-    commands:
-      - |
-        if [[ "${ROCM_CI_ARTIFACT_ONLY:-0}" == "1" ]]; then
-          echo "ROCM_CI_ARTIFACT_ONLY=1; building ROCm wheel artifact only"
-          IMAGE_TAG="" bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-artifacts
-        else
-          bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-wheel
+    - >
+      docker build
+      --build-arg max_jobs=16
+      --build-arg REMOTE_VLLM=1
+      --build-arg ARG_PYTORCH_ROCM_ARCH='gfx90a;gfx942;gfx950'
+      --build-arg VLLM_BRANCH=$BUILDKITE_COMMIT
+      --tag "rocm/vllm-ci:${BUILDKITE_COMMIT}"
+      -f docker/Dockerfile.rocm
+      --target test
+      --no-cache
+      --progress plain .
+    - |
+      docker run --rm --network=none --entrypoint /bin/bash "rocm/vllm-ci:${BUILDKITE_COMMIT}" -ec '
+        if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi
+        if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi
+        if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi
+        if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi
+        command -v python3
+        command -v uv
+        command -v pytest
+        if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then
+          echo No ROCm CLI found in image >&2
+          exit 1
         fi
-      - |
-        docker run --rm --network=none --entrypoint /bin/bash "rocm/vllm-ci:${BUILDKITE_COMMIT}" -ec '
-          if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi
-          if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi
-          if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi
-          if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi
-          command -v python3
-          command -v uv
-          command -v pytest
-          if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then
-            echo No ROCm CLI found in image >&2
-            exit 1
-          fi
-          python3 - <<PY
-        import torch, vllm
-        print(torch.__version__)
-        print(vllm.__version__)
-        PY
-          echo AMD image smoke OK
-        '
+        python3 - <<PY
+      import torch, vllm
+      print(torch.__version__)
+      print(vllm.__version__)
+      PY
+        echo AMD image smoke OK
+      '
+    - docker push "rocm/vllm-ci:${BUILDKITE_COMMIT}"
     env:
       DOCKER_BUILDKIT: "1"
-      VLLM_BAKE_FILE: "docker/docker-bake-rocm.hcl"
-      PYTORCH_ROCM_ARCH: "gfx90a;gfx942;gfx950"
-      IMAGE_TAG: "rocm/vllm-ci:$BUILDKITE_COMMIT"
-      REMOTE_VLLM: "1"
-      VLLM_BRANCH: "$BUILDKITE_COMMIT"
-    retry:
-      automatic:
-        - exit_status: -1  # Agent was lost
-          limit: 1
-        - exit_status: -10  # Agent was lost
-          limit: 1

.buildkite/hardware_tests/cpu.yaml

@@ -28,19 +28,18 @@ steps:
       pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
       pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
 
-# Note: SDE can't be downloaded from CI host because of AWS WAF
-# - label: CPU-Compatibility Tests
-#   depends_on: []
-#   device: intel_cpu
-#   no_plugin: true
-#   source_file_dependencies:
-#   - cmake/cpu_extension.cmake
-#   - setup.py
-#   - vllm/platforms/cpu.py
-#   commands:
-#     - |
-#       bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
-#       bash .buildkite/scripts/hardware_ci/run-cpu-compatibility-test.sh"
+- label: CPU-Compatibility Tests
+  depends_on: []
+  device: intel_cpu
+  no_plugin: true
+  source_file_dependencies:
+  - cmake/cpu_extension.cmake
+  - setup.py
+  - vllm/platforms/cpu.py
+  commands:
+    - |
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
+      bash .buildkite/scripts/hardware_ci/run-cpu-compatibility-test.sh"
 
 - label: CPU-Language Generation and Pooling Model Tests
   depends_on: []

.buildkite/intel_jobs/basic_correctness.yaml

@@ -1,22 +0,0 @@
-group: Basic Correctness
-depends_on:
-  - image-build-xpu
-steps:
-- label: XPU Sleep Mode
-  timeout_in_minutes: 30
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-  - vllm/
-  - tests/basic_correctness/test_cumem.py
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'cd tests &&
-      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
-      pytest -v -s basic_correctness/test_mem.py::test_end_to_end'

.buildkite/intel_jobs/misc_intel.yaml

@@ -53,26 +53,3 @@ steps:
       pytest -v -s v1/test_oracle.py &&
       pytest -v -s v1/test_request.py &&
       pytest -v -s v1/test_outputs.py'
-
-- label: XPU CPU Offload
-  timeout_in_minutes: 60
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-    - vllm/
-    - vllm/v1/kv_offload/
-    - vllm/v1/kv_connector/
-    - tests/v1/kv_offload/
-    - tests/v1/kv_connector/unit/test_offloading_connector.py
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
-      cd tests &&
-      pytest -v -s v1/kv_offload &&
-      pytest -v -s v1/kv_connector/unit/test_offloading_connector.py'

.buildkite/intel_jobs/test-intel.yaml

@@ -40,9 +40,7 @@ steps:
         python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --block-size 64 --enforce-eager --max-model-len 8192 &&
         python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 &&
         python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel &&
-        python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192 &&
-        VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192 &&
-        python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP4-LLMC --enforce-eager -tp 2 --max-model-len 8192
+        python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192
         '
   - label: "XPU V1 test"
     depends_on:

.buildkite/performance-benchmarks/tests/latency-tests-arm64-cpu.json

@@ -2,6 +2,7 @@
     {
         "test_name": "latency_llama8B_tp1",
         "environment_variables": {
+            "VLLM_RPC_TIMEOUT": 100000,
             "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
             "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
             "VLLM_CPU_KVCACHE_SPACE": 40

.buildkite/performance-benchmarks/tests/latency-tests-cpu.json

@@ -2,6 +2,7 @@
     {
         "test_name": "latency_llama8B_tp2",
         "environment_variables": {
+            "VLLM_RPC_TIMEOUT": 100000,
 	    "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
 	    "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
 	    "VLLM_CPU_SGL_KERNEL": 1,

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

@@ -13,6 +13,7 @@
       200
     ],
     "server_environment_variables": {
+      "VLLM_RPC_TIMEOUT": 100000,
       "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
       "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
       "VLLM_CPU_SGL_KERNEL": 1,

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

@@ -5,6 +5,7 @@
     ],
     "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200],
     "server_environment_variables": {
+      "VLLM_RPC_TIMEOUT": 100000,
       "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120
     },
     "server_parameters": {

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

@@ -9,6 +9,7 @@
       128
     ],
     "server_environment_variables": {
+      "VLLM_RPC_TIMEOUT": 100000,
       "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
       "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
       "VLLM_CPU_SGL_KERNEL": 1,

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

@@ -5,6 +5,7 @@
     ],
     "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200],
     "server_environment_variables": {
+      "VLLM_RPC_TIMEOUT": 100000,
       "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
       "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
       "VLLM_CPU_SGL_KERNEL": 1,

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

@@ -5,6 +5,7 @@
     ],
     "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200],
     "server_environment_variables": {
+      "VLLM_RPC_TIMEOUT": 100000,
       "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
       "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
       "VLLM_CPU_SGL_KERNEL": 1,

.buildkite/performance-benchmarks/tests/throughput-tests-arm64-cpu.json

@@ -2,6 +2,7 @@
     {
         "test_name": "throughput_llama8B_tp1",
         "environment_variables": {
+            "VLLM_RPC_TIMEOUT": 100000,
             "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
             "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
             "VLLM_CPU_KVCACHE_SPACE": 40

.buildkite/performance-benchmarks/tests/throughput-tests-cpu.json

@@ -2,6 +2,7 @@
     {
         "test_name": "throughput_llama8B_tp2",
         "environment_variables": {
+            "VLLM_RPC_TIMEOUT": 100000,
 	    "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
 	    "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
 	    "VLLM_CPU_SGL_KERNEL": 1,

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

@@ -28,10 +28,8 @@
 ###############################################################################
 set -o pipefail
 
-# Export Python path for commands that run directly on the host. Containerized
-# tests set this to /vllm-workspace below so spawned Python processes do not
-# depend on their current working directory.
-export PYTHONPATH="${PYTHONPATH:-..}"
+# Export Python path
+export PYTHONPATH=".."
 
 ###############################################################################
 # Helper Functions
@@ -54,108 +52,6 @@ cleanup_network() {
   fi
 }
 
-prepare_artifact_image() {
-  if [[ "${VLLM_CI_USE_ARTIFACTS:-0}" != "1" ]]; then
-    return 1
-  fi
-  if ! command -v buildkite-agent >/dev/null 2>&1; then
-    echo "buildkite-agent not found; cannot download ROCm wheel artifact"
-    return 1
-  fi
-
-  local artifact_glob="${VLLM_CI_ARTIFACT_GLOB:-artifacts/vllm-rocm-install/vllm-rocm-install.tar.gz}"
-  local archive=""
-  local metadata_file=""
-  local base_image="${VLLM_CI_BASE_IMAGE:-rocm/vllm-dev:ci_base}"
-  local artifact_image=""
-  local artifact_key=""
-  local base_digest=""
-  local wheel_dir=""
-  local context_dir=""
-  local workspace_dir=""
-
-  artifact_work_dir=$(mktemp -d -t vllm-rocm-artifact.XXXXXX)
-  wheel_dir="${artifact_work_dir}/wheels"
-  context_dir="${artifact_work_dir}/context"
-  workspace_dir="${context_dir}/workspace"
-  mkdir -p "${wheel_dir}" "${context_dir}/wheels" "${workspace_dir}"
-
-  echo "--- Downloading ROCm wheel artifact"
-  if ! buildkite-agent artifact download "${artifact_glob}" "${artifact_work_dir}"; then
-    echo "Failed to download ${artifact_glob}"
-    return 1
-  fi
-  buildkite-agent artifact download \
-    "artifacts/vllm-rocm-install/ci-base-image.txt" \
-    "${artifact_work_dir}" >/dev/null 2>&1 || true
-
-  archive=$(find "${artifact_work_dir}" -name "vllm-rocm-install.tar.gz" -type f | head -1)
-  if [[ -z "${archive}" || ! -f "${archive}" ]]; then
-    echo "ROCm wheel artifact archive was not found"
-    return 1
-  fi
-
-  metadata_file=$(find "${artifact_work_dir}" -name "ci-base-image.txt" -type f | head -1)
-  if [[ -n "${metadata_file}" && -s "${metadata_file}" ]]; then
-    base_image=$(tr -d '[:space:]' < "${metadata_file}")
-  fi
-
-  echo "--- Preparing local ROCm test image"
-  echo "Base image: ${base_image}"
-  docker pull "${base_image}" || return 1
-  base_digest=$(
-    docker image inspect \
-      --format='{{if .RepoDigests}}{{index .RepoDigests 0}}{{else}}{{.Id}}{{end}}' \
-      "${base_image}" 2>/dev/null || printf '%s' "${base_image}"
-  )
-
-  artifact_key=$(
-    {
-      printf 'base-image:%s\n' "${base_digest}"
-      sha256sum "${archive}"
-    } | sha256sum | cut -c1-24
-  )
-  artifact_image="rocm/vllm-ci-artifact:${artifact_key}"
-
-  if docker image inspect "${artifact_image}" >/dev/null 2>&1; then
-    echo "Using existing local ROCm artifact image: ${artifact_image}"
-    image_name="${artifact_image}"
-    return 0
-  fi
-
-  tar -xzf "${archive}" -C "${wheel_dir}" || return 1
-  if ! ls "${wheel_dir}"/*.whl >/dev/null 2>&1; then
-    echo "ROCm wheel artifact did not contain a wheel"
-    return 1
-  fi
-  if [[ ! -d "${wheel_dir}/tests" ]]; then
-    echo "ROCm wheel artifact did not contain the test workspace"
-    return 1
-  fi
-
-  cp "${wheel_dir}"/*.whl "${context_dir}/wheels/" || return 1
-  tar -C "${wheel_dir}" --exclude='*.whl' -cf - . \
-    | tar -C "${workspace_dir}" -xf - || return 1
-  cat > "${context_dir}/Dockerfile" <<'EOF'
-ARG BASE_IMAGE
-FROM ${BASE_IMAGE}
-COPY wheels/ /tmp/vllm-wheels/
-COPY workspace/ /vllm-workspace/
-RUN python3 -m pip install --no-deps --force-reinstall /tmp/vllm-wheels/*.whl \
-    && rm -rf /tmp/vllm-wheels
-WORKDIR /vllm-workspace
-EOF
-
-  echo "--- Building local ROCm test image"
-  docker build \
-    --pull=false \
-    --build-arg "BASE_IMAGE=${base_image}" \
-    -t "${artifact_image}" \
-    "${context_dir}" || return 1
-  image_name="${artifact_image}"
-  return 0
-}
-
 is_multi_node() {
   local cmds="$1"
   # Primary signal: NUM_NODES environment variable set by the pipeline
@@ -347,30 +243,22 @@ report_docker_usage
 
 # --- Pull test image ---
 echo "--- Pulling container"
-image_name="${VLLM_CI_FALLBACK_IMAGE:-rocm/vllm-ci:${BUILDKITE_COMMIT:-local}}"
-artifact_work_dir=""
+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() {
-  if docker container inspect "${container_name}" >/dev/null 2>&1; then
-    docker rm -f "${container_name}" || true
-  fi
-  if [[ "${VLLM_CI_REMOVE_TEST_IMAGE:-0}" == "1" ]]; then
-    docker image rm -f "${image_name}" || true
-  else
-    # Keep images by default so later jobs on the same AMD node can reuse layers.
-    echo "Keeping ROCm test image locally: ${image_name}"
-  fi
-  if [[ -n "${artifact_work_dir}" ]]; then
-    rm -rf "${artifact_work_dir}"
-  fi
+  # docker run uses --rm, so the container is normally already gone when the
+  # EXIT trap runs. Cleanup is best-effort and must not affect the test result.
+  docker rm -f "${container_name}" >/dev/null 2>&1 || true
 }
-trap remove_docker_container EXIT
 
-if ! prepare_artifact_image; then
-  echo "Using full ROCm CI image: ${image_name}"
-  docker pull "${image_name}" || exit 1
-fi
+on_exit() {
+  local exit_code=$?
+  remove_docker_container
+  exit "$exit_code"
+}
+trap on_exit EXIT
 
 # --- Prepare commands ---
 echo "--- Running container"
@@ -379,14 +267,6 @@ HF_CACHE="$(realpath ~)/huggingface"
 mkdir -p "${HF_CACHE}"
 HF_MOUNT="/root/.cache/huggingface"
 
-# Hugging Face Hub defaults to 10s request/download timeouts, while the ROCm
-# CI image currently raises downloads to 60s. AMD model-test jobs routinely
-# start from a cold or partially-populated shared cache, and the 60s read cap
-# has still timed out before pytest reached the vLLM behavior under test.
-# Keep the CI default explicit and overridable from the Buildkite environment.
-: "${HF_HUB_DOWNLOAD_TIMEOUT:=300}"
-: "${HF_HUB_ETAG_TIMEOUT:=60}"
-
 # ---- Command source selection ----
 # Prefer VLLM_TEST_COMMANDS (preserves all inner quoting intact).
 # Fall back to $* for backward compatibility, but warn that inner
@@ -426,14 +306,7 @@ fi
 
 echo "Final commands: $commands"
 
-MYPYTHONPATH="/vllm-workspace"
-
-container_job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-0}}"
-container_job_id="${container_job_id//[^A-Za-z0-9_.-]/_}"
-container_job_id_short="${container_job_id:0:8}"
-CONTAINER_TMPDIR="/tmp/vllm-${container_job_id_short}"
-CONTAINER_CACHE_ROOT="/tmp/vllm-buildkite-${container_job_id}/cache"
-CONTAINER_PREFLIGHT="mkdir -p \"\$TMPDIR\" \"\$TORCHINDUCTOR_CACHE_DIR\" \"\$TRITON_CACHE_DIR\" \"\$VLLM_CACHE_ROOT\" \"\$XDG_CACHE_HOME\" && python -c \"import encodings, importlib.metadata as im, importlib.util as iu; [im.version(d) for d in ('transformers', 'torch', 'ray', 'sympy', 'markupsafe', 'vllm')]; missing=[m for m in ('torch.utils.model_zoo', 'transformers.models.nomic_bert', 'ray.dag', 'sympy.physics', 'markupsafe._speedups') if iu.find_spec(m) is None]; assert not missing, missing\""
+MYPYTHONPATH=".."
 
 # Verify GPU access
 render_gid=$(getent group render | cut -d: -f3)
@@ -510,8 +383,6 @@ else
     --group-add "$render_gid" \
     --rm \
     -e HF_TOKEN \
-    -e "HF_HUB_DOWNLOAD_TIMEOUT=${HF_HUB_DOWNLOAD_TIMEOUT}" \
-    -e "HF_HUB_ETAG_TIMEOUT=${HF_HUB_ETAG_TIMEOUT}" \
     -e AWS_ACCESS_KEY_ID \
     -e AWS_SECRET_ACCESS_KEY \
     -e BUILDKITE_PARALLEL_JOB \
@@ -519,15 +390,10 @@ else
     -v "${HF_CACHE}:${HF_MOUNT}" \
     -e "HF_HOME=${HF_MOUNT}" \
     -e "PYTHONPATH=${MYPYTHONPATH}" \
-    -e "TMPDIR=${CONTAINER_TMPDIR}/tmp" \
-    -e "TORCHINDUCTOR_CACHE_DIR=${CONTAINER_CACHE_ROOT}/torchinductor" \
-    -e "TRITON_CACHE_DIR=${CONTAINER_CACHE_ROOT}/triton" \
-    -e "VLLM_CACHE_ROOT=${CONTAINER_CACHE_ROOT}/vllm" \
-    -e "XDG_CACHE_HOME=${CONTAINER_CACHE_ROOT}/xdg" \
     -e "PYTORCH_ROCM_ARCH=" \
     --name "${container_name}" \
     "${image_name}" \
-    /bin/bash -c "${CONTAINER_PREFLIGHT} && ${commands}"
+    /bin/bash -c "${commands}"
 
   exit_code=$?
   handle_pytest_exit "$exit_code"

.buildkite/test-amd.yaml

@@ -315,6 +315,24 @@ steps:
   - pytest -v -s distributed/test_pp_cudagraph.py
   - pytest -v -s distributed/test_pipeline_parallel.py
 
+#----------------------------------------------------------  mi250 · engine  -----------------------------------------------------------#
+
+- label: Engine # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  agent_pool: mi250_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/engine
+  - tests/test_sequence
+  - tests/test_config
+  - tests/test_logger
+  - tests/test_vllm_port
+  commands:
+  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
+
 #-----------------------------------------------------------  mi250 · evals  -----------------------------------------------------------#
 
 - label: Multi-Modal Accuracy Eval (Small Models) # TBD
@@ -431,6 +449,29 @@ steps:
   commands:
   - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
 
+#------------------------------------------------------  mi250 · model_executor  -------------------------------------------------------#
+
+- label: Model Executor # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  agent_pool: mi250_1
+  optional: true
+  torch_nightly: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/engine/arg_utils.py
+  - vllm/config/model.py
+  - vllm/model_executor
+  - tests/model_executor
+  - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  commands:
+  - apt-get update && apt-get install -y curl libsodium23
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s model_executor -m '(not slow_test)'
+  - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
+
 #------------------------------------------------------  mi250 · models / basic  -------------------------------------------------------#
 
 - label: Basic Models Test (Other CPU) # TBD
@@ -890,10 +931,10 @@ steps:
   - vllm/
   - tests/basic_correctness/test_basic_correctness
   - tests/basic_correctness/test_cpu_offload
-  - tests/basic_correctness/test_mem.py
+  - tests/basic_correctness/test_cumem.py
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s basic_correctness/test_mem.py
+  - pytest -v -s basic_correctness/test_cumem.py
   - pytest -v -s basic_correctness/test_basic_correctness.py
   - pytest -v -s basic_correctness/test_cpu_offload.py
 
@@ -1281,11 +1322,12 @@ steps:
   source_file_dependencies:
   - vllm/
   - tests/entrypoints/llm
+  - tests/entrypoints/offline_mode
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py --ignore=entrypoints/llm/offline_mode
-  - pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
-  - pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
+  - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py
+  - pytest -v -s entrypoints/llm/test_generate.py
+  - pytest -v -s entrypoints/offline_mode
 
 - label: Entrypoints Integration (Pooling) # TBD
   timeout_in_minutes: 180
@@ -1327,13 +1369,12 @@ steps:
   - vllm/platforms/rocm.py
   commands:
   - pytest -v -s entrypoints/openai/tool_parsers
-  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
+  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
 
 - label: OpenAI API correctness # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_1
-  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - csrc/
@@ -1392,7 +1433,6 @@ steps:
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_1
-  optional: true
   working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
   source_file_dependencies:
   - csrc/
@@ -1699,29 +1739,6 @@ steps:
   - pytest -v -s -x lora/test_gptoss_tp.py
   - pytest -v -s -x lora/test_qwen35_densemodel_lora.py
 
-#------------------------------------------------------  mi300 · model_executor  -------------------------------------------------------#
-
-- label: Model Executor # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  optional: true
-  torch_nightly: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/engine/arg_utils.py
-  - vllm/config/model.py
-  - vllm/model_executor
-  - tests/model_executor
-  - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
-  - vllm/_aiter_ops.py
-  - vllm/platforms/rocm.py
-  commands:
-  - apt-get update && apt-get install -y curl libsodium23
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s model_executor -m '(not slow_test)'
-  - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
-
 #-----------------------------------------------------  mi300 · models / language  -----------------------------------------------------#
 
 - label: Language Models Test (Extended Pooling)  # TBD
@@ -2167,72 +2184,10 @@ steps:
   commands:
   - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
 
-- label: Speculators Correctness # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/config/speculative.py
-  - vllm/engine/arg_utils.py
-  - vllm/transformers_utils/config.py
-  - vllm/transformers_utils/configs/speculators/
-  - vllm/v1/spec_decode/
-  - vllm/v1/worker/gpu/spec_decode/
-  - vllm/v1/worker/gpu_model_runner.py
-  - vllm/v1/sample/
-  - vllm/v1/attention/backends/
-  - vllm/v1/attention/selector.py
-  - vllm/model_executor/model_loader/
-  - vllm/model_executor/layers/
-  - vllm/model_executor/models/llama_eagle3.py
-  - vllm/model_executor/models/qwen3.py
-  - vllm/model_executor/models/qwen3_dflash.py
-  - vllm/model_executor/models/registry.py
-  - vllm/_aiter_ops.py
-  - tests/evals/gsm8k/
-  - tests/v1/spec_decode/test_speculators_correctness.py
-  - vllm/platforms/rocm.py
-  commands:
-  - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
-  - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
-
-- label: Extract Hidden States Integration # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/config/speculative.py
-  - vllm/distributed/kv_transfer/kv_connector/
-  - vllm/model_executor/layers/attention/
-  - vllm/model_executor/layers/mamba/
-  - vllm/model_executor/model_loader/
-  - vllm/model_executor/models/extract_hidden_states.py
-  - vllm/model_executor/models/llama.py
-  - vllm/model_executor/models/qwen3_5.py
-  - vllm/model_executor/models/qwen3_next.py
-  - vllm/model_executor/models/registry.py
-  - vllm/transformers_utils/configs/extract_hidden_states.py
-  - vllm/transformers_utils/configs/qwen3_5.py
-  - vllm/v1/attention/backends/
-  - vllm/v1/attention/selector.py
-  - vllm/v1/kv_cache_interface.py
-  - vllm/v1/spec_decode/extract_hidden_states.py
-  - vllm/v1/worker/gpu_model_runner.py
-  - vllm/_aiter_ops.py
-  - tests/v1/kv_connector/extract_hidden_states_integration/
-  - vllm/platforms/rocm.py
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s v1/kv_connector/extract_hidden_states_integration
-
 - label: V1 attention (H100-MI300) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_1
-  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/config/attention.py
@@ -2591,7 +2546,6 @@ steps:
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
   agent_pool: mi325_1
-  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
@@ -2601,7 +2555,7 @@ steps:
   - tests/test_logger
   - tests/test_vllm_port
   commands:
-  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
+  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
 
 #-----------------------------------------------------------  mi325 · evals  -----------------------------------------------------------#
 
@@ -2683,7 +2637,6 @@ steps:
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
   agent_pool: mi325_1
-  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/

.buildkite/test_areas/attention.yaml

@@ -2,8 +2,8 @@ group: Attention
 depends_on: 
   - image-build
 steps:
-- label: V1 attention (H100-MI300)
-  key: v1-attention-h100-mi300
+- label: V1 attention (H100)
+  key: v1-attention-h100
   timeout_in_minutes: 30
   device: h100
   source_file_dependencies:
@@ -13,20 +13,6 @@ steps:
     - tests/v1/attention
   commands:
     - pytest -v -s v1/attention
-  mirror:
-    amd:
-      device: mi300_1
-      timeout_in_minutes: 70
-      depends_on:
-      - image-build-amd
-      source_file_dependencies:
-      - vllm/config/attention.py
-      - vllm/model_executor/layers/attention
-      - vllm/v1/attention
-      - tests/v1/attention
-      - vllm/_aiter_ops.py
-      - vllm/envs.py
-      - vllm/platforms/rocm.py
 
 - label: V1 attention (B200)
   key: v1-attention-b200

.buildkite/test_areas/basic_correctness.yaml

@@ -10,9 +10,9 @@ steps:
   - vllm/
   - tests/basic_correctness/test_basic_correctness
   - tests/basic_correctness/test_cpu_offload
-  - tests/basic_correctness/test_mem.py
+  - tests/basic_correctness/test_cumem.py
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s basic_correctness/test_mem.py
+  - pytest -v -s basic_correctness/test_cumem.py
   - pytest -v -s basic_correctness/test_basic_correctness.py
   - pytest -v -s basic_correctness/test_cpu_offload.py

.buildkite/test_areas/engine.yaml

@@ -26,12 +26,6 @@ steps:
   - tests/test_jit_monitor.py
   commands:
   - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
-  mirror:
-    amd:
-      device: mi300_1
-      timeout_in_minutes: 60
-      depends_on:
-      - image-build-amd
 
 - label: Engine (1 GPU)
   key: engine-1-gpu
@@ -44,7 +38,7 @@ steps:
     - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 40
       depends_on:
       - image-build-amd

.buildkite/test_areas/entrypoints.yaml

@@ -11,7 +11,7 @@ steps:
   - tests/entrypoints/
   commands:
   - pytest -v -s entrypoints/openai/tool_parsers
-  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
+  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
 
 - label: Entrypoints Integration (LLM)
   key: entrypoints-integration-llm
@@ -20,14 +20,16 @@ steps:
   source_file_dependencies:
   - vllm/
   - tests/entrypoints/llm
+  - tests/entrypoints/offline_mode
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py --ignore=entrypoints/llm/offline_mode
+  - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py
   - pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
-  - pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
+  - pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
   mirror:
     amd:
       device: mi325_1
+      soft_fail: true
       depends_on:
       - image-build-amd
 
@@ -44,7 +46,8 @@ steps:
   - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
+      soft_fail: true
       timeout_in_minutes: 80
       depends_on:
       - image-build-amd
@@ -66,7 +69,8 @@ steps:
   - pytest -v -s tool_use
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
+      soft_fail: true
       timeout_in_minutes: 60
       depends_on:
       - image-build-amd
@@ -86,6 +90,7 @@ steps:
   mirror:
     amd:
       device: mi325_1
+      soft_fail: true
       timeout_in_minutes: 60
       depends_on:
       - image-build-amd
@@ -104,7 +109,8 @@ steps:
   - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
+      soft_fail: true
       depends_on:
       - image-build-amd
 
@@ -150,21 +156,3 @@ steps:
   - vllm/entrypoints/openai/
   commands: # LMEval
   - pytest -s entrypoints/openai/correctness/
-  mirror:
-    amd:
-      device: mi300_1
-      depends_on:
-      - image-build-amd
-      source_file_dependencies:
-      - csrc/
-      - vllm/entrypoints/openai/
-      - vllm/model_executor/models/whisper.py
-      - vllm/model_executor/layers/
-      - vllm/v1/attention/backends/
-      - vllm/v1/attention/selector.py
-      - vllm/_aiter_ops.py
-      - vllm/platforms/rocm.py
-      - vllm/model_executor/model_loader/
-      commands:
-      - bash ../tools/install_torchcodec_rocm.sh || exit 1
-      - pytest -s entrypoints/openai/correctness/

.buildkite/test_areas/kernels.yaml

@@ -86,7 +86,7 @@ steps:
   parallelism: 2
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       source_file_dependencies:
       - csrc/quantization/
       - vllm/model_executor/layers/quantization

.buildkite/test_areas/lm_eval.yaml

@@ -12,21 +12,6 @@ steps:
   autorun_on_main: true
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
-  mirror:
-    amd:
-      device: mi300_1
-      timeout_in_minutes: 55
-      depends_on:
-      - image-build-amd
-      source_file_dependencies:
-      - csrc/
-      - vllm/model_executor/layers/quantization
-      - vllm/model_executor/models/
-      - vllm/model_executor/model_loader/
-      - vllm/v1/attention/backends/
-      - vllm/v1/attention/selector.py
-      - vllm/_aiter_ops.py
-      - vllm/platforms/rocm.py
 
 # - label: LM Eval Large Models (4 GPUs)(A100)
 #   device: a100

.buildkite/test_areas/misc.yaml

@@ -281,7 +281,6 @@ steps:
   - vllm/model_executor/layers/quantization/quark/
   - vllm/multimodal/
   - vllm/outputs.py
-  - vllm/parser/
   - vllm/platforms/
   - vllm/pooling_params.py
   - vllm/ray/

.buildkite/test_areas/model_runner_v2.yaml

@@ -94,13 +94,11 @@ steps:
     - vllm/v1/worker/gpu_worker.py
     - tests/distributed/test_pipeline_parallel.py
     - tests/distributed/test_pp_cudagraph.py
-    - tests/v1/distributed/test_pp_dp_v2.py
   commands:
     - set -x
     - export VLLM_USE_V2_MODEL_RUNNER=1
     - pytest -v -s distributed/test_pipeline_parallel.py -k "not ray and not Jamba"
     - pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
-    - pytest -v -s v1/distributed/test_pp_dp_v2.py
 
 - label: Model Runner V2 Spec Decode
   device: h200_35gb

.buildkite/test_areas/models_language.yaml

@@ -51,7 +51,6 @@ steps:
     torch_nightly: {}
     amd:
       device: mi325_1
-      timeout_in_minutes: 90
       depends_on:
       - image-build-amd
       commands:

.buildkite/test_areas/models_multimodal.yaml

@@ -15,7 +15,7 @@ steps:
     - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
 
@@ -33,7 +33,7 @@ steps:
     - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
 
@@ -50,7 +50,7 @@ steps:
     - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
 
@@ -153,12 +153,3 @@ steps:
   - tests/models/multimodal/pooling
   commands:
     - pytest -v -s models/multimodal/pooling -m 'not core_model'
-  mirror:
-    amd:
-      device: mi300_1
-      timeout_in_minutes: 60
-      depends_on:
-      - image-build-amd
-      source_file_dependencies:
-      - vllm/
-      - tests/models/multimodal/pooling

.buildkite/test_areas/spec_decode.yaml

@@ -37,21 +37,6 @@ steps:
     - tests/v1/e2e/spec_decode/
   commands:
     - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
-  mirror:
-    amd:
-      device: mi300_1
-      timeout_in_minutes: 65
-      depends_on:
-      - image-build-amd
-      source_file_dependencies:
-      - vllm/v1/spec_decode/
-      - vllm/v1/worker/gpu/spec_decode/
-      - vllm/model_executor/model_loader/
-      - vllm/v1/sample/
-      - vllm/model_executor/layers/
-      - vllm/transformers_utils/configs/speculators/
-      - tests/v1/e2e/spec_decode/
-      - vllm/platforms/rocm.py
 
 - label: Spec Decode Speculators + MTP Nightly B200
   key: spec-decode-speculators-mtp-nightly-b200
@@ -76,20 +61,6 @@ steps:
     - tests/v1/e2e/spec_decode/
   commands:
     - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
-  mirror:
-    amd:
-      device: mi300_1
-      timeout_in_minutes: 65
-      depends_on:
-      - image-build-amd
-      source_file_dependencies:
-      - vllm/v1/spec_decode/
-      - vllm/v1/worker/gpu/spec_decode/
-      - vllm/model_executor/model_loader/
-      - vllm/v1/sample/
-      - vllm/model_executor/layers/
-      - tests/v1/e2e/spec_decode/
-      - vllm/platforms/rocm.py
 
 - label: Spec Decode Draft Model
   key: spec-decode-draft-model
@@ -101,20 +72,6 @@ steps:
     - tests/v1/e2e/spec_decode/
   commands:
     - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
-  mirror:
-    amd:
-      device: mi300_1
-      timeout_in_minutes: 50
-      depends_on:
-      - image-build-amd
-      source_file_dependencies:
-      - vllm/v1/spec_decode/
-      - vllm/v1/worker/gpu/spec_decode/
-      - vllm/model_executor/model_loader/
-      - vllm/v1/sample/
-      - vllm/model_executor/layers/
-      - tests/v1/e2e/spec_decode/
-      - vllm/platforms/rocm.py
 
 - label: Spec Decode Draft Model Nightly B200
   key: spec-decode-draft-model-nightly-b200

.github/CODEOWNERS

@@ -34,11 +34,10 @@
 /vllm/entrypoints/speech_to_text/realtime @njhill
 /vllm/entrypoints/speech_to_text @NickLucche
 /vllm/entrypoints/pooling @noooop
-/vllm/entrypoints/serve/sagemaker @DarkLight1337
+/vllm/entrypoints/sagemaker @DarkLight1337
 /vllm/entrypoints/serve @njhill
 /vllm/entrypoints/*.py @njhill
 /vllm/entrypoints/chat_utils.py @DarkLight1337
-/vllm/entrypoints/offline_utils.py @DarkLight1337
 /vllm/entrypoints/llm.py @DarkLight1337
 
 # Rust Frontend

.github/mergify.yml

@@ -21,9 +21,6 @@ pull_request_rules:
     - check-failure=pre-commit
     - -closed
     - -draft
-    - or:
-      - label=ready
-      - label=verified
   actions:
     comment:
       message: |
@@ -39,6 +36,18 @@ pull_request_rules:
 
         For future commits, `pre-commit` will run automatically on changed files before each commit.
 
+        > [!TIP]
+        > <details>
+        > <summary>Is <code>mypy</code> failing?</summary>
+        > <br/>
+        > <code>mypy</code> is run differently in CI. If the failure is related to this check, please use the following command to run it locally:
+        >
+        > ```bash
+        > # For mypy (substitute "3.10" with the failing version if needed)
+        > pre-commit run --hook-stage manual mypy-3.10
+        > ```
+        > </details>
+
 - name: comment-dco-failure
   description: Comment on PR when DCO check fails
   conditions:

.github/workflows/pre-commit.yml

@@ -49,7 +49,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
-    - uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
+    - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
       with:
         python-version: "3.12"
     - run: echo "::add-matcher::.github/workflows/matchers/actionlint.json"

.github/workflows/stale.yml

@@ -15,7 +15,7 @@ jobs:
       actions: write
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/stale@eb5cf3af3ac0a1aa4c9c45633dd1ae542a27a899 # v10.3.0
+      - uses: actions/stale@997185467fa4f803885201cee163a9f38240193d # v10.1.1
         with:
           # Increasing this value ensures that changes to this workflow
           # propagate to all issues and PRs in days rather than months

.pre-commit-config.yaml

@@ -148,27 +148,33 @@ repos:
     language: python
     entry: python tools/pre_commit/generate_nightly_torch_test.py
     files: ^requirements/test/cuda\.(in|txt)$
-  - id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
-    name: Run mypy for Python 3.10
-    entry: python tools/pre_commit/mypy.py "3.10"
+  - id: mypy-local
+    name: Run mypy locally for lowest supported Python version
+    entry: python tools/pre_commit/mypy.py 0 "3.10"
+    stages: [pre-commit] # Don't run in CI
     <<: &mypy_common
       language: python
       types_or: [python, pyi]
       require_serial: true
-      additional_dependencies: ["mypy==1.20.2", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
+      additional_dependencies: ["mypy[faster-cache]==1.19.1", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
+  - id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
+    name: Run mypy for Python 3.10
+    entry: python tools/pre_commit/mypy.py 1 "3.10"
+    <<: *mypy_common
+    stages: [manual] # Only run in CI
   - id: mypy-3.11 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
     name: Run mypy for Python 3.11
-    entry: python tools/pre_commit/mypy.py "3.11"
+    entry: python tools/pre_commit/mypy.py 1 "3.11"
     <<: *mypy_common
     stages: [manual] # Only run in CI
   - id: mypy-3.12 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
     name: Run mypy for Python 3.12
-    entry: python tools/pre_commit/mypy.py "3.12"
+    entry: python tools/pre_commit/mypy.py 1 "3.12"
     <<: *mypy_common
     stages: [manual] # Only run in CI
   - id: mypy-3.13 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
     name: Run mypy for Python 3.13
-    entry: python tools/pre_commit/mypy.py "3.13"
+    entry: python tools/pre_commit/mypy.py 1 "3.13"
     <<: *mypy_common
     stages: [manual] # Only run in CI
   - id: shellcheck

AGENTS.md

@@ -98,13 +98,11 @@ pre-commit run --all-files
 pre-commit run ruff-check --all-files
 
 # Run mypy as it is in CI:
-pre-commit run mypy-3.12 --all-files --hook-stage manual
+pre-commit run mypy-3.10 --all-files --hook-stage manual
 ```
 
 The line length limit for Python code is 88 characters. If you are not sure, use pre-commit to check.
 
-Use [Google-style docstrings](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings) (`Args:`/`Returns:`/`Raises:` sections), not reStructuredText/Sphinx fields (`:param:`, `:return:`, `:rtype:`).
-
 ### Commit messages
 
 Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:

CMakeLists.txt

@@ -112,8 +112,6 @@ endif()
 #
 # spinloop extension (pure CXX; must stay above the non-CUDA device branch so
 # CPU builds define the target before the early return)
-# This extension requires SABI 3.11 since it relies on Py_buffer support. Loading
-# failure is handled gracefully on vLLM side for lower Python versions.
 #
 set(VLLM_SPINLOOP_EXT_SRC "csrc/spinloop.cpp")
 set(SPINLOOP_COMPILE_FLAGS "")
@@ -307,10 +305,18 @@ endif()
 #
 
 set(VLLM_EXT_SRC
+  "csrc/cuda_view.cu"
+  "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
   "csrc/quantization/activation_kernels.cu"
-  "csrc/torch_bindings.cpp")
+  "csrc/cuda_utils_kernels.cu"
+  "csrc/custom_all_reduce.cu"
+  "csrc/torch_bindings.cpp"
+  "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
 
 if(VLLM_GPU_LANG STREQUAL "CUDA")
+  list(APPEND VLLM_EXT_SRC
+    "csrc/minimax_reduce_rms_kernel.cu")
+
   SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
 
   # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
@@ -343,6 +349,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   endif()
   FetchContent_MakeAvailable(cutlass)
 
+  list(APPEND VLLM_EXT_SRC
+    "csrc/cutlass_extensions/common.cpp")
+
   set_gencode_flags_for_srcs(
     SRCS "${VLLM_EXT_SRC}"
     CUDA_ARCHS "${CUDA_ARCHS}")
@@ -494,12 +503,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   endif()
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND ES_MXFP8_GROUPED_MM_ARCHS)
     set(SRCS
-      "csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu"
-      "csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu")
+      "csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu"
+      "csrc/moe/mxfp8_moe/mxfp8_experts_quant.cu")
     set_gencode_flags_for_srcs(
       SRCS "${SRCS}"
       CUDA_ARCHS "${ES_MXFP8_GROUPED_MM_ARCHS}")
-    list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
+    list(APPEND VLLM_EXT_SRC "${SRCS}")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_ES_MXFP8_GROUPED_MM_SM100=1")
     message(STATUS "Building ES MXFP8 grouped kernels for archs: ${ES_MXFP8_GROUPED_MM_ARCHS}")
   else()
@@ -589,12 +598,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
 endif()
 
 if (VLLM_GPU_LANG STREQUAL "HIP")
-  # Add QuickReduce kernels (ROCm-only; not part of stable ABI migration).
-  # TODO: Remove the cuda_view when ROCm upgrade to torch 2.11.
+  # Add QuickReduce kernels
   list(APPEND VLLM_EXT_SRC
     "csrc/custom_quickreduce.cu"
-    "csrc/cuda_view.cu"
-    "csrc/libtorch_stable/cuda_utils_kernels.cu"
   )
 # if ROCM endif
 endif()
@@ -627,9 +633,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/activation_kernels.cu"
     "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
     "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
-    "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
-    "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
-    "csrc/libtorch_stable/permute_cols.cu"
     "csrc/libtorch_stable/quantization/gptq/q_gemm.cu"
     "csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu"
     "csrc/libtorch_stable/pos_encoding_kernels.cu"
@@ -637,7 +640,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/layernorm_kernels.cu"
     "csrc/libtorch_stable/layernorm_quant_kernels.cu"
     "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
-    "csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
     "csrc/libtorch_stable/attention/merge_attn_states.cu"
     "csrc/libtorch_stable/sampler.cu"
     "csrc/libtorch_stable/topk.cu"
@@ -645,21 +647,19 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/attention/paged_attention_v1.cu"
     "csrc/libtorch_stable/attention/paged_attention_v2.cu"
     "csrc/libtorch_stable/cache_kernels.cu"
-    "csrc/libtorch_stable/cache_kernels.cu"
-    "csrc/libtorch_stable/cache_kernels_fused.cu"
-    "csrc/libtorch_stable/custom_all_reduce.cu"
-    "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
+    "csrc/libtorch_stable/cache_kernels_fused.cu")
 
   if(VLLM_GPU_LANG STREQUAL "CUDA")
     list(APPEND VLLM_STABLE_EXT_SRC
-      "csrc/libtorch_stable/cuda_view.cu"
-      "csrc/libtorch_stable/cuda_utils_kernels.cu"
-      "csrc/libtorch_stable/cutlass_extensions/common.cpp"
+      "csrc/cuda_utils_kernels.cu"
+      "csrc/cutlass_extensions/common.cpp"
       "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
       "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
       "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu"
-      "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu"
-      "csrc/libtorch_stable/minimax_reduce_rms_kernel.cu")
+      "csrc/libtorch_stable/permute_cols.cu"
+      "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
+      "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
+      "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu")
 
     set_gencode_flags_for_srcs(
       SRCS "${VLLM_STABLE_EXT_SRC}"
@@ -1072,25 +1072,20 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     USE_SABI 3
     WITH_SOABI)
 
+  # Set TORCH_TARGET_VERSION for stable ABI compatibility.
+  # This ensures we only use C-shim APIs available in PyTorch 2.10.
+  # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
+  # which is currently set to 2.10.
+  target_compile_definitions(_C_stable_libtorch PRIVATE
+    TORCH_TARGET_VERSION=0x020A000000000000ULL)
+
   # Needed to use cuda/hip APIs from C-shim
   if(VLLM_GPU_LANG STREQUAL "CUDA")
-    # Set TORCH_TARGET_VERSION for stable ABI compatibility.
-    # This ensures we only use C-shim APIs available in PyTorch 2.11.
-    # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
-    # which is currently set to 2.11.
-    target_compile_definitions(_C_stable_libtorch PRIVATE
-      TORCH_TARGET_VERSION=0x020B000000000000ULL)
     target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
     # Needed by CUTLASS kernels
     target_compile_definitions(_C_stable_libtorch PRIVATE
       CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
   elseif(VLLM_GPU_LANG STREQUAL "HIP")
-    # Set TORCH_TARGET_VERSION for stable ABI compatibility.
-    # This ensures we only use C-shim APIs available in PyTorch 2.10.
-    # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
-    # which is currently set to 2.10.
-    target_compile_definitions(_C_stable_libtorch PRIVATE
-      TORCH_TARGET_VERSION=0x020A000000000000ULL)
     target_compile_definitions(_C_stable_libtorch PRIVATE USE_ROCM)
   endif()
 
@@ -1308,8 +1303,7 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
     set(VLLM_ROCM_HAS_GFX1100 ON)
     list(APPEND VLLM_ROCM_EXT_SRC
       "csrc/rocm/q_gemm_rdna3.cu"
-      "csrc/rocm/q_gemm_rdna3_wmma.cu"
-      "csrc/rocm/moe_q_gemm_rdna3.cu")
+      "csrc/rocm/q_gemm_rdna3_wmma.cu")
   endif()
 
   define_extension_target(

benchmarks/kernels/benchmark_moe.py

@@ -271,6 +271,7 @@ def benchmark_config(
                     moe_config=moe_config,
                     quant_config=quant_config,
                 ),
+                inplace=not disable_inplace(),
             )
 
         with override_config(config):
@@ -278,6 +279,7 @@ def benchmark_config(
                 x, input_gating, topk, renormalize=not use_deep_gemm
             )
 
+            inplace = not disable_inplace()
             if use_deep_gemm:
                 return deep_gemm_experts.apply(
                     x,
@@ -296,6 +298,7 @@ def benchmark_config(
                 w2,
                 topk_weights,
                 topk_ids,
+                inplace=inplace,
                 quant_config=quant_config,
             )
 

cmake/hipify.py

@@ -14,18 +14,7 @@ import argparse
 import os
 import shutil
 
-from torch.utils.hipify.hipify_python import get_hip_file_path, hipify
-
-
-def _expected_hip_build_path(source_abs: str, output_directory: str) -> str:
-    """Match torch.utils.hipify.hipify_python.preprocessor fout_path naming."""
-    rel = os.path.relpath(source_abs, output_directory)
-    return os.path.abspath(
-        os.path.join(
-            output_directory, get_hip_file_path(rel, is_pytorch_extension=True)
-        )
-    )
-
+from torch.utils.hipify.hipify_python import hipify
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -64,11 +53,7 @@ if __name__ == "__main__":
     hipify_result = hipify(
         project_directory=args.project_dir,
         output_directory=args.output_dir,
-        # Hipify resolves quoted includes next to the including file first; vLLM
-        # uses paths relative to csrc/ (e.g. "libtorch_stable/torch_utils.h"
-        # from quantization/w8a8/fp8/*.cu). Without an include root here, those
-        # headers are never found and are not hipified or rewritten in dependents.
-        header_include_dirs=["."],
+        header_include_dirs=[],
         includes=includes,
         extra_files=extra_files,
         show_detailed=True,
@@ -79,20 +64,14 @@ if __name__ == "__main__":
     hipified_sources = []
     for source in args.sources:
         s_abs = os.path.abspath(source)
-        if s_abs in hipify_result and hipify_result[s_abs].hipified_path is not None:
-            path = hipify_result[s_abs].hipified_path
-            # PyTorch skips writing when is_pytorch_extension and text unchanged;
-            # hipified_path then stays *.cu. CMake expects *.hip under output_dir.
-            if s_abs.endswith(".cu") and path.endswith(".cu"):
-                dest = _expected_hip_build_path(s_abs, args.output_dir)
-                if os.path.normpath(path) != os.path.normpath(dest):
-                    os.makedirs(os.path.dirname(dest), exist_ok=True)
-                    shutil.copy2(path, dest)
-                hipified_s_abs = dest
-            else:
-                hipified_s_abs = path
-        else:
-            hipified_s_abs = s_abs
+        hipified_s_abs = (
+            hipify_result[s_abs].hipified_path
+            if (
+                s_abs in hipify_result
+                and hipify_result[s_abs].hipified_path is not None
+            )
+            else s_abs
+        )
         hipified_sources.append(hipified_s_abs)
 
     assert len(hipified_sources) == len(args.sources)

cmake/utils.cmake

@@ -81,14 +81,6 @@ function (hipify_sources_target OUT_SRCS NAME ORIG_SRCS)
   set_property(GLOBAL APPEND PROPERTY VLLM_HIPIFY_ALL_SRCS ${SRCS})
   set_property(GLOBAL APPEND PROPERTY VLLM_HIPIFY_ALL_BYPRODUCTS ${HIP_SRCS})
 
-  # Chain hipify targets so they run sequentially. Parallel hipify
-  # invocations race on shutil.copytree, overwriting .hip files
-  # produced by another target back to .cu originals.
-  if (DEFINED _VLLM_LAST_HIPIFY_TARGET)
-    add_dependencies(hipify${NAME} ${_VLLM_LAST_HIPIFY_TARGET})
-  endif()
-  set(_VLLM_LAST_HIPIFY_TARGET "hipify${NAME}" PARENT_SCOPE)
-
   # Swap out original extension sources with hipified sources.
   list(APPEND HIP_SRCS ${CXX_SRCS})
   set(${OUT_SRCS} ${HIP_SRCS} PARENT_SCOPE)

csrc/cpu/cpu_fused_moe.cpp

@@ -30,12 +30,7 @@
   }()
 
 namespace {
-enum class FusedMOEAct {
-  SiluAndMul,
-  SwigluOAIAndMul,
-  GeluAndMul,
-  GeluTanhAndMul,
-};
+enum class FusedMOEAct { SiluAndMul, SwigluOAIAndMul, GeluAndMul };
 
 FusedMOEAct get_act_type(const std::string& act) {
   if (act == "silu") {
@@ -44,8 +39,6 @@ FusedMOEAct get_act_type(const std::string& act) {
     return FusedMOEAct::SwigluOAIAndMul;
   } else if (act == "gelu") {
     return FusedMOEAct::GeluAndMul;
-  } else if (act == "gelu_tanh") {
-    return FusedMOEAct::GeluTanhAndMul;
   } else {
     TORCH_CHECK(false, "Invalid act type: " + act);
   }
@@ -150,44 +143,6 @@ void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
   }
 }
 
-template <typename scalar_t>
-void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
-                       const int32_t m_size, const int32_t n_size,
-                       const int32_t input_stride,
-                       const int32_t output_stride) {
-  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
-  const int32_t dim = n_size / 2;
-  float* __restrict__ gate = input;
-  float* __restrict__ up = input + dim;
-  vec_op::FP32Vec16 one_vec(1.0);
-  vec_op::FP32Vec16 w1_vec(0.7978845608028654);
-  vec_op::FP32Vec16 w2_vec(0.5);
-  vec_op::FP32Vec16 w3_vec(0.044715);
-  alignas(64) float temp[16];
-
-  for (int32_t m = 0; m < m_size; ++m) {
-    for (int32_t n = 0; n < dim; n += 16) {
-      vec_op::FP32Vec16 gate_vec(gate + n);
-      vec_op::FP32Vec16 up_vec(up + n);
-      auto gate_pow3_vec = gate_vec * gate_vec * gate_vec;
-      auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec);
-
-      inner_vec.save(temp);
-      for (int32_t i = 0; i < 16; ++i) {
-        temp[i] = std::tanh(temp[i]);
-      }
-      vec_op::FP32Vec16 tanh_vec(temp);
-      auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec);
-      auto gated_output_fp32 = up_vec * gelu_tanh;
-      scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
-      gated_output.save(output + n);
-    }
-    gate += input_stride;
-    up += input_stride;
-    output += output_stride;
-  }
-}
-
 template <typename scalar_t>
 FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
                                   float* __restrict__ input,
@@ -205,9 +160,6 @@ FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
     case FusedMOEAct::GeluAndMul:
       gelu_and_mul(input, output, m, n, input_stride, output_stride);
       return;
-    case FusedMOEAct::GeluTanhAndMul:
-      gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride);
-      return;
     default:
       TORCH_CHECK(false, "Unsupported act type.");
   }

csrc/cuda_view.cu

@@ -1,4 +1,3 @@
-// TODO: Remove this once ROCm upgrade to torch 2.11.
 #include <torch/all.h>
 #include <torch/cuda.h>
 #include <cuda_runtime.h>

csrc/custom_all_reduce.cu

@@ -1,11 +1,7 @@
-#include "torch_utils.h"
-
-#include <torch/csrc/stable/macros.h>
-#include <torch/csrc/stable/accelerator.h>
-#include <torch/csrc/stable/tensor.h>
-#include <torch/csrc/stable/ops.h>
-#include <torch/headeronly/core/ScalarType.h>
-#include <torch/csrc/stable/device.h>
+#include <ATen/cuda/Exceptions.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#include <torch/all.h>
 
 #include "custom_all_reduce.cuh"
 
@@ -15,7 +11,7 @@ using fptr_t = int64_t;
 static_assert(sizeof(void*) == sizeof(fptr_t));
 
 fptr_t init_custom_ar(const std::vector<fptr_t>& fake_ipc_ptrs,
-                      torch::stable::Tensor& rank_data, int64_t rank,
+                      torch::Tensor& rank_data, int64_t rank,
                       bool fully_connected) {
   int world_size = fake_ipc_ptrs.size();
   if (world_size > 8)
@@ -29,9 +25,9 @@ fptr_t init_custom_ar(const std::vector<fptr_t>& fake_ipc_ptrs,
   for (int i = 0; i < world_size; i++) {
     ipc_ptrs[i] = reinterpret_cast<vllm::Signal*>(fake_ipc_ptrs[i]);
   }
-  return (fptr_t) new vllm::CustomAllreduce(
-      ipc_ptrs, rank_data.mutable_data_ptr(), rank_data.numel(), rank,
-      world_size, fully_connected);
+  return (fptr_t) new vllm::CustomAllreduce(ipc_ptrs, rank_data.data_ptr(),
+                                            rank_data.numel(), rank, world_size,
+                                            fully_connected);
 }
 
 /**
@@ -50,14 +46,10 @@ fptr_t init_custom_ar(const std::vector<fptr_t>& fake_ipc_ptrs,
  * 5. A[None].expand(2, -1, -1, -1): Not OK
  * 6. A[:, 1:, 1:]: Not OK
  */
-bool _is_weak_contiguous(torch::stable::Tensor& t) {
-  if (t.is_contiguous()) {
-    return true;
-  }
-  int64_t storage_nbytes = 0;
-  TORCH_ERROR_CODE_CHECK(aoti_torch_get_storage_size(t.get(), &storage_nbytes));
-  return storage_nbytes - t.storage_offset() * t.element_size() ==
-         static_cast<int64_t>(t.numel() * t.element_size());
+bool _is_weak_contiguous(torch::Tensor& t) {
+  return t.is_contiguous() ||
+         (t.storage().nbytes() - t.storage_offset() * t.element_size() ==
+          t.numel() * t.element_size());
 }
 
 /**
@@ -67,45 +59,42 @@ bool _is_weak_contiguous(torch::stable::Tensor& t) {
  * Otherwise, _reg_buffer is assumed to be IPC-registered and inp is first
  * copied into _reg_buffer.
  */
-void all_reduce(fptr_t _fa, torch::stable::Tensor& inp,
-                torch::stable::Tensor& out, fptr_t _reg_buffer,
-                int64_t reg_buffer_sz_bytes) {
+void all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
+                fptr_t _reg_buffer, int64_t reg_buffer_sz_bytes) {
   auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      inp.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index());
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(inp));
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
 
-  STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type()));
-  STD_TORCH_CHECK((inp.numel()) == (out.numel()));
-  STD_TORCH_CHECK(_is_weak_contiguous(out));
-  STD_TORCH_CHECK(_is_weak_contiguous(inp));
+  TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
+  TORCH_CHECK_EQ(inp.numel(), out.numel());
+  TORCH_CHECK(_is_weak_contiguous(out));
+  TORCH_CHECK(_is_weak_contiguous(inp));
   auto input_size = inp.numel() * inp.element_size();
   auto reg_buffer = reinterpret_cast<void*>(_reg_buffer);
   if (reg_buffer) {
-    STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes));
-    STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size,
-                                   cudaMemcpyDeviceToDevice, stream));
+    TORCH_CHECK_LE(input_size, reg_buffer_sz_bytes);
+    AT_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.data_ptr(), input_size,
+                                  cudaMemcpyDeviceToDevice, stream));
   } else {
-    reg_buffer = inp.mutable_data_ptr();
+    reg_buffer = inp.data_ptr();
   }
   switch (out.scalar_type()) {
-    case torch::headeronly::ScalarType::Float: {
+    case at::ScalarType::Float: {
       fa->allreduce<float>(stream, reinterpret_cast<float*>(reg_buffer),
-                           reinterpret_cast<float*>(out.mutable_data_ptr()),
+                           reinterpret_cast<float*>(out.data_ptr()),
                            out.numel());
       break;
     }
-    case torch::headeronly::ScalarType::Half: {
+    case at::ScalarType::Half: {
       fa->allreduce<half>(stream, reinterpret_cast<half*>(reg_buffer),
-                          reinterpret_cast<half*>(out.mutable_data_ptr()),
-                          out.numel());
+                          reinterpret_cast<half*>(out.data_ptr()), out.numel());
       break;
     }
 #if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
-    case torch::headeronly::ScalarType::BFloat16: {
+    case at::ScalarType::BFloat16: {
       fa->allreduce<nv_bfloat16>(
           stream, reinterpret_cast<nv_bfloat16*>(reg_buffer),
-          reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), out.numel());
+          reinterpret_cast<nv_bfloat16*>(out.data_ptr()), out.numel());
       break;
     }
 #endif
@@ -123,7 +112,7 @@ int64_t meta_size() { return sizeof(vllm::Signal); }
 
 void register_buffer(fptr_t _fa, const std::vector<fptr_t>& fake_ipc_ptrs) {
   auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
-  STD_TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_);
+  TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_);
   void* ipc_ptrs[8];
   for (int i = 0; i < fake_ipc_ptrs.size(); i++) {
     ipc_ptrs[i] = reinterpret_cast<void*>(fake_ipc_ptrs[i]);
@@ -154,49 +143,47 @@ void register_graph_buffers(fptr_t _fa,
   fa->register_graph_buffers(bytes, offsets);
 }
 
-std::tuple<fptr_t, torch::stable::Tensor> allocate_shared_buffer_and_handle(
+std::tuple<fptr_t, torch::Tensor> allocate_shared_buffer_and_handle(
     int64_t size) {
-  int device_index;
-  STD_CUDA_CHECK(cudaGetDevice(&device_index));
-  const torch::stable::accelerator::DeviceGuard device_guard(device_index);
+  auto device_index = c10::cuda::current_device();
+  at::DeviceGuard device_guard(at::Device(at::DeviceType::CUDA, device_index));
   void* buffer;
   cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed;
-  const cudaStream_t stream = get_current_cuda_stream(device_index);
-  STD_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  AT_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));
 
   // Allocate buffer
 #if defined(USE_ROCM)
   // data buffers need to be "uncached" for signal on MI200
-  STD_CUDA_CHECK(
+  AT_CUDA_CHECK(
       hipExtMallocWithFlags((void**)&buffer, size, hipDeviceMallocUncached));
 #else
-  STD_CUDA_CHECK(cudaMalloc((void**)&buffer, size));
+  AT_CUDA_CHECK(cudaMalloc((void**)&buffer, size));
 #endif
-  STD_CUDA_CHECK(cudaMemsetAsync(buffer, 0, size, stream));
-  STD_CUDA_CHECK(cudaStreamSynchronize(stream));
-  STD_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));
+  AT_CUDA_CHECK(cudaMemsetAsync(buffer, 0, size, stream));
+  AT_CUDA_CHECK(cudaStreamSynchronize(stream));
+  AT_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));
 
   // Create IPC memhandle for the allocated buffer.
   // Will use it in open_mem_handle.
-  auto handle = torch::stable::empty(
-      {static_cast<int64_t>(sizeof(cudaIpcMemHandle_t))},
-      torch::headeronly::ScalarType::Byte, std::nullopt,
-      torch::stable::Device(torch::stable::DeviceType::CPU));
-  STD_CUDA_CHECK(cudaIpcGetMemHandle(
-      (cudaIpcMemHandle_t*)handle.mutable_data_ptr(), buffer));
+  auto options =
+      torch::TensorOptions().dtype(torch::kUInt8).device(torch::kCPU);
+  auto handle =
+      torch::empty({static_cast<int64_t>(sizeof(cudaIpcMemHandle_t))}, options);
+  AT_CUDA_CHECK(
+      cudaIpcGetMemHandle((cudaIpcMemHandle_t*)handle.data_ptr(), buffer));
 
   return std::make_tuple(reinterpret_cast<fptr_t>(buffer), handle);
 }
 
-fptr_t open_mem_handle(torch::stable::Tensor& mem_handle) {
+fptr_t open_mem_handle(torch::Tensor& mem_handle) {
   void* ipc_ptr;
-  STD_CUDA_CHECK(cudaIpcOpenMemHandle(
-      (void**)&ipc_ptr,
-      *((const cudaIpcMemHandle_t*)mem_handle.const_data_ptr()),
+  AT_CUDA_CHECK(cudaIpcOpenMemHandle(
+      (void**)&ipc_ptr, *((const cudaIpcMemHandle_t*)mem_handle.data_ptr()),
       cudaIpcMemLazyEnablePeerAccess));
   return reinterpret_cast<fptr_t>(ipc_ptr);
 }
 
 void free_shared_buffer(fptr_t buffer) {
-  STD_CUDA_CHECK(cudaFree(reinterpret_cast<void*>(buffer)));
+  AT_CUDA_CHECK(cudaFree(reinterpret_cast<void*>(buffer)));
 }

csrc/cutlass_extensions/common.cpp

@@ -1,4 +1,4 @@
-#include "common.hpp"
+#include "cutlass_extensions/common.hpp"
 
 int32_t get_sm_version_num() {
   int32_t major_capability, minor_capability;

csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu

@@ -28,20 +28,7 @@
  *   [bs*576,       bs*576 + bs*8):   UE8M0 scales, 7 real + 1 pad per token
  */
 
-#include "torch_utils.h"
-
-#include <torch/csrc/stable/macros.h>
-#include <torch/csrc/stable/accelerator.h>
-#include <torch/csrc/stable/tensor.h>
-#include <torch/csrc/stable/ops.h>
-#include <torch/headeronly/core/ScalarType.h>
-#include <torch/csrc/stable/device.h>
-
 #include <cmath>
-#include "cuda_compat.h"
-#include "dispatch_utils.h"
-#include "type_convert.cuh"
-
 #ifndef USE_ROCM
   #include <cuda_fp8.h>
 #else
@@ -50,6 +37,14 @@
 #include <cuda_runtime.h>
 #include <type_traits>
 
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/cuda.h>
+
+#include "cuda_compat.h"
+#include "dispatch_utils.h"
+#include "type_convert.cuh"
+
 #ifndef FINAL_MASK
   #ifdef USE_ROCM
     #define FINAL_MASK 0xffffffffffffffffULL
@@ -75,7 +70,7 @@ namespace deepseek_v4_fused_ops {
 
 namespace {
 inline int getSMVersion() {
-  auto* props = get_device_prop();
+  auto* props = at::cuda::getCurrentDeviceProperties();
   return props->major * 10 + props->minor;
 }
 }  // namespace
@@ -102,35 +97,6 @@ constexpr float NUM_TOKEN_CUTOFF = 1024;
 constexpr int kNumLanes = 32;
 constexpr int kElemsPerLane = kHeadDim / kNumLanes;  // 16
 
-// Pack this lane's 16 fp32 elements into per-tensor E4M3 FP8 (one uint4 = 16
-// B), scaling by `scale` (a reciprocal scale) and saturating to ±448.  Used by
-// the FlashInfer full-cache path for both the Q and KV stores.
-__device__ __forceinline__ uint4 packFp8E4M3x16(float const* values,
-                                                float const scale) {
-#ifndef USE_ROCM
-  uint4 out;
-  auto* out2 = reinterpret_cast<__nv_fp8x2_storage_t*>(&out);
-  #pragma unroll
-  for (int i = 0; i < kElemsPerLane / 2; i++) {
-    float2 scaled =
-        make_float2(values[2 * i] * scale, values[2 * i + 1] * scale);
-    scaled.x = fminf(fmaxf(scaled.x, -kFp8Max), kFp8Max);
-    scaled.y = fminf(fmaxf(scaled.y, -kFp8Max), kFp8Max);
-    out2[i] = __nv_cvt_float2_to_fp8x2(scaled, __NV_SATFINITE, __NV_E4M3);
-  }
-  return out;
-#else
-  uint8_t out_bytes[kElemsPerLane];
-  #pragma unroll
-  for (int i = 0; i < kElemsPerLane; i++) {
-    float scaled = values[i] * scale;
-    scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
-    out_bytes[i] = rocm_cvt_float_to_fp8_e4m3(scaled);
-  }
-  return *reinterpret_cast<uint4 const*>(out_bytes);
-#endif
-}
-
 // ────────────────────────────────────────────────────────────────────────────
 // Small inline helpers
 // ────────────────────────────────────────────────────────────────────────────
@@ -598,7 +564,7 @@ static void launchFusedDeepseekV4Templated(
   // bf16 on pre-Ampere (sm_70/sm_75) because _typeConvert<BFloat16> is
   // unavailable there.  Refuse the launch loudly instead of silently
   // skipping the work.
-  STD_TORCH_CHECK(
+  TORCH_CHECK(
       sm_version >= 80,
       "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert requires sm_80+ "
       "(Ampere or newer); got sm_",
@@ -669,7 +635,7 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
     DISPATCH(64)
     DISPATCH(128)
     default:
-      STD_TORCH_CHECK(false,
+      TORCH_CHECK(false,
                   "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert: "
                   "unsupported num_heads_q_padded=",
                   num_heads_q_padded,
@@ -678,504 +644,80 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
 #undef DISPATCH
 }
 
-// ────────────────────────────────────────────────────────────────────────────
-// FlashInfer full-cache kernel
-// ────────────────────────────────────────────────────────────────────────────
-//
-// Sibling to the FlashMLA kernel above, used by the FlashInfer V4 sparse-MLA
-// backend.  Differences from the legacy path:
-//   * No Q head padding — output Q layout matches the input num_heads_q.
-//   * KV is written as a *contiguous* 512-wide row per token (token-strided),
-//     not the legacy UE8M0 paged layout with a separate scale tail.
-//   * Q/KV are stored either as bf16 or as per-tensor E4M3 FP8 (one global
-//     scale), selected by the STORE_Q_FP8 / STORE_KV_FP8 template flags.
-//
-// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) / warps).
-// Each warp handles one (token, slot): slot < num_heads_q → Q, slot ==
-// num_heads_q → KV.
-template <typename scalar_t_in, bool STORE_Q_FP8, bool STORE_KV_FP8>
-__global__ void fusedDeepseekV4FullCacheKernel(
-    scalar_t_in* __restrict__ q_inout,          // [N, H, 512], in place (bf16)
-    uint8_t* __restrict__ q_fp8_out,            // [N, H, 512] fp8, optional
-    int64_t const q_fp8_stride0,                // elements (fp8 == bytes)
-    int64_t const q_fp8_stride1,                // elements (fp8 == bytes)
-    scalar_t_in const* __restrict__ kv_in,      // [N, 512] bf16
-    uint8_t* __restrict__ k_cache,              // contiguous bf16 or fp8 cache
-    int64_t const* __restrict__ slot_mapping,   // [num_tokens_insert] i64
-    int64_t const* __restrict__ position_ids,   // [N] i64
-    float const* __restrict__ cos_sin_cache,    // [max_pos, 64] fp32
-    float const* __restrict__ fp8_scale_ptr,    // scalar, KV fp8 only
-    float const* __restrict__ q_fp8_scale_inv,  // scalar, Q fp8 only
-    float const eps,
-    int const num_tokens_full,      // = q.size(0) = kv.size(0)
-    int const num_tokens_insert,    // = slot_mapping.size(0)
-    int const num_heads_q,          // H (no padding)
-    int const cache_block_size,     // tokens per cache block
-    int64_t const kv_block_stride,  // bytes per cache block
-    int64_t const kv_token_stride) {  // bytes per cache token
-#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
-  if constexpr (std::is_same_v<scalar_t_in, c10::BFloat16>) {
-    return;
-  } else {
-#endif
-    using Converter = vllm::_typeConvert<scalar_t_in>;
-    int const warpsPerBlock = blockDim.x / 32;
-    int const warpId = threadIdx.x / 32;
-    int const laneId = threadIdx.x % 32;
-    int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
-
-    int const slotsPerToken = num_heads_q + 1;
-    int const tokenIdx = globalWarpIdx / slotsPerToken;
-    int const slotIdx = globalWarpIdx % slotsPerToken;
-    if (tokenIdx >= num_tokens_full) return;
-    bool const isKV = (slotIdx == num_heads_q);
-    // KV branch: skip DP-padded tokens (no slot reserved for them).
-    if (isKV && tokenIdx >= num_tokens_insert) return;
-
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-    cudaGridDependencySynchronize();
-#endif
-
-    int const dim_base = laneId * kElemsPerLane;  // in [0, 512) step 16
-    scalar_t_in const* src_ptr;
-    if (isKV) {
-      src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
-    } else {
-      src_ptr = q_inout +
-                (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) *
-                    kHeadDim +
-                dim_base;
-    }
-    uint4 const v0 = *reinterpret_cast<uint4 const*>(src_ptr);
-    uint4 const v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
-
-    // ── Decode bf16 → 16 fp32 registers ───────────────────────────────────
-    float elements[kElemsPerLane];
-    {
-      auto const* p0 =
-          reinterpret_cast<typename Converter::packed_hip_type const*>(&v0);
-      auto const* p1 =
-          reinterpret_cast<typename Converter::packed_hip_type const*>(&v1);
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        float2 f2 = Converter::convert(p0[i]);
-        elements[2 * i] = f2.x;
-        elements[2 * i + 1] = f2.y;
-      }
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        float2 f2 = Converter::convert(p1[i]);
-        elements[8 + 2 * i] = f2.x;
-        elements[8 + 2 * i + 1] = f2.y;
-      }
-    }
-
-    // ── Q branch: RMSNorm (no weight) ─────────────────────────────────────
-    if (!isKV) {
-      float sumOfSquares = 0.0f;
-#pragma unroll
-      for (int i = 0; i < kElemsPerLane; i++) {
-        sumOfSquares += elements[i] * elements[i];
-      }
-      sumOfSquares = warpSum<float>(sumOfSquares);
-      float const rms_rcp =
-          rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
-#pragma unroll
-      for (int i = 0; i < kElemsPerLane; i++) {
-        elements[i] = elements[i] * rms_rcp;
-      }
-    }
-
-    // ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ───────────────────────────
-    bool const is_rope_lane = dim_base >= kNopeDim;
-    if (is_rope_lane) {
-      int64_t const pos = position_ids[tokenIdx];
-      constexpr int kHalfRope = kRopeDim / 2;
-      float const* cos_ptr = cos_sin_cache + pos * kRopeDim;
-      float const* sin_ptr = cos_ptr + kHalfRope;
-      int const rope_local_base = dim_base - kNopeDim;
-      int const half_base = rope_local_base >> 1;
-      float4 const c0 = *reinterpret_cast<float4 const*>(cos_ptr + half_base);
-      float4 const c1 = *reinterpret_cast<float4 const*>(cos_ptr + half_base + 4);
-      float4 const s0 = *reinterpret_cast<float4 const*>(sin_ptr + half_base);
-      float4 const s1 = *reinterpret_cast<float4 const*>(sin_ptr + half_base + 4);
-      float const cos_arr[8] = {c0.x, c0.y, c0.z, c0.w, c1.x, c1.y, c1.z, c1.w};
-      float const sin_arr[8] = {s0.x, s0.y, s0.z, s0.w, s1.x, s1.y, s1.z, s1.w};
-#pragma unroll
-      for (int p = 0; p < kElemsPerLane / 2; p++) {
-        float const x_even = elements[2 * p];
-        float const x_odd = elements[2 * p + 1];
-        elements[2 * p] = x_even * cos_arr[p] - x_odd * sin_arr[p];
-        elements[2 * p + 1] = x_even * sin_arr[p] + x_odd * cos_arr[p];
-      }
-    }
-
-    // ── Store ─────────────────────────────────────────────────────────────
-    if (!isKV) {
-      if constexpr (STORE_Q_FP8) {
-        float const scale_inv = VLLM_LDG(q_fp8_scale_inv);
-        uint4 const out = packFp8E4M3x16(elements, scale_inv);
-        uint8_t* dst = q_fp8_out +
-                       static_cast<int64_t>(tokenIdx) * q_fp8_stride0 +
-                       static_cast<int64_t>(slotIdx) * q_fp8_stride1 + dim_base;
-        *reinterpret_cast<uint4*>(dst) = out;
-      } else {
-        uint4 out0, out1;
-        auto* po0 = reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
-        auto* po1 = reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
-#pragma unroll
-        for (int i = 0; i < 4; i++) {
-          po0[i] = Converter::convert(
-              make_float2(elements[2 * i], elements[2 * i + 1]));
-        }
-#pragma unroll
-        for (int i = 0; i < 4; i++) {
-          po1[i] = Converter::convert(
-              make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
-        }
-        scalar_t_in* dst =
-            q_inout +
-            (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
-            dim_base;
-        *reinterpret_cast<uint4*>(dst) = out0;
-        *reinterpret_cast<uint4*>(dst + 8) = out1;
-      }
-    } else {
-      int64_t const slot_id = slot_mapping[tokenIdx];
-      if (slot_id >= 0) {
-        int64_t const block_idx = slot_id / cache_block_size;
-        int64_t const pos_in_block = slot_id % cache_block_size;
-        uint8_t* cache_row =
-            k_cache + block_idx * kv_block_stride + pos_in_block * kv_token_stride;
-        if constexpr (STORE_KV_FP8) {
-          float const inv_scale = 1.0f / VLLM_LDG(fp8_scale_ptr);
-          uint4 const out = packFp8E4M3x16(elements, inv_scale);
-          *reinterpret_cast<uint4*>(cache_row + dim_base) = out;
-        } else {
-          uint4 out0, out1;
-          auto* po0 =
-              reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
-          auto* po1 =
-              reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
-#pragma unroll
-          for (int i = 0; i < 4; i++) {
-            po0[i] = Converter::convert(
-                make_float2(elements[2 * i], elements[2 * i + 1]));
-          }
-#pragma unroll
-          for (int i = 0; i < 4; i++) {
-            po1[i] = Converter::convert(
-                make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
-          }
-          scalar_t_in* dst = reinterpret_cast<scalar_t_in*>(cache_row) + dim_base;
-          *reinterpret_cast<uint4*>(dst) = out0;
-          *reinterpret_cast<uint4*>(dst + 8) = out1;
-        }
-      }
-    }
-
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-    cudaTriggerProgrammaticLaunchCompletion();
-#endif
-#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
-  }
-#endif
-}
-
-// Configure + launch helper shared by the bf16 and fp8 full-cache launchers.
-template <typename scalar_t_in, bool STORE_Q_FP8, bool STORE_KV_FP8>
-static void launchFullCacheKernel(
-    scalar_t_in* q_inout, uint8_t* q_fp8_out, int64_t q_fp8_stride0,
-    int64_t q_fp8_stride1, scalar_t_in const* kv_in, uint8_t* k_cache,
-    int64_t const* slot_mapping, int64_t const* position_ids,
-    float const* cos_sin_cache, float const* fp8_scale,
-    float const* q_fp8_scale_inv, float const eps, int const num_tokens_full,
-    int const num_tokens_insert, int const num_heads_q,
-    int const cache_block_size, int64_t const kv_block_stride,
-    int64_t const kv_token_stride, char const* op_name, cudaStream_t stream) {
-  constexpr int kBlockSize = 256;
-  constexpr int kWarpsPerBlock = kBlockSize / 32;
-  int64_t const total_warps =
-      static_cast<int64_t>(num_tokens_full) * (num_heads_q + 1);
-  int const grid =
-      static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
-  auto* kernel =
-      fusedDeepseekV4FullCacheKernel<scalar_t_in, STORE_Q_FP8, STORE_KV_FP8>;
-#ifndef USE_ROCM
-  static int const sm_version = getSMVersion();
-  STD_TORCH_CHECK(sm_version >= 80, op_name,
-                  " requires sm_80+ (Ampere or newer); got sm_", sm_version);
-  cudaLaunchConfig_t config;
-  config.gridDim = dim3(grid);
-  config.blockDim = dim3(kBlockSize);
-  config.dynamicSmemBytes = 0;
-  config.stream = stream;
-  cudaLaunchAttribute attrs[1];
-  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-  attrs[0].val.programmaticStreamSerializationAllowed = 1;
-  config.attrs = attrs;
-  config.numAttrs = (sm_version >= 90) ? 1 : 0;
-  cudaLaunchKernelEx(&config, kernel, q_inout, q_fp8_out, q_fp8_stride0,
-                     q_fp8_stride1, kv_in, k_cache, slot_mapping, position_ids,
-                     cos_sin_cache, fp8_scale, q_fp8_scale_inv, eps,
-                     num_tokens_full, num_tokens_insert, num_heads_q,
-                     cache_block_size, kv_block_stride, kv_token_stride);
-#else
-  kernel<<<grid, kBlockSize, 0, stream>>>(
-      q_inout, q_fp8_out, q_fp8_stride0, q_fp8_stride1, kv_in, k_cache,
-      slot_mapping, position_ids, cos_sin_cache, fp8_scale, q_fp8_scale_inv,
-      eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
-      kv_block_stride, kv_token_stride);
-#endif
-}
-
 }  // namespace deepseek_v4_fused_ops
 }  // namespace vllm
 
 // ────────────────────────────────────────────────────────────────────────────
 // Torch op wrapper
 // ────────────────────────────────────────────────────────────────────────────
-torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::stable::Tensor const& q_in,           // [N, num_heads_q, 512] bf16
-    torch::stable::Tensor const& kv,             // [N, 512] bf16 (read-only)
-    torch::stable::Tensor& k_cache,              // [num_blocks, block_bytes] uint8
-    torch::stable::Tensor const& slot_mapping,   // [N] int64
-    torch::stable::Tensor const& position_ids,   // [N] int64
-    torch::stable::Tensor const& cos_sin_cache,  // [max_pos, rope_dim] bf16
-    int64_t q_head_padded,                       // padded Q head count for output
+torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+    torch::Tensor const& q_in,           // [N, num_heads_q, 512] bf16
+    torch::Tensor const& kv,             // [N, 512] bf16 (read-only)
+    torch::Tensor& k_cache,              // [num_blocks, block_bytes] uint8
+    torch::Tensor const& slot_mapping,   // [N] int64
+    torch::Tensor const& position_ids,   // [N] int64
+    torch::Tensor const& cos_sin_cache,  // [max_pos, rope_dim] bf16
+    int64_t q_head_padded,               // padded Q head count for output
     double eps, int64_t cache_block_size) {
-  STD_TORCH_CHECK(q_in.device().is_cuda() && q_in.is_contiguous(),
-                  "q_in must be contiguous CUDA");
-  STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(),
-                  "kv must be contiguous CUDA");
-  STD_TORCH_CHECK(k_cache.device().is_cuda(), "k_cache must be CUDA");
-  STD_TORCH_CHECK(slot_mapping.device().is_cuda() &&
-                      slot_mapping.scalar_type() ==
-                          torch::headeronly::ScalarType::Long,
-                  "slot_mapping must be int64 CUDA");
-  STD_TORCH_CHECK(position_ids.device().is_cuda() &&
-                      position_ids.scalar_type() ==
-                          torch::headeronly::ScalarType::Long,
-                  "position_ids must be int64 CUDA");
-  STD_TORCH_CHECK(cos_sin_cache.device().is_cuda(), "cos_sin_cache must be CUDA");
-  STD_TORCH_CHECK(q_in.dim() == 3 && q_in.size(2) == 512,
-                  "q_in shape [N, num_heads_q, 512]");
-  STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
-  STD_TORCH_CHECK(q_in.scalar_type() == kv.scalar_type(),
-                  "q_in and kv dtype must match");
-  STD_TORCH_CHECK(q_head_padded >= q_in.size(1),
-                  "q_head_padded must be >= q_in.size(1) (num_heads_q)");
-  STD_TORCH_CHECK(k_cache.scalar_type() == torch::headeronly::ScalarType::Byte,
-                  "k_cache must be uint8");
-  STD_TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
-                  "cos_sin_cache shape [max_pos, 64]");
-  STD_TORCH_CHECK(cos_sin_cache.scalar_type() ==
-                      torch::headeronly::ScalarType::Float,
-                  "cos_sin_cache must be float32");
+  TORCH_CHECK(q_in.is_cuda() && q_in.is_contiguous(),
+              "q_in must be contiguous CUDA");
+  TORCH_CHECK(kv.is_cuda() && kv.is_contiguous(), "kv must be contiguous CUDA");
+  TORCH_CHECK(k_cache.is_cuda(), "k_cache must be CUDA");
+  TORCH_CHECK(slot_mapping.is_cuda() && slot_mapping.dtype() == torch::kInt64,
+              "slot_mapping must be int64 CUDA");
+  TORCH_CHECK(position_ids.is_cuda() && position_ids.dtype() == torch::kInt64,
+              "position_ids must be int64 CUDA");
+  TORCH_CHECK(cos_sin_cache.is_cuda(), "cos_sin_cache must be CUDA");
+  TORCH_CHECK(q_in.dim() == 3 && q_in.size(2) == 512,
+              "q_in shape [N, num_heads_q, 512]");
+  TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
+  TORCH_CHECK(q_in.dtype() == kv.dtype(), "q_in and kv dtype must match");
+  TORCH_CHECK(q_head_padded >= q_in.size(1),
+              "q_head_padded must be >= q_in.size(1) (num_heads_q)");
+  TORCH_CHECK(k_cache.dtype() == torch::kUInt8, "k_cache must be uint8");
+  TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
+              "cos_sin_cache shape [max_pos, 64]");
+  TORCH_CHECK(cos_sin_cache.dtype() == torch::kFloat32,
+              "cos_sin_cache must be float32");
 
   // With DP padding, slot_mapping can be shorter than q/kv/positions.
   // Q-norm+RoPE runs on all q.size(0) rows (downstream attention uses them);
   // KV quant+insert runs only on the first slot_mapping.size(0) rows.
   int const num_tokens_full = static_cast<int>(q_in.size(0));
   int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
-  STD_TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
-                      static_cast<int>(position_ids.size(0)) == num_tokens_full,
-                  "q/kv/position_ids row counts must match");
-  STD_TORCH_CHECK(num_tokens_insert <= num_tokens_full,
-                  "slot_mapping must not exceed q row count");
+  TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
+                  static_cast<int>(position_ids.size(0)) == num_tokens_full,
+              "q/kv/position_ids row counts must match");
+  TORCH_CHECK(num_tokens_insert <= num_tokens_full,
+              "slot_mapping must not exceed q row count");
   int const num_heads_q = static_cast<int>(q_in.size(1));
   int const num_heads_q_padded = static_cast<int>(q_head_padded);
   int const cache_block_size_i = static_cast<int>(cache_block_size);
   int const kv_block_stride = static_cast<int>(k_cache.stride(0));
 
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      q_in.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream(q_in.get_device_index());
+  at::cuda::OptionalCUDAGuard device_guard(device_of(q_in));
+  auto stream = at::cuda::getCurrentCUDAStream();
 
   // Allocate the padded q output.  The kernel writes every element (live
   // region gets RMSNorm+RoPE; pad region gets zeros), so `empty` is safe.
-  auto q_out = torch::stable::new_empty(
-      q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type());
+  torch::Tensor q_out = torch::empty(
+      {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.options());
 
-  VLLM_STABLE_DISPATCH_HALF_TYPES(
+  VLLM_DISPATCH_HALF_TYPES(
       q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
         using qkv_scalar_t = scalar_t;
         vllm::deepseek_v4_fused_ops::
             launchFusedDeepseekV4QNormRopeKVRopeQuantInsert<qkv_scalar_t>(
-                reinterpret_cast<qkv_scalar_t const*>(q_in.const_data_ptr()),
-                reinterpret_cast<qkv_scalar_t*>(q_out.mutable_data_ptr()),
-                reinterpret_cast<qkv_scalar_t const*>(kv.const_data_ptr()),
-                reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr()),
-                slot_mapping.const_data_ptr<int64_t>(),
-                position_ids.const_data_ptr<int64_t>(),
-                cos_sin_cache.const_data_ptr<float>(), static_cast<float>(eps),
+                reinterpret_cast<qkv_scalar_t const*>(q_in.data_ptr()),
+                reinterpret_cast<qkv_scalar_t*>(q_out.data_ptr()),
+                reinterpret_cast<qkv_scalar_t const*>(kv.data_ptr()),
+                reinterpret_cast<uint8_t*>(k_cache.data_ptr()),
+                reinterpret_cast<int64_t const*>(slot_mapping.data_ptr()),
+                reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
+                cos_sin_cache.data_ptr<float>(), static_cast<float>(eps),
                 num_tokens_full, num_tokens_insert, num_heads_q,
                 num_heads_q_padded, cache_block_size_i, kv_block_stride,
                 stream);
       });
   return q_out;
 }
-
-// ────────────────────────────────────────────────────────────────────────────
-// FlashInfer full-cache torch ops
-// ────────────────────────────────────────────────────────────────────────────
-void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
-    torch::stable::Tensor& q,                    // [N, H, 512] bf16, in place
-    torch::stable::Tensor const& kv,             // [N, 512] bf16, read-only
-    torch::stable::Tensor& k_cache,              // [num_blocks, bs, 512] bf16
-    torch::stable::Tensor const& slot_mapping,   // [num_tokens_insert] int64
-    torch::stable::Tensor const& position_ids,   // [N] int64
-    torch::stable::Tensor const& cos_sin_cache,  // [max_pos, 64] float32
-    double eps, int64_t cache_block_size) {
-  using torch::headeronly::ScalarType;
-  STD_TORCH_CHECK(q.device().is_cuda() && q.is_contiguous(),
-                  "q must be contiguous CUDA");
-  STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(),
-                  "kv must be contiguous CUDA");
-  STD_TORCH_CHECK(k_cache.device().is_cuda(), "k_cache must be CUDA");
-  STD_TORCH_CHECK(slot_mapping.device().is_cuda() &&
-                      slot_mapping.scalar_type() == ScalarType::Long,
-                  "slot_mapping must be int64 CUDA");
-  STD_TORCH_CHECK(position_ids.device().is_cuda() &&
-                      position_ids.scalar_type() == ScalarType::Long,
-                  "position_ids must be int64 CUDA");
-  STD_TORCH_CHECK(cos_sin_cache.device().is_cuda() &&
-                      cos_sin_cache.scalar_type() == ScalarType::Float &&
-                      cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
-                  "cos_sin_cache shape [max_pos, 64] float32");
-  STD_TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
-  STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
-  STD_TORCH_CHECK(q.scalar_type() == ScalarType::BFloat16 &&
-                      kv.scalar_type() == ScalarType::BFloat16,
-                  "q and kv must be bfloat16");
-  STD_TORCH_CHECK(k_cache.dim() == 3 && k_cache.size(1) == cache_block_size &&
-                      k_cache.size(2) == 512 && k_cache.stride(2) == 1,
-                  "k_cache shape [num_blocks, cache_block_size, 512] contiguous");
-  STD_TORCH_CHECK(k_cache.scalar_type() == ScalarType::BFloat16,
-                  "k_cache must be bfloat16");
-
-  int const num_tokens_full = static_cast<int>(q.size(0));
-  int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
-  STD_TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
-                      static_cast<int>(position_ids.size(0)) == num_tokens_full,
-                  "q/kv/position_ids row counts must match");
-  STD_TORCH_CHECK(num_tokens_insert <= num_tokens_full,
-                  "slot_mapping must not exceed q row count");
-  int const num_heads_q = static_cast<int>(q.size(1));
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      q.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream(q.get_device_index());
-
-  // bf16 cache: 2 bytes/element -> byte strides for the uint8-addressed kernel.
-  int64_t const kv_block_stride = k_cache.stride(0) * 2;
-  int64_t const kv_token_stride = k_cache.stride(1) * 2;
-
-  VLLM_STABLE_DISPATCH_HALF_TYPES(
-      q.scalar_type(),
-      "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert", [&] {
-        vllm::deepseek_v4_fused_ops::launchFullCacheKernel<scalar_t, false,
-                                                           false>(
-            reinterpret_cast<scalar_t*>(q.mutable_data_ptr()), nullptr, 0, 0,
-            reinterpret_cast<scalar_t const*>(kv.const_data_ptr()),
-            reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr()),
-            slot_mapping.const_data_ptr<int64_t>(),
-            position_ids.const_data_ptr<int64_t>(),
-            cos_sin_cache.const_data_ptr<float>(), nullptr, nullptr,
-            static_cast<float>(eps), num_tokens_full, num_tokens_insert,
-            num_heads_q, static_cast<int>(cache_block_size), kv_block_stride,
-            kv_token_stride,
-            "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert",
-            stream);
-      });
-}
-
-void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
-    torch::stable::Tensor const& q,                // [N, H, 512] bf16, read-only
-    torch::stable::Tensor const& kv,               // [N, 512] bf16, read-only
-    torch::stable::Tensor& q_fp8,                  // [N, H, 512] fp8 e4m3
-    torch::stable::Tensor& k_cache,                // [num_blocks, bs, 512] fp8
-    torch::stable::Tensor const& slot_mapping,     // [num_tokens_insert] int64
-    torch::stable::Tensor const& position_ids,     // [N] int64
-    torch::stable::Tensor const& cos_sin_cache,    // [max_pos, 64] float32
-    torch::stable::Tensor const& fp8_scale,        // scalar float32 (KV scale)
-    torch::stable::Tensor const& q_fp8_scale_inv,  // scalar float32 (1 / Q scale)
-    double eps, int64_t cache_block_size) {
-  using torch::headeronly::ScalarType;
-  STD_TORCH_CHECK(q.device().is_cuda() && q.is_contiguous(),
-                  "q must be contiguous CUDA");
-  STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(),
-                  "kv must be contiguous CUDA");
-  STD_TORCH_CHECK(q_fp8.device().is_cuda() && q_fp8.is_contiguous() &&
-                      q_fp8.scalar_type() == ScalarType::Float8_e4m3fn &&
-                      q_fp8.dim() == 3 && q_fp8.size(0) == q.size(0) &&
-                      q_fp8.size(1) == q.size(1) && q_fp8.size(2) == q.size(2),
-                  "q_fp8 must be a contiguous float8_e4m3fn tensor matching q");
-  STD_TORCH_CHECK(k_cache.device().is_cuda(), "k_cache must be CUDA");
-  STD_TORCH_CHECK(slot_mapping.device().is_cuda() &&
-                      slot_mapping.scalar_type() == ScalarType::Long,
-                  "slot_mapping must be int64 CUDA");
-  STD_TORCH_CHECK(position_ids.device().is_cuda() &&
-                      position_ids.scalar_type() == ScalarType::Long,
-                  "position_ids must be int64 CUDA");
-  STD_TORCH_CHECK(cos_sin_cache.device().is_cuda() &&
-                      cos_sin_cache.scalar_type() == ScalarType::Float &&
-                      cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
-                  "cos_sin_cache shape [max_pos, 64] float32");
-  STD_TORCH_CHECK(fp8_scale.device().is_cuda() &&
-                      fp8_scale.scalar_type() == ScalarType::Float &&
-                      fp8_scale.size(0) == 1,
-                  "fp8_scale must be a scalar float32 CUDA tensor");
-  STD_TORCH_CHECK(q_fp8_scale_inv.device().is_cuda() &&
-                      q_fp8_scale_inv.scalar_type() == ScalarType::Float &&
-                      q_fp8_scale_inv.size(0) == 1,
-                  "q_fp8_scale_inv must be a scalar float32 CUDA tensor");
-  STD_TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
-  STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
-  STD_TORCH_CHECK(q.scalar_type() == kv.scalar_type(),
-                  "q and kv dtype must match");
-  STD_TORCH_CHECK(k_cache.dim() == 3 && k_cache.size(1) == cache_block_size &&
-                      k_cache.size(2) == 512 && k_cache.stride(2) == 1,
-                  "k_cache shape [num_blocks, cache_block_size, 512] contiguous");
-  STD_TORCH_CHECK(k_cache.scalar_type() == ScalarType::Float8_e4m3fn,
-                  "k_cache must be float8_e4m3fn");
-
-  int const num_tokens_full = static_cast<int>(q.size(0));
-  int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
-  STD_TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
-                      static_cast<int>(position_ids.size(0)) == num_tokens_full,
-                  "q/kv/position_ids row counts must match");
-  STD_TORCH_CHECK(num_tokens_insert <= num_tokens_full,
-                  "slot_mapping must not exceed q row count");
-  int const num_heads_q = static_cast<int>(q.size(1));
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      q.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream(q.get_device_index());
-
-  VLLM_STABLE_DISPATCH_HALF_TYPES(
-      q.scalar_type(),
-      "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert", [&] {
-        vllm::deepseek_v4_fused_ops::launchFullCacheKernel<scalar_t, true,
-                                                           true>(
-            // q is read-only in the fp8 path (the kernel writes q_fp8); the
-            // launcher signature is non-const, so cast away const on the ptr.
-            reinterpret_cast<scalar_t*>(
-                const_cast<void*>(q.const_data_ptr())),
-            reinterpret_cast<uint8_t*>(q_fp8.mutable_data_ptr()),
-            q_fp8.stride(0), q_fp8.stride(1),
-            reinterpret_cast<scalar_t const*>(kv.const_data_ptr()),
-            reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr()),
-            slot_mapping.const_data_ptr<int64_t>(),
-            position_ids.const_data_ptr<int64_t>(),
-            cos_sin_cache.const_data_ptr<float>(),
-            fp8_scale.const_data_ptr<float>(),
-            q_fp8_scale_inv.const_data_ptr<float>(), static_cast<float>(eps),
-            num_tokens_full, num_tokens_insert, num_heads_q,
-            static_cast<int>(cache_block_size),
-            // fp8 cache: 1 byte/element -> stride already in bytes.
-            k_cache.stride(0), k_cache.stride(1),
-            "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert",
-            stream);
-      });
-}

csrc/libtorch_stable/concat_mla_q.cuh

@@ -1,6 +1,9 @@
 #ifndef CONCAT_MLA_Q_CUH_
 #define CONCAT_MLA_Q_CUH_
 
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
 #include "cuda_vec_utils.cuh"
 
 namespace vllm {

csrc/libtorch_stable/cuda_view.cu

@@ -1,76 +0,0 @@
-#include <torch/csrc/stable/tensor.h>
-#include <torch/csrc/stable/ops.h>
-#include <torch/csrc/stable/accelerator.h>
-#include <torch/headeronly/core/ScalarType.h>
-#include <torch/csrc/stable/device.h>
-#include <torch/csrc/stable/c/shim.h>
-#include <torch/headeronly/version.h>
-#include <cuda_runtime.h>
-
-#include <array>
-#include <optional>
-
-// This function assumes that `cpu_tensor` is a CPU tensor,
-// and that UVA (Unified Virtual Addressing) is enabled.
-torch::stable::Tensor get_cuda_view_from_cpu_tensor(
-    torch::stable::Tensor& cpu_tensor) {
-  STD_TORCH_CHECK(cpu_tensor.device().is_cpu(), "Input tensor must be on CPU");
-
-  const auto dtype = cpu_tensor.scalar_type();
-  const auto layout = cpu_tensor.layout();
-  const torch::stable::Device cuda_dev(torch::headeronly::DeviceType::CUDA);
-
-  // handle empty tensor
-  if (cpu_tensor.numel() == 0) {
-    return torch::stable::empty(cpu_tensor.sizes(), dtype, layout, cuda_dev);
-  }
-
-  std::array<StableIValue, 2> is_pinned_stack{
-      torch::stable::detail::from(cpu_tensor),
-      torch::stable::detail::from(std::nullopt)};
-  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
-      "aten::is_pinned", "", is_pinned_stack.data(), TORCH_ABI_VERSION));
-  if (torch::stable::detail::to<bool>(is_pinned_stack[0])) {
-    // If CPU tensor is pinned, directly get the device pointer.
-    void* host_ptr = const_cast<void*>(cpu_tensor.mutable_data_ptr());
-    void* device_ptr = nullptr;
-    cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
-    STD_TORCH_CHECK(err == cudaSuccess, "cudaHostGetDevicePointer failed: ",
-                    cudaGetErrorString(err));
-
-    return torch::stable::from_blob(
-        device_ptr, cpu_tensor.sizes(), cpu_tensor.strides(), cuda_dev, dtype,
-        [base = cpu_tensor](void*) {});  // keep cpu tensor alive
-  }
-
-  // If CPU tensor is not pinned, allocate a new pinned memory buffer.
-  torch::stable::Tensor contiguous_cpu = torch::stable::contiguous(cpu_tensor);
-  size_t nbytes = contiguous_cpu.numel() * contiguous_cpu.element_size();
-
-  void* host_ptr = nullptr;
-  cudaError_t err = cudaHostAlloc(&host_ptr, nbytes, cudaHostAllocMapped);
-  if (err != cudaSuccess) {
-    STD_TORCH_CHECK(false, "cudaHostAlloc failed: ", cudaGetErrorString(err));
-  }
-
-  err = cudaMemcpy(host_ptr, contiguous_cpu.const_data_ptr(), nbytes,
-                   cudaMemcpyDefault);
-  if (err != cudaSuccess) {
-    cudaFreeHost(host_ptr);
-    STD_TORCH_CHECK(false, "cudaMemcpy failed: ", cudaGetErrorString(err));
-  }
-
-  void* device_ptr = nullptr;
-  err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
-  if (err != cudaSuccess) {
-    cudaFreeHost(host_ptr);
-    STD_TORCH_CHECK(
-        false, "cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
-  }
-
-  auto deleter = [host_ptr](void*) { cudaFreeHost(host_ptr); };
-
-  return torch::stable::from_blob(device_ptr, contiguous_cpu.sizes(),
-                                  contiguous_cpu.strides(), cuda_dev,
-                                  contiguous_cpu.scalar_type(), deleter);
-}

csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu

@@ -1,69 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-// Adapted from SGLang:
-// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled.cu
-
-#include <torch/csrc/stable/library.h>
-#include <torch/csrc/stable/tensor.h>
-#include "libtorch_stable/torch_utils.h"
-
-#include "cutlass_mxfp8_grouped_mm_launcher.cuh"
-
-void cutlass_mxfp8_grouped_mm(const torch::stable::Tensor& a,
-                              const torch::stable::Tensor& b,
-                              const torch::stable::Tensor& sfa,
-                              const torch::stable::Tensor& sfb,
-                              torch::stable::Tensor& d,
-                              const torch::stable::Tensor& problem_sizes,
-                              const torch::stable::Tensor& expert_offsets,
-                              const torch::stable::Tensor& blockscale_offsets) {
-#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
-  STD_TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
-  STD_TORCH_CHECK(problem_sizes.size(1) == 3,
-                  "problem_sizes must have shape (num_experts, 3)");
-  STD_TORCH_CHECK(
-      problem_sizes.size(0) == expert_offsets.size(0),
-      "Number of experts in problem_sizes must match expert_offsets");
-  STD_TORCH_CHECK(
-      problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
-      "problem_sizes must be int32");
-  STD_TORCH_CHECK(
-      expert_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
-      "expert_offsets must be int32");
-  STD_TORCH_CHECK(
-      blockscale_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
-      "blockscale_offsets must be int32");
-  STD_TORCH_CHECK(a.dim() == 2,
-                  "a must be a 2D tensor of shape (num_tokens, k)");
-  STD_TORCH_CHECK(b.dim() == 3,
-                  "b must be a 3D tensor of shape (num_experts, k, n)");
-  STD_TORCH_CHECK(a.size(1) == b.size(1) && a.size(1) % 128 == 0,
-                  "k should align 128");
-  STD_TORCH_CHECK(b.size(2) % 128 == 0, "n should align 128");
-  STD_TORCH_CHECK(a.stride(1) == 1, "a must be row major");
-  STD_TORCH_CHECK(b.stride(1) == 1, "b must be column major");
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      a.get_device_index());
-  auto stream = get_current_cuda_stream(a.get_device_index());
-  if (d.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
-    expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
-        cutlass::bfloat16_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
-                             blockscale_offsets, stream);
-  } else if (d.scalar_type() == torch::headeronly::ScalarType::Half) {
-    expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
-        cutlass::half_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
-                         blockscale_offsets, stream);
-  } else {
-    STD_TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
-  }
-#else
-  STD_TORCH_CHECK(false,
-                  "No implemented cutlass_mxfp8_grouped_mm for "
-                  "current device");
-#endif
-}
-
-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
-  m.impl("cutlass_mxfp8_grouped_mm", TORCH_BOX(&cutlass_mxfp8_grouped_mm));
-}

csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu

@@ -1,66 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-// Adapted from SGLang:
-// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cu
-
-#include <torch/csrc/stable/library.h>
-#include <torch/csrc/stable/tensor.h>
-#include "libtorch_stable/torch_utils.h"
-
-#include "mxfp8_experts_quant.cuh"
-
-void mxfp8_experts_quant(const torch::stable::Tensor& input,
-                         const torch::stable::Tensor& problem_sizes,
-                         const torch::stable::Tensor& expert_offsets,
-                         const torch::stable::Tensor& blockscale_offsets,
-                         torch::stable::Tensor& quant_output,
-                         torch::stable::Tensor& scale_factor) {
-#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
-  STD_TORCH_CHECK(input.dim() == 2, "input must be 2D tensor");
-  STD_TORCH_CHECK(input.size(1) % 128 == 0, "k must align to 128");
-  STD_TORCH_CHECK(input.stride(1) == 1, "input must be row major");
-  STD_TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
-  STD_TORCH_CHECK(
-      problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
-      "problem_sizes must be int32");
-  STD_TORCH_CHECK(
-      expert_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
-      "expert_offsets must be int32");
-  STD_TORCH_CHECK(
-      blockscale_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
-      "blockscale_offsets must be int32");
-
-  auto groups = problem_sizes.size(0);
-  STD_TORCH_CHECK(
-      expert_offsets.dim() == 1 && expert_offsets.size(0) == groups,
-      "expert_offsets must be 1D and have size equal to the number of groups");
-  STD_TORCH_CHECK(
-      blockscale_offsets.dim() == 1 && blockscale_offsets.size(0) == groups,
-      "blockscale_offsets must be 1D and have size equal to the number of "
-      "groups");
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      input.get_device_index());
-  if (input.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
-    expert_specialization::launch_mxfp8_experts_quant<__nv_bfloat16>(
-        input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
-        scale_factor);
-  } else if (input.scalar_type() == torch::headeronly::ScalarType::Half) {
-    expert_specialization::launch_mxfp8_experts_quant<__half>(
-        input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
-        scale_factor);
-  } else {
-    STD_TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
-  }
-#else
-  STD_TORCH_CHECK(false,
-                  "No implemented mxfp8_experts_quant for "
-                  "current device");
-#endif
-}
-
-// Registered here (not torch_bindings.cpp) because ENABLE_ES_MXFP8_GROUPED_MM
-// is applied only under COMPILE_LANGUAGE:CUDA.
-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
-  m.impl("mxfp8_experts_quant", TORCH_BOX(&mxfp8_experts_quant));
-}

csrc/libtorch_stable/ops.h

@@ -3,6 +3,10 @@
 #include <torch/csrc/stable/library.h>
 #include <torch/csrc/stable/tensor.h>
 
+#ifndef USE_ROCM
+torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
+                                   torch::stable::Tensor const& perm);
+
 void per_token_group_quant_fp8(const torch::stable::Tensor& input,
                                torch::stable::Tensor& output_q,
                                torch::stable::Tensor& output_s,
@@ -24,10 +28,6 @@ void per_token_group_quant_int8(const torch::stable::Tensor& input,
                                 int64_t group_size, double eps, double int8_min,
                                 double int8_max);
 
-torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
-                                   torch::stable::Tensor const& perm);
-
-#ifndef USE_ROCM
 bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
 bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
 bool cutlass_group_gemm_supported(int64_t cuda_device_capability);
@@ -162,11 +162,6 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
 // AllSpark ops: declarations are in the source files
 // (allspark_repack.cu and allspark_qgemm_w8a16.cu)
 
-// TODO: Move this out once ROCm upgrade their torch to 2.11.
-// CPU tensor -> CUDA UVA view (shared CUDA)
-torch::stable::Tensor get_cuda_view_from_cpu_tensor(
-    torch::stable::Tensor& cpu_tensor);
-
 #endif
 
 // Attention kernels (shared CUDA/ROCm)
@@ -220,13 +215,6 @@ void rms_norm_per_block_quant(torch::stable::Tensor& out,
                               std::optional<torch::stable::Tensor> residual,
                               int64_t group_size, bool is_scale_transposed);
 
-void silu_and_mul_per_block_quant(torch::stable::Tensor& out,
-                                  torch::stable::Tensor const& input,
-                                  torch::stable::Tensor& scales,
-                                  int64_t group_size,
-                                  std::optional<torch::stable::Tensor> scale_ub,
-                                  bool is_scale_transposed);
-
 // Positional encoding kernels (shared CUDA/ROCm)
 void rotary_embedding(torch::stable::Tensor& positions,
                       torch::stable::Tensor& query,
@@ -243,44 +231,6 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q,
                         torch::stable::Tensor& position_ids,
                         int64_t forced_token_heads_per_warp);
 
-torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv,
-    torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
-    torch::stable::Tensor const& position_ids,
-    torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
-    double eps, int64_t cache_block_size);
-
-void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
-    torch::stable::Tensor& q, torch::stable::Tensor const& kv,
-    torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
-    torch::stable::Tensor const& position_ids,
-    torch::stable::Tensor const& cos_sin_cache, double eps,
-    int64_t cache_block_size);
-
-void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
-    torch::stable::Tensor const& q, torch::stable::Tensor const& kv,
-    torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache,
-    torch::stable::Tensor const& slot_mapping,
-    torch::stable::Tensor const& position_ids,
-    torch::stable::Tensor const& cos_sin_cache,
-    torch::stable::Tensor const& fp8_scale,
-    torch::stable::Tensor const& q_fp8_scale_inv, double eps,
-    int64_t cache_block_size);
-
-#ifndef USE_ROCM
-torch::stable::Tensor minimax_allreduce_rms(
-    torch::stable::Tensor const& input,
-    torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace,
-    int64_t const rank, int64_t const nranks, double const eps);
-std::tuple<torch::stable::Tensor, torch::stable::Tensor>
-minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
-                         torch::stable::Tensor const& norm_weight_q,
-                         torch::stable::Tensor const& norm_weight_k,
-                         torch::stable::Tensor workspace, int64_t const q_size,
-                         int64_t const kv_size, int64_t const rank,
-                         int64_t const nranks, double const eps);
-#endif
-
 // Sampler kernels (shared CUDA/ROCm)
 void apply_repetition_penalties_(
     torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
@@ -323,26 +273,6 @@ void selective_scan_fwd(
     const std::optional<torch::stable::Tensor>& cu_chunk_seqlen,
     const std::optional<torch::stable::Tensor>& last_chunk_indices);
 
-using fptr_t = int64_t;
-fptr_t init_custom_ar(const std::vector<int64_t>& fake_ipc_ptrs,
-                      torch::stable::Tensor& rank_data, int64_t rank,
-                      bool fully_connected);
-void all_reduce(fptr_t _fa, torch::stable::Tensor& inp,
-                torch::stable::Tensor& out, fptr_t reg_buffer,
-                int64_t reg_buffer_sz_bytes);
-void dispose(fptr_t _fa);
-int64_t meta_size();
-void register_buffer(fptr_t _fa, const std::vector<int64_t>& fake_ipc_ptrs);
-std::tuple<std::vector<int64_t>, std::vector<int64_t>>
-get_graph_buffer_ipc_meta(fptr_t _fa);
-void register_graph_buffers(fptr_t _fa,
-                            const std::vector<std::vector<int64_t>>& handles,
-                            const std::vector<std::vector<int64_t>>& offsets);
-std::tuple<int64_t, torch::stable::Tensor> allocate_shared_buffer_and_handle(
-    int64_t size);
-int64_t open_mem_handle(torch::stable::Tensor& mem_handle);
-void free_shared_buffer(int64_t buffer);
-
 // Activation kernels (shared CUDA/ROCm)
 void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void silu_and_mul_clamp(torch::stable::Tensor& out,

csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu

@@ -18,7 +18,7 @@
 #include <torch/csrc/stable/tensor.h>
 #include "libtorch_stable/torch_utils.h"
 #include "cutlass_extensions/torch_utils.hpp"
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 
 #include "get_group_starts.cuh"
 #include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu

@@ -21,7 +21,7 @@
 #include "cutlass/util/packed_stride.hpp"
 #include "cutlass/util/mixed_dtype_utils.hpp"
 
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 #include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
 
 #include <cuda_runtime.h>

csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu

@@ -12,7 +12,7 @@
 
 #include <cutlass/arch/arch.h>
 
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 
 #include "cute/tensor.hpp"
 #include "cutlass/tensor_ref.h"

csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu

@@ -20,7 +20,7 @@
 
 #include <cutlass/arch/arch.h>
 
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 
 #include "cute/tensor.hpp"
 #include "cutlass/tensor_ref.h"

csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu

@@ -18,7 +18,7 @@
 
 #include "libtorch_stable/torch_utils.h"
 
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 #include "nvfp4_utils.cuh"
 
 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \

csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu

@@ -18,7 +18,7 @@
 
 #include "libtorch_stable/torch_utils.h"
 
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 
 #if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
 void cutlass_scaled_fp4_mm_sm100a(torch::stable::Tensor& D,

csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu

@@ -18,7 +18,7 @@
 
 #include "libtorch_stable/torch_utils.h"
 
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 
 #include "cutlass/cutlass.h"
 

csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu

@@ -18,7 +18,7 @@
 
 #include "libtorch_stable/torch_utils.h"
 
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 
 #include "cutlass/cutlass.h"
 

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh

@@ -20,7 +20,7 @@
 #include "cutlass/util/packed_stride.hpp"
 
 #include "core/math.hpp"
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 // clang-format on
 
 namespace vllm::c3x {

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh

@@ -15,7 +15,7 @@
 #include "cutlass/gemm/collective/collective_builder.hpp"
 
 #include "core/math.hpp"
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 // clang-format on
 
 /*

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp

@@ -1,7 +1,7 @@
 #include <torch/csrc/stable/tensor.h>
 #include <torch/headeronly/core/ScalarType.h>
 #include "cuda_utils.h"
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 
 template <typename Fp8Func, typename Int8Func, typename BlockwiseFunc>
 void dispatch_scaled_mm(torch::stable::Tensor& c,

csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh

@@ -8,7 +8,7 @@
 
 #include <torch/csrc/stable/ops.h>
 #include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 #include "get_group_starts.cuh"
 
 using namespace cute;

csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh

@@ -23,7 +23,7 @@
 #include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h"
 
 #include "core/math.hpp"
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 // clang-format on
 
 using namespace cute;

csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu

@@ -4,7 +4,7 @@
 
 #include "libtorch_stable/torch_utils.h"
 
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 
 void cutlass_scaled_mm_sm75(torch::stable::Tensor& c,
                             torch::stable::Tensor const& a,

csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu

@@ -7,11 +7,7 @@
 
 #include <cmath>
 
-#ifdef USE_ROCM
-  #include <hip/hip_fp8.h>
-#else
-  #include <cuda_fp8.h>
-#endif
+#include <cuda_fp8.h>
 
 #include "libtorch_stable/quantization/vectorization.cuh"
 #include "libtorch_stable/quantization/vectorization_utils.cuh"
@@ -19,23 +15,12 @@
 #include "libtorch_stable/torch_utils.h"
 
 __device__ __forceinline__ float GroupReduceMax(float val) {
-#ifdef USE_ROCM
-  // 16-thread logical groups may pack up to four per 64-lane wavefront; use a
-  // 64-bit mask and explicit width so shuffles stay within each group.
-  const int lane_in_wave = threadIdx.x % warpSize;
-  const unsigned long long mask = 0xFFFFull << ((lane_in_wave / 16) * 16);
-  val = fmaxf(val, __shfl_xor_sync(mask, val, 8, 16));
-  val = fmaxf(val, __shfl_xor_sync(mask, val, 4, 16));
-  val = fmaxf(val, __shfl_xor_sync(mask, val, 2, 16));
-  val = fmaxf(val, __shfl_xor_sync(mask, val, 1, 16));
-#else
   unsigned mask = threadIdx.x % 32 >= 16 ? 0xffff0000 : 0x0000ffff;
 
   val = fmaxf(val, __shfl_xor_sync(mask, val, 8));
   val = fmaxf(val, __shfl_xor_sync(mask, val, 4));
   val = fmaxf(val, __shfl_xor_sync(mask, val, 2));
   val = fmaxf(val, __shfl_xor_sync(mask, val, 1));
-#endif
   return val;
 }
 
@@ -252,12 +237,10 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
 
   VLLM_STABLE_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "per_token_group_quant_8bit", ([&] {
-        if (dst_type == torch::headeronly::ScalarType::Char) {
+        if (dst_type == torch::headeronly::ScalarType::Float8_e4m3fn) {
+          LAUNCH_KERNEL(scalar_t, __nv_fp8_e4m3);
+        } else if (dst_type == torch::headeronly::ScalarType::Char) {
           LAUNCH_KERNEL(scalar_t, int8_t);
-        } else {
-          VLLM_STABLE_DISPATCH_FP8_TYPES(
-              dst_type, "per_token_group_quant_8bit_fp8",
-              ([&] { LAUNCH_KERNEL(scalar_t, fp8_t); }));
         }
       }));
 
@@ -334,18 +317,10 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
 
   // 8-lane subgroup shuffle reduce (octet of the warp). The mask selects the
   // 8 lanes within the warp that share a group.
-#ifdef USE_ROCM
-  const int lane_in_wave = threadIdx.x % warpSize;
-  const unsigned long long mask = 0xFFull << (lane_in_wave & ~7);
-  local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4, 8));
-  local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2, 8));
-  local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1, 8));
-#else
   unsigned mask = 0xffu << (threadIdx.x & 24u);
   local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4));
   local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2));
   local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1));
-#endif
 
   float y_s = local_absmax / max_8bit;
   y_s = fmaxf(y_s, 1e-10f);
@@ -528,12 +503,15 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
 
   VLLM_STABLE_DISPATCH_HALF_TYPES(
       input.scalar_type(), "per_token_group_quant_8bit_packed_register", ([&] {
-        if (dst_type == torch::headeronly::ScalarType::Char) {
+        if (dst_type == torch::headeronly::ScalarType::Float8_e4m3fn) {
+          LAUNCH_REG_KERNEL(scalar_t, __nv_fp8_e4m3);
+        } else if (dst_type == torch::headeronly::ScalarType::Char) {
           LAUNCH_REG_KERNEL(scalar_t, int8_t);
         } else {
-          VLLM_STABLE_DISPATCH_FP8_TYPES(
-              dst_type, "per_token_group_quant_8bit_packed_fp8",
-              ([&] { LAUNCH_REG_KERNEL(scalar_t, fp8_t); }));
+          STD_TORCH_CHECK(
+              false,
+              "per_token_group_quant_8bit_packed only supports FP8/INT8 "
+              "outputs.");
         }
       }));
 

csrc/libtorch_stable/torch_bindings.cpp

@@ -1,5 +1,4 @@
 #include "ops.h"
-#include "cuda_utils.h"
 #include "core/registration.h"
 
 #include <torch/csrc/stable/library.h>
@@ -8,6 +7,11 @@
 // Note: We register under namespace "_C" so ops are accessible as
 // torch.ops._C.<op_name> for compatibility with existing code.
 STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
+#ifndef USE_ROCM
+  ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
+#endif
+
+#ifndef USE_ROCM
   // Compute per-token-group FP8 quantized tensor and scaling factor.
   // The dummy arguments are here so we can correctly fuse with RMSNorm.
   ops.def(
@@ -27,15 +31,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! "
       "output_s, int group_size, float eps, float int8_min, float int8_max) -> "
       "()");
-  ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
 
-#ifndef USE_ROCM
-
-  // TODO: Remove this once ROCm upgrade to torch 2.11.
-  ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
-#endif
-
-#ifndef USE_ROCM
   // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column
   // quantization, as well as bias
   ops.def(
@@ -325,16 +321,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "Tensor? scale_ub, Tensor!? residual, int group_size, "
       "bool is_scale_transposed) -> ()");
 
-  // Fused SiLU+Mul + per-block quantization
-  ops.def(
-      "silu_and_mul_per_block_quant("
-      "Tensor! out, "
-      "Tensor input, "
-      "Tensor! scales, "
-      "int group_size, "
-      "Tensor? scale_ub=None, "
-      "bool is_scale_transposed=False) -> ()");
-
   // Rotary embedding
   // Apply GPT-NeoX or GPT-J style rotary embedding to query and key.
   ops.def(
@@ -351,38 +337,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "bool is_neox, Tensor position_ids, "
       "int forced_token_heads_per_warp=-1) -> ()");
 
-  ops.def(
-      "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
-      "Tensor q_in, Tensor kv, Tensor! k_cache, "
-      "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
-      "int q_head_padded, float eps, int cache_block_size) -> Tensor");
-
-  // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate
-  // FP8 tensor, and KV into a contiguous 512-wide token-strided cache.
-  ops.def(
-      "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert("
-      "Tensor! q, Tensor kv, Tensor! k_cache, Tensor slot_mapping, "
-      "Tensor position_ids, Tensor cos_sin_cache, float eps, "
-      "int cache_block_size) -> ()");
-  ops.def(
-      "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert("
-      "Tensor q, Tensor kv, Tensor! q_fp8, Tensor! k_cache, "
-      "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
-      "Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, "
-      "int cache_block_size) -> ()");
-
-#ifndef USE_ROCM
-  ops.def(
-      "minimax_allreduce_rms("
-      "Tensor input, Tensor norm_weight, Tensor workspace, "
-      "int rank, int nranks, float eps) -> Tensor");
-  ops.def(
-      "minimax_allreduce_rms_qk("
-      "Tensor qkv, Tensor norm_weight_q, Tensor norm_weight_k, "
-      "Tensor workspace, int q_size, int kv_size, int rank, int nranks, "
-      "float eps) -> (Tensor, Tensor)");
-#endif
-
   // Apply repetition penalties to logits in-place.
   ops.def(
       "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -554,6 +508,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
 }
 
 STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
+#ifndef USE_ROCM
+  ops.impl("permute_cols", TORCH_BOX(&permute_cols));
+#endif
+
+#ifndef USE_ROCM
   // Per-token group quantization
   ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8));
   ops.impl("per_token_group_fp8_quant_packed",
@@ -561,9 +520,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("per_token_group_quant_int8",
            TORCH_BOX(&per_token_group_quant_int8));
 
-  ops.impl("permute_cols", TORCH_BOX(&permute_cols));
-
-#ifndef USE_ROCM
   // CUTLASS scaled_mm ops
   ops.impl("cutlass_scaled_mm", TORCH_BOX(&cutlass_scaled_mm));
   ops.impl("cutlass_scaled_mm_azp", TORCH_BOX(&cutlass_scaled_mm_azp));
@@ -611,24 +567,10 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("rms_norm_dynamic_per_token_quant",
            TORCH_BOX(&rms_norm_dynamic_per_token_quant));
   ops.impl("rms_norm_per_block_quant", TORCH_BOX(&rms_norm_per_block_quant));
-  ops.impl("silu_and_mul_per_block_quant",
-           TORCH_BOX(&silu_and_mul_per_block_quant));
 
   // Positional encoding kernels (shared CUDA/ROCm)
   ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding));
   ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));
-  ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert",
-           TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert));
-  ops.impl(
-      "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert",
-      TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert));
-  ops.impl(
-      "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert",
-      TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert));
-#ifndef USE_ROCM
-  ops.impl("minimax_allreduce_rms", TORCH_BOX(&minimax_allreduce_rms));
-  ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk));
-#endif
 
   // Sampler kernels (shared CUDA/ROCm)
   ops.impl("apply_repetition_penalties_",
@@ -675,28 +617,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2));
 }
 
-// TODO: Remove this once ROCm upgrade to torch 2.11.
-#ifndef USE_ROCM
-STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) {
-  ops.impl("get_cuda_view_from_cpu_tensor",
-           TORCH_BOX(&get_cuda_view_from_cpu_tensor));
-}
-
-STABLE_TORCH_LIBRARY_FRAGMENT(_C_cuda_utils, cuda_utils) {
-  cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int");
-  cuda_utils.def(
-      "get_max_shared_memory_per_block_device_attribute(int device_id) -> int");
-}
-
-STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd,
-                          cuda_utils) {
-  cuda_utils.impl("get_device_attribute", TORCH_BOX(&get_device_attribute));
-  cuda_utils.impl("get_max_shared_memory_per_block_device_attribute",
-                  TORCH_BOX(&get_max_shared_memory_per_block_device_attribute));
-}
-
-#endif
-
 // These capability-check functions take only primitive args (no tensors), so
 // there is no device to dispatch on. CompositeExplicitAutograd makes them
 // available for all backends. This is the stable ABI equivalent of calling
@@ -805,45 +725,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
       "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()");
 }
 
-STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ar) {
-  custom_ar.def(
-      "init_custom_ar(int[] ipc_tensors, Tensor rank_data, "
-      "int rank, bool fully_connected) -> int");
-  custom_ar.def(
-      "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, "
-      "int reg_buffer_sz_bytes) -> ()");
-  custom_ar.def("dispose(int fa) -> ()");
-  custom_ar.def("meta_size() -> int");
-  custom_ar.def("register_buffer(int fa, int[] ipc_tensors) -> ()");
-  custom_ar.def("get_graph_buffer_ipc_meta(int fa) -> (int[], int[])");
-  custom_ar.def(
-      "register_graph_buffers(int fa, int[][] handles, int[][] offsets) -> ()");
-  custom_ar.def("allocate_shared_buffer_and_handle(int size) -> (int, Tensor)");
-  custom_ar.def("open_mem_handle(Tensor mem_handle) -> int");
-  custom_ar.def("free_shared_buffer(int ptr) -> ()");
-}
-
-STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ar) {
-  custom_ar.impl("init_custom_ar", TORCH_BOX(&init_custom_ar));
-  custom_ar.impl("all_reduce", TORCH_BOX(&all_reduce));
-}
-
-STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CPU, custom_ar) {
-  custom_ar.impl("open_mem_handle", TORCH_BOX(&open_mem_handle));
-}
-
-STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CompositeExplicitAutograd, custom_ar) {
-  custom_ar.impl("dispose", TORCH_BOX(&dispose));
-  custom_ar.impl("meta_size", TORCH_BOX(&meta_size));
-  custom_ar.impl("register_buffer", TORCH_BOX(&register_buffer));
-  custom_ar.impl("get_graph_buffer_ipc_meta",
-                 TORCH_BOX(&get_graph_buffer_ipc_meta));
-  custom_ar.impl("register_graph_buffers", TORCH_BOX(&register_graph_buffers));
-  custom_ar.impl("allocate_shared_buffer_and_handle",
-                 TORCH_BOX(&allocate_shared_buffer_and_handle));
-  custom_ar.impl("free_shared_buffer", TORCH_BOX(&free_shared_buffer));
-}
-
 STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CPU, ops) {
   ops.impl("swap_blocks_batch", TORCH_BOX(&swap_blocks_batch));
 }

csrc/libtorch_stable/torch_utils.h

@@ -6,7 +6,11 @@
 #include <torch/csrc/stable/tensor.h>
 #include <torch/headeronly/util/shim_utils.h>
 
-#include <cuda_runtime.h>
+#ifndef USE_ROCM
+  #include <cuda_runtime.h>
+#else
+  #include <hip/hip_runtime.h>
+#endif
 #include <cublas_v2.h>
 
 #include <deque>

csrc/minimax_reduce_rms_kernel.cu

@@ -15,19 +15,16 @@
  * limitations under the License.
  */
 
-#include "torch_utils.h"
-
-#include <torch/csrc/stable/macros.h>
-#include <torch/csrc/stable/accelerator.h>
-#include <torch/csrc/stable/tensor.h>
-#include <torch/csrc/stable/ops.h>
-#include <torch/headeronly/core/ScalarType.h>
-#include <torch/csrc/stable/device.h>
-
 #include <cooperative_groups.h>
 #include <cuda_runtime.h>
 
+#include <torch/cuda.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+
 #include "cuda_compat.h"
+#include "cuda_utils.h"
+#include "core/registration.h"
 #include "minimax_reduce_rms_kernel.h"
 
 #include <algorithm>
@@ -614,7 +611,7 @@ int get_sm_count() {
   static int sm_count = 0;
   if (sm_count == 0) {
     int device_id;
-    STD_CUDA_CHECK(cudaGetDevice(&device_id));
+    CUDA_CHECK(cudaGetDevice(&device_id));
     cudaDeviceProp device_prop;
     cudaGetDeviceProperties(&device_prop, device_id);
     sm_count = device_prop.multiProcessorCount;
@@ -624,13 +621,13 @@ int get_sm_count() {
 
 inline int getSMVersion(bool queryRealSmArch = false) {
   int device{-1};
-  STD_CUDA_CHECK(cudaGetDevice(&device));
+  CUDA_CHECK(cudaGetDevice(&device));
   int sm_major = 0;
   int sm_minor = 0;
-  STD_CUDA_CHECK(cudaDeviceGetAttribute(
-      &sm_major, cudaDevAttrComputeCapabilityMajor, device));
-  STD_CUDA_CHECK(cudaDeviceGetAttribute(
-      &sm_minor, cudaDevAttrComputeCapabilityMinor, device));
+  CUDA_CHECK(cudaDeviceGetAttribute(&sm_major,
+                                    cudaDevAttrComputeCapabilityMajor, device));
+  CUDA_CHECK(cudaDeviceGetAttribute(&sm_minor,
+                                    cudaDevAttrComputeCapabilityMinor, device));
   int sm = sm_major * 10 + sm_minor;
   if (sm == 121 && !queryRealSmArch) {
     return 120;
@@ -642,7 +639,7 @@ template <typename KernelFunc>
 int get_max_active_blocks(KernelFunc kernel, int block_size,
                           int dynamic_smem = 0) {
   int max_active = 0;
-  STD_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+  CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
       &max_active, kernel, block_size, dynamic_smem));
   return std::max(max_active, 1);
 }
@@ -681,27 +678,27 @@ void minimax_reduce_rms_kernel_launcher(MiniMaxReduceRMSParams const& params) {
   cfg.attrs = attribute;
   cfg.numAttrs = SM >= 90 ? 2 : 0;
 
-  STD_CUDA_CHECK(cudaLaunchKernelEx(
+  CUDA_CHECK(cudaLaunchKernelEx(
       &cfg, minimax_reduce_rms_kernel_lamport<DType, NRanks>, params));
 }
 
 template <typename DType, int NRanks, int OriginQDim, int OriginKDim>
 void minimax_reduce_rms_kernel_launcher_float4(
     MiniMaxReduceRMSParams const& params) {
-  STD_TORCH_CHECK(params.size_q % params.hidden_dim == 0);
-  STD_TORCH_CHECK(params.hidden_dim % kElemsPerAccess<DType> == 0);
+  TORCH_CHECK(params.size_q % params.hidden_dim == 0);
+  TORCH_CHECK(params.hidden_dim % kElemsPerAccess<DType> == 0);
   if (params.stride_q > 0) {
-    STD_TORCH_CHECK(params.stride_q % kElemsPerAccess<DType> == 0);
+    TORCH_CHECK(params.stride_q % kElemsPerAccess<DType> == 0);
   }
-  STD_TORCH_CHECK(params.allreduce_in_k != nullptr,
-                  "float4 QK kernel requires K input");
-  STD_TORCH_CHECK(params.hidden_dim >= params.hidden_dim_k);
-  STD_TORCH_CHECK(params.size_k % params.hidden_dim_k == 0);
-  STD_TORCH_CHECK(params.hidden_dim_k % kElemsPerAccess<DType> == 0);
-  STD_TORCH_CHECK(params.size_q / params.hidden_dim ==
-                  params.size_k / params.hidden_dim_k);
+  TORCH_CHECK(params.allreduce_in_k != nullptr,
+              "float4 QK kernel requires K input");
+  TORCH_CHECK(params.hidden_dim >= params.hidden_dim_k);
+  TORCH_CHECK(params.size_k % params.hidden_dim_k == 0);
+  TORCH_CHECK(params.hidden_dim_k % kElemsPerAccess<DType> == 0);
+  TORCH_CHECK(params.size_q / params.hidden_dim ==
+              params.size_k / params.hidden_dim_k);
   if (params.stride_k > 0) {
-    STD_TORCH_CHECK(params.stride_k % kElemsPerAccess<DType> == 0);
+    TORCH_CHECK(params.stride_k % kElemsPerAccess<DType> == 0);
   }
 
   int token_num = params.size_q / params.hidden_dim;
@@ -749,7 +746,7 @@ void minimax_reduce_rms_kernel_launcher_float4(
   cfg.attrs = attribute;
   cfg.numAttrs = SM >= 90 ? 2 : 0;
 
-  STD_CUDA_CHECK(cudaLaunchKernelEx(&cfg, kfn, params));
+  CUDA_CHECK(cudaLaunchKernelEx(&cfg, kfn, params));
 }
 
 template <int NRanks>
@@ -762,21 +759,21 @@ void dispatch_dtype(MiniMaxReduceRMSParams const& params) {
                     (params.hidden_dim * params.nranks == 6144) &&
                     (params.hidden_dim_k * params.nranks == 1024);
 
-  if (params.dtype == torch::headeronly::ScalarType::Half) {
+  if (params.dtype == at::ScalarType::Half) {
     if (use_float4) {
       minimax_reduce_rms_kernel_launcher_float4<half, NRanks, 6144, 1024>(
           params);
     } else {
       minimax_reduce_rms_kernel_launcher<half, NRanks>(params);
     }
-  } else if (params.dtype == torch::headeronly::ScalarType::BFloat16) {
+  } else if (params.dtype == at::ScalarType::BFloat16) {
     if (use_float4) {
       minimax_reduce_rms_kernel_launcher_float4<__nv_bfloat16, NRanks, 6144,
                                                 1024>(params);
     } else {
       minimax_reduce_rms_kernel_launcher<__nv_bfloat16, NRanks>(params);
     }
-  } else if (params.dtype == torch::headeronly::ScalarType::Float) {
+  } else if (params.dtype == at::ScalarType::Float) {
     if (use_float4) {
       minimax_reduce_rms_kernel_launcher_float4<float, NRanks, 6144, 1024>(
           params);
@@ -784,7 +781,7 @@ void dispatch_dtype(MiniMaxReduceRMSParams const& params) {
       minimax_reduce_rms_kernel_launcher<float, NRanks>(params);
     }
   } else {
-    STD_TORCH_CHECK(false, "Unsupported data type for minimax_reduce_rms_op");
+    TORCH_CHECK(false, "Unsupported data type for minimax_reduce_rms_op");
   }
 }
 
@@ -798,18 +795,16 @@ void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params) {
   } else if (params.nranks == 16) {
     dispatch_dtype<16>(params);
   } else {
-    STD_TORCH_CHECK(false, "minimax_reduce_rms_op: unsupported ranks number!");
+    TORCH_CHECK(false, "minimax_reduce_rms_op: unsupported ranks number!");
   }
 }
 }  // namespace tensorrt_llm
 }  // namespace vllm
 
-torch::stable::Tensor minimax_allreduce_rms(
-    torch::stable::Tensor const& input,
-    torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace,
-    int64_t const rank, int64_t const nranks, double const eps) {
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      input.get_device_index());
+torch::Tensor minimax_allreduce_rms(torch::Tensor const& input,
+                                    torch::Tensor const& norm_weight,
+                                    torch::Tensor workspace, int64_t const rank,
+                                    int64_t const nranks, double const eps) {
   auto allreduce_params = vllm::tensorrt_llm::MiniMaxReduceRMSParams();
 
   allreduce_params.nranks = static_cast<int>(nranks);
@@ -820,12 +815,12 @@ torch::stable::Tensor minimax_allreduce_rms(
   allreduce_params.stride_q = allreduce_params.hidden_dim;
   allreduce_params.workspace =
       reinterpret_cast<void**>(workspace.mutable_data_ptr());
-  allreduce_params.allreduce_in = const_cast<void*>(input.const_data_ptr());
-  allreduce_params.rms_gamma = const_cast<void*>(norm_weight.const_data_ptr());
+  allreduce_params.allreduce_in = input.data_ptr();
+  allreduce_params.rms_gamma = norm_weight.data_ptr();
   allreduce_params.rms_eps = static_cast<float>(eps);
-  allreduce_params.stream = get_current_cuda_stream(input.get_device_index());
+  allreduce_params.stream = at::cuda::getCurrentCUDAStream(input.get_device());
 
-  torch::stable::Tensor rms_norm_out = torch::stable::empty_like(input);
+  torch::Tensor rms_norm_out = torch::empty_like(input);
   allreduce_params.rms_norm_out = rms_norm_out.mutable_data_ptr();
 
   vllm::tensorrt_llm::minimax_reduce_rms_op(allreduce_params);
@@ -833,33 +828,26 @@ torch::stable::Tensor minimax_allreduce_rms(
   return rms_norm_out;
 }
 
-std::tuple<torch::stable::Tensor, torch::stable::Tensor>
-minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
-                         torch::stable::Tensor const& norm_weight_q,
-                         torch::stable::Tensor const& norm_weight_k,
-                         torch::stable::Tensor workspace, int64_t const q_size,
-                         int64_t const kv_size, int64_t const rank,
-                         int64_t const nranks, double const eps) {
-  STD_TORCH_CHECK(qkv.dim() == 2, "minimax_allreduce_rms_qk: qkv must be 2D");
-  STD_TORCH_CHECK(qkv.is_contiguous(),
-                  "minimax_allreduce_rms_qk: qkv must be contiguous");
+std::tuple<torch::Tensor, torch::Tensor> minimax_allreduce_rms_qk(
+    torch::Tensor qkv, torch::Tensor const& norm_weight_q,
+    torch::Tensor const& norm_weight_k, torch::Tensor workspace,
+    int64_t const q_size, int64_t const kv_size, int64_t const rank,
+    int64_t const nranks, double const eps) {
+  TORCH_CHECK(qkv.dim() == 2, "minimax_allreduce_rms_qk: qkv must be 2D");
+  TORCH_CHECK(qkv.is_contiguous(),
+              "minimax_allreduce_rms_qk: qkv must be contiguous");
   int64_t qkv_dim = qkv.size(-1);
-  STD_TORCH_CHECK(qkv_dim == q_size + 2 * kv_size,
-                  "minimax_allreduce_rms_qk: qkv last dim must equal "
-                  "q_size + 2 * kv_size");
-  STD_TORCH_CHECK(rank < nranks,
-                  "minimax_allreduce_rms_qk: rank must be less than nranks");
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      qkv.get_device_index());
+  TORCH_CHECK(qkv_dim == q_size + 2 * kv_size,
+              "minimax_allreduce_rms_qk: qkv last dim must equal "
+              "q_size + 2 * kv_size");
+  TORCH_CHECK(rank < nranks,
+              "minimax_allreduce_rms_qk: rank must be less than nranks");
 
   int64_t num_tokens = qkv.size(0);
   int elem_bytes = qkv.element_size();
 
-  torch::stable::Tensor q_out =
-      torch::stable::new_empty(qkv, {num_tokens, q_size}, qkv.scalar_type());
-  torch::stable::Tensor k_out =
-      torch::stable::new_empty(qkv, {num_tokens, kv_size}, qkv.scalar_type());
+  torch::Tensor q_out = torch::empty({num_tokens, q_size}, qkv.options());
+  torch::Tensor k_out = torch::empty({num_tokens, kv_size}, qkv.options());
 
   auto params = vllm::tensorrt_llm::MiniMaxReduceRMSParams();
   params.nranks = static_cast<int>(nranks);
@@ -875,14 +863,13 @@ minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
   params.stride_k_out = 0;  // k_out is contiguous; kernel uses hidden_dim_k
   params.workspace = reinterpret_cast<void**>(workspace.mutable_data_ptr());
 
-  uint8_t* base =
-      const_cast<uint8_t*>(static_cast<const uint8_t*>(qkv.const_data_ptr()));
+  uint8_t* base = static_cast<uint8_t*>(qkv.data_ptr());
   params.allreduce_in = base;
   params.allreduce_in_k = base + q_size * elem_bytes;
-  params.rms_gamma = const_cast<void*>(norm_weight_q.const_data_ptr());
-  params.rms_gamma_k = const_cast<void*>(norm_weight_k.const_data_ptr());
+  params.rms_gamma = norm_weight_q.data_ptr();
+  params.rms_gamma_k = norm_weight_k.data_ptr();
   params.rms_eps = static_cast<float>(eps);
-  params.stream = get_current_cuda_stream(qkv.get_device_index());
+  params.stream = at::cuda::getCurrentCUDAStream(qkv.get_device());
 
   params.rms_norm_out = q_out.mutable_data_ptr();
   params.rms_norm_out_k = k_out.mutable_data_ptr();

csrc/minimax_reduce_rms_kernel.h

@@ -19,7 +19,7 @@
 #include <cuda_bf16.h>
 #include <cuda_fp16.h>
 
-#include <torch/headeronly/core/ScalarType.h>
+#include <torch/types.h>
 
 namespace vllm {
 namespace tensorrt_llm {
@@ -51,7 +51,7 @@ static constexpr int kElemsPerAccess = ElemsPerAccess<DType>::value;
 struct MiniMaxReduceRMSParams {
   int nranks{};
   int rank{};
-  torch::headeronly::ScalarType dtype{torch::headeronly::ScalarType::Undefined};
+  at::ScalarType dtype{at::ScalarType::Undefined};
   int size_q{};
   int hidden_dim{};
   int size_k{};

csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu

@@ -0,0 +1,60 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+// Adapted from SGLang:
+// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled.cu
+
+#include <torch/all.h>
+
+#include "cutlass_mxfp8_grouped_mm_launcher.cuh"
+
+void cutlass_mxfp8_grouped_mm(const torch::Tensor& a, const torch::Tensor& b,
+                              const torch::Tensor& sfa,
+                              const torch::Tensor& sfb, torch::Tensor& d,
+                              const torch::Tensor& problem_sizes,
+                              const torch::Tensor& expert_offsets,
+                              const torch::Tensor& blockscale_offsets) {
+#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
+  TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
+  TORCH_CHECK(problem_sizes.size(1) == 3,
+              "problem_sizes must have shape (num_experts, 3)");
+  TORCH_CHECK(problem_sizes.size(0) == expert_offsets.size(0),
+              "Number of experts in problem_sizes must match expert_offsets");
+  TORCH_CHECK(problem_sizes.dtype() == torch::kInt32,
+              "problem_sizes must be int32");
+  TORCH_CHECK(expert_offsets.dtype() == torch::kInt32,
+              "expert_offsets must be int32");
+  TORCH_CHECK(blockscale_offsets.dtype() == torch::kInt32,
+              "blockscale_offsets must be int32");
+  TORCH_CHECK(a.dim() == 2, "a must be a 2D tensor of shape (num_tokens, k)");
+  TORCH_CHECK(b.dim() == 3,
+              "b must be a 3D tensor of shape (num_experts, k, n)");
+  TORCH_CHECK(a.size(1) == b.size(1) && a.size(1) % 128 == 0,
+              "k should align 128");
+  TORCH_CHECK(b.size(2) % 128 == 0, "n should align 128");
+  TORCH_CHECK(a.strides()[1] == 1, "a must be row major");
+  TORCH_CHECK(b.strides()[1] == 1, "b must be column major");
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+  if (d.dtype() == torch::kBFloat16) {
+    expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
+        cutlass::bfloat16_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
+                             blockscale_offsets, stream);
+  } else if (d.dtype() == torch::kFloat16) {
+    expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
+        cutlass::half_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
+                         blockscale_offsets, stream);
+  } else {
+    TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
+  }
+#else
+  TORCH_CHECK(false,
+              "No implemented cutlass_mxfp8_grouped_mm for "
+              "current device");
+#endif
+}
+
+#include "core/registration.h"
+
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
+  m.impl("cutlass_mxfp8_grouped_mm", cutlass_mxfp8_grouped_mm);
+}

csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_launcher.cuh

@@ -4,9 +4,9 @@
 // https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_launcher.cuh
 
 #pragma once
-
-#include <torch/csrc/stable/tensor.h>
-#include <torch/headeronly/util/Exception.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/all.h>
 
 #include <cassert>
 #include <iostream>
@@ -15,22 +15,18 @@
 #include "cute/tensor.hpp"
 #include "cutlass_mxfp8_grouped_mm_functor.cuh"
 #include "cutlass_mxfp8_grouped_mm_traits.cuh"
-#include "libtorch_stable/torch_utils.h"
 
 namespace expert_specialization {
 
 template <typename GemmTraits>
 void cutlass_mxfp8_grouped_mm_pre_compute(
-    torch::stable::Tensor& a_ptrs, torch::stable::Tensor& b_ptrs,
-    torch::stable::Tensor& sfa_ptrs, torch::stable::Tensor& sfb_ptrs,
-    torch::stable::Tensor& d_ptrs, torch::stable::Tensor& stride_a,
-    torch::stable::Tensor& stride_b, torch::stable::Tensor& stride_d,
-    torch::stable::Tensor& layout_sfa, torch::stable::Tensor& layout_sfb,
-    const torch::stable::Tensor& a, const torch::stable::Tensor& b,
-    const torch::stable::Tensor& sfa, const torch::stable::Tensor& sfb,
-    const torch::stable::Tensor& d, const torch::stable::Tensor& problem_sizes,
-    const torch::stable::Tensor& expert_offsets,
-    const torch::stable::Tensor& blockscale_offsets, cudaStream_t stream) {
+    torch::Tensor& a_ptrs, torch::Tensor& b_ptrs, torch::Tensor& sfa_ptrs,
+    torch::Tensor& sfb_ptrs, torch::Tensor& d_ptrs, torch::Tensor& stride_a,
+    torch::Tensor& stride_b, torch::Tensor& stride_d, torch::Tensor& layout_sfa,
+    torch::Tensor& layout_sfb, const torch::Tensor& a, const torch::Tensor& b,
+    const torch::Tensor& sfa, const torch::Tensor& sfb, const torch::Tensor& d,
+    const torch::Tensor& problem_sizes, const torch::Tensor& expert_offsets,
+    const torch::Tensor& blockscale_offsets, cudaStream_t stream) {
   using OffsetFunctor = CutlassMxfp8GroupedMmOffsetFunctor<GemmTraits>;
   using ElementA = typename OffsetFunctor::ElementA;
   using ElementB = typename OffsetFunctor::ElementB;
@@ -46,10 +42,10 @@ void cutlass_mxfp8_grouped_mm_pre_compute(
   using StrideB = typename StrideFunctor::StrideB;
   using StrideD = typename StrideFunctor::StrideD;
 
-  int num_experts = static_cast<int>(expert_offsets.size(0));
-  STD_TORCH_CHECK(num_experts <= 1024,
-                  "Number of experts cannot exceed 1024, the maximum number of "
-                  "threads per block.");
+  int num_experts = (int)expert_offsets.size(0);
+  TORCH_CHECK(num_experts <= 1024,
+              "Number of experts cannot exceed 1024, the maximum number of "
+              "threads per block.");
 
   OffsetFunctor offset_functor(
       reinterpret_cast<int*>(expert_offsets.data_ptr()),
@@ -76,18 +72,13 @@ void cutlass_mxfp8_grouped_mm_pre_compute(
 }
 
 template <typename GemmTraits>
-void cutlass_mxfp8_grouped_mm(const torch::stable::Tensor& a_ptrs,
-                              const torch::stable::Tensor& b_ptrs,
-                              const torch::stable::Tensor& sfa_ptrs,
-                              const torch::stable::Tensor& sfb_ptrs,
-                              const torch::stable::Tensor& d_ptrs,
-                              const torch::stable::Tensor& stride_a,
-                              const torch::stable::Tensor& stride_b,
-                              const torch::stable::Tensor& stride_d,
-                              const torch::stable::Tensor& layout_sfa,
-                              const torch::stable::Tensor& layout_sfb,
-                              const torch::stable::Tensor& problem_sizes,
-                              cudaStream_t stream) {
+void cutlass_mxfp8_grouped_mm(
+    const torch::Tensor& a_ptrs, const torch::Tensor& b_ptrs,
+    const torch::Tensor& sfa_ptrs, const torch::Tensor& sfb_ptrs,
+    const torch::Tensor& d_ptrs, const torch::Tensor& stride_a,
+    const torch::Tensor& stride_b, const torch::Tensor& stride_d,
+    const torch::Tensor& layout_sfa, const torch::Tensor& layout_sfb,
+    const torch::Tensor& problem_sizes, cudaStream_t stream) {
   using Gemm = typename GemmTraits::Gemm;
   using ElementA = typename Gemm::ElementA;
   using ElementB = typename Gemm::ElementB;
@@ -102,12 +93,13 @@ void cutlass_mxfp8_grouped_mm(const torch::stable::Tensor& a_ptrs,
       typename GemmTraits::ProblemShape::UnderlyingProblemShape;
 
   cutlass::KernelHardwareInfo hw_info;
-  hw_info.device_id = d_ptrs.get_device_index();
-  hw_info.sm_count = get_device_prop()->multiProcessorCount;
+  hw_info.device_id = c10::cuda::current_device();
+  hw_info.sm_count =
+      at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
   hw_info.cluster_shape = GemmTraits::MMAConfig::preferred_cluster;
   hw_info.cluster_shape_fallback = GemmTraits::MMAConfig::fallback_cluster;
 
-  int num_experts = static_cast<int>(problem_sizes.size(0));
+  int num_experts = (int)problem_sizes.size(0);
 
   UnderlyingProblemShape* underlying_problem_shape =
       reinterpret_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
@@ -135,55 +127,44 @@ void cutlass_mxfp8_grouped_mm(const torch::stable::Tensor& a_ptrs,
   Gemm gemm;
 
   auto can_implement_status = gemm.can_implement(arguments);
-  STD_TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
-                  "Failed to implement GEMM");
+  TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
+              "Failed to implement GEMM");
 
+  torch::TensorOptions options_uint8 =
+      torch::TensorOptions().dtype(torch::kUInt8).device(d_ptrs.device());
   size_t workspace_size = gemm.get_workspace_size(arguments);
-  torch::stable::Tensor workspace = torch::stable::empty(
-      {static_cast<int64_t>(workspace_size)},
-      torch::headeronly::ScalarType::Byte, std::nullopt, d_ptrs.device());
+  torch::Tensor workspace = torch::empty(workspace_size, options_uint8);
 
   auto status = gemm.initialize(arguments, workspace.data_ptr(), stream);
-  STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
-                  "Failed to initialize GEMM");
+  TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to initialize GEMM");
 
   status = gemm.run(stream, nullptr, true);  // Enable PDL
-  STD_TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
+  TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
 }
 
 template <typename OutType>
 void cutlass_mxfp8_grouped_mm_dispatch_out_dtype(
-    const torch::stable::Tensor& a, const torch::stable::Tensor& b,
-    const torch::stable::Tensor& sfa, const torch::stable::Tensor& sfb,
-    torch::stable::Tensor& d, const torch::stable::Tensor& problem_sizes,
-    const torch::stable::Tensor& expert_offsets,
-    const torch::stable::Tensor& blockscale_offsets, cudaStream_t stream) {
-  int num_experts = static_cast<int>(problem_sizes.size(0));
-  auto device = a.device();
+    const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& sfa,
+    const torch::Tensor& sfb, torch::Tensor& d,
+    const torch::Tensor& problem_sizes, const torch::Tensor& expert_offsets,
+    const torch::Tensor& blockscale_offsets, cudaStream_t stream) {
+  int num_experts = (int)problem_sizes.size(0);
+  torch::TensorOptions options_int64 =
+      torch::TensorOptions().dtype(torch::kInt64).device(a.device());
+  torch::TensorOptions options_int32 =
+      torch::TensorOptions().dtype(torch::kInt32).device(a.device());
 
-  torch::stable::Tensor a_ptrs = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor b_ptrs = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor sfa_ptrs = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor sfb_ptrs = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor d_ptrs = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
+  torch::Tensor a_ptrs = torch::empty(num_experts, options_int64);
+  torch::Tensor b_ptrs = torch::empty(num_experts, options_int64);
+  torch::Tensor sfa_ptrs = torch::empty(num_experts, options_int64);
+  torch::Tensor sfb_ptrs = torch::empty(num_experts, options_int64);
+  torch::Tensor d_ptrs = torch::empty(num_experts, options_int64);
 
-  torch::stable::Tensor stride_a = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor stride_b = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor stride_d = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor layout_sfa =
-      torch::stable::empty({num_experts, 5}, torch::headeronly::ScalarType::Int,
-                           std::nullopt, device);
-  torch::stable::Tensor layout_sfb =
-      torch::stable::empty({num_experts, 5}, torch::headeronly::ScalarType::Int,
-                           std::nullopt, device);
+  torch::Tensor stride_a = torch::empty(num_experts, options_int64);
+  torch::Tensor stride_b = torch::empty(num_experts, options_int64);
+  torch::Tensor stride_d = torch::empty(num_experts, options_int64);
+  torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int32);
+  torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int32);
 
   using GemmTraits = CutlassMxfp8GroupedMmGemmTraits<MMA1SMConfig, OutType>;
   cutlass_mxfp8_grouped_mm_pre_compute<GemmTraits>(
@@ -195,4 +176,4 @@ void cutlass_mxfp8_grouped_mm_dispatch_out_dtype(
       layout_sfa, layout_sfb, problem_sizes, stream);
 }
 
-}  // namespace expert_specialization
+}  // namespace expert_specialization

csrc/moe/mxfp8_moe/mxfp8_experts_quant.cu

@@ -0,0 +1,60 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+// Adapted from SGLang:
+// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cu
+
+#include <torch/all.h>
+
+#include "mxfp8_experts_quant.cuh"
+
+void mxfp8_experts_quant(const torch::Tensor& input,
+                         const torch::Tensor& problem_sizes,
+                         const torch::Tensor& expert_offsets,
+                         const torch::Tensor& blockscale_offsets,
+                         torch::Tensor& quant_output,
+                         torch::Tensor& scale_factor) {
+#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
+  TORCH_CHECK(input.dim() == 2, "input must be 2D tensor");
+  TORCH_CHECK(input.size(1) % 128 == 0, "k must align to 128");
+  TORCH_CHECK(input.strides()[1] == 1, "input must be row major");
+  TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
+  TORCH_CHECK(problem_sizes.dtype() == torch::kInt32,
+              "problem_sizes must be int32");
+  TORCH_CHECK(expert_offsets.dtype() == torch::kInt32,
+              "expert_offsets must be int32");
+  TORCH_CHECK(blockscale_offsets.dtype() == torch::kInt32,
+              "blockscale_offsets must be int32");
+
+  auto groups = problem_sizes.size(0);
+  TORCH_CHECK(
+      expert_offsets.dim() == 1 && expert_offsets.size(0) == groups,
+      "expert_offsets must be 1D and have size equal to the number of groups");
+  TORCH_CHECK(
+      blockscale_offsets.dim() == 1 && blockscale_offsets.size(0) == groups,
+      "blockscale_offsets must be 1D and have size equal to the number of "
+      "groups");
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+  if (input.dtype() == torch::kBFloat16) {
+    expert_specialization::launch_mxfp8_experts_quant<__nv_bfloat16>(
+        input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
+        scale_factor);
+  } else if (input.dtype() == torch::kFloat16) {
+    expert_specialization::launch_mxfp8_experts_quant<__half>(
+        input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
+        scale_factor);
+  } else {
+    TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
+  }
+#else
+  TORCH_CHECK(false,
+              "No implemented mxfp8_experts_quant for "
+              "current device");
+#endif
+}
+
+#include "core/registration.h"
+
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
+  m.impl("mxfp8_experts_quant", mxfp8_experts_quant);
+}

csrc/moe/mxfp8_moe/mxfp8_experts_quant.cuh

@@ -4,19 +4,16 @@
 // https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh
 
 #pragma once
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
 #include <cuda.h>
 #include <cuda_bf16.h>
 #include <cuda_fp16.h>
-
-#include <torch/csrc/inductor/aoti_torch/c/shim.h>
-#include <torch/csrc/stable/macros.h>
-#include <torch/csrc/stable/tensor.h>
-#include <torch/headeronly/util/Exception.h>
+#include <torch/all.h>
 
 #include <cuda/ptx>
 
 #include "cute/tensor.hpp"
-#include "libtorch_stable/torch_utils.h"
 
 namespace expert_specialization {
 
@@ -359,12 +356,12 @@ __global__ void mxfp8_experts_quant_kernel(
 }
 
 template <typename T_IN>
-void launch_mxfp8_experts_quant(const torch::stable::Tensor& input,
-                                const torch::stable::Tensor& problem_sizes,
-                                const torch::stable::Tensor& expert_offsets,
-                                const torch::stable::Tensor& blockscale_offsets,
-                                torch::stable::Tensor& quant_output,
-                                torch::stable::Tensor& scale_factor) {
+void launch_mxfp8_experts_quant(const torch::Tensor& input,
+                                const torch::Tensor& problem_sizes,
+                                const torch::Tensor& expert_offsets,
+                                const torch::Tensor& blockscale_offsets,
+                                torch::Tensor& quant_output,
+                                torch::Tensor& scale_factor) {
   ThrLayout thr_layout{};
   ValLayout val_layout{};
   SfR2SThrLayout r2s_thr_layout{};
@@ -389,18 +386,19 @@ void launch_mxfp8_experts_quant(const torch::stable::Tensor& input,
       CopyAtomR2S{}, r2s_thr_layout, r2s_val_layout);  // Tiler_MN: (16, 4)
 
   int max_active_blocks_per_sm = -1;
-  STD_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+  AT_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
       &max_active_blocks_per_sm,
       mxfp8_experts_quant_kernel<T_IN, decltype(tiled_copy_g2r),
                                  decltype(tiled_copy_r2g),
                                  decltype(tiled_copy_r2s)>,
       THREAD_BLOCK_SIZE, 0));
 
-  dim3 grid(get_device_prop()->multiProcessorCount * max_active_blocks_per_sm,
+  dim3 grid(at::cuda::getCurrentDeviceProperties()->multiProcessorCount *
+                max_active_blocks_per_sm,
             1, 1);
   dim3 block(THREAD_BLOCK_SIZE, 1, 1);
-  int num_experts = static_cast<int>(problem_sizes.size(0));
-  auto stream = get_current_cuda_stream(input.get_device_index());
+  int num_experts = (int)problem_sizes.size(0);
+  auto stream = at::cuda::getCurrentCUDAStream();
   mxfp8_experts_quant_kernel<T_IN, decltype(tiled_copy_g2r),
                              decltype(tiled_copy_r2g), decltype(tiled_copy_r2s)>
       <<<grid, block, 0, stream>>>(

csrc/ops.h

@@ -40,6 +40,18 @@ void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
 void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
                         torch::Tensor& weight, double epsilon);
 
+torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+    torch::Tensor const& q_in, torch::Tensor const& kv, torch::Tensor& k_cache,
+    torch::Tensor const& slot_mapping, torch::Tensor const& position_ids,
+    torch::Tensor const& cos_sin_cache, int64_t q_head_padded, double eps,
+    int64_t cache_block_size);
+
+void silu_and_mul_per_block_quant(torch::Tensor& out,
+                                  torch::Tensor const& input,
+                                  torch::Tensor& scales, int64_t group_size,
+                                  std::optional<torch::Tensor> scale_ub,
+                                  bool is_scale_transposed);
+
 // rotary_embedding also exist in csrc/libtorch_stable/ops.h (torch::stable
 // ABI for CUDA). It remains here because the CPU build still uses these
 // torch::Tensor declarations.
@@ -78,6 +90,8 @@ void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,
                         torch::Tensor const& seq_lens,
                         torch::Tensor const& page_table, double scale);
 
+torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
+
 void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
                               torch::Tensor const& scale,
                               std::optional<torch::Tensor> const& azp);
@@ -93,6 +107,24 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
     int64_t activation_kind);
 
 using fptr_t = int64_t;
+fptr_t init_custom_ar(const std::vector<int64_t>& fake_ipc_ptrs,
+                      torch::Tensor& rank_data, int64_t rank,
+                      bool fully_connected);
+void all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
+                fptr_t reg_buffer, int64_t reg_buffer_sz_bytes);
+void dispose(fptr_t _fa);
+int64_t meta_size();
+void register_buffer(fptr_t _fa, const std::vector<int64_t>& fake_ipc_ptrs);
+std::tuple<std::vector<int64_t>, std::vector<int64_t>>
+get_graph_buffer_ipc_meta(fptr_t _fa);
+void register_graph_buffers(fptr_t _fa,
+                            const std::vector<std::vector<int64_t>>& handles,
+                            const std::vector<std::vector<int64_t>>& offsets);
+std::tuple<int64_t, torch::Tensor> allocate_shared_buffer_and_handle(
+    int64_t size);
+int64_t open_mem_handle(torch::Tensor& mem_handle);
+void free_shared_buffer(int64_t buffer);
+
 #ifdef USE_ROCM
 fptr_t init_custom_qr(int64_t rank, int64_t world_size,
                       std::optional<int64_t> qr_max_size = std::nullopt);
@@ -102,7 +134,16 @@ void qr_open_handles(fptr_t _fa, const std::vector<torch::Tensor>& handles);
 void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
                    int64_t quant_level, bool cast_bf2half = false);
 int64_t qr_max_size();
+#endif
 
-// TODO: Remove this once ROCm upgrade to torch 2.11.
-torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
+#ifndef USE_ROCM
+torch::Tensor minimax_allreduce_rms(torch::Tensor const& input,
+                                    torch::Tensor const& norm_weight,
+                                    torch::Tensor workspace, int64_t const rank,
+                                    int64_t const nranks, double const eps);
+std::tuple<torch::Tensor, torch::Tensor> minimax_allreduce_rms_qk(
+    torch::Tensor qkv, torch::Tensor const& norm_weight_q,
+    torch::Tensor const& norm_weight_k, torch::Tensor workspace,
+    int64_t const q_size, int64_t const kv_size, int64_t const rank,
+    int64_t const nranks, double const eps);
 #endif

csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu

@@ -1,10 +1,11 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-#include "../../torch_utils.h"
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
 
 #include "../../dispatch_utils.h"
-#include "quant_conversions.cuh"
+#include "libtorch_stable/quantization/fused_kernels/quant_conversions.cuh"
 
 namespace vllm {
 
@@ -104,70 +105,64 @@ __global__ void silu_and_mul_per_block_quant_kernel(
 
 }  // namespace vllm
 
-void silu_and_mul_per_block_quant(torch::stable::Tensor& out,
-                                  torch::stable::Tensor const& input,
-                                  torch::stable::Tensor& scales,
-                                  int64_t group_size,
-                                  std::optional<torch::stable::Tensor> scale_ub,
+void silu_and_mul_per_block_quant(torch::Tensor& out,
+                                  torch::Tensor const& input,
+                                  torch::Tensor& scales, int64_t group_size,
+                                  std::optional<torch::Tensor> scale_ub,
                                   bool is_scale_transposed) {
-  static torch::headeronly::ScalarType kFp8Type =
-      is_fp8_ocp() ? torch::headeronly::ScalarType::Float8_e4m3fn
-                   : torch::headeronly::ScalarType::Float8_e4m3fnuz;
+  static c10::ScalarType kFp8Type = is_fp8_ocp()
+                                        ? c10::ScalarType::Float8_e4m3fn
+                                        : c10::ScalarType::Float8_e4m3fnuz;
 
-  STD_TORCH_CHECK(out.scalar_type() == kFp8Type ||
-                  out.scalar_type() == torch::headeronly::ScalarType::Char);
-  STD_TORCH_CHECK(out.is_contiguous() && input.is_contiguous());
-  STD_TORCH_CHECK(
-      input.scalar_type() == torch::headeronly::ScalarType::Half ||
-          input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+  TORCH_CHECK(out.dtype() == kFp8Type || out.dtype() == torch::kInt8);
+  TORCH_CHECK(out.is_contiguous() && input.is_contiguous());
+  TORCH_CHECK(
+      input.dtype() == torch::kFloat16 || input.dtype() == torch::kBFloat16,
       "Input must be FP16 or BF16");
-  STD_TORCH_CHECK(scales.scalar_type() == torch::headeronly::ScalarType::Float);
-  STD_TORCH_CHECK(group_size == 128 || group_size == 64,
-                  "Unsupported group size: ", group_size);
+  TORCH_CHECK(scales.dtype() == torch::kFloat32, "Scales must be FP32");
+  TORCH_CHECK(group_size == 128 || group_size == 64,
+              "Unsupported group size: ", group_size);
 
   if (scale_ub.has_value()) {
-    STD_TORCH_CHECK(out.scalar_type() == kFp8Type);
+    TORCH_CHECK(out.dtype() == kFp8Type);
   }
 
   int32_t hidden_size = out.size(-1);
   auto num_tokens = input.size(0);
   int32_t num_groups = hidden_size / group_size;
 
-  STD_TORCH_CHECK(input.size(-1) == hidden_size * 2,
-                  "input last dim must be 2x output hidden_size");
-  STD_TORCH_CHECK(hidden_size % group_size == 0,
-                  "hidden_size must be divisible by group_size");
+  TORCH_CHECK(input.size(-1) == hidden_size * 2,
+              "input last dim must be 2x output hidden_size");
+  TORCH_CHECK(hidden_size % group_size == 0,
+              "hidden_size must be divisible by group_size");
 
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      input.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
   dim3 grid(num_tokens, num_groups);
   dim3 block(group_size);
 
-  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+  VLLM_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "silu_and_mul_per_block_quant", [&] {
         using scalar_in_t = scalar_t;
 
-        VLLM_STABLE_DISPATCH_QUANT_TYPES(
+        VLLM_DISPATCH_QUANT_TYPES(
             out.scalar_type(), "silu_and_mul_per_block_quant", [&] {
               using scalar_out_t = scalar_t;
 
-              VLLM_STABLE_DISPATCH_GROUP_SIZE(group_size, gs, [&] {
-                VLLM_STABLE_DISPATCH_BOOL(
-                    is_scale_transposed, transpose_scale, [&] {
-                      vllm::silu_and_mul_per_block_quant_kernel<
-                          scalar_in_t, scalar_out_t, transpose_scale, gs>
-                          <<<grid, block, 0, stream>>>(
-                              out.mutable_data_ptr<scalar_out_t>(),
-                              scales.mutable_data_ptr<float>(),
-                              input.const_data_ptr<scalar_in_t>(),
-                              scale_ub.has_value()
-                                  ? scale_ub->const_data_ptr<float>()
-                                  : nullptr,
-                              hidden_size);
-                    });
+              VLLM_DISPATCH_GROUP_SIZE(group_size, gs, [&] {
+                VLLM_DISPATCH_BOOL(is_scale_transposed, transpose_scale, [&] {
+                  vllm::silu_and_mul_per_block_quant_kernel<
+                      scalar_in_t, scalar_out_t, transpose_scale, gs>
+                      <<<grid, block, 0, stream>>>(
+                          out.data_ptr<scalar_out_t>(),
+                          scales.data_ptr<float>(),
+                          input.data_ptr<scalar_in_t>(),
+                          scale_ub.has_value() ? scale_ub->data_ptr<float>()
+                                               : nullptr,
+                          hidden_size);
+                });
               });
             });
       });
-}
+}

csrc/rocm/moe_q_gemm_rdna3.cu

@@ -1,639 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-//
-// Fused MoE W4A16 GPTQ kernel for RDNA3 (gfx1100).
-//
-// Combines expert routing (sorted_token_ids / expert_ids) with the RDNA3
-// W4A16 dequant+dot from q_gemm_rdna3.cu into a single kernel launch.
-// Each block processes BLOCK_SIZE_M tokens assigned to one expert, covering
-// a tile of N output columns and K input positions.
-//
-// Weight format: same as the dense kernel — [E, K/8, N] uint32 shuffled,
-// [E, groups, N] scales, [E, groups, N/8] packed zeros.
-//
-// Design: THREADS_X=256 (8 waves on wave32), BLOCK_KN_SIZE=256, each thread
-// handles 4 N columns. Output via 64-bit packed CAS atomic-add directly to
-// the pre-zeroed output tensor (no FP32 scratch buffer).
-
-#include <cstdint>
-
-#include <torch/all.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <ATen/cuda/CUDAContext.h>
-
-#include <hip/hip_runtime.h>
-#include <hip/hip_bf16.h>
-#include <hip/hip_fp16.h>
-
-#include "qdq_4_rdna3.cuh"
-
-#if defined(__HIPCC__) && defined(__gfx1100__)
-  #define __HIP__RDNA3__
-#endif
-
-namespace vllm {
-namespace moe_gptq_rdna3 {
-
-#define BLOCK_KN_SIZE 256
-#define THREADS_X 256
-
-#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
-
-using gptq_rdna3::bf162_t;
-using gptq_rdna3::bf16_t;
-
-// --- Helpers (same as q_gemm_rdna3.cu) ---
-
-template <typename T>
-__forceinline__ __device__ T tzero();
-
-template <>
-__forceinline__ __device__ half tzero<half>() {
-  return __float2half_rn(0.0f);
-}
-
-template <>
-__forceinline__ __device__ bf16_t tzero<bf16_t>() {
-  return __float2bfloat16(0.0f);
-}
-
-__forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) {
-  float result = 0.0f;
-  const half2* a2_ptr = (const half2*)a_ptr;
-  #pragma unroll
-  for (int i = 0; i < 4; i++) {
-    result = __builtin_amdgcn_fdot2(dq[i], *a2_ptr++, result, /*clamp=*/false);
-  }
-  return result;
-}
-
-__forceinline__ __device__ float dot22_8_f(float (&dq)[8],
-                                           const bf16_t* a_ptr) {
-  float result = 0.0f;
-  #pragma unroll
-  for (int i = 0; i < 4; i++) {
-    uint32_t aw;
-    __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t));
-    float a_x = __uint_as_float((aw & 0xFFFFu) << 16);
-    float a_y = __uint_as_float(aw & 0xFFFF0000u);
-    result = __fmaf_rn(dq[2 * i + 0], a_x, result);
-    result = __fmaf_rn(dq[2 * i + 1], a_y, result);
-  }
-  return result;
-}
-
-__forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01,
-                                                   half2 v23) {
-  unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
-  unsigned long long old = *addr_u;
-  while (true) {
-    union {
-      unsigned long long u;
-      half2 h2[2];
-    } cur, sum;
-    cur.u = old;
-    sum.h2[0] = __hadd2(cur.h2[0], v01);
-    sum.h2[1] = __hadd2(cur.h2[1], v23);
-    unsigned long long prev = atomicCAS(addr_u, old, sum.u);
-    if (prev == old) break;
-    old = prev;
-  }
-}
-
-__forceinline__ __device__ void atomic_add_pk4_bf16(bf16_t* addr, bf162_t v01,
-                                                    bf162_t v23) {
-  unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
-  unsigned long long old = *addr_u;
-  while (true) {
-    union {
-      unsigned long long u;
-      bf162_t b2[2];
-    } cur, sum;
-    cur.u = old;
-    sum.b2[0] = __hadd2(cur.b2[0], v01);
-    sum.b2[1] = __hadd2(cur.b2[1], v23);
-    unsigned long long prev = atomicCAS(addr_u, old, sum.u);
-    if (prev == old) break;
-    old = prev;
-  }
-}
-
-__forceinline__ __device__ void load4_zeros(const uint32_t* qzeros_row, int n,
-                                            int (&zeros)[4]) {
-  int qcol = n / 8;
-  int shift = (n & 0x07) * 4;
-  uint32_t d = qzeros_row[qcol] >> shift;
-  zeros[0] = (int)(d & 0xF);
-  zeros[1] = (int)((d >> 4) & 0xF);
-  zeros[2] = (int)((d >> 8) & 0xF);
-  zeros[3] = (int)((d >> 12) & 0xF);
-}
-
-template <typename T>
-__forceinline__ __device__ void load4_scales(const T* scales_row, int n,
-                                             T (&scales)[4]) {
-  scales[0] = scales_row[n + 0];
-  scales[1] = scales_row[n + 1];
-  scales[2] = scales_row[n + 2];
-  scales[3] = scales_row[n + 3];
-}
-
-// ---------------------------------------------------------------------------
-// Fused MoE kernel.
-// ---------------------------------------------------------------------------
-
-template <typename T, int BLOCK_SIZE_M>
-__global__ void moe_gemm_q4_kernel_rdna3(
-    const T* __restrict__ a,                  // [size_m, size_k] or [M*topk, K]
-    T* __restrict__ c,                        // [M*topk, size_n] pre-zeroed
-    const uint32_t* __restrict__ b_q_weight,  // [E, K/8, N] packed
-    const T* __restrict__ b_scales,           // [E, groups, N]
-    const uint32_t* __restrict__ b_qzeros,    // [E, groups, N/8] packed
-    const float* __restrict__ topk_weights,   // [M*topk] or nullptr
-    const int32_t* __restrict__ sorted_token_ids,
-    const int32_t* __restrict__ expert_ids,
-    const int32_t* __restrict__ num_tokens_post_padded,
-    const int size_m,  // total tokens (original M, or M*topk for w2)
-    const int size_n,  // output features per expert
-    const int size_k,  // input features
-    const int groups,  // K / group_size
-    const int top_k,   // routing top-k (1 for w2 pass)
-    // Per-expert strides (in elements, not bytes)
-    const int expert_weight_stride,  // (K/8) * N
-    const int expert_scales_stride,  // groups * N
-    const int expert_zeros_stride,   // groups * (N/8)
-    const bool mul_topk_weight,
-    const int output_topk) {  // >0: reduce output by token_id/output_topk
-  const int t = threadIdx.x;
-  const int token_block = blockIdx.x;
-  const int offset_n = blockIdx.y * BLOCK_KN_SIZE * 4;
-  const int offset_k = blockIdx.z * BLOCK_KN_SIZE;
-  const int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
-  const int n = offset_n + t * 4;
-
-  // Early exit for padding blocks or invalid experts (expert_map = -1)
-  if (token_block * BLOCK_SIZE_M >= num_tokens_post_padded[0]) return;
-
-  const int expert_id = expert_ids[token_block];
-  if (expert_id == -1) return;
-
-  // Expert-specific pointers
-  const uint32_t* expert_weights =
-      b_q_weight + (int64_t)expert_id * expert_weight_stride;
-  const T* expert_scales = b_scales + (int64_t)expert_id * expert_scales_stride;
-  const uint32_t* expert_qzeros =
-      b_qzeros + (int64_t)expert_id * expert_zeros_stride;
-
-  // LDS for activations
-  constexpr int LDS_PAD = 8;
-  __shared__ T block_a[BLOCK_SIZE_M][BLOCK_KN_SIZE + LDS_PAD];
-
-  static_assert(BLOCK_KN_SIZE == THREADS_X,
-                "BLOCK_KN_SIZE must equal THREADS_X");
-
-  // For bf16 M=1, we can skip LDS and read A from global (same as dense).
-  // fp16 always needs LDS due to the dot22_8_f indexing pattern.
-  constexpr bool USE_LDS_A = (BLOCK_SIZE_M > 1) || std::is_same<T, half>::value;
-
-  const int offset_m_base = token_block * BLOCK_SIZE_M;
-
-  if constexpr (USE_LDS_A) {
-    if (offset_k + t < end_k) {
-  #pragma unroll
-      for (int m = 0; m < BLOCK_SIZE_M; ++m) {
-        int32_t token_id = sorted_token_ids[offset_m_base + m];
-        int token_row = token_id / top_k;
-        T av;
-        if (token_row < size_m) {
-          av = a[(int64_t)token_row * size_k + offset_k + t];
-        } else {
-          av = tzero<T>();
-        }
-        block_a[m][t] = av;
-      }
-    }
-    __syncthreads();
-  }
-
-  if (n >= size_n) return;
-
-  // Group bookkeeping
-  const int groupsize = size_k / groups;
-  int group = offset_k / groupsize;
-  int nextgroup = (group + 1) * groupsize;
-
-  // Weight pointer for this expert
-  int qk = offset_k / 8;
-  const uint32_t* b_ptr = expert_weights + qk * size_n + n;
-
-  // Per-column dequant constants (4 columns per thread)
-  half2 z1z16_h[4][2], y1y16_h[4][2];
-  float z_b_f[4], y_b_f[4];
-
-  // GPTQv1: zero_offset = 1
-  constexpr int zero_offset = 1;
-
-  auto refresh_group = [&](int g) {
-    const uint32_t* qz_row = expert_qzeros + g * (size_n / 8);
-    const T* sc_row = expert_scales + g * size_n;
-    int zeros[4];
-    T scales[4];
-    load4_zeros(qz_row, n, zeros);
-    load4_scales<T>(sc_row, n, scales);
-    if constexpr (std::is_same<T, half>::value) {
-  #pragma unroll
-      for (int i = 0; i < 4; ++i) {
-        gptq_rdna3::prep_zero_scale_fp16((uint32_t)(zeros[i] + zero_offset),
-                                         scales[i], z1z16_h[i], y1y16_h[i]);
-      }
-    } else {
-  #pragma unroll
-      for (int i = 0; i < 4; ++i) {
-        gptq_rdna3::prep_zero_scale_bf16_f32((uint32_t)(zeros[i] + zero_offset),
-                                             scales[i], z_b_f[i], y_b_f[i]);
-      }
-    }
-  };
-
-  refresh_group(group);
-
-  float block_c[BLOCK_SIZE_M][4];
-  #pragma unroll
-  for (int m = 0; m < BLOCK_SIZE_M; ++m) {
-  #pragma unroll
-    for (int j = 0; j < 4; ++j) block_c[m][j] = 0.0f;
-  }
-
-  // --- Main K-loop ---
-  int k = offset_k;
-  while (k < end_k) {
-    if (k == nextgroup) {
-      group++;
-      nextgroup += groupsize;
-      refresh_group(group);
-    }
-
-    // Prefetch 4 weight words (128 bytes)
-    int4 b_w[4];
-  #pragma unroll
-    for (int j = 0; j < 4; ++j) {
-      b_w[j] = *(const int4*)(b_ptr + j * size_n);
-    }
-    b_ptr += 4 * size_n;
-
-  #pragma unroll
-    for (int j = 0; j < 4; ++j) {
-      const int a_off = (k - offset_k) + 8 * j;
-
-      if constexpr (std::is_same<T, half>::value) {
-        // fp16 path: dequant via bit-trick, dot via v_dot2_f32_f16
-        half2 dq[4][4];
-        gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].x, dq[0], z1z16_h[0],
-                                        y1y16_h[0]);
-        gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].y, dq[1], z1z16_h[1],
-                                        y1y16_h[1]);
-        gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].z, dq[2], z1z16_h[2],
-                                        y1y16_h[2]);
-        gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].w, dq[3], z1z16_h[3],
-                                        y1y16_h[3]);
-
-  #pragma unroll
-        for (int m = 0; m < BLOCK_SIZE_M; ++m) {
-          const half* a_ptr = reinterpret_cast<const half*>(&block_a[m][a_off]);
-          block_c[m][0] += dot22_8_f(dq[0], a_ptr);
-          block_c[m][1] += dot22_8_f(dq[1], a_ptr);
-          block_c[m][2] += dot22_8_f(dq[2], a_ptr);
-          block_c[m][3] += dot22_8_f(dq[3], a_ptr);
-        }
-      } else if constexpr (BLOCK_SIZE_M == 1) {
-        // bf16 M=1: v_dot2_f32_bf16 with InstCombine-defeating opacity
-        typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
-        constexpr uint32_t BF16_MAGIC = 0x43004300u;
-        constexpr uint32_t BF16_ONES = 0x3F803F80u;
-        union pack4 {
-          float f[4];
-          uint32_t u[4];
-        };
-
-        uint32_t w[4];
-        __builtin_memcpy(w, &b_w[j], sizeof(int4));
-
-        // Load activations — read from global (no LDS for bf16 M=1)
-        pack4 a_pack;
-        {
-          int32_t token_id = sorted_token_ids[offset_m_base];
-          int token_row = token_id / top_k;
-          if (token_row < size_m) {
-            const uint32_t* a_words = reinterpret_cast<const uint32_t*>(
-                a + (int64_t)token_row * size_k + offset_k + a_off);
-            a_pack.u[0] = a_words[0];
-            a_pack.u[1] = a_words[1];
-            a_pack.u[2] = a_words[2];
-            a_pack.u[3] = a_words[3];
-          } else {
-            a_pack.u[0] = 0;
-            a_pack.u[1] = 0;
-            a_pack.u[2] = 0;
-            a_pack.u[3] = 0;
-          }
-        }
-
-        // sum_a for bias correction
-        float sum_a = 0.0f;
-  #pragma unroll
-        for (int b = 0; b < 4; ++b) {
-          sum_a = __builtin_amdgcn_fdot2_f32_bf16(
-              *((bf16x2_t*)(&a_pack.f[b])), *((const bf16x2_t*)&BF16_ONES),
-              sum_a, /*clamp=*/false);
-        }
-
-  #pragma unroll 1
-        for (int col = 0; col < 4; ++col) {
-          pack4 q_pack;
-          const uint32_t qa = w[col];
-          q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
-          q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
-          q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
-          q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
-
-          float partial = 0.0f;
-  #pragma unroll
-          for (int b = 0; b < 4; ++b) {
-            partial = __builtin_amdgcn_fdot2_f32_bf16(
-                *((bf16x2_t*)(&a_pack.f[b])), *((bf16x2_t*)(&q_pack.f[b])),
-                partial, /*clamp=*/false);
-          }
-
-          block_c[0][col] =
-              __fmaf_rn(y_b_f[col], partial,
-                        __fmaf_rn(z_b_f[col], sum_a, block_c[0][col]));
-        }
-      } else {
-        // bf16 M>1: v_dot2_f32_bf16 with LDS-staged activations
-        typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
-        constexpr uint32_t BF16_MAGIC = 0x43004300u;
-        constexpr uint32_t BF16_ONES = 0x3F803F80u;
-        union pack4 {
-          float f[4];
-          uint32_t u[4];
-        };
-
-        uint32_t w[4];
-        __builtin_memcpy(w, &b_w[j], sizeof(int4));
-
-        pack4 a_pack[BLOCK_SIZE_M];
-  #pragma unroll
-        for (int m = 0; m < BLOCK_SIZE_M; ++m) {
-          const uint32_t* a_words =
-              reinterpret_cast<const uint32_t*>(&block_a[m][a_off]);
-          a_pack[m].u[0] = a_words[0];
-          a_pack[m].u[1] = a_words[1];
-          a_pack[m].u[2] = a_words[2];
-          a_pack[m].u[3] = a_words[3];
-        }
-
-        float sum_a[BLOCK_SIZE_M];
-  #pragma unroll
-        for (int m = 0; m < BLOCK_SIZE_M; ++m) {
-          float s = 0.0f;
-  #pragma unroll
-          for (int b = 0; b < 4; ++b) {
-            s = __builtin_amdgcn_fdot2_f32_bf16(*((bf16x2_t*)(&a_pack[m].f[b])),
-                                                *((const bf16x2_t*)&BF16_ONES),
-                                                s, /*clamp=*/false);
-          }
-          sum_a[m] = s;
-        }
-
-  #pragma unroll 1
-        for (int col = 0; col < 4; ++col) {
-          pack4 q_pack;
-          const uint32_t qa = w[col];
-          q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
-          q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
-          q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
-          q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
-
-  #pragma unroll
-          for (int m = 0; m < BLOCK_SIZE_M; ++m) {
-            float partial = 0.0f;
-  #pragma unroll
-            for (int b = 0; b < 4; ++b) {
-              partial = __builtin_amdgcn_fdot2_f32_bf16(
-                  *((bf16x2_t*)(&a_pack[m].f[b])), *((bf16x2_t*)(&q_pack.f[b])),
-                  partial, /*clamp=*/false);
-            }
-            block_c[m][col] =
-                __fmaf_rn(y_b_f[col], partial,
-                          __fmaf_rn(z_b_f[col], sum_a[m], block_c[m][col]));
-          }
-        }
-      }
-    }
-    k += 32;
-  }
-
-  // --- Epilogue: apply topk_weight and atomic-add to output ---
-  #pragma unroll
-  for (int m = 0; m < BLOCK_SIZE_M; ++m) {
-    int32_t token_id = sorted_token_ids[offset_m_base + m];
-    if (token_id / top_k >= size_m) continue;
-
-    // Apply router weight
-    if (mul_topk_weight && topk_weights != nullptr) {
-      float tw = topk_weights[token_id];
-  #pragma unroll
-      for (int j = 0; j < 4; ++j) block_c[m][j] *= tw;
-    }
-
-    // output_topk > 0: reduce by mapping token_id back to original token
-    // (multiple experts write to the same row via atomics)
-    int64_t out_row = (output_topk > 0) ? (int64_t)(token_id / output_topk)
-                                        : (int64_t)token_id;
-    T* out = c + out_row * size_n + n;
-    if constexpr (std::is_same<T, half>::value) {
-      half2 r01 = __halves2half2(__float2half_rn(block_c[m][0]),
-                                 __float2half_rn(block_c[m][1]));
-      half2 r23 = __halves2half2(__float2half_rn(block_c[m][2]),
-                                 __float2half_rn(block_c[m][3]));
-      atomic_add_pk4_f16(out, r01, r23);
-    } else {
-      bf162_t r01;
-      r01.x = __float2bfloat16(block_c[m][0]);
-      r01.y = __float2bfloat16(block_c[m][1]);
-      bf162_t r23;
-      r23.x = __float2bfloat16(block_c[m][2]);
-      r23.y = __float2bfloat16(block_c[m][3]);
-      atomic_add_pk4_bf16(out, r01, r23);
-    }
-  }
-}
-
-#else  // non-RDNA3: empty stub for symbol parity
-
-template <typename T, int BLOCK_SIZE_M>
-__global__ void moe_gemm_q4_kernel_rdna3(
-    const T*, T*, const uint32_t*, const T*, const uint32_t*, const float*,
-    const int32_t*, const int32_t*, const int32_t*, const int, const int,
-    const int, const int, const int, const int, const int, const int,
-    const bool, const int) {}
-
-#endif  // __HIP__RDNA3__ || !__HIP_DEVICE_COMPILE__
-
-// ---------------------------------------------------------------------------
-// Launcher
-// ---------------------------------------------------------------------------
-
-template <typename T, int BLOCK_SIZE_M>
-void launch_moe_gemm_q4(
-    const T* a, T* c, const uint32_t* b_q_weight, const T* b_scales,
-    const uint32_t* b_qzeros, const float* topk_weights,
-    const int32_t* sorted_token_ids, const int32_t* expert_ids,
-    const int32_t* num_tokens_post_padded, int num_token_blocks, int size_m,
-    int size_n, int size_k, int groups, int top_k, int expert_weight_stride,
-    int expert_scales_stride, int expert_zeros_stride, bool mul_topk_weight,
-    int output_topk, cudaStream_t stream) {
-  dim3 block(THREADS_X);
-  dim3 grid(num_token_blocks,
-            (size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4),
-            (size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE);
-
-  moe_gemm_q4_kernel_rdna3<T, BLOCK_SIZE_M><<<grid, block, 0, stream>>>(
-      a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
-      expert_ids, num_tokens_post_padded, size_m, size_n, size_k, groups, top_k,
-      expert_weight_stride, expert_scales_stride, expert_zeros_stride,
-      mul_topk_weight, output_topk);
-}
-
-template <typename T>
-void dispatch_moe_gemm_q4(
-    const T* a, T* c, const uint32_t* b_q_weight, const T* b_scales,
-    const uint32_t* b_qzeros, const float* topk_weights,
-    const int32_t* sorted_token_ids, const int32_t* expert_ids,
-    const int32_t* num_tokens_post_padded, int num_token_blocks, int size_m,
-    int size_n, int size_k, int groups, int top_k, int block_size_m,
-    int expert_weight_stride, int expert_scales_stride, int expert_zeros_stride,
-    bool mul_topk_weight, int output_topk, cudaStream_t stream) {
-  // Dispatch to template instantiation based on block_size_m
-  switch (block_size_m) {
-    case 1:
-      launch_moe_gemm_q4<T, 1>(
-          a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
-          expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
-          size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
-          expert_zeros_stride, mul_topk_weight, output_topk, stream);
-      break;
-    case 2:
-      launch_moe_gemm_q4<T, 2>(
-          a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
-          expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
-          size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
-          expert_zeros_stride, mul_topk_weight, output_topk, stream);
-      break;
-    case 4:
-      launch_moe_gemm_q4<T, 4>(
-          a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
-          expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
-          size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
-          expert_zeros_stride, mul_topk_weight, output_topk, stream);
-      break;
-    case 8:
-      launch_moe_gemm_q4<T, 8>(
-          a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
-          expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
-          size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
-          expert_zeros_stride, mul_topk_weight, output_topk, stream);
-      break;
-    default:
-      TORCH_CHECK(false,
-                  "moe_gptq_gemm_rdna3: block_size_m must be 1, 2, 4, or 8, "
-                  "got ",
-                  block_size_m);
-  }
-}
-
-}  // namespace moe_gptq_rdna3
-}  // namespace vllm
-
-// ---------------------------------------------------------------------------
-// Public entry point
-// ---------------------------------------------------------------------------
-//
-// Inputs:
-//   a                      [M, K] or [M*top_k, K]  half or bfloat16
-//   c                      [M*top_k, N]             same dtype (pre-zeroed!)
-//   b_q_weight             [E, K/8, N]              uint32 (shuffled)
-//   b_scales               [E, groups, N]           same dtype as a
-//   b_qzeros               [E, groups, N/8]         uint32 (packed 4-bit)
-//   topk_weights           [M*top_k] or empty       float32
-//   sorted_token_ids       [num_blocks * block_m]   int32
-//   expert_ids             [num_blocks]              int32
-//   num_tokens_post_padded [1]                       int32
-//   top_k                  int
-//   block_size_m           int (1, 2, 4, or 8)
-//   mul_topk_weight        bool
-
-void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c,
-                         torch::Tensor b_q_weight, torch::Tensor b_scales,
-                         torch::Tensor b_qzeros, torch::Tensor topk_weights,
-                         torch::Tensor sorted_token_ids,
-                         torch::Tensor expert_ids,
-                         torch::Tensor num_tokens_post_padded, int64_t top_k,
-                         int64_t block_size_m, bool mul_topk_weight,
-                         int64_t output_topk) {
-  TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor");
-  TORCH_CHECK(c.is_cuda(), "c must be a CUDA/HIP tensor");
-  TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor");
-  TORCH_CHECK(a.dim() == 2, "a must be 2D");
-  TORCH_CHECK(c.dim() == 2, "c must be 2D");
-  TORCH_CHECK(b_q_weight.dim() == 3, "b_q_weight must be 3D [E, K/8, N]");
-  TORCH_CHECK(b_scales.dim() == 3, "b_scales must be 3D [E, groups, N]");
-  TORCH_CHECK(b_qzeros.dim() == 3, "b_qzeros must be 3D [E, groups, N/8]");
-  TORCH_CHECK(
-      a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16,
-      "a must be half or bfloat16");
-  TORCH_CHECK(a.scalar_type() == b_scales.scalar_type(),
-              "b_scales dtype must match a");
-
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
-  auto stream = at::cuda::getCurrentCUDAStream();
-
-  int size_m = (int)a.size(0);
-  int size_k = (int)a.size(1);
-  int size_n = (int)b_q_weight.size(2);
-  int groups = (int)b_scales.size(1);
-
-  // Per-expert strides
-  int expert_weight_stride = (int)(b_q_weight.size(1) * b_q_weight.size(2));
-  int expert_scales_stride = (int)(b_scales.size(1) * b_scales.size(2));
-  int expert_zeros_stride = (int)(b_qzeros.size(1) * b_qzeros.size(2));
-
-  int num_token_blocks = (int)(sorted_token_ids.size(0) / block_size_m);
-
-  const float* topk_w_ptr =
-      (topk_weights.numel() > 0) ? topk_weights.data_ptr<float>() : nullptr;
-
-  // Manual dtype dispatch using HIP native types (c10::Half/BFloat16 don't
-  // implicitly convert to half/__hip_bfloat16 in device code).
-  using vllm::gptq_rdna3::bf16_t;
-
-  auto dispatch = [&](auto* a_ptr, auto* c_ptr, const auto* s_ptr) {
-    using T = std::remove_const_t<std::remove_pointer_t<decltype(a_ptr)>>;
-    vllm::moe_gptq_rdna3::dispatch_moe_gemm_q4<T>(
-        a_ptr, c_ptr, (const uint32_t*)b_q_weight.data_ptr<int32_t>(), s_ptr,
-        (const uint32_t*)b_qzeros.data_ptr<int32_t>(), topk_w_ptr,
-        sorted_token_ids.data_ptr<int32_t>(), expert_ids.data_ptr<int32_t>(),
-        num_tokens_post_padded.data_ptr<int32_t>(), num_token_blocks, size_m,
-        size_n, size_k, groups, (int)top_k, (int)block_size_m,
-        expert_weight_stride, expert_scales_stride, expert_zeros_stride,
-        mul_topk_weight, (int)output_topk, stream);
-  };
-
-  if (a.scalar_type() == torch::kHalf) {
-    dispatch((const half*)a.data_ptr(), (half*)c.data_ptr(),
-             (const half*)b_scales.data_ptr());
-  } else {
-    dispatch((const bf16_t*)a.data_ptr(), (bf16_t*)c.data_ptr(),
-             (const bf16_t*)b_scales.data_ptr());
-  }
-}

csrc/rocm/ops.h

@@ -27,15 +27,6 @@ torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
                                    torch::Tensor b_scales,
                                    torch::Tensor b_g_idx, bool use_v2_format);
 
-void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c,
-                         torch::Tensor b_q_weight, torch::Tensor b_scales,
-                         torch::Tensor b_qzeros, torch::Tensor topk_weights,
-                         torch::Tensor sorted_token_ids,
-                         torch::Tensor expert_ids,
-                         torch::Tensor num_tokens_post_padded, int64_t top_k,
-                         int64_t block_size_m, bool mul_topk_weight,
-                         int64_t output_topk);
-
 void paged_attention(
     torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
     torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,

csrc/rocm/torch_bindings.cpp

@@ -50,15 +50,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
       "gptq_gemm_rdna3_wmma(Tensor a, Tensor b_q_weight, Tensor b_qzeros, "
       "Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor");
   rocm_ops.impl("gptq_gemm_rdna3_wmma", torch::kCUDA, &gptq_gemm_rdna3_wmma);
-
-  rocm_ops.def(
-      "moe_gptq_gemm_rdna3(Tensor a, Tensor! c, Tensor b_q_weight, "
-      "Tensor b_scales, Tensor b_qzeros, Tensor topk_weights, "
-      "Tensor sorted_token_ids, Tensor expert_ids, "
-      "Tensor num_tokens_post_padded, "
-      "int top_k, int block_size_m, bool mul_topk_weight, "
-      "int output_topk) -> ()");
-  rocm_ops.impl("moe_gptq_gemm_rdna3", torch::kCUDA, &moe_gptq_gemm_rdna3);
 #endif
 
   // Custom attention op

csrc/torch_bindings.cpp

@@ -32,24 +32,37 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("weak_ref_tensor(Tensor input) -> Tensor");
   ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
 
-#ifdef USE_ROCM
-  // TODO: Remove this once we upgrade to torch 2.11.
-  // ROCm still uses torch 2.10,
-  // So we still need to use unstable torch ABI for now.
   ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
   ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
            &get_cuda_view_from_cpu_tensor);
-#endif
 
   // Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
   ops.def(
       "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
   ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
 
+  // Fused SiLU+Mul + per-block quantization
+  ops.def(
+      "silu_and_mul_per_block_quant("
+      "Tensor! out, "
+      "Tensor input, "
+      "Tensor! scales, "
+      "int group_size, "
+      "Tensor? scale_ub=None, "
+      "bool is_scale_transposed=False) -> ()");
+  ops.impl("silu_and_mul_per_block_quant", torch::kCUDA,
+           &silu_and_mul_per_block_quant);
+
   // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
   // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
-  // kernel launch. Registered in _C_stable_libtorch (incl. the FlashInfer V4
-  // full-cache bf16/fp8 variants).
+  // kernel launch.
+  ops.def(
+      "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
+      "Tensor q_in, Tensor kv, Tensor! k_cache, "
+      "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
+      "int q_head_padded, float eps, int cache_block_size) -> Tensor");
+  ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
+           &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
 
   // Quantization ops
 #ifndef USE_ROCM
@@ -150,29 +163,37 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   // conditionally compiled so impl registration is in source file
 
 #endif
+
+#ifndef USE_ROCM
+  ops.def(
+      "minimax_allreduce_rms("
+      "Tensor input,"
+      "Tensor norm_weight,"
+      "Tensor workspace,"
+      "int rank,"
+      "int nranks,"
+      "float eps) -> Tensor");
+  ops.impl("minimax_allreduce_rms", torch::kCUDA, &minimax_allreduce_rms);
+  ops.def(
+      "minimax_allreduce_rms_qk("
+      "Tensor qkv,"
+      "Tensor norm_weight_q,"
+      "Tensor norm_weight_k,"
+      "Tensor workspace,"
+      "int q_size,"
+      "int kv_size,"
+      "int rank,"
+      "int nranks,"
+      "float eps) -> (Tensor, Tensor)");
+  ops.impl("minimax_allreduce_rms_qk", torch::kCUDA, &minimax_allreduce_rms_qk);
+
+  //  conditionally compiled so impl in source file
+#endif
 }
 
-#ifdef USE_ROCM
-TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
-  // Quick Reduce all-reduce kernels (ROCm-only; stays on legacy _C).
-  custom_ar.def(
-      "qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool "
-      "cast_bf2half) -> ()");
-  custom_ar.impl("qr_all_reduce", torch::kCUDA, &qr_all_reduce);
-
-  custom_ar.def("init_custom_qr", &init_custom_qr);
-  custom_ar.def("qr_destroy", &qr_destroy);
-  custom_ar.def("qr_get_handle", &qr_get_handle);
-
-  custom_ar.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()");
-  custom_ar.impl("qr_open_handles", torch::kCPU, &qr_open_handles);
-
-  custom_ar.def("qr_max_size", &qr_max_size);
-}
-
-// TODO: Remove this once ROCm upgrade to torch 2.11.
 TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
   // Cuda utils
+
   // Gets the specified device attribute.
   cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int");
   cuda_utils.impl("get_device_attribute", &get_device_attribute);
@@ -183,6 +204,49 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
   cuda_utils.impl("get_max_shared_memory_per_block_device_attribute",
                   &get_max_shared_memory_per_block_device_attribute);
 }
+
+TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
+  // Custom all-reduce kernels
+  custom_ar.def(
+      "init_custom_ar(int[] ipc_tensors, Tensor rank_data, "
+      "int rank, bool fully_connected) -> int");
+  custom_ar.impl("init_custom_ar", torch::kCUDA, &init_custom_ar);
+  custom_ar.def(
+      "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, "
+      "int reg_buffer_sz_bytes) -> ()");
+  custom_ar.impl("all_reduce", torch::kCUDA, &all_reduce);
+
+  custom_ar.def("dispose", &dispose);
+  custom_ar.def("meta_size", &meta_size);
+
+  custom_ar.def("register_buffer", &register_buffer);
+  custom_ar.def("get_graph_buffer_ipc_meta", &get_graph_buffer_ipc_meta);
+  custom_ar.def("register_graph_buffers", &register_graph_buffers);
+
+  custom_ar.def("allocate_shared_buffer_and_handle",
+                &allocate_shared_buffer_and_handle);
+  custom_ar.def("open_mem_handle(Tensor mem_handle) -> int", &open_mem_handle);
+  custom_ar.impl("open_mem_handle", torch::kCPU, &open_mem_handle);
+
+  custom_ar.def("free_shared_buffer", &free_shared_buffer);
+#ifdef USE_ROCM
+  // Quick Reduce all-reduce kernels
+  custom_ar.def(
+      "qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool "
+      "cast_bf2half) -> ()");
+  custom_ar.impl("qr_all_reduce", torch::kCUDA, &qr_all_reduce);
+
+  custom_ar.def("init_custom_qr", &init_custom_qr);
+  custom_ar.def("qr_destroy", &qr_destroy);
+
+  custom_ar.def("qr_get_handle", &qr_get_handle);
+
+  custom_ar.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()");
+  custom_ar.impl("qr_open_handles", torch::kCPU, &qr_open_handles);
+
+  // Max input size in bytes
+  custom_ar.def("qr_max_size", &qr_max_size);
 #endif
+}
 
 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

csrc/type_convert.cuh

@@ -50,7 +50,7 @@ struct _typeConvert<float> {
 #if defined(USE_ROCM) || (defined(CUDA_VERSION) && (CUDA_VERSION >= 12000))
 // CUDA < 12.0 runs into issues with packed type conversion
 template <>
-struct _typeConvert<torch::headeronly::Half> {
+struct _typeConvert<c10::Half> {
   static constexpr bool exists = true;
   using hip_type = __half;
   using packed_hip_type = __half2;
@@ -73,7 +73,7 @@ struct _typeConvert<torch::headeronly::Half> {
 // CUDA_ARCH < 800 does not have BF16 support
 // ROCm 7.0+ supports bfloat16
 template <>
-struct _typeConvert<torch::headeronly::BFloat16> {
+struct _typeConvert<c10::BFloat16> {
   static constexpr bool exists = true;
   using hip_type = __nv_bfloat16;
   using packed_hip_type = __nv_bfloat162;

docker/Dockerfile

@@ -218,10 +218,6 @@ COPY requirements/common.txt requirements/common.txt
 COPY requirements/cuda.txt requirements/cuda.txt
 COPY use_existing_torch.py use_existing_torch.py
 COPY pyproject.toml pyproject.toml
-# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that
-# share paths with different content. uv can extract them in either order,
-# leaving base files that break CUDA 13 CuTe DSL JIT.
-# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed.
 RUN --mount=type=cache,target=/opt/uv/cache \
     if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
         sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' requirements/cuda.txt; \
@@ -238,13 +234,6 @@ RUN --mount=type=cache,target=/opt/uv/cache \
     else \
         uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \
         --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
-    fi \
-    && if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \
-        CUTLASS_DSL_VERSION=$(uv pip show --python /opt/venv/bin/python3 nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \
-        if [ -n "$CUTLASS_DSL_VERSION" ]; then \
-            uv pip install --python /opt/venv/bin/python3 --force-reinstall --no-deps \
-                "nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \
-        fi; \
     fi
 
 # Track PyTorch lib versions used during build and match in downstream instances.
@@ -756,10 +745,6 @@ ENV VLLM_ENABLE_CUDA_COMPATIBILITY=0
 ARG PYTORCH_CUDA_INDEX_BASE_URL
 COPY requirements/common.txt /tmp/common.txt
 COPY requirements/cuda.txt /tmp/requirements-cuda.txt
-# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that
-# share paths with different content. uv can extract them in either order,
-# leaving base files that break CUDA 13 CuTe DSL JIT.
-# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed.
 RUN --mount=type=cache,target=/opt/uv/cache \
     if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
         sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' /tmp/requirements-cuda.txt; \
@@ -767,22 +752,15 @@ RUN --mount=type=cache,target=/opt/uv/cache \
     fi && \
     uv pip install --system -r /tmp/requirements-cuda.txt \
         --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') && \
-    if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \
-        CUTLASS_DSL_VERSION=$(uv pip show --system nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \
-        if [ -n "$CUTLASS_DSL_VERSION" ]; then \
-            uv pip install --system --force-reinstall --no-deps \
-                "nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \
-        fi; \
-    fi && \
     rm /tmp/requirements-cuda.txt /tmp/common.txt
 
 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
 # https://docs.flashinfer.ai/installation.html
 # From versions.json: .flashinfer.version
-ARG FLASHINFER_VERSION=0.6.12
+ARG FLASHINFER_VERSION=0.6.11.post2
 RUN --mount=type=cache,target=/opt/uv/cache \
     uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
-        --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
+        --extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
 
 # ============================================================
 # OPENAI API SERVER DEPENDENCIES
@@ -864,19 +842,6 @@ RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm
     uv pip install --system ep_kernels/dist/*.whl --verbose \
         --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
 
-# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that
-# share paths with different content. Force -libs-cu13 last after runtime
-# dependency installs so uv cannot leave base files behind.
-# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed.
-RUN --mount=type=cache,target=/opt/uv/cache \
-    if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \
-        CUTLASS_DSL_VERSION=$(uv pip show --system nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \
-        if [ -n "$CUTLASS_DSL_VERSION" ]; then \
-            uv pip install --system --force-reinstall --no-deps \
-                "nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \
-        fi; \
-    fi
-
 # Download FlashInfer precompiled cubins AFTER all pip installs are done.
 # This must run after the vLLM wheel and EP kernels installs above, because
 # those can reinstall/touch flashinfer packages. Downloading cubins earlier

docker/Dockerfile.cpu

@@ -168,12 +168,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 ######################### TRITON-CPU BUILD IMAGE #########################
 FROM base AS vllm-triton-cpu-build
 
-# Support for cross-compilation with x86 ISA including AVX2 and AVX512: docker build --build-arg VLLM_CPU_X86="true" ...
-# Re-declared here because this stage is `FROM base` (not `vllm-build`), so it
-# does not inherit the ARG/ENV defined there. Without it, the guard below would
-# see an empty value and build triton-cpu on non-x86 targets (e.g. arm64).
-ARG VLLM_CPU_X86=0
-
 WORKDIR /vllm-workspace
 
 RUN mkdir dist
@@ -275,11 +269,6 @@ ENV HF_HUB_DOWNLOAD_TIMEOUT 60
 ######################### RELEASE IMAGE #########################
 FROM base AS vllm-openai
 
-# Re-declared here because this stage is `FROM base` (not `vllm-build`), so the
-# RUN below that gates the triton-cpu wheel install on $VLLM_CPU_X86 would
-# otherwise see an empty value and try to install it on non-x86 targets.
-ARG VLLM_CPU_X86=0
-
 WORKDIR /vllm-workspace
 
 RUN --mount=type=cache,target=/root/.cache/uv \

docker/Dockerfile.nightly_torch

@@ -256,13 +256,13 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.
 
 
 # build flashinfer for torch nightly from source around 10 mins
-# release version: v0.6.12
+# release version: v0.6.11.post2
 # todo(elainewy): cache flashinfer build result for faster build
 ENV CCACHE_DIR=/root/.cache/ccache
 RUN --mount=type=cache,target=/root/.cache/ccache \
     --mount=type=cache,target=/root/.cache/uv \
     echo "git clone flashinfer..." \
-    && git clone --depth 1 --branch v0.6.12 --recursive https://github.com/flashinfer-ai/flashinfer.git \
+    && git clone --depth 1 --branch v0.6.11.post2 --recursive https://github.com/flashinfer-ai/flashinfer.git \
     && cd flashinfer \
     && git submodule update --init --recursive \
     && echo "finish git clone flashinfer..." \

docker/Dockerfile.rocm

@@ -2,7 +2,6 @@
 ARG REMOTE_VLLM="0"
 ARG COMMON_WORKDIR=/app
 ARG BASE_IMAGE=rocm/vllm-dev:base
-ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base
 # NIC backend for MoRI RDMA support.
 # By default (all), drivers and userspace libraries for all supported NIC types
 # (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime.
@@ -17,8 +16,7 @@ ARG NIC_BACKEND=all
 ARG AINIC_VERSION=1.117.3-hydra
 ARG UBUNTU_CODENAME=jammy
 
-# Sccache configuration. Release builds use this today; CI can opt in when a
-# shared S3-compatible cache backend is available.
+# Sccache configuration (only used in release pipeline)
 ARG USE_SCCACHE
 ARG SCCACHE_DOWNLOAD_URL
 ARG SCCACHE_ENDPOINT
@@ -31,16 +29,12 @@ FROM ${BASE_IMAGE} AS base
 ARG ARG_PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
 
-# Install build dependencies and utilities
+# Install some basic utilities
 RUN apt-get update -q -y && apt-get install -q -y \
     sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \
     apt-transport-https ca-certificates wget curl \
-    libnuma-dev ccache mold
-RUN --mount=type=cache,target=/root/.cache/pip \
-    python3 -m pip install --upgrade pip
-# Note: mold is installed but not set as the system default linker because
-# some packages use JIT compilation at runtime with flags mold does not support.
-# Build stages opt in via LDFLAGS="-fuse-ld=mold".
+    libnuma-dev
+RUN python3 -m pip install --upgrade pip
 # Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_base)
 ARG USE_SCCACHE
 RUN if [ "$USE_SCCACHE" != "1" ]; then \
@@ -61,12 +55,6 @@ ENV UV_HTTP_TIMEOUT=500
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 # Use copy mode to avoid hardlink failures with Docker cache mounts
 ENV UV_LINK_MODE=copy
-# ccache directory - persisted across layer rebuilds via cache mounts.
-ENV CCACHE_DIR=/root/.cache/ccache
-ENV CCACHE_COMPILERCHECK=content
-# Empty by default so build steps fall back to $(nproc); CI can override.
-ARG max_jobs
-ENV MAX_JOBS=${max_jobs}
 
 # Install sccache if USE_SCCACHE is enabled (for release builds)
 ARG USE_SCCACHE
@@ -126,7 +114,8 @@ FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm
 # -----------------------
 # Rust build stage
 # Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages
-# don't need the rust toolchain or protoc.
+# don't need the rust toolchain or protoc. Runs in parallel with the main wheel
+# build for faster end-to-end builds.
 FROM fetch_vllm AS rust-build
 ARG COMMON_WORKDIR
 
@@ -155,74 +144,24 @@ ENV RUSTUP_MAX_RETRIES=10
 # layer for later COPY --from=rust-build.
 RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \
     --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \
-    --mount=type=cache,id=vllm-rocm-cargo-target,target=${COMMON_WORKDIR}/vllm/rust/target,sharing=locked \
     cd ${COMMON_WORKDIR}/vllm \
     && VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh \
     && test -x /tmp/vllm-rs
 
 # -----------------------
-# vLLM native build stages
-#
-# csrc-build intentionally copies only files that affect ROCm native extension
-# compilation. That keeps unrelated CI/test/docs edits from invalidating the
-# expensive HIP/C++ build layer.
-FROM base AS csrc-build
-ARG COMMON_WORKDIR
-WORKDIR ${COMMON_WORKDIR}/vllm
-
-COPY requirements/rocm.txt requirements/rocm.txt
-COPY requirements/common.txt requirements/common.txt
-RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
-    uv pip install --system -r requirements/rocm.txt
-
-# pyproject.toml is bind-mounted in the RUN step so metadata-only changes do
-# not invalidate the expensive native build layer.
-COPY setup.py CMakeLists.txt ./
-COPY cmake cmake/
-COPY csrc csrc/
-COPY vllm/envs.py vllm/envs.py
-COPY vllm/__init__.py vllm/__init__.py
-
-ENV VLLM_TARGET_DEVICE=rocm
-ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+rocm.csrc.build"
-
-RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproject.toml \
-    --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
-    export CCACHE_BASEDIR="$PWD" \
-    && echo "=== ccache stats before ROCm native build ===" \
-    && (ccache --show-stats || true) \
-    && (ccache --zero-stats || true) \
-    && EFFECTIVE_MAX_JOBS="${MAX_JOBS:-$(nproc)}" \
-    && echo "Building ROCm native extension wheel with MAX_JOBS=${EFFECTIVE_MAX_JOBS}" \
-    && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${EFFECTIVE_MAX_JOBS}" python3 setup.py bdist_wheel --dist-dir=dist \
-    && test -d dist \
-    && ls dist/*.whl >/dev/null \
-    && echo "=== ccache stats after ROCm native build ===" \
-    && (ccache --show-stats || true)
-
-# Build the full vLLM ROCm wheel by reusing the native extension wheel from
-# csrc-build. This stage still rebuilds for Python/package changes, but skips
-# the expensive HIP/C++ compile when native inputs are unchanged.
+# vLLM build stages
 FROM fetch_vllm AS build_vllm
 ARG COMMON_WORKDIR
-ENV VLLM_TARGET_DEVICE=rocm
-
-COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels
 
 # Drop the pre-built rust frontend binary into the source tree. setup.py
 # detects it and ships it as-is, skipping the local cargo build.
 COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
 
-RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
-    cd vllm \
-    && uv pip install --system -r requirements/rocm.txt \
-    && export VLLM_USE_PRECOMPILED=1 \
-    && export VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled-wheels/*.whl)" \
-    && export VLLM_DOCKER_BUILD_CONTEXT=1 \
-    && echo "Packaging vLLM ROCm wheel using precompiled extensions from ${VLLM_PRECOMPILED_WHEEL_LOCATION}" \
-    && python3 setup.py bdist_wheel --dist-dir=dist \
-    && test -d dist \
-    && ls dist/*.whl >/dev/null
+# Build vLLM (setup.py auto-detects sccache in PATH)
+RUN cd vllm \
+    && python3 -m pip install -r requirements/rocm.txt \
+    && python3 setup.py clean --all  \
+    && python3 setup.py bdist_wheel --dist-dir=dist
 FROM scratch AS export_vllm
 ARG COMMON_WORKDIR
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/dist/*.whl /
@@ -230,10 +169,8 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/requirements /requirements
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
 
 # RIXL/UCX build stages
@@ -264,17 +201,14 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
     ibverbs-providers \
     && rm -rf /var/lib/apt/lists/*
 
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system meson auditwheel patchelf tomlkit
+RUN uv pip install --system meson auditwheel patchelf tomlkit
 
-RUN --mount=type=cache,target=/root/.cache/ccache \
-    cd /usr/local/src && \
+RUN cd /usr/local/src && \
     git clone ${UCX_REPO} &&  \
     cd ucx  && \
     git checkout ${UCX_BRANCH} && \
     ./autogen.sh && \
     mkdir build && cd build && \
-    CC="ccache gcc" CXX="ccache g++" \
     ../configure \
         --prefix=/usr/local/ucx \
         --enable-shared \
@@ -286,22 +220,20 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
         --with-verbs \
         --with-dm \
         --enable-mt && \
-    make -j$(nproc) && \
+    make -j && \
     make install
 
 ENV PATH=/usr/local/ucx/bin:$PATH
 ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
 
-RUN --mount=type=cache,target=/root/.cache/ccache \
-    git clone ${RIXL_REPO} /opt/rixl && \
+RUN git clone ${RIXL_REPO} /opt/rixl && \
     cd /opt/rixl && \
     git checkout ${RIXL_BRANCH} && \
-    CC="ccache gcc" CXX="ccache g++" \
     meson setup build --prefix=${RIXL_HOME} \
                      -Ducx_path=${UCX_HOME} \
                      -Drocm_path=${ROCM_PATH} && \
     cd build && \
-    ninja -j$(nproc) && \
+    ninja && \
     ninja install
 
 # Generate RIXL wheel
@@ -318,44 +250,30 @@ RUN cd /opt/rixl && \
         --ucx-plugins-dir ${UCX_HOME}/lib/ucx \
         --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
 
-# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not
-# invalidate the slow ROCShmem build.
-FROM base AS build_rocshmem
+# DeepEP build stage
+FROM base AS build_deep
 ARG ROCSHMEM_BRANCH="f0acb0c6"
 ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
-# DeepEP only supports gfx942 and gfx950; build ROCShmem for the same set so
-# it can be linked against DeepEP without arch mismatches.
-ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
-ENV ROCM_PATH=/opt/rocm
-ENV ROCSHMEM_DIR=/opt/rocshmem
-
-RUN --mount=type=cache,target=/root/.cache/ccache \
-    git clone --no-checkout --filter=blob:none ${ROCSHMEM_REPO} \
- && cd rocm-systems \
- && git sparse-checkout set --cone projects/rocshmem \
- && git checkout ${ROCSHMEM_BRANCH} \
- && mkdir -p projects/rocshmem/build \
- && cd projects/rocshmem/build \
- && CC="ccache gcc" CXX="ccache g++" INSTALL_PREFIX=${ROCSHMEM_DIR} \
-    bash ../scripts/build_configs/all_backends \
-      -DROCM_PATH=${ROCM_PATH} \
-      -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \
-      -DUSE_EXTERNAL_MPI=OFF
-
-# DeepEP build stage - depends on ROCShmem, builds the HIP kernel wheel.
-FROM build_rocshmem AS build_deepep
 ARG DEEPEP_BRANCH="a9ea9774"
 ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
 ARG DEEPEP_NIC="cx7"
+ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
+ENV ROCSHMEM_DIR=/opt/rocshmem
 
-# Build DeepEP wheel. DeepEP looks for rocshmem at ROCSHMEM_DIR.
-# DeepEP only supports gfx942 and gfx950, so avoid gfx90a in the default list.
-RUN --mount=type=cache,target=/root/.cache/ccache \
-    export PYTORCH_ROCM_ARCH="gfx942;gfx950" \
- && git clone ${DEEPEP_REPO} \
+RUN git clone ${ROCSHMEM_REPO} \
+ && cd rocm-systems \
+ && git checkout ${ROCSHMEM_BRANCH} \
+ && mkdir -p projects/rocshmem/build \
+ && cd projects/rocshmem/build \
+ && INSTALL_PREFIX=${ROCSHMEM_DIR} \
+    ../scripts/build_configs/all_backends -DUSE_EXTERNAL_MPI=OFF
+
+# Build DeepEP wheel.
+# DeepEP looks for rocshmem at ROCSHMEM_DIR.
+RUN git clone ${DEEPEP_REPO} \
  && cd DeepEP \
  && git checkout ${DEEPEP_BRANCH} \
- && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
+ && python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
 
 # MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do
 # not force users to rebuild the long-lived Dockerfile.rocm_base image.
@@ -454,9 +372,8 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \
 # Extract version from git BEFORE any modifications (pin_rocm_dependencies.py modifies requirements/rocm.txt)
 # This ensures setuptools_scm sees clean repo state for version detection
 RUN --mount=type=bind,source=.git,target=vllm/.git \
-    --mount=type=cache,target=/root/.cache/uv \
     cd vllm \
-    && uv pip install --system setuptools_scm regex \
+    && pip install setuptools_scm regex \
     && VLLM_VERSION=$(python3 -c "import setuptools_scm; print(setuptools_scm.get_version())") \
     && echo "Detected vLLM version: ${VLLM_VERSION}" \
     && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt
@@ -492,20 +409,18 @@ RUN echo "Pinning vLLM dependencies to custom wheel versions..." \
     && python3 /tmp/pin_rocm_dependencies.py /install ${COMMON_WORKDIR}/vllm/requirements/rocm.txt
 
 # Install dependencies using custom wheels from /install
-RUN --mount=type=cache,target=/root/.cache/uv \
-    cd vllm \
+RUN cd vllm \
     && echo "Building vLLM with custom wheels from /install" \
-    && uv pip install --system --find-links /install -r requirements/rocm.txt
+    && python3 -m pip install --find-links /install -r requirements/rocm.txt \
+    && python3 setup.py clean --all
 
 # Build wheel using pre-extracted version to avoid dirty state from modified requirements/rocm.txt
-# (setup.py auto-detects ccache/sccache in PATH)
+# (setup.py auto-detects sccache in PATH)
 RUN --mount=type=bind,source=.git,target=vllm/.git \
-    --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
     cd vllm \
-    && export CCACHE_BASEDIR="$PWD" \
     && export SETUPTOOLS_SCM_PRETEND_VERSION=$(cat /tmp/vllm_version.txt) \
     && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \
-    && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist
+    && python3 setup.py bdist_wheel --dist-dir=dist
 
 FROM scratch AS export_vllm_wheel_release
 ARG COMMON_WORKDIR
@@ -514,122 +429,115 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/requirements /requir
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
 
 # -----------------------
-# CI base image (Tier 1) - stable, rarely changing CI dependencies.
-# Per-PR test builds pull this as CI_BASE_IMAGE so the test stage only layers
-# in the vLLM artifacts for the current commit.
-FROM mori_base AS ci_base
-ARG COMMON_WORKDIR
+# Test vLLM image
+FROM mori_base AS test
 
-# Update rdma-core to support latest rocshmem.
+RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
+
+# Install vLLM using uv (inherited from base stage)
+# Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version
+RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
+    --mount=type=cache,target=/root/.cache/uv \
+    cd /install \
+    && uv pip install --system -r requirements/rocm.txt \
+    && uv pip install --system -r requirements/test/rocm.txt \
+    && pip uninstall -y vllm \
+    && uv pip install --system *.whl
+
+# Persist the built wheel in the image so python_only_compile_rocm.sh can
+# reinstall it after removing compilers. The bind-mounted /install contents
+# above are not available once that RUN step completes.
+COPY --from=export_vllm /*.whl /opt/vllm-wheels/
+
+# Update rdma-core to support latest rocshmem
 ARG DEEPEP_NIC
 RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
     git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \
     cd /tmp/rdma-core && \
     mkdir -p build && cd build && \
     cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \
-    ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
+    ninja &&     ninja install &&     ldconfig &&     rm -rf /tmp/rdma-core; \
 fi
 
-# Install RIXL + DeepEP wheels.
+# Install RIXL wheel
 RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
-    --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
-    uv pip install --system /rixl_install/*.whl /deep_install/*.whl
+    uv pip install --system /rixl_install/*.whl
 
-# Copy ROCShmem runtime libraries.
-COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem
+# Install DeepEP wheel
+RUN --mount=type=bind,from=build_deep,src=/app/deep_install,target=/deep_install \
+    uv pip install --system /deep_install/*.whl
+COPY --from=build_deep /opt/rocshmem /opt/rocshmem
 
-# RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec.
-RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
+# RIXL/MoRIIO runtime dependencies (RDMA userspace libraries)
+RUN apt-get update -q -y && apt-get install -q -y \
     librdmacm1 \
     libibverbs1 \
     ibverbs-providers \
     ibverbs-utils \
-    pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \
-    libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \
     && rm -rf /var/lib/apt/lists/*
 
-# Install torchcodec from source for ROCm/torch ABI compatibility.
-COPY tools/install_torchcodec_rocm.sh /tmp/install_torchcodec.sh
-RUN --mount=type=cache,target=/root/.cache/uv \
-    --mount=type=cache,target=/root/.cache/pip \
-    --mount=type=cache,target=/root/.cache/torchcodec-wheels \
-    bash /tmp/install_torchcodec.sh \
-    && rm /tmp/install_torchcodec.sh \
-    && apt-get clean && rm -rf /var/lib/apt/lists/*
+WORKDIR /vllm-workspace
+ARG COMMON_WORKDIR
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace
 
-# Pre-install shared ROCm runtime dependencies.
-COPY requirements/common.txt requirements/rocm.txt /tmp/ci-base-requirements/
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system -r /tmp/ci-base-requirements/rocm.txt \
-    && rm -rf /tmp/ci-base-requirements
+# install development dependencies (for testing)
+RUN cd /vllm-workspace \
+    && python3 -m pip install -e tests/vllm_test_utils \
+    && python3 -m pip install pytest-shard
 
-# Enable fast and less brittle model downloads in tests.
+# enable fast downloads from hf (for testing)
 ENV HF_XET_HIGH_PERFORMANCE=1
-ENV HF_HUB_DOWNLOAD_TIMEOUT=60
 
-# Pre-install vLLM test dependencies.
-COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system -r /tmp/rocm-test-reqs.txt
+# increase timeout for hf downloads (for testing)
+ENV HF_HUB_DOWNLOAD_TIMEOUT 60
 
-# Rebuild fastsafetensors from source so its C++ extension is compiled with
-# USE_ROCM and can detect libamdhip64.so at runtime.
-RUN --mount=type=cache,target=/root/.cache/pip \
-    FASTSAFETENSORS_REQ="$(grep -E '^fastsafetensors(==| @ )' /tmp/rocm-test-reqs.txt | head -1)" \
-    && test -n "${FASTSAFETENSORS_REQ}" \
-    && python3 -m pip install --force-reinstall --no-deps \
-        --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \
-    && rm /tmp/rocm-test-reqs.txt
+# install audio decode package `torchcodec` from source (required due to 
+# ROCm and torch version mismatch) for tests with datasets package
+COPY tools/install_torchcodec_rocm.sh /tmp/install_torchcodec.sh
+RUN bash /tmp/install_torchcodec.sh \
+    && rm /tmp/install_torchcodec.sh \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
 
-# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel.
+# Copy in the v1 package (for python-only install test group)
+COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
+
+# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel
 # See: https://github.com/pytorch/pytorch/issues/169857
 ENV MIOPEN_DEBUG_CONV_DIRECT=0
 ENV MIOPEN_DEBUG_CONV_GEMM=0
 
-# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc.
+# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
 # See: https://github.com/ROCm/rocm-libraries/issues/6266
 ENV HSA_ENABLE_IPC_MODE_LEGACY=1
 
-# ROCm profiler limits workaround.
-RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
-ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
-
-# Install vllm_test_utils in ci_base for ci_base + wheel parity.
-COPY tests/vllm_test_utils /tmp/vllm_test_utils
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system /tmp/vllm_test_utils \
-    && rm -rf /tmp/vllm_test_utils
-
-# -----------------------
-# Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base.
-FROM ${CI_BASE_IMAGE} AS test
-ARG COMMON_WORKDIR
-
-# Install the vLLM wheel (--no-deps: all deps already in ci_base).
-RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
-    --mount=type=cache,target=/root/.cache/uv \
-    cd /install \
-    && uv pip install --system --no-deps *.whl
-
-# Store the vLLM wheel in the image for python-only install tests.
-COPY --from=export_vllm /*.whl /opt/vllm-wheels/
-
-WORKDIR /vllm-workspace
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace
-
-# Copy in the v1 package (for python-only install test group).
-COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
-
-# Hide source under src/ so it won't shadow the installed package in tests.
+# Source code is used in the `python_only_compile.sh` test
+# We hide it inside `src/` so that this source code
+# will not be imported by other tests
 RUN mkdir src && mv vllm src/vllm
 
+# This is a workaround to ensure pytest exits with the correct status code in CI tests.
+RUN printf '%s\n' \
+    'import os' \
+    '' \
+    '_exit_code = 1' \
+    '' \
+    'def pytest_sessionfinish(session, exitstatus):' \
+    '    global _exit_code' \
+    '    _exit_code = int(exitstatus)' \
+    '' \
+    'def pytest_unconfigure(config):' \
+    '    import sys' \
+    '    sys.stdout.flush()' \
+    '    sys.stderr.flush()' \
+    '    os._exit(_exit_code)' \
+    > /vllm-workspace/conftest.py
+
 # -----------------------
 # Final vLLM image
 FROM mori_base AS final
@@ -645,7 +553,6 @@ RUN rm -f /usr/bin/sccache || true \
 # This prevents S3 bucket config from leaking into production images
 ENV SCCACHE_BUCKET=
 ENV SCCACHE_REGION=
-ENV SCCACHE_ENDPOINT=
 ENV SCCACHE_S3_NO_CREDENTIALS=
 ENV SCCACHE_IDLE_TIMEOUT=
 

docker/ci-rocm.hcl

@@ -1,376 +0,0 @@
-# ci-rocm.hcl - CI-specific configuration for vLLM ROCm Docker builds
-#
-# This file lives in the vLLM repo at docker/ci-rocm.hcl so ROCm Docker
-# build mechanics can evolve with Dockerfile.rocm and docker-bake-rocm.hcl.
-# Used with: docker buildx bake -f docker/docker-bake-rocm.hcl -f docker/ci-rocm.hcl test-rocm-ci
-#
-# Registry cache: Docker Hub (rocm/vllm-ci-cache) is used exclusively.
-# AMD build agents already have Docker Hub credentials (they push the test
-# image to rocm/vllm-ci), so no additional credential setup is required.
-# ROCm CI uses Docker Hub for BuildKit layer cache by default. A separate
-# compiler cache can be enabled with USE_SCCACHE=1 when AMD provides a shared
-# S3-compatible cache endpoint.
-
-# CI metadata
-
-variable "BUILDKITE_COMMIT" {
-  default = ""
-}
-
-variable "BUILDKITE_BUILD_NUMBER" {
-  default = ""
-}
-
-variable "BUILDKITE_BUILD_ID" {
-  default = ""
-}
-
-variable "PARENT_COMMIT" {
-  default = ""
-}
-
-# Merge-base of HEAD with main - provides a more stable cache fallback than
-# parent commit for long-lived PRs. Mirrors the VLLM_MERGE_BASE_COMMIT
-# pattern used in the shared ci.hcl file. Auto-computed by ci-bake-rocm.sh
-# when unset.
-variable "VLLM_MERGE_BASE_COMMIT" {
-  default = ""
-}
-
-# Bridge to vLLM's COMMIT variable for OCI labels
-variable "COMMIT" {
-  default = BUILDKITE_COMMIT
-}
-
-# Image tags (set by CI)
-
-variable "IMAGE_TAG" {
-  default = ""
-}
-
-variable "IMAGE_TAG_LATEST" {
-  default = ""
-}
-
-# ROCm-specific GPU architecture targets
-
-variable "PYTORCH_ROCM_ARCH" {
-  default = "gfx90a;gfx942;gfx950"
-}
-
-# Pre-built CI base image (Tier 1). Per-PR builds pull this instead of
-# rebuilding RIXL/DeepEP/torchcodec from scratch. The ci_base stage in
-# Dockerfile.rocm inherits from base, so CI_BASE_IMAGE only affects the test
-# stage and is irrelevant when building --target ci_base itself.
-variable "CI_BASE_IMAGE" {
-  default = "rocm/vllm-dev:ci_base"
-}
-
-# Leave CI_MAX_JOBS empty so the Dockerfile falls back to $(nproc) and uses
-# the full builder parallelism. Operators can still override this per build.
-variable "CI_MAX_JOBS" {
-  default = ""
-}
-
-# Upstream dependency commit pins -- extracted from Dockerfile.rocm by
-# ci-bake-rocm.sh at build time. Empty defaults are safe: the cache
-# functions produce no entries when the variable is empty.
-variable "RIXL_BRANCH" {
-  default = ""
-}
-
-variable "UCX_BRANCH" {
-  default = ""
-}
-
-variable "ROCSHMEM_BRANCH" {
-  default = ""
-}
-
-variable "DEEPEP_BRANCH" {
-  default = ""
-}
-
-variable "RIXL_CACHE_KEY" {
-  default = ""
-}
-
-variable "ROCSHMEM_CACHE_KEY" {
-  default = ""
-}
-
-variable "DEEPEP_CACHE_KEY" {
-  default = ""
-}
-
-# Docker Hub registry cache for AMD builds.
-#
-# A separate repo (rocm/vllm-ci-cache) is used for BuildKit layer cache.
-# Final-image cache exports use mode=min to reduce the volume of data pushed.
-# Source-scoped csrc cache exports default to mode=max so fresh workers can
-# recover more of the native build graph when ROCm extension inputs change.
-# NOTE: mode=min still includes all layers referenced by the final image
-# manifest, including inherited base layers (~7.25GB ROCm runtime).
-# Docker Hub auto-creates the repo on first push.
-#
-# Final-image cache stays commit-scoped. Branch-to-branch reuse for the test
-# image comes from importing the parent and merge-base commit cache refs.
-#
-# The source-scoped native cache is exported both per-commit and per-branch so
-# ROCm extension rebuilds are shareable within the same commit reruns and across
-# consecutive commits on the same branch without depending on a single global
-# latest tag.
-
-variable "DOCKERHUB_CACHE_REPO" {
-  default = "rocm/vllm-ci-cache"
-}
-
-variable "DOCKERHUB_CACHE_TO" {
-  default = ""
-}
-
-variable "ROCM_CACHE_BRANCH_TAG" {
-  default = ""
-}
-
-variable "ROCM_CACHE_UPSTREAM_BRANCH_TAG" {
-  default = ""
-}
-
-variable "ROCM_CSRC_CACHE_TO_MODE" {
-  default = "max"
-}
-
-variable "ROCM_FINAL_CACHE_TO_MODE" {
-  default = "min"
-}
-
-# Functions
-
-function "get_cache_from_rocm" {
-  params = []
-  result = compact([
-    # Exact commit hit - fastest cache on re-runs of the same commit
-    BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${BUILDKITE_COMMIT}" : "",
-    # Parent commit - useful cache for incremental changes
-    PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${PARENT_COMMIT}" : "",
-    # Merge-base with main - stable fallback for long-lived or rebased PRs;
-    # maps to a real main-branch commit whose cache layers are likely warm
-    VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${VLLM_MERGE_BASE_COMMIT}" : "",
-    # Import the source-scoped native build cache as well so builds whose
-    # Python/package layers changed can still reuse compiled ROCm objects.
-    BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${BUILDKITE_COMMIT}" : "",
-    PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${PARENT_COMMIT}" : "",
-    VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${VLLM_MERGE_BASE_COMMIT}" : "",
-    ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "",
-    ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "",
-    # Branch-scoped full image cache - fallback when parent-commit cache is evicted
-    ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "",
-    ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "",
-  ])
-}
-
-function "get_cache_to_rocm" {
-  params = []
-  result = compact([
-    # Commit-scoped cache for exact re-runs.
-    BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${BUILDKITE_COMMIT},mode=${ROCM_FINAL_CACHE_TO_MODE}" : "",
-    # Branch-scoped cache so later commits on the same branch can reuse the full
-    # image layers when the parent-commit cache is evicted. Unlike the old
-    # rocm-latest tag (which caused duplicate exporter 400s), this is per-branch.
-    ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${ROCM_FINAL_CACHE_TO_MODE}" : "",
-  ])
-}
-
-function "get_cache_from_rocm_csrc" {
-  params = []
-  result = compact([
-    BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${BUILDKITE_COMMIT}" : "",
-    PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${PARENT_COMMIT}" : "",
-    VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${VLLM_MERGE_BASE_COMMIT}" : "",
-    ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "",
-    ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "",
-  ])
-}
-
-function "get_cache_to_rocm_csrc" {
-  params = []
-  result = compact([
-    # Export the exact-commit native cache for same-commit reruns.
-    BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${BUILDKITE_COMMIT},mode=${ROCM_CSRC_CACHE_TO_MODE}" : "",
-    # Export the branch-scoped native cache so later commits on the same branch
-    # can reuse compiled ROCm objects even when the exact parent cache is absent.
-    ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${ROCM_CSRC_CACHE_TO_MODE}" : "",
-  ])
-}
-
-# Cache functions for upstream dependency stages (RIXL/UCX, ROCShmem, DeepEP).
-# These stages are pinned to specific upstream commit hashes, so cache keys use
-# those hashes rather than the Buildkite commit. This means the cache persists
-# across all vLLM commits as long as the upstream dependency pins don't change.
-
-function "get_cache_from_rocm_deps" {
-  params = []
-  result = compact([
-    RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY}" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
-    ROCSHMEM_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY}" : (ROCSHMEM_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_BRANCH}" : ""),
-    DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY}" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH}" : ""),
-  ])
-}
-
-function "get_cache_to_rocm_rixl" {
-  params = []
-  result = compact([
-    RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY},mode=min" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
-  ])
-}
-
-function "get_cache_to_rocm_rocshmem" {
-  params = []
-  result = compact([
-    ROCSHMEM_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY},mode=min" : (ROCSHMEM_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_BRANCH},mode=min" : ""),
-  ])
-}
-
-function "get_cache_to_rocm_deepep" {
-  params = []
-  result = compact([
-    DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY},mode=min" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH},mode=min" : ""),
-  ])
-}
-
-# CI targets
-
-target "_ci-rocm" {
-  annotations = [
-    "manifest:vllm.buildkite.build_number=${BUILDKITE_BUILD_NUMBER}",
-    "manifest:vllm.buildkite.build_id=${BUILDKITE_BUILD_ID}",
-  ]
-  args = {
-    ARG_PYTORCH_ROCM_ARCH = PYTORCH_ROCM_ARCH
-    CI_BASE_IMAGE         = CI_BASE_IMAGE
-    max_jobs              = CI_MAX_JOBS
-  }
-}
-
-target "test-rocm-ci" {
-  inherits   = ["_common-rocm", "_ci-rocm", "_labels"]
-  target     = "test"
-  cache-from = get_cache_from_rocm()
-  cache-to   = get_cache_to_rocm()
-  tags = compact([
-    IMAGE_TAG,
-    IMAGE_TAG_LATEST,
-  ])
-  output = ["type=registry"]
-}
-
-# Cache-only target for the source-scoped ROCm native build stage.
-# This persists the csrc-build stage in the registry cache even though the
-# final test image only consumes it indirectly while packaging the wheel.
-target "csrc-rocm-ci" {
-  inherits   = ["_common-rocm", "_ci-rocm"]
-  target     = "csrc-build"
-  cache-from = get_cache_from_rocm_csrc()
-  cache-to   = get_cache_to_rocm_csrc()
-  output     = ["type=cacheonly"]
-}
-
-# Keep wheel export on the same CI graph as the test image build so the
-# shared build_vllm/export_vllm stages resolve identically within one bake
-# invocation. Without this, export-wheel-rocm uses the plain local target
-# args while test-rocm-ci uses CI-only args, which can lead to separate
-# cache lineages and inconsistent export_vllm results.
-target "export-wheel-rocm" {
-  inherits   = ["_common-rocm", "_ci-rocm"]
-  target     = "export_vllm"
-  cache-from = get_cache_from_rocm()
-  cache-to   = get_cache_to_rocm()
-  output     = ["type=local,dest=./wheel-export"]
-}
-
-# Artifact-only vLLM build. GPU test jobs consume this artifact on top of
-# ci_base, avoiding a per-commit multi-GB image push/pull.
-group "test-rocm-ci-with-artifacts" {
-  targets = ["csrc-rocm-ci", "export-wheel-rocm"]
-}
-
-# Full test image + wheel export. Kept for fallback/debugging when a pushed
-# per-commit image is useful.
-group "test-rocm-ci-with-wheel" {
-  targets = ["csrc-rocm-ci", "test-rocm-ci", "export-wheel-rocm"]
-}
-
-# Image tags for the ci_base build. ci-bake-rocm.sh rewrites CI_BASE_IMAGE_TAG
-# to the primary tag for this build. Non-nightly builds use a commit-scoped tag
-# and also publish a content tag for reuse. NIGHTLY=1 builds on the stable branch
-# can additionally set CI_BASE_IMAGE_TAG_STABLE to refresh rocm/vllm-dev:ci_base.
-variable "CI_BASE_IMAGE_TAG" {
-  default = "rocm/vllm-dev:ci_base"
-}
-
-variable "CI_BASE_IMAGE_TAG_CONTENT" {
-  default = ""
-}
-
-variable "CI_BASE_IMAGE_TAG_STABLE" {
-  default = ""
-}
-
-# Cache-only targets for upstream dependency stages. These persist each stage
-# in the registry cache keyed by its upstream commit hash. When ci_base rebuilds
-# (e.g., requirements change), these stages are cache hits if their upstream
-# pins haven't changed -- saving ~35min of compilation.
-target "rixl-rocm-ci" {
-  inherits   = ["_common-rocm", "_ci-rocm"]
-  target     = "build_rixl"
-  cache-from = get_cache_from_rocm_deps()
-  cache-to   = get_cache_to_rocm_rixl()
-  output     = ["type=cacheonly"]
-}
-
-target "rocshmem-rocm-ci" {
-  inherits   = ["_common-rocm", "_ci-rocm"]
-  target     = "build_rocshmem"
-  cache-from = get_cache_from_rocm_deps()
-  cache-to   = get_cache_to_rocm_rocshmem()
-  output     = ["type=cacheonly"]
-}
-
-target "deepep-rocm-ci" {
-  inherits   = ["_common-rocm", "_ci-rocm"]
-  target     = "build_deepep"
-  cache-from = get_cache_from_rocm_deps()
-  cache-to   = get_cache_to_rocm_deepep()
-  output     = ["type=cacheonly"]
-}
-
-# Builds only the ci_base stage (RIXL, DeepEP, torchcodec, etc.)
-# Invoked by the ensure-ci-base step when the content hash of ci_base-affecting
-# files drifts from the remote image label. Per-PR builds then pull the result
-# as CI_BASE_IMAGE instead of rebuilding those slow layers on every commit.
-# Uses inline cache metadata on the ci_base image itself instead of exporting a
-# separate registry cache artifact.
-target "ci-base-rocm-ci" {
-  inherits   = ["_common-rocm", "_ci-rocm", "_labels"]
-  target     = "ci_base"
-  cache-from = concat(
-    compact([
-      CI_BASE_IMAGE_TAG != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG}" : "",
-      CI_BASE_IMAGE_TAG_CONTENT != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_CONTENT}" : "",
-      CI_BASE_IMAGE_TAG_STABLE != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_STABLE}" : "",
-    ]),
-    # Import upstream dependency caches so RIXL/ROCShmem/DeepEP stages
-    # are cache hits even when ci_base itself needs rebuilding.
-    get_cache_from_rocm_deps(),
-  )
-  cache-to = ["type=inline"]
-  tags     = compact([CI_BASE_IMAGE_TAG, CI_BASE_IMAGE_TAG_CONTENT, CI_BASE_IMAGE_TAG_STABLE])
-  output   = ["type=registry"]
-}
-
-# Group for ci_base builds -- exports dependency stage caches alongside the
-# ci_base image so future rebuilds can reuse them independently.
-group "ci-base-rocm-ci-with-deps" {
-  targets = ["rixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
-}

docker/docker-bake-rocm.hcl

@@ -1,143 +0,0 @@
-# docker-bake-rocm.hcl - vLLM ROCm Docker build configuration
-#
-# This file lives in the vLLM repo at docker/docker-bake-rocm.hcl
-# Equivalent of docker-bake.hcl for ROCm builds.
-#
-# Usage:
-#   docker buildx bake -f docker/docker-bake-rocm.hcl              # Build test (default)
-#   docker buildx bake -f docker/docker-bake-rocm.hcl final-rocm   # Build final image
-#   docker buildx bake -f docker/docker-bake-rocm.hcl --print      # Show resolved config
-#
-# CI usage (with the vLLM-owned CI overlay):
-#   docker buildx bake -f docker/docker-bake-rocm.hcl -f docker/ci-rocm.hcl test-rocm-ci
-
-variable "MAX_JOBS" {
-  # Empty string lets the Dockerfile fall back to $(nproc) via
-  # MAX_JOBS="${MAX_JOBS:-$(nproc)}" in each RUN step, which uses all
-  # available cores on whatever machine the build runs on.
-  # Override with --set '*.args.max_jobs=8' for local builds on small machines.
-  default = ""
-}
-
-variable "PYTORCH_ROCM_ARCH" {
-  default = "gfx90a;gfx942;gfx950"
-}
-
-variable "COMMIT" {
-  default = ""
-}
-
-# Content hash of ci_base-affecting files. Computed by ci-bake-rocm.sh and
-# embedded as a label so future builds can compare without rebuilding.
-variable "CI_BASE_CONTENT_HASH" {
-  default = ""
-}
-
-# REMOTE_VLLM=0: use local source via Docker build context (ONBUILD COPY ./ vllm/)
-# REMOTE_VLLM=1: clone from GitHub at VLLM_BRANCH (standalone builds without local source)
-variable "REMOTE_VLLM" {
-  default = "0"
-}
-
-variable "VLLM_BRANCH" {
-  default = "main"
-}
-
-# CI_BASE_IMAGE: pre-built ci_base image for per-PR test builds.
-# Defaults to the local "ci_base" stage for standalone/local builds.
-# CI overrides this to "rocm/vllm-dev:ci_base" via environment variable.
-variable "CI_BASE_IMAGE" {
-  default = "rocm/vllm-dev:ci_base"
-}
-
-# Upstream dependency commit pins. Plain local bake builds use the Dockerfile
-# ARG defaults. ci-bake-rocm.sh resolves those defaults (plus any env
-# overrides) and writes a small HCL override before invoking CI targets.
-variable "RIXL_BRANCH" {
-  default = ""
-}
-
-variable "UCX_BRANCH" {
-  default = ""
-}
-
-variable "ROCSHMEM_BRANCH" {
-  default = ""
-}
-
-variable "DEEPEP_BRANCH" {
-  default = ""
-}
-
-group "default" {
-  targets = ["test-rocm"]
-}
-
-target "_common-rocm" {
-  dockerfile = "docker/Dockerfile.rocm"
-  context    = "."
-  args = {
-    max_jobs                        = MAX_JOBS
-    ARG_PYTORCH_ROCM_ARCH           = PYTORCH_ROCM_ARCH
-    REMOTE_VLLM                     = REMOTE_VLLM
-    VLLM_BRANCH                     = VLLM_BRANCH
-    CI_BASE_IMAGE                   = CI_BASE_IMAGE
-  }
-}
-
-target "_labels" {
-  labels = {
-    "org.opencontainers.image.source"      = "https://github.com/vllm-project/vllm"
-    "org.opencontainers.image.vendor"      = "vLLM"
-    "org.opencontainers.image.title"       = "vLLM ROCm"
-    "org.opencontainers.image.description" = "vLLM: A high-throughput and memory-efficient inference and serving engine for LLMs (ROCm)"
-    "org.opencontainers.image.licenses"    = "Apache-2.0"
-    "org.opencontainers.image.revision"    = COMMIT
-  }
-  annotations = [
-    "manifest:org.opencontainers.image.revision=${COMMIT}",
-  ]
-}
-
-target "test-rocm" {
-  inherits = ["_common-rocm", "_labels"]
-  target   = "test"
-  tags     = ["rocm/vllm:test"]
-  output   = ["type=docker"]
-}
-
-# CI base image target - builds only the ci_base stage (RIXL, DeepEP,
-# torchcodec, requirements, etc.). Used by the weekly scheduled build and
-# the auto-rebuild trigger when requirements change in a PR.
-target "ci-base-rocm" {
-  inherits = ["_common-rocm", "_labels"]
-  target   = "ci_base"
-  labels   = {
-    "vllm.ci_base.content_hash" = CI_BASE_CONTENT_HASH
-  }
-  tags     = ["rocm/vllm-dev:ci_base"]
-  output   = ["type=docker"]
-}
-
-# Wheel export target - extracts the built vLLM wheel + test workspace
-# to local disk. Used by CI to upload the wheel as a Buildkite artifact
-# so test jobs can assemble images locally from ci_base + wheel instead
-# of pulling the full large image from Docker Hub.
-#
-# Usage:
-#   docker buildx bake -f docker/docker-bake-rocm.hcl export-wheel-rocm
-#   # Creates ./wheel-export/*.whl, ./wheel-export/requirements/, etc.
-#
-# After a full bake build, BuildKit cache makes this nearly instant.
-target "export-wheel-rocm" {
-  inherits = ["_common-rocm"]
-  target   = "export_vllm"
-  output   = ["type=local,dest=./wheel-export"]
-}
-
-target "final-rocm" {
-  inherits = ["_common-rocm", "_labels"]
-  target   = "final"
-  tags     = ["rocm/vllm:latest"]
-  output   = ["type=docker"]
-}

docker/versions.json

@@ -68,7 +68,7 @@
       "default": "true"
     },
     "FLASHINFER_VERSION": {
-      "default": "0.6.12"
+      "default": "0.6.11.post2"
     },
     "GDRCOPY_CUDA_VERSION": {
       "default": "12.8"

docs/configuration/optimization.md

@@ -296,7 +296,7 @@ llm = LLM(model="Qwen/Qwen3-8B")
 The `fastokens` Python package (>= 0.2.0) must be installed; if it isn't,
 vLLM raises a clear `ImportError` at tokenizer load. The override applies to
 any `--tokenizer-mode` that ends up loading an HF fast tokenizer (`hf`,
-`deepseek_v32`, `deepseek_v4`, `qwen_vl`, …). Models that don't use the HF
+`deepseek_v32`, `deepseek_v4`, `qwen_vl`, …). Modes that don't use the HF
 fast tokenizer (`mistral`, `grok2`, `kimi_audio`) ignore the flag.
 
 Tokenizer-bound workloads — long shared prefixes, bursty short prompts,

docs/contributing/README.md

@@ -101,7 +101,7 @@ vLLM's `pre-commit` hooks will now run automatically every time you commit.
     Some `pre-commit` hooks only run in CI. If you need to, you can run them locally with:
 
     ```bash
-    pre-commit run --hook-stage manual mypy-3.11
+    pre-commit run --hook-stage manual mypy-3.10
     ```
 
 ### Documentation

docs/contributing/profiling.md

@@ -35,7 +35,8 @@ Traces can be visualized using <https://ui.perfetto.dev/>.
 
 !!! tip
     To stop the profiler - it flushes out all the profile trace files to the directory. This takes time, for example for about 100 requests worth of data for a llama 70b, it takes about 10 minutes to flush out on a H100.
-    The engine client waits for this flush to complete without timing out, so simply allow the stop call to run to completion.
+    Set the env variable VLLM_RPC_TIMEOUT to a big number before you start the server. Say something like 30 minutes.
+    `export VLLM_RPC_TIMEOUT=1800000`
 
 ### Example commands and usage
 

docs/design/attention_backends.md

@@ -228,17 +228,3 @@ MLA decode backends are selected using the standard
 | `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
 | `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
 | `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
-
-### DeepSeek V4 Decode Backends
-
-DeepSeek V4 sparse MLA uses its own decode backends, selected via
-`--attention-backend=<BACKEND>` (e.g., `FLASHMLA_SPARSE_DSV4`,
-`FLASHINFER_MLA_SPARSE_DSV4`). They share the V4 sparse-index
-pipeline (compressor + SWA + indexer, 256-token blocks, head 512);
-default on NVIDIA is `FLASHMLA_SPARSE_DSV4`.
-
-| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
-| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
-| `FLASHINFER_MLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
-| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla`, `fp8` | 256 | 512 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
-| `ROCM_FLASHMLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |

docs/design/cuda_graphs_multimodal.md

@@ -17,7 +17,6 @@ The encoder CUDA Graph system uses a **budget-based capture/replay** strategy, m
 
 * [EncoderCudaGraphManager][vllm.v1.worker.encoder_cudagraph.EncoderCudaGraphManager]: orchestrates capture, replay, greedy packing, and data-parallel execution for encoder CUDA Graphs.
 * [SupportsEncoderCudaGraph][vllm.model_executor.models.interfaces.SupportsEncoderCudaGraph]: a runtime-checkable protocol that models implement to opt-in to encoder CUDA Graphs.
-* [EncoderItemSpec][vllm.v1.worker.encoder_cudagraph_defs.EncoderItemSpec]: describes a single encoder input item (image or video) with its input size and output token count.
 * [BudgetGraphMetadata][vllm.v1.worker.encoder_cudagraph.BudgetGraphMetadata]: holds the captured CUDA Graph and its associated I/O buffers for a single token budget level.
 
 ### Budget-based graph capture
@@ -31,7 +30,8 @@ class BudgetGraphMetadata:
     max_batch_size: int
     max_frames_per_batch: int
     graph: torch.cuda.CUDAGraph
-    input_buffers: dict[str, torch.Tensor]  # e.g. pixel_values, embeddings, seq metadata
+    input_buffer: torch.Tensor       # e.g. pixel_values
+    metadata_buffers: dict[str, torch.Tensor]  # e.g. embeddings, seq metadata
     output_buffer: torch.Tensor      # encoder hidden states
 ```
 
@@ -43,8 +43,8 @@ When a batch of images arrives, the manager sorts images by output token count (
 
 For each graph replay:
 
-1. Call `prepare_encoder_cudagraph_replay_buffers()` to compute buffer values (including `pixel_values` and precomputed metadata) from actual batch inputs.
-2. Zero the pre-allocated `input_buffers`, then slice-copy the replay values into them.
+1. Zero the pre-allocated `input_buffer`, then copy input tensors (e.g., `pixel_values`) into it.
+2. Zero `metadata_buffers`, then slice-copy precomputed values (e.g., rotary embeddings, sequence metadata).
 3. Replay the CUDA Graph.
 4. Clone outputs from `output_buffer` (cloning is necessary since the buffer is reused across replays).
 
@@ -65,15 +65,19 @@ Following <https://github.com/vllm-project/vllm/pull/35963> (ViT full CUDA graph
 
 Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGraph][vllm.model_executor.models.interfaces.SupportsEncoderCudaGraph] protocol. This protocol encapsulates all model-specific logic so that the manager remains model-agnostic. The protocol defines the following methods:
 
-* `get_encoder_cudagraph_config()` — returns static configuration (supported modalities, buffer keys, output hidden size, padding logics, max frames per video).
+* `get_encoder_cudagraph_config()` — returns static configuration (supported modalities, input key, buffer keys, output hidden size).
 * `get_encoder_cudagraph_budget_range(vllm_config)` — returns `(min_budget, max_budget)` for auto-inference of token budgets.
-* `get_encoder_cudagraph_item_specs(mm_kwargs)` — returns `list[EncoderItemSpec]` describing each item with its input size and output token count. Replaces the former three separate methods (`get_num_items`, `get_per_item_output_tokens`, `get_per_item_input_sizes`).
+* `get_encoder_cudagraph_num_items(mm_kwargs)` — returns the number of items (e.g. images) in the batch.
+* `get_encoder_cudagraph_per_item_output_tokens(mm_kwargs)` — returns per-item output token counts, used for greedy packing.
+* `get_encoder_cudagraph_per_item_input_sizes(mm_kwargs)` — returns per-item input sizes (e.g. patch counts), used for DP load balancing.
 * `select_encoder_cudagraph_items(mm_kwargs, indices)` — extracts a sub-batch of items by index, used during greedy packing and DP sharding.
-* `prepare_encoder_cudagraph_capture_inputs(...)` — creates dummy inputs for graph capture. Returns `EncoderCudaGraphCaptureInputs` with a single `values: dict[str, torch.Tensor]` that contains all buffers to be recorded into the graph.
-* `prepare_encoder_cudagraph_replay_buffers(mm_kwargs, max_batch_size, max_frames_per_batch)` — computes buffer values from actual batch inputs. Returns `EncoderCudaGraphReplayBuffers` with a `values` dict whose keys match `buffer_keys` in the config.
-* `encoder_cudagraph_forward(inputs: dict[str, torch.Tensor])` — forward pass accepting only fixed-shaped input tensors (the captured `values` dict). Called during both capture and replay. The `pixel_values` tensor is included in `inputs` alongside metadata buffers.
-* `encoder_eager_forward(mm_kwargs)` — fallback eager forward when no graph fits.
-* `postprocess_encoder_output(...)` — post-process encoder output, delegates to `scatter_output_slices` by default.
+* `prepare_encoder_cudagraph_capture_inputs(...)` — creates dummy inputs for graph capture.
+* `prepare_encoder_cudagraph_replay_buffers(...)` — computes new buffer values from actual batch inputs before replay.
+* `encoder_cudagraph_forward(...)` — forward pass using precomputed buffers (called during capture and replay).
+* `encoder_eager_forward(...)` — fallback eager forward when no graph fits.
+* `get_input_modality(...)` - return the modality of the inputs.
+* `get_max_frames_per_video()` - return model-specific max frames per video.
+* `postprocess_encoder_output(...)` - post process encoder output, directly call scatter_output_slices by default
 
 !!! note
     The `SupportsEncoderCudaGraph` protocol is designed to be model-agnostic. New vision encoder models can opt-in by implementing the protocol methods without modifying the manager.
@@ -82,7 +86,6 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
 
 | Architecture | Models | CG for Image | CG for Video |
 | ------------ | ------ | ------------ | ------------ |
-| `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ |
 | `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ |
 | `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ |
 | `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
@@ -100,7 +103,7 @@ Three fields in `CompilationConfig` control encoder CUDA Graphs:
 * `cudagraph_mm_encoder` (`bool`, default `False`) — enable CUDA Graph capture for multimodal encoder. When enabled, captures the full encoder forward as a CUDA Graph for each token budget level.
 * `encoder_cudagraph_token_budgets` (`list[int]`, default `[]`) — token budget levels for capture. If empty (default), auto-inferred from model architecture as power-of-2 levels. User-provided values override auto-inference.
 * `encoder_cudagraph_max_vision_items_per_batch` (`int`, default `0`) — maximum number of images/videos per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`.
-* `encoder_cudagraph_max_frames_per_batch` (`int`, default `None`) — maximum number of video frames per batch during capture. If `None` (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * max_frames_per_video` (`max_frames_per_video` is a model-specific value from `EncoderCudaGraphConfig`, computed by `get_max_frames_per_video()` on the model). If we limit the video count per prompt to `0`, it will also be set to `0` (i.e., fall back to image-only mode).
+* `encoder_cudagraph_max_frames_per_batch` (`int`, default `None`) — maximum number of video frames per batch during capture. If `None` (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * max_frames_per_video` (`max_frames_per_video` is a model-specific value according to its `processing_info`). If we limit the video count per prompt to `0`, it will also be set to `0` (i.e., fall back to image-only mode).
 
 ## Usage guide
 

docs/design/nixl_kv_cache_lease.md

@@ -128,7 +128,7 @@ The lease mechanism is controlled through `kv_connector_extra_config` in `--kv-t
 vllm serve <MODEL> \
   --kv-transfer-config '{
     "kv_connector": "NixlConnector",
-    "kv_role": "kv_producer",
+    "kv_role": "kv_both",
     "kv_connector_extra_config": {"kv_lease_duration": 60}
   }'
 ```

docs/design/torch_compile_multimodal.md

@@ -88,7 +88,7 @@ If compilation fails for a multimodal model:
 
 1. **Disable and test**: First verify the model works without compilation:
    ```bash
-   vllm serve <model> --compilation-config='{"mode":0,"compile_mm_encoder":"false"}'
+   VLLM_TORCH_COMPILE_LEVEL=0 vllm serve <model> --compilation-config='{"compile_mm_encoder":"false"}'
    ```
 
 2. **Check logs**: Enable debug logging to see compilation details:

docs/features/batch_invariance.md

@@ -17,7 +17,10 @@ Batch invariance is crucial for several use cases:
 
 ## Hardware Requirements
 
-Batch invariance requires NVIDIA GPUs with compute capability 8.0 or higher.
+Batch invariance currently requires NVIDIA GPUs with compute capability 9.0 or higher:
+
+- **H-series**: H100, H200
+- **B-series**: B100, B200
 
 ## Enabling Batch Invariance
 
@@ -104,7 +107,7 @@ Batch invariance has been tested and verified on the following models:
 - **Qwen3 (Dense)**: `Qwen/Qwen3-1.7B`, `Qwen/Qwen3-8B`, `Qwen/Qwen3-4B-AWQ`, `Qwen/Qwen3-8B-AWQ`
 - **Qwen3 (MoE)**: `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-Next-80B-A3B-Instruct`, `Qwen/Qwen3-30B-A3B-Thinking-2507-FP8`
 - **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct`
-- **Llama 3**: Llama3.1 and 3.2 series, `meta-llama/Llama-3.2-3B-Instruct` for example
+- **Llama 3**: `meta-llama/Llama-3.1-8B-Instruct`, `meta-llama/Llama-3.2-1B-Instruct`
 - **GPT-OSS**: `openai/gpt-oss-20b`, `openai/gpt-oss-120b`
 - **Mistral**: `mistralai/Mistral-7B-v0.3`
 

docs/features/nixl_connector_usage.md

@@ -50,7 +50,7 @@ To select a different backend, set `kv_connector_extra_config.backends` in `--kv
 vllm serve <MODEL> \
   --kv-transfer-config '{
     "kv_connector":"NixlConnector",
-    "kv_role":"kv_producer",
+    "kv_role":"kv_both",
     "kv_connector_extra_config":{"backends":["LIBFABRIC"]}
   }'
 ```
@@ -60,7 +60,7 @@ You can also pass JSON keys individually using dotted arguments, and you can app
 ```bash
 vllm serve <MODEL> \
   --kv-transfer-config.kv_connector NixlConnector \
-  --kv-transfer-config.kv_role kv_producer \
+  --kv-transfer-config.kv_role kv_both \
   --kv-transfer-config.kv_connector_extra_config.backends+ LIBFABRIC
 ```
 
@@ -81,7 +81,7 @@ VLLM_NIXL_SIDE_CHANNEL_PORT=5600 \
 vllm serve Qwen/Qwen3-0.6B \
   --port 8100 \
   --enforce-eager \
-  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_producer","kv_load_failure_policy":"fail"}'
+  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail"}'
 ```
 
 ### Consumer (Decoder) Configuration
@@ -96,7 +96,7 @@ VLLM_NIXL_SIDE_CHANNEL_PORT=5601 \
 vllm serve Qwen/Qwen3-0.6B \
   --port 8200 \
   --enforce-eager \
-  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer","kv_load_failure_policy":"fail"}'
+  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail"}'
 ```
 
 ### Proxy Server
@@ -212,21 +212,10 @@ sequenceDiagram
 Enable bidirectional KV transfer by setting `bidirectional_kv_xfer` in `kv_connector_extra_config` on **both** P and D instances:
 
 ```bash
-# Prefill instance
 vllm serve <MODEL> \
   --kv-transfer-config '{
     "kv_connector": "NixlConnector",
-    "kv_role": "kv_producer",
-    "kv_connector_extra_config": {
-      "bidirectional_kv_xfer": true
-    }
-  }'
-
-# Decode instance
-vllm serve <MODEL> \
-  --kv-transfer-config '{
-    "kv_connector": "NixlConnector",
-    "kv_role": "kv_consumer",
+    "kv_role": "kv_both",
     "kv_connector_extra_config": {
       "bidirectional_kv_xfer": true
     }
@@ -370,10 +359,11 @@ For multi-host DP deployment, only need to provide the host/port of the head ins
 
 - **kv_producer**: For prefiller instances that generate KV caches
 - **kv_consumer**: For decoder instances that consume KV caches from prefiller
-- **kv_both** (deprecated): Previously used as a catch-all when the role was not predetermined. This value is now deprecated for NixlConnector and will be removed in a future release.
+- **kv_both**: Enables symmetric functionality where the connector can act as both producer and consumer. This provides flexibility for experimental setups and scenarios where the role distinction is not predetermined.
 
-!!! warning
-    `kv_role="kv_both"` is deprecated for NixlConnector. Please set `kv_role="kv_producer"` for prefill instances and `kv_role="kv_consumer"` for decode instances. See [#33702](https://github.com/vllm-project/vllm/issues/33702) for details.
+!!! tip
+    NixlConnector currently does not distinguish `kv_role`; the actual prefiller/decoder roles are determined by the upper-level proxy (e.g., `toy_proxy_server.py` using `--prefiller-hosts` and `--decoder-hosts`).
+    Therefore, `kv_role` in `--kv-transfer-config` is effectively a placeholder and does not affect NixlConnector's behavior.
 
 ### KV Load Failure Policy