[Distributed] Add UCX one-shot AllReduce for DP metadata sync

Replace Gloo TCP AllReduce with UCX tag-matching over IB RDMA for the
per-iteration DP metadata synchronization in data-parallel mode.

The one-shot allreduce sends each rank's full tensor to every other rank
and locally sums, staying well within the RDMA eager threshold (no
rendezvous handshake). This reduces P99 DP sync latency from ~100ms
(Gloo/TCP) to ~0.1ms (UCX/RDMA) on InfiniBand clusters.

Enabled by setting VLLM_DP_SYNC_BACKEND=ucx. Falls back to Gloo
automatically if UCX is unavailable or initialization fails.

New files:
- ucx_dp_communicator.py: Python wrapper with auto-compile of C library
- ucx_dp_sync.c: C-level UCX one-shot allreduce primitive
- test_ucx_allreduce.py: Latency benchmark (Gloo vs UCX)
- test_ucx_correctness.py: Multi-pattern correctness tests

Modified files:
- vllm/v1/worker/dp_utils.py: UCX path for V1 model runner
- vllm/v1/worker/gpu/dp_utils.py: UCX path for V2 model runner

Co-authored-by: Claude <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

.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/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

@@ -52,6 +52,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 +345,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"

.buildkite/test-amd.yaml

@@ -449,29 +449,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
@@ -1739,6 +1716,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

.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

.dockerignore

@@ -33,3 +33,10 @@ share/python-wheels/
 *.egg
 MANIFEST
 rust/target/
+# Not needed in Docker builds
+docs/
+.github/
+.pre-commit-config.yaml
+.clang-format
+.gitattributes
+format.sh

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 "")
@@ -309,14 +311,9 @@ set(VLLM_EXT_SRC
   "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.
@@ -503,12 +500,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,7 +595,7 @@ 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).
   list(APPEND VLLM_EXT_SRC
     "csrc/custom_quickreduce.cu"
   )
@@ -633,6 +630,8 @@ 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/quantization/gptq/q_gemm.cu"
     "csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu"
     "csrc/libtorch_stable/pos_encoding_kernels.cu"
@@ -647,7 +646,10 @@ 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
@@ -659,7 +661,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
       "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}"

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/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/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
@@ -564,7 +569,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 +640,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,
@@ -650,71 +655,80 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
 // ────────────────────────────────────────────────────────────────────────────
 // 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);

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);
 
+#ifndef USE_ROCM
+torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
+                                   torch::stable::Tensor const& perm);
+
 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);
@@ -231,6 +231,27 @@ 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);
+
+#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 +294,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/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

@@ -7,11 +7,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(
@@ -32,6 +27,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "output_s, int group_size, float eps, float int8_min, float int8_max) -> "
       "()");
 
+#ifndef USE_ROCM
+  ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
+#endif
+
+#ifndef USE_ROCM
   // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column
   // quantization, as well as bias
   ops.def(
@@ -337,6 +337,24 @@ 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");
+
+#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 +526,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 +533,11 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("per_token_group_quant_int8",
            TORCH_BOX(&per_token_group_quant_int8));
 
+#ifndef USE_ROCM
+  ops.impl("permute_cols", TORCH_BOX(&permute_cols));
+#endif
+
+#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));
@@ -571,6 +589,12 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   // 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));
+#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_",
@@ -725,6 +749,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,12 +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,
@@ -107,24 +101,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);
@@ -135,15 +111,3 @@ 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);
-#endif

csrc/torch_bindings.cpp

@@ -55,14 +55,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 
   // 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.
 
   // Quantization ops
 #ifndef USE_ROCM
@@ -163,34 +156,27 @@ 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);
+}
+#endif
+
 TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
   // Cuda utils
 
@@ -205,48 +191,4 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
                   &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

@@ -757,10 +757,10 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 # 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

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
@@ -86,6 +98,7 @@ RUN if [ "$USE_SCCACHE" = "1" ]; then \
 ARG USE_SCCACHE
 ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}}
 ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
+ENV SCCACHE_ENDPOINT=${USE_SCCACHE:+${SCCACHE_ENDPOINT}}
 ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
 ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
 
@@ -114,8 +127,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 +156,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 /
@@ -171,6 +233,7 @@ 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/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
@@ -431,112 +516,118 @@ 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/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 +644,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/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/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.
@@ -103,7 +99,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/training/weight_transfer/base.md

@@ -156,5 +156,6 @@ from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
 engine = WeightTransferEngineFactory.create_engine(
     config=weight_transfer_config,
     parallel_config=parallel_config,
+    model=model,
 )
 ```

requirements/cuda.txt

@@ -9,8 +9,8 @@ torchaudio==2.11.0
 # These must be updated alongside torch
 torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
 # FlashInfer should be updated together with the Dockerfile
-flashinfer-python==0.6.11.post2
-flashinfer-cubin==0.6.11.post2
+flashinfer-python==0.6.12
+flashinfer-cubin==0.6.12
 apache-tvm-ffi==0.1.9
 tilelang==0.1.9
 # Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to

rust/Cargo.lock

@@ -5622,6 +5622,7 @@ dependencies = [
  "expect-test",
  "futures",
  "half",
+ "indexmap 2.13.0",
  "itertools 0.14.0",
  "llm-multimodal",
  "minijinja",

rust/Cargo.toml

@@ -43,6 +43,7 @@ half = { version = "2.7.1", features = ["bytemuck"] }
 hex = "0.4.3"
 hf-hub = { version = "0.5.0", features = ["tokio"] }
 http-body = "1.0.1"
+indexmap = "2.13.0"
 itertools = "0.14.0"
 libc = "0.2.177"
 llm-multimodal = { git = "https://github.com/vllm-project/llm-multimodal", rev = "5b558989844d1c7af3e43d0f604069ffd9c06320" }
@@ -69,7 +70,7 @@ rustc-hash = "1.1.0"
 serde = { version = "1.0.228", features = ["derive"] }
 serde-json-fmt = "0.1.0"
 serde_default = "0.2.0"
-serde_json = { version = "1.0.145", features = ["arbitrary_precision", "preserve_order"] }
+serde_json = { version = "1.0.145", features = ["preserve_order"] }
 serde_repr = "0.1.20"
 serde_tuple = "1.1.3"
 serde_with = "3.18.0"

rust/src/chat/Cargo.toml

@@ -10,6 +10,7 @@ asynk-strim-attr.workspace = true
 easy-ext.workspace = true
 futures.workspace = true
 half.workspace = true
+indexmap.workspace = true
 itertools.workspace = true
 llm-multimodal.workspace = true
 minijinja.workspace = true

rust/src/chat/src/lib.rs

@@ -189,6 +189,7 @@ impl ChatLlm {
             cache_salt: request.cache_salt,
             add_special_tokens: request.add_special_tokens,
             data_parallel_rank: request.data_parallel_rank,
+            lora_request: request.lora_request,
         };
         let decoded_stream = self.text.generate(text_request).await?.map_err(Error::from).boxed();
 

rust/src/chat/src/parser/mod.rs

@@ -6,10 +6,10 @@ use std::convert::Infallible;
 use std::fmt;
 use std::str::FromStr;
 
-use serde_with::DeserializeFromStr;
+use serde_with::{DeserializeFromStr, SerializeDisplay};
 
 /// Specify which reasoning or tool-call parser implementation to use.
-#[derive(Debug, Clone, PartialEq, Eq, Default, DeserializeFromStr)]
+#[derive(Debug, Clone, PartialEq, Eq, Default, DeserializeFromStr, SerializeDisplay)]
 pub enum ParserSelection {
     /// Use model-based auto-detection.
     #[default]

rust/src/chat/src/renderer/hf/format.rs

@@ -5,7 +5,7 @@ use std::str::FromStr;
 use minijinja::machinery::ast::{Expr, ForLoop, Set, Stmt};
 use minijinja::machinery::{WhitespaceConfig, parse};
 use minijinja::syntax::SyntaxConfig;
-use serde_with::DeserializeFromStr;
+use serde_with::{DeserializeFromStr, SerializeDisplay};
 
 /// Chat template content format.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
@@ -18,7 +18,7 @@ pub enum ChatTemplateContentFormat {
 }
 
 /// Configurable chat-template content format selection.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, DeserializeFromStr)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, DeserializeFromStr, SerializeDisplay)]
 pub enum ChatTemplateContentFormatOption {
     /// Detect the format from the template source.
     #[default]

rust/src/chat/src/renderer/hf/mod.rs

@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 
 use serde::Serialize;
-use serde_json::Value;
+use serde_json::Value as JsonValue;
 use thiserror_ext::AsReport as _;
 use tracing::{info, trace, warn};
 use vllm_text::Prompt;
@@ -13,6 +13,7 @@ use self::format::{
     ChatTemplateContentFormat, ChatTemplateContentFormatOption as ContentFormatOption,
 };
 use self::template::{CompiledChatTemplate, TemplateContext};
+use self::value::{TemplateValue, to_template_value};
 use super::{ChatRenderer, RenderedPrompt};
 use crate::error::Result;
 use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest};
@@ -24,6 +25,7 @@ mod error;
 mod format;
 mod template;
 mod tojson;
+mod value;
 
 pub use template::{load_chat_template, resolve_chat_template};
 
@@ -38,7 +40,7 @@ pub struct MultimodalRenderInfo {
 /// state.
 pub struct HfChatRenderer {
     default_template: Option<CompiledChatTemplate>,
-    default_template_kwargs: HashMap<String, Value>,
+    default_template_kwargs: HashMap<String, JsonValue>,
     content_format: ContentFormatOption,
     special_tokens: Option<HfSpecialTokens>,
     multimodal: Option<MultimodalRenderInfo>,
@@ -48,7 +50,7 @@ impl HfChatRenderer {
     /// Create a renderer from the given template string.
     pub fn new(
         template: Option<String>,
-        default_template_kwargs: HashMap<String, Value>,
+        default_template_kwargs: HashMap<String, JsonValue>,
         content_format: ContentFormatOption,
     ) -> Result<Self> {
         Ok(Self {
@@ -245,7 +247,7 @@ struct TemplateToolCall {
 #[derive(Debug, Serialize)]
 struct TemplateToolFunction {
     name: String,
-    arguments: Value,
+    arguments: TemplateValue,
 }
 
 #[derive(Debug, Serialize)]
@@ -259,7 +261,7 @@ pub(super) struct TemplateTool {
 struct TemplateToolDefinition {
     name: String,
     description: Option<String>,
-    parameters: Value,
+    parameters: TemplateValue,
     strict: Option<bool>,
 }
 
@@ -345,13 +347,14 @@ fn to_template_tool_calls(
     let mut tool_calls = Vec::new();
 
     for tool_call in content.tool_calls() {
-        let arguments = serde_json::from_str::<Value>(&tool_call.arguments).map_err(|error| {
+        let arguments = serde_json::from_str(&tool_call.arguments).map_err(|error| {
             Error::ChatTemplate(format!(
                 "assistant tool call `{}` has invalid JSON arguments: {}",
                 tool_call.id,
                 error.as_report()
             ))
         })?;
+        let arguments = to_template_value(arguments);
 
         tool_calls.push(TemplateToolCall {
             id: tool_call.id.clone(),
@@ -434,7 +437,7 @@ fn to_template_tools(tools: &[ChatTool]) -> Vec<TemplateTool> {
             function: TemplateToolDefinition {
                 name: tool.name.clone(),
                 description: tool.description.clone(),
-                parameters: tool.parameters.clone(),
+                parameters: to_template_value(tool.parameters.clone()),
                 strict: tool.strict,
             },
         })
@@ -909,6 +912,29 @@ mod tests {
         assert_eq!(rendered, "get_weather|Paris|call_1|Sunny");
     }
 
+    #[test]
+    fn chat_template_tool_call_argument_items_method_is_not_shadowed_by_field() {
+        let request = sample_request(vec![ChatMessage::assistant_blocks(vec![
+            AssistantContentBlock::ToolCall(crate::AssistantToolCall {
+                id: "call_1".to_string(),
+                name: "add".to_string(),
+                arguments: r#"{"items":"operands","x":2,"y":1.0}"#.to_string(),
+            }),
+        ])]);
+
+        let rendered = render(
+            Some(
+                "{%- set arguments = messages[0].tool_calls[0].function.arguments -%}
+{%- for key, value in arguments.items() -%}{{ key }}={{ value }};{%- endfor -%}
+|{{ arguments['items'] }}",
+            ),
+            &request,
+        )
+        .unwrap();
+
+        assert_eq!(rendered, "items=operands;x=2;y=1.0;|operands");
+    }
+
     #[test]
     fn qwen35_template_renders_prefilled_reasoning_start_when_thinking_enabled() {
         let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);

rust/src/chat/src/renderer/hf/tojson.rs

@@ -208,11 +208,27 @@ mod tests {
     }
 
     #[test]
-    fn tojson_preserves_arbitrary_precision_number_spelling() {
+    fn tojson_uses_standard_serde_json_number_spelling() {
         let payload = serde_json::from_str(r#"{"x":2,"y":1.00}"#).unwrap();
         let rendered = render("{{ payload|tojson }}", payload);
 
-        assert_eq!(rendered, "{\"x\": 2, \"y\": 1.00}");
+        // TODO: we cannot preserve the original number precision by enabling `serde_json`'s
+        // `arbitrary_precision` feature, otherwise the following test
+        // `serialized_json_numbers_do_not_leak_serde_private_representation` will fail.
+        // See issue: https://github.com/mitsuhiko/minijinja/issues/641
+        assert_eq!(rendered, "{\"x\": 2, \"y\": 1.0}");
+    }
+
+    #[test]
+    fn serialized_json_numbers_do_not_leak_serde_private_representation() {
+        let payload: serde_json::Value = serde_json::from_str(r#"{"x":2,"y":1.00}"#).unwrap();
+        let rendered = render("{{ payload }}", payload);
+
+        // TODO: we cannot preserve the original number precision by enabling `serde_json`'s
+        // `arbitrary_precision` feature, otherwise this will fail.
+        // See issue: https://github.com/mitsuhiko/minijinja/issues/641
+        assert!(!rendered.contains("$serde_json::private::Number"));
+        assert_eq!(rendered, r#"{"x": 2, "y": 1.0}"#);
     }
 
     #[test]

rust/src/chat/src/renderer/hf/value.rs

@@ -0,0 +1,77 @@
+use std::sync::Arc;
+
+use indexmap::IndexMap;
+use minijinja::value::{Enumerator, Object, ObjectExt, ObjectRepr};
+use minijinja::{Error as TemplateError, ErrorKind as TemplateErrorKind, State};
+use serde::Serialize;
+use serde_json::Value as JsonValue;
+
+/// A wrapper around `minijinja::Value` that can be constructed with `to_template_value` and used
+/// as a value in the chat template.
+#[derive(Debug, Serialize)]
+#[serde(transparent)]
+pub(super) struct TemplateValue(minijinja::Value);
+
+pub(super) fn to_template_value(value: JsonValue) -> TemplateValue {
+    TemplateValue(match value {
+        JsonValue::Array(values) => values
+            .into_iter()
+            .map(to_template_value)
+            .map(|value| value.0)
+            .collect::<minijinja::Value>(),
+        JsonValue::Object(values) => minijinja::Value::from_object(TemplateMap(
+            values
+                .into_iter()
+                .map(|(key, value)| (key, to_template_value(value).0))
+                .collect(),
+        )),
+        // For primitive values, directly convert them to `minijinja::Value` using `from_serialize`.
+        value => minijinja::Value::from_serialize(value),
+    })
+}
+
+/// A custom map type that always returns `UnknownMethod` for method calls, so that pycompat can
+/// always handle dict methods through the unknown-method callback.
+///
+/// Use `IndexMap` to preserve the original key order when iterating.
+///
+/// MiniJinja's default map can resolve a same-named field before Python dict methods. HF templates
+/// commonly call `dict.items()`, which would fail if the map had an `items` field.
+/// See issue: https://github.com/mitsuhiko/minijinja/issues/903
+#[derive(Debug)]
+struct TemplateMap(IndexMap<String, minijinja::Value>);
+
+impl Object for TemplateMap {
+    fn repr(self: &Arc<Self>) -> ObjectRepr {
+        ObjectRepr::Map
+    }
+
+    fn get_value(self: &Arc<Self>, key: &minijinja::Value) -> Option<minijinja::Value> {
+        self.0.get(key.as_str()?).cloned()
+    }
+
+    fn get_value_by_str(self: &Arc<Self>, key: &str) -> Option<minijinja::Value> {
+        self.0.get(key).cloned()
+    }
+
+    fn enumerate(self: &Arc<Self>) -> Enumerator {
+        self.mapped_rev_enumerator(|this| {
+            Box::new(this.0.keys().map(|key| minijinja::Value::from(key.as_str())))
+        })
+    }
+
+    fn enumerator_len(self: &Arc<Self>) -> Option<usize> {
+        Some(self.0.len())
+    }
+
+    fn call_method(
+        self: &Arc<Self>,
+        _state: &State<'_, '_>,
+        _method: &str,
+        _args: &[minijinja::Value],
+    ) -> std::result::Result<minijinja::Value, TemplateError> {
+        // Always return `UnknownMethod` for method calls,
+        // so that pycompat can handle dict methods through the unknown-method callback.
+        Err(TemplateError::from(TemplateErrorKind::UnknownMethod))
+    }
+}

rust/src/chat/src/renderer/selection.rs

@@ -1,10 +1,10 @@
 use std::fmt;
 use std::str::FromStr;
 
-use serde_with::DeserializeFromStr;
+use serde_with::{DeserializeFromStr, SerializeDisplay};
 
 /// Specify which chat renderer implementation to use.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, DeserializeFromStr)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, DeserializeFromStr, SerializeDisplay)]
 pub enum RendererSelection {
     /// Use model-based auto-detection.
     #[default]

rust/src/chat/src/request.rs

@@ -3,6 +3,7 @@ use std::collections::HashMap;
 use llm_multimodal::ImageDetail;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
+use vllm_engine_core_client::protocol::lora::LoraRequest;
 pub use vllm_text::SamplingParams;
 use vllm_text::TextDecodeOptions;
 pub use vllm_tool_parser::Tool as ChatTool;
@@ -426,6 +427,9 @@ pub struct ChatRequest {
     /// Override data parallel rank.
     #[serde(default)]
     pub data_parallel_rank: Option<u32>,
+    /// LoRA adapter selected for this request.
+    #[serde(default)]
+    pub lora_request: Option<LoraRequest>,
 }
 
 impl ChatRequest {
@@ -445,6 +449,7 @@ impl ChatRequest {
             cache_salt: None,
             add_special_tokens: false,
             data_parallel_rank: None,
+            lora_request: None,
         }
     }
 

rust/src/chat/tests/roundtrip.rs

@@ -233,7 +233,7 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
         "roundtrip-reasoning-tools",
         vec![ChatMessage::text(
             ChatRole::User,
-            "Check Shanghai weather and add 1.00 plus 2.",
+            "Check Shanghai weather and add 1.0 plus 2.",
         )],
         test_tools(),
         Some(true), // always enable thinking in this fixture
@@ -261,9 +261,10 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
                 AssistantContentBlock::ToolCall(AssistantToolCall {
                     id: "functions.add:1".to_string(),
                     name: "add".to_string(),
-                    // Intentionally use a non-lexical order of keys and a different number
-                    // formatting style to verify text-level fidelity of the roundtrip.
-                    arguments: r#"{"y":1.00,"x":2}"#.to_string(),
+                    // Intentionally use a non-lexical order of keys to verify text-level
+                    // fidelity of the roundtrip where JSON formatting remains stable. The
+                    // `items` key also exercises templates that call `arguments.items()`.
+                    arguments: r#"{"y":1.0,"x":2,"items":["left","right"]}"#.to_string(),
                 }),
             ],
         },
@@ -291,7 +292,7 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
     assert_eq!(tool_calls[1].name, "add");
     assert_eq!(
         tool_calls[1].arguments,
-        expected_arguments(&case, r#"{"y": 1.00, "x": 2}"#)?,
+        expected_arguments(&case, r#"{"y": 1.0, "x": 2, "items": ["left", "right"]}"#)?,
     );
 
     assert_eq!(
@@ -585,9 +586,13 @@ fn test_tools() -> Vec<ChatTool> {
                 "type": "object",
                 "properties": {
                     "y": { "type": "number" },
-                    "x": { "type": "number" }
+                    "x": { "type": "number" },
+                    "items": {
+                        "type": "array",
+                        "items": { "type": "string" }
+                    }
                 },
-                "required": ["y", "x"]
+                "required": ["y", "x", "items"]
             }),
             strict: None,
         },

rust/src/cmd/src/cli/tests.rs

@@ -87,6 +87,17 @@ fn serve_args_auto_forward_python_flags_without_separator() {
     );
 }
 
+#[test]
+fn serve_args_auto_forward_enable_lora_to_python() {
+    let cli =
+        Cli::try_parse_from(["vllm-rs", "serve", "Qwen/Qwen3-0.6B", "--enable-lora"]).unwrap();
+
+    let Command::Serve(args) = cli.command else {
+        panic!("expected serve args");
+    };
+    assert_eq!(args.managed_engine.python_args, vec!["--enable-lora"]);
+}
+
 #[test]
 fn serve_args_auto_forward_python_multi_char_alias_without_separator() {
     let cli = Cli::try_parse_from(["vllm-rs", "serve", "Qwen/Qwen3-0.6B", "-tp", "2"]).unwrap();

rust/src/cmd/src/cli/unsupported.rs

@@ -326,15 +326,6 @@ pub struct EngineUnsupportedArgs {
     #[arg(long)]
     pub mm_processor_cache_type: Option<Unsupported>,
 
-    /// If True, enable handling of LoRA adapters.
-    #[arg(
-        long,
-        visible_alias = "no-enable-lora",
-        default_missing_value = "true",
-        num_args = 0..=1
-    )]
-    pub enable_lora: Option<Unsupported>,
-
     /// Dictionary mapping specific modalities to LoRA model paths.
     #[arg(long)]
     pub default_mm_loras: Option<Unsupported>,

rust/src/engine-core-client/src/client.rs

@@ -2,6 +2,7 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use futures::future::{join_all, try_join_all};
+use serde::Serialize;
 use tokio::sync::mpsc;
 use tokio_util::task::AbortOnDropHandle;
 use tracing::{debug, info, trace};
@@ -10,6 +11,7 @@ use crate::client::imp::{ClientInner, run_abort_loop, run_output_dispatcher_loop
 use crate::coordinator::CoordinatorHandle;
 use crate::error::{Error, Result};
 use crate::protocol::handshake::EngineCoreReadyResponse;
+use crate::protocol::lora::LoraRequest;
 use crate::protocol::utility::EngineCoreUtilityRequest;
 use crate::protocol::{EngineCoreRequest, EngineCoreRequestType, ModelDtype};
 use crate::transport::{self, ConnectedEngine};
@@ -22,7 +24,7 @@ pub use stream::{EngineCoreOutputStream, EngineCoreStreamOutput};
 
 /// How the frontend acquires its request/response transport with Python
 /// `EngineCoreProc`s.
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
 pub enum TransportMode {
     /// The Rust process owns the startup handshake and allocates or binds the
     /// frontend transport addresses itself before replying to engine
@@ -659,6 +661,24 @@ impl EngineCoreClient {
         Ok(results.into_iter().all(|ok| ok))
     }
 
+    /// Load or refresh one LoRA adapter on every connected engine.
+    pub async fn add_lora(&self, lora_request: &LoraRequest) -> Result<bool> {
+        Ok(self
+            .call_utility::<bool, _>("add_lora", (lora_request,))
+            .await?
+            .into_iter()
+            .all(|loaded| loaded))
+    }
+
+    /// Remove one LoRA adapter from every connected engine.
+    pub async fn remove_lora(&self, lora_id: u64) -> Result<bool> {
+        Ok(self
+            .call_utility::<bool, _>("remove_lora", (lora_id,))
+            .await?
+            .into_iter()
+            .all(|removed| removed))
+    }
+
     /// Put the engine to sleep.
     pub async fn sleep(&self, level: u32, mode: &str) -> Result<()> {
         self.call_utility::<(), _>("sleep", (level, mode)).await?;

rust/src/engine-core-client/src/protocol/lora.rs

@@ -0,0 +1,42 @@
+use serde_tuple::{Deserialize_tuple, Serialize_tuple};
+
+use crate::protocol::OpaqueValue;
+
+/// Request for a LoRA adapter.
+///
+/// Mirrors Python `vllm.lora.request.LoRARequest`, which is a msgspec
+/// `array_like=True` struct. Keep the field order aligned with Python.
+#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple)]
+pub struct LoraRequest {
+    pub lora_name: String,
+    pub lora_int_id: u64,
+    pub lora_path: String,
+    #[serde(default)]
+    pub base_model_name: Option<String>,
+    #[serde(default)]
+    pub tensorizer_config_dict: Option<OpaqueValue>,
+    #[serde(default)]
+    pub load_inplace: bool,
+    #[serde(default)]
+    pub is_3d_lora_weight: bool,
+}
+
+impl LoraRequest {
+    pub fn new(
+        lora_name: String,
+        lora_int_id: u64,
+        lora_path: String,
+        load_inplace: bool,
+        is_3d_lora_weight: bool,
+    ) -> Self {
+        Self {
+            lora_name,
+            lora_int_id,
+            lora_path,
+            base_model_name: None,
+            tensorizer_config_dict: None,
+            load_inplace,
+            is_3d_lora_weight,
+        }
+    }
+}

rust/src/engine-core-client/src/protocol/mod.rs

@@ -48,6 +48,7 @@ mod classified_outputs;
 pub mod dtype;
 pub mod handshake;
 pub mod logprobs;
+pub mod lora;
 pub mod multimodal;
 pub mod stats;
 pub mod tensor;
@@ -349,7 +350,7 @@ pub struct EngineCoreRequest {
     pub pooling_params: Option<OpaqueValue>,
     pub arrival_time: f64,
     #[serde(default)]
-    pub lora_request: Option<OpaqueValue>,
+    pub lora_request: Option<lora::LoraRequest>,
     #[serde(default)]
     pub cache_salt: Option<String>,
     #[serde(default)]

rust/src/llm/src/request.rs

@@ -2,8 +2,9 @@ use std::collections::BTreeMap;
 use std::time::{SystemTime, UNIX_EPOCH};
 
 use uuid::Uuid;
+use vllm_engine_core_client::protocol::lora::LoraRequest;
 use vllm_engine_core_client::protocol::multimodal::MmFeatures;
-use vllm_engine_core_client::protocol::{EngineCoreRequest, EngineCoreSamplingParams, OpaqueValue};
+use vllm_engine_core_client::protocol::{EngineCoreRequest, EngineCoreSamplingParams};
 
 use crate::error::{Error, Result};
 
@@ -34,7 +35,7 @@ pub struct GenerateRequest {
     pub priority: i32,
     pub data_parallel_rank: Option<u32>,
     pub reasoning_ended: Option<bool>,
-    pub lora_request: Option<OpaqueValue>,
+    pub lora_request: Option<LoraRequest>,
 }
 
 #[derive(Debug)]

rust/src/server/src/config.rs

@@ -2,12 +2,13 @@ use std::collections::HashMap;
 use std::time::Duration;
 
 use anyhow::Result;
+use serde::Serialize;
 use serde_json::Value;
 use vllm_chat::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection};
 use vllm_engine_core_client::{CoordinatorMode as EngineCoreCoordinatorMode, TransportMode};
 
 /// How the HTTP server obtains its listening socket.
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
 pub enum HttpListenerMode {
     /// Bind a fresh TCP listener on the given host/port.
     BindTcp { host: String, port: u16 },
@@ -20,7 +21,7 @@ pub enum HttpListenerMode {
 
 /// Which coordinator implementation should be active when one is present for a
 /// frontend client.
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
 pub enum CoordinatorMode {
     /// Do not run a coordinator at all.
     None,
@@ -32,7 +33,7 @@ pub enum CoordinatorMode {
 }
 
 /// Normalized runtime configuration for the minimal OpenAI-compatible server.
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
 pub struct Config {
     /// Frontend-to-engine transport setup.
     pub transport_mode: TransportMode,

rust/src/server/src/grpc/convert.rs

@@ -91,6 +91,7 @@ pub fn to_text_request(
         cache_salt: kv.map(|k| &k.cache_salt).filter(|s| !s.is_empty()).cloned(),
         add_special_tokens: true,
         data_parallel_rank: None,
+        lora_request: None,
     })
 }
 

rust/src/server/src/lib.rs

@@ -4,15 +4,17 @@ mod config;
 mod error;
 mod grpc;
 mod listener;
+mod lora;
 mod middleware;
 mod routes;
+mod server_info;
 mod state;
 mod utils;
 
 use std::sync::{Arc, OnceLock};
 
 use anyhow::{Context as _, Result};
-use axum::serve::ListenerExt as _;
+use axum::{Router, serve::ListenerExt as _};
 pub use config::{Config, CoordinatorMode, HttpListenerMode};
 use tokio::net::TcpListener;
 use tokio::time::{Instant, sleep_until};
@@ -29,6 +31,7 @@ use vllm_text::TextLlm;
 
 use crate::listener::Listener;
 use crate::routes::build_router;
+use crate::server_info::ServerInfoSnapshot;
 use crate::state::AppState;
 
 /// Build the shared application state for one configured model and one engine
@@ -87,7 +90,8 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
     Ok(Arc::new(
         AppState::new(served_model_names, chat)
             .with_log_requests(config.enable_log_requests)
-            .with_request_id_headers(config.enable_request_id_headers),
+            .with_request_id_headers(config.enable_request_id_headers)
+            .with_server_info(ServerInfoSnapshot::from_config(config)),
     ))
 }
 
@@ -97,6 +101,21 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
 /// The server owns one `vllm-chat` facade, which in turn owns the lower
 /// `vllm-text` and `vllm-llm` layers, and shuts them down before returning.
 pub async fn serve(config: Config, shutdown: CancellationToken) -> Result<()> {
+    serve_with_router_extension(config, shutdown, |router| router).await
+}
+
+/// Run the OpenAI-compatible HTTP server with an opt-in router extension.
+///
+/// The extension receives the finalized vLLM router and can merge additional
+/// routes before the server starts accepting requests.
+pub async fn serve_with_router_extension<F>(
+    config: Config,
+    shutdown: CancellationToken,
+    extend_router: F,
+) -> Result<()>
+where
+    F: FnOnce(Router) -> Router,
+{
     config.validate().context("invalid OpenAI frontend configuration")?;
 
     // Also check shutdown during the (potentially long) startup handshake.
@@ -109,7 +128,7 @@ pub async fn serve(config: Config, shutdown: CancellationToken) -> Result<()> {
         .context("failed to bind listener for OpenAI server")?;
     let bind_address = listener.local_addr()?;
     let model = state.primary_model_name().to_owned();
-    let app = build_router(state.clone());
+    let app = extend_router(build_router(state.clone()));
 
     // Optionally bind the gRPC Generate server on a separate port. Bind
     // synchronously here so bind errors (port in use, permission denied, ...)

rust/src/server/src/lora.rs

@@ -0,0 +1,168 @@
+use std::collections::BTreeMap;
+use std::sync::atomic::{AtomicU64, Ordering};
+
+use tokio::sync::{Mutex, RwLock};
+use vllm_engine_core_client::EngineCoreClient;
+use vllm_engine_core_client::protocol::lora::LoraRequest;
+
+/// Snapshot of the currently served model names plus the requested LoRA, if
+/// the model name resolves to a dynamic adapter.
+#[derive(Debug, Clone)]
+pub(crate) struct LoraModelResolution {
+    pub model_names: Vec<String>,
+    pub lora_request: Option<LoraRequest>,
+}
+
+/// Runtime registry for dynamically loaded LoRA adapters.
+pub(crate) struct LoraManager {
+    /// Dynamically loaded LoRA adapters keyed by public model name.
+    requests: RwLock<BTreeMap<String, LoraRequest>>,
+    /// Monotonic adapter id allocator. LoRA ids are one-indexed.
+    id_counter: AtomicU64,
+    /// Serialize dynamic LoRA registry updates around engine utility calls.
+    update_lock: Mutex<()>,
+}
+
+#[derive(Debug)]
+pub(crate) enum LoadLoraError {
+    AlreadyLoaded { lora_name: String },
+    BaseModelName { lora_name: String },
+    Engine(vllm_engine_core_client::Error),
+    NotLoaded { lora_name: String },
+}
+
+#[derive(Debug)]
+pub(crate) enum UnloadLoraError {
+    NotFound {
+        lora_name: String,
+    },
+    IntIdMismatch {
+        lora_name: String,
+        expected: u64,
+        actual: u64,
+    },
+    Engine(vllm_engine_core_client::Error),
+    NotRemoved {
+        lora_name: String,
+        lora_int_id: u64,
+    },
+}
+
+impl LoraManager {
+    pub fn new() -> Self {
+        Self {
+            requests: RwLock::new(BTreeMap::new()),
+            id_counter: AtomicU64::new(0),
+            update_lock: Mutex::new(()),
+        }
+    }
+
+    /// Return base served model names plus dynamically loaded LoRA adapter
+    /// names.
+    pub async fn served_model_names(&self, base_model_names: &[String]) -> Vec<String> {
+        let mut names = base_model_names.to_vec();
+        names.extend(self.requests.read().await.keys().cloned());
+        names
+    }
+
+    /// Resolve the requested model against one consistent LoRA registry
+    /// snapshot.
+    pub async fn resolve_model(
+        &self,
+        base_model_names: &[String],
+        model_name: Option<&str>,
+    ) -> LoraModelResolution {
+        let requests = self.requests.read().await;
+        let mut model_names = base_model_names.to_vec();
+        model_names.extend(requests.keys().cloned());
+        let lora_request = model_name.and_then(|name| requests.get(name).cloned());
+
+        LoraModelResolution {
+            model_names,
+            lora_request,
+        }
+    }
+
+    /// Load one dynamic LoRA adapter and register it as a public model name.
+    pub async fn load_lora(
+        &self,
+        engine_core_client: &EngineCoreClient,
+        base_model_names: &[String],
+        lora_name: String,
+        lora_path: String,
+        load_inplace: bool,
+        is_3d_lora_weight: bool,
+    ) -> Result<LoraRequest, LoadLoraError> {
+        let _guard = self.update_lock.lock().await;
+        if base_model_names.iter().any(|name| name == &lora_name) {
+            return Err(LoadLoraError::BaseModelName { lora_name });
+        }
+        if !load_inplace && self.requests.read().await.contains_key(&lora_name) {
+            return Err(LoadLoraError::AlreadyLoaded { lora_name });
+        }
+
+        let lora_int_id = self
+            .requests
+            .read()
+            .await
+            .get(&lora_name)
+            .map(|request| request.lora_int_id)
+            .unwrap_or_else(|| self.id_counter.fetch_add(1, Ordering::Relaxed) + 1);
+        let lora_request = LoraRequest::new(
+            lora_name.clone(),
+            lora_int_id,
+            lora_path,
+            load_inplace,
+            is_3d_lora_weight,
+        );
+
+        let loaded = engine_core_client
+            .add_lora(&lora_request)
+            .await
+            .map_err(LoadLoraError::Engine)?;
+        if !loaded {
+            return Err(LoadLoraError::NotLoaded { lora_name });
+        }
+        self.requests.write().await.insert(lora_name, lora_request.clone());
+        Ok(lora_request)
+    }
+
+    /// Remove one dynamic LoRA adapter from the engine and public model
+    /// registry.
+    pub async fn unload_lora(
+        &self,
+        engine_core_client: &EngineCoreClient,
+        lora_name: &str,
+        requested_lora_int_id: Option<u64>,
+    ) -> Result<LoraRequest, UnloadLoraError> {
+        let _guard = self.update_lock.lock().await;
+        let lora_request = self.requests.read().await.get(lora_name).cloned().ok_or_else(|| {
+            UnloadLoraError::NotFound {
+                lora_name: lora_name.to_string(),
+            }
+        })?;
+
+        if let Some(actual) = requested_lora_int_id
+            && actual != lora_request.lora_int_id
+        {
+            return Err(UnloadLoraError::IntIdMismatch {
+                lora_name: lora_name.to_string(),
+                expected: lora_request.lora_int_id,
+                actual,
+            });
+        }
+
+        let removed = engine_core_client
+            .remove_lora(lora_request.lora_int_id)
+            .await
+            .map_err(UnloadLoraError::Engine)?;
+        if !removed {
+            return Err(UnloadLoraError::NotRemoved {
+                lora_name: lora_request.lora_name,
+                lora_int_id: lora_request.lora_int_id,
+            });
+        }
+
+        Ok(self.requests.write().await.remove(lora_name).unwrap_or(lora_request))
+    }
+}

rust/src/server/src/routes.rs

@@ -3,8 +3,10 @@ mod collective_rpc;
 mod health;
 mod inference;
 mod load;
+mod lora;
 mod metrics;
 pub(crate) mod openai;
+mod server_info;
 mod sleep;
 mod version;
 
@@ -25,12 +27,40 @@ fn server_dev_mode_enabled() -> bool {
         .is_some_and(|value| value != 0)
 }
 
-/// Build the minimal OpenAI-compatible router for one configured model.
-pub fn build_router(state: Arc<AppState>) -> Router {
-    build_router_with_dev_mode(state, server_dev_mode_enabled())
+fn runtime_lora_updating_enabled() -> bool {
+    std::env::var("VLLM_ALLOW_RUNTIME_LORA_UPDATING")
+        .ok()
+        .is_some_and(|value| matches!(value.trim().to_lowercase().as_str(), "1" | "true"))
 }
 
+/// Build the minimal OpenAI-compatible router for one configured model.
+pub fn build_router(state: Arc<AppState>) -> Router {
+    build_router_with_options(
+        state,
+        server_dev_mode_enabled(),
+        runtime_lora_updating_enabled(),
+    )
+}
+
+#[cfg(test)]
 fn build_router_with_dev_mode(state: Arc<AppState>, dev_mode_enabled: bool) -> Router {
+    build_router_with_dev_mode_and_lora(state, dev_mode_enabled, false)
+}
+
+#[cfg(test)]
+fn build_router_with_dev_mode_and_lora(
+    state: Arc<AppState>,
+    dev_mode_enabled: bool,
+    runtime_lora_updating_enabled: bool,
+) -> Router {
+    build_router_with_options(state, dev_mode_enabled, runtime_lora_updating_enabled)
+}
+
+fn build_router_with_options(
+    state: Arc<AppState>,
+    dev_mode_enabled: bool,
+    runtime_lora_updating_enabled: bool,
+) -> Router {
     let mut router = Router::new()
         // Health & monitoring
         .route("/health", get(health::health))
@@ -44,6 +74,12 @@ fn build_router_with_dev_mode(state: Arc<AppState>, dev_mode_enabled: bool) -> R
         // vLLM specific inference endpoints
         .route("/inference/v1/generate", post(inference::generate));
 
+    if runtime_lora_updating_enabled {
+        router = router
+            .route("/v1/load_lora_adapter", post(lora::load_lora_adapter))
+            .route("/v1/unload_lora_adapter", post(lora::unload_lora_adapter));
+    }
+
     if dev_mode_enabled {
         // Development-only
         router = router
@@ -54,6 +90,7 @@ fn build_router_with_dev_mode(state: Arc<AppState>, dev_mode_enabled: bool) -> R
             .route("/sleep", post(sleep::sleep))
             .route("/wake_up", post(sleep::wake_up))
             .route("/is_sleeping", get(sleep::is_sleeping))
+            .route("/server_info", get(server_info::server_info))
     }
 
     let enable_request_id_headers = state.enable_request_id_headers;

rust/src/server/src/routes/inference/generate.rs

@@ -42,8 +42,8 @@ pub async fn generate(
     ValidatedJson(body): ValidatedJson<GenerateRequest>,
 ) -> Response {
     let request_context = resolve_request_context(&headers, body.request_id.as_deref());
-    let prepared = match prepare_generate_request(body, state.served_model_names(), request_context)
-    {
+    let lora_resolution = state.resolve_model_with_loras(body.model.as_deref()).await;
+    let prepared = match prepare_generate_request(body, &lora_resolution, request_context) {
         Ok(prepared) => prepared,
         Err(error) => return error.into_response(),
     };

rust/src/server/src/routes/inference/generate/convert.rs

@@ -3,6 +3,7 @@ use vllm_text::{Prompt, TextDecodeOptions, TextRequest};
 use super::types::GenerateRequest;
 use super::validate;
 use crate::error::ApiError;
+use crate::lora::LoraModelResolution;
 use crate::utils::{ResolvedRequestContext, merge_kv_transfer_params};
 
 /// Lowered generate request plus the response request ID.
@@ -21,10 +22,10 @@ pub struct PreparedRequest {
 /// text-generation format.
 pub fn prepare_generate_request(
     request: GenerateRequest,
-    served_model_names: &[String],
+    lora_resolution: &LoraModelResolution,
     ctx: ResolvedRequestContext,
 ) -> Result<PreparedRequest, ApiError> {
-    validate::validate_request_compat(&request, served_model_names)?;
+    validate::validate_request_compat(&request, &lora_resolution.model_names)?;
 
     let stream = request.stream;
     let include_usage = request
@@ -57,6 +58,7 @@ pub fn prepare_generate_request(
         cache_salt: request.cache_salt,
         add_special_tokens: false,
         data_parallel_rank: ctx.data_parallel_rank,
+        lora_request: lora_resolution.lora_request.clone(),
     };
 
     Ok(PreparedRequest {
@@ -76,9 +78,17 @@ mod tests {
     use vllm_text::Prompt;
 
     use super::prepare_generate_request;
+    use crate::lora::LoraModelResolution;
     use crate::routes::inference::generate::types::GenerateRequest;
     use crate::utils::ResolvedRequestContext;
 
+    fn served(names: &[&str]) -> LoraModelResolution {
+        LoraModelResolution {
+            model_names: names.iter().map(|s| s.to_string()).collect(),
+            lora_request: None,
+        }
+    }
+
     #[test]
     fn prepare_generate_request_maps_token_prompt_and_sampling_params() {
         let request: GenerateRequest = serde_json::from_value(json!({
@@ -100,7 +110,7 @@ mod tests {
 
         let prepared = prepare_generate_request(
             request,
-            &["Qwen/Qwen1.5-0.5B-Chat".to_string()],
+            &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
             ResolvedRequestContext::default(),
         )
         .expect("prepare");
@@ -143,7 +153,7 @@ mod tests {
 
         let prepared = prepare_generate_request(
             request,
-            &["Qwen/Qwen1.5-0.5B-Chat".to_string()],
+            &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
             ResolvedRequestContext::default(),
         )
         .expect("prepare");

rust/src/server/src/routes/lora.rs

@@ -0,0 +1,296 @@
+use std::path::{Component, Path, PathBuf};
+use std::sync::Arc;
+
+use axum::extract::State;
+use serde::Deserialize;
+use thiserror_ext::AsReport;
+use validator::Validate;
+
+use crate::error::ApiError;
+use crate::lora::{LoadLoraError, UnloadLoraError};
+use crate::routes::openai::utils::types::Normalizable;
+use crate::routes::openai::utils::validated_json::ValidatedJson;
+use crate::state::AppState;
+
+const RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV: &str = "VLLM_RUNTIME_LORA_ALLOWED_PATH_PREFIXES";
+
+#[derive(Debug, Deserialize, Validate)]
+pub(crate) struct LoadLoraAdapterRequest {
+    lora_name: String,
+    lora_path: String,
+    #[serde(default)]
+    load_inplace: bool,
+    #[serde(default)]
+    is_3d_lora_weight: bool,
+}
+
+impl Normalizable for LoadLoraAdapterRequest {}
+
+#[derive(Debug, Deserialize, Validate)]
+pub(crate) struct UnloadLoraAdapterRequest {
+    lora_name: String,
+    #[serde(default)]
+    lora_int_id: Option<u64>,
+}
+
+impl Normalizable for UnloadLoraAdapterRequest {}
+
+fn runtime_lora_allowed_path_prefixes() -> Option<Vec<PathBuf>> {
+    let prefixes = std::env::var_os(RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV)?;
+    let prefixes: Vec<_> = std::env::split_paths(&prefixes)
+        .filter(|path| !path.as_os_str().is_empty())
+        .collect();
+    (!prefixes.is_empty()).then_some(prefixes)
+}
+
+fn looks_like_local_lora_path(lora_path: &str) -> bool {
+    let path = Path::new(lora_path);
+    path.is_absolute()
+        || lora_path.starts_with('~')
+        || lora_path.starts_with('.')
+        || path.components().any(|component| matches!(component, Component::ParentDir))
+}
+
+fn validate_lora_path_access(
+    lora_path: &str,
+    allowed_prefixes: Option<&[PathBuf]>,
+) -> Result<Option<String>, ApiError> {
+    let path = Path::new(lora_path);
+    if !looks_like_local_lora_path(lora_path) && !path.exists() {
+        return Ok(None);
+    }
+
+    let Some(allowed_prefixes) = allowed_prefixes else {
+        return Err(ApiError::invalid_request(
+            format!(
+                "Local LoRA adapter paths require {RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV} to be configured."
+            ),
+            Some("lora_path"),
+        ));
+    };
+
+    if !path.is_absolute() {
+        return Err(ApiError::invalid_request(
+            format!(
+                "Local LoRA adapter paths must be absolute and under one of the prefixes configured by {RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV}."
+            ),
+            Some("lora_path"),
+        ));
+    }
+
+    let canonical_path = path.canonicalize().map_err(|_| {
+        ApiError::invalid_request(
+            "Local LoRA adapter path must exist and be accessible.".to_string(),
+            Some("lora_path"),
+        )
+    })?;
+    let canonical_prefixes = allowed_prefixes
+        .iter()
+        .map(|prefix| {
+            prefix.canonicalize().map_err(|_| {
+                ApiError::server_error(format!(
+                    "configured {RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV} path prefix must exist and be accessible"
+                ))
+            })
+        })
+        .collect::<Result<Vec<_>, _>>()?;
+
+    if !canonical_prefixes.iter().any(|prefix| canonical_path.starts_with(prefix)) {
+        return Err(ApiError::invalid_request(
+            "Local LoRA adapter path is outside the configured allowed prefixes.".to_string(),
+            Some("lora_path"),
+        ));
+    }
+
+    Ok(Some(canonical_path.to_string_lossy().into_owned()))
+}
+
+/// Dynamically load one LoRA adapter and expose it as an OpenAI model id.
+pub async fn load_lora_adapter(
+    State(state): State<Arc<AppState>>,
+    ValidatedJson(request): ValidatedJson<LoadLoraAdapterRequest>,
+) -> Result<String, ApiError> {
+    if request.lora_name.is_empty() || request.lora_path.is_empty() {
+        return Err(ApiError::invalid_request(
+            "Both 'lora_name' and 'lora_path' must be provided.".to_string(),
+            None,
+        ));
+    }
+    let allowed_prefixes = runtime_lora_allowed_path_prefixes();
+    let lora_path = validate_lora_path_access(&request.lora_path, allowed_prefixes.as_deref())?
+        .unwrap_or(request.lora_path);
+
+    let lora_name = request.lora_name;
+    state
+        .load_lora(
+            lora_name.clone(),
+            lora_path,
+            request.load_inplace,
+            request.is_3d_lora_weight,
+        )
+        .await
+        .map_err(|error| match error {
+            LoadLoraError::AlreadyLoaded { lora_name } => ApiError::invalid_request(
+                format!(
+                    "The lora adapter '{lora_name}' has already been loaded. If you want to load the adapter in place, set 'load_inplace' to true."
+                ),
+                Some("lora_name"),
+            ),
+            LoadLoraError::BaseModelName { lora_name } => ApiError::invalid_request(
+                format!("The lora adapter name '{lora_name}' conflicts with a served base model."),
+                Some("lora_name"),
+            ),
+            LoadLoraError::Engine(error) => ApiError::server_error(format!(
+                "failed to load LoRA adapter '{lora_name}': {}",
+                error.to_report_string()
+            )),
+            LoadLoraError::NotLoaded { lora_name } => ApiError::server_error(format!(
+                "failed to load LoRA adapter '{lora_name}': engine rejected the adapter"
+            )),
+        })?;
+
+    Ok(format!(
+        "Success: LoRA adapter '{lora_name}' added successfully."
+    ))
+}
+
+/// Remove one LoRA adapter from the engine and frontend registry.
+pub async fn unload_lora_adapter(
+    State(state): State<Arc<AppState>>,
+    ValidatedJson(request): ValidatedJson<UnloadLoraAdapterRequest>,
+) -> Result<String, ApiError> {
+    if request.lora_name.is_empty() {
+        return Err(ApiError::invalid_request(
+            "'lora_name' needs to be provided to unload a LoRA adapter.".to_string(),
+            Some("lora_name"),
+        ));
+    }
+
+    let lora_request = state
+        .unload_lora(&request.lora_name, request.lora_int_id)
+        .await
+        .map_err(|error| match error {
+            UnloadLoraError::NotFound { lora_name } => ApiError::model_not_found(lora_name),
+            UnloadLoraError::IntIdMismatch {
+                lora_name,
+                expected,
+                actual,
+            } => ApiError::invalid_request(
+                format!(
+                    "The requested lora_int_id {actual} does not match loaded adapter '{lora_name}' with id {expected}."
+                ),
+                Some("lora_int_id"),
+            ),
+            UnloadLoraError::Engine(error) => ApiError::server_error(format!(
+                "failed to unload LoRA adapter '{}': {}",
+                request.lora_name,
+                error.to_report_string()
+            )),
+            UnloadLoraError::NotRemoved {
+                lora_name,
+                lora_int_id,
+            } => ApiError::server_error(format!(
+                "failed to unload LoRA adapter '{lora_name}' with id {lora_int_id}"
+            )),
+        })?;
+
+    Ok(format!(
+        "Success: LoRA adapter '{}' removed successfully.",
+        lora_request.lora_name
+    ))
+}
+
+#[cfg(test)]
+mod tests {
+    use std::fs;
+    use std::path::PathBuf;
+    use std::time::{SystemTime, UNIX_EPOCH};
+
+    use super::validate_lora_path_access;
+
+    fn temp_lora_dir(test_name: &str) -> PathBuf {
+        let suffix = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .expect("clock should be after unix epoch")
+            .as_nanos();
+        let path = std::env::temp_dir().join(format!(
+            "vllm-lora-{test_name}-{}-{suffix}",
+            std::process::id()
+        ));
+        fs::create_dir_all(&path).expect("create temp lora dir");
+        path
+    }
+
+    #[test]
+    fn lora_path_allows_hf_repo_ids_without_prefixes() {
+        assert_eq!(
+            validate_lora_path_access("org/adapter-a", None).expect("hf repo id should be allowed"),
+            None
+        );
+    }
+
+    #[test]
+    fn lora_path_rejects_local_paths_without_prefixes() {
+        assert!(validate_lora_path_access("/tmp/adapter-a", None).is_err());
+        assert!(validate_lora_path_access("./adapter-a", None).is_err());
+        assert!(validate_lora_path_access("~/adapter-a", None).is_err());
+        assert!(validate_lora_path_access("subdir/../../../etc/sensitive", None).is_err());
+    }
+
+    #[test]
+    fn lora_path_rejects_existing_bare_relative_paths_without_prefixes() {
+        let root =
+            PathBuf::from("target").join(format!("vllm-lora-relative-{}", std::process::id()));
+        let adapter = root.join("adapter-a");
+        fs::create_dir_all(&adapter).expect("create relative adapter dir");
+
+        assert!(
+            validate_lora_path_access(adapter.to_str().expect("utf-8 temp path"), None).is_err()
+        );
+
+        fs::remove_dir_all(root).ok();
+    }
+
+    #[test]
+    fn lora_path_allows_absolute_paths_under_configured_prefixes() {
+        let root = temp_lora_dir("allowed-prefix");
+        let allowed = root.join("allowed");
+        let adapter = allowed.join("adapter-a");
+        fs::create_dir_all(&adapter).expect("create adapter dir");
+
+        let prefixes = [allowed];
+        let resolved =
+            validate_lora_path_access(adapter.to_str().expect("utf-8 temp path"), Some(&prefixes))
+                .expect("path under configured prefix should be allowed");
+        assert_eq!(
+            resolved.as_deref(),
+            Some(
+                adapter
+                    .canonicalize()
+                    .expect("canonical adapter")
+                    .to_str()
+                    .expect("utf-8 temp path")
+            )
+        );
+
+        fs::remove_dir_all(root).ok();
+    }
+
+    #[test]
+    fn lora_path_rejects_parent_escape_from_configured_prefixes() {
+        let root = temp_lora_dir("parent-escape");
+        let allowed = root.join("allowed");
+        let private_adapter = root.join("private").join("adapter-a");
+        fs::create_dir_all(&allowed).expect("create allowed dir");
+        fs::create_dir_all(&private_adapter).expect("create private adapter dir");
+
+        let escaped = allowed.join("../private/adapter-a");
+        let prefixes = [allowed];
+        assert!(
+            validate_lora_path_access(escaped.to_str().expect("utf-8 temp path"), Some(&prefixes))
+                .is_err()
+        );
+
+        fs::remove_dir_all(root).ok();
+    }
+}

rust/src/server/src/routes/openai/chat_completions.rs

@@ -49,8 +49,9 @@ pub async fn chat_completions(
 ) -> Response {
     let stream = body.stream;
     let request_context = resolve_request_context(&headers, body.request_id.as_deref());
+    let lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await;
 
-    let prepared = match prepare_chat_request(body, state.served_model_names(), request_context) {
+    let prepared = match prepare_chat_request(body, &lora_resolution, request_context) {
         Ok(prepared) => prepared,
         Err(error) => return error.into_response(),
     };

rust/src/server/src/routes/openai/chat_completions/convert.rs

@@ -8,6 +8,7 @@ use vllm_chat::{
 use super::types::ChatCompletionRequest;
 use super::validate;
 use crate::error::{ApiError, bail_invalid_request};
+use crate::lora::LoraModelResolution;
 use crate::routes::openai::utils::structured_outputs::convert_from_response_format;
 use crate::routes::openai::utils::types::{
     ChatMessage, ContentPart, MessageContent, Tool, ToolChoice, ToolChoiceValue,
@@ -41,16 +42,21 @@ pub struct PreparedRequest {
 /// Validate and lower one OpenAI chat completion request into the internal chat
 /// format.
 ///
-/// `served_model_names` must be non-empty; the first entry is used as the
-/// `model` field in responses.
+/// `lora_resolution.model_names` must be non-empty; the first entry is used as
+/// the base `model` field in responses when no LoRA adapter is selected.
 pub(crate) fn prepare_chat_request(
     request: ChatCompletionRequest,
-    served_model_names: &[String],
+    lora_resolution: &LoraModelResolution,
     ctx: ResolvedRequestContext,
 ) -> Result<PreparedRequest, ApiError> {
-    validate::validate_request_compat(&request, served_model_names)?;
+    validate::validate_request_compat(&request, &lora_resolution.model_names)?;
 
     let request_id = format!("chatcmpl-{}", ctx.request_id);
+    let response_model = lora_resolution
+        .lora_request
+        .as_ref()
+        .map(|request| request.lora_name.clone())
+        .unwrap_or_else(|| lora_resolution.model_names.first().cloned().unwrap_or_default());
     let echo = request
         .echo
         .then(|| extract_last_assistant_content(&request.messages))
@@ -131,11 +137,12 @@ pub(crate) fn prepare_chat_request(
         cache_salt: request.cache_salt,
         add_special_tokens: request.add_special_tokens,
         data_parallel_rank: ctx.data_parallel_rank,
+        lora_request: lora_resolution.lora_request.clone(),
     };
 
     Ok(PreparedRequest {
         request_id,
-        response_model: served_model_names.first().cloned().unwrap_or_default(),
+        response_model,
         include_usage,
         requested_logprobs,
         include_prompt_logprobs,
@@ -352,6 +359,7 @@ mod tests {
     use vllm_text::output::TextDecodeOptions;
 
     use super::prepare_chat_request;
+    use crate::lora::LoraModelResolution;
     use crate::routes::openai::chat_completions::types::{
         AssistantRole, ChatCompletionMessage, ChatCompletionRequest,
     };
@@ -365,8 +373,11 @@ mod tests {
         resolve_request_context(headers, request_id)
     }
 
-    fn served(names: &[&str]) -> Vec<String> {
-        names.iter().map(|s| s.to_string()).collect()
+    fn served(names: &[&str]) -> LoraModelResolution {
+        LoraModelResolution {
+            model_names: names.iter().map(|s| s.to_string()).collect(),
+            lora_request: None,
+        }
     }
 
     fn base_request() -> ChatCompletionRequest {

rust/src/server/src/routes/openai/completions.rs

@@ -44,12 +44,12 @@ pub async fn completions(
     let stream = body.stream;
     let logprobs = body.logprobs;
     let request_context = resolve_request_context(&headers, body.request_id.as_deref());
+    let lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await;
 
-    let prepared =
-        match prepare_completion_request(body, state.served_model_names(), request_context) {
-            Ok(prepared) => prepared,
-            Err(error) => return error.into_response(),
-        };
+    let prepared = match prepare_completion_request(body, &lora_resolution, request_context) {
+        Ok(prepared) => prepared,
+        Err(error) => return error.into_response(),
+    };
     let request_span = tracing::info_span!(
         "completions",
         request_id = %prepared.request_id,

rust/src/server/src/routes/openai/completions/convert.rs

@@ -2,6 +2,7 @@ use vllm_text::{SamplingParams, TextDecodeOptions, TextRequest};
 
 use super::types::CompletionRequest;
 use crate::error::ApiError;
+use crate::lora::LoraModelResolution;
 use crate::routes::openai::completions::validate;
 use crate::routes::openai::utils::structured_outputs::convert_from_response_format_value;
 use crate::utils::{ResolvedRequestContext, convert_logit_bias, merge_kv_transfer_params};
@@ -30,16 +31,21 @@ pub struct PreparedRequest {
 /// Validate and lower one OpenAI completions request into the internal
 /// text-generation format.
 ///
-/// `served_model_names` must be non-empty; the first entry is used as the
-/// `model` field in responses.
+/// `lora_resolution.model_names` must be non-empty; the first entry is used as
+/// the base `model` field in responses when no LoRA adapter is selected.
 pub(crate) fn prepare_completion_request(
     request: CompletionRequest,
-    served_model_names: &[String],
+    lora_resolution: &LoraModelResolution,
     ctx: ResolvedRequestContext,
 ) -> Result<PreparedRequest, ApiError> {
-    validate::validate_request_compat(&request, served_model_names)?;
+    validate::validate_request_compat(&request, &lora_resolution.model_names)?;
 
     let request_id = format!("cmpl-{}", ctx.request_id);
+    let response_model = lora_resolution
+        .lora_request
+        .as_ref()
+        .map(|request| request.lora_name.clone())
+        .unwrap_or_else(|| lora_resolution.model_names.first().cloned().unwrap_or_default());
 
     let logprobs = match request.logprobs {
         Some(logprobs) => Some(i32::try_from(logprobs).map_err(|_| {
@@ -104,11 +110,12 @@ pub(crate) fn prepare_completion_request(
         cache_salt: request.cache_salt,
         add_special_tokens: request.add_special_tokens,
         data_parallel_rank: ctx.data_parallel_rank,
+        lora_request: lora_resolution.lora_request.clone(),
     };
 
     Ok(PreparedRequest {
         request_id,
-        response_model: served_model_names.first().cloned().unwrap_or_default(),
+        response_model,
         include_usage,
         text_request,
         echo,
@@ -124,6 +131,7 @@ mod tests {
     use vllm_text::Prompt;
 
     use super::prepare_completion_request;
+    use crate::lora::LoraModelResolution;
     use crate::routes::openai::completions::types::CompletionRequest;
     use crate::utils::{ResolvedRequestContext, resolve_request_context};
 
@@ -131,8 +139,11 @@ mod tests {
         resolve_request_context(headers, request_id)
     }
 
-    fn served(names: &[&str]) -> Vec<String> {
-        names.iter().map(|s| s.to_string()).collect()
+    fn served(names: &[&str]) -> LoraModelResolution {
+        LoraModelResolution {
+            model_names: names.iter().map(|s| s.to_string()).collect(),
+            lora_request: None,
+        }
     }
 
     fn base_request_json() -> serde_json::Value {

rust/src/server/src/routes/openai/models.rs

@@ -8,13 +8,13 @@ use crate::state::AppState;
 
 /// Return all configured served model names in OpenAI `list models` format.
 pub async fn list_models(State(state): State<Arc<AppState>>) -> Json<ListModelsResponse> {
+    let model_names = state.served_model_names_with_loras().await;
     Json(ListModelsResponse {
         object: "list".to_string(),
-        data: state
-            .served_model_names()
-            .iter()
+        data: model_names
+            .into_iter()
             .map(|name| ModelObject {
-                id: name.clone(),
+                id: name,
                 object: "model".to_string(),
                 created: 0,
                 owned_by: "vllm-frontend-rs".to_string(),

rust/src/server/src/routes/server_info.rs

@@ -0,0 +1,47 @@
+use std::sync::Arc;
+
+use axum::Json;
+use axum::extract::{Query, State};
+use axum::http::StatusCode;
+use axum::response::{IntoResponse, Response};
+use serde::Deserialize;
+
+use crate::server_info::ServerInfoConfigFormat;
+use crate::state::AppState;
+
+#[derive(Debug, Clone, Copy, Deserialize)]
+#[serde(rename_all = "lowercase")]
+enum ConfigFormat {
+    Text,
+    Json,
+}
+
+impl From<ConfigFormat> for ServerInfoConfigFormat {
+    fn from(value: ConfigFormat) -> Self {
+        match value {
+            ConfigFormat::Text => Self::Text,
+            ConfigFormat::Json => Self::Json,
+        }
+    }
+}
+
+fn default_config_format() -> ConfigFormat {
+    ConfigFormat::Text
+}
+
+#[derive(Debug, Deserialize)]
+pub(crate) struct ServerInfoParams {
+    #[serde(default = "default_config_format")]
+    config_format: ConfigFormat,
+}
+
+/// Get server configuration and environment metadata.
+pub async fn server_info(
+    State(state): State<Arc<AppState>>,
+    Query(params): Query<ServerInfoParams>,
+) -> Response {
+    match state.server_info_response(params.config_format.into()) {
+        Some(response) => Json(response).into_response(),
+        None => StatusCode::NOT_FOUND.into_response(),
+    }
+}

rust/src/server/src/routes/tests.rs

@@ -43,7 +43,8 @@ use vllm_text::{Prompt, TextBackend};
 use zeromq::prelude::{SocketRecv, SocketSend};
 use zeromq::{DealerSocket, PushSocket, ZmqMessage};
 
-use super::{build_router, build_router_with_dev_mode};
+use super::{build_router, build_router_with_dev_mode, build_router_with_dev_mode_and_lora};
+use crate::lora::LoraModelResolution;
 use crate::routes::openai::chat_completions::convert::prepare_chat_request;
 use crate::state::AppState;
 
@@ -141,6 +142,13 @@ fn default_stream_output_specs() -> Vec<(Vec<u32>, Option<EngineCoreFinishReason
     ]
 }
 
+fn assert_adapter_a_lora_request(request: &EngineCoreRequest) {
+    let lora = request.lora_request.as_ref().expect("lora request");
+    assert_eq!(lora.lora_name, "adapter-a");
+    assert_eq!(lora.lora_int_id, 1);
+    assert_eq!(lora.lora_path, "org/adapter-a");
+}
+
 fn sse_data_payloads(text: &str) -> Vec<&str> {
     text.lines().filter_map(|line| line.strip_prefix("data: ")).collect()
 }
@@ -734,16 +742,23 @@ async fn test_chat_with_engine_outputs(
 }
 
 async fn test_app() -> axum::Router {
+    test_app_with_dev_mode(false).await
+}
+
+async fn test_app_with_dev_mode(dev_mode_enabled: bool) -> axum::Router {
     let (chat, _engine_task) = test_models_with_engine_outputs_and_backend(
         b"engine-openai",
         default_stream_output_specs(),
         Arc::new(FakeChatBackend::new()),
     )
     .await;
-    build_router(Arc::new(AppState::new(
-        vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
-        chat,
-    )))
+    build_router_with_dev_mode(
+        Arc::new(AppState::new(
+            vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
+            chat,
+        )),
+        dev_mode_enabled,
+    )
 }
 
 async fn test_app_with_request_id_headers() -> (axum::Router, MockEngineTask) {
@@ -822,12 +837,13 @@ where
 
     let chat = ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new()));
     (
-        build_router_with_dev_mode(
+        build_router_with_dev_mode_and_lora(
             Arc::new(AppState::new(
                 vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
                 chat,
             )),
             true,
+            true,
         ),
         engine_task,
     )
@@ -1078,6 +1094,401 @@ async fn version_returns_engine_vllm_version() {
     );
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn server_info_endpoint_is_dev_mode_only() {
+    let mut app = test_app().await;
+    let response = app
+        .call(
+            Request::builder()
+                .uri("/server_info")
+                .body(Body::empty())
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+
+    assert_eq!(response.status(), StatusCode::NOT_FOUND);
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn load_lora_adapter_registers_model_and_forwards_lora_request() {
+    let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
+        boxed_test_future(async move {
+            let utility = recv_engine_message(dealer).await;
+            assert_eq!(utility[0].as_ref(), &[0x03]);
+
+            let payload = decode_value(&utility[1]).expect("decode utility payload");
+            let array = payload.as_array().expect("utility payload array");
+            let call_id = array[1].as_u64().expect("call id");
+            assert_eq!(array[2], Value::from("add_lora"));
+
+            let args = array[3].as_array().expect("utility args");
+            let lora = args[0].as_array().expect("lora request tuple");
+            assert_eq!(lora[0], Value::from("adapter-a"));
+            assert_eq!(lora[1], Value::from(1));
+            assert_eq!(lora[2], Value::from("org/adapter-a"));
+
+            send_outputs(push, utility_outputs(call_id, utility_result_value(true))).await;
+
+            let add = recv_engine_message(dealer).await;
+            assert_eq!(add[0].as_ref(), &[0x00]);
+            let request: EngineCoreRequest =
+                rmp_serde::from_slice(&add[1]).expect("decode engine request");
+            assert_adapter_a_lora_request(&request);
+
+            send_outputs(
+                push,
+                engine_outputs_for_request(&request.request_id, default_stream_output_specs()),
+            )
+            .await;
+
+            let add = recv_engine_message(dealer).await;
+            assert_eq!(add[0].as_ref(), &[0x00]);
+            let request: EngineCoreRequest =
+                rmp_serde::from_slice(&add[1]).expect("decode engine request");
+            assert_adapter_a_lora_request(&request);
+
+            send_outputs(
+                push,
+                engine_outputs_for_request(&request.request_id, default_stream_output_specs()),
+            )
+            .await;
+
+            let add = recv_engine_message(dealer).await;
+            assert_eq!(add[0].as_ref(), &[0x00]);
+            let request: EngineCoreRequest =
+                rmp_serde::from_slice(&add[1]).expect("decode engine request");
+            assert_eq!(request.prompt_token_ids.as_deref(), Some(&[11, 22][..]));
+            assert_adapter_a_lora_request(&request);
+
+            send_outputs(
+                push,
+                engine_outputs_for_request(&request.request_id, default_stream_output_specs()),
+            )
+            .await;
+
+            let utility = recv_engine_message(dealer).await;
+            assert_eq!(utility[0].as_ref(), &[0x03]);
+
+            let payload = decode_value(&utility[1]).expect("decode utility payload");
+            let array = payload.as_array().expect("utility payload array");
+            let call_id = array[1].as_u64().expect("call id");
+            assert_eq!(array[2], Value::from("remove_lora"));
+
+            let args = array[3].as_array().expect("utility args");
+            assert_eq!(args[0], Value::from(1));
+
+            send_outputs(push, utility_outputs(call_id, utility_result_value(true))).await;
+        })
+    })
+    .await;
+
+    let response = app
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/v1/load_lora_adapter")
+                .header("content-type", "application/json")
+                .body(Body::from(
+                    json!({
+                        "lora_name": "adapter-a",
+                        "lora_path": "org/adapter-a"
+                    })
+                    .to_string(),
+                ))
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+    assert_eq!(response.status(), StatusCode::OK);
+
+    let models = app
+        .call(Request::builder().uri("/v1/models").body(Body::empty()).expect("build request"))
+        .await
+        .expect("call app");
+    let body = to_bytes(models.into_body(), usize::MAX).await.expect("read body");
+    let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
+    assert_eq!(json["data"][1]["id"], "adapter-a");
+
+    let response = app
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/v1/completions")
+                .header("content-type", "application/json")
+                .body(Body::from(
+                    json!({
+                        "model": "adapter-a",
+                        "prompt": "hello",
+                        "max_tokens": 2
+                    })
+                    .to_string(),
+                ))
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+    assert_eq!(response.status(), StatusCode::OK);
+
+    let response = app
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/v1/chat/completions")
+                .header("content-type", "application/json")
+                .body(Body::from(
+                    json!({
+                        "model": "adapter-a",
+                        "stream": false,
+                        "messages": [{"role": "user", "content": "hello"}]
+                    })
+                    .to_string(),
+                ))
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+    assert_eq!(response.status(), StatusCode::OK);
+
+    let response = app
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/inference/v1/generate")
+                .header("content-type", "application/json")
+                .body(Body::from(
+                    json!({
+                        "model": "adapter-a",
+                        "token_ids": [11, 22],
+                        "stream": false,
+                        "sampling_params": {
+                            "max_tokens": 2
+                        }
+                    })
+                    .to_string(),
+                ))
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+    assert_eq!(response.status(), StatusCode::OK);
+
+    let response = app
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/v1/unload_lora_adapter")
+                .header("content-type", "application/json")
+                .body(Body::from(
+                    json!({
+                        "lora_name": "adapter-a",
+                        "lora_int_id": 1
+                    })
+                    .to_string(),
+                ))
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+    assert_eq!(response.status(), StatusCode::OK);
+
+    let models = app
+        .call(Request::builder().uri("/v1/models").body(Body::empty()).expect("build request"))
+        .await
+        .expect("call app");
+    let body = to_bytes(models.into_body(), usize::MAX).await.expect("read body");
+    let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
+    assert_eq!(json["data"].as_array().expect("model data").len(), 1);
+
+    let response = app
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/v1/completions")
+                .header("content-type", "application/json")
+                .body(Body::from(
+                    json!({
+                        "model": "adapter-a",
+                        "prompt": "hello",
+                        "max_tokens": 2
+                    })
+                    .to_string(),
+                ))
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+    assert_eq!(response.status(), StatusCode::NOT_FOUND);
+
+    drop(app);
+    engine_task.finish().await;
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn server_info_endpoint_returns_not_found_without_snapshot() {
+    let mut app = test_app_with_dev_mode(true).await;
+    let response = app
+        .call(
+            Request::builder()
+                .uri("/server_info")
+                .body(Body::empty())
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+
+    assert_eq!(response.status(), StatusCode::NOT_FOUND);
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn unload_lora_adapter_rejects_mismatched_lora_int_id() {
+    let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
+        boxed_test_future(async move {
+            let utility = recv_engine_message(dealer).await;
+            assert_eq!(utility[0].as_ref(), &[0x03]);
+
+            let payload = decode_value(&utility[1]).expect("decode utility payload");
+            let array = payload.as_array().expect("utility payload array");
+            let call_id = array[1].as_u64().expect("call id");
+            assert_eq!(array[2], Value::from("add_lora"));
+
+            send_outputs(push, utility_outputs(call_id, utility_result_value(true))).await;
+        })
+    })
+    .await;
+
+    let response = app
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/v1/load_lora_adapter")
+                .header("content-type", "application/json")
+                .body(Body::from(
+                    json!({
+                        "lora_name": "adapter-a",
+                        "lora_path": "org/adapter-a"
+                    })
+                    .to_string(),
+                ))
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+    assert_eq!(response.status(), StatusCode::OK);
+
+    let response = app
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/v1/unload_lora_adapter")
+                .header("content-type", "application/json")
+                .body(Body::from(
+                    json!({
+                        "lora_name": "adapter-a",
+                        "lora_int_id": 99
+                    })
+                    .to_string(),
+                ))
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+    assert_eq!(response.status(), StatusCode::BAD_REQUEST);
+
+    let models = app
+        .call(Request::builder().uri("/v1/models").body(Body::empty()).expect("build request"))
+        .await
+        .expect("call app");
+    let body = to_bytes(models.into_body(), usize::MAX).await.expect("read body");
+    let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
+    assert_eq!(json["data"][1]["id"], "adapter-a");
+
+    drop(app);
+    engine_task.finish().await;
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn load_lora_adapter_rejects_engine_false_result() {
+    let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
+        boxed_test_future(async move {
+            let utility = recv_engine_message(dealer).await;
+            assert_eq!(utility[0].as_ref(), &[0x03]);
+
+            let payload = decode_value(&utility[1]).expect("decode utility payload");
+            let array = payload.as_array().expect("utility payload array");
+            let call_id = array[1].as_u64().expect("call id");
+            assert_eq!(array[2], Value::from("add_lora"));
+
+            send_outputs(push, utility_outputs(call_id, utility_result_value(false))).await;
+        })
+    })
+    .await;
+
+    let response = app
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/v1/load_lora_adapter")
+                .header("content-type", "application/json")
+                .body(Body::from(
+                    json!({
+                        "lora_name": "adapter-a",
+                        "lora_path": "org/adapter-a"
+                    })
+                    .to_string(),
+                ))
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+    assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
+
+    let models = app
+        .call(Request::builder().uri("/v1/models").body(Body::empty()).expect("build request"))
+        .await
+        .expect("call app");
+    let body = to_bytes(models.into_body(), usize::MAX).await.expect("read body");
+    let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
+    assert_eq!(json["data"].as_array().expect("model data").len(), 1);
+
+    drop(app);
+    engine_task.finish().await;
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn load_lora_adapter_rejects_base_model_name_collision() {
+    let (mut app, engine_task) =
+        test_admin_app_with_engine_script(|_, _| boxed_test_future(async move {})).await;
+
+    let response = app
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/v1/load_lora_adapter")
+                .header("content-type", "application/json")
+                .body(Body::from(
+                    json!({
+                        "lora_name": "Qwen/Qwen1.5-0.5B-Chat",
+                        "lora_path": "org/adapter-a"
+                    })
+                    .to_string(),
+                ))
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+    assert_eq!(response.status(), StatusCode::BAD_REQUEST);
+
+    drop(app);
+    engine_task.finish().await;
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn http_metrics_record_list_models_requests() {
@@ -3000,7 +3411,10 @@ async fn prepared_openai_request_streams_text_events() {
             "messages": [{"role": "user", "content": "hello"}]
         }))
         .expect("decode request"),
-        &["Qwen/Qwen1.5-0.5B-Chat".to_string()],
+        &LoraModelResolution {
+            model_names: vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
+            lora_request: None,
+        },
         crate::utils::ResolvedRequestContext::default(),
     )
     .expect("prepare request");

rust/src/server/src/server_info.rs

@@ -0,0 +1,144 @@
+use std::collections::BTreeMap;
+
+use serde_json::{Value, json};
+
+use crate::config::Config;
+
+const SENSITIVE_VLLM_ENV_PATTERNS: &[&str] =
+    &["KEY", "SECRET", "TOKEN", "PASSWORD", "CREDENTIAL", "AUTH"];
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum ServerInfoConfigFormat {
+    Text,
+    Json,
+}
+
+/// Snapshot returned by `/server_info`.
+#[derive(Debug, Clone)]
+pub(crate) struct ServerInfoSnapshot {
+    vllm_config_text: String,
+    vllm_config_json: Value,
+    vllm_env: BTreeMap<String, String>,
+    system_env: BTreeMap<String, String>,
+}
+
+impl ServerInfoSnapshot {
+    /// Capture the runtime configuration fields available to the Rust frontend.
+    pub(crate) fn from_config(config: &Config) -> Self {
+        let vllm_config_json =
+            serde_json::to_value(config).expect("server info value must serialize");
+
+        Self {
+            vllm_config_text: render_config_text(&vllm_config_json),
+            vllm_config_json,
+            vllm_env: collect_vllm_env(),
+            system_env: collect_system_env(),
+        }
+    }
+
+    pub(crate) fn response(&self, config_format: ServerInfoConfigFormat) -> Value {
+        let vllm_config = match config_format {
+            ServerInfoConfigFormat::Text => Value::String(self.vllm_config_text.clone()),
+            ServerInfoConfigFormat::Json => self.vllm_config_json.clone(),
+        };
+
+        json!({
+            "vllm_config": vllm_config,
+            "vllm_env": self.vllm_env.clone(),
+            "system_env": self.system_env.clone(),
+        })
+    }
+}
+
+fn render_config_text(config: &Value) -> String {
+    match config {
+        Value::Object(fields) => fields
+            .iter()
+            .map(|(key, value)| format!("{key}={}", render_config_text_value(value)))
+            .collect::<Vec<_>>()
+            .join("\n"),
+        _ => render_config_text_value(config),
+    }
+}
+
+fn render_config_text_value(value: &Value) -> String {
+    match value {
+        Value::Null => "None".to_string(),
+        Value::String(value) => value.clone(),
+        _ => value.to_string(),
+    }
+}
+
+fn collect_vllm_env() -> BTreeMap<String, String> {
+    std::env::vars().filter(|(key, _)| is_public_vllm_env_key(key)).collect()
+}
+
+fn is_public_vllm_env_key(key: &str) -> bool {
+    let key = key.to_ascii_uppercase();
+    key.starts_with("VLLM_")
+        && !SENSITIVE_VLLM_ENV_PATTERNS.iter().any(|pattern| key.contains(pattern))
+}
+
+fn collect_system_env() -> BTreeMap<String, String> {
+    BTreeMap::from([
+        ("arch".to_string(), std::env::consts::ARCH.to_string()),
+        ("family".to_string(), std::env::consts::FAMILY.to_string()),
+        ("os".to_string(), std::env::consts::OS.to_string()),
+    ])
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::BTreeSet;
+
+    use serde_json::{Value, json};
+
+    use super::{is_public_vllm_env_key, render_config_text};
+
+    #[test]
+    fn render_config_text_formats_config_snapshot() {
+        let rendered = render_config_text(&json!({
+            "model": "test-model",
+            "served_model_name": ["served-model"],
+            "chat_template": null,
+            "enable_log_requests": true,
+        }));
+        let lines = rendered.lines().collect::<BTreeSet<_>>();
+
+        assert_eq!(
+            lines,
+            BTreeSet::from([
+                "chat_template=None",
+                "enable_log_requests=true",
+                "model=test-model",
+                "served_model_name=[\"served-model\"]",
+            ])
+        );
+        assert_eq!(
+            render_config_text(&Value::String("inline".to_string())),
+            "inline"
+        );
+        assert_eq!(render_config_text(&Value::Null), "None");
+    }
+
+    #[test]
+    fn server_info_env_filter_excludes_sensitive_vllm_keys() {
+        for key in [
+            "VLLM_API_KEY",
+            "VLLM_AUTH_TOKEN",
+            "VLLM_SECRET",
+            "VLLM_PASSWORD",
+            "VLLM_CREDENTIAL_FILE",
+            "vllm_token",
+        ] {
+            assert!(!is_public_vllm_env_key(key), "{key}");
+        }
+    }
+
+    #[test]
+    fn server_info_env_filter_includes_public_vllm_keys() {
+        assert!(is_public_vllm_env_key("VLLM_LOGGING_LEVEL"));
+        assert!(is_public_vllm_env_key("VLLM_USE_MODELSCOPE"));
+        assert!(!is_public_vllm_env_key("OTHER_ENV"));
+    }
+}

rust/src/server/src/state.rs

@@ -1,10 +1,16 @@
 use std::sync::Arc;
 use std::sync::atomic::{AtomicU64, Ordering};
 
+use serde_json::Value;
 use tokio::time::{Duration, Instant, sleep_until};
 use tracing::warn;
 use vllm_chat::ChatLlm;
 use vllm_engine_core_client::EngineCoreClient;
+use vllm_engine_core_client::protocol::lora::LoraRequest;
+
+use crate::lora::{LoadLoraError, LoraManager, LoraModelResolution, UnloadLoraError};
+
+use crate::server_info::{ServerInfoConfigFormat, ServerInfoSnapshot};
 
 const SHUTDOWN_REFCOUNT_POLL_INTERVAL: Duration = Duration::from_millis(100);
 
@@ -19,8 +25,12 @@ pub struct AppState {
     pub enable_log_requests: bool,
     /// Whether to set X-Request-Id on every HTTP response.
     pub enable_request_id_headers: bool,
+    /// Runtime server information returned by `/server_info`, when available.
+    server_info: Option<ServerInfoSnapshot>,
     /// Number of in-flight inference requests currently owned by this frontend.
     server_load: AtomicU64,
+    /// Dynamic LoRA adapter registry.
+    lora_manager: LoraManager,
 }
 
 impl AppState {
@@ -42,7 +52,9 @@ impl AppState {
             chat,
             enable_log_requests: false,
             enable_request_id_headers: false,
+            server_info: None,
             server_load: AtomicU64::new(0),
+            lora_manager: LoraManager::new(),
         }
     }
 
@@ -58,6 +70,20 @@ impl AppState {
         self
     }
 
+    /// Attach the runtime server information snapshot used by `/server_info`.
+    pub(crate) fn with_server_info(mut self, server_info: ServerInfoSnapshot) -> Self {
+        self.server_info = Some(server_info);
+        self
+    }
+
+    /// Build a `/server_info` response payload.
+    pub(crate) fn server_info_response(
+        &self,
+        config_format: ServerInfoConfigFormat,
+    ) -> Option<Value> {
+        self.server_info.as_ref().map(|server_info| server_info.response(config_format))
+    }
+
     /// The primary model name echoed back in API responses (the first served
     /// name).
     pub fn primary_model_name(&self) -> &str {
@@ -69,6 +95,49 @@ impl AppState {
         &self.served_model_names
     }
 
+    /// Return base served model names plus dynamically loaded LoRA adapter
+    /// names.
+    pub async fn served_model_names_with_loras(&self) -> Vec<String> {
+        self.lora_manager.served_model_names(&self.served_model_names).await
+    }
+
+    /// Resolve the requested model against one dynamic LoRA registry snapshot.
+    pub async fn resolve_model_with_loras(&self, model_name: Option<&str>) -> LoraModelResolution {
+        self.lora_manager.resolve_model(&self.served_model_names, model_name).await
+    }
+
+    /// Load one dynamic LoRA adapter and register it as a public model name.
+    pub async fn load_lora(
+        &self,
+        lora_name: String,
+        lora_path: String,
+        load_inplace: bool,
+        is_3d_lora_weight: bool,
+    ) -> Result<LoraRequest, LoadLoraError> {
+        self.lora_manager
+            .load_lora(
+                self.engine_core_client(),
+                &self.served_model_names,
+                lora_name,
+                lora_path,
+                load_inplace,
+                is_3d_lora_weight,
+            )
+            .await
+    }
+
+    /// Remove one dynamic LoRA adapter from the engine and public model
+    /// registry.
+    pub async fn unload_lora(
+        &self,
+        lora_name: &str,
+        lora_int_id: Option<u64>,
+    ) -> Result<LoraRequest, UnloadLoraError> {
+        self.lora_manager
+            .unload_lora(self.engine_core_client(), lora_name, lora_int_id)
+            .await
+    }
+
     /// Return a reference to the underlying engine core client for utility
     /// calls.
     pub(crate) fn engine_core_client(&self) -> &EngineCoreClient {

rust/src/text/src/lower.rs

@@ -40,11 +40,11 @@ pub fn lower_text_request(
         cache_salt: request.cache_salt.clone(),
         priority: request.priority,
         data_parallel_rank: request.data_parallel_rank,
+        lora_request: request.lora_request.clone(),
         // Fields below are currently placeholders.
         arrival_time: None,
         trace_headers: None,
         reasoning_ended: None,
-        lora_request: None,
     };
 
     Ok(PreparedTextRequest {

rust/src/text/src/request.rs

@@ -4,6 +4,7 @@ use enum_as_inner::EnumAsInner;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use vllm_engine_core_client::protocol::StructuredOutputsParams;
+use vllm_engine_core_client::protocol::lora::LoraRequest;
 use vllm_engine_core_client::protocol::multimodal::MmFeatures;
 
 use crate::error::{Error, Result};
@@ -166,6 +167,9 @@ pub struct TextRequest {
     /// Override data parallel rank.
     #[serde(default)]
     pub data_parallel_rank: Option<u32>,
+    /// LoRA adapter selected for this request.
+    #[serde(default)]
+    pub lora_request: Option<LoraRequest>,
 }
 
 impl TextRequest {
@@ -182,6 +186,7 @@ impl TextRequest {
             cache_salt: None,
             add_special_tokens: false,
             data_parallel_rank: None,
+            lora_request: None,
         }
     }
 

rust/src/tool-parser/src/json/internlm2.rs

@@ -34,26 +34,23 @@ const INTERNLM2_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
 /// This Rust port intentionally diverges from
 /// `vllm/tool_parsers/internlm2_tool_parser.py` in two user-visible ways:
 ///
-/// - **Parallel tool calls are supported.** Python silently drops every
-///   `<|action_start|>` block after the first (`current_tool_id > 0` returns
-///   an empty delta); this parser emits every well-formed block with
-///   incrementing `tool_index`. Models that legitimately emit multiple action
-///   blocks therefore produce more tool calls under Rust than under Python.
-/// - **End-marker bytes inside JSON string values are preserved.** Python
-///   does `action.split("<|action_end|>")[0]` which truncates regardless of
-///   JSON context; this parser scans matched braces and quotes so a literal
-///   `<|action_end|>` inside an arguments string is forwarded intact.
+/// - **Parallel tool calls are supported.** Python silently drops every `<|action_start|>` block
+///   after the first (`current_tool_id > 0` returns an empty delta); this parser emits every
+///   well-formed block with incrementing `tool_index`. Models that legitimately emit multiple
+///   action blocks therefore produce more tool calls under Rust than under Python.
+/// - **End-marker bytes inside JSON string values are preserved.** Python does
+///   `action.split("<|action_end|>")[0]` which truncates regardless of JSON context; this parser
+///   scans matched braces and quotes so a literal `<|action_end|>` inside an arguments string is
+///   forwarded intact.
 /// - **Only whitespace is allowed before the `{`.** Python's non-streaming
-///   `action[action.find("{"):]` drops any bytes before the first `{`, but
-///   its streaming path has no equivalent and the model format always emits
-///   `<|plugin|>{...`; this parser allows only whitespace there, matching the
-///   other JSON parsers in this crate.
-/// - **Truncated tool calls error rather than silently dropping.** Python's
-///   streaming wrapper swallows mid-stream errors with `except Exception:
-///   return None` (logging a traceback) while its non-streaming path raises
-///   `JSONDecodeError`; this parser returns an `incomplete InternLM2 tool
-///   call` error from `finish()`, matching the other JSON parsers and Python's
-///   non-streaming behavior.
+///   `action[action.find("{"):]` drops any bytes before the first `{`, but its streaming path has
+///   no equivalent and the model format always emits `<|plugin|>{...`; this parser allows only
+///   whitespace there, matching the other JSON parsers in this crate.
+/// - **Truncated tool calls error rather than silently dropping.** Python's streaming wrapper
+///   swallows mid-stream errors with `except Exception: return None` (logging a traceback) while
+///   its non-streaming path raises `JSONDecodeError`; this parser returns an `incomplete InternLM2
+///   tool call` error from `finish()`, matching the other JSON parsers and Python's non-streaming
+///   behavior.
 ///
 /// # Known unaddressed divergences (TODO)
 ///
@@ -63,21 +60,18 @@ const INTERNLM2_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
 /// Qwen as well. If a real-world InternLM2 deployment hits one of these,
 /// prioritize the corresponding fix.
 ///
-/// - **Arguments value type.** The shared core requires the arguments value
-///   to be a JSON object (`take_json_object` rejects anything not starting
-///   with `{`). Python's `json.dumps(action_dict.get("parameters", ...))`
-///   accepts `null`, arrays, strings, and numbers and round-trips them
-///   verbatim. Models that legitimately emit `"parameters":null` will hard-
+/// - **Arguments value type.** The shared core requires the arguments value to be a JSON object
+///   (`take_json_object` rejects anything not starting with `{`). Python's
+///   `json.dumps(action_dict.get("parameters", ...))` accepts `null`, arrays, strings, and numbers
+///   and round-trips them verbatim. Models that legitimately emit `"parameters":null` will hard-
 ///   fail under Rust.
-/// - **Unknown arguments key.** Python falls back to `{}` via
-///   `action_dict.get("parameters", action_dict.get("arguments", {}))` when
-///   neither key is present; the Rust header parser raises
-///   `parsing failed: invalid InternLM2` for any unrecognized key. A model
-///   that emits a typo (e.g. `"params"`) breaks the whole response.
-/// - **Field order independence.** The header parser requires the JSON keys
-///   to appear in the order `name` then arguments key. Python's
-///   `json.loads` + `dict.get` is order-independent, so a model emitting
-///   `{"parameters":{...},"name":"foo"}` parses in Python but fails in Rust.
+/// - **Unknown arguments key.** Python falls back to `{}` via `action_dict.get("parameters",
+///   action_dict.get("arguments", {}))` when neither key is present; the Rust header parser raises
+///   `parsing failed: invalid InternLM2` for any unrecognized key. A model that emits a typo (e.g.
+///   `"params"`) breaks the whole response.
+/// - **Field order independence.** The header parser requires the JSON keys to appear in the order
+///   `name` then arguments key. Python's `json.loads` + `dict.get` is order-independent, so a model
+///   emitting `{"parameters":{...},"name":"foo"}` parses in Python but fails in Rust.
 pub struct Internlm2ToolParser {
     inner: JsonToolCallParser,
 }

rust/src/tool-parser/src/parameters.rs

@@ -501,15 +501,21 @@ mod tests {
 
         assert_eq!(converted_number_text(&params, "5"), "5");
         assert_eq!(converted_number_text(&params, "5.0"), "5.0");
-        assert_eq!(converted_number_text(&params, "5.00"), "5.00");
-        assert_eq!(converted_number_text(&params, "1e0"), "1e+0");
         assert_eq!(converted_number_text(&params, "5."), "5.0");
         assert_eq!(converted_number_text(&params, "+1"), "1");
         assert_eq!(converted_number_text(&params, "+1.0"), "1.0");
-        assert_eq!(
-            converted_number_text(&params, "9223372036854775807.5"),
-            "9223372036854775807.5"
-        );
+
+        // TODO: we cannot preserve the original number precision by enabling `serde_json`'s
+        // `arbitrary_precision` feature, otherwise the test
+        // `serialized_json_numbers_do_not_leak_serde_private_representation` will fail.
+        // See issue: https://github.com/mitsuhiko/minijinja/issues/641
+
+        // assert_eq!(converted_number_text(&params, "5.00"), "5.00");
+        // assert_eq!(converted_number_text(&params, "1e0"), "1e+0");
+        // assert_eq!(
+        //     converted_number_text(&params, "9223372036854775807.5"),
+        //     "9223372036854775807.5"
+        // );
     }
 
     fn converted_number_text(params: &ToolSchema, value: &str) -> String {

tests/benchmarks/test_custom_dataset_chat_template_kwargs.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import argparse
+import json
+from pathlib import Path
+
+import pytest
+
+from vllm.benchmarks.datasets import get_samples
+
+
+class _RecordingTokenizer:
+    """Minimal tokenizer stub that records the kwargs forwarded to
+    apply_chat_template, so we can assert chat_template_kwargs propagation
+    without loading a real model/template."""
+
+    def __init__(self) -> None:
+        self.captured_kwargs: dict | None = None
+        self.chat_template = "dummy-template"
+
+    def apply_chat_template(
+        self,
+        conversation,
+        add_generation_prompt: bool = True,
+        tokenize: bool = False,
+        **kwargs,
+    ) -> str:
+        self.captured_kwargs = kwargs
+        return conversation[0]["content"]
+
+    def __call__(self, text: str):
+        return argparse.Namespace(input_ids=list(range(len(text.split()))))
+
+
+def _args(dataset_path: str, chat_template_kwargs) -> argparse.Namespace:
+    return argparse.Namespace(
+        dataset_name="custom",
+        dataset_path=dataset_path,
+        disable_shuffle=True,
+        num_prompts=1,
+        custom_output_len=32,
+        skip_chat_template=False,
+        chat_template_kwargs=chat_template_kwargs,
+        no_oversample=False,
+        seed=0,
+        request_id_prefix="",
+    )
+
+
+def _write_one(path: Path) -> None:
+    path.write_text(json.dumps({"prompt": "hello world"}) + "\n")
+
+
+@pytest.mark.benchmark
+def test_chat_template_kwargs_forwarded(tmp_path: Path) -> None:
+    """--chat-template-kwargs must reach the client-side apply_chat_template."""
+    jsonl = tmp_path / "data.jsonl"
+    _write_one(jsonl)
+
+    tok = _RecordingTokenizer()
+    get_samples(_args(str(jsonl), {"thinking": True}), tok)
+
+    assert tok.captured_kwargs == {"thinking": True}
+
+
+@pytest.mark.benchmark
+def test_chat_template_kwargs_default_is_noop(tmp_path: Path) -> None:
+    """When not provided, no extra kwargs are passed (existing behavior)."""
+    jsonl = tmp_path / "data.jsonl"
+    _write_one(jsonl)
+
+    tok = _RecordingTokenizer()
+    get_samples(_args(str(jsonl), None), tok)
+
+    assert tok.captured_kwargs == {}

tests/distributed/test_weight_transfer.py

@@ -291,7 +291,9 @@ def test_nccl_receive_sparse_weights_without_init_raises():
 
     config = WeightTransferConfig(backend="nccl")
     parallel_config = create_mock_parallel_config()
-    engine = NCCLWeightTransferEngine(config, parallel_config)
+    engine = NCCLWeightTransferEngine(
+        config, parallel_config, MagicMock(spec=torch.nn.Module)
+    )
 
     update_info = NCCLWeightTransferUpdateInfo(
         names=["w"],
@@ -526,7 +528,9 @@ def inference_receive_sparse_tensor(
     parallel_config.data_parallel_rank = 0
     parallel_config.data_parallel_index = 0
 
-    engine = NCCLWeightTransferEngine(config, parallel_config)
+    engine = NCCLWeightTransferEngine(
+        config, parallel_config, MagicMock(spec=torch.nn.Module)
+    )
     engine.init_transfer_engine(
         NCCLWeightTransferInitInfo(
             master_address=master_address,

tests/entrypoints/test_chat_utils.py

@@ -13,6 +13,8 @@ from vllm.assets.image import ImageAsset
 from vllm.assets.video import VideoAsset
 from vllm.config import ModelConfig
 from vllm.entrypoints.chat_utils import (
+    ConversationMessage,
+    _postprocess_messages,
     parse_chat_messages,
     parse_chat_messages_async,
 )
@@ -2714,3 +2716,29 @@ async def test_parse_chat_messages_video_vision_chunk_with_uuid_async(
     assert conversation == expected_conversation
     _assert_mm_data_is_vision_chunk_input(mm_data, 1)
     _assert_mm_uuids(mm_uuids, 1, expected_uuids=[video_uuid], modality="vision_chunk")
+
+
+def test_postprocess_messages_null_arguments_string():
+    """arguments="null" must not reach the chat template as Python None.
+
+    json.loads("null") returns None, which causes Jinja2 templates that call
+    tc.arguments.items() to raise 'None' has no attribute 'items'.
+    The function should coerce it to {} instead.
+    """
+    messages: list[ConversationMessage] = [
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [
+                {
+                    "id": "call_1",
+                    "type": "function",
+                    "function": {"name": "get_current_time", "arguments": "null"},
+                }
+            ],
+        }
+    ]
+    _postprocess_messages(messages)
+    tool_calls = messages[0]["tool_calls"]
+    assert tool_calls is not None
+    assert tool_calls[0]["function"]["arguments"] == {}

tests/kernels/attention/test_cpu_attn.py

@@ -258,10 +258,13 @@ def varlen_with_paged_kv(
 
     # KV cache for CPU attention
     cache_dtype = torch.uint8 if is_fp8 else dtype
-    packed_key_cache = torch.empty(
-        num_blocks, num_kv_heads, block_size, head_size, dtype=cache_dtype
+    packed_key_value_cache = torch.empty(
+        num_blocks, num_kv_heads, block_size, head_size * 2, dtype=cache_dtype
     )
-    packed_value_cache = torch.empty_like(packed_key_cache)
+    packed_key_value_cache = packed_key_value_cache.view(
+        (num_blocks, num_kv_heads, block_size * 2, -1)
+    )
+    packed_key_cache, packed_value_cache = packed_key_value_cache.chunk(2, dim=2)
 
     cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
         dim=0, dtype=torch.int32

tests/kernels/moe/test_cpu_fused_moe.py

@@ -20,7 +20,12 @@ EXPERT_NUM = [
 HIDDEN_DIM = [128, 2880]
 INTERMEDIATE_DIM = [128, 2880]
 BATCH_SIZE = [1, 64, 256]
-ACT = [MoEActivation.SILU, MoEActivation.SWIGLUOAI, MoEActivation.GELU]
+ACT = [
+    MoEActivation.SILU,
+    MoEActivation.SWIGLUOAI,
+    MoEActivation.GELU,
+    MoEActivation.GELU_TANH,
+]
 USE_BIAS = [True, False]
 ISA = ["amx", "vec"] if torch.cpu._is_amx_tile_supported() else ["vec"]
 DTYPE = [torch.bfloat16]

tests/kernels/moe/test_cutlass_moe.py

@@ -22,6 +22,7 @@ from vllm.model_executor.layers.fused_moe.config import (
     fp8_w8a8_moe_quant_config,
 )
 from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
+    CutlassExpertsFp4,
     CutlassExpertsFp8,
     run_cutlass_moe_fp8,
 )
@@ -52,6 +53,12 @@ MNK_FACTORS = [
 vllm_config = VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
 
 
+def test_cutlass_moe_supports_gelu_tanh_activation_metadata():
+    assert CutlassExpertsFp8._supports_activation(MoEActivation.GELU_TANH)
+    assert CutlassExpertsFp4._supports_activation(MoEActivation.GELU_TANH)
+    assert CutlassExpertsFp4._supports_activation(MoEActivation.GELU_TANH_NO_MUL)
+
+
 @dataclasses.dataclass
 class MOETensors:
     a: torch.Tensor

tests/kernels/quantization/test_per_token_group_quant.py

@@ -6,6 +6,7 @@ import pytest
 import torch
 
 from vllm.model_executor.layers.quantization.utils import fp8_utils, int8_utils
+from vllm.model_executor.layers.quantization.utils.quant_utils import get_fp8_min_max
 from vllm.platforms import current_platform
 
 
@@ -16,7 +17,9 @@ from vllm.platforms import current_platform
 @pytest.mark.parametrize("tma_aligned", [False, True])
 @pytest.mark.parametrize("scale_ue8m0", [False, True])
 @pytest.mark.parametrize("group_size", [64, 128])
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="Only test on CUDA/ROCm."
+)
 def test_per_token_group_quant_fp8(
     shape, column_major: bool, tma_aligned: bool, scale_ue8m0: bool, group_size: int
 ):
@@ -37,7 +40,7 @@ def test_per_token_group_quant_fp8(
     )
 
     # triton ref
-    with patch("vllm.platforms.current_platform.is_cuda", return_value=False):
+    with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
         ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(
             x,
             group_size,
@@ -77,7 +80,8 @@ def test_per_token_group_quant_fp8(
 )
 @pytest.mark.parametrize("poisoned_scales", [False, True])
 @pytest.mark.skipif(
-    not current_platform.is_cuda(), reason="DeepGEMM not available on this platform"
+    not current_platform.is_cuda_alike(),
+    reason="DeepGEMM not available on this platform",
 )
 def test_per_token_group_quant_fp8_packed(
     num_tokens, hidden_dim, group_size, poisoned_scales
@@ -99,8 +103,8 @@ def test_per_token_group_quant_fp8_packed(
     if poisoned_scales:
         # Call the kernel with poisoned scale buffer to
         # ensure padded indices are correctly zeroed.
-        fp8_dtype = torch.float8_e4m3fn
-        finfo = torch.finfo(fp8_dtype)
+        fp8_dtype = current_platform.fp8_dtype()
+        fp8_min, fp8_max = get_fp8_min_max()
         out_q = torch.empty_like(x, dtype=fp8_dtype)
         out_s_packed = torch.empty_strided(
             (mn, k_num_packed),
@@ -115,8 +119,8 @@ def test_per_token_group_quant_fp8_packed(
             out_s_packed,
             group_size,
             1e-10,
-            finfo.min,
-            finfo.max,
+            fp8_min,
+            fp8_max,
         )
     else:
         out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
@@ -126,7 +130,7 @@ def test_per_token_group_quant_fp8_packed(
         )
 
     # Triton reference (row-major float32 scales, UE8M0)
-    with patch("vllm.platforms.current_platform.is_cuda", return_value=False):
+    with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
         ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(
             x,
             group_size,
@@ -157,7 +161,8 @@ def test_per_token_group_quant_fp8_packed(
 
 
 @pytest.mark.skipif(
-    not current_platform.is_cuda(), reason="DeepGEMM not available on this platform"
+    not current_platform.is_cuda_alike(),
+    reason="DeepGEMM not available on this platform",
 )
 def test_per_token_group_quant_fp8_packed_all_zero():
     """All-zero input must produce well-defined UE8M0 scale bytes via the eps
@@ -216,7 +221,8 @@ def test_per_token_group_quant_fp8_packed_all_zero():
 
 
 @pytest.mark.skipif(
-    not current_platform.is_cuda(), reason="DeepGEMM not available on this platform"
+    not current_platform.is_cuda_alike(),
+    reason="DeepGEMM not available on this platform",
 )
 def test_per_token_group_quant_fp8_packed_mantissa_rounds_up():
     """Inputs whose absmax/max_8bit produces a non-power-of-2 force the
@@ -244,7 +250,7 @@ def test_per_token_group_quant_fp8_packed_mantissa_rounds_up():
         use_ue8m0=True,
     )
 
-    with patch("vllm.platforms.current_platform.is_cuda", return_value=False):
+    with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
         ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(
             x,
             group_size,
@@ -286,7 +292,8 @@ def test_per_token_group_quant_fp8_packed_mantissa_rounds_up():
     ],
 )
 @pytest.mark.skipif(
-    not current_platform.is_cuda(), reason="DeepGEMM not available on this platform"
+    not current_platform.is_cuda_alike(),
+    reason="DeepGEMM not available on this platform",
 )
 def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q(
     num_tokens, hidden_dim
@@ -305,8 +312,8 @@ def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q(
     k_num_packed = (groups_per_row + 3) // 4
     tma_aligned_mn = ((mn + 3) // 4) * 4
 
-    fp8_dtype = torch.float8_e4m3fn
-    finfo = torch.finfo(fp8_dtype)
+    fp8_dtype = current_platform.fp8_dtype()
+    fp8_min, fp8_max = get_fp8_min_max()
     # Allocate output_q with the padded mn extent and pre-fill with 0xFF
     # so the kernel cannot rely on a clean buffer.
     out_q = torch.empty((tma_aligned_mn, hidden_dim), device=device, dtype=fp8_dtype)
@@ -320,11 +327,11 @@ def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q(
     )
 
     torch.ops._C.per_token_group_fp8_quant_packed(
-        x, out_q, out_s_packed, group_size, 1e-10, finfo.min, finfo.max
+        x, out_q, out_s_packed, group_size, 1e-10, fp8_min, fp8_max
     )
 
     # Live rows must match the Triton reference.
-    with patch("vllm.platforms.current_platform.is_cuda", return_value=False):
+    with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
         ref_q, _ = fp8_utils.per_token_group_quant_fp8(x, group_size, use_ue8m0=True)
     assert torch.equal(out_q[:mn], ref_q), "Live region mismatch"
 
@@ -356,7 +363,7 @@ def test_per_token_group_quant_int8(shape, group_size: int):
     )
 
     # triton ref
-    with patch("vllm.platforms.current_platform.is_cuda", return_value=False):
+    with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
         ref_q, ref_s = int8_utils.per_token_group_quant_int8(
             x,
             group_size,

tests/model_executor/layers/test_rocm_unquantized_gemm.py

@@ -41,8 +41,10 @@ def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch):
     assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
 
 
-def test_rocm_unquantized_gemm_gfx1x_n_gt_4_falls_back(monkeypatch):
-    x = torch.randn(5, 64, dtype=torch.float16)
+def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch):
+    # wvSplitK skinny GEMM handles n in [1, 5] (see PR #40687); n > 5 must
+    # fall back to torch.nn.functional.linear.
+    x = torch.randn(6, 64, dtype=torch.float16)
     weight = torch.randn(128, 64, dtype=torch.float16)
 
     monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)

tests/model_executor/model_loader/test_reload.py

@@ -28,6 +28,22 @@ from vllm.model_executor.model_loader.reload.utils import get_layer_tensors
 from vllm.platforms import current_platform
 
 
+def _fp8_reload_unsupported() -> bool:
+    """Whether the FP8 reload/online-quantize tests should be skipped.
+
+    ``supports_fp8()`` returns True on MI250 (gfx90a) because the general
+    quantization paths upcast FP8 weights, but gfx90a has no native FP8 and
+    cannot run these reload models, so treat it as unsupported here.
+    """
+    if not current_platform.supports_fp8():
+        return True
+    if current_platform.is_rocm():
+        from vllm.platforms.rocm import on_gfx90a
+
+        return on_gfx90a()
+    return False
+
+
 class _AliasedBufferLayer(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -284,7 +300,7 @@ def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner):
     if current_platform.device_count() < tp_size:
         pytest.skip(reason="Not enough CUDA devices")
 
-    if "FP8" in base_model and not current_platform.supports_fp8():
+    if "FP8" in base_model and _fp8_reload_unsupported():
         pytest.skip(reason="Requires FP8 support")
 
     with vllm_runner(
@@ -308,7 +324,7 @@ def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner):
 
 def test_kv_scale_reload(vllm_runner):
     """Test reloading a checkpoint that contains k_scale/v_scale weights."""
-    if not current_platform.supports_fp8():
+    if _fp8_reload_unsupported():
         pytest.skip(reason="Requires FP8 support")
 
     model = "nm-testing/Llama-3.2-1B-Instruct-FP8-KV"
@@ -378,7 +394,7 @@ def test_online_quantize_reload(
     if current_platform.device_count() < tp_size:
         pytest.skip(reason="Not enough GPU devices")
 
-    if quantization == "fp8" and not current_platform.supports_fp8():
+    if quantization == "fp8" and _fp8_reload_unsupported():
         pytest.skip(reason="Requires FP8 support")
 
     with vllm_runner(

tests/parser/test_parse.py

@@ -6,7 +6,7 @@ import json
 import pytest
 
 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
-from vllm.parser.abstract_parser import _WrappedParser
+from vllm.parser.abstract_parser import DelegatingParser
 from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser
 from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser
 
@@ -65,9 +65,11 @@ TOOLS = [
 
 
 def make_parser(tokenizer, reasoning=False, tool=False):
-    _WrappedParser.reasoning_parser_cls = ThinkReasoningParser if reasoning else None
-    _WrappedParser.tool_parser_cls = Hermes2ProToolParser if tool else None
-    return _WrappedParser(tokenizer)
+    class TestParser(DelegatingParser):
+        reasoning_parser_cls = ThinkReasoningParser if reasoning else None
+        tool_parser_cls = Hermes2ProToolParser if tool else None
+
+    return TestParser(tokenizer)
 
 
 @pytest.mark.parametrize(