[CI] Make Model Executor test hangs fail fast with a traceback (#43971)

Signed-off-by: khluu <khluu000@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>

.buildkite/hardware_tests/cpu.yaml

@@ -62,11 +62,16 @@ steps:
   source_file_dependencies:
   - vllm/v1/worker/cpu/
   - vllm/v1/worker/gpu/
+  - vllm/v1/sample/ops/topk_topp_triton.py
+  - vllm/v1/sample/ops/topk_topp_sampler.py
+  - tests/v1/sample/test_topk_topp_sampler.py
   commands:
     - |
       bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
       uv pip install git+https://github.com/triton-lang/triton-cpu.git@270e696d
-      VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model"
+      VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model
+      # TODO: move to CPU-Kernel Tests once triton-cpu has a pre-built wheel
+      pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp"
 
 - label: CPU-Quantization Model Tests
   depends_on: []

.buildkite/image_build/image_build.yaml

@@ -6,14 +6,26 @@ steps:
     timeout_in_minutes: 600
     commands:
     - if [[ "$BUILDKITE_BRANCH" == "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $IMAGE_TAG $IMAGE_TAG_LATEST; else .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $IMAGE_TAG; fi
-    # Non-root smoke 1: the default (root) image must still be importable
+    retry:
+      automatic:
+        - exit_status: -1  # Agent was lost
+          limit: 2
+        - exit_status: -10  # Agent was lost
+          limit: 2
+
+  - label: ":docker: :smoking: Non-root smoke tests"
+    key: image-smoke-test
+    depends_on:
+      - image-build
+    commands:
+    # Smoke 1: the default (root) image must still be importable
     # under a non-root UID via `--user 2000:0`. Validates the `vllm` passwd
     # entry + group-0-writable /home/vllm + uv path cleanup from #31959.
     # Uses `import vllm` rather than `vllm serve --help` because the latter
     # instantiates `VllmConfig` which requires a GPU attached to the
     # container.
     - docker run --rm --user 2000:0 --entrypoint python3 "$IMAGE_TAG" -c "import vllm; print(vllm.__version__)"
-    # Non-root smoke 2: assert the non-root enabling invariants are baked
+    # Smoke 2: assert the non-root enabling invariants are baked
     # into the image. Runs as UID 2000:0 via a shell so we can verify
     # filesystem perms + passwd/group file state + wrapper presence without
     # triggering vLLM's GPU-requiring config-init path. The opt-in

.buildkite/scripts/ci-fetch-log.sh

@@ -9,6 +9,13 @@
 # Find <build_number> and <job_uuid> via:
 #   gh pr checks <PR> --repo vllm-project/vllm
 # Each failing row's URL is .../builds/<build_number>#<job_uuid>.
+#
+# Default output path: ci-<build>-<uuid_first_13_chars>.log (e.g.
+# ci-68478-019e6b07-daae.log). Jobs in the same build share the UUID's
+# first 8 chars, so the second segment is needed for uniqueness when
+# fetching multiple jobs in parallel. The script refuses to overwrite an
+# existing output file; pass an explicit path or set CI_FETCH_LOG_FORCE=1
+# to override.
 
 set -euo pipefail
 
@@ -26,12 +33,12 @@ if [ $# -lt 1 ]; then usage; fi
 if [[ "$1" == https://* ]]; then
     BUILD=$(echo "$1" | sed -nE 's#.*/builds/([0-9]+).*#\1#p')
     JOB=$(echo "$1" | grep -oE '[0-9a-f]{8}-[0-9a-f-]+' | head -n 1)
-    OUT="${2:-ci-${BUILD}-${JOB:0:8}.log}"
+    OUT="${2:-}"
 else
     if [ $# -lt 2 ]; then usage; fi
     BUILD="$1"
     JOB="$2"
-    OUT="${3:-ci-${BUILD}-${JOB:0:8}.log}"
+    OUT="${3:-}"
 fi
 
 if [ -z "$BUILD" ] || [ -z "$JOB" ]; then
@@ -39,6 +46,18 @@ if [ -z "$BUILD" ] || [ -z "$JOB" ]; then
     usage
 fi
 
+# Jobs in the same build share the UUID's first segment, so include the
+# second segment (chars 9-13, e.g. "019e6b07-daae") to keep default filenames
+# unique when fetching multiple jobs from one build in parallel.
+if [ -z "$OUT" ]; then
+    OUT="ci-${BUILD}-${JOB:0:13}.log"
+fi
+
+if [ -e "$OUT" ] && [ -z "${CI_FETCH_LOG_FORCE:-}" ]; then
+    echo "Refusing to overwrite existing $OUT (set CI_FETCH_LOG_FORCE=1 or pass an explicit output path)." >&2
+    exit 1
+fi
+
 COOKIES=$(mktemp)
 trap 'rm -f "$COOKIES"' EXIT
 

.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh

@@ -49,6 +49,7 @@ for BACK in "${BACKENDS[@]}"; do
     --data-parallel-size 2 \
     --enable-expert-parallel \
     --enable-eplb \
+    --eplb-config '{"use_async": false}' \
     --trust-remote-code \
     --max-model-len 2048 \
     --all2all-backend "$BACK" \

.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh

@@ -48,7 +48,7 @@ for BACK in "${BACKENDS[@]}"; do
     --enforce-eager \
     --enable-eplb \
     --all2all-backend "$BACK" \
-    --eplb-config '{"window_size":10, "step_interval":100, "num_redundant_experts":0, "log_balancedness":true}' \
+    --eplb-config '{"window_size":10, "step_interval":100, "num_redundant_experts":0, "log_balancedness":true, "use_async":false}' \
     --tensor-parallel-size "${TENSOR_PARALLEL_SIZE}" \
     --data-parallel-size "${DATA_PARALLEL_SIZE}" \
     --enable-expert-parallel \

.buildkite/scripts/tool_call/run-bfcl-eval.sh

@@ -70,7 +70,7 @@ echo "============================================"
 # ---- Install bfcl-eval if missing ----
 if ! python3 -c "import bfcl_eval" 2>/dev/null; then
     echo "Installing bfcl-eval..."
-    pip install "bfcl-eval>=2025.10.20.1,<2026"
+    uv pip install "bfcl-eval>=2025.10.20.1,<2026"
 fi
 
 # ---- Cleanup handler ----
@@ -100,7 +100,7 @@ SERVE_ARGS=(
     --tensor-parallel-size "$TP_SIZE"
     --max-model-len "$MAX_MODEL_LEN"
     --enforce-eager
-    --no-enable-prefix-caching
+    --enable-prefix-caching
 )
 
 # Append reasoning parser if specified

.buildkite/test-amd.yaml

@@ -1484,7 +1484,7 @@ steps:
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt
 
-- label: DeepSeek V2-Lite Accuracy (4xH100-4xMI300) # TBD
+- label: DeepSeek V2-Lite Sync EPLB Accuracy (4xH100-4xMI300) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_4
@@ -1526,7 +1526,7 @@ steps:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
 
-- label: Qwen3-30B-A3B-FP8-block Accuracy (4xH100-4xMI300) # TBD
+- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (4xH100-4xMI300) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_4
@@ -2895,7 +2895,7 @@ steps:
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt
 
-- label: Qwen3-30B-A3B-FP8-block Accuracy (B200-MI355) # TBD
+- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200-MI355) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   agent_pool: mi355_2

.buildkite/test_areas/e2e_integration.yaml

@@ -2,8 +2,8 @@ group: E2E Integration
 depends_on: 
   - image-build
 steps:
-- label: DeepSeek V2-Lite Accuracy
-  key: deepseek-v2-lite-accuracy
+- label: DeepSeek V2-Lite Sync EPLB Accuracy
+  key: deepseek-v2-lite-sync-eplb-accuracy
   timeout_in_minutes: 60
   device: h100
   optional: true
@@ -12,8 +12,8 @@ steps:
   commands:
   - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010
 
-- label: Qwen3-30B-A3B-FP8-block Accuracy
-  key: qwen3-30b-a3b-fp8-block-accuracy
+- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy
+  key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy
   timeout_in_minutes: 60
   device: h100
   optional: true
@@ -22,8 +22,8 @@ steps:
   commands:
   - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020
 
-- label: Qwen3-30B-A3B-FP8-block Accuracy (B200)
-  key: qwen3-30b-a3b-fp8-block-accuracy-b200
+- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200)
+  key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-b200
   timeout_in_minutes: 60
   device: b200-k8s
   optional: true

.buildkite/test_areas/model_executor.yaml

@@ -14,5 +14,12 @@ steps:
   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
+    # Dump tracebacks of all threads if a test hangs, so a wedged GPU/CUDA
+    # init surfaces a stack instead of silently stalling.
+    - export PYTHONFAULTHANDLER=1
+    # Per-test watchdog: a single hung test (e.g. stuck during engine/CUDA
+    # init) fails fast with a traceback instead of running until the global
+    # build timeout. The `thread` method also handles hangs inside C/CUDA
+    # calls that the signal method cannot interrupt.
+    - pytest -v -s model_executor -m '(not slow_test)' --timeout=900 --timeout-method=thread
+    - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread

.readthedocs.yaml

@@ -9,8 +9,8 @@ build:
     python: "3.12"
   jobs:
     post_checkout:
-      - bash docs/pre_run_check.sh
       - git fetch origin main --unshallow --no-tags --filter=blob:none || true
+      - bash docs/pre_run_check.sh
     pre_create_environment:
       - pip install uv
     create_environment:

CMakeLists.txt

@@ -305,10 +305,6 @@ endif()
 #
 
 set(VLLM_EXT_SRC
-  "csrc/cache_kernels.cu"
-  "csrc/cache_kernels_fused.cu"
-  "csrc/attention/paged_attention_v1.cu"
-  "csrc/attention/paged_attention_v2.cu"
   "csrc/cuda_view.cu"
   "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
   "csrc/quantization/activation_kernels.cu"
@@ -647,7 +643,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/attention/merge_attn_states.cu"
     "csrc/libtorch_stable/sampler.cu"
     "csrc/libtorch_stable/topk.cu"
-    "csrc/libtorch_stable/mamba/selective_scan_fwd.cu")
+    "csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
+    "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")
 
   if(VLLM_GPU_LANG STREQUAL "CUDA")
     list(APPEND VLLM_STABLE_EXT_SRC
@@ -924,13 +924,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
       SRCS "${SRCS}"
       CUDA_ARCHS "${FP4_ARCHS}")
     list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
-    # nvfp4_kv_cache_kernels uses non-stable torch API and is called directly
-    # from cache_kernels.cu, so it belongs in _C rather than _C_stable.
-    set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
+    set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
     set_gencode_flags_for_srcs(
       SRCS "${NVFP4_KV_SRC}"
       CUDA_ARCHS "${FP4_ARCHS}")
-    target_sources(_C PRIVATE ${NVFP4_KV_SRC})
+    list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}")
     target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
@@ -960,11 +958,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
       SRCS "${SRCS}"
       CUDA_ARCHS "${FP4_ARCHS}")
     list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
-    set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
+    set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
     set_gencode_flags_for_srcs(
       SRCS "${NVFP4_KV_SRC}"
       CUDA_ARCHS "${FP4_ARCHS}")
-    target_sources(_C PRIVATE ${NVFP4_KV_SRC})
+    list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}")
     target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
@@ -1286,6 +1284,14 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/rocm/skinny_gemms.cu"
     "csrc/rocm/attention.cu")
 
+  set(VLLM_ROCM_HAS_GFX1100 OFF)
+  if(VLLM_GPU_ARCHES MATCHES "gfx1100")
+    set(VLLM_ROCM_HAS_GFX1100 ON)
+    list(APPEND VLLM_ROCM_EXT_SRC
+      "csrc/rocm/q_gemm_rdna3.cu"
+      "csrc/rocm/q_gemm_rdna3_wmma.cu")
+  endif()
+
   define_extension_target(
     _rocm_C
     DESTINATION vllm
@@ -1295,6 +1301,10 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
     ARCHITECTURES ${VLLM_GPU_ARCHES}
     USE_SABI 3
     WITH_SOABI)
+
+  if(VLLM_ROCM_HAS_GFX1100)
+    target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GFX1100)
+  endif()
 endif()
 
 # Must run after the last HIP `define_extension_target` so every extension

csrc/libtorch_stable/activation_kernels.cu

@@ -4,7 +4,7 @@
 #include <cmath>
 
 #include "../cuda_compat.h"
-#include "../cuda_vec_utils.cuh"
+#include "cuda_vec_utils.cuh"
 #include "dispatch_utils.h"
 #include "torch_utils.h"
 

csrc/libtorch_stable/attention/attention_kernels.cuh

@@ -17,21 +17,18 @@
  * limitations under the License.
  */
 
-#include <torch/all.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
 #include <algorithm>
 
-#include "attention_dtypes.h"
+#include "../../attention/attention_dtypes.h"
 #include "attention_utils.cuh"
-#include "../cuda_compat.h"
+#include "../../cuda_compat.h"
 
 #ifdef USE_ROCM
   #include <hip/hip_bf16.h>
-  #include "../quantization/w8a8/fp8/amd/quant_utils.cuh"
+  #include "../../quantization/w8a8/fp8/amd/quant_utils.cuh"
 typedef __hip_bfloat16 __nv_bfloat16;
 #else
-  #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
+  #include "../../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
 #endif
 
 #define MAX(a, b) ((a) > (b) ? (a) : (b))

csrc/libtorch_stable/attention/attention_utils.cuh

@@ -18,8 +18,8 @@
  */
 #pragma once
 
-#include "../cuda_compat.h"
-#include "attention_dtypes.h"
+#include "../../cuda_compat.h"
+#include "../../attention/attention_dtypes.h"
 
 #include <float.h>
 #include <type_traits>

csrc/libtorch_stable/attention/merge_attn_states.cu

@@ -7,7 +7,7 @@
 #include <torch/headeronly/core/ScalarType.h>
 
 #include "../../attention/attention_dtypes.h"
-#include "../../attention/attention_utils.cuh"
+#include "attention_utils.cuh"
 #include "../../quantization/w8a8/fp8/common.cuh"
 
 namespace vllm {

csrc/libtorch_stable/attention/paged_attention_v1.cu

@@ -16,8 +16,9 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+#include "../torch_utils.h"
 #include "attention_kernels.cuh"
-#include "../cuda_compat.h"
+#include "../../cuda_compat.h"
 
 #define MAX(a, b) ((a) > (b) ? (a) : (b))
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -44,13 +45,15 @@ template <typename T, typename CACHE_T, int BLOCK_SIZE,
           vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
           int NUM_THREADS = 128>
 void paged_attention_v1_launcher(
-    torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache,
-    torch::Tensor& value_cache, int num_kv_heads, float scale,
-    torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len,
-    const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,
-    torch::Tensor& v_scale, const int tp_rank,
-    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
-    const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
+    torch::stable::Tensor& out, torch::stable::Tensor& query,
+    torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
+    int num_kv_heads, float scale, torch::stable::Tensor& block_tables,
+    torch::stable::Tensor& seq_lens, int max_seq_len,
+    const std::optional<torch::stable::Tensor>& alibi_slopes,
+    torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale,
+    const int tp_rank, const int blocksparse_local_blocks,
+    const int blocksparse_vert_stride, const int blocksparse_block_size,
+    const int blocksparse_head_sliding_step) {
   int num_seqs = query.size(0);
   int num_heads = query.size(1);
   int head_size = query.size(2);
@@ -69,8 +72,8 @@ void paged_attention_v1_launcher(
   T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
   CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
   CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
-  int* block_tables_ptr = block_tables.data_ptr<int>();
-  int* seq_lens_ptr = seq_lens.data_ptr<int>();
+  int* block_tables_ptr = block_tables.mutable_data_ptr<int>();
+  int* seq_lens_ptr = seq_lens.mutable_data_ptr<int>();
   const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
   const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
 
@@ -85,8 +88,9 @@ void paged_attention_v1_launcher(
 
   dim3 grid(num_heads, num_seqs, 1);
   dim3 block(NUM_THREADS);
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      query.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
   switch (head_size) {
     // NOTE(woosuk): To reduce the compilation time, we only compile for the
     // head sizes that we use in the model. However, we can easily extend this
@@ -119,7 +123,7 @@ void paged_attention_v1_launcher(
       LAUNCH_PAGED_ATTENTION_V1(256);
       break;
     default:
-      TORCH_CHECK(false, "Unsupported head size: ", head_size);
+      STD_TORCH_CHECK(false, "Unsupported head size: ", head_size);
       break;
   }
 }
@@ -141,43 +145,43 @@ void paged_attention_v1_launcher(
 
 // NOTE(woosuk): To reduce the compilation time, we omitted block sizes
 // 1, 2, 4, 64, 128, 256.
-#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)         \
-  switch (block_size) {                                           \
-    case 8:                                                       \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);         \
-      break;                                                      \
-    case 16:                                                      \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);        \
-      break;                                                      \
-    case 32:                                                      \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);        \
-      break;                                                      \
-    default:                                                      \
-      TORCH_CHECK(false, "Unsupported block size: ", block_size); \
-      break;                                                      \
+#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)             \
+  switch (block_size) {                                               \
+    case 8:                                                           \
+      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);             \
+      break;                                                          \
+    case 16:                                                          \
+      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);            \
+      break;                                                          \
+    case 32:                                                          \
+      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);            \
+      break;                                                          \
+    default:                                                          \
+      STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \
+      break;                                                          \
   }
 
 void paged_attention_v1(
-    torch::Tensor& out,    // [num_seqs, num_heads, head_size]
-    torch::Tensor& query,  // [num_seqs, num_heads, head_size]
-    torch::Tensor&
+    torch::stable::Tensor& out,    // [num_seqs, num_heads, head_size]
+    torch::stable::Tensor& query,  // [num_seqs, num_heads, head_size]
+    torch::stable::Tensor&
         key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
-    torch::Tensor&
+    torch::stable::Tensor&
         value_cache,       // [num_blocks, num_heads, head_size, block_size]
     int64_t num_kv_heads,  // [num_heads]
     double scale,
-    torch::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
-    torch::Tensor& seq_lens,      // [num_seqs]
+    torch::stable::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
+    torch::stable::Tensor& seq_lens,      // [num_seqs]
     int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::Tensor& k_scale,
-    torch::Tensor& v_scale, const int64_t tp_rank,
+    const std::optional<torch::stable::Tensor>& alibi_slopes,
+    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
+    torch::stable::Tensor& v_scale, const int64_t tp_rank,
     const int64_t blocksparse_local_blocks,
     const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
     const int64_t blocksparse_head_sliding_step) {
   const bool is_block_sparse = (blocksparse_vert_stride > 1);
 
-  DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype,
+  DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype,
                              CALL_V1_LAUNCHER_BLOCK_SIZE)
 }
 

csrc/libtorch_stable/attention/paged_attention_v2.cu

@@ -16,8 +16,9 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+#include "../torch_utils.h"
 #include "attention_kernels.cuh"
-#include "../cuda_compat.h"
+#include "../../cuda_compat.h"
 
 #define MAX(a, b) ((a) > (b) ? (a) : (b))
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -44,14 +45,16 @@ template <typename T, typename CACHE_T, int BLOCK_SIZE,
           vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
           int NUM_THREADS = 128, int PARTITION_SIZE = 512>
 void paged_attention_v2_launcher(
-    torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
-    torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
-    torch::Tensor& value_cache, int num_kv_heads, float scale,
-    torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len,
-    const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,
-    torch::Tensor& v_scale, const int tp_rank,
-    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
-    const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
+    torch::stable::Tensor& out, torch::stable::Tensor& exp_sums,
+    torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out,
+    torch::stable::Tensor& query, torch::stable::Tensor& key_cache,
+    torch::stable::Tensor& value_cache, int num_kv_heads, float scale,
+    torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens,
+    int max_seq_len, const std::optional<torch::stable::Tensor>& alibi_slopes,
+    torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale,
+    const int tp_rank, const int blocksparse_local_blocks,
+    const int blocksparse_vert_stride, const int blocksparse_block_size,
+    const int blocksparse_head_sliding_step) {
   int num_seqs = query.size(0);
   int num_heads = query.size(1);
   int head_size = query.size(2);
@@ -73,8 +76,8 @@ void paged_attention_v2_launcher(
   T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
   CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
   CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
-  int* block_tables_ptr = block_tables.data_ptr<int>();
-  int* seq_lens_ptr = seq_lens.data_ptr<int>();
+  int* block_tables_ptr = block_tables.mutable_data_ptr<int>();
+  int* seq_lens_ptr = seq_lens.mutable_data_ptr<int>();
   const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
   const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
 
@@ -91,8 +94,9 @@ void paged_attention_v2_launcher(
   int reduce_shared_mem_size = 2 * max_num_partitions * sizeof(float);
 
   dim3 block(NUM_THREADS);
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      query.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
   switch (head_size) {
     // NOTE(woosuk): To reduce the compilation time, we only compile for the
     // head sizes that we use in the model. However, we can easily extend this
@@ -125,7 +129,7 @@ void paged_attention_v2_launcher(
       LAUNCH_PAGED_ATTENTION_V2(256);
       break;
     default:
-      TORCH_CHECK(false, "Unsupported head size: ", head_size);
+      STD_TORCH_CHECK(false, "Unsupported head size: ", head_size);
       break;
   }
 }
@@ -148,46 +152,48 @@ void paged_attention_v2_launcher(
 
 // NOTE(woosuk): To reduce the compilation time, we omitted block sizes
 // 1, 2, 4, 64, 128, 256.
-#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)         \
-  switch (block_size) {                                           \
-    case 8:                                                       \
-      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);         \
-      break;                                                      \
-    case 16:                                                      \
-      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);        \
-      break;                                                      \
-    case 32:                                                      \
-      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);        \
-      break;                                                      \
-    default:                                                      \
-      TORCH_CHECK(false, "Unsupported block size: ", block_size); \
-      break;                                                      \
+#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)             \
+  switch (block_size) {                                               \
+    case 8:                                                           \
+      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);             \
+      break;                                                          \
+    case 16:                                                          \
+      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);            \
+      break;                                                          \
+    case 32:                                                          \
+      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);            \
+      break;                                                          \
+    default:                                                          \
+      STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \
+      break;                                                          \
   }
 
 void paged_attention_v2(
-    torch::Tensor& out,         // [num_seqs, num_heads, head_size]
-    torch::Tensor& exp_sums,    // [num_seqs, num_heads, max_num_partitions]
-    torch::Tensor& max_logits,  // [num_seqs, num_heads, max_num_partitions]
-    torch::Tensor&
+    torch::stable::Tensor& out,  // [num_seqs, num_heads, head_size]
+    torch::stable::Tensor&
+        exp_sums,  // [num_seqs, num_heads, max_num_partitions]
+    torch::stable::Tensor&
+        max_logits,  // [num_seqs, num_heads, max_num_partitions]
+    torch::stable::Tensor&
         tmp_out,  // [num_seqs, num_heads, max_num_partitions, head_size]
-    torch::Tensor& query,  // [num_seqs, num_heads, head_size]
-    torch::Tensor&
+    torch::stable::Tensor& query,  // [num_seqs, num_heads, head_size]
+    torch::stable::Tensor&
         key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
-    torch::Tensor&
+    torch::stable::Tensor&
         value_cache,       // [num_blocks, num_heads, head_size, block_size]
     int64_t num_kv_heads,  // [num_heads]
     double scale,
-    torch::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
-    torch::Tensor& seq_lens,      // [num_seqs]
+    torch::stable::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
+    torch::stable::Tensor& seq_lens,      // [num_seqs]
     int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::Tensor& k_scale,
-    torch::Tensor& v_scale, const int64_t tp_rank,
+    const std::optional<torch::stable::Tensor>& alibi_slopes,
+    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
+    torch::stable::Tensor& v_scale, const int64_t tp_rank,
     const int64_t blocksparse_local_blocks,
     const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
     const int64_t blocksparse_head_sliding_step) {
   const bool is_block_sparse = (blocksparse_vert_stride > 1);
-  DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype,
+  DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype,
                              CALL_V2_LAUNCHER_BLOCK_SIZE)
 }
 

csrc/libtorch_stable/cache_kernels_fused.cu

@@ -1,15 +1,13 @@
-#include <torch/all.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-
-#include "cuda_compat.h"
+#include "torch_utils.h"
 #include "dispatch_utils.h"
 
-#include "quantization/w8a8/fp8/common.cuh"
+#include "../cuda_compat.h"
+
+#include "../quantization/w8a8/fp8/common.cuh"
 #ifdef USE_ROCM
-  #include "quantization/w8a8/fp8/amd/quant_utils.cuh"
+  #include "../quantization/w8a8/fp8/amd/quant_utils.cuh"
 #else
-  #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh"
+  #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
 #endif
 
 #ifdef USE_ROCM
@@ -164,43 +162,52 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(
 
 }  // namespace vllm
 
-#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE)     \
-  do {                                                                        \
-    VLLM_DISPATCH_FLOATING_TYPES(q_pe.scalar_type(), "qk_scalar_type", [&] {  \
-      using qk_t = scalar_t;                                                  \
-      VLLM_DISPATCH_FLOATING_TYPES(                                           \
-          rope_cos_sin_cache.scalar_type(), "rope_cos_sin_cache_scalar_type", \
-          [&] {                                                               \
-            using cos_sin_t = scalar_t;                                       \
-            if (rope_is_neox) {                                               \
-              vllm::concat_and_cache_mla_rope_fused_kernel<                   \
-                  qk_t, cos_sin_t, true, RAW_KV_T, CACHE_T, KV_DTYPE>         \
-                  <<<grid, block, 0, stream>>>(                               \
-                      positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(),   \
-                      k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(),           \
-                      rope_cos_sin_cache.data_ptr<cos_sin_t>(), rot_dim,      \
-                      q_pe_stride_token, q_pe_stride_head, k_pe_stride,       \
-                      kv_c_stride, num_q_heads,                               \
-                      reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()),        \
-                      slot_mapping.data_ptr<int64_t>(), block_stride,         \
-                      entry_stride, kv_lora_rank, block_size,                 \
-                      kv_cache_quant_scale.data_ptr<float>());                \
-            } else {                                                          \
-              vllm::concat_and_cache_mla_rope_fused_kernel<                   \
-                  qk_t, cos_sin_t, false, RAW_KV_T, CACHE_T, KV_DTYPE>        \
-                  <<<grid, block, 0, stream>>>(                               \
-                      positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(),   \
-                      k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(),           \
-                      rope_cos_sin_cache.data_ptr<cos_sin_t>(), rot_dim,      \
-                      q_pe_stride_token, q_pe_stride_head, k_pe_stride,       \
-                      kv_c_stride, num_q_heads,                               \
-                      reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()),        \
-                      slot_mapping.data_ptr<int64_t>(), block_stride,         \
-                      entry_stride, kv_lora_rank, block_size,                 \
-                      kv_cache_quant_scale.data_ptr<float>());                \
-            }                                                                 \
-          });                                                                 \
-    });                                                                       \
+#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE)  \
+  do {                                                                     \
+    VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                   \
+        q_pe.scalar_type(), "qk_scalar_type", [&] {                        \
+          using qk_t = scalar_t;                                           \
+          VLLM_STABLE_DISPATCH_FLOATING_TYPES(                             \
+              rope_cos_sin_cache.scalar_type(),                            \
+              "rope_cos_sin_cache_scalar_type", [&] {                      \
+                using cos_sin_t = scalar_t;                                \
+                if (rope_is_neox) {                                        \
+                  vllm::concat_and_cache_mla_rope_fused_kernel<            \
+                      qk_t, cos_sin_t, true, RAW_KV_T, CACHE_T, KV_DTYPE>  \
+                      <<<grid, block, 0, stream>>>(                        \
+                          positions.const_data_ptr<int64_t>(),             \
+                          q_pe.mutable_data_ptr<qk_t>(),                   \
+                          k_pe.mutable_data_ptr<qk_t>(),                   \
+                          kv_c.const_data_ptr<qk_t>(),                     \
+                          rope_cos_sin_cache.const_data_ptr<cos_sin_t>(),  \
+                          rot_dim, q_pe_stride_token, q_pe_stride_head,    \
+                          k_pe_stride, kv_c_stride, num_q_heads,           \
+                          reinterpret_cast<CACHE_T*>(                      \
+                              kv_cache.mutable_data_ptr()),                \
+                          slot_mapping.const_data_ptr<int64_t>(),          \
+                          block_stride, entry_stride, kv_lora_rank,        \
+                          block_size,                                      \
+                          kv_cache_quant_scale.const_data_ptr<float>());   \
+                } else {                                                   \
+                  vllm::concat_and_cache_mla_rope_fused_kernel<            \
+                      qk_t, cos_sin_t, false, RAW_KV_T, CACHE_T, KV_DTYPE> \
+                      <<<grid, block, 0, stream>>>(                        \
+                          positions.const_data_ptr<int64_t>(),             \
+                          q_pe.mutable_data_ptr<qk_t>(),                   \
+                          k_pe.mutable_data_ptr<qk_t>(),                   \
+                          kv_c.const_data_ptr<qk_t>(),                     \
+                          rope_cos_sin_cache.const_data_ptr<cos_sin_t>(),  \
+                          rot_dim, q_pe_stride_token, q_pe_stride_head,    \
+                          k_pe_stride, kv_c_stride, num_q_heads,           \
+                          reinterpret_cast<CACHE_T*>(                      \
+                              kv_cache.mutable_data_ptr()),                \
+                          slot_mapping.const_data_ptr<int64_t>(),          \
+                          block_stride, entry_stride, kv_lora_rank,        \
+                          block_size,                                      \
+                          kv_cache_quant_scale.const_data_ptr<float>());   \
+                }                                                          \
+              });                                                          \
+        });                                                                \
   } while (false)
 
 // Executes RoPE on q_pe and k_pe, then writes k_pe and kv_c in the kv cache.
@@ -208,64 +215,69 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(
 // Replaces DeepseekScalingRotaryEmbedding.self.rotary_emb and
 // concat_and_cache_mla.
 void concat_and_cache_mla_rope_fused(
-    torch::Tensor& positions,           // [num_tokens]
-    torch::Tensor& q_pe,                // [num_tokens, num_q_heads, rot_dim]
-    torch::Tensor& k_pe,                // [num_tokens, rot_dim]
-    torch::Tensor& kv_c,                // [num_tokens, kv_lora_rank]
-    torch::Tensor& rope_cos_sin_cache,  // [max_position, rot_dim]
+    torch::stable::Tensor& positions,  // [num_tokens]
+    torch::stable::Tensor& q_pe,       // [num_tokens, num_q_heads, rot_dim]
+    torch::stable::Tensor& k_pe,       // [num_tokens, rot_dim]
+    torch::stable::Tensor& kv_c,       // [num_tokens, kv_lora_rank]
+    torch::stable::Tensor& rope_cos_sin_cache,  // [max_position, rot_dim]
     bool rope_is_neox,
-    torch::Tensor& slot_mapping,  // [num_tokens] or [num_actual_tokens]
-    torch::Tensor&
+    torch::stable::Tensor& slot_mapping,  // [num_tokens] or [num_actual_tokens]
+    torch::stable::Tensor&
         kv_cache,  // [num_blocks, block_size, (kv_lora_rank + rot_dim)]
-    const std::string& kv_cache_dtype, torch::Tensor& kv_cache_quant_scale) {
+    const std::string& kv_cache_dtype,
+    torch::stable::Tensor& kv_cache_quant_scale) {
   // NOTE(woosuk): In vLLM V1, query/key/position.size(0) can be different from
   // slot_mapping.size(0) because of padding for CUDA graphs.
-  // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
-  // both include padding.
-  // In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0)
-  // since key includes padding for CUDA graphs, while slot_mapping does not.
-  // In this case, slot_mapping.size(0) represents the actual number of tokens
+  // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0)
+  // because both include padding.
+  // In vLLM V1, however, key.size(0) can be larger than
+  // slot_mapping.size(0) since key includes padding for CUDA graphs,
+  // while slot_mapping does not. In this case,
+  // slot_mapping.size(0) represents the actual number of tokens
   // before padding.
-  // For compatibility with both cases, we use slot_mapping.size(0) as the
-  // number of tokens.
-  int num_tokens = slot_mapping.size(0);
-  int num_padded_tokens = q_pe.size(0);
-  TORCH_CHECK_GE(num_padded_tokens, num_tokens);
+  // For compatibility with both cases, we use slot_mapping.size(0) as
+  // the number of tokens.
+  const int64_t num_tokens = slot_mapping.size(0);
+  const int64_t num_padded_tokens = q_pe.size(0);
+  STD_TORCH_CHECK(num_padded_tokens >= num_tokens);
 
   const int num_q_heads = q_pe.size(1);
   const int rot_dim = q_pe.size(2);
   const int kv_lora_rank = kv_c.size(1);
 
-  TORCH_CHECK_EQ(positions.size(0), num_padded_tokens);
-  TORCH_CHECK_EQ(positions.dim(), 1);
-  TORCH_CHECK_EQ(positions.scalar_type(), c10::ScalarType::Long);
+  STD_TORCH_CHECK(positions.size(0) == num_padded_tokens);
+  STD_TORCH_CHECK(positions.dim() == 1);
+  STD_TORCH_CHECK(positions.scalar_type() ==
+                  torch::headeronly::ScalarType::Long);
 
-  TORCH_CHECK_EQ(q_pe.dim(), 3);
-  TORCH_CHECK_EQ(q_pe.size(0), num_padded_tokens);
-  TORCH_CHECK_EQ(q_pe.size(1), num_q_heads);
-  TORCH_CHECK_EQ(q_pe.size(2), rot_dim);
+  STD_TORCH_CHECK(q_pe.dim() == 3);
+  STD_TORCH_CHECK(q_pe.size(0) == num_padded_tokens);
+  STD_TORCH_CHECK(q_pe.size(1) == num_q_heads);
+  STD_TORCH_CHECK(q_pe.size(2) == rot_dim);
 
-  TORCH_CHECK_EQ(k_pe.dim(), 2);
-  TORCH_CHECK_EQ(k_pe.size(0), num_padded_tokens);
-  TORCH_CHECK_EQ(k_pe.size(1), rot_dim);
-  TORCH_CHECK_EQ(k_pe.scalar_type(), q_pe.scalar_type());
+  STD_TORCH_CHECK(k_pe.dim() == 2);
+  STD_TORCH_CHECK(k_pe.size(0) == num_padded_tokens);
+  STD_TORCH_CHECK(k_pe.size(1) == rot_dim);
+  STD_TORCH_CHECK(k_pe.scalar_type() == q_pe.scalar_type());
 
-  TORCH_CHECK_EQ(kv_c.dim(), 2);
-  TORCH_CHECK_EQ(kv_c.size(0), num_padded_tokens);
-  TORCH_CHECK_EQ(kv_c.size(1), kv_lora_rank);
-  TORCH_CHECK_EQ(kv_c.scalar_type(), q_pe.scalar_type());
-  TORCH_CHECK_EQ(kv_c.dtype(), q_pe.dtype());
+  STD_TORCH_CHECK(kv_c.dim() == 2);
+  STD_TORCH_CHECK(kv_c.size(0) == num_padded_tokens);
+  STD_TORCH_CHECK(kv_c.size(1) == kv_lora_rank);
+  STD_TORCH_CHECK(kv_c.scalar_type() == q_pe.scalar_type());
 
-  TORCH_CHECK_EQ(rope_cos_sin_cache.size(1), rot_dim);
+  STD_TORCH_CHECK(rope_cos_sin_cache.size(1) == rot_dim);
+  STD_TORCH_CHECK(rope_cos_sin_cache.scalar_type() == q_pe.scalar_type());
 
-  TORCH_CHECK_EQ(slot_mapping.size(0), num_tokens);
-  TORCH_CHECK_EQ(slot_mapping.scalar_type(), c10::ScalarType::Long);
+  STD_TORCH_CHECK(slot_mapping.size(0) == num_tokens);
+  STD_TORCH_CHECK(slot_mapping.scalar_type() ==
+                  torch::headeronly::ScalarType::Long);
 
-  TORCH_CHECK_EQ(kv_cache.size(2), kv_lora_rank + rot_dim);
-  TORCH_CHECK_EQ(kv_cache.dim(), 3);
+  STD_TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + rot_dim);
+  STD_TORCH_CHECK(kv_cache.dim() == 3);
 
-  TORCH_CHECK_EQ(kv_cache_quant_scale.numel(), 1);
-  TORCH_CHECK_EQ(kv_cache_quant_scale.scalar_type(), c10::ScalarType::Float);
+  STD_TORCH_CHECK(kv_cache_quant_scale.numel() == 1);
+  STD_TORCH_CHECK(kv_cache_quant_scale.scalar_type() ==
+                  torch::headeronly::ScalarType::Float);
 
   int64_t q_pe_stride_token = q_pe.stride(0);
   int64_t q_pe_stride_head = q_pe.stride(1);
@@ -286,9 +298,10 @@ void concat_and_cache_mla_rope_fused(
   dim3 grid(num_tokens, 1, 1);
   dim3 block(thread_block_size, 1, 1);
 
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(positions));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      positions.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
 
-  DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype,
+  DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype,
                              CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED);
 }

csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu

@@ -17,11 +17,8 @@
 #define NVFP4_ENABLE_ELTS16 1
 #include "libtorch_stable/quantization/fp4/nvfp4_utils.cuh"
 
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <torch/all.h>
-
-#include "dispatch_utils.h"
+#include "libtorch_stable/dispatch_utils.h"
+#include "libtorch_stable/torch_utils.h"
 
 namespace vllm {
 
@@ -184,12 +181,13 @@ __global__ void reshape_and_cache_nvfp4_kernel(
 // Receives key_cache/value_cache as kv_cache[:, 0] and kv_cache[:, 1].
 // Each KV side contains both data and scale:
 //   page = [K_data | K_scale | V_data | V_scale]
-void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
-                                      torch::Tensor& key_cache,
-                                      torch::Tensor& value_cache,
-                                      torch::Tensor& slot_mapping,
-                                      torch::Tensor& k_scale,
-                                      torch::Tensor& v_scale) {
+void reshape_and_cache_nvfp4_dispatch(torch::stable::Tensor& key,
+                                      torch::stable::Tensor& value,
+                                      torch::stable::Tensor& key_cache,
+                                      torch::stable::Tensor& value_cache,
+                                      torch::stable::Tensor& slot_mapping,
+                                      torch::stable::Tensor& k_scale,
+                                      torch::stable::Tensor& v_scale) {
   int num_tokens = slot_mapping.size(0);
   int num_heads = key.size(1);
   int head_size = key.size(2);
@@ -200,17 +198,18 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
   // key_cache is kv_cache[:, 0] with shape
   // [num_blocks, block_size, num_heads, full_dim] in logical order.
   // Strides encode the physical layout (HND or NHD).
-  TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D");
-  TORCH_CHECK(key_cache.size(3) == full_dim,
-              "key_cache last dim must be data_dim + scale_dim, got ",
-              key_cache.size(3), " expected ", full_dim);
+  STD_TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D");
+  STD_TORCH_CHECK(key_cache.size(3) == full_dim,
+                  "key_cache last dim must be data_dim + scale_dim, got ",
+                  key_cache.size(3), " expected ", full_dim);
 
   int block_size = key_cache.size(1);
 
-  TORCH_CHECK(head_size % 16 == 0,
-              "head_size must be divisible by 16 for NVFP4 KV cache");
-  TORCH_CHECK(block_size % 4 == 0,
-              "block_size must be divisible by 4 for NVFP4 KV cache swizzle");
+  STD_TORCH_CHECK(head_size % 16 == 0,
+                  "head_size must be divisible by 16 for NVFP4 KV cache");
+  STD_TORCH_CHECK(block_size % 4 == 0,
+                  "block_size must be divisible by 4 for NVFP4 KV cache "
+                  "swizzle");
 
   // Detect physical layout from strides (based on full_dim).
   // HND: head stride > block_offset stride.
@@ -230,8 +229,9 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
   // Scale follows data within each KV side.
   int64_t data_per_kv = (int64_t)num_heads * block_size * data_dim;
 
-  uint8_t* key_scale_ptr = key_cache.data_ptr<uint8_t>() + data_per_kv;
-  uint8_t* value_scale_ptr = value_cache.data_ptr<uint8_t>() + data_per_kv;
+  uint8_t* key_scale_ptr = key_cache.mutable_data_ptr<uint8_t>() + data_per_kv;
+  uint8_t* value_scale_ptr =
+      value_cache.mutable_data_ptr<uint8_t>() + data_per_kv;
 
   // Scale strides: same page stride, inner strides from layout.
   int64_t scale_block_stride = data_block_stride;
@@ -244,8 +244,8 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
     scale_block_offset_stride = (int64_t)num_heads * scale_dim;
   }
 
-  const float* k_scale_ptr = k_scale.data_ptr<float>();
-  const float* v_scale_ptr = v_scale.data_ptr<float>();
+  const float* k_scale_ptr = k_scale.const_data_ptr<float>();
+  const float* v_scale_ptr = v_scale.const_data_ptr<float>();
 
   int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
   int total_groups = num_heads * groups_per_head;
@@ -256,20 +256,22 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
   dim3 grid(num_tokens);
   dim3 block(num_threads);
 
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      key.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
 
-  AT_DISPATCH_REDUCED_FLOATING_TYPES(
+  VLLM_STABLE_DISPATCH_HALF_TYPES(
       key.scalar_type(), "reshape_and_cache_nvfp4", [&] {
         vllm::reshape_and_cache_nvfp4_kernel<scalar_t>
             <<<grid, block, 0, stream>>>(
-                key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
-                key_cache.data_ptr<uint8_t>(), value_cache.data_ptr<uint8_t>(),
-                key_scale_ptr, value_scale_ptr,
-                slot_mapping.data_ptr<int64_t>(), k_scale_ptr, v_scale_ptr,
-                key.stride(0), value.stride(0), num_heads, head_size,
-                block_size, data_block_stride, data_head_stride,
-                data_block_offset_stride, scale_block_stride, scale_head_stride,
-                scale_block_offset_stride);
+                key.const_data_ptr<scalar_t>(),
+                value.const_data_ptr<scalar_t>(),
+                key_cache.mutable_data_ptr<uint8_t>(),
+                value_cache.mutable_data_ptr<uint8_t>(), key_scale_ptr,
+                value_scale_ptr, slot_mapping.const_data_ptr<int64_t>(),
+                k_scale_ptr, v_scale_ptr, key.stride(0), value.stride(0),
+                num_heads, head_size, block_size, data_block_stride,
+                data_head_stride, data_block_offset_stride, scale_block_stride,
+                scale_head_stride, scale_block_offset_stride);
       });
 }

csrc/libtorch_stable/ops.h

@@ -355,3 +355,132 @@ torch::stable::Tensor ggml_moe_a8_vec(torch::stable::Tensor X,
                                       int64_t tokens);
 
 int64_t ggml_moe_get_block_size(int64_t type);
+
+void paged_attention_v1(
+    torch::stable::Tensor& out, torch::stable::Tensor& query,
+    torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
+    int64_t num_kv_heads, double scale, torch::stable::Tensor& block_tables,
+    torch::stable::Tensor& seq_lens, int64_t block_size, int64_t max_seq_len,
+    const std::optional<torch::stable::Tensor>& alibi_slopes,
+    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
+    torch::stable::Tensor& v_scale, const int64_t tp_rank,
+    const int64_t blocksparse_local_blocks,
+    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
+    const int64_t blocksparse_head_sliding_step);
+
+void paged_attention_v2(
+    torch::stable::Tensor& out, torch::stable::Tensor& exp_sums,
+    torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out,
+    torch::stable::Tensor& query, torch::stable::Tensor& key_cache,
+    torch::stable::Tensor& value_cache, int64_t num_kv_heads, double scale,
+    torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens,
+    int64_t block_size, int64_t max_seq_len,
+    const std::optional<torch::stable::Tensor>& alibi_slopes,
+    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
+    torch::stable::Tensor& v_scale, const int64_t tp_rank,
+    const int64_t blocksparse_local_blocks,
+    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
+    const int64_t blocksparse_head_sliding_step);
+
+// Cache ops (shared CUDA/ROCm)
+void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst,
+                 int64_t block_size_in_bytes,
+                 const torch::stable::Tensor& block_mapping);
+
+// Batch swap: submit all block copies in a single driver call.
+void swap_blocks_batch(const torch::stable::Tensor& src_ptrs,
+                       const torch::stable::Tensor& dst_ptrs,
+                       const torch::stable::Tensor& sizes,
+                       bool is_src_access_order_any);
+
+void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value,
+                       torch::stable::Tensor& key_cache,
+                       torch::stable::Tensor& value_cache,
+                       torch::stable::Tensor& slot_mapping,
+                       const std::string& kv_cache_dtype,
+                       torch::stable::Tensor& k_scale,
+                       torch::stable::Tensor& v_scale);
+
+void reshape_and_cache_flash(
+    torch::stable::Tensor& key, torch::stable::Tensor& value,
+    torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
+    torch::stable::Tensor& slot_mapping, const std::string& kv_cache_dtype,
+    torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale);
+
+void concat_and_cache_mla(torch::stable::Tensor& kv_c,
+                          torch::stable::Tensor& k_pe,
+                          torch::stable::Tensor& kv_cache,
+                          torch::stable::Tensor& slot_mapping,
+                          const std::string& kv_cache_dtype,
+                          torch::stable::Tensor& scale);
+
+// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla
+void concat_and_cache_mla_rope_fused(
+    torch::stable::Tensor& positions, torch::stable::Tensor& q_pe,
+    torch::stable::Tensor& k_pe, torch::stable::Tensor& kv_c,
+    torch::stable::Tensor& rope_cos_sin_cache, bool rope_is_neox,
+    torch::stable::Tensor& slot_mapping, torch::stable::Tensor& kv_cache,
+    const std::string& kv_cache_dtype,
+    torch::stable::Tensor& kv_cache_quant_scale);
+
+// Just for unittest
+void convert_fp8(torch::stable::Tensor& dst_cache,
+                 torch::stable::Tensor& src_cache, const double scale,
+                 const std::string& kv_cache_dtype);
+
+void gather_and_maybe_dequant_cache(
+    torch::stable::Tensor const& src_cache,     // [NUM_BLOCKS, BLOCK_SIZE,
+                                                // ENTRIES...]
+    torch::stable::Tensor const& dst,           // [TOT_TOKENS, ENTRIES...]
+    torch::stable::Tensor const& block_table,   // [BATCH, BLOCK_INDICES]
+    torch::stable::Tensor const& cu_seq_lens,   // [BATCH+1]
+    torch::stable::Tensor const& token_to_seq,  // [MAX_TOKEN_ACROSS_CHUNKS]
+    int64_t num_tokens, const std::string& kv_cache_dtype,
+    torch::stable::Tensor const& scale,
+    std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
+
+// TODO(hc): cp_gather_cache need support scaled kvcahe in the future.
+void cp_gather_cache(
+    torch::stable::Tensor const& src_cache,    // [NUM_BLOCKS, BLOCK_SIZE,
+                                               // ENTRIES...]
+    torch::stable::Tensor const& dst,          // [TOT_TOKENS, ENTRIES...]
+    torch::stable::Tensor const& block_table,  // [BATCH, BLOCK_INDICES]
+    torch::stable::Tensor const& cu_seq_lens,  // [BATCH+1]
+    int64_t batch_size,
+    std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
+
+// Gather and upconvert FP8 KV cache to BF16 workspace
+void cp_gather_and_upconvert_fp8_kv_cache(
+    torch::stable::Tensor const& src_cache,         // [NUM_BLOCKS, BLOCK_SIZE,
+                                                    // 656]
+    torch::stable::Tensor const& dst,               // [TOT_TOKENS, 576]
+    torch::stable::Tensor const& block_table,       // [BATCH, BLOCK_INDICES]
+    torch::stable::Tensor const& seq_lens,          // [BATCH]
+    torch::stable::Tensor const& workspace_starts,  // [BATCH]
+    int64_t batch_size);
+
+// Indexer K quantization and cache function
+void indexer_k_quant_and_cache(
+    torch::stable::Tensor& k,             // [num_tokens, head_dim]
+    torch::stable::Tensor& kv_cache,      // [num_blocks, block_size,
+                                          // cache_stride]
+    torch::stable::Tensor& slot_mapping,  // [num_tokens]
+    int64_t quant_block_size,             // quantization block size
+    const std::string& scale_fmt);
+
+// Concatenate query nope and rope for MLA/DSA attention
+void concat_mla_q(
+    torch::stable::Tensor& ql_nope,  // [num_tokens, num_heads, nope_dim]
+    torch::stable::Tensor& q_pe,     // [num_tokens, num_heads, rope_dim]
+    torch::stable::Tensor& q_out);   // [num_tokens, num_heads, nope_dim +
+                                     // rope_dim]
+
+// Extract function to gather quantized K cache
+void cp_gather_indexer_k_quant_cache(
+    const torch::stable::Tensor& kv_cache,      // [num_blocks, block_size,
+                                                // cache_stride]
+    torch::stable::Tensor& dst_k,               // [num_tokens, head_dim]
+    torch::stable::Tensor& dst_scale,           // [num_tokens, head_dim /
+                                                // quant_block_size * 4]
+    const torch::stable::Tensor& block_table,   // [batch_size, num_blocks]
+    const torch::stable::Tensor& cu_seq_lens);  // [batch_size + 1]

csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu

@@ -17,7 +17,7 @@
 #include <torch/csrc/stable/tensor.h>
 #include "libtorch_stable/torch_utils.h"
 #include "libtorch_stable/dispatch_utils.h"
-#include "cuda_vec_utils.cuh"
+#include "../../cuda_vec_utils.cuh"
 
 #include <cuda_runtime_api.h>
 #include <cuda_runtime.h>

csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu

@@ -27,7 +27,7 @@
 #include <torch/csrc/stable/tensor.h>
 #include "libtorch_stable/torch_utils.h"
 #include "libtorch_stable/dispatch_utils.h"
-#include "cuda_vec_utils.cuh"
+#include "../../cuda_vec_utils.cuh"
 #include "cuda_utils.h"
 
 #include "nvfp4_utils.cuh"

csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu

@@ -17,7 +17,7 @@
 #include <torch/csrc/stable/tensor.h>
 #include "libtorch_stable/torch_utils.h"
 #include "libtorch_stable/dispatch_utils.h"
-#include "cuda_vec_utils.cuh"
+#include "../../cuda_vec_utils.cuh"
 
 #include <cuda_runtime_api.h>
 #include <cuda_runtime.h>

csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu

@@ -23,7 +23,7 @@
 
 #include "libtorch_stable/torch_utils.h"
 #include "libtorch_stable/dispatch_utils.h"
-#include "cuda_vec_utils.cuh"
+#include "../../cuda_vec_utils.cuh"
 
 #include "cuda_utils.h"
 #include "launch_bounds_utils.h"

csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh

@@ -20,7 +20,7 @@
 #include <cuda_fp8.h>
 #include <utility>
 
-#include "cuda_vec_utils.cuh"
+#include "../../cuda_vec_utils.cuh"
 
 #if defined(NVFP4_ENABLE_ELTS16) && defined(CUDA_VERSION) && \
     CUDA_VERSION >= 12090

csrc/libtorch_stable/torch_bindings.cpp

@@ -474,6 +474,33 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "Tensor? initial_state_idx,"
       "Tensor? cu_chunk_seqlen,"
       "Tensor? last_chunk_indices) -> ()");
+
+  // Attention ops
+  // Compute the attention between an input query and the cached
+  // keys/values using PagedAttention.
+  ops.def(
+      "paged_attention_v1("
+      "    Tensor! out, Tensor query, Tensor key_cache,"
+      "    Tensor value_cache, int num_kv_heads, float scale,"
+      "    Tensor block_tables, Tensor seq_lens, int block_size,"
+      "    int max_seq_len, Tensor? alibi_slopes,"
+      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
+      "    int tp_rank, int blocksparse_local_blocks,"
+      "    int blocksparse_vert_stride, int blocksparse_block_size,"
+      "    int blocksparse_head_sliding_step) -> ()");
+
+  // PagedAttention V2.
+  ops.def(
+      "paged_attention_v2("
+      "    Tensor! out, Tensor! exp_sums, Tensor! max_logits,"
+      "    Tensor! tmp_out, Tensor query, Tensor key_cache,"
+      "    Tensor value_cache, int num_kv_heads, float scale,"
+      "    Tensor block_tables, Tensor seq_lens, int block_size,"
+      "    int max_seq_len, Tensor? alibi_slopes,"
+      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
+      "    int tp_rank, int blocksparse_local_blocks,"
+      "    int blocksparse_vert_stride, int blocksparse_block_size,"
+      "    int blocksparse_head_sliding_step) -> ()");
 }
 
 STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
@@ -581,6 +608,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("ggml_moe_a8", TORCH_BOX(&ggml_moe_a8));
   ops.impl("ggml_moe_a8_vec", TORCH_BOX(&ggml_moe_a8_vec));
   ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd));
+
+  ops.impl("paged_attention_v1", TORCH_BOX(&paged_attention_v1));
+  ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2));
 }
 
 // These capability-check functions take only primitive args (no tensors), so
@@ -603,4 +633,115 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) {
   ops.impl("ggml_moe_get_block_size", TORCH_BOX(&ggml_moe_get_block_size));
 }
 
+// Cache ops
+STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
+  // Swap in (out) the cache blocks from src to dst.
+  ops.def(
+      "swap_blocks(Tensor src, Tensor! dst,"
+      "            int block_size_in_bytes, Tensor block_mapping) -> ()");
+
+  // Batch swap: submit all block copies in a single driver call.
+  ops.def(
+      "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs,"
+      "                  Tensor sizes,"
+      "                  bool is_src_access_order_any=False) -> ()");
+
+  // Reshape the key and value tensors and cache them.
+  ops.def(
+      "reshape_and_cache(Tensor key, Tensor value,"
+      "                  Tensor! key_cache, Tensor! value_cache,"
+      "                  Tensor slot_mapping,"
+      "                  str kv_cache_dtype,"
+      "                  Tensor k_scale, Tensor v_scale) -> ()");
+
+  // Reshape the key and value tensors and cache them.
+  ops.def(
+      "reshape_and_cache_flash(Tensor key, Tensor value,"
+      "                        Tensor! key_cache,"
+      "                        Tensor! value_cache,"
+      "                        Tensor slot_mapping,"
+      "                        str kv_cache_dtype,"
+      "                        Tensor k_scale, Tensor v_scale) -> ()");
+
+  // Concat kv_c and k_pe and cache them.
+  ops.def(
+      "concat_and_cache_mla(Tensor kv_c, Tensor k_pe,"
+      "                     Tensor! kv_cache,"
+      "                     Tensor slot_mapping,"
+      "                     str kv_cache_dtype,"
+      "                     Tensor scale) -> ()");
+
+  // Rotate Q and K, then write to kv cache for MLA
+  ops.def(
+      "concat_and_cache_mla_rope_fused("
+      "                     Tensor positions,"
+      "                     Tensor! q_pe,"
+      "                     Tensor! k_pe,"
+      "                     Tensor kv_c,"
+      "                     Tensor cos_sin_cache,"
+      "                     bool is_neox,"
+      "                     Tensor slot_mapping,"
+      "                     Tensor! kv_cache,"
+      "                     str kv_cache_dtype,"
+      "                     Tensor kv_cache_scale) -> ()");
+
+  // Convert the key and value cache to fp8 data type.
+  ops.def(
+      "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, "
+      "str kv_cache_dtype) -> ()");
+
+  // Gather cache blocks from src_cache to dst, dequantizing from
+  // src_cache's dtype to dst's dtype if necessary.
+  ops.def(
+      "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, "
+      "                               Tensor block_table, Tensor cu_seq_lens, "
+      "                               Tensor token_to_seq, "
+      "                               int num_tokens, "
+      "                               str kv_cache_dtype, "
+      "                               Tensor scale, Tensor? seq_starts) -> ()");
+
+  ops.def(
+      "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, "
+      "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()");
+
+  ops.def(
+      "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, "
+      "Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int "
+      "batch_size) -> ()");
+
+  ops.def(
+      "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor "
+      "slot_mapping, "
+      "int quant_block_size, str kv_cache_dtype) -> ()");
+
+  ops.def("concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()");
+
+  ops.def(
+      "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! "
+      "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()");
+}
+
+STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CPU, ops) {
+  ops.impl("swap_blocks_batch", TORCH_BOX(&swap_blocks_batch));
+}
+
+STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) {
+  ops.impl("swap_blocks", TORCH_BOX(&swap_blocks));
+  ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache));
+  ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash));
+  ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla));
+  ops.impl("concat_and_cache_mla_rope_fused",
+           TORCH_BOX(&concat_and_cache_mla_rope_fused));
+  ops.impl("convert_fp8", TORCH_BOX(&convert_fp8));
+  ops.impl("gather_and_maybe_dequant_cache",
+           TORCH_BOX(&gather_and_maybe_dequant_cache));
+  ops.impl("cp_gather_cache", TORCH_BOX(&cp_gather_cache));
+  ops.impl("cp_gather_and_upconvert_fp8_kv_cache",
+           TORCH_BOX(&cp_gather_and_upconvert_fp8_kv_cache));
+  ops.impl("indexer_k_quant_and_cache", TORCH_BOX(&indexer_k_quant_and_cache));
+  ops.impl("concat_mla_q", TORCH_BOX(&concat_mla_q));
+  ops.impl("cp_gather_indexer_k_quant_cache",
+           TORCH_BOX(&cp_gather_indexer_k_quant_cache));
+}
+
 REGISTER_EXTENSION(_C_stable_libtorch)

csrc/ops.h

@@ -31,29 +31,6 @@ torch::Tensor weak_ref_tensor(torch::Tensor& tensor) {
   return new_tensor;
 }
 
-void paged_attention_v1(
-    torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache,
-    torch::Tensor& value_cache, int64_t num_kv_heads, double scale,
-    torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size,
-    int64_t max_seq_len, const std::optional<torch::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::Tensor& k_scale,
-    torch::Tensor& v_scale, const int64_t tp_rank,
-    const int64_t blocksparse_local_blocks,
-    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
-    const int64_t blocksparse_head_sliding_step);
-
-void paged_attention_v2(
-    torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
-    torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
-    torch::Tensor& value_cache, int64_t num_kv_heads, double scale,
-    torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size,
-    int64_t max_seq_len, const std::optional<torch::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::Tensor& k_scale,
-    torch::Tensor& v_scale, const int64_t tp_rank,
-    const int64_t blocksparse_local_blocks,
-    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
-    const int64_t blocksparse_head_sliding_step);
-
 // rms_norm and fused_add_rms_norm declarations also exist in
 // csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
 // because the CPU build still uses these torch::Tensor declarations.

csrc/quantization/w8a8/fp8/amd/quant_utils.cuh

@@ -6,6 +6,7 @@
 #include <hip/hip_bfloat16.h>
 
 #include "../../../../attention/attention_dtypes.h"
+#include <torch/headeronly/core/ScalarType.h>
 
 namespace vllm {
 #ifdef USE_ROCM
@@ -642,27 +643,29 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) {
     vllm::Fp8KVCacheDataType KV_CACHE_DTYPE =                                  \
         vllm::get_fp8_kv_cache_data_type(KV_DTYPE);                            \
     if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kAuto) {                   \
-      if (SRC_DTYPE == at::ScalarType::Float) {                                \
+      if (SRC_DTYPE == torch::headeronly::ScalarType::Float) {                 \
         FN(float, float, vllm::Fp8KVCacheDataType::kAuto);                     \
-      } else if (SRC_DTYPE == at::ScalarType::Half) {                          \
+      } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) {           \
         FN(uint16_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto);               \
-      } else if (SRC_DTYPE == at::ScalarType::BFloat16) {                      \
+      } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) {       \
         FN(__nv_bfloat16, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto);     \
       } else {                                                                 \
-        TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
+        STD_TORCH_CHECK(false,                                                 \
+                        "Unsupported input type of kv cache: ", SRC_DTYPE);    \
       }                                                                        \
     } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E4M3) {         \
-      if (SRC_DTYPE == at::ScalarType::Float) {                                \
+      if (SRC_DTYPE == torch::headeronly::ScalarType::Float) {                 \
         FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3);                \
-      } else if (SRC_DTYPE == at::ScalarType::Half) {                          \
+      } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) {           \
         FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3);             \
-      } else if (SRC_DTYPE == at::ScalarType::BFloat16) {                      \
+      } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) {       \
         FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3);        \
       } else {                                                                 \
-        TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
+        STD_TORCH_CHECK(false,                                                 \
+                        "Unsupported input type of kv cache: ", SRC_DTYPE);    \
       }                                                                        \
     } else {                                                                   \
-      TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE);     \
+      STD_TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \
     }
 
 }  // namespace fp8

csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "../../../../attention/attention_dtypes.h"
+#include <torch/headeronly/core/ScalarType.h>
 #include <assert.h>
 #include <float.h>
 #include <stdint.h>
@@ -546,37 +547,40 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) {
     vllm::Fp8KVCacheDataType KV_CACHE_DTYPE =                                  \
         vllm::get_fp8_kv_cache_data_type(KV_DTYPE);                            \
     if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kAuto) {                   \
-      if (SRC_DTYPE == at::ScalarType::Float) {                                \
+      if (SRC_DTYPE == torch::headeronly::ScalarType::Float) {                 \
         FN(float, float, vllm::Fp8KVCacheDataType::kAuto);                     \
-      } else if (SRC_DTYPE == at::ScalarType::Half) {                          \
+      } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) {           \
         FN(uint16_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto);               \
-      } else if (SRC_DTYPE == at::ScalarType::BFloat16) {                      \
+      } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) {       \
         FN(__nv_bfloat16, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto);     \
       } else {                                                                 \
-        TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
+        STD_TORCH_CHECK(false,                                                 \
+                        "Unsupported input type of kv cache: ", SRC_DTYPE);    \
       }                                                                        \
     } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E4M3) {         \
-      if (SRC_DTYPE == at::ScalarType::Float) {                                \
+      if (SRC_DTYPE == torch::headeronly::ScalarType::Float) {                 \
         FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3);                \
-      } else if (SRC_DTYPE == at::ScalarType::Half) {                          \
+      } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) {           \
         FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3);             \
-      } else if (SRC_DTYPE == at::ScalarType::BFloat16) {                      \
+      } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) {       \
         FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3);        \
       } else {                                                                 \
-        TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
+        STD_TORCH_CHECK(false,                                                 \
+                        "Unsupported input type of kv cache: ", SRC_DTYPE);    \
       }                                                                        \
     } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E5M2) {         \
-      if (SRC_DTYPE == at::ScalarType::Float) {                                \
+      if (SRC_DTYPE == torch::headeronly::ScalarType::Float) {                 \
         FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2);                \
-      } else if (SRC_DTYPE == at::ScalarType::Half) {                          \
+      } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) {           \
         FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2);             \
-      } else if (SRC_DTYPE == at::ScalarType::BFloat16) {                      \
+      } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) {       \
         FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2);        \
       } else {                                                                 \
-        TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
+        STD_TORCH_CHECK(false,                                                 \
+                        "Unsupported input type of kv cache: ", SRC_DTYPE);    \
       }                                                                        \
     } else {                                                                   \
-      TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE);     \
+      STD_TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \
     }
 
 }  // namespace fp8

csrc/rocm/ops.h

@@ -18,6 +18,15 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
                const at::Tensor& scale_a, const at::Tensor& scale_b,
                const int64_t CuCount);
 
+torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight,
+                              torch::Tensor b_qzeros, torch::Tensor b_scales,
+                              torch::Tensor b_g_idx, bool use_v2_format);
+
+torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
+                                   torch::Tensor b_qzeros,
+                                   torch::Tensor b_scales,
+                                   torch::Tensor b_g_idx, bool use_v2_format);
+
 void paged_attention(
     torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
     torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,

csrc/rocm/q_gemm_rdna3.cu

@@ -0,0 +1,780 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// W4A16 GPTQ kernel for RDNA3 (gfx1100 / RX 7900 XTX class), templated on the
+// activation dtype (half or __hip_bfloat16). Adapted from exllamav2's 4-bit
+// kernel (csrc/quantization/gptq/q_gemm.cu) with the following changes:
+//
+//   1. Direct write to the T-typed output via packed CAS-loop on a 64-bit
+//      word (atomic_add_pk4_{f16,bf16}). gfx11 has no native
+//      v_global_atomic_pk_add_{f16,bf16}, so the kernel emulates one with
+//      global_atomic_cmpswap_b64. This avoids the M*N*4-byte FP32 scratch
+//      buffer + memset + cast-pass that an fp32-accumulator design would
+//      need; the caller passes a zero-initialised T-typed output tensor
+//      and every block atomically adds its partial sum into it.
+//
+//   2. The bf16 path uses a dedicated bit-trick that avoids the fp16-only
+//      "upper nibble * 16" trick, which would overflow the 7-bit bf16
+//      mantissa. See qdq_4_rdna3.cuh for details.
+//
+//   3. Wave32 geometry sized for high CU saturation: THREADS_X=256
+//      (8 waves per block) and BLOCK_KN_SIZE=256, with each thread
+//      computing 4 N output columns. gridDim.z = K / BLOCK_KN_SIZE
+//      splits K and the output is atomically accumulated. fp16 uses
+//      v_dot2_f32_f16 (__builtin_amdgcn_fdot2) for the inner dot;
+//      bf16 widens to fp32 (no v_pk_fma_bf16 on gfx11) and accumulates
+//      with v_fma_f32. M_COUNT ∈ {1,2,4,8} is selected at launch
+//      based on size_m.
+//
+//   4. The bf16 dispatch with M >= 16 forwards to the WMMA kernel in
+//      q_gemm_rdna3_wmma.cu (separate translation unit) where
+//      v_wmma_f32_16x16x16_bf16_w32 wins. The fp16 path always stays
+//      scalar (the bit-trick dequant beats WMMA below M=64).
+
+#include <cstdint>
+#include <cstdio>
+
+#include <torch/all.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+#include <hip/hip_fp16.h>
+
+#include "qdq_4_rdna3.cuh"
+
+#if defined(__HIPCC__) && defined(__gfx1100__)
+  #define __HIP__RDNA3__
+#endif
+
+namespace vllm {
+namespace gptq_rdna3 {
+
+// BLOCK_KN_SIZE = 256 (was 128 in exllama). Each block covers 256 K
+// elements and THREADS_X*4 = 1024 N columns. For Qwen-class K=4096 this
+// halves gridDim.z (32 → 16) and therefore halves the atomic count per
+// output position vs the exllama default. THREADS_X=256 = 8 waves on RDNA3
+// wave32; with ~32 wave slots per CU we still fit 4 blocks per CU at peak.
+//
+// We tried BLOCK_KN_SIZE=512 (microbench on Qwen3.6-27B): bf16 improved
+// 5-10% at large M (atomic CAS halved), but fp16 decode regressed up to
+// +40% on qkv-square (32 → 45 μs at M=1). Cause: 16 waves/block × 16
+// total blocks for [M=1, K=N=4096] only saturates ~8 of the 96 CUs,
+// breaking memory-latency hiding for the fp16 path which is already
+// memory-bound. Reverted to 256; bf16 keeps most of its gains from the
+// fp32 dequant rewrite alone.
+#define BLOCK_KN_SIZE 256
+#define THREADS_X 256
+
+// Device code below is RDNA3-only; non-RDNA3 device passes fall through to
+// the empty __global__ stub at the #else below for symbol parity.
+#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
+
+// ---------------------------------------------------------------------------
+// Per-dtype helpers. We avoid heavy template metaprogramming and just provide
+// overloaded inline functions; the kernel below selects via `if constexpr`.
+// ---------------------------------------------------------------------------
+
+// Type-generic zero — both half and bf16_t in HIP/ROCm have a converting
+// constructor from float, but going through __float2half_rn / __float2bfloat16
+// is the unambiguously correct path on every ROCm version.
+template <typename T>
+__forceinline__ __device__ T tzero();
+
+template <>
+__forceinline__ __device__ half tzero<half>() {
+  return __float2half_rn(0.0f);
+}
+
+template <>
+__forceinline__ __device__ bf16_t tzero<bf16_t>() {
+  return __float2bfloat16(0.0f);
+}
+
+__forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) {
+  // RDNA3 has v_dot2_f32_f16 (`__builtin_amdgcn_fdot2`) which computes
+  // fp32 += a.x*b.x + a.y*b.y in a single instruction with the accumulator
+  // staying in fp32 throughout. hipcc 7.2 does NOT peephole the obvious
+  // `__hfma2 + cast + add` pattern into v_dot2 (verified by ISA
+  // disassembly: 0 v_dot2_f32_f16 vs 256 v_cvt_f32_f16 + 218 v_add_f32 in
+  // the M_COUNT=8 kernel before this change), so we issue the builtin
+  // explicitly. Saves the trailing 2× v_cvt_f32_f16 + v_add_f32 (3 ops)
+  // per dot22_8_f call vs the half2-accumulator form. With 128 calls per
+  // K=32 step that's ~384 ops/K-step less issue pressure on the VALU.
+  //
+  // Numerical bonus: accumulator stays fp32 throughout the dot. The old
+  // form accumulated 8 muladds in fp16 (10-bit mantissa) before casting,
+  // which could lose ~3 bits of precision on borderline magnitudes.
+  float result = 0.0f;
+  const half2* a2_ptr = (const half2*)a_ptr;
+  #pragma unroll
+  for (int i = 0; i < 4; i++) {
+    result = __builtin_amdgcn_fdot2(dq[i], *a2_ptr++, result, /*clamp=*/false);
+  }
+  return result;
+}
+
+__forceinline__ __device__ float dot22_8_f(bf162_t (&dq)[4],
+                                           const bf16_t* a_ptr) {
+  // RDNA3 (gfx1100) lacks a packed bf16 FMA: there is no v_pk_fma_bf16 in
+  // the gfx11 ISA (it only landed on CDNA3+ / gfx94x and later). hipcc
+  // therefore lowers __hfma2(bf162_t, bf162_t, bf162_t) to a serialised
+  // fallback (single-element FMAs or fp32 round-trips), which empirically
+  // runs ~2× the cycle count of v_pk_fma_f16 on the same VALU. The bf16
+  // decode path was paying that tax in full, scaling linearly with M (the
+  // fp16 path scales sub-linearly because its v_pk_fma_f16 is full rate
+  // and the kernel becomes memory-bound).
+  //
+  // Fix: widen bf16 → fp32 explicitly (a left-shift by 16, free in VGPRs)
+  // and accumulate with v_fma_f32, which IS full rate on RDNA3. Same FMA
+  // count, but each FMA is fast. Bonus: the accumulator is now fp32
+  // throughout instead of bf16, which is also numerically more accurate
+  // (no compounding bf16-rounding inside the dot loop).
+  float result = 0.0f;
+  #pragma unroll
+  for (int i = 0; i < 4; i++) {
+    uint32_t aw, dw;
+    __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t));
+    __builtin_memcpy(&dw, &dq[i], sizeof(uint32_t));
+    // bf16 in low 16 bits  → fp32 by left-shifting into the upper half.
+    // bf16 in high 16 bits → already aligned with fp32's upper half.
+    float a_x = __uint_as_float((aw & 0xFFFFu) << 16);
+    float a_y = __uint_as_float(aw & 0xFFFF0000u);
+    float d_x = __uint_as_float((dw & 0xFFFFu) << 16);
+    float d_y = __uint_as_float(dw & 0xFFFF0000u);
+    result = __fmaf_rn(d_x, a_x, result);
+    result = __fmaf_rn(d_y, a_y, result);
+  }
+  return result;
+}
+
+// fp32-input dot product: paired with dequant_4bit_8_bf16_f32 which already
+// produces fp32 dq[8]. Saves the bf16→fp32 widening that the bf162_t
+// overload above does for dq (still need to widen A from bf16). Wins more
+// at high N: the bf162_t version's per-call widening cost scales with the
+// number of dequants × M_COUNT × 4 dot calls; the fp32 version pays only
+// for A widening (M_COUNT × 4 × 4 widens, half as many).
+__forceinline__ __device__ float dot22_8_f(float (&dq)[8],
+                                           const bf16_t* a_ptr) {
+  float result = 0.0f;
+  #pragma unroll
+  for (int i = 0; i < 4; i++) {
+    uint32_t aw;
+    __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t));
+    float a_x = __uint_as_float((aw & 0xFFFFu) << 16);
+    float a_y = __uint_as_float(aw & 0xFFFF0000u);
+    result = __fmaf_rn(dq[2 * i + 0], a_x, result);
+    result = __fmaf_rn(dq[2 * i + 1], a_y, result);
+  }
+  return result;
+}
+
+// ---------------------------------------------------------------------------
+// Packed atomic-add via CAS-loop on a 64-bit word (4 fp16/bf16 lanes per CAS).
+// RDNA3 (gfx11) does NOT have native v_global_atomic_pk_add_f16 / _bf16 (those
+// landed on gfx940 / gfx1250 respectively), so this lowers to
+// global_atomic_cmpswap_b64 plus retry. We use this in the kernel epilogue to
+// write 4 output columns per row in a single atomic operation — half the
+// atomic instruction count and half the contention vs two 32-bit CAS calls.
+//
+// Writing directly to fp16/bf16 (instead of through an FP32 scratch buffer +
+// cast pass) saves M*N*4 bytes of allocation, the memset, and the epilogue
+// cast pass that an fp32-accumulator design would need.
+//
+// 64-bit alignment: the kernel writes at `out + n` where n = offset_n + t*4
+// (always multiple of 4), and partition_weight_shape[1] is required to be a
+// multiple of 8 by can_implement(), so every (m, n) write target is 8-byte
+// aligned. Required by global_atomic_cmpswap_b64.
+// ---------------------------------------------------------------------------
+
+__forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01,
+                                                   half2 v23) {
+  unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
+  unsigned long long old = *addr_u;
+  while (true) {
+    union {
+      unsigned long long u;
+      half2 h2[2];
+    } cur, sum;
+    cur.u = old;
+    sum.h2[0] = __hadd2(cur.h2[0], v01);
+    sum.h2[1] = __hadd2(cur.h2[1], v23);
+    unsigned long long prev = atomicCAS(addr_u, old, sum.u);
+    if (prev == old) break;
+    old = prev;
+  }
+}
+
+__forceinline__ __device__ void atomic_add_pk4_bf16(bf16_t* addr, bf162_t v01,
+                                                    bf162_t v23) {
+  unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
+  unsigned long long old = *addr_u;
+  while (true) {
+    union {
+      unsigned long long u;
+      bf162_t b2[2];
+    } cur, sum;
+    cur.u = old;
+    sum.b2[0] = __hadd2(cur.b2[0], v01);
+    sum.b2[1] = __hadd2(cur.b2[1], v23);
+    unsigned long long prev = atomicCAS(addr_u, old, sum.u);
+    if (prev == old) break;
+    old = prev;
+  }
+}
+
+// Load one row's worth of 4 packed zeros (column n..n+3) from a [groups, N/8]
+// uint32 tensor. n is a multiple of 4 by construction (n = offset_n + t*4 with
+// offset_n = blockIdx.x * 512), so the 4 nibbles always live within one or two
+// uint32 words; in practice within one because n & 7 is 0 or 4.
+__forceinline__ __device__ void load4_zeros(const uint32_t* qzeros_row, int n,
+                                            int (&zeros)[4]) {
+  int qcol = n / 8;
+  int shift = (n & 0x07) * 4;
+  uint32_t d = qzeros_row[qcol] >> shift;
+  zeros[0] = (int)(d & 0xF);
+  zeros[1] = (int)((d >> 4) & 0xF);
+  zeros[2] = (int)((d >> 8) & 0xF);
+  zeros[3] = (int)((d >> 12) & 0xF);
+}
+
+template <typename T>
+__forceinline__ __device__ void load4_scales(const T* scales_row, int n,
+                                             T (&scales)[4]) {
+  scales[0] = scales_row[n + 0];
+  scales[1] = scales_row[n + 1];
+  scales[2] = scales_row[n + 2];
+  scales[3] = scales_row[n + 3];
+}
+
+// ---------------------------------------------------------------------------
+// Main kernel.
+// ---------------------------------------------------------------------------
+
+template <typename T, int M_COUNT>
+__global__ void gemm_q4_kernel_rdna3(
+    const T* __restrict__ a, const uint32_t* __restrict__ b_q_weight,
+    const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
+    T* __restrict__ c, const int size_m, const int size_n, const int size_k,
+    const int groups, const int zero_offset, const int* __restrict__ b_q_perm) {
+  const int t = threadIdx.x;
+  const int offset_n = blockIdx.x * BLOCK_KN_SIZE * 4;
+  const int offset_m = blockIdx.y * M_COUNT;
+  const int offset_k = blockIdx.z * BLOCK_KN_SIZE;
+  const int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
+  const int n = offset_n + t * 4;
+
+  // LDS layout: [M_COUNT][BLOCK_KN_SIZE + LDS_PAD]. The PAD=8 elements per M
+  // row break the natural 256-element/512-byte alignment that would otherwise
+  // collide on the same LDS bank when a thread reads block_a[0..M_COUNT-1][k]
+  // (same k, different m). Row stride becomes 264 elements * 2B = 528B = 132
+  // 4-byte banks, so m-stride hits banks (m*132)%32 = (m*4)%32 — distinct for
+  // all M_COUNT ≤ 8. Cost: 16B LDS per block, irrelevant.
+  constexpr int LDS_PAD = 8;
+  __shared__ T block_a[M_COUNT][BLOCK_KN_SIZE + LDS_PAD];
+
+  // Stage A: each thread loads 1 K element per M row into LDS (with optional
+  // act-order permutation). THREADS_X == BLOCK_KN_SIZE so this is a 1:1 map.
+  // For M_COUNT > 1 with size_m not a multiple of M_COUNT, slots past size_m
+  // are zero-padded so the dot product contribution is 0 (we then skip the
+  // atomic write for those rows below).
+  //
+  // M=1 fast path: skip LDS staging + __syncthreads entirely. All 256 threads
+  // read the SAME 8-element A window per inner step (a_off is uniform across
+  // the block), so the cache-line broadcast through L1 makes global reads as
+  // cheap as LDS reads. Measured: ~1% on 4B b=1, ~6% on 27B b=1 in=128.
+  static_assert(BLOCK_KN_SIZE == THREADS_X,
+                "BLOCK_KN_SIZE must equal THREADS_X (1 K element per thread)");
+  // The M=1 fast path (skip LDS) only has a global-read code path for bf16
+  // (the v_dot2_f32_bf16 branch).  The fp16 inner loop still indexes
+  // block_a[m][a_off] unconditionally, so for fp16 we MUST stage A through
+  // LDS even at M=1 to avoid reading uninitialized shared memory.
+  constexpr bool USE_LDS_A = (M_COUNT > 1) || std::is_same<T, half>::value;
+  if constexpr (USE_LDS_A) {
+    if (offset_k + t < end_k) {
+  #pragma unroll
+      for (int m = 0; m < M_COUNT; ++m) {
+        T av;
+        if (offset_m + m < size_m) {
+          const T* a_row = a + (offset_m + m) * size_k;
+          if (b_q_perm)
+            av = a_row[b_q_perm[offset_k + t]];
+          else
+            av = a_row[offset_k + t];
+        } else {
+          av = tzero<T>();  // zero-pad invalid M rows
+        }
+        block_a[m][t] = av;
+      }
+    }
+
+    // Threads beyond the right edge of N have nothing to do. Note: we must NOT
+    // return before __syncthreads() if any thread in the block participates in
+    // the LDS load above — but here all THREADS_X (=256) threads always do,
+    // regardless of whether their `n` is in bounds.
+    __syncthreads();
+  } else if (b_q_perm) {
+    // bf16 M=1 fast path skips LDS, but its global read below is sequential
+    // and cannot apply act-order. When a permutation is present, stage the
+    // single A row through LDS (as fp16 / M>1 do) so the read picks it up.
+    // b_q_perm is block-uniform, so the __syncthreads is non-divergent.
+    if (offset_k + t < end_k)
+      block_a[0][t] = a[offset_m * size_k + b_q_perm[offset_k + t]];
+    __syncthreads();
+  }
+  if (n >= size_n) return;
+
+  // Group bookkeeping. We require size_k % groups == 0 (groupsize divides K).
+  const int groupsize = size_k / groups;
+  int group = offset_k / groupsize;
+  int nextgroup = (group + 1) * groupsize;
+
+  // qweight stride: weights are [K/8, N] uint32 with K packed at dim 0.
+  int qk = offset_k / 8;
+  const uint32_t* b_ptr = b_q_weight + qk * size_n + n;
+
+  // Per-column dequant constants. We hold one set of (z, y) pairs per column.
+  // fp16 uses the exllama (z1z16, y1y16) double-pair to enable the upper-
+  // nibble-*16 trick. bf16 uses fp32 scalars (z, y) because the dequant
+  // produces fp32 directly — see prep_zero_scale_bf16_f32 / the FMA
+  // bypass for the missing v_pk_fma_bf16 on gfx11.
+  half2 z1z16_h[4][2], y1y16_h[4][2];
+  float z_b_f[4], y_b_f[4];
+
+  auto refresh_group = [&](int g) {
+    const uint32_t* qz_row = b_qzeros + g * (size_n / 8);
+    const T* sc_row = b_scales + g * size_n;
+    int zeros[4];
+    T scales[4];
+    load4_zeros(qz_row, n, zeros);
+    load4_scales<T>(sc_row, n, scales);
+    if constexpr (std::is_same<T, half>::value) {
+  #pragma unroll
+      for (int i = 0; i < 4; ++i) {
+        prep_zero_scale_fp16((uint32_t)(zeros[i] + zero_offset), scales[i],
+                             z1z16_h[i], y1y16_h[i]);
+      }
+    } else {
+  #pragma unroll
+      for (int i = 0; i < 4; ++i) {
+        prep_zero_scale_bf16_f32((uint32_t)(zeros[i] + zero_offset), scales[i],
+                                 z_b_f[i], y_b_f[i]);
+      }
+    }
+  };
+
+  refresh_group(group);
+
+  float block_c[M_COUNT][4];
+  #pragma unroll
+  for (int m = 0; m < M_COUNT; ++m) {
+  #pragma unroll
+    for (int j = 0; j < 4; ++j) block_c[m][j] = 0.0f;
+  }
+
+  // Note on group-transition granularity: we check `k == nextgroup` at the
+  // start of each outer iteration (which advances K by 32). This is correct
+  // when group_size >= 32 OR group_size divides 32 evenly (groupsize is one
+  // of {1,2,4,8,16,32,64,128,...}). For group_size in {16, 8, 4, ...} the
+  // inner loop would cross a group boundary between j-iterations; we require
+  // group_size >= 32 here, mirroring exllama's assumption.
+  //
+  // Software pipelining: we issue all 4 vectorized weight loads up front
+  // before any dequant/FMA depends on them. This gives the AMDGPU backend
+  // freedom to schedule the global_loads early and overlap their latency
+  // with dequant + v_pk_fma_f16 of earlier iterations. Cost: 4×int4 = 16
+  // VGPRs in flight per thread, plenty of headroom on RDNA3.
+  int k = offset_k;
+  while (k < end_k) {
+    if (k == nextgroup) {
+      group++;
+      nextgroup += groupsize;
+      refresh_group(group);
+    }
+
+    // Prefetch all four j-iterations' weight words. The compiler emits 4
+    // global_load_b128 instructions back-to-back; the dependent dequant +
+    // FMA work below hides their latency.
+    int4 b_w[4];
+  #pragma unroll
+    for (int j = 0; j < 4; ++j) {
+      b_w[j] = *(const int4*)(b_ptr + j * size_n);
+    }
+    b_ptr += 4 * size_n;
+
+  #pragma unroll
+    for (int j = 0; j < 4; ++j) {
+      const int a_off = (k - offset_k) + 8 * j;
+
+      if constexpr (std::is_same<T, half>::value) {
+        half2 dq[4][4];
+        dequant_4bit_8_fp16((uint32_t)b_w[j].x, dq[0], z1z16_h[0], y1y16_h[0]);
+        dequant_4bit_8_fp16((uint32_t)b_w[j].y, dq[1], z1z16_h[1], y1y16_h[1]);
+        dequant_4bit_8_fp16((uint32_t)b_w[j].z, dq[2], z1z16_h[2], y1y16_h[2]);
+        dequant_4bit_8_fp16((uint32_t)b_w[j].w, dq[3], z1z16_h[3], y1y16_h[3]);
+
+  #pragma unroll
+        for (int m = 0; m < M_COUNT; ++m) {
+          const half* a_ptr = reinterpret_cast<const half*>(&block_a[m][a_off]);
+          block_c[m][0] += dot22_8_f(dq[0], a_ptr);
+          block_c[m][1] += dot22_8_f(dq[1], a_ptr);
+          block_c[m][2] += dot22_8_f(dq[2], a_ptr);
+          block_c[m][3] += dot22_8_f(dq[3], a_ptr);
+        }
+      } else if constexpr (M_COUNT == 1) {
+        // bf16 decode (M=1), v_dot2_f32_bf16 path. Mirrors the data-flow of
+        // Hybrid PR #40977's wvSplitK_int4 kernel exactly so clang's
+        // InstCombine cannot fold the bf16→fp32 widening (LLVM #76000):
+        //   * activations and magic-value weights share a fp32-aliased
+        //     union (bytes written as uint32, read as bf16x2_t for the
+        //     dot — pointer-cast opacity defeats the fold)
+        //   * sum_a computed via a *second* v_dot2 with bf162(1,1) as the
+        //     second operand, avoiding any explicit bf16→fp32 widen of A
+        //   * bias correction y_b_f * partial + z_b_f * sum_a, identical
+        //     to the previous fp32-FMA-chain path
+        //
+        // Net: 20 v_dot2_f32_bf16 + 8 fp32 FMA per int32 weight vs the
+        // previous 40 fp32 FMA. v_dot2 runs at full rate on gfx1100, so
+        // the substitution is ~2× cheaper for the inner accumulator.
+        typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
+        constexpr uint32_t BF16_MAGIC = 0x43004300u;  // bf162(128, 128)
+        constexpr uint32_t BF16_ONES = 0x3F803F80u;   // bf162(1.0, 1.0)
+        union pack4 {
+          float f[4];
+          uint32_t u[4];
+        };
+
+        uint32_t w[4];
+        __builtin_memcpy(w, &b_w[j], sizeof(int4));
+
+        // Load 8 bf16 activations as 4 uint32s (= 4 bf16x2 pairs) into a
+        // fp32-aliased union. Storing as uint32 keeps the IR-level type
+        // opaque so the inner v_dot2 cannot be folded to fp32 widening.
+        //
+        // A is read direct from global (no LDS staging — see USE_LDS_A above),
+        // except under act-order, where it comes from the permuted LDS copy.
+        pack4 a_pack;
+        {
+          const uint32_t* a_words =
+              b_q_perm
+                  ? reinterpret_cast<const uint32_t*>(&block_a[0][a_off])
+                  : reinterpret_cast<const uint32_t*>(a + offset_k + a_off);
+          a_pack.u[0] = a_words[0];
+          a_pack.u[1] = a_words[1];
+          a_pack.u[2] = a_words[2];
+          a_pack.u[3] = a_words[3];
+        }
+
+        // sum_a = Σ a[i]. Computed via 4× v_dot2_f32_bf16 with bf162(1,1) as
+        // the second operand — every bf16 pair contributes 1·a_lo + 1·a_hi.
+        // No fp32 widening of activations: the bytes go straight from LDS
+        // through v_dot2 into the fp32 accumulator.
+        float sum_a = 0.0f;
+  #pragma unroll
+        for (int b = 0; b < 4; ++b) {
+          sum_a = __builtin_amdgcn_fdot2_f32_bf16(
+              *((bf16x2_t*)(&a_pack.f[b])), *((const bf16x2_t*)&BF16_ONES),
+              sum_a, /*clamp=*/false);
+        }
+
+        // unroll 1 keeps q_pack alive only one col at a time (8 fp32 VGPRs
+        // recycled across cols), avoiding straight-line expansion that
+        // would inflate live-range to 32 VGPRs.
+  #pragma unroll 1
+        for (int col = 0; col < 4; ++col) {
+          // Build dequant magic values bf16(128 + nibble) directly into a
+          // fp32-aliased union via uint32 stores. No fp32 in the data flow
+          // until v_dot2 consumes the bytes.
+          pack4 q_pack;
+          const uint32_t qa = w[col];
+          q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
+          q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
+          q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
+          q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
+
+          // partial = Σ (128 + nibble[i]) · a[i], via 4× v_dot2_f32_bf16.
+          float partial = 0.0f;
+  #pragma unroll
+          for (int b = 0; b < 4; ++b) {
+            partial = __builtin_amdgcn_fdot2_f32_bf16(
+                *((bf16x2_t*)(&a_pack.f[b])), *((bf16x2_t*)(&q_pack.f[b])),
+                partial, /*clamp=*/false);
+          }
+
+          // block_c += y_b_f * partial + z_b_f * sum_a
+          //   y_b_f = scale, z_b_f = -(128+zero)*scale
+          // partial holds (128 + nibble) · a; subtracting (128+zero)·sum_a
+          // and scaling yields scale · (nibble - zero) · a as required.
+          block_c[0][col] =
+              __fmaf_rn(y_b_f[col], partial,
+                        __fmaf_rn(z_b_f[col], sum_a, block_c[0][col]));
+        }
+      } else {
+        // bf16 M_COUNT > 1 path with v_dot2_f32_bf16. Same opacity trick as
+        // the M=1 branch: activations + magic-value weights stored in
+        // fp32-aliased unions, dot via __builtin_amdgcn_fdot2_f32_bf16 with
+        // pointer-cast to bf16x2_t. sum_a[m] computed via second v_dot2
+        // with BF16_ONES; bias correction (y_b_f * partial + z_b_f * sum_a)
+        // applied after the dot. Magic values built once per col and reused
+        // across all M rows — amortizes dequant cost across M_COUNT.
+        typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
+        constexpr uint32_t BF16_MAGIC = 0x43004300u;  // bf162(128, 128)
+        constexpr uint32_t BF16_ONES = 0x3F803F80u;   // bf162(1.0, 1.0)
+        union pack4 {
+          float f[4];
+          uint32_t u[4];
+        };
+
+        uint32_t w[4];
+        __builtin_memcpy(w, &b_w[j], sizeof(int4));
+
+        // Load M_COUNT × 8 bf16 activations as 4 uint32s each into pack4
+        // unions. Stored as uint32 to keep IR-level types opaque (defeats
+        // InstCombine fold). At M_COUNT=8 this is 32 fp32 VGPRs — within RDNA3
+        // budget.
+        pack4 a_pack[M_COUNT];
+  #pragma unroll
+        for (int m = 0; m < M_COUNT; ++m) {
+          const uint32_t* a_words =
+              reinterpret_cast<const uint32_t*>(&block_a[m][a_off]);
+          a_pack[m].u[0] = a_words[0];
+          a_pack[m].u[1] = a_words[1];
+          a_pack[m].u[2] = a_words[2];
+          a_pack[m].u[3] = a_words[3];
+        }
+
+        // sum_a[m] = Σ a[m][i] via 4× v_dot2 with bf162(1,1) — no fp32 widen.
+        float sum_a[M_COUNT];
+  #pragma unroll
+        for (int m = 0; m < M_COUNT; ++m) {
+          float s = 0.0f;
+  #pragma unroll
+          for (int b = 0; b < 4; ++b) {
+            s = __builtin_amdgcn_fdot2_f32_bf16(*((bf16x2_t*)(&a_pack[m].f[b])),
+                                                *((const bf16x2_t*)&BF16_ONES),
+                                                s, /*clamp=*/false);
+          }
+          sum_a[m] = s;
+        }
+
+        // Per col: build magic-value pack, dot against all M activations.
+        // unroll 1 keeps q_pack live one col at a time (8 fp32 VGPRs recycled)
+        // — same register-pressure trick as the previous fp32 path.
+  #pragma unroll 1
+        for (int col = 0; col < 4; ++col) {
+          pack4 q_pack;
+          const uint32_t qa = w[col];
+          q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
+          q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
+          q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
+          q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
+
+  #pragma unroll
+          for (int m = 0; m < M_COUNT; ++m) {
+            float partial = 0.0f;
+  #pragma unroll
+            for (int b = 0; b < 4; ++b) {
+              partial = __builtin_amdgcn_fdot2_f32_bf16(
+                  *((bf16x2_t*)(&a_pack[m].f[b])), *((bf16x2_t*)(&q_pack.f[b])),
+                  partial, /*clamp=*/false);
+            }
+            // block_c += y_b_f * partial + z_b_f * sum_a (same correction as
+            // M=1)
+            block_c[m][col] =
+                __fmaf_rn(y_b_f[col], partial,
+                          __fmaf_rn(z_b_f[col], sum_a[m], block_c[m][col]));
+          }
+        }
+      }
+    }
+    k += 32;  // 4 weight words * 8 nibbles = 32 K elements
+  }
+
+  // Pack the 4 FP32 partial sums into 2 packed pairs and atomically add all
+  // four lanes in a single 64-bit CAS write directly to the T-typed output
+  // (caller pre-zeros it). On gfx11 the packed atomic is a CAS-loop, but with
+  // a single b64 op we halve the atomic instruction count vs two b32 CAS
+  // calls, AND save the FP32 buffer + memset + cast pass entirely.
+  #pragma unroll
+  for (int m = 0; m < M_COUNT; ++m) {
+    if (offset_m + m >= size_m) continue;  // skip padding rows past size_m
+    T* out = c + (offset_m + m) * size_n + n;
+    if constexpr (std::is_same<T, half>::value) {
+      half2 r01 = __halves2half2(__float2half_rn(block_c[m][0]),
+                                 __float2half_rn(block_c[m][1]));
+      half2 r23 = __halves2half2(__float2half_rn(block_c[m][2]),
+                                 __float2half_rn(block_c[m][3]));
+      atomic_add_pk4_f16(out, r01, r23);
+    } else {
+      bf162_t r01;
+      r01.x = __float2bfloat16(block_c[m][0]);
+      r01.y = __float2bfloat16(block_c[m][1]);
+      bf162_t r23;
+      r23.x = __float2bfloat16(block_c[m][2]);
+      r23.y = __float2bfloat16(block_c[m][3]);
+      atomic_add_pk4_bf16(out, r01, r23);
+    }
+  }
+}
+
+#else  // non-RDNA3 device pass: empty __global__ for symbol parity.
+
+template <typename T, int M_COUNT>
+__global__ void gemm_q4_kernel_rdna3(const T*, const uint32_t*, const uint32_t*,
+                                     const T*, T*, const int, const int,
+                                     const int, const int, const int,
+                                     const int*) {}
+
+#endif  // __HIP__RDNA3__ || !__HIP_DEVICE_COMPILE__
+
+// ---------------------------------------------------------------------------
+// Launcher.
+// ---------------------------------------------------------------------------
+
+template <typename T, int M_COUNT>
+void launch_gemm_q4_for_mcount(const T* a, const uint32_t* b_q_weight,
+                               const uint32_t* b_qzeros, const T* b_scales,
+                               const int* b_q_perm, T* c, int size_m,
+                               int size_n, int size_k, int groups,
+                               int zero_offset, cudaStream_t stream) {
+  dim3 block(THREADS_X);
+  dim3 grid((size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4),
+            (size_m + M_COUNT - 1) / M_COUNT,
+            (size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE);
+
+  gemm_q4_kernel_rdna3<T, M_COUNT><<<grid, block, 0, stream>>>(
+      a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
+      zero_offset, b_q_perm);
+}
+
+// Dispatch to the largest M_COUNT template that doesn't waste more than
+// half a tile. Caps at 8: above that, the WMMA-prefill kernel (M >= 16) is
+// the right tool, not bigger M_COUNT in the scalar dot-product path.
+//
+// Tile-waste table:
+//   M=1   -> M_COUNT=1   (no waste)
+//   M=2,3 -> M_COUNT=2   (M=3 wastes 1/2 of last tile)
+//   M=4-7 -> M_COUNT=4   (worst case M=5: wastes 3/4 of last tile)
+//   M=8-15-> M_COUNT=8   (worst case M=9: wastes 7/8 of last tile)
+// "Wasted" rows are zero-padded in LDS and skip the atomic write, so they
+// only burn instructions on the last block, never affect correctness.
+template <typename T>
+void launch_gemm_q4(const T* a, const uint32_t* b_q_weight,
+                    const uint32_t* b_qzeros, const T* b_scales,
+                    const int* b_q_perm, T* c, int size_m, int size_n,
+                    int size_k, int groups, bool use_v2_format,
+                    cudaStream_t stream) {
+  const int zero_offset = use_v2_format ? 0 : 1;
+
+  if (size_m == 1) {
+    launch_gemm_q4_for_mcount<T, 1>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
+                                    c, size_m, size_n, size_k, groups,
+                                    zero_offset, stream);
+  } else if (size_m <= 3) {
+    launch_gemm_q4_for_mcount<T, 2>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
+                                    c, size_m, size_n, size_k, groups,
+                                    zero_offset, stream);
+  } else if (size_m <= 7) {
+    launch_gemm_q4_for_mcount<T, 4>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
+                                    c, size_m, size_n, size_k, groups,
+                                    zero_offset, stream);
+  } else {
+    // M_COUNT=8 covers M up to 15 here; M >= 16 should ideally take the
+    // WMMA path, but if it falls through we still produce correct output —
+    // just leaving 3-5× of throughput on the table for prefill workloads.
+    launch_gemm_q4_for_mcount<T, 8>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
+                                    c, size_m, size_n, size_k, groups,
+                                    zero_offset, stream);
+  }
+}
+
+}  // namespace gptq_rdna3
+}  // namespace vllm
+
+// ---------------------------------------------------------------------------
+// Public entry point.
+// ---------------------------------------------------------------------------
+//
+// Inputs:
+//   a         [M, K]            half or bfloat16
+//   b_q_weight[K/8, N]          uint32 (already shuffled via gptq_shuffle)
+//   b_qzeros  [groups, N/8]     uint32 (packed 4-bit zeros)
+//   b_scales  [groups, N]       half or bfloat16
+//   b_g_idx   [K] or empty      int32 (act-order permutation; empty=identity)
+//   use_v2_format                bool   (true = GPTQv2, no +1 zero offset)
+//
+// Output:
+//   c         [M, N]            same dtype as a
+
+torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
+                                   torch::Tensor b_qzeros,
+                                   torch::Tensor b_scales,
+                                   torch::Tensor b_g_idx, bool use_v2_format);
+
+torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight,
+                              torch::Tensor b_qzeros, torch::Tensor b_scales,
+                              torch::Tensor b_g_idx, bool use_v2_format) {
+  if (a.dim() == 2 && b_q_weight.dim() == 2 && a.size(1) % 16 == 0 &&
+      b_q_weight.size(1) % 16 == 0 &&
+      ((a.scalar_type() == torch::kBFloat16 && a.size(0) >= 16) ||
+       (a.scalar_type() == torch::kHalf && a.size(0) >= 64))) {
+    return gptq_gemm_rdna3_wmma(a, b_q_weight, b_qzeros, b_scales, b_g_idx,
+                                use_v2_format);
+  }
+
+  TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor");
+  TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor");
+  TORCH_CHECK(b_qzeros.is_cuda(), "b_qzeros must be a CUDA/HIP tensor");
+  TORCH_CHECK(b_scales.is_cuda(), "b_scales must be a CUDA/HIP tensor");
+  TORCH_CHECK(a.dim() == 2, "a must be 2D [M, K]");
+  TORCH_CHECK(b_q_weight.dim() == 2, "b_q_weight must be 2D [K/8, N]");
+  TORCH_CHECK(
+      a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16,
+      "a must be half or bfloat16");
+  TORCH_CHECK(a.scalar_type() == b_scales.scalar_type(),
+              "b_scales dtype must match a");
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  int size_m = (int)a.size(0);
+  int size_k = (int)a.size(1);
+  int size_n = (int)b_q_weight.size(1);
+  int groups = (int)b_qzeros.size(0);
+
+  TORCH_CHECK(b_q_weight.size(0) * 8 == size_k,
+              "b_q_weight first dim must be K/8");
+  TORCH_CHECK(b_scales.size(0) == groups,
+              "b_scales must have same group count as qzeros");
+  TORCH_CHECK(b_scales.size(1) == size_n, "b_scales last dim must be N");
+  TORCH_CHECK(size_n % 8 == 0, "N must be a multiple of 8 (64-bit atomic CAS)");
+
+  auto opts = torch::TensorOptions().dtype(a.dtype()).device(a.device());
+  at::Tensor c = torch::zeros({size_m, size_n}, opts);
+
+  const int* g_idx_ptr = nullptr;
+  if (!b_g_idx.device().is_meta() && b_g_idx.numel() > 0) {
+    TORCH_CHECK(b_g_idx.scalar_type() == torch::kInt32,
+                "b_g_idx must be int32");
+    g_idx_ptr = (const int*)b_g_idx.data_ptr();
+  }
+
+  if (a.scalar_type() == torch::kHalf) {
+    vllm::gptq_rdna3::launch_gemm_q4<half>(
+        (const half*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(),
+        (const uint32_t*)b_qzeros.data_ptr(), (const half*)b_scales.data_ptr(),
+        g_idx_ptr, (half*)c.data_ptr(), size_m, size_n, size_k, groups,
+        use_v2_format, stream);
+  } else {
+    vllm::gptq_rdna3::launch_gemm_q4<vllm::gptq_rdna3::bf16_t>(
+        (const vllm::gptq_rdna3::bf16_t*)a.data_ptr(),
+        (const uint32_t*)b_q_weight.data_ptr(),
+        (const uint32_t*)b_qzeros.data_ptr(),
+        (const vllm::gptq_rdna3::bf16_t*)b_scales.data_ptr(), g_idx_ptr,
+        (vllm::gptq_rdna3::bf16_t*)c.data_ptr(), size_m, size_n, size_k, groups,
+        use_v2_format, stream);
+  }
+
+  return c;
+}

csrc/rocm/qdq_4_rdna3.cuh

@@ -0,0 +1,239 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// W4A16 dequant primitives for RDNA3 (gfx1100/gfx1101/gfx1102), templated on
+// the activation/scale dtype (half or __hip_bfloat16). The fp16 path reuses
+// the classic exllamav2 bit-trick:
+//
+//   (qa & 0x000F000F) | 0x64006400  ->  half2(1024+q_lo, 1024+q_hi)
+//   (qa & 0x00F000F0) | 0x64006400  ->  half2(1024+q_lo*16, 1024+q_hi*16)
+//
+// The "*16 then divide by 16 in the FMA" trick for the upper-nibble pairs
+// works in fp16 because the mantissa (10 bits) is wide enough to hold a value
+// shifted by 4 bits. In bf16 the mantissa is only 7 bits, so shifting an upper
+// nibble into bits [7:4] would spill into the exponent. To avoid that, the
+// bf16 path shifts each pair of nibbles down to bits [3:0]/[19:16] with a
+// single right-shift before the OR with 0x43004300 (= bf162(128, 128)).
+
+#ifndef _qdq_4_rdna3_cuh
+#define _qdq_4_rdna3_cuh
+
+#include <cstdint>
+
+#include <hip/hip_bf16.h>
+#include <hip/hip_fp16.h>
+
+namespace vllm {
+namespace gptq_rdna3 {
+
+using bf16_t = __hip_bfloat16;
+using bf162_t = __hip_bfloat162;
+
+// Bit-shuffle for an int32 holding 8 sequential 4-bit weights q[0..7]:
+//   in:  q[7] q[6] q[5] q[4] q[3] q[2] q[1] q[0]   (LSB first)
+//   out: q[7] q[5] q[3] q[1] q[6] q[4] q[2] q[0]   (even/odd interleaved)
+//
+// After shuffle, q[2k]   sits at bits [4k   : 4k+3]   (lower 16)
+//                q[2k+1] sits at bits [16+4k: 16+4k+3] (upper 16)
+// so a single mask 0x000F000F selects the matching even/odd pair, ready to
+// bitcast to half2 / bfloat162 after OR-ing with the magic constant.
+__forceinline__ __device__ void shuffle_4bit_8(uint32_t* q) {
+  uint32_t qa = q[0];
+  uint32_t qb = 0;
+#pragma unroll
+  for (int i = 0; i < 4; i++) {
+    uint32_t qa0 = qa & 0x0F;
+    uint32_t qa1 = (qa & 0xF0) >> 4;
+    qa >>= 8;
+    qb |= (qa1 << (i * 4 + 16));
+    qb |= (qa0 << (i * 4));
+  }
+  q[0] = qb;
+}
+
+// ---------------------------------------------------------------------------
+// fp16 path
+// ---------------------------------------------------------------------------
+
+// Precompute scale-baked constants for a single zero/scale pair.
+//   z1z16[0] = scale * (-1024 - zero)            (used for "low" pairs)
+//   z1z16[1] = scale * (-64   - zero)            (used for "high" pairs)
+//   y1y16[0] = scale * 1                          (low pairs are q + 1024)
+//   y1y16[1] = scale * (1/16)                     (high pairs are q*16 + 1024)
+__forceinline__ __device__ void prep_zero_scale_fp16(uint32_t zero, half scale,
+                                                     half2 (&z1z16)[2],
+                                                     half2 (&y1y16)[2]) {
+  // half(-1024 - zero) via the exllamav2 bit-trick:
+  //   half bits 0xE400 == -1024.0 ; ORing the zero into mantissa subtracts it.
+  union {
+    uint16_t u;
+    half h;
+  } z1u;
+  z1u.u = (uint16_t)(0xE400 | zero);
+  half z1 = z1u.h;
+  half z16 = __hsub(__int2half_rn(-64), __int2half_rn((int)zero));
+
+  half2 scale2 = __half2half2(scale);
+  z1z16[0] = __hmul2(scale2, __half2half2(z1));
+  z1z16[1] = __hmul2(scale2, __half2half2(z16));
+
+  half y1 = __float2half_rn(1.0f);
+  half y16 = __float2half_rn(1.0f / 16.0f);
+  y1y16[0] = __hmul2(scale2, __half2half2(y1));
+  y1y16[1] = __hmul2(scale2, __half2half2(y16));
+}
+
+// Dequantize one int32 (8 shuffled 4-bit weights) into 4 half2 pairs:
+//   dq[0] = (q[0], q[1]) * scale - zero*scale
+//   dq[1] = (q[2], q[3]) * scale - zero*scale
+//   dq[2] = (q[4], q[5]) * scale - zero*scale
+//   dq[3] = (q[6], q[7]) * scale - zero*scale
+__forceinline__ __device__ void dequant_4bit_8_fp16(uint32_t qa, half2 (&dq)[4],
+                                                    half2 (&z1z16)[2],
+                                                    half2 (&y1y16)[2]) {
+  const uint32_t c0 = 0x64006400;
+
+  union {
+    uint32_t u;
+    half2 h2;
+  } q0, q1, q2, q3;
+  q0.u = (qa & 0x000F000F) | c0;  // half2(q[0]+1024, q[1]+1024)
+  q1.u = (qa & 0x00F000F0) | c0;  // half2(q[2]*16+1024, q[3]*16+1024)
+  uint32_t qa_hi = qa >> 8;
+  q2.u = (qa_hi & 0x000F000F) | c0;  // half2(q[4]+1024, q[5]+1024)
+  q3.u = (qa_hi & 0x00F000F0) | c0;  // half2(q[6]*16+1024, q[7]*16+1024)
+
+  dq[0] = __hfma2(q0.h2, y1y16[0], z1z16[0]);
+  dq[1] = __hfma2(q1.h2, y1y16[1], z1z16[1]);
+  dq[2] = __hfma2(q2.h2, y1y16[0], z1z16[0]);
+  dq[3] = __hfma2(q3.h2, y1y16[1], z1z16[1]);
+}
+
+// ---------------------------------------------------------------------------
+// bf16 path
+// ---------------------------------------------------------------------------
+
+// Bit-trick magic for bf16:
+//   bf16(128) == 0x4300 (sign 0, exp 134, mantissa 0).
+//   For nibble n in [0..15], bits [3:0] of mantissa hold n exactly because
+//   bf16's ULP at 128 is 1 (mantissa step = 2^(7-7) = 1). So
+//   ((qa & 0x000F000F) | 0x43004300) bitcasts to bfloat162(128+n_lo, 128+n_hi).
+//
+// Because bf16's mantissa is only 7 bits, we cannot use the fp16 "upper nibble
+// * 16" trick. Instead each pair of nibbles is shifted down to [3:0]/[19:16]
+// via a single 4/8/12-bit right-shift before the OR. That costs one extra
+// shift per pair vs fp16, but keeps the FMA structure identical.
+__forceinline__ __device__ void prep_zero_scale_bf16(uint32_t zero,
+                                                     bf16_t scale,
+                                                     bf162_t& z_prep,
+                                                     bf162_t& y_prep) {
+  // z = scale * -(128 + zero); y = scale.
+  float scale_f = __bfloat162float(scale);
+  float zf = -(128.0f + (float)zero) * scale_f;
+  bf16_t zb = __float2bfloat16(zf);
+  z_prep = __bfloat162bfloat162(zb);
+  y_prep = __bfloat162bfloat162(scale);
+}
+
+__forceinline__ __device__ void dequant_4bit_8_bf16(uint32_t qa,
+                                                    bf162_t (&dq)[4],
+                                                    bf162_t z_prep,
+                                                    bf162_t y_prep) {
+  const uint32_t c0 = 0x43004300;
+
+  union {
+    uint32_t u;
+    bf162_t b2;
+  } q0, q1, q2, q3;
+  q0.u = ((qa >> 0) & 0x000F000F) | c0;   // bf162(128+q[0], 128+q[1])
+  q1.u = ((qa >> 4) & 0x000F000F) | c0;   // bf162(128+q[2], 128+q[3])
+  q2.u = ((qa >> 8) & 0x000F000F) | c0;   // bf162(128+q[4], 128+q[5])
+  q3.u = ((qa >> 12) & 0x000F000F) | c0;  // bf162(128+q[6], 128+q[7])
+
+  // dq = q_b * scale + (-(128+zero)*scale) = (q - zero) * scale
+  dq[0] = __hfma2(q0.b2, y_prep, z_prep);
+  dq[1] = __hfma2(q1.b2, y_prep, z_prep);
+  dq[2] = __hfma2(q2.b2, y_prep, z_prep);
+  dq[3] = __hfma2(q3.b2, y_prep, z_prep);
+}
+
+// ---------------------------------------------------------------------------
+// bf16-input → fp32-output dequant (RDNA3 scalar path).
+//
+// RDNA3 (gfx1100) has no v_pk_fma_bf16; packed bf16 FMA lowers to a slow
+// fallback. Rather than computing dq in bf16 and widening at FMA time in
+// the dot product, we widen to fp32 here once (a free left-shift by 16) and
+// emit the (q - zero) * scale FMA directly in fp32. This:
+//   * Replaces 4× slow bf16 packed FMA with 8× fast fp32 FMA per int32.
+//   * Eliminates 4× bf16→fp32 widens that the dot product would do.
+//   * Keeps the dot product accumulator in fp32 without a roundtrip.
+//
+// Output: fp32 dq[8], one element per K position (consumed by the
+// fp32-overload of dot22_8_f in q_gemm_rdna3.cu).
+__forceinline__ __device__ void prep_zero_scale_bf16_f32(uint32_t zero,
+                                                         bf16_t scale,
+                                                         float& z_prep,
+                                                         float& y_prep) {
+  float scale_f = __bfloat162float(scale);
+  z_prep = -(128.0f + (float)zero) * scale_f;
+  y_prep = scale_f;
+}
+
+// Pure-q dequant for the M_COUNT=1 factored path: outputs the unscaled fp32
+// values 128+nibble, without folding scale/zero. The caller folds scale/zb
+// into the accumulator outside the inner loop using a precomputed sum_a,
+// which saves ~27% of the FMA count vs the per-col-dequant approach above
+// (only beneficial at M_COUNT=1; break-even at M_COUNT=2).
+//
+// Cost: 0 FMAs (pure bit-trick + as_float reinterprets).
+__forceinline__ __device__ void dequant_4bit_8_bf16_q_only(uint32_t qa,
+                                                           float (&q_f32)[8]) {
+  const uint32_t c0 = 0x43004300;
+  const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0;
+  const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0;
+  const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0;
+  const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0;
+  q_f32[0] = __uint_as_float((q0 & 0xFFFFu) << 16);
+  q_f32[1] = __uint_as_float(q0 & 0xFFFF0000u);
+  q_f32[2] = __uint_as_float((q1 & 0xFFFFu) << 16);
+  q_f32[3] = __uint_as_float(q1 & 0xFFFF0000u);
+  q_f32[4] = __uint_as_float((q2 & 0xFFFFu) << 16);
+  q_f32[5] = __uint_as_float(q2 & 0xFFFF0000u);
+  q_f32[6] = __uint_as_float((q3 & 0xFFFFu) << 16);
+  q_f32[7] = __uint_as_float(q3 & 0xFFFF0000u);
+}
+
+__forceinline__ __device__ void dequant_4bit_8_bf16_f32(uint32_t qa,
+                                                        float (&dq)[8],
+                                                        float z_prep,
+                                                        float y_prep) {
+  const uint32_t c0 = 0x43004300;
+  const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0;
+  const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0;
+  const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0;
+  const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0;
+  // bf16(128+nibble) bits → fp32(128+nibble) bits via left-shift by 16
+  // (just zero-extends the mantissa from 7 to 23 bits; exponent preserved).
+  const float q0x = __uint_as_float((q0 & 0xFFFFu) << 16);
+  const float q0y = __uint_as_float(q0 & 0xFFFF0000u);
+  const float q1x = __uint_as_float((q1 & 0xFFFFu) << 16);
+  const float q1y = __uint_as_float(q1 & 0xFFFF0000u);
+  const float q2x = __uint_as_float((q2 & 0xFFFFu) << 16);
+  const float q2y = __uint_as_float(q2 & 0xFFFF0000u);
+  const float q3x = __uint_as_float((q3 & 0xFFFFu) << 16);
+  const float q3y = __uint_as_float(q3 & 0xFFFF0000u);
+  // dq[i] = q_f32 * scale + (-(128+zero)*scale) = (nibble - zero) * scale
+  dq[0] = __fmaf_rn(q0x, y_prep, z_prep);
+  dq[1] = __fmaf_rn(q0y, y_prep, z_prep);
+  dq[2] = __fmaf_rn(q1x, y_prep, z_prep);
+  dq[3] = __fmaf_rn(q1y, y_prep, z_prep);
+  dq[4] = __fmaf_rn(q2x, y_prep, z_prep);
+  dq[5] = __fmaf_rn(q2y, y_prep, z_prep);
+  dq[6] = __fmaf_rn(q3x, y_prep, z_prep);
+  dq[7] = __fmaf_rn(q3y, y_prep, z_prep);
+}
+
+}  // namespace gptq_rdna3
+}  // namespace vllm
+
+#endif  // _qdq_4_rdna3_cuh

csrc/rocm/torch_bindings.cpp

@@ -39,6 +39,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
       "          Tensor scale_b, int CuCount) -> ()");
   rocm_ops.impl("wvSplitKQ", torch::kCUDA, &wvSplitKQ);
 
+#ifdef VLLM_ROCM_GFX1100
+  // W4A16 GPTQ kernels for AMD RDNA3 (gfx1100).
+  rocm_ops.def(
+      "gptq_gemm_rdna3(Tensor a, Tensor b_q_weight, Tensor b_qzeros, "
+      "Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor");
+  rocm_ops.impl("gptq_gemm_rdna3", torch::kCUDA, &gptq_gemm_rdna3);
+
+  rocm_ops.def(
+      "gptq_gemm_rdna3_wmma(Tensor a, Tensor b_q_weight, Tensor b_qzeros, "
+      "Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor");
+  rocm_ops.impl("gptq_gemm_rdna3_wmma", torch::kCUDA, &gptq_gemm_rdna3_wmma);
+#endif
+
   // Custom attention op
   // Compute the attention between an input query and the cached
   // keys/values using PagedAttention.

csrc/torch_bindings.cpp

@@ -1,4 +1,7 @@
-#include "cache.h"
+// Provides torch::Tensor for ops.h (previously included transitively via
+// cache.h, which is no longer included here after cache ops moved to
+// _C_stable_libtorch).
+#include <torch/all.h>
 #include "cuda_utils.h"
 #include "ops.h"
 #include "core/registration.h"
@@ -33,35 +36,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
            &get_cuda_view_from_cpu_tensor);
 
-  // Attention ops
-  // Compute the attention between an input query and the cached
-  // keys/values using PagedAttention.
-  ops.def(
-      "paged_attention_v1("
-      "    Tensor! out, Tensor query, Tensor key_cache,"
-      "    Tensor value_cache, int num_kv_heads, float scale,"
-      "    Tensor block_tables, Tensor seq_lens, int block_size,"
-      "    int max_seq_len, Tensor? alibi_slopes,"
-      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
-      "    int tp_rank, int blocksparse_local_blocks,"
-      "    int blocksparse_vert_stride, int blocksparse_block_size,"
-      "    int blocksparse_head_sliding_step) -> ()");
-  ops.impl("paged_attention_v1", torch::kCUDA, &paged_attention_v1);
-
-  // PagedAttention V2.
-  ops.def(
-      "paged_attention_v2("
-      "    Tensor! out, Tensor! exp_sums, Tensor! max_logits,"
-      "    Tensor! tmp_out, Tensor query, Tensor key_cache,"
-      "    Tensor value_cache, int num_kv_heads, float scale,"
-      "    Tensor block_tables, Tensor seq_lens, int block_size,"
-      "    int max_seq_len, Tensor? alibi_slopes,"
-      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
-      "    int tp_rank, int blocksparse_local_blocks,"
-      "    int blocksparse_vert_stride, int blocksparse_block_size,"
-      "    int blocksparse_head_sliding_step) -> ()");
-  ops.impl("paged_attention_v2", torch::kCUDA, &paged_attention_v2);
-
   // Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
   ops.def(
       "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
@@ -217,114 +191,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 #endif
 }
 
-TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
-  // Cache ops
-  // Swap in (out) the cache blocks from src to dst.
-  cache_ops.def(
-      "swap_blocks(Tensor src, Tensor! dst,"
-      "            int block_size_in_bytes, Tensor block_mapping) -> ()");
-  cache_ops.impl("swap_blocks", torch::kCUDA, &swap_blocks);
-
-  // Batch swap: submit all block copies in a single driver call.
-  cache_ops.def(
-      "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs,"
-      "                  Tensor sizes,"
-      "                  bool is_src_access_order_any=False) -> ()");
-  cache_ops.impl("swap_blocks_batch", torch::kCPU, &swap_blocks_batch);
-
-  // Reshape the key and value tensors and cache them.
-  cache_ops.def(
-      "reshape_and_cache(Tensor key, Tensor value,"
-      "                  Tensor! key_cache, Tensor! value_cache,"
-      "                  Tensor slot_mapping,"
-      "                  str kv_cache_dtype,"
-      "                  Tensor k_scale, Tensor v_scale) -> ()");
-  cache_ops.impl("reshape_and_cache", torch::kCUDA, &reshape_and_cache);
-
-  // Reshape the key and value tensors and cache them.
-  cache_ops.def(
-      "reshape_and_cache_flash(Tensor key, Tensor value,"
-      "                        Tensor! key_cache,"
-      "                        Tensor! value_cache,"
-      "                        Tensor slot_mapping,"
-      "                        str kv_cache_dtype,"
-      "                        Tensor k_scale, Tensor v_scale) -> ()");
-  cache_ops.impl("reshape_and_cache_flash", torch::kCUDA,
-                 &reshape_and_cache_flash);
-
-  // Concat kv_c and k_pe and cache them.
-  cache_ops.def(
-      "concat_and_cache_mla(Tensor kv_c, Tensor k_pe,"
-      "                     Tensor! kv_cache,"
-      "                     Tensor slot_mapping,"
-      "                     str kv_cache_dtype,"
-      "                     Tensor scale) -> ()");
-  cache_ops.impl("concat_and_cache_mla", torch::kCUDA, &concat_and_cache_mla);
-
-  // Rotate Q and K, then write to kv cache for MLA
-  cache_ops.def(
-      "concat_and_cache_mla_rope_fused("
-      "                     Tensor positions,"
-      "                     Tensor! q_pe,"
-      "                     Tensor! k_pe,"
-      "                     Tensor kv_c,"
-      "                     Tensor cos_sin_cache,"
-      "                     bool is_neox,"
-      "                     Tensor slot_mapping,"
-      "                     Tensor! kv_cache,"
-      "                     str kv_cache_dtype,"
-      "                     Tensor kv_cache_scale) -> ()");
-  cache_ops.impl("concat_and_cache_mla_rope_fused", torch::kCUDA,
-                 &concat_and_cache_mla_rope_fused);
-
-  // Convert the key and value cache to fp8 data type.
-  cache_ops.def(
-      "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, "
-      "str kv_cache_dtype) -> ()");
-  cache_ops.impl("convert_fp8", torch::kCUDA, &convert_fp8);
-
-  // Gather cache blocks from src_cache to dst, dequantizing from
-  // src_cache's dtype to dst's dtype if necessary.
-  cache_ops.def(
-      "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, "
-      "                               Tensor block_table, Tensor cu_seq_lens, "
-      "                               Tensor token_to_seq, "
-      "                               int num_tokens, "
-      "                               str kv_cache_dtype, "
-      "                               Tensor scale, Tensor? seq_starts) -> ()");
-  cache_ops.impl("gather_and_maybe_dequant_cache", torch::kCUDA,
-                 &gather_and_maybe_dequant_cache);
-
-  cache_ops.def(
-      "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, "
-      "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()");
-  cache_ops.impl("cp_gather_cache", torch::kCUDA, &cp_gather_cache);
-
-  cache_ops.def(
-      "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, "
-      "Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int "
-      "batch_size) -> ()");
-  cache_ops.impl("cp_gather_and_upconvert_fp8_kv_cache", torch::kCUDA,
-                 &cp_gather_and_upconvert_fp8_kv_cache);
-
-  cache_ops.def(
-      "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor "
-      "slot_mapping, "
-      "int quant_block_size, str kv_cache_dtype) -> ()");
-  cache_ops.impl("indexer_k_quant_and_cache", torch::kCUDA,
-                 &indexer_k_quant_and_cache);
-
-  cache_ops.def(
-      "concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()");
-  cache_ops.impl("concat_mla_q", torch::kCUDA, &concat_mla_q);
-
-  cache_ops.def(
-      "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! "
-      "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()");
-  cache_ops.impl("cp_gather_indexer_k_quant_cache", torch::kCUDA,
-                 &cp_gather_indexer_k_quant_cache);
-}
-
 TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
   // Cuda utils
 

docker/Dockerfile.cpu

@@ -165,6 +165,23 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=cache,target=/vllm-workspace/.deps,sharing=locked \
     VLLM_TARGET_DEVICE=cpu python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38
 
+######################### TRITON-CPU BUILD IMAGE #########################
+FROM base AS vllm-triton-cpu-build
+
+WORKDIR /vllm-workspace
+
+RUN mkdir dist
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=cache,target=/root/.cache/ccache \
+    --mount=type=cache,target=/vllm-workspace/.deps,sharing=locked \
+    if [ "$TARGETARCH" = "amd64" ] || [ "$VLLM_CPU_X86" != "0" ]; then \
+        git clone --recurse-submodules "https://github.com/triton-lang/triton-cpu.git"; \
+        cd triton-cpu; \
+        git checkout "270e696d"; \
+        uv build --wheel --out-dir=../dist; \
+    fi
+
 ######################### TEST DEPS #########################
 FROM base AS vllm-test-deps
 
@@ -257,7 +274,14 @@ WORKDIR /vllm-workspace
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=cache,target=/root/.cache/ccache \
     --mount=type=bind,from=vllm-build,src=/vllm-workspace/dist,target=dist \
-    uv pip install "$(realpath dist/*.whl)[audio,triton-cpu]"
+    uv pip install "$(realpath dist/*.whl)[audio]"
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=cache,target=/root/.cache/ccache \
+    --mount=type=bind,from=vllm-triton-cpu-build,src=/vllm-workspace/dist,target=dist \
+    if [ "$TARGETARCH" = "amd64" ] || [ "$VLLM_CPU_X86" != "0" ]; then \
+        uv pip install "$(realpath dist/*.whl)"; \
+    fi
 
 # Add labels to document build configuration
 LABEL org.opencontainers.image.title="vLLM CPU"

docs/configuration/optimization.md

@@ -46,14 +46,14 @@ In V1, **chunked prefill is enabled by default whenever possible**. With chunked
 
 This policy has two benefits:
 
-- It improves ITL and generation decode because decode requests are prioritized.
+- It improves inter-token latency (ITL) and generation decode because decode requests are prioritized.
 - It helps achieve better GPU utilization by locating compute-bound (prefill) and memory-bound (decode) requests to the same batch.
 
 ### Performance Tuning with Chunked Prefill
 
 You can tune the performance by adjusting `max_num_batched_tokens`:
 
-- Smaller values (e.g., 2048) achieve better inter-token latency (ITL) because there are fewer prefills slowing down decodes.
+- Smaller values (e.g., 2048) achieve better ITL because there are fewer prefills slowing down decodes.
 - Higher values achieve better time to first token (TTFT) as you can process more prefill tokens in a batch.
 - For optimal throughput, we recommend setting `max_num_batched_tokens > 8192` especially for smaller models on large GPUs.
 - If `max_num_batched_tokens` is the same as `max_model_len`, that's almost the equivalent to the V0 default scheduling policy (except that it still prioritizes decodes).

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

@@ -41,7 +41,7 @@ pip install -v -r requirements/xpu.txt
 
     ```bash
     pip uninstall -y triton triton-xpu
-    pip install triton-xpu==3.6.0 --extra-index-url https://download.pytorch.org/whl/xpu
+    pip install triton-xpu==3.7.0 --extra-index-url https://download.pytorch.org/whl/xpu
     ```
 
     !!! note

docs/models/supported_models.md

@@ -633,6 +633,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `SmolVLMForConditionalGeneration` | SmolVLM2 | T + I | `SmolVLM2-2.2B-Instruct` | ✅︎ | |
 | `Step3VLForConditionalGeneration` | Step3-VL | T + I<sup>+</sup> | `stepfun-ai/step3` | | ✅︎ |
 | `StepVLForConditionalGeneration` | Step3-VL-10B | T + I<sup>+</sup> | `stepfun-ai/Step3-VL-10B` | | ✅︎ |
+| `Step3p7ForConditionalGeneration` | Step-3.7-Flash | T + I<sup>+</sup> | `stepfun-ai/Step-3.7-Flash` | | ✅︎ |
 | `TarsierForConditionalGeneration` | Tarsier | T + I<sup>E+</sup> | `omni-search/Tarsier-7b`, `omni-search/Tarsier-34b` | | ✅︎ |
 | `Tarsier2ForConditionalGeneration`<sup>^</sup> | Tarsier2 | T + I<sup>E+</sup> + V<sup>E+</sup> | `omni-research/Tarsier2-Recap-7b`, `omni-research/Tarsier2-7b-0115` | | ✅︎ |
 | `UltravoxModel` | Ultravox | T + A<sup>E+</sup> | `fixie-ai/ultravox-v0_5-llama-3_2-1b` | ✅︎ | ✅︎ |

docs/pre_run_check.sh

@@ -1,41 +1,60 @@
-if [ "$READTHEDOCS_VERSION_TYPE" = "external" ]; then
-  MAX_WAIT=300
-  INTERVAL=60
-  ELAPSED=0
-  while :; do
-    RAW=$(curl -sS -w "\n%{http_code}" "https://api.github.com/repos/vllm-project/vllm/commits/${READTHEDOCS_GIT_COMMIT_HASH}/check-runs?check_name=pre-run-check&filter=latest")
-    HTTP_CODE=$(printf %s "$RAW" | tail -n1)
-    BODY=$(printf %s "$RAW" | sed '$d')
-    if [ "$HTTP_CODE" != "200" ]; then
-      echo "GitHub API returned HTTP $HTTP_CODE (likely rate-limited); skipping pre-run-check gate."
-      break
-    fi
-    STATUS=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"status\") or \"\") if r else \"none\")")
-    CONCLUSION=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"conclusion\") or \"\") if r else \"\")")
-    CHECK_URL=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"html_url\") or \"\") if r else \"\")")
-    if [ "$STATUS" = "none" ]; then
-      echo "no pre-run-check found for this commit; skipping gate."
-      break
-    fi
-    if [ -n "$CONCLUSION" ]; then
-      echo "pre-run-check conclusion: $CONCLUSION"
-      if [ "$CONCLUSION" = "failure" ] || [ "$CONCLUSION" = "cancelled" ] || [ "$CONCLUSION" = "timed_out" ]; then
-        echo "pre-run-check did not pass; skipping docs build."
-        if [ -n "$CHECK_URL" ]; then
-          echo "pre-run-check failure reason: $CHECK_URL"
-        fi
-        exit 1
-      fi
-      break
-    fi
-    if [ "$ELAPSED" -ge "$MAX_WAIT" ]; then
-      echo "pre-run-check status=$STATUS after ${MAX_WAIT}s; skipping gate."
-      break
-    fi
-    echo "pre-run-check status=$STATUS; waiting ${INTERVAL}s..."
-    sleep "$INTERVAL"
-    ELAPSED=$((ELAPSED + INTERVAL))
-  done
-else
+if [ "$READTHEDOCS_VERSION_TYPE" != "external" ]; then
   echo "Not a PR build (version type=$READTHEDOCS_VERSION_TYPE); skipping pre-run-check gate."
-fi
+  exit 0
+fi
+
+echo "Checking for changes to docs-affecting files vs origin/main..."
+DOCS_PATHS=(
+  docs/                       # Actual docs content
+  examples/                   # Examples are rendered in docs
+  vllm/                       # API & CLI reference
+  requirements/test/cuda.txt  # CLI reference (see docs/mkdocs/hooks/generate_argparse.py)
+  mkdocs.yaml                 # Affects build process
+  .readthedocs.yaml           # Affects build process
+  requirements/docs.txt       # Affects build process
+  requirements/docs.in        # Affects build process
+)
+if git diff --quiet origin/main -- "${DOCS_PATHS[@]}"; then
+  echo "No docs-affecting files changed vs origin/main; cancelling build."
+  # See https://docs.readthedocs.com/platform/latest/guides/build/skip-build.html for info on exit code
+  exit 183
+fi
+echo "Docs-affecting files changed; continuing pre-run-check."
+echo "Checking pre-commit/pre-run-check status..."
+MAX_WAIT=300
+INTERVAL=60
+ELAPSED=0
+while :; do
+  RAW=$(curl -sS -w "\n%{http_code}" "https://api.github.com/repos/vllm-project/vllm/commits/${READTHEDOCS_GIT_COMMIT_HASH}/check-runs?check_name=pre-run-check&filter=latest")
+  HTTP_CODE=$(printf %s "$RAW" | tail -n1)
+  BODY=$(printf %s "$RAW" | sed '$d')
+  if [ "$HTTP_CODE" != "200" ]; then
+    echo "GitHub API returned HTTP $HTTP_CODE (likely rate-limited); skipping pre-commit/pre-run-check gate."
+    break
+  fi
+  STATUS=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"status\") or \"\") if r else \"none\")")
+  CONCLUSION=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"conclusion\") or \"\") if r else \"\")")
+  CHECK_URL=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"html_url\") or \"\") if r else \"\")")
+  if [ "$STATUS" = "none" ]; then
+    echo "no pre-commit/pre-run-check found for this commit; skipping gate."
+    break
+  fi
+  if [ -n "$CONCLUSION" ]; then
+    echo "pre-commit/pre-run-check conclusion: $CONCLUSION"
+    if [ "$CONCLUSION" = "failure" ] || [ "$CONCLUSION" = "cancelled" ] || [ "$CONCLUSION" = "timed_out" ]; then
+      echo "pre-commit/pre-run-check did not pass; skipping docs build."
+      if [ -n "$CHECK_URL" ]; then
+        echo "pre-commit/pre-run-check failure reason: $CHECK_URL"
+      fi
+      exit 1
+    fi
+    break
+  fi
+  if [ "$ELAPSED" -ge "$MAX_WAIT" ]; then
+    echo "pre-commit/pre-run-check status=$STATUS after ${MAX_WAIT}s; skipping gate."
+    break
+  fi
+  echo "pre-commit/pre-run-check status=$STATUS; waiting ${INTERVAL}s..."
+  sleep "$INTERVAL"
+  ELAPSED=$((ELAPSED + INTERVAL))
+done

docs/serving/expert_parallel_deployment.md

@@ -151,7 +151,7 @@ Configure EPLB with the `--eplb-config` argument, which accepts a JSON string. T
 | `step_interval` | Frequency of rebalancing (every N engine steps) | 3000 |
 | `log_balancedness` | Log balancedness metrics (avg tokens per expert ÷ max tokens per expert) | `false` |
 | `num_redundant_experts` | Additional global experts per EP rank beyond equal distribution | `0` |
-| `use_async` | Use non-blocking EPLB for reduced latency overhead | `false` |
+| `use_async` | Use non-blocking EPLB for reduced latency overhead | `true` |
 | `policy` | The policy type for expert parallel load balancing | `"default"` |
 | `communicator` | Backend for expert weight transfers: `"torch_nccl"`, `"torch_gloo"`, `"pynccl"`, `"nixl"`,  or `null` (auto) | `null` |
 

examples/tool_chat_template_gemma4.jinja

@@ -295,6 +295,15 @@
                                 {%- endif -%}
                             {%- endfor -%}
                             {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}
+                            {%- for part in tool_body -%}
+                                {%- if part.get('type') == 'image' -%}
+                                    {{- '<|image|>' -}}
+                                {%- elif part.get('type') == 'audio' -%}
+                                    {{- '<|audio|>' -}}
+                                {%- elif part.get('type') == 'video' -%}
+                                    {{- '<|video|>' -}}
+                                {%- endif -%}
+                            {%- endfor -%}
                         {%- else -%}
                             {{- format_tool_response_block(ns_tname.name, tool_body) -}}
                         {%- endif -%}

requirements/xpu.txt

@@ -17,4 +17,4 @@ torchaudio
 torchvision
 
 auto_round_lib>=0.13.0
-vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.8/vllm_xpu_kernels-0.1.8-cp38-abi3-manylinux_2_28_x86_64.whl
+vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.9/vllm_xpu_kernels-0.1.9-cp38-abi3-manylinux_2_28_x86_64.whl

rust/src/chat/src/lib.rs

@@ -93,7 +93,7 @@ pub struct ChatLlm {
     text: TextLlm,
     backend: DynChatBackend,
     /// Effective model dtype reported by the engine.
-    model_dtype: Option<ModelDtype>,
+    model_dtype: ModelDtype,
     /// Tool-call parser selection.
     tool_call_parser: ParserSelection,
     /// Reasoning parser selection.
@@ -135,7 +135,7 @@ impl ChatLlm {
     }
 
     /// Override the effective model dtype used for multimodal tensor encoding.
-    pub fn with_model_dtype(mut self, model_dtype: Option<ModelDtype>) -> Self {
+    pub fn with_model_dtype(mut self, model_dtype: ModelDtype) -> Self {
         self.model_dtype = model_dtype;
         self
     }

rust/src/chat/src/multimodal.rs

@@ -12,7 +12,7 @@
 use std::collections::{HashMap, HashSet};
 use std::fs;
 use std::path::Path;
-use std::sync::{Arc, LazyLock, Once};
+use std::sync::{Arc, LazyLock};
 
 use itertools::izip;
 use llm_multimodal::{
@@ -239,7 +239,7 @@ pub(crate) async fn finalize_rendered_prompt(
     request: &ChatRequest,
     rendered: RenderedPrompt,
     info: Option<&MultimodalModelInfo>,
-    model_dtype: Option<ModelDtype>,
+    model_dtype: ModelDtype,
 ) -> Result<(Prompt, Option<MmFeatures>)> {
     if !request.has_multimodal() {
         return Ok((rendered.prompt, None));
@@ -249,16 +249,6 @@ pub(crate) async fn finalize_rendered_prompt(
         bail_multimodal!("multimodal chat renderer must return a text prompt before expansion");
     };
     let media_parts = extract_media_parts(request)?;
-    let model_dtype = model_dtype.unwrap_or_else(|| {
-        static WARN_ONCE: Once = Once::new();
-        WARN_ONCE.call_once(|| {
-            warn!(
-                "engine handshake did not report model dtype; \
-                 falling back to float32 for multimodal tensor encoding"
-            );
-        });
-        ModelDtype::Float32
-    });
 
     let mut prompt_token_ids = info
         .context

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

@@ -290,10 +290,8 @@ impl EngineCoreClient {
 
         // If any engine reported a dp_stats_address in its ready response, use it
         // as the external coordinator address.
-        let dp_stats_address: Option<String> = engines
-            .iter()
-            .filter_map(|e| e.ready_response.as_ref())
-            .find_map(|r| r.dp_stats_address.clone());
+        let dp_stats_address: Option<String> =
+            engines.iter().find_map(|engine| engine.ready_response.dp_stats_address.clone());
 
         let (coordinator, coordinator_output_task, coordinator_task) =
             if let Some(coordinator_transport) = connected.coordinator {
@@ -368,40 +366,44 @@ impl EngineCoreClient {
     /// Return the ready responses received from all engines on the input
     /// socket.
     pub fn ready_responses(&self) -> Vec<&EngineCoreReadyResponse> {
-        self.engines
-            .iter()
-            .filter_map(|engine| engine.ready_response.as_ref())
-            .collect()
+        self.engines.iter().map(|engine| &engine.ready_response).collect()
     }
 
-    /// Return the engine-reported effective model dtype, when available.
-    pub fn model_dtype(&self) -> Option<ModelDtype> {
+    /// Return the engine-reported effective model dtype.
+    pub fn model_dtype(&self) -> ModelDtype {
         self.engines
-            .iter()
-            .filter_map(|engine| engine.ready_response.as_ref())
-            .find_map(|response| response.dtype)
+            .first()
+            .expect("engine core client requires at least one engine")
+            .ready_response
+            .dtype
+    }
+
+    /// Return the engine-reported Python vLLM version.
+    pub fn vllm_version(&self) -> &str {
+        self.engines
+            .first()
+            .expect("engine core client requires at least one engine")
+            .ready_response
+            .vllm_version
+            .as_str()
     }
 
     /// Return the total number of GPU blocks summed across all connected
     /// engines.
     pub fn total_num_gpu_blocks(&self) -> u64 {
-        self.engines
-            .iter()
-            .filter_map(|engine| engine.ready_response.as_ref())
-            .map(|r| r.num_gpu_blocks)
-            .sum()
+        self.engines.iter().map(|engine| engine.ready_response.num_gpu_blocks).sum()
     }
 
     /// Return the minimum engine-reported `max_model_len` across all engines.
     ///
     /// This is the auto-fitted value after KV cache profiling and may differ
     /// from the originally configured value.
-    pub fn max_model_len(&self) -> Option<u32> {
+    pub fn max_model_len(&self) -> u32 {
         self.engines
             .iter()
-            .filter_map(|e| e.ready_response.as_ref())
-            .map(|r| r.max_model_len as u32)
+            .map(|engine| engine.ready_response.max_model_len as u32)
             .min()
+            .expect("engine core client requires at least one engine")
     }
 
     /// Get the model name associated with this client used for metrics

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

@@ -382,6 +382,7 @@ mod tests {
     use zeromq::{RouterSocket, Socket};
 
     use super::*;
+    use crate::mock_engine::default_ready_response;
 
     async fn test_inner() -> ClientInner {
         let mut socket = RouterSocket::new();
@@ -392,7 +393,7 @@ mod tests {
             "test-model".to_string(),
             &[ConnectedEngine {
                 engine_id: EngineId::from(b"engine-0"),
-                ready_response: None,
+                ready_response: default_ready_response(),
             }],
         )
     }

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

@@ -339,15 +339,20 @@ mod tests {
     use super::{EngineRoutingState, RequestRegistry, UtilityRegistry};
     use crate::EngineId;
     use crate::client::state::EngineLoadSnapshot;
+    use crate::mock_engine::default_ready_response;
     use crate::protocol::{EngineCoreFinishReason, EngineCoreOutput};
     use crate::transport::ConnectedEngine;
 
+    fn connected_engine(engine_id: EngineId) -> ConnectedEngine {
+        ConnectedEngine {
+            engine_id,
+            ready_response: default_ready_response(),
+        }
+    }
+
     #[test]
     fn registry_rejects_duplicate_request_ids() {
-        let mut registry = RequestRegistry::new(&[ConnectedEngine {
-            engine_id: EngineId::from(b"engine-0"),
-            ready_response: None,
-        }]);
+        let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]);
         registry.register("req-1".to_string(), None).unwrap();
         let error = registry.register("req-1".to_string(), None).unwrap_err();
         assert!(matches!(
@@ -358,10 +363,7 @@ mod tests {
 
     #[test]
     fn registry_removes_finished_request_on_output() {
-        let mut registry = RequestRegistry::new(&[ConnectedEngine {
-            engine_id: EngineId::from(b"engine-0"),
-            ready_response: None,
-        }]);
+        let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]);
         registry.register("req-1".to_string(), None).unwrap();
 
         let sender = registry.sender_for_output(&EngineCoreOutput {
@@ -376,10 +378,7 @@ mod tests {
 
     #[test]
     fn registry_closes_all_requests_on_failure() {
-        let mut registry = RequestRegistry::new(&[ConnectedEngine {
-            engine_id: EngineId::from(b"engine-0"),
-            ready_response: None,
-        }]);
+        let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]);
         registry.register("req-1".to_string(), None).unwrap();
         registry.register("req-2".to_string(), None).unwrap();
 
@@ -394,14 +393,8 @@ mod tests {
         let engine_0 = EngineId::from_engine_index(0);
         let engine_1 = EngineId::from_engine_index(1);
         let mut registry = RequestRegistry::new(&[
-            ConnectedEngine {
-                engine_id: engine_0.clone(),
-                ready_response: None,
-            },
-            ConnectedEngine {
-                engine_id: engine_1.clone(),
-                ready_response: None,
-            },
+            connected_engine(engine_0.clone()),
+            connected_engine(engine_1.clone()),
         ]);
         let (chosen_0, _) = registry.register("req-1".to_string(), None).unwrap();
         let (chosen_1, _) = registry.register("req-2".to_string(), None).unwrap();
@@ -428,14 +421,8 @@ mod tests {
         let engine_0 = EngineId::from_engine_index(0);
         let engine_1 = EngineId::from_engine_index(1);
         let mut registry = RequestRegistry::new(&[
-            ConnectedEngine {
-                engine_id: engine_0.clone(),
-                ready_response: None,
-            },
-            ConnectedEngine {
-                engine_id: engine_1.clone(),
-                ready_response: None,
-            },
+            connected_engine(engine_0.clone()),
+            connected_engine(engine_1.clone()),
         ]);
 
         let (chosen_0, _) = registry.register("req-1".to_string(), None).unwrap();
@@ -488,14 +475,8 @@ mod tests {
         let engine_0 = EngineId::from_engine_index(0);
         let engine_1 = EngineId::from_engine_index(1);
         let mut registry = RequestRegistry::new(&[
-            ConnectedEngine {
-                engine_id: engine_0.clone(),
-                ready_response: None,
-            },
-            ConnectedEngine {
-                engine_id: engine_1.clone(),
-                ready_response: None,
-            },
+            connected_engine(engine_0.clone()),
+            connected_engine(engine_1.clone()),
         ]);
 
         assert!(registry.apply_scheduler_counts(
@@ -523,18 +504,9 @@ mod tests {
         let engine_1 = EngineId::from_engine_index(1);
         let engine_2 = EngineId::from_engine_index(2);
         let mut registry = RequestRegistry::new(&[
-            ConnectedEngine {
-                engine_id: engine_0.clone(),
-                ready_response: None,
-            },
-            ConnectedEngine {
-                engine_id: engine_1.clone(),
-                ready_response: None,
-            },
-            ConnectedEngine {
-                engine_id: engine_2.clone(),
-                ready_response: None,
-            },
+            connected_engine(engine_0.clone()),
+            connected_engine(engine_1.clone()),
+            connected_engine(engine_2.clone()),
         ]);
 
         // Explicitly target rank 2 (third engine).
@@ -555,14 +527,8 @@ mod tests {
         let engine_0 = EngineId::from_engine_index(0);
         let engine_1 = EngineId::from_engine_index(1);
         let mut registry = RequestRegistry::new(&[
-            ConnectedEngine {
-                engine_id: engine_0.clone(),
-                ready_response: None,
-            },
-            ConnectedEngine {
-                engine_id: engine_1.clone(),
-                ready_response: None,
-            },
+            connected_engine(engine_0.clone()),
+            connected_engine(engine_1.clone()),
         ]);
 
         // Load-balance: first two go to engine_0 and engine_1.
@@ -577,14 +543,8 @@ mod tests {
     #[test]
     fn register_with_out_of_range_rank_returns_error() {
         let mut registry = RequestRegistry::new(&[
-            ConnectedEngine {
-                engine_id: EngineId::from_engine_index(0),
-                ready_response: None,
-            },
-            ConnectedEngine {
-                engine_id: EngineId::from_engine_index(1),
-                ready_response: None,
-            },
+            connected_engine(EngineId::from_engine_index(0)),
+            connected_engine(EngineId::from_engine_index(1)),
         ]);
 
         let error = registry.register("req-1".to_string(), Some(2)).unwrap_err();
@@ -600,10 +560,7 @@ mod tests {
     #[test]
     fn register_with_rank_on_single_engine_only_accepts_zero() {
         let engine_0 = EngineId::from_engine_index(0);
-        let mut registry = RequestRegistry::new(&[ConnectedEngine {
-            engine_id: engine_0.clone(),
-            ready_response: None,
-        }]);
+        let mut registry = RequestRegistry::new(&[connected_engine(engine_0.clone())]);
 
         let (chosen, _) = registry.register("req-ok".to_string(), Some(0)).unwrap();
         assert_eq!(chosen, engine_0);

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

@@ -47,7 +47,8 @@ pub fn default_ready_response() -> EngineCoreReadyResponse {
         max_model_len: DEFAULT_MOCK_MAX_MODEL_LEN,
         num_gpu_blocks: DEFAULT_MOCK_NUM_GPU_BLOCKS,
         dp_stats_address: None,
-        dtype: Some(ModelDtype::Float32),
+        dtype: ModelDtype::Float32,
+        vllm_version: "test-vllm-version".to_string(),
     }
 }
 

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

@@ -39,10 +39,9 @@ pub struct EngineCoreReadyResponse {
     /// DP coordinator stats publish address, if applicable.
     pub dp_stats_address: Option<String>,
     /// Effective model dtype after Python vLLM resolves `--dtype`.
-    // TODO: This is currently not wired up on the engine side. After it's added, remove `Option`
-    // and `serde(default)`.
-    #[serde(default)]
-    pub dtype: Option<ModelDtype>,
+    pub dtype: ModelDtype,
+    /// Python vLLM version reported by the engine process.
+    pub vllm_version: String,
 }
 
 /// Frontend-owned ZMQ addresses that are sent to the engine during startup
@@ -69,22 +68,3 @@ pub struct HandshakeInitMessage {
     pub addresses: HandshakeAddresses,
     pub parallel_config: BTreeMap<String, OpaqueValue>,
 }
-
-#[cfg(test)]
-mod tests {
-    use super::EngineCoreReadyResponse;
-    use crate::protocol::ModelDtype;
-
-    #[test]
-    fn ready_response_accepts_effective_dtype() {
-        let response: EngineCoreReadyResponse = serde_json::from_value(serde_json::json!({
-            "max_model_len": 4096,
-            "num_gpu_blocks": 2,
-            "dp_stats_address": null,
-            "dtype": "bfloat16"
-        }))
-        .unwrap();
-
-        assert_eq!(response.dtype, Some(ModelDtype::BFloat16));
-    }
-}

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

@@ -10,9 +10,9 @@ use crate::EngineId;
 pub use crate::mock_engine::{MockCoordinatorSockets, MockEngineSockets};
 use crate::mock_engine::{
     MockEngineConfig, MockEngineDataSockets, connect_to_bootstrapped_frontend, connect_to_frontend,
+    default_ready_response,
 };
-use crate::protocol::ModelDtype;
-use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage};
+use crate::protocol::handshake::HandshakeInitMessage;
 
 /// Per-test IPC endpoint namespace backed by a unique temporary directory.
 ///
@@ -57,12 +57,7 @@ fn test_mock_engine_config() -> MockEngineConfig {
     MockEngineConfig {
         local: true,
         headless: true,
-        ready_response: EngineCoreReadyResponse {
-            max_model_len: 4096,
-            num_gpu_blocks: 0,
-            dp_stats_address: None,
-            dtype: Some(ModelDtype::Float32),
-        },
+        ready_response: default_ready_response(),
         ..Default::default()
     }
 }

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

@@ -925,6 +925,7 @@ async fn client_fail_closes_when_main_output_path_receives_dp_control() {
     .await;
     assert_eq!(client.engine_identities()[0], b"engine-0");
     assert!(client.ready_responses()[0].max_model_len > 0);
+    assert_eq!(client.vllm_version(), "test-vllm-version");
 
     let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap();
     let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap();

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

@@ -104,8 +104,8 @@ pub struct ConnectedEngine {
     /// The identity of the connected engine.
     pub engine_id: EngineId,
     /// Post-initialization configuration received from the engine on the input
-    /// socket registration message. `None` until the registration is received.
-    pub ready_response: Option<EngineCoreReadyResponse>,
+    /// socket registration message.
+    pub ready_response: EngineCoreReadyResponse,
 }
 
 /// Represents the connected shared transport plus all registered engines after
@@ -295,18 +295,9 @@ pub async fn connect_handshake(
         }
     }
 
-    // 4. Wait for every engine to connect to the shared input socket and register itself. The
-    //    `ready_response` is a placeholder; it is populated for each engine by
-    //    `wait_for_input_registrations` below.
-    let mut engines: Vec<_> = engines
-        .into_keys()
-        .map(|engine_id| ConnectedEngine {
-            engine_id,
-            ready_response: None,
-        })
-        .collect();
-
-    wait_for_input_registrations(&mut input_socket, &mut engines, ready_timeout).await?;
+    // 6. Wait for every engine to connect to the shared input socket and register itself.
+    let engines =
+        wait_for_input_registrations(&mut input_socket, engines.into_keys(), ready_timeout).await?;
     debug!(
         engine_count = engines.len(),
         "all engines registered on shared input socket"
@@ -349,15 +340,13 @@ pub async fn connect_bootstrapped(
     let mut output_socket = PullSocket::new();
     let output_address = output_socket.bind(output_address).await?.to_string();
 
-    // TODO: follow start rank
-    let mut engines = (0..engine_count)
-        .map(|index| ConnectedEngine {
-            engine_id: EngineId::from((index as u16).to_le_bytes().to_vec()),
-            ready_response: None,
-        })
-        .collect::<Vec<_>>();
-
-    wait_for_input_registrations(&mut input_socket, &mut engines, ready_timeout).await?;
+    let engines = wait_for_input_registrations(
+        &mut input_socket,
+        // TODO: follow start rank
+        (0..engine_count).map(|index| EngineId::from((index as u16).to_le_bytes().to_vec())),
+        ready_timeout,
+    )
+    .await?;
     info!(
         engine_count = engines.len(),
         "bootstrapped engines connected"
@@ -455,17 +444,14 @@ async fn send_init_message(
 /// Simplify API server handshake"), the payload is a msgpack-encoded
 /// [`EngineCoreReadyResponse`] carrying post-initialization values such as
 /// `max_model_len`.
-///
-/// Older engines sent an empty second frame here just to establish the
-/// ROUTER/DEALER backchannel, with no structured payload on the input socket.
-/// We continue to tolerate that legacy shape so the frontend can still connect
-/// to slightly older local engine checkouts.
 async fn wait_for_input_registrations(
     input_socket: &mut RouterSocket,
-    engines: &mut [ConnectedEngine],
+    expected_engines: impl IntoIterator<Item = EngineId>,
     ready_timeout: Duration,
-) -> Result<()> {
-    let mut pending = engines.iter().map(|e| e.engine_id.clone()).collect::<BTreeSet<_>>();
+) -> Result<Vec<ConnectedEngine>> {
+    let expected_engines = expected_engines.into_iter().collect::<Vec<_>>();
+    let mut pending = expected_engines.iter().cloned().collect::<BTreeSet<_>>();
+    let mut ready_responses = BTreeMap::new();
 
     while !pending.is_empty() {
         let registration = timeout(ready_timeout, input_socket.recv()).await.map_err(|_| {
@@ -489,29 +475,33 @@ async fn wait_for_input_registrations(
             );
         }
 
-        let ready_response = if frames[1].is_empty() {
-            debug!(
-                ?actual_id,
-                "received legacy empty input registration from engine"
+        if frames[1].is_empty() {
+            bail_unexpected_handshake_message!(
+                "expected msgpack EngineCoreReadyResponse for engine input registration, got empty payload from engine id {actual_id:?}"
             );
-            None
-        } else {
-            let ready_response: EngineCoreReadyResponse = decode_msgpack(&frames[1])?;
-            debug!(
-                ?actual_id,
-                ?ready_response,
-                "received input registration from engine"
-            );
-            Some(ready_response)
-        };
-
-        // Store the ready response in the corresponding engine entry.
-        if let Some(engine) = engines.iter_mut().find(|e| e.engine_id == actual_id) {
-            engine.ready_response = ready_response;
         }
+
+        let ready_response: EngineCoreReadyResponse = decode_msgpack(&frames[1])?;
+        debug!(
+            ?actual_id,
+            ?ready_response,
+            "received input registration from engine"
+        );
+        ready_responses.insert(actual_id, ready_response);
     }
 
-    Ok(())
+    Ok(expected_engines
+        .into_iter()
+        .map(|engine_id| {
+            let ready_response = ready_responses
+                .remove(&engine_id)
+                .expect("every expected engine id has a decoded ready response");
+            ConnectedEngine {
+                engine_id,
+                ready_response,
+            }
+        })
+        .collect())
 }
 
 /// Send an encoded message to the engine through the input socket.

rust/src/mock-engine/src/tests.rs

@@ -98,7 +98,8 @@ async fn mock_engine_connects_over_tcp() {
     let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await;
     assert_eq!(client.engine_count(), 1);
     assert_eq!(client.engine_identities()[0], &[0, 0]);
-    assert_eq!(client.max_model_len(), Some(1024 * 1024));
+    assert_eq!(client.max_model_len(), 1024 * 1024);
+    assert_eq!(client.vllm_version(), "test-vllm-version");
     shutdown_mock(client, shutdown, task).await;
 }
 

rust/src/server/src/routes.rs

@@ -6,6 +6,7 @@ mod load;
 mod metrics;
 pub(crate) mod openai;
 mod sleep;
+mod version;
 
 use std::sync::Arc;
 
@@ -35,6 +36,7 @@ fn build_router_with_dev_mode(state: Arc<AppState>, dev_mode_enabled: bool) -> R
         .route("/health", get(health::health))
         .route("/metrics", get(metrics::scrape))
         .route("/load", get(load::load))
+        .route("/version", get(version::version))
         // OpenAI-compatible endpoints
         .route("/v1/models", get(openai::list_models))
         .route("/v1/completions", post(openai::completions))

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

@@ -994,6 +994,27 @@ async fn list_models_returns_configured_model() {
     assert_eq!(json["data"][0]["id"], "Qwen/Qwen1.5-0.5B-Chat");
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn version_returns_engine_vllm_version() {
+    let mut app = test_app().await;
+    let response = app
+        .call(Request::builder().uri("/version").body(Body::empty()).expect("build request"))
+        .await
+        .expect("call app");
+
+    assert_eq!(response.status(), StatusCode::OK);
+    let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
+    let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
+    assert_eq!(
+        json,
+        json!({
+            "version": "test-vllm-version",
+            "rust_frontend_version": env!("CARGO_PKG_VERSION"),
+        })
+    );
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn http_metrics_record_list_models_requests() {

rust/src/server/src/routes/version.rs

@@ -0,0 +1,23 @@
+use std::sync::Arc;
+
+use axum::Json;
+use axum::extract::State;
+use serde::Serialize;
+
+use crate::state::AppState;
+
+#[derive(Serialize)]
+pub(crate) struct VersionResponse {
+    version: String,
+    rust_frontend_version: &'static str,
+}
+
+/// Get engine and Rust frontend version metadata.
+pub async fn version(State(state): State<Arc<AppState>>) -> Json<VersionResponse> {
+    let version = state.engine_core_client().vllm_version().to_string();
+
+    Json(VersionResponse {
+        version,
+        rust_frontend_version: env!("CARGO_PKG_VERSION"),
+    })
+}

rust/src/text/src/lib.rs

@@ -45,9 +45,9 @@ pub struct TextLlm {
     /// Tokenizer/model metadata backend responsible for prompt encode/decode
     /// and sampling hints.
     backend: DynTextBackend,
-    /// Context window size derived by the backend or from engine startup
-    /// handshake, with optional override from config.
-    max_model_len: Option<u32>,
+    /// Context window size reported by the engine startup handshake, with
+    /// optional override from config.
+    max_model_len: u32,
 }
 
 impl TextLlm {
@@ -71,7 +71,7 @@ impl TextLlm {
     /// This takes priority over both the engine-reported default and any
     /// tokenizer/model metadata exposed by the backend.
     pub fn with_max_model_len(mut self, max_model_len: u32) -> Self {
-        self.max_model_len = Some(max_model_len);
+        self.max_model_len = max_model_len;
         self
     }
 
@@ -129,9 +129,7 @@ impl TextLlm {
         };
 
         let mut sampling_hints = self.backend.sampling_hints()?;
-        if let Some(max_model_len) = self.max_model_len {
-            sampling_hints.max_model_len = Some(max_model_len);
-        }
+        sampling_hints.max_model_len = Some(self.max_model_len);
         let PreparedTextRequest {
             text_request,
             generate_request,

setup.py

@@ -1195,11 +1195,6 @@ setup(
             "opentelemetry-exporter-otlp>=1.26.0",
             "opentelemetry-semantic-conventions-ai>=0.4.1",
         ],
-        "triton-cpu": [
-            "triton @ "
-            "git+https://github.com/triton-lang/triton-cpu.git@270e696d ; "
-            "platform_machine == 'x86_64'",
-        ],  # Remove after stable release
     },
     cmdclass=cmdclass,
     package_data=package_data,

tests/compile/fusions_e2e/common.py

@@ -20,6 +20,7 @@ class Matches(NamedTuple):
     attn_quant_fusion: int = 0
     # distributed
     ar_rms_fusion: int = 0
+    aiter_ar_rms_fusion: int = 0
     sequence_parallel: int = 0
     async_tp: int = 0
 
@@ -97,6 +98,9 @@ FUSION_LOG_PATTERNS: dict[str, re.Pattern] = {
     "ar_rms_fusion": re.compile(
         r"allreduce_rms_fusion.py:\d+] Replaced (\d+) patterns"
     ),
+    "aiter_ar_rms_fusion": re.compile(
+        r"RocmAiterAllReduceFusionPass Replaced (\d+) patterns"
+    ),
     "sequence_parallel": re.compile(
         r"sequence_parallelism.py:\d+] Replaced (\d+) patterns"
     ),

tests/compile/fusions_e2e/conftest.py

@@ -97,6 +97,11 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
                 f"attention backend '{attn_backend.backend.name}'"
             )
 
+        if backend_name == "rocm_attn" and model_name == "openai/gpt-oss-20b":
+            pytest.skip(
+                "ROCM_ATTN does not support attention sinks (required by gpt-oss-20b)"
+            )
+
         if attn_backend.backend.name == "FLASHINFER":
             from vllm.utils.flashinfer import supports_trtllm_attention
 

tests/compile/fusions_e2e/models.py

@@ -75,6 +75,7 @@ llama3_8b = ModelFusionInfo(
     model_name="meta-llama/Llama-3.1-8B-Instruct",
     matches=lambda n_layers: Matches(
         ar_rms_fusion=n_layers * 2 + 1,
+        aiter_ar_rms_fusion=n_layers * 2,
         sequence_parallel=n_layers * 2 + 1,
         async_tp=n_layers * 4,
     ),
@@ -136,6 +137,7 @@ qwen3_a3b = ModelFusionInfo(
     matches=lambda n_layers: Matches(
         norm_rope_fusion=n_layers,
         ar_rms_fusion=n_layers * 2 + 1,
+        aiter_ar_rms_fusion=n_layers * 2,
         sequence_parallel=n_layers * 2 + 1,
         async_tp=n_layers * 2,
     ),
@@ -211,6 +213,7 @@ gpt_oss_20b = ModelFusionInfo(
     model_name="openai/gpt-oss-20b",
     matches=lambda n_layers: Matches(
         ar_rms_fusion=n_layers * 2 + 1,
+        aiter_ar_rms_fusion=n_layers + 1,
         sequence_parallel=n_layers * 2 + 1,
         async_tp=n_layers * 2,
     ),

tests/compile/fusions_e2e/test_tp2_ar_rms.py

@@ -84,6 +84,7 @@ def test_tp2_ar_rms_fp8_fusions(
     model_kwargs["load_format"] = "dummy"
     model_kwargs["max_model_len"] = 1024
     model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
+    model_kwargs["disable_custom_all_reduce"] = False
 
     compilation_config = dict(
         use_inductor_graph_partition=inductor_graph_partition,
@@ -149,6 +150,7 @@ def test_tp2_ar_rms_fp4_fusions(
     model_kwargs["load_format"] = "dummy"
     model_kwargs["max_model_len"] = 1024
     model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
+    model_kwargs["disable_custom_all_reduce"] = False
 
     compilation_config = dict(
         use_inductor_graph_partition=inductor_graph_partition,
@@ -213,6 +215,7 @@ def test_tp2_ar_rms_fusions(
     model_kwargs["load_format"] = "dummy"
     model_kwargs["max_model_len"] = 1024
     model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
+    model_kwargs["disable_custom_all_reduce"] = False
 
     compilation_config = dict(
         use_inductor_graph_partition=inductor_graph_partition,
@@ -225,9 +228,13 @@ def test_tp2_ar_rms_fusions(
 
     matches_check = [
         "norm_rope_fusion",
-        "ar_rms_fusion",
     ]
 
+    if current_platform.is_rocm():
+        matches_check.append("aiter_ar_rms_fusion")
+    else:
+        matches_check.append("ar_rms_fusion")
+
     run_e2e_fusion_test(
         model_name,
         matches,

tests/compile/passes/distributed/test_async_tp.py

@@ -29,6 +29,7 @@ from vllm.distributed.parallel_state import (
     initialize_model_parallel,
 )
 from vllm.platforms import current_platform
+from vllm.utils.network_utils import get_open_port
 from vllm.utils.system_utils import update_environment_variables
 from vllm.utils.torch_utils import set_random_seed
 
@@ -234,8 +235,20 @@ class TestAGCutlassScaledMMModel(_BaseScaledMMModel):
         TestAGMMModel,
         TestScaledMMRSModel,
         TestAGScaledMMModel,
-        TestCutlassScaledMMRSModel,
-        TestAGCutlassScaledMMModel,
+        pytest.param(
+            TestCutlassScaledMMRSModel,
+            marks=pytest.mark.skipif(
+                not hasattr(torch.ops._C, "cutlass_scaled_mm"),
+                reason="Requires cutlass_scaled_mm",
+            ),
+        ),
+        pytest.param(
+            TestAGCutlassScaledMMModel,
+            marks=pytest.mark.skipif(
+                not hasattr(torch.ops._C, "cutlass_scaled_mm"),
+                reason="Requires cutlass_scaled_mm",
+            ),
+        ),
     ],
 )
 @pytest.mark.parametrize("batch_size", [8])
@@ -268,6 +281,7 @@ def test_async_tp_pass_replace(
         )
 
     num_processes = 2
+    master_port = str(get_open_port())
 
     def run_torch_spawn(fn, nprocs):
         # need to use torch.mp.spawn otherwise will have problems with
@@ -282,6 +296,7 @@ def test_async_tp_pass_replace(
                 hidden_size,
                 dtype,
                 dynamic,
+                master_port,
             ),
             nprocs=nprocs,
         )
@@ -314,6 +329,7 @@ def async_tp_pass_on_test_model(
     hidden_size: int,
     dtype: torch.dtype,
     dynamic: bool,
+    master_port: str = "0",
 ):
     set_random_seed(0)
 
@@ -328,7 +344,7 @@ def async_tp_pass_on_test_model(
             "LOCAL_RANK": str(local_rank),
             "WORLD_SIZE": str(world_size),
             "MASTER_ADDR": "localhost",
-            "MASTER_PORT": "12345",
+            "MASTER_PORT": master_port,
         }
     )
 

tests/distributed/test_comm_ops.py

@@ -292,7 +292,7 @@ def test_async_intermediate_tensors_lazy_wait() -> None:
     )
 
     # accessing non-tensor attributes should not trigger wait.
-    assert it.kv_connector_output is None
+    assert it._comm_handles is not None
     assert work.wait_calls == 0
     assert post_calls["n"] == 0
 

tests/distributed/test_eplb_spec_decode.py

@@ -15,7 +15,7 @@ def get_model_args(
     spec_method: str,
     tp_size: int,
     model_max_len: int,
-    use_async: bool = False,
+    use_async: bool = True,
 ) -> dict:
     speculative_config = {
         "method": spec_method,
@@ -28,9 +28,8 @@ def get_model_args(
         "window_size": 128,
         "step_interval": 1024,
         "log_balancedness": False,
+        "use_async": use_async,
     }
-    if use_async:
-        eplb_config["use_async"] = True
     model_args = {
         "pretrained": model_name,
         "dtype": "auto",

tests/kernels/quantization/test_rdna3_w4a16.py

@@ -0,0 +1,278 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Correctness tests for the ROCm RDNA3 W4A16 GPTQ kernel (gfx1100).
+
+Exercises ``RDNA3W4A16LinearKernel`` end-to-end: it builds a layer with
+GPTQ-format checkpoint parameters, runs ``process_weights_after_loading``
+(weight shuffle + zero-point synthesis), then ``apply_weights``, and compares
+the result against an fp32 reference dequant-and-matmul.
+
+The kernel is exposed via ``torch.ops._rocm_C.gptq_gemm_rdna3`` and is only
+built for gfx11; tests are skipped elsewhere.
+
+Run `pytest tests/kernels/quantization/test_rdna3_w4a16.py`.
+"""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+if not current_platform.is_rocm():
+    pytest.skip("RDNA3 W4A16 kernel is ROCm-only", allow_module_level=True)
+
+from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (  # noqa: E402
+    MPLinearLayerConfig,
+)
+from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import (  # noqa: E402
+    RDNA3W4A16LinearKernel,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (  # noqa: E402
+    pack_quantized_values_into_int32,
+)
+from vllm.model_executor.parameter import (  # noqa: E402
+    GroupQuantScaleParameter,
+    PackedvLLMParameter,
+)
+from vllm.platforms.rocm import on_gfx1100  # noqa: E402
+from vllm.scalar_type import scalar_types  # noqa: E402
+from vllm.utils.torch_utils import set_random_seed  # noqa: E402
+
+device = "cuda"
+
+WEIGHT_TYPE = scalar_types.uint4b8  # symmetric int4, bias = 8
+PACK_FACTOR = 8  # 8 x 4-bit nibbles per int32
+
+# Skip everything in this module unless we are on the only architecture the
+# kernel is built/registered for.
+gfx1100_only = pytest.mark.skipif(
+    not (
+        on_gfx1100()
+        and hasattr(torch.ops, "_rocm_C")
+        and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3")
+    ),
+    reason="requires gfx1100 with the _rocm_C.gptq_gemm_rdna3 op built in",
+)
+
+
+# ---------------------------------------------------------------------------
+# Reference implementation
+# ---------------------------------------------------------------------------
+
+
+def _reference(
+    x_mk: torch.Tensor,
+    q_int4_kn: torch.Tensor,
+    scales_gn: torch.Tensor,
+    zeros_gn: torch.Tensor | None,
+    group_size: int,
+    bias: torch.Tensor | None,
+) -> torch.Tensor:
+    """fp32 reference for the RDNA3 W4A16 op.
+
+    x_mk:       [M, K] fp16/bf16 activations.
+    q_int4_kn:  [K, N] int32 raw stored nibbles in [0, 15].
+    scales_gn:  [K//G, N] per-group scales (act dtype).
+    zeros_gn:   [K//G, N] int32 raw stored zero points in [0, 15], or None
+                for the symmetric path (kernel synthesizes stored zero = 7).
+    group_size: G.
+
+    The kernel applies the GPTQv1 "+1" zero-point quirk, so the effective
+    zero is ``stored_zero + 1`` (symmetric path: 7 + 1 == bias == 8).
+    """
+    K, N = q_int4_kn.shape
+    s_full = scales_gn.repeat_interleave(group_size, dim=0).to(torch.float32)  # [K,N]
+    if zeros_gn is None:
+        z_full = torch.full(
+            (K, N), float(WEIGHT_TYPE.bias), device=x_mk.device, dtype=torch.float32
+        )
+    else:
+        z_full = (zeros_gn + 1).repeat_interleave(group_size, dim=0).to(torch.float32)
+    w_fp = (q_int4_kn.to(torch.float32) - z_full) * s_full  # [K, N]
+    out = x_mk.to(torch.float32) @ w_fp  # [M, N]
+    if bias is not None:
+        out = out + bias.to(torch.float32)
+    return out.to(x_mk.dtype)
+
+
+# ---------------------------------------------------------------------------
+# Layer construction (GPTQ checkpoint format)
+# ---------------------------------------------------------------------------
+
+
+def _build_layer(
+    q_int4_kn: torch.Tensor,
+    scales_gn: torch.Tensor,
+    zeros_gn: torch.Tensor | None,
+    dtype: torch.dtype,
+) -> torch.nn.Module:
+    """Build a dummy layer carrying GPTQ-format params, as the loader would."""
+    no_loader = lambda *args, **kwargs: None  # noqa: E731
+
+    # qweight: int4 packed along K into int32 -> [K//8, N].
+    qweight = pack_quantized_values_into_int32(q_int4_kn, WEIGHT_TYPE, packed_dim=0)
+
+    class DummyLayer(torch.nn.Module):
+        pass
+
+    layer = DummyLayer()
+    layer.register_parameter(
+        "qweight",
+        PackedvLLMParameter(
+            data=qweight,
+            weight_loader=no_loader,
+            input_dim=0,
+            output_dim=1,
+            packed_dim=0,
+            packed_factor=PACK_FACTOR,
+        ),
+    )
+    layer.register_parameter(
+        "scales",
+        GroupQuantScaleParameter(
+            data=scales_gn.to(dtype),
+            weight_loader=no_loader,
+            input_dim=0,
+            output_dim=1,
+        ),
+    )
+    if zeros_gn is not None:
+        # qzeros: int4 packed along N into int32 -> [K//G, N//8].
+        qzeros = pack_quantized_values_into_int32(zeros_gn, WEIGHT_TYPE, packed_dim=1)
+        layer.register_parameter(
+            "qzeros",
+            PackedvLLMParameter(
+                data=qzeros,
+                weight_loader=no_loader,
+                input_dim=0,
+                output_dim=1,
+                packed_dim=1,
+                packed_factor=PACK_FACTOR,
+            ),
+        )
+    return layer
+
+
+def _run_kernel(
+    x_mk: torch.Tensor,
+    q_int4_kn: torch.Tensor,
+    scales_gn: torch.Tensor,
+    zeros_gn: torch.Tensor | None,
+    group_size: int,
+    bias: torch.Tensor | None,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    K, N = q_int4_kn.shape
+    has_zp = zeros_gn is not None
+
+    config = MPLinearLayerConfig(
+        full_weight_shape=(K, N),
+        partition_weight_shape=(K, N),
+        weight_type=WEIGHT_TYPE,
+        act_type=dtype,
+        group_size=group_size,
+        zero_points=has_zp,
+        has_g_idx=False,
+    )
+    ok, reason = RDNA3W4A16LinearKernel.can_implement(config)
+    assert ok, f"can_implement rejected a supported config: {reason}"
+
+    layer = _build_layer(q_int4_kn, scales_gn, zeros_gn, dtype)
+    kernel = RDNA3W4A16LinearKernel(
+        config,
+        w_q_param_name="qweight",
+        w_s_param_name="scales",
+        w_zp_param_name="qzeros" if has_zp else None,
+        w_gidx_param_name=None,
+    )
+    kernel.process_weights_after_loading(layer)
+    return kernel.apply_weights(layer, x_mk, bias=bias)
+
+
+# Relative-L2 tolerance per dtype. The bf16 path widens dequantized weights
+# to fp32 and accumulates in fp32, so it matches the reference almost exactly
+# (<0.4% incl. the WMMA prefill path). The fp16 path uses the exllamav2
+# "+1024" bit-trick (see qdq_4_rdna3.cuh): the dequantized weight is recovered
+# as the fp16 difference of two ~1024*scale magnitudes, which sheds low-order
+# mantissa bits and leaves ~2-3% relative noise that accumulates over K. We
+# compare on the relative Frobenius norm rather than elementwise, since the
+# bit-trick noise produces large *relative* errors on individual near-zero
+# outputs that carry negligible absolute weight.
+_REL_L2_TOL = {torch.float16: 5e-2, torch.bfloat16: 1e-2}
+
+
+def _assert_close(out: torch.Tensor, ref: torch.Tensor, dtype: torch.dtype):
+    rel_l2 = (out.to(torch.float32) - ref.to(torch.float32)).norm() / ref.to(
+        torch.float32
+    ).norm()
+    tol = _REL_L2_TOL[dtype]
+    assert rel_l2 < tol, f"relative L2 error {rel_l2:.4f} exceeds {tol} for {dtype}"
+
+
+# ---------------------------------------------------------------------------
+# Forward correctness
+# ---------------------------------------------------------------------------
+
+
+# (M, K, N, group_size). M spans the scalar decode path (small M) and the
+# WMMA prefill path (M >= 16 on the bf16 dispatch). K/N satisfy the kernel's
+# divisibility constraints (K % G == 0, K % 8 == 0, N % 8 == 0).
+MKNG_SHAPES = [
+    (1, 128, 128, 128),  # single group, decode
+    (2, 256, 256, 128),  # two groups
+    (8, 256, 512, 64),  # M=8 scalar, smaller group
+    (16, 512, 256, 128),  # M=16 -> WMMA path for bf16
+    (32, 512, 512, 64),  # larger prefill
+]
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("has_zp", [False, True], ids=["no_zp", "with_zp"])
+@pytest.mark.parametrize(
+    "M,K,N,G", MKNG_SHAPES, ids=[f"m{m}_k{k}_n{n}_g{g}" for m, k, n, g in MKNG_SHAPES]
+)
+def test_rdna3_w4a16_matches_reference(dtype, has_zp, M, K, N, G, dist_init):
+    set_random_seed(0)
+    assert K % G == 0 and K % PACK_FACTOR == 0 and N % PACK_FACTOR == 0
+
+    groups = K // G
+    x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
+    q_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
+    scales_gn = (
+        0.05 * torch.rand((groups, N), device=device, dtype=torch.float32) + 0.01
+    ).to(dtype)
+    zeros_gn = (
+        torch.randint(0, 16, (groups, N), device=device, dtype=torch.int32)
+        if has_zp
+        else None
+    )
+
+    out = _run_kernel(x_mk, q_int4_kn, scales_gn, zeros_gn, G, None, dtype)
+    ref = _reference(x_mk, q_int4_kn, scales_gn, zeros_gn, G, None)
+
+    assert out.shape == (M, N) and out.dtype == dtype
+    _assert_close(out, ref, dtype)
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("M", [1, 32], ids=["decode", "prefill"])
+def test_rdna3_w4a16_bias(dtype, M, dist_init):
+    """Bias is added on both the scalar (M=1) and WMMA (M=32) paths."""
+    set_random_seed(0)
+    K, N, G = 512, 256, 128
+    groups = K // G
+
+    x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
+    q_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
+    scales_gn = (
+        0.05 * torch.rand((groups, N), device=device, dtype=torch.float32) + 0.01
+    ).to(dtype)
+    bias = (0.1 * torch.randn(N, device=device, dtype=torch.float32)).to(dtype)
+
+    out = _run_kernel(x_mk, q_int4_kn, scales_gn, None, G, bias, dtype)
+    ref = _reference(x_mk, q_int4_kn, scales_gn, None, G, bias)
+
+    _assert_close(out, ref, dtype)

tests/kernels/quantization/test_rdna3_w4a16_selection.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Kernel-selection / gating tests for the ROCm RDNA3 W4A16 GPTQ kernel.
+
+Verifies that ``choose_mp_linear_kernel`` resolves a supported W4A16 GPTQ
+config to ``RDNA3W4A16LinearKernel`` on gfx1100 (it is registered ahead of
+``TritonW4A16LinearKernel`` in the ROCm priority list), and that
+``RDNA3W4A16LinearKernel.can_implement`` rejects the configs it does not
+support so selection falls through to the next kernel.
+
+Run `pytest tests/kernels/quantization/test_rdna3_w4a16_selection.py`.
+"""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+if not current_platform.is_rocm():
+    pytest.skip("RDNA3 W4A16 kernel is ROCm-only", allow_module_level=True)
+
+from vllm.model_executor.kernels.linear import (  # noqa: E402
+    choose_mp_linear_kernel,
+)
+from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (  # noqa: E402
+    MPLinearLayerConfig,
+)
+from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import (  # noqa: E402
+    RDNA3W4A16LinearKernel,
+)
+from vllm.platforms.rocm import on_gfx1100  # noqa: E402
+from vllm.scalar_type import scalar_types  # noqa: E402
+
+WEIGHT_TYPE = scalar_types.uint4b8  # symmetric int4, bias = 8
+
+# The kernel is only selectable when running on gfx1100 with the custom op
+# compiled in; otherwise can_implement rejects and selection falls through.
+gfx1100_only = pytest.mark.skipif(
+    not (
+        on_gfx1100()
+        and hasattr(torch.ops, "_rocm_C")
+        and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3")
+    ),
+    reason="requires gfx1100 with the _rocm_C.gptq_gemm_rdna3 op built in",
+)
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_selection_prefers_rdna3(dtype):
+    """A supported W4A16 GPTQ config resolves to the RDNA3 kernel on gfx1100."""
+    config = MPLinearLayerConfig(
+        full_weight_shape=(1024, 256),
+        partition_weight_shape=(1024, 256),
+        weight_type=WEIGHT_TYPE,
+        act_type=dtype,
+        group_size=128,
+        zero_points=False,
+        has_g_idx=False,
+    )
+    assert choose_mp_linear_kernel(config).__name__ == "RDNA3W4A16LinearKernel"
+
+
+@gfx1100_only
+@pytest.mark.parametrize(
+    "weight_type,group_size,N,full_k,expected_ok",
+    [
+        (scalar_types.uint4b8, 128, 256, 1024, True),  # nominal: supported
+        (scalar_types.uint4b8, -1, 256, 1024, False),  # channelwise unsupported
+        (scalar_types.uint4b8, 128, 252, 1024, False),  # N not a multiple of 8
+        (scalar_types.uint4b8, 96, 256, 1024, False),  # group does not divide K
+        (scalar_types.uint8b128, 128, 256, 1024, False),  # wrong quant type
+    ],
+    ids=["ok", "channelwise", "bad_n", "group_ndiv_k", "wrong_qtype"],
+)
+def test_can_implement(weight_type, group_size, N, full_k, expected_ok):
+    """can_implement gates on quant type, group size, and N divisibility."""
+    config = MPLinearLayerConfig(
+        full_weight_shape=(full_k, N),
+        partition_weight_shape=(full_k, N),
+        weight_type=weight_type,
+        act_type=torch.float16,
+        group_size=group_size,
+        zero_points=False,
+        has_g_idx=False,
+    )
+    ok, reason = RDNA3W4A16LinearKernel.can_implement(config)
+    assert ok is expected_ok, reason

tests/kernels/test_mhc_kernels.py

@@ -340,22 +340,17 @@ def test_hc_head_tilelang(num_tokens, hidden_size, hc_mult):
     hc_base = torch.randn((hc_mult,), dtype=torch.float32) * 0.1
     rms_eps = hc_eps = 1e-6
 
-    out = torch.empty((num_tokens, hidden_size), dtype=torch.bfloat16)
-    out.fill_(float("nan"))
-
-    result = torch.ops.vllm.hc_head_fused_kernel_tilelang(
+    out = torch.ops.vllm.hc_head_fused_kernel_tilelang(
         residual,
         fn,
         hc_scale,
         hc_base,
-        out,
-        hidden_size,
         rms_eps,
         hc_eps,
-        hc_mult,
     )
 
-    assert result is None
+    assert out.shape == (num_tokens, hidden_size)
+    assert out.dtype == torch.bfloat16
     assert not torch.isnan(out).any()
 
     out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps)

tests/lora/test_utils.py

@@ -175,7 +175,7 @@ def test_get_adapter_absolute_path_local_existing(mock_abspath, mock_exist):
     assert get_adapter_absolute_path(path) == absolute_path
 
 
-@patch("huggingface_hub.snapshot_download")
+@patch("huggingface_hub.HfApi.snapshot_download")
 @patch("os.path.exists")
 def test_get_adapter_absolute_path_huggingface(mock_exist, mock_snapshot_download):
     # Hugging Face model identifier
@@ -186,7 +186,7 @@ def test_get_adapter_absolute_path_huggingface(mock_exist, mock_snapshot_downloa
     assert get_adapter_absolute_path(path) == absolute_path
 
 
-@patch("huggingface_hub.snapshot_download")
+@patch("huggingface_hub.HfApi.snapshot_download")
 @patch("os.path.exists")
 def test_get_adapter_absolute_path_huggingface_error(
     mock_exist, mock_snapshot_download

tests/model_executor/layers/test_pooler_methods.py

@@ -0,0 +1,499 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for sequence and token pooling methods and their factories."""
+
+from dataclasses import dataclass
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.pooler.seqwise.methods import (
+    CLSPool,
+    LastPool,
+    MeanPool,
+    get_seq_pooling_method,
+)
+from vllm.model_executor.layers.pooler.tokwise.methods import (
+    AllPool,
+    StepPool,
+    get_tok_pooling_method,
+)
+from vllm.pooling_params import PoolingParams
+from vllm.v1.pool.metadata import PoolingCursor, PoolingMetadata, PoolingStates
+
+_CPU = torch.device("cpu")
+
+
+def _make_pooling_cursor(
+    prompt_lens: list[int],
+    *,
+    num_scheduled_tokens: list[int] | None = None,
+    seq_lens: list[int] | None = None,
+    device: torch.device = _CPU,
+) -> PoolingCursor:
+    """Build a PoolingCursor from a list of per-sequence prompt lengths."""
+    prompt_lens_cpu = torch.tensor(prompt_lens, dtype=torch.long)
+    if num_scheduled_tokens is None:
+        num_scheduled_tokens_cpu = prompt_lens_cpu.clone()
+    else:
+        num_scheduled_tokens_cpu = torch.tensor(num_scheduled_tokens, dtype=torch.long)
+    if seq_lens is None:
+        seq_lens_cpu = prompt_lens_cpu.clone()
+    else:
+        seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.long)
+
+    cumsum = torch.zeros(len(prompt_lens) + 1, dtype=torch.long, device=device)
+    torch.cumsum(num_scheduled_tokens_cpu, dim=0, out=cumsum[1:])
+
+    return PoolingCursor(
+        first_token_indices_gpu=cumsum[: len(prompt_lens)].to(device),
+        last_token_indices_gpu=(cumsum[1:] - 1).to(device),
+        prompt_lens_cpu=prompt_lens_cpu,
+        seq_lens_cpu=seq_lens_cpu,
+        num_scheduled_tokens_cpu=num_scheduled_tokens_cpu,
+    )
+
+
+def _make_metadata(
+    prompt_lens: list[int],
+    *,
+    tasks: list[str] | None = None,
+    token_ids: list[list[int]] | None = None,
+    pooling_params: list[PoolingParams] | None = None,
+    num_scheduled_tokens: list[int] | None = None,
+    seq_lens: list[int] | None = None,
+    device: torch.device = _CPU,
+) -> PoolingMetadata:
+    """Build a minimal PoolingMetadata for testing pooling methods."""
+    n_seqs = len(prompt_lens)
+    if tasks is None:
+        tasks = ["embed"] * n_seqs
+    if pooling_params is None:
+        pooling_params = [PoolingParams(task=t) for t in tasks]
+
+    prompt_lens_tensor = torch.tensor(prompt_lens, dtype=torch.long)
+
+    prompt_token_ids_cpu = None
+    prompt_token_ids = None
+    if token_ids is not None:
+        max_len = max(len(t) for t in token_ids)
+        padded = [t + [0] * (max_len - len(t)) for t in token_ids]
+        prompt_token_ids_cpu = torch.tensor(padded, dtype=torch.long)
+        prompt_token_ids = prompt_token_ids_cpu.to(device)
+
+    cursor = _make_pooling_cursor(
+        prompt_lens,
+        num_scheduled_tokens=num_scheduled_tokens,
+        seq_lens=seq_lens,
+        device=device,
+    )
+
+    pooling_states = [PoolingStates() for _ in range(n_seqs)]
+
+    return PoolingMetadata(
+        prompt_lens=prompt_lens_tensor,
+        prompt_token_ids=prompt_token_ids,
+        prompt_token_ids_cpu=prompt_token_ids_cpu,
+        pooling_params=pooling_params,
+        pooling_states=pooling_states,
+        pooling_cursor=cursor,
+    )
+
+
+# ---------------------------------------------------------------------------
+# CLSPool
+# ---------------------------------------------------------------------------
+class TestCLSPool:
+    def test_extracts_first_token(self):
+        hidden = torch.tensor(
+            [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]
+        )
+        metadata = _make_metadata([2, 3])
+        pooler = CLSPool()
+        out = pooler(hidden, metadata)
+        expected = torch.tensor([[1.0, 2.0], [5.0, 6.0]])
+        assert torch.equal(out, expected)
+
+    def test_rejects_partial_prefill(self):
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
+        metadata = _make_metadata([3], num_scheduled_tokens=[2])
+        pooler = CLSPool()
+        with pytest.raises(AssertionError, match="partial prefill"):
+            pooler(hidden, metadata)
+
+
+# ---------------------------------------------------------------------------
+# LastPool
+# ---------------------------------------------------------------------------
+class TestLastPool:
+    def test_extracts_last_token(self):
+        hidden = torch.tensor(
+            [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]
+        )
+        metadata = _make_metadata([2, 3])
+        pooler = LastPool()
+        out = pooler(hidden, metadata)
+        expected = torch.tensor([[3.0, 4.0], [9.0, 10.0]])
+        assert torch.equal(out, expected)
+
+    def test_partial_prefill_extracts_last_scheduled(self):
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+        metadata = _make_metadata([4], num_scheduled_tokens=[2])
+        pooler = LastPool()
+        out = pooler(hidden, metadata)
+        expected = torch.tensor([[3.0, 4.0]])
+        assert torch.equal(out, expected)
+
+
+# ---------------------------------------------------------------------------
+# MeanPool
+# ---------------------------------------------------------------------------
+class TestMeanPool:
+    def test_computes_mean(self):
+        hidden = torch.tensor(
+            [[1.0, 2.0], [3.0, 4.0], [10.0, 20.0]], dtype=torch.float32
+        )
+        metadata = _make_metadata([2, 1])
+        pooler = MeanPool()
+        out = pooler(hidden, metadata)
+        expected = torch.tensor([[2.0, 3.0], [10.0, 20.0]], dtype=torch.float32)
+        assert torch.allclose(out, expected, atol=1e-5)
+
+    def test_single_token_is_identity(self):
+        hidden = torch.tensor([[5.0, 10.0]], dtype=torch.float32)
+        metadata = _make_metadata([1])
+        pooler = MeanPool()
+        out = pooler(hidden, metadata)
+        assert torch.allclose(out, hidden, atol=1e-5)
+
+    def test_uniform_values_return_same(self):
+        hidden = torch.full((4, 3), 7.0, dtype=torch.float32)
+        metadata = _make_metadata([4])
+        pooler = MeanPool()
+        out = pooler(hidden, metadata)
+        expected = torch.full((1, 3), 7.0, dtype=torch.float32)
+        assert torch.allclose(out, expected, atol=1e-5)
+
+    def test_multiple_sequences(self):
+        hidden = torch.tensor(
+            [
+                [0.0, 0.0],
+                [2.0, 4.0],
+                [4.0, 8.0],
+                [10.0, 10.0],
+            ],
+            dtype=torch.float32,
+        )
+        metadata = _make_metadata([3, 1])
+        pooler = MeanPool()
+        out = pooler(hidden, metadata)
+        expected = torch.tensor([[2.0, 4.0], [10.0, 10.0]], dtype=torch.float32)
+        assert torch.allclose(out, expected, atol=1e-5)
+
+    def test_empty_batch(self):
+        hidden = torch.empty((0, 8), dtype=torch.float32)
+        metadata = _make_metadata([])
+        pooler = MeanPool()
+        out = pooler(hidden, metadata)
+        assert out.shape == (0, 8)
+
+    def test_rejects_partial_prefill(self):
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32)
+        metadata = _make_metadata([3], num_scheduled_tokens=[2])
+        pooler = MeanPool()
+        with pytest.raises(AssertionError, match="partial prefill"):
+            pooler(hidden, metadata)
+
+    def test_chunked_accumulation(self):
+        hidden = torch.arange(20, dtype=torch.float32).reshape(5, 4)
+        metadata = _make_metadata([3, 2])
+        pooler = MeanPool()
+        with patch(
+            "vllm.model_executor.layers.pooler.seqwise.methods"
+            "._MEAN_POOL_ACCUMULATION_CHUNK_BYTES",
+            16,
+        ):
+            out = pooler(hidden, metadata)
+        expected_seq0 = hidden[:3].float().mean(dim=0, keepdim=True)
+        expected_seq1 = hidden[3:].float().mean(dim=0, keepdim=True)
+        expected = torch.cat([expected_seq0, expected_seq1], dim=0)
+        assert torch.allclose(out, expected, atol=1e-5)
+
+    def test_upcasts_to_float32(self):
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float16)
+        metadata = _make_metadata([2])
+        pooler = MeanPool()
+        out = pooler(hidden, metadata)
+        assert out.dtype == torch.float32
+        expected = torch.tensor([[2.0, 3.0]], dtype=torch.float32)
+        assert torch.allclose(out, expected, atol=1e-2)
+
+
+# ---------------------------------------------------------------------------
+# get_seq_pooling_method factory
+# ---------------------------------------------------------------------------
+class TestGetSeqPoolingMethod:
+    def test_cls(self):
+        assert isinstance(get_seq_pooling_method("CLS"), CLSPool)
+
+    def test_last(self):
+        assert isinstance(get_seq_pooling_method("LAST"), LastPool)
+
+    def test_mean(self):
+        assert isinstance(get_seq_pooling_method("MEAN"), MeanPool)
+
+    def test_unknown_raises(self):
+        with pytest.raises(NotImplementedError, match="UNKNOWN"):
+            get_seq_pooling_method("UNKNOWN")
+
+
+# ---------------------------------------------------------------------------
+# AllPool
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class _FakeSchedulerConfig:
+    enable_chunked_prefill: bool = False
+
+
+@dataclass
+class _FakeVllmConfig:
+    scheduler_config: _FakeSchedulerConfig
+
+
+class TestAllPool:
+    @staticmethod
+    def _make_all_pool(*, chunked: bool = False) -> AllPool:
+        fake_config = _FakeVllmConfig(
+            scheduler_config=_FakeSchedulerConfig(
+                enable_chunked_prefill=chunked,
+            ),
+        )
+        with patch(
+            "vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config",
+            return_value=fake_config,
+        ):
+            return AllPool()
+
+    def test_splits_by_sequence(self):
+        pooler = self._make_all_pool()
+        hidden = torch.tensor(
+            [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]
+        )
+        metadata = _make_metadata([2, 3])
+        out = pooler(hidden, metadata)
+        assert len(out) == 2
+        assert torch.equal(out[0], hidden[:2])
+        assert torch.equal(out[1], hidden[2:])
+
+    def test_single_sequence(self):
+        pooler = self._make_all_pool()
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
+        metadata = _make_metadata([3])
+        out = pooler(hidden, metadata)
+        assert len(out) == 1
+        assert torch.equal(out[0], hidden)
+
+    def test_chunked_prefill_returns_none_for_unfinished(self):
+        pooler = self._make_all_pool(chunked=True)
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+        metadata = _make_metadata(
+            [4],
+            num_scheduled_tokens=[2],
+            seq_lens=[2],
+        )
+        out = pooler(hidden, metadata)
+        assert len(out) == 1
+        assert out[0] is None
+
+    def test_chunked_prefill_returns_concat_when_finished(self):
+        pooler = self._make_all_pool(chunked=True)
+
+        chunk1 = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+        metadata1 = _make_metadata(
+            [4],
+            num_scheduled_tokens=[2],
+            seq_lens=[2],
+        )
+        out1 = pooler(chunk1, metadata1)
+        assert out1[0] is None
+
+        chunk2 = torch.tensor([[5.0, 6.0], [7.0, 8.0]])
+        metadata2 = _make_metadata(
+            [4],
+            num_scheduled_tokens=[2],
+            seq_lens=[4],
+        )
+        metadata2.pooling_states = metadata1.pooling_states
+        out2 = pooler(chunk2, metadata2)
+        assert out2[0] is not None
+        expected = torch.cat([chunk1, chunk2], dim=0)
+        assert torch.equal(out2[0], expected)
+
+    def test_chunked_prefill_single_shot_matches_non_chunked(self):
+        pooler = self._make_all_pool(chunked=True)
+        hidden = torch.tensor(
+            [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]
+        )
+        metadata = _make_metadata([2, 3])
+        out = pooler(hidden, metadata)
+        assert len(out) == 2
+        assert torch.equal(out[0], hidden[:2])
+        assert torch.equal(out[1], hidden[2:])
+
+    def test_chunked_prefill_mixed_finished_unfinished(self):
+        pooler = self._make_all_pool(chunked=True)
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
+        metadata = _make_metadata(
+            [2, 4],
+            num_scheduled_tokens=[2, 1],
+            seq_lens=[2, 1],
+        )
+        out = pooler(hidden, metadata)
+        assert len(out) == 2
+        assert torch.equal(out[0], hidden[:2])
+        assert out[1] is None
+
+
+# ---------------------------------------------------------------------------
+# StepPool
+# ---------------------------------------------------------------------------
+class TestStepPool:
+    @staticmethod
+    def _make_step_pool(*, chunked: bool = False) -> StepPool:
+        fake_config = _FakeVllmConfig(
+            scheduler_config=_FakeSchedulerConfig(
+                enable_chunked_prefill=chunked,
+            ),
+        )
+        with patch(
+            "vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config",
+            return_value=fake_config,
+        ):
+            return StepPool()
+
+    def test_filters_by_step_tag_id(self):
+        pooler = self._make_step_pool()
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
+        token_ids = [[10, 99, 10, 20]]
+        params = [PoolingParams(task="token_classify", step_tag_id=10)]
+        metadata = _make_metadata([4], token_ids=token_ids, pooling_params=params)
+        out = pooler(hidden, metadata)
+        assert len(out) == 1
+        expected = torch.tensor([[1.0, 2.0], [5.0, 6.0]])
+        assert torch.equal(out[0], expected)
+
+    def test_filters_by_returned_token_ids(self):
+        pooler = self._make_step_pool()
+        hidden = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+        token_ids = [[10, 20]]
+        params = [PoolingParams(task="token_classify", returned_token_ids=[0, 2])]
+        metadata = _make_metadata([2], token_ids=token_ids, pooling_params=params)
+        out = pooler(hidden, metadata)
+        assert len(out) == 1
+        expected = torch.tensor([[1.0, 3.0], [4.0, 6.0]])
+        assert torch.equal(out[0], expected)
+
+    def test_no_filtering_without_params(self):
+        pooler = self._make_step_pool()
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+        token_ids = [[10, 20]]
+        params = [PoolingParams(task="token_classify")]
+        metadata = _make_metadata([2], token_ids=token_ids, pooling_params=params)
+        out = pooler(hidden, metadata)
+        assert len(out) == 1
+        assert torch.equal(out[0], hidden)
+
+    def test_combined_step_tag_and_returned_token_ids(self):
+        pooler = self._make_step_pool()
+        hidden = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]])
+        token_ids = [[99, 10, 99]]
+        params = [
+            PoolingParams(
+                task="token_classify",
+                step_tag_id=10,
+                returned_token_ids=[0, 2],
+            )
+        ]
+        metadata = _make_metadata([3], token_ids=token_ids, pooling_params=params)
+        out = pooler(hidden, metadata)
+        assert len(out) == 1
+        expected = torch.tensor([[4.0, 6.0]])
+        assert torch.equal(out[0], expected)
+
+    def test_step_tag_id_no_match_returns_empty(self):
+        pooler = self._make_step_pool()
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+        token_ids = [[10, 20]]
+        params = [PoolingParams(task="token_classify", step_tag_id=999)]
+        metadata = _make_metadata([2], token_ids=token_ids, pooling_params=params)
+        out = pooler(hidden, metadata)
+        assert len(out) == 1
+        assert out[0].shape == (0, 2)
+
+    def test_chunked_prefill_propagates_none_for_unfinished(self):
+        pooler = self._make_step_pool(chunked=True)
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+        token_ids = [[10, 20, 30, 40]]
+        params = [PoolingParams(task="token_classify", step_tag_id=10)]
+        metadata = _make_metadata(
+            [4],
+            token_ids=token_ids,
+            pooling_params=params,
+            num_scheduled_tokens=[2],
+            seq_lens=[2],
+        )
+        out = pooler(hidden, metadata)
+        assert len(out) == 1
+        assert out[0] is None
+
+    def test_chunked_prefill_filters_when_finished(self):
+        pooler = self._make_step_pool(chunked=True)
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
+        token_ids = [[10, 99, 10, 20]]
+        params = [PoolingParams(task="token_classify", step_tag_id=10)]
+        metadata = _make_metadata([4], token_ids=token_ids, pooling_params=params)
+        out = pooler(hidden, metadata)
+        assert len(out) == 1
+        expected = torch.tensor([[1.0, 2.0], [5.0, 6.0]])
+        assert torch.equal(out[0], expected)
+
+    def test_requires_token_ids_update(self):
+        pooler = self._make_step_pool()
+        update = pooler.get_pooling_updates("token_classify")
+        assert update.requires_token_ids is True
+
+
+# ---------------------------------------------------------------------------
+# get_tok_pooling_method factory
+# ---------------------------------------------------------------------------
+class TestGetTokPoolingMethod:
+    def test_all(self):
+        fake_config = _FakeVllmConfig(
+            scheduler_config=_FakeSchedulerConfig(
+                enable_chunked_prefill=False,
+            ),
+        )
+        with patch(
+            "vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config",
+            return_value=fake_config,
+        ):
+            assert isinstance(get_tok_pooling_method("ALL"), AllPool)
+
+    def test_step(self):
+        fake_config = _FakeVllmConfig(
+            scheduler_config=_FakeSchedulerConfig(
+                enable_chunked_prefill=False,
+            ),
+        )
+        with patch(
+            "vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config",
+            return_value=fake_config,
+        ):
+            assert isinstance(get_tok_pooling_method("STEP"), StepPool)
+
+    def test_unknown_raises(self):
+        with pytest.raises(NotImplementedError, match="UNKNOWN"):
+            get_tok_pooling_method("UNKNOWN")

tests/models/registry.py

@@ -935,7 +935,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {
         trust_remote_code=True,
         hf_overrides={"architectures": ["GLM4VForCausalLM"]},
     ),
-    "Glm4vForConditionalGeneration": _HfExamplesInfo("zai-org/GLM-4.1V-9B-Thinking"),
+    "Glm4vForConditionalGeneration": _HfExamplesInfo(
+        "zai-org/GLM-4.1V-9B-Thinking",
+        extras={"4.6V": "zai-org/GLM-4.6V-Flash"},
+    ),
     "Glm4vMoeForConditionalGeneration": _HfExamplesInfo("zai-org/GLM-4.5V"),
     "GlmOcrForConditionalGeneration": _HfExamplesInfo(
         "zai-org/GLM-OCR",
@@ -1368,6 +1371,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {
     "StepVLForConditionalGeneration": _HfExamplesInfo(
         "stepfun-ai/Step3-VL-10B", trust_remote_code=True
     ),
+    "Step3p7ForConditionalGeneration": _HfExamplesInfo(
+        "stepfun-ai/Step-3.7-Flash", is_available_online=False, trust_remote_code=True
+    ),
     "UltravoxModel": _HfExamplesInfo(
         "fixie-ai/ultravox-v0_5-llama-3_2-1b",
         trust_remote_code=True,

tests/renderers/test_gemma4_chat_template.py

@@ -343,3 +343,72 @@ class TestGemma4ChatTemplate:
         assert '<|"|>Alice<|"|>' in result
         assert "active:true" in result
         assert "count:42" in result
+
+    def test_tool_response_with_multimodal_content(self, gemma4_template):
+        """Multimodal placeholders in tool messages are emitted after the
+        tool_response block."""
+        messages = [
+            {"role": "user", "content": "Download the image and describe it."},
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "id": "call_1",
+                        "type": "function",
+                        "function": {
+                            "name": "download_image",
+                            "arguments": '{"url": "https://example.com/x.png"}',
+                        },
+                    },
+                ],
+            },
+            {
+                "role": "tool",
+                "tool_call_id": "call_1",
+                "content": [
+                    {"type": "text", "text": "Image downloaded successfully."},
+                    {"type": "image"},
+                ],
+            },
+        ]
+        result = _render(gemma4_template, messages, add_generation_prompt=True)
+        assert "<|tool_response>" in result
+        assert "response:download_image{" in result
+        assert "<tool_response|>" in result
+        assert "<|image|>" in result
+
+    def test_tool_response_with_all_modalities(self, gemma4_template):
+        """All multimodal types (image, audio, video) in a single tool
+        response are rendered."""
+        messages = [
+            {"role": "user", "content": "Process media"},
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "id": "c1",
+                        "type": "function",
+                        "function": {
+                            "name": "process",
+                            "arguments": "{}",
+                        },
+                    },
+                ],
+            },
+            {
+                "role": "tool",
+                "tool_call_id": "c1",
+                "content": [
+                    {"type": "text", "text": "Results."},
+                    {"type": "image"},
+                    {"type": "audio"},
+                    {"type": "video"},
+                ],
+            },
+        ]
+        result = _render(gemma4_template, messages, add_generation_prompt=True)
+        assert "<|image|>" in result
+        assert "<|audio|>" in result
+        assert "<|video|>" in result

tests/utils_/test_numa_utils.py

@@ -10,6 +10,33 @@ from vllm.config import ParallelConfig
 from vllm.utils import numa_utils
 
 
+@pytest.fixture(autouse=True)
+def _disable_pct_by_default(monkeypatch):
+    """Force PCT detection OFF unless a test opts in via ``_patch_pct_gates``.
+
+    The CI / dev machines themselves can be Xeon 6776P with PCT enabled, so a
+    plain ``cache_clear`` would let the gate auto-detect ``True`` from the
+    live filesystem and silently re-route "baseline" tests through the PCT
+    path. Stub ``/proc/cpuinfo`` and ``acpi_cppc/highest_perf`` to a state
+    that fails the gate; ``_patch_pct_gates`` re-stubs on top when needed.
+    """
+    from io import StringIO
+
+    real_open = open
+
+    def _no_pct_open(path, *args, **kwargs):
+        if path == numa_utils._PROC_CPUINFO_PATH:
+            return StringIO("processor\t: 0\nmodel name\t: Generic Test CPU\n")
+        if path == numa_utils._PCT_HIGHEST_PERF_PATH:
+            raise OSError("PCT disabled by autouse fixture")
+        return real_open(path, *args, **kwargs)
+
+    monkeypatch.setattr("builtins.open", _no_pct_open)
+    numa_utils._pct_sku_config.cache_clear()
+    yield
+    numa_utils._pct_sku_config.cache_clear()
+
+
 def _make_config(**parallel_kwargs):
     parallel_defaults = dict(
         numa_bind=False,
@@ -31,7 +58,7 @@ def _make_config(**parallel_kwargs):
 def test_get_numactl_args_with_node_binding():
     vllm_config = _make_config(numa_bind=True, numa_bind_nodes=[0, 1])
     assert (
-        numa_utils._get_numactl_args(vllm_config, local_rank=1)
+        numa_utils._get_numactl_worker_args(vllm_config.parallel_config, local_rank=1)
         == "--cpunodebind=1 --membind=1"
     )
 
@@ -43,7 +70,300 @@ def test_get_numactl_args_with_cpu_binding():
         numa_bind_cpus=["0-3", "4-7"],
     )
     assert (
-        numa_utils._get_numactl_args(vllm_config, local_rank=1)
+        numa_utils._get_numactl_worker_args(vllm_config.parallel_config, local_rank=1)
+        == "--physcpubind=4-7 --membind=1"
+    )
+
+
+def _patch_pct_gates(
+    monkeypatch,
+    *,
+    model_match: bool,
+    highest_perf: int | None,
+    cpulist: str | None = "0-31,64-95",
+    cpulist_by_node: dict[int, str | None] | None = None,
+    sku: str = "6776P",
+):
+    """Force `_pct_sku_config` and node cpulist read to deterministic state.
+
+    ``cpulist`` is the default returned for any node not present in
+    ``cpulist_by_node``. ``cpulist_by_node`` lets a test return different
+    cpulists for different NUMA nodes (e.g. node 0 vs node 1). ``sku`` lets
+    the test pick which Granite Rapids SKU appears in the fake
+    ``/proc/cpuinfo`` ``model name`` (only used when ``model_match=True``).
+    """
+    import pathlib
+    from io import StringIO
+
+    import regex as re
+
+    cpuinfo = (
+        f"processor\t: 0\nmodel name\t: Intel(R) Xeon(R) Platinum {sku} CPU @ 2.40GHz\n"
+        if model_match
+        else "processor\t: 0\nmodel name\t: Intel(R) Xeon(R) Platinum 8480+\n"
+    )
+
+    real_open = open
+
+    def fake_open(path, *args, **kwargs):
+        if path == numa_utils._PROC_CPUINFO_PATH:
+            return StringIO(cpuinfo)
+        if path == numa_utils._PCT_HIGHEST_PERF_PATH:
+            if highest_perf is None:
+                raise OSError("missing")
+            return StringIO(f"{highest_perf}\n")
+        return real_open(path, *args, **kwargs)
+
+    real_read_text = pathlib.Path.read_text
+    cpulist_by_node = cpulist_by_node or {}
+
+    def fake_read_text(self, *args, **kwargs):
+        path_str = str(self)
+        if path_str.endswith("/cpulist") and "/sys/devices/system/node" in path_str:
+            match = re.search(r"/node(\d+)/cpulist$", path_str)
+            if match:
+                node_id = int(match.group(1))
+                if node_id in cpulist_by_node:
+                    val = cpulist_by_node[node_id]
+                    if val is None:
+                        raise OSError(f"missing cpulist for node{node_id}")
+                    return val
+            if cpulist is None:
+                raise OSError("missing cpulist")
+            return cpulist
+        return real_read_text(self, *args, **kwargs)
+
+    monkeypatch.setattr("builtins.open", fake_open)
+    monkeypatch.setattr("pathlib.Path.read_text", fake_read_text)
+    numa_utils._pct_sku_config.cache_clear()
+
+
+def test_pct_binding_filters_cpus(monkeypatch):
+    _patch_pct_gates(monkeypatch, model_match=True, highest_perf=46)
+    assert numa_utils._maybe_get_pct_cpu_binding([0]) == [0, 1, 16, 17, 64, 65, 80, 81]
+
+
+@pytest.mark.parametrize(
+    "sku,expected_cpus",
+    [
+        # 64-core SKUs (stride 16): cpus from "0-31,64-95" with cpu_id % 16
+        # in (0, 1) -> 0, 1, 16, 17, 64, 65, 80, 81.
+        ("6776P", [0, 1, 16, 17, 64, 65, 80, 81]),
+        ("6774P", [0, 1, 16, 17, 64, 65, 80, 81]),
+        # 72-core SKU (stride 18): cpus from "0-31,64-95" with cpu_id % 18
+        # in (0, 1) -> 0, 1, 18, 19, 72, 73, 90, 91.
+        ("6962P", [0, 1, 18, 19, 72, 73, 90, 91]),
+    ],
+)
+def test_pct_binding_fires_on_every_capable_sku(monkeypatch, sku, expected_cpus):
+    """Each SKU in ``_PCT_CAPABLE_SKUS`` engages the gate at its own
+    expected ``highest_perf`` and uses its own priority-core stride."""
+    sku_config = numa_utils._PCT_CAPABLE_SKUS[sku]
+    _patch_pct_gates(
+        monkeypatch,
+        model_match=True,
+        highest_perf=sku_config.highest_perf,
+        sku=sku,
+    )
+    assert numa_utils._maybe_get_pct_cpu_binding([0]) == expected_cpus
+
+
+def test_pct_binding_fails_closed_when_sku_perf_mismatch(monkeypatch):
+    """6962P with 6776P's highest_perf (46 vs expected 44) must fail closed."""
+    _patch_pct_gates(monkeypatch, model_match=True, highest_perf=46, sku="6962P")
+    assert numa_utils._maybe_get_pct_cpu_binding([0]) is None
+
+
+def test_pct_binding_disabled_when_cpu_model_mismatch(monkeypatch):
+    _patch_pct_gates(monkeypatch, model_match=False, highest_perf=46)
+    assert numa_utils._maybe_get_pct_cpu_binding([0]) is None
+
+
+def test_pct_binding_disabled_when_highest_perf_does_not_match(monkeypatch):
+    _patch_pct_gates(monkeypatch, model_match=True, highest_perf=42)
+    assert numa_utils._maybe_get_pct_cpu_binding([0]) is None
+
+
+def test_pct_binding_disabled_when_files_missing(monkeypatch):
+    _patch_pct_gates(monkeypatch, model_match=True, highest_perf=None)
+    assert numa_utils._maybe_get_pct_cpu_binding([0]) is None
+
+
+def test_pct_binding_returns_none_when_node_cpulist_filter_empty(monkeypatch):
+    _patch_pct_gates(
+        monkeypatch,
+        model_match=True,
+        highest_perf=46,
+        cpulist="2-15,18-31",
+    )
+    assert numa_utils._maybe_get_pct_cpu_binding([0]) is None
+
+
+def test_pct_binding_returns_none_when_node_cpulist_missing(monkeypatch):
+    _patch_pct_gates(monkeypatch, model_match=True, highest_perf=46, cpulist=None)
+    assert numa_utils._maybe_get_pct_cpu_binding([0]) is None
+
+
+def test_get_numactl_args_uses_pct_when_user_did_not_specify_cpus(monkeypatch):
+    _patch_pct_gates(monkeypatch, model_match=True, highest_perf=46)
+    vllm_config = _make_config(numa_bind=True, numa_bind_nodes=[0, 1])
+    assert (
+        numa_utils._get_numactl_worker_args(vllm_config.parallel_config, local_rank=1)
+        == "--physcpubind=0,1,16,17,64,65,80,81 --membind=1"
+    )
+
+
+def test_get_numactl_args_engine_core_baseline_single_node_shard():
+    """Baseline (no PCT): single-NUMA shard -> single-node bind."""
+    vllm_config = _make_config(numa_bind=True, numa_bind_nodes=[0, 1])
+    assert (
+        numa_utils._get_numactl_enginecore_args(
+            vllm_config.parallel_config, local_rank=0
+        )
+        == "--cpunodebind=0 --membind=0"
+    )
+
+
+def test_get_numactl_args_engine_core_baseline_spans_shard_numa_nodes():
+    """Baseline (no PCT): a TP=4 shard spanning both NUMA nodes -> bind to both."""
+    vllm_config = _make_config(
+        numa_bind=True,
+        numa_bind_nodes=[0, 0, 1, 1],
+        tensor_parallel_size=4,
+    )
+    assert (
+        numa_utils._get_numactl_enginecore_args(
+            vllm_config.parallel_config, local_rank=0
+        )
+        == "--cpunodebind=0,1 --membind=0,1"
+    )
+
+
+def test_get_numactl_args_engine_core_pct_spans_shard_numa_nodes(monkeypatch):
+    """PCT: EngineCore for a multi-NUMA shard binds to the union of priority
+    cores across all shard nodes, so worker `--physcpubind` is always a
+    subset of EngineCore's `cpus_allowed`."""
+    _patch_pct_gates(
+        monkeypatch,
+        model_match=True,
+        highest_perf=46,
+        cpulist_by_node={0: "0-31,128-159", 1: "64-95,192-223"},
+    )
+    vllm_config = _make_config(
+        numa_bind=True,
+        numa_bind_nodes=[0, 0, 1, 1],
+        tensor_parallel_size=4,
+    )
+    assert numa_utils._get_numactl_enginecore_args(
+        vllm_config.parallel_config, local_rank=0
+    ) == (
+        "--physcpubind="
+        "0,1,16,17,64,65,80,81,128,129,144,145,192,193,208,209"
+        " --membind=0,1"
+    )
+
+
+def test_get_numactl_args_engine_core_pct_dp_shard_picks_local_nodes(monkeypatch):
+    """With DP=2, each shard's EngineCore binds only to its own NUMA nodes."""
+    _patch_pct_gates(
+        monkeypatch,
+        model_match=True,
+        highest_perf=46,
+        cpulist_by_node={0: "0-31,128-159", 1: "64-95,192-223"},
+    )
+    vllm_config = _make_config(
+        numa_bind=True,
+        numa_bind_nodes=[0, 0, 1, 1],
+        tensor_parallel_size=2,
+        data_parallel_rank_local=1,
+    )
+    # Shard 1 owns gpu_indices 2 and 3 -> nodes [1, 1] -> {1}.
+    assert (
+        numa_utils._get_numactl_enginecore_args(
+            vllm_config.parallel_config, local_rank=0
+        )
+        == "--physcpubind=64,65,80,81,192,193,208,209 --membind=1"
+    )
+
+
+def test_get_numactl_args_engine_core_pct_external_launcher_spans_local_nodes(
+    monkeypatch,
+):
+    """external_launcher (or multi-node-within-DP, or Ray) hits the
+    fallback branch. EngineCore must still span every local NUMA node
+    so it can mp-spawn its local workers without ``--physcpubind``
+    strict-validation failures."""
+    _patch_pct_gates(
+        monkeypatch,
+        model_match=True,
+        highest_perf=46,
+        cpulist_by_node={0: "0-31,128-159", 1: "64-95,192-223"},
+    )
+    vllm_config = _make_config(
+        numa_bind=True,
+        numa_bind_nodes=[0, 0, 0, 0, 1, 1, 1, 1],
+        distributed_executor_backend="external_launcher",
+        tensor_parallel_size=8,
+    )
+    assert numa_utils._get_numactl_enginecore_args(
+        vllm_config.parallel_config, local_rank=0
+    ) == (
+        "--physcpubind="
+        "0,1,16,17,64,65,80,81,128,129,144,145,192,193,208,209"
+        " --membind=0,1"
+    )
+
+
+def test_get_numactl_args_engine_core_baseline_multi_node_within_dp_spans_locals():
+    """Multi-node-within-DP fallback: bind EngineCore to all local NUMA
+    nodes that the visible ``numa_bind_nodes`` reference."""
+    vllm_config = _make_config(
+        numa_bind=True,
+        numa_bind_nodes=[0, 0, 0, 0, 1, 1, 1, 1],
+        nnodes_within_dp=2,
+        tensor_parallel_size=8,
+    )
+    assert (
+        numa_utils._get_numactl_enginecore_args(
+            vllm_config.parallel_config, local_rank=0
+        )
+        == "--cpunodebind=0,1 --membind=0,1"
+    )
+
+
+def test_get_numactl_args_engine_core_skips_user_cpu_list(monkeypatch):
+    """EngineCore ignores ``--numa-bind-cpus``.
+
+    Those are per-worker lists; binding EngineCore to any of them would
+    shrink its ``cpus_allowed`` below the strict-superset workers'
+    ``--physcpubind`` spawns need. We fall back to ``--cpunodebind`` over
+    the shard's NUMA nodes instead. PCT auto-detect is also bypassed when
+    the user is explicit (its priority-core union may not be a superset
+    of the user's per-worker cores)."""
+    _patch_pct_gates(monkeypatch, model_match=True, highest_perf=46)
+    vllm_config = _make_config(
+        numa_bind=True,
+        numa_bind_nodes=[0, 0, 1, 1],
+        numa_bind_cpus=["0-3", "4-7", "64-67", "68-71"],
+        tensor_parallel_size=4,
+    )
+    assert (
+        numa_utils._get_numactl_enginecore_args(
+            vllm_config.parallel_config, local_rank=0
+        )
+        == "--cpunodebind=0,1 --membind=0,1"
+    )
+
+
+def test_get_numactl_args_user_cpus_override_pct(monkeypatch):
+    _patch_pct_gates(monkeypatch, model_match=True, highest_perf=46)
+    vllm_config = _make_config(
+        numa_bind=True,
+        numa_bind_nodes=[0, 1],
+        numa_bind_cpus=["0-3", "4-7"],
+    )
+    assert (
+        numa_utils._get_numactl_worker_args(vllm_config.parallel_config, local_rank=1)
         == "--physcpubind=4-7 --membind=1"
     )
 
@@ -57,7 +377,7 @@ def test_get_numactl_args_uses_dp_offset():
         tensor_parallel_size=2,
     )
     assert (
-        numa_utils._get_numactl_args(vllm_config, local_rank=1)
+        numa_utils._get_numactl_worker_args(vllm_config.parallel_config, local_rank=1)
         == "--cpunodebind=1 --membind=1"
     )
 
@@ -66,7 +386,21 @@ def test_get_numactl_args_requires_detectable_nodes(monkeypatch):
     vllm_config = _make_config(numa_bind=True)
     monkeypatch.setattr(numa_utils, "get_auto_numa_nodes", lambda: None)
     with pytest.raises(RuntimeError):
-        numa_utils._get_numactl_args(vllm_config, local_rank=0)
+        numa_utils._get_numactl_worker_args(vllm_config.parallel_config, local_rank=0)
+
+
+def test_configure_subprocess_rejects_unknown_process_kind():
+    """configure_subprocess only knows 'worker' and 'EngineCore'; anything
+    else must raise ValueError instead of silently routing to the worker
+    path."""
+    vllm_config = _make_config(numa_bind=True, numa_bind_nodes=[0])
+    with (
+        pytest.raises(ValueError, match="process_kind"),
+        numa_utils.configure_subprocess(
+            vllm_config, local_rank=0, process_kind="bogus"
+        ),
+    ):
+        pass
 
 
 def test_log_numactl_show(monkeypatch):

tests/v1/cudagraph/test_encoder_cudagraph.py

@@ -87,8 +87,10 @@ class _MockModel(SupportsEncoderCudaGraph):
     def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
         return EncoderCudaGraphConfig(
             modalities=["image"],
-            input_key_by_modality={"image": "pixel_values"},
-            buffer_keys=["dummy_buf"],
+            buffer_keys=[
+                "pixel_values",
+                "dummy_buf",
+            ],
             out_hidden_size=32,
         )
 
@@ -269,9 +271,6 @@ class SimpleMockViTModel(torch.nn.Module, SupportsEncoderCudaGraph):
     def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
         return EncoderCudaGraphConfig(
             modalities=["image"],
-            input_key_by_modality={
-                "image": "pixel_values",
-            },
             buffer_keys=["dummy_buf"],
             out_hidden_size=_HIDDEN,
         )
@@ -350,11 +349,10 @@ class SimpleMockViTModel(torch.nn.Module, SupportsEncoderCudaGraph):
         n_out = _count_output_tokens(grid_config, _SPATIAL_MERGE)
         dummy_buf = torch.zeros(n_out, _HIDDEN, device=device, dtype=dtype)
         return EncoderCudaGraphCaptureInputs(
-            mm_kwargs={
+            values={
                 "pixel_values": dummy_pixel_values,
-                "image_grid_thw": grid_config,
+                "dummy_buf": dummy_buf,
             },
-            buffers={"dummy_buf": dummy_buf},
         )
 
     def prepare_encoder_cudagraph_replay_buffers(
@@ -367,14 +365,18 @@ class SimpleMockViTModel(torch.nn.Module, SupportsEncoderCudaGraph):
         n_out = _count_output_tokens(grid_thw, _SPATIAL_MERGE)
         p = next(self.parameters())
         dummy_buf = torch.zeros(n_out, _HIDDEN, device=p.device, dtype=p.dtype)
-        return EncoderCudaGraphReplayBuffers(buffers={"dummy_buf": dummy_buf})
+        return EncoderCudaGraphReplayBuffers(
+            values={
+                "pixel_values": mm_kwargs["pixel_values"],
+                "dummy_buf": dummy_buf,
+            }
+        )
 
     def encoder_cudagraph_forward(
         self,
-        mm_kwargs: dict[str, Any],
-        buffers: dict[str, torch.Tensor],
+        values: dict[str, torch.Tensor],
     ) -> torch.Tensor:
-        return self._forward(mm_kwargs["pixel_values"])
+        return self._forward(values["pixel_values"])
 
     def encoder_eager_forward(
         self,
@@ -551,10 +553,6 @@ class SimpleMockViTVideoModel(SimpleMockViTModel):
     def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
         return EncoderCudaGraphConfig(
             modalities=["image", "video"],
-            input_key_by_modality={
-                "image": "pixel_values",
-                "video": "pixel_values_videos",
-            },
             buffer_keys=["dummy_buf"],
             out_hidden_size=_HIDDEN,
         )
@@ -654,11 +652,10 @@ class SimpleMockViTVideoModel(SimpleMockViTModel):
         n_out = _count_output_tokens(grid_config, _SPATIAL_MERGE)
         dummy_buf = torch.zeros(n_out, _HIDDEN, device=device, dtype=dtype)
         return EncoderCudaGraphCaptureInputs(
-            mm_kwargs={
+            values={
                 "pixel_values": dummy_pixel_values,
-                "image_grid_thw": grid_config,
+                "dummy_buf": dummy_buf,
             },
-            buffers={"dummy_buf": dummy_buf},
         )
 
     def prepare_encoder_cudagraph_replay_buffers(
@@ -670,14 +667,18 @@ class SimpleMockViTVideoModel(SimpleMockViTModel):
         n_out = _count_output_tokens(self._get_grid_thw(mm_kwargs), _SPATIAL_MERGE)
         p = next(self.parameters())
         dummy_buf = torch.zeros(n_out, _HIDDEN, device=p.device, dtype=p.dtype)
-        return EncoderCudaGraphReplayBuffers(buffers={"dummy_buf": dummy_buf})
+        return EncoderCudaGraphReplayBuffers(
+            values={
+                "pixel_values": self._get_pixel_values(mm_kwargs),
+                "dummy_buf": dummy_buf,
+            }
+        )
 
     def encoder_cudagraph_forward(
         self,
-        mm_kwargs: dict[str, Any],
-        buffers: dict[str, torch.Tensor],
+        values: dict[str, torch.Tensor],
     ) -> torch.Tensor:
-        return self._forward(self._get_pixel_values(mm_kwargs))
+        return self._forward(values["pixel_values"])
 
     def encoder_eager_forward(
         self,
@@ -718,14 +719,6 @@ class TestGetInputModality:
         }
         assert model.get_input_modality(mm_kwargs) == "video"
 
-    def test_video_model_config_has_both_modalities(self):
-        model = SimpleMockViTVideoModel()
-        cfg = model.get_encoder_cudagraph_config()
-        assert "image" in cfg.modalities
-        assert "video" in cfg.modalities
-        assert cfg.input_key_by_modality["image"] == "pixel_values"
-        assert cfg.input_key_by_modality["video"] == "pixel_values_videos"
-
 
 # ---------------------------------------------------------------------------
 # GPU tests — video capture, replay, fallback, and mixed image+video

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

@@ -24,7 +24,9 @@ from vllm.v1.kv_cache_interface import (
 from vllm.v1.kv_offload.base import (
     OffloadingEvent,
     OffloadingManager,
+    OffloadPolicy,
     ReqContext,
+    RequestOffloadingContext,
     get_offload_block_hash,
 )
 from vllm.v1.request import RequestStatus
@@ -748,12 +750,8 @@ class TestSlidingWindowLookup:
 
 
 @pytest.mark.parametrize("async_scheduling", [True, False])
-def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bool):
-    """With do_remote_decode=True, after loading prefix blocks from CPU,
-    all blocks must be re-stored — not just the newly computed ones.
-
-    This supports P/D disaggregation where the prefill instance offloads the
-    complete KV cache so a remote decode node can consume it."""
+def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling: bool):
+    """With REQUEST_LEVEL policy, all blocks are stored — including prefix hits."""
     gpu_block_size = 4
     block_size_factor = 3
     offloaded_block_size = gpu_block_size * block_size_factor
@@ -780,12 +778,13 @@ def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bo
     # Reset GPU prefix cache so the next request must load from CPU.
     runner.scheduler.reset_prefix_cache()
 
-    # New request with do_remote_decode=True and 2 offloaded blocks.
-    # The first offloaded block matches what we stored in CPU.
-    runner.new_request(
-        token_ids=[0] * offloaded_block_size * 2,
-        kv_transfer_params={"do_remote_decode": True},
+    # Manager returns REQUEST_LEVEL for the next request.
+    runner.manager.on_new_request.return_value = RequestOffloadingContext(
+        policy=OffloadPolicy.REQUEST_LEVEL
     )
+
+    # New request with 2 offloaded blocks; first matches what's in CPU.
+    runner.new_request(token_ids=[0] * offloaded_block_size * 2)
     runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output(keys)
@@ -957,9 +956,10 @@ def test_max_offload_tokens_validation(request_runner, async_scheduling: bool):
         )
 
     def setup(r, max_offload_tokens):
-        r.new_request(token_ids=[0] * offloaded_block_size * 3)
-        req = r.scheduler.requests[str(r.req_id)]
-        req.kv_transfer_params = {"max_offload_tokens": max_offload_tokens}
+        r.new_request(
+            token_ids=[0] * offloaded_block_size * 3,
+            kv_transfer_params={"max_offload_tokens": max_offload_tokens},
+        )
         r.manager.prepare_store.side_effect = (
             lambda keys, req_context: generate_store_output(keys)
         )
@@ -1029,6 +1029,56 @@ def test_max_offload_tokens_validation(request_runner, async_scheduling: bool):
     )
 
 
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_offload_prompt_only(request_runner, async_scheduling: bool):
+    """offload_prompt_only=True offloads prompt blocks but never decode blocks.
+
+    Setup: a 2-offloaded-block prompt followed by enough decode tokens to fill
+    4 more offloaded blocks. The flag clamps the offloadable token count to the
+    prompt length, so only the prompt's blocks (GPU offsets 0-5) are ever
+    eligible for store; the decode blocks (offsets >= 6) are skipped.
+
+    The request is intentionally not terminated (no EOS): a store is only
+    flushed when a request finishes (or is preempted), so without a finish
+    there is nothing to flush and the assertion stays free of flush-timing
+    subtleties. The decode steps are still enough for the prompt store to
+    complete and show up in expected_stored.
+    """
+    gpu_block_size = 4
+    block_size_factor = 3
+    offloaded_block_size = gpu_block_size * block_size_factor  # 12
+    num_prompt_blocks = 2
+    num_decode_blocks = 4
+    prompt_offsets = (0, 1, 2, 3, 4, 5)
+
+    runner = request_runner(
+        block_size=gpu_block_size,
+        num_gpu_blocks=100,
+        async_scheduling=async_scheduling,
+        block_size_factor=block_size_factor,
+        extra_config_overrides={"offload_prompt_only": True},
+    )
+
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output(keys)
+    )
+
+    runner.new_request(token_ids=[0] * offloaded_block_size * num_prompt_blocks)
+    runner.run(
+        decoded_tokens=[0] * (offloaded_block_size * num_decode_blocks),
+        expected_stored=prompt_offsets,
+    )
+
+    # Timing-independent guard: only the prompt's blocks were ever offered for
+    # store. If decode blocks leaked through, more keys would appear here.
+    offered_keys = {
+        key
+        for call in runner.manager.prepare_store.call_args_list
+        for key in call.args[0]
+    }
+    assert len(offered_keys) == num_prompt_blocks
+
+
 def test_flush_all_jobs_when_no_requests_remain(request_runner):
     """When all tracked requests are finished, build_connector_meta flushes
     all pending jobs since there will be no future step to complete them."""

tests/v1/kv_connector/unit/offloading_connector/utils.py

@@ -43,6 +43,7 @@ from vllm.v1.kv_offload.base import (
     OffloadingSpec,
     OffloadKey,
     PrepareStoreOutput,
+    RequestOffloadingContext,
     make_offload_key,
 )
 from vllm.v1.kv_offload.worker.worker import (
@@ -119,6 +120,7 @@ class MockOffloadingSpec(OffloadingSpec):
         self.manager.lookup.return_value = 0
         self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys)
         self.manager.lookup.return_value = False
+        self.manager.on_new_request.return_value = RequestOffloadingContext()
         self.handler = MockOffloadingHandler()
 
     def get_manager(self) -> OffloadingManager:
@@ -169,6 +171,7 @@ class RequestRunner:
         block_size_factor: int = 1,
         async_scheduling: bool = True,
         kv_cache_groups: list[KVCacheGroupSpec] | None = None,
+        extra_config_overrides: dict[str, Any] | None = None,
     ):
         assert block_size_factor == 1 or kv_cache_groups is None, (
             "block_size_factor > 1 requires all groups to have the same "
@@ -192,9 +195,13 @@ class RequestRunner:
         extra_config: dict[str, Any] = {
             "spec_name": "MockOffloadingSpec",
             "spec_module_path": "tests.v1.kv_connector.unit.offloading_connector.utils",  # noqa: E501
+            # Preserve legacy behavior for tests; new opt-in tests override.
+            "offload_prompt_only": False,
         }
         if block_size_factor > 1:
             extra_config["block_size"] = block_size * block_size_factor
+        if extra_config_overrides:
+            extra_config.update(extra_config_overrides)
 
         vllm_config.kv_transfer_config = KVTransferConfig(
             kv_connector="OffloadingConnector",
@@ -591,6 +598,7 @@ def request_runner():
         async_scheduling,
         block_size_factor=1,
         kv_cache_groups=None,
+        extra_config_overrides=None,
     ):
         runner = RequestRunner(
             block_size=block_size,
@@ -598,6 +606,7 @@ def request_runner():
             block_size_factor=block_size_factor,
             async_scheduling=async_scheduling,
             kv_cache_groups=kv_cache_groups,
+            extra_config_overrides=extra_config_overrides,
         )
         runners.append(runner)
         return runner

tests/v1/kv_connector/unit/test_moriio_connector.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import importlib.util
-import subprocess
 import uuid
 from unittest.mock import MagicMock, patch
 
@@ -48,18 +47,6 @@ def _make_test_kv_cache_config() -> KVCacheConfig:
 aiter_available = importlib.util.find_spec("aiter") is not None
 mori_available = importlib.util.find_spec("mori") is not None
 
-
-def _rdma_available() -> bool:
-    """Check if RDMA devices are available."""
-    try:
-        result = subprocess.run(["ibv_devinfo"], capture_output=True, text=True)
-        return "No IB devices found" not in result.stderr
-    except FileNotFoundError:
-        return False
-
-
-rdma_available = _rdma_available()
-
 pytestmark = pytest.mark.skipif(
     not (current_platform.is_rocm() and mori_available),
     reason="MoRIIOs are only available on ROCm with aiter package installed",
@@ -224,11 +211,14 @@ def create_vllm_config(
         cache_dtype="auto",
         enable_prefix_caching=True,
     )
+    # These tests exercise connector setup, not real RDMA transfer (MoRI wrapper is
+    # mocked), so we can use any backend without affecting test validity. Use xGMI to
+    # avoid requiring RNICs in CI.
     kv_transfer_config = KVTransferConfig(
         kv_connector="MoRIIOConnector",
         kv_role=role,
         enable_permute_local_kv=enable_permute_local_kv,
-        kv_connector_extra_config={"read_mode": read_mode},
+        kv_connector_extra_config={"read_mode": read_mode, "backend": "xgmi"},
     )
     return VllmConfig(
         scheduler_config=scheduler_config,
@@ -417,7 +407,6 @@ def test_read_mode_loads_remote_block_ids():
 @pytest.mark.skipif(
     not aiter_available, reason="Requires aiter package for ROCm FlashAttention backend"
 )
-@pytest.mark.skipif(not rdma_available, reason="No RDMA devices available")
 def test_register_kv_caches(mock_parallel_groups):
     """Test that MoRIIOConnector.register_kv_caches correctly registers kv caches."""
     ROLE = "kv_consumer"
@@ -517,7 +506,6 @@ def test_register_kv_caches(mock_parallel_groups):
 @pytest.mark.skipif(
     not aiter_available, reason="Requires aiter package for ROCm FlashAttention backend"
 )
-@pytest.mark.skipif(not rdma_available, reason="No RDMA devices available")
 def test_moriio_handshake_returns_metadata(mock_parallel_groups):
     """MoRIIO handshake socket returns valid agent metadata over ZMQ."""
 

tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py

@@ -0,0 +1,325 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for MultiConnector wrapping NixlConnector + SimpleCPUOffloadConnector.
+
+Verifies scheduler-level behavior: HMA support detection, load delegation
+(first-wins), store-to-all, and connector metadata aggregation.
+"""
+
+import uuid
+from collections import defaultdict
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from tests.v1.kv_connector.unit.utils import (
+    create_model_runner_output,
+    create_request,
+    create_scheduler,
+    create_vllm_config,
+)
+from vllm import SamplingParams
+from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
+    MultiConnector,
+    MultiKVConnectorMetadata,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl import (
+    NixlConnectorMetadata,
+)
+from vllm.utils.hashing import sha256
+from vllm.v1.core.kv_cache_utils import get_request_block_hasher
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheTensor,
+    SlidingWindowSpec,
+)
+from vllm.v1.outputs import KVConnectorOutput
+from vllm.v1.request import Request
+from vllm.v1.simple_kv_offload.metadata import (
+    SimpleCPUOffloadMetadata,
+    SimpleCPUOffloadWorkerMetadata,
+)
+
+NIXL_WRAPPER_PATCH = (
+    "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper"
+)
+
+
+class FakeNixlWrapper:
+    """Minimal mock of NixlWrapper for testing without NIXL hardware.
+
+    Duplicated from test_nixl_connector.py to avoid importing that module
+    (which has heavy dependencies like ray).
+    """
+
+    AGENT_METADATA = b"fake_agent_metadata"
+    REMOTE_AGENT_NAME = "remote_agent"
+
+    def __init__(self, agent_name: str, *args, **kwargs):
+        self._cycles_before_xfer_done = 0
+        self._check_xfer_state_cycles: defaultdict[int, int] = defaultdict(lambda: 0)
+
+    def get_reg_descs(self, caches_data, memory_type: str) -> list:
+        return [str(uuid.uuid4()) for _ in caches_data]
+
+    def register_memory(self, descs, backends) -> None:
+        pass
+
+    def deregister_memory(self, descs) -> None:
+        pass
+
+    def get_xfer_descs(self, blocks_data, memory_type: str) -> list:
+        return [str(uuid.uuid4()) for _ in blocks_data]
+
+    def prep_xfer_dlist(self, agent_name: str, descs: list) -> int:
+        return uuid.uuid4().int
+
+    def get_agent_metadata(self) -> bytes:
+        return self.AGENT_METADATA
+
+    def add_remote_agent(self, agent_metadata: bytes) -> str:
+        return self.REMOTE_AGENT_NAME
+
+    def get_new_notifs(self) -> dict[str, list[bytes]]:
+        return {}
+
+    def check_xfer_state(self, handle: int) -> str:
+        if self._check_xfer_state_cycles[handle] >= self._cycles_before_xfer_done:
+            return "DONE"
+        self._check_xfer_state_cycles[handle] += 1
+        return "PROC"
+
+    def release_xfer_handle(self, handle: int) -> None:
+        pass
+
+    def release_dlist_handle(self, handle: int) -> None:
+        pass
+
+    def remove_remote_agent(self, agent: str) -> None:
+        pass
+
+    def send_notif(self, agent_name: str, notif_msg: bytes) -> None:
+        pass
+
+    def make_prepped_xfer(self, *args, **kwargs) -> int:
+        return uuid.uuid4().int
+
+    def transfer(self, handle: int) -> str:
+        return "PROC"
+
+    def get_xfer_telemetry(self, handle: int) -> dict:
+        return {}
+
+    def set_cycles_before_xfer_done(self, cycles: int):
+        pass
+
+
+BLOCK_SIZE = 16
+NUM_KV_HEADS = 1
+HEAD_SIZE = 16
+DTYPE = torch.float16
+_BYTES_PER_BLOCK = BLOCK_SIZE * NUM_KV_HEADS * HEAD_SIZE * 2 * DTYPE.itemsize
+
+
+def _make_kv_cache_config(
+    num_blocks: int = 100,
+    swa_enabled: bool = False,
+) -> KVCacheConfig:
+    """Build KVCacheConfig with non-empty kv_cache_tensors.
+
+    SimpleCPUOffloadConnector requires kv_cache_tensors with real sizes
+    (used by _derive_cpu_config).
+    """
+    fa_layers = ["layer0", "layer2"]
+    groups = [
+        KVCacheGroupSpec(
+            fa_layers,
+            FullAttentionSpec(
+                block_size=BLOCK_SIZE,
+                num_kv_heads=NUM_KV_HEADS,
+                head_size=HEAD_SIZE,
+                dtype=DTYPE,
+            ),
+        )
+    ]
+    tensors = [
+        KVCacheTensor(
+            size=_BYTES_PER_BLOCK * num_blocks,
+            shared_by=fa_layers,
+        )
+    ]
+    if swa_enabled:
+        sw_layers = ["layer1", "layer3"]
+        groups.append(
+            KVCacheGroupSpec(
+                sw_layers,
+                SlidingWindowSpec(
+                    block_size=BLOCK_SIZE,
+                    num_kv_heads=NUM_KV_HEADS,
+                    head_size=HEAD_SIZE,
+                    dtype=DTYPE,
+                    sliding_window=128,
+                ),
+            )
+        )
+        tensors.append(
+            KVCacheTensor(
+                size=_BYTES_PER_BLOCK * num_blocks,
+                shared_by=sw_layers,
+            )
+        )
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=tensors,
+        kv_cache_groups=groups,
+    )
+
+
+def _multi_connector_config(swa_enabled: bool = False):
+    """Return (vllm_config, kv_cache_config) for a MultiConnector test."""
+    kv_cache_config = _make_kv_cache_config(swa_enabled=swa_enabled)
+    vllm_config = create_vllm_config(
+        kv_connector="MultiConnector",
+        kv_connector_extra_config={
+            "connectors": [
+                {
+                    "kv_connector": "NixlConnector",
+                    "kv_role": "kv_both",
+                },
+                {
+                    "kv_connector": "SimpleCPUOffloadConnector",
+                    "kv_role": "kv_both",
+                    "kv_connector_extra_config": {
+                        "cpu_bytes_to_use": 1 << 30,
+                    },
+                },
+            ],
+        },
+    )
+    return vllm_config, kv_cache_config
+
+
+@patch(NIXL_WRAPPER_PATCH, FakeNixlWrapper)
+def test_nixl_wins_load_over_cpu_offload():
+    """When NixlConnector (index 0) has matched tokens from a remote prefill, it should
+    win the load: Nixl metadata tracks the recv while CPU offload metadata has no load
+     scheduled."""
+    vllm_config, kv_cache_config = _multi_connector_config()
+    scheduler = create_scheduler(vllm_config, kv_cache_config=kv_cache_config)
+    mc = scheduler.connector
+    assert isinstance(mc, MultiConnector)
+
+    request = create_request(
+        request_id=1,
+        num_tokens=BLOCK_SIZE * 3,
+        do_remote_prefill=True,
+        block_size=BLOCK_SIZE,
+    )
+    scheduler.add_request(request)
+    sched_out = scheduler.schedule()
+
+    assert mc._requests_to_connector[request.request_id] == 0
+
+    meta = sched_out.kv_connector_metadata
+    assert isinstance(meta, MultiKVConnectorMetadata)
+    assert len(meta.metadata) == 2
+
+    nixl_meta = meta.metadata[0]
+    assert isinstance(nixl_meta, NixlConnectorMetadata)
+    # nixl is tracking the request
+    assert request.request_id in nixl_meta.reqs_to_recv
+
+    cpu_meta = meta.metadata[1]
+    assert isinstance(cpu_meta, SimpleCPUOffloadMetadata)
+    assert not cpu_meta.load_gpu_blocks
+
+
+@patch(NIXL_WRAPPER_PATCH, FakeNixlWrapper)
+def test_cpu_offload_wins_when_nixl_has_no_match():
+    """When NixlConnector returns 0 matched tokens and SimpleCPUOffloadConnector has a
+    CPU cache hit, the CPU offload connector (index 1) wins the load."""
+    vllm_config, kv_cache_config = _multi_connector_config()
+    scheduler = create_scheduler(vllm_config, kv_cache_config=kv_cache_config)
+    mc = scheduler.connector
+    assert isinstance(mc, MultiConnector)
+
+    req1 = create_request(
+        request_id=10,
+        num_tokens=BLOCK_SIZE * 3,
+        block_size=BLOCK_SIZE,
+    )
+    scheduler.add_request(req1)
+    sched_out1 = scheduler.schedule()
+
+    # build_connector_meta runs before _update_after_schedule, so num_computed_tokens
+    # is still 0 during the first schedule and the store sees no confirmed blocks.
+    # Simulate one model step so the second schedule triggers the store.
+    model_output = create_model_runner_output(reqs=[req1])
+    scheduler.update_from_output(sched_out1, model_output)
+
+    sched_out2 = scheduler.schedule()
+    meta2 = sched_out2.kv_connector_metadata
+    assert isinstance(meta2, MultiKVConnectorMetadata)
+    cpu_meta = meta2.metadata[1]
+    assert isinstance(cpu_meta, SimpleCPUOffloadMetadata)
+    assert cpu_meta.store_event >= 0, "Expected a store event on the second schedule"
+
+    cpu_connector = mc._connectors[1]
+    worker_meta = SimpleCPUOffloadWorkerMetadata(
+        completed_store_events={
+            cpu_meta.store_event: cpu_connector.scheduler_manager._expected_worker_count
+        },
+    )
+    output = KVConnectorOutput(
+        finished_recving=set(),
+        kv_connector_worker_meta=worker_meta,
+    )
+    cpu_connector.update_connector_output(output)
+
+    req2 = Request(
+        request_id="id-cpu-offload-hit",
+        prompt_token_ids=req1.prompt_token_ids,
+        sampling_params=SamplingParams(max_tokens=16),
+        pooling_params=None,
+        mm_features=None,
+        block_hasher=get_request_block_hasher(BLOCK_SIZE, sha256),
+    )
+
+    hit_tokens, is_async = mc.get_num_new_matched_tokens(req2, num_computed_tokens=0)
+    assert hit_tokens is not None and hit_tokens > 0
+    assert mc._requests_to_connector[req2.request_id] == 1
+    assert is_async is True
+
+
+@pytest.mark.parametrize("swa_enabled", [False, True], ids=["fa_only", "fa_sw"])
+@patch(NIXL_WRAPPER_PATCH, FakeNixlWrapper)
+def test_request_finished_no_async_save(swa_enabled: bool):
+    """A normal request (no P/D) produces no async save from either connector.
+    MultiConnector returns (False, None) via both request_finished and
+    request_finished_all_groups, and cleans up _requests_to_connector."""
+    from vllm.v1.request import RequestStatus
+
+    vllm_config, kv_cache_config = _multi_connector_config(swa_enabled=swa_enabled)
+    scheduler = create_scheduler(vllm_config, kv_cache_config=kv_cache_config)
+    mc = scheduler.connector
+    assert isinstance(mc, MultiConnector)
+
+    request = create_request(
+        request_id=40,
+        num_tokens=BLOCK_SIZE * 2,
+        block_size=BLOCK_SIZE,
+    )
+    scheduler.add_request(request)
+    scheduler.schedule()
+    request.status = RequestStatus.FINISHED_STOPPED
+
+    block_ids = (list(range(2)), list(range(2))) if swa_enabled else (list(range(2)),)
+    async_save, txfer_params = mc.request_finished_all_groups(request, block_ids)
+
+    assert async_save is False
+    assert txfer_params is None
+    assert len(mc._extra_async_saves) == 0
+    assert request.request_id not in mc._requests_to_connector

tests/v1/kv_offload/test_tiering_offloading.py

@@ -19,7 +19,9 @@ import torch
 
 from vllm.v1.kv_offload.base import (
     OffloadKey,
+    OffloadPolicy,
     ReqContext,
+    RequestOffloadingContext,
     make_offload_key,
 )
 from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager
@@ -412,6 +414,78 @@ class TestTieringOffloadingManager:
         job_metadata = self.secondary_tier1.submit_store.call_args.args[0]
         assert job_metadata.req_context is ctx
 
+    def test_on_new_request_lifecycle(self, manager_setup):
+        """Policy defaults to BLOCK_LEVEL, escalates when a tier requests it,
+        and is cleaned up on on_request_finished."""
+        # Default: all tiers return BLOCK_LEVEL
+        ctx = ReqContext(req_id="req_policy_lifecycle")
+        result = self.manager.on_new_request(ctx)
+        assert result.policy == OffloadPolicy.BLOCK_LEVEL
+        self.manager.on_request_finished(ctx)
+
+        # Escalate: tier1 requests REQUEST_LEVEL
+        self.secondary_tier1.on_new_request = (
+            lambda req_context: RequestOffloadingContext(
+                policy=OffloadPolicy.REQUEST_LEVEL
+            )
+        )
+
+        ctx = ReqContext(req_id="req_policy_lifecycle_2")
+        result = self.manager.on_new_request(ctx)
+        assert result.policy == OffloadPolicy.REQUEST_LEVEL
+        assert ctx.req_id in self.manager._request_level_tiers
+
+        # Cleanup
+        self.manager.on_request_finished(ctx)
+        assert ctx.req_id not in self.manager._request_level_tiers
+
+    def test_prepare_store_cascades_existing_blocks_to_request_level_tiers(
+        self, manager_setup
+    ):
+        """prepare_store cascades hit blocks to request-level tiers only."""
+        # Store some blocks to primary first
+        existing_blocks = to_keys(range(3))
+        result = self.manager.prepare_store(existing_blocks, _CTX)
+        assert result is not None
+        self.manager.complete_store(existing_blocks, _CTX, success=True)
+        # Drain cascade completions
+        list(self.manager.take_events())
+        list(self.manager.take_events())
+
+        # Make tier1 request-level, tier2 stays block-level
+        self.secondary_tier1.on_new_request = (
+            lambda req_context: RequestOffloadingContext(
+                policy=OffloadPolicy.REQUEST_LEVEL
+            )
+        )
+
+        ctx = ReqContext(req_id="req_cascade")
+        self.manager.on_new_request(ctx)
+
+        # Spy on submit_store
+        self.secondary_tier1.submit_store = MagicMock(
+            wraps=self.secondary_tier1.submit_store
+        )
+        self.secondary_tier2.submit_store = MagicMock(
+            wraps=self.secondary_tier2.submit_store
+        )
+
+        # Call prepare_store with existing + new blocks
+        new_blocks = to_keys(range(3, 5))
+        all_blocks = existing_blocks + new_blocks
+        result = self.manager.prepare_store(all_blocks, ctx)
+        assert result is not None
+        assert set(result.keys_to_store) == set(new_blocks)
+
+        # Only tier1 (request-level) should get existing blocks cascaded now.
+        # New blocks are cascaded to ALL tiers later via complete_store().
+        self.secondary_tier1.submit_store.assert_called_once()
+        job_metadata = self.secondary_tier1.submit_store.call_args.args[0]
+        assert set(job_metadata.keys) == set(existing_blocks)
+
+        # tier2 (block-level) does not get existing blocks here.
+        self.secondary_tier2.submit_store.assert_not_called()
+
 
 class TestTieringOffloadingWithoutSecondaryTiers:
     """Test TieringOffloadingManager with no secondary tiers (backward compat)."""

tests/v1/sample/test_topk_topp_sampler.py

@@ -5,6 +5,7 @@ import torch
 from torch import Generator
 
 from vllm.platforms import current_platform
+from vllm.triton_utils import HAS_TRITON
 from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch
 
 DEVICE_TYPE = current_platform.device_type
@@ -151,7 +152,7 @@ def test_flashinfer_sampler():
 # =============================================================================
 
 
-@pytest.mark.skipif("cpu" in DEVICE_TYPE, reason="CUDA/XPU not available")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available on this platform")
 class TestTritonTopkTopp:
     """Tests for the Triton top-k/top-p kernel."""
 

tests/v1/worker/test_gpu_model_runner.py

@@ -39,6 +39,7 @@ from vllm.v1.kv_cache_interface import (
     KVCacheGroupSpec,
     KVCacheTensor,
 )
+from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
 from vllm.v1.worker.gpu_input_batch import InputBatch
@@ -279,7 +280,8 @@ def test_sample_tokens_receives_pp_sampled_ids_only_on_non_last_rank(
         lambda: SimpleNamespace(world_size=world_size, is_last_rank=is_last_rank),
     )
 
-    assert GPUModelRunner.sample_tokens(runner, None) is None
+    output = GPUModelRunner.sample_tokens(runner, None)
+    assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None)
     assert receive_calls == expected_calls
 
 
@@ -298,7 +300,8 @@ def test_sample_tokens_skips_pp_group_lookup_without_async_scheduling(
         pytest.fail,
     )
 
-    assert GPUModelRunner.sample_tokens(runner, None) is None
+    output = GPUModelRunner.sample_tokens(runner, None)
+    assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None)
 
 
 def test_select_common_block_size_no_valid_option():

tests/v1/worker/test_gpu_model_runner_v2_eplb.py

@@ -7,6 +7,7 @@ from typing import Any
 
 import torch
 
+from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
 from vllm.v1.worker.gpu import eplb_utils as eplb
 from vllm.v1.worker.gpu import model_runner as mrv2
 
@@ -76,7 +77,10 @@ def _make_runner(**overrides: Any) -> Any:
     runner.max_num_reqs = 8
     runner.max_num_tokens = 16
     runner.decode_query_len = 1
-    runner.kv_connector = SimpleNamespace(set_disabled=lambda *_: None)
+    runner.kv_connector = SimpleNamespace(
+        set_disabled=lambda *_: None,
+        post_forward=lambda *_, **__: None,
+    )
     runner.eplb = eplb.EPLBController(runner.parallel_config, runner.device)
     runner.pooling_runner = None
     runner.execute_model_state = None
@@ -183,5 +187,6 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):
         ),
     )
 
-    assert mrv2.GPUModelRunner.sample_tokens(runner, None) is None
+    output = mrv2.GPUModelRunner.sample_tokens(runner, None)
+    assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None)
     assert events == ["postprocess", "eplb"]

vllm/_custom_ops.py

@@ -650,6 +650,51 @@ def gptq_shuffle(q_weight: torch.Tensor, q_perm: torch.Tensor, bit: int) -> None
     torch.ops._C.gptq_shuffle(q_weight, q_perm, bit)
 
 
+def gptq_gemm_rdna3(
+    a: torch.Tensor,
+    b_q_weight: torch.Tensor,
+    b_qzeros: torch.Tensor,
+    b_scales: torch.Tensor,
+    b_g_idx: torch.Tensor,
+    use_v2_format: bool,
+) -> torch.Tensor:
+    return torch.ops._rocm_C.gptq_gemm_rdna3(
+        a, b_q_weight, b_qzeros, b_scales, b_g_idx, use_v2_format
+    )
+
+
+if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3"):
+
+    @register_fake("_rocm_C::gptq_gemm_rdna3")
+    def _gptq_gemm_rdna3_fake(
+        a: torch.Tensor,
+        b_q_weight: torch.Tensor,
+        b_qzeros: torch.Tensor,
+        b_scales: torch.Tensor,
+        b_g_idx: torch.Tensor,
+        use_v2_format: bool,
+    ) -> torch.Tensor:
+        return torch.empty(
+            (a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device
+        )
+
+
+if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3_wmma"):
+
+    @register_fake("_rocm_C::gptq_gemm_rdna3_wmma")
+    def _gptq_gemm_rdna3_wmma_fake(
+        a: torch.Tensor,
+        b_q_weight: torch.Tensor,
+        b_qzeros: torch.Tensor,
+        b_scales: torch.Tensor,
+        b_g_idx: torch.Tensor,
+        use_v2_format: bool,
+    ) -> torch.Tensor:
+        return torch.empty(
+            (a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device
+        )
+
+
 if hasattr(torch.ops._C, "allspark_w8a16_gemm"):
 
     @register_fake("_C::allspark_w8a16_gemm")
@@ -3719,29 +3764,6 @@ def matmul_mxf4_bf16_tn(
     return torch.ops._qutlass_C.matmul_mxf4_bf16_tn(a, b, a_sf, b_sf, alpha)
 
 
-if hasattr(torch.ops._qutlass_C, "matmul_ada_mxf4_bf16_tn"):
-
-    @register_fake("_qutlass_C::matmul_ada_mxf4_bf16_tn")
-    def _fake_matmul_ada_mxf4_bf16_tn(
-        a: torch.Tensor,
-        b: torch.Tensor,
-        a_sf: torch.Tensor,
-        b_sf: torch.Tensor,
-        alpha: torch.Tensor,
-    ):
-        return a.new_empty(*a.shape[:-1], b.shape[0], dtype=torch.bfloat16)
-
-
-def matmul_ada_mxf4_bf16_tn(
-    a: torch.Tensor,
-    b: torch.Tensor,
-    a_sf: torch.Tensor,
-    b_sf: torch.Tensor,
-    alpha: torch.Tensor,
-) -> torch.Tensor:
-    return torch.ops._qutlass_C.matmul_ada_mxf4_bf16_tn(a, b, a_sf, b_sf, alpha)
-
-
 if hasattr(torch.ops._qutlass_C, "fusedQuantizeMxQuest"):
 
     @register_fake("_qutlass_C::fusedQuantizeMxQuest")

vllm/_xpu_ops.py

@@ -114,8 +114,34 @@ def _gdn_attention_core_xpu_impl(
     attn_metadata = attn_metadata_raw[self.prefix]
     assert isinstance(attn_metadata, GDNAttentionMetadata)
 
-    # TODO: xpu does not support speculative decoding yet
-    assert attn_metadata.spec_sequence_masks is None  # type: ignore[attr-defined]
+    num_actual_tokens = attn_metadata.num_actual_tokens
+    num_accepted_tokens = attn_metadata.num_accepted_tokens
+
+    num_prefills = attn_metadata.num_prefills
+    num_decodes = attn_metadata.num_decodes
+    num_spec_decodes = attn_metadata.num_spec_decodes
+
+    has_initial_state = attn_metadata.has_initial_state
+
+    non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc
+    non_spec_token_indx = attn_metadata.non_spec_token_indx
+    non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor  # noqa: E501
+    non_spec_state_indices_tensor = (
+        non_spec_state_indices_tensor.contiguous()
+        if non_spec_state_indices_tensor is not None
+        else None
+    )
+
+    spec_query_start_loc = attn_metadata.spec_query_start_loc
+    spec_token_indx = attn_metadata.spec_token_indx
+    spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor  # noqa: E501
+
+    spec_sequence_masks = attn_metadata.spec_sequence_masks
+    if spec_sequence_masks is not None:
+        if non_spec_token_indx is not None:
+            non_spec_token_indx = non_spec_token_indx.to(torch.int32)
+        if spec_token_indx is not None:
+            spec_token_indx = spec_token_indx.to(torch.int32)
 
     conv_weights = self.conv1d.weight.view(
         self.conv1d.weight.size(0), self.conv1d.weight.size(2)
@@ -137,12 +163,18 @@ def _gdn_attention_core_xpu_impl(
         activation=self.activation,
         A_log=self.A_log,
         dt_bias=self.dt_bias,
-        num_prefills=attn_metadata.num_prefills,  # type: ignore[attr-defined]
-        num_decodes=attn_metadata.num_decodes,  # type: ignore[attr-defined]
-        has_initial_state=attn_metadata.has_initial_state,  # type: ignore[attr-defined]
-        non_spec_query_start_loc=attn_metadata.non_spec_query_start_loc,  # type: ignore[attr-defined]
-        non_spec_state_indices_tensor=attn_metadata.non_spec_state_indices_tensor,  # type: ignore[attr-defined]
-        num_actual_tokens=attn_metadata.num_actual_tokens,  # type: ignore[attr-defined]
+        num_prefills=num_prefills,  # type: ignore[attr-defined]
+        num_decodes=num_decodes,  # type: ignore[attr-defined]
+        num_spec_decodes=num_spec_decodes,  # type: ignore[attr-defined]
+        has_initial_state=has_initial_state,  # type: ignore[attr-defined]
+        non_spec_query_start_loc=non_spec_query_start_loc,  # type: ignore[attr-defined]
+        non_spec_token_indx=non_spec_token_indx,  # type: ignore[attr-defined]
+        non_spec_state_indices_tensor=non_spec_state_indices_tensor,  # type: ignore[attr-defined]
+        spec_query_start_loc=spec_query_start_loc,  # type: ignore[attr-defined]
+        spec_token_indx=spec_token_indx,  # type: ignore[attr-defined]
+        spec_state_indices_tensor=spec_state_indices_tensor,
+        num_accepted_tokens=num_accepted_tokens,  # type: ignore[attr-defined]
+        num_actual_tokens=num_actual_tokens,  # type: ignore[attr-defined]
         tp_size=self.tp_size,
         reorder_input=not self.gqa_interleaved_layout,
     )

vllm/assets/video.py

@@ -7,10 +7,10 @@ from typing import Any, ClassVar, Literal
 
 import numpy as np
 import numpy.typing as npt
-from huggingface_hub import hf_hub_download
 from PIL import Image
 
 from vllm.multimodal.media.audio import load_audio_pyav
+from vllm.transformers_utils.repo_utils import hf_api
 
 from .base import get_cache_dir
 
@@ -27,7 +27,7 @@ def download_video_asset(filename: str) -> str:
     video_path = video_directory / filename
     video_path_str = str(video_path)
     if not video_path.exists():
-        video_path_str = hf_hub_download(
+        video_path_str = hf_api().hf_hub_download(
             repo_id="raushan-testing-hf/videos-test",
             filename=filename,
             repo_type="dataset",

vllm/benchmarks/datasets/datasets.py

@@ -31,7 +31,6 @@ from typing import Any, cast
 
 import numpy as np
 import pybase64 as base64
-from huggingface_hub import snapshot_download
 from PIL import Image
 from typing_extensions import deprecated
 
@@ -46,6 +45,7 @@ from vllm.lora.utils import get_adapter_absolute_path
 from vllm.multimodal.audio import get_audio_duration
 from vllm.multimodal.image import convert_image_mode
 from vllm.tokenizers import TokenizerLike
+from vllm.transformers_utils.repo_utils import hf_api
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.utils.import_utils import PlaceholderModule
 from vllm.utils.mistral import is_mistral_tokenizer
@@ -3334,7 +3334,10 @@ class MMVUDataset(HuggingFaceDataset):
         self._remote_path_root = (
             f"https://huggingface.co/datasets/{self.hf_name}/resolve/main"
         )
-        self._local_path_root = snapshot_download(self.hf_name, repo_type="dataset")
+        self._local_path_root = hf_api().snapshot_download(
+            self.hf_name,
+            repo_type="dataset",
+        )
 
     def sample(
         self,

vllm/compilation/passes/fusion/allreduce_rms_fusion.py

@@ -1101,6 +1101,14 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass):
             return False
         return bool(compile_range.end <= self.max_token_num)
 
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: fx.Graph) -> None:
+        self.matched_count = self.pm_pass.apply(graph)
+        VllmPatternMatcherPass.match_table[self.pass_name] += self.matched_count
+        logger.debug(
+            "%s Replaced %s patterns", self.__class__.__name__, self.matched_count
+        )
+
     def __del__(self) -> None:
         if getattr(self, "disabled", True):
             return

vllm/compilation/passes/fusion/collective_fusion.py

@@ -917,13 +917,13 @@ class AsyncTPPass(VllmFusionPatternMatcherPass):
             AllGatherScaledMMPattern(self.model_dtype, self.device).register(
                 self.pm_pass
             )
-
-            CutlassScaledMMReduceScatterPattern(self.model_dtype, self.device).register(
-                self.pm_pass
-            )
-            AllGatherCutlassScaledMMPattern(self.model_dtype, self.device).register(
-                self.pm_pass
-            )
+            if hasattr(torch.ops._C, "cutlass_scaled_mm"):
+                CutlassScaledMMReduceScatterPattern(
+                    self.model_dtype, self.device
+                ).register(self.pm_pass)
+                AllGatherCutlassScaledMMPattern(self.model_dtype, self.device).register(
+                    self.pm_pass
+                )
             with suppress(ImportError):
                 import vllm.utils.flashinfer  # noqa: F401
             if hasattr(torch.ops.vllm, "bmm_fp8"):

vllm/config/parallel.py

@@ -78,7 +78,7 @@ class EPLBConfig:
     """
     Interval for logging the balancedness.
     """
-    use_async: bool = False
+    use_async: bool = True
     """
     Whether to use non-blocking EPLB.
     """
@@ -267,7 +267,12 @@ class ParallelConfig:
     """num of nodes for multi-node distributed
     inference when distributed_executor_backend is mp."""
     numa_bind: bool = False
-    """Enable NUMA binding for GPU worker subprocesses."""
+    """Enable NUMA binding for GPU worker subprocesses.
+
+    By default, workers are pinned to their GPU's NUMA-local CPUs and
+    memory; on PCT-capable Xeons they also auto-bind to the SKU's
+    PCT priority cores.
+    """
     numa_bind_nodes: list[int] | None = None
     """NUMA node to bind each GPU worker to.
 

vllm/config/speculative.py

@@ -485,7 +485,16 @@ class SpeculativeConfig:
                 {"n_predict": n_predict, "architectures": ["LongCatFlashMTPModel"]}
             )
 
-        if hf_config.model_type == "step3p5":
+        if hf_config.model_type in ("step3p5", "step3p7") or hf_config.architectures[
+            0
+        ] in ("Step3p5ForCausalLM", "Step3p7ForConditionalGeneration"):
+            quantization_config = getattr(hf_config, "quantization_config", None)
+            hf_config = getattr(hf_config, "text_config", hf_config)
+            if (
+                quantization_config is not None
+                and getattr(hf_config, "quantization_config", None) is None
+            ):
+                hf_config.update({"quantization_config": quantization_config})
             hf_config.model_type = "step3p5_mtp"
             n_predict = getattr(hf_config, "num_nextn_predict_layers", 1)
             hf_config.update({"n_predict": n_predict, "architectures": ["Step3p5MTP"]})
@@ -705,22 +714,16 @@ class SpeculativeConfig:
                     MTPModelTypes
                 ):
                     self.method = "mtp"
-                    if self.num_speculative_tokens > 1:
+                    if (
+                        self.num_speculative_tokens > 1
+                        and self.draft_model_config.hf_config.model_type
+                        != "step3p5_mtp"
+                    ):
                         logger.warning(
                             "Enabling num_speculative_tokens > 1 will run "
                             "multiple times of forward on same MTP layer"
                             ",which may result in lower acceptance rate"
                         )
-                elif self.draft_model_config.hf_config.model_type in (
-                    "longcat_flash_mtp"
-                ):
-                    self.method = "longcat_flash_mtp"
-                    if self.num_speculative_tokens > 1:
-                        logger.warning(
-                            "LongCat MTP models only have "
-                            "one layer. Might need some code changes "
-                            "to support multiple layers."
-                        )
                 elif self.method == "draft_model":
                     pass
                 else:
@@ -1056,6 +1059,14 @@ class SpeculativeConfig:
             == "gemma4_mtp"
         )
 
+    def use_step3p5_mtp(self) -> bool:
+        return (
+            self.method == "mtp"
+            and self.draft_model_config is not None
+            and getattr(self.draft_model_config.hf_config, "model_type", None)
+            == "step3p5_mtp"
+        )
+
     def use_eagle(self) -> bool:
         return self.method in ("eagle", "eagle3", "mtp", "dflash")
 

vllm/config/vllm.py

@@ -121,7 +121,13 @@ def enable_act_fusion(cfg: "VllmConfig") -> bool:
 
 
 def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool:
-    """Enable if TP > 1 and Hopper/Blackwell and flashinfer installed."""
+    """Enable if TP > 1, PP == 1, Hopper/Blackwell, and flashinfer installed.
+
+    Gated off for PP > 1: the fused op's GPU-side peer-signal spin-wait
+    assumes byte-identical kernel launches across TP peers, but concurrent
+    independent warmup of multiple TP subgroups lets ranks pick divergent
+    FlashInfer launch configs and deadlock.
+    """
     from vllm.platforms import current_platform
     from vllm.utils.flashinfer import has_flashinfer
 
@@ -134,6 +140,7 @@ def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool:
 
     return (
         cfg.parallel_config.tensor_parallel_size > 1
+        and cfg.parallel_config.pipeline_parallel_size == 1
         and current_platform.is_cuda()
         and has_flashinfer()
         and (

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

@@ -29,7 +29,9 @@ from vllm.v1.kv_offload.base import (
     OffloadingManager,
     OffloadingSpec,
     OffloadKey,
+    OffloadPolicy,
     ReqContext,
+    RequestOffloadingContext,
     get_offload_block_hash,
     make_offload_key,
 )
@@ -92,6 +94,7 @@ class SchedulerOffloadConfig(NamedTuple):
     kv_group_configs: tuple[GroupOffloadConfig, ...]
     block_size_factor: int
     num_workers: int
+    offload_prompt_only: bool
 
     @classmethod
     def from_spec(cls, spec: OffloadingSpec) -> "SchedulerOffloadConfig":
@@ -154,6 +157,7 @@ class SchedulerOffloadConfig(NamedTuple):
                 for idx, gpu_block_size in enumerate(spec.gpu_block_size)
             ),
             block_size_factor=spec.block_size_factor,
+            offload_prompt_only=spec.offload_prompt_only,
         )
 
 
@@ -172,8 +176,9 @@ class RequestGroupState:
 class RequestOffloadState:
     config: SchedulerOffloadConfig
     req: Request
+    req_context: ReqContext
+    offloading_context: RequestOffloadingContext
     group_states: tuple[RequestGroupState, ...] = field(init=False)
-    req_context: ReqContext = field(init=False)
     # upper bound on tokens to offload for this request; None means no cap
     max_offload_tokens: int | None = None
     # number of hits in the GPU cache
@@ -186,10 +191,6 @@ class RequestOffloadState:
         self.group_states = tuple(
             RequestGroupState() for _ in self.config.kv_group_configs
         )
-        self.req_context = ReqContext(
-            req_id=self.req.request_id,
-            kv_transfer_params=self.req.kv_transfer_params,
-        )
         params = self.req.kv_transfer_params
 
         # NOTE: This field is experimental and subject to change in the future.
@@ -248,6 +249,13 @@ class RequestOffloadState:
             )
 
 
+def _create_req_context(req: Request) -> ReqContext:
+    return ReqContext(
+        req_id=req.request_id,
+        kv_transfer_params=req.kv_transfer_params,
+    )
+
+
 class OffloadingConnectorScheduler:
     """Implementation of Scheduler side methods"""
 
@@ -512,6 +520,18 @@ class OffloadingConnectorScheduler:
 
         return num_hit_tokens
 
+    def on_new_request(self, request: Request) -> None:
+        """Called when a new request is added to the scheduler."""
+        req_context = _create_req_context(request)
+        offloading_context = self.manager.on_new_request(req_context)
+        req_status = RequestOffloadState(
+            config=self.config,
+            req=request,
+            req_context=req_context,
+            offloading_context=offloading_context,
+        )
+        self._req_status[request.request_id] = req_status
+
     def get_num_new_matched_tokens(
         self, request: Request, num_computed_tokens: int
     ) -> tuple[int | None, bool]:
@@ -534,24 +554,15 @@ class OffloadingConnectorScheduler:
                 - `True` if tokens will be loaded asynchronously
                   (between scheduler steps).
         """
-        is_new_request = False
-        if req_status := self._req_status.get(request.request_id):
-            # make sure block IDs are cleared
-            for group_state in req_status.group_states:
-                group_state.block_ids.clear()
-        else:
-            is_new_request = True
-            req_status = RequestOffloadState(config=self.config, req=request)
-            self._req_status[request.request_id] = req_status
+        req_status = self._req_status[request.request_id]
+        for group_state in req_status.group_states:
+            group_state.block_ids.clear()
 
         req_status.update_offload_keys()
         req_status.num_locally_computed_tokens = num_computed_tokens
 
         num_hit_tokens = self._lookup(req_status)
-        if is_new_request:
-            req_status.update_num_hit_blocks(
-                num_computed_tokens + (num_hit_tokens or 0)
-            )
+        req_status.update_num_hit_blocks(num_computed_tokens + (num_hit_tokens or 0))
 
         self._touch(req_status)
 
@@ -568,9 +579,6 @@ class OffloadingConnectorScheduler:
         num_locally_computed_tokens = req_status.num_locally_computed_tokens
         num_cached_tokens = num_locally_computed_tokens + num_external_tokens
 
-        params = req_status.req_context.kv_transfer_params
-        do_remote_decode = params is not None and params.get("do_remote_decode")
-
         keys_to_load: list[OffloadKey] = []
         dst_block_ids: list[int] = []
         # per group
@@ -624,10 +632,10 @@ class OffloadingConnectorScheduler:
             group_sizes.append(num_pending_gpu_blocks)
             block_indices.append(num_locally_computed_gpu_blocks)
 
-            if not do_remote_decode:
-                # For P/D prefill requests (do_remote_decode=True), we do
-                # NOT skip saving the hit prefix, as we need to stream the
-                # entire KV cache so a remote decode node can consume it.
+            # Skip prefix-hit blocks for block-level policy; for
+            # request-level, next_stored_block_idx stays at 0 so all
+            # blocks (including hits) are offloaded.
+            if req_status.offloading_context.policy == OffloadPolicy.BLOCK_LEVEL:
                 group_state.next_stored_block_idx = num_blocks
 
         # Fence dst blocks against finished-request pending stores.
@@ -704,6 +712,15 @@ class OffloadingConnectorScheduler:
             if max_offload_tokens is not None:
                 num_offloadable_tokens = min(num_offloadable_tokens, max_offload_tokens)
 
+            # Skip decode-phase blocks: clamp to the prompt length so only
+            # prefill (prompt) blocks become eligible for store. next_stored_idx
+            # never advances past this boundary, so decode blocks are never
+            # queued in this or any later step.
+            if self.config.offload_prompt_only:
+                num_offloadable_tokens = min(
+                    num_offloadable_tokens, req.num_prompt_tokens
+                )
+
             # Filter out blocks skipped due to sliding window attention / SSM
             # or unreachable by the load path's alignment constraints.
             new_offload_keys: list[OffloadKey] = []
@@ -950,6 +967,12 @@ class OffloadingConnectorScheduler:
         # TODO(orozery): possibly kickoff offload for last block
         # which may have been deferred due to async scheduling
         req_status = self._req_status.get(request.request_id)
+
+        req_context = (
+            req_status.req_context if req_status else _create_req_context(request)
+        )
+        self.manager.on_request_finished(req_context)
+
         if req_status is None:
             return False, None
         if not req_status.transfer_jobs:

vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py

@@ -124,6 +124,10 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
             return self.connector_worker.build_connector_worker_meta()
         return None
 
+    def on_new_request(self, request: "Request") -> None:
+        assert self.connector_scheduler is not None
+        self.connector_scheduler.on_new_request(request)
+
     def get_num_new_matched_tokens(
         self, request: "Request", num_computed_tokens: int
     ) -> tuple[int | None, bool]:

vllm/entrypoints/openai/responses/serving.py

@@ -1041,7 +1041,10 @@ class OpenAIServingResponses(OpenAIServing):
 
         # Use parser to extract and create response output items
         if self.parser:
-            parser = self.parser(tokenizer, request.tools)
+            chat_template_kwargs = self._effective_chat_template_kwargs(request)
+            parser = self.parser(
+                tokenizer, request.tools, chat_template_kwargs=chat_template_kwargs
+            )
             return parser.extract_response_outputs(
                 model_output=final_output.text,
                 model_output_token_ids=final_output.token_ids,
@@ -1378,7 +1381,15 @@ class OpenAIServingResponses(OpenAIServing):
         ],
     ) -> AsyncGenerator[StreamingResponsesResponse, None]:
         processor = SimpleStreamingEventProcessor()
-        parser = self.parser(tokenizer, request.tools) if self.parser else None
+        parser = (
+            self.parser(
+                tokenizer,
+                request.tools,
+                chat_template_kwargs=self._effective_chat_template_kwargs(request),
+            )
+            if self.parser
+            else None
+        )
 
         def _get_logprobs(
             output: CompletionOutput,

vllm/lora/utils.py

@@ -4,7 +4,6 @@
 import os
 from typing import TYPE_CHECKING
 
-import huggingface_hub
 import regex as re
 from huggingface_hub.utils import HfHubHTTPError, HFValidationError
 from torch import nn
@@ -37,6 +36,7 @@ from vllm.lora.layers import (
 from vllm.model_executor.layers.fused_moe import FusedMoE
 from vllm.model_executor.layers.linear import LinearBase
 from vllm.model_executor.utils import get_moe_expert_mapping, get_packed_modules_mapping
+from vllm.transformers_utils.repo_utils import hf_api
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -340,7 +340,9 @@ def get_adapter_absolute_path(lora_path: str) -> str:
         error_log = "Error downloading the ModelScope model"
     else:
         # Otherwise, we assume the path is a Hugging Face Hub repo.
-        download_fn = lambda: huggingface_hub.snapshot_download(repo_id=lora_path)
+        download_fn = lambda: hf_api().snapshot_download(
+            repo_id=lora_path,
+        )
         download_exceptions = (HfHubHTTPError, HFValidationError)
         error_log = "Error downloading the HuggingFace model"
 

vllm/model_executor/kernels/linear/__init__.py

@@ -51,6 +51,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.machete import (
 from vllm.model_executor.kernels.linear.mixed_precision.marlin import (
     MarlinLinearKernel,
 )
+from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import (
+    RDNA3W4A16LinearKernel,
+)
 from vllm.model_executor.kernels.linear.mixed_precision.triton_w4a16 import (
     TritonW4A16LinearKernel,
 )
@@ -339,6 +342,7 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
         TritonW4A16LinearKernel,
     ],
     PlatformEnum.ROCM: [
+        RDNA3W4A16LinearKernel,
         TritonW4A16LinearKernel,
         ConchLinearKernel,
         ExllamaLinearKernel,