[Docker] Fix CUTLASS DSL cu13 install order in Dockerfile (#45204)

Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
(cherry picked from commit 40e065e86a)

.buildkite/ci_config_rocm.yaml

@@ -0,0 +1,23 @@
+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,42 +1,73 @@
-group: Hardware - AMD Build 
+group: Hardware - AMD Build
 steps:
-  - label: "AMD: :docker: build image"
-    key: image-build-amd
+  # 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
     depends_on: []
     device: amd_cpu
     no_plugin: true
     commands:
-    - >
-      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
-        python3 - <<PY
-      import torch, vllm
-      print(torch.__version__)
-      print(vllm.__version__)
-      PY
-        echo AMD image smoke OK
-      '
-    - docker push "rocm/vllm-ci:${BUILDKITE_COMMIT}"
+      - 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
+        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
+        '
+    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,18 +28,19 @@ 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"
 
-- 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"
+# 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-Language Generation and Pooling Model Tests
   depends_on: []

.buildkite/intel_jobs/basic_correctness.yaml

@@ -0,0 +1,22 @@
+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,3 +53,26 @@ 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,7 +40,9 @@ 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
+        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
         '
   - label: "XPU V1 test"
     depends_on:

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

@@ -2,7 +2,6 @@
     {
         "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,7 +2,6 @@
     {
         "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,7 +13,6 @@
       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,7 +5,6 @@
     ],
     "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,7 +9,6 @@
       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,7 +5,6 @@
     ],
     "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,7 +5,6 @@
     ],
     "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,7 +2,6 @@
     {
         "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,7 +2,6 @@
     {
         "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,8 +28,10 @@
 ###############################################################################
 set -o pipefail
 
-# Export Python path
-export PYTHONPATH=".."
+# 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:-..}"
 
 ###############################################################################
 # Helper Functions
@@ -52,6 +54,108 @@ 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
@@ -243,22 +347,30 @@ report_docker_usage
 
 # --- Pull test image ---
 echo "--- Pulling container"
-image_name="rocm/vllm-ci:${BUILDKITE_COMMIT}"
+image_name="${VLLM_CI_FALLBACK_IMAGE:-rocm/vllm-ci:${BUILDKITE_COMMIT:-local}}"
+artifact_work_dir=""
 container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
-docker pull "${image_name}"
 
 remove_docker_container() {
-  # docker 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
+  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
 }
+trap remove_docker_container EXIT
 
-on_exit() {
-  local exit_code=$?
-  remove_docker_container
-  exit "$exit_code"
-}
-trap on_exit EXIT
+if ! prepare_artifact_image; then
+  echo "Using full ROCm CI image: ${image_name}"
+  docker pull "${image_name}" || exit 1
+fi
 
 # --- Prepare commands ---
 echo "--- Running container"
@@ -267,6 +379,14 @@ 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
@@ -306,7 +426,14 @@ fi
 
 echo "Final commands: $commands"
 
-MYPYTHONPATH=".."
+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\""
 
 # Verify GPU access
 render_gid=$(getent group render | cut -d: -f3)
@@ -383,6 +510,8 @@ 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 \
@@ -390,10 +519,15 @@ 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 "${commands}"
+    /bin/bash -c "${CONTAINER_PREFLIGHT} && ${commands}"
 
   exit_code=$?
   handle_pytest_exit "$exit_code"

.buildkite/test-amd.yaml

@@ -315,24 +315,6 @@ 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
@@ -449,29 +431,6 @@ 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
@@ -931,10 +890,10 @@ steps:
   - vllm/
   - tests/basic_correctness/test_basic_correctness
   - tests/basic_correctness/test_cpu_offload
-  - tests/basic_correctness/test_cumem.py
+  - tests/basic_correctness/test_mem.py
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s basic_correctness/test_cumem.py
+  - pytest -v -s basic_correctness/test_mem.py
   - pytest -v -s basic_correctness/test_basic_correctness.py
   - pytest -v -s basic_correctness/test_cpu_offload.py
 
@@ -1322,12 +1281,11 @@ 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
-  - pytest -v -s entrypoints/llm/test_generate.py
-  - pytest -v -s entrypoints/offline_mode
+  - 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
 
 - label: Entrypoints Integration (Pooling) # TBD
   timeout_in_minutes: 180
@@ -1369,12 +1327,13 @@ steps:
   - vllm/platforms/rocm.py
   commands:
   - pytest -v -s entrypoints/openai/tool_parsers
-  - 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
+  - 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
 
 - 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/
@@ -1433,6 +1392,7 @@ 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/
@@ -1739,6 +1699,29 @@ 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
@@ -2184,10 +2167,72 @@ 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
@@ -2546,6 +2591,7 @@ 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/
@@ -2555,7 +2601,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
+  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
 
 #-----------------------------------------------------------  mi325 · evals  -----------------------------------------------------------#
 
@@ -2637,6 +2683,7 @@ 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)
-  key: v1-attention-h100
+- label: V1 attention (H100-MI300)
+  key: v1-attention-h100-mi300
   timeout_in_minutes: 30
   device: h100
   source_file_dependencies:
@@ -13,6 +13,20 @@ 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_cumem.py
+  - tests/basic_correctness/test_mem.py
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s basic_correctness/test_cumem.py
+  - pytest -v -s basic_correctness/test_mem.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,6 +26,12 @@ 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
@@ -38,7 +44,7 @@ steps:
     - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
   mirror:
     amd:
-      device: mi325_1
+      device: mi300_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/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
+  - 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
 
 - label: Entrypoints Integration (LLM)
   key: entrypoints-integration-llm
@@ -20,16 +20,14 @@ 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
+  - 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/offline_mode # Needs to avoid interference with other tests
+  - pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
   mirror:
     amd:
       device: mi325_1
-      soft_fail: true
       depends_on:
       - image-build-amd
 
@@ -46,8 +44,7 @@ steps:
   - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
   mirror:
     amd:
-      device: mi325_1
-      soft_fail: true
+      device: mi300_1
       timeout_in_minutes: 80
       depends_on:
       - image-build-amd
@@ -69,8 +66,7 @@ steps:
   - pytest -v -s tool_use
   mirror:
     amd:
-      device: mi325_1
-      soft_fail: true
+      device: mi300_1
       timeout_in_minutes: 60
       depends_on:
       - image-build-amd
@@ -90,7 +86,6 @@ steps:
   mirror:
     amd:
       device: mi325_1
-      soft_fail: true
       timeout_in_minutes: 60
       depends_on:
       - image-build-amd
@@ -109,8 +104,7 @@ steps:
   - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
   mirror:
     amd:
-      device: mi325_1
-      soft_fail: true
+      device: mi300_1
       depends_on:
       - image-build-amd
 
@@ -156,3 +150,21 @@ 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: mi325_1
+      device: mi300_1
       source_file_dependencies:
       - csrc/quantization/
       - vllm/model_executor/layers/quantization

.buildkite/test_areas/lm_eval.yaml

@@ -12,6 +12,21 @@ 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,6 +281,7 @@ 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,11 +94,13 @@ 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,6 +51,7 @@ 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: mi325_1
+      device: mi300_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: mi325_1
+      device: mi300_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: mi325_1
+      device: mi300_1
       depends_on:
       - image-build-amd
 
@@ -153,3 +153,12 @@ 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,6 +37,21 @@ 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
@@ -61,6 +76,20 @@ 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
@@ -72,6 +101,20 @@ 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,10 +34,11 @@
 /vllm/entrypoints/speech_to_text/realtime @njhill
 /vllm/entrypoints/speech_to_text @NickLucche
 /vllm/entrypoints/pooling @noooop
-/vllm/entrypoints/sagemaker @DarkLight1337
+/vllm/entrypoints/serve/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,6 +21,9 @@ pull_request_rules:
     - check-failure=pre-commit
     - -closed
     - -draft
+    - or:
+      - label=ready
+      - label=verified
   actions:
     comment:
       message: |
@@ -36,18 +39,6 @@ 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/stale.yml

@@ -15,7 +15,7 @@ jobs:
       actions: write
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/stale@997185467fa4f803885201cee163a9f38240193d # v10.1.1
+      - uses: actions/stale@eb5cf3af3ac0a1aa4c9c45633dd1ae542a27a899 # v10.3.0
         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,33 +148,27 @@ repos:
     language: python
     entry: python tools/pre_commit/generate_nightly_torch_test.py
     files: ^requirements/test/cuda\.(in|txt)$
-  - 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
+  - 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"
     <<: &mypy_common
       language: python
       types_or: [python, pyi]
       require_serial: true
-      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
+      additional_dependencies: ["mypy==1.20.2", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
   - 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 1 "3.11"
+    entry: python tools/pre_commit/mypy.py "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 1 "3.12"
+    entry: python tools/pre_commit/mypy.py "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 1 "3.13"
+    entry: python tools/pre_commit/mypy.py "3.13"
     <<: *mypy_common
     stages: [manual] # Only run in CI
   - id: shellcheck

AGENTS.md

@@ -98,11 +98,13 @@ pre-commit run --all-files
 pre-commit run ruff-check --all-files
 
 # Run mypy as it is in CI:
-pre-commit run mypy-3.10 --all-files --hook-stage manual
+pre-commit run mypy-3.12 --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,6 +112,8 @@ 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 "")
@@ -305,18 +307,10 @@ 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/cuda_utils_kernels.cu"
-  "csrc/custom_all_reduce.cu"
-  "csrc/torch_bindings.cpp"
-  "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
+  "csrc/torch_bindings.cpp")
 
 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.
@@ -349,9 +343,6 @@ 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}")
@@ -503,12 +494,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/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu"
-      "csrc/moe/mxfp8_moe/mxfp8_experts_quant.cu")
+      "csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu"
+      "csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu")
     set_gencode_flags_for_srcs(
       SRCS "${SRCS}"
       CUDA_ARCHS "${ES_MXFP8_GROUPED_MM_ARCHS}")
-    list(APPEND VLLM_EXT_SRC "${SRCS}")
+    list(APPEND VLLM_STABLE_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()
@@ -598,9 +589,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
 endif()
 
 if (VLLM_GPU_LANG STREQUAL "HIP")
-  # Add QuickReduce kernels
+  # Add QuickReduce kernels (ROCm-only; not part of stable ABI migration).
+  # TODO: Remove the cuda_view when ROCm upgrade to torch 2.11.
   list(APPEND VLLM_EXT_SRC
     "csrc/custom_quickreduce.cu"
+    "csrc/cuda_view.cu"
+    "csrc/libtorch_stable/cuda_utils_kernels.cu"
   )
 # if ROCM endif
 endif()
@@ -633,6 +627,9 @@ 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"
@@ -640,6 +637,7 @@ 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"
@@ -647,19 +645,21 @@ 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_fused.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")
 
   if(VLLM_GPU_LANG STREQUAL "CUDA")
     list(APPEND VLLM_STABLE_EXT_SRC
-      "csrc/cuda_utils_kernels.cu"
-      "csrc/cutlass_extensions/common.cpp"
+      "csrc/libtorch_stable/cuda_view.cu"
+      "csrc/libtorch_stable/cuda_utils_kernels.cu"
+      "csrc/libtorch_stable/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/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")
+      "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu"
+      "csrc/libtorch_stable/minimax_reduce_rms_kernel.cu")
 
     set_gencode_flags_for_srcs(
       SRCS "${VLLM_STABLE_EXT_SRC}"
@@ -1072,20 +1072,25 @@ 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()
 
@@ -1303,7 +1308,8 @@ 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/q_gemm_rdna3_wmma.cu"
+      "csrc/rocm/moe_q_gemm_rdna3.cu")
   endif()
 
   define_extension_target(

benchmarks/kernels/benchmark_moe.py

@@ -271,7 +271,6 @@ def benchmark_config(
                     moe_config=moe_config,
                     quant_config=quant_config,
                 ),
-                inplace=not disable_inplace(),
             )
 
         with override_config(config):
@@ -279,7 +278,6 @@ 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,
@@ -298,7 +296,6 @@ def benchmark_config(
                 w2,
                 topk_weights,
                 topk_ids,
-                inplace=inplace,
                 quant_config=quant_config,
             )
 

cmake/hipify.py

@@ -14,7 +14,18 @@ import argparse
 import os
 import shutil
 
-from torch.utils.hipify.hipify_python import hipify
+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)
+        )
+    )
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -53,7 +64,11 @@ if __name__ == "__main__":
     hipify_result = hipify(
         project_directory=args.project_dir,
         output_directory=args.output_dir,
-        header_include_dirs=[],
+        # 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=["."],
         includes=includes,
         extra_files=extra_files,
         show_detailed=True,
@@ -64,14 +79,20 @@ if __name__ == "__main__":
     hipified_sources = []
     for source in args.sources:
         s_abs = os.path.abspath(source)
-        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
-        )
+        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_sources.append(hipified_s_abs)
 
     assert len(hipified_sources) == len(args.sources)

cmake/utils.cmake

@@ -81,6 +81,14 @@ 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,7 +30,12 @@
   }()
 
 namespace {
-enum class FusedMOEAct { SiluAndMul, SwigluOAIAndMul, GeluAndMul };
+enum class FusedMOEAct {
+  SiluAndMul,
+  SwigluOAIAndMul,
+  GeluAndMul,
+  GeluTanhAndMul,
+};
 
 FusedMOEAct get_act_type(const std::string& act) {
   if (act == "silu") {
@@ -39,6 +44,8 @@ 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);
   }
@@ -143,6 +150,44 @@ 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,
@@ -160,6 +205,9 @@ 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,3 +1,4 @@
+// TODO: Remove this once ROCm upgrade to torch 2.11.
 #include <torch/all.h>
 #include <torch/cuda.h>
 #include <cuda_runtime.h>

csrc/libtorch_stable/concat_mla_q.cuh

@@ -1,9 +1,6 @@
 #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

@@ -0,0 +1,76 @@
+#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/custom_all_reduce.cu

@@ -1,7 +1,11 @@
-#include <ATen/cuda/Exceptions.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAStream.h>
-#include <torch/all.h>
+#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 "custom_all_reduce.cuh"
 
@@ -11,7 +15,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::Tensor& rank_data, int64_t rank,
+                      torch::stable::Tensor& rank_data, int64_t rank,
                       bool fully_connected) {
   int world_size = fake_ipc_ptrs.size();
   if (world_size > 8)
@@ -25,9 +29,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.data_ptr(),
-                                            rank_data.numel(), rank, world_size,
-                                            fully_connected);
+  return (fptr_t) new vllm::CustomAllreduce(
+      ipc_ptrs, rank_data.mutable_data_ptr(), rank_data.numel(), rank,
+      world_size, fully_connected);
 }
 
 /**
@@ -46,10 +50,14 @@ 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::Tensor& t) {
-  return t.is_contiguous() ||
-         (t.storage().nbytes() - t.storage_offset() * t.element_size() ==
-          t.numel() * t.element_size());
+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());
 }
 
 /**
@@ -59,42 +67,45 @@ bool _is_weak_contiguous(torch::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::Tensor& inp, torch::Tensor& out,
-                fptr_t _reg_buffer, int64_t reg_buffer_sz_bytes) {
+void all_reduce(fptr_t _fa, torch::stable::Tensor& inp,
+                torch::stable::Tensor& out, fptr_t _reg_buffer,
+                int64_t reg_buffer_sz_bytes) {
   auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(inp));
-  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      inp.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index());
 
-  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));
+  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));
   auto input_size = inp.numel() * inp.element_size();
   auto reg_buffer = reinterpret_cast<void*>(_reg_buffer);
   if (reg_buffer) {
-    TORCH_CHECK_LE(input_size, reg_buffer_sz_bytes);
-    AT_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.data_ptr(), input_size,
-                                  cudaMemcpyDeviceToDevice, stream));
+    STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes));
+    STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size,
+                                   cudaMemcpyDeviceToDevice, stream));
   } else {
-    reg_buffer = inp.data_ptr();
+    reg_buffer = inp.mutable_data_ptr();
   }
   switch (out.scalar_type()) {
-    case at::ScalarType::Float: {
+    case torch::headeronly::ScalarType::Float: {
       fa->allreduce<float>(stream, reinterpret_cast<float*>(reg_buffer),
-                           reinterpret_cast<float*>(out.data_ptr()),
+                           reinterpret_cast<float*>(out.mutable_data_ptr()),
                            out.numel());
       break;
     }
-    case at::ScalarType::Half: {
+    case torch::headeronly::ScalarType::Half: {
       fa->allreduce<half>(stream, reinterpret_cast<half*>(reg_buffer),
-                          reinterpret_cast<half*>(out.data_ptr()), out.numel());
+                          reinterpret_cast<half*>(out.mutable_data_ptr()),
+                          out.numel());
       break;
     }
 #if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
-    case at::ScalarType::BFloat16: {
+    case torch::headeronly::ScalarType::BFloat16: {
       fa->allreduce<nv_bfloat16>(
           stream, reinterpret_cast<nv_bfloat16*>(reg_buffer),
-          reinterpret_cast<nv_bfloat16*>(out.data_ptr()), out.numel());
+          reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), out.numel());
       break;
     }
 #endif
@@ -112,7 +123,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);
-  TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_);
+  STD_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]);
@@ -143,47 +154,49 @@ void register_graph_buffers(fptr_t _fa,
   fa->register_graph_buffers(bytes, offsets);
 }
 
-std::tuple<fptr_t, torch::Tensor> allocate_shared_buffer_and_handle(
+std::tuple<fptr_t, torch::stable::Tensor> allocate_shared_buffer_and_handle(
     int64_t size) {
-  auto device_index = c10::cuda::current_device();
-  at::DeviceGuard device_guard(at::Device(at::DeviceType::CUDA, device_index));
+  int device_index;
+  STD_CUDA_CHECK(cudaGetDevice(&device_index));
+  const torch::stable::accelerator::DeviceGuard device_guard(device_index);
   void* buffer;
   cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed;
-  auto stream = c10::cuda::getCurrentCUDAStream().stream();
-  AT_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));
+  const cudaStream_t stream = get_current_cuda_stream(device_index);
+  STD_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));
 
   // Allocate buffer
 #if defined(USE_ROCM)
   // data buffers need to be "uncached" for signal on MI200
-  AT_CUDA_CHECK(
+  STD_CUDA_CHECK(
       hipExtMallocWithFlags((void**)&buffer, size, hipDeviceMallocUncached));
 #else
-  AT_CUDA_CHECK(cudaMalloc((void**)&buffer, size));
+  STD_CUDA_CHECK(cudaMalloc((void**)&buffer, size));
 #endif
-  AT_CUDA_CHECK(cudaMemsetAsync(buffer, 0, size, stream));
-  AT_CUDA_CHECK(cudaStreamSynchronize(stream));
-  AT_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));
+  STD_CUDA_CHECK(cudaMemsetAsync(buffer, 0, size, stream));
+  STD_CUDA_CHECK(cudaStreamSynchronize(stream));
+  STD_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));
 
   // Create IPC memhandle for the allocated buffer.
   // Will use it in open_mem_handle.
-  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));
+  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));
 
   return std::make_tuple(reinterpret_cast<fptr_t>(buffer), handle);
 }
 
-fptr_t open_mem_handle(torch::Tensor& mem_handle) {
+fptr_t open_mem_handle(torch::stable::Tensor& mem_handle) {
   void* ipc_ptr;
-  AT_CUDA_CHECK(cudaIpcOpenMemHandle(
-      (void**)&ipc_ptr, *((const cudaIpcMemHandle_t*)mem_handle.data_ptr()),
+  STD_CUDA_CHECK(cudaIpcOpenMemHandle(
+      (void**)&ipc_ptr,
+      *((const cudaIpcMemHandle_t*)mem_handle.const_data_ptr()),
       cudaIpcMemLazyEnablePeerAccess));
   return reinterpret_cast<fptr_t>(ipc_ptr);
 }
 
 void free_shared_buffer(fptr_t buffer) {
-  AT_CUDA_CHECK(cudaFree(reinterpret_cast<void*>(buffer)));
+  STD_CUDA_CHECK(cudaFree(reinterpret_cast<void*>(buffer)));
 }

csrc/libtorch_stable/cutlass_extensions/common.cpp

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

csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu

@@ -28,7 +28,20 @@
  *   [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
@@ -37,14 +50,6 @@
 #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
@@ -70,7 +75,7 @@ namespace deepseek_v4_fused_ops {
 
 namespace {
 inline int getSMVersion() {
-  auto* props = at::cuda::getCurrentDeviceProperties();
+  auto* props = get_device_prop();
   return props->major * 10 + props->minor;
 }
 }  // namespace
@@ -97,6 +102,35 @@ 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
 // ────────────────────────────────────────────────────────────────────────────
@@ -564,7 +598,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.
-  TORCH_CHECK(
+  STD_TORCH_CHECK(
       sm_version >= 80,
       "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert requires sm_80+ "
       "(Ampere or newer); got sm_",
@@ -635,7 +669,7 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
     DISPATCH(64)
     DISPATCH(128)
     default:
-      TORCH_CHECK(false,
+      STD_TORCH_CHECK(false,
                   "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert: "
                   "unsupported num_heads_q_padded=",
                   num_heads_q_padded,
@@ -644,80 +678,504 @@ 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::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
+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
     double eps, int64_t cache_block_size) {
-  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");
+  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");
 
   // 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));
-  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");
+  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_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));
 
-  at::cuda::OptionalCUDAGuard device_guard(device_of(q_in));
-  auto stream = at::cuda::getCurrentCUDAStream();
+  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());
 
   // Allocate the padded q output.  The kernel writes every element (live
   // region gets RMSNorm+RoPE; pad region gets zeros), so `empty` is safe.
-  torch::Tensor q_out = torch::empty(
-      {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.options());
+  auto q_out = torch::stable::new_empty(
+      q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type());
 
-  VLLM_DISPATCH_HALF_TYPES(
+  VLLM_STABLE_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.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),
+                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),
                 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/minimax_reduce_rms_kernel.cu

@@ -15,16 +15,19 @@
  * 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>
@@ -611,7 +614,7 @@ int get_sm_count() {
   static int sm_count = 0;
   if (sm_count == 0) {
     int device_id;
-    CUDA_CHECK(cudaGetDevice(&device_id));
+    STD_CUDA_CHECK(cudaGetDevice(&device_id));
     cudaDeviceProp device_prop;
     cudaGetDeviceProperties(&device_prop, device_id);
     sm_count = device_prop.multiProcessorCount;
@@ -621,13 +624,13 @@ int get_sm_count() {
 
 inline int getSMVersion(bool queryRealSmArch = false) {
   int device{-1};
-  CUDA_CHECK(cudaGetDevice(&device));
+  STD_CUDA_CHECK(cudaGetDevice(&device));
   int sm_major = 0;
   int sm_minor = 0;
-  CUDA_CHECK(cudaDeviceGetAttribute(&sm_major,
-                                    cudaDevAttrComputeCapabilityMajor, device));
-  CUDA_CHECK(cudaDeviceGetAttribute(&sm_minor,
-                                    cudaDevAttrComputeCapabilityMinor, device));
+  STD_CUDA_CHECK(cudaDeviceGetAttribute(
+      &sm_major, cudaDevAttrComputeCapabilityMajor, device));
+  STD_CUDA_CHECK(cudaDeviceGetAttribute(
+      &sm_minor, cudaDevAttrComputeCapabilityMinor, device));
   int sm = sm_major * 10 + sm_minor;
   if (sm == 121 && !queryRealSmArch) {
     return 120;
@@ -639,7 +642,7 @@ template <typename KernelFunc>
 int get_max_active_blocks(KernelFunc kernel, int block_size,
                           int dynamic_smem = 0) {
   int max_active = 0;
-  CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+  STD_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
       &max_active, kernel, block_size, dynamic_smem));
   return std::max(max_active, 1);
 }
@@ -678,27 +681,27 @@ void minimax_reduce_rms_kernel_launcher(MiniMaxReduceRMSParams const& params) {
   cfg.attrs = attribute;
   cfg.numAttrs = SM >= 90 ? 2 : 0;
 
-  CUDA_CHECK(cudaLaunchKernelEx(
+  STD_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) {
-  TORCH_CHECK(params.size_q % params.hidden_dim == 0);
-  TORCH_CHECK(params.hidden_dim % kElemsPerAccess<DType> == 0);
+  STD_TORCH_CHECK(params.size_q % params.hidden_dim == 0);
+  STD_TORCH_CHECK(params.hidden_dim % kElemsPerAccess<DType> == 0);
   if (params.stride_q > 0) {
-    TORCH_CHECK(params.stride_q % kElemsPerAccess<DType> == 0);
+    STD_TORCH_CHECK(params.stride_q % kElemsPerAccess<DType> == 0);
   }
-  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);
+  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);
   if (params.stride_k > 0) {
-    TORCH_CHECK(params.stride_k % kElemsPerAccess<DType> == 0);
+    STD_TORCH_CHECK(params.stride_k % kElemsPerAccess<DType> == 0);
   }
 
   int token_num = params.size_q / params.hidden_dim;
@@ -746,7 +749,7 @@ void minimax_reduce_rms_kernel_launcher_float4(
   cfg.attrs = attribute;
   cfg.numAttrs = SM >= 90 ? 2 : 0;
 
-  CUDA_CHECK(cudaLaunchKernelEx(&cfg, kfn, params));
+  STD_CUDA_CHECK(cudaLaunchKernelEx(&cfg, kfn, params));
 }
 
 template <int NRanks>
@@ -759,21 +762,21 @@ void dispatch_dtype(MiniMaxReduceRMSParams const& params) {
                     (params.hidden_dim * params.nranks == 6144) &&
                     (params.hidden_dim_k * params.nranks == 1024);
 
-  if (params.dtype == at::ScalarType::Half) {
+  if (params.dtype == torch::headeronly::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 == at::ScalarType::BFloat16) {
+  } else if (params.dtype == torch::headeronly::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 == at::ScalarType::Float) {
+  } else if (params.dtype == torch::headeronly::ScalarType::Float) {
     if (use_float4) {
       minimax_reduce_rms_kernel_launcher_float4<float, NRanks, 6144, 1024>(
           params);
@@ -781,7 +784,7 @@ void dispatch_dtype(MiniMaxReduceRMSParams const& params) {
       minimax_reduce_rms_kernel_launcher<float, NRanks>(params);
     }
   } else {
-    TORCH_CHECK(false, "Unsupported data type for minimax_reduce_rms_op");
+    STD_TORCH_CHECK(false, "Unsupported data type for minimax_reduce_rms_op");
   }
 }
 
@@ -795,16 +798,18 @@ void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params) {
   } else if (params.nranks == 16) {
     dispatch_dtype<16>(params);
   } else {
-    TORCH_CHECK(false, "minimax_reduce_rms_op: unsupported ranks number!");
+    STD_TORCH_CHECK(false, "minimax_reduce_rms_op: unsupported ranks number!");
   }
 }
 }  // namespace tensorrt_llm
 }  // namespace vllm
 
-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) {
+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());
   auto allreduce_params = vllm::tensorrt_llm::MiniMaxReduceRMSParams();
 
   allreduce_params.nranks = static_cast<int>(nranks);
@@ -815,12 +820,12 @@ torch::Tensor minimax_allreduce_rms(torch::Tensor const& input,
   allreduce_params.stride_q = allreduce_params.hidden_dim;
   allreduce_params.workspace =
       reinterpret_cast<void**>(workspace.mutable_data_ptr());
-  allreduce_params.allreduce_in = input.data_ptr();
-  allreduce_params.rms_gamma = norm_weight.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.rms_eps = static_cast<float>(eps);
-  allreduce_params.stream = at::cuda::getCurrentCUDAStream(input.get_device());
+  allreduce_params.stream = get_current_cuda_stream(input.get_device_index());
 
-  torch::Tensor rms_norm_out = torch::empty_like(input);
+  torch::stable::Tensor rms_norm_out = torch::stable::empty_like(input);
   allreduce_params.rms_norm_out = rms_norm_out.mutable_data_ptr();
 
   vllm::tensorrt_llm::minimax_reduce_rms_op(allreduce_params);
@@ -828,26 +833,33 @@ torch::Tensor minimax_allreduce_rms(torch::Tensor const& input,
   return rms_norm_out;
 }
 
-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");
+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");
   int64_t qkv_dim = qkv.size(-1);
-  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");
+  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());
 
   int64_t num_tokens = qkv.size(0);
   int elem_bytes = qkv.element_size();
 
-  torch::Tensor q_out = torch::empty({num_tokens, q_size}, qkv.options());
-  torch::Tensor k_out = torch::empty({num_tokens, kv_size}, qkv.options());
+  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());
 
   auto params = vllm::tensorrt_llm::MiniMaxReduceRMSParams();
   params.nranks = static_cast<int>(nranks);
@@ -863,13 +875,14 @@ std::tuple<torch::Tensor, torch::Tensor> minimax_allreduce_rms_qk(
   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 = static_cast<uint8_t*>(qkv.data_ptr());
+  uint8_t* base =
+      const_cast<uint8_t*>(static_cast<const uint8_t*>(qkv.const_data_ptr()));
   params.allreduce_in = base;
   params.allreduce_in_k = base + q_size * elem_bytes;
-  params.rms_gamma = norm_weight_q.data_ptr();
-  params.rms_gamma_k = norm_weight_k.data_ptr();
+  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_eps = static_cast<float>(eps);
-  params.stream = at::cuda::getCurrentCUDAStream(qkv.get_device());
+  params.stream = get_current_cuda_stream(qkv.get_device_index());
 
   params.rms_norm_out = q_out.mutable_data_ptr();
   params.rms_norm_out_k = k_out.mutable_data_ptr();

csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu

@@ -0,0 +1,69 @@
+// 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/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 <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <torch/all.h>
+
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/util/Exception.h>
 
 #include <cassert>
 #include <iostream>
@@ -15,18 +15,22 @@
 #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::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) {
+    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) {
   using OffsetFunctor = CutlassMxfp8GroupedMmOffsetFunctor<GemmTraits>;
   using ElementA = typename OffsetFunctor::ElementA;
   using ElementB = typename OffsetFunctor::ElementB;
@@ -42,10 +46,10 @@ void cutlass_mxfp8_grouped_mm_pre_compute(
   using StrideB = typename StrideFunctor::StrideB;
   using StrideD = typename StrideFunctor::StrideD;
 
-  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.");
+  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.");
 
   OffsetFunctor offset_functor(
       reinterpret_cast<int*>(expert_offsets.data_ptr()),
@@ -72,13 +76,18 @@ void cutlass_mxfp8_grouped_mm_pre_compute(
 }
 
 template <typename GemmTraits>
-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) {
+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) {
   using Gemm = typename GemmTraits::Gemm;
   using ElementA = typename Gemm::ElementA;
   using ElementB = typename Gemm::ElementB;
@@ -93,13 +102,12 @@ void cutlass_mxfp8_grouped_mm(
       typename GemmTraits::ProblemShape::UnderlyingProblemShape;
 
   cutlass::KernelHardwareInfo hw_info;
-  hw_info.device_id = c10::cuda::current_device();
-  hw_info.sm_count =
-      at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+  hw_info.device_id = d_ptrs.get_device_index();
+  hw_info.sm_count = get_device_prop()->multiProcessorCount;
   hw_info.cluster_shape = GemmTraits::MMAConfig::preferred_cluster;
   hw_info.cluster_shape_fallback = GemmTraits::MMAConfig::fallback_cluster;
 
-  int num_experts = (int)problem_sizes.size(0);
+  int num_experts = static_cast<int>(problem_sizes.size(0));
 
   UnderlyingProblemShape* underlying_problem_shape =
       reinterpret_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
@@ -127,44 +135,55 @@ void cutlass_mxfp8_grouped_mm(
   Gemm gemm;
 
   auto can_implement_status = gemm.can_implement(arguments);
-  TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
-              "Failed to implement GEMM");
+  STD_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::Tensor workspace = torch::empty(workspace_size, options_uint8);
+  torch::stable::Tensor workspace = torch::stable::empty(
+      {static_cast<int64_t>(workspace_size)},
+      torch::headeronly::ScalarType::Byte, std::nullopt, d_ptrs.device());
 
   auto status = gemm.initialize(arguments, workspace.data_ptr(), stream);
-  TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to initialize GEMM");
+  STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
+                  "Failed to initialize GEMM");
 
   status = gemm.run(stream, nullptr, true);  // Enable PDL
-  TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
+  STD_TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
 }
 
 template <typename OutType>
 void cutlass_mxfp8_grouped_mm_dispatch_out_dtype(
-    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());
+    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();
 
-  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 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 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);
+  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);
 
   using GemmTraits = CutlassMxfp8GroupedMmGemmTraits<MMA1SMConfig, OutType>;
   cutlass_mxfp8_grouped_mm_pre_compute<GemmTraits>(
@@ -176,4 +195,4 @@ void cutlass_mxfp8_grouped_mm_dispatch_out_dtype(
       layout_sfa, layout_sfb, problem_sizes, stream);
 }
 
-}  // namespace expert_specialization
+}  // namespace expert_specialization

csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu

@@ -0,0 +1,66 @@
+// 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/moe/mxfp8_moe/mxfp8_experts_quant.cuh

@@ -4,16 +4,19 @@
 // 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/all.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 <cuda/ptx>
 
 #include "cute/tensor.hpp"
+#include "libtorch_stable/torch_utils.h"
 
 namespace expert_specialization {
 
@@ -356,12 +359,12 @@ __global__ void mxfp8_experts_quant_kernel(
 }
 
 template <typename T_IN>
-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) {
+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) {
   ThrLayout thr_layout{};
   ValLayout val_layout{};
   SfR2SThrLayout r2s_thr_layout{};
@@ -386,19 +389,18 @@ void launch_mxfp8_experts_quant(const torch::Tensor& input,
       CopyAtomR2S{}, r2s_thr_layout, r2s_val_layout);  // Tiler_MN: (16, 4)
 
   int max_active_blocks_per_sm = -1;
-  AT_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+  STD_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(at::cuda::getCurrentDeviceProperties()->multiProcessorCount *
-                max_active_blocks_per_sm,
+  dim3 grid(get_device_prop()->multiProcessorCount * max_active_blocks_per_sm,
             1, 1);
   dim3 block(THREAD_BLOCK_SIZE, 1, 1);
-  int num_experts = (int)problem_sizes.size(0);
-  auto stream = at::cuda::getCurrentCUDAStream();
+  int num_experts = static_cast<int>(problem_sizes.size(0));
+  auto stream = get_current_cuda_stream(input.get_device_index());
   mxfp8_experts_quant_kernel<T_IN, decltype(tiled_copy_g2r),
                              decltype(tiled_copy_r2g), decltype(tiled_copy_r2s)>
       <<<grid, block, 0, stream>>>(

csrc/libtorch_stable/ops.h

@@ -3,10 +3,6 @@
 #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,
@@ -28,6 +24,10 @@ 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,6 +162,11 @@ 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)
@@ -215,6 +220,13 @@ 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,
@@ -231,6 +243,44 @@ 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,
@@ -273,6 +323,26 @@ 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 "cutlass_extensions/common.hpp"
+#include "libtorch_stable/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 "cutlass_extensions/common.hpp"
+#include "libtorch_stable/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 "cutlass_extensions/common.hpp"
+#include "libtorch_stable/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 "cutlass_extensions/common.hpp"
+#include "libtorch_stable/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 "cutlass_extensions/common.hpp"
+#include "libtorch_stable/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 "cutlass_extensions/common.hpp"
+#include "libtorch_stable/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 "cutlass_extensions/common.hpp"
+#include "libtorch_stable/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 "cutlass_extensions/common.hpp"
+#include "libtorch_stable/cutlass_extensions/common.hpp"
 
 #include "cutlass/cutlass.h"
 

csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu

@@ -1,11 +1,10 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
+#include "../../torch_utils.h"
 
 #include "../../dispatch_utils.h"
-#include "libtorch_stable/quantization/fused_kernels/quant_conversions.cuh"
+#include "quant_conversions.cuh"
 
 namespace vllm {
 
@@ -105,64 +104,70 @@ __global__ void silu_and_mul_per_block_quant_kernel(
 
 }  // namespace vllm
 
-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,
+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) {
-  static c10::ScalarType kFp8Type = is_fp8_ocp()
-                                        ? c10::ScalarType::Float8_e4m3fn
-                                        : c10::ScalarType::Float8_e4m3fnuz;
+  static torch::headeronly::ScalarType kFp8Type =
+      is_fp8_ocp() ? torch::headeronly::ScalarType::Float8_e4m3fn
+                   : torch::headeronly::ScalarType::Float8_e4m3fnuz;
 
-  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,
+  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,
       "Input must be FP16 or BF16");
-  TORCH_CHECK(scales.dtype() == torch::kFloat32, "Scales must be FP32");
-  TORCH_CHECK(group_size == 128 || group_size == 64,
-              "Unsupported group size: ", group_size);
+  STD_TORCH_CHECK(scales.scalar_type() == torch::headeronly::ScalarType::Float);
+  STD_TORCH_CHECK(group_size == 128 || group_size == 64,
+                  "Unsupported group size: ", group_size);
 
   if (scale_ub.has_value()) {
-    TORCH_CHECK(out.dtype() == kFp8Type);
+    STD_TORCH_CHECK(out.scalar_type() == kFp8Type);
   }
 
   int32_t hidden_size = out.size(-1);
   auto num_tokens = input.size(0);
   int32_t num_groups = hidden_size / 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");
+  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");
 
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
 
   dim3 grid(num_tokens, num_groups);
   dim3 block(group_size);
 
-  VLLM_DISPATCH_FLOATING_TYPES(
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "silu_and_mul_per_block_quant", [&] {
         using scalar_in_t = scalar_t;
 
-        VLLM_DISPATCH_QUANT_TYPES(
+        VLLM_STABLE_DISPATCH_QUANT_TYPES(
             out.scalar_type(), "silu_and_mul_per_block_quant", [&] {
               using scalar_out_t = scalar_t;
 
-              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);
-                });
+              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);
+                    });
               });
             });
       });
-}
+}

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 "cutlass_extensions/common.hpp"
+#include "libtorch_stable/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 "cutlass_extensions/common.hpp"
+#include "libtorch_stable/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 "cutlass_extensions/common.hpp"
+#include "libtorch_stable/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 "cutlass_extensions/common.hpp"
+#include "libtorch_stable/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 "cutlass_extensions/common.hpp"
+#include "libtorch_stable/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 "cutlass_extensions/common.hpp"
+#include "libtorch_stable/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,7 +7,11 @@
 
 #include <cmath>
 
-#include <cuda_fp8.h>
+#ifdef USE_ROCM
+  #include <hip/hip_fp8.h>
+#else
+  #include <cuda_fp8.h>
+#endif
 
 #include "libtorch_stable/quantization/vectorization.cuh"
 #include "libtorch_stable/quantization/vectorization_utils.cuh"
@@ -15,12 +19,23 @@
 #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;
 }
 
@@ -237,10 +252,12 @@ 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::Float8_e4m3fn) {
-          LAUNCH_KERNEL(scalar_t, __nv_fp8_e4m3);
-        } else if (dst_type == torch::headeronly::ScalarType::Char) {
+        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); }));
         }
       }));
 
@@ -317,10 +334,18 @@ __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);
@@ -503,15 +528,12 @@ 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::Float8_e4m3fn) {
-          LAUNCH_REG_KERNEL(scalar_t, __nv_fp8_e4m3);
-        } else if (dst_type == torch::headeronly::ScalarType::Char) {
+        if (dst_type == torch::headeronly::ScalarType::Char) {
           LAUNCH_REG_KERNEL(scalar_t, int8_t);
         } else {
-          STD_TORCH_CHECK(
-              false,
-              "per_token_group_quant_8bit_packed only supports FP8/INT8 "
-              "outputs.");
+          VLLM_STABLE_DISPATCH_FP8_TYPES(
+              dst_type, "per_token_group_quant_8bit_packed_fp8",
+              ([&] { LAUNCH_REG_KERNEL(scalar_t, fp8_t); }));
         }
       }));
 

csrc/libtorch_stable/torch_bindings.cpp

@@ -1,4 +1,5 @@
 #include "ops.h"
+#include "cuda_utils.h"
 #include "core/registration.h"
 
 #include <torch/csrc/stable/library.h>
@@ -7,11 +8,6 @@
 // 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(
@@ -31,7 +27,15 @@ 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(
@@ -321,6 +325,16 @@ 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(
@@ -337,6 +351,38 @@ 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, "
@@ -508,11 +554,6 @@ 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",
@@ -520,6 +561,9 @@ 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));
@@ -567,10 +611,24 @@ 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_",
@@ -617,6 +675,28 @@ 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
@@ -725,6 +805,45 @@ 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,11 +6,7 @@
 #include <torch/csrc/stable/tensor.h>
 #include <torch/headeronly/util/shim_utils.h>
 
-#ifndef USE_ROCM
-  #include <cuda_runtime.h>
-#else
-  #include <hip/hip_runtime.h>
-#endif
+#include <cuda_runtime.h>
 #include <cublas_v2.h>
 
 #include <deque>

csrc/minimax_reduce_rms_kernel.h

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

csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu

@@ -1,60 +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/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/mxfp8_experts_quant.cu

@@ -1,60 +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/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/ops.h

@@ -40,18 +40,6 @@ 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.
@@ -90,8 +78,6 @@ 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);
@@ -107,24 +93,6 @@ 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);
@@ -134,16 +102,7 @@ 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
 
-#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);
+// TODO: Remove this once ROCm upgrade to torch 2.11.
+torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
 #endif

csrc/rocm/moe_q_gemm_rdna3.cu

@@ -0,0 +1,639 @@
+// 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,6 +27,15 @@ 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,6 +50,15 @@ 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,37 +32,24 @@ 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.
-  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);
+  // kernel launch. Registered in _C_stable_libtorch (incl. the FlashInfer V4
+  // full-cache bf16/fp8 variants).
 
   // Quantization ops
 #ifndef USE_ROCM
@@ -163,37 +150,29 @@ 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);
@@ -204,49 +183,6 @@ 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<c10::Half> {
+struct _typeConvert<torch::headeronly::Half> {
   static constexpr bool exists = true;
   using hip_type = __half;
   using packed_hip_type = __half2;
@@ -73,7 +73,7 @@ struct _typeConvert<c10::Half> {
 // CUDA_ARCH < 800 does not have BF16 support
 // ROCm 7.0+ supports bfloat16
 template <>
-struct _typeConvert<c10::BFloat16> {
+struct _typeConvert<torch::headeronly::BFloat16> {
   static constexpr bool exists = true;
   using hip_type = __nv_bfloat16;
   using packed_hip_type = __nv_bfloat162;

docker/Dockerfile

@@ -218,6 +218,10 @@ 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; \
@@ -234,6 +238,13 @@ 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.
@@ -745,6 +756,10 @@ 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; \
@@ -752,15 +767,22 @@ 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.11.post2
+ARG FLASHINFER_VERSION=0.6.12
 RUN --mount=type=cache,target=/opt/uv/cache \
     uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
-        --extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
+        --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
 
 # ============================================================
 # OPENAI API SERVER DEPENDENCIES
@@ -842,6 +864,19 @@ 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.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.11.post2
+# release version: v0.6.12
 # 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.11.post2 --recursive https://github.com/flashinfer-ai/flashinfer.git \
+    && git clone --depth 1 --branch v0.6.12 --recursive https://github.com/flashinfer-ai/flashinfer.git \
     && cd flashinfer \
     && git submodule update --init --recursive \
     && echo "finish git clone flashinfer..." \

docker/Dockerfile.rocm

@@ -2,6 +2,7 @@
 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.
@@ -16,7 +17,8 @@ ARG NIC_BACKEND=all
 ARG AINIC_VERSION=1.117.3-hydra
 ARG UBUNTU_CODENAME=jammy
 
-# Sccache configuration (only used in release pipeline)
+# Sccache configuration. Release builds use this today; CI can opt in when a
+# shared S3-compatible cache backend is available.
 ARG USE_SCCACHE
 ARG SCCACHE_DOWNLOAD_URL
 ARG SCCACHE_ENDPOINT
@@ -29,12 +31,16 @@ FROM ${BASE_IMAGE} AS base
 ARG ARG_PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
 
-# Install some basic utilities
+# Install build dependencies and 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
-RUN python3 -m pip install --upgrade pip
+    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".
 # 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 \
@@ -55,6 +61,12 @@ 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
@@ -114,8 +126,7 @@ 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. Runs in parallel with the main wheel
-# build for faster end-to-end builds.
+# don't need the rust toolchain or protoc.
 FROM fetch_vllm AS rust-build
 ARG COMMON_WORKDIR
 
@@ -144,24 +155,74 @@ 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 build stages
+# 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.
 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
 
-# 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
+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
 FROM scratch AS export_vllm
 ARG COMMON_WORKDIR
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/dist/*.whl /
@@ -169,8 +230,10 @@ 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
@@ -201,14 +264,17 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
     ibverbs-providers \
     && rm -rf /var/lib/apt/lists/*
 
-RUN uv pip install --system meson auditwheel patchelf tomlkit
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system meson auditwheel patchelf tomlkit
 
-RUN cd /usr/local/src && \
+RUN --mount=type=cache,target=/root/.cache/ccache \
+    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 \
@@ -220,20 +286,22 @@ RUN cd /usr/local/src && \
         --with-verbs \
         --with-dm \
         --enable-mt && \
-    make -j && \
+    make -j$(nproc) && \
     make install
 
 ENV PATH=/usr/local/ucx/bin:$PATH
 ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
 
-RUN git clone ${RIXL_REPO} /opt/rixl && \
+RUN --mount=type=cache,target=/root/.cache/ccache \
+    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 && \
+    ninja -j$(nproc) && \
     ninja install
 
 # Generate RIXL wheel
@@ -250,30 +318,44 @@ RUN cd /opt/rixl && \
         --ucx-plugins-dir ${UCX_HOME}/lib/ucx \
         --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
 
-# DeepEP build stage
-FROM base AS build_deep
+# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not
+# invalidate the slow ROCShmem build.
+FROM base AS build_rocshmem
 ARG ROCSHMEM_BRANCH="f0acb0c6"
 ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
-ARG DEEPEP_BRANCH="a9ea9774"
-ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
-ARG DEEPEP_NIC="cx7"
+# 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 git clone ${ROCSHMEM_REPO} \
+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 \
- && INSTALL_PREFIX=${ROCSHMEM_DIR} \
-    ../scripts/build_configs/all_backends -DUSE_EXTERNAL_MPI=OFF
+ && 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
 
-# Build DeepEP wheel.
-# DeepEP looks for rocshmem at ROCSHMEM_DIR.
-RUN git clone ${DEEPEP_REPO} \
+# 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"
+
+# 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} \
  && cd DeepEP \
  && git checkout ${DEEPEP_BRANCH} \
- && python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
+ && 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
 
 # MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do
 # not force users to rebuild the long-lived Dockerfile.rocm_base image.
@@ -372,8 +454,9 @@ 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 \
-    && pip install setuptools_scm regex \
+    && uv pip install --system 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
@@ -409,18 +492,20 @@ 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 cd vllm \
+RUN --mount=type=cache,target=/root/.cache/uv \
+    cd vllm \
     && echo "Building vLLM with custom wheels from /install" \
-    && python3 -m pip install --find-links /install -r requirements/rocm.txt \
-    && python3 setup.py clean --all
+    && uv pip install --system --find-links /install -r requirements/rocm.txt
 
 # Build wheel using pre-extracted version to avoid dirty state from modified requirements/rocm.txt
-# (setup.py auto-detects sccache in PATH)
+# (setup.py auto-detects ccache/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}" \
-    && python3 setup.py bdist_wheel --dist-dir=dist
+    && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist
 
 FROM scratch AS export_vllm_wheel_release
 ARG COMMON_WORKDIR
@@ -429,114 +514,121 @@ 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
 
 # -----------------------
-# Test vLLM image
-FROM mori_base AS test
+# 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
 
-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
+# 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 wheel
+# Install RIXL + DeepEP wheels.
 RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
-    uv pip install --system /rixl_install/*.whl
+    --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
+    uv pip install --system /rixl_install/*.whl /deep_install/*.whl
 
-# 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
+# Copy ROCShmem runtime libraries.
+COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem
 
-# RIXL/MoRIIO runtime dependencies (RDMA userspace libraries)
-RUN apt-get update -q -y && apt-get install -q -y \
+# RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec.
+RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
     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/*
 
-WORKDIR /vllm-workspace
-ARG COMMON_WORKDIR
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace
-
-# 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 downloads from hf (for testing)
-ENV HF_XET_HIGH_PERFORMANCE=1
-
-# increase timeout for hf downloads (for testing)
-ENV HF_HUB_DOWNLOAD_TIMEOUT 60
-
-# install audio decode package `torchcodec` from source (required due to 
-# ROCm and torch version mismatch) for tests with datasets package
+# Install torchcodec from source for ROCm/torch ABI compatibility.
 COPY tools/install_torchcodec_rocm.sh /tmp/install_torchcodec.sh
-RUN bash /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/*
+    && apt-get clean && rm -rf /var/lib/apt/lists/*
 
-# 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
+# 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
 
-# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel
+# Enable fast and less brittle model downloads in tests.
+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
+
+# 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
+
+# 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
 
-# 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
+# ROCm profiler limits workaround.
+RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
+ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
 
-# 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
+# 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.
+RUN mkdir src && mv vllm src/vllm
 
 # -----------------------
 # Final vLLM image
@@ -553,6 +645,7 @@ 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

@@ -0,0 +1,376 @@
+# 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

@@ -0,0 +1,143 @@
+# 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.11.post2"
+      "default": "0.6.12"
     },
     "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`, …). Modes that don't use the HF
+`deepseek_v32`, `deepseek_v4`, `qwen_vl`, …). Models 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.10
+    pre-commit run --hook-stage manual mypy-3.11
     ```
 
 ### Documentation

docs/contributing/profiling.md

@@ -35,8 +35,7 @@ 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.
-    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`
+    The engine client waits for this flush to complete without timing out, so simply allow the stop call to run to completion.
 
 ### Example commands and usage
 

docs/design/attention_backends.md

@@ -228,3 +228,17 @@ 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,6 +17,7 @@ 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
@@ -30,8 +31,7 @@ class BudgetGraphMetadata:
     max_batch_size: int
     max_frames_per_batch: int
     graph: torch.cuda.CUDAGraph
-    input_buffer: torch.Tensor       # e.g. pixel_values
-    metadata_buffers: dict[str, torch.Tensor]  # e.g. embeddings, seq metadata
+    input_buffers: dict[str, torch.Tensor]  # e.g. pixel_values, 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. 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).
+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.
 3. Replay the CUDA Graph.
 4. Clone outputs from `output_buffer` (cloning is necessary since the buffer is reused across replays).
 
@@ -65,19 +65,15 @@ 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, input key, buffer keys, output hidden size).
+* `get_encoder_cudagraph_config()` — returns static configuration (supported modalities, buffer keys, output hidden size, padding logics, max frames per video).
 * `get_encoder_cudagraph_budget_range(vllm_config)` — returns `(min_budget, max_budget)` for auto-inference of token budgets.
-* `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.
+* `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`).
 * `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.
-* `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
+* `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.
 
 !!! 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.
@@ -86,6 +82,7 @@ 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` | ✅︎ | ✅︎ |
@@ -103,7 +100,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 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).
+* `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).
 
 ## 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_both",
+    "kv_role": "kv_producer",
     "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_TORCH_COMPILE_LEVEL=0 vllm serve <model> --compilation-config='{"compile_mm_encoder":"false"}'
+   vllm serve <model> --compilation-config='{"mode":0,"compile_mm_encoder":"false"}'
    ```
 
 2. **Check logs**: Enable debug logging to see compilation details:

docs/features/batch_invariance.md

@@ -17,10 +17,7 @@ Batch invariance is crucial for several use cases:
 
 ## Hardware Requirements
 
-Batch invariance currently requires NVIDIA GPUs with compute capability 9.0 or higher:
-
-- **H-series**: H100, H200
-- **B-series**: B100, B200
+Batch invariance requires NVIDIA GPUs with compute capability 8.0 or higher.
 
 ## Enabling Batch Invariance
 
@@ -107,7 +104,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**: `meta-llama/Llama-3.1-8B-Instruct`, `meta-llama/Llama-3.2-1B-Instruct`
+- **Llama 3**: Llama3.1 and 3.2 series, `meta-llama/Llama-3.2-3B-Instruct` for example
 - **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_both",
+    "kv_role":"kv_producer",
     "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_both \
+  --kv-transfer-config.kv_role kv_producer \
   --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_both","kv_load_failure_policy":"fail"}'
+  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_producer","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_both","kv_load_failure_policy":"fail"}'
+  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer","kv_load_failure_policy":"fail"}'
 ```
 
 ### Proxy Server
@@ -212,10 +212,21 @@ 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_both",
+    "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_connector_extra_config": {
       "bidirectional_kv_xfer": true
     }
@@ -359,11 +370,10 @@ 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**: 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.
+- **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.
 
-!!! 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.
+!!! 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.
 
 ### KV Load Failure Policy
 

docs/features/quantization/README.md

@@ -3,7 +3,7 @@
 Quantization trades off model precision for smaller memory footprint, allowing large models to be run on a wider range of devices.
 
 !!! tip
-    To get started with quantization, see [LLM Compressor](llm_compressor.md), a library for optimizing models for deployment with vLLM that supports FP8, INT8, INT4, and other quantization formats.
+    To get started with quantization, see [LLM Compressor](llm_compressor/README.md), a library for optimizing models for deployment with vLLM that supports FP8, INT8, INT4, and other quantization formats.
 
 The following are the supported quantization formats for vLLM:
 
@@ -12,9 +12,11 @@ The following are the supported quantization formats for vLLM:
 - [GGUF](gguf.md)
 - [GPTQModel](gptqmodel.md)
 - [Intel Neural Compressor](inc.md)
-- [INT4 W4A16](int4.md)
-- [INT8 W8A8](int8.md)
-- [FP8 W8A8](fp8.md)
+- [LLM Compressor](llm_compressor/README.md)
+    - [FP8 W8A8](llm_compressor/fp8.md)
+    - [INT4 W4A16](llm_compressor/int4.md)
+    - [INT8 W4A8](llm_compressor/int8_w4a8.md)
+    - [INT8 W8A8](llm_compressor/int8_w8a8.md)
 - [NVIDIA Model Optimizer](modelopt.md)
 - [Online Quantization](online.md)
 - [AMD Quark](quark.md)
@@ -46,16 +48,17 @@ th:not(:first-child) {
 }
 </style>
 
-| Implementation            | Volta | Turing | Ampere | Ada | Hopper | AMD GPU | Intel GPU | x86 CPU |
-| ------------------------- | ----- | ------ | ------ | --- | ------ | ------- | --------- | ------- |
-| AWQ                       | ❌    | ✅︎     | ✅︎     | ✅︎  | ✅︎     | ❌      | ✅︎        | ✅︎      |
-| GPTQ                      | ✅︎    | ✅︎     | ✅︎     | ✅︎  | ✅︎     | ❌      | ✅︎        | ✅︎      |
-| Marlin (GPTQ/AWQ/FP8/FP4) | ❌    | ✅︎*    | ✅︎     | ✅︎  | ✅︎     | ❌      | ❌        | ❌      |
-| INT8 (W8A8)               | ❌    | ✅︎     | ✅︎     | ✅︎  | ✅︎     | ❌      | ❌        | ✅︎      |
-| FP8 (W8A8)                | ❌    | ❌     | ❌     | ✅︎  | ✅︎     | ✅︎      | ❌        | ❌      |
-| bitsandbytes              | ✅︎    | ✅︎     | ✅︎     | ✅︎  | ✅︎     | ❌      | ❌        | ❌      |
-| DeepSpeedFP               | ✅︎    | ✅︎     | ✅︎     | ✅︎  | ✅︎     | ❌      | ❌        | ❌      |
-| GGUF                      | ✅︎    | ✅︎     | ✅︎     | ✅︎  | ✅︎     | ✅︎      | ❌        | ❌      |
+| Implementation            | Volta | Turing | Ampere | Ada | Hopper | AMD GPU | Intel GPU | x86 CPU | Arm CPU |
+| ------------------------- | ----- | ------ | ------ | --- | ------ | ------- | --------- | ------- | ------- |
+| AWQ                       | ❌    | ✅︎     | ✅︎     | ✅︎  | ✅︎     | ❌      | ✅︎        | ✅︎      | ❌      |
+| GPTQ                      | ✅︎    | ✅︎     | ✅︎     | ✅︎  | ✅︎     | ❌      | ✅︎        | ✅︎      | ❌      |
+| Marlin (GPTQ/AWQ/FP8/FP4) | ❌    | ✅︎*    | ✅︎     | ✅︎  | ✅︎     | ❌      | ❌        | ❌      | ❌      |
+| llm-compressor INT8 (W8A8)| ❌    | ✅︎     | ✅︎     | ✅︎  | ✅︎     | ❌      | ❌        | ✅︎      | ✅︎      |
+| llm-compressor INT8 (W4A8)| ❌    | ❌     | ❌     | ❌  | ❌     | ❌      | ❌        | ❌      | ✅︎      |
+| llm-compressor FP8 (W8A8) | ❌    | ❌     | ❌     | ✅︎  | ✅︎     | ✅︎      | ❌        | ❌      | ❌      |
+| bitsandbytes              | ✅︎    | ✅︎     | ✅︎     | ✅︎  | ✅︎     | ❌      | ❌        | ❌      | ❌      |
+| DeepSpeedFP               | ✅︎    | ✅︎     | ✅︎     | ✅︎  | ✅︎     | ❌      | ❌        | ❌      | ❌      |
+| GGUF                      | ✅︎    | ✅︎     | ✅︎     | ✅︎  | ✅︎     | ✅︎      | ❌        | ❌      | ❌      |
 
 - Volta refers to SM 7.0, Turing to SM 7.5, Ampere to SM 8.0/8.6, Ada to SM 8.9, and Hopper to SM 9.0.
 - ✅︎ indicates that the quantization method is supported on the specified hardware.

docs/features/quantization/llm_compressor/fp8.md

@@ -21,9 +21,17 @@ The FP8 types typically supported in hardware have two distinct representations,
 To produce performant FP8 quantized models with vLLM, you'll need to install the [llm-compressor](https://github.com/vllm-project/llm-compressor/) library:
 
 ```bash
-pip install llmcompressor
+(venv-llm-compressor) pip install llmcompressor
 ```
 
+Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:
+
+```bash
+(venv-vllm) pip install vllm "lm-eval[api]>=0.4.12"
+```
+
+Please use separate environments for vLLM and llm-compressor as they might not work together.
+
 ## Quantization Process
 
 The quantization process involves three main steps:
@@ -57,36 +65,28 @@ For FP8 quantization, we can recover accuracy with simple RTN quantization. We r
 
 Since simple RTN does not require data for weight quantization and the activations are quantized dynamically, we do not need any calibration data for this quantization flow.
 
-??? code
+```python
+from llmcompressor import oneshot
+from llmcompressor.modifiers.quantization import QuantizationModifier
 
-    ```python
-    from llmcompressor import oneshot
-    from llmcompressor.modifiers.quantization import QuantizationModifier
+# Configure the simple PTQ quantization
+recipe = QuantizationModifier(
+    targets="Linear",
+    scheme="FP8_DYNAMIC",
+    ignore=["lm_head"],
+)
 
-    # Configure the simple PTQ quantization
-    recipe = QuantizationModifier(
-        targets="Linear",
-        scheme="FP8_DYNAMIC",
-        ignore=["lm_head"],
-    )
+# Apply the quantization algorithm.
+oneshot(model=model, recipe=recipe)
 
-    # Apply the quantization algorithm.
-    oneshot(model=model, recipe=recipe)
-
-    # Save the model: Meta-Llama-3-8B-Instruct-FP8-Dynamic
-    SAVE_DIR = MODEL_ID.split("/")[1] + "-FP8-Dynamic"
-    model.save_pretrained(SAVE_DIR)
-    tokenizer.save_pretrained(SAVE_DIR)
-    ```
+# Save the model: Meta-Llama-3-8B-Instruct-FP8-Dynamic
+SAVE_DIR = MODEL_ID.split("/")[1] + "-FP8-Dynamic"
+model.save_pretrained(SAVE_DIR)
+tokenizer.save_pretrained(SAVE_DIR)
+```
 
 ### 3. Evaluating Accuracy
 
-Install `vllm` and `lm-evaluation-harness` for evaluation:
-
-```bash
-pip install vllm "lm-eval[api]>=0.4.12"
-```
-
 Load and run the model in `vllm`:
 
 ```python