refactor & simplify

Signed-off-by: Bugen Zhao <i@bugenzhao.com>

.buildkite/ci_config_rocm.yaml

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

.buildkite/hardware_tests/amd.yaml

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

.buildkite/hardware_tests/cpu.yaml

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

.buildkite/intel_jobs/basic_correctness.yaml

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

.buildkite/intel_jobs/misc_intel.yaml

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

.buildkite/intel_jobs/test-intel.yaml

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -37,8 +37,7 @@ function cpu_tests() {
     pytest -x -v -s tests/kernels/test_onednn.py
     pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
     pytest -x -v -s tests/kernels/core/test_cpu_activation.py
-    pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic
-    pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
+    pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic"
 
   # skip tests requiring model downloads if HF_TOKEN is not set
   # due to rate-limits

.buildkite/scripts/install-kv-connectors.sh

@@ -1,39 +0,0 @@
-#!/bin/bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-set -euo pipefail
-
-REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}"
-
-uv pip install --system -r "${REQUIREMENTS_FILE}"
-
-NIXL_METADATA=$(python3 - <<'PY'
-import importlib.metadata as metadata
-
-import torch
-
-cuda_version = torch.version.cuda
-if cuda_version is None:
-    raise SystemExit("torch.version.cuda is not set")
-
-print(cuda_version.split(".", 1)[0], metadata.version("nixl"))
-PY
-)
-read -r CUDA_MAJOR NIXL_VERSION <<<"${NIXL_METADATA}"
-
-# nixl>=1.1.0 can install multiple CUDA wheel variants. Keep only the variant
-# matching this CI image so nixl_ep_cpp links against the available libcudart.
-uv pip uninstall --system nixl-cu12 nixl-cu13 2>/dev/null || true
-uv pip install --system --no-deps "nixl-cu${CUDA_MAJOR}==${NIXL_VERSION}"
-
-python3 - <<'PY'
-import importlib.metadata as metadata
-
-for package_name in ("nixl", "nixl-cu12", "nixl-cu13"):
-    try:
-        version = metadata.version(package_name)
-    except metadata.PackageNotFoundError:
-        version = "not installed"
-    print(f"{package_name}: {version}")
-PY

.buildkite/test-amd.yaml

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

.buildkite/test_areas/attention.yaml

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

.buildkite/test_areas/basic_correctness.yaml

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

.buildkite/test_areas/disaggregated.yaml

@@ -11,7 +11,7 @@ steps:
     - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
     - tests/v1/kv_connector/nixl_integration/
   commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
     - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 - label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs)
   key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus
@@ -22,7 +22,7 @@ steps:
     - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
     - tests/v1/kv_connector/nixl_integration/
   commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
     - FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
 - label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
@@ -34,7 +34,7 @@ steps:
     - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
     - tests/v1/kv_connector/nixl_integration/
   commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
     - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
 - label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
@@ -46,7 +46,7 @@ steps:
     - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
     - tests/v1/kv_connector/nixl_integration/
   commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
     - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
 - label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
@@ -58,7 +58,7 @@ steps:
     - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
     - tests/v1/kv_connector/nixl_integration/
   commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
     - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
 - label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
@@ -73,7 +73,7 @@ steps:
     - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
     - tests/v1/kv_connector/nixl_integration/
   commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
     - bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
 
 - label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
@@ -87,7 +87,7 @@ steps:
     - vllm/v1/worker/kv_connector_model_runner_mixin.py
     - tests/v1/kv_connector/nixl_integration/
   commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
     - bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
 
 - label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
@@ -102,5 +102,5 @@ steps:
     - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
     - tests/v1/kv_connector/nixl_integration/
   commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
-    - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh

.buildkite/test_areas/engine.yaml

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

.buildkite/test_areas/entrypoints.yaml

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

.buildkite/test_areas/kernels.yaml

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

.buildkite/test_areas/lm_eval.yaml

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

.buildkite/test_areas/misc.yaml

@@ -86,7 +86,7 @@ steps:
     - tests/v1/metrics
     - tests/entrypoints/openai/correctness/test_lmeval.py
   commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     # split the test to avoid interference
     - pytest -v -s -m 'not cpu_test' v1/core
@@ -281,7 +281,6 @@ steps:
   - vllm/model_executor/layers/quantization/quark/
   - vllm/multimodal/
   - vllm/outputs.py
-  - vllm/parser/
   - vllm/platforms/
   - vllm/pooling_params.py
   - vllm/ray/

.buildkite/test_areas/model_runner_v2.yaml

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

.buildkite/test_areas/models_language.yaml

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

.buildkite/test_areas/models_multimodal.yaml

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

.buildkite/test_areas/rust_frontend.yaml

@@ -45,19 +45,19 @@ steps:
   - vllm/entrypoints/serve/
   - vllm/v1/engine/
   - tests/utils.py
-  # - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py
+  # - tests/entrypoints/rpc/test_collective_rpc.py
   - tests/entrypoints/serve/disagg/test_serving_tokens.py
   - tests/entrypoints/serve/instrumentator/test_basic.py
   - tests/entrypoints/serve/instrumentator/test_metrics.py
-  # - tests/entrypoints/serve/dev/test_sleep.py
+  # - tests/entrypoints/serve/instrumentator/test_sleep.py
   commands:
   - export VLLM_USE_RUST_FRONTEND=1
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  # - pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
+  # - pytest -v -s entrypoints/rpc/test_collective_rpc.py
   - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
   - pytest -v -s entrypoints/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
   - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
-  # - pytest -v -s entrypoints/serve/dev/test_sleep.py
+  # - pytest -v -s entrypoints/serve/instrumentator/test_sleep.py
 
 - label: Rust Frontend Core Correctness
   timeout_in_minutes: 30

.buildkite/test_areas/spec_decode.yaml

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

.github/CODEOWNERS

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

.github/mergify.yml

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

.github/workflows/add_label_automerge.yml

@@ -10,7 +10,7 @@ jobs:
         runs-on: ubuntu-latest
         steps:
             -   name: Add label
-                uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
+                uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
                 with:
                     script: |
                         github.rest.issues.addLabels({

.github/workflows/issue_autolabel.yml

@@ -14,7 +14,7 @@ jobs:
     steps:
       - name: Label issues based on keywords
         id: label-step
-        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3  # v9.0.0
+        uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd  # v8.0.0
         with:
           script: |
             // Configuration: Add new labels and keywords here
@@ -315,7 +315,7 @@ jobs:
 
       - name: CC users for labeled issues
         if: steps.label-step.outputs.labels_added != '[]'
-        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3  # v9.0.0
+        uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd  # v8.0.0
         with:
           script: |
             // Configuration: Map labels to GitHub users to CC
@@ -392,7 +392,7 @@ jobs:
 
       - name: Request missing ROCm info from issue author
         if: contains(steps.label-step.outputs.labels_added, 'rocm') && contains(toJSON(github.event.issue.labels.*.name), 'bug')
-        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3  # v9.0.0
+        uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd  # v8.0.0
         with:
           script: |
             const body = (context.payload.issue.body || '').toLowerCase();

.github/workflows/new_pr_bot.yml

@@ -12,7 +12,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Update PR description
-        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
+        uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
         with:
           script: |
             const { owner, repo } = context.repo;
@@ -55,7 +55,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Post welcome comment for first-time contributors
-        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
+        uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
         with:
           script: |
             const { owner, repo } = context.repo;

.github/workflows/pre-commit.yml

@@ -20,7 +20,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - name: Check PR label and author merge count
-      uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
+      uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
       with:
         script: |
           const { data: pr } = await github.rest.pulls.get({

.github/workflows/stale.yml

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

.pre-commit-config.yaml

@@ -21,7 +21,7 @@ repos:
   rev: v21.1.2
   hooks:
   - id: clang-format
-    exclude: 'csrc/(moe/topk_softmax_kernels.cu|libtorch_stable/quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*'
+    exclude: 'csrc/(moe/topk_softmax_kernels.cu|quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*'
     types_or: [c++, cuda]
     args: [--style=file, --verbose]
 - repo: https://github.com/DavidAnson/markdownlint-cli2
@@ -148,27 +148,33 @@ repos:
     language: python
     entry: python tools/pre_commit/generate_nightly_torch_test.py
     files: ^requirements/test/cuda\.(in|txt)$
-  - id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
-    name: Run mypy for Python 3.10
-    entry: python tools/pre_commit/mypy.py "3.10"
+  - id: mypy-local
+    name: Run mypy locally for lowest supported Python version
+    entry: python tools/pre_commit/mypy.py 0 "3.10"
+    stages: [pre-commit] # Don't run in CI
     <<: &mypy_common
       language: python
       types_or: [python, pyi]
       require_serial: true
-      additional_dependencies: ["mypy==1.20.2", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
+      additional_dependencies: ["mypy[faster-cache]==1.19.1", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
+  - id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
+    name: Run mypy for Python 3.10
+    entry: python tools/pre_commit/mypy.py 1 "3.10"
+    <<: *mypy_common
+    stages: [manual] # Only run in CI
   - id: mypy-3.11 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
     name: Run mypy for Python 3.11
-    entry: python tools/pre_commit/mypy.py "3.11"
+    entry: python tools/pre_commit/mypy.py 1 "3.11"
     <<: *mypy_common
     stages: [manual] # Only run in CI
   - id: mypy-3.12 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
     name: Run mypy for Python 3.12
-    entry: python tools/pre_commit/mypy.py "3.12"
+    entry: python tools/pre_commit/mypy.py 1 "3.12"
     <<: *mypy_common
     stages: [manual] # Only run in CI
   - id: mypy-3.13 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
     name: Run mypy for Python 3.13
-    entry: python tools/pre_commit/mypy.py "3.13"
+    entry: python tools/pre_commit/mypy.py 1 "3.13"
     <<: *mypy_common
     stages: [manual] # Only run in CI
   - id: shellcheck

AGENTS.md

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

CMakeLists.txt

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

benchmarks/kernels/benchmark_moe.py

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

cmake/cpu_extension.cmake

@@ -369,18 +369,6 @@ else()
     add_compile_definitions(-DVLLM_NUMA_DISABLED)
 endif()
 
-# check if the pytorch wheel ships libopenblas.so.
-set(VLLM_OPENBLAS_LIB "")
-if (NOT ENABLE_X86_ISA)
-    file(GLOB _VLLM_TORCH_OPENBLAS_LIBS
-        "${TORCH_INSTALL_PREFIX}/lib/libopenblas*.so*")
-    # Note: we don't link openblas directly to _C extension, as it's available through libtorch.so 
-    if (_VLLM_TORCH_OPENBLAS_LIBS)
-        list(GET _VLLM_TORCH_OPENBLAS_LIBS 0 VLLM_OPENBLAS_LIB)
-        message(STATUS "CPU OpenBLAS library: ${VLLM_OPENBLAS_LIB}")
-    endif()
-endif()
-
 #
 # Generate CPU attention dispatch header
 #
@@ -399,7 +387,6 @@ endif()
 #
 set(VLLM_EXT_SRC
     "csrc/cpu/activation.cpp"
-    "csrc/cpu/sgl-kernels/fla.cpp"
     "csrc/cpu/utils.cpp"
     "csrc/cpu/spec_decode_utils.cpp"
     "csrc/cpu/layernorm.cpp"
@@ -409,13 +396,6 @@ set(VLLM_EXT_SRC
     "csrc/cpu/cpu_attn.cpp"
     "csrc/cpu/torch_bindings.cpp")
 
-if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64" AND VLLM_RVV_VLEN AND
-        VLLM_RVV_VLEN GREATER 0 AND (RVV_FP16_FOUND OR RVV_BF16_FOUND))
-    set(VLLM_EXT_SRC
-        "csrc/cpu/cpu_wna16.cpp"
-        ${VLLM_EXT_SRC})
-endif()
-
 if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
     set(VLLM_EXT_SRC
         "csrc/cpu/shm.cpp"
@@ -423,12 +403,6 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
         ${VLLM_EXT_SRC})
 endif()
 
-if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)	
-    set(VLLM_EXT_SRC
-        "csrc/cpu/shm.cpp"
-        ${VLLM_EXT_SRC})
-endif()
-
 if(USE_ONEDNN)
     set(VLLM_EXT_SRC
         "csrc/cpu/dnnl_kernels.cpp"
@@ -437,6 +411,7 @@ endif()
 
 if (ENABLE_X86_ISA)
     set(VLLM_EXT_SRC_SGL
+        "csrc/cpu/sgl-kernels/fla.cpp"
         "csrc/cpu/sgl-kernels/conv.cpp"
         "csrc/cpu/sgl-kernels/gemm.cpp"
         "csrc/cpu/sgl-kernels/gemm_int8.cpp"
@@ -448,7 +423,6 @@ if (ENABLE_X86_ISA)
         "csrc/cpu/sgl-kernels/moe_fp8.cpp")
 
     set(VLLM_EXT_SRC_AVX512
-        "csrc/cpu/sgl-kernels/fla.cpp"
         "csrc/cpu/shm.cpp"
         "csrc/cpu/cpu_wna16.cpp"
         "csrc/cpu/cpu_fused_moe.cpp"
@@ -465,7 +439,6 @@ if (ENABLE_X86_ISA)
         "csrc/moe/dynamic_4bit_int_moe_cpu.cpp") 
 
     set(VLLM_EXT_SRC_AVX2
-        "csrc/cpu/sgl-kernels/fla.cpp"
         "csrc/cpu/utils.cpp"
         "csrc/cpu/spec_decode_utils.cpp"
         "csrc/cpu/cpu_attn.cpp"
@@ -539,9 +512,6 @@ else()
         USE_SABI 3
         WITH_SOABI
     )
-    if (VLLM_OPENBLAS_LIB)
-        target_compile_definitions(_C PRIVATE VLLM_HAS_OPENBLAS)
-    endif()
 endif()
 
 message(STATUS "Enabling C extension.")

cmake/external_projects/vllm_flash_attn.cmake

@@ -31,7 +31,7 @@ endif()
 
 if(VLLM_FLASH_ATTN_SRC_DIR)
   FetchContent_Declare(
-          vllm-flash-attn SOURCE_DIR
+          vllm-flash-attn SOURCE_DIR 
           ${VLLM_FLASH_ATTN_SRC_DIR}
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
   )
@@ -39,7 +39,7 @@ else()
   FetchContent_Declare(
           vllm-flash-attn
           GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG dd62dac706b1cf7895bd99b18c6cb7e7e117ee25
+          GIT_TAG bce29425653ec0fbc579d329883030e832d15ada
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

cmake/hipify.py

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

cmake/utils.cmake

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

csrc/cpu/cpu_fused_moe.cpp

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

csrc/cpu/cpu_types_riscv_impl.hpp

@@ -94,10 +94,6 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
       : reg(RVVI(__riscv_vle16_v_f16, LMUL_256)(
             static_cast<const _Float16*>(ptr), VEC_ELEM_NUM)) {};
 
-  explicit FP16Vec16(const c10::Half v)
-      : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _f16, LMUL_256)(
-            RVVI(__riscv_vmv_v_x_u16, LMUL_256)(v.x, VEC_ELEM_NUM))) {};
-
   explicit FP16Vec16(const FP32Vec16& vec);
 
   void save(void* ptr) const {
@@ -169,9 +165,6 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
             reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
 
   explicit BF16Vec16(fixed_bf16x16_t data) : reg(data) {};
-  explicit BF16Vec16(const c10::BFloat16 v)
-      : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _bf16, LMUL_256)(
-            RVVI(__riscv_vmv_v_x_u16, LMUL_256)(v.x, VEC_ELEM_NUM))) {};
   explicit BF16Vec16(const FP32Vec16&);
 
   void save(void* ptr) const {
@@ -297,9 +290,6 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
     }
     reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16);
   }
-  explicit BF16Vec16(const c10::BFloat16 v)
-      : reg_fp32(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(static_cast<float>(v),
-                                                      VEC_ELEM_NUM)) {}
   explicit BF16Vec16(const FP32Vec16&);
   void save(void* ptr) const {
     float tmp[16];
@@ -639,19 +629,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
       : reg(RVVI4(__riscv_vcreate_v_f32, LMUL_256, _f32, LMUL_512)(
             data.reg, data.reg)) {};
   explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
-  explicit FP32Vec16(int64_t value, const FP32Vec16& lut) {
-    const uint64_t q_values = static_cast<uint64_t>(value);
-    auto packed = RVVI(__riscv_vmv_v_x_u64, LMUL_1024)(q_values, VEC_ELEM_NUM);
-    auto lane_ids = RVVI(__riscv_vid_v_u64, LMUL_1024)(VEC_ELEM_NUM);
-    auto shifts =
-        RVVI(__riscv_vsll_vx_u64, LMUL_1024)(lane_ids, 2, VEC_ELEM_NUM);
-    auto shifted =
-        RVVI(__riscv_vsrl_vv_u64, LMUL_1024)(packed, shifts, VEC_ELEM_NUM);
-    auto idx64 =
-        RVVI(__riscv_vand_vx_u64, LMUL_1024)(shifted, 0xF, VEC_ELEM_NUM);
-    auto idx32 = RVVI(__riscv_vnsrl_wx_u32, LMUL_512)(idx64, 0, VEC_ELEM_NUM);
-    reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx32, VEC_ELEM_NUM);
-  }
   explicit FP32Vec16(const FP16Vec16& v);
 
 #ifdef __riscv_zvfbfmin
@@ -664,10 +641,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
   explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
 #endif
 
-  // FP8 stub: dead code on RISC-V (fp8 KV cache is x86-only), needed for
-  // load_b_pair_vec template to compile on all platforms.
-  explicit FP32Vec16(const BF16Vec32&, int) : FP32Vec16() {}
-
   FP32Vec16 operator+(const FP32Vec16& b) const {
     return FP32Vec16(
         RVVI(__riscv_vfadd_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
@@ -918,30 +891,6 @@ inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
   acc = acc.fma(a, b);
 }
 
-template <typename VecT>
-static void interleave_save_16b(const VecT& vec0, const VecT& vec1, void* ptr) {
-  alignas(64) uint16_t values0[VecT::VEC_ELEM_NUM];
-  alignas(64) uint16_t values1[VecT::VEC_ELEM_NUM];
-  vec0.save(values0);
-  vec1.save(values1);
-
-  auto* packed = reinterpret_cast<uint32_t*>(ptr);
-  for (int32_t i = 0; i < VecT::VEC_ELEM_NUM; ++i) {
-    packed[i] = static_cast<uint32_t>(values0[i]) |
-                (static_cast<uint32_t>(values1[i]) << 16);
-  }
-}
-
-static void interleave_save(const FP16Vec16& vec0, const FP16Vec16& vec1,
-                            void* ptr) {
-  interleave_save_16b(vec0, vec1, ptr);
-}
-
-static void interleave_save(const BF16Vec16& vec0, const BF16Vec16& vec1,
-                            void* ptr) {
-  interleave_save_16b(vec0, vec1, ptr);
-}
-
 #ifdef __riscv_zvfbfmin
 template <>
 inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {

csrc/cpu/cpu_types_vsx.hpp

@@ -89,35 +89,6 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
   }
 };
 
-struct FP16Vec16 : public Vec<FP16Vec16> {
-  constexpr static int VEC_ELEM_NUM = 16;
-  ss16x8x2_t reg;
-
-  explicit FP16Vec16(const void* ptr) {
-    reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)ptr);
-    reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
-  }
-
-  explicit FP16Vec16(bool, const void* ptr) : FP16Vec16(ptr) {}
-
-  explicit FP16Vec16(const FP32Vec16&);
-
-  void save(void* ptr) const {
-    vec_xst(reg.val[0], 0, (signed short*)ptr);
-    vec_xst(reg.val[1], 16, (signed short*)ptr);
-  }
-
-  void save(void* ptr, int elem_num) const {
-    int num = std::max(0, std::min(elem_num, VEC_ELEM_NUM));
-    if (num <= 8) {
-      vec_xst_len(reg.val[0], (signed short*)ptr, num * 2);
-    } else {
-      vec_xst(reg.val[0], 0, (signed short*)ptr);
-      vec_xst_len(reg.val[1], (signed short*)ptr + 8, (num - 8) * 2);
-    }
-  }
-};
-
 struct BF16Vec16 : public Vec<BF16Vec16> {
   constexpr static int VEC_ELEM_NUM = 16;
 
@@ -129,8 +100,6 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
     reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
   }
 
-  explicit BF16Vec16(bool, const void* ptr) : BF16Vec16(ptr) {}
-
   explicit BF16Vec16(const FP32Vec16&);
 
   void save(void* ptr) const {
@@ -410,8 +379,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
     reg.val[3] = vec_xl(48, ptr);
   }
 
-  explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {}
-
   explicit FP32Vec16(f32x4x4_t data) : reg(data) {}
 
   explicit FP32Vec16(const FP32Vec16& data) {
@@ -435,7 +402,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
     reg.val[3] = data.reg.val[1];
   }
 
-  explicit FP32Vec16(const FP16Vec16& v);
   explicit FP32Vec16(const BF16Vec16& v) {
     reg.val[0] = (__vector float)vec_mergeh(zero, v.reg.val[0]);
     reg.val[1] = (__vector float)vec_mergel(zero, v.reg.val[0]);
@@ -769,40 +735,6 @@ inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
 #endif
 }
 
-inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
-  alignas(16) float temp_fp32[16];
-  alignas(16) c10::Half temp_fp16[16];
-
-  vec_xst(v.reg.val[0], 0, temp_fp32);
-  vec_xst(v.reg.val[1], 16, temp_fp32);
-  vec_xst(v.reg.val[2], 32, temp_fp32);
-  vec_xst(v.reg.val[3], 48, temp_fp32);
-
-  for (int i = 0; i < 16; i++) {
-    temp_fp16[i] = c10::Half(temp_fp32[i]);
-  }
-
-  reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)temp_fp16);
-  reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)temp_fp16);
-}
-
-inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
-  alignas(16) c10::Half temp_fp16[16];
-  alignas(16) float temp_fp32[16];
-
-  vec_xst(v.reg.val[0], 0, (signed short*)temp_fp16);
-  vec_xst(v.reg.val[1], 16, (signed short*)temp_fp16);
-
-  for (int i = 0; i < 16; i++) {
-    temp_fp32[i] = float(temp_fp16[i]);
-  }
-
-  reg.val[0] = vec_xl(0, temp_fp32);
-  reg.val[1] = vec_xl(16, temp_fp32);
-  reg.val[2] = vec_xl(32, temp_fp32);
-  reg.val[3] = vec_xl(48, temp_fp32);
-}
-
 inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
 #ifdef _ARCH_PWR10
   __vector signed short ret[4];
@@ -862,43 +794,6 @@ inline void prefetch(const void* addr) {
   __asm__ __volatile__("dcbt 0, %0" : : "r"(addr) : "memory");
 }
 
-struct INT8Vec64 {
-  __vector signed char data[4];
-
-  INT8Vec64() = default;
-
-  explicit INT8Vec64(const int8_t* ptr) {
-    data[0] = vec_xl(0, ptr);
-    data[1] = vec_xl(16, ptr);
-    data[2] = vec_xl(32, ptr);
-    data[3] = vec_xl(48, ptr);
-  }
-
-  explicit INT8Vec64(bool, const int8_t* ptr) : INT8Vec64(ptr) {}
-
-  void save(int8_t* ptr) const {
-    vec_xst(data[0], 0, ptr);
-    vec_xst(data[1], 16, ptr);
-    vec_xst(data[2], 32, ptr);
-    vec_xst(data[3], 48, ptr);
-  }
-
-  void save(int8_t* ptr, int elem_num) const {
-    if (elem_num <= 0) return;
-
-    int full_vecs = elem_num / 16;
-    for (int i = 0; i < full_vecs && i < 4; i++) {
-      vec_xst(data[i], i * 16, ptr);
-    }
-
-    int remaining = elem_num % 16;
-    if (remaining > 0 && full_vecs < 4) {
-      vec_xst_len(data[full_vecs], ptr + full_vecs * 16, remaining);
-    }
-  }
-
-  void nt_save(int8_t* ptr) const { save(ptr); }
-};
-}  // namespace vec_op
+};  // namespace vec_op
 
 #endif

csrc/cpu/sgl-kernels/blas_gemm.h

@@ -1,82 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-#include <ATen/native/CPUBlas.h>
-
-// Unlike brgemm, PyTorch does not publicly expose at::native::cpublas::gemm
-// If OpenBLS is available in the PyTorch wheel, we rely on it for fast
-// bf16:bf16->fp32 GEMMs Otherwise, we fall back to PyTorch reference BLAS path.
-#if defined(VLLM_HAS_OPENBLAS)
-extern "C" void sbgemm_(char* transa, char* transb, int* m, int* n, int* k,
-                        float* alpha, const at::BFloat16* a, int* lda,
-                        const at::BFloat16* b, int* ldb, float* beta, float* c,
-                        int* ldc);
-
-extern "C" void sgemm_(char* transa, char* transb, int* m, int* n, int* k,
-                       float* alpha, const float* a, int* lda, const float* b,
-                       int* ldb, float* beta, float* c, int* ldc);
-
-inline char blas_transpose(at::native::TransposeType trans) {
-  switch (trans) {
-    case at::native::TransposeType::NoTranspose:
-      return 'n';
-    case at::native::TransposeType::Transpose:
-      return 't';
-    case at::native::TransposeType::ConjTranspose:
-      return 'c';
-  }
-  return 'n';
-}
-
-inline void blas_gemm(at::native::TransposeType transa,
-                      at::native::TransposeType transb, int64_t m, int64_t n,
-                      int64_t k, float alpha, const at::BFloat16* a,
-                      int64_t lda, const at::BFloat16* b, int64_t ldb,
-                      float beta, float* c, int64_t ldc) {
-  char transa_ = blas_transpose(transa);
-  char transb_ = blas_transpose(transb);
-  int m_ = static_cast<int>(m);
-  int n_ = static_cast<int>(n);
-  int k_ = static_cast<int>(k);
-  int lda_ = static_cast<int>(lda);
-  int ldb_ = static_cast<int>(ldb);
-  int ldc_ = static_cast<int>(ldc);
-  sbgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta,
-          c, &ldc_);
-}
-
-inline void blas_gemm(at::native::TransposeType transa,
-                      at::native::TransposeType transb, int64_t m, int64_t n,
-                      int64_t k, float alpha, const float* a, int64_t lda,
-                      const float* b, int64_t ldb, float beta, float* c,
-                      int64_t ldc) {
-  char transa_ = blas_transpose(transa);
-  char transb_ = blas_transpose(transb);
-  int m_ = static_cast<int>(m);
-  int n_ = static_cast<int>(n);
-  int k_ = static_cast<int>(k);
-  int lda_ = static_cast<int>(lda);
-  int ldb_ = static_cast<int>(ldb);
-  int ldc_ = static_cast<int>(ldc);
-  sgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta,
-         c, &ldc_);
-}
-
-inline void blas_gemm(at::native::TransposeType, at::native::TransposeType,
-                      int64_t, int64_t, int64_t, float, const at::Half*,
-                      int64_t, const at::Half*, int64_t, float, float*,
-                      int64_t) {
-  TORCH_CHECK(false, "CPU OpenBLAS hgemm is not available.");
-}
-#else
-template <typename scalar_t>
-inline void blas_gemm(at::native::TransposeType transa,
-                      at::native::TransposeType transb, int64_t m, int64_t n,
-                      int64_t k, float alpha, const scalar_t* a, int64_t lda,
-                      const scalar_t* b, int64_t ldb, float beta, float* c,
-                      int64_t ldc) {
-  auto gemm = at::native::cpublas::gemm_no_downcast_stub.DEFAULT;
-  gemm(c10::CppTypeToScalarType<scalar_t>::value, transa, transb, m, n, k,
-       at::Scalar(alpha), a, lda, b, ldb, at::Scalar(beta), c, ldc);
-}
-#endif

csrc/cpu/sgl-kernels/fla.cpp

@@ -301,42 +301,25 @@ void chunk_gated_delta_rule_kernel_impl(
       // attn = k_beta @ key.transpose(-1, -2)
       // attn: [B, HV, num_chunk, chunk_size, chunk_size]
       // transpose and pack for key
-      if constexpr (brgemm_supported()) {
-        pack_vnni<scalar_t>(
-            /*    dst */ k_transpose,
-            /*    src */ curr_k_pad,
-            /*     N  */ chunk_size,
-            /*     K  */ qk_head_size,
-            /* ld_src */ qk_head_size,
-            /* ld_dst */ chunk_size);
-        // k_beta @ key.transpose(-1, -2)
-        at::native::cpublas::brgemm(
-            /*     M */ chunk_size,
-            /*     N */ chunk_size,
-            /*     K */ qk_head_size,
-            /*   lda */ qk_head_size,
-            /*   ldb */ chunk_size,
-            /*   ldc */ chunk_size,
-            /* add_C */ false,
-            /*     A */ curr_k_beta,
-            /*     B */ k_transpose,
-            /*     C */ curr_attn);
-      } else {
-        blas_gemm(
-            at::native::TransposeType::Transpose,
-            at::native::TransposeType::NoTranspose,
-            chunk_size,
-            chunk_size,
-            qk_head_size,
-            1.0f,
-            curr_k_pad,
-            qk_head_size,
-            curr_k_beta,
-            qk_head_size,
-            0.0f,
-            curr_attn,
-            chunk_size);
-      }
+      pack_vnni<scalar_t>(
+          /*    dst */ k_transpose,
+          /*    src */ curr_k_pad,
+          /*     N  */ chunk_size,
+          /*     K  */ qk_head_size,
+          /* ld_src */ qk_head_size,
+          /* ld_dst */ chunk_size);
+      // k_beta @ key.transpose(-1, -2)
+      at::native::cpublas::brgemm(
+          /*     M */ chunk_size,
+          /*     N */ chunk_size,
+          /*     K */ qk_head_size,
+          /*   lda */ qk_head_size,
+          /*   ldb */ chunk_size,
+          /*   ldc */ chunk_size,
+          /* add_C */ false,
+          /*     A */ curr_k_beta,
+          /*     B */ k_transpose,
+          /*     C */ curr_attn);
       // attn = attn * decay_mask
       for (int64_t m = 0; m < chunk_size; m++) {
         at::vec::map2<float>(
@@ -430,42 +413,25 @@ void chunk_gated_delta_rule_kernel_impl(
       // k_beta_g = k_beta * g: [B, HV, num_chunk, chunk_size, EK]
       // k_cumdecay: [B, HV, num_chunk, chunk_size, EK]
       // pack for value
-      if constexpr (brgemm_supported()) {
-        pack_vnni2<scalar_t>(
-            /*    dst */ v_pack,
-            /*    src */ curr_v_beta,
-            /*     N  */ chunk_size,
-            /*     K  */ v_head_size,
-            /* ld_src */ v_head_size,
-            /* ld_dst */ v_head_size);
-        // value = attn @ v_beta
-        at::native::cpublas::brgemm(
-            /*     M */ chunk_size,
-            /*     N */ v_head_size,
-            /*     K */ chunk_size,
-            /*   lda */ chunk_size,
-            /*   ldb */ v_head_size,
-            /*   ldc */ v_head_size,
-            /* add_C */ false,
-            /*     A */ curr_attn_reduced,
-            /*     B */ v_pack,
-            /*     C */ curr_value);
-      } else {
-        blas_gemm(
-            at::native::TransposeType::NoTranspose,
-            at::native::TransposeType::NoTranspose,
-            v_head_size,
-            chunk_size,
-            chunk_size,
-            1.0f,
-            curr_v_beta,
-            v_head_size,
-            curr_attn_reduced,
-            chunk_size,
-            0.0f,
-            curr_value,
-            v_head_size);
-      }
+      pack_vnni2<scalar_t>(
+          /*    dst */ v_pack,
+          /*    src */ curr_v_beta,
+          /*     N  */ chunk_size,
+          /*     K  */ v_head_size,
+          /* ld_src */ v_head_size,
+          /* ld_dst */ v_head_size);
+      // value = attn @ v_beta
+      at::native::cpublas::brgemm(
+          /*     M */ chunk_size,
+          /*     N */ v_head_size,
+          /*     K */ chunk_size,
+          /*   lda */ chunk_size,
+          /*   ldb */ v_head_size,
+          /*   ldc */ v_head_size,
+          /* add_C */ false,
+          /*     A */ curr_attn_reduced,
+          /*     B */ v_pack,
+          /*     C */ curr_value);
       // k_beta_g = k_beta * g.exp().unsqueeze(-1)
       for (int64_t j = 0; j < chunk_size; j++) {
         int64_t i = 0;
@@ -479,42 +445,25 @@ void chunk_gated_delta_rule_kernel_impl(
         }
       }
       // pack for k_beta_g
-      if constexpr (brgemm_supported()) {
-        pack_vnni2<scalar_t>(
-            /*    dst */ k_beta_g_pack,
-            /*    src */ k_beta_g,
-            /*     N  */ chunk_size,
-            /*     K  */ qk_head_size,
-            /* ld_src */ qk_head_size,
-            /* ld_dst */ qk_head_size);
-        // k_cumdecay = attn @ k_beta_g
-        at::native::cpublas::brgemm(
-            /*     M */ chunk_size,
-            /*     N */ qk_head_size,
-            /*     K */ chunk_size,
-            /*   lda */ chunk_size,
-            /*   ldb */ qk_head_size,
-            /*   ldc */ qk_head_size,
-            /* add_C */ false,
-            /*     A */ curr_attn_reduced,
-            /*     B */ k_beta_g_pack,
-            /*     C */ k_cumdecay);
-      } else {
-        blas_gemm(
-            at::native::TransposeType::NoTranspose,
-            at::native::TransposeType::NoTranspose,
-            qk_head_size,
-            chunk_size,
-            chunk_size,
-            1.0f,
-            k_beta_g,
-            qk_head_size,
-            curr_attn_reduced,
-            chunk_size,
-            0.0f,
-            k_cumdecay,
-            qk_head_size);
-      }
+      pack_vnni2<scalar_t>(
+          /*    dst */ k_beta_g_pack,
+          /*    src */ k_beta_g,
+          /*     N  */ chunk_size,
+          /*     K  */ qk_head_size,
+          /* ld_src */ qk_head_size,
+          /* ld_dst */ qk_head_size);
+      // k_cumdecay = attn @ k_beta_g
+      at::native::cpublas::brgemm(
+          /*     M */ chunk_size,
+          /*     N */ qk_head_size,
+          /*     K */ chunk_size,
+          /*   lda */ chunk_size,
+          /*   ldb */ qk_head_size,
+          /*   ldc */ qk_head_size,
+          /* add_C */ false,
+          /*     A */ curr_attn_reduced,
+          /*     B */ k_beta_g_pack,
+          /*     C */ k_cumdecay);
       for (int i = 0; i < chunk_size; i++) {
         at::vec::map<scalar_t>(
             [](fVec x) { return x; },
@@ -602,42 +551,25 @@ void chunk_gated_delta_rule_kernel_impl(
 
         // attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0)
         // k_transpose_i = k_i.transpose(-1, -2)
-        if constexpr (brgemm_supported()) {
-          pack_vnni<scalar_t>(
-              /*    dst */ k_transpose_i,
-              /*    src */ k_i,
-              /*     N  */ chunk_size,
-              /*     K  */ qk_head_size,
-              /* ld_src */ qk_head_size,
-              /* ld_dst */ chunk_size);
-          // attn_i = q_i @ k_transpose_i
-          at::native::cpublas::brgemm(
-              /* M */ chunk_size,
-              /* N */ chunk_size,
-              /* K */ qk_head_size,
-              /* lda */ qk_head_size,
-              /* ldb */ chunk_size,
-              /* ldc */ chunk_size,
-              /* add_C */ false,
-              /* A */ q_i,
-              /* B */ k_transpose_i,
-              /* C */ attn_i);
-        } else {
-          blas_gemm(
-              at::native::TransposeType::Transpose,
-              at::native::TransposeType::NoTranspose,
-              chunk_size,
-              chunk_size,
-              qk_head_size,
-              1.0f,
-              k_i,
-              qk_head_size,
-              q_i,
-              qk_head_size,
-              0.0f,
-              attn_i,
-              chunk_size);
-        }
+        pack_vnni<scalar_t>(
+            /*    dst */ k_transpose_i,
+            /*    src */ k_i,
+            /*     N  */ chunk_size,
+            /*     K  */ qk_head_size,
+            /* ld_src */ qk_head_size,
+            /* ld_dst */ chunk_size);
+        // attn_i = q_i @ k_transpose_i
+        at::native::cpublas::brgemm(
+            /* M */ chunk_size,
+            /* N */ chunk_size,
+            /* K */ qk_head_size,
+            /* lda */ qk_head_size,
+            /* ldb */ chunk_size,
+            /* ldc */ chunk_size,
+            /* add_C */ false,
+            /* A */ q_i,
+            /* B */ k_transpose_i,
+            /* C */ attn_i);
         // attn_i = attn_i * decay_mask_i
         for (int64_t m = 0; m < chunk_size; m++) {
           auto attn_i_m = attn_i + m * chunk_size;
@@ -677,45 +609,28 @@ void chunk_gated_delta_rule_kernel_impl(
         }
 
         // pack for curr_last_recurrent_state
-        if constexpr (brgemm_supported()) {
-          pack_vnni2<scalar_t>(
-              /*    dst */ curr_last_recurrent_state_pack_reduced,
-              /*    src */ curr_last_recurrent_state_reduced,
-              /*     N  */ qk_head_size,
-              /*     K  */ v_head_size,
-              /* ld_src */ v_head_size,
-              /* ld_dst */ v_head_size);
+        pack_vnni2<scalar_t>(
+            /*    dst */ curr_last_recurrent_state_pack_reduced,
+            /*    src */ curr_last_recurrent_state_reduced,
+            /*     N  */ qk_head_size,
+            /*     K  */ v_head_size,
+            /* ld_src */ v_head_size,
+            /* ld_dst */ v_head_size);
 
-          // v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV]
-          // k_cumdecay_i: [chunk_size, EK]
-          // curr_last_recurrent_state: [EK, EV]
-          at::native::cpublas::brgemm(
-              /*     M */ chunk_size,
-              /*     N */ v_head_size,
-              /*     K */ qk_head_size,
-              /*   lda */ qk_head_size,
-              /*   ldb */ v_head_size,
-              /*   ldc */ v_head_size,
-              /* add_C */ false,
-              /*     A */ k_cumdecay_i_reduced,
-              /*     B */ curr_last_recurrent_state_pack_reduced,
-              /*     C */ v_prime);
-        } else {
-          blas_gemm(
-              at::native::TransposeType::NoTranspose,
-              at::native::TransposeType::NoTranspose,
-              v_head_size,
-              chunk_size,
-              qk_head_size,
-              1.0f,
-              curr_last_recurrent_state_reduced,
-              v_head_size,
-              k_cumdecay_i_reduced,
-              qk_head_size,
-              0.0f,
-              v_prime,
-              v_head_size);
-        }
+        // v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV]
+        // k_cumdecay_i: [chunk_size, EK]
+        // curr_last_recurrent_state: [EK, EV]
+        at::native::cpublas::brgemm(
+            /*     M */ chunk_size,
+            /*     N */ v_head_size,
+            /*     K */ qk_head_size,
+            /*   lda */ qk_head_size,
+            /*   ldb */ v_head_size,
+            /*   ldc */ v_head_size,
+            /* add_C */ false,
+            /*     A */ k_cumdecay_i_reduced,
+            /*     B */ curr_last_recurrent_state_pack_reduced,
+            /*     C */ v_prime);
 
         // v_new = v_prime = v_i - v_prime
         // v_i: [chunk_size, EV]
@@ -748,75 +663,41 @@ void chunk_gated_delta_rule_kernel_impl(
         }
         // attn_inter = qg @ curr_last_recurrent_state: [chunk_size, EV]
         // curr_last_recurrent_state: [EK, EV]
-        if constexpr (brgemm_supported()) {
-          at::native::cpublas::brgemm(
-              /* M */ chunk_size,
-              /* N */ v_head_size,
-              /* K */ qk_head_size,
-              /* lda */ qk_head_size,
-              /* ldb */ v_head_size,
-              /* ldc */ v_head_size,
-              /* add_C */ false,
-              /* A */ qg,
-              /* B */ curr_last_recurrent_state_pack_reduced,
-              /* C */ attn_inter);
-        } else {
-          blas_gemm(
-              at::native::TransposeType::NoTranspose,
-              at::native::TransposeType::NoTranspose,
-              v_head_size,
-              chunk_size,
-              qk_head_size,
-              1.0f,
-              curr_last_recurrent_state_reduced,
-              v_head_size,
-              qg,
-              qk_head_size,
-              0.0f,
-              attn_inter,
-              v_head_size);
-        }
+        at::native::cpublas::brgemm(
+            /* M */ chunk_size,
+            /* N */ v_head_size,
+            /* K */ qk_head_size,
+            /* lda */ qk_head_size,
+            /* ldb */ v_head_size,
+            /* ldc */ v_head_size,
+            /* add_C */ false,
+            /* A */ qg,
+            /* B */ curr_last_recurrent_state_pack_reduced,
+            /* C */ attn_inter);
 
         // core_attn_out[:, :, i] = attn_inter + attn_i @ v_new
         // pack for v_prime
-        if constexpr (brgemm_supported()) {
-          pack_vnni2<scalar_t>(
-              /*    dst */ v_prime_pack_reduced,
-              /*    src */ v_prime_reduced,
-              /*     N  */ chunk_size,
-              /*     K  */ v_head_size,
-              /* ld_src */ v_head_size,
-              /* ld_dst */ v_head_size);
-          // attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV]
-          // attn_i: [chunk_size, chunk_size]
-          // v_new: [chunk_size, EV]
-          at::native::cpublas::brgemm(
-              /* M */ chunk_size,
-              /* N */ v_head_size,
-              /* K */ chunk_size,
-              /* lda */ chunk_size,
-              /* ldb */ v_head_size,
-              /* ldc */ v_head_size,
-              /* add_C */ true,
-              /* A */ attn_i_reduced,
-              /* B */ v_prime_pack_reduced,
-              /* C */ attn_inter);
-        } else {
-          blas_gemm(
-              at::native::TransposeType::NoTranspose,
-              at::native::TransposeType::NoTranspose,
-              v_head_size,
-              chunk_size,
-              chunk_size,
-              1.0f,
-              v_prime_reduced,
-              v_head_size,
-              attn_i_reduced,
-              chunk_size,
-              1.0f,
-              attn_inter,
-              v_head_size);
-        }
+        pack_vnni2<scalar_t>(
+            /*    dst */ v_prime_pack_reduced,
+            /*    src */ v_prime_reduced,
+            /*     N  */ chunk_size,
+            /*     K  */ v_head_size,
+            /* ld_src */ v_head_size,
+            /* ld_dst */ v_head_size);
+        // attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV]
+        // attn_i: [chunk_size, chunk_size]
+        // v_new: [chunk_size, EV]
+        at::native::cpublas::brgemm(
+            /* M */ chunk_size,
+            /* N */ v_head_size,
+            /* K */ chunk_size,
+            /* lda */ chunk_size,
+            /* ldb */ v_head_size,
+            /* ldc */ v_head_size,
+            /* add_C */ true,
+            /* A */ attn_i_reduced,
+            /* B */ v_prime_pack_reduced,
+            /* C */ attn_inter);
 
         // core_attn_out[:, :, i] = attn_inter
         for (int64_t m = 0; m < chunk_size; m++) {
@@ -881,34 +762,17 @@ void chunk_gated_delta_rule_kernel_impl(
             /* ld_dst */ chunk_size);
         // kgv = kg.transpose(-1, -2) @ v_new
         // v_new: [chunk_size, EV]
-        if constexpr (brgemm_supported()) {
-          at::native::cpublas::brgemm(
-              /* M */ qk_head_size,
-              /* N */ v_head_size,
-              /* K */ chunk_size,
-              /* lda */ chunk_size,
-              /* ldb */ v_head_size,
-              /* ldc */ v_head_size,
-              /* add_C */ false,
-              /* A */ kg_transpose,
-              /* B */ v_prime_pack_reduced,
-              /* C */ kgv);
-        } else {
-          blas_gemm(
-              at::native::TransposeType::NoTranspose,
-              at::native::TransposeType::NoTranspose,
-              v_head_size,
-              qk_head_size,
-              chunk_size,
-              1.0f,
-              v_prime_reduced,
-              v_head_size,
-              kg_transpose,
-              chunk_size,
-              0.0f,
-              kgv,
-              v_head_size);
-        }
+        at::native::cpublas::brgemm(
+            /* M */ qk_head_size,
+            /* N */ v_head_size,
+            /* K */ chunk_size,
+            /* lda */ chunk_size,
+            /* ldb */ v_head_size,
+            /* ldc */ v_head_size,
+            /* add_C */ false,
+            /* A */ kg_transpose,
+            /* B */ v_prime_pack_reduced,
+            /* C */ kgv);
         // last_recurrent_state = 1) + 2)
         for (int64_t m = 0; m < qk_head_size; m++) {
           at::vec::map2<float>(
@@ -1057,8 +921,7 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
       float k_scale = use_qk_l2norm_in_kernel ? qk_scale_buf[k_scale_offset] : 1.0f;
       int64_t v_offset = si * v_strideS + bi * v_strideB + ni * v_strideH;
       int64_t o_offset = ((bi * seq_len + si) * v_num_heads + ni) * v_head_dim;
-      // See: https://github.com/sgl-project/sglang/pull/26634
-      float beta_val = 1 / (1 + std::exp(-b_ptr[bi * v_num_heads + ni]));
+      float beta_val = 1 / (1 + std::exp(-b_ptr[ni]));
       fVec beta_vec = fVec(beta_val);
       int64_t dvi = 0;
       for (; dvi <= v_head_dim - VecSize; dvi += VecSize) {

csrc/cpu/sgl-kernels/gemm.h

@@ -4,12 +4,9 @@
 // clang-format off
 
 #pragma once
-#include "common.h"
-#include "blas_gemm.h"
+#include <ATen/native/CPUBlas.h>
 
-#if defined(__AVX512F__) && defined(__AVX512BF16__) && defined(__AMX_BF16__)
-#define CPU_CAPABILITY_AVX512
-#endif
+#include "common.h"
 
 // amx-bf16
 #define TILE_M 16
@@ -24,39 +21,31 @@ constexpr int block_size_n() {
   return 2 * TILE_N;
 }
 
-constexpr bool brgemm_supported() {
-#if defined(CPU_CAPABILITY_AVX512)
-  return true;
-#else
-  return false;
-#endif
-}
-
 // define threshold using brgemm (intel AMX)
 template <typename T>
 inline bool can_use_brgemm(int M);
 template <>
 inline bool can_use_brgemm<at::BFloat16>(int M) {
-  return brgemm_supported() && M > 4;
+  return M > 4;
 }
 template <>
 inline bool can_use_brgemm<at::Half>(int M) {
-  return brgemm_supported();
+  return true;
 }
 // this requires PyTorch 2.7 or above
 template <>
 inline bool can_use_brgemm<int8_t>(int M) {
-  return brgemm_supported() && M > 4;
+  return M > 4;
 }
 
 template <>
 inline bool can_use_brgemm<uint8_t>(int M) {
-  return brgemm_supported() && M > 4;
+  return M > 4;
 }
 
 template <>
 inline bool can_use_brgemm<at::Float8_e4m3fn>(int M) {
-  return brgemm_supported() && M > 4;
+  return M > 4;
 }
 
 // work around compiler internal error

csrc/cpu/sgl-kernels/vec.h

@@ -11,9 +11,7 @@
 
 #include <ATen/cpu/vec/functional.h>
 #include <ATen/cpu/vec/vec.h>
-#if defined(CPU_CAPABILITY_AVX512)
 #include <immintrin.h>
-#endif
 namespace {
 
 using namespace at::vec;

csrc/cpu/shm.cpp

@@ -5,7 +5,7 @@
 #include <sys/stat.h>
 #include <unistd.h>
 
-#if defined(__aarch64__) || defined(__powerpc64__)
+#ifdef __aarch64__
   #include <atomic>
 #endif
 
@@ -38,7 +38,7 @@ struct KernelVecType<c10::Half> {
 };
 
 struct ThreadSHMContext {
-#if defined(__aarch64__) || defined(__powerpc64__)
+#ifdef __aarch64__
   // memory model is weaker on AArch64, so we use atomic variables for
   // consumer (load-acquire) and producer (store-release) to make sure
   // that a stamp cannot be ready before the corresponding data is ready.
@@ -75,7 +75,7 @@ struct ThreadSHMContext {
     TORCH_CHECK(group_size <= MAX_SHM_RANK_NUM);
     TORCH_CHECK((size_t)this % 64 == 0);
     TORCH_CHECK((size_t)thread_shm_ptr % 64 == 0);
-#if defined(__aarch64__) || defined(__powerpc64__)
+#ifdef __aarch64__
     _curr_thread_stamp[0].store(1, std::memory_order_relaxed);
     _curr_thread_stamp[1].store(1, std::memory_order_relaxed);
     _ready_thread_stamp[0].store(0, std::memory_order_relaxed);
@@ -124,7 +124,7 @@ struct ThreadSHMContext {
   }
 
   char get_curr_stamp(int idx) const {
-#if defined(__aarch64__) || defined(__powerpc64__)
+#ifdef __aarch64__
     return _curr_thread_stamp[idx].load(std::memory_order_acquire);
 #else
     return _curr_thread_stamp[idx];
@@ -132,7 +132,7 @@ struct ThreadSHMContext {
   }
 
   char get_ready_stamp(int idx) const {
-#if defined(__aarch64__) || defined(__powerpc64__)
+#ifdef __aarch64__
     return _ready_thread_stamp[idx].load(std::memory_order_acquire);
 #else
     return _ready_thread_stamp[idx];
@@ -140,7 +140,7 @@ struct ThreadSHMContext {
   }
 
   void next_stamp() {
-#if defined(__aarch64__) || defined(__powerpc64__)
+#ifdef __aarch64__
     _curr_thread_stamp[local_stamp_buffer_idx].fetch_add(
         1, std::memory_order_release);
 #else
@@ -150,7 +150,7 @@ struct ThreadSHMContext {
   }
 
   void commit_ready_stamp() {
-#if defined(__aarch64__) || defined(__powerpc64__)
+#ifdef __aarch64__
     _ready_thread_stamp[local_stamp_buffer_idx].store(
         _curr_thread_stamp[local_stamp_buffer_idx].load(
             std::memory_order_relaxed),
@@ -186,10 +186,8 @@ struct ThreadSHMContext {
         break;
       }
       ++_spinning_count;
-#if defined(__aarch64__)
+#ifdef __aarch64__
       __asm__ __volatile__("yield");
-#elif defined(__powerpc64__)
-      __asm__ __volatile__("or 1,1,1");
 #else
       _mm_pause();
 #endif  // __aarch64__

csrc/cpu/torch_bindings.cpp

@@ -378,8 +378,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 #endif
 
 // SHM CCL
-#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__)) || \
-    defined(__powerpc64__)
+#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__))
   ops.def(
       "init_shm_manager(str name, int group_size, int rank, int thread_num) -> "
       "int",
@@ -448,25 +447,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "bool is_vnni) -> Tensor");
   ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu);
 
-  // Adapted from sglang: casual_conv1d kernels
-  ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
-  ops.impl("causal_conv1d_weight_pack", torch::kCPU,
-           &causal_conv1d_weight_pack);
-  ops.def(
-      "causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "
-      "conv_states, Tensor? query_start_loc,"
-      "Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "
-      "int pad_slot_id, bool is_vnni) -> "
-      "Tensor");
-  ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
-  ops.def(
-      "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
-      "weight, Tensor? bias, bool silu_activation,"
-      "Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
-      "bool is_vnni) -> Tensor");
-  ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
-#endif
-
   // Adapted from sglang: GDN kernels
   ops.def(
       "chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, "
@@ -490,6 +470,25 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "-> (Tensor, Tensor)");
   ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu);
 
+  // Adapted from sglang: casual_conv1d kernels
+  ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
+  ops.impl("causal_conv1d_weight_pack", torch::kCPU,
+           &causal_conv1d_weight_pack);
+  ops.def(
+      "causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "
+      "conv_states, Tensor? query_start_loc,"
+      "Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "
+      "int pad_slot_id, bool is_vnni) -> "
+      "Tensor");
+  ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
+  ops.def(
+      "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
+      "weight, Tensor? bias, bool silu_activation,"
+      "Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
+      "bool is_vnni) -> Tensor");
+  ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
+#endif
+
   // CPU attention kernels
   ops.def(
       "get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, "
@@ -519,7 +518,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("dynamic_per_token_scaled_fp8_quant() -> ()", placeholder_op);
 
   // WNA16
-#if defined(__AVX512F__) || defined(__riscv_v)
+#if defined(__AVX512F__)
   ops.def(
       "cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, "
       "Tensor scales, Tensor? zeros, Tensor? g_idx, Tensor? bias, SymInt "

csrc/cuda_view.cu

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

csrc/custom_all_reduce.cu

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

csrc/cutlass_extensions/common.cpp

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

csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu

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

csrc/libtorch_stable/activation_kernels.cu

@@ -10,11 +10,20 @@
 
 namespace vllm {
 
-template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
+// `alpha` and `beta` are applied to opposite operands:
+//   - alpha lives INSIDE the activation (the activated half): the gated
+//     activation computes act_half * sigmoid(alpha * act_half).
+//   - beta is added to the OTHER (non-activated) half before the multiply.
+// So the result is always ACT(act_half, alpha) * (other_half + beta).
+// Which half is which depends on `act_first` (see below). Defaults
+// alpha=1.0, beta=0.0 reproduce the plain SwiGLU/GeGLU behavior.
+template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&, const float),
           bool act_first, bool HAS_CLAMP>
 __device__ __forceinline__ scalar_t compute(const scalar_t& x,
                                             const scalar_t& y,
-                                            const float limit) {
+                                            const float limit,
+                                            const float alpha,
+                                            const float beta) {
   if constexpr (act_first) {
     scalar_t gate = x;
     scalar_t up = y;
@@ -22,7 +31,9 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
       gate = (scalar_t)fminf((float)gate, limit);
       up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
     }
-    return ACT_FN(gate) * up;
+    // act_first: gate is the activated half -> alpha applies to gate;
+    // beta is added to up (the non-activated half).
+    return ACT_FN(gate, alpha) * (scalar_t)((float)up + beta);
   } else {
     scalar_t gate = x;
     scalar_t up = y;
@@ -30,55 +41,66 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
       gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
       up = (scalar_t)fminf((float)up, limit);
     }
-    return gate * ACT_FN(up);
+    // !act_first: up is the activated half -> alpha applies to up;
+    // beta is added to gate (the non-activated half).
+    return (scalar_t)((float)gate + beta) * ACT_FN(up, alpha);
   }
 }
 
-template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
+template <typename packed_t,
+          packed_t (*PACKED_ACT_FN)(const packed_t&, const float),
           bool act_first, bool HAS_CLAMP>
 __device__ __forceinline__ packed_t packed_compute(const packed_t& x,
                                                    const packed_t& y,
-                                                   const float limit) {
+                                                   const float limit,
+                                                   const float alpha,
+                                                   const float beta) {
   if constexpr (act_first) {
     packed_t gate = x;
     packed_t up = y;
+    float2 u = cast_to_float2(up);
     if constexpr (HAS_CLAMP) {
       float2 g = cast_to_float2(gate);
-      float2 u = cast_to_float2(up);
       g.x = fminf(g.x, limit);
       g.y = fminf(g.y, limit);
       u.x = fmaxf(fminf(u.x, limit), -limit);
       u.y = fmaxf(fminf(u.y, limit), -limit);
       gate = cast_to_packed<packed_t>(g);
-      up = cast_to_packed<packed_t>(u);
     }
-    return packed_mul(PACKED_ACT_FN(gate), up);
+    // act_first: gate is the activated half -> alpha applies to gate;
+    // beta is added to up (the non-activated half).
+    u.x += beta;
+    u.y += beta;
+    return packed_mul(PACKED_ACT_FN(gate, alpha), cast_to_packed<packed_t>(u));
   } else {
     packed_t gate = x;
     packed_t up = y;
+    float2 g = cast_to_float2(gate);
     if constexpr (HAS_CLAMP) {
-      float2 g = cast_to_float2(gate);
       float2 u = cast_to_float2(up);
       g.x = fmaxf(fminf(g.x, limit), -limit);
       g.y = fmaxf(fminf(g.y, limit), -limit);
       u.x = fminf(u.x, limit);
       u.y = fminf(u.y, limit);
-      gate = cast_to_packed<packed_t>(g);
       up = cast_to_packed<packed_t>(u);
     }
-    return packed_mul(gate, PACKED_ACT_FN(up));
+    // !act_first: up is the activated half -> alpha applies to up;
+    // beta is added to gate (the non-activated half).
+    g.x += beta;
+    g.y += beta;
+    return packed_mul(cast_to_packed<packed_t>(g), PACKED_ACT_FN(up, alpha));
   }
 }
 
 // Activation and gating kernel template.
 template <typename scalar_t, typename packed_t,
-          scalar_t (*ACT_FN)(const scalar_t&),
-          packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
-          bool use_vec, bool HAS_CLAMP, bool use_256b = false>
+          scalar_t (*ACT_FN)(const scalar_t&, const float),
+          packed_t (*PACKED_ACT_FN)(const packed_t&, const float),
+          bool act_first, bool use_vec, bool HAS_CLAMP, bool use_256b = false>
 __global__ void act_and_mul_kernel(
     scalar_t* __restrict__ out,          // [..., d]
     const scalar_t* __restrict__ input,  // [..., 2, d]
-    const int d, const float limit) {
+    const int d, const float limit, const float alpha, const float beta) {
   const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
   const scalar_t* y_ptr = x_ptr + d;
   scalar_t* out_ptr = out + blockIdx.x * d;
@@ -105,7 +127,7 @@ __global__ void act_and_mul_kernel(
       for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
         x.elts[j] =
             packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
-                x.elts[j], y.elts[j], limit);
+                x.elts[j], y.elts[j], limit, alpha, beta);
       }
       if constexpr (use_256b) {
         st256(x, &out_vec[i]);
@@ -118,29 +140,34 @@ __global__ void act_and_mul_kernel(
     for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
       const scalar_t x = VLLM_LDG(&x_ptr[idx]);
       const scalar_t y = VLLM_LDG(&y_ptr[idx]);
-      out_ptr[idx] =
-          compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
+      out_ptr[idx] = compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(
+          x, y, limit, alpha, beta);
     }
   }
 }
 
+// Gated activations take an `alpha` argument that scales the sigmoid input
+// (`x * sigmoid(alpha * x)`). alpha defaults to 1.0 at all call sites, which
+// is exactly SiLU; only the clamp path (silu_and_mul_with_clamp) passes a
+// non-default alpha. Activations that do not use alpha simply ignore it.
 template <typename T>
-__device__ __forceinline__ T silu_kernel(const T& x) {
-  // x * sigmoid(x)
-  return (T)(((float)x) / (1.0f + expf((float)-x)));
+__device__ __forceinline__ T silu_kernel(const T& x, const float alpha) {
+  // x * sigmoid(alpha * x)
+  return (T)(((float)x) / (1.0f + expf((float)-x * alpha)));
 }
 
 template <typename packed_t>
-__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val) {
-  // x * sigmoid(x)
+__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val,
+                                                       const float alpha) {
+  // x * sigmoid(alpha * x)
   float2 fval = cast_to_float2(val);
-  fval.x = fval.x / (1.0f + expf(-fval.x));
-  fval.y = fval.y / (1.0f + expf(-fval.y));
+  fval.x = fval.x / (1.0f + expf(-fval.x * alpha));
+  fval.y = fval.y / (1.0f + expf(-fval.y * alpha));
   return cast_to_packed<packed_t>(fval);
 }
 
 template <typename T>
-__device__ __forceinline__ T gelu_kernel(const T& x) {
+__device__ __forceinline__ T gelu_kernel(const T& x, const float /*alpha*/) {
   // Equivalent to PyTorch GELU with 'none' approximation.
   // Refer to:
   // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
@@ -150,7 +177,8 @@ __device__ __forceinline__ T gelu_kernel(const T& x) {
 }
 
 template <typename packed_t>
-__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) {
+__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val,
+                                                       const float /*alpha*/) {
   // Equivalent to PyTorch GELU with 'none' approximation.
   // Refer to:
   // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
@@ -162,7 +190,8 @@ __device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) {
 }
 
 template <typename T>
-__device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
+__device__ __forceinline__ T gelu_tanh_kernel(const T& x,
+                                              const float /*alpha*/) {
   // Equivalent to PyTorch GELU with 'tanh' approximation.
   // Refer to:
   // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
@@ -176,7 +205,7 @@ __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
 
 template <typename packed_t>
 __device__ __forceinline__ packed_t
-packed_gelu_tanh_kernel(const packed_t& val) {
+packed_gelu_tanh_kernel(const packed_t& val, const float /*alpha*/) {
   // Equivalent to PyTorch GELU with 'tanh' approximation.
   // Refer to:
   // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
@@ -202,7 +231,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
 // clamped (max only) and up input is clamped (both sides) before the
 // activation function is applied.
 #define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST,        \
-                                      HAS_CLAMP, LIMIT)                        \
+                                      HAS_CLAMP, LIMIT, ALPHA, BETA)           \
   auto dtype = input.scalar_type();                                            \
   int d = input.size(-1) / 2;                                                  \
   int64_t num_tokens = input.numel() / input.size(-1);                         \
@@ -230,7 +259,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
             ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>(     \
             out.mutable_data_ptr<scalar_t>(),                                  \
-            input.const_data_ptr<scalar_t>(), d, LIMIT);                       \
+            input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA);          \
       });                                                                      \
     } else {                                                                   \
       VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {   \
@@ -240,7 +269,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
             ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>(    \
             out.mutable_data_ptr<scalar_t>(),                                  \
-            input.const_data_ptr<scalar_t>(), d, LIMIT);                       \
+            input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA);          \
       });                                                                      \
     }                                                                          \
   } else {                                                                     \
@@ -252,7 +281,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
           PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
           ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>(            \
           out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),  \
-          d, LIMIT);                                                           \
+          d, LIMIT, ALPHA, BETA);                                              \
     });                                                                        \
   }
 
@@ -260,14 +289,18 @@ void silu_and_mul(torch::stable::Tensor& out,    // [..., d]
                   torch::stable::Tensor& input)  // [..., 2 * d]
 {
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true, false, 0.0f);
+                                true, false, 0.0f, 1.0f, 0.0f);
 }
 
 void silu_and_mul_clamp(torch::stable::Tensor& out,    // [..., d]
                         torch::stable::Tensor& input,  // [..., 2 * d]
-                        double limit) {
+                        double limit, double alpha, double beta) {
+  // out = (gate.clamp(max=limit) * sigmoid(alpha * gate.clamp(max=limit)))
+  //       * (up.clamp(+-limit) + beta)
+  // alpha=1.0, beta=0.0 reduce this to silu(gate) * up.
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true, true, (float)limit);
+                                true, true, (float)limit, (float)alpha,
+                                (float)beta);
 }
 
 void mul_and_silu(torch::stable::Tensor& out,    // [..., d]
@@ -276,21 +309,22 @@ void mul_and_silu(torch::stable::Tensor& out,    // [..., d]
   // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
   // applies the silu to the latter half of the input.
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                false, false, 0.0f);
+                                false, false, 0.0f, 1.0f, 0.0f);
 }
 
 void gelu_and_mul(torch::stable::Tensor& out,    // [..., d]
                   torch::stable::Tensor& input)  // [..., 2 * d]
 {
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
-                                true, false, 0.0f);
+                                true, false, 0.0f, 1.0f, 0.0f);
 }
 
 void gelu_tanh_and_mul(torch::stable::Tensor& out,    // [..., d]
                        torch::stable::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(
-      vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
+                                vllm::packed_gelu_tanh_kernel, true, false,
+                                0.0f, 1.0f, 0.0f);
 }
 
 namespace vllm {

csrc/libtorch_stable/concat_mla_q.cuh

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

csrc/libtorch_stable/cuda_view.cu

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

csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp

@@ -2,7 +2,7 @@
 
 #include <torch/csrc/stable/tensor.h>
 
-#include "broadcast_load_epilogue_c2x.hpp"
+#include "cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp"
 
 /*
    This file defines custom epilogues for fusing channel scales, token scales,

csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu

@@ -0,0 +1,615 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+ *
+ * Horizontally-fused MiniMax-M3 attention pre-processing kernel.
+ *
+ * Replaces the per-token Python sequence in
+ * ``MiniMaxM3SparseAttention.forward`` / ``MiniMaxM3Attention.forward``:
+ *
+ *     q  = q_norm(q);  k = k_norm(k);  q, k = rotary_emb(pos, q, k)
+ *     index_q = index_q_norm(index_q);  index_k = index_k_norm(index_k)
+ *     index_q, index_k = rotary_emb(pos, index_q, index_k)
+ *     _insert_kv(k, v, index_k)
+ *
+ * All branches share head_dim=128 and the *same* partial-NeoX RoPE table
+ * (``rotary_dim`` rotated, the trailing dims pass through).  The four norms
+ * are Gemma-style RMSNorm (``x * rsqrt(mean(x^2)+eps) * (1 + weight)``) with
+ * independent weights.
+ *
+ * Everything lives in a single fused ``qkv`` tensor.  The sparse layer's
+ * fused projection (MinimaxM3QKVParallelLinearWithIndexer) emits, per token::
+ *
+ *     [ q | k | v | index_q | index_k ]   (the "5 results")
+ *
+ * while the dense layer emits just ``[ q | k | v ]``.  The kernel reads the
+ * index branch straight out of that packed row -- no separate index tensors.
+ *
+ * One kernel, one grid; each warp owns one (token, head-slot) pair.  Slot
+ * enumeration per token:
+ *     [0, nq)                         Q  heads   -> norm(q_w)  + RoPE, write
+ * qkv [nq, nq+nkv)                    K  heads   -> norm(k_w)  + RoPE, write
+ * qkv
+ *                                                  (+ insert into key cache)
+ *     [nq+nkv, nq+2*nkv)              V  heads   -> insert into value cache
+ *     IQ heads (niq)                             -> norm(iq_w) + RoPE, write iq
+ *     IK       (1)                               -> norm(ik_w) + RoPE
+ *                                                  (+ insert into index cache)
+ *
+ * The IQ/IK warps address the index_q/index_k sub-blocks *inside* qkv at the
+ * fixed physical offsets (nq+2*nkv)*128 and (nq+2*nkv+niq)*128.
+ *
+ * Dense vs sparse is a compile-time choice via the ``kIsSparse``/``kInsertKV``
+ * template bools (3 instantiations: dense <false,false>, sparse-profiling
+ * <true,false>, sparse-serving <true,true>), so the index slots, the V slots
+ * and the cache inserts fold away entirely on paths that don't use them. The
+ * dense layer passes no caches/index: norm+RoPE happens in place and the
+ * generic ``Attention`` layer owns the cache write.
+ *
+ * Q/K and (sparse) index_q/index_k are all rewritten in place inside the fused
+ * ``qkv`` tensor.  Caches (bf16) are scatter-written by slot.
+ */
+
+#include <cmath>
+#include <cuda_runtime.h>
+#include <type_traits>
+
+#include "torch_utils.h"
+
+#include "../cuda_compat.h"
+#include "../type_convert.cuh"
+#include "dispatch_utils.h"
+
+#ifndef FINAL_MASK
+  #ifdef USE_ROCM
+    #define FINAL_MASK 0xffffffffffffffffULL
+  #else
+    #define FINAL_MASK 0xffffffffu
+  #endif
+#endif
+
+namespace vllm {
+namespace minimax_m3_fused_ops {
+
+namespace {
+inline int getSMVersion() {
+  auto* props = get_device_prop();
+  return props->major * 10 + props->minor;
+}
+}  // namespace
+
+// ────────────────────────────────────────────────────────────────────────────
+// Constants (hard-coded for MiniMax-M3-preview).
+// ────────────────────────────────────────────────────────────────────────────
+constexpr int kHeadDim = 128;
+constexpr int kNumLanes = 32;
+constexpr int kElemsPerLane = kHeadDim / kNumLanes;  // 4
+
+// ────────────────────────────────────────────────────────────────────────────
+// Helpers
+// ────────────────────────────────────────────────────────────────────────────
+__device__ __forceinline__ float warpReduceSum(float val) {
+#pragma unroll
+  for (int mask = 16; mask > 0; mask >>= 1) {
+    val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
+  }
+  return val;
+}
+
+// Gemma RMSNorm over the full head (no-op when ``weight == nullptr``) followed
+// by partial NeoX RoPE on the leading ``rotary_dim`` dims, all in fp32.  Each
+// lane owns ``kElemsPerLane`` contiguous dims [laneId*4, laneId*4+4).
+template <typename scalar_t>
+__device__ __forceinline__ void normAndRope(
+    float (&elems)[kElemsPerLane], int const laneId, float const eps,
+    scalar_t const* __restrict__ weight,  // [kHeadDim] or nullptr (no norm)
+    bool const do_rope, int const rotary_dim,
+    scalar_t const* __restrict__ cos_ptr,  // cos_sin_cache + pos*rotary_dim
+    bool const apply_norm) {
+  // ── Gemma RMSNorm: x * rsqrt(mean(x^2)+eps) * (1 + w) ──────────────────
+  if (apply_norm) {
+    float sumsq = 0.0f;
+#pragma unroll
+    for (int i = 0; i < kElemsPerLane; i++) sumsq += elems[i] * elems[i];
+    sumsq = warpReduceSum(sumsq);
+    float const rms_rcp = rsqrtf(sumsq / static_cast<float>(kHeadDim) + eps);
+#pragma unroll
+    for (int i = 0; i < kElemsPerLane; i++) {
+      int const dim = laneId * kElemsPerLane + i;
+      float const w = 1.0f + static_cast<float>(weight[dim]);
+      elems[i] = elems[i] * rms_rcp * w;
+    }
+  }
+
+  // ── Partial NeoX RoPE on dims [0, rotary_dim) ──────────────────────────
+  // half = rotary_dim/2.  Pair (i, i+half) for i in [0, half).  Lane L owns
+  // dims [4L, 4L+4); since half is a multiple of 4, a lane lies wholly in the
+  // first half (own=x[i]) or second half (own=x[i+half]); its partner lives
+  // ``half/4`` lanes away (XOR with that distance).
+  if (do_rope) {
+    int const half = rotary_dim / 2;
+    int const dim0 = laneId * kElemsPerLane;
+    bool const in_rope = dim0 < rotary_dim;
+    int const lane_xor = half / kElemsPerLane;  // partner-lane distance
+
+    float partner[kElemsPerLane];
+#pragma unroll
+    for (int i = 0; i < kElemsPerLane; i++) {
+      partner[i] = __shfl_xor_sync(FINAL_MASK, elems[i], lane_xor, 32);
+    }
+    if (in_rope) {
+      bool const first_half = dim0 < half;
+      int const i_base = first_half ? dim0 : (dim0 - half);  // cos/sin index
+      scalar_t const* sin_ptr = cos_ptr + half;
+#pragma unroll
+      for (int i = 0; i < kElemsPerLane; i++) {
+        float const c = static_cast<float>(cos_ptr[i_base + i]);
+        float const s = static_cast<float>(sin_ptr[i_base + i]);
+        if (first_half) {
+          elems[i] = elems[i] * c - partner[i] * s;
+        } else {
+          elems[i] = elems[i] * c + partner[i] * s;
+        }
+      }
+    }
+  }
+}
+
+// Load 4 contiguous bf16 -> 4 fp32 registers.
+template <typename scalar_t>
+__device__ __forceinline__ void loadElems(scalar_t const* __restrict__ src,
+                                          float (&elems)[kElemsPerLane]) {
+  using Converter = vllm::_typeConvert<scalar_t>;
+  uint2 v = *reinterpret_cast<uint2 const*>(src);
+  auto const* p =
+      reinterpret_cast<typename Converter::packed_hip_type const*>(&v);
+#pragma unroll
+  for (int i = 0; i < kElemsPerLane / 2; i++) {
+    float2 f2 = Converter::convert(p[i]);
+    elems[2 * i] = f2.x;
+    elems[2 * i + 1] = f2.y;
+  }
+}
+
+// Store 4 fp32 registers -> 4 contiguous bf16.
+template <typename scalar_t>
+__device__ __forceinline__ void storeElems(
+    scalar_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
+  using Converter = vllm::_typeConvert<scalar_t>;
+  uint2 v;
+  auto* p = reinterpret_cast<typename Converter::packed_hip_type*>(&v);
+#pragma unroll
+  for (int i = 0; i < kElemsPerLane / 2; i++) {
+    p[i] = Converter::convert(make_float2(elems[2 * i], elems[2 * i + 1]));
+  }
+  *reinterpret_cast<uint2*>(dst) = v;
+}
+
+// ────────────────────────────────────────────────────────────────────────────
+// Kernel
+// ────────────────────────────────────────────────────────────────────────────
+// Grid: 1D, ceil(num_tokens * slots_per_token / warps_per_block).
+// Each warp = one (token, slot).
+//
+// `kIsSparse` and `kInsertKV` are compile-time template bools, so all the
+// branch decisions that distinguish the dense layer from the sparse layer
+// (index slots, KV/index inserts, V slots) fold away per instantiation.
+// Three instantiations are built: dense <false,false>, sparse-profiling
+// <true,false> and sparse-serving <true,true>.  Slots per token:
+//     Q : nq                            (always — norm+RoPE)
+//     K : nkv                           (always — norm+RoPE; +K-cache insert)
+//     V : nkv  only if kInsertKV        (V-cache insert; no warps in dense)
+//     IQ: niq  only if kIsSparse        (norm+RoPE)
+//     IK: 1    only if kIsSparse        (norm+RoPE; +index-cache insert)
+template <typename scalar_t, bool kIsSparse, bool kInsertKV>
+__global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
+    scalar_t* __restrict__ qkv,  // [N, qkv_row] in/out (packs index if sparse)
+    scalar_t* __restrict__ q_out,        // [N, nq*128] contiguous, or nullptr
+    scalar_t* __restrict__ index_q_out,  // [N, niq*128] contiguous, or nullptr
+    scalar_t const* __restrict__ q_norm_w,
+    scalar_t const* __restrict__ k_norm_w,
+    scalar_t const* __restrict__ iq_norm_w,
+    scalar_t const* __restrict__ ik_norm_w,
+    scalar_t const* __restrict__ cos_sin_cache,  // [max_pos, rotary_dim]
+    int64_t const* __restrict__ positions,       // [N] i64
+    int64_t const* __restrict__ slot_mapping,    // [N] i64 or nullptr
+    scalar_t* __restrict__ kv_cache,             // [nb,2,bs,nkv,128] or nullptr
+    scalar_t* __restrict__ index_cache,          // [nb*bs, 128] or nullptr
+    float const eps, int const rotary_dim, int const num_tokens, int const nq,
+    int const nkv, int const niq, int const block_size,
+    // kv_cache strides (in elements) for logical shape [nb, 2, bs, nkv, 128].
+    // The head_dim (last) dim is always innermost-contiguous (stride 1), so the
+    // NHD/HND layout choice is fully captured by these four strides: NHD keeps
+    // s_token < s_head, HND swaps them. dim_base addresses head_dim directly.
+    int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token,
+    int64_t const kv_s_head) {
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
+  // _typeConvert<BFloat16> is unavailable on pre-Ampere; the M3 kernel only
+  // runs with bf16/fp16 inputs in practice.  Discard the bf16 body there.
+  if constexpr (std::is_same_v<scalar_t, c10::BFloat16>) {
+    return;
+  } else {
+#endif
+    int const warpsPerBlock = blockDim.x / 32;
+    int const laneId = threadIdx.x % 32;
+    int const globalWarpIdx = blockIdx.x * warpsPerBlock + (threadIdx.x / 32);
+
+    // Slot layout (compile-time gated: dense has neither V nor index slots).
+    int const v_slots = kInsertKV ? nkv : 0;
+    int const idx_slots = kIsSparse ? niq + 1 : 0;
+    int const slots_per_token = nq + nkv + v_slots + idx_slots;
+
+    int const tokenIdx = globalWarpIdx / slots_per_token;
+    int const slot = globalWarpIdx % slots_per_token;
+    if (tokenIdx >= num_tokens) return;
+
+    // Slot boundaries.
+    int const k_begin = nq;
+    int const v_begin = nq + nkv;             // valid only when kInsertKV
+    int const iq_begin = nq + nkv + v_slots;  // index block start
+    int const ik_slot = iq_begin + niq;       // valid only when kIsSparse
+
+    bool const isQ = slot < k_begin;
+    bool const isK = slot >= k_begin && slot < v_begin;
+    bool isV = false;
+    if constexpr (kInsertKV) isV = slot >= v_begin && slot < v_begin + nkv;
+    bool isIQ = false, isIK = false;
+    if constexpr (kIsSparse) {
+      isIQ = slot >= iq_begin && slot < ik_slot;
+      isIK = slot == ik_slot;
+    }
+
+    int const dim_base = laneId * kElemsPerLane;
+    // Physical row width of qkv: the dense layer packs [q|k|v]; the sparse
+    // layer additionally packs [index_q (niq heads) | index_k (1 head)].
+    int const qkv_row = (nq + 2 * nkv + (kIsSparse ? (niq + 1) : 0)) * kHeadDim;
+
+    // ── Resolve source pointer + per-branch parameters. ────────────────────
+    scalar_t* row_ptr = nullptr;       // in-place output location
+    scalar_t const* norm_w = nullptr;  // nullptr -> skip norm (V)
+    bool do_rope = true;
+    int head = 0;  // kv head index for inserts
+
+    if (isQ) {
+      row_ptr =
+          qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
+      norm_w = q_norm_w;
+    } else if (isK) {
+      head = slot - k_begin;
+      row_ptr =
+          qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
+      norm_w = k_norm_w;
+    } else if (isV) {
+      // qkv V section starts at slot index (nq + nkv): slot * kHeadDim is the
+      // correct in-tensor offset.
+      head = slot - v_begin;
+      row_ptr =
+          qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
+      norm_w = nullptr;  // V: no norm, no rope
+      do_rope = false;
+    } else if (isIQ) {
+      // index_q sub-block lives at physical offset (nq+2*nkv)*128 in qkv.
+      int const ih = slot - iq_begin;
+      row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
+                (nq + 2 * nkv + ih) * kHeadDim;
+      norm_w = iq_norm_w;
+    } else {  // isIK -- single shared index key at (nq+2*nkv+niq)*128.
+      row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
+                (nq + 2 * nkv + niq) * kHeadDim;
+      norm_w = ik_norm_w;
+    }
+
+    // Store destination.  Q and index_q are gathered into dedicated contiguous
+    // output buffers (when provided) so the downstream SM100 sparse kernel's
+    // flat TMA descriptor can address them as [tokens*heads, head_dim]; this
+    // folds the de-interleaving into the store the kernel already does, instead
+    // of a separate q.contiguous() copy.  Everything else stays in place.
+    scalar_t* store_ptr = row_ptr;
+    if (isQ && q_out != nullptr) {
+      store_ptr = q_out + static_cast<int64_t>(tokenIdx) * nq * kHeadDim +
+                  slot * kHeadDim;
+    } else if (isIQ && index_q_out != nullptr) {
+      store_ptr = index_q_out +
+                  static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
+                  (slot - iq_begin) * kHeadDim;
+    }
+
+    // PDL: wait for the predecessor kernel (the qkv-projection GEMM that
+    // produces ``qkv``) to finish before touching any global memory.  No-op
+    // when PDL is not enabled on the launch.  The CUDA runtime wrapper emits
+    // the griddepcontrol.wait PTX with the required memory clobber internally.
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    cudaGridDependencySynchronize();
+#endif
+
+    // ── Load -> norm+rope (fp32) -> store back in place. ───────────────────
+    float elems[kElemsPerLane];
+    loadElems<scalar_t>(row_ptr + dim_base, elems);
+
+    if (!isV) {
+      int64_t const pos = positions[tokenIdx];
+      scalar_t const* cos_ptr = cos_sin_cache + pos * rotary_dim;
+      normAndRope<scalar_t>(elems, laneId, eps, norm_w, do_rope, rotary_dim,
+                            cos_ptr, /*apply_norm=*/norm_w != nullptr);
+      storeElems<scalar_t>(store_ptr + dim_base, elems);
+    }
+
+    // ── Cache inserts (sparse serving only). ───────────────────────────────
+    if constexpr (kInsertKV) {
+      // Guard (not early-return) so every thread reaches the PDL trigger below.
+      int64_t const sm = (isK || isV || isIK) ? slot_mapping[tokenIdx] : -1;
+      if (sm >= 0) {  // skip padded / unscheduled tokens
+        if (isIK) {
+          scalar_t* dst = index_cache + sm * kHeadDim + dim_base;
+          storeElems<scalar_t>(dst, elems);
+        } else if (isK || isV) {
+          // kv_cache logical shape [num_blocks, 2, block_size, nkv, head_dim].
+          // Paging is logical (block = sm/block_size, token = sm%block_size);
+          // the physical NHD/HND layout is honoured via the passed strides.
+          int64_t const b = sm / block_size;
+          int64_t const t = sm % block_size;
+          int const kv = isK ? 0 : 1;
+          int64_t const off =
+              b * kv_s_block + kv * kv_s_kv + t * kv_s_token + head * kv_s_head;
+          storeElems<scalar_t>(kv_cache + off + dim_base, elems);
+        }
+      }
+    }
+
+    // PDL: signal that this kernel is done so a dependent successor may launch
+    // early.  No-op when PDL is not enabled on the launch.
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    cudaTriggerProgrammaticLaunchCompletion();
+#endif
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
+  }
+#endif
+}
+
+// ────────────────────────────────────────────────────────────────────────────
+// Launch wrapper
+// ────────────────────────────────────────────────────────────────────────────
+template <typename scalar_t>
+void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
+                          scalar_t const* q_norm_w, scalar_t const* k_norm_w,
+                          scalar_t const* iq_norm_w, scalar_t const* ik_norm_w,
+                          scalar_t const* cos_sin_cache,
+                          int64_t const* positions, int64_t const* slot_mapping,
+                          scalar_t* kv_cache, scalar_t* index_cache,
+                          float const eps, int const rotary_dim,
+                          int const num_tokens, int const nq, int const nkv,
+                          int const niq, int const block_size,
+                          int64_t const kv_s_block, int64_t const kv_s_kv,
+                          int64_t const kv_s_token, int64_t const kv_s_head,
+                          bool const has_index, bool const insert_kv,
+                          cudaStream_t stream) {
+  // Slot count must match the kernel's compile-time gating.
+  int const v_slots = insert_kv ? nkv : 0;
+  int const idx_slots = has_index ? niq + 1 : 0;
+  int const slots_per_token = nq + nkv + v_slots + idx_slots;
+
+  constexpr int kBlockSize = 256;
+  constexpr int kWarpsPerBlock = kBlockSize / 32;
+  int64_t const total_warps =
+      static_cast<int64_t>(num_tokens) * slots_per_token;
+  int const grid =
+      static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
+  if (grid == 0) return;
+
+#ifndef USE_ROCM
+  // PDL: enable programmatic stream serialization whenever the hardware
+  // supports it (SM90+).  On pre-Hopper GPUs the attribute is unavailable, so
+  // leave numAttrs = 0 and launch as a regular kernel via cudaLaunchKernelEx.
+  static int const sm_version = getSMVersion();
+  cudaLaunchConfig_t config;
+  config.gridDim = dim3(grid);
+  config.blockDim = dim3(kBlockSize);
+  config.dynamicSmemBytes = 0;
+  config.stream = stream;
+  cudaLaunchAttribute attrs[1];
+  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attrs[0].val.programmaticStreamSerializationAllowed = 1;
+  config.attrs = attrs;
+  config.numAttrs = (sm_version >= 90) ? 1 : 0;
+
+  #define LAUNCH(IS_SPARSE, INSERT)                                            \
+    cudaLaunchKernelEx(                                                        \
+        &config,                                                               \
+        fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, IS_SPARSE, INSERT>,    \
+        qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, ik_norm_w,     \
+        cos_sin_cache, positions, slot_mapping, kv_cache, index_cache, eps,    \
+        rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_kv, \
+        kv_s_token, kv_s_head)
+#else
+  // ROCm: standard kernel launch syntax (no PDL/stream serialization).
+  // clang-format off
+  #define LAUNCH(IS_SPARSE, INSERT)                                          \
+    fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, IS_SPARSE, INSERT>       \
+        <<<grid, kBlockSize, 0, stream>>>(                                   \
+            qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w,          \
+            ik_norm_w, cos_sin_cache, positions, slot_mapping, kv_cache,     \
+            index_cache, eps, rotary_dim, num_tokens, nq, nkv, niq,          \
+            block_size, kv_s_block, kv_s_kv, kv_s_token, kv_s_head)
+  // clang-format on
+#endif
+
+  if (has_index) {
+    if (insert_kv) {
+      LAUNCH(true, true);  // sparse serving
+    } else {
+      LAUNCH(true, false);  // sparse profiling
+    }
+  } else {
+    // Dense layer: never has an index branch and never inserts here (the
+    // generic Attention layer owns the KV insert).
+    LAUNCH(false, false);
+  }
+#undef LAUNCH
+}
+
+}  // namespace minimax_m3_fused_ops
+}  // namespace vllm
+
+// ────────────────────────────────────────────────────────────────────────────
+// Torch op wrapper
+// ────────────────────────────────────────────────────────────────────────────
+void fused_minimax_m3_qknorm_rope_kv_insert(
+    torch::stable::Tensor& qkv,  // [N, qkv_row] (packs index if sparse)
+    torch::stable::Tensor const& q_norm_weight,  // [128]
+    torch::stable::Tensor const& k_norm_weight,  // [128]
+    torch::stable::Tensor const& cos_sin_cache,  // [max_pos, rotary_dim]
+    torch::stable::Tensor const& positions,      // [N] i64
+    int64_t num_heads, int64_t num_kv_heads, int64_t rotary_dim, double eps,
+    std::optional<torch::stable::Tensor> index_q_norm_weight,  // [128]
+    std::optional<torch::stable::Tensor> index_k_norm_weight,  // [128]
+    int64_t num_index_heads,                            // niq; 0 => dense
+    std::optional<torch::stable::Tensor> slot_mapping,  // [N] i64
+    std::optional<torch::stable::Tensor> kv_cache,      // [nb,2,bs,nkv,128]
+    std::optional<torch::stable::Tensor> index_cache,   // [nb,bs,128]
+    int64_t block_size,
+    std::optional<torch::stable::Tensor> q_out,  // [N, nq*128] contiguous
+    std::optional<torch::stable::Tensor>
+        index_q_out) {  // [N, niq*128] contiguous
+  STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(),
+                  "qkv must be contiguous CUDA");
+  STD_TORCH_CHECK(
+      positions.is_cuda() &&
+          positions.scalar_type() == torch::headeronly::ScalarType::Long,
+      "positions must be int64 CUDA");
+  STD_TORCH_CHECK(cos_sin_cache.is_cuda() && cos_sin_cache.is_contiguous(),
+                  "cos_sin_cache must be contiguous CUDA");
+  STD_TORCH_CHECK(cos_sin_cache.scalar_type() == qkv.scalar_type(),
+                  "cos_sin_cache dtype must match qkv");
+  STD_TORCH_CHECK(
+      cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == rotary_dim,
+      "cos_sin_cache shape [max_pos, rotary_dim]");
+
+  STD_TORCH_CHECK(q_norm_weight.scalar_type() == qkv.scalar_type() &&
+                      k_norm_weight.scalar_type() == qkv.scalar_type(),
+                  "q/k norm weight dtype must match qkv");
+  STD_TORCH_CHECK(
+      q_norm_weight.numel() == vllm::minimax_m3_fused_ops::kHeadDim &&
+          k_norm_weight.numel() == vllm::minimax_m3_fused_ops::kHeadDim,
+      "q/k norm weight must have 128 elements");
+  STD_TORCH_CHECK(rotary_dim > 0 && rotary_dim % 8 == 0 &&
+                      rotary_dim <= vllm::minimax_m3_fused_ops::kHeadDim,
+                  "rotary_dim must be a positive multiple of 8 and <= 128");
+
+  int const num_tokens = static_cast<int>(qkv.size(0));
+  int const nq = static_cast<int>(num_heads);
+  int const nkv = static_cast<int>(num_kv_heads);
+  int const niq = static_cast<int>(num_index_heads);
+
+  // The sparse layer packs the index branch ([index_q (niq heads) | index_k
+  // (1 head)]) right after [q|k|v] in the same row; the dense layer does not.
+  bool const has_index = niq > 0;
+  bool const insert_kv = kv_cache.has_value();
+  int const kHeadDim = vllm::minimax_m3_fused_ops::kHeadDim;
+  int const expected_row =
+      (nq + 2 * nkv + (has_index ? niq + 1 : 0)) * kHeadDim;
+  STD_TORCH_CHECK(qkv.size(1) == expected_row,
+                  "qkv last dim must be (num_heads + 2*num_kv_heads"
+                  " + num_index_heads + 1) * 128 for sparse, "
+                  "(num_heads + 2*num_kv_heads) * 128 for dense");
+
+  // Only the sparse layer inserts here (dense lets the generic Attention layer
+  // own the KV write); there is no dense+insert kernel instantiation.
+  STD_TORCH_CHECK(
+      !insert_kv || has_index,
+      "insert mode (kv_cache) requires the index branch (sparse layer)");
+  if (has_index) {
+    STD_TORCH_CHECK(
+        index_q_norm_weight.has_value() && index_k_norm_weight.has_value(),
+        "index branch requires both index norm weights");
+    STD_TORCH_CHECK(index_q_norm_weight->scalar_type() == qkv.scalar_type() &&
+                        index_k_norm_weight->scalar_type() == qkv.scalar_type(),
+                    "index norm weights dtype must match qkv");
+    STD_TORCH_CHECK(index_q_norm_weight->numel() == kHeadDim &&
+                        index_k_norm_weight->numel() == kHeadDim,
+                    "index norm weights must have 128 elements");
+  }
+  // kv_cache strides (logical shape [nb, 2, bs, nkv, head_dim]). Read straight
+  // off the tensor so the kernel honours whatever physical layout the attention
+  // backend allocated (NHD: stride order (0,1,2,3,4); HND: (0,1,3,2,4)). No new
+  // op argument is needed -- the strides ride along with the tensor itself.
+  int64_t kv_s_block = 0, kv_s_kv = 0, kv_s_token = 0, kv_s_head = 0;
+  if (insert_kv) {
+    STD_TORCH_CHECK(
+        slot_mapping.has_value() &&
+            slot_mapping->scalar_type() == torch::headeronly::ScalarType::Long,
+        "insert mode requires int64 slot_mapping");
+    STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
+                    "kv_cache dtype must match qkv (bf16 cache only)");
+    STD_TORCH_CHECK(index_cache.has_value() &&
+                        index_cache->scalar_type() == qkv.scalar_type(),
+                    "insert mode requires matching index_cache");
+    STD_TORCH_CHECK(kv_cache->dim() == 5 && kv_cache->stride(4) == 1,
+                    "kv_cache must be [nb,2,bs,nkv,head_dim] with contiguous "
+                    "head_dim (stride(4)==1)");
+    kv_s_block = kv_cache->stride(0);
+    kv_s_kv = kv_cache->stride(1);
+    kv_s_token = kv_cache->stride(2);
+    kv_s_head = kv_cache->stride(3);
+  }
+  // Optional contiguous gather targets: when given, the normed/roped q (and
+  // index_q) are written here instead of in place, so callers avoid a separate
+  // .contiguous() copy.  index_q_out only makes sense on the sparse path.
+  if (q_out.has_value()) {
+    STD_TORCH_CHECK(
+        q_out->is_cuda() && q_out->is_contiguous() &&
+            q_out->scalar_type() == qkv.scalar_type(),
+        "q_out must be a contiguous CUDA tensor matching qkv dtype");
+    STD_TORCH_CHECK(
+        q_out->numel() == static_cast<int64_t>(num_tokens) * nq * kHeadDim,
+        "q_out must have num_tokens * num_heads * 128 elements");
+  }
+  if (index_q_out.has_value()) {
+    STD_TORCH_CHECK(
+        has_index,
+        "index_q_out requires the index branch (num_index_heads > 0)");
+    STD_TORCH_CHECK(
+        index_q_out->is_cuda() && index_q_out->is_contiguous() &&
+            index_q_out->scalar_type() == qkv.scalar_type(),
+        "index_q_out must be a contiguous CUDA tensor matching qkv dtype");
+    STD_TORCH_CHECK(index_q_out->numel() ==
+                        static_cast<int64_t>(num_tokens) * niq * kHeadDim,
+                    "index_q_out must have num_tokens * num_index_heads * 128 "
+                    "elements");
+  }
+
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      qkv.get_device_index());
+  auto stream = get_current_cuda_stream(qkv.get_device_index());
+
+  VLLM_STABLE_DISPATCH_HALF_TYPES(
+      qkv.scalar_type(), "fused_minimax_m3_qknorm_rope_kv_insert", [&] {
+        using st = scalar_t;
+        vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st>(
+            reinterpret_cast<st*>(qkv.data_ptr()),
+            q_out.has_value() ? reinterpret_cast<st*>(q_out->data_ptr())
+                              : nullptr,
+            index_q_out.has_value()
+                ? reinterpret_cast<st*>(index_q_out->data_ptr())
+                : nullptr,
+            reinterpret_cast<st const*>(q_norm_weight.data_ptr()),
+            reinterpret_cast<st const*>(k_norm_weight.data_ptr()),
+            has_index
+                ? reinterpret_cast<st const*>(index_q_norm_weight->data_ptr())
+                : nullptr,
+            has_index
+                ? reinterpret_cast<st const*>(index_k_norm_weight->data_ptr())
+                : nullptr,
+            reinterpret_cast<st const*>(cos_sin_cache.data_ptr()),
+            reinterpret_cast<int64_t const*>(positions.data_ptr()),
+            insert_kv
+                ? reinterpret_cast<int64_t const*>(slot_mapping->data_ptr())
+                : nullptr,
+            insert_kv ? reinterpret_cast<st*>(kv_cache->data_ptr()) : nullptr,
+            (insert_kv && has_index)
+                ? reinterpret_cast<st*>(index_cache->data_ptr())
+                : nullptr,
+            static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens,
+            nq, nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_kv,
+            kv_s_token, kv_s_head, has_index, insert_kv, stream);
+      });
+}

csrc/libtorch_stable/fused_qknorm_rope_kernel.cu

@@ -20,7 +20,7 @@
 
 #include "torch_utils.h"
 
-#include "async_util.cuh"
+#include "../async_util.cuh"
 #include "../cuda_compat.h"
 #include "../type_convert.cuh"
 #include "dispatch_utils.h"

csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu

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

csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu

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

csrc/libtorch_stable/ops.h

@@ -3,6 +3,10 @@
 #include <torch/csrc/stable/library.h>
 #include <torch/csrc/stable/tensor.h>
 
+#ifndef USE_ROCM
+torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
+                                   torch::stable::Tensor const& perm);
+
 void per_token_group_quant_fp8(const torch::stable::Tensor& input,
                                torch::stable::Tensor& output_q,
                                torch::stable::Tensor& output_s,
@@ -24,10 +28,6 @@ void per_token_group_quant_int8(const torch::stable::Tensor& input,
                                 int64_t group_size, double eps, double int8_min,
                                 double int8_max);
 
-torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
-                                   torch::stable::Tensor const& perm);
-
-#ifndef USE_ROCM
 bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
 bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
 bool cutlass_group_gemm_supported(int64_t cuda_device_capability);
@@ -162,11 +162,6 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
 // AllSpark ops: declarations are in the source files
 // (allspark_repack.cu and allspark_qgemm_w8a16.cu)
 
-// TODO: Move this out once ROCm upgrade their torch to 2.11.
-// CPU tensor -> CUDA UVA view (shared CUDA)
-torch::stable::Tensor get_cuda_view_from_cpu_tensor(
-    torch::stable::Tensor& cpu_tensor);
-
 #endif
 
 // Attention kernels (shared CUDA/ROCm)
@@ -220,13 +215,6 @@ void rms_norm_per_block_quant(torch::stable::Tensor& out,
                               std::optional<torch::stable::Tensor> residual,
                               int64_t group_size, bool is_scale_transposed);
 
-void silu_and_mul_per_block_quant(torch::stable::Tensor& out,
-                                  torch::stable::Tensor const& input,
-                                  torch::stable::Tensor& scales,
-                                  int64_t group_size,
-                                  std::optional<torch::stable::Tensor> scale_ub,
-                                  bool is_scale_transposed);
-
 // Positional encoding kernels (shared CUDA/ROCm)
 void rotary_embedding(torch::stable::Tensor& positions,
                       torch::stable::Tensor& query,
@@ -243,43 +231,22 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q,
                         torch::stable::Tensor& position_ids,
                         int64_t forced_token_heads_per_warp);
 
-torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv,
-    torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
-    torch::stable::Tensor const& position_ids,
-    torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
-    double eps, int64_t cache_block_size);
-
-void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
-    torch::stable::Tensor& q, torch::stable::Tensor const& kv,
-    torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
-    torch::stable::Tensor const& position_ids,
-    torch::stable::Tensor const& cos_sin_cache, double eps,
-    int64_t cache_block_size);
-
-void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
-    torch::stable::Tensor const& q, torch::stable::Tensor const& kv,
-    torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache,
-    torch::stable::Tensor const& slot_mapping,
-    torch::stable::Tensor const& position_ids,
+// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV /
+// index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs
+// the index branch and scatters k/v/index_k into their paged caches.
+void fused_minimax_m3_qknorm_rope_kv_insert(
+    torch::stable::Tensor& qkv, torch::stable::Tensor const& q_norm_weight,
+    torch::stable::Tensor const& k_norm_weight,
     torch::stable::Tensor const& cos_sin_cache,
-    torch::stable::Tensor const& fp8_scale,
-    torch::stable::Tensor const& q_fp8_scale_inv, double eps,
-    int64_t cache_block_size);
-
-#ifndef USE_ROCM
-torch::stable::Tensor minimax_allreduce_rms(
-    torch::stable::Tensor const& input,
-    torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace,
-    int64_t const rank, int64_t const nranks, double const eps);
-std::tuple<torch::stable::Tensor, torch::stable::Tensor>
-minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
-                         torch::stable::Tensor const& norm_weight_q,
-                         torch::stable::Tensor const& norm_weight_k,
-                         torch::stable::Tensor workspace, int64_t const q_size,
-                         int64_t const kv_size, int64_t const rank,
-                         int64_t const nranks, double const eps);
-#endif
+    torch::stable::Tensor const& positions, int64_t num_heads,
+    int64_t num_kv_heads, int64_t rotary_dim, double eps,
+    std::optional<torch::stable::Tensor> index_q_norm_weight,
+    std::optional<torch::stable::Tensor> index_k_norm_weight,
+    int64_t num_index_heads, std::optional<torch::stable::Tensor> slot_mapping,
+    std::optional<torch::stable::Tensor> kv_cache,
+    std::optional<torch::stable::Tensor> index_cache, int64_t block_size,
+    std::optional<torch::stable::Tensor> q_out,
+    std::optional<torch::stable::Tensor> index_q_out);
 
 // Sampler kernels (shared CUDA/ROCm)
 void apply_repetition_penalties_(
@@ -323,30 +290,11 @@ void selective_scan_fwd(
     const std::optional<torch::stable::Tensor>& cu_chunk_seqlen,
     const std::optional<torch::stable::Tensor>& last_chunk_indices);
 
-using fptr_t = int64_t;
-fptr_t init_custom_ar(const std::vector<int64_t>& fake_ipc_ptrs,
-                      torch::stable::Tensor& rank_data, int64_t rank,
-                      bool fully_connected);
-void all_reduce(fptr_t _fa, torch::stable::Tensor& inp,
-                torch::stable::Tensor& out, fptr_t reg_buffer,
-                int64_t reg_buffer_sz_bytes);
-void dispose(fptr_t _fa);
-int64_t meta_size();
-void register_buffer(fptr_t _fa, const std::vector<int64_t>& fake_ipc_ptrs);
-std::tuple<std::vector<int64_t>, std::vector<int64_t>>
-get_graph_buffer_ipc_meta(fptr_t _fa);
-void register_graph_buffers(fptr_t _fa,
-                            const std::vector<std::vector<int64_t>>& handles,
-                            const std::vector<std::vector<int64_t>>& offsets);
-std::tuple<int64_t, torch::stable::Tensor> allocate_shared_buffer_and_handle(
-    int64_t size);
-int64_t open_mem_handle(torch::stable::Tensor& mem_handle);
-void free_shared_buffer(int64_t buffer);
-
 // Activation kernels (shared CUDA/ROCm)
 void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void silu_and_mul_clamp(torch::stable::Tensor& out,
-                        torch::stable::Tensor& input, double limit);
+                        torch::stable::Tensor& input, double limit,
+                        double alpha = 1.0, double beta = 0.0);
 void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_tanh_and_mul(torch::stable::Tensor& out,

csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu

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

csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu

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

csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu

@@ -25,7 +25,7 @@
 #include <cuda_fp8.h>
 
 #include "cuda_utils.h"
-#include "libtorch_stable/launch_bounds_utils.h"
+#include "launch_bounds_utils.h"
 
 // Define before including nvfp4_utils.cuh so the header
 // can use this macro during compilation.

csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu

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

csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu

@@ -34,7 +34,7 @@
 static_assert(CVT_FP4_ELTS_PER_THREAD == 16,
               "MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)");
 
-#include "libtorch_stable/launch_bounds_utils.h"
+#include "launch_bounds_utils.h"
 
 namespace vllm {
 

csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu

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

csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu

@@ -26,7 +26,7 @@
 
 #include "cuda_utils.h"
 #include "nvfp4_utils.cuh"
-#include "libtorch_stable/launch_bounds_utils.h"
+#include "launch_bounds_utils.h"
 
 namespace vllm {
 

csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu

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

csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu

@@ -26,7 +26,7 @@
 #include "../../cuda_vec_utils.cuh"
 
 #include "cuda_utils.h"
-#include "libtorch_stable/launch_bounds_utils.h"
+#include "launch_bounds_utils.h"
 
 // Define before including nvfp4_utils.cuh so the header
 // can use this macro during compilation.

csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu

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

csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu

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

csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu

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

csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu

@@ -7,11 +7,14 @@
 
 #include <torch/csrc/stable/ops.h>
 
-#include "ggml-common.h"
-#include "vecdotq.cuh"
-#include "dequantize.cuh"
-#include "mmvq.cuh"
-#include "mmq.cuh"
+// NOTE: These headers are intentionally kept in csrc/quantization/gguf/ (not
+// moved to libtorch_stable) to avoid unnecessary reformatting that would break
+// git rename detection and pollute blame history.
+#include "../../../quantization/gguf/ggml-common.h"
+#include "../../../quantization/gguf/vecdotq.cuh"
+#include "../../../quantization/gguf/dequantize.cuh"
+#include "../../../quantization/gguf/mmvq.cuh"
+#include "../../../quantization/gguf/mmq.cuh"
 #include "moe.cuh"
 #include "moe_vec.cuh"
 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

csrc/libtorch_stable/topk.cu

@@ -7,7 +7,7 @@
 #include "torch_utils.h"
 
 #ifndef USE_ROCM
-  #include "persistent_topk.cuh"
+  #include "../persistent_topk.cuh"
 #endif
 
 namespace {

csrc/libtorch_stable/torch_bindings.cpp

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

csrc/libtorch_stable/torch_utils.h

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

csrc/minimax_m3_build_k2q_csr.cu

@@ -0,0 +1,742 @@
+// CUDA C++ q2k -> k2q CSR builder.
+//
+// Five-stage pipeline. q-ascending order within each CSR row is preserved
+// by partitioning q across (CTA, warp_in_CTA) units; each unit owns a
+// contiguous q-sub-range and reserves a contiguous slot range per row via
+// a precomputed exclusive prefix scan.
+//
+//   M:  build_row_map      -- round-robin packing of rows across batches
+//   H:  histogram + tile_counts
+//   PR: row prefix         -- single block per head, row_counts -> row_ptr
+//   PT: tile prefix        -- multi-block, scan tile_counts along (c, w) axis
+//   S:  scatter (sorted)   -- per-warp slot range, q-sequential within warp
+//
+// Per-warp partitioning: each CTA has kWarps warps; warp w of CTA c owns
+// q-range [c*q_per_cta + w*q_per_warp, c*q_per_cta + (w+1)*q_per_warp).
+// tile_counts is shaped [G * kWarps, H, total_rows]; the "row" dimension
+// of the prefix scan is the flattened (c * kWarps + w) index, scanned in
+// lexicographic order so that warp-local slot ranges concatenate to the
+// global q-sorted output.
+
+#include <torch/all.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include <algorithm>
+
+#define CHECK_CUDA(x) TORCH_CHECK((x).is_cuda(), #x " must be CUDA")
+#define CHECK_CONTIGUOUS(x) \
+  TORCH_CHECK((x).is_contiguous(), #x " must be contiguous")
+#define CHECK_INT(x) \
+  TORCH_CHECK((x).scalar_type() == at::kInt, #x " must be int32")
+#define CHECK_INPUT(x) \
+  CHECK_CUDA(x);       \
+  CHECK_CONTIGUOUS(x); \
+  CHECK_INT(x)
+
+namespace {
+
+constexpr int kWarpSize = 32;
+
+__device__ __forceinline__ void advance_batch_only(int const* __restrict__ cu_q,
+                                                   int B, int q_abs, int& bi) {
+  while (bi < B && cu_q[bi + 1] <= q_abs) ++bi;
+}
+
+// Atomic increment of a 16-bit half within a 32-bit SMEM word; returns the
+// OLD 16-bit value (slot). Per-warp count must stay < 32768 so the low
+// half does not carry into the high half.
+//   base_int32 : int32 pointer; element i holds rows 2*i (low) and 2*i+1
+//   (high).
+__device__ __forceinline__ int atomic_inc_int16_packed(int* base_int32,
+                                                       int row) {
+  int idx = row >> 1;
+  int shift = (row & 1) << 4;  // 0 or 16
+  int delta = 1 << shift;
+  int old = atomicAdd(&base_int32[idx], delta);
+  return (old >> shift) & 0xFFFF;
+}
+
+// Read 16-bit half from packed int32 storage.
+__device__ __forceinline__ int read_int16_packed(int const* base_int32,
+                                                 int row) {
+  int v = base_int32[row >> 1];
+  int shift = (row & 1) << 4;
+  return (v >> shift) & 0xFFFF;
+}
+
+// ---------------------------------------------------------------------------
+// M: round-robin row map.
+// ---------------------------------------------------------------------------
+template <int kBlockK>
+__global__ void k2q_build_row_map_kernel(int const* __restrict__ cu_k,
+                                         int* __restrict__ row_map,
+                                         int* __restrict__ row_coords, int B,
+                                         int max_kv_blocks) {
+  int level = blockIdx.x;
+  if (level >= max_kv_blocks) return;
+  if (threadIdx.x != 0) return;
+  int rows_before = 0;
+  for (int b = 0; b < B; ++b) {
+    int rb = (cu_k[b + 1] - cu_k[b] + kBlockK - 1) / kBlockK;
+    rows_before += (rb < level ? rb : level);
+  }
+  int active_before = 0;
+  for (int b = 0; b < B; ++b) {
+    int rb = (cu_k[b + 1] - cu_k[b] + kBlockK - 1) / kBlockK;
+    if (rb > level) {
+      int row_linear = rows_before + active_before;
+      row_map[(size_t)b * max_kv_blocks + level] = row_linear;
+      if (row_coords != nullptr) {
+        row_coords[(size_t)row_linear * 2] = b;
+        row_coords[(size_t)row_linear * 2 + 1] = level;
+      }
+      ++active_before;
+    } else {
+      row_map[(size_t)b * max_kv_blocks + level] = -1;
+    }
+  }
+}
+
+// ---------------------------------------------------------------------------
+// H: per-warp histogram + tile_counts.
+// kWarps warps per CTA, each owns q-sub-range = q_per_cta / kWarps.
+// SMEM hist[kWarps, total_rows] int32 (stored as packed int16 cursor:
+// 2 entries per int32 word). Each warp counts to its own row.
+// At end-of-CTA, write tile_counts[c*kWarps + w, h, r] = smem_hist[w, r]
+// and atomicAdd(row_counts[h, r], sum over w of smem_hist[w, r]).
+// ---------------------------------------------------------------------------
+template <int kTopK, int kBlockK, int kWarps>
+__global__ void k2q_hist_kernel(int const* __restrict__ q2k,
+                                int const* __restrict__ cu_q,
+                                int const* __restrict__ row_map,
+                                int* __restrict__ row_counts,
+                                int* __restrict__ tile_counts, int H, int B,
+                                int S_Q, int total_rows, int max_kv_blocks,
+                                int q_per_cta, int q_per_warp) {
+  constexpr int kThreads = kWarps * kWarpSize;
+  extern __shared__ int smem_hist_int[];
+  int* smem_hist = smem_hist_int;
+  int tid = threadIdx.x;
+  int warp_id = tid >> 5;
+  int lane = tid & 31;
+  int c = blockIdx.x;
+  int q_start_cta = c * q_per_cta;
+  int q_end_cta = min(q_start_cta + q_per_cta, S_Q);
+  int q_start_warp = min(q_start_cta + warp_id * q_per_warp, q_end_cta);
+  int q_end_warp = min(q_start_warp + q_per_warp, q_end_cta);
+
+  constexpr int kInt4PerToken = kTopK / 4;
+  int packed_per_warp = (total_rows + 1) >> 1;
+  int* my_hist = smem_hist + warp_id * packed_per_warp;
+
+  for (int h = 0; h < H; ++h) {
+    for (int i = lane; i < packed_per_warp; i += kWarpSize) my_hist[i] = 0;
+    __syncthreads();
+
+    if (q_start_warp < q_end_warp) {
+      int bi = 0;
+      int qi = q_start_warp + lane;
+      advance_batch_only(cu_q, B, qi, bi);
+
+      int4 const* head_topk4 =
+          reinterpret_cast<int4 const*>(q2k + (size_t)h * S_Q * kTopK);
+
+      for (; qi < q_end_warp; qi += kWarpSize) {
+        advance_batch_only(cu_q, B, qi, bi);
+        int const* my_row_map = row_map + (size_t)bi * max_kv_blocks;
+
+        int4 buf[kInt4PerToken];
+#pragma unroll
+        for (int v = 0; v < kInt4PerToken; ++v) {
+          buf[v] = head_topk4[(size_t)qi * kInt4PerToken + v];
+        }
+#pragma unroll
+        for (int t = 0; t < kTopK; ++t) {
+          int kvb_local = reinterpret_cast<int const*>(buf)[t];
+          if (kvb_local >= 0 && kvb_local < max_kv_blocks) {
+            int row = my_row_map[kvb_local];
+            if (row >= 0 && row < total_rows) {
+              atomic_inc_int16_packed(my_hist, row);
+            }
+          }
+        }
+      }
+    }
+    __syncthreads();
+
+    int* head_row_counts = row_counts + (size_t)h * total_rows;
+    // Each warp writes its own slice of tile_counts (full int32) by
+    // unpacking int16 entries from SMEM.
+    int* my_tile =
+        tile_counts + ((size_t)(c * kWarps + warp_id) * H + h) * total_rows;
+    for (int i = lane; i < total_rows; i += kWarpSize) {
+      my_tile[i] = read_int16_packed(my_hist, i);
+    }
+    __syncthreads();
+
+    // Sum across warps (int32 accumulator), atomicAdd to row_counts.
+    for (int i = tid; i < total_rows; i += kThreads) {
+      int sum = 0;
+#pragma unroll
+      for (int w = 0; w < kWarps; ++w) {
+        sum += read_int16_packed(smem_hist + w * packed_per_warp, i);
+      }
+      if (sum > 0) atomicAdd(&head_row_counts[i], sum);
+    }
+    if (h + 1 < H) __syncthreads();
+  }
+}
+
+// ---------------------------------------------------------------------------
+// PR: row prefix. One block per head.
+// ---------------------------------------------------------------------------
+template <int kThreads>
+__global__ void k2q_row_prefix_kernel(int const* __restrict__ row_counts,
+                                      int* __restrict__ row_ptr,
+                                      int const* __restrict__ row_coords,
+                                      int* __restrict__ scheduler_metadata,
+                                      int* __restrict__ work_count,
+                                      int total_rows, int target_q_per_cta,
+                                      int work_capacity) {
+  int h = blockIdx.x;
+  int tid = threadIdx.x;
+  __shared__ int scan_buf[kThreads];
+
+  int const* head_counts = row_counts + (size_t)h * total_rows;
+  int* head_rowptr = row_ptr + (size_t)h * (total_rows + 1);
+  int chunk = (total_rows + kThreads - 1) / kThreads;
+  int lo = tid * chunk;
+  int hi = min(lo + chunk, total_rows);
+
+  int local_sum = 0;
+  for (int i = lo; i < hi; ++i) local_sum += head_counts[i];
+  scan_buf[tid] = local_sum;
+  __syncthreads();
+
+  for (int off = 1; off < kThreads; off <<= 1) {
+    int add = (tid >= off) ? scan_buf[tid - off] : 0;
+    __syncthreads();
+    scan_buf[tid] += add;
+    __syncthreads();
+  }
+  int running = scan_buf[tid] - local_sum;
+  for (int i = lo; i < hi; ++i) {
+    int row_count = head_counts[i];
+    running += row_count;
+    head_rowptr[i + 1] = running;
+    if (scheduler_metadata != nullptr && work_count != nullptr &&
+        row_count > 0) {
+      int num_chunks = (row_count + target_q_per_cta - 1) / target_q_per_cta;
+      int base = atomicAdd(work_count, num_chunks);
+      int batch_idx = row_coords[(size_t)i * 2];
+      int kv_block_idx = row_coords[(size_t)i * 2 + 1];
+      for (int c = 0; c < num_chunks; ++c) {
+        int work_idx = base + c;
+        if (work_idx < work_capacity) {
+          int q_begin = c * target_q_per_cta;
+          int q_count = min(target_q_per_cta, row_count - q_begin);
+          int* meta = scheduler_metadata + (size_t)work_idx * 6;
+          meta[0] = h;
+          meta[1] = i;
+          meta[2] = q_begin;
+          meta[3] = q_count;
+          meta[4] = batch_idx;
+          meta[5] = kv_block_idx;
+        }
+      }
+    }
+  }
+}
+
+// ---------------------------------------------------------------------------
+// PT_smem: SMEM-staged tile prefix scan.
+// Each block handles kRowsPerBlock rows for one head h. Cooperative load
+// of tile_counts[*, h, base_r..base_r+M) into SMEM (better coalescing
+// than per-warp uncoalesced stride reads), then per-warp scan in SMEM,
+// then cooperative store back. Fuses row_ptr into the base.
+// ---------------------------------------------------------------------------
+template <int kThreads, int kRowsPerBlock>
+__global__ void k2q_tile_prefix_smem_kernel(int* __restrict__ tile_counts,
+                                            int const* __restrict__ row_ptr,
+                                            int H, int total_rows,
+                                            int G_total) {
+  static_assert(kRowsPerBlock > 0, "kRowsPerBlock must be positive");
+  extern __shared__ int smem_tprefix[];
+  // smem layout: smem[r_off][g] for r_off in [0, M), g in [0, G_total).
+
+  int tid = threadIdx.x;
+  int lane = tid & 31;
+  int warp_id = tid >> 5;
+
+  // Grid: H * blocks_per_h. Each block stays within a single head h
+  // and processes kRowsPerBlock contiguous rows starting at b_in_h *
+  // kRowsPerBlock. (Earlier flat-grid mapping `h = block_job /
+  // total_rows; base_r = block_job - h*total_rows` skipped rows when
+  // total_rows was not a multiple of kRowsPerBlock and H > 1, because
+  // the last partial block of head h-1 left blocks of head h starting
+  // at a non-zero row offset.)
+  int blocks_per_h = (total_rows + kRowsPerBlock - 1) / kRowsPerBlock;
+  int h = blockIdx.x / blocks_per_h;
+  int b_in_h = blockIdx.x - h * blocks_per_h;
+  if (h >= H) return;
+  int base_r = b_in_h * kRowsPerBlock;
+  if (base_r >= total_rows) return;
+  int actual_M = min(kRowsPerBlock, total_rows - base_r);
+
+  size_t stride_g = (size_t)H * total_rows;
+  int* base_ptr = tile_counts + (size_t)h * total_rows + base_r;
+  int total_elems = G_total * actual_M;
+
+  // Cooperative load. Pattern: thread tid -> (r_off=tid%M, g=tid/M),
+  // then strided. 32 lanes hit M r's × (32/M) g's, giving 32/M cache
+  // lines per warp (vs 32 in the naive stride-along-g pattern).
+  for (int i = tid; i < total_elems; i += kThreads) {
+    int r_off = i % actual_M;
+    int g = i / actual_M;
+    smem_tprefix[r_off * G_total + g] = base_ptr[g * stride_g + r_off];
+  }
+  __syncthreads();
+
+  // Per-warp scan: warp w scans row (base_r + w) if w < actual_M.
+  if (warp_id < actual_M) {
+    int abs_r = base_r + warp_id;
+    int rp = row_ptr[(size_t)h * (total_rows + 1) + abs_r];
+    int* my_smem = smem_tprefix + warp_id * G_total;
+    int running = rp;
+    for (int g0 = 0; g0 < G_total; g0 += kWarpSize) {
+      int g = g0 + lane;
+      int v = (g < G_total) ? my_smem[g] : 0;
+      int x = v;
+#pragma unroll
+      for (int off = 1; off < kWarpSize; off <<= 1) {
+        int nbr = __shfl_up_sync(0xFFFFFFFF, x, off);
+        if (lane >= off) x += nbr;
+      }
+      int excl = running + x - v;
+      if (g < G_total) my_smem[g] = excl;
+      int chunk_sum = __shfl_sync(0xFFFFFFFF, x, 31);
+      running += chunk_sum;
+    }
+  }
+  __syncthreads();
+
+  // Cooperative store back.
+  for (int i = tid; i < total_elems; i += kThreads) {
+    int r_off = i % actual_M;
+    int g = i / actual_M;
+    base_ptr[g * stride_g + r_off] = smem_tprefix[r_off * G_total + g];
+  }
+}
+
+// ---------------------------------------------------------------------------
+// S: scatter. kWarps warps per CTA, each owns q-sub-range. Per-warp SMEM
+// cursor and per-warp tile_offset slot range. Within a warp, q's are
+// processed sequentially; lanes 0..kTopK-1 handle the topK slots in
+// lockstep. Across distinct q's in the same warp, the lockstep ordering
+// guarantees q-monotonic atomicAdd on smem_cursor[r].
+// ---------------------------------------------------------------------------
+// kQPerIter * kTopK lanes are active per warp iter; remaining lanes idle.
+// For kTopK=16, kQPerIter=2 uses all 32 lanes; for kTopK=8, kQPerIter=4.
+// CORRECTNESS NOTE: relies on lane-ordered SMEM atomicAdd return values
+// within a single warp instruction (verified on B200; tests pass).
+//
+// SMEM cursor stored as packed int16 (two cursors per int32). Per-warp
+// row count must stay < 32768 (~q_per_warp * kTopK at max sink), which
+// holds for all task.md sizes up to 1024K.
+template <int kTopK, int kBlockK, int kWarps>
+__global__ void k2q_scatter_kernel(
+    int const* __restrict__ q2k, int const* __restrict__ cu_q,
+    int const* __restrict__ row_map, int const* __restrict__ abs_base,
+    int* __restrict__ q_idx, int* __restrict__ qsplit_idx,
+    int* __restrict__ split_counts, int H, int B, int S_Q, int total_rows,
+    int max_kv_blocks, int q_per_cta, int q_per_warp, int max_seqlen_q) {
+  constexpr int kQPerIter = kWarpSize / kTopK > 0 ? kWarpSize / kTopK : 1;
+  extern __shared__ int smem_cursor_int[];
+  int* smem_cursor = smem_cursor_int;
+  int tid = threadIdx.x;
+  int warp_id = tid >> 5;
+  int lane = tid & 31;
+  int c = blockIdx.x;
+  int q_start_cta = c * q_per_cta;
+  int q_end_cta = min(q_start_cta + q_per_cta, S_Q);
+  int q_start_warp = min(q_start_cta + warp_id * q_per_warp, q_end_cta);
+  int q_end_warp = min(q_start_warp + q_per_warp, q_end_cta);
+
+  int q_in_iter = lane / kTopK;
+  int slot_in_q = lane % kTopK;
+  bool lane_active = (lane < kQPerIter * kTopK);
+
+  // Per-warp packed cursor: total_rows int16 entries -> ceil(total_rows/2)
+  // int32.
+  int packed_per_warp = (total_rows + 1) >> 1;
+  int* my_cursor = smem_cursor + warp_id * packed_per_warp;
+
+  for (int h = 0; h < H; ++h) {
+    for (int i = lane; i < packed_per_warp; i += kWarpSize) my_cursor[i] = 0;
+    __syncwarp();
+
+    if (q_start_warp < q_end_warp) {
+      int bi = 0;
+      advance_batch_only(cu_q, B, q_start_warp, bi);
+
+      int const* head_q2k = q2k + (size_t)h * S_Q * kTopK;
+      int const* my_abs_base =
+          abs_base + ((size_t)(c * kWarps + warp_id) * H + h) * total_rows;
+      int* head_qidx = q_idx + (size_t)h * S_Q * kTopK;
+
+      // (Hot-row register cache experiment showed no measurable
+      // benefit; relying on L1 to keep row 0 / row total_rows-1
+      // hot since they're hit every iteration in sink workloads.)
+
+      constexpr int kUnroll = 16;
+      int qi_base = q_start_warp;
+      for (; qi_base + kUnroll * kQPerIter <= q_end_warp;
+           qi_base += kUnroll * kQPerIter) {
+        int kvb[kUnroll];
+        int qloc[kUnroll];
+        int batch[kUnroll];
+        int const* rmap[kUnroll];
+
+#pragma unroll
+        for (int u = 0; u < kUnroll; ++u) {
+          int qi_u = qi_base + u * kQPerIter + q_in_iter;
+          kvb[u] = -1;
+          qloc[u] = 0;
+          batch[u] = 0;
+          if (lane_active) {
+            advance_batch_only(cu_q, B, qi_u, bi);
+            qloc[u] = qi_u - cu_q[bi];
+            batch[u] = bi;
+            kvb[u] = head_q2k[(size_t)qi_u * kTopK + slot_in_q];
+          }
+          rmap[u] = row_map + (size_t)bi * max_kv_blocks;
+        }
+
+        int row[kUnroll];
+#pragma unroll
+        for (int u = 0; u < kUnroll; ++u) {
+          row[u] = -1;
+          if (lane_active && kvb[u] >= 0 && kvb[u] < max_kv_blocks)
+            row[u] = rmap[u][kvb[u]];
+        }
+
+        // Pre-issue all kUnroll abs_base loads in parallel before
+        // the atomic chain so memory pipeline runs concurrently
+        // with SMEM atomic-adds.
+        int abs_v[kUnroll];
+#pragma unroll
+        for (int u = 0; u < kUnroll; ++u) {
+          abs_v[u] =
+              (row[u] >= 0 && row[u] < total_rows) ? my_abs_base[row[u]] : 0;
+        }
+
+#pragma unroll
+        for (int u = 0; u < kUnroll; ++u) {
+          int r = row[u];
+          bool valid_edge = r >= 0 && r < total_rows;
+          unsigned int valid_mask = __ballot_sync(0xFFFFFFFFu, valid_edge);
+          unsigned int group_mask =
+              (kTopK == 32) ? 0xFFFFFFFFu
+                            : (((1u << kTopK) - 1u) << (q_in_iter * kTopK));
+          unsigned int lower_lane_mask = lane == 0 ? 0u : ((1u << lane) - 1u);
+          int split_slot = __popc(valid_mask & group_mask & lower_lane_mask);
+          int valid_count = __popc(valid_mask & group_mask);
+          if (split_counts != nullptr && slot_in_q == 0) {
+            split_counts[((size_t)batch[u] * max_seqlen_q + qloc[u]) * H + h] =
+                valid_count;
+          }
+          if (valid_edge) {
+            int slot = atomic_inc_int16_packed(my_cursor, r);
+            int out_pos = abs_v[u] + slot;
+            head_qidx[out_pos] = qloc[u];
+            if (qsplit_idx != nullptr) {
+              qsplit_idx[(size_t)h * S_Q * kTopK + out_pos] =
+                  qloc[u] | ((split_slot & 0xFF) << 24);
+            }
+          }
+        }
+      }
+      // Tail: 1-3 iters left.
+      for (; qi_base < q_end_warp; qi_base += kQPerIter) {
+        int my_qi = qi_base + q_in_iter;
+        bool valid_q = (my_qi < q_end_warp) && lane_active;
+        int kvb_local = -1;
+        int q_local = 0;
+        int batch_local = 0;
+        if (valid_q) {
+          advance_batch_only(cu_q, B, my_qi, bi);
+          batch_local = bi;
+          q_local = my_qi - cu_q[bi];
+          kvb_local = head_q2k[(size_t)my_qi * kTopK + slot_in_q];
+        }
+        int const* my_row_map = row_map + (size_t)bi * max_kv_blocks;
+        int row = -1;
+        if (valid_q && kvb_local >= 0 && kvb_local < max_kv_blocks) {
+          row = my_row_map[kvb_local];
+        }
+        bool valid_edge = row >= 0 && row < total_rows;
+        unsigned int valid_mask = __ballot_sync(0xFFFFFFFFu, valid_edge);
+        unsigned int group_mask =
+            (kTopK == 32) ? 0xFFFFFFFFu
+                          : (((1u << kTopK) - 1u) << (q_in_iter * kTopK));
+        unsigned int lower_lane_mask = lane == 0 ? 0u : ((1u << lane) - 1u);
+        int split_slot = __popc(valid_mask & group_mask & lower_lane_mask);
+        int valid_count = __popc(valid_mask & group_mask);
+        if (split_counts != nullptr && valid_q && slot_in_q == 0) {
+          split_counts[((size_t)batch_local * max_seqlen_q + q_local) * H + h] =
+              valid_count;
+        }
+        if (valid_edge) {
+          int slot = atomic_inc_int16_packed(my_cursor, row);
+          int out_pos = my_abs_base[row] + slot;
+          head_qidx[out_pos] = q_local;
+          if (qsplit_idx != nullptr) {
+            qsplit_idx[(size_t)h * S_Q * kTopK + out_pos] =
+                q_local | ((split_slot & 0xFF) << 24);
+          }
+        }
+      }
+    }
+    if (h + 1 < H) __syncthreads();
+  }
+}
+
+}  // anonymous namespace
+
+// ===========================================================================
+// Host orchestration
+// ===========================================================================
+
+template <int kTopK, int kBlockK>
+static void launch_pipeline(torch::Tensor q2k, torch::Tensor cu_q,
+                            torch::Tensor cu_k, torch::Tensor row_ptr,
+                            torch::Tensor q_idx, int total_rows,
+                            int max_kv_blocks,
+                            torch::Tensor scheduler_metadata = torch::Tensor(),
+                            torch::Tensor work_count = torch::Tensor(),
+                            torch::Tensor qsplit_idx = torch::Tensor(),
+                            torch::Tensor split_counts = torch::Tensor(),
+                            int target_q_per_cta = 1, int work_capacity = 0,
+                            int max_seqlen_q = 0) {
+  int H = (int)q2k.size(0);
+  int S_Q = (int)q2k.size(1);
+  int topK = (int)q2k.size(2);
+  TORCH_CHECK(topK == kTopK, "topK runtime != template kTopK");
+  int B = (int)cu_q.size(0) - 1;
+  auto device = q2k.device();
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  AT_CUDA_CHECK(cudaMemsetAsync(row_ptr.data_ptr<int>(), 0,
+                                (size_t)H * (total_rows + 1) * sizeof(int),
+                                stream));
+  AT_CUDA_CHECK(cudaMemsetAsync(q_idx.data_ptr<int>(), 0xFF,
+                                (size_t)H * S_Q * kTopK * sizeof(int), stream));
+
+  auto opts = torch::TensorOptions().dtype(torch::kInt32).device(device);
+  auto row_counts = torch::zeros({H, total_rows}, opts);
+  auto row_map = torch::empty({B, max_kv_blocks}, opts);
+  bool emit_schedule = scheduler_metadata.defined();
+  auto row_coords =
+      emit_schedule ? torch::empty({total_rows, 2}, opts) : torch::Tensor();
+  int* scheduler_metadata_ptr =
+      emit_schedule ? scheduler_metadata.data_ptr<int>() : nullptr;
+  int* work_count_ptr = emit_schedule ? work_count.data_ptr<int>() : nullptr;
+  int* qsplit_idx_ptr = emit_schedule ? qsplit_idx.data_ptr<int>() : nullptr;
+  int* split_counts_ptr =
+      emit_schedule ? split_counts.data_ptr<int>() : nullptr;
+  int* row_coords_ptr = emit_schedule ? row_coords.data_ptr<int>() : nullptr;
+  if (emit_schedule) {
+    AT_CUDA_CHECK(cudaMemsetAsync(work_count_ptr, 0, sizeof(int), stream));
+    AT_CUDA_CHECK(cudaMemsetAsync(scheduler_metadata_ptr, 0,
+                                  (size_t)work_capacity * 6 * sizeof(int),
+                                  stream));
+  }
+
+  int dev = q2k.get_device();
+  int num_sms = 0;
+  AT_CUDA_CHECK(
+      cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, dev));
+
+  // -- Pick kWarps per CTA based on SMEM budget for cursor/hist ---------
+  // SMEM per CTA = kWarps * total_rows * sizeof(int) (for both H and S).
+  // Want at least 2 CTAs/SM for memory parallelism. SM100 SMEM = 228KB.
+  // Pick the largest kWarps that fits two CTAs/SM, capped at 4.
+  // SMEM cursor packed as int16 (2 entries per int32 word):
+  int per_warp_smem = ((total_rows + 1) >> 1) * (int)sizeof(int);
+  int kWarps_pick = 4;
+  while (kWarps_pick > 1 && (kWarps_pick * per_warp_smem) * 2 > 228 * 1024) {
+    kWarps_pick >>= 1;
+  }
+  if (kWarps_pick < 1) kWarps_pick = 1;
+
+  // -- Pick G (CTAs) ----------------------------------------------------
+  // For each (kWarps, per_warp_smem) pair, the SMEM-bound occupancy is
+  // 228KB / (kWarps*per_warp_smem) CTAs/SM. We size G as
+  // num_sms * occupancy so a single resident wave covers all CTAs and
+  // the memory pipeline runs at peak.
+  int per_cta_smem_bytes = kWarps_pick * per_warp_smem;
+  int max_ctas_per_sm =
+      std::max(1, (228 * 1024) / std::max(1, per_cta_smem_bytes));
+  if (max_ctas_per_sm > 8) max_ctas_per_sm = 8;
+  constexpr int kMinQPerCta = 256;
+  // Cap target_g at num_sms * 3 — empirically this balances
+  // per-CTA work-size against parallelism. Higher caps regress
+  // mid-size cases due to row_counts atomicAdd contention and
+  // smaller q_per_cta. SMEM-bound configurations naturally cap
+  // lower if max_ctas_per_sm < 3.
+  int target_g = num_sms * std::min(max_ctas_per_sm, 3);
+  int max_g_for_q = (S_Q + kMinQPerCta - 1) / kMinQPerCta;
+  int G = std::min({target_g, max_g_for_q, S_Q});
+  if (G < 1) G = 1;
+  int q_per_cta = (S_Q + G - 1) / G;
+  G = (S_Q + q_per_cta - 1) / q_per_cta;
+  int q_per_warp = (q_per_cta + kWarps_pick - 1) / kWarps_pick;
+  int G_total = G * kWarps_pick;
+
+  auto tile_counts = torch::empty({G_total, H, total_rows}, opts);
+
+  // -- Compile-time switch on kWarps for the templated kernels ---------
+  auto rmap_fn = k2q_build_row_map_kernel<kBlockK>;
+  auto rprefix_fn = k2q_row_prefix_kernel<1024>;
+  constexpr int kPtRowsPerBlock = 8;
+  constexpr int kPtThreads = 256;
+  auto tprefix_smem_fn =
+      k2q_tile_prefix_smem_kernel<kPtThreads, kPtRowsPerBlock>;
+
+  if (max_kv_blocks > 0) {
+    rmap_fn<<<max_kv_blocks, 32, 0, stream>>>(cu_k.data_ptr<int>(),
+                                              row_map.data_ptr<int>(),
+                                              row_coords_ptr, B, max_kv_blocks);
+  }
+
+  auto launch_hist_scatter = [&](auto kWarps_const) {
+    constexpr int W = decltype(kWarps_const)::value;
+    size_t smem_bytes = (size_t)W * per_warp_smem;
+    auto hist_fn = k2q_hist_kernel<kTopK, kBlockK, W>;
+    auto scat_fn = k2q_scatter_kernel<kTopK, kBlockK, W>;
+    AT_CUDA_CHECK(cudaFuncSetAttribute(
+        hist_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem_bytes));
+    AT_CUDA_CHECK(cudaFuncSetAttribute(
+        scat_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem_bytes));
+
+    hist_fn<<<G, W * kWarpSize, smem_bytes, stream>>>(
+        q2k.data_ptr<int>(), cu_q.data_ptr<int>(), row_map.data_ptr<int>(),
+        row_counts.data_ptr<int>(), tile_counts.data_ptr<int>(), H, B, S_Q,
+        total_rows, max_kv_blocks, q_per_cta, q_per_warp);
+
+    rprefix_fn<<<H, 1024, 0, stream>>>(
+        row_counts.data_ptr<int>(), row_ptr.data_ptr<int>(),
+        emit_schedule ? row_coords.data_ptr<int>() : nullptr,
+        scheduler_metadata_ptr, work_count_ptr, total_rows, target_q_per_cta,
+        work_capacity);
+
+    // Grid is H * blocks_per_h so each block stays within a single
+    // head; flat (H*total_rows) grid would skip rows when total_rows
+    // is not a multiple of kPtRowsPerBlock.
+    int blocks_per_h = (total_rows + kPtRowsPerBlock - 1) / kPtRowsPerBlock;
+    int pt_grid = H * blocks_per_h;
+    if (pt_grid < 1) pt_grid = 1;
+    size_t pt_smem = (size_t)kPtRowsPerBlock * G_total * sizeof(int);
+    AT_CUDA_CHECK(cudaFuncSetAttribute(
+        tprefix_smem_fn, cudaFuncAttributeMaxDynamicSharedMemorySize,
+        (int)pt_smem));
+    tprefix_smem_fn<<<pt_grid, kPtThreads, pt_smem, stream>>>(
+        tile_counts.data_ptr<int>(), row_ptr.data_ptr<int>(), H, total_rows,
+        G_total);
+
+    scat_fn<<<G, W * kWarpSize, smem_bytes, stream>>>(
+        q2k.data_ptr<int>(), cu_q.data_ptr<int>(), row_map.data_ptr<int>(),
+        tile_counts.data_ptr<int>(), q_idx.data_ptr<int>(), qsplit_idx_ptr,
+        split_counts_ptr, H, B, S_Q, total_rows, max_kv_blocks, q_per_cta,
+        q_per_warp, max_seqlen_q);
+  };
+
+  if (kWarps_pick == 4) {
+    launch_hist_scatter(std::integral_constant<int, 4>{});
+  } else if (kWarps_pick == 2) {
+    launch_hist_scatter(std::integral_constant<int, 2>{});
+  } else {
+    launch_hist_scatter(std::integral_constant<int, 1>{});
+  }
+}
+
+void run_minimax_m3_build_k2q_csr_with_schedule(
+    torch::Tensor q2k, torch::Tensor cu_q, torch::Tensor cu_k,
+    torch::Tensor row_ptr, torch::Tensor q_idx,
+    torch::Tensor scheduler_metadata, torch::Tensor work_count,
+    torch::Tensor qsplit_idx, torch::Tensor split_counts, int64_t topk,
+    int64_t blk_kv, int64_t total_rows, int64_t max_kv_blocks,
+    int64_t target_q_per_cta, int64_t work_capacity, int64_t max_seqlen_q) {
+  CHECK_INPUT(q2k);
+  CHECK_INPUT(cu_q);
+  CHECK_INPUT(cu_k);
+  CHECK_INPUT(row_ptr);
+  CHECK_INPUT(q_idx);
+  CHECK_INPUT(scheduler_metadata);
+  CHECK_INPUT(work_count);
+  CHECK_INPUT(qsplit_idx);
+  CHECK_INPUT(split_counts);
+  TORCH_CHECK(blk_kv == 128, "build_k2q_csr only supports blk_kv == 128");
+  int H = (int)q2k.size(0);
+  int S_Q = (int)q2k.size(1);
+  int tr = (int)total_rows;
+  int mkv = (int)max_kv_blocks;
+  int target = (int)target_q_per_cta;
+  int capacity = (int)work_capacity;
+  int max_sq = (int)max_seqlen_q;
+  TORCH_CHECK(tr >= 0 && mkv >= 0 && target > 0 && capacity > 0 && max_sq >= 0,
+              "invalid schedule sizing arguments");
+  TORCH_CHECK(row_ptr.size(0) == H && row_ptr.size(1) == tr + 1,
+              "row_ptr shape mismatch");
+  TORCH_CHECK(q_idx.size(0) == H && q_idx.size(1) == (int64_t)S_Q * (int)topk,
+              "q_idx shape mismatch");
+  TORCH_CHECK(qsplit_idx.sizes() == q_idx.sizes(), "qsplit_idx shape mismatch");
+  TORCH_CHECK(
+      scheduler_metadata.size(0) == capacity && scheduler_metadata.size(1) == 6,
+      "scheduler_metadata shape mismatch");
+  TORCH_CHECK(work_count.numel() == 1,
+              "work_count must have one int32 element");
+  TORCH_CHECK(split_counts.dim() == 3 &&
+                  split_counts.size(0) == cu_q.size(0) - 1 &&
+                  split_counts.size(1) == max_sq && split_counts.size(2) == H,
+              "split_counts shape mismatch");
+  if (S_Q == 0 || tr == 0 || H == 0 || mkv == 0) {
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    AT_CUDA_CHECK(cudaMemsetAsync(row_ptr.data_ptr<int>(), 0,
+                                  (size_t)H * (tr + 1) * sizeof(int), stream));
+    AT_CUDA_CHECK(cudaMemsetAsync(q_idx.data_ptr<int>(), 0xFF,
+                                  (size_t)H * S_Q * (int)topk * sizeof(int),
+                                  stream));
+    AT_CUDA_CHECK(
+        cudaMemsetAsync(work_count.data_ptr<int>(), 0, sizeof(int), stream));
+    if (split_counts.numel() > 0) {
+      AT_CUDA_CHECK(cudaMemsetAsync(split_counts.data_ptr<int>(), 0,
+                                    (size_t)split_counts.numel() * sizeof(int),
+                                    stream));
+    }
+    return;
+  }
+
+  if (topk == 16) {
+    launch_pipeline<16, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
+                             scheduler_metadata, work_count, qsplit_idx,
+                             split_counts, target, capacity, max_sq);
+  } else if (topk == 8) {
+    launch_pipeline<8, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
+                            scheduler_metadata, work_count, qsplit_idx,
+                            split_counts, target, capacity, max_sq);
+  } else if (topk == 32) {
+    launch_pipeline<32, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
+                             scheduler_metadata, work_count, qsplit_idx,
+                             split_counts, target, capacity, max_sq);
+  } else if (topk == 4) {
+    launch_pipeline<4, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
+                            scheduler_metadata, work_count, qsplit_idx,
+                            split_counts, target, capacity, max_sq);
+  } else {
+    TORCH_CHECK(false, "unsupported topK ", topk,
+                " (expected 4, 8, 16, or 32)");
+  }
+}

csrc/minimax_reduce_rms_kernel.cu

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