refactor & simplify

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

.buildkite/hardware_tests/amd.yaml

@@ -17,6 +17,26 @@ steps:
       --target test
       --no-cache
       --progress plain .
+    - |
+      docker run --rm --network=none --entrypoint /bin/bash "rocm/vllm-ci:${BUILDKITE_COMMIT}" -ec '
+        if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi
+        if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi
+        if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi
+        if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi
+        command -v python3
+        command -v uv
+        command -v pytest
+        if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then
+          echo No ROCm CLI found in image >&2
+          exit 1
+        fi
+        python3 - <<PY
+      import torch, vllm
+      print(torch.__version__)
+      print(vllm.__version__)
+      PY
+        echo AMD image smoke OK
+      '
     - docker push "rocm/vllm-ci:${BUILDKITE_COMMIT}"
     env:
       DOCKER_BUILDKIT: "1"

.buildkite/hardware_tests/cpu.yaml

@@ -54,17 +54,35 @@ steps:
       pytest -x -v -s tests/models/language/generation -m cpu_model
       pytest -x -v -s tests/models/language/pooling -m cpu_model"
 
+- label: CPU-ModelRunnerV2 Tests
+  depends_on: []
+  device: intel_cpu
+  no_plugin: true
+  soft_fail: true
+  source_file_dependencies:
+  - vllm/v1/worker/cpu/
+  - vllm/v1/worker/gpu/
+  - vllm/v1/sample/ops/topk_topp_triton.py
+  - vllm/v1/sample/ops/topk_topp_sampler.py
+  - tests/v1/sample/test_topk_topp_sampler.py
+  commands:
+    - |
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
+      uv pip install git+https://github.com/triton-lang/triton-cpu.git@270e696d
+      VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model
+      # TODO: move to CPU-Kernel Tests once triton-cpu has a pre-built wheel
+      pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp"
+
 - label: CPU-Quantization Model Tests
   depends_on: []
   device: intel_cpu
   no_plugin: true
   source_file_dependencies:
   - csrc/cpu/
-  - vllm/model_executor/layers/quantization/cpu_wna16.py
   - vllm/model_executor/layers/quantization/auto_gptq.py
   - vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py
-  - vllm/model_executor/layers/quantization/kernels/scaled_mm/cpu.py
-  - vllm/model_executor/layers/quantization/kernels/mixed_precision/cpu.py
+  - vllm/model_executor/kernels/linear/mixed_precision/cpu.py
+  - vllm/model_executor/kernels/linear/scaled_mm/cpu.py
   - vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
   - tests/quantization/test_compressed_tensors.py
   - tests/quantization/test_cpu_wna16.py

.buildkite/image_build/image_build.yaml

@@ -13,6 +13,60 @@ steps:
         - exit_status: -10  # Agent was lost
           limit: 2
 
+  - label: ":docker: :smoking: Non-root smoke tests"
+    key: image-build-smoke-test
+    depends_on:
+      - image-build
+    commands:
+    # Smoke 1: the default (root) image must still be importable
+    # under a non-root UID via `--user 2000:0`. Validates the `vllm` passwd
+    # entry + group-0-writable /home/vllm + uv path cleanup from #31959.
+    # Uses `import vllm` rather than `vllm serve --help` because the latter
+    # instantiates `VllmConfig` which requires a GPU attached to the
+    # container.
+    - docker run --rm --user 2000:0 --entrypoint python3 "$IMAGE_TAG" -c "import vllm; print(vllm.__version__)"
+    # Smoke 2: assert the non-root enabling invariants are baked
+    # into the image. Runs as UID 2000:0 via a shell so we can verify
+    # filesystem perms + passwd/group file state + wrapper presence without
+    # triggering vLLM's GPU-requiring config-init path. The opt-in
+    # `vllm-openai-nonroot` target adds only `USER vllm`, `WORKDIR
+    # /home/vllm`, and an `ENTRYPOINT` override on top of these invariants;
+    # its build correctness is reviewed at the Dockerfile level. Wrapper
+    # logic is covered separately by the pre-commit hook
+    # `test-nonroot-entrypoint` (see .pre-commit-config.yaml).
+    - |
+      docker run --rm --user 2000:0 --entrypoint /bin/sh "$IMAGE_TAG" -ec '
+        if ! getent passwd 2000 | grep -q ^vllm:; then
+          echo FAIL: UID 2000 != vllm
+          exit 1
+        fi
+        if ! id -gn 2>/dev/null | grep -qx root; then
+          echo FAIL: GID 0 not root group
+          exit 1
+        fi
+        touch /home/vllm/.smoke && rm /home/vllm/.smoke
+        touch /opt/uv/cache/.smoke && rm /opt/uv/cache/.smoke
+        if ! test -x /usr/local/bin/vllm-nonroot-entrypoint.sh; then
+          echo FAIL: wrapper missing
+          exit 1
+        fi
+        if ! test -w /etc/passwd; then
+          echo FAIL: /etc/passwd not group-writable
+          exit 1
+        fi
+        if ! test -w /etc/group; then
+          echo FAIL: /etc/group not group-writable
+          exit 1
+        fi
+        echo non-root invariants OK
+      '
+    retry:
+      automatic:
+        - exit_status: -1  # Agent was lost
+          limit: 2
+        - exit_status: -10  # Agent was lost
+          limit: 2
+
   - label: ":docker: Build CPU image"
     key: image-build-cpu
     depends_on: []
@@ -56,3 +110,21 @@ steps:
           limit: 2
         - exit_status: -10  # Agent was lost
           limit: 2
+
+  - label: ":docker: Build arm64 image"
+    key: arm64-image-build
+    depends_on: []
+    source_file_dependencies:
+      - ".buildkite/image_build/image_build.yaml"
+      - ".buildkite/image_build/image_build_arm64.sh"
+      - "docker/Dockerfile"
+    commands:
+    - .buildkite/image_build/image_build_arm64.sh $REGISTRY $REPO $BUILDKITE_COMMIT
+    env:
+      DOCKER_BUILDKIT: "1"
+    retry:
+      automatic:
+        - exit_status: -1  # Agent was lost
+          limit: 2
+        - exit_status: -10  # Agent was lost
+          limit: 2

.buildkite/image_build/image_build_arm64.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+set -e
+
+if [[ $# -lt 3 ]]; then
+  echo "Usage: $0 <registry> <repo> <commit>"
+  exit 1
+fi
+
+REGISTRY=$1
+REPO=$2
+BUILDKITE_COMMIT=$3
+
+# authenticate with AWS ECR
+aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
+
+# skip build if image already exists
+if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64) ]]; then
+  echo "Image not found, proceeding with build..."
+else
+  echo "Image found"
+  exit 0
+fi
+
+# build (Grace/GH200 is the arm64 GPU target; sm_90)
+docker build --file docker/Dockerfile \
+  --platform linux/arm64 \
+  --build-arg max_jobs=16 \
+  --build-arg nvcc_threads=4 \
+  --build-arg torch_cuda_arch_list="9.0" \
+  --build-arg USE_SCCACHE=1 \
+  --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
+  --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64 \
+  --target test \
+  --progress plain .
+
+# push
+docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64

.buildkite/image_build/image_build_hpu.sh

@@ -11,7 +11,7 @@ REPO=$2
 BUILDKITE_COMMIT=$3
 
 # authenticate with AWS ECR
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
+aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
 
 # skip build if image already exists
 if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-hpu) ]]; then

.buildkite/image_build/image_build_xpu.sh

@@ -11,8 +11,8 @@ REPO=$2
 BUILDKITE_COMMIT=$3
 
 # authenticate with AWS ECR
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
-aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
+aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
+aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true
 
 # skip build if image already exists
 if ! docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-xpu &> /dev/null; then

.buildkite/intel_jobs/test-intel.yaml

@@ -39,7 +39,9 @@ steps:
         python3 examples/basic/offline_inference/generate.py --model nvidia/Llama-3.1-8B-Instruct-FP8 --block-size 64 --enforce-eager --quantization modelopt --kv-cache-dtype fp8 --attention-backend TRITON_ATTN --max-model-len 4096 &&
         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 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
+        '
   - label: "XPU V1 test"
     depends_on:
       - image-build-xpu
@@ -64,3 +66,22 @@ steps:
         pytest -v -s v1/test_serial_utils.py &&
         pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py &&
         pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py'
+  - label: "XPU server test"
+    depends_on:
+      - image-build-xpu
+    timeout_in_minutes: 30
+    device: intel_gpu
+    no_plugin: true
+    env:
+      REGISTRY: "public.ecr.aws/q9t5s3a7"
+      REPO: "vllm-ci-test-repo"
+    source_file_dependencies:
+      - vllm/
+      - .buildkite/intel_jobs/test-intel.yaml 
+    commands:
+      - >-
+        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+        'pip install av &&
+        cd tests &&
+        pytest -v -s entrypoints/openai/chat_completion/test_audio_in_video.py &&
+        pytest -v -s benchmarks/test_serve_cli.py'

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

@@ -1,77 +1,112 @@
-[
-    {
-        "test_name": "serving_llama8B_tp1_sharegpt",
-        "qps_list": [1, 4, 16, "inf"],
-        "server_parameters": {
-            "model": "meta-llama/Meta-Llama-3.1-8B-Instruct",
-            "tensor_parallel_size": 1,
-            "disable_log_stats": "",
-            "load_format": "dummy"
-        },
-        "client_parameters": {
-            "model": "meta-llama/Meta-Llama-3.1-8B-Instruct",
-            "backend": "vllm",
-            "dataset_name": "sharegpt",
-            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
-            "temperature": 0,
-            "num_prompts": 200
-        }
+{
+  "defaults": {
+    "qps_list": [
+      "inf"
+    ],
+    "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200],
+    "server_parameters": {
+      "model": "meta-llama/Llama-3.1-8B-Instruct",
+      "tensor_parallel_size": 1,
+      "dtype": "bfloat16"
     },
-    {
-        "test_name": "serving_llama70B_tp4_sharegpt",
-        "qps_list": [1, 4, 16, "inf"],
-        "server_parameters": {
-            "model": "meta-llama/Meta-Llama-3.1-70B-Instruct",
-            "tensor_parallel_size": 4,
-            "disable_log_stats": "",
-            "load_format": "dummy"
-        },
-        "client_parameters": {
-            "model": "meta-llama/Meta-Llama-3.1-70B-Instruct",
-            "backend": "vllm",
-            "dataset_name": "sharegpt",
-            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
-            "temperature": 0,
-            "num_prompts": 200
-        }
-    },
-    {
-        "test_name": "serving_mixtral8x7B_tp2_sharegpt",
-        "qps_list": [1, 4, 16, "inf"],
-        "server_parameters": {
-            "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
-            "tensor_parallel_size": 2,
-            "disable_log_stats": "",
-            "load_format": "dummy"
-        },
-        "client_parameters": {
-            "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
-            "backend": "vllm",
-            "dataset_name": "sharegpt",
-            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
-            "temperature": 0,
-            "num_prompts": 200
-        }
-    },
-    {
-        "test_name": "serving_llama70B_tp4_sharegpt_specdecode",
-        "qps_list": [2],
-        "server_parameters": {
-            "model": "meta-llama/Meta-Llama-3.1-70B-Instruct", 
-            "tensor_parallel_size": 4,
-            "speculative_config": {
-                "model": "turboderp/Qwama-0.5B-Instruct",
-                "num_speculative_tokens": 4,
-                "draft_tensor_parallel_size": 1
-            }
-        },
-        "client_parameters": {
-            "model": "meta-llama/Meta-Llama-3.1-70B-Instruct",
-            "backend": "vllm",
-            "dataset_name": "sharegpt",
-            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
-            "temperature": 0,
-            "num_prompts": 200 
-        }
+    "client_parameters": {
+      "model": "meta-llama/Llama-3.1-8B-Instruct",
+      "backend": "vllm",
+      "ignore-eos": "",
+      "temperature": 0,
+      "num_prompts": 200
     }
-]
+  },
+  "tests": [
+    {
+      "test_name": "serving_llama8B_tp1_sharegpt",
+      "server_parameters": {
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "dataset_name": "sharegpt",
+        "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
+      }
+    },
+    {
+        "dataset_name": "sharegpt",
+        "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
+      }
+    },
+    {
+      "test_name": "serving_llama8B_tp1_random_128_128",
+      "server_parameters": {
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_llama8B_tp1_random_128_2048",
+      "server_parameters": {
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 2048
+      }
+    },
+    {
+      "test_name": "serving_llama8B_tp1_random_2048_128",
+      "server_parameters": {
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "dataset_name": "random",
+        "random-input-len": 2048,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_llama8B_tp1_random_2048_2048",
+      "server_parameters": {
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "dataset_name": "random",
+        "random-input-len": 2048,
+        "random-output-len": 2048
+      }
+    },
+    {
+      "test_name": "serving_llama70B_tp4_random_128_128",
+      "server_parameters": {
+        "model": "meta-llama/Llama-3.3-70B-Instruct",
+        "async_scheduling": "",
+	"no_enable_prefix_caching": "",
+        "max_num_batched_tokens": 8192
+      },
+      "client_parameters": {
+        "model": "meta-llama/Llama-3.3-70B-Instruct",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_gemma4-e4b_tp1_random_128_128",
+      "server_parameters": {
+        "model": "google/gemma-4-E4B-it",
+	"enable_auto_tool_choice": "",
+	"tool_call_parser": "gemma4",
+	"chat_template": "examples/tool_chat_template_gemma4.jinja",
+	"reasoning_parser": "gemma4"
+      },
+      "client_parameters": {
+        "model": "google/gemma-4-E4B-it",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    }
+  ]
+}

.buildkite/release-pipeline.yaml

@@ -737,7 +737,7 @@ steps:
       - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
     env:
       S3_BUCKET: "vllm-wheels"
-      VARIANT: "rocm722"
+      VARIANT: "rocm723"
 
   # ROCm Job 6: Build ROCm Release Docker Image
   - label: ":docker: Build release image - x86_64 - ROCm"

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

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

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

@@ -35,25 +35,9 @@ export PYTHONPATH=".."
 # Helper Functions
 ###############################################################################
 
-cleanup_docker() {
-  # Get Docker's root directory
-  docker_root=$(docker info -f '{{.DockerRootDir}}')
-  if [ -z "$docker_root" ]; then
-    echo "Failed to determine Docker root directory."
-    exit 1
-  fi
-  echo "Docker root directory: $docker_root"
-
-  disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
-  threshold=70
-  if [ "$disk_usage" -gt "$threshold" ]; then
-    echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
-    docker image prune -f
-    docker volume prune -f && docker system prune --force --filter "until=72h" --all
-    echo "Docker images and volumes cleanup completed."
-  else
-    echo "Disk usage is below $threshold%. No cleanup needed."
-  fi
+report_docker_usage() {
+  echo "--- Docker usage"
+  docker system df || true
 }
 
 cleanup_network() {
@@ -254,8 +238,8 @@ re_quote_pytest_markers() {
 echo "--- ROCm info"
 rocminfo
 
-# --- Docker housekeeping ---
-cleanup_docker
+# --- Docker status ---
+report_docker_usage
 
 # --- Pull test image ---
 echo "--- Pulling container"
@@ -264,9 +248,17 @@ container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | hea
 docker pull "${image_name}"
 
 remove_docker_container() {
-  docker rm -f "${container_name}" || docker image rm -f "${image_name}" || true
+  # 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
+
+on_exit() {
+  local exit_code=$?
+  remove_docker_container
+  exit "$exit_code"
+}
+trap on_exit EXIT
 
 # --- Prepare commands ---
 echo "--- Running container"

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

@@ -352,17 +352,31 @@ if [[ -z "${ZE_AFFINITY_MASK:-}" ]]; then
   echo "Warning: ZE_AFFINITY_MASK is not set. Proceeding without device affinity." >&2
 fi
 
-docker run \
+export CMDS="${commands}"
+export HF_TOKEN ZE_AFFINITY_MASK
+
+{
+  flock 9
+  if ! docker image inspect "${IMAGE}" >/dev/null 2>&1; then
+    echo 'Image missing before container creation, pulling again...'
+    timeout 900 docker pull "${IMAGE}"
+  fi
+
+  docker create \
     --device /dev/dri:/dev/dri \
     --net=host \
     --ipc=host \
     --privileged \
     -v /dev/dri/by-path:/dev/dri/by-path \
-    -v ${HOME}/.cache/huggingface:/root/.cache/huggingface \
-    --entrypoint="" \
-    -e "HF_TOKEN=${HF_TOKEN:-}" \
-    -e "ZE_AFFINITY_MASK=${ZE_AFFINITY_MASK:-}" \
-    -e "CMDS=${commands}" \
+    -v "${HOME}/.cache/huggingface:/root/.cache/huggingface" \
+    --entrypoint='' \
+    -e HF_TOKEN \
+    -e ZE_AFFINITY_MASK \
+    -e CMDS \
     --name "${container_name}" \
-    "${image_name}" \
-    bash -c 'set -e; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"'
+    "${IMAGE}" \
+    bash -c 'set -e; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \
+    >/dev/null
+} 9>/tmp/docker-pull.lock
+
+docker start -a "${container_name}"

.buildkite/scripts/run-rust-frontend-cargo-ci.sh

@@ -0,0 +1,156 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+MODE="${1:-}"
+
+if [[ "$MODE" != "style-clippy" && "$MODE" != "test" ]]; then
+  echo "Usage: $0 {style-clippy|test}" >&2
+  exit 2
+fi
+
+ROOT_DIR="$(git rev-parse --show-toplevel)"
+cd "$ROOT_DIR"
+
+export CARGO_TERM_COLOR="${CARGO_TERM_COLOR:-always}"
+export CARGO_HOME="${CARGO_HOME:-$HOME/.cargo}"
+export RUSTUP_HOME="${RUSTUP_HOME:-$HOME/.rustup}"
+export PATH="$CARGO_HOME/bin:$PATH"
+
+log_section() {
+  echo "--- $*"
+}
+
+install_protoc() {
+  if command -v protoc >/dev/null 2>&1; then
+    return
+  fi
+
+  local version="${PROTOC_VERSION:-31.1}"
+  local arch
+  case "$(uname -m)" in
+    x86_64)
+      arch="x86_64"
+      ;;
+    aarch64|arm64)
+      arch="aarch_64"
+      ;;
+    *)
+      echo "Unsupported protoc architecture: $(uname -m)" >&2
+      return 1
+      ;;
+  esac
+
+  local url="https://github.com/protocolbuffers/protobuf/releases/download/v${version}/protoc-${version}-linux-${arch}.zip"
+  local tmp_dir
+  tmp_dir="$(mktemp -d)"
+
+  log_section "Installing protoc ${version}"
+  curl -L --proto '=https' --tlsv1.2 -sSf "$url" -o "$tmp_dir/protoc.zip"
+  mkdir -p "$CARGO_HOME/bin"
+  unzip -q "$tmp_dir/protoc.zip" bin/protoc 'include/*' -d "$CARGO_HOME"
+  chmod +x "$CARGO_HOME/bin/protoc"
+  rm -rf "$tmp_dir"
+}
+
+rust_toolchain() {
+  awk -F '"' '/channel[[:space:]]*=/ { print $2; exit }' rust-toolchain.toml
+}
+
+install_rust_toolchain() {
+  log_section "Installing Rust toolchain"
+  if ! command -v rustup >/dev/null 2>&1; then
+    curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \
+      | sh -s -- -y --profile minimal --default-toolchain none
+  fi
+
+  local toolchain
+  toolchain="$(rust_toolchain)"
+  rustup toolchain install "$toolchain" --profile minimal --component rustfmt,clippy
+  rustup component add --toolchain "$toolchain" rustfmt clippy
+}
+
+install_cargo_binstall() {
+  if command -v cargo-binstall >/dev/null 2>&1; then
+    return
+  fi
+
+  log_section "Installing cargo-binstall"
+  curl -L --proto '=https' --tlsv1.2 -sSf \
+    https://raw.githubusercontent.com/cargo-bins/cargo-binstall/main/install-from-binstall-release.sh \
+    | bash
+}
+
+install_cargo_sort() {
+  if command -v cargo-sort >/dev/null 2>&1; then
+    return
+  fi
+
+  log_section "Installing cargo-sort"
+  install_cargo_binstall
+  cargo binstall --no-confirm cargo-sort
+}
+
+install_cargo_nextest() {
+  if command -v cargo-nextest >/dev/null 2>&1; then
+    return
+  fi
+
+  log_section "Installing cargo-nextest"
+  install_cargo_binstall
+  cargo binstall --no-confirm --secure cargo-nextest
+}
+
+install_uv() {
+  if command -v uv >/dev/null 2>&1; then
+    return
+  fi
+
+  log_section "Installing uv"
+  curl -LsSf --proto '=https' --tlsv1.2 https://astral.sh/uv/install.sh \
+    | env UV_INSTALL_DIR="$CARGO_HOME/bin" sh
+}
+
+run_style_clippy() {
+  install_cargo_sort
+
+  log_section "Checking Rust formatting"
+  cargo fmt --manifest-path rust/Cargo.toml --all -- --check
+
+  log_section "Checking Cargo.toml ordering"
+  cargo sort --workspace --check rust
+
+  log_section "Running clippy"
+  cargo clippy \
+    --manifest-path rust/Cargo.toml \
+    --workspace \
+    --all-targets \
+    --all-features \
+    --locked \
+    -- \
+    -D warnings
+}
+
+run_tests() {
+  install_uv
+  install_cargo_nextest
+
+  log_section "Running cargo nextest"
+  cargo nextest run \
+    --manifest-path rust/Cargo.toml \
+    --workspace \
+    --all-features \
+    --locked \
+    --no-fail-fast
+}
+
+install_protoc
+install_rust_toolchain
+
+case "$MODE" in
+  style-clippy)
+    run_style_clippy
+    ;;
+  test)
+    run_tests
+    ;;
+esac

.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh

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

.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh

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

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

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

.buildkite/test-amd.yaml

@@ -139,19 +139,6 @@ steps:
   - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py
   - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)'
 
-#--------------------------------------------------------  mi250 · benchmarks  ---------------------------------------------------------#
-
-- label: Benchmarks # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  working_dir: "/vllm-workspace/.buildkite"
-  source_file_dependencies:
-  - benchmarks/
-  - vllm/platforms/rocm.py
-  commands:
-  - bash scripts/run-benchmarks.sh
-
 #----------------------------------------------------------  mi250 · compile  ----------------------------------------------------------#
 
 - label: PyTorch Compilation Unit Tests # TBD
@@ -485,7 +472,7 @@ steps:
   - pytest -v -s model_executor -m '(not slow_test)'
   - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
 
-#----------------------------------------------------------  mi250 · models  -----------------------------------------------------------#
+#------------------------------------------------------  mi250 · models / basic  -------------------------------------------------------#
 
 - label: Basic Models Test (Other CPU) # TBD
   timeout_in_minutes: 180
@@ -546,6 +533,8 @@ steps:
   commands:
   - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
 
+#-----------------------------------------------------  mi250 · models / language  -----------------------------------------------------#
+
 - label: Language Models Test (MTEB) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -590,6 +579,8 @@ steps:
   - pip freeze | grep -E 'torch'
   - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
 
+#----------------------------------------------------  mi250 · models / multimodal  ----------------------------------------------------#
+
 - label: Multi-Modal Models (Extended Generation 2) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -976,18 +967,6 @@ steps:
 
 #--------------------------------------------------------  mi300 · benchmarks  ---------------------------------------------------------#
 
-- label: Benchmarks # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  optional: true
-  working_dir: "/vllm-workspace/.buildkite"
-  source_file_dependencies:
-  - benchmarks/
-  - vllm/platforms/rocm.py
-  commands:
-  - bash scripts/run-benchmarks.sh
-
 - label: Benchmarks CLI Test # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1282,7 +1261,7 @@ steps:
   - tests/entrypoints/test_chat_utils
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
+  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
 
 - label: Entrypoints Integration (API Server openai - Part 2) # TBD
   timeout_in_minutes: 180
@@ -1505,7 +1484,7 @@ steps:
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt
 
-- label: DeepSeek V2-Lite Accuracy (4xH100-4xMI300) # TBD
+- label: DeepSeek V2-Lite Sync EPLB Accuracy (4xH100-4xMI300) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_4
@@ -1547,7 +1526,7 @@ steps:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
 
-- label: Qwen3-30B-A3B-FP8-block Accuracy (4xH100-4xMI300) # TBD
+- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (4xH100-4xMI300) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_4
@@ -1759,7 +1738,7 @@ steps:
   - pytest -v -s -x lora/test_gptoss_tp.py
   - pytest -v -s -x lora/test_qwen35_densemodel_lora.py
 
-#----------------------------------------------------------  mi300 · models  -----------------------------------------------------------#
+#-----------------------------------------------------  mi300 · models / language  -----------------------------------------------------#
 
 - label: Language Models Test (Extended Pooling)  # TBD
   timeout_in_minutes: 180
@@ -1787,6 +1766,8 @@ steps:
   - pip freeze | grep -E 'torch'
   - pytest -v -s models/language -m 'core_model and (not slow_test)'
 
+#----------------------------------------------------  mi300 · models / multimodal  ----------------------------------------------------#
+
 - label: Multi-Modal Models (Extended Generation 1) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1892,10 +1873,11 @@ steps:
   - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
   - pytest -v -s models/multimodal/processing/test_tensor_schema.py
 
-- label: Multi-Modal Processor (CPU) # TBD
+- label: Multi-Modal Processor (CPU) %N # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_1
+  parallelism: 4
   no_gpu: true
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -1905,7 +1887,9 @@ steps:
   - tests/models/registry.py
   commands:
   - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-  - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py
+  - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+
+#-----------------------------------------------------  mi300 · models / quantized  -----------------------------------------------------#
 
 - label: Quantized Models Test # TBD
   timeout_in_minutes: 180
@@ -1921,7 +1905,31 @@ steps:
   commands:
   - pytest -v -s models/quantization
 
-- label: Transformers Nightly Models # TBD
+#--------------------------------------------------  mi300 · models / transformers  ---------------------------------------------------#
+
+- label: Transformers Nightly Models (Shardable) %N # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  parallelism: 4
+  optional: true
+  working_dir: "/vllm-workspace/"
+  source_file_dependencies:
+  - vllm/model_executor/models/
+  - vllm/model_executor/model_loader/
+  - vllm/multimodal/
+  - vllm/model_executor/layers/
+  - vllm/v1/attention/backends/
+  - vllm/v1/attention/selector.py
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  - tests/models/
+  commands:
+  - pip install --upgrade git+https://github.com/huggingface/transformers
+  - pytest -v -s tests/models/test_initialization.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+  - pytest -v -s tests/models/multimodal/processing/ --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+
+- label: Transformers Nightly Models (Single) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_1
@@ -1940,9 +1948,7 @@ steps:
   - examples/
   commands:
   - pip install --upgrade git+https://github.com/huggingface/transformers
-  - pytest -v -s tests/models/test_initialization.py
   - pytest -v -s tests/models/test_transformers.py
-  - pytest -v -s tests/models/multimodal/processing/
   - pytest -v -s tests/models/multimodal/test_mapping.py
   - python3 examples/basic/offline_inference/chat.py
   - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
@@ -2003,7 +2009,7 @@ steps:
   - vllm/model_executor/layers
   - vllm/sampling_metadata.py
   - vllm/v1/sample/
-  - vllm/beam_search.py
+  - vllm/entrypoints/generate/beam_search/
   - tests/samplers
   - tests/conftest.py
   - vllm/_aiter_ops.py
@@ -2391,7 +2397,7 @@ steps:
   - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
   - DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
-- label: Hyrbid SSM NixlConnector PD accuracy tests (4 GPUs) # TBD
+- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_4
@@ -2593,7 +2599,7 @@ steps:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8
 
-#----------------------------------------------------------  mi325 · models  -----------------------------------------------------------#
+#-----------------------------------------------------  mi325 · models / language  -----------------------------------------------------#
 
 - label: Language Models Test (Extended Generation) # TBD
   timeout_in_minutes: 180
@@ -2624,6 +2630,8 @@ steps:
   - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
   - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
 
+#----------------------------------------------------  mi325 · models / multimodal  ----------------------------------------------------#
+
 - label: Multi-Modal Models (Extended Pooling) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
@@ -2695,19 +2703,35 @@ steps:
   optional: true
   working_dir: "/vllm-workspace/"
   source_file_dependencies:
+  - csrc/custom_quickreduce.cu
+  - csrc/ops.h
+  - csrc/torch_bindings.cpp
   - vllm/distributed/
-  - vllm/v1/distributed/
+  - vllm/model_executor/layers/
+  - vllm/entrypoints/llm.py
+  - vllm/config/parallel.py
   - vllm/model_executor/layers/fused_moe/
+  - vllm/v1/engine/
+  - vllm/v1/executor/
+  - vllm/v1/worker/
+  - vllm/v1/distributed/
   - vllm/v1/attention/backends/
   - vllm/v1/attention/selector.py
-  - tests/distributed/test_context_parallel.py
-  - tests/v1/distributed/test_dbo.py
-  - examples/features/data_parallel/data_parallel_offline.py
   - vllm/_aiter_ops.py
+  - vllm/_custom_ops.py
   - vllm/platforms/rocm.py
+  - vllm/envs.py
+  - examples/offline_inference/data_parallel.py
+  - tests/distributed/test_context_parallel.py
+  - tests/distributed/test_rocm_quick_reduce.py
+  - tests/distributed/test_quick_all_reduce.py
+  - tests/v1/distributed/test_dbo.py
+  - tests/utils.py
   commands:
   - pytest -v -s tests/distributed/test_context_parallel.py
   - pytest -v -s tests/v1/distributed/test_dbo.py
+  - pytest -v -s tests/distributed/test_rocm_quick_reduce.py
+  - pytest -v -s tests/distributed/test_quick_all_reduce.py
 
 #--------------------------------------------------------  mi355 · entrypoints  --------------------------------------------------------#
 
@@ -2744,7 +2768,7 @@ steps:
   - tests/entrypoints/test_chat_utils
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
+  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
 
 - label: Entrypoints Integration (API Server openai - Part 2) # TBD
   timeout_in_minutes: 180
@@ -2871,7 +2895,7 @@ steps:
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt
 
-- label: Qwen3-30B-A3B-FP8-block Accuracy (B200-MI355) # TBD
+- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200-MI355) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   agent_pool: mi355_2
@@ -3043,7 +3067,7 @@ steps:
   commands:
     - pytest -v -s kernels/moe/test_deepep_moe.py
 
-#----------------------------------------------------------  mi355 · models  -----------------------------------------------------------#
+#-----------------------------------------------------  mi355 · models / language  -----------------------------------------------------#
 
 - label: Language Models Test (Extended Generation) # TBD
   timeout_in_minutes: 180
@@ -3111,6 +3135,8 @@ steps:
   - pip freeze | grep -E 'torch'
   - pytest -v -s models/language -m 'core_model and (not slow_test)'
 
+#----------------------------------------------------  mi355 · models / multimodal  ----------------------------------------------------#
+
 - label: Multi-Modal Models (Extended Generation 1) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
@@ -3182,6 +3208,8 @@ steps:
   - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model
   - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model
 
+#-----------------------------------------------------  mi355 · models / quantized  -----------------------------------------------------#
+
 - label: Quantized Models Test # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]

.buildkite/test_areas/disaggregated.yaml

@@ -88,7 +88,7 @@ steps:
     - tests/v1/kv_connector/nixl_integration/
   commands:
     - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
+    - bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
 
 - label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
   key: multiconnector-nixl-offloading-pd-edge-cases-2-gpus

.buildkite/test_areas/e2e_integration.yaml

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

.buildkite/test_areas/engine.yaml

@@ -38,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
@@ -60,6 +60,7 @@ steps:
       - image-build-amd
 
 - label: e2e Core (1 GPU)
+  device: h200_35gb
   key: e2e-core-1-gpu
   timeout_in_minutes: 30
   source_file_dependencies:

.buildkite/test_areas/entrypoints.yaml

@@ -28,7 +28,8 @@ steps:
   - pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
+      soft_fail: true
       depends_on:
       - image-build-amd
 
@@ -42,10 +43,11 @@ steps:
   - tests/entrypoints/test_chat_utils
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
+  - 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
@@ -63,7 +65,8 @@ steps:
   - pytest -v -s entrypoints/test_chat_utils.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
+      soft_fail: true
       timeout_in_minutes: 60
       depends_on:
       - image-build-amd
@@ -82,12 +85,14 @@ steps:
   - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
+      soft_fail: true
       timeout_in_minutes: 60
       depends_on:
       - image-build-amd
 
 - label: Entrypoints Integration (API Server 2)
+  device: h200_35gb
   key: entrypoints-integration-api-server-2
   timeout_in_minutes: 130
   working_dir: "/vllm-workspace/tests"
@@ -103,11 +108,13 @@ steps:
   - pytest -v -s tool_use
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
+      soft_fail: true
       depends_on:
       - image-build-amd
 
 - label: Entrypoints Integration (Speech to Text)
+  device: h200_35gb
   key: entrypoints-integration-speech_to_text
   timeout_in_minutes: 50
   working_dir: "/vllm-workspace/tests"

.buildkite/test_areas/kernels.yaml

@@ -38,6 +38,28 @@ steps:
   commands:
     - pytest -v -s kernels/core/test_minimax_reduce_rms.py
 
+- label: Deepseek V4 Kernel Test (H100)
+  key: deepseek-v4-kernel-test-h100
+  timeout_in_minutes: 15
+  device: h100
+  source_file_dependencies:
+  - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+  - vllm/models/deepseek_v4/common/ops/
+  - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+  commands:
+    - pytest -v -s kernels/test_fused_deepseek_v4_*.py
+
+- label: Deepseek V4 Kernel Test (B200)
+  key: deepseek-v4-kernel-test-b200
+  timeout_in_minutes: 15
+  device: b200-k8s
+  source_file_dependencies:
+  - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+  - vllm/models/deepseek_v4/common/ops/
+  - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+  commands:
+    - pytest -v -s kernels/test_fused_deepseek_v4_*.py
+
 - label: Kernels Attention Test %N
   key: kernels-attention-test
   timeout_in_minutes: 35
@@ -64,7 +86,7 @@ steps:
   parallelism: 2
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       source_file_dependencies:
       - csrc/quantization/
       - vllm/model_executor/layers/quantization
@@ -150,9 +172,12 @@ steps:
   - csrc/quantization/fp4/
   - csrc/attention/mla/
   - csrc/quantization/cutlass_w8a8/moe/
-  - vllm/model_executor/layers/fused_moe/cutlass_moe.py
-  - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
-  - vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py
+  - vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
+  - vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py
+  - vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+  - vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+  - vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py
+  - vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_two_sided.py
   - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
   - vllm/v1/attention/backends/flashinfer.py
   - vllm/v1/attention/backends/mla/cutlass_mla.py
@@ -185,6 +210,7 @@ steps:
     - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
     - pytest -v -s tests/kernels/moe/test_flashinfer.py
     - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
+    - pytest -v -s tests/kernels/moe/test_trtllm_nvfp4_moe.py
     - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
     # e2e
     - pytest -v -s tests/models/quantization/test_nvfp4.py

.buildkite/test_areas/lm_eval.yaml

@@ -3,6 +3,7 @@ depends_on:
   - image-build
 steps:
 - label: LM Eval Small Models
+  device: h200_35gb
   key: lm-eval-small-models
   timeout_in_minutes: 75
   source_file_dependencies:
@@ -152,6 +153,7 @@ steps:
     - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt
 
 - label: MRCR Eval Small Models
+  device: h200_35gb
   timeout_in_minutes: 30
   source_file_dependencies:
   - tests/evals/mrcr/

.buildkite/test_areas/lora.yaml

@@ -3,6 +3,7 @@ depends_on:
   - image-build
 steps:
 - label: LoRA %N
+  device: h200_35gb
   key: lora
   timeout_in_minutes: 30
   source_file_dependencies:

.buildkite/test_areas/misc.yaml

@@ -3,6 +3,7 @@ depends_on:
   - image-build
 steps:
 - label: V1 Spec Decode
+  device: h200_35gb
   key: v1-spec-decode
   timeout_in_minutes: 30
   source_file_dependencies:
@@ -51,7 +52,7 @@ steps:
     - pytest -v -s v1/test_outputs.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
 
@@ -165,6 +166,7 @@ steps:
   working_dir: "/vllm-workspace/tests" # optional
 
 - label: Examples
+  device: h200_35gb
   key: examples
   timeout_in_minutes: 45
   working_dir: "/vllm-workspace/examples"
@@ -235,6 +237,7 @@ steps:
   - bash standalone_tests/python_only_compile.sh
 
 - label: Async Engine, Inputs, Utils, Worker
+  device: h200_35gb
   key: async-engine-inputs-utils-worker
   timeout_in_minutes: 50
   source_file_dependencies:
@@ -364,8 +367,10 @@ steps:
     - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
     - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
     - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
+    - pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py
   
 - label: Acceptance Length Test (Large Models) # optional
+  device: h200_35gb
   key: acceptance-length-test-large-models
   timeout_in_minutes: 25
   gpu: h100

.buildkite/test_areas/model_executor.yaml

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

.buildkite/test_areas/model_runner_v2.yaml

@@ -3,6 +3,7 @@ depends_on:
   - image-build
 steps:
 - label: Model Runner V2 Core Tests
+  device: h200_35gb
   key: model-runner-v2-core-tests
   timeout_in_minutes: 45
   source_file_dependencies:
@@ -26,6 +27,7 @@ steps:
   - pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"
 
 - label: Model Runner V2 Examples
+  device: h200_35gb
   key: model-runner-v2-examples
   timeout_in_minutes: 45
   working_dir: "/vllm-workspace/examples"
@@ -99,6 +101,7 @@ steps:
     - pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
 
 - label: Model Runner V2 Spec Decode
+  device: h200_35gb
   key: model-runner-v2-spec-decode
   timeout_in_minutes: 30
   working_dir: "/vllm-workspace/tests"

.buildkite/test_areas/models_basic.yaml

@@ -18,6 +18,7 @@ steps:
     torch_nightly: {}
 
 - label: Basic Models Tests (Extra Initialization) %N
+  device: h200_35gb
   key: basic-models-tests-extra-initialization
   timeout_in_minutes: 45
   source_file_dependencies:
@@ -34,6 +35,7 @@ steps:
     torch_nightly: {}
 
 - label: Basic Models Tests (Other)
+  device: h200_35gb
   key: basic-models-tests-other
   timeout_in_minutes: 45
   source_file_dependencies:
@@ -56,35 +58,3 @@ steps:
   device: cpu-small
   commands:
     - pytest -v -s models/test_utils.py models/test_vision.py
-
-- label: Transformers Nightly Models
-  key: transformers-nightly-models
-  working_dir: "/vllm-workspace/"
-  optional: true
-  soft_fail: true
-  commands:
-    - pip install --upgrade git+https://github.com/huggingface/transformers
-    - pytest -v -s tests/models/test_initialization.py
-    - pytest -v -s tests/models/test_transformers.py
-    - pytest -v -s tests/models/multimodal/processing/
-    - pytest -v -s tests/models/multimodal/test_mapping.py
-    - python3 examples/basic/offline_inference/chat.py
-    - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
-    # Whisper needs spawn method to avoid deadlock
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
-
-- label: Transformers Backward Compatibility Models Test
-  key: transformers-backward-compatibility-models-test
-  working_dir: "/vllm-workspace/"
-  optional: true
-  soft_fail: true
-  commands:
-    - pip install transformers==4.57.5
-    - pytest -v -s tests/models/test_initialization.py
-    - pytest -v -s tests/models/test_transformers.py
-    - pytest -v -s tests/models/multimodal/processing/
-    - pytest -v -s tests/models/multimodal/test_mapping.py
-    - python3 examples/basic/offline_inference/chat.py
-    - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
-    # Whisper needs spawn method to avoid deadlock
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper

.buildkite/test_areas/models_language.yaml

@@ -50,7 +50,7 @@ steps:
   mirror:
     torch_nightly: {}
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
       commands:
@@ -59,6 +59,7 @@ steps:
       - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
 
 - label: Language Models Test (Extended Generation) # 80min
+  device: h200_35gb
   key: language-models-test-extended-generation
   timeout_in_minutes: 110
   optional: true
@@ -84,6 +85,7 @@ steps:
     - pytest -v -s models/language/generation_ppl_test
 
 - label: Language Models Test (Extended Pooling)  # 36min
+  device: h200_35gb
   key: language-models-test-extended-pooling
   timeout_in_minutes: 50
   optional: true
@@ -94,7 +96,7 @@ steps:
     - pytest -v -s models/language/pooling -m 'not core_model'
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 100
       depends_on:
       - image-build-amd

.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,11 +33,12 @@ 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
 
 - label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl"
+  device: h200_35gb
   key: multi-modal-models-standard-3-llava-qwen2-vl
   timeout_in_minutes: 45
   source_file_dependencies:
@@ -49,11 +50,12 @@ 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
 
 - label: "Multi-Modal Models (Standard) 4: other + whisper"
+  device: h200_35gb
   key: multi-modal-models-standard-4-other-whisper
   timeout_in_minutes: 45
   source_file_dependencies:
@@ -92,6 +94,7 @@ steps:
     - pytest -v -s models/multimodal/processing/test_tensor_schema.py
 
 - label: Multi-Modal Accuracy Eval (Small Models) # 50min
+  device: h200_35gb
   key: multi-modal-accuracy-eval-small-models
   timeout_in_minutes: 70
   working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
@@ -115,11 +118,12 @@ steps:
     - pytest -v -s models/multimodal/test_mapping.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
 
 - label: Multi-Modal Models (Extended Generation 2)
+  device: h200_35gb
   key: multi-modal-models-extended-generation-2
   optional: true
   source_file_dependencies:
@@ -130,6 +134,7 @@ steps:
     - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'
 
 - label: Multi-Modal Models (Extended Generation 3)
+  device: h200_35gb
   key: multi-modal-models-extended-generation-3
   optional: true
   source_file_dependencies:

.buildkite/test_areas/pytorch.yaml

@@ -3,6 +3,7 @@ depends_on:
   - image-build
 steps:
 - label: PyTorch Compilation Unit Tests
+  device: h200_35gb
   key: pytorch-compilation-unit-tests
   timeout_in_minutes: 10
   source_file_dependencies:

.buildkite/test_areas/rust_frontend.yaml

@@ -0,0 +1,107 @@
+group: Rust Frontend E2E
+depends_on:
+  - image-build
+steps:
+- label: Rust Frontend OpenAI Coverage
+  timeout_in_minutes: 90
+  device: h200_18gb
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - rust/
+  - vllm/benchmarks/
+  - vllm/entrypoints/openai/
+  - vllm/entrypoints/serve/
+  - vllm/v1/sample/
+  - tests/utils.py
+  - tests/benchmarks/test_serve_cli.py
+  - tests/entrypoints/openai/chat_completion/test_chat_completion.py
+  # - tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py
+
+  # - tests/entrypoints/openai/completion/test_prompt_validation.py
+  - tests/entrypoints/openai/completion/test_shutdown.py
+  # - tests/entrypoints/openai/test_return_token_ids.py
+  # - tests/entrypoints/openai/test_uds.py
+  - tests/v1/sample/test_logprobs_e2e.py
+  commands:
+  - export VLLM_USE_RUST_FRONTEND=1
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
+  - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py
+  # - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not invalid"
+
+  # - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
+  - pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
+  # - pytest -v -s entrypoints/openai/test_return_token_ids.py
+  # - pytest -v -s entrypoints/openai/test_uds.py
+  - pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server"
+
+- label: Rust Frontend Serve/Admin Coverage
+  timeout_in_minutes: 60
+  device: h200_18gb
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - rust/
+  - vllm/entrypoints/openai/
+  - vllm/entrypoints/serve/
+  - vllm/v1/engine/
+  - tests/utils.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/instrumentator/test_sleep.py
+  commands:
+  - export VLLM_USE_RUST_FRONTEND=1
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  # - 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/instrumentator/test_sleep.py
+
+- label: Rust Frontend Core Correctness
+  timeout_in_minutes: 30
+  device: h200_18gb
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - rust/
+  - vllm/entrypoints/openai/
+  - tests/utils.py
+  - tests/entrypoints/openai/correctness/test_lmeval.py
+  commands:
+  - export VLLM_USE_RUST_FRONTEND=1
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
+
+- label: Rust Frontend Tool Use
+  timeout_in_minutes: 60
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - rust/
+  - vllm/entrypoints/openai/
+  - vllm/tool_parsers/
+  - tests/utils.py
+  - tests/tool_use/
+  commands:
+  - export VLLM_USE_RUST_FRONTEND=1
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice"
+
+- label: Rust Frontend Distributed
+  timeout_in_minutes: 30
+  num_devices: 4
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - rust/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/executor/
+  - vllm/v1/engine/
+  - vllm/v1/worker/
+  - tests/utils.py
+  - tests/v1/distributed/test_internal_lb_dp.py
+  commands:
+  - export VLLM_USE_RUST_FRONTEND=1
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - export NCCL_CUMEM_HOST_ENABLE=0
+  - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info"

.buildkite/test_areas/rust_frontend_cargo.yaml

@@ -0,0 +1,30 @@
+group: Rust Frontend Cargo
+depends_on: []
+steps:
+- label: Rust Frontend Cargo Style + Clippy
+  key: rust-frontend-cargo-style-clippy
+  depends_on: []
+  timeout_in_minutes: 30
+  device: cpu-medium
+  no_plugin: true
+  source_file_dependencies:
+  - rust/
+  - rust-toolchain.toml
+  - .buildkite/test_areas/rust_frontend_cargo.yaml
+  - .buildkite/scripts/run-rust-frontend-cargo-ci.sh
+  commands:
+  - .buildkite/scripts/run-rust-frontend-cargo-ci.sh style-clippy
+
+- label: Rust Frontend Cargo Tests
+  key: rust-frontend-cargo-tests
+  depends_on: []
+  timeout_in_minutes: 30
+  device: cpu-medium
+  no_plugin: true
+  source_file_dependencies:
+  - rust/
+  - rust-toolchain.toml
+  - .buildkite/test_areas/rust_frontend_cargo.yaml
+  - .buildkite/scripts/run-rust-frontend-cargo-ci.sh
+  commands:
+  - .buildkite/scripts/run-rust-frontend-cargo-ci.sh test

.buildkite/test_areas/samplers.yaml

@@ -3,6 +3,7 @@ depends_on:
   - image-build
 steps:
 - label: Samplers Test
+  device: h200_35gb
   key: samplers-test
   timeout_in_minutes: 75
   source_file_dependencies:
@@ -10,6 +11,7 @@ steps:
   - vllm/sampling_metadata.py
   - tests/samplers
   - tests/conftest.py
+  - vllm/entrypoints/generate/beam_search
   commands:
     # VLLM_USE_FLASHINFER_SAMPLER defaults to 1 now, so we need to pin both
     # values explicitly to still cover the PyTorch-native (Triton) path.

.buildkite/test_areas/spec_decode.yaml

@@ -32,6 +32,7 @@ steps:
   source_file_dependencies:
     - vllm/v1/spec_decode/
     - vllm/v1/worker/gpu/spec_decode/
+    - vllm/v1/attention/backends/
     - vllm/transformers_utils/configs/speculators/
     - tests/v1/e2e/spec_decode/
   commands:

.dockerignore

@@ -2,6 +2,7 @@
 /build
 dist
 vllm/*.so
+vllm/vllm-rs
 
 # Byte-compiled / optimized / DLL files
 __pycache__/
@@ -31,3 +32,4 @@ share/python-wheels/
 .installed.cfg
 *.egg
 MANIFEST
+rust/target/

.github/CODEOWNERS

@@ -40,6 +40,12 @@
 /vllm/entrypoints/chat_utils.py @DarkLight1337
 /vllm/entrypoints/llm.py @DarkLight1337
 
+# Rust Frontend
+/rust/ @BugenZhao @njhill
+/build_rust.sh @BugenZhao @njhill
+/rust-toolchain.toml @BugenZhao @njhill
+/.buildkite/test_areas/rust* @BugenZhao @njhill
+
 # Input/Output Processing
 /vllm/sampling_params.py @njhill @NickLucche
 /vllm/pooling_params.py @noooop @DarkLight1337
@@ -72,21 +78,23 @@
 /vllm/v1/worker/gpu/kv_connector.py @orozery
 
 # CI & building
-/.buildkite @Harry-Chen
-/docker/Dockerfile @Harry-Chen
+/.buildkite @Harry-Chen @khluu
+/docker/Dockerfile @Harry-Chen @khluu
+/pyproject.toml @khluu
+/setup.py @khluu
 
 # Test ownership
-/.buildkite/lm-eval-harness @mgoin 
+/.buildkite/lm-eval-harness @mgoin
 /tests/distributed/test_multi_node_assignment.py @youkaichao
 /tests/distributed/test_pipeline_parallel.py @youkaichao
 /tests/distributed/test_same_node.py @youkaichao
-/tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche
-/tests/evals @mgoin @vadiklyutiy
-/tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256 @zyongye
+/tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche @AndreasKaratzas
+/tests/evals @mgoin @vadiklyutiy @AndreasKaratzas
+/tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256 @zyongye @AndreasKaratzas
 /tests/kernels/ir @ProExpertProg @tjtanaa
-/tests/models @DarkLight1337 @ywang96
+/tests/models @DarkLight1337 @ywang96 @AndreasKaratzas
 /tests/multimodal @DarkLight1337 @ywang96 @NickLucche
-/tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye
+/tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye @AndreasKaratzas
 /tests/test_inputs.py @DarkLight1337 @ywang96
 /tests/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm
 /tests/v1/structured_output @mgoin @russellb @aarnphm
@@ -153,9 +161,7 @@ mkdocs.yaml @hmellor
 /vllm/model_executor/models/deepseek_mtp.py @luccafong
 
 # DeepseekV4-specific files
-/vllm/v1/attention/ops/deepseek_v4_ops @zyongye
-/vllm/model_executor/layers/deepseek_compressor.py @zyongye
-/vllm/model_executor/layers/deepseek_v4_attention.py @zyongye
+/vllm/models/deepseek_v4 @zyongye
 /vllm/model_executor/layers/sparse_attn_indexer.py @zyongye
 
 # Mistral-specific files
@@ -172,21 +178,21 @@ mkdocs.yaml @hmellor
 /vllm/model_executor/layers/fla @ZJY0516 @vadiklyutiy
 
 # ROCm related: specify owner with write access to notify AMD folks for careful code review
-/vllm/**/*rocm* @tjtanaa
-/docker/Dockerfile.rocm* @gshtras @tjtanaa
-/vllm/v1/attention/backends/rocm*.py @gshtras @tjtanaa
-/vllm/v1/attention/backends/mla/rocm*.py @gshtras @tjtanaa
-/vllm/v1/attention/ops/rocm*.py @gshtras @tjtanaa
-/vllm/model_executor/layers/fused_moe/rocm*.py @gshtras @tjtanaa
-/csrc/rocm @gshtras @tjtanaa
-/requirements/*rocm* @tjtanaa
-/tests/**/*rocm* @tjtanaa
+/vllm/**/*rocm* @tjtanaa @dllehr-amd
+/docker/Dockerfile.rocm* @tjtanaa @dllehr-amd @AndreasKaratzas
+/vllm/v1/attention/backends/rocm*.py @tjtanaa @dllehr-amd
+/vllm/v1/attention/backends/mla/rocm*.py @tjtanaa @dllehr-amd
+/vllm/v1/attention/ops/rocm*.py @tjtanaa @dllehr-amd
+/vllm/model_executor/layers/fused_moe/rocm*.py @tjtanaa @dllehr-amd
+/csrc/rocm @tjtanaa @dllehr-amd
+/requirements/*rocm* @tjtanaa @AndreasKaratzas
+/tests/**/*rocm* @tjtanaa @AndreasKaratzas
 /docs/**/*rocm* @tjtanaa
 /vllm/**/*quark* @tjtanaa
-/tests/**/*quark* @tjtanaa
+/tests/**/*quark* @tjtanaa @AndreasKaratzas
 /docs/**/*quark* @tjtanaa
-/vllm/**/*aiter* @tjtanaa
-/tests/**/*aiter* @tjtanaa
+/vllm/**/*aiter* @tjtanaa @AndreasKaratzas
+/tests/**/*aiter* @tjtanaa @AndreasKaratzas
 
 # TPU
 /vllm/v1/worker/tpu* @NickLucche

.github/mergify.yml

@@ -103,6 +103,19 @@ pull_request_rules:
       add:
         - frontend
 
+- name: label-rust
+  description: Automatically apply rust label
+  conditions:
+    - label != stale
+    - or:
+      - files~=(?i)rust
+      - title~=(?i)rust
+      - title~=(?i)vllm-rs
+  actions:
+    label:
+      add:
+        - rust
+
 - name: label-llama
   description: Automatically apply llama label
   conditions:

.gitignore

@@ -26,6 +26,9 @@ __pycache__/
 # C extensions
 *.so
 
+# Rust binaries
+vllm/vllm-rs
+
 # Distribution / packaging
 .Python
 build/

.pre-commit-config.yaml

@@ -222,6 +222,12 @@ repos:
     name: Update Dockerfile dependency graph
     entry: tools/pre_commit/update-dockerfile-graph.sh
     language: script
+  - id: test-nonroot-entrypoint
+    name: Test non-root entrypoint wrapper
+    entry: bash docker/entrypoints/test_vllm_nonroot_entrypoint.sh
+    language: system
+    pass_filenames: false
+    files: ^docker/entrypoints/(vllm-nonroot-entrypoint|test_vllm_nonroot_entrypoint)\.sh$
   - id: check-forbidden-imports
     name: Check for forbidden imports
     entry: python tools/pre_commit/check_forbidden_imports.py
@@ -256,6 +262,32 @@ repos:
     entry: python tools/pre_commit/check_boolean_context_manager.py
     language: python
     types: [python]
+  # Rust hooks. These shell out to `cargo`; tools/pre_commit/rust-check.sh
+  # skips with a warning when cargo is not installed.
+  - id: rust-cargo-autoinherit
+    name: Rust - Normalize Cargo manifests with autoinherit
+    entry: tools/pre_commit/rust-check.sh autoinherit --prefer-simple-dotted
+    language: script
+    pass_filenames: false
+    require_serial: true
+    stages: [pre-commit] # Only run locally as Buildkite will cover this
+    files: ^rust/(Cargo\.toml|src/.*/Cargo\.toml)$
+  - id: rust-cargo-sort
+    name: Rust - Sort Cargo manifest sections
+    entry: tools/pre_commit/rust-check.sh sort --workspace
+    language: script
+    pass_filenames: false
+    require_serial: true
+    stages: [pre-commit] # Only run locally as Buildkite will cover this
+    files: ^rust/(Cargo\.toml|src/.*/Cargo\.toml)$
+  - id: rust-cargo-fmt
+    name: Rust - Format code
+    entry: tools/pre_commit/rust-check.sh fmt
+    language: script
+    pass_filenames: false
+    require_serial: true
+    stages: [pre-commit] # Only run locally as Buildkite will cover this
+    files: ^rust/.*(\.rs|Cargo\.toml|rustfmt\.toml)$
   # Keep `suggestion` last
   - id: suggestion
     name: Suggestion

.readthedocs.yaml

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

AGENTS.md

@@ -101,6 +101,8 @@ pre-commit run ruff-check --all-files
 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.
+
 ### Commit messages
 
 Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:

CMakeLists.txt

@@ -144,14 +144,14 @@ endif()
 # Set up GPU language and check the torch version and warn if it isn't
 # what is expected.
 #
-if (NOT HIP_FOUND AND CUDA_FOUND)
+if (NOT HIP_FOUND AND NOT PYTORCH_FOUND_HIP AND CUDA_FOUND)
   set(VLLM_GPU_LANG "CUDA")
 
   if (NOT Torch_VERSION VERSION_EQUAL ${TORCH_SUPPORTED_VERSION_CUDA})
     message(WARNING "Pytorch version ${TORCH_SUPPORTED_VERSION_CUDA} "
       "expected for CUDA build, saw ${Torch_VERSION} instead.")
   endif()
-elseif(HIP_FOUND)
+elseif(HIP_FOUND OR PYTORCH_FOUND_HIP)
   set(VLLM_GPU_LANG "HIP")
 
   # Importing torch recognizes and sets up some HIP/ROCm configuration but does
@@ -305,26 +305,8 @@ endif()
 #
 
 set(VLLM_EXT_SRC
-  "csrc/mamba/mamba_ssm/selective_scan_fwd.cu"
-  "csrc/cache_kernels.cu"
-  "csrc/cache_kernels_fused.cu"
-  "csrc/attention/paged_attention_v1.cu"
-  "csrc/attention/paged_attention_v2.cu"
-  "csrc/attention/merge_attn_states.cu"
-  "csrc/pos_encoding_kernels.cu"
-  "csrc/activation_kernels.cu"
-  "csrc/layernorm_kernels.cu"
-  "csrc/fused_qknorm_rope_kernel.cu"
-  "csrc/layernorm_quant_kernels.cu"
-  "csrc/sampler.cu"
-  "csrc/topk.cu"
   "csrc/cuda_view.cu"
-  "csrc/quantization/gptq/q_gemm.cu"
-  "csrc/quantization/w8a8/int8/scaled_quant.cu"
-  "csrc/quantization/w8a8/fp8/common.cu"
-  "csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
   "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
-  "csrc/quantization/gguf/gguf_kernel.cu"
   "csrc/quantization/activation_kernels.cu"
   "csrc/cuda_utils_kernels.cu"
   "csrc/custom_all_reduce.cu"
@@ -333,7 +315,8 @@ set(VLLM_EXT_SRC
 
 if(VLLM_GPU_LANG STREQUAL "CUDA")
   list(APPEND VLLM_EXT_SRC
-    "csrc/minimax_reduce_rms_kernel.cu")
+    "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")
 
@@ -379,16 +362,30 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   # are not supported by Machete yet.
 
   # marlin arches for fp16 output
-  cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX" "${CUDA_ARCHS}")
+  # Family-conditional 12.0f (one cubin for SM12x family) requires CUDA >= 13.0;
+  # fall back to architecture-specific 12.0a;12.1a on CUDA < 13.0 (e.g. 12.8).
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
+  endif()
   # marlin has limited support for turing
   cuda_archs_loose_intersection(MARLIN_SM75_ARCHS "7.5" "${CUDA_ARCHS}")
   # marlin arches for bf16 output (we need 9.0 for bf16 atomicAdd PTX)
-  cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
+  endif()
   # marlin arches for fp8 input
   # - sm80 doesn't support fp8 computation
   # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
   # so we only enable fp8 computation for SM89 (e.g. RTX 40x0)  and 12.0 (e.g. RTX 50x0)
-  cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0;12.1" "${CUDA_ARCHS}")
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0a;12.1a" "${CUDA_ARCHS}")
+  endif()
   # marlin arches for other files
   cuda_archs_loose_intersection(MARLIN_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}")
 
@@ -628,33 +625,47 @@ define_extension_target(
 # Setting this variable sidesteps the issue by calling the driver directly.
 target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
 
-# add OR VLLM_GPU_LANG STREQUAL "HIP" here once
-# https://github.com/vllm-project/vllm/issues/35163 is resolved
-if(VLLM_GPU_LANG STREQUAL "CUDA")
+if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
   #
   # _C_stable_libtorch extension (ops registered via STABLE_TORCH_LIBRARY)
   #
   set(VLLM_STABLE_EXT_SRC
     "csrc/libtorch_stable/torch_bindings.cpp"
-    "csrc/cutlass_extensions/common.cpp"
-    "csrc/cuda_utils_kernels.cu"
-    "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/activation_kernels.cu"
+    "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
+    "csrc/libtorch_stable/quantization/w8a8/fp8/common.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/attention/merge_attn_states.cu"
+    "csrc/libtorch_stable/sampler.cu"
+    "csrc/libtorch_stable/topk.cu"
+    "csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
+    "csrc/libtorch_stable/attention/paged_attention_v1.cu"
+    "csrc/libtorch_stable/attention/paged_attention_v2.cu"
+    "csrc/libtorch_stable/cache_kernels.cu"
+    "csrc/libtorch_stable/cache_kernels_fused.cu")
 
   if(VLLM_GPU_LANG STREQUAL "CUDA")
     list(APPEND VLLM_STABLE_EXT_SRC
+      "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/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")
-  endif()
 
-  if(VLLM_GPU_LANG STREQUAL "CUDA")
     set_gencode_flags_for_srcs(
       SRCS "${VLLM_STABLE_EXT_SRC}"
       CUDA_ARCHS "${CUDA_ARCHS}")
-  endif()
 
   # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
@@ -674,6 +685,22 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
                    "in CUDA target architectures.")
   endif()
 
+  # FP32 router GEMM (H=3072, E=256, M<=32). Requires SM90+ and CUDA >= 12.0.
+  cuda_archs_sm90plus(FP32_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}")
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND FP32_ROUTER_GEMM_ARCHS)
+    set(SRCS
+      "csrc/libtorch_stable/fp32_router_gemm_entry.cu"
+      "csrc/libtorch_stable/fp32_router_gemm.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${SRCS}"
+      CUDA_ARCHS "${FP32_ROUTER_GEMM_ARCHS}")
+    list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
+    message(STATUS "Building fp32_router_gemm for archs: ${FP32_ROUTER_GEMM_ARCHS}")
+  else()
+    message(STATUS "Not building fp32_router_gemm as no compatible archs found "
+                   "(requires SM90+ and CUDA >= 12.0).")
+  endif()
+
   # Only build AllSpark kernels if we are building for at least some compatible archs.
   cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
   if (ALLSPARK_ARCHS)
@@ -915,13 +942,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       SRCS "${SRCS}"
       CUDA_ARCHS "${FP4_ARCHS}")
     list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
-    # nvfp4_kv_cache_kernels uses non-stable torch API and is called directly
-    # from cache_kernels.cu, so it belongs in _C rather than _C_stable.
-    set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
+    set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
     set_gencode_flags_for_srcs(
       SRCS "${NVFP4_KV_SRC}"
       CUDA_ARCHS "${FP4_ARCHS}")
-    target_sources(_C PRIVATE ${NVFP4_KV_SRC})
+    list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}")
     target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
@@ -951,11 +976,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       SRCS "${SRCS}"
       CUDA_ARCHS "${FP4_ARCHS}")
     list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
-    set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
+    set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
     set_gencode_flags_for_srcs(
       SRCS "${NVFP4_KV_SRC}"
       CUDA_ARCHS "${FP4_ARCHS}")
-    target_sources(_C PRIVATE ${NVFP4_KV_SRC})
+    list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}")
     target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
@@ -1034,6 +1059,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     message(STATUS "Building hadacore")
   endif()
 
+  # if CUDA endif
+  endif()
+
   message(STATUS "Enabling C_stable extension.")
   define_extension_target(
     _C_stable_libtorch
@@ -1053,13 +1081,34 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   target_compile_definitions(_C_stable_libtorch PRIVATE
     TORCH_TARGET_VERSION=0x020A000000000000ULL)
 
-  # Needed to use cuda APIs from C-shim
-  target_compile_definitions(_C_stable_libtorch PRIVATE
-    USE_CUDA)
+  # Needed to use cuda/hip APIs from C-shim
+  if(VLLM_GPU_LANG STREQUAL "CUDA")
+    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")
+    target_compile_definitions(_C_stable_libtorch PRIVATE USE_ROCM)
+  endif()
 
-  # Needed by CUTLASS kernels
-  target_compile_definitions(_C_stable_libtorch PRIVATE
-    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
+  # On ROCm, _C_stable_libtorch calls raw HIP APIs (e.g. hipGetDevice in
+  # get_device_prop()) which must resolve to the same libamdhip64.so that
+  # PyTorch uses.  When PyTorch bundles its own copy (pip/conda wheels),
+  # the raw HIP calls would otherwise resolve to the system ROCm copy,
+  # initializing a second HIP runtime that corrupts device state (wrong
+  # device on DeviceGuard, core dumps on multi-GPU tests).
+  #
+  # If PyTorch doesn't bundle libamdhip64 (built from source against system
+  # ROCm), there is only one copy in the process and no action is needed —
+  # the HIP compiler already links the system libamdhip64 automatically.
+  if(VLLM_GPU_LANG STREQUAL "HIP")
+    find_library(_STABLE_TORCH_AMDHIP64 amdhip64
+      PATHS "${TORCH_INSTALL_PREFIX}/lib" NO_DEFAULT_PATH)
+    if(_STABLE_TORCH_AMDHIP64)
+      message(STATUS "Found PyTorch-bundled libamdhip64 at ${_STABLE_TORCH_AMDHIP64}")
+      target_link_libraries(_C_stable_libtorch PRIVATE ${_STABLE_TORCH_AMDHIP64})
+    endif()
+  endif()
 endif()
 
 #
@@ -1102,7 +1151,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   # moe marlin arches
   # note that we always set `use_atomic_add=False` for moe marlin now,
   # so we don't need 9.0 for bf16 atomicAdd PTX
-  cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX" "${CUDA_ARCHS}")
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX;12.0f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
+  endif()
   # moe marlin has limited support for turing
   cuda_archs_loose_intersection(MARLIN_MOE_SM75_ARCHS "7.5" "${CUDA_ARCHS}")
   # moe marlin arches for fp8 input
@@ -1205,34 +1258,22 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
                    " in CUDA target architectures")
   endif()
 
-  # DeepSeek V3 router GEMM kernel - requires SM90+
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
-  endif()
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_ROUTER_GEMM_ARCHS)
+  # DeepSeek V3 router GEMM kernel requires SM90+ and CUDA >= 12.0.
+  # (fp32_router_gemm has been migrated to _C_stable_libtorch above.)
+  cuda_archs_sm90plus(SM90PLUS_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}")
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SM90PLUS_ROUTER_GEMM_ARCHS)
     set(DSV3_ROUTER_GEMM_SRC
       "csrc/moe/dsv3_router_gemm_entry.cu"
       "csrc/moe/dsv3_router_gemm_float_out.cu"
       "csrc/moe/dsv3_router_gemm_bf16_out.cu")
     set_gencode_flags_for_srcs(
       SRCS "${DSV3_ROUTER_GEMM_SRC}"
-      CUDA_ARCHS "${DSV3_ROUTER_GEMM_ARCHS}")
+      CUDA_ARCHS "${SM90PLUS_ROUTER_GEMM_ARCHS}")
     list(APPEND VLLM_MOE_EXT_SRC "${DSV3_ROUTER_GEMM_SRC}")
-    message(STATUS "Building DSV3 router GEMM kernel for archs: ${DSV3_ROUTER_GEMM_ARCHS}")
 
-    # DeepSeek V4 fused RMSNorm + router GEMV - same arch gating as DSV3.
-    set(DSV4_NORM_ROUTER_GEMM_SRC
-      "csrc/moe/dsv4_norm_router_gemm_entry.cu"
-      "csrc/moe/dsv4_norm_router_gemm_kernel.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${DSV4_NORM_ROUTER_GEMM_SRC}"
-      CUDA_ARCHS "${DSV3_ROUTER_GEMM_ARCHS}")
-    list(APPEND VLLM_MOE_EXT_SRC "${DSV4_NORM_ROUTER_GEMM_SRC}")
-    message(STATUS "Building DSV4 norm+router GEMV kernel for archs: ${DSV3_ROUTER_GEMM_ARCHS}")
+    message(STATUS "Building DSV3 router GEMM kernels for archs: ${SM90PLUS_ROUTER_GEMM_ARCHS}")
   else()
-    message(STATUS "Not building DSV3 router GEMM kernel as no compatible archs found"
+    message(STATUS "Not building DSV3 router GEMM kernels as no compatible archs found"
                    " (requires SM90+ and CUDA >= 12.0)")
   endif()
 endif()
@@ -1259,6 +1300,14 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/rocm/skinny_gemms.cu"
     "csrc/rocm/attention.cu")
 
+  set(VLLM_ROCM_HAS_GFX1100 OFF)
+  if(VLLM_GPU_ARCHES MATCHES "gfx1100")
+    set(VLLM_ROCM_HAS_GFX1100 ON)
+    list(APPEND VLLM_ROCM_EXT_SRC
+      "csrc/rocm/q_gemm_rdna3.cu"
+      "csrc/rocm/q_gemm_rdna3_wmma.cu")
+  endif()
+
   define_extension_target(
     _rocm_C
     DESTINATION vllm
@@ -1268,6 +1317,16 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
     ARCHITECTURES ${VLLM_GPU_ARCHES}
     USE_SABI 3
     WITH_SOABI)
+
+  if(VLLM_ROCM_HAS_GFX1100)
+    target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GFX1100)
+  endif()
+endif()
+
+# Must run after the last HIP `define_extension_target` so every extension
+# has registered its sources.
+if (VLLM_GPU_LANG STREQUAL "HIP")
+  vllm_finalize_hipify_target()
 endif()
 
 # For CUDA and HIP builds also build the triton_kernels external package.

benchmarks/benchmark_hidden_state_extraction.py

@@ -0,0 +1,415 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Benchmark hidden state extraction throughput.
+
+Measures two modes:
+  1. Baseline: bulk inference with max_tokens=1, no extraction.
+  2. Extract:  async hidden state extraction via ExampleHiddenStatesConnector
+               with N concurrent clients, each consuming hidden states as
+               soon as their request finishes (overlapping I/O with generation).
+
+Reports tokens/s and prompts/s for each mode.
+
+Usage:
+  python benchmarks/benchmark_hidden_state_extraction.py \
+      --model Qwen/Qwen3-0.6B \
+      --num-prompts 64 \
+      --num-clients 8 \
+      --prompt-len 8192 \
+      --layers 1 2 3 4
+"""
+
+import argparse
+import asyncio
+import time
+from concurrent.futures import ThreadPoolExecutor
+
+import torch
+from transformers import AutoConfig
+
+from vllm import LLM, SamplingParams
+from vllm.config.kv_transfer import KVTransferConfig
+from vllm.distributed.kv_transfer.kv_connector.v1 import (
+    example_hidden_states_connector,
+)
+from vllm.engine.arg_utils import AsyncEngineArgs
+from vllm.sampling_params import RequestOutputKind
+from vllm.v1.engine.async_llm import AsyncLLM
+
+
+def _make_profiler_config(profile_dir: str) -> dict:
+    """Build a profiler_config dict for torch profiling."""
+    return {
+        "profiler": "torch",
+        "torch_profiler_dir": profile_dir,
+        "torch_profiler_with_stack": True,
+    }
+
+
+def make_random_prompts(
+    num_prompts: int, prompt_len: int, vocab_size: int, seed: int = 42
+) -> list[list[int]]:
+    """Generate lists of random token IDs."""
+    # Set seed for reproducibility
+    torch.manual_seed(seed)
+    return [
+        torch.randint(0, vocab_size, (prompt_len,)).tolist() for _ in range(num_prompts)
+    ]
+
+
+def consume_hidden_states(path: str) -> float:
+    """Load hidden states from disk and compute per-position mean.
+
+    Returns a single float: the grand mean of all hidden state values.
+    This forces the benchmark to actually read and reduce the data.
+
+    Uses :func:`load_hidden_states` which acquires a shared flock,
+    blocking (without polling) until the async writer releases its
+    exclusive lock.
+    """
+    obj = example_hidden_states_connector.load_hidden_states(path)
+    hs = obj["hidden_states"]
+    total = hs.mean().item()
+
+    example_hidden_states_connector.cleanup_hidden_states(path)
+
+    return total
+
+
+def run_baseline(
+    model: str,
+    prompts: list[list[int]],
+    extra_args: dict,
+    profile_dir: str | None = None,
+) -> dict:
+    """Baseline: bulk inference, no hidden state extraction."""
+    if profile_dir:
+        extra_args = {
+            **extra_args,
+            "profiler_config": _make_profiler_config(profile_dir),
+        }
+    llm = LLM(
+        model=model,
+        enable_prefix_caching=False,
+        enable_chunked_prefill=False,
+        **extra_args,
+    )
+    sampling_params = SamplingParams(max_tokens=1)
+    prompt_inputs = [{"prompt_token_ids": p} for p in prompts]
+
+    # Warmup
+    llm.generate(prompt_inputs[:4], sampling_params, use_tqdm=False)
+
+    if profile_dir:
+        llm.start_profile()
+
+    t0 = time.perf_counter()
+    outputs = llm.generate(prompt_inputs, sampling_params, use_tqdm=True)
+    elapsed = time.perf_counter() - t0
+
+    if profile_dir:
+        llm.stop_profile()
+
+    total_prompt_tokens = sum(len(o.prompt_token_ids) for o in outputs)
+    num_prompts = len(outputs)
+
+    del llm
+    torch.accelerator.empty_cache()
+
+    return {
+        "mode": "baseline",
+        "elapsed_s": elapsed,
+        "num_prompts": num_prompts,
+        "total_prompt_tokens": total_prompt_tokens,
+        "tokens_per_s": total_prompt_tokens / elapsed,
+        "prompts_per_s": num_prompts / elapsed,
+    }
+
+
+# ---- Async extraction benchmark ----
+
+
+async def _client_loop(
+    engine: AsyncLLM,
+    prompt_queue: asyncio.Queue,
+    consume_pool: ThreadPoolExecutor,
+    results: list[dict],
+    client_id: int,
+):
+    """A single async client: pulls prompts, submits to engine, consumes
+    hidden states as soon as each request finishes."""
+    loop = asyncio.get_event_loop()
+    while True:
+        item = await prompt_queue.get()
+        if item is None:
+            prompt_queue.task_done()
+            break
+        idx, token_ids = item
+
+        request_id = f"req-{idx}"
+        sampling_params = SamplingParams(
+            max_tokens=1,
+            output_kind=RequestOutputKind.FINAL_ONLY,
+        )
+
+        final_output = None
+        async for output in engine.generate(
+            request_id=request_id,
+            prompt={"prompt_token_ids": token_ids},
+            sampling_params=sampling_params,
+        ):
+            if output.finished:
+                final_output = output
+
+        # Consume hidden states on a thread (disk I/O)
+        path = final_output.kv_transfer_params["hidden_states_path"]
+        mean_val = await loop.run_in_executor(consume_pool, consume_hidden_states, path)
+        num_tokens = len(final_output.prompt_token_ids)
+
+        results.append(
+            {
+                "request_id": request_id,
+                "num_prompt_tokens": num_tokens,
+                "mean_hidden_value": mean_val,
+            }
+        )
+        prompt_queue.task_done()
+
+
+async def _run_extraction_async(
+    model: str,
+    prompts: list[list[int]],
+    num_clients: int,
+    layers: list[int],
+    tmpdir: str,
+    extra_args: dict,
+    profile_dir: str | None = None,
+) -> dict:
+    if profile_dir:
+        extra_args = {
+            **extra_args,
+            "profiler_config": _make_profiler_config(profile_dir),
+        }
+    engine_args = AsyncEngineArgs(
+        model=model,
+        enable_prefix_caching=False,
+        enable_chunked_prefill=False,
+        max_num_batched_tokens=40960,
+        max_model_len=40960,
+        speculative_config={
+            "method": "extract_hidden_states",
+            "num_speculative_tokens": 1,
+            "draft_model_config": {
+                "hf_config": {
+                    "eagle_aux_hidden_state_layer_ids": layers,
+                },
+            },
+        },
+        kv_transfer_config=KVTransferConfig(
+            kv_connector="ExampleHiddenStatesConnector",
+            kv_role="kv_producer",
+            kv_connector_extra_config={
+                "shared_storage_path": tmpdir,
+            },
+        ),
+        **extra_args,
+    )
+    engine = AsyncLLM.from_engine_args(engine_args)
+
+    try:
+        # Warmup: run a few prompts sequentially, cleaning up generated files
+        for i in range(min(4, len(prompts))):
+            sp = SamplingParams(max_tokens=1, output_kind=RequestOutputKind.FINAL_ONLY)
+            final_output = None
+            async for output in engine.generate(
+                request_id=f"warmup-{i}",
+                prompt={"prompt_token_ids": prompts[i]},
+                sampling_params=sp,
+            ):
+                if output.finished:
+                    final_output = output
+            if final_output and final_output.kv_transfer_params:
+                path = final_output.kv_transfer_params.get("hidden_states_path")
+                if path:
+                    example_hidden_states_connector.cleanup_hidden_states(path)
+
+        if profile_dir:
+            await engine.start_profile()
+
+        # Fill prompt queue
+        prompt_queue: asyncio.Queue = asyncio.Queue()
+        for idx, token_ids in enumerate(prompts):
+            prompt_queue.put_nowait((idx, token_ids))
+        # Sentinel per client
+        for _ in range(num_clients):
+            prompt_queue.put_nowait(None)
+
+        results: list[dict] = []
+        consume_pool = ThreadPoolExecutor(max_workers=num_clients)
+
+        t0 = time.perf_counter()
+        tasks = [
+            asyncio.create_task(
+                _client_loop(engine, prompt_queue, consume_pool, results, i)
+            )
+            for i in range(num_clients)
+        ]
+        await asyncio.gather(*tasks)
+        elapsed = time.perf_counter() - t0
+
+        consume_pool.shutdown(wait=True)
+
+        if profile_dir:
+            await engine.stop_profile()
+
+        total_prompt_tokens = sum(r["num_prompt_tokens"] for r in results)
+        num_prompts = len(results)
+        mean_hidden = sum(r["mean_hidden_value"] for r in results) / max(
+            len(results), 1
+        )
+
+        return {
+            "mode": "extract",
+            "elapsed_s": elapsed,
+            "num_prompts": num_prompts,
+            "total_prompt_tokens": total_prompt_tokens,
+            "tokens_per_s": total_prompt_tokens / elapsed,
+            "prompts_per_s": num_prompts / elapsed,
+            "mean_hidden_value": mean_hidden,
+        }
+    finally:
+        engine.shutdown()
+
+
+def run_extraction(
+    model: str,
+    prompts: list[list[int]],
+    num_clients: int,
+    layers: list[int],
+    extra_args: dict,
+    profile_dir: str | None = None,
+) -> dict:
+    return asyncio.run(
+        _run_extraction_async(
+            model,
+            prompts,
+            num_clients,
+            layers,
+            "/dev/shm",
+            extra_args,
+            profile_dir=profile_dir,
+        )
+    )
+
+
+def print_results(results: dict):
+    mode = results["mode"]
+    print(f"\n{'=' * 60}")
+    print(f"  {mode.upper()} RESULTS")
+    print(f"{'=' * 60}")
+    print(f"  Prompts:             {results['num_prompts']}")
+    print(f"  Total prompt tokens: {results['total_prompt_tokens']:,}")
+    print(f"  Wall time:           {results['elapsed_s']:.2f}s")
+    print(f"  Tokens/s:            {results['tokens_per_s']:,.0f}")
+    print(f"  Prompts/s:           {results['prompts_per_s']:.2f}")
+    if mode == "extract":
+        print(f"  Mean hidden value:   {results['mean_hidden_value']:.6f}")
+    print(f"{'=' * 60}\n")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Benchmark hidden state extraction throughput"
+    )
+    parser.add_argument("--model", type=str, required=True)
+    parser.add_argument("--num-prompts", type=int, default=64)
+    parser.add_argument("--num-clients", type=int, default=8)
+    parser.add_argument("--prompt-len", type=int, default=8192)
+    parser.add_argument("--layers", type=int, nargs="+", default=[1, 2, 3, 4])
+    parser.add_argument("--skip-baseline", action="store_true")
+    parser.add_argument("--skip-extract", action="store_true")
+    parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
+    parser.add_argument("--max-num-batched-tokens", type=int, default=None)
+    parser.add_argument("--max-cudagraph-capture-size", type=int, default=None)
+    parser.add_argument("--max-model-len", type=int, default=None)
+    parser.add_argument("--enforce-eager", action="store_true")
+    parser.add_argument("--load-format", type=str, default=None)
+    parser.add_argument(
+        "--profile",
+        action="store_true",
+        help="Enable torch profiler for both baseline and extraction runs.",
+    )
+    parser.add_argument(
+        "--torch-profiler-dir",
+        type=str,
+        default="./vllm_profile",
+        help="Directory to save torch profiler traces (default: ./vllm_profile).",
+    )
+    parser.add_argument(
+        "--enable-flashinfer-autotune",
+        action="store_true",
+        default=False,
+        help="Enable FlashInfer autotuning (can be slow).",
+    )
+    args = parser.parse_args()
+
+    extra_args = {
+        "gpu_memory_utilization": args.gpu_memory_utilization,
+    }
+    if args.max_model_len is not None:
+        extra_args["max_model_len"] = args.max_model_len
+    if args.max_num_batched_tokens is not None:
+        extra_args["max_num_batched_tokens"] = args.max_num_batched_tokens
+        if args.max_model_len and args.max_num_batched_tokens < args.max_model_len:
+            raise ValueError(
+                "max_num_batched_tokens must be >= max_model_len since chunked prefill"
+                " is not supported by hidden state extraction."
+            )
+    if args.enforce_eager:
+        extra_args["enforce_eager"] = True
+    if args.load_format is not None:
+        extra_args["load_format"] = args.load_format
+    if args.max_cudagraph_capture_size is not None:
+        extra_args["max_cudagraph_capture_size"] = args.max_cudagraph_capture_size
+    extra_args["enable_flashinfer_autotune"] = args.enable_flashinfer_autotune
+
+    # Get vocab size from HF config without loading the full model
+    hf_config = AutoConfig.from_pretrained(args.model, trust_remote_code=True)
+    vocab_size = hf_config.vocab_size
+    prompts = make_random_prompts(args.num_prompts, args.prompt_len, vocab_size)
+    print(
+        f"Generated {args.num_prompts} prompts, "
+        f"{args.prompt_len} tokens each (vocab {vocab_size})"
+    )
+
+    profile_dir = args.torch_profiler_dir if args.profile else None
+    if profile_dir:
+        print(f"Torch profiler enabled, traces will be saved to {profile_dir}/")
+
+    if not args.skip_baseline:
+        baseline_profile_dir = f"{profile_dir}/baseline" if profile_dir else None
+        baseline = run_baseline(
+            args.model, prompts, extra_args, profile_dir=baseline_profile_dir
+        )
+        print_results(baseline)
+
+    if not args.skip_extract:
+        extract_profile_dir = f"{profile_dir}/extract" if profile_dir else None
+        extract = run_extraction(
+            args.model,
+            prompts,
+            args.num_clients,
+            args.layers,
+            extra_args,
+            profile_dir=extract_profile_dir,
+        )
+        print_results(extract)
+
+    if not args.skip_baseline and not args.skip_extract:
+        slowdown = baseline["tokens_per_s"] / extract["tokens_per_s"]
+        print("Extraction slowdown factor: {:.2f}x".format(slowdown))
+
+
+if __name__ == "__main__":
+    main()

benchmarks/kernels/benchmark_fused_moe_lora_one_shot.py

@@ -0,0 +1,465 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark the fused MoE-LoRA fast path (one-shot) vs two-kernel baseline.
+
+The "one_shot" provider goes through `vllm.lora.ops.triton_ops.fused_moe_lora`
+which dispatches to the single-kernel one-shot implementation when
+fully_sharded=False (the prefill default).
+
+The "two_kernel" provider drives `fused_moe_lora_shrink` + `fused_moe_lora_expand`
+directly, bypassing the dispatch and matching the legacy two-kernel path's
+work distribution. This isolates the win from kernel fusion.
+
+Run:
+    .venv/bin/python -m benchmarks.kernels.benchmark_fused_moe_lora_one_shot
+    .venv/bin/python -m benchmarks.kernels.benchmark_fused_moe_lora_one_shot \\
+        --model qwen3moe
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import random
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.lora.ops.triton_ops import (
+    fused_moe_lora,
+    fused_moe_lora_expand,
+    fused_moe_lora_shrink,
+)
+from vllm.triton_utils import triton
+
+DTYPE = torch.bfloat16
+DEVICE = "cuda"
+
+
+# ----- input fabrication -----------------------------------------------------
+
+
+def _round_up(x: int, base: int) -> int:
+    return ((x + base - 1) // base) * base
+
+
+def _ceildiv(x: int, y: int) -> int:
+    return (x + y - 1) // y
+
+
+def _assign_loras(num_tokens: int, num_sequences: int, max_loras: int) -> torch.Tensor:
+    tokens_per_seq = num_tokens // num_sequences
+    rem = num_tokens % num_sequences
+    out = torch.empty(num_tokens, dtype=torch.int32)
+    start = 0
+    for i in range(num_sequences):
+        end = start + tokens_per_seq + (1 if i < rem else 0)
+        out[start:end] = random.randint(0, max_loras - 1)
+        start = end
+    return out
+
+
+def _assign_experts(num_tokens: int, num_experts: int, top_k: int):
+    expert_indices = torch.empty((num_tokens, top_k), dtype=torch.int32)
+    for i in range(num_tokens):
+        expert_indices[i] = torch.randperm(num_experts)[:top_k]
+    weights = torch.rand((num_tokens, top_k), dtype=torch.float32)
+    weights = weights / weights.sum(dim=1, keepdim=True)
+    return expert_indices, weights
+
+
+def _make_inputs(
+    M: int,
+    K: int,
+    N_per_slice: int,
+    rank: int,
+    num_experts: int,
+    top_k: int,
+    max_loras: int,
+    num_slices: int,
+    block_size_m: int,
+):
+    """Mirrors the production caller's tensor layout."""
+    torch.manual_seed(0)
+    random.seed(0)
+
+    num_sequences = max(1, min(M, 8))
+    topk_ids_cpu, topk_weights_cpu = _assign_experts(M, num_experts, top_k)
+    token_lora_cpu = _assign_loras(M, num_sequences, max_loras)
+    lora_ids_cpu = torch.full((max_loras + 1,), -1, dtype=torch.int32)
+    uniq = torch.unique(token_lora_cpu, sorted=True)
+    lora_ids_cpu[: uniq.size(0)].copy_(uniq)
+
+    topk_ids = topk_ids_cpu.to(DEVICE)
+    topk_weights = topk_weights_cpu.to(device=DEVICE, dtype=DTYPE)
+    token_lora_mapping = token_lora_cpu.to(DEVICE)
+    lora_ids = lora_ids_cpu.to(DEVICE)
+    adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32, device=DEVICE)
+
+    lora_a = [
+        torch.randn((max_loras, num_experts, rank, K), dtype=DTYPE, device=DEVICE)
+        / max(K, 1) ** 0.5
+        for _ in range(num_slices)
+    ]
+    lora_b = [
+        torch.randn(
+            (max_loras, num_experts, N_per_slice, rank),
+            dtype=DTYPE,
+            device=DEVICE,
+        )
+        / max(rank, 1) ** 0.5
+        for _ in range(num_slices)
+    ]
+    hidden = torch.randn((M, K), dtype=DTYPE, device=DEVICE)
+    out_template = torch.zeros(
+        (M, top_k, num_slices * N_per_slice), dtype=DTYPE, device=DEVICE
+    )
+
+    # Sorted-path metadata (the prefill default).
+    max_pad = topk_ids.numel() + num_experts * (block_size_m - 1)
+    max_pad = _round_up(max_pad, block_size_m)
+    max_blocks = _ceildiv(max_pad, block_size_m)
+    sorted_token_ids = torch.empty(
+        (max_loras * max_pad,), dtype=torch.int32, device=DEVICE
+    )
+    expert_ids = torch.empty(
+        (max_loras * max_blocks,), dtype=torch.int32, device=DEVICE
+    )
+    num_post = torch.empty((max_loras,), dtype=torch.int32, device=DEVICE)
+    ops.moe_lora_align_block_size(
+        topk_ids,
+        token_lora_mapping,
+        num_experts,
+        block_size_m,
+        max_loras,
+        max_pad,
+        max_blocks,
+        sorted_token_ids,
+        expert_ids,
+        num_post,
+        adapter_enabled,
+        lora_ids,
+    )
+    expert_ids = expert_ids.view(max_loras, -1).contiguous()
+    sorted_token_ids = sorted_token_ids.view(max_loras, -1).contiguous()
+    num_active = torch.tensor([max_loras + 1], dtype=torch.int32, device="cpu")
+
+    return dict(
+        hidden=hidden,
+        lora_a=lora_a,
+        lora_b=lora_b,
+        topk_weights=topk_weights,
+        sorted_token_ids=sorted_token_ids,
+        expert_ids=expert_ids,
+        num_post=num_post,
+        token_lora_mapping=token_lora_mapping,
+        lora_ids=lora_ids,
+        num_active=num_active,
+        adapter_enabled=adapter_enabled,
+        out_template=out_template,
+        # bookkeeping
+        M=M,
+        K=K,
+        N_per_slice=N_per_slice,
+        rank=rank,
+        num_experts=num_experts,
+        top_k=top_k,
+        max_loras=max_loras,
+        num_slices=num_slices,
+        block_size_m=block_size_m,
+    )
+
+
+# ----- providers -------------------------------------------------------------
+
+
+def _run_one_shot(inp: dict):
+    """Drive `fused_moe_lora` with fully_sharded=False -> one-shot fast path."""
+    out = inp["out_template"].clone()
+    fused_moe_lora(
+        out,
+        inp["hidden"],
+        inp["lora_a"],
+        inp["lora_b"],
+        inp["topk_weights"],
+        inp["sorted_token_ids"],
+        inp["expert_ids"],
+        inp["num_post"],
+        inp["token_lora_mapping"],
+        inp["rank"],
+        inp["top_k"],
+        inp["lora_ids"],
+        inp["num_active"],
+        inp["adapter_enabled"],
+        inp["block_size_m"],
+        64,
+        32,
+        8,
+        4,
+        3,
+        1,
+        inp["block_size_m"],
+        64,
+        32,
+        8,
+        4,
+        3,
+        1,
+        False,
+        False,
+        0,
+    )
+    return out
+
+
+def _run_two_kernel(inp: dict):
+    """Drive `fused_moe_lora_shrink` + `fused_moe_lora_expand` directly,
+    bypassing the dispatch. Matches the legacy two-kernel work distribution.
+    """
+    M = inp["M"]
+    top_k = inp["top_k"]
+    rank = inp["rank"]
+    num_slices = inp["num_slices"]
+    N_per_slice = inp["N_per_slice"]
+    K = inp["K"]
+    num_experts = inp["num_experts"]
+    block_m = inp["block_size_m"]
+
+    intermediate = torch.zeros((num_slices, M, top_k, rank), dtype=DTYPE, device=DEVICE)
+    out = inp["out_template"].clone()
+    EM = inp["sorted_token_ids"].shape[1]
+    num_tokens = M * top_k
+
+    fused_moe_lora_shrink(
+        intermediate,
+        inp["hidden"],
+        inp["lora_a"],
+        inp["topk_weights"],
+        inp["sorted_token_ids"],
+        inp["expert_ids"],
+        inp["num_post"],
+        inp["token_lora_mapping"],
+        top_k,
+        inp["lora_ids"],
+        inp["adapter_enabled"],
+        torch.device(DEVICE),
+        rank,
+        M,
+        EM,
+        K,
+        num_tokens,
+        num_experts,
+        num_slices,
+        block_m,
+        64,
+        32,
+        8,
+        4,
+        3,
+        1,
+        inp["num_active"],
+        False,
+    )
+    fused_moe_lora_expand(
+        out,
+        intermediate,
+        inp["lora_b"],
+        inp["topk_weights"],
+        inp["sorted_token_ids"],
+        inp["expert_ids"],
+        inp["num_post"],
+        inp["token_lora_mapping"],
+        top_k,
+        inp["lora_ids"],
+        inp["adapter_enabled"],
+        torch.device(DEVICE),
+        rank,
+        M,
+        EM,
+        K,
+        num_tokens,
+        num_experts,
+        num_slices,
+        rank,
+        N_per_slice,
+        block_m,
+        64,
+        32,
+        8,
+        4,
+        3,
+        1,
+        inp["num_active"],
+        False,
+        0,
+    )
+    return out
+
+
+PROVIDER_FNS = {
+    "one_shot": _run_one_shot,
+    "two_kernel": _run_two_kernel,
+}
+
+
+# ----- model presets ---------------------------------------------------------
+
+
+MODEL_PRESETS: dict[str, dict] = {
+    # Mixtral-8x7B style: E=8, top_k=2, hidden=4096, intermediate=14336
+    "mixtral": dict(
+        K=4096,
+        N_per_slice=7168,
+        num_experts=8,
+        top_k=2,
+        max_loras=4,
+        num_slices=2,
+        block_size_m=64,
+    ),
+    # Qwen3-MoE / DeepSeek-V2 style: E=64, top_k=8, hidden=2048, inter=1408
+    "qwen3moe": dict(
+        K=2048,
+        N_per_slice=1408,
+        num_experts=64,
+        top_k=8,
+        max_loras=4,
+        num_slices=2,
+        block_size_m=64,
+    ),
+    # GLM-5.1 (zai-org/GLM-5.1-FP8): E=256, top_k=8, hidden=6144,
+    # moe_intermediate=2048
+    "glm5_1": dict(
+        K=6144,
+        N_per_slice=2048,
+        num_experts=256,
+        top_k=8,
+        max_loras=4,
+        num_slices=2,
+        block_size_m=64,
+    ),
+}
+
+
+M_RANGE = [16, 64, 256, 1024, 4096, 16384]
+RANK_RANGE = [8, 16, 32, 64]
+
+
+def get_benchmark(model: str, max_loras: int | None = None):
+    preset = dict(MODEL_PRESETS[model])
+    if max_loras is not None:
+        preset["max_loras"] = max_loras
+
+    @triton.testing.perf_report(
+        triton.testing.Benchmark(
+            x_names=["M", "rank"],
+            x_vals=[(M, R) for M in M_RANGE for R in RANK_RANGE],
+            line_arg="provider",
+            line_vals=list(PROVIDER_FNS.keys()),
+            line_names=["one_shot (fused)", "two_kernel (legacy)"],
+            styles=[("red", "-"), ("blue", "-")],
+            ylabel="ms",
+            plot_name=f"fused_moe_lora-{model}-loras{preset['max_loras']}",
+            args={"preset": preset},
+        )
+    )
+    def benchmark(M, rank, provider, preset):
+        inp = _make_inputs(
+            M=M,
+            K=preset["K"],
+            N_per_slice=preset["N_per_slice"],
+            rank=rank,
+            num_experts=preset["num_experts"],
+            top_k=preset["top_k"],
+            max_loras=preset["max_loras"],
+            num_slices=preset["num_slices"],
+            block_size_m=preset["block_size_m"],
+        )
+        fn = PROVIDER_FNS[provider]
+        quantiles = [0.5, 0.2, 0.8]
+        ms, min_ms, max_ms = triton.testing.do_bench(
+            lambda: fn(inp), quantiles=quantiles
+        )
+        return ms, max_ms, min_ms
+
+    return benchmark
+
+
+# ----- correctness sanity ---------------------------------------------------
+
+
+def calculate_diff(model: str, M: int, rank: int, max_loras: int | None = None):
+    preset = dict(MODEL_PRESETS[model])
+    if max_loras is not None:
+        preset["max_loras"] = max_loras
+    inp = _make_inputs(
+        M=M,
+        K=preset["K"],
+        N_per_slice=preset["N_per_slice"],
+        rank=rank,
+        num_experts=preset["num_experts"],
+        top_k=preset["top_k"],
+        max_loras=preset["max_loras"],
+        num_slices=preset["num_slices"],
+        block_size_m=preset["block_size_m"],
+    )
+    out_one = _run_one_shot(inp)
+    out_two = _run_two_kernel(inp)
+    max_abs = (out_one.float() - out_two.float()).abs().max().item()
+    print(
+        f"  model={model:<9} M={M:<6} rank={rank:<3}  "
+        f"max|one_shot - two_kernel|={max_abs:.4g}  "
+        f"ref|max|={out_two.float().abs().max().item():.3g}"
+    )
+    if max_abs <= 5e-2:
+        print("  ✅ outputs match within bf16 tolerance")
+    else:
+        print("  ❌ outputs differ beyond expected bf16 noise")
+
+
+# ----- main ------------------------------------------------------------------
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="mixtral",
+        choices=list(MODEL_PRESETS.keys()),
+        help="Model preset to sweep",
+    )
+    parser.add_argument(
+        "--save-path",
+        type=str,
+        default="./configs/fused_moe_lora_one_shot/",
+        help="Directory to save benchmark results",
+    )
+    parser.add_argument(
+        "--check-only",
+        action="store_true",
+        help="Run correctness sanity check only, no perf sweep",
+    )
+    parser.add_argument(
+        "--max-loras",
+        type=int,
+        default=None,
+        help="Override max_loras in the model preset (number of LoRA adapters "
+        "active in the batch). Defaults to the preset's value.",
+    )
+    args = parser.parse_args()
+
+    print(f"Correctness check ({args.model}):")
+    calculate_diff(args.model, M=256, rank=32, max_loras=args.max_loras)
+    if args.check_only:
+        raise SystemExit(0)
+
+    effective_max_loras = (
+        args.max_loras
+        if args.max_loras is not None
+        else MODEL_PRESETS[args.model]["max_loras"]
+    )
+    print(f"\nGPU: {torch.cuda.get_device_name()}")
+    print(f"Model preset: {args.model}  max_loras={effective_max_loras}\n")
+    benchmark = get_benchmark(args.model, max_loras=args.max_loras)
+    os.makedirs(args.save_path, exist_ok=True)
+    benchmark.run(print_data=True, save_path=args.save_path)

benchmarks/kernels/benchmark_moe_permute_unpermute.py

@@ -10,6 +10,7 @@ from transformers import AutoConfig
 
 from vllm.model_executor.layers.fused_moe import fused_topk
 from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
+    MoEPermuteScratch,
     moe_permute,
     moe_unpermute,
 )
@@ -54,6 +55,15 @@ def benchmark_permute(
     topk_weights, topk_ids, token_expert_indices = fused_topk(
         qhidden_states, input_gating, topk, False
     )
+    scratch = MoEPermuteScratch(
+        max_num_tokens=num_tokens,
+        topk=topk,
+        num_experts=num_experts,
+        num_local_experts=num_experts,
+        device=qhidden_states.device,
+        hidden_size=hidden_size,
+        hidden_dtype=qhidden_states.dtype,
+    )
 
     def prepare(i: int):
         input_gating.copy_(gating_output[i])
@@ -65,6 +75,7 @@ def benchmark_permute(
             topk_ids=topk_ids,
             n_expert=num_experts,
             expert_map=None,
+            scratch=scratch,
         )
 
     # JIT compilation & warmup
@@ -123,6 +134,15 @@ def benchmark_unpermute(
     topk_weights, topk_ids, token_expert_indices = fused_topk(
         qhidden_states, input_gating, topk, False
     )
+    scratch = MoEPermuteScratch(
+        max_num_tokens=num_tokens,
+        topk=topk,
+        num_experts=num_experts,
+        num_local_experts=num_experts,
+        device=qhidden_states.device,
+        hidden_size=hidden_size,
+        hidden_dtype=qhidden_states.dtype,
+    )
 
     def prepare():
         (
@@ -137,6 +157,7 @@ def benchmark_unpermute(
             topk_ids=topk_ids,
             n_expert=num_experts,
             expert_map=None,
+            scratch=scratch,
         )
         # convert to fp16/bf16 as gemm output
         return (

benchmarks/kernels/benchmark_norm_router_gemm.py

@@ -1,183 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Benchmark and correctness check for ``ops.dsv4_norm_router_gemm``.
-
-Two implementations are compared:
-
-  1. ``unfused``   — ``vllm_ops.rms_norm`` then ``ops.dsv3_router_gemm``,
-                     i.e. the current vLLM hot path (two kernel launches).
-  2. ``fused``     — ``ops.dsv4_norm_router_gemm``, the new single-kernel
-                     fused path.
-
-Both produce ``(normed_x: bf16, router_logits: fp32)``.  The correctness
-check verifies that ``fused`` and ``unfused`` agree to within ~1 bf16
-ULP — that is the precision floor for this op.
-"""
-
-import argparse
-
-import torch
-
-from vllm import _custom_ops as vllm_ops
-from vllm.triton_utils import triton
-
-# The fused dsv4_norm_router_gemm kernel is templated only for DSV4-Pro
-# (hidden_size=7168, num_experts=384).  Other shapes fall back to the
-# unfused path on the Python side (NormGatedLinear), so benchmark only
-# the configuration that the fused kernel actually targets.
-HIDDEN_SIZE = 7168
-NUM_EXPERTS_CHOICES = (384,)
-RMS_EPS = 1e-6
-
-
-def unfused_norm_router_gemm(
-    x: torch.Tensor,
-    norm_weight: torch.Tensor,
-    gate_weight: torch.Tensor,
-    eps: float,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    # Call ``_C::rms_norm`` directly (mirroring ``_dsv4_pro_norm_gate``'s
-    # fallback path) so the benchmarked baseline doesn't inherit any
-    # Python wrapper overhead or risk falling through to the native
-    # eager-primitive ``RMSNorm.forward_native`` path.
-    normed = torch.empty_like(x)
-    torch.ops._C.rms_norm(normed, x, norm_weight, eps)
-    logits = vllm_ops.dsv3_router_gemm(normed, gate_weight, torch.float32)
-    return normed, logits
-
-
-def fused_norm_router_gemm(
-    x: torch.Tensor,
-    norm_weight: torch.Tensor,
-    gate_weight: torch.Tensor,
-    eps: float,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    return vllm_ops.dsv4_norm_router_gemm(x, norm_weight, gate_weight, eps)
-
-
-def _make_inputs(num_tokens: int, num_experts: int, hidden_size: int, seed: int = 0):
-    torch.manual_seed(seed)
-    device = "cuda"
-    x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
-    norm_w = torch.randn(hidden_size, dtype=torch.bfloat16, device=device)
-    gate_w = torch.randn(num_experts, hidden_size, dtype=torch.bfloat16, device=device)
-    # Down-scale gate_w so the GEMV output stays in a representable range.
-    gate_w = gate_w / float(hidden_size) ** 0.5
-    norm_w = (norm_w * 0.1) + 1.0
-    return x, norm_w, gate_w
-
-
-def calculate_diff(
-    num_tokens: int,
-    num_experts: int,
-    hidden_size: int = HIDDEN_SIZE,
-    normed_atol: float = 2e-3,
-    logits_atol: float = 1e-2,
-    rtol: float = 1e-2,
-) -> None:
-    x, norm_w, gate_w = _make_inputs(num_tokens, num_experts, hidden_size)
-
-    normed_unfused, logits_unfused = unfused_norm_router_gemm(
-        x.clone(), norm_w, gate_w, RMS_EPS
-    )
-    normed_fused, logits_fused = fused_norm_router_gemm(
-        x.clone(), norm_w, gate_w, RMS_EPS
-    )
-
-    def _max_abs(a, b):
-        return (a.float() - b.float()).abs().max().item()
-
-    print(f"\n=== M={num_tokens} E={num_experts} H={hidden_size} ===")
-    print(f"normed_x  |fused - unfused| = {_max_abs(normed_fused, normed_unfused):.3e}")
-    print(f"logits    |fused - unfused| = {_max_abs(logits_fused, logits_unfused):.3e}")
-
-    ok_normed = torch.allclose(
-        normed_fused.float(),
-        normed_unfused.float(),
-        atol=normed_atol,
-        rtol=rtol,
-    )
-    ok_logits = torch.allclose(
-        logits_fused.float(),
-        logits_unfused.float(),
-        atol=logits_atol,
-        rtol=rtol,
-    )
-    if ok_normed and ok_logits:
-        print(
-            f"OK   fused vs unfused within "
-            f"normed_atol={normed_atol:.0e} logits_atol={logits_atol:.0e} "
-            f"rtol={rtol:.0e}"
-        )
-    else:
-        print(
-            f"FAIL normed_ok={ok_normed} logits_ok={ok_logits}; "
-            f"see max-abs values above"
-        )
-
-
-def get_benchmark():
-    # Only num_tokens varies (DSV4-Pro hard-codes E=384); single-axis
-    # sweep yields a clean line plot with M on the x-axis.
-    num_experts = NUM_EXPERTS_CHOICES[0]
-
-    @triton.testing.perf_report(
-        triton.testing.Benchmark(
-            x_names=["num_tokens"],
-            x_vals=list(range(1, 17)),
-            line_arg="provider",
-            line_vals=["unfused", "fused"],
-            line_names=["unfused (rms+dsv3)", "fused (dsv4)"],
-            styles=[("green", "-"), ("red", "-")],
-            ylabel="us",
-            plot_name=f"norm-router-gemm-E{num_experts}-H{HIDDEN_SIZE}",
-            args={},
-        )
-    )
-    def benchmark(num_tokens, provider):
-        x, norm_w, gate_w = _make_inputs(num_tokens, num_experts, HIDDEN_SIZE)
-
-        quantiles = [0.5, 0.2, 0.8]
-        if provider == "unfused":
-            fn = lambda: unfused_norm_router_gemm(  # noqa: E731
-                x, norm_w, gate_w, RMS_EPS
-            )
-        else:
-            fn = lambda: fused_norm_router_gemm(  # noqa: E731
-                x, norm_w, gate_w, RMS_EPS
-            )
-
-        ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=quantiles)
-        return 1000 * ms, 1000 * max_ms, 1000 * min_ms
-
-    return benchmark
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--save-path",
-        type=str,
-        default="./configs/norm_router_gemm/",
-    )
-    parser.add_argument(
-        "--skip-bench",
-        action="store_true",
-        help="Run only the correctness check, not the perf sweep.",
-    )
-    args = parser.parse_args()
-
-    # Correctness sweep over the full fast-path range M=1..16.
-    for m in range(1, 17):
-        for e in NUM_EXPERTS_CHOICES:
-            calculate_diff(num_tokens=m, num_experts=e, hidden_size=HIDDEN_SIZE)
-
-    if args.skip_bench:
-        return
-
-    benchmark = get_benchmark()
-    benchmark.run(print_data=True, save_path=args.save_path)
-
-
-if __name__ == "__main__":
-    main()

benchmarks/kernels/benchmark_router_gemm.py

@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+import torch.nn.functional as F
+
+from vllm import _custom_ops as ops
+from vllm.platforms import current_platform
+from vllm.transformers_utils.config import get_config
+from vllm.triton_utils import triton
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+# Dimensions supported by the DSV3 specialized kernel
+DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
+DSV3_SUPPORTED_HIDDEN_SIZES = [7168]
+
+# Dimensions supported by the gpt-oss specialized kernel
+GPT_OSS_SUPPORTED_NUM_EXPERTS = [32, 128]
+GPT_OSS_SUPPORTED_HIDDEN_SIZES = [2880]
+
+# Dimensions supported by the fp32 specialized kernel (MiniMax-M2)
+FP32_SUPPORTED_NUM_EXPERTS = [256]
+FP32_SUPPORTED_HIDDEN_SIZES = [3072]
+FP32_MAX_TOKENS = 32
+
+
+def get_batch_size_range(max_batch_size):
+    return [2**x for x in range(14) if 2**x <= max_batch_size]
+
+
+def get_model_params(config):
+    if config.architectures[0] in (
+        "DeepseekV2ForCausalLM",
+        "DeepseekV3ForCausalLM",
+        "DeepseekV32ForCausalLM",
+    ):
+        num_experts = config.n_routed_experts
+        hidden_size = config.hidden_size
+    elif config.architectures[0] in ("GptOssForCausalLM",) or config.architectures[
+        0
+    ] in ("MiniMaxM2ForCausalLM",):
+        num_experts = config.num_local_experts
+        hidden_size = config.hidden_size
+    else:
+        raise ValueError(f"Unsupported architecture: {config.architectures}")
+    return num_experts, hidden_size
+
+
+def get_benchmark(model, max_batch_size, trust_remote_code):
+    @triton.testing.perf_report(
+        triton.testing.Benchmark(
+            x_names=["batch_size"],
+            x_vals=get_batch_size_range(max_batch_size),
+            x_log=False,
+            line_arg="provider",
+            line_vals=[
+                "torch",
+                "vllm",
+            ],
+            line_names=["PyTorch", "vLLM"],
+            styles=([("blue", "-"), ("red", "-")]),
+            ylabel="TFLOPs",
+            plot_name=f"{model} router gemm throughput",
+            args={},
+        )
+    )
+    def benchmark(batch_size, provider):
+        config = get_config(model=model, trust_remote_code=trust_remote_code)
+        num_experts, hidden_size = get_model_params(config)
+
+        is_hopper_or_blackwell = current_platform.is_device_capability(
+            90
+        ) or current_platform.is_device_capability_family(100)
+        allow_dsv3_router_gemm = (
+            is_hopper_or_blackwell
+            and num_experts in DSV3_SUPPORTED_NUM_EXPERTS
+            and hidden_size in DSV3_SUPPORTED_HIDDEN_SIZES
+        )
+        allow_gpt_oss_router_gemm = (
+            is_hopper_or_blackwell
+            and num_experts in GPT_OSS_SUPPORTED_NUM_EXPERTS
+            and hidden_size in GPT_OSS_SUPPORTED_HIDDEN_SIZES
+        )
+        is_fp32_router_model = (
+            is_hopper_or_blackwell
+            and num_experts in FP32_SUPPORTED_NUM_EXPERTS
+            and hidden_size in FP32_SUPPORTED_HIDDEN_SIZES
+        )
+        allow_fp32_router_gemm = is_fp32_router_model and batch_size <= FP32_MAX_TOKENS
+
+        # Weight dtype: fp32 kernel requires fp32 weights; others use bf16.
+        weight_dtype = torch.float32 if is_fp32_router_model else torch.bfloat16
+        mat_a = torch.randn(
+            (batch_size, hidden_size), dtype=torch.bfloat16, device="cuda"
+        ).contiguous()
+        mat_b = torch.randn(
+            (num_experts, hidden_size), dtype=weight_dtype, device="cuda"
+        ).contiguous()
+        bias = torch.randn(
+            num_experts, dtype=torch.bfloat16, device="cuda"
+        ).contiguous()
+
+        has_bias = allow_gpt_oss_router_gemm
+
+        quantiles = [0.5, 0.2, 0.8]
+
+        if provider == "torch":
+
+            def runner():
+                if allow_fp32_router_gemm:
+                    F.linear(mat_a.float(), mat_b)
+                elif has_bias:
+                    F.linear(mat_a, mat_b, bias)
+                else:
+                    F.linear(mat_a, mat_b)
+        elif provider == "vllm":
+
+            def runner():
+                if allow_dsv3_router_gemm:
+                    ops.dsv3_router_gemm(mat_a, mat_b, torch.bfloat16)
+                elif allow_fp32_router_gemm:
+                    ops.fp32_router_gemm(mat_a, mat_b)
+                elif allow_gpt_oss_router_gemm:
+                    ops.gpt_oss_router_gemm(mat_a, mat_b, bias)
+                elif is_fp32_router_model:
+                    # batch_size > FP32_MAX_TOKENS: fall back to F.linear
+                    F.linear(mat_a.float(), mat_b)
+                else:
+                    F.linear(mat_a, mat_b)
+
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            runner, quantiles=quantiles
+        )
+
+        def tflops(t_ms):
+            flops = 2 * batch_size * hidden_size * num_experts
+            return flops / (t_ms * 1e-3) / 1e12
+
+        return tflops(ms), tflops(max_ms), tflops(min_ms)
+
+    return benchmark
+
+
+if __name__ == "__main__":
+    parser = FlexibleArgumentParser()
+    parser.add_argument("--model", type=str, default="openai/gpt-oss-20b")
+    parser.add_argument("--max-batch-size", default=16, type=int)
+    parser.add_argument("--trust-remote-code", action="store_true")
+    args = parser.parse_args()
+
+    # Get the benchmark function
+    benchmark = get_benchmark(args.model, args.max_batch_size, args.trust_remote_code)
+    # Run performance benchmark
+    benchmark.run(print_data=True)

benchmarks/kernels/benchmark_selective_state_update.py

@@ -0,0 +1,774 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Benchmark and tuning script for the Mamba selective_state_update kernel.
+
+Mirrors the fused MoE tuning workflow: sweeps (BLOCK_SIZE_M, num_warps) across
+an effective_batch grid for a given (headdim, dstate, ngroups, cache_dtype) and
+saves the best config per effective_batch to JSON. Generated configs are picked
+up by selective_state_update at runtime.
+
+Usage:
+    python -m benchmarks.kernels.benchmark_selective_state_update \
+        --all-dstates --save-configs --compare
+"""
+
+import argparse
+import json
+import os
+import sys
+from io import StringIO
+from itertools import product
+from typing import Any
+
+import torch
+
+from tests.kernels.mamba.utils import selective_state_update_ref
+from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
+    _CONFIGS_DIR,
+    _canonical_cache_dtype,
+    _get_default_ssm_launch_config,
+    get_ssm_config_file_name,
+    get_ssm_device_name,
+    override_ssm_config,
+    selective_state_update,
+)
+from vllm.triton_utils import triton
+
+# bf16 shares configs with fp16 - same bit width.
+_SSM_CACHE_DTYPE_MAP: dict[str, torch.dtype] = {
+    "float32": torch.float32,
+    "float16": torch.float16,
+    "bfloat16": torch.float16,
+}
+
+_RESULTS_DIR = os.path.dirname(os.path.realpath(__file__))
+
+# ---------------------------------------------------------------------------
+# Tuning search space
+# ---------------------------------------------------------------------------
+
+_BSM_CHOICES_ALL = [4, 8, 16, 32, 64, 128, 256]
+
+NUM_WARPS_CHOICES = [1, 2, 4, 8]
+
+
+def _block_size_m_choices(headdim: int) -> list[int]:
+    """BLOCK_SIZE_M candidates worth sweeping for a given headdim.
+
+    BLOCK_SIZE_M > next_pow2(headdim) wastes >=50% of each tile via masking
+    (offs_m >= dim rows are zeroed out), so we cap the sweep there.
+    """
+    ceiling = 1
+    while ceiling < headdim:
+        ceiling <<= 1
+    return [b for b in _BSM_CHOICES_ALL if b <= ceiling]
+
+
+# Default deployment shapes. effective_batch = batch * nheads scales the
+# kernel grid, so configs transfer across (model, TP) combos sharing
+# (headdim, dstate, cache_dtype).
+DEFAULT_BATCH_SIZES = [1, 8, 16, 32, 64, 128, 256, 512, 1024, 1536, 2048]
+DEFAULT_NHEADS = [128, 256]
+
+ALL_DSTATES = [16, 32, 64, 128, 256]
+
+# Default tuning shape — matches Nemotron-3-Super and Nemotron-3-Nano Mamba layers.
+# Override with CLI flags for other architectures.
+DEFAULT_HEADDIM = 64
+DEFAULT_NGROUPS = 8
+
+
+# ---------------------------------------------------------------------------
+# Benchmark helper
+# ---------------------------------------------------------------------------
+
+
+def _make_inputs(
+    batch: int,
+    nheads: int,
+    dim: int,
+    dstate: int,
+    ngroups: int,
+    dtype: torch.dtype,
+    state_dtype: torch.dtype | None = None,
+    device: str = "cuda",
+):
+    if state_dtype is None:
+        state_dtype = dtype
+    state = torch.randn(batch, nheads, dim, dstate, dtype=state_dtype, device=device)
+    x = torch.randn(batch, nheads, dim, dtype=dtype, device=device)
+    dt = torch.randn(batch, nheads, dim, dtype=dtype, device=device)
+    A = -torch.rand(nheads, dim, dstate, dtype=torch.float32, device=device)
+    B = torch.randn(batch, ngroups, dstate, dtype=dtype, device=device)
+    C = torch.randn(batch, ngroups, dstate, dtype=dtype, device=device)
+    D = torch.randn(nheads, dim, dtype=dtype, device=device)
+    dt_bias = torch.randn(nheads, dim, dtype=dtype, device=device)
+    out = torch.zeros(batch, nheads, dim, dtype=dtype, device=device)
+    return state, x, dt, A, B, C, D, dt_bias, out
+
+
+def benchmark_config(
+    batch: int,
+    nheads: int,
+    dim: int,
+    dstate: int,
+    ngroups: int,
+    block_size_m: int,
+    num_warps_val: int,
+    dtype: torch.dtype,
+    state_dtype: torch.dtype | None = None,
+    num_iters: int = 100,
+    num_warmup: int = 20,
+    graph_batch_size: int = 10,
+) -> float | None:
+    """
+    Time one (BLOCK_SIZE_M, num_warps) config for selective_state_update.
+    Returns elapsed time in microseconds, or None on error.
+
+    Uses CUDA graph capture-and-replay to isolate kernel time from Python
+    eager-mode dispatch / kwarg-resolution overhead, mirroring the timing
+    methodology in benchmarks/kernels/benchmark_moe.py.
+    """
+    state, x, dt, A, B, C, D, dt_bias, out = _make_inputs(
+        batch, nheads, dim, dstate, ngroups, dtype, state_dtype=state_dtype
+    )
+
+    def _call_kernel() -> None:
+        selective_state_update(
+            state,
+            x,
+            dt,
+            A,
+            B,
+            C,
+            D=D,
+            z=None,
+            dt_bias=dt_bias,
+            dt_softplus=True,
+            out=out,
+        )
+
+    try:
+        with override_ssm_config((block_size_m, num_warps_val)):
+            # Eager-mode warmup: triggers Triton autotune / JIT, primes caches.
+            for _ in range(num_warmup):
+                _call_kernel()
+            torch.accelerator.synchronize()
+
+            # Capture graph_batch_size invocations into a CUDA graph so the
+            # timed region runs without Python dispatch overhead per call.
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                for _ in range(graph_batch_size):
+                    _call_kernel()
+            torch.accelerator.synchronize()
+
+            # Warmup graph replays (let the runtime stabilize).
+            for _ in range(5):
+                graph.replay()
+            torch.accelerator.synchronize()
+
+            start = torch.cuda.Event(enable_timing=True)
+            end = torch.cuda.Event(enable_timing=True)
+            latencies: list[float] = []
+            for _ in range(num_iters):
+                start.record()
+                graph.replay()
+                end.record()
+                end.synchronize()
+                latencies.append(start.elapsed_time(end))
+            graph.reset()
+        # elapsed_time returns ms; each replay runs graph_batch_size kernels,
+        # so divide by (num_iters * graph_batch_size) and convert ms -> us.
+        return sum(latencies) / (num_iters * graph_batch_size) * 1000
+    except Exception as e:
+        if "OutOfResources" not in str(e):
+            print(
+                f"    Warning: config M={block_size_m},w={num_warps_val} "
+                f"raised {type(e).__name__}: {e}"
+            )
+        return None
+
+
+# ---------------------------------------------------------------------------
+# Tuning loop
+# ---------------------------------------------------------------------------
+
+
+# CUDA grid Y/Z dim limit — both `batch` and `nheads` must fit individually.
+_CUDA_MAX_GRID_DIM = 65535
+
+# Above this, kernel state-offset arithmetic (batch * nheads * headdim * dstate)
+# overflows int32 and the launch raises cudaErrorIllegalAddress.
+# 262144 covers Nemotron Super TP1 BS=2048.
+_MAX_EFFECTIVE_BATCH = 262144
+
+
+def expand_batch_x_nheads(
+    batch_sizes: list[int],
+    nheads_list: list[int],
+    ngroups: int,
+) -> list[tuple[int, int, int]]:
+    """Cross-product batch_sizes × nheads_list → sorted [(effective_batch,
+    batch, nheads)], deduped by effective_batch. Filters pairs that exceed
+    the CUDA grid dim limit, the effective_batch ceiling, or where nheads is
+    not a positive multiple of ngroups.
+    """
+    seen: dict[int, tuple[int, int]] = {}
+    skipped_grid: list[tuple[int, int]] = []
+    skipped_ngroups: list[tuple[int, int]] = []
+    skipped_eb: list[tuple[int, int]] = []
+    for b, n in product(batch_sizes, nheads_list):
+        if b <= 0 or n <= 0:
+            continue
+        if b > _CUDA_MAX_GRID_DIM or n > _CUDA_MAX_GRID_DIM:
+            skipped_grid.append((b, n))
+            continue
+        if n % ngroups != 0:
+            skipped_ngroups.append((b, n))
+            continue
+        if b * n > _MAX_EFFECTIVE_BATCH:
+            skipped_eb.append((b, n))
+            continue
+        seen.setdefault(b * n, (b, n))
+    if skipped_grid:
+        print(
+            f"  Note: skipping (batch, nheads) pairs exceeding CUDA grid dim "
+            f"{_CUDA_MAX_GRID_DIM}: {skipped_grid}"
+        )
+    if skipped_ngroups:
+        print(
+            f"  Note: skipping (batch, nheads) pairs where nheads % ngroups != 0 "
+            f"for ngroups={ngroups}: {skipped_ngroups}"
+        )
+    if skipped_eb:
+        print(
+            f"  Note: skipping (batch, nheads) pairs whose effective_batch "
+            f"exceeds {_MAX_EFFECTIVE_BATCH}: {skipped_eb}"
+        )
+    return sorted((eb, b, n) for eb, (b, n) in seen.items())
+
+
+def tune_dstate(
+    dstate: int,
+    headdim: int,
+    ngroups: int,
+    dtype: torch.dtype,
+    num_iters: int,
+    verbose: bool,
+    active: list[tuple[int, int, int]],
+    state_dtype: torch.dtype | None = None,
+) -> tuple[dict[int, dict], dict[int, dict[tuple[int, int], float]]]:
+    """For each (effective_batch, batch, nheads) in *active*, sweep
+    (BLOCK_SIZE_M, num_warps) and return
+    ({effective_batch: best_config}, {effective_batch: {(bsm, nw): us}}).
+    The second map is the full timing grid, used downstream so we don't
+    re-measure the same config in the comparison phase.
+    """
+    best_per_eb: dict[int, dict] = {}
+    timings: dict[int, dict[tuple[int, int], float]] = {}
+
+    print(f"\n{'=' * 74}")
+    effective_state_dtype = state_dtype if state_dtype is not None else dtype
+    print(
+        f"Tuning  headdim={headdim}  dstate={dstate}  ngroups={ngroups}  "
+        f"dtype={dtype}  ssm_cache_dtype={effective_state_dtype}"
+    )
+    print(f"{'=' * 74}")
+
+    bsm_choices = _block_size_m_choices(headdim)
+    print(f"BSM candidates (capped at next_pow2(headdim={headdim})): {bsm_choices}")
+
+    hdr = f"{'EffBatch':>8} | {'BLOCK_M':>7} | {'warps':>5} | {'us':>10} | note"
+    print(hdr)
+    print("-" * 52)
+
+    for eb, batch, nheads in active:
+        best_time = float("inf")
+        best_cfg: dict = {}
+        eb_timings: dict[tuple[int, int], float] = {}
+
+        for bsm, nw in product(bsm_choices, NUM_WARPS_CHOICES):
+            t = benchmark_config(
+                batch=batch,
+                nheads=nheads,
+                dim=headdim,
+                dstate=dstate,
+                ngroups=ngroups,
+                block_size_m=bsm,
+                num_warps_val=nw,
+                dtype=dtype,
+                state_dtype=state_dtype,
+                num_iters=num_iters,
+            )
+            if t is None:
+                continue
+            eb_timings[(bsm, nw)] = t
+            is_best = t < best_time
+            if is_best:
+                best_time = t
+                best_cfg = {"BLOCK_SIZE_M": bsm, "num_warps": nw}
+            if verbose:
+                marker = " <-- best" if is_best else ""
+                print(f"{eb:>8} | {bsm:>7} | {nw:>5} | {t:>10.2f} |{marker}")
+
+        timings[eb] = eb_timings
+
+        if not best_cfg:
+            print(
+                f"{eb:>8} | {'-':>7} | {'-':>5} | {'-':>10} | "
+                f"no working config (skipped)"
+            )
+            continue
+
+        if not verbose:
+            print(
+                f"{eb:>8} | {best_cfg['BLOCK_SIZE_M']:>7} | "
+                f"{best_cfg['num_warps']:>5} | {best_time:>10.2f} | best"
+            )
+
+        best_per_eb[eb] = best_cfg
+
+    return best_per_eb, timings
+
+
+# ---------------------------------------------------------------------------
+# Correctness validation
+# ---------------------------------------------------------------------------
+
+
+def validate_configs(
+    dstate: int,
+    headdim: int,
+    ngroups: int,
+    tuned: dict[int, dict],
+    active: list[tuple[int, int, int]],
+    dtype: torch.dtype,
+    atol: float = 1e-2,
+    rtol: float = 1e-2,
+    state_dtype: torch.dtype | None = None,
+) -> dict[int, bool]:
+    """
+    For every (effective_batch, batch, nheads) in *active* that has a tuned
+    config, run the kernel with that config and compare against the reference.
+    Returns {effective_batch: passed}.
+    """
+    # Disable TF32 so the reference's matmul matches the Triton kernel's
+    # fp32 accumulation; otherwise large ebs show bf16 rounding mismatches.
+    torch.set_float32_matmul_precision("highest")
+
+    print(f"\n{'=' * 74}")
+    effective_state_dtype = state_dtype if state_dtype is not None else dtype
+    print(
+        f"Validation  headdim={headdim}  dstate={dstate}  ngroups={ngroups}  "
+        f"dtype={dtype}  ssm_cache_dtype={effective_state_dtype}  atol={atol}"
+    )
+    print(f"{'=' * 74}")
+    print(f"{'EffBatch':>8} | {'MaxAbsErr':>12} | {'Status':>8}")
+    print("-" * 36)
+
+    results: dict[int, bool] = {}
+
+    for eb, batch, nheads in active:
+        cfg = tuned.get(eb)
+        if cfg is None:
+            continue
+        state, x, dt, A, B, C, D, dt_bias, out = _make_inputs(
+            batch=batch,
+            nheads=nheads,
+            dim=headdim,
+            dstate=dstate,
+            ngroups=ngroups,
+            dtype=dtype,
+            state_dtype=state_dtype,
+        )
+        # Clone state before GPU kernel modifies it in-place
+        state_ref = state.clone()
+
+        with override_ssm_config((cfg["BLOCK_SIZE_M"], cfg["num_warps"])):
+            selective_state_update(
+                state,
+                x,
+                dt,
+                A,
+                B,
+                C,
+                D=D,
+                z=None,
+                dt_bias=dt_bias,
+                dt_softplus=True,
+                out=out,
+            )
+        torch.accelerator.synchronize()
+        gpu_out = out.detach().cpu()
+
+        # Reference uses the original (unmodified) state
+        # Upcast to fp32 so the reference sums in fp32 (matches the Triton
+        # kernel); summing in bf16 over `dstate` blows up the error.
+        ref_out = (
+            selective_state_update_ref(
+                state_ref.float(),
+                x.float(),
+                dt.float(),
+                A.float(),
+                B.float(),
+                C.float(),
+                D=D.float(),
+                dt_bias=dt_bias.float(),
+                dt_softplus=True,
+            )
+            .to(out.dtype)
+            .cpu()
+        )
+
+        passed = torch.allclose(gpu_out.float(), ref_out.float(), atol=atol, rtol=rtol)
+        max_err = (gpu_out.float() - ref_out.float()).abs().max().item()
+        status = "PASS" if passed else "FAIL"
+        results[eb] = passed
+        print(f"{eb:>8} | {max_err:>12.6f} | {status:>8}")
+
+    n_pass = sum(results.values())
+    n_total = len(results)
+    print(f"\n  {n_pass}/{n_total} configs passed validation for dstate={dstate}")
+    return results
+
+
+# ---------------------------------------------------------------------------
+# Save configs
+# ---------------------------------------------------------------------------
+
+
+def save_configs(
+    headdim: int,
+    dstate: int,
+    cache_dtype: str,
+    configs: dict[int, dict],
+    save_dir: str | None = None,
+) -> str:
+    # bf16 shares configs with fp16, use common filename for both
+    cache_dtype = _canonical_cache_dtype(cache_dtype)
+
+    base_dir = save_dir if save_dir else _CONFIGS_DIR
+    os.makedirs(base_dir, exist_ok=True)
+    file_path = os.path.join(
+        base_dir,
+        get_ssm_config_file_name(headdim, dstate, cache_dtype, get_ssm_device_name()),
+    )
+    # triton_version is informational only, the loader ignores it
+    payload: dict[str, Any] = {
+        "triton_version": triton.__version__,
+        **{str(k): v for k, v in sorted(configs.items())},
+    }
+    with open(file_path, "w") as f:
+        json.dump(payload, f, indent=4)
+    return file_path
+
+
+# ---------------------------------------------------------------------------
+# Comparison table
+# ---------------------------------------------------------------------------
+
+
+def current_heuristic(dstate: int, is_blackwell: bool = False) -> dict:
+    """Return the current hard-coded BLOCK_SIZE_M / num_warps for dstate."""
+    bsm, nw = _get_default_ssm_launch_config(dstate, is_blackwell)
+    return {"BLOCK_SIZE_M": bsm, "num_warps": nw}
+
+
+def compare_heuristic_vs_tuned(
+    dstate: int,
+    headdim: int,
+    ngroups: int,
+    tuned: dict[int, dict],
+    timings: dict[int, dict[tuple[int, int], float]],
+    active: list[tuple[int, int, int]],
+    dtype: torch.dtype,
+    num_iters: int,
+    is_blackwell: bool,
+    state_dtype: torch.dtype | None = None,
+):
+    heur_cfg = current_heuristic(dstate, is_blackwell)
+    heur_key = (heur_cfg["BLOCK_SIZE_M"], heur_cfg["num_warps"])
+
+    print(f"\n{'=' * 74}")
+    print(
+        f"Comparison  headdim={headdim}  dstate={dstate}  "
+        f"ngroups={ngroups}  —  heuristic vs tuned"
+    )
+    print(
+        f"Heuristic: BLOCK_SIZE_M={heur_cfg['BLOCK_SIZE_M']}, "
+        f"num_warps={heur_cfg['num_warps']}"
+    )
+    print(f"{'=' * 74}")
+    hdr = (
+        f"{'EffBatch':>8} | {'Heur(us)':>10} | {'Tuned(us)':>10} | "
+        f"{'Speedup':>8} | Best config"
+    )
+    print(hdr)
+    print("-" * len(hdr))
+
+    for eb, batch, nheads in active:
+        eb_timings = timings.get(eb, {})
+
+        # Heuristic timing: reuse the tuning measurement if the heuristic
+        # config was in the swept grid; otherwise measure it once.
+        t_h = eb_timings.get(heur_key)
+        if t_h is None:
+            t_h = benchmark_config(
+                batch=batch,
+                nheads=nheads,
+                dim=headdim,
+                dstate=dstate,
+                ngroups=ngroups,
+                block_size_m=heur_cfg["BLOCK_SIZE_M"],
+                num_warps_val=heur_cfg["num_warps"],
+                dtype=dtype,
+                state_dtype=state_dtype,
+                num_iters=num_iters,
+            )
+
+        # `tuned[eb]` may be missing if all configs failed in tune_dstate;
+        # in that case fall back to the heuristic so the table still prints.
+        best = tuned.get(eb) or heur_cfg
+        t_t = eb_timings.get((best["BLOCK_SIZE_M"], best["num_warps"]))
+
+        if t_h is None or t_t is None:
+            print(f"{eb:>8} | {'N/A':>10} | {'N/A':>10} | {'N/A':>8} |")
+            continue
+        speedup = t_h / t_t
+        marker = " <--" if speedup > 1.05 else ""
+        print(
+            f"{eb:>8} | {t_h:>10.2f} | {t_t:>10.2f} | "
+            f"{speedup:>7.2f}x | "
+            f"M={best['BLOCK_SIZE_M']},w={best['num_warps']}{marker}"
+        )
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+
+def save_results(device_name: str, output: str, results_file: str | None = None) -> str:
+    """Save the full benchmark output to a results text file."""
+    if results_file is None:
+        results_file = os.path.join(
+            _RESULTS_DIR, f"ssm_benchmark_results_{device_name}.txt"
+        )
+    with open(results_file, "w") as f:
+        f.write(output)
+    return results_file
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Tune selective_state_update kernel for Mamba SSM"
+    )
+    parser.add_argument(
+        "--dstate",
+        type=int,
+        default=128,
+        help="SSM state size to tune for (default: 128)",
+    )
+    parser.add_argument(
+        "--all-dstates",
+        action="store_true",
+        help="Tune all common dstate values: " + str(ALL_DSTATES),
+    )
+    parser.add_argument(
+        "--dtype",
+        type=str,
+        default="bfloat16",
+        choices=["float16", "bfloat16"],
+        help="Activation / input data type (default: bfloat16)",
+    )
+    parser.add_argument(
+        "--mamba-ssm-cache-dtype",
+        type=str,
+        default="float32",
+        choices=list(_SSM_CACHE_DTYPE_MAP.keys()),
+        help="SSM state cache dtype (default: float32)",
+    )
+    parser.add_argument(
+        "--num-iters",
+        type=int,
+        default=100,
+        help="Number of timing iterations (default: 100)",
+    )
+    parser.add_argument(
+        "--save-configs",
+        action="store_true",
+        help=f"Save best configs to JSON in {_CONFIGS_DIR}",
+    )
+    parser.add_argument(
+        "--compare",
+        action="store_true",
+        help="Show comparison table: heuristic vs tuned",
+    )
+    parser.add_argument(
+        "--verbose",
+        action="store_true",
+        help="Print every (BLOCK_SIZE_M, num_warps) result, not just best",
+    )
+    parser.add_argument(
+        "--results-file",
+        type=str,
+        default=None,
+        help="Path to save the benchmark results text file "
+        "(default: ssm_benchmark_results_<device>.txt alongside this script)",
+    )
+    parser.add_argument(
+        "--save-dir",
+        type=str,
+        default=None,
+        help=f"Directory to save JSON configs (default: {_CONFIGS_DIR})",
+    )
+    parser.add_argument(
+        "--headdim",
+        type=int,
+        default=DEFAULT_HEADDIM,
+        help=f"Per-head feature dim (default: {DEFAULT_HEADDIM})",
+    )
+    parser.add_argument(
+        "--ngroups",
+        type=int,
+        default=DEFAULT_NGROUPS,
+        help=f"Number of B/C groups (default: {DEFAULT_NGROUPS})",
+    )
+    parser.add_argument(
+        "--batch-sizes",
+        type=int,
+        nargs="+",
+        default=DEFAULT_BATCH_SIZES,
+        metavar="B",
+        help=f"Decoder batch sizes to sweep (default: {DEFAULT_BATCH_SIZES})",
+    )
+    parser.add_argument(
+        "--nheads",
+        type=int,
+        nargs="+",
+        default=DEFAULT_NHEADS,
+        metavar="N",
+        help=f"Number of heads per rank to sweep (default: {DEFAULT_NHEADS}). "
+        "effective_batch = batch * nheads; cross-product is deduped by eb.",
+    )
+    parser.add_argument(
+        "--validate",
+        action="store_true",
+        help="After tuning, verify each best config against a CPU reference "
+        "implementation. Configs that fail are flagged in the output.",
+    )
+    parser.add_argument(
+        "--atol",
+        type=float,
+        default=1e-2,
+        help="Absolute tolerance for --validate (default: 1e-2)",
+    )
+    args = parser.parse_args()
+
+    dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+    state_dtype = _SSM_CACHE_DTYPE_MAP[args.mamba_ssm_cache_dtype]
+    device_name = get_ssm_device_name()
+    cap = torch.cuda.get_device_capability()
+    is_blackwell = cap[0] >= 10
+
+    # Mirror all output to a results file (like Unix tee).
+    buf = StringIO()
+
+    class _Tee:
+        """Writes to both the original stdout and an in-memory buffer."""
+
+        def write(self, s):
+            buf.write(s)
+            sys.__stdout__.write(s)
+
+        def flush(self):
+            sys.__stdout__.flush()
+
+    sys.stdout = _Tee()  # type: ignore[assignment]
+
+    try:
+        print(f"Device : {device_name}  (sm_{cap[0]}{cap[1]})")
+        print(f"Blackwell: {is_blackwell}")
+        print(f"dtype  : {args.dtype}")
+        print(f"ssm_cache_dtype: {args.mamba_ssm_cache_dtype}")
+        print(f"headdim: {args.headdim}")
+        print(f"ngroups: {args.ngroups}")
+        print(f"triton : {triton.__version__}")
+
+        dstates = ALL_DSTATES if args.all_dstates else [args.dstate]
+        active = expand_batch_x_nheads(args.batch_sizes, args.nheads, args.ngroups)
+
+        for dstate in dstates:
+            tuned, timings = tune_dstate(
+                dstate=dstate,
+                headdim=args.headdim,
+                ngroups=args.ngroups,
+                dtype=dtype,
+                num_iters=args.num_iters,
+                verbose=args.verbose,
+                active=active,
+                state_dtype=state_dtype,
+            )
+
+            if args.compare:
+                compare_heuristic_vs_tuned(
+                    dstate=dstate,
+                    headdim=args.headdim,
+                    ngroups=args.ngroups,
+                    tuned=tuned,
+                    timings=timings,
+                    active=active,
+                    dtype=dtype,
+                    num_iters=args.num_iters,
+                    is_blackwell=is_blackwell,
+                    state_dtype=state_dtype,
+                )
+
+            if args.validate:
+                validity = validate_configs(
+                    dstate=dstate,
+                    headdim=args.headdim,
+                    ngroups=args.ngroups,
+                    tuned=tuned,
+                    active=active,
+                    dtype=dtype,
+                    atol=args.atol,
+                    state_dtype=state_dtype,
+                )
+                # Filter out any configs that failed correctness check
+                failed = [eb for eb, ok in validity.items() if not ok]
+                if failed:
+                    print(
+                        f"\n  WARNING: {len(failed)} config(s) failed validation "
+                        f"for dstate={dstate}: effective_batches {failed}"
+                    )
+                    print("  These will NOT be saved even with --save-configs.")
+                    tuned = {
+                        eb: cfg for eb, cfg in tuned.items() if validity.get(eb, True)
+                    }
+
+            if args.save_configs:
+                path = save_configs(
+                    headdim=args.headdim,
+                    dstate=dstate,
+                    cache_dtype=args.mamba_ssm_cache_dtype,
+                    configs=tuned,
+                    save_dir=args.save_dir,
+                )
+                print(f"\nSaved: {path}")
+            else:
+                print(f"\nBest configs for dstate={dstate}:")
+                for eb, cfg in sorted(tuned.items()):
+                    print(f"  effective_batch={eb:>6}: {cfg}")
+                print("\n(Re-run with --save-configs to persist to JSON)")
+    finally:
+        sys.stdout = sys.__stdout__
+        results_path = save_results(device_name, buf.getvalue(), args.results_file)
+        print(f"\nResults saved to: {results_path}")
+
+
+if __name__ == "__main__":
+    main()

build_rust.sh

@@ -0,0 +1,44 @@
+#!/bin/bash
+# Build the vllm-rs Rust frontend binary and install it into the vllm package.
+# Usage: ./build_rust.sh [--debug]
+#
+# By default builds in release mode. Pass --debug for faster compile times
+# during development.
+
+set -euo pipefail
+
+REPO_ROOT="$(cd "$(dirname "$0")" && pwd)"
+RUST_DIR="$REPO_ROOT/rust"
+TARGET_PATH="${VLLM_RS_TARGET_PATH:-$REPO_ROOT/vllm/vllm-rs}"
+
+# Read the required toolchain from rust-toolchain.toml.
+TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/')
+
+# Ensure rustup and the required toolchain are available.
+if ! command -v rustup &>/dev/null; then
+    echo "rustup not found, installing..."
+    curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain none
+    source "$HOME/.cargo/env"
+fi
+
+if ! rustup run "$TOOLCHAIN" rustc --version &>/dev/null; then
+    echo "Installing Rust toolchain: $TOOLCHAIN"
+    rustup toolchain install "$TOOLCHAIN"
+fi
+
+if [[ "${1:-}" == "--debug" ]]; then
+    PROFILE_ARGS=()
+    PROFILE_DIR="debug"
+else
+    PROFILE_ARGS=(--release)
+    PROFILE_DIR="release"
+fi
+
+cargo +"$TOOLCHAIN" build "${PROFILE_ARGS[@]}" \
+    --manifest-path "$RUST_DIR/Cargo.toml" \
+    --bin vllm-rs \
+    --features native-tls-vendored
+
+mkdir -p "$(dirname "$TARGET_PATH")"
+cp "$RUST_DIR/target/$PROFILE_DIR/vllm-rs" "$TARGET_PATH"
+echo "Installed vllm-rs to $TARGET_PATH"

cmake/external_projects/triton_kernels.cmake

@@ -1,6 +1,6 @@
 # Install OpenAI triton_kernels from https://github.com/triton-lang/triton/tree/main/python/triton_kernels
 
-set(DEFAULT_TRITON_KERNELS_TAG "v3.6.0")
+set(DEFAULT_TRITON_KERNELS_TAG "v3.5.1")
 
 # Set TRITON_KERNELS_SRC_DIR for use with local development with vLLM. We expect TRITON_KERNELS_SRC_DIR to
 # be directly set to the triton_kernels python directory.

cmake/utils.cmake

@@ -47,12 +47,17 @@ macro (append_cmake_prefix_path PKG EXPR)
   list(APPEND CMAKE_PREFIX_PATH ${_PREFIX_PATH})
 endmacro()
 
-#
-# Add a target named `hipify${NAME}` that runs the hipify preprocessor on a set
-# of CUDA source files. The names of the corresponding "hipified" sources are
-# stored in `OUT_SRCS`.
-#
+# Resolve hipified output paths for `NAME` into `OUT_SRCS` and register the
+# `.cu` sources with the shared `hipify_all` target. Per-extension hipify
+# targets are unsafe to run in parallel against a shared csrc/ output dir, so
+# accumulation here is paired with a single finalize step.
 function (hipify_sources_target OUT_SRCS NAME ORIG_SRCS)
+  if (TARGET hipify_all)
+    message(FATAL_ERROR
+      "hipify_sources_target(${NAME}) called after vllm_finalize_hipify_target. "
+      "Add the new HIP extension before the finalizer call in CMakeLists.txt.")
+  endif()
+
   #
   # Split into C++ and non-C++ (i.e. CUDA) sources.
   #
@@ -73,19 +78,41 @@ function (hipify_sources_target OUT_SRCS NAME ORIG_SRCS)
     list(APPEND HIP_SRCS "${CMAKE_CURRENT_BINARY_DIR}/${SRC}")
   endforeach()
 
-  set(CSRC_BUILD_DIR ${CMAKE_CURRENT_BINARY_DIR}/csrc)
-  add_custom_target(
-    hipify${NAME}
-    COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/cmake/hipify.py -p ${CMAKE_SOURCE_DIR}/csrc -o ${CSRC_BUILD_DIR} ${SRCS}
-    DEPENDS ${CMAKE_SOURCE_DIR}/cmake/hipify.py ${SRCS}
-    BYPRODUCTS ${HIP_SRCS}
-    COMMENT "Running hipify on ${NAME} extension source files.")
+  set_property(GLOBAL APPEND PROPERTY VLLM_HIPIFY_ALL_SRCS ${SRCS})
+  set_property(GLOBAL APPEND PROPERTY VLLM_HIPIFY_ALL_BYPRODUCTS ${HIP_SRCS})
 
   # Swap out original extension sources with hipified sources.
   list(APPEND HIP_SRCS ${CXX_SRCS})
   set(${OUT_SRCS} ${HIP_SRCS} PARENT_SCOPE)
 endfunction()
 
+# Define the single shared `hipify_all` custom target that runs hipify once
+# on the union of every HIP extension's sources. Call after the last HIP
+# `define_extension_target`.
+function (vllm_finalize_hipify_target)
+  if (TARGET hipify_all)
+    return()
+  endif()
+
+  get_property(ALL_SRCS GLOBAL PROPERTY VLLM_HIPIFY_ALL_SRCS)
+  get_property(ALL_BYPRODUCTS GLOBAL PROPERTY VLLM_HIPIFY_ALL_BYPRODUCTS)
+
+  if (NOT ALL_SRCS)
+    return()
+  endif()
+
+  list(REMOVE_DUPLICATES ALL_SRCS)
+  list(REMOVE_DUPLICATES ALL_BYPRODUCTS)
+
+  set(CSRC_BUILD_DIR ${CMAKE_CURRENT_BINARY_DIR}/csrc)
+  add_custom_target(
+    hipify_all
+    COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/cmake/hipify.py -p ${CMAKE_SOURCE_DIR}/csrc -o ${CSRC_BUILD_DIR} ${ALL_SRCS}
+    DEPENDS ${CMAKE_SOURCE_DIR}/cmake/hipify.py ${ALL_SRCS}
+    BYPRODUCTS ${ALL_BYPRODUCTS}
+    COMMENT "Running hipify on all extension source files.")
+endfunction()
+
 #
 # Get additional GPU compiler flags from torch.
 #
@@ -449,6 +476,16 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
   set(${OUT_CUDA_ARCHS} ${_CUDA_ARCHS} PARENT_SCOPE)
 endfunction()
 
+
+function(cuda_archs_sm90plus OUT_CUDA_ARCHS TGT_CUDA_ARCHS)
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f" "${TGT_CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a" "${TGT_CUDA_ARCHS}")
+  endif()
+  set(${OUT_CUDA_ARCHS} ${_archs} PARENT_SCOPE)
+endfunction()
+
 #
 # Override the GPU architectures detected by cmake/torch and filter them by
 # `GPU_SUPPORTED_ARCHES`. Sets the final set of architectures in
@@ -551,7 +588,7 @@ function (define_extension_target MOD_NAME)
 
   if (ARG_LANGUAGE STREQUAL "HIP")
     # Make this target dependent on the hipify preprocessor step.
-    add_dependencies(${MOD_NAME} hipify${MOD_NAME})
+    add_dependencies(${MOD_NAME} hipify_all)
     # Make sure we include the hipified versions of the headers, and avoid conflicts with the ones in the original source folder
     target_include_directories(${MOD_NAME} PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/csrc
       ${ARG_INCLUDE_DIRECTORIES})

csrc/attention/dtype_fp8.cuh

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "attention_generic.cuh"
+#include "torch_utils.h"
 
 #include <stdint.h>
 #ifdef ENABLE_FP8
@@ -30,7 +31,7 @@ inline Fp8KVCacheDataType get_fp8_kv_cache_data_type(
   } else if (dtype_str == "fp8_e5m2") {
     return Fp8KVCacheDataType::kFp8E5M2;
   }
-  TORCH_CHECK(false, "Unsupported fp8 kv cache data type: ", dtype_str);
+  TORCH_UTILS_CHECK(false, "Unsupported fp8 kv cache data type: ", dtype_str);
 }
 
 // fp8 vector types for quantization of kv cache

csrc/cpu/cpu_attn_impl.hpp

@@ -408,9 +408,19 @@ class AttentionScheduler {
     const int64_t cache_size = cpu_utils::get_available_l2_size();
     const int32_t max_num_q_per_iter = input.max_num_q_per_iter;
     const int32_t kv_len_alignment = input.kv_block_alignment;
+    bool has_decode_request = false;
+    bool decode_only_batch = true;
+    for (int32_t req_id = 0; req_id < input.num_reqs; ++req_id) {
+      const int32_t q_token_num =
+          input.query_start_loc[req_id + 1] - input.query_start_loc[req_id];
+      has_decode_request = has_decode_request || (q_token_num == 1);
+      decode_only_batch = decode_only_batch && (q_token_num == 1);
+    }
     int32_t q_head_per_kv = input.num_heads_q / input.num_heads_kv;
-    const bool use_gqa = (max_num_q_per_iter % q_head_per_kv == 0);
-    if (!use_gqa) {
+    const bool supports_gqa = q_head_per_kv <= max_num_q_per_iter;
+    const bool use_gqa_fast_path = supports_gqa && decode_only_batch;
+    const bool use_gqa_scratchpad = supports_gqa && has_decode_request;
+    if (!use_gqa_scratchpad) {
       q_head_per_kv = 1;  // fallback to MHA
     }
     const int32_t min_split_kv_len =
@@ -680,7 +690,7 @@ class AttentionScheduler {
         metadata_ptr->attention_scratchpad_size_per_thread *
             metadata_ptr->thread_num +
         metadata_ptr->reduction_scratchpad_size_per_kv_head *
-            (use_gqa ? input.num_heads_kv : input.num_heads_q);
+            (use_gqa_fast_path ? input.num_heads_kv : input.num_heads_q);
     cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(
         scratchpad_size);
 
@@ -1409,13 +1419,24 @@ class AttentionMainLoop {
       const int32_t q_head_num = input->num_heads;
       const int32_t kv_head_num = input->num_kv_heads;
       const int32_t q_heads_per_kv = q_head_num / kv_head_num;
-      const bool use_gqa =
-          (max_q_head_num_per_iter % q_heads_per_kv == 0) ? true : false;
-      const int32_t actual_kv_head_num = use_gqa ? kv_head_num : q_head_num;
-      const int32_t actual_q_heads_per_kv = use_gqa ? q_heads_per_kv : 1;
+      AttentionWorkItemGroup* const workitem_groups =
+          metadata.workitem_groups_ptr;
+      const int32_t* cu_workitem_num_per_thread =
+          metadata.cu_workitem_num_per_thread;
+      ReductionWorkItemGroup* const reduction_items =
+          metadata.reduction_items_ptr;
+      const bool supports_gqa = q_heads_per_kv <= max_q_head_num_per_iter;
+      bool decode_only_batch = true;
+      for (int32_t i = 0; i < metadata.workitem_group_num; ++i) {
+        decode_only_batch =
+            decode_only_batch && (workitem_groups[i].q_token_num == 1);
+      }
+      const bool use_gqa_fast_path = supports_gqa && decode_only_batch;
+      const int32_t actual_kv_head_num =
+          use_gqa_fast_path ? kv_head_num : q_head_num;
+      const int32_t actual_q_heads_per_kv =
+          use_gqa_fast_path ? q_heads_per_kv : 1;
       TORCH_CHECK_LE(actual_q_heads_per_kv, max_q_head_num_per_iter);
-      const int32_t max_q_token_num_per_iter =
-          max_q_head_num_per_iter / actual_q_heads_per_kv;
       const int64_t q_token_num_stride = input->query_num_tokens_stride;
       const int64_t q_head_num_stride = input->query_num_heads_stride;
       const int64_t kv_cache_head_num_stride = input->cache_num_kv_heads_stride;
@@ -1461,15 +1482,6 @@ class AttentionMainLoop {
               sizeof(q_buffer_t), sizeof(logits_buffer_t),
               sizeof(partial_output_buffer_t), max_q_head_num_per_iter,
               max_q_head_num_per_iter);
-      const int32_t default_q_tile_token_num =
-          default_tile_size / actual_q_heads_per_kv;
-
-      AttentionWorkItemGroup* const workitem_groups =
-          metadata.workitem_groups_ptr;
-      const int32_t* cu_workitem_num_per_thread =
-          metadata.cu_workitem_num_per_thread;
-      ReductionWorkItemGroup* const reduction_items =
-          metadata.reduction_items_ptr;
 
       const int32_t effective_thread_num = metadata.effective_thread_num;
       const int32_t reduction_item_num = metadata.reduction_item_num;
@@ -1513,8 +1525,6 @@ class AttentionMainLoop {
               cu_workitem_num_per_thread[thread_offset + 1] -
               cu_workitem_num_per_thread[thread_offset];
 
-          const int32_t q_head_start_idx = kv_head_idx * actual_q_heads_per_kv;
-
           for (int32_t workitem_group_idx = 0;
                workitem_group_idx < curr_workitem_groups_num;
                ++workitem_group_idx) {
@@ -1529,6 +1539,21 @@ class AttentionMainLoop {
             const int32_t q_token_id_start =
                 current_workitem_group->q_token_id_start;
             const int32_t q_token_num = current_workitem_group->q_token_num;
+            const bool curr_use_gqa =
+                use_gqa_fast_path || (supports_gqa && q_token_num == 1);
+            if (!use_gqa_fast_path && curr_use_gqa &&
+                kv_head_idx % q_heads_per_kv != 0) {
+              continue;
+            }
+            const int32_t curr_q_heads_per_kv =
+                curr_use_gqa ? q_heads_per_kv : 1;
+            const int32_t curr_max_q_token_num_per_iter =
+                max_q_head_num_per_iter / curr_q_heads_per_kv;
+            const int32_t curr_default_q_tile_token_num =
+                default_tile_size / curr_q_heads_per_kv;
+            const int32_t q_head_start_idx =
+                use_gqa_fast_path ? (kv_head_idx * q_heads_per_kv)
+                                  : kv_head_idx;
 
             // taskgroup general information
             const int32_t q_end = input->query_start_loc[current_group_idx + 1];
@@ -1542,7 +1567,7 @@ class AttentionMainLoop {
                              current_workitem_group->local_split_id == 0);
 
             for (int32_t q_token_offset = 0; q_token_offset < q_token_num;
-                 q_token_offset += default_q_tile_token_num) {
+                 q_token_offset += curr_default_q_tile_token_num) {
               bool first_iter_flag[AttentionScheduler::MaxQTileIterNum];
               for (int32_t i = 0; i < AttentionScheduler::MaxQTileIterNum;
                    ++i) {
@@ -1552,9 +1577,9 @@ class AttentionMainLoop {
               const int32_t q_token_start_idx =
                   q_start + q_token_offset + q_token_id_start;
               const int32_t actual_q_token_num = std::min(
-                  default_q_tile_token_num, q_token_num - q_token_offset);
+                  curr_default_q_tile_token_num, q_token_num - q_token_offset);
               const int32_t q_head_tile_size =
-                  actual_q_token_num * actual_q_heads_per_kv;
+                  actual_q_token_num * curr_q_heads_per_kv;
               const int32_t rounded_q_head_tile_size =
                   ((q_head_tile_size + max_q_head_num_per_iter - 1) /
                    max_q_head_num_per_iter) *
@@ -1591,10 +1616,9 @@ class AttentionMainLoop {
                   AttentionScheduler::align_kv_tile_pos(
                       kv_tile_start_pos, kv_tile_end_pos, blocksize_alignment);
 
-              int32_t curr_kv_head_idx =
-                  use_gqa ? kv_head_idx
-                          : (kv_head_idx /
-                             q_heads_per_kv);  // for GQA disabled case
+              const int32_t curr_kv_head_idx =
+                  use_gqa_fast_path ? kv_head_idx
+                                    : (kv_head_idx / q_heads_per_kv);
 
               // std::printf("thread_id: %d, req_id: %d, q_token_start: %d,
               // q_token_end: %d, q_head_start: %d, q_head_end: %d, kv_head_idx:
@@ -1629,12 +1653,12 @@ class AttentionMainLoop {
                   (s_aux != nullptr ? s_aux + q_head_start_idx : nullptr);
 
               // copy the Q tile to q_buffer, the logical layout of q_buffer is
-              // [actual_q_token_num, actual_q_heads_per_kv, head_dim]
+              // [actual_q_token_num, curr_q_heads_per_kv, head_dim]
               {
                 attn_impl.copy_q_heads_tile(
                     q_tile_ptr, q_buffer, actual_q_token_num,
-                    actual_q_heads_per_kv, q_token_num_stride,
-                    q_head_num_stride, scale);
+                    curr_q_heads_per_kv, q_token_num_stride, q_head_num_stride,
+                    scale);
               }
 
               if (use_sink) {
@@ -1648,29 +1672,29 @@ class AttentionMainLoop {
                 float* __restrict__ curr_max_buffer = max_buffer;
                 for (int32_t token_idx = 0; token_idx < actual_q_token_num;
                      ++token_idx) {
-                  for (int32_t head_idx = 0; head_idx < actual_q_heads_per_kv;
+                  for (int32_t head_idx = 0; head_idx < curr_q_heads_per_kv;
                        ++head_idx) {
                     curr_sum_buffer[head_idx] = 1.0f;
                     curr_max_buffer[head_idx] = s_aux_fp32[head_idx];
                   }
 
-                  curr_sum_buffer += actual_q_heads_per_kv;
-                  curr_max_buffer += actual_q_heads_per_kv;
+                  curr_sum_buffer += curr_q_heads_per_kv;
+                  curr_max_buffer += curr_q_heads_per_kv;
                 }
               } else {
                 float* __restrict__ curr_sum_buffer = sum_buffer;
                 float* __restrict__ curr_max_buffer = max_buffer;
                 for (int32_t token_idx = 0; token_idx < actual_q_token_num;
                      ++token_idx) {
-                  for (int32_t head_idx = 0; head_idx < actual_q_heads_per_kv;
+                  for (int32_t head_idx = 0; head_idx < curr_q_heads_per_kv;
                        ++head_idx) {
                     curr_sum_buffer[head_idx] = 0.0f;
                     curr_max_buffer[head_idx] =
                         std::numeric_limits<float>::lowest();
                   }
 
-                  curr_sum_buffer += actual_q_heads_per_kv;
-                  curr_max_buffer += actual_q_heads_per_kv;
+                  curr_sum_buffer += curr_q_heads_per_kv;
+                  curr_max_buffer += curr_q_heads_per_kv;
                 }
               }
 
@@ -1683,16 +1707,17 @@ class AttentionMainLoop {
                     kv_tile_pos_left + kv_tile_size, rounded_kv_tile_end_pos);
                 for (int32_t q_head_tile_token_offset = 0;
                      q_head_tile_token_offset < actual_q_token_num;
-                     q_head_tile_token_offset += max_q_token_num_per_iter) {
+                     q_head_tile_token_offset +=
+                     curr_max_q_token_num_per_iter) {
                   const int32_t q_tile_pos_left =
                       q_tile_start_pos + q_head_tile_token_offset;
                   const int32_t q_tile_token_num =
-                      std::min(max_q_token_num_per_iter,
+                      std::min(curr_max_q_token_num_per_iter,
                                actual_q_token_num - q_head_tile_token_offset);
                   const int32_t q_tile_head_offset =
-                      q_head_tile_token_offset * actual_q_heads_per_kv;
+                      q_head_tile_token_offset * curr_q_heads_per_kv;
                   const int32_t q_tile_head_num =
-                      q_tile_token_num * actual_q_heads_per_kv;
+                      q_tile_token_num * curr_q_heads_per_kv;
                   const int32_t q_tile_pos_right =
                       q_tile_pos_left + q_tile_token_num;
                   const auto [actual_kv_tile_pos_left,
@@ -1702,7 +1727,7 @@ class AttentionMainLoop {
                           q_tile_pos_right, sliding_window_left,
                           sliding_window_right);
                   const int32_t q_iter_idx =
-                      q_head_tile_token_offset / max_q_token_num_per_iter;
+                      q_head_tile_token_offset / curr_max_q_token_num_per_iter;
 
                   if (actual_kv_tile_pos_right <= actual_kv_tile_pos_left) {
                     continue;
@@ -1768,7 +1793,7 @@ class AttentionMainLoop {
                       aligned_actual_kv_tile_pos_left,
                       aligned_actual_kv_tile_pos_right, actual_kv_token_num,
                       kv_cache_block_num_stride, q_tile_head_num,
-                      q_tile_token_num, q_tile_pos_left, actual_q_heads_per_kv,
+                      q_tile_token_num, q_tile_pos_left, curr_q_heads_per_kv,
                       block_size, sliding_window_left, sliding_window_right,
                       scale, softcap_scale, curr_alibi_slopes,
                       first_iter_flag[q_iter_idx], use_sink, debug_info);
@@ -1782,11 +1807,11 @@ class AttentionMainLoop {
                   final_output(partial_q_buffer,
                                reinterpret_cast<query_t*>(input->output) +
                                    output_buffer_offset,
-                               sum_buffer, actual_q_heads_per_kv,
+                               sum_buffer, curr_q_heads_per_kv,
                                actual_q_token_num, q_head_num, output_v_scale);
                 } else {
                   const int32_t stride =
-                      actual_q_heads_per_kv * split_kv_q_token_num_threshold;
+                      curr_q_heads_per_kv * split_kv_q_token_num_threshold;
                   buffer_manager.update(kv_head_idx, total_reduction_split_num,
                                         head_dim, stride, sizeof(float));
                   volatile bool* split_flag_buffer =
@@ -1822,18 +1847,26 @@ class AttentionMainLoop {
           const int32_t curr_split_id = curr_workitem_groups->split_start_id;
           const int32_t curr_split_num = curr_workitem_groups->split_num;
           const int32_t current_group_idx = curr_workitem_groups->req_id;
+          const bool curr_use_gqa =
+              use_gqa_fast_path || (supports_gqa && curr_output_token_num == 1);
+          if (!use_gqa_fast_path && curr_use_gqa &&
+              kv_head_idx % q_heads_per_kv != 0) {
+            continue;
+          }
+          const int32_t curr_q_heads_per_kv = curr_use_gqa ? q_heads_per_kv : 1;
           const int32_t curr_output_head_num =
-              curr_output_token_num * actual_q_heads_per_kv;
+              curr_output_token_num * curr_q_heads_per_kv;
 
           const int32_t q_start = input->query_start_loc[current_group_idx];
           const int32_t q_token_start_idx = q_start + curr_output_token_idx;
-          const int32_t q_head_start_idx = kv_head_idx * actual_q_heads_per_kv;
+          const int32_t q_head_start_idx =
+              use_gqa_fast_path ? (kv_head_idx * q_heads_per_kv) : kv_head_idx;
           size_t output_buffer_offset =
               q_token_start_idx * q_head_num * head_dim +
               q_head_start_idx * head_dim;
 
           const int32_t stride =
-              actual_q_heads_per_kv * split_kv_q_token_num_threshold;
+              curr_q_heads_per_kv * split_kv_q_token_num_threshold;
           buffer_manager.update(kv_head_idx, total_reduction_split_num,
                                 head_dim, stride, sizeof(float));
           volatile bool* split_flag_buffer =
@@ -1852,7 +1885,7 @@ class AttentionMainLoop {
           final_output(
               split_output_buffer,
               reinterpret_cast<query_t*>(input->output) + output_buffer_offset,
-              split_sum_buffer, actual_q_heads_per_kv, curr_output_token_num,
+              split_sum_buffer, curr_q_heads_per_kv, curr_output_token_num,
               q_head_num, output_v_scale);
         }
       }

csrc/cpu/cpu_attn_rvv.hpp

@@ -4,17 +4,18 @@
 #ifndef CPU_ATTN_RVV_HPP
 #define CPU_ATTN_RVV_HPP
 
-// This kernel is currently hardcoded to VLEN=128 (m1/m2 intrinsics, vl=8).
-// The fixed-width typedefs below use `riscv_rvv_vector_bits(128)`, which
-// only matches `vfloat16m1_t`/`vuint16m1_t` register layout when VLEN==128;
-// at VLEN>=256 those typedefs fail to compile.  Scalar RISC-V builds
-// (-march=rv64gc) additionally don't have <riscv_vector.h>.  For both
-// cases we omit the file entirely and let the dispatcher fall back to the
-// scalar VEC / VEC16 implementations.  TODO: migrate to RVVI() macros +
-// semantic names in cpu_types_riscv_defs.hpp to support VLEN>=256 natively.
-#if defined(__riscv_v_min_vlen) && __riscv_v_min_vlen == 128
+// RVV attention kernel using VLEN-agnostic RVVI() macros from
+// cpu_types_riscv_defs.hpp.  The Mx8 tile GEMM uses 8 FP32 elements
+// per vector (LMUL_256 bits of FP32 data), which maps to:
+//   VLEN=128: m2 (256 bits = 8 x FP32)
+//   VLEN=256: m1 (256 bits = 8 x FP32)
+// Only VLEN=128 and VLEN=256 are supported; other VLENs (512, 1024)
+// and scalar RISC-V builds fall back to VEC/VEC16.
+#if defined(__riscv_v_min_vlen) && \
+    (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)
 
   #include "cpu_attn_impl.hpp"
+  #include "cpu_types_riscv_defs.hpp"
   #include <riscv_vector.h>
   #include <type_traits>
 
@@ -22,73 +23,50 @@ namespace cpu_attention {
 
 namespace {
 
-// File-local concrete-LMUL typedefs.  The shared _defs.hpp exposes
-// VLEN-independent semantic names (fixed_fp32x8_t, fixed_fp16x8_t, ...),
-// but this kernel is currently hardcoded to VLEN=128 (m1/m2 intrinsics),
-// so keep the legacy concrete aliases scoped to this file.
-typedef vfloat16m1_t fixed_vfloat16m1_t
-    __attribute__((riscv_rvv_vector_bits(128)));
-typedef vfloat32m2_t fixed_vfloat32m2_t
-    __attribute__((riscv_rvv_vector_bits(256)));
-typedef vuint16m1_t fixed_vuint16m1_t
-    __attribute__((riscv_rvv_vector_bits(128)));
-typedef vuint32m2_t fixed_vuint32m2_t
-    __attribute__((riscv_rvv_vector_bits(256)));
-  #ifdef __riscv_zvfbfmin
-typedef vbfloat16m1_t fixed_vbfloat16m1_t
-    __attribute__((riscv_rvv_vector_bits(128)));
-  #endif
-
   #define BLOCK_SIZE_ALIGNMENT 32
   #define HEAD_SIZE_ALIGNMENT 32
   #define MAX_Q_HEAD_NUM_PER_ITER 16
 
 // ============================================================================
-// B-matrix row loading: load 8 elements as FP32 (using m2 LMUL at VLEN=128)
+// B-matrix row loading: load 8 elements as FP32
 // ============================================================================
 
 template <typename kv_cache_t>
-FORCE_INLINE fixed_vfloat32m2_t load_row8_B_as_f32(const kv_cache_t* p);
+FORCE_INLINE fixed_fp32x8_t load_row8_B_as_f32(const kv_cache_t* p);
 
 template <>
-FORCE_INLINE fixed_vfloat32m2_t load_row8_B_as_f32<float>(const float* p) {
-  return __riscv_vle32_v_f32m2(p, 8);
+FORCE_INLINE fixed_fp32x8_t load_row8_B_as_f32<float>(const float* p) {
+  return RVVI(__riscv_vle32_v_f32, LMUL_256)(p, 8);
 }
 
 template <>
-FORCE_INLINE fixed_vfloat32m2_t
-load_row8_B_as_f32<c10::Half>(const c10::Half* p) {
+FORCE_INLINE fixed_fp32x8_t load_row8_B_as_f32<c10::Half>(const c10::Half* p) {
   #ifdef __riscv_zvfh
-  fixed_vfloat16m1_t h =
-      __riscv_vle16_v_f16m1(reinterpret_cast<const _Float16*>(p), 8);
-  return __riscv_vfwcvt_f_f_v_f32m2(h, 8);
+  fixed_fp16x8_t h = RVVI(__riscv_vle16_v_f16, LMUL_128)(
+      reinterpret_cast<const _Float16*>(p), 8);
+  return RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(h, 8);
   #else
-  // Fallback for hardware without Zvfh: scalar half->float conversion.
-  // c10::Half provides operator float() so this is correct on any RVV CPU
-  // that has only the base V extension.  Slower than the Zvfh path, but
-  // keeps the kernel buildable on Zvfhmin-only / no-fp16 hardware.
   alignas(16) float tmp[8];
   for (int i = 0; i < 8; ++i) {
     tmp[i] = static_cast<float>(p[i]);
   }
-  return __riscv_vle32_v_f32m2(tmp, 8);
+  return RVVI(__riscv_vle32_v_f32, LMUL_256)(tmp, 8);
   #endif
 }
 
 template <>
-FORCE_INLINE fixed_vfloat32m2_t
+FORCE_INLINE fixed_fp32x8_t
 load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p) {
   #ifdef __riscv_zvfbfmin
-  fixed_vbfloat16m1_t bf =
-      __riscv_vle16_v_bf16m1(reinterpret_cast<const __bf16*>(p), 8);
-  return __riscv_vfwcvtbf16_f_f_v_f32m2(bf, 8);
+  fixed_bf16x8_t bf = RVVI(__riscv_vle16_v_bf16, LMUL_128)(
+      reinterpret_cast<const __bf16*>(p), 8);
+  return RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(bf, 8);
   #else
-  // Fallback: load as uint16, zero-extend to uint32, shift left by 16
-  fixed_vuint16m1_t raw =
-      __riscv_vle16_v_u16m1(reinterpret_cast<const uint16_t*>(p), 8);
-  fixed_vuint32m2_t wide = __riscv_vzext_vf2_u32m2(raw, 8);
-  fixed_vuint32m2_t shifted = __riscv_vsll_vx_u32m2(wide, 16, 8);
-  return __riscv_vreinterpret_v_u32m2_f32m2(shifted);
+  fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)(
+      reinterpret_cast<const uint16_t*>(p), 8);
+  fixed_u32x8_t wide = RVVI(__riscv_vzext_vf2_u32, LMUL_256)(raw, 8);
+  fixed_u32x8_t shifted = RVVI(__riscv_vsll_vx_u32, LMUL_256)(wide, 16, 8);
+  return RVVI4(__riscv_vreinterpret_v_u32, LMUL_256, _f32, LMUL_256)(shifted);
   #endif
 }
 
@@ -96,14 +74,12 @@ load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p) {
 // Micro kernel: Mx8 tile, K unrolled by 4, RVV scalar-broadcast FMA
 // ============================================================================
 //
-// NEON uses vfmaq_laneq_f32 (lane-indexed FMA from a preloaded A vector).
 // RVV has no lane-indexed FMA; instead we load A elements as scalars and
-// use __riscv_vfmacc_vf (scalar * vector + accumulator), which is equally
-// efficient and avoids the need for vrgather/vslidedown.
+// use vfmacc_vf (scalar * vector + accumulator).
 //
-// At VLEN=128, m2 holds 8 x FP32, matching the 8-column tile width.
-// Register budget: M accumulators (m2 each) + 1 B temp = 2M+2 regs.
-// M=8 => 18 regs out of 32 available — no spills.
+// The 8-column tile uses LMUL_256 bits of FP32 data:
+//   VLEN=128: m2 (2 regs per accumulator), M=8 => 18 of 32 regs
+//   VLEN=256: m1 (1 reg  per accumulator), M=8 =>  9 of 32 regs
 
 template <int32_t M, typename kv_cache_t>
 FORCE_INLINE void gemm_micro_rvv_fma_Mx8_Ku4(
@@ -115,94 +91,90 @@ FORCE_INLINE void gemm_micro_rvv_fma_Mx8_Ku4(
 
   constexpr size_t vl = 8;
 
-  // helpers for per-M codegen
   #define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
   #define IF_M(i) if constexpr (M > (i))
 
-    // A row base pointers
   #define DECL_A(i) const float* a##i = A + (i) * lda;
   ROWS_APPLY(DECL_A)
   #undef DECL_A
 
-    // declare one m2 accumulator per row
-  #define DECL_ACC(i) fixed_vfloat32m2_t acc##i;
+  #define DECL_ACC(i) fixed_fp32x8_t acc##i;
   ROWS_APPLY(DECL_ACC)
   #undef DECL_ACC
 
-    // initialize accumulators
-  #define INIT_ACC(i)                                      \
-    IF_M(i) {                                              \
-      if (accumulate) {                                    \
-        acc##i = __riscv_vle32_v_f32m2(C + (i) * ldc, vl); \
-      } else {                                             \
-        acc##i = __riscv_vfmv_v_f_f32m2(0.f, vl);          \
-      }                                                    \
+  #define INIT_ACC(i)                                                    \
+    IF_M(i) {                                                            \
+      if (accumulate) {                                                  \
+        acc##i = RVVI(__riscv_vle32_v_f32, LMUL_256)(C + (i) * ldc, vl); \
+      } else {                                                           \
+        acc##i = RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.f, vl);          \
+      }                                                                  \
     }
   ROWS_APPLY(INIT_ACC)
   #undef INIT_ACC
 
   int32_t k = 0;
 
-  // K unrolled by 4
   for (; k + 3 < K; k += 4) {
-    // k + 0
     {
-      fixed_vfloat32m2_t b =
+      fixed_fp32x8_t b =
           load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 0) * ldb);
-  #define STEP_K0(i)                                                    \
-    IF_M(i) {                                                           \
-      acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k + 0), b, vl); \
+  #define STEP_K0(i)                                                          \
+    IF_M(i) {                                                                 \
+      acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)(acc##i, *(a##i + k + 0), \
+                                                     b, vl);                  \
     }
       ROWS_APPLY(STEP_K0)
   #undef STEP_K0
     }
-    // k + 1
     {
-      fixed_vfloat32m2_t b =
+      fixed_fp32x8_t b =
           load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 1) * ldb);
-  #define STEP_K1(i)                                                    \
-    IF_M(i) {                                                           \
-      acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k + 1), b, vl); \
+  #define STEP_K1(i)                                                          \
+    IF_M(i) {                                                                 \
+      acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)(acc##i, *(a##i + k + 1), \
+                                                     b, vl);                  \
     }
       ROWS_APPLY(STEP_K1)
   #undef STEP_K1
     }
-    // k + 2
     {
-      fixed_vfloat32m2_t b =
+      fixed_fp32x8_t b =
           load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 2) * ldb);
-  #define STEP_K2(i)                                                    \
-    IF_M(i) {                                                           \
-      acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k + 2), b, vl); \
+  #define STEP_K2(i)                                                          \
+    IF_M(i) {                                                                 \
+      acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)(acc##i, *(a##i + k + 2), \
+                                                     b, vl);                  \
     }
       ROWS_APPLY(STEP_K2)
   #undef STEP_K2
     }
-    // k + 3
     {
-      fixed_vfloat32m2_t b =
+      fixed_fp32x8_t b =
           load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 3) * ldb);
-  #define STEP_K3(i)                                                    \
-    IF_M(i) {                                                           \
-      acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k + 3), b, vl); \
+  #define STEP_K3(i)                                                          \
+    IF_M(i) {                                                                 \
+      acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)(acc##i, *(a##i + k + 3), \
+                                                     b, vl);                  \
     }
       ROWS_APPLY(STEP_K3)
   #undef STEP_K3
     }
   }
 
-  // K tail
   for (; k < K; ++k) {
-    fixed_vfloat32m2_t b = load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb);
-  #define TAIL_ROW(i) \
-    IF_M(i) { acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k), b, vl); }
+    fixed_fp32x8_t b = load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb);
+  #define TAIL_ROW(i)                                                        \
+    IF_M(i) {                                                                \
+      acc##i =                                                               \
+          RVVI(__riscv_vfmacc_vf_f32, LMUL_256)(acc##i, *(a##i + k), b, vl); \
+    }
     ROWS_APPLY(TAIL_ROW)
   #undef TAIL_ROW
   }
 
-    // store accumulators to C
   #define STORE_ROW(i) \
-    IF_M(i) { __riscv_vse32_v_f32m2(C + (i) * ldc, acc##i, vl); }
+    IF_M(i) { RVVI(__riscv_vse32_v_f32, LMUL_256)(C + (i) * ldc, acc##i, vl); }
   ROWS_APPLY(STORE_ROW)
   #undef STORE_ROW
 
@@ -381,7 +353,6 @@ class AttentionImpl<ISA::RVV, scalar_t, head_dim, kv_cache_scalar_t> {
         const int64_t block_idx = pos / block_size;
         const int64_t block_offset = pos % block_size;
         {
-          // Write Key (transpose to column-major: [head_dim, block_size])
           const scalar_t* key_start_ptr = key +
                                           token_idx * key_token_num_stride +
                                           head_idx * key_head_num_stride;
@@ -389,8 +360,6 @@ class AttentionImpl<ISA::RVV, scalar_t, head_dim, kv_cache_scalar_t> {
               key_cache + block_idx * num_blocks_stride +
               head_idx * cache_head_num_stride + block_offset;
 
-          // Strided vector store for efficient transpose.
-          // Load contiguous key elements, store with stride = block_size.
           {
             const ptrdiff_t byte_stride = block_size * sizeof(scalar_t);
             int64_t i = 0;
@@ -405,7 +374,6 @@ class AttentionImpl<ISA::RVV, scalar_t, head_dim, kv_cache_scalar_t> {
                                              i * block_size),
                     byte_stride, v, vl);
               } else {
-                // Half and BFloat16 are both 16-bit types
                 vl = __riscv_vsetvl_e16m1(head_dim - i);
                 vuint16m1_t v = __riscv_vle16_v_u16m1(
                     reinterpret_cast<const uint16_t*>(key_start_ptr + i), vl);
@@ -419,7 +387,6 @@ class AttentionImpl<ISA::RVV, scalar_t, head_dim, kv_cache_scalar_t> {
           }
         }
         {
-          // Write Value (row-major: [block_size, head_dim])
           const scalar_t* value_start_ptr = value +
                                             token_idx * value_token_num_stride +
                                             head_idx * value_head_num_stride;
@@ -440,6 +407,6 @@ class AttentionImpl<ISA::RVV, scalar_t, head_dim, kv_cache_scalar_t> {
   #undef HEAD_SIZE_ALIGNMENT
   #undef MAX_Q_HEAD_NUM_PER_ITER
 
-#endif  // __riscv_v_min_vlen == 128
+#endif  // __riscv_v_min_vlen == 128 || 256
 
 #endif  // CPU_ATTN_RVV_HPP

csrc/cpu/cpu_types_riscv_defs.hpp

@@ -71,6 +71,10 @@ typedef RVVTYPE(vuint16, LMUL_256, _t) fixed_u16x16_t
 typedef RVVTYPE(vuint16, LMUL_512, _t) fixed_u16x32_t
     __attribute__((riscv_rvv_vector_bits(512)));
 
+// uint32
+typedef RVVTYPE(vuint32, LMUL_256, _t) fixed_u32x8_t
+    __attribute__((riscv_rvv_vector_bits(256)));
+
 // bfloat16
 #ifdef __riscv_zvfbfmin
 typedef RVVTYPE(vbfloat16, LMUL_128, _t) fixed_bf16x8_t

csrc/cpu/generate_cpu_attn_dispatch.py

@@ -150,12 +150,10 @@ def generate_header_file() -> str:
   #include "cpu_attn_vxe.hpp"
 #endif
 
-// cpu_attn_rvv.hpp is hardcoded to VLEN==128 (m1/m2 intrinsics, vl=8) and
-// itself includes <riscv_vector.h>, which is unavailable on scalar
-// (-march=rv64gc) builds. Gate the include the same way as the dispatch
-// macro below, so non-128 / scalar RISC-V builds skip it entirely.
+// cpu_attn_rvv.hpp supports VLEN=128 and VLEN=256 via RVVI() macros.
+// Other VLENs and scalar RISC-V builds skip it entirely.
 #if defined(__riscv) && defined(__riscv_v_min_vlen) && \
-    __riscv_v_min_vlen == 128
+    (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)
   #include "cpu_attn_rvv.hpp"
 #endif
 
@@ -222,15 +220,12 @@ def generate_header_file() -> str:
         ["VXE", "VEC", "VEC16"],
         fp8=False,
     )
-    # RISC-V with RVV.  cpu_attn_rvv.hpp is hardcoded to VLEN==128
-    # (riscv_rvv_vector_bits(128) typedefs + vl=8 m1/m2 intrinsics), so
-    # we split the dispatch into two top-level branches: VLEN==128 builds
-    # get the full RVV+VEC+VEC16 case set, other VLEN builds get a
-    # VEC/VEC16-only fallback.  Preprocessor directives cannot appear
-    # inside a #define body, so this duplication is necessary.
+    # RISC-V with RVV.  cpu_attn_rvv.hpp supports VLEN=128 and VLEN=256
+    # via RVVI() macros.  Builds with a supported VLEN get
+    # RVV+VEC+VEC16; other RISC-V builds fall back to VEC/VEC16 only.
     header += _macro_block(
         "#elif defined(__riscv) && defined(__riscv_v_min_vlen) "
-        "&& __riscv_v_min_vlen == 128",
+        "&& (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)",
         ["RVV", "VEC", "VEC16"],
         fp8=False,
     )

csrc/cutlass_extensions/torch_utils.hpp

@@ -1,5 +1,7 @@
 #pragma once
 
+#include "torch_utils.h"
+
 // This header is shared between _C (unstable ABI, used by machete) and
 // _C_stable_libtorch (stable ABI, used by W4A8/sparse). TORCH_TARGET_VERSION
 // is defined only for the stable target, so we switch includes and types
@@ -8,13 +10,9 @@
   #include <torch/csrc/stable/tensor.h>
   #include <torch/headeronly/util/BFloat16.h>
   #include <torch/headeronly/util/Half.h>
-  #include <torch/headeronly/util/shim_utils.h>  // for STD_TORCH_CHECK
 using TorchTensor = torch::stable::Tensor;
-  #define TORCH_UTILS_CHECK STD_TORCH_CHECK
 #else
-  #include <torch/all.h>
 using TorchTensor = torch::Tensor;
-  #define TORCH_UTILS_CHECK TORCH_CHECK
 #endif
 
 #include "cute/layout.hpp"

csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu

@@ -87,6 +87,12 @@ constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1;  // 8 (7 real + 1 pad)
 constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2;  // 448 + 128 = 576
 constexpr float kFp8Max = 448.0f;
 
+#ifndef USE_ROCM
+// When num_tokens is less than this threshold,
+// run the reduced grid variant on cuda
+constexpr float NUM_TOKEN_CUTOFF = 1024;
+#endif
+
 // Per-warp layout:  32 lanes × 16 elems/lane = 512 elems = HEAD_DIM.
 constexpr int kNumLanes = 32;
 constexpr int kElemsPerLane = kHeadDim / kNumLanes;  // 16
@@ -112,24 +118,255 @@ __device__ __forceinline__ float warpSum(float val) {
   return val;
 }
 
+// ────────────────────────────────────────────────────────────────────────────
+// Per-slot inner pipeline
+// ────────────────────────────────────────────────────────────────────────────
+// Shared by both kernel variants: 1 CTA per (token, head) pair vs. 1 CTA per
+// token.  Templated on `kNumHeadsQPadded` so the KV-sentinel comparison and
+// q_out stride fold to compile-time constants.
+//
+// Slot layout (per token):
+//   slot < num_heads_q                          → live-Q   (RMSNorm + RoPE,
+//                                                           read q_in →
+//                                                           write q_out)
+//   num_heads_q <= slot < kNumHeadsQPadded      → pad-Q    (zero-fill q_out;
+//                                                           v0/v1 unused)
+//   slot == kNumHeadsQPadded                    → KV       (RoPE + UE8M0 quant
+//                                                           + paged-cache
+//                                                           insert)
+template <typename scalar_t_in, int kNumHeadsQPadded>
+__device__ __forceinline__ void processDeepseekV4Slot(
+    uint4 v0, uint4 v1, int const tokenIdx, int const slotIdx,
+    int const dim_base, int const laneId, int const num_heads_q,
+    float const eps, scalar_t_in* __restrict__ q_out,
+    uint8_t* __restrict__ k_cache, int64_t const* __restrict__ slot_mapping,
+    int64_t const* __restrict__ position_ids,
+    float const* __restrict__ cos_sin_cache, int const cache_block_size,
+    int const kv_block_stride) {
+  using Converter = vllm::_typeConvert<scalar_t_in>;
+  bool const isKV = (slotIdx == kNumHeadsQPadded);
+  bool const isPadQ = !isKV && (slotIdx >= num_heads_q);
+
+  // ── Pad-Q branch: write 32 B of zeros and exit. ─────────────────────────
+  // FlashMLA reads these slots; bf16 +0.0 is bit pattern 0x0000, so a uint4
+  // zero literal is correct.  Matches the live-Q branch's vectorized store.
+  if (isPadQ) {
+    scalar_t_in* dst =
+        q_out +
+        (static_cast<int64_t>(tokenIdx) * kNumHeadsQPadded + slotIdx) *
+            kHeadDim +
+        dim_base;
+    uint4 const zero4 = {0u, 0u, 0u, 0u};
+    *reinterpret_cast<uint4*>(dst) = zero4;
+    *reinterpret_cast<uint4*>(dst + 8) = zero4;
+    return;
+  }
+
+  // ── Decode the bf16 → 16 fp32 registers ─────────────────────────────
+  float elements[kElemsPerLane];
+  {
+    typename Converter::packed_hip_type const* p0 =
+        reinterpret_cast<typename Converter::packed_hip_type const*>(&v0);
+    typename Converter::packed_hip_type 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) ─────────────────────────────
+  // All math in fp32.  cos_sin_cache is loaded as fp32 (its native storage).
+  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;
+
+    // Load phase: 4 vectorized LDGs issue back-to-back.
+    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];
+    }
+  }
+
+  // ═══════════════════════════════════════════════════════════════════
+  // Q / KV branch dispatch. Restructured as if/else (no early `return`)
+  // so every code path lands at the same exit point — callers own PDL
+  // triggering and per-iteration buffer rotation.
+  // ═══════════════════════════════════════════════════════════════════
+  if (!isKV) {
+    // ── Live-Q: cast back to bf16 and store into the padded q_out. ─────
+    uint4 out0, out1;
+    typename Converter::packed_hip_type* po0 =
+        reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
+    typename Converter::packed_hip_type* 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_out +
+        (static_cast<int64_t>(tokenIdx) * kNumHeadsQPadded + slotIdx) *
+            kHeadDim +
+        dim_base;
+    *reinterpret_cast<uint4*>(dst) = out0;
+    *reinterpret_cast<uint4*>(dst + 8) = out1;
+  } else {
+    // ── KV: FP8 quant on NoPE + bf16 store on RoPE + cache insert.
+    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* block_base =
+          k_cache + block_idx * static_cast<int64_t>(kv_block_stride);
+      uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes;
+      uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim;
+      uint8_t* token_scale_ptr =
+          block_base +
+          static_cast<int64_t>(cache_block_size) * kTokenDataBytes +
+          pos_in_block * kScaleBytesPerToken;
+
+#pragma unroll
+      for (int i = 0; i < kElemsPerLane; i++) {
+        elements[i] = Converter::convert(Converter::convert(elements[i]));
+      }
+
+      float local_absmax = 0.0f;
+#pragma unroll
+      for (int i = 0; i < kElemsPerLane; i++) {
+        local_absmax = fmaxf(local_absmax, fabsf(elements[i]));
+      }
+      float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f);
+      float const exponent = ceilf(log2f(absmax / kFp8Max));
+      float const inv_scale = exp2f(-exponent);
+
+      if (!is_rope_lane) {
+        uint8_t out_bytes[kElemsPerLane];
+#pragma unroll
+        for (int i = 0; i < kElemsPerLane; i++) {
+          float scaled = elements[i] * inv_scale;
+          scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
+#ifndef USE_ROCM
+          __nv_fp8_storage_t s =
+              __nv_cvt_float_to_fp8(scaled, __NV_SATFINITE, __NV_E4M3);
+          out_bytes[i] = static_cast<uint8_t>(s);
+#else
+          out_bytes[i] = rocm_cvt_float_to_fp8_e4m3(scaled);
+#endif
+        }
+        *reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
+            *reinterpret_cast<uint4 const*>(out_bytes);
+
+        if ((laneId & 3) == 0) {
+          int const q_block_idx = laneId >> 2;
+          float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f);
+          token_scale_ptr[q_block_idx] = static_cast<uint8_t>(encoded);
+        }
+        if (laneId == 0) {
+          token_scale_ptr[kNumQuantBlocks] = 0;
+        }
+      } else {
+        uint4 out0, out1;
+        typename Converter::packed_hip_type* po0 =
+            reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
+        typename Converter::packed_hip_type* 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]));
+        }
+        int const rope_local_base = dim_base - kNopeDim;
+        scalar_t_in* bf16_dst =
+            reinterpret_cast<scalar_t_in*>(token_bf16_ptr) + rope_local_base;
+        *reinterpret_cast<uint4*>(bf16_dst) = out0;
+        *reinterpret_cast<uint4*>(bf16_dst + 8) = out1;
+      }
+    }
+  }
+}
+
 // ────────────────────────────────────────────────────────────────────────────
 // Kernel
 // ────────────────────────────────────────────────────────────────────────────
 //
-// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) /
+// Grid: 1D, gridDim.x = ceil(num_tokens_full * (kNumHeadsQPadded + 1) /
 // warps_per_block) Block: blockDim.x = 256 threads (8 warps per block) Each
-// warp handles one (token, head_slot) pair. head_slot < num_heads_q          →
-// Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q         → KV
-// branch (RoPE + UE8M0 quant + insert)
+// warp handles one (token, head_slot) pair.
+//   slot < num_heads_q                              → live-Q branch
+//                                                     (RMSNorm + RoPE,
+//                                                      read q_in → write q_out)
+//   num_heads_q <= slot < kNumHeadsQPadded          → pad-Q branch
+//                                                     (zero-fill q_out)
+//   slot == kNumHeadsQPadded                        → KV branch
+//                                                     (RoPE + UE8M0 quant +
+//                                                      paged-cache insert)
+//
+// `kNumHeadsQPadded` is a template parameter (compile-time constant) so the
+// divisions in the grid math and the KV-sentinel comparison fold to fast
+// constant operations.  The launch wrapper dispatches the runtime value to
+// the matching instantiation.
 //
 // With DP padding, q/kv/position_ids can have more rows than slot_mapping.
-// The Q branch covers all `num_tokens_full` rows (downstream attention uses
-// them).  The KV branch only inserts the first `num_tokens_insert` tokens
-// (= slot_mapping length) into the paged cache.
+// The live-Q and pad-Q branches cover all `num_tokens_full` rows (downstream
+// attention uses them).  The KV branch only inserts the first
+// `num_tokens_insert` tokens (= slot_mapping length) into the paged cache.
 //
-template <typename scalar_t_in>
+template <typename scalar_t_in, int kNumHeadsQPadded>
 __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
-    scalar_t_in* __restrict__ q_inout,         // [N, H, 512] bf16, in place
+    scalar_t_in const* __restrict__ q_in,      // [N, num_heads_q,      512]
+    scalar_t_in* __restrict__ q_out,           // [N, kNumHeadsQPadded, 512]
     scalar_t_in const* __restrict__ kv_in,     // [N, 512] bf16
     uint8_t* __restrict__ k_cache,             // [num_blocks, block_stride]
     int64_t const* __restrict__ slot_mapping,  // [num_tokens_insert] i64
@@ -138,7 +375,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
     float const eps,
     int const num_tokens_full,    // = q.size(0) = kv.size(0)
     int const num_tokens_insert,  // = slot_mapping.size(0), ≤ num_tokens_full
-    int const num_heads_q,        // H
+    int const num_heads_q,        // live Q heads (input layout)
     int const cache_block_size,   // tokens per paged-cache block
     int const kv_block_stride) {  // bytes per paged-cache block
 #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
@@ -149,19 +386,18 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
     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 total_slots_per_token = num_heads_q + 1;
-    int const tokenIdx = globalWarpIdx / total_slots_per_token;
-    int const slotIdx = globalWarpIdx % total_slots_per_token;
+    constexpr int kTotalSlotsPerToken = kNumHeadsQPadded + 1;
+    int const tokenIdx = globalWarpIdx / kTotalSlotsPerToken;
+    int const slotIdx = globalWarpIdx % kTotalSlotsPerToken;
     if (tokenIdx >= num_tokens_full) return;
 
-    bool const isKV = (slotIdx == num_heads_q);
+    bool const isKV = (slotIdx == kNumHeadsQPadded);
+    bool const isPadQ = !isKV && (slotIdx >= num_heads_q);
     // KV branch: skip DP-padded tokens (no slot reserved for them).
     if (isKV && tokenIdx >= num_tokens_insert) return;
 
@@ -176,209 +412,123 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
     // Dim range this lane owns within the 512-wide head.
     int const dim_base = laneId * kElemsPerLane;  // in [0, 512) step 16
 
-    // ── Load 16 bf16 → 16 fp32 registers (one 16-byte + one 16-byte LDG) ────
-    float elements[kElemsPerLane];
-    float sumOfSquares = 0.0f;
-
-    scalar_t_in const* src_ptr;
-    if (isKV) {
-      src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
-    } else {
-      int64_t const q_row_offset =
-          (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
-          dim_base;
-      src_ptr = q_inout + q_row_offset;
+    // Load only for live-Q and KV slots; pad-Q skips the read (q_in beyond
+    // num_heads_q is out of bounds) and the helper zero-fills its output.
+    uint4 v0, v1;
+    if (!isPadQ) {
+      scalar_t_in const* src_ptr;
+      if (isKV) {
+        src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
+      } else {
+        int64_t const q_row_offset =
+            (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) *
+                kHeadDim +
+            dim_base;
+        src_ptr = q_in + q_row_offset;
+      }
+      v0 = *reinterpret_cast<uint4 const*>(src_ptr);
+      v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
     }
 
-    // Two 16-byte loads per thread (8 bf16 each).  Use uint4 as the vector
-    // type and bitcast to scalar_t_in packed pairs for conversion.
-    uint4 v0 = *reinterpret_cast<uint4 const*>(src_ptr);
-    uint4 v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
+    processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded>(
+        v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_out,
+        k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size,
+        kv_block_stride);
 
-    {
-      typename Converter::packed_hip_type const* p0 =
-          reinterpret_cast<typename Converter::packed_hip_type const*>(&v0);
-      typename Converter::packed_hip_type const* p1 =
-          reinterpret_cast<typename Converter::packed_hip_type const*>(&v1);
-// Each packed_hip_type holds 2 bf16 → 4 packed = 8 elems per uint4.
-#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 with no weight (has_weight=False) ─────────────────
-    // Variance + rsqrt + multiply all in fp32, no intermediate bf16 round.
-    // The downstream bf16 round only happens at the final store.
-    if (!isKV) {
-#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) ─────────────────────────────
-    // All math in fp32.  cos_sin_cache is loaded as fp32 (its native storage).
-    bool const is_rope_lane = dim_base >= kNopeDim;
-    if (is_rope_lane) {
-      int64_t const pos = position_ids[tokenIdx];
-      constexpr int kHalfRope = kRopeDim / 2;  // 32
-      float const* cos_ptr = cos_sin_cache + pos * kRopeDim;
-      float const* sin_ptr = cos_ptr + kHalfRope;
-
-      int const rope_local_base = dim_base - kNopeDim;  // in [0, 64) step 16
-#pragma unroll
-      for (int p = 0; p < kElemsPerLane / 2; p++) {
-        int const pair_dim = rope_local_base + 2 * p;
-        int const half_idx = pair_dim / 2;
-        float const cos_v = VLLM_LDG(cos_ptr + half_idx);
-        float const sin_v = VLLM_LDG(sin_ptr + half_idx);
-        float const x_even = elements[2 * p];
-        float const x_odd = elements[2 * p + 1];
-        elements[2 * p] = x_even * cos_v - x_odd * sin_v;
-        elements[2 * p + 1] = x_even * sin_v + x_odd * cos_v;
-      }
-    }
-
-    // ═══════════════════════════════════════════════════════════════════════
-    // Q branch: cast to bf16 and store back in place.
-    // ═══════════════════════════════════════════════════════════════════════
-    if (!isKV) {
-      uint4 out0, out1;
-      typename Converter::packed_hip_type* po0 =
-          reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
-      typename Converter::packed_hip_type* 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;
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-      cudaTriggerProgrammaticLaunchCompletion();
+    cudaTriggerProgrammaticLaunchCompletion();
 #endif
-      return;
-    }
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
+  }
+#endif
+}
+
+// ────────────────────────────────────────────────────────────────────────────
+// Kernel
+// ────────────────────────────────────────────────────────────────────────────
+//
+// Grid: 1D, gridDim.x = num_tokens_full
+// Block: blockDim.x = 256 threads (8 warps per block) Each
+// warp handles one token, iterating over each head.
+// Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q
+// KV branch (RoPE + UE8M0 quant + insert)
+//
+template <typename scalar_t_in, int kNumHeadsQPadded>
+__global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
+    scalar_t_in const* __restrict__ q_in, scalar_t_in* __restrict__ q_out,
+    scalar_t_in const* __restrict__ kv_in, uint8_t* __restrict__ k_cache,
+    int64_t const* __restrict__ slot_mapping,
+    int64_t const* __restrict__ position_ids,
+    float const* __restrict__ cos_sin_cache, float const eps,
+    int const num_tokens_full, int const num_tokens_insert,
+    int const num_heads_q, int const cache_block_size,
+    int const kv_block_stride) {
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
+  if constexpr (std::is_same_v<scalar_t_in, c10::BFloat16>) {
+    return;
+  } else {
+#endif
+    int const warpsPerBlock = blockDim.x / 32;
+    int const warpId = threadIdx.x / 32;
+    int const laneId = threadIdx.x % 32;
+
+    int const tokenIdx = blockIdx.x;
+    if (tokenIdx >= num_tokens_full) return;
 
-    // ═══════════════════════════════════════════════════════════════════════
-    // KV branch.
-    // ═══════════════════════════════════════════════════════════════════════
-    int64_t const slot_id = slot_mapping[tokenIdx];
-    if (slot_id < 0) {
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-      cudaTriggerProgrammaticLaunchCompletion();
+    cudaGridDependencySynchronize();
 #endif
-      return;
-    }
 
-    int64_t const block_idx = slot_id / cache_block_size;
-    int64_t const pos_in_block = slot_id % cache_block_size;
-    uint8_t* block_base =
-        k_cache + block_idx * static_cast<int64_t>(kv_block_stride);
-    uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes;
-    uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim;
-    uint8_t* token_scale_ptr =
-        block_base + static_cast<int64_t>(cache_block_size) * kTokenDataBytes +
-        pos_in_block * kScaleBytesPerToken;
+    int const dim_base = laneId * kElemsPerLane;  // in [0, 512) step 16
+    // Slot enumeration: live-Q + pad-Q + (KV if this token has a slot).
+    int const slot_end = (tokenIdx >= num_tokens_insert)
+                             ? kNumHeadsQPadded
+                             : (kNumHeadsQPadded + 1);
 
-    // Round K to bf16 first, matching the unfused reference path where K is
-    // materialized as bf16 before K quantization.  absmax, clamp, and FP8
-    // quant below all run on these bf16-rounded values.
-#pragma unroll
-    for (int i = 0; i < kElemsPerLane; i++) {
-      elements[i] = Converter::convert(Converter::convert(elements[i]));
-    }
+    auto load_slot = [&](int s, uint4& va, uint4& vb) {
+      // pad-Q slots skip the load — q_in beyond num_heads_q is OOB.
+      if (s >= num_heads_q && s < kNumHeadsQPadded) return;
+      scalar_t_in const* src;
+      if (s == kNumHeadsQPadded) {
+        src = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
+      } else {
+        src = q_in +
+              (static_cast<int64_t>(tokenIdx) * num_heads_q +
+               static_cast<int64_t>(s)) *
+                  kHeadDim +
+              dim_base;
+      }
+      va = *reinterpret_cast<uint4 const*>(src);
+      vb = *reinterpret_cast<uint4 const*>(src + 8);
+    };
 
-    // Per-quant-block absmax must be computed by ALL 32 lanes (warp-collective
-    // shuffle requires full participation).  RoPE lanes contribute garbage,
-    // but their values are gated out below via `!is_rope_lane`.
-    float local_absmax = 0.0f;
-#pragma unroll
-    for (int i = 0; i < kElemsPerLane; i++) {
-      local_absmax = fmaxf(local_absmax, fabsf(elements[i]));
-    }
-    float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f);
-    float const exponent = ceilf(log2f(absmax / kFp8Max));
-    float const inv_scale = exp2f(-exponent);
+    if (warpId < slot_end) {
+      int curr_slot = warpId;
+      uint4 v0_curr, v1_curr;
+      load_slot(curr_slot, v0_curr, v1_curr);
 
-    if (!is_rope_lane) {
-      // ── NoPE lane: UE8M0 FP8 quant ───────────────────────────────────────
-      uint8_t out_bytes[kElemsPerLane];
-#pragma unroll
-      for (int i = 0; i < kElemsPerLane; i++) {
-        float scaled = elements[i] * inv_scale;
-        scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
-#ifndef USE_ROCM
-        __nv_fp8_storage_t s =
-            __nv_cvt_float_to_fp8(scaled, __NV_SATFINITE, __NV_E4M3);
-        out_bytes[i] = static_cast<uint8_t>(s);
-#else
-      out_bytes[i] = rocm_cvt_float_to_fp8_e4m3(scaled);
-#endif
-      }
-      // One 16-byte STG per lane.
-      *reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
-          *reinterpret_cast<uint4 const*>(out_bytes);
+      while (curr_slot < slot_end) {
+        int const next_slot = curr_slot + warpsPerBlock;
+        bool const has_next = (next_slot < slot_end);
+
+        // Prefetch src for the next slot
+        uint4 v0_next, v1_next;
+        if (has_next) {
+          load_slot(next_slot, v0_next, v1_next);
+        }
+
+        processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded>(
+            v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId,
+            num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids,
+            cos_sin_cache, cache_block_size, kv_block_stride);
+
+        // ── Buffer rotation: hand the prefetched LDGs to the next iter.
+        v0_curr = v0_next;
+        v1_curr = v1_next;
+        curr_slot = next_slot;
+      }  // while
+    }  // if (warpId < slot_end)
 
-      // Lane (4k) of each 4-lane group writes the scale byte for block k<7.
-      if ((laneId & 3) == 0) {
-        int const q_block_idx = laneId >> 2;  // 0..6 for NoPE lanes
-        float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f);
-        token_scale_ptr[q_block_idx] = static_cast<uint8_t>(encoded);
-      }
-      // Lane 0 also writes the padding byte at index 7.
-      if (laneId == 0) {
-        token_scale_ptr[kNumQuantBlocks] = 0;  // pad
-      }
-    } else {
-      // ── RoPE lane: cast back to bf16 and store to cache bf16 tail ────────
-      uint4 out0, out1;
-      typename Converter::packed_hip_type* po0 =
-          reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
-      typename Converter::packed_hip_type* 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]));
-      }
-      int const rope_local_base = dim_base - kNopeDim;  // in [0, 64)
-      scalar_t_in* bf16_dst =
-          reinterpret_cast<scalar_t_in*>(token_bf16_ptr) + rope_local_base;
-      *reinterpret_cast<uint4*>(bf16_dst) = out0;
-      *reinterpret_cast<uint4*>(bf16_dst + 8) = out1;
-    }
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
     cudaTriggerProgrammaticLaunchCompletion();
 #endif
@@ -390,10 +540,10 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
 // ────────────────────────────────────────────────────────────────────────────
 // Launch wrapper
 // ────────────────────────────────────────────────────────────────────────────
-template <typename scalar_t_in>
-void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
-    scalar_t_in* q_inout, scalar_t_in const* kv_in, uint8_t* k_cache,
-    int64_t const* slot_mapping, int64_t const* position_ids,
+template <typename scalar_t_in, int kNumHeadsQPadded>
+static void launchFusedDeepseekV4Templated(
+    scalar_t_in const* q_in, scalar_t_in* q_out, 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 eps, int const num_tokens_full,
     int const num_tokens_insert, int const num_heads_q,
     int const cache_block_size, int const kv_block_stride,
@@ -401,7 +551,7 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
   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);
+      static_cast<int64_t>(num_tokens_full) * (kNumHeadsQPadded + 1);
   int const grid =
       static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
 
@@ -430,37 +580,87 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
   config.attrs = attrs;
   config.numAttrs = (sm_version >= 90) ? 1 : 0;
 
-  cudaLaunchKernelEx(
-      &config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>,
-      q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
-      num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
-      kv_block_stride);
+  if (num_tokens_full < NUM_TOKEN_CUTOFF) {
+    cudaLaunchKernelEx(
+        &config,
+        fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in,
+                                                        kNumHeadsQPadded>,
+        q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
+        eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
+        kv_block_stride);
+  } else {
+    config.gridDim = dim3(num_tokens_full);
+    cudaLaunchKernelEx(
+        &config,
+        fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid<
+            scalar_t_in, kNumHeadsQPadded>,
+        q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
+        eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
+        kv_block_stride);
+  }
 #else
   // ROCm: use standard kernel launch syntax (no PDL/stream serialization)
   // clang-format off
-  fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>
+  fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in, kNumHeadsQPadded>
       <<<grid, kBlockSize, 0, stream>>>(
-          q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
-          eps, num_tokens_full, num_tokens_insert, num_heads_q,
+          q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,
+          cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q,
           cache_block_size, kv_block_stride);
 #endif
 }
 
+// Runtime dispatch into one of the precompiled `kNumHeadsQPadded`
+// instantiations.  Supported padded head counts: 8, 16, 32, 64, 128.
+template <typename scalar_t_in>
+void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
+    scalar_t_in const* q_in, scalar_t_in* q_out, 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 eps,
+    int const num_tokens_full, int const num_tokens_insert,
+    int const num_heads_q, int const num_heads_q_padded,
+    int const cache_block_size, int const kv_block_stride,
+    cudaStream_t stream) {
+#define DISPATCH(N)                                                         \
+  case N:                                                                   \
+    launchFusedDeepseekV4Templated<scalar_t_in, N>(                         \
+        q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,            \
+        cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, \
+        cache_block_size, kv_block_stride, stream);                         \
+    return;
+
+  switch (num_heads_q_padded) {
+    DISPATCH(8)
+    DISPATCH(16)
+    DISPATCH(32)
+    DISPATCH(64)
+    DISPATCH(128)
+    default:
+      TORCH_CHECK(false,
+                  "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert: "
+                  "unsupported num_heads_q_padded=",
+                  num_heads_q_padded,
+                  " (compiled instantiations: 8, 16, 32, 64, 128).");
+  }
+#undef DISPATCH
+}
+
 }  // namespace deepseek_v4_fused_ops
 }  // namespace vllm
 
 // ────────────────────────────────────────────────────────────────────────────
 // Torch op wrapper
 // ────────────────────────────────────────────────────────────────────────────
-void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::Tensor& q,                    // [N, H, 512] bf16, in place
+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) {
-  TORCH_CHECK(q.is_cuda() && q.is_contiguous(), "q must be contiguous CUDA");
+  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,
@@ -468,9 +668,12 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
   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.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
+  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.dtype() == kv.dtype(), "q and kv dtype must match");
+  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]");
@@ -480,32 +683,41 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
   // 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.size(0));
+  int const num_tokens_full = static_cast<int>(q_in.size(0));
   int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
   TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
                   static_cast<int>(position_ids.size(0)) == num_tokens_full,
               "q/kv/position_ids row counts must match");
   TORCH_CHECK(num_tokens_insert <= num_tokens_full,
               "slot_mapping must not exceed q row count");
-  int const num_heads_q = static_cast<int>(q.size(1));
+  int const num_heads_q = static_cast<int>(q_in.size(1));
+  int const num_heads_q_padded = static_cast<int>(q_head_padded);
   int const cache_block_size_i = static_cast<int>(cache_block_size);
   int const kv_block_stride = static_cast<int>(k_cache.stride(0));
 
-  at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+  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.
+  torch::Tensor q_out = torch::empty(
+      {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.options());
+
   VLLM_DISPATCH_HALF_TYPES(
-      q.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
+      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*>(q.data_ptr()),
+                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,
-                cache_block_size_i, kv_block_stride, stream);
+                num_heads_q_padded, cache_block_size_i, kv_block_stride,
+                stream);
       });
+  return q_out;
 }

csrc/libtorch_stable/activation_kernels.cu

@@ -1,20 +1,29 @@
-#include <ATen/cuda/CUDAContext.h>
-#include <torch/all.h>
-#include <c10/cuda/CUDAGuard.h>
+#include <cuda.h>
+#include <torch/csrc/stable/tensor.h>
 
 #include <cmath>
 
-#include "cuda_compat.h"
+#include "../cuda_compat.h"
 #include "cuda_vec_utils.cuh"
 #include "dispatch_utils.h"
+#include "torch_utils.h"
 
 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);                         \
@@ -210,83 +239,92 @@ packed_gelu_tanh_kernel(const packed_t& val) {
     return;                                                                    \
   }                                                                            \
   dim3 grid(num_tokens);                                                       \
-  int cc_major = at::cuda::getCurrentDeviceProperties()->major;                \
+  int cc_major = get_device_prop()->major;                                     \
   int support_vec =                                                            \
       (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128)            \
           ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE                           \
           : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE;                         \
-  int vec_size = support_vec / at::elementSize(dtype);                         \
+  int vec_size = support_vec / input.element_size();                           \
   const bool use_vec = (d % vec_size == 0);                                    \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
+  const torch::stable::accelerator::DeviceGuard device_guard(                  \
+      input.get_device_index());                                               \
+  const cudaStream_t stream = get_current_cuda_stream();                       \
   if (use_vec) {                                                               \
     dim3 block(std::min(d / vec_size, 1024));                                  \
     if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) {         \
-      VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {          \
+      VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {   \
         vllm::act_and_mul_kernel<                                              \
             scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
             KERNEL<scalar_t>,                                                  \
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
             ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>(     \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
+            out.mutable_data_ptr<scalar_t>(),                                  \
+            input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA);          \
       });                                                                      \
     } else {                                                                   \
-      VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {          \
+      VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {   \
         vllm::act_and_mul_kernel<                                              \
             scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
             KERNEL<scalar_t>,                                                  \
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
             ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>(    \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
+            out.mutable_data_ptr<scalar_t>(),                                  \
+            input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA);          \
       });                                                                      \
     }                                                                          \
   } else {                                                                     \
     dim3 block(std::min(d, 1024));                                             \
-    VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {            \
+    VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {     \
       vllm::act_and_mul_kernel<                                                \
           scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,        \
           KERNEL<scalar_t>,                                                    \
           PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
           ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>(            \
-          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);     \
+          out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),  \
+          d, LIMIT, ALPHA, BETA);                                              \
     });                                                                        \
   }
 
-void silu_and_mul(torch::Tensor& out,    // [..., d]
-                  torch::Tensor& input)  // [..., 2 * d]
+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::Tensor& out,    // [..., d]
-                        torch::Tensor& input,  // [..., 2 * d]
-                        double limit) {
+void silu_and_mul_clamp(torch::stable::Tensor& out,    // [..., d]
+                        torch::stable::Tensor& input,  // [..., 2 * d]
+                        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::Tensor& out,    // [..., d]
-                  torch::Tensor& input)  // [..., 2 * d]
+void mul_and_silu(torch::stable::Tensor& out,    // [..., d]
+                  torch::stable::Tensor& input)  // [..., 2 * 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::Tensor& out,    // [..., d]
-                  torch::Tensor& input)  // [..., 2 * d]
+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::Tensor& out,    // [..., d]
-                       torch::Tensor& input)  // [..., 2 * d]
+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 {
@@ -434,19 +472,20 @@ __global__ void swigluoai_and_mul_kernel(
     return;                                                                    \
   }                                                                            \
   dim3 grid(num_tokens);                                                       \
-  int cc_major = at::cuda::getCurrentDeviceProperties()->major;                \
+  int cc_major = get_device_prop()->major;                                     \
   int support_vec =                                                            \
       (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128)            \
           ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE                           \
           : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE;                         \
-  int vec_size = support_vec / at::elementSize(dtype);                         \
+  int vec_size = support_vec / input.element_size();                           \
   const bool use_vec = (d % vec_size == 0);                                    \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
+  const torch::stable::accelerator::DeviceGuard device_guard(                  \
+      input.get_device_index());                                               \
+  const cudaStream_t stream = get_current_cuda_stream();                       \
   if (use_vec) {                                                               \
     dim3 block(std::min(d / vec_size, 1024));                                  \
     if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) {         \
-      VLLM_DISPATCH_FLOATING_TYPES(                                            \
+      VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                     \
           dtype, "act_and_mul_kernel_with_param", [&] {                        \
             vllm::act_and_mul_kernel_with_param<                               \
                 scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,  \
@@ -454,11 +493,11 @@ __global__ void swigluoai_and_mul_kernel(
                 PACKED_KERNEL<                                                 \
                     typename vllm::PackedTypeConverter<scalar_t>::Type>,       \
                 true, true><<<grid, block, 0, stream>>>(                       \
-                out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d,       \
-                PARAM);                                                        \
+                out.mutable_data_ptr<scalar_t>(),                              \
+                input.const_data_ptr<scalar_t>(), d, PARAM);                   \
           });                                                                  \
     } else {                                                                   \
-      VLLM_DISPATCH_FLOATING_TYPES(                                            \
+      VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                     \
           dtype, "act_and_mul_kernel_with_param", [&] {                        \
             vllm::act_and_mul_kernel_with_param<                               \
                 scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,  \
@@ -466,45 +505,49 @@ __global__ void swigluoai_and_mul_kernel(
                 PACKED_KERNEL<                                                 \
                     typename vllm::PackedTypeConverter<scalar_t>::Type>,       \
                 true, false><<<grid, block, 0, stream>>>(                      \
-                out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d,       \
-                PARAM);                                                        \
+                out.mutable_data_ptr<scalar_t>(),                              \
+                input.const_data_ptr<scalar_t>(), d, PARAM);                   \
           });                                                                  \
     }                                                                          \
   } else {                                                                     \
     dim3 block(std::min(d, 1024));                                             \
-    VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel_with_param", [&] { \
-      vllm::act_and_mul_kernel_with_param<                                     \
-          scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,        \
-          KERNEL<scalar_t>,                                                    \
-          PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
-          false><<<grid, block, 0, stream>>>(                                  \
-          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, PARAM);     \
-    });                                                                        \
+    VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                       \
+        dtype, "act_and_mul_kernel_with_param", [&] {                          \
+          vllm::act_and_mul_kernel_with_param<                                 \
+              scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,    \
+              KERNEL<scalar_t>,                                                \
+              PACKED_KERNEL<                                                   \
+                  typename vllm::PackedTypeConverter<scalar_t>::Type>,         \
+              false><<<grid, block, 0, stream>>>(                              \
+              out.mutable_data_ptr<scalar_t>(),                                \
+              input.const_data_ptr<scalar_t>(), d, PARAM);                     \
+        });                                                                    \
   }
 
-#define LAUNCH_SIGLUOAI_AND_MUL(KERNEL, ALPHA, LIMIT)                          \
-  int d = input.size(-1) / 2;                                                  \
-  int64_t num_tokens = input.numel() / input.size(-1);                         \
-  dim3 grid(num_tokens);                                                       \
-  dim3 block(std::min(d, 1024));                                               \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
-  VLLM_DISPATCH_FLOATING_TYPES(                                                \
-      input.scalar_type(), "clamp_swiglu_kernel_with_params", [&] {            \
-        vllm::swigluoai_and_mul_kernel<scalar_t, KERNEL<scalar_t>>             \
-            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),             \
-                                         input.data_ptr<scalar_t>(), d, ALPHA, \
-                                         LIMIT);                               \
+#define LAUNCH_SIGLUOAI_AND_MUL(KERNEL, ALPHA, LIMIT)                         \
+  int d = input.size(-1) / 2;                                                 \
+  int64_t num_tokens = input.numel() / input.size(-1);                        \
+  dim3 grid(num_tokens);                                                      \
+  dim3 block(std::min(d, 1024));                                              \
+  const torch::stable::accelerator::DeviceGuard device_guard(                 \
+      input.get_device_index());                                              \
+  const cudaStream_t stream = get_current_cuda_stream();                      \
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                        \
+      input.scalar_type(), "clamp_swiglu_kernel_with_params", [&] {           \
+        vllm::swigluoai_and_mul_kernel<scalar_t, KERNEL<scalar_t>>            \
+            <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(),    \
+                                         input.const_data_ptr<scalar_t>(), d, \
+                                         ALPHA, LIMIT);                       \
       });
 
-void fatrelu_and_mul(torch::Tensor& out,    // [..., d],
-                     torch::Tensor& input,  // [..., 2 * d]
+void fatrelu_and_mul(torch::stable::Tensor& out,    // [..., d],
+                     torch::stable::Tensor& input,  // [..., 2 * d]
                      double threshold) {
   LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(
       vllm::fatrelu_kernel, vllm::packed_fatrelu_kernel, threshold);
 }
-void swigluoai_and_mul(torch::Tensor& out,    // [..., d]
-                       torch::Tensor& input,  // [..., 2 * d]
+void swigluoai_and_mul(torch::stable::Tensor& out,    // [..., d]
+                       torch::stable::Tensor& input,  // [..., 2 * d]
                        double alpha, double limit) {
   LAUNCH_SIGLUOAI_AND_MUL(vllm::swigluoai_and_mul, alpha, limit);
 }
@@ -559,45 +602,46 @@ __global__ void activation_kernel(
 }  // namespace vllm
 
 // Launch element-wise activation kernel.
-#define LAUNCH_ACTIVATION_KERNEL(KERNEL)                                 \
-  auto dtype = input.scalar_type();                                      \
-  int d = input.size(-1);                                                \
-  int64_t num_tokens = input.numel() / input.size(-1);                   \
-  if (num_tokens == 0) {                                                 \
-    return;                                                              \
-  }                                                                      \
-  dim3 grid(num_tokens);                                                 \
-  int cc_major = at::cuda::getCurrentDeviceProperties()->major;          \
-  int support_vec =                                                      \
-      (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128)      \
-          ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE                     \
-          : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE;                   \
-  int vec_size = support_vec / at::elementSize(dtype);                   \
-  const bool use_vec = (d % vec_size == 0);                              \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));      \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();          \
-  if (use_vec) {                                                         \
-    dim3 block(std::min(d / vec_size, 1024));                            \
-    if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) {   \
-      VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {     \
-        vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, true>  \
-            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \
-                                         input.data_ptr<scalar_t>(), d); \
-      });                                                                \
-    } else {                                                             \
-      VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {     \
-        vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, false> \
-            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \
-                                         input.data_ptr<scalar_t>(), d); \
-      });                                                                \
-    }                                                                    \
-  } else {                                                               \
-    dim3 block(std::min(d, 1024));                                       \
-    VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {       \
-      vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, false>         \
-          <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),         \
-                                       input.data_ptr<scalar_t>(), d);   \
-    });                                                                  \
+#define LAUNCH_ACTIVATION_KERNEL(KERNEL)                                       \
+  auto dtype = input.scalar_type();                                            \
+  int d = input.size(-1);                                                      \
+  int64_t num_tokens = input.numel() / input.size(-1);                         \
+  if (num_tokens == 0) {                                                       \
+    return;                                                                    \
+  }                                                                            \
+  dim3 grid(num_tokens);                                                       \
+  int cc_major = get_device_prop()->major;                                     \
+  int support_vec =                                                            \
+      (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128)            \
+          ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE                           \
+          : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE;                         \
+  int vec_size = support_vec / input.element_size();                           \
+  const bool use_vec = (d % vec_size == 0);                                    \
+  const torch::stable::accelerator::DeviceGuard device_guard(                  \
+      input.get_device_index());                                               \
+  const cudaStream_t stream = get_current_cuda_stream();                       \
+  if (use_vec) {                                                               \
+    dim3 block(std::min(d / vec_size, 1024));                                  \
+    if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) {         \
+      VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {    \
+        vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, true>        \
+            <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(),     \
+                                         input.const_data_ptr<scalar_t>(), d); \
+      });                                                                      \
+    } else {                                                                   \
+      VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {    \
+        vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, false>       \
+            <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(),     \
+                                         input.const_data_ptr<scalar_t>(), d); \
+      });                                                                      \
+    }                                                                          \
+  } else {                                                                     \
+    dim3 block(std::min(d, 1024));                                             \
+    VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {      \
+      vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, false>               \
+          <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(),       \
+                                       input.const_data_ptr<scalar_t>(), d);   \
+    });                                                                        \
   }
 
 namespace vllm {
@@ -625,20 +669,20 @@ __device__ __forceinline__ T gelu_quick_kernel(const T& x) {
 
 }  // namespace vllm
 
-void gelu_new(torch::Tensor& out,    // [..., d]
-              torch::Tensor& input)  // [..., d]
+void gelu_new(torch::stable::Tensor& out,    // [..., d]
+              torch::stable::Tensor& input)  // [..., d]
 {
   LAUNCH_ACTIVATION_KERNEL(vllm::gelu_new_kernel);
 }
 
-void gelu_fast(torch::Tensor& out,    // [..., d]
-               torch::Tensor& input)  // [..., d]
+void gelu_fast(torch::stable::Tensor& out,    // [..., d]
+               torch::stable::Tensor& input)  // [..., d]
 {
   LAUNCH_ACTIVATION_KERNEL(vllm::gelu_fast_kernel);
 }
 
-void gelu_quick(torch::Tensor& out,    // [..., d]
-                torch::Tensor& input)  // [..., d]
+void gelu_quick(torch::stable::Tensor& out,    // [..., d]
+                torch::stable::Tensor& input)  // [..., d]
 {
   LAUNCH_ACTIVATION_KERNEL(vllm::gelu_quick_kernel);
 }

csrc/libtorch_stable/attention/attention_kernels.cuh

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

csrc/libtorch_stable/attention/attention_utils.cuh

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

csrc/libtorch_stable/attention/merge_attn_states.cu

@@ -1,14 +1,14 @@
 #include <optional>
-#include <torch/all.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
 #include <algorithm>
 #include <limits>
 
-#include "attention_dtypes.h"
-#include "attention_utils.cuh"
-#include "../quantization/w8a8/fp8/common.cuh"
+#include "../torch_utils.h"
 #include "../dispatch_utils.h"
+#include <torch/headeronly/core/ScalarType.h>
+
+#include "../../attention/attention_dtypes.h"
+#include "attention_utils.cuh"
+#include "../../quantization/w8a8/fp8/common.cuh"
 
 namespace vllm {
 
@@ -196,17 +196,17 @@ __global__ void merge_attn_states_kernel(
 // The following macro is used to dispatch the conversion function based on
 // the output data type. The FN is a macro that calls a function with
 // template<typename scalar_t>.
-#define DISPATCH_BY_SCALAR_DTYPE(scalar_dtype, fn)                      \
-  {                                                                     \
-    if (scalar_dtype == at::ScalarType::Float) {                        \
-      fn(float);                                                        \
-    } else if (scalar_dtype == at::ScalarType::Half) {                  \
-      fn(uint16_t);                                                     \
-    } else if (scalar_dtype == at::ScalarType::BFloat16) {              \
-      fn(__nv_bfloat16);                                                \
-    } else {                                                            \
-      TORCH_CHECK(false, "Unsupported data type of O: ", scalar_dtype); \
-    }                                                                   \
+#define DISPATCH_BY_SCALAR_DTYPE(scalar_dtype, fn)                          \
+  {                                                                         \
+    if (scalar_dtype == torch::headeronly::ScalarType::Float) {             \
+      fn(float);                                                            \
+    } else if (scalar_dtype == torch::headeronly::ScalarType::Half) {       \
+      fn(uint16_t);                                                         \
+    } else if (scalar_dtype == torch::headeronly::ScalarType::BFloat16) {   \
+      fn(__nv_bfloat16);                                                    \
+    } else {                                                                \
+      STD_TORCH_CHECK(false, "Unsupported data type of O: ", scalar_dtype); \
+    }                                                                       \
   }
 
 #define LAUNCH_MERGE_ATTN_STATES(scalar_t, output_t, NUM_THREADS,           \
@@ -245,11 +245,14 @@ __global__ void merge_attn_states_kernel(
  */
 template <typename scalar_t>
 void merge_attn_states_launcher(
-    torch::Tensor& output, std::optional<torch::Tensor> output_lse,
-    const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse,
-    const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,
+    torch::stable::Tensor& output,
+    std::optional<torch::stable::Tensor> output_lse,
+    const torch::stable::Tensor& prefix_output,
+    const torch::stable::Tensor& prefix_lse,
+    const torch::stable::Tensor& suffix_output,
+    const torch::stable::Tensor& suffix_lse,
     const std::optional<int64_t> prefill_tokens_with_context,
-    const std::optional<torch::Tensor>& output_scale) {
+    const std::optional<torch::stable::Tensor>& output_scale) {
   constexpr uint NUM_THREADS = 128;
   const uint num_tokens = output.size(0);
   const uint num_heads = output.size(1);
@@ -258,23 +261,23 @@ void merge_attn_states_launcher(
   const uint output_head_stride = output.stride(1);
   // Thread mapping is based on input BF16 pack_size
   const uint pack_size = 16 / sizeof(scalar_t);
-  TORCH_CHECK(head_size % pack_size == 0,
-              "headsize must be multiple of pack_size:", pack_size);
+  STD_TORCH_CHECK(head_size % pack_size == 0,
+                  "headsize must be multiple of pack_size:", pack_size);
 
   const uint prefix_num_tokens =
       prefill_tokens_with_context.has_value()
           ? static_cast<uint>(prefill_tokens_with_context.value())
           : num_tokens;
-  TORCH_CHECK(prefix_num_tokens <= num_tokens,
-              "prefix_num_tokens must be <= num_tokens");
+  STD_TORCH_CHECK(prefix_num_tokens <= num_tokens,
+                  "prefix_num_tokens must be <= num_tokens");
 
   float* output_lse_ptr = nullptr;
   if (output_lse.has_value()) {
-    output_lse_ptr = output_lse.value().data_ptr<float>();
+    output_lse_ptr = output_lse.value().mutable_data_ptr<float>();
   }
   float* output_scale_ptr = nullptr;
   if (output_scale.has_value()) {
-    output_scale_ptr = output_scale.value().data_ptr<float>();
+    output_scale_ptr = output_scale.value().mutable_data_ptr<float>();
   }
   // Process one pack elements per thread. for float, the
   // pack_size is 4 for half/bf16, the pack_size is 8.
@@ -284,14 +287,15 @@ void merge_attn_states_launcher(
   dim3 block(NUM_THREADS);
   dim3 grid((total_threads + NUM_THREADS - 1) / NUM_THREADS);
 
-  const c10::cuda::OptionalCUDAGuard device_guard(prefix_output.device());
-  auto stream = at::cuda::getCurrentCUDAStream();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      prefix_output.get_device_index());
+  auto stream = get_current_cuda_stream();
 
   if (output_scale.has_value()) {
     // FP8 output path - dispatch on output FP8 type
-    VLLM_DISPATCH_FP8_TYPES(output.scalar_type(), "merge_attn_states_fp8", [&] {
-      LAUNCH_MERGE_ATTN_STATES(scalar_t, fp8_t, NUM_THREADS, true);
-    });
+    VLLM_STABLE_DISPATCH_FP8_TYPES(
+        output.scalar_type(), "merge_attn_states_fp8",
+        [&] { LAUNCH_MERGE_ATTN_STATES(scalar_t, fp8_t, NUM_THREADS, true); });
   } else {
     // Original BF16/FP16/FP32 output path
     LAUNCH_MERGE_ATTN_STATES(scalar_t, scalar_t, NUM_THREADS, false);
@@ -305,26 +309,29 @@ void merge_attn_states_launcher(
         suffix_lse, prefill_tokens_with_context, output_scale);       \
   }
 
-void merge_attn_states(torch::Tensor& output,
-                       std::optional<torch::Tensor> output_lse,
-                       const torch::Tensor& prefix_output,
-                       const torch::Tensor& prefix_lse,
-                       const torch::Tensor& suffix_output,
-                       const torch::Tensor& suffix_lse,
-                       std::optional<int64_t> prefill_tokens_with_context,
-                       const std::optional<torch::Tensor>& output_scale) {
+void merge_attn_states(
+    torch::stable::Tensor& output,
+    std::optional<torch::stable::Tensor> output_lse,
+    const torch::stable::Tensor& prefix_output,
+    const torch::stable::Tensor& prefix_lse,
+    const torch::stable::Tensor& suffix_output,
+    const torch::stable::Tensor& suffix_lse,
+    const std::optional<int64_t> prefill_tokens_with_context,
+    const std::optional<torch::stable::Tensor>& output_scale) {
   if (output_scale.has_value()) {
-    TORCH_CHECK(output.scalar_type() == at::ScalarType::Float8_e4m3fn ||
-                    output.scalar_type() == at::ScalarType::Float8_e4m3fnuz,
-                "output must be FP8 when output_scale is provided, got: ",
-                output.scalar_type());
+    STD_TORCH_CHECK(
+        output.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
+            output.scalar_type() ==
+                torch::headeronly::ScalarType::Float8_e4m3fnuz,
+        "output must be FP8 when output_scale is provided, got: ",
+        output.scalar_type());
   } else {
-    TORCH_CHECK(output.scalar_type() == prefix_output.scalar_type(),
-                "output dtype (", output.scalar_type(),
-                ") must match prefix_output dtype (",
-                prefix_output.scalar_type(), ") when output_scale is not set");
+    STD_TORCH_CHECK(
+        output.scalar_type() == prefix_output.scalar_type(), "output dtype (",
+        output.scalar_type(), ") must match prefix_output dtype (",
+        prefix_output.scalar_type(), ") when output_scale is not set");
   }
   // Always dispatch on prefix_output (input) dtype
-  DISPATCH_BY_SCALAR_DTYPE(prefix_output.dtype(),
+  DISPATCH_BY_SCALAR_DTYPE(prefix_output.scalar_type(),
                            CALL_MERGE_ATTN_STATES_LAUNCHER);
 }

csrc/libtorch_stable/attention/paged_attention_v1.cu

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

csrc/libtorch_stable/attention/paged_attention_v2.cu

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

csrc/libtorch_stable/cache_kernels_fused.cu

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

csrc/libtorch_stable/cuda_vec_utils.cuh

@@ -9,6 +9,8 @@
 
 #ifdef USE_ROCM
   #include <hip/hip_runtime.h>
+  #include <hip/hip_bf16.h>
+  #include <hip/hip_fp16.h>
 #else
   #include <cuda_bf16.h>
   #include <cuda_fp16.h>

csrc/libtorch_stable/dispatch_utils.h

@@ -58,6 +58,35 @@
   THO_DISPATCH_SWITCH(TYPE, NAME,                        \
                       VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(__VA_ARGS__))
 
+// Quant type dispatch (FP8 + INT8)
+#ifdef USE_ROCM
+  #define VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(...)                  \
+    THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn,   \
+                      __VA_ARGS__)                                    \
+    THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fnuz, \
+                      __VA_ARGS__)                                    \
+    THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__)
+#else
+  #define VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(...)                \
+    THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, \
+                      __VA_ARGS__)                                  \
+    THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__)
+#endif
+
+#define VLLM_STABLE_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \
+  THO_DISPATCH_SWITCH(TYPE, NAME,                         \
+                      VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__))
+
+// Group size dispatch (pure C++ if/else, no ATen dependency)
+#define VLLM_STABLE_DISPATCH_GROUP_SIZE(group_size, const_group_size, ...) \
+  if (group_size == 128) {                                                 \
+    constexpr int const_group_size = 128;                                  \
+    __VA_ARGS__();                                                         \
+  } else if (group_size == 64) {                                           \
+    constexpr int const_group_size = 64;                                   \
+    __VA_ARGS__();                                                         \
+  }
+
 // Boolean dispatch
 #define VLLM_STABLE_DISPATCH_BOOL(expr, const_expr, ...) \
   if (expr) {                                            \
@@ -67,3 +96,56 @@
     constexpr bool const_expr = false;                   \
     __VA_ARGS__();                                       \
   }
+
+// Vec size dispatch (pure C++ switch, no ATen dependency)
+#define VLLM_STABLE_DISPATCH_VEC_SIZE(VEC_SIZE, ...) \
+  switch (VEC_SIZE) {                                \
+    case 16: {                                       \
+      constexpr int vec_size = 16;                   \
+      __VA_ARGS__();                                 \
+      break;                                         \
+    }                                                \
+    case 8: {                                        \
+      constexpr int vec_size = 8;                    \
+      __VA_ARGS__();                                 \
+      break;                                         \
+    }                                                \
+    case 4: {                                        \
+      constexpr int vec_size = 4;                    \
+      __VA_ARGS__();                                 \
+      break;                                         \
+    }                                                \
+    case 2: {                                        \
+      constexpr int vec_size = 2;                    \
+      __VA_ARGS__();                                 \
+      break;                                         \
+    }                                                \
+    default: {                                       \
+      constexpr int vec_size = 1;                    \
+      __VA_ARGS__();                                 \
+      break;                                         \
+    }                                                \
+  }
+
+// Tensor rank dispatch (2D, 3D, 4D)
+#define VLLM_STABLE_DISPATCH_RANK234(NUM_DIMS, ...)                          \
+  switch (NUM_DIMS) {                                                        \
+    case 2: {                                                                \
+      constexpr int tensor_rank = 2;                                         \
+      __VA_ARGS__();                                                         \
+      break;                                                                 \
+    }                                                                        \
+    case 3: {                                                                \
+      constexpr int tensor_rank = 3;                                         \
+      __VA_ARGS__();                                                         \
+      break;                                                                 \
+    }                                                                        \
+    case 4: {                                                                \
+      constexpr int tensor_rank = 4;                                         \
+      __VA_ARGS__();                                                         \
+      break;                                                                 \
+    }                                                                        \
+    default:                                                                 \
+      STD_TORCH_CHECK(                                                       \
+          false, "Expects rank 2, 3 or 4 tensors but got unsupported rank"); \
+  }

csrc/libtorch_stable/fp32_router_gemm.cu

@@ -0,0 +1,223 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Router GEMM: activation(T) x weight(fp32) -> fp32, H=3072, E=256, M<=32.
+// Supports bf16 or fp32 activation; weight is always fp32.
+// Adapted from dsv3_router_gemm_float_out.cu.
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+// ---------------------------------------------------------------------------
+// Load helpers
+// ---------------------------------------------------------------------------
+
+// Load VPT fp32 values from the weight matrix (always fp32).
+//   VPT=4 when activation is fp32 (one float4 load)
+//   VPT=8 when activation is bf16 (two float4 loads)
+template <int VPT>
+__device__ __forceinline__ void load_weight(float const* ptr, float* dst);
+
+template <>
+__device__ __forceinline__ void load_weight<4>(float const* ptr, float* dst) {
+  float4 v = *reinterpret_cast<float4 const*>(ptr);
+  dst[0] = v.x;
+  dst[1] = v.y;
+  dst[2] = v.z;
+  dst[3] = v.w;
+}
+
+template <>
+__device__ __forceinline__ void load_weight<8>(float const* ptr, float* dst) {
+  float4 v0 = *reinterpret_cast<float4 const*>(ptr);
+  float4 v1 = *reinterpret_cast<float4 const*>(ptr + 4);
+  dst[0] = v0.x;
+  dst[1] = v0.y;
+  dst[2] = v0.z;
+  dst[3] = v0.w;
+  dst[4] = v1.x;
+  dst[5] = v1.y;
+  dst[6] = v1.z;
+  dst[7] = v1.w;
+}
+
+// Load VPT activation values and convert to fp32.
+template <typename T, int VPT>
+__device__ __forceinline__ void load_activation(T const* ptr, float* dst);
+
+// fp32 activation: one float4 load, no conversion needed.
+template <>
+__device__ __forceinline__ void load_activation<float, 4>(float const* ptr,
+                                                          float* dst) {
+  float4 v = *reinterpret_cast<float4 const*>(ptr);
+  dst[0] = v.x;
+  dst[1] = v.y;
+  dst[2] = v.z;
+  dst[3] = v.w;
+}
+
+// bf16 activation: one uint4 load (8 × bf16) + element-wise conversion.
+template <>
+__device__ __forceinline__ void load_activation<__nv_bfloat16, 8>(
+    __nv_bfloat16 const* ptr, float* dst) {
+  uint4 v = *reinterpret_cast<uint4 const*>(ptr);
+  __nv_bfloat16 const* bf16_ptr = reinterpret_cast<__nv_bfloat16 const*>(&v);
+#pragma unroll
+  for (int i = 0; i < 8; i++) dst[i] = __bfloat162float(bf16_ptr[i]);
+}
+
+// ---------------------------------------------------------------------------
+// Kernel
+// ---------------------------------------------------------------------------
+
+// InputT : type of activation (float or __nv_bfloat16)
+// Weight is always fp32; output is always fp32.
+// VPT = 16 / sizeof(InputT):  4 for fp32, 8 for bf16
+template <typename InputT, int kBlockSize, int kNumTokens, int kNumExperts,
+          int kHiddenDim>
+__global__ __launch_bounds__(128, 1) void fp32_router_gemm_kernel(
+    float* out, InputT const* mat_a, float const* mat_b) {
+  constexpr int VPT = 16 / sizeof(InputT);
+  constexpr int k_elems_per_k_iteration = VPT * kBlockSize;
+  constexpr int k_iterations = kHiddenDim / k_elems_per_k_iteration;
+  constexpr int kWarpSize = 32;
+  constexpr int kNumWarps = kBlockSize / kWarpSize;
+
+  int const n_idx = blockIdx.x;
+  int const tid = threadIdx.x;
+  int const warpId = tid / kWarpSize;
+  int const laneId = tid % kWarpSize;
+
+  float acc[kNumTokens] = {};
+  __shared__ float sm_reduction[kNumTokens][kNumWarps];
+
+  float const* b_col = mat_b + n_idx * kHiddenDim;
+
+  int k_bases[k_iterations];
+#pragma unroll
+  for (int ki = 0; ki < k_iterations; ki++) {
+    k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT;
+  }
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  asm volatile("griddepcontrol.wait;");
+#endif
+
+  for (int ki = 0; ki < k_iterations; ki++) {
+    int const k_base = k_bases[ki];
+
+    float b_float[VPT];
+    load_weight<VPT>(b_col + k_base, b_float);
+
+#pragma unroll
+    for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
+      float a_float[VPT];
+      load_activation<InputT, VPT>(mat_a + m_idx * kHiddenDim + k_base,
+                                   a_float);
+#pragma unroll
+      for (int k = 0; k < VPT; k++) {
+        acc[m_idx] += a_float[k] * b_float[k];
+      }
+    }
+  }
+
+  // Warp-level butterfly reduction
+#pragma unroll
+  for (int m = 0; m < kNumTokens; m++) {
+    float sum = acc[m];
+    sum += __shfl_xor_sync(0xffffffff, sum, 16);
+    sum += __shfl_xor_sync(0xffffffff, sum, 8);
+    sum += __shfl_xor_sync(0xffffffff, sum, 4);
+    sum += __shfl_xor_sync(0xffffffff, sum, 2);
+    sum += __shfl_xor_sync(0xffffffff, sum, 1);
+    if (laneId == 0) sm_reduction[m][warpId] = sum;
+  }
+
+  __syncthreads();
+
+  if (tid == 0) {
+#pragma unroll
+    for (int m = 0; m < kNumTokens; m++) {
+      float final_sum = 0.0f;
+#pragma unroll
+      for (int w = 0; w < kNumWarps; w++) final_sum += sm_reduction[m][w];
+      out[m * kNumExperts + n_idx] = final_sum;
+    }
+  }
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  asm volatile("griddepcontrol.launch_dependents;");
+#endif
+}
+
+// ---------------------------------------------------------------------------
+// Launcher
+// ---------------------------------------------------------------------------
+
+template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
+void invokeFp32RouterGemm(float* output, InputT const* mat_a,
+                          float const* mat_b, cudaStream_t stream) {
+  constexpr int kBlockSize = 128;
+  cudaLaunchConfig_t config;
+  config.gridDim = kNumExperts;
+  config.blockDim = kBlockSize;
+  config.dynamicSmemBytes = 0;
+  config.stream = stream;
+  cudaLaunchAttribute attrs[1];
+  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attrs[0].val.programmaticStreamSerializationAllowed = 1;
+  config.numAttrs = 1;
+  config.attrs = attrs;
+  cudaLaunchKernelEx(&config,
+                     fp32_router_gemm_kernel<InputT, kBlockSize, kNumTokens,
+                                             kNumExperts, kHiddenDim>,
+                     output, mat_a, mat_b);
+}
+
+// ---------------------------------------------------------------------------
+// Explicit instantiations: M=1..32, E=256, H=3072, for both input types
+// ---------------------------------------------------------------------------
+
+#define INSTANTIATE(T, M)                              \
+  template void invokeFp32RouterGemm<T, M, 256, 3072>( \
+      float*, T const*, float const*, cudaStream_t);
+
+#define INSTANTIATE_ALL(T) \
+  INSTANTIATE(T, 1)        \
+  INSTANTIATE(T, 2)        \
+  INSTANTIATE(T, 3)        \
+  INSTANTIATE(T, 4)        \
+  INSTANTIATE(T, 5)        \
+  INSTANTIATE(T, 6)        \
+  INSTANTIATE(T, 7)        \
+  INSTANTIATE(T, 8)        \
+  INSTANTIATE(T, 9)        \
+  INSTANTIATE(T, 10)       \
+  INSTANTIATE(T, 11)       \
+  INSTANTIATE(T, 12)       \
+  INSTANTIATE(T, 13)       \
+  INSTANTIATE(T, 14)       \
+  INSTANTIATE(T, 15)       \
+  INSTANTIATE(T, 16)       \
+  INSTANTIATE(T, 17)       \
+  INSTANTIATE(T, 18)       \
+  INSTANTIATE(T, 19)       \
+  INSTANTIATE(T, 20)       \
+  INSTANTIATE(T, 21)       \
+  INSTANTIATE(T, 22)       \
+  INSTANTIATE(T, 23)       \
+  INSTANTIATE(T, 24)       \
+  INSTANTIATE(T, 25)       \
+  INSTANTIATE(T, 26)       \
+  INSTANTIATE(T, 27)       \
+  INSTANTIATE(T, 28)       \
+  INSTANTIATE(T, 29)       \
+  INSTANTIATE(T, 30)       \
+  INSTANTIATE(T, 31)       \
+  INSTANTIATE(T, 32)
+
+INSTANTIATE_ALL(float)
+INSTANTIATE_ALL(__nv_bfloat16)
+
+#undef INSTANTIATE_ALL
+#undef INSTANTIATE

csrc/libtorch_stable/fp32_router_gemm_entry.cu

@@ -0,0 +1,127 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+
+#include "core/registration.h"
+#include "libtorch_stable/torch_utils.h"
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#include <stdexcept>
+
+namespace {
+
+inline int getSMVersion() {
+  auto* props = get_device_prop();
+  return props->major * 10 + props->minor;
+}
+
+}  // namespace
+
+static constexpr int FP32_NUM_EXPERTS = 256;
+static constexpr int FP32_HIDDEN_DIM = 3072;
+static constexpr int FP32_MAX_TOKENS = 32;
+
+// Forward declarations — 4 template params must match fp32_router_gemm.cu
+template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
+void invokeFp32RouterGemm(float* output, InputT const* mat_a,
+                          float const* mat_b, cudaStream_t stream);
+
+// LoopUnroller templated on InputT
+template <typename InputT, int kBegin, int kEnd>
+struct Fp32LoopUnroller {
+  static void unroll(int num_tokens, float* output, InputT const* mat_a,
+                     float const* mat_b, cudaStream_t stream) {
+    if (num_tokens == kBegin) {
+      invokeFp32RouterGemm<InputT, kBegin, FP32_NUM_EXPERTS, FP32_HIDDEN_DIM>(
+          output, mat_a, mat_b, stream);
+    } else {
+      Fp32LoopUnroller<InputT, kBegin + 1, kEnd>::unroll(num_tokens, output,
+                                                         mat_a, mat_b, stream);
+    }
+  }
+};
+
+template <typename InputT, int kEnd>
+struct Fp32LoopUnroller<InputT, kEnd, kEnd> {
+  static void unroll(int num_tokens, float* output, InputT const* mat_a,
+                     float const* mat_b, cudaStream_t stream) {
+    if (num_tokens == kEnd) {
+      invokeFp32RouterGemm<InputT, kEnd, FP32_NUM_EXPERTS, FP32_HIDDEN_DIM>(
+          output, mat_a, mat_b, stream);
+    } else {
+      throw std::invalid_argument(
+          "fp32_router_gemm: num_tokens must be in [1, 32]");
+    }
+  }
+};
+
+void fp32_router_gemm(
+    torch::stable::Tensor& output,       // [num_tokens, num_experts]
+    torch::stable::Tensor const& mat_a,  // [num_tokens, hidden_dim]
+    torch::stable::Tensor const& mat_b   // [num_experts, hidden_dim]
+) {
+  STD_TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2);
+  STD_TORCH_CHECK(output.is_cuda() && mat_a.is_cuda() && mat_b.is_cuda(),
+                  "fp32_router_gemm: all tensors must be CUDA tensors");
+  STD_TORCH_CHECK(output.get_device_index() == mat_a.get_device_index() &&
+                      output.get_device_index() == mat_b.get_device_index(),
+                  "fp32_router_gemm: all tensors must be on the same device");
+  STD_TORCH_CHECK(
+      output.is_contiguous() && mat_a.is_contiguous() && mat_b.is_contiguous(),
+      "fp32_router_gemm: all tensors must be contiguous");
+
+  const int num_tokens = mat_a.size(0);
+  const int num_experts = mat_b.size(0);
+  const int hidden_dim = mat_a.size(1);
+
+  STD_TORCH_CHECK(output.size(0) == num_tokens && output.size(1) == num_experts,
+                  "fp32_router_gemm: output must have shape [num_tokens, "
+                  "num_experts]");
+  STD_TORCH_CHECK(
+      mat_a.size(1) == mat_b.size(1),
+      "fp32_router_gemm: mat_a and mat_b must have the same hidden_dim");
+  STD_TORCH_CHECK(hidden_dim == FP32_HIDDEN_DIM,
+                  "fp32_router_gemm: expected hidden_dim=3072");
+  STD_TORCH_CHECK(num_experts == FP32_NUM_EXPERTS,
+                  "fp32_router_gemm: expected num_experts=256");
+  STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS,
+                  "fp32_router_gemm: num_tokens must be in [0, 32]");
+  STD_TORCH_CHECK(
+      mat_a.scalar_type() == torch::headeronly::ScalarType::Float ||
+          mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+      "fp32_router_gemm: mat_a must be float32 or bfloat16");
+  STD_TORCH_CHECK(mat_b.scalar_type() == torch::headeronly::ScalarType::Float,
+                  "fp32_router_gemm: mat_b (weight) must be float32");
+  STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Float,
+                  "fp32_router_gemm: output must be float32");
+
+  if (num_tokens == 0) {
+    return;
+  }
+
+  STD_TORCH_CHECK(getSMVersion() >= 90, "fp32_router_gemm: requires SM90+");
+
+  auto stream = get_current_cuda_stream(mat_a.get_device_index());
+  float* out_ptr = reinterpret_cast<float*>(output.mutable_data_ptr());
+  float const* mat_b_ptr = reinterpret_cast<float const*>(mat_b.data_ptr());
+
+  if (mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
+    auto const* mat_a_ptr =
+        reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
+    Fp32LoopUnroller<__nv_bfloat16, 1, FP32_MAX_TOKENS>::unroll(
+        num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream);
+  } else {
+    auto const* mat_a_ptr = reinterpret_cast<float const*>(mat_a.data_ptr());
+    Fp32LoopUnroller<float, 1, FP32_MAX_TOKENS>::unroll(
+        num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream);
+  }
+}
+
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
+  m.impl("fp32_router_gemm", TORCH_BOX(&fp32_router_gemm));
+}

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

@@ -18,21 +18,20 @@
 #include <cuda_runtime.h>
 #include <type_traits>
 
-#include <torch/cuda.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
+#include "torch_utils.h"
 
-#include "async_util.cuh"
-#include "cuda_compat.h"
+#include "../async_util.cuh"
+#include "../cuda_compat.h"
+#include "../type_convert.cuh"
 #include "dispatch_utils.h"
-#include "type_convert.cuh"
 
-#define CHECK_TYPE(x, st)                                              \
-  TORCH_CHECK(x.scalar_type() == st, #x " dtype is ", x.scalar_type(), \
-              ", while ", st, " is expected")
-#define CHECK_TH_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_TYPE(x, st)                                                  \
+  STD_TORCH_CHECK(x.scalar_type() == st, #x " dtype is ", x.scalar_type(), \
+                  ", while ", st, " is expected")
+#define CHECK_TH_CUDA(x) \
+  STD_TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
 #define CHECK_CONTIGUOUS(x) \
-  TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+  STD_TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
 #define CHECK_INPUT(x) \
   CHECK_TH_CUDA(x);    \
   CHECK_CONTIGUOUS(x)
@@ -589,8 +588,8 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens,
       });
       break;
     default:
-      TORCH_CHECK(false,
-                  "Unsupported head dimension for fusedQKNormRope: ", head_dim);
+      STD_TORCH_CHECK(
+          false, "Unsupported head dimension for fusedQKNormRope: ", head_dim);
   }
 }
 
@@ -604,10 +603,10 @@ void launchFusedQKNormRopeNTokenHeads(
     void const* k_weight, void const* cos_sin_cache, bool const interleave,
     int64_t const* position_ids, int const token_heads_per_warp,
     cudaStream_t stream) {
-  TORCH_CHECK(token_heads_per_warp == 1 || token_heads_per_warp == 2 ||
-                  token_heads_per_warp == 4 || token_heads_per_warp == 8,
-              "token_heads_per_warp must be 1, 2, 4, or 8, got ",
-              token_heads_per_warp);
+  STD_TORCH_CHECK(token_heads_per_warp == 1 || token_heads_per_warp == 2 ||
+                      token_heads_per_warp == 4 || token_heads_per_warp == 8,
+                  "token_heads_per_warp must be 1, 2, 4, or 8, got ",
+                  token_heads_per_warp);
 
   // token_heads_per_warp == 1: delegate to the 1-head baseline kernel.
   if (token_heads_per_warp == 1) {
@@ -691,7 +690,7 @@ void launchFusedQKNormRopeNTokenHeads(
         });                                                                  \
         break;                                                               \
       default:                                                               \
-        TORCH_CHECK(false, "Unsupported head dimension: ", head_dim);        \
+        STD_TORCH_CHECK(false, "Unsupported head dimension: ", head_dim);    \
     }                                                                        \
   } while (0)
 
@@ -708,19 +707,21 @@ void launchFusedQKNormRopeNTokenHeads(
 }  // namespace tensorrt_llm::kernels
 
 void fused_qk_norm_rope(
-    torch::Tensor& qkv,       // Combined QKV tensor [num_tokens,
-                              // (num_heads_q+num_heads_k+num_heads_v)*head_dim]
-    int64_t num_heads_q,      // Number of query heads
-    int64_t num_heads_k,      // Number of key heads
-    int64_t num_heads_v,      // Number of value heads
-    int64_t head_dim,         // Dimension per head
-    double eps,               // Epsilon for RMS normalization
-    torch::Tensor& q_weight,  // RMSNorm weights for query [head_dim]
-    torch::Tensor& k_weight,  // RMSNorm weights for key [head_dim]
-    torch::Tensor& cos_sin_cache,  // Cos/sin cache [max_position, head_dim]
-    bool is_neox,                  // Whether RoPE is applied in Neox style
-    torch::Tensor& position_ids,   // Position IDs for RoPE [num_tokens]
-    int64_t forced_token_heads_per_warp  // -1 = auto-select, >0 = forced value
+    torch::stable::Tensor&
+        qkv,              // Combined QKV tensor [num_tokens,
+                          // (num_heads_q+num_heads_k+num_heads_v)*head_dim]
+    int64_t num_heads_q,  // Number of query heads
+    int64_t num_heads_k,  // Number of key heads
+    int64_t num_heads_v,  // Number of value heads
+    int64_t head_dim,     // Dimension per head
+    double eps,           // Epsilon for RMS normalization
+    torch::stable::Tensor& q_weight,  // RMSNorm weights for query [head_dim]
+    torch::stable::Tensor& k_weight,  // RMSNorm weights for key [head_dim]
+    torch::stable::Tensor& cos_sin_cache,  // Cos/sin cache [max_position,
+                                           // head_dim]
+    bool is_neox,  // Whether RoPE is applied in Neox style
+    torch::stable::Tensor& position_ids,  // Position IDs for RoPE [num_tokens]
+    int64_t forced_token_heads_per_warp   // -1 = auto-select, >0 = forced value
 ) {
   // Input validation
   CHECK_INPUT(qkv);
@@ -728,40 +729,42 @@ void fused_qk_norm_rope(
   CHECK_INPUT(q_weight);
   CHECK_INPUT(k_weight);
   CHECK_INPUT(cos_sin_cache);
-  CHECK_TYPE(position_ids, torch::kInt64);
+  CHECK_TYPE(position_ids, torch::headeronly::ScalarType::Long);
 
-  TORCH_CHECK(qkv.dim() == 2,
-              "QKV tensor must be 2D: [num_tokens, "
-              "(num_heads_q+num_heads_k+num_heads_v)*head_dim]");
-  TORCH_CHECK(position_ids.dim() == 1, "Position IDs must be 1D: [num_tokens]");
-  TORCH_CHECK(q_weight.dim() == 1, "Query weights must be 1D: [head_dim]");
-  TORCH_CHECK(k_weight.dim() == 1, "Key weights must be 1D: [head_dim]");
-  TORCH_CHECK(cos_sin_cache.dim() == 2,
-              "Cos/sin cache must be 2D: [max_position, head_dim]");
-  TORCH_CHECK(q_weight.size(0) == head_dim,
-              "Query weights size must match head dimension");
-  TORCH_CHECK(k_weight.size(0) == head_dim,
-              "Key weights size must match head dimension");
+  STD_TORCH_CHECK(qkv.dim() == 2,
+                  "QKV tensor must be 2D: [num_tokens, "
+                  "(num_heads_q+num_heads_k+num_heads_v)*head_dim]");
+  STD_TORCH_CHECK(position_ids.dim() == 1,
+                  "Position IDs must be 1D: [num_tokens]");
+  STD_TORCH_CHECK(q_weight.dim() == 1, "Query weights must be 1D: [head_dim]");
+  STD_TORCH_CHECK(k_weight.dim() == 1, "Key weights must be 1D: [head_dim]");
+  STD_TORCH_CHECK(cos_sin_cache.dim() == 2,
+                  "Cos/sin cache must be 2D: [max_position, head_dim]");
+  STD_TORCH_CHECK(q_weight.size(0) == head_dim,
+                  "Query weights size must match head dimension");
+  STD_TORCH_CHECK(k_weight.size(0) == head_dim,
+                  "Key weights size must match head dimension");
 
-  TORCH_CHECK(cos_sin_cache.size(1) % 2 == 0, "rotary_dim must be even");
-  TORCH_CHECK(cos_sin_cache.size(1) <= head_dim,
-              "rotary_dim must be less than or equal to head_dim");
+  STD_TORCH_CHECK(cos_sin_cache.size(1) % 2 == 0, "rotary_dim must be even");
+  STD_TORCH_CHECK(cos_sin_cache.size(1) <= head_dim,
+                  "rotary_dim must be less than or equal to head_dim");
 
-  TORCH_CHECK(qkv.scalar_type() == q_weight.scalar_type() &&
-                  qkv.scalar_type() == k_weight.scalar_type(),
-              "qkv, q_weight and k_weight must have the same dtype");
+  STD_TORCH_CHECK(qkv.scalar_type() == q_weight.scalar_type() &&
+                      qkv.scalar_type() == k_weight.scalar_type(),
+                  "qkv, q_weight and k_weight must have the same dtype");
 
   int64_t num_tokens = qkv.size(0);
-  TORCH_CHECK(position_ids.size(0) == num_tokens,
-              "Number of tokens in position_ids must match QKV");
+  STD_TORCH_CHECK(position_ids.size(0) == num_tokens,
+                  "Number of tokens in position_ids must match QKV");
 
   int64_t total_heads = num_heads_q + num_heads_k + num_heads_v;
-  TORCH_CHECK(
+  STD_TORCH_CHECK(
       qkv.size(1) == total_heads * head_dim,
       "QKV tensor size must match total number of heads and head dimension");
 
-  auto device_id = qkv.get_device();
-  auto stream = at::cuda::getCurrentCUDAStream(device_id);
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      qkv.get_device_index());
+  auto stream = get_current_cuda_stream(qkv.get_device_index());
 
   // Select token_heads_per_warp: forced value if >0, else auto-select.
   // Auto thresholds are calibrated on SM 9.0 (H100). On other architectures,
@@ -771,8 +774,7 @@ void fused_qk_norm_rope(
     token_heads_per_warp = static_cast<int>(forced_token_heads_per_warp);
   } else {
     token_heads_per_warp = 1;
-    auto* dev_prop = at::cuda::getDeviceProperties(device_id);
-    int sm_version = dev_prop->major * 10 + dev_prop->minor;
+    int sm_version = get_device_prop()->major * 10 + get_device_prop()->minor;
     int64_t total_qk_units = num_tokens * (num_heads_q + num_heads_k);
     if (sm_version == 90) {
       if (head_dim >= 256) {
@@ -795,21 +797,22 @@ void fused_qk_norm_rope(
     }
   }
 
-  VLLM_DISPATCH_HALF_TYPES(qkv.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
-    using qkv_scalar_t = scalar_t;
-    VLLM_DISPATCH_FLOATING_TYPES(
-        cos_sin_cache.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
-          using cache_scalar_t = scalar_t;
-          tensorrt_llm::kernels::launchFusedQKNormRopeNTokenHeads<
-              qkv_scalar_t, cache_scalar_t>(
-              qkv.data_ptr(), static_cast<int>(num_tokens),
-              static_cast<int>(num_heads_q), static_cast<int>(num_heads_k),
-              static_cast<int>(num_heads_v), static_cast<int>(head_dim),
-              static_cast<int>(cos_sin_cache.size(1)), static_cast<float>(eps),
-              q_weight.data_ptr(), k_weight.data_ptr(),
-              cos_sin_cache.data_ptr(), !is_neox,
-              reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
-              token_heads_per_warp, stream);
-        });
-  });
+  VLLM_STABLE_DISPATCH_HALF_TYPES(
+      qkv.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
+        using qkv_scalar_t = scalar_t;
+        VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+            cos_sin_cache.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
+              using cache_scalar_t = scalar_t;
+              tensorrt_llm::kernels::launchFusedQKNormRopeNTokenHeads<
+                  qkv_scalar_t, cache_scalar_t>(
+                  qkv.data_ptr(), static_cast<int>(num_tokens),
+                  static_cast<int>(num_heads_q), static_cast<int>(num_heads_k),
+                  static_cast<int>(num_heads_v), static_cast<int>(head_dim),
+                  static_cast<int>(cos_sin_cache.size(1)),
+                  static_cast<float>(eps), q_weight.data_ptr(),
+                  k_weight.data_ptr(), cos_sin_cache.data_ptr(), !is_neox,
+                  reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
+                  token_heads_per_warp, stream);
+            });
+      });
 }

csrc/libtorch_stable/layernorm_kernels.cu

@@ -1,11 +1,12 @@
-#include "type_convert.cuh"
-#include "dispatch_utils.h"
-#include "cub_helpers.h"
-#include "core/batch_invariant.hpp"
-#include "libtorch_stable/quantization/vectorization_utils.cuh"
+#include <numeric>
 
-#include <torch/cuda.h>
-#include <c10/cuda/CUDAGuard.h>
+#include "torch_utils.h"
+
+#include "../cub_helpers.h"
+#include "../core/batch_invariant.hpp"
+#include "../type_convert.cuh"
+#include "dispatch_utils.h"
+#include "quantization/vectorization_utils.cuh"
 
 namespace vllm {
 
@@ -77,8 +78,7 @@ __global__ void rms_norm_kernel(
 #pragma unroll
     for (int j = 0; j < VEC_SIZE; j++) {
       float x = static_cast<float>(src1.val[j]);
-      float w = static_cast<float>(src2.val[j]);
-      dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
+      dst.val[j] = static_cast<scalar_t>(x * s_variance) * src2.val[j];
     }
     v_out[i] = dst;
   }
@@ -142,8 +142,7 @@ fused_add_rms_norm_kernel(
 #pragma unroll
     for (int j = 0; j < width; ++j) {
       float x = Converter::convert(res.data[j]);
-      float wf = Converter::convert(w.data[j]);
-      out.data[j] = Converter::convert(x * s_variance * wf);
+      out.data[j] = Converter::convert(x * s_variance) * w.data[j];
     }
     input_v[strided_id] = out;
   }
@@ -182,23 +181,23 @@ fused_add_rms_norm_kernel(
 
   for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
     float x = (float)residual[blockIdx.x * hidden_size + idx];
-    float w = (float)weight[idx];
-    input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
+    input[blockIdx.x * input_stride + idx] =
+        (scalar_t)(x * s_variance) * weight[idx];
   }
 }
 
 }  // namespace vllm
 
-void rms_norm(torch::Tensor& out,     // [..., hidden_size]
-              torch::Tensor& input,   // [..., hidden_size]
-              torch::Tensor& weight,  // [hidden_size]
+void rms_norm(torch::stable::Tensor& out,     // [..., hidden_size]
+              torch::stable::Tensor& input,   // [..., hidden_size]
+              torch::stable::Tensor& weight,  // [hidden_size]
               double epsilon) {
-  TORCH_CHECK(out.is_contiguous());
+  STD_TORCH_CHECK(out.is_contiguous());
   if (input.stride(-1) != 1) {
-    input = input.contiguous();
+    input = torch::stable::contiguous(input);
   }
-  TORCH_CHECK(input.stride(-1) == 1);
-  TORCH_CHECK(weight.is_contiguous());
+  STD_TORCH_CHECK(input.stride(-1) == 1);
+  STD_TORCH_CHECK(weight.is_contiguous());
 
   int hidden_size = input.size(-1);
 
@@ -213,45 +212,49 @@ void rms_norm(torch::Tensor& out,     // [..., hidden_size]
   // For large num_tokens, use smaller blocks to increase SM concurrency.
   const int max_block_size = (num_tokens < 256) ? 1024 : 256;
   dim3 grid(num_tokens);
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  VLLM_DISPATCH_RANK234(num_dims, [&] {
-    VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_kernel", [&] {
-      const int calculated_vec_size =
-          std::gcd(16 / sizeof(scalar_t), hidden_size);
-      const int block_size =
-          std::min(hidden_size / calculated_vec_size, max_block_size);
-      dim3 block(block_size);
-      VLLM_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
-        vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank>
-            <<<grid, block, 0, stream>>>(
-                out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
-                input_stride_d2, input_stride_d3, input_stride_d4,
-                input_shape_d2, input_shape_d3, weight.data_ptr<scalar_t>(),
-                epsilon, num_tokens, hidden_size);
-      });
-    });
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+  VLLM_STABLE_DISPATCH_RANK234(num_dims, [&] {
+    VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+        input.scalar_type(), "rms_norm_kernel", [&] {
+          const int calculated_vec_size =
+              std::gcd(16 / sizeof(scalar_t), hidden_size);
+          const int block_size =
+              std::min(hidden_size / calculated_vec_size, max_block_size);
+          dim3 block(block_size);
+          VLLM_STABLE_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
+            vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank>
+                <<<grid, block, 0, stream>>>(
+                    out.mutable_data_ptr<scalar_t>(),
+                    input.const_data_ptr<scalar_t>(), input_stride_d2,
+                    input_stride_d3, input_stride_d4, input_shape_d2,
+                    input_shape_d3, weight.const_data_ptr<scalar_t>(), epsilon,
+                    num_tokens, hidden_size);
+          });
+        });
   });
 }
 
-#define LAUNCH_FUSED_ADD_RMS_NORM(width)                                    \
-  VLLM_DISPATCH_FLOATING_TYPES(                                             \
-      input.scalar_type(), "fused_add_rms_norm_kernel", [&] {               \
-        vllm::fused_add_rms_norm_kernel<scalar_t, width>                    \
-            <<<grid, block, 0, stream>>>(                                   \
-                input.data_ptr<scalar_t>(), input_stride,                   \
-                residual.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(), \
-                epsilon, num_tokens, hidden_size);                          \
+#define LAUNCH_FUSED_ADD_RMS_NORM(width)                                \
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                  \
+      input.scalar_type(), "fused_add_rms_norm_kernel", [&] {           \
+        vllm::fused_add_rms_norm_kernel<scalar_t, width>                \
+            <<<grid, block, 0, stream>>>(                               \
+                input.mutable_data_ptr<scalar_t>(), input_stride,       \
+                residual.mutable_data_ptr<scalar_t>(),                  \
+                weight.const_data_ptr<scalar_t>(), epsilon, num_tokens, \
+                hidden_size);                                           \
       });
 
-void fused_add_rms_norm(torch::Tensor& input,     // [..., hidden_size]
-                        torch::Tensor& residual,  // [..., hidden_size]
-                        torch::Tensor& weight,    // [hidden_size]
+void fused_add_rms_norm(torch::stable::Tensor& input,     // [..., hidden_size]
+                        torch::stable::Tensor& residual,  // [..., hidden_size]
+                        torch::stable::Tensor& weight,    // [hidden_size]
                         double epsilon) {
-  TORCH_CHECK(weight.scalar_type() == input.scalar_type());
-  TORCH_CHECK(input.scalar_type() == residual.scalar_type());
-  TORCH_CHECK(residual.is_contiguous());
-  TORCH_CHECK(weight.is_contiguous());
+  STD_TORCH_CHECK(weight.scalar_type() == input.scalar_type());
+  STD_TORCH_CHECK(input.scalar_type() == residual.scalar_type());
+  STD_TORCH_CHECK(residual.is_contiguous());
+  STD_TORCH_CHECK(weight.is_contiguous());
   int hidden_size = input.size(-1);
   int64_t input_stride = input.stride(-2);
   int num_tokens = input.numel() / hidden_size;
@@ -263,8 +266,9 @@ void fused_add_rms_norm(torch::Tensor& input,     // [..., hidden_size]
      hiding on global mem ops. */
   const int max_block_size = (num_tokens < 256) ? 1024 : 256;
   dim3 block(std::min(hidden_size, max_block_size));
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
   /*If the tensor types are FP16/BF16, try to use the optimized kernel
     with packed + vectorized ops.
     Max optimization is achieved with a width-8 vector of FP16/BF16s

csrc/libtorch_stable/layernorm_quant_kernels.cu

@@ -5,15 +5,16 @@
  * Currently, only static fp8 quantization is supported.
  */
 
-#include "type_convert.cuh"
-#include "quantization/w8a8/fp8/common.cuh"
-#include "dispatch_utils.h"
-#include "cub_helpers.h"
-#include "core/batch_invariant.hpp"
-#include "libtorch_stable/quantization/vectorization_utils.cuh"
+#include <numeric>
 
-#include <torch/cuda.h>
-#include <c10/cuda/CUDAGuard.h>
+#include "torch_utils.h"
+
+#include "../cub_helpers.h"
+#include "../core/batch_invariant.hpp"
+#include "../quantization/w8a8/fp8/common.cuh"
+#include "../type_convert.cuh"
+#include "dispatch_utils.h"
+#include "quantization/vectorization_utils.cuh"
 
 namespace vllm {
 
@@ -65,13 +66,8 @@ __global__ void rms_norm_static_fp8_quant_kernel(
 #pragma unroll
     for (int j = 0; j < VEC_SIZE; j++) {
       float x = static_cast<float>(src1.val[j]);
-      float w = static_cast<float>(src2.val[j]);
-      // Round normalized result through scalar_t to match the precision of the
-      // unfused composite (rms_norm writes scalar_t, then
-      // static_scaled_fp8_quant re-loads it as float before FP8 conversion).
-      // Without this round, the fused path is strictly more accurate and
-      // disagrees with the composite at exact E4M3 quantization tie boundaries.
-      scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
+      // Multiply in weight's native dtype to match rms_norm_kernel.
+      scalar_t out_norm = static_cast<scalar_t>(x * s_variance) * src2.val[j];
       out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
           scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
                                                 scale_inv);
@@ -141,12 +137,8 @@ fused_add_rms_norm_static_fp8_quant_kernel(
 #pragma unroll
     for (int i = 0; i < width; ++i) {
       float x = Converter::convert(res.data[i]);
-      float wf = Converter::convert(w.data[i]);
-      // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
-      // to match the unfused composite path at FP8 boundaries. We use the
-      // backend's hip_type for the intermediate since c10::Half/BFloat16 has
-      // ambiguous conversions on CUDA and no implicit conversion on ROCm.
-      HipT out_norm_h = Converter::convert(x * s_variance * wf);
+      // Multiply in weight's native dtype to match fused_add_rms_norm_kernel.
+      HipT out_norm_h = Converter::convert(x * s_variance) * w.data[i];
       out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
           Converter::convert(out_norm_h), scale_inv);
     }
@@ -191,10 +183,8 @@ fused_add_rms_norm_static_fp8_quant_kernel(
 
   for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
     float x = (float)residual[blockIdx.x * hidden_size + idx];
-    float w = (float)weight[idx];
-    // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
-    // to match the unfused composite path at FP8 boundaries.
-    scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
+    // Multiply in weight's native dtype to match fused_add_rms_norm_kernel.
+    scalar_t out_norm = static_cast<scalar_t>(x * s_variance) * weight[idx];
     out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
         static_cast<float>(out_norm), scale_inv);
   }
@@ -202,12 +192,13 @@ fused_add_rms_norm_static_fp8_quant_kernel(
 
 }  // namespace vllm
 
-void rms_norm_static_fp8_quant(torch::Tensor& out,     // [..., hidden_size]
-                               torch::Tensor& input,   // [..., hidden_size]
-                               torch::Tensor& weight,  // [hidden_size]
-                               torch::Tensor& scale,   // [1]
-                               double epsilon) {
-  TORCH_CHECK(out.is_contiguous());
+void rms_norm_static_fp8_quant(
+    torch::stable::Tensor& out,     // [..., hidden_size]
+    torch::stable::Tensor& input,   // [..., hidden_size]
+    torch::stable::Tensor& weight,  // [hidden_size]
+    torch::stable::Tensor& scale,   // [1]
+    double epsilon) {
+  STD_TORCH_CHECK(out.is_contiguous());
   int hidden_size = input.size(-1);
   int input_stride = input.stride(-2);
   int num_tokens = input.numel() / hidden_size;
@@ -215,24 +206,26 @@ void rms_norm_static_fp8_quant(torch::Tensor& out,     // [..., hidden_size]
   // For large num_tokens, use smaller blocks to increase SM concurrency.
   const int max_block_size = (num_tokens < 256) ? 1024 : 256;
   dim3 grid(num_tokens);
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  VLLM_DISPATCH_FLOATING_TYPES(
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "rms_norm_kernel_scalar_type", [&] {
-        VLLM_DISPATCH_FP8_TYPES(
+        VLLM_STABLE_DISPATCH_FP8_TYPES(
             out.scalar_type(), "rms_norm_kernel_fp8_type", [&] {
               const int calculated_vec_size =
                   std::gcd(16 / sizeof(scalar_t), hidden_size);
               const int block_size =
                   std::min(hidden_size / calculated_vec_size, max_block_size);
               dim3 block(block_size);
-              VLLM_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
+              VLLM_STABLE_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
                 vllm::rms_norm_static_fp8_quant_kernel<scalar_t, fp8_t,
                                                        vec_size>
                     <<<grid, block, 0, stream>>>(
-                        out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
-                        input_stride, weight.data_ptr<scalar_t>(),
-                        scale.data_ptr<float>(), epsilon, num_tokens,
+                        out.mutable_data_ptr<fp8_t>(),
+                        input.const_data_ptr<scalar_t>(), input_stride,
+                        weight.const_data_ptr<scalar_t>(),
+                        scale.const_data_ptr<float>(), epsilon, num_tokens,
                         hidden_size);
               });
             });
@@ -240,30 +233,32 @@ void rms_norm_static_fp8_quant(torch::Tensor& out,     // [..., hidden_size]
 }
 
 #define LAUNCH_FUSED_ADD_RMS_NORM(width)                                     \
-  VLLM_DISPATCH_FLOATING_TYPES(                                              \
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                       \
       input.scalar_type(), "fused_add_rms_norm_kernel_scalar_type", [&] {    \
-        VLLM_DISPATCH_FP8_TYPES(                                             \
+        VLLM_STABLE_DISPATCH_FP8_TYPES(                                      \
             out.scalar_type(), "fused_add_rms_norm_kernel_fp8_type", [&] {   \
               vllm::fused_add_rms_norm_static_fp8_quant_kernel<scalar_t,     \
                                                                width, fp8_t> \
                   <<<grid, block, 0, stream>>>(                              \
-                      out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),     \
-                      input_stride, residual.data_ptr<scalar_t>(),           \
-                      weight.data_ptr<scalar_t>(), scale.data_ptr<float>(),  \
-                      epsilon, num_tokens, hidden_size);                     \
+                      out.mutable_data_ptr<fp8_t>(),                         \
+                      input.mutable_data_ptr<scalar_t>(), input_stride,      \
+                      residual.mutable_data_ptr<scalar_t>(),                 \
+                      weight.const_data_ptr<scalar_t>(),                     \
+                      scale.const_data_ptr<float>(), epsilon, num_tokens,    \
+                      hidden_size);                                          \
             });                                                              \
       });
 void fused_add_rms_norm_static_fp8_quant(
-    torch::Tensor& out,       // [..., hidden_size],
-    torch::Tensor& input,     // [..., hidden_size]
-    torch::Tensor& residual,  // [..., hidden_size]
-    torch::Tensor& weight,    // [hidden_size]
-    torch::Tensor& scale,     // [1]
+    torch::stable::Tensor& out,       // [..., hidden_size],
+    torch::stable::Tensor& input,     // [..., hidden_size]
+    torch::stable::Tensor& residual,  // [..., hidden_size]
+    torch::stable::Tensor& weight,    // [hidden_size]
+    torch::stable::Tensor& scale,     // [1]
     double epsilon) {
-  TORCH_CHECK(out.is_contiguous());
-  TORCH_CHECK(residual.is_contiguous());
-  TORCH_CHECK(residual.scalar_type() == input.scalar_type());
-  TORCH_CHECK(weight.scalar_type() == input.scalar_type());
+  STD_TORCH_CHECK(out.is_contiguous());
+  STD_TORCH_CHECK(residual.is_contiguous());
+  STD_TORCH_CHECK(residual.scalar_type() == input.scalar_type());
+  STD_TORCH_CHECK(weight.scalar_type() == input.scalar_type());
   int hidden_size = input.size(-1);
   int input_stride = input.stride(-2);
   int num_tokens = input.numel() / hidden_size;
@@ -275,8 +270,9 @@ void fused_add_rms_norm_static_fp8_quant(
      hiding on global mem ops. */
   const int max_block_size = (num_tokens < 256) ? 1024 : 256;
   dim3 block(std::min(hidden_size, max_block_size));
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
   /*If the tensor types are FP16/BF16, try to use the optimized kernel
     with packed + vectorized ops.
     Max optimization is achieved with a width-8 vector of FP16/BF16s

csrc/libtorch_stable/mamba/selective_scan.h

@@ -12,6 +12,9 @@
     #include <hip/hip_bf16.h>
 #endif
 #include <cuda_fp16.h>
+
+#include <torch/headeronly/util/Half.h>
+#include <torch/headeronly/util/BFloat16.h>
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 struct SSMParamsBase {
@@ -159,8 +162,8 @@ struct Converter{
 };
 
 template<int N>
-struct Converter<at::Half, N>{
-    static inline __device__ void to_float(const at::Half (&src)[N], float (&dst)[N]) {
+struct Converter<torch::headeronly::Half, N>{
+    static inline __device__ void to_float(const torch::headeronly::Half (&src)[N], float (&dst)[N]) {
         static_assert(N % 2 == 0);
         auto &src2 = reinterpret_cast<const half2 (&)[N / 2]>(src);
         auto &dst2 = reinterpret_cast<float2 (&)[N / 2]>(dst);
@@ -171,8 +174,8 @@ struct Converter<at::Half, N>{
 
 #if __CUDA_ARCH__ >= 800
 template<int N>
-struct Converter<at::BFloat16, N>{
-    static inline __device__ void to_float(const at::BFloat16 (&src)[N], float (&dst)[N]) {
+struct Converter<torch::headeronly::BFloat16, N>{
+    static inline __device__ void to_float(const torch::headeronly::BFloat16 (&src)[N], float (&dst)[N]) {
         static_assert(N % 2 == 0);
         auto &src2 = reinterpret_cast<const nv_bfloat162 (&)[N / 2]>(src);
         auto &dst2 = reinterpret_cast<float2 (&)[N / 2]>(dst);

csrc/libtorch_stable/mamba/selective_scan_fwd.cu

@@ -1,18 +1,9 @@
 // clang-format off
 // adapted from https://github.com/state-spaces/mamba/blob/main/csrc/selective_scan/selective_scan_fwd_kernel.cuh
-#include <torch/all.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
+#include "../torch_utils.h"
+#include <torch/csrc/stable/macros.h>
 #include "selective_scan.h"
 
-#include <c10/util/BFloat16.h>
-#include <c10/util/Half.h>
-#ifdef USE_ROCM
-    #include <c10/hip/HIPException.h>  // For C10_HIP_CHECK and C10_HIP_KERNEL_LAUNCH_CHECK
-#else
-    #include <c10/cuda/CUDAException.h>  // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
-#endif
-
 #ifndef USE_ROCM
     #include <cub/block/block_load.cuh>
     #include <cub/block/block_store.cuh>
@@ -416,15 +407,15 @@ void selective_scan_fwd_launch(SSMParamsBase &params, cudaStream_t stream) {
                 auto kernel = &selective_scan_fwd_kernel<Ktraits>;
                 if (kSmemSize >= 48 * 1024) {
 #ifdef USE_ROCM
-                    C10_HIP_CHECK(hipFuncSetAttribute(
+                    STD_CUDA_CHECK(hipFuncSetAttribute(
                         reinterpret_cast<const void*>(kernel), hipFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
 #else
-                    C10_CUDA_CHECK(cudaFuncSetAttribute(
+                    STD_CUDA_CHECK(cudaFuncSetAttribute(
                         kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
 #endif
                 }
                 kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
-                C10_CUDA_KERNEL_LAUNCH_CHECK();
+                STD_CUDA_KERNEL_LAUNCH_CHECK();
             });
         });
     });
@@ -462,46 +453,46 @@ void selective_scan_fwd_cuda(SSMParamsBase &params, cudaStream_t stream) {
     #endif
 }
 
-template void selective_scan_fwd_cuda<at::BFloat16, float, at::BFloat16>(SSMParamsBase &params, cudaStream_t stream);
-template void selective_scan_fwd_cuda<at::BFloat16, float, float>(SSMParamsBase &params, cudaStream_t stream);
-template void selective_scan_fwd_cuda<at::Half, float, at::Half>(SSMParamsBase &params, cudaStream_t stream);
-template void selective_scan_fwd_cuda<at::Half, float, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<torch::headeronly::BFloat16, float, torch::headeronly::BFloat16>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<torch::headeronly::BFloat16, float, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<torch::headeronly::Half, float, torch::headeronly::Half>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<torch::headeronly::Half, float, float>(SSMParamsBase &params, cudaStream_t stream);
 template void selective_scan_fwd_cuda<float, float, float>(SSMParamsBase &params, cudaStream_t stream);
 
-#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
+#define CHECK_SHAPE(x, ...) STD_TORCH_CHECK(x.sizes().equals(torch::headeronly::IntHeaderOnlyArrayRef({__VA_ARGS__})), #x " must have shape (" #__VA_ARGS__ ")")
 
 #define DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(ITYPE, STYPE, NAME, ...)       \
-    if (ITYPE == at::ScalarType::Half) {                                            \
-        using input_t = at::Half;                                                   \
+    if (ITYPE == torch::headeronly::ScalarType::Half) {                             \
+        using input_t = torch::headeronly::Half;                                                  \
         using weight_t = float;                                                     \
-        if (STYPE == at::ScalarType::Half) {                                        \
-            using state_t = at::Half;                                               \
+        if (STYPE == torch::headeronly::ScalarType::Half) {                         \
+            using state_t = torch::headeronly::Half;                                              \
             __VA_ARGS__();                                                          \
-        } else if (STYPE == at::ScalarType::Float) {                                \
+        } else if (STYPE == torch::headeronly::ScalarType::Float) {                 \
             using state_t = float;                                                  \
             __VA_ARGS__();                                                          \
         } else {                                                                    \
-            AT_ERROR(#NAME, " not implemented for state type '", toString(STYPE), "'"); \
+            STD_TORCH_CHECK(false, #NAME " not implemented for state type '", STYPE, "'"); \
         }                                                                           \
-    } else if (ITYPE == at::ScalarType::BFloat16) {                                 \
-        using input_t = at::BFloat16;                                               \
+    } else if (ITYPE == torch::headeronly::ScalarType::BFloat16) {                  \
+        using input_t = torch::headeronly::BFloat16;                                              \
         using weight_t = float;                                                     \
-        if (STYPE == at::ScalarType::BFloat16) {                                    \
-            using state_t = at::BFloat16;                                           \
+        if (STYPE == torch::headeronly::ScalarType::BFloat16) {                     \
+            using state_t = torch::headeronly::BFloat16;                                          \
             __VA_ARGS__();                                                          \
-        } else if (STYPE == at::ScalarType::Float) {                                \
+        } else if (STYPE == torch::headeronly::ScalarType::Float) {                 \
             using state_t = float;                                                  \
             __VA_ARGS__();                                                          \
         } else {                                                                    \
-            AT_ERROR(#NAME, " not implemented for state type '", toString(STYPE), "'"); \
+            STD_TORCH_CHECK(false, #NAME " not implemented for state type '", STYPE, "'"); \
         }                                                                           \
-    } else if (ITYPE == at::ScalarType::Float)  {                                   \
+    } else if (ITYPE == torch::headeronly::ScalarType::Float)  {                    \
         using input_t = float;                                                      \
         using weight_t = float;                                                     \
         using state_t = float;                                                      \
         __VA_ARGS__();                                                              \
     } else {                                                                        \
-        AT_ERROR(#NAME, " not implemented for input type '", toString(ITYPE), "'"); \
+        STD_TORCH_CHECK(false, #NAME " not implemented for input type '", ITYPE, "'"); \
     }
 
 
@@ -518,30 +509,30 @@ void set_ssm_params_fwd(SSMParamsBase &params,
                         const bool is_variable_B,
                         const bool is_variable_C,
                         // device pointers
-                        const torch::Tensor u,
-                        const torch::Tensor delta,
-                        const torch::Tensor A,
-                        const torch::Tensor B,
-                        const torch::Tensor C,
-                        const torch::Tensor out,
-                        const torch::Tensor z,
-                        const torch::Tensor out_z,
-                        const std::optional<at::Tensor>& D,
-                        const std::optional<at::Tensor>& delta_bias,
-                        const torch::Tensor ssm_states,
+                        const torch::stable::Tensor u,
+                        const torch::stable::Tensor delta,
+                        const torch::stable::Tensor A,
+                        const torch::stable::Tensor B,
+                        const torch::stable::Tensor C,
+                        const torch::stable::Tensor out,
+                        const torch::stable::Tensor z,
+                        const torch::stable::Tensor out_z,
+                        const std::optional<torch::stable::Tensor>& D,
+                        const std::optional<torch::stable::Tensor>& delta_bias,
+                        const torch::stable::Tensor ssm_states,
                         bool has_z,
                         bool delta_softplus,
-                        const std::optional<at::Tensor>& query_start_loc,
-                        const std::optional<at::Tensor>& cache_indices,
-                        const std::optional<at::Tensor>& has_initial_state,
+                        const std::optional<torch::stable::Tensor>& query_start_loc,
+                        const std::optional<torch::stable::Tensor>& cache_indices,
+                        const std::optional<torch::stable::Tensor>& has_initial_state,
                         bool varlen,
                         int64_t null_block_id,
                         int64_t block_size,
-                        const std::optional<torch::Tensor> &block_idx_first_scheduled_token,
-                        const std::optional<torch::Tensor> &block_idx_last_scheduled_token,
-                        const std::optional<torch::Tensor> &initial_state_idx,
-                        const std::optional<torch::Tensor> &cu_chunk_seqlen,
-                        const std::optional<torch::Tensor> &last_chunk_indices) {
+                        const std::optional<torch::stable::Tensor> &block_idx_first_scheduled_token,
+                        const std::optional<torch::stable::Tensor> &block_idx_last_scheduled_token,
+                        const std::optional<torch::stable::Tensor> &initial_state_idx,
+                        const std::optional<torch::stable::Tensor> &cu_chunk_seqlen,
+                        const std::optional<torch::stable::Tensor> &last_chunk_indices) {
 
     // Reset the parameters
     memset(&params, 0, sizeof(params));
@@ -654,45 +645,45 @@ void set_ssm_params_fwd(SSMParamsBase &params,
     }
 }
 
-void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
-                  const torch::Tensor &A, const torch::Tensor &B, const torch::Tensor &C,
-                  const std::optional<torch::Tensor> &D_,
-                  const std::optional<torch::Tensor> &z_,
-                  const std::optional<torch::Tensor> &delta_bias_,
+void selective_scan_fwd(const torch::stable::Tensor &u, const torch::stable::Tensor &delta,
+                  const torch::stable::Tensor &A, const torch::stable::Tensor &B, const torch::stable::Tensor &C,
+                  const std::optional<torch::stable::Tensor> &D_,
+                  const std::optional<torch::stable::Tensor> &z_,
+                  const std::optional<torch::stable::Tensor> &delta_bias_,
                   bool delta_softplus,
-                  const std::optional<torch::Tensor> &query_start_loc,
-                  const std::optional<torch::Tensor> &cache_indices,
-                  const std::optional<torch::Tensor> &has_initial_state,
-                  const torch::Tensor &ssm_states,
+                  const std::optional<torch::stable::Tensor> &query_start_loc,
+                  const std::optional<torch::stable::Tensor> &cache_indices,
+                  const std::optional<torch::stable::Tensor> &has_initial_state,
+                  const torch::stable::Tensor &ssm_states,
                   // used to identify padding entries if cache_indices provided
                   // in case of padding, the kernel will return early
                   int64_t null_block_id,
                   int64_t block_size,
-                  const std::optional<torch::Tensor> &block_idx_first_scheduled_token,
-                  const std::optional<torch::Tensor> &block_idx_last_scheduled_token,
-                  const std::optional<torch::Tensor> &initial_state_idx,
-                  const std::optional<torch::Tensor> &cu_chunk_seqlen,
-                  const std::optional<torch::Tensor> &last_chunk_indices) {
+                  const std::optional<torch::stable::Tensor> &block_idx_first_scheduled_token,
+                  const std::optional<torch::stable::Tensor> &block_idx_last_scheduled_token,
+                  const std::optional<torch::stable::Tensor> &initial_state_idx,
+                  const std::optional<torch::stable::Tensor> &cu_chunk_seqlen,
+                  const std::optional<torch::stable::Tensor> &last_chunk_indices) {
     auto input_type = u.scalar_type();
     auto weight_type = A.scalar_type();
-    TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
-    TORCH_CHECK(weight_type == at::ScalarType::Float);
+    STD_TORCH_CHECK(input_type == torch::headeronly::ScalarType::Float || input_type == torch::headeronly::ScalarType::Half || input_type == torch::headeronly::ScalarType::BFloat16);
+    STD_TORCH_CHECK(weight_type == torch::headeronly::ScalarType::Float);
 
     const bool is_variable_B = B.dim() >= 3;
     const bool is_variable_C = C.dim() >= 3;
 
-    TORCH_CHECK(delta.scalar_type() == input_type);
-    TORCH_CHECK(B.scalar_type() == (!is_variable_B ? weight_type : input_type));
-    TORCH_CHECK(C.scalar_type() == (!is_variable_C ? weight_type : input_type));
+    STD_TORCH_CHECK(delta.scalar_type() == input_type);
+    STD_TORCH_CHECK(B.scalar_type() == (!is_variable_B ? weight_type : input_type));
+    STD_TORCH_CHECK(C.scalar_type() == (!is_variable_C ? weight_type : input_type));
 
-    TORCH_CHECK(u.is_cuda());
-    TORCH_CHECK(delta.is_cuda());
-    TORCH_CHECK(A.is_cuda());
-    TORCH_CHECK(B.is_cuda());
-    TORCH_CHECK(C.is_cuda());
+    STD_TORCH_CHECK(u.is_cuda());
+    STD_TORCH_CHECK(delta.is_cuda());
+    STD_TORCH_CHECK(A.is_cuda());
+    STD_TORCH_CHECK(B.is_cuda());
+    STD_TORCH_CHECK(C.is_cuda());
 
-    TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1);
-    TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1);
+    STD_TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1);
+    STD_TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1);
 
     const auto sizes = u.sizes();
     const bool varlen = query_start_loc.has_value();
@@ -702,7 +693,7 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
     const int dstate = A.size(1);
     const int n_groups = varlen ? B.size(0) : B.size(1);
 
-    TORCH_CHECK(dstate <= 256, "selective_scan only supports state dimension <= 256");
+    STD_TORCH_CHECK(dstate <= 256, "selective_scan only supports state dimension <= 256");
 
     if (varlen) {
         CHECK_SHAPE(u, dim, seqlen);
@@ -712,94 +703,94 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
         CHECK_SHAPE(delta, batch_size, dim, seqlen);
     }
     CHECK_SHAPE(A, dim, dstate);
-    TORCH_CHECK(is_variable_B, "is_variable_B = False is disabled in favor of reduced binary size")
+    STD_TORCH_CHECK(is_variable_B, "is_variable_B = False is disabled in favor of reduced binary size");
     if (varlen) {
         CHECK_SHAPE(B, n_groups, dstate, seqlen);
     } else {
-        CHECK_SHAPE(B, batch_size, n_groups, dstate, seqlen); 
+        CHECK_SHAPE(B, batch_size, n_groups, dstate, seqlen);
     }
-    TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1);
+    STD_TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1);
 
-    TORCH_CHECK(is_variable_C, "is_variable_C = False is disabled in favor of reduced binary size")
+    STD_TORCH_CHECK(is_variable_C, "is_variable_C = False is disabled in favor of reduced binary size");
     if (varlen) {
         CHECK_SHAPE(C, n_groups, dstate, seqlen);
     } else {
-        CHECK_SHAPE(C, batch_size, n_groups, dstate, seqlen); 
+        CHECK_SHAPE(C, batch_size, n_groups, dstate, seqlen);
     }
-    TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1);
+    STD_TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1);
 
     if (D_.has_value()) {
         auto D = D_.value();
-        TORCH_CHECK(D.scalar_type() == at::ScalarType::Float);
-        TORCH_CHECK(D.is_cuda());
-        TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1);
+        STD_TORCH_CHECK(D.scalar_type() == torch::headeronly::ScalarType::Float);
+        STD_TORCH_CHECK(D.is_cuda());
+        STD_TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1);
         CHECK_SHAPE(D, dim);
     }
 
     if (delta_bias_.has_value()) {
         auto delta_bias = delta_bias_.value();
-        TORCH_CHECK(delta_bias.scalar_type() == at::ScalarType::Float);
-        TORCH_CHECK(delta_bias.is_cuda());
-        TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1);
+        STD_TORCH_CHECK(delta_bias.scalar_type() == torch::headeronly::ScalarType::Float);
+        STD_TORCH_CHECK(delta_bias.is_cuda());
+        STD_TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1);
         CHECK_SHAPE(delta_bias, dim);
     }
 
 
     if (has_initial_state.has_value()) {
         auto has_initial_state_ = has_initial_state.value();
-        TORCH_CHECK(has_initial_state_.scalar_type() == at::ScalarType::Bool);
-        TORCH_CHECK(has_initial_state_.is_cuda());
+        STD_TORCH_CHECK(has_initial_state_.scalar_type() == torch::headeronly::ScalarType::Bool);
+        STD_TORCH_CHECK(has_initial_state_.is_cuda());
         CHECK_SHAPE(has_initial_state_, batch_size);
     }
 
 
     if (query_start_loc.has_value()) {
         auto query_start_loc_ = query_start_loc.value();
-        TORCH_CHECK(query_start_loc_.scalar_type() == at::ScalarType::Int);
-        TORCH_CHECK(query_start_loc_.is_cuda());
+        STD_TORCH_CHECK(query_start_loc_.scalar_type() == torch::headeronly::ScalarType::Int);
+        STD_TORCH_CHECK(query_start_loc_.is_cuda());
     }
 
 
     if (cache_indices.has_value()) {
         auto cache_indices_ = cache_indices.value();
-        TORCH_CHECK(cache_indices_.scalar_type() == at::ScalarType::Int);
-        TORCH_CHECK(cache_indices_.is_cuda());
+        STD_TORCH_CHECK(cache_indices_.scalar_type() == torch::headeronly::ScalarType::Int);
+        STD_TORCH_CHECK(cache_indices_.is_cuda());
 
         // cache_indices can be either 1D (batch_size,) for non-APC mode
         // or 2D (batch_size, max_positions) for APC mode
         const bool is_apc_mode = block_idx_first_scheduled_token.has_value();
         if (is_apc_mode) {
-            TORCH_CHECK(cache_indices_.dim() == 2, "cache_indices must be 2D for APC mode");
-            TORCH_CHECK(cache_indices_.size(0) == batch_size, "cache_indices first dimension must match batch_size");
+            STD_TORCH_CHECK(cache_indices_.dim() == 2, "cache_indices must be 2D for APC mode");
+            STD_TORCH_CHECK(cache_indices_.size(0) == batch_size, "cache_indices first dimension must match batch_size");
         } else {
             CHECK_SHAPE(cache_indices_, batch_size);
         }
     }
-   
 
-    at::Tensor z, out_z;
+
+    torch::stable::Tensor z, out_z;
     const bool has_z = z_.has_value();
     if (has_z) {
         z = z_.value();
-        TORCH_CHECK(z.scalar_type() == input_type);
-        TORCH_CHECK(z.is_cuda());
-        TORCH_CHECK(z.stride(-1) == 1 || z.size(-1) == 1);
+        STD_TORCH_CHECK(z.scalar_type() == input_type);
+        STD_TORCH_CHECK(z.is_cuda());
+        STD_TORCH_CHECK(z.stride(-1) == 1 || z.size(-1) == 1);
         if (varlen){
             CHECK_SHAPE(z, dim, seqlen);
         } else {
             CHECK_SHAPE(z, batch_size, dim, seqlen);
         }
-        
+
         out_z = z;
     }
 
     // Right now u has BHL layout and delta has HBL layout, and we want out to have HBL layout
-    at::Tensor out = delta;
+    torch::stable::Tensor out = delta;
     // ssm_states can now be either the same as input_type or float32
     auto state_type = ssm_states.scalar_type();
-    TORCH_CHECK(state_type == input_type || state_type == at::ScalarType::Float);
-    TORCH_CHECK(ssm_states.is_cuda());
-    TORCH_CHECK(ssm_states.stride(-1) == 1);
+    STD_TORCH_CHECK(state_type == input_type || state_type == torch::headeronly::ScalarType::Float);
+    STD_TORCH_CHECK(ssm_states.is_cuda());
+    STD_TORCH_CHECK(ssm_states.stride(-1) == 1);
 
     SSMParamsBase params;
     set_ssm_params_fwd(params, batch_size, dim, seqlen, dstate, n_groups, is_variable_B, is_variable_C,
@@ -823,8 +814,8 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
                        );
 
     
-    const at::cuda::OptionalCUDAGuard device_guard(device_of(u));
-    auto stream = at::cuda::getCurrentCUDAStream().stream();
+    const torch::stable::accelerator::DeviceGuard device_guard(u.get_device_index());
+    auto stream = get_current_cuda_stream();
     DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), ssm_states.scalar_type(), "selective_scan_fwd", [&] {
         selective_scan_fwd_cuda<input_t, weight_t, state_t>(params, stream);
     });

csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu

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

csrc/libtorch_stable/ops.h

@@ -164,5 +164,341 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
 
 #endif
 
+// Attention kernels (shared CUDA/ROCm)
+void merge_attn_states(
+    torch::stable::Tensor& output,
+    std::optional<torch::stable::Tensor> output_lse,
+    const torch::stable::Tensor& prefix_output,
+    const torch::stable::Tensor& prefix_lse,
+    const torch::stable::Tensor& suffix_output,
+    const torch::stable::Tensor& suffix_lse,
+    const std::optional<int64_t> prefill_tokens_with_context,
+    const std::optional<torch::stable::Tensor>& output_scale = std::nullopt);
+
 torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x,
                                          bool inplace);
+
+// Layernorm kernels (shared CUDA/ROCm)
+void rms_norm(torch::stable::Tensor& out, torch::stable::Tensor& input,
+              torch::stable::Tensor& weight, double epsilon);
+
+void fused_add_rms_norm(torch::stable::Tensor& input,
+                        torch::stable::Tensor& residual,
+                        torch::stable::Tensor& weight, double epsilon);
+
+// Layernorm-quant kernels (shared CUDA/ROCm)
+void rms_norm_static_fp8_quant(torch::stable::Tensor& out,
+                               torch::stable::Tensor& input,
+                               torch::stable::Tensor& weight,
+                               torch::stable::Tensor& scale, double epsilon);
+
+void fused_add_rms_norm_static_fp8_quant(torch::stable::Tensor& out,
+                                         torch::stable::Tensor& input,
+                                         torch::stable::Tensor& residual,
+                                         torch::stable::Tensor& weight,
+                                         torch::stable::Tensor& scale,
+                                         double epsilon);
+
+// Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm)
+void rms_norm_dynamic_per_token_quant(
+    torch::stable::Tensor& out, torch::stable::Tensor const& input,
+    torch::stable::Tensor const& weight, torch::stable::Tensor& scales,
+    double const var_epsilon, std::optional<torch::stable::Tensor> scale_ub,
+    std::optional<torch::stable::Tensor> residual);
+
+void rms_norm_per_block_quant(torch::stable::Tensor& out,
+                              torch::stable::Tensor const& input,
+                              torch::stable::Tensor const& weight,
+                              torch::stable::Tensor& scales,
+                              double const var_epsilon,
+                              std::optional<torch::stable::Tensor> scale_ub,
+                              std::optional<torch::stable::Tensor> residual,
+                              int64_t group_size, bool is_scale_transposed);
+
+// Positional encoding kernels (shared CUDA/ROCm)
+void rotary_embedding(torch::stable::Tensor& positions,
+                      torch::stable::Tensor& query,
+                      std::optional<torch::stable::Tensor> key,
+                      int64_t head_size, torch::stable::Tensor& cos_sin_cache,
+                      bool is_neox, int64_t rope_dim_offset, bool inverse);
+
+void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q,
+                        int64_t num_heads_k, int64_t num_heads_v,
+                        int64_t head_dim, double eps,
+                        torch::stable::Tensor& q_weight,
+                        torch::stable::Tensor& k_weight,
+                        torch::stable::Tensor& cos_sin_cache, bool is_neox,
+                        torch::stable::Tensor& position_ids,
+                        int64_t forced_token_heads_per_warp);
+
+// 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& 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_(
+    torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
+    const torch::stable::Tensor& output_mask,
+    const torch::stable::Tensor& repetition_penalties);
+
+void top_k_per_row_prefill(const torch::stable::Tensor& logits,
+                           const torch::stable::Tensor& rowStarts,
+                           const torch::stable::Tensor& rowEnds,
+                           torch::stable::Tensor& indices, int64_t numRows,
+                           int64_t stride0, int64_t stride1, int64_t topK);
+
+void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n,
+                          const torch::stable::Tensor& seqLens,
+                          torch::stable::Tensor& indices, int64_t numRows,
+                          int64_t stride0, int64_t stride1, int64_t topK);
+
+void persistent_topk(const torch::stable::Tensor& logits,
+                     const torch::stable::Tensor& lengths,
+                     torch::stable::Tensor& output,
+                     torch::stable::Tensor& workspace, int64_t k,
+                     int64_t max_seq_len);
+
+void selective_scan_fwd(
+    const torch::stable::Tensor& u, const torch::stable::Tensor& delta,
+    const torch::stable::Tensor& A, const torch::stable::Tensor& B,
+    const torch::stable::Tensor& C,
+    const std::optional<torch::stable::Tensor>& D_,
+    const std::optional<torch::stable::Tensor>& z_,
+    const std::optional<torch::stable::Tensor>& delta_bias_,
+    bool delta_softplus,
+    const std::optional<torch::stable::Tensor>& query_start_loc,
+    const std::optional<torch::stable::Tensor>& cache_indices,
+    const std::optional<torch::stable::Tensor>& has_initial_state,
+    const torch::stable::Tensor& ssm_states, int64_t null_block_id,
+    int64_t block_size,
+    const std::optional<torch::stable::Tensor>& block_idx_first_scheduled_token,
+    const std::optional<torch::stable::Tensor>& block_idx_last_scheduled_token,
+    const std::optional<torch::stable::Tensor>& initial_state_idx,
+    const std::optional<torch::stable::Tensor>& cu_chunk_seqlen,
+    const std::optional<torch::stable::Tensor>& last_chunk_indices);
+
+// 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,
+                        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,
+                       torch::stable::Tensor& input);
+void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
+                     double threshold);
+void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
+                       double alpha = 1.702, double limit = 7.0);
+void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input);
+void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input);
+void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input);
+
+// INT8 quantization kernels (shared CUDA/ROCm)
+void static_scaled_int8_quant(torch::stable::Tensor& out,
+                              torch::stable::Tensor const& input,
+                              torch::stable::Tensor const& scale,
+                              std::optional<torch::stable::Tensor> const& azp);
+
+void dynamic_scaled_int8_quant(torch::stable::Tensor& out,
+                               torch::stable::Tensor const& input,
+                               torch::stable::Tensor& scales,
+                               std::optional<torch::stable::Tensor> const& azp);
+
+// FP8 quantization kernels (shared CUDA/ROCm)
+void static_scaled_fp8_quant(
+    torch::stable::Tensor& out, torch::stable::Tensor const& input,
+    torch::stable::Tensor const& scale,
+    std::optional<torch::headeronly::IntHeaderOnlyArrayRef> group_shape =
+        std::nullopt);
+
+void dynamic_scaled_fp8_quant(torch::stable::Tensor& out,
+                              torch::stable::Tensor const& input,
+                              torch::stable::Tensor& scale);
+
+void dynamic_per_token_scaled_fp8_quant(
+    torch::stable::Tensor& out, torch::stable::Tensor const& input,
+    torch::stable::Tensor& scale,
+    std::optional<torch::stable::Tensor> const& scale_ub);
+
+// GPTQ kernels (shared CUDA/ROCm)
+torch::stable::Tensor gptq_gemm(torch::stable::Tensor a,
+                                torch::stable::Tensor b_q_weight,
+                                torch::stable::Tensor b_gptq_qzeros,
+                                torch::stable::Tensor b_gptq_scales,
+                                torch::stable::Tensor b_g_idx, bool use_exllama,
+                                bool use_v2_format, int64_t bit);
+
+void gptq_shuffle(torch::stable::Tensor q_weight, torch::stable::Tensor q_perm,
+                  int64_t bit);
+
+// GGML kernels (shared CUDA/ROCm)
+torch::stable::Tensor ggml_dequantize(
+    torch::stable::Tensor W, int64_t type, int64_t m, int64_t n,
+    std::optional<torch::headeronly::ScalarType> const& dtype);
+
+torch::stable::Tensor ggml_mul_mat_vec_a8(torch::stable::Tensor W,
+                                          torch::stable::Tensor X, int64_t type,
+                                          int64_t row);
+
+torch::stable::Tensor ggml_mul_mat_a8(torch::stable::Tensor W,
+                                      torch::stable::Tensor X, int64_t type,
+                                      int64_t row);
+
+torch::stable::Tensor ggml_moe_a8(torch::stable::Tensor X,
+                                  torch::stable::Tensor W,
+                                  torch::stable::Tensor sorted_token_ids,
+                                  torch::stable::Tensor expert_ids,
+                                  torch::stable::Tensor num_tokens_post_padded,
+                                  int64_t type, int64_t row, int64_t top_k,
+                                  int64_t tokens);
+
+torch::stable::Tensor ggml_moe_a8_vec(torch::stable::Tensor X,
+                                      torch::stable::Tensor W,
+                                      torch::stable::Tensor topk_ids,
+                                      int64_t top_k, int64_t type, int64_t row,
+                                      int64_t tokens);
+
+int64_t ggml_moe_get_block_size(int64_t type);
+
+void paged_attention_v1(
+    torch::stable::Tensor& out, torch::stable::Tensor& query,
+    torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
+    int64_t num_kv_heads, double scale, torch::stable::Tensor& block_tables,
+    torch::stable::Tensor& seq_lens, int64_t block_size, int64_t max_seq_len,
+    const std::optional<torch::stable::Tensor>& alibi_slopes,
+    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
+    torch::stable::Tensor& v_scale, const int64_t tp_rank,
+    const int64_t blocksparse_local_blocks,
+    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
+    const int64_t blocksparse_head_sliding_step);
+
+void paged_attention_v2(
+    torch::stable::Tensor& out, torch::stable::Tensor& exp_sums,
+    torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out,
+    torch::stable::Tensor& query, torch::stable::Tensor& key_cache,
+    torch::stable::Tensor& value_cache, int64_t num_kv_heads, double scale,
+    torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens,
+    int64_t block_size, int64_t max_seq_len,
+    const std::optional<torch::stable::Tensor>& alibi_slopes,
+    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
+    torch::stable::Tensor& v_scale, const int64_t tp_rank,
+    const int64_t blocksparse_local_blocks,
+    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
+    const int64_t blocksparse_head_sliding_step);
+
+// Cache ops (shared CUDA/ROCm)
+void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst,
+                 int64_t block_size_in_bytes,
+                 const torch::stable::Tensor& block_mapping);
+
+// Batch swap: submit all block copies in a single driver call.
+void swap_blocks_batch(const torch::stable::Tensor& src_ptrs,
+                       const torch::stable::Tensor& dst_ptrs,
+                       const torch::stable::Tensor& sizes,
+                       bool is_src_access_order_any);
+
+void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value,
+                       torch::stable::Tensor& key_cache,
+                       torch::stable::Tensor& value_cache,
+                       torch::stable::Tensor& slot_mapping,
+                       const std::string& kv_cache_dtype,
+                       torch::stable::Tensor& k_scale,
+                       torch::stable::Tensor& v_scale);
+
+void reshape_and_cache_flash(
+    torch::stable::Tensor& key, torch::stable::Tensor& value,
+    torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
+    torch::stable::Tensor& slot_mapping, const std::string& kv_cache_dtype,
+    torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale);
+
+void concat_and_cache_mla(torch::stable::Tensor& kv_c,
+                          torch::stable::Tensor& k_pe,
+                          torch::stable::Tensor& kv_cache,
+                          torch::stable::Tensor& slot_mapping,
+                          const std::string& kv_cache_dtype,
+                          torch::stable::Tensor& scale);
+
+// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla
+void concat_and_cache_mla_rope_fused(
+    torch::stable::Tensor& positions, torch::stable::Tensor& q_pe,
+    torch::stable::Tensor& k_pe, torch::stable::Tensor& kv_c,
+    torch::stable::Tensor& rope_cos_sin_cache, bool rope_is_neox,
+    torch::stable::Tensor& slot_mapping, torch::stable::Tensor& kv_cache,
+    const std::string& kv_cache_dtype,
+    torch::stable::Tensor& kv_cache_quant_scale);
+
+// Just for unittest
+void convert_fp8(torch::stable::Tensor& dst_cache,
+                 torch::stable::Tensor& src_cache, const double scale,
+                 const std::string& kv_cache_dtype);
+
+void gather_and_maybe_dequant_cache(
+    torch::stable::Tensor const& src_cache,     // [NUM_BLOCKS, BLOCK_SIZE,
+                                                // ENTRIES...]
+    torch::stable::Tensor const& dst,           // [TOT_TOKENS, ENTRIES...]
+    torch::stable::Tensor const& block_table,   // [BATCH, BLOCK_INDICES]
+    torch::stable::Tensor const& cu_seq_lens,   // [BATCH+1]
+    torch::stable::Tensor const& token_to_seq,  // [MAX_TOKEN_ACROSS_CHUNKS]
+    int64_t num_tokens, const std::string& kv_cache_dtype,
+    torch::stable::Tensor const& scale,
+    std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
+
+// TODO(hc): cp_gather_cache need support scaled kvcahe in the future.
+void cp_gather_cache(
+    torch::stable::Tensor const& src_cache,    // [NUM_BLOCKS, BLOCK_SIZE,
+                                               // ENTRIES...]
+    torch::stable::Tensor const& dst,          // [TOT_TOKENS, ENTRIES...]
+    torch::stable::Tensor const& block_table,  // [BATCH, BLOCK_INDICES]
+    torch::stable::Tensor const& cu_seq_lens,  // [BATCH+1]
+    int64_t batch_size,
+    std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
+
+// Gather and upconvert FP8 KV cache to BF16 workspace
+void cp_gather_and_upconvert_fp8_kv_cache(
+    torch::stable::Tensor const& src_cache,         // [NUM_BLOCKS, BLOCK_SIZE,
+                                                    // 656]
+    torch::stable::Tensor const& dst,               // [TOT_TOKENS, 576]
+    torch::stable::Tensor const& block_table,       // [BATCH, BLOCK_INDICES]
+    torch::stable::Tensor const& seq_lens,          // [BATCH]
+    torch::stable::Tensor const& workspace_starts,  // [BATCH]
+    int64_t batch_size);
+
+// Indexer K quantization and cache function
+void indexer_k_quant_and_cache(
+    torch::stable::Tensor& k,             // [num_tokens, head_dim]
+    torch::stable::Tensor& kv_cache,      // [num_blocks, block_size,
+                                          // cache_stride]
+    torch::stable::Tensor& slot_mapping,  // [num_tokens]
+    int64_t quant_block_size,             // quantization block size
+    const std::string& scale_fmt);
+
+// Concatenate query nope and rope for MLA/DSA attention
+void concat_mla_q(
+    torch::stable::Tensor& ql_nope,  // [num_tokens, num_heads, nope_dim]
+    torch::stable::Tensor& q_pe,     // [num_tokens, num_heads, rope_dim]
+    torch::stable::Tensor& q_out);   // [num_tokens, num_heads, nope_dim +
+                                     // rope_dim]
+
+// Extract function to gather quantized K cache
+void cp_gather_indexer_k_quant_cache(
+    const torch::stable::Tensor& kv_cache,      // [num_blocks, block_size,
+                                                // cache_stride]
+    torch::stable::Tensor& dst_k,               // [num_tokens, head_dim]
+    torch::stable::Tensor& dst_scale,           // [num_tokens, head_dim /
+                                                // quant_block_size * 4]
+    const torch::stable::Tensor& block_table,   // [batch_size, num_blocks]
+    const torch::stable::Tensor& cu_seq_lens);  // [batch_size + 1]

csrc/libtorch_stable/pos_encoding_kernels.cu

@@ -1,8 +1,6 @@
-#include <torch/all.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
+#include "torch_utils.h"
 
-#include "cuda_compat.h"
+#include "../cuda_compat.h"
 #include "dispatch_utils.h"
 
 namespace vllm {
@@ -103,35 +101,37 @@ __global__ void rotary_embedding_kernel(
 }  // namespace vllm
 
 void rotary_embedding(
-    torch::Tensor& positions,  // [batch_size, seq_len] or [num_tokens]
-    torch::Tensor& query,  // [batch_size, seq_len, num_heads * head_size] or
-                           // [num_tokens, num_heads * head_size] or
-                           // [batch_size, seq_len, num_heads, head_size] or
-                           // [num_tokens, num_heads, head_size]
-    std::optional<torch::Tensor> key,
+    torch::stable::Tensor& positions,  // [batch_size, seq_len] or [num_tokens]
+    torch::stable::Tensor&
+        query,  // [batch_size, seq_len, num_heads * head_size] or
+                // [num_tokens, num_heads * head_size] or
+                // [batch_size, seq_len, num_heads, head_size] or
+                // [num_tokens, num_heads, head_size]
+    std::optional<torch::stable::Tensor> key,
     // null or
     // [batch_size, seq_len, num_kv_heads * head_size] or
     // [num_tokens, num_kv_heads * head_size] or
     // [batch_size, seq_len, num_heads, head_size] or
     // [num_tokens, num_heads, head_size]
     int64_t head_size,
-    torch::Tensor& cos_sin_cache,  // [max_position, rot_dim]
+    torch::stable::Tensor& cos_sin_cache,  // [max_position, rot_dim]
     bool is_neox, int64_t rope_dim_offset, bool inverse) {
   // num_tokens = batch_size * seq_len
   int64_t num_tokens = positions.numel();
   int positions_ndim = positions.dim();
 
   // Make sure num_tokens dim is consistent across positions, query, and key
-  TORCH_CHECK(
+  STD_TORCH_CHECK(
       positions_ndim == 1 || positions_ndim == 2,
       "positions must have shape [num_tokens] or [batch_size, seq_len]");
   if (positions_ndim == 1) {
-    TORCH_CHECK(query.size(0) == positions.size(0) &&
-                    (!key.has_value() || key->size(0) == positions.size(0)),
-                "query, key and positions must have the same number of tokens");
+    STD_TORCH_CHECK(
+        query.size(0) == positions.size(0) &&
+            (!key.has_value() || key->size(0) == positions.size(0)),
+        "query, key and positions must have the same number of tokens");
   }
   if (positions_ndim == 2) {
-    TORCH_CHECK(
+    STD_TORCH_CHECK(
         query.size(0) == positions.size(0) &&
             (!key.has_value() || key->size(0) == positions.size(0)) &&
             query.size(1) == positions.size(1) &&
@@ -143,20 +143,20 @@ void rotary_embedding(
   // hidden_size = num_heads * head_size
   int query_hidden_size = query.numel() / num_tokens;
   int key_hidden_size = key.has_value() ? key->numel() / num_tokens : 0;
-  TORCH_CHECK(query_hidden_size % head_size == 0);
-  TORCH_CHECK(key_hidden_size % head_size == 0);
+  STD_TORCH_CHECK(query_hidden_size % head_size == 0);
+  STD_TORCH_CHECK(key_hidden_size % head_size == 0);
 
   // Make sure query and key have consistent number of heads
   int num_heads = query_hidden_size / head_size;
   int num_kv_heads = key.has_value() ? key_hidden_size / head_size : num_heads;
-  TORCH_CHECK(num_heads % num_kv_heads == 0);
+  STD_TORCH_CHECK(num_heads % num_kv_heads == 0);
 
   int rot_dim = cos_sin_cache.size(1);
   int seq_dim_idx = positions_ndim - 1;
   int64_t query_stride = query.stride(seq_dim_idx);
   int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0;
 
-  TORCH_CHECK((rot_dim + rope_dim_offset) <= head_size);
+  STD_TORCH_CHECK((rot_dim + rope_dim_offset) <= head_size);
   // Determine head stride: for [*, heads, head_size] use stride of last dim;
   // for flat [*, heads*head_size], heads blocks are contiguous of size
   // head_size
@@ -166,30 +166,36 @@ void rotary_embedding(
 
   dim3 grid(num_tokens);
   dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "rotary_embedding", [&] {
-    using query_t = scalar_t;
-    VLLM_DISPATCH_FLOATING_TYPES(
-        cos_sin_cache.scalar_type(), "rotary_embedding_cache", [&] {
-          using cache_t = scalar_t;
-          if (is_neox) {
-            vllm::rotary_embedding_kernel<query_t, cache_t, true>
-                <<<grid, block, 0, stream>>>(
-                    positions.data_ptr<int64_t>(), query.data_ptr<query_t>(),
-                    key.has_value() ? key->data_ptr<query_t>() : nullptr,
-                    cos_sin_cache.data_ptr<cache_t>(), rot_dim, query_stride,
-                    key_stride, head_stride, num_heads, num_kv_heads, head_size,
-                    rope_dim_offset, inverse);
-          } else {
-            vllm::rotary_embedding_kernel<query_t, cache_t, false>
-                <<<grid, block, 0, stream>>>(
-                    positions.data_ptr<int64_t>(), query.data_ptr<query_t>(),
-                    key.has_value() ? key->data_ptr<query_t>() : nullptr,
-                    cos_sin_cache.data_ptr<cache_t>(), rot_dim, query_stride,
-                    key_stride, head_stride, num_heads, num_kv_heads, head_size,
-                    rope_dim_offset, inverse);
-          }
-        });
-  });
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      query.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+      query.scalar_type(), "rotary_embedding", [&] {
+        using query_t = scalar_t;
+        VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+            cos_sin_cache.scalar_type(), "rotary_embedding_cache", [&] {
+              using cache_t = scalar_t;
+              if (is_neox) {
+                vllm::rotary_embedding_kernel<query_t, cache_t, true>
+                    <<<grid, block, 0, stream>>>(
+                        positions.const_data_ptr<int64_t>(),
+                        query.mutable_data_ptr<query_t>(),
+                        key.has_value() ? key->mutable_data_ptr<query_t>()
+                                        : nullptr,
+                        cos_sin_cache.const_data_ptr<cache_t>(), rot_dim,
+                        query_stride, key_stride, head_stride, num_heads,
+                        num_kv_heads, head_size, rope_dim_offset, inverse);
+              } else {
+                vllm::rotary_embedding_kernel<query_t, cache_t, false>
+                    <<<grid, block, 0, stream>>>(
+                        positions.const_data_ptr<int64_t>(),
+                        query.mutable_data_ptr<query_t>(),
+                        key.has_value() ? key->mutable_data_ptr<query_t>()
+                                        : nullptr,
+                        cos_sin_cache.const_data_ptr<cache_t>(), rot_dim,
+                        query_stride, key_stride, head_stride, num_heads,
+                        num_kv_heads, head_size, rope_dim_offset, inverse);
+              }
+            });
+      });
 }

csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu

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

csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu

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

csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu

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

csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu

@@ -23,7 +23,7 @@
 
 #include "libtorch_stable/torch_utils.h"
 #include "libtorch_stable/dispatch_utils.h"
-#include "cuda_vec_utils.cuh"
+#include "../../cuda_vec_utils.cuh"
 
 #include "cuda_utils.h"
 #include "launch_bounds_utils.h"
@@ -38,8 +38,8 @@ namespace vllm {
 // Use UE4M3 by default.
 template <class Type, bool UE8M0_SF = false>
 __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
-    cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, int32_t num_padded_cols,
-                    Type const* __restrict__ in,
+    cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, int32_t outputCols,
+                    int32_t num_padded_cols, Type const* __restrict__ in,
                     float const* __restrict__ SFScale,
                     uint32_t* __restrict__ out, uint32_t* __restrict__ SFout) {
   using PackedVec = vllm::PackedVec<Type, CVT_FP4_PACK16>;
@@ -50,7 +50,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
                 "Vec size is not matched.");
 
   // Precompute SF layout parameter (constant for entire kernel).
-  int32_t const numKTiles = (numCols + 63) / 64;
+  int32_t const numKTiles = (outputCols + 63) / 64;
 
   int sf_m = round_up<int>(numRows, 128);
   int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x;
@@ -68,16 +68,17 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
       PackedVec in_vec;
       int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
 
-      // If we are outside valid rows OR outside valid columns -> Use Zeros
-      bool valid = (rowIdx < numRows) && (elem_idx < numCols);
+      // If we are outside valid columns, feed zeros
+      bool valid_input = (rowIdx < numRows) && (elem_idx < numCols);
+      bool valid_output = (rowIdx < numRows) && (elem_idx < outputCols);
       if constexpr (CVT_FP4_PACK16) {
         ld256_cg_or_zero(reinterpret_cast<u32x8_t&>(in_vec),
                          &reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
-                         valid);
+                         valid_input);
       } else {
         ld128_cg_or_zero(reinterpret_cast<uint4&>(in_vec),
                          &reinterpret_cast<const uint32_t*>(in)[inOffset * 4],
-                         valid);
+                         valid_input);
       }
 
       auto sf_out =
@@ -89,16 +90,16 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
           cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
               in_vec, global_scale, sf_out);
 
-      // We do NOT write output for padding because the 'out' tensor is not
-      // padded.
-      if (valid) {
+      if (valid_output) {
         if constexpr (CVT_FP4_PACK16) {
-          int64_t outOffset = rowIdx * (numCols / 8) + colIdx * 2;
+          int64_t outOffset = rowIdx * (outputCols / 8) + colIdx * 2;
           uint64_t packed64 =
               (uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo);
           reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64;
         } else {
-          out[inOffset] = out_val;
+          int64_t outOffset =
+              rowIdx * (outputCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
+          out[outOffset] = out_val;
         }
       }
     }
@@ -109,7 +110,8 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
 template <class Type, bool UE8M0_SF = false>
 __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
     cvt_fp16_to_fp4_sf_major(int32_t numRows, int32_t numCols,
-                             int32_t sf_n_unpadded, int32_t num_packed_cols,
+                             int32_t outputCols, int32_t sf_n_unpadded,
+                             int32_t num_packed_cols,
                              Type const* __restrict__ in,
                              float const* __restrict__ SFScale,
                              uint32_t* __restrict__ out,
@@ -136,7 +138,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
       PackedVec in_vec;
       int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
 
-      // If we are outside valid rows OR outside valid columns -> Use Zeros
+      // If we are outside valid columns, feed zeros
       bool valid = (rowIdx < numRows) && (elem_idx < numCols);
       if constexpr (CVT_FP4_PACK16) {
         ld256_cg_or_zero(reinterpret_cast<u32x8_t&>(in_vec),
@@ -155,16 +157,16 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
           cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
               in_vec, global_scale, sf_out);
 
-      // We do NOT write output for padding because the 'out' tensor is not
-      // padded.
-      if (valid) {
+      if (rowIdx < numRows) {
         if constexpr (CVT_FP4_PACK16) {
-          int64_t outOffset = rowIdx * (numCols / 8) + colIdx * 2;
+          int64_t outOffset = rowIdx * (outputCols / 8) + colIdx * 2;
           uint64_t packed64 =
               (uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo);
           reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64;
         } else {
-          out[inOffset] = out_val;
+          int64_t outOffset =
+              rowIdx * (outputCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
+          out[outOffset] = out_val;
         }
       }
     }
@@ -180,8 +182,11 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
                              bool is_sf_swizzled_layout) {
   int32_t m = input.size(0);
   int32_t n = input.size(1);
+  int32_t output_n = output.size(1) * 2;
 
   STD_TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
+  STD_TORCH_CHECK(output_n % 16 == 0,
+                  "The output tensor width must be a multiple of 16.");
   STD_TORCH_CHECK(
       input.scalar_type() == torch::headeronly::ScalarType::Half ||
           input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
@@ -197,7 +202,7 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
       input.get_device_index());
   auto stream = get_current_cuda_stream(input.get_device_index());
 
-  int sf_n_unpadded = int(n / CVT_FP4_SF_VEC_SIZE);
+  int output_sf_n_unpadded = int(output_n / CVT_FP4_SF_VEC_SIZE);
 
   // Grid, Block size. Each thread converts 8 values.
   dim3 block(std::min(int(n / ELTS_PER_THREAD), 512));
@@ -205,7 +210,7 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
       vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
 
   if (is_sf_swizzled_layout) {
-    int sf_n_int = int(vllm::round_up(sf_n_unpadded, 4) / 4);
+    int sf_n_int = int(vllm::round_up(output_sf_n_unpadded, 4) / 4);
     int32_t num_padded_cols =
         sf_n_int * 4 * CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
 
@@ -220,12 +225,12 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
           using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
           auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
           vllm::cvt_fp16_to_fp4<cuda_type, false><<<grid, block, 0, stream>>>(
-              m, n, num_padded_cols, input_ptr, input_sf_ptr,
+              m, n, output_n, num_padded_cols, input_ptr, input_sf_ptr,
               reinterpret_cast<uint32_t*>(output_ptr),
               reinterpret_cast<uint32_t*>(sf_out));
         });
   } else {
-    int num_packed_cols = n / CVT_FP4_ELTS_PER_THREAD;
+    int num_packed_cols = output_n / CVT_FP4_ELTS_PER_THREAD;
     int grid_y = vllm::div_round_up(num_packed_cols, static_cast<int>(block.x));
     int grid_x = std::min(
         m, std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
@@ -237,7 +242,8 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
           auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
           vllm::cvt_fp16_to_fp4_sf_major<cuda_type, false>
               <<<grid, block, 0, stream>>>(
-                  m, n, sf_n_unpadded, num_packed_cols, input_ptr, input_sf_ptr,
+                  m, n, output_n, output_sf_n_unpadded, num_packed_cols,
+                  input_ptr, input_sf_ptr,
                   reinterpret_cast<uint32_t*>(output_ptr),
                   reinterpret_cast<uint32_t*>(sf_out));
         });

csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu

@@ -22,6 +22,8 @@
 
 #include "cutlass/cutlass.h"
 
+#include <type_traits>
+
 #include "cutlass/gemm/collective/collective_builder.hpp"
 #include "cutlass/epilogue/collective/collective_builder.hpp"
 #include "cutlass/gemm/device/gemm_universal_adapter.h"
@@ -30,15 +32,21 @@
 #include "cutlass/util/packed_stride.hpp"
 
 #include "core/math.hpp"
+#include "core/batch_invariant.hpp"
 
 using namespace cute;
 
 #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
 
-// Configuration for M in (256, inf)
+// Configuration for M in (256, inf), also reused for batch-invariant mode
+// to keep a fixed large-M tiling across all batch sizes.
+// Do not change the tile K or tile scheduler here unless you are also
+// updating the batch-invariant behavior; if batch-invariant mode needs a
+// different schedule, add a dedicated batch-invariant config/path instead.
 struct sm100_fp4_config_default {
   using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
   using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
+  using TileScheduler = cutlass::gemm::PersistentScheduler;
   using TileShape = Shape<_256, _256, _256>;
   using ClusterShape = Shape<_2, _1, _1>;
   using PerSmTileShape_MNK = Shape<_128, _256, _256>;
@@ -48,6 +56,7 @@ struct sm100_fp4_config_default {
 struct sm100_fp4_config_M256 {
   using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
   using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
+  using TileScheduler = void;
   using TileShape = Shape<_256, _128, _256>;
   using ClusterShape = Shape<_2, _1, _1>;
   using PerSmTileShape_MNK = Shape<_128, _128, _256>;
@@ -57,6 +66,7 @@ struct sm100_fp4_config_M256 {
 struct sm100_fp4_config_M16 {
   using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
   using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
+  using TileScheduler = void;
   using TileShape = Shape<_128, _128, _256>;
   using ClusterShape = Shape<_1, _1, _1>;
   using PerSmTileShape_MNK = Shape<_128, _128, _256>;
@@ -98,7 +108,7 @@ struct Fp4GemmSm100 {
           cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator,
           ElementAccumulator, ElementC, LayoutCTag, AlignmentC, ElementD,
           LayoutDTag, AlignmentD,
-          cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp;
+          typename Config::EpilogueSchedule>::CollectiveOp;
 
   using CollectiveMainloop =
       typename cutlass::gemm::collective::CollectiveBuilder<
@@ -107,10 +117,13 @@ struct Fp4GemmSm100 {
           ClusterShape,
           cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
               sizeof(typename CollectiveEpilogue::SharedStorage))>,
-          cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
+          typename Config::KernelSchedule>::CollectiveOp;
 
-  using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
-      Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
+  using TileScheduler = typename Config::TileScheduler;
+  using GemmKernel =
+      cutlass::gemm::kernel::GemmUniversal<Shape<int, int, int, int>,
+                                           CollectiveMainloop,
+                                           CollectiveEpilogue, TileScheduler>;
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
   using StrideA = typename Gemm::GemmKernel::StrideA;
   using LayoutA = decltype(cute::make_layout(make_shape(0, 0, 0), StrideA{}));
@@ -205,6 +218,17 @@ void cutlass_fp4_gemm_dispatch(torch::stable::Tensor& D,
                                torch::stable::Tensor const& B_sf,
                                torch::stable::Tensor const& alpha, int64_t m,
                                int64_t n, int64_t k, cudaStream_t stream) {
+  if (vllm::vllm_is_batch_invariant()) {
+    using BiGemm = Fp4GemmSm100<sm100_fp4_config_default, OutType>;
+    static_assert(
+        cute::is_same_v<typename BiGemm::TileScheduler,
+                        cutlass::gemm::PersistentScheduler>,
+        "batch_invariant requires a persistent tile scheduler; stream-K or "
+        "split-K would break numerical invariance");
+    runGemm<BiGemm>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+    return;
+  }
+
   uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
 
   if (mp2 <= 16) {

csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu

@@ -22,6 +22,8 @@
 
 #include "cutlass/cutlass.h"
 
+#include <type_traits>
+
 #include "cutlass/gemm/collective/collective_builder.hpp"
 #include "cutlass/epilogue/collective/collective_builder.hpp"
 #include "cutlass/gemm/device/gemm_universal_adapter.h"
@@ -30,6 +32,7 @@
 #include "cutlass/util/packed_stride.hpp"
 
 #include "core/math.hpp"
+#include "core/batch_invariant.hpp"
 
 using namespace cute;
 
@@ -49,12 +52,22 @@ constexpr auto FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte;
 constexpr auto SF_DTYPE = torch::headeronly::ScalarType::Float8_e4m3fn;
 
 struct sm120_fp4_config_M256 {
+  using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
+  using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
+  using TileScheduler = void;
   using ClusterShape = Shape<_1, _1, _1>;
   using MmaTileShape = Shape<_128, _128, _128>;
   using PerSmTileShape_MNK = Shape<_128, _128, _128>;
 };
 
 struct sm120_fp4_config_default {
+  // Also used for batch-invariant mode.
+  // Do not change the tile K or tile scheduler here unless you are also
+  // updating the batch-invariant behavior; if batch-invariant mode needs a
+  // different schedule, add a dedicated batch-invariant config/path instead.
+  using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
+  using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
+  using TileScheduler = cutlass::gemm::PersistentScheduler;
   using ClusterShape = Shape<_1, _1, _1>;
   using MmaTileShape = Shape<_256, _128, _128>;
   using PerSmTileShape_MNK = Shape<_256, _128, _128>;
@@ -91,7 +104,7 @@ struct Fp4GemmSm120 {
           cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator,
           ElementAccumulator, ElementC, LayoutCTag, AlignmentC, ElementD,
           LayoutDTag, AlignmentD,
-          cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp;
+          typename Config::EpilogueSchedule>::CollectiveOp;
 
   using CollectiveMainloop =
       typename cutlass::gemm::collective::CollectiveBuilder<
@@ -100,10 +113,13 @@ struct Fp4GemmSm120 {
           ClusterShape,
           cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
               sizeof(typename CollectiveEpilogue::SharedStorage))>,
-          cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
+          typename Config::KernelSchedule>::CollectiveOp;
 
-  using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
-      Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
+  using TileScheduler = typename Config::TileScheduler;
+  using GemmKernel =
+      cutlass::gemm::kernel::GemmUniversal<Shape<int, int, int, int>,
+                                           CollectiveMainloop,
+                                           CollectiveEpilogue, TileScheduler>;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
 };
@@ -180,39 +196,41 @@ void runGemm(torch::stable::Tensor& D, torch::stable::Tensor const& A,
   CUTLASS_CHECK(gemm.run(arguments, workspace.data_ptr(), stream));
 }
 
-void cutlass_fp4_bf16_gemm_dispatch(torch::stable::Tensor& D,
-                                    torch::stable::Tensor const& A,
-                                    torch::stable::Tensor const& B,
-                                    torch::stable::Tensor const& A_sf,
-                                    torch::stable::Tensor const& B_sf,
-                                    torch::stable::Tensor const& alpha, int m,
-                                    int n, int k, cudaStream_t stream) {
+namespace {
+
+// Dispatch function to select appropriate config based on M (file-local;
+// internal linkage avoids clashing with SM100's cutlass_fp4_gemm_dispatch in
+// nvfp4_scaled_mm_kernels.cu).
+template <typename OutType>
+void cutlass_fp4_gemm_dispatch(torch::stable::Tensor& D,
+                               torch::stable::Tensor const& A,
+                               torch::stable::Tensor const& B,
+                               torch::stable::Tensor const& A_sf,
+                               torch::stable::Tensor const& B_sf,
+                               torch::stable::Tensor const& alpha, int m, int n,
+                               int k, cudaStream_t stream) {
+  if (vllm::vllm_is_batch_invariant()) {
+    using BiGemm = Fp4GemmSm120<sm120_fp4_config_default, OutType>;
+    static_assert(
+        cute::is_same_v<typename BiGemm::TileScheduler,
+                        cutlass::gemm::PersistentScheduler>,
+        "batch_invariant requires a persistent tile scheduler; stream-K or "
+        "split-K would break numerical invariance");
+    runGemm<typename BiGemm::Gemm>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+    return;
+  }
+
   uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
   if (mp2 <= 256) {
-    runGemm<Fp4GemmSm120<sm120_fp4_config_M256, cutlass::bfloat16_t>::Gemm>(
+    runGemm<typename Fp4GemmSm120<sm120_fp4_config_M256, OutType>::Gemm>(
         D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
   } else {
-    runGemm<Fp4GemmSm120<sm120_fp4_config_default, cutlass::bfloat16_t>::Gemm>(
+    runGemm<typename Fp4GemmSm120<sm120_fp4_config_default, OutType>::Gemm>(
         D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
   }
 }
 
-void cutlass_fp4_f16_gemm_dispatch(torch::stable::Tensor& D,
-                                   torch::stable::Tensor const& A,
-                                   torch::stable::Tensor const& B,
-                                   torch::stable::Tensor const& A_sf,
-                                   torch::stable::Tensor const& B_sf,
-                                   torch::stable::Tensor const& alpha, int m,
-                                   int n, int k, cudaStream_t stream) {
-  uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
-  if (mp2 <= 256) {
-    runGemm<Fp4GemmSm120<sm120_fp4_config_M256, cutlass::half_t>::Gemm>(
-        D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
-  } else {
-    runGemm<Fp4GemmSm120<sm120_fp4_config_default, cutlass::half_t>::Gemm>(
-        D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
-  }
-}
+}  // namespace
 
 void cutlass_scaled_fp4_mm_sm120a(torch::stable::Tensor& D,
                                   torch::stable::Tensor const& A,
@@ -275,11 +293,11 @@ void cutlass_scaled_fp4_mm_sm120a(torch::stable::Tensor& D,
   const cudaStream_t stream = get_current_cuda_stream(A.get_device_index());
 
   if (out_dtype == torch::headeronly::ScalarType::BFloat16) {
-    return cutlass_fp4_bf16_gemm_dispatch(D, A, B, A_sf, B_sf, alpha, m, n, k,
-                                          stream);
+    return cutlass_fp4_gemm_dispatch<cutlass::bfloat16_t>(
+        D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
   } else if (out_dtype == torch::headeronly::ScalarType::Half) {
-    return cutlass_fp4_f16_gemm_dispatch(D, A, B, A_sf, B_sf, alpha, m, n, k,
-                                         stream);
+    return cutlass_fp4_gemm_dispatch<cutlass::half_t>(D, A, B, A_sf, B_sf,
+                                                      alpha, m, n, k, stream);
   } else {
     STD_TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm sm120 (",
                     out_dtype, ")");

csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh

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

csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu

@@ -1,6 +1,5 @@
 
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
+#include "../../torch_utils.h"
 
 #include "../../dispatch_utils.h"
 #include "layernorm_utils.cuh"
@@ -134,63 +133,71 @@ __global__ void rms_norm_per_block_quant_kernel(
 // Residual add + RMS norm + dynamic per token
 template <typename scalar_in_t>
 void rms_norm_dynamic_per_token_quant_dispatch(
-    torch::Tensor& out,           // [..., hidden_size]
-    torch::Tensor const& input,   // [..., hidden_size]
-    torch::Tensor const& weight,  // [hidden_size]
-    torch::Tensor& scales,        // [num_tokens]
-    double const var_epsilon,     // Variance epsilon used in norm calculation
-    std::optional<at::Tensor> const& scale_ub,
-    std::optional<at::Tensor>& residual) {
+    torch::stable::Tensor& out,           // [..., hidden_size]
+    torch::stable::Tensor const& input,   // [..., hidden_size]
+    torch::stable::Tensor const& weight,  // [hidden_size]
+    torch::stable::Tensor& scales,        // [num_tokens]
+    double const var_epsilon,  // Variance epsilon used in norm calculation
+    std::optional<torch::stable::Tensor> const& scale_ub,
+    std::optional<torch::stable::Tensor>& residual) {
   int32_t hidden_size = input.size(-1);
-  int32_t input_stride = input.view({-1, hidden_size}).stride(0);
+  int32_t input_stride =
+      torch::stable::view(input, {-1, hidden_size}).stride(0);
   auto num_tokens = input.numel() / hidden_size;
 
   dim3 grid(num_tokens);
   dim3 block(std::min(hidden_size, 1024));
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
 
-  VLLM_DISPATCH_BOOL(residual.has_value(), has_residual, [&] {
-    VLLM_DISPATCH_QUANT_TYPES(
+  VLLM_STABLE_DISPATCH_BOOL(residual.has_value(), has_residual, [&] {
+    VLLM_STABLE_DISPATCH_QUANT_TYPES(
         out.scalar_type(), "rms_norm_dynamic_per_token_quant_kernel", [&] {
           vllm::rms_norm_dynamic_per_token_quant_kernel<scalar_in_t, scalar_t,
                                                         has_residual>
               <<<grid, block, 0, stream>>>(
-                  out.data_ptr<scalar_t>(), scales.data_ptr<float>(),
-                  input.data_ptr<scalar_in_t>(), weight.data_ptr<scalar_in_t>(),
-                  scale_ub.has_value() ? scale_ub->data_ptr<float>() : nullptr,
+                  out.mutable_data_ptr<scalar_t>(),
+                  scales.mutable_data_ptr<float>(),
+                  input.const_data_ptr<scalar_in_t>(),
+                  weight.const_data_ptr<scalar_in_t>(),
+                  scale_ub.has_value() ? scale_ub->const_data_ptr<float>()
+                                       : nullptr,
                   var_epsilon, hidden_size, input_stride,
-                  has_residual ? residual->data_ptr<scalar_in_t>() : nullptr);
+                  has_residual ? residual->mutable_data_ptr<scalar_in_t>()
+                               : nullptr);
         });
   });
 }
 
 void rms_norm_dynamic_per_token_quant(
-    torch::Tensor& out,           // [..., hidden_size]
-    torch::Tensor const& input,   // [..., hidden_size]
-    torch::Tensor const& weight,  // [hidden_size]
-    torch::Tensor& scales,        // [num_tokens]
-    double const var_epsilon,     // Variance epsilon used in norm calculation
-    std::optional<at::Tensor> scale_ub, std::optional<at::Tensor> residual) {
-  static c10::ScalarType kFp8Type = is_fp8_ocp()
-                                        ? c10::ScalarType::Float8_e4m3fn
-                                        : c10::ScalarType::Float8_e4m3fnuz;
-  TORCH_CHECK(out.dtype() == kFp8Type || out.dtype() == torch::kInt8);
-  TORCH_CHECK(out.is_contiguous());
-  TORCH_CHECK(input.stride(-1) == 1,
-              "Input must be contiguous in the last dimension");
+    torch::stable::Tensor& out,           // [..., hidden_size]
+    torch::stable::Tensor const& input,   // [..., hidden_size]
+    torch::stable::Tensor const& weight,  // [hidden_size]
+    torch::stable::Tensor& scales,        // [num_tokens]
+    double const var_epsilon,  // Variance epsilon used in norm calculation
+    std::optional<torch::stable::Tensor> scale_ub,
+    std::optional<torch::stable::Tensor> residual) {
+  static torch::headeronly::ScalarType kFp8Type =
+      is_fp8_ocp() ? torch::headeronly::ScalarType::Float8_e4m3fn
+                   : torch::headeronly::ScalarType::Float8_e4m3fnuz;
+  STD_TORCH_CHECK(out.scalar_type() == kFp8Type ||
+                  out.scalar_type() == torch::headeronly::ScalarType::Char);
+  STD_TORCH_CHECK(out.is_contiguous());
+  STD_TORCH_CHECK(input.stride(-1) == 1,
+                  "Input must be contiguous in the last dimension");
 
   if (scale_ub.has_value()) {
-    TORCH_CHECK(out.dtype() == kFp8Type);
+    STD_TORCH_CHECK(out.scalar_type() == kFp8Type);
   }
-  TORCH_CHECK(weight.dtype() == input.dtype());
-  TORCH_CHECK(scales.dtype() == torch::kFloat32);
+  STD_TORCH_CHECK(weight.scalar_type() == input.scalar_type());
+  STD_TORCH_CHECK(scales.scalar_type() == torch::headeronly::ScalarType::Float);
   if (residual) {
-    TORCH_CHECK(residual->scalar_type() == input.scalar_type());
-    TORCH_CHECK(residual->is_contiguous());
+    STD_TORCH_CHECK(residual->scalar_type() == input.scalar_type());
+    STD_TORCH_CHECK(residual->is_contiguous());
   }
 
-  VLLM_DISPATCH_FLOATING_TYPES(
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "rms_norm_dynamic_per_token_quant_dispatch", [&] {
         rms_norm_dynamic_per_token_quant_dispatch<scalar_t>(
             out, input, weight, scales, var_epsilon, scale_ub, residual);
@@ -199,103 +206,115 @@ void rms_norm_dynamic_per_token_quant(
 
 // Residual add + RMS norm + dynamic per token
 void rms_norm_per_block_quant_dispatch(
-    torch::Tensor& out,           // [..., hidden_size]
-    torch::Tensor const& input,   // [..., hidden_size]
-    torch::Tensor const& weight,  // [hidden_size]
-    torch::Tensor& scales,        // [num_tokens, hidden_size / group_size] or
-                                  // [hidden_size / group_size, num_tokens]
+    torch::stable::Tensor& out,           // [..., hidden_size]
+    torch::stable::Tensor const& input,   // [..., hidden_size]
+    torch::stable::Tensor const& weight,  // [hidden_size]
+    torch::stable::Tensor& scales,        // [num_tokens, hidden_size /
+                                          // group_size] or
+                                          // [hidden_size / group_size,
+                                          // num_tokens]
     int32_t group_size,
     double const var_epsilon,  // Variance epsilon used in norm calculation
-    std::optional<at::Tensor> const& scale_ub,
-    std::optional<at::Tensor>& residual, bool is_scale_transposed) {
+    std::optional<torch::stable::Tensor> const& scale_ub,
+    std::optional<torch::stable::Tensor>& residual, bool is_scale_transposed) {
   int32_t hidden_size = input.size(-1);
-  int32_t input_stride = input.view({-1, hidden_size}).stride(0);
+  int32_t input_stride =
+      torch::stable::view(input, {-1, hidden_size}).stride(0);
 
-  TORCH_CHECK(hidden_size % 4 == 0,
-              "Hidden size must be divisible by 4 for vectorized access");
-  TORCH_CHECK(input_stride % 4 == 0,
-              "Input stride must be divisible by 4 for vectorized access");
-  TORCH_CHECK(group_size % 4 == 0,
-              "Group size must be divisible by 4 for vectorized access");
+  STD_TORCH_CHECK(hidden_size % 4 == 0,
+                  "Hidden size must be divisible by 4 for vectorized access");
+  STD_TORCH_CHECK(input_stride % 4 == 0,
+                  "Input stride must be divisible by 4 for vectorized access");
+  STD_TORCH_CHECK(group_size % 4 == 0,
+                  "Group size must be divisible by 4 for vectorized access");
 
   auto num_tokens = input.numel() / hidden_size;
 
   dim3 grid(num_tokens);
   const int max_block_size = (num_tokens <= 256) ? 512 : 256;
   dim3 block(std::min(hidden_size, max_block_size));
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
 
-  VLLM_DISPATCH_FLOATING_TYPES(
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "rms_norm_per_block_quant_fp_dispatch", [&] {
         using scalar_in_t = scalar_t;
-        VLLM_DISPATCH_GROUP_SIZE(group_size, gs, [&] {
-          VLLM_DISPATCH_BOOL(residual.has_value(), has_residual, [&] {
-            VLLM_DISPATCH_BOOL(is_scale_transposed, transpose_scale, [&] {
-              VLLM_DISPATCH_QUANT_TYPES(
-                  out.scalar_type(), "rms_norm_per_block_quant_kernel", [&] {
-                    vllm::rms_norm_per_block_quant_kernel<scalar_in_t, scalar_t,
-                                                          has_residual,
-                                                          transpose_scale, gs>
-                        <<<grid, block, 0, stream>>>(
-                            out.data_ptr<scalar_t>(), scales.data_ptr<float>(),
-                            input.data_ptr<scalar_in_t>(),
-                            weight.data_ptr<scalar_in_t>(),
-                            scale_ub.has_value() ? scale_ub->data_ptr<float>()
-                                                 : nullptr,
+        VLLM_STABLE_DISPATCH_GROUP_SIZE(group_size, gs, [&] {
+          VLLM_STABLE_DISPATCH_BOOL(residual.has_value(), has_residual, [&] {
+            VLLM_STABLE_DISPATCH_BOOL(
+                is_scale_transposed, transpose_scale, [&] {
+                  VLLM_STABLE_DISPATCH_QUANT_TYPES(
+                      out.scalar_type(), "rms_norm_per_block_quant_kernel",
+                      [&] {
+                        vllm::rms_norm_per_block_quant_kernel<
+                            scalar_in_t, scalar_t, has_residual,
+                            transpose_scale, gs><<<grid, block, 0, stream>>>(
+                            out.mutable_data_ptr<scalar_t>(),
+                            scales.mutable_data_ptr<float>(),
+                            input.const_data_ptr<scalar_in_t>(),
+                            weight.const_data_ptr<scalar_in_t>(),
+                            scale_ub.has_value()
+                                ? scale_ub->const_data_ptr<float>()
+                                : nullptr,
                             var_epsilon, hidden_size, input_stride,
-                            has_residual ? residual->data_ptr<scalar_in_t>()
-                                         : nullptr,
+                            has_residual
+                                ? residual->mutable_data_ptr<scalar_in_t>()
+                                : nullptr,
                             scales.stride(1));
-                  });
-            });
+                      });
+                });
           });
         });
       });
 }
 
-void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
-                              torch::Tensor const& weight,
-                              torch::Tensor& scales, double const var_epsilon,
-                              std::optional<torch::Tensor> scale_ub,
-                              std::optional<torch::Tensor> residual,
+void rms_norm_per_block_quant(torch::stable::Tensor& out,
+                              torch::stable::Tensor const& input,
+                              torch::stable::Tensor const& weight,
+                              torch::stable::Tensor& scales,
+                              double const var_epsilon,
+                              std::optional<torch::stable::Tensor> scale_ub,
+                              std::optional<torch::stable::Tensor> residual,
                               int64_t group_size, bool is_scale_transposed) {
-  static c10::ScalarType kFp8Type = is_fp8_ocp()
-                                        ? c10::ScalarType::Float8_e4m3fn
-                                        : c10::ScalarType::Float8_e4m3fnuz;
-  TORCH_CHECK(out.dtype() == kFp8Type || out.dtype() == torch::kInt8);
-  TORCH_CHECK(out.is_contiguous());
-  TORCH_CHECK(input.stride(-1) == 1,
-              "Input must be contiguous in the last dimension");
+  static torch::headeronly::ScalarType kFp8Type =
+      is_fp8_ocp() ? torch::headeronly::ScalarType::Float8_e4m3fn
+                   : torch::headeronly::ScalarType::Float8_e4m3fnuz;
+  STD_TORCH_CHECK(out.scalar_type() == kFp8Type ||
+                  out.scalar_type() == torch::headeronly::ScalarType::Char);
+  STD_TORCH_CHECK(out.is_contiguous());
+  STD_TORCH_CHECK(input.stride(-1) == 1,
+                  "Input must be contiguous in the last dimension");
 
   if (scale_ub.has_value()) {
-    TORCH_CHECK(out.dtype() == kFp8Type);
+    STD_TORCH_CHECK(out.scalar_type() == kFp8Type);
   }
-  TORCH_CHECK(weight.dtype() == input.dtype());
-  TORCH_CHECK(scales.dtype() == torch::kFloat32);
+  STD_TORCH_CHECK(weight.scalar_type() == input.scalar_type());
+  STD_TORCH_CHECK(scales.scalar_type() == torch::headeronly::ScalarType::Float);
   if (residual) {
-    TORCH_CHECK(residual->scalar_type() == input.scalar_type());
-    TORCH_CHECK(residual->is_contiguous());
+    STD_TORCH_CHECK(residual->scalar_type() == input.scalar_type());
+    STD_TORCH_CHECK(residual->is_contiguous());
   }
 
-  TORCH_CHECK(group_size == 128 || group_size == 64,
-              "Unsupported group size: ", group_size);
+  STD_TORCH_CHECK(group_size == 128 || group_size == 64,
+                  "Unsupported group size: ", group_size);
 
   if (scales.stride(1) > 1) {
-    TORCH_CHECK(is_scale_transposed,
-                "Outer scale stride must be 1 when scales are not transposed");
+    STD_TORCH_CHECK(
+        is_scale_transposed,
+        "Outer scale stride must be 1 when scales are not transposed");
   }
 
   int64_t hidden_size = input.size(-1);
-  TORCH_CHECK(hidden_size > 0 && hidden_size % group_size == 0,
-              "hidden_size must be a positive multiple of group_size");
+  STD_TORCH_CHECK(hidden_size > 0 && hidden_size % group_size == 0,
+                  "hidden_size must be a positive multiple of group_size");
   int64_t num_tokens = input.numel() / hidden_size;
   int64_t num_groups = hidden_size / group_size;
-  TORCH_CHECK(scales.numel() >= num_tokens * num_groups,
-              "scales buffer too small: need ", num_tokens * num_groups,
-              " elements, got ", scales.numel());
+  STD_TORCH_CHECK(scales.numel() >= num_tokens * num_groups,
+                  "scales buffer too small: need ", num_tokens * num_groups,
+                  " elements, got ", scales.numel());
 
   rms_norm_per_block_quant_dispatch(out, input, weight, scales, group_size,
                                     var_epsilon, scale_ub, residual,
                                     is_scale_transposed);
-}
+}

csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh

@@ -8,8 +8,8 @@
 #include "quantization/utils.cuh"
 #include "quant_conversions.cuh"
 
-#include "../../cub_helpers.h"
-#include "../../cuda_compat.h"
+#include "../../../cub_helpers.h"
+#include "../../../cuda_compat.h"
 
 namespace vllm {
 

csrc/libtorch_stable/quantization/fused_kernels/quant_conversions.cuh

@@ -6,7 +6,7 @@
 
 #include "libtorch_stable/quantization/vectorization.cuh"
 // TODO(luka/varun):refactor common.cuh to use this file instead
-#include "../w8a8/fp8/common.cuh"
+#include "../../../quantization/w8a8/fp8/common.cuh"
 
 namespace vllm {
 

csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu

@@ -1,17 +1,20 @@
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
 
-#include <torch/all.h>
-#include <c10/cuda/CUDAGuard.h>
+#include "../../../cuda_compat.h"
+#include "../../dispatch_utils.h"
+#include "../../torch_utils.h"
 
-#include "../../cuda_compat.h"
-#include "dispatch_utils.h"
+#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"
 
@@ -71,16 +74,17 @@ static void quantize_row_q8_1_cuda(const scalar_t* x, void* vy, const int kx,
   }
 }
 
-torch::Tensor ggml_dequantize(torch::Tensor W,  // quant weight
-                              int64_t type, int64_t m, int64_t n,
-                              std::optional<at::ScalarType> const& dtype) {
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(W));
-  auto dtype_ = dtype.value_or(torch::kFloat16);
-  auto options = torch::TensorOptions().dtype(dtype_).device(W.device());
-  at::Tensor DW = torch::empty({m, n}, options);
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
+torch::stable::Tensor ggml_dequantize(
+    torch::stable::Tensor W,  // quant weight
+    int64_t type, int64_t m, int64_t n,
+    std::optional<torch::headeronly::ScalarType> const& dtype) {
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      W.get_device_index());
+  auto dtype_ = dtype.value_or(torch::headeronly::ScalarType::Half);
+  auto DW = torch::stable::empty({m, n}, dtype_, std::nullopt, W.device());
+  cudaStream_t stream = get_current_cuda_stream();
 
-  VLLM_DISPATCH_FLOATING_TYPES(DW.scalar_type(), "ggml_dequantize", [&] {
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(DW.scalar_type(), "ggml_dequantize", [&] {
     auto to_cuda = ggml_get_to_cuda<scalar_t>(type);
     to_cuda((void*)W.data_ptr(), (scalar_t*)DW.data_ptr(), m * n, stream);
   });
@@ -88,135 +92,142 @@ torch::Tensor ggml_dequantize(torch::Tensor W,  // quant weight
   return DW;
 }
 
-torch::Tensor ggml_mul_mat_vec_a8(torch::Tensor W,  // quant weight
-                                  torch::Tensor X,  // input
-                                  int64_t type, int64_t row) {
+torch::stable::Tensor ggml_mul_mat_vec_a8(
+    torch::stable::Tensor W,  // quant weight
+    torch::stable::Tensor X,  // input
+    int64_t type, int64_t row) {
   int col = X.sizes()[1];
   int vecs = X.sizes()[0];
   const int padded = (col + 512 - 1) / 512 * 512;
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
-  auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device());
-  at::Tensor Y = torch::empty({vecs, row}, options);
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
-  options = torch::TensorOptions().dtype(torch::kInt32).device(W.device());
-  at::Tensor quant_X = torch::empty({vecs, padded / 32 * 9}, options);
-  VLLM_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_mul_mat_vec_a8", [&] {
-    quantize_row_q8_1_cuda<scalar_t>(
-        (scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, vecs, stream);
-    switch (type) {
-      case 2:
-        mul_mat_vec_q4_0_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 3:
-        mul_mat_vec_q4_1_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 6:
-        mul_mat_vec_q5_0_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 7:
-        mul_mat_vec_q5_1_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 8:
-        mul_mat_vec_q8_0_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 10:
-        mul_mat_vec_q2_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 11:
-        mul_mat_vec_q3_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 12:
-        mul_mat_vec_q4_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 13:
-        mul_mat_vec_q5_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 14:
-        mul_mat_vec_q6_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 16:
-        mul_mat_vec_iq2_xxs_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 17:
-        mul_mat_vec_iq2_xs_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 18:
-        mul_mat_vec_iq3_xxs_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 19:
-        mul_mat_vec_iq1_s_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 20:
-        mul_mat_vec_iq4_nl_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 21:
-        mul_mat_vec_iq3_s_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 22:
-        mul_mat_vec_iq2_s_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 23:
-        mul_mat_vec_iq4_xs_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-      case 29:
-        mul_mat_vec_iq1_m_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
-        break;
-    }
-  });
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      X.get_device_index());
+  auto Y = torch::stable::empty({vecs, row}, X.scalar_type(), std::nullopt,
+                                W.device());
+  cudaStream_t stream = get_current_cuda_stream();
+  auto quant_X = torch::stable::empty({vecs, padded / 32 * 9},
+                                      torch::headeronly::ScalarType::Int,
+                                      std::nullopt, W.device());
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+      X.scalar_type(), "ggml_mul_mat_vec_a8", [&] {
+        quantize_row_q8_1_cuda<scalar_t>((scalar_t*)X.data_ptr(),
+                                         (void*)quant_X.data_ptr(), col, vecs,
+                                         stream);
+        switch (type) {
+          case 2:
+            mul_mat_vec_q4_0_q8_1_cuda<scalar_t>(
+                (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
+                (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
+            break;
+          case 3:
+            mul_mat_vec_q4_1_q8_1_cuda<scalar_t>(
+                (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
+                (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
+            break;
+          case 6:
+            mul_mat_vec_q5_0_q8_1_cuda<scalar_t>(
+                (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
+                (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
+            break;
+          case 7:
+            mul_mat_vec_q5_1_q8_1_cuda<scalar_t>(
+                (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
+                (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
+            break;
+          case 8:
+            mul_mat_vec_q8_0_q8_1_cuda<scalar_t>(
+                (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
+                (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
+            break;
+          case 10:
+            mul_mat_vec_q2_K_q8_1_cuda<scalar_t>(
+                (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
+                (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
+            break;
+          case 11:
+            mul_mat_vec_q3_K_q8_1_cuda<scalar_t>(
+                (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
+                (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
+            break;
+          case 12:
+            mul_mat_vec_q4_K_q8_1_cuda<scalar_t>(
+                (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
+                (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
+            break;
+          case 13:
+            mul_mat_vec_q5_K_q8_1_cuda<scalar_t>(
+                (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
+                (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
+            break;
+          case 14:
+            mul_mat_vec_q6_K_q8_1_cuda<scalar_t>(
+                (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
+                (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
+            break;
+          case 16:
+            mul_mat_vec_iq2_xxs_q8_1_cuda<scalar_t>(
+                (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
+                (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
+            break;
+          case 17:
+            mul_mat_vec_iq2_xs_q8_1_cuda<scalar_t>(
+                (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
+                (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
+            break;
+          case 18:
+            mul_mat_vec_iq3_xxs_q8_1_cuda<scalar_t>(
+                (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
+                (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
+            break;
+          case 19:
+            mul_mat_vec_iq1_s_q8_1_cuda<scalar_t>(
+                (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
+                (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
+            break;
+          case 20:
+            mul_mat_vec_iq4_nl_q8_1_cuda<scalar_t>(
+                (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
+                (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
+            break;
+          case 21:
+            mul_mat_vec_iq3_s_q8_1_cuda<scalar_t>(
+                (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
+                (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
+            break;
+          case 22:
+            mul_mat_vec_iq2_s_q8_1_cuda<scalar_t>(
+                (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
+                (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
+            break;
+          case 23:
+            mul_mat_vec_iq4_xs_q8_1_cuda<scalar_t>(
+                (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
+                (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
+            break;
+          case 29:
+            mul_mat_vec_iq1_m_q8_1_cuda<scalar_t>(
+                (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
+                (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
+            break;
+        }
+      });
   return Y;
 }
 
-torch::Tensor ggml_mul_mat_a8(torch::Tensor W,  // quant weight
-                              torch::Tensor X,  // input
-                              int64_t type, int64_t row) {
+torch::stable::Tensor ggml_mul_mat_a8(torch::stable::Tensor W,  // quant weight
+                                      torch::stable::Tensor X,  // input
+                                      int64_t type, int64_t row) {
   int col = X.sizes()[1];
   int padded = (col + 512 - 1) / 512 * 512;
   int batch = X.sizes()[0];
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
-  auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device());
-  at::Tensor Y = torch::empty({batch, row}, options);
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
-  options = torch::TensorOptions().dtype(torch::kInt32).device(W.device());
-  at::Tensor quant_X = torch::empty({batch, padded / 32 * 9}, options);
-  VLLM_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_mul_mat_a8", [&] {
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      X.get_device_index());
+  auto Y = torch::stable::empty({batch, row}, X.scalar_type(), std::nullopt,
+                                W.device());
+  cudaStream_t stream = get_current_cuda_stream();
+  auto quant_X = torch::stable::empty({batch, padded / 32 * 9},
+                                      torch::headeronly::ScalarType::Int,
+                                      std::nullopt, W.device());
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_mul_mat_a8", [&] {
     quantize_row_q8_1_cuda((scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(),
                            col, batch, stream);
 
@@ -276,21 +287,24 @@ torch::Tensor ggml_mul_mat_a8(torch::Tensor W,  // quant weight
   return Y;
 }
 
-torch::Tensor ggml_moe_a8(torch::Tensor X,  // input
-                          torch::Tensor W,  // expert weights
-                          torch::Tensor sorted_token_ids,
-                          torch::Tensor expert_ids,
-                          torch::Tensor num_tokens_post_padded, int64_t type,
-                          int64_t row, int64_t top_k, int64_t tokens) {
+torch::stable::Tensor ggml_moe_a8(torch::stable::Tensor X,  // input
+                                  torch::stable::Tensor W,  // expert weights
+                                  torch::stable::Tensor sorted_token_ids,
+                                  torch::stable::Tensor expert_ids,
+                                  torch::stable::Tensor num_tokens_post_padded,
+                                  int64_t type, int64_t row, int64_t top_k,
+                                  int64_t tokens) {
   int col = X.sizes()[1];
   int padded = (col + 512 - 1) / 512 * 512;
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
-  auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device());
-  at::Tensor Y = torch::empty({tokens * top_k, row}, options);
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
-  options = torch::TensorOptions().dtype(torch::kInt32).device(W.device());
-  at::Tensor quant_X = torch::empty({tokens, padded / 32 * 9}, options);
-  VLLM_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_moe_a8", [&] {
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      X.get_device_index());
+  auto Y = torch::stable::empty({tokens * top_k, row}, X.scalar_type(),
+                                std::nullopt, W.device());
+  cudaStream_t stream = get_current_cuda_stream();
+  auto quant_X = torch::stable::empty({tokens, padded / 32 * 9},
+                                      torch::headeronly::ScalarType::Int,
+                                      std::nullopt, W.device());
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_moe_a8", [&] {
     quantize_row_q8_1_cuda((scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(),
                            col, tokens, stream);
     switch (type) {
@@ -379,19 +393,23 @@ torch::Tensor ggml_moe_a8(torch::Tensor X,  // input
   return Y;
 }
 
-torch::Tensor ggml_moe_a8_vec(torch::Tensor X,  // input
-                              torch::Tensor W,  // expert weights
-                              torch::Tensor topk_ids, int64_t top_k,
-                              int64_t type, int64_t row, int64_t tokens) {
+torch::stable::Tensor ggml_moe_a8_vec(
+    torch::stable::Tensor X,  // input
+    torch::stable::Tensor W,  // expert weights
+    torch::stable::Tensor topk_ids, int64_t top_k, int64_t type, int64_t row,
+    int64_t tokens) {
   int col = X.sizes()[1];
   const int padded = (col + 512 - 1) / 512 * 512;
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
-  auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device());
-  at::Tensor Y = torch::zeros({tokens * top_k, row}, options);
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
-  options = torch::TensorOptions().dtype(torch::kInt32).device(W.device());
-  at::Tensor quant_X = torch::empty({tokens, padded / 32 * 9}, options);
-  VLLM_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_moe_vec_a8", [&] {
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      X.get_device_index());
+  auto Y = torch::stable::empty({tokens * top_k, row}, X.scalar_type(),
+                                std::nullopt, W.device());
+  torch::stable::fill_(Y, 0.0);
+  cudaStream_t stream = get_current_cuda_stream();
+  auto quant_X = torch::stable::empty({tokens, padded / 32 * 9},
+                                      torch::headeronly::ScalarType::Int,
+                                      std::nullopt, W.device());
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_moe_vec_a8", [&] {
     quantize_row_q8_1_cuda<scalar_t>((scalar_t*)X.data_ptr(),
                                      (void*)quant_X.data_ptr(), col, tokens,
                                      stream);

csrc/libtorch_stable/quantization/gptq/q_gemm.cu

@@ -6,9 +6,8 @@ https://github.com/qwopqwop200/GPTQ-for-LLaMa
 #include <cstdint>
 #include <cstdio>
 
-#include <torch/all.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <ATen/cuda/CUDAContext.h>
+#include "../../torch_utils.h"
+#include <torch/csrc/stable/ops.h>
 #include <cuda_runtime.h>
 #include <cuda_fp16.h>
 
@@ -735,7 +734,7 @@ void gemm_half_q_half_cuda_part(const half* a, const uint32_t* b_q_weight,
   fp_gemm_half_q_half_gptq_kernel kernel =
       pick_gemm_half_q_half_gptq_kernel(true, m_count, bit);
 
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream();
   kernel<<<gridDim, blockDim, 0, stream>>>(
       a, b_q_weight, b_gptq_qzeros, b_gptq_scales, c, size_m, size_n, size_k,
       groups, use_v2_format, b_q_perm);
@@ -1164,7 +1163,7 @@ void reconstruct_exllama(const uint32_t* b_q_weight,
     reconstruct_exllama_kernel = reconstruct_exllama_8bit_kernel;
   }
 
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream();
   reconstruct_exllama_kernel<<<gridDim, blockDim, 0, stream>>>(
       b_q_weight, b_q_perm, b_gptq_qzeros, b_gptq_scales, height, width, groups,
       use_v2_format, out);
@@ -1376,7 +1375,7 @@ void gemm_half_q_half_alt(const half* a, const uint32_t* b_q_weight,
     kernel = gemm_half_q_half_alt_8bit_kernel;
   }
 
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream();
   kernel<<<gridDim, blockDim, 0, stream>>>(
       (const half2*)a, b_q_weight, c, b_gptq_scales, b_gptq_qzeros, b_g_idx,
       size_m, size_k / 32 * bit, size_n, use_v2_format);
@@ -1485,7 +1484,7 @@ void reconstruct_gptq(const uint32_t* b_q_weight, const uint32_t* b_gptq_qzeros,
     gridDim.y = DIVIDE(height, 32);
   }
 
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream();
   kernel<<<gridDim, blockDim, 0, stream>>>(b_q_weight, b_gptq_scales,
                                            b_gptq_qzeros, b_g_idx, height,
                                            width, groups, use_v2_format, out);
@@ -1794,7 +1793,7 @@ void shuffle_exllama_weight(uint32_t* q_weight, int* q_perm, int height,
     } else if (bit == 8) {
       kernel = make_sequential_8bit_kernel;
     }
-    const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    const cudaStream_t stream = get_current_cuda_stream();
     kernel<<<gridDim, blockDim, 0, stream>>>(q_weight, new_qweight, q_perm,
                                              width);
     // Replace qweights
@@ -1818,29 +1817,34 @@ void shuffle_exllama_weight(uint32_t* q_weight, int* q_perm, int height,
   } else if (bit == 8) {
     shuffle_kernel = shuffle_8bit_kernel;
   }
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream();
   shuffle_kernel<<<gridDim, blockDim, 0, stream>>>(q_weight, height, width);
 }
 
 }  // namespace gptq
 }  // namespace vllm
 
-torch::Tensor gptq_gemm(torch::Tensor a, torch::Tensor b_q_weight,
-                        torch::Tensor b_gptq_qzeros,
-                        torch::Tensor b_gptq_scales, torch::Tensor b_g_idx,
-                        bool use_exllama, bool use_v2_format, int64_t bit) {
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
-  auto options = torch::TensorOptions().dtype(a.dtype()).device(a.device());
-  at::Tensor c = torch::zeros({a.size(0), b_q_weight.size(1)}, options);
-  at::Tensor temp_dq = torch::empty(
-      {b_q_weight.size(0) * 32 / bit, b_q_weight.size(1)}, options);
+torch::stable::Tensor gptq_gemm(torch::stable::Tensor a,
+                                torch::stable::Tensor b_q_weight,
+                                torch::stable::Tensor b_gptq_qzeros,
+                                torch::stable::Tensor b_gptq_scales,
+                                torch::stable::Tensor b_g_idx, bool use_exllama,
+                                bool use_v2_format, int64_t bit) {
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      a.get_device_index());
+  auto c = torch::stable::new_zeros(a, {a.size(0), b_q_weight.size(1)});
+  auto temp_dq =
+      torch::stable::empty({b_q_weight.size(0) * 32 / bit, b_q_weight.size(1)},
+                           a.scalar_type(), std::nullopt, a.device());
 
   vllm::gptq::gemm_half_q_half_cuda(
-      at::cuda::getCurrentCUDABlasHandle(), (const half*)a.data_ptr(),
+      get_current_cuda_blas_handle(), (const half*)a.data_ptr(),
       (const uint32_t*)b_q_weight.data_ptr(),
       (const uint32_t*)b_gptq_qzeros.data_ptr(),
       (const half*)b_gptq_scales.data_ptr(),
-      b_g_idx.device().is_meta() ? NULL : (const int*)b_g_idx.data_ptr(),
+      b_g_idx.device().type() == torch::stable::DeviceType::Meta
+          ? NULL
+          : (const int*)b_g_idx.data_ptr(),
       (half*)c.data_ptr(), (half*)temp_dq.data_ptr(),
       c.size(0),              // m
       c.size(1),              // n
@@ -1850,11 +1854,14 @@ torch::Tensor gptq_gemm(torch::Tensor a, torch::Tensor b_q_weight,
   return c;
 }
 
-void gptq_shuffle(torch::Tensor q_weight, torch::Tensor q_perm, int64_t bit) {
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(q_weight));
+void gptq_shuffle(torch::stable::Tensor q_weight, torch::stable::Tensor q_perm,
+                  int64_t bit) {
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      q_weight.get_device_index());
   vllm::gptq::shuffle_exllama_weight(
       (uint32_t*)q_weight.data_ptr(),
-      q_perm.device().is_meta() || q_perm.numel() == 0
+      q_perm.device().type() == torch::stable::DeviceType::Meta ||
+              q_perm.numel() == 0
           ? NULL
           : (int*)q_perm.data_ptr(),
       q_weight.size(0) * 32 / bit, q_weight.size(1), bit);