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

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

@@ -6,14 +6,26 @@ steps:
     timeout_in_minutes: 600
     commands:
     - if [[ "$BUILDKITE_BRANCH" == "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $IMAGE_TAG $IMAGE_TAG_LATEST; else .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $IMAGE_TAG; fi
-    # Non-root smoke 1: the default (root) image must still be importable
+    retry:
+      automatic:
+        - exit_status: -1  # Agent was lost
+          limit: 2
+        - exit_status: -10  # Agent was lost
+          limit: 2
+
+  - label: ":docker: :smoking: Non-root smoke tests"
+    key: image-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__)"
-    # Non-root smoke 2: assert the non-root enabling invariants are baked
+    # Smoke 2: assert the non-root enabling invariants are baked
     # into the image. Runs as UID 2000:0 via a shell so we can verify
     # filesystem perms + passwd/group file state + wrapper presence without
     # triggering vLLM's GPU-requiring config-init path. The opt-in
@@ -98,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/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/coverage/aggregate-coverage.py

@@ -1,208 +0,0 @@
-#!/usr/bin/env python3
-"""Aggregate per-step coverage JSON files into a test-selection mapping.
-
-Downloads all coverage_*.json artifacts from the current Buildkite build,
-then produces two output files:
-
-1. coverage_map.json — inverted index: {source_file: [step_keys]}
-   Used by the pipeline generator to determine which steps to trigger.
-
-2. step_coverage.json — forward index: {step_key: [source_files]}
-   Useful for debugging and understanding test coverage.
-
-Usage:
-    # Run as a Buildkite step at the end of nightly CI
-    python3 .buildkite/scripts/coverage/aggregate-coverage.py
-
-    # Or locally with downloaded artifacts
-    python3 .buildkite/scripts/coverage/aggregate-coverage.py --local-dir ./artifacts/
-"""
-
-import argparse
-import json
-import os
-import subprocess
-import sys
-import tempfile
-from collections import defaultdict
-from pathlib import Path
-
-
-def download_artifacts(dest_dir: str) -> list[str]:
-    """Download all coverage_*.json artifacts from the current build."""
-    try:
-        subprocess.run(
-            ["buildkite-agent", "artifact", "download", "coverage_*.json", dest_dir],
-            check=True,
-            capture_output=True,
-            text=True,
-        )
-    except FileNotFoundError:
-        print("buildkite-agent not found, skipping download", file=sys.stderr)
-        return []
-    except subprocess.CalledProcessError as e:
-        print(f"Artifact download failed: {e.stderr}", file=sys.stderr)
-        return []
-
-    return list(Path(dest_dir).glob("coverage_*.json"))
-
-
-def load_coverage_files(files: list[Path]) -> dict[str, list[str]]:
-    """Load coverage JSON files and extract source files per step.
-
-    Returns: {step_key: [source_files]}
-    """
-    step_coverage = {}
-
-    for filepath in files:
-        filename = filepath.name
-        # coverage_<step_key>.json -> step_key
-        step_key = filename.removeprefix("coverage_").removesuffix(".json")
-
-        try:
-            with open(filepath) as f:
-                data = json.load(f)
-        except (json.JSONDecodeError, OSError) as e:
-            print(f"Warning: skipping {filename}: {e}", file=sys.stderr)
-            continue
-
-        source_files = []
-        for fpath, fdata in data.get("files", {}).items():
-            # Skip files with zero executed lines — coverage.py reports
-            # all files in the source tree, not just those actually run.
-            # Supports both full format (summary.covered_lines) and
-            # stripped format (covered_lines directly).
-            covered = fdata.get("covered_lines") or fdata.get("summary", {}).get("covered_lines", 0)
-            if covered == 0:
-                continue
-            # If function-level data is available, skip import-only files
-            # (files where only module-level code ran but no named functions
-            # were actually called).
-            funcs_called = fdata.get("functions_called")
-            if funcs_called is not None and funcs_called == 0:
-                continue
-            # Normalize paths to be relative to the vllm package root.
-            # coverage.py may report absolute paths or paths relative to
-            # the installed package location. We only care about files
-            # under the vllm/ directory.
-            normalized = _normalize_path(fpath)
-            if normalized:
-                source_files.append(normalized)
-
-        if source_files:
-            step_coverage[step_key] = sorted(set(source_files))
-            print(f"  {step_key}: {len(source_files)} source files")
-
-    return step_coverage
-
-
-def _normalize_path(path: str) -> str | None:
-    """Normalize a coverage path to a vllm-relative path.
-
-    Returns None for paths outside the vllm package (tests, third-party, etc).
-    """
-    # Strip common prefixes from installed package paths
-    markers = ["/site-packages/", "/dist-packages/", "/vllm-workspace/src/"]
-    for marker in markers:
-        idx = path.find(marker)
-        if idx != -1:
-            path = path[idx + len(marker):]
-            break
-
-    # Also handle paths that are already relative
-    if path.startswith("vllm/"):
-        return path
-
-    # Handle absolute paths that contain /vllm/
-    idx = path.find("/vllm/")
-    if idx != -1:
-        return path[idx + 1:]
-
-    return None
-
-
-def build_inverted_index(
-    step_coverage: dict[str, list[str]],
-) -> dict[str, list[str]]:
-    """Build {source_file: [step_keys]} from {step_key: [source_files]}."""
-    inverted = defaultdict(list)
-    for step_key, source_files in step_coverage.items():
-        for src_file in source_files:
-            inverted[src_file].append(step_key)
-
-    # Sort step lists for deterministic output
-    return {k: sorted(v) for k, v in sorted(inverted.items())}
-
-
-def main():
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument(
-        "--local-dir",
-        help="Directory containing coverage_*.json files (skip artifact download)",
-    )
-    parser.add_argument(
-        "--output-dir",
-        default=".",
-        help="Directory to write output files (default: cwd)",
-    )
-    args = parser.parse_args()
-
-    if args.local_dir:
-        artifact_dir = args.local_dir
-        files = list(Path(artifact_dir).glob("coverage_*.json"))
-    else:
-        artifact_dir = tempfile.mkdtemp(prefix="coverage_artifacts_")
-        files = download_artifacts(artifact_dir)
-
-    if not files:
-        print("No coverage files found. Nothing to aggregate.")
-        sys.exit(0)
-
-    print(f"Found {len(files)} coverage files:")
-
-    # Build the forward index: step -> source files
-    step_coverage = load_coverage_files(files)
-
-    if not step_coverage:
-        print("No valid coverage data found.")
-        sys.exit(0)
-
-    # Build the inverted index: source file -> steps
-    coverage_map = build_inverted_index(step_coverage)
-
-    # Write outputs
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    step_coverage_path = output_dir / "step_coverage.json"
-    with open(step_coverage_path, "w") as f:
-        json.dump(step_coverage, f, indent=2)
-    print(f"\nWrote {step_coverage_path} ({len(step_coverage)} steps)")
-
-    coverage_map_path = output_dir / "coverage_map.json"
-    with open(coverage_map_path, "w") as f:
-        json.dump(coverage_map, f, indent=2)
-    print(f"Wrote {coverage_map_path} ({len(coverage_map)} source files)")
-
-    # Summary stats
-    total_files = len(coverage_map)
-    total_mappings = sum(len(v) for v in coverage_map.values())
-    print(f"\nSummary: {total_files} source files mapped to "
-          f"{len(step_coverage)} steps ({total_mappings} total mappings)")
-
-    # Upload aggregated files as artifacts
-    for output_file in [step_coverage_path, coverage_map_path]:
-        try:
-            subprocess.run(
-                ["buildkite-agent", "artifact", "upload", str(output_file)],
-                check=True,
-                capture_output=True,
-                text=True,
-            )
-            print(f"Uploaded {output_file}")
-        except (FileNotFoundError, subprocess.CalledProcessError):
-            pass  # Not in Buildkite or upload failed — that's fine for local runs
-
-
-if __name__ == "__main__":
-    main()

.buildkite/scripts/coverage/upload-step-coverage.sh

@@ -1,42 +0,0 @@
-#!/bin/bash
-# Upload coverage data for the current Buildkite step.
-# Called automatically at the end of each step when COLLECT_COVERAGE=1.
-#
-# Expects:
-#   - .coverage.${BUILDKITE_STEP_KEY} data file from coverage run --append
-#   - BUILDKITE_STEP_KEY, BUILDKITE_BUILD_NUMBER env vars
-#
-# Produces:
-#   - coverage_${BUILDKITE_STEP_KEY}.json uploaded as a Buildkite artifact
-
-set -euo pipefail
-
-STEP_KEY="${BUILDKITE_STEP_KEY:-unknown}"
-DATA_FILE=".coverage.${STEP_KEY}"
-OUTPUT_JSON="coverage_${STEP_KEY}.json"
-
-if [ ! -f "$DATA_FILE" ]; then
-    echo "~~~ No coverage data file found ($DATA_FILE), skipping upload"
-    exit 0
-fi
-
-echo "~~~ :bar_chart: Exporting coverage data for step: ${STEP_KEY}"
-
-coverage json \
-    --data-file="$DATA_FILE" \
-    -o "$OUTPUT_JSON" \
-    --omit='*/tests/*,*/test_*,*/__pycache__/*' \
-    2>&1 || {
-        echo "Warning: coverage json export failed, skipping"
-        exit 0
-    }
-
-FILE_COUNT=$(python3 -c "import json; d=json.load(open('$OUTPUT_JSON')); print(len(d.get('files', {})))" 2>/dev/null || echo "?")
-echo "Coverage captured ${FILE_COUNT} source files for step ${STEP_KEY}"
-
-buildkite-agent artifact upload "$OUTPUT_JSON" 2>&1 || {
-    echo "Warning: artifact upload failed"
-    exit 0
-}
-
-echo "Uploaded $OUTPUT_JSON"

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

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

.buildkite/test_areas/e2e_integration.yaml

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

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

.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,7 +85,8 @@ 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
@@ -104,7 +108,8 @@ steps:
   - pytest -v -s tool_use
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
+      soft_fail: true
       depends_on:
       - image-build-amd
 

.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

.buildkite/test_areas/misc.yaml

@@ -52,7 +52,7 @@ steps:
     - pytest -v -s v1/test_outputs.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
 

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

@@ -58,37 +58,3 @@ steps:
   device: cpu-small
   commands:
     - pytest -v -s models/test_utils.py models/test_vision.py
-
-- label: Transformers Nightly Models
-  device: h200_35gb
-  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
-  device: h200_35gb
-  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:
@@ -96,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,7 +33,7 @@ steps:
     - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
 
@@ -50,7 +50,7 @@ steps:
     - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
 
@@ -118,7 +118,7 @@ steps:
     - pytest -v -s models/multimodal/test_mapping.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
 

.buildkite/test_areas/rust_frontend.yaml

@@ -16,7 +16,7 @@ steps:
   - 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/chat_completion/test_chat_with_tool_reasoning.py
+
   # - tests/entrypoints/openai/completion/test_prompt_validation.py
   - tests/entrypoints/openai/completion/test_shutdown.py
   # - tests/entrypoints/openai/test_return_token_ids.py
@@ -28,7 +28,7 @@ steps:
   - 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/chat_completion/test_chat_with_tool_reasoning.py
+
   # - 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

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

.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
@@ -171,20 +179,20 @@ mkdocs.yaml @hmellor
 
 # ROCm related: specify owner with write access to notify AMD folks for careful code review
 /vllm/**/*rocm* @tjtanaa @dllehr-amd
-/docker/Dockerfile.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
-/tests/**/*rocm* @tjtanaa
+/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:

.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,14 +305,6 @@ 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/sampler.cu"
-  "csrc/topk.cu"
   "csrc/cuda_view.cu"
   "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
   "csrc/quantization/activation_kernels.cu"
@@ -323,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")
 
@@ -369,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}")
 
@@ -631,9 +638,18 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu"
     "csrc/libtorch_stable/pos_encoding_kernels.cu"
     "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu"
+    "csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu"
     "csrc/libtorch_stable/layernorm_kernels.cu"
     "csrc/libtorch_stable/layernorm_quant_kernels.cu"
-    "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu")
+    "csrc/libtorch_stable/quantization/fused_kernels/fused_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
@@ -669,6 +685,22 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
                    "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)
@@ -910,13 +942,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
       SRCS "${SRCS}"
       CUDA_ARCHS "${FP4_ARCHS}")
     list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
-    # nvfp4_kv_cache_kernels uses non-stable torch API and is called directly
-    # from cache_kernels.cu, so it belongs in _C rather than _C_stable.
-    set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
+    set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
     set_gencode_flags_for_srcs(
       SRCS "${NVFP4_KV_SRC}"
       CUDA_ARCHS "${FP4_ARCHS}")
-    target_sources(_C PRIVATE ${NVFP4_KV_SRC})
+    list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}")
     target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
@@ -946,11 +976,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
       SRCS "${SRCS}"
       CUDA_ARCHS "${FP4_ARCHS}")
     list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
-    set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
+    set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
     set_gencode_flags_for_srcs(
       SRCS "${NVFP4_KV_SRC}"
       CUDA_ARCHS "${FP4_ARCHS}")
-    target_sources(_C PRIVATE ${NVFP4_KV_SRC})
+    list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}")
     target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
@@ -1121,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
@@ -1224,24 +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}")
+
+    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()
@@ -1268,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
@@ -1277,6 +1317,10 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
     ARCHITECTURES ${VLLM_GPU_ARCHES}
     USE_SABI 3
     WITH_SOABI)
+
+  if(VLLM_ROCM_HAS_GFX1100)
+    target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GFX1100)
+  endif()
 endif()
 
 # Must run after the last HIP `define_extension_target` so every extension

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_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)

cmake/utils.cmake

@@ -476,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

csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu

@@ -122,18 +122,45 @@ __device__ __forceinline__ float warpSum(float val) {
 // Per-slot inner pipeline
 // ────────────────────────────────────────────────────────────────────────────
 // Shared by both kernel variants: 1 CTA per (token, head) pair vs. 1 CTA per
-// token
-template <typename scalar_t_in>
+// 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_inout,
+    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 == num_heads_q);
+  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];
@@ -207,7 +234,7 @@ __device__ __forceinline__ void processDeepseekV4Slot(
   // triggering and per-iteration buffer rotation.
   // ═══════════════════════════════════════════════════════════════════
   if (!isKV) {
-    // ── Q: cast back to bf16 and store. ────────────────────────────
+    // ── 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);
@@ -224,8 +251,9 @@ __device__ __forceinline__ void processDeepseekV4Slot(
           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 +
+        q_out +
+        (static_cast<int64_t>(tokenIdx) * kNumHeadsQPadded + slotIdx) *
+            kHeadDim +
         dim_base;
     *reinterpret_cast<uint4*>(dst) = out0;
     *reinterpret_cast<uint4*>(dst + 8) = out1;
@@ -313,20 +341,32 @@ __device__ __forceinline__ void processDeepseekV4Slot(
 // 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
@@ -335,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)
@@ -351,12 +391,13 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
     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;
 
@@ -371,22 +412,26 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
     // Dim range this lane owns within the 512-wide head.
     int const dim_base = laneId * kElemsPerLane;  // in [0, 512) step 16
 
-    // Two 16-byte loads per thread (8 bf16 each).  Use uint4 as the vector
-    // type; the shared per-slot helper bitcasts to scalar_t_in packed pairs.
-    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);
     }
-    uint4 const v0 = *reinterpret_cast<uint4 const*>(src_ptr);
-    uint4 const v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
 
-    processDeepseekV4Slot<scalar_t_in>(
-        v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_inout,
+    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);
 
@@ -408,9 +453,9 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
 // Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q
 // KV branch (RoPE + UE8M0 quant + insert)
 //
-template <typename scalar_t_in>
+template <typename scalar_t_in, int kNumHeadsQPadded>
 __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
-    scalar_t_in* __restrict__ q_inout,  // [N, H, 512] bf16, in place
+    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,
@@ -435,25 +480,32 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
 #endif
 
     int const dim_base = laneId * kElemsPerLane;  // in [0, 512) step 16
-    int const slot_end =
-        (tokenIdx >= num_tokens_insert) ? num_heads_q : (num_heads_q + 1);
+    // Slot enumeration: live-Q + pad-Q + (KV if this token has a slot).
+    int const slot_end = (tokenIdx >= num_tokens_insert)
+                             ? kNumHeadsQPadded
+                             : (kNumHeadsQPadded + 1);
 
-    auto src_for_slot = [&](int s) -> scalar_t_in const* {
-      if (s == num_heads_q) {
-        return kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
+    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;
       }
-      return q_inout +
-             (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);
     };
 
     if (warpId < slot_end) {
       int curr_slot = warpId;
-      scalar_t_in const* src_curr = src_for_slot(curr_slot);
-      uint4 v0_curr = *reinterpret_cast<uint4 const*>(src_curr);
-      uint4 v1_curr = *reinterpret_cast<uint4 const*>(src_curr + 8);
+      uint4 v0_curr, v1_curr;
+      load_slot(curr_slot, v0_curr, v1_curr);
 
       while (curr_slot < slot_end) {
         int const next_slot = curr_slot + warpsPerBlock;
@@ -462,14 +514,12 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
         // Prefetch src for the next slot
         uint4 v0_next, v1_next;
         if (has_next) {
-          scalar_t_in const* src_next = src_for_slot(next_slot);
-          v0_next = *reinterpret_cast<uint4 const*>(src_next);
-          v1_next = *reinterpret_cast<uint4 const*>(src_next + 8);
+          load_slot(next_slot, v0_next, v1_next);
         }
 
-        processDeepseekV4Slot<scalar_t_in>(
+        processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded>(
             v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId,
-            num_heads_q, eps, q_inout, k_cache, slot_mapping, position_ids,
+            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.
@@ -490,10 +540,10 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
 // ────────────────────────────────────────────────────────────────────────────
 // 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,
@@ -501,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);
 
@@ -532,46 +582,85 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
 
   if (num_tokens_full < NUM_TOKEN_CUTOFF) {
     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,
+        &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>,
-        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,
+        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,
@@ -579,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]");
@@ -591,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

@@ -4,17 +4,26 @@
 #include <cmath>
 
 #include "../cuda_compat.h"
-#include "../cuda_vec_utils.cuh"
+#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);                         \
@@ -230,7 +259,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
             ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>(     \
             out.mutable_data_ptr<scalar_t>(),                                  \
-            input.const_data_ptr<scalar_t>(), d, LIMIT);                       \
+            input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA);          \
       });                                                                      \
     } else {                                                                   \
       VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {   \
@@ -240,7 +269,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
             ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>(    \
             out.mutable_data_ptr<scalar_t>(),                                  \
-            input.const_data_ptr<scalar_t>(), d, LIMIT);                       \
+            input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA);          \
       });                                                                      \
     }                                                                          \
   } else {                                                                     \
@@ -252,7 +281,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
           PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
           ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>(            \
           out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),  \
-          d, LIMIT);                                                           \
+          d, LIMIT, ALPHA, BETA);                                              \
     });                                                                        \
   }
 
@@ -260,14 +289,18 @@ void silu_and_mul(torch::stable::Tensor& out,    // [..., d]
                   torch::stable::Tensor& input)  // [..., 2 * d]
 {
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true, false, 0.0f);
+                                true, false, 0.0f, 1.0f, 0.0f);
 }
 
 void silu_and_mul_clamp(torch::stable::Tensor& out,    // [..., d]
                         torch::stable::Tensor& input,  // [..., 2 * d]
-                        double limit) {
+                        double limit, double alpha, double beta) {
+  // out = (gate.clamp(max=limit) * sigmoid(alpha * gate.clamp(max=limit)))
+  //       * (up.clamp(+-limit) + beta)
+  // alpha=1.0, beta=0.0 reduce this to silu(gate) * up.
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true, true, (float)limit);
+                                true, true, (float)limit, (float)alpha,
+                                (float)beta);
 }
 
 void mul_and_silu(torch::stable::Tensor& out,    // [..., d]
@@ -276,21 +309,22 @@ void mul_and_silu(torch::stable::Tensor& out,    // [..., d]
   // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
   // applies the silu to the latter half of the input.
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                false, false, 0.0f);
+                                false, false, 0.0f, 1.0f, 0.0f);
 }
 
 void gelu_and_mul(torch::stable::Tensor& out,    // [..., d]
                   torch::stable::Tensor& input)  // [..., 2 * d]
 {
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
-                                true, false, 0.0f);
+                                true, false, 0.0f, 1.0f, 0.0f);
 }
 
 void gelu_tanh_and_mul(torch::stable::Tensor& out,    // [..., d]
                        torch::stable::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(
-      vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
+                                vllm::packed_gelu_tanh_kernel, true, false,
+                                0.0f, 1.0f, 0.0f);
 }
 
 namespace vllm {

csrc/libtorch_stable/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/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/layernorm_kernels.cu

@@ -78,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;
   }
@@ -143,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;
   }
@@ -183,8 +181,8 @@ 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];
   }
 }
 

csrc/libtorch_stable/layernorm_quant_kernels.cu

@@ -66,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);
@@ -142,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);
     }
@@ -192,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);
   }

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

csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu

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

csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu

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

csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu

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

csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu

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

csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh

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

csrc/libtorch_stable/sampler.cu

@@ -1,8 +1,6 @@
-#include "cuda_compat.h"
+#include "../cuda_compat.h"
 #include "dispatch_utils.h"
-
-#include <torch/cuda.h>
-#include <c10/cuda/CUDAGuard.h>
+#include "torch_utils.h"
 
 #ifndef USE_ROCM
   #include <cub/cub.cuh>
@@ -618,14 +616,14 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(
 }  // namespace vllm
 
 void apply_repetition_penalties_(
-    torch::Tensor& logits,             // [num_seqs, vocab_size], in-place
-    const torch::Tensor& prompt_mask,  // [num_seqs, vocab_size]
-    const torch::Tensor& output_mask,  // [num_seqs, vocab_size]
-    const torch::Tensor& repetition_penalties) {  // [num_seqs]
-  TORCH_CHECK(logits.is_contiguous());
-  TORCH_CHECK(prompt_mask.is_contiguous());
-  TORCH_CHECK(output_mask.is_contiguous());
-  TORCH_CHECK(repetition_penalties.is_contiguous());
+    torch::stable::Tensor& logits,  // [num_seqs, vocab_size], in-place
+    const torch::stable::Tensor& prompt_mask,  // [num_seqs, vocab_size]
+    const torch::stable::Tensor& output_mask,  // [num_seqs, vocab_size]
+    const torch::stable::Tensor& repetition_penalties) {  // [num_seqs]
+  STD_TORCH_CHECK(logits.is_contiguous());
+  STD_TORCH_CHECK(prompt_mask.is_contiguous());
+  STD_TORCH_CHECK(output_mask.is_contiguous());
+  STD_TORCH_CHECK(repetition_penalties.is_contiguous());
 
   int vocab_size = logits.size(-1);
   int num_seqs = logits.size(0);
@@ -635,7 +633,7 @@ void apply_repetition_penalties_(
   // Get number of SMs on the current device
   int sms = 0;
   cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount,
-                         logits.get_device());
+                         logits.get_device_index());
 
   // Compute tile_num and tile_size
   int tile_num =
@@ -645,27 +643,29 @@ void apply_repetition_penalties_(
   // Each block handles one sequence and a tile of vocab
   dim3 grid(num_seqs, tile_num);
   dim3 block(std::min(tile_size, 1024));
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(logits));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  VLLM_DISPATCH_FLOATING_TYPES(
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      logits.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
       logits.scalar_type(), "apply_repetition_penalties_kernel", [&] {
         vllm::apply_repetition_penalties_kernel<scalar_t>
             <<<grid, block, 0, stream>>>(
-                logits.data_ptr<scalar_t>(), prompt_mask.data_ptr<bool>(),
-                output_mask.data_ptr<bool>(),
-                repetition_penalties.data_ptr<scalar_t>(), num_seqs, vocab_size,
-                tile_size);
+                logits.mutable_data_ptr<scalar_t>(),
+                prompt_mask.const_data_ptr<bool>(),
+                output_mask.const_data_ptr<bool>(),
+                repetition_penalties.const_data_ptr<scalar_t>(), num_seqs,
+                vocab_size, tile_size);
       });
 }
 
-void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,
-                          const torch::Tensor& seqLens, torch::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) {
   constexpr int kSortingAlgorithmThreshold = 12288;
   constexpr int kSplitWorkThreshold = 200 * 1000;
   constexpr int kNumThreadsPerBlock = 512;
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream();
   const auto numColumns = logits.size(1);
 
   // True if seqLens is 2D (B, next_n): each logit row has its own pre-computed
@@ -677,73 +677,76 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,
     // Use insertion sort
     vllm::topKPerRowDecode<kNumThreadsPerBlock, false>
         <<<numRows, kNumThreadsPerBlock, topK * sizeof(int32_t), stream>>>(
-            logits.data_ptr<float>(), seqLens.data_ptr<int>(),
-            indices.data_ptr<int>(), static_cast<int>(stride0),
+            logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
+            indices.mutable_data_ptr<int>(), static_cast<int>(stride0),
             static_cast<int>(stride1), static_cast<int>(topK),
             static_cast<int>(next_n), seqLensIs2D);
   } else if (numColumns < kSplitWorkThreshold) {
     // From this threshold, use radix sort instead
     vllm::topKPerRowDecode<kNumThreadsPerBlock, true>
         <<<numRows, kNumThreadsPerBlock, topK * sizeof(int32_t), stream>>>(
-            logits.data_ptr<float>(), seqLens.data_ptr<int>(),
-            indices.data_ptr<int>(), static_cast<int>(stride0),
+            logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
+            indices.mutable_data_ptr<int>(), static_cast<int>(stride0),
             static_cast<int>(stride1), static_cast<int>(topK),
             static_cast<int>(next_n), seqLensIs2D);
   } else {
     // Long sequences are run in two steps
     constexpr auto multipleBlocksPerRowConfig = 10;
 
-    const auto outIndicesAux =
-        torch::empty({numRows, multipleBlocksPerRowConfig, topK},
-                     torch::dtype(torch::kInt32).device(logits.device()));
-    const auto outLogitsAux =
-        torch::empty({numRows, multipleBlocksPerRowConfig, topK},
-                     torch::dtype(torch::kFloat).device(logits.device()));
+    const auto outIndicesAux = torch::stable::empty(
+        {numRows, multipleBlocksPerRowConfig, topK},
+        torch::headeronly::ScalarType::Int, std::nullopt, logits.device());
+    const auto outLogitsAux = torch::stable::empty(
+        {numRows, multipleBlocksPerRowConfig, topK},
+        torch::headeronly::ScalarType::Float, std::nullopt, logits.device());
 
     vllm::topKPerRowDecode<kNumThreadsPerBlock, true, true>
         <<<dim3(numRows, multipleBlocksPerRowConfig), kNumThreadsPerBlock,
            2 * topK * sizeof(int32_t), stream>>>(
-            logits.data_ptr<float>(), seqLens.data_ptr<int>(),
-            outIndicesAux.data_ptr<int>(), static_cast<int>(stride0),
+            logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
+            outIndicesAux.mutable_data_ptr<int>(), static_cast<int>(stride0),
             static_cast<int>(stride1), static_cast<int>(topK),
             static_cast<int>(next_n), seqLensIs2D,
-            outLogitsAux.data_ptr<float>());
+            outLogitsAux.mutable_data_ptr<float>());
 
     constexpr int kNumThreadsPerBlockMerge = 1024;
     vllm::topKPerRowDecode<kNumThreadsPerBlockMerge, true, false, true>
         <<<numRows, kNumThreadsPerBlockMerge, topK * sizeof(int32_t), stream>>>(
-            outLogitsAux.data_ptr<float>(), seqLens.data_ptr<int>(),
-            indices.data_ptr<int>(), multipleBlocksPerRowConfig * topK, 1,
-            static_cast<int>(topK), static_cast<int>(next_n), seqLensIs2D,
-            nullptr, multipleBlocksPerRowConfig, outIndicesAux.data_ptr<int>());
+            outLogitsAux.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
+            indices.mutable_data_ptr<int>(), multipleBlocksPerRowConfig * topK,
+            1, static_cast<int>(topK), static_cast<int>(next_n), seqLensIs2D,
+            nullptr, multipleBlocksPerRowConfig,
+            outIndicesAux.const_data_ptr<int>());
   }
 }
 
-void top_k_per_row_prefill(const torch::Tensor& logits,
-                           const torch::Tensor& rowStarts,
-                           const torch::Tensor& rowEnds, torch::Tensor& indices,
-                           int64_t numRows, int64_t stride0, int64_t stride1,
-                           int64_t topK) {
+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) {
   constexpr int kSortingAlgorithmThreshold = 12288;
   constexpr int kNumThreadsPerBlock = 512;
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream();
 
   int numInsertionBlocks =
       std::min(static_cast<int>(numRows), kSortingAlgorithmThreshold);
   vllm::topKPerRowPrefill<kNumThreadsPerBlock, false>
       <<<numInsertionBlocks, kNumThreadsPerBlock, topK * sizeof(int32_t),
-         stream>>>(logits.data_ptr<float>(), rowStarts.data_ptr<int>(),
-                   rowEnds.data_ptr<int>(), indices.data_ptr<int>(),
-                   static_cast<int>(stride0), static_cast<int>(stride1),
-                   static_cast<int>(topK), 0);
+         stream>>>(logits.const_data_ptr<float>(),
+                   rowStarts.const_data_ptr<int>(),
+                   rowEnds.const_data_ptr<int>(),
+                   indices.mutable_data_ptr<int>(), static_cast<int>(stride0),
+                   static_cast<int>(stride1), static_cast<int>(topK), 0);
 
   if (numRows > kSortingAlgorithmThreshold) {
     int numRadixBlocks = numRows - kSortingAlgorithmThreshold;
     vllm::topKPerRowPrefill<kNumThreadsPerBlock, true>
         <<<numRadixBlocks, kNumThreadsPerBlock, topK * sizeof(int32_t),
-           stream>>>(logits.data_ptr<float>(), rowStarts.data_ptr<int>(),
-                     rowEnds.data_ptr<int>(), indices.data_ptr<int>(),
-                     static_cast<int>(stride0), static_cast<int>(stride1),
-                     static_cast<int>(topK), kSortingAlgorithmThreshold);
+           stream>>>(
+            logits.const_data_ptr<float>(), rowStarts.const_data_ptr<int>(),
+            rowEnds.const_data_ptr<int>(), indices.mutable_data_ptr<int>(),
+            static_cast<int>(stride0), static_cast<int>(stride1),
+            static_cast<int>(topK), kSortingAlgorithmThreshold);
   }
 }

csrc/libtorch_stable/topk.cu

@@ -1,49 +1,51 @@
 // Persistent TopK kernel for DeepSeek V3 sparse attention indexer.
 // See persistent_topk.cuh for kernel implementation.
 
-#include <torch/all.h>
-#include <ATen/cuda/CUDAContext.h>
 #include <cuda_runtime.h>
 #include <algorithm>
 
+#include "torch_utils.h"
+
 #ifndef USE_ROCM
-  #include "persistent_topk.cuh"
+  #include "../persistent_topk.cuh"
 #endif
 
 namespace {
 
 #ifndef USE_ROCM
 template <int TopK>
-void launch_persistent_topk(const torch::Tensor& logits,
-                            const torch::Tensor& lengths, torch::Tensor& output,
-                            torch::Tensor& workspace, int64_t max_seq_len) {
+void launch_persistent_topk(const torch::stable::Tensor& logits,
+                            const torch::stable::Tensor& lengths,
+                            torch::stable::Tensor& output,
+                            torch::stable::Tensor& workspace,
+                            int64_t max_seq_len) {
   namespace P = vllm::persistent;
 
   const int64_t num_rows = logits.size(0);
   const int64_t stride = logits.stride(0);
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream();
 
   static int num_sms = 0;
   static int max_smem_per_block = 0;
   if (num_sms == 0) {
-    int device;
-    cudaGetDevice(&device);
-    cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, device);
-    cudaDeviceGetAttribute(&max_smem_per_block,
-                           cudaDevAttrMaxSharedMemoryPerBlockOptin, device);
+    const cudaDeviceProp* device_prop = get_device_prop();
+    num_sms = device_prop->multiProcessorCount;
+    max_smem_per_block = device_prop->sharedMemPerBlockOptin;
   }
 
   if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
     cudaError_t status =
         vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
-            logits.data_ptr<float>(), output.data_ptr<int32_t>(),
-            lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
+            logits.const_data_ptr<float>(), output.mutable_data_ptr<int32_t>(),
+            lengths.const_data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
             static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride), stream);
-    TORCH_CHECK(status == cudaSuccess,
-                "FilteredTopK failed: ", cudaGetErrorString(status));
+    STD_TORCH_CHECK(status == cudaSuccess,
+                    "FilteredTopK failed: ", cudaGetErrorString(status));
   } else {
-    TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
-    TORCH_CHECK(workspace.dtype() == torch::kUInt8, "workspace must be uint8");
+    STD_TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
+    STD_TORCH_CHECK(
+        workspace.scalar_type() == torch::headeronly::ScalarType::Byte,
+        "workspace must be uint8");
 
     int effective_max_smem;
     if (num_rows <= 4) {
@@ -99,9 +101,9 @@ void launch_persistent_topk(const torch::Tensor& logits,
           &occupancy, P::persistent_topk_kernel<TopK, 1>, P::kThreadsPerBlock,
           smem_size);
     }
-    TORCH_CHECK(occ_err == cudaSuccess,
-                "persistent_topk occupancy query failed: ",
-                cudaGetErrorString(occ_err));
+    STD_TORCH_CHECK(occ_err == cudaSuccess,
+                    "persistent_topk occupancy query failed: ",
+                    cudaGetErrorString(occ_err));
     if (occupancy < 1) occupancy = 1;
 
     // The cooperative spin-wait barrier only runs when at least one row hits
@@ -131,27 +133,29 @@ void launch_persistent_topk(const torch::Tensor& logits,
     // If the cooperative launch wouldn't fit, fall back to FilteredTopK
     // instead of deadlocking. Only relevant when needs_cooperative.
     if (needs_cooperative && total_ctas > hw_resident_cap) {
-      TORCH_CHECK(max_smem_per_block >= 128 * 1024,
-                  "persistent_topk would oversubscribe and the FilteredTopK "
-                  "fallback requires >=128KB smem per block (have ",
-                  max_smem_per_block, "). total_ctas=", total_ctas,
-                  " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
-                  ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
-                  ", smem=", smem_size, ").");
+      STD_TORCH_CHECK(
+          max_smem_per_block >= 128 * 1024,
+          "persistent_topk would oversubscribe and the FilteredTopK "
+          "fallback requires >=128KB smem per block (have ",
+          max_smem_per_block, "). total_ctas=", total_ctas,
+          " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
+          ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
+          ", smem=", smem_size, ").");
       cudaError_t status =
           vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
-              logits.data_ptr<float>(), output.data_ptr<int32_t>(),
-              lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
-              static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride),
-              stream);
-      TORCH_CHECK(status == cudaSuccess,
-                  "FilteredTopK fallback failed: ", cudaGetErrorString(status));
+              logits.const_data_ptr<float>(),
+              output.mutable_data_ptr<int32_t>(),
+              lengths.const_data_ptr<int32_t>(),
+              static_cast<uint32_t>(num_rows), static_cast<uint32_t>(TopK),
+              static_cast<uint32_t>(stride), stream);
+      STD_TORCH_CHECK(status == cudaSuccess, "FilteredTopK fallback failed: ",
+                      cudaGetErrorString(status));
       return;
     }
 
     size_t state_bytes = num_groups * sizeof(P::RadixRowState);
-    TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
-                "workspace too small, need ", state_bytes, " bytes");
+    STD_TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
+                    "workspace too small, need ", state_bytes, " bytes");
 
     // Zero the per-group RadixRowState region before launch.
     //
@@ -179,22 +183,22 @@ void launch_persistent_topk(const torch::Tensor& logits,
     //      first red_release. cudaMemsetAsync is stream-ordered: the zero
     //      is globally visible before any CTA runs.
     {
-      cudaError_t mz_err = cudaMemsetAsync(workspace.data_ptr<uint8_t>(), 0,
-                                           state_bytes, stream);
-      TORCH_CHECK(mz_err == cudaSuccess,
-                  "row_states memset failed: ", cudaGetErrorString(mz_err));
+      cudaError_t mz_err = cudaMemsetAsync(
+          workspace.mutable_data_ptr<uint8_t>(), 0, state_bytes, stream);
+      STD_TORCH_CHECK(mz_err == cudaSuccess,
+                      "row_states memset failed: ", cudaGetErrorString(mz_err));
     }
 
     P::PersistentTopKParams params;
-    params.input = logits.data_ptr<float>();
-    params.output = output.data_ptr<int32_t>();
-    params.lengths = lengths.data_ptr<int32_t>();
+    params.input = logits.const_data_ptr<float>();
+    params.output = output.mutable_data_ptr<int32_t>();
+    params.lengths = lengths.const_data_ptr<int32_t>();
     params.num_rows = static_cast<uint32_t>(num_rows);
     params.stride = static_cast<uint32_t>(stride);
     params.top_k = static_cast<uint32_t>(TopK);
     params.chunk_size = chunk_size;
-    params.row_states =
-        reinterpret_cast<P::RadixRowState*>(workspace.data_ptr<uint8_t>());
+    params.row_states = reinterpret_cast<P::RadixRowState*>(
+        workspace.mutable_data_ptr<uint8_t>());
     params.ctas_per_group = ctas_per_group;
     params.max_seq_len = static_cast<uint32_t>(max_seq_len);
 
@@ -203,8 +207,8 @@ void launch_persistent_topk(const torch::Tensor& logits,
       auto kernel = &P::persistent_topk_kernel<TOPK_VAL, VS>;                 \
       cudaError_t err = cudaFuncSetAttribute(                                 \
           kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);    \
-      TORCH_CHECK(err == cudaSuccess,                                         \
-                  "Failed to set smem: ", cudaGetErrorString(err));           \
+      STD_TORCH_CHECK(err == cudaSuccess,                                     \
+                      "Failed to set smem: ", cudaGetErrorString(err));       \
       kernel<<<total_ctas, P::kThreadsPerBlock, smem_size, stream>>>(params); \
     } while (0)
 
@@ -219,37 +223,42 @@ void launch_persistent_topk(const torch::Tensor& logits,
   }
 
   cudaError_t err = cudaGetLastError();
-  TORCH_CHECK(err == cudaSuccess,
-              "persistent_topk failed: ", cudaGetErrorString(err));
+  STD_TORCH_CHECK(err == cudaSuccess,
+                  "persistent_topk failed: ", cudaGetErrorString(err));
 }
 #endif
 
 }  // anonymous namespace
 
-void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
-                     torch::Tensor& output, torch::Tensor& workspace, int64_t k,
+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) {
 #ifndef USE_ROCM
-  TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
-  TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
-  TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
-  TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
-  TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
-  TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
-  TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
-  TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
-              "lengths must be 1D or 2D");
-  TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
-  TORCH_CHECK(output.dim() == 2, "output must be 2D");
+  STD_TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
+  STD_TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
+  STD_TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
+  STD_TORCH_CHECK(logits.scalar_type() == torch::headeronly::ScalarType::Float,
+                  "Only float32 supported");
+  STD_TORCH_CHECK(lengths.scalar_type() == torch::headeronly::ScalarType::Int,
+                  "lengths must be int32");
+  STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Int,
+                  "output must be int32");
+  STD_TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
+  STD_TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
+                  "lengths must be 1D or 2D");
+  STD_TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
+  STD_TORCH_CHECK(output.dim() == 2, "output must be 2D");
 
   const int64_t num_rows = logits.size(0);
-  const int64_t stride = logits.stride(0);
 
-  TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
-  TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
-              "output size mismatch");
-  TORCH_CHECK(k == 512 || k == 1024 || k == 2048,
-              "persistent_topk supports k=512, k=1024, or k=2048, got k=", k);
+  STD_TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
+  STD_TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
+                  "output size mismatch");
+  STD_TORCH_CHECK(
+      k == 512 || k == 1024 || k == 2048,
+      "persistent_topk supports k=512, k=1024, or k=2048, got k=", k);
 
   if (k == 512) {
     launch_persistent_topk<512>(logits, lengths, output, workspace,
@@ -262,6 +271,6 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
                                  max_seq_len);
   }
 #else
-  TORCH_CHECK(false, "persistent_topk is not supported on ROCm");
+  STD_TORCH_CHECK(false, "persistent_topk is not supported on ROCm");
 #endif
 }

csrc/libtorch_stable/torch_bindings.cpp

@@ -247,6 +247,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
   ops.def(
       "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
 
+  // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+).
+  // conditionally compiled so impl registration is in source file
+  ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
+
   // reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
   ops.def(
       "rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "
@@ -263,6 +267,20 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor");
 #endif
 
+  // Merge attn states
+  // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
+  // can be used to combine partial attention results (in the split-KV case)
+  ops.def(
+      "merge_attn_states("
+      "    Tensor! output,"
+      "    Tensor!? output_lse,"
+      "    Tensor prefix_output,"
+      "    Tensor prefix_lse,"
+      "    Tensor suffix_output,"
+      "    Tensor suffix_lse,"
+      "    int!? prefill_tokens_with_context,"
+      "    Tensor? output_scale=None) -> ()");
+
   // Hadamard transforms
   // conditionally compiled so impl registration is in source file
   ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor");
@@ -319,6 +337,37 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "bool is_neox, Tensor position_ids, "
       "int forced_token_heads_per_warp=-1) -> ()");
 
+  // Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE + KV-insert.
+  ops.def(
+      "fused_minimax_m3_qknorm_rope_kv_insert("
+      "Tensor! qkv, Tensor q_norm_weight, Tensor k_norm_weight, "
+      "Tensor cos_sin_cache, Tensor positions, int num_heads, "
+      "int num_kv_heads, int rotary_dim, float eps, "
+      "Tensor? index_q_norm_weight, Tensor? index_k_norm_weight, "
+      "int num_index_heads, "
+      "Tensor? slot_mapping, Tensor!? kv_cache, Tensor!? index_cache, "
+      "int block_size, Tensor!? q_out, Tensor!? index_q_out) -> ()");
+
+  // Apply repetition penalties to logits in-place.
+  ops.def(
+      "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
+      "Tensor output_mask, Tensor repetition_penalties) -> ()");
+
+  // Optimized top-k per row operations.
+  ops.def(
+      "top_k_per_row_prefill(Tensor logits, Tensor rowStarts, Tensor rowEnds, "
+      "Tensor! indices, int numRows, int stride0, "
+      "int stride1, int topK) -> ()");
+
+  ops.def(
+      "top_k_per_row_decode(Tensor logits, int next_n, "
+      "Tensor seq_lens, Tensor! indices, "
+      "int numRows, int stride0, int stride1, int topK) -> ()");
+
+  ops.def(
+      "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, "
+      "Tensor workspace, int k, int max_seq_len) -> ()");
+
   // Activation ops
   // Activation function used in SwiGLU.
   ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
@@ -326,9 +375,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
   ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
 
   // SwiGLU activation with input clamping.
+  // alpha scales the sigmoid (gate * sigmoid(alpha * gate)); beta is added to
+  // the up half (up + beta). Defaults alpha=1.0, beta=0.0 give silu(gate)*up.
   ops.def(
-      "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
-      "-> ()");
+      "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, "
+      "float alpha=1.0, float beta=0.0) -> ()");
 
   // Activation function used in GeGLU with `none` approximation.
   ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
@@ -422,6 +473,51 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "int type, SymInt row, SymInt tokens) -> Tensor");
 
   ops.def("ggml_moe_get_block_size(int type) -> int");
+
+  // Mamba selective scan kernel
+  ops.def(
+      "selective_scan_fwd(Tensor! u, Tensor! delta,"
+      "Tensor! A, Tensor! B, Tensor! C,"
+      "Tensor? D_, Tensor!? z_, Tensor? delta_bias_,"
+      "bool delta_softplus,"
+      "Tensor? query_start_loc,"
+      "Tensor? cache_indices,"
+      "Tensor? has_initial_state,"
+      "Tensor! ssm_states,"
+      "int null_block_id,"
+      "int block_size,"
+      "Tensor? block_idx_first_scheduled_token,"
+      "Tensor? block_idx_last_scheduled_token,"
+      "Tensor? initial_state_idx,"
+      "Tensor? cu_chunk_seqlen,"
+      "Tensor? last_chunk_indices) -> ()");
+
+  // Attention ops
+  // Compute the attention between an input query and the cached
+  // keys/values using PagedAttention.
+  ops.def(
+      "paged_attention_v1("
+      "    Tensor! out, Tensor query, Tensor key_cache,"
+      "    Tensor value_cache, int num_kv_heads, float scale,"
+      "    Tensor block_tables, Tensor seq_lens, int block_size,"
+      "    int max_seq_len, Tensor? alibi_slopes,"
+      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
+      "    int tp_rank, int blocksparse_local_blocks,"
+      "    int blocksparse_vert_stride, int blocksparse_block_size,"
+      "    int blocksparse_head_sliding_step) -> ()");
+
+  // PagedAttention V2.
+  ops.def(
+      "paged_attention_v2("
+      "    Tensor! out, Tensor! exp_sums, Tensor! max_logits,"
+      "    Tensor! tmp_out, Tensor query, Tensor key_cache,"
+      "    Tensor value_cache, int num_kv_heads, float scale,"
+      "    Tensor block_tables, Tensor seq_lens, int block_size,"
+      "    int max_seq_len, Tensor? alibi_slopes,"
+      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
+      "    int tp_rank, int blocksparse_local_blocks,"
+      "    int blocksparse_vert_stride, int blocksparse_block_size,"
+      "    int blocksparse_head_sliding_step) -> ()");
 }
 
 STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
@@ -469,6 +565,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   // files (allspark_repack.cu and allspark_qgemm_w8a16.cu)
 #endif
 
+  ops.impl("merge_attn_states", TORCH_BOX(&merge_attn_states));
+
   // Layernorm kernels (shared CUDA/ROCm)
   ops.impl("rms_norm", TORCH_BOX(&rms_norm));
   ops.impl("fused_add_rms_norm", TORCH_BOX(&fused_add_rms_norm));
@@ -486,6 +584,15 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   // Positional encoding kernels (shared CUDA/ROCm)
   ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding));
   ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));
+  ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
+           TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert));
+
+  // Sampler kernels (shared CUDA/ROCm)
+  ops.impl("apply_repetition_penalties_",
+           TORCH_BOX(&apply_repetition_penalties_));
+  ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill));
+  ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode));
+  ops.impl("persistent_topk", TORCH_BOX(&persistent_topk));
 
   // Activation kernels (shared CUDA/ROCm)
   ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul));
@@ -519,6 +626,10 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("ggml_mul_mat_a8", TORCH_BOX(&ggml_mul_mat_a8));
   ops.impl("ggml_moe_a8", TORCH_BOX(&ggml_moe_a8));
   ops.impl("ggml_moe_a8_vec", TORCH_BOX(&ggml_moe_a8_vec));
+  ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd));
+
+  ops.impl("paged_attention_v1", TORCH_BOX(&paged_attention_v1));
+  ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2));
 }
 
 // These capability-check functions take only primitive args (no tensors), so
@@ -541,4 +652,115 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) {
   ops.impl("ggml_moe_get_block_size", TORCH_BOX(&ggml_moe_get_block_size));
 }
 
+// Cache ops
+STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
+  // Swap in (out) the cache blocks from src to dst.
+  ops.def(
+      "swap_blocks(Tensor src, Tensor! dst,"
+      "            int block_size_in_bytes, Tensor block_mapping) -> ()");
+
+  // Batch swap: submit all block copies in a single driver call.
+  ops.def(
+      "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs,"
+      "                  Tensor sizes,"
+      "                  bool is_src_access_order_any=False) -> ()");
+
+  // Reshape the key and value tensors and cache them.
+  ops.def(
+      "reshape_and_cache(Tensor key, Tensor value,"
+      "                  Tensor! key_cache, Tensor! value_cache,"
+      "                  Tensor slot_mapping,"
+      "                  str kv_cache_dtype,"
+      "                  Tensor k_scale, Tensor v_scale) -> ()");
+
+  // Reshape the key and value tensors and cache them.
+  ops.def(
+      "reshape_and_cache_flash(Tensor key, Tensor value,"
+      "                        Tensor! key_cache,"
+      "                        Tensor! value_cache,"
+      "                        Tensor slot_mapping,"
+      "                        str kv_cache_dtype,"
+      "                        Tensor k_scale, Tensor v_scale) -> ()");
+
+  // Concat kv_c and k_pe and cache them.
+  ops.def(
+      "concat_and_cache_mla(Tensor kv_c, Tensor k_pe,"
+      "                     Tensor! kv_cache,"
+      "                     Tensor slot_mapping,"
+      "                     str kv_cache_dtype,"
+      "                     Tensor scale) -> ()");
+
+  // Rotate Q and K, then write to kv cache for MLA
+  ops.def(
+      "concat_and_cache_mla_rope_fused("
+      "                     Tensor positions,"
+      "                     Tensor! q_pe,"
+      "                     Tensor! k_pe,"
+      "                     Tensor kv_c,"
+      "                     Tensor cos_sin_cache,"
+      "                     bool is_neox,"
+      "                     Tensor slot_mapping,"
+      "                     Tensor! kv_cache,"
+      "                     str kv_cache_dtype,"
+      "                     Tensor kv_cache_scale) -> ()");
+
+  // Convert the key and value cache to fp8 data type.
+  ops.def(
+      "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, "
+      "str kv_cache_dtype) -> ()");
+
+  // Gather cache blocks from src_cache to dst, dequantizing from
+  // src_cache's dtype to dst's dtype if necessary.
+  ops.def(
+      "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, "
+      "                               Tensor block_table, Tensor cu_seq_lens, "
+      "                               Tensor token_to_seq, "
+      "                               int num_tokens, "
+      "                               str kv_cache_dtype, "
+      "                               Tensor scale, Tensor? seq_starts) -> ()");
+
+  ops.def(
+      "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, "
+      "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()");
+
+  ops.def(
+      "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, "
+      "Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int "
+      "batch_size) -> ()");
+
+  ops.def(
+      "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor "
+      "slot_mapping, "
+      "int quant_block_size, str kv_cache_dtype) -> ()");
+
+  ops.def("concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()");
+
+  ops.def(
+      "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! "
+      "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()");
+}
+
+STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CPU, ops) {
+  ops.impl("swap_blocks_batch", TORCH_BOX(&swap_blocks_batch));
+}
+
+STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) {
+  ops.impl("swap_blocks", TORCH_BOX(&swap_blocks));
+  ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache));
+  ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash));
+  ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla));
+  ops.impl("concat_and_cache_mla_rope_fused",
+           TORCH_BOX(&concat_and_cache_mla_rope_fused));
+  ops.impl("convert_fp8", TORCH_BOX(&convert_fp8));
+  ops.impl("gather_and_maybe_dequant_cache",
+           TORCH_BOX(&gather_and_maybe_dequant_cache));
+  ops.impl("cp_gather_cache", TORCH_BOX(&cp_gather_cache));
+  ops.impl("cp_gather_and_upconvert_fp8_kv_cache",
+           TORCH_BOX(&cp_gather_and_upconvert_fp8_kv_cache));
+  ops.impl("indexer_k_quant_and_cache", TORCH_BOX(&indexer_k_quant_and_cache));
+  ops.impl("concat_mla_q", TORCH_BOX(&concat_mla_q));
+  ops.impl("cp_gather_indexer_k_quant_cache",
+           TORCH_BOX(&cp_gather_indexer_k_quant_cache));
+}
+
 REGISTER_EXTENSION(_C_stable_libtorch)

csrc/minimax_m3_build_k2q_csr.cu

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

csrc/moe/moe_ops.h

@@ -62,6 +62,9 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
 
 bool moe_permute_unpermute_supported();
 
+int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
+                                        int64_t num_experts);
+
 void shuffle_rows(const torch::Tensor& input_tensor,
                   const torch::Tensor& dst2src_map,
                   torch::Tensor& output_tensor);

csrc/moe/moe_permute_unpermute_op.cu

@@ -8,6 +8,108 @@
 // moe_permute kernels require at least CUDA 12.0
 #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)
 
+namespace {
+
+torch::Tensor maybe_allocate_tensor(
+    const std::optional<torch::Tensor>& maybe_tensor,
+    at::IntArrayRef expected_sizes, torch::ScalarType dtype, c10::Device device,
+    char const* name) {
+  auto expected_numel = c10::multiply_integers(expected_sizes);
+  if (maybe_tensor.has_value()) {
+    auto tensor = maybe_tensor.value();
+    TORCH_CHECK(tensor.device() == device, name, " must be on the same device");
+    TORCH_CHECK(tensor.scalar_type() == dtype, name, " has incorrect dtype");
+    TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+    TORCH_CHECK(tensor.numel() >= expected_numel, name,
+                " is too small for the requested shape");
+    auto flat_tensor = tensor.view({tensor.numel()});
+    return flat_tensor.narrow(0, 0, expected_numel).view(expected_sizes);
+  }
+  return torch::empty(expected_sizes, torch::dtype(dtype).device(device));
+}
+
+}  // namespace
+
+int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
+                                        int64_t n_expert) {
+  return static_cast<int64_t>(
+      CubKeyValueSorter::getWorkspaceSize(num_expanded_rows, n_expert));
+}
+
+void moe_permute_impl(
+    const torch::Tensor& input,                      // [n_token, hidden]
+    const torch::Tensor& topk_ids,                   // [n_token, topk]
+    const torch::Tensor& token_expert_indices,       // [n_token, topk]
+    const std::optional<torch::Tensor>& expert_map,  // [n_expert]
+    int64_t n_expert, int64_t n_local_expert, int64_t topk,
+    torch::Tensor& permuted_input,             // [permuted_size, hidden]
+    torch::Tensor& expert_first_token_offset,  // [n_local_expert + 1]
+    torch::Tensor& inv_permuted_idx,           // [n_token, topk]
+    torch::Tensor& permuted_idx,               // [permute_size]
+    const std::optional<torch::Tensor>& maybe_sort_workspace,
+    const std::optional<torch::Tensor>& maybe_permuted_experts_id,
+    const std::optional<torch::Tensor>& maybe_sorted_row_idx,
+    const std::optional<torch::Tensor>& maybe_topk_ids_for_sort) {
+  TORCH_CHECK(expert_first_token_offset.scalar_type() == at::ScalarType::Long,
+              "expert_first_token_offset must be int64");
+  TORCH_CHECK(topk_ids.scalar_type() == at::ScalarType::Int,
+              "topk_ids must be int32");
+  TORCH_CHECK(token_expert_indices.scalar_type() == at::ScalarType::Int,
+              "token_expert_indices must be int32");
+  TORCH_CHECK(inv_permuted_idx.scalar_type() == at::ScalarType::Int,
+              "inv_permuted_idx must be int32");
+  TORCH_CHECK(expert_first_token_offset.size(0) == n_local_expert + 1,
+              "expert_first_token_offset shape != n_local_expert+1");
+  TORCH_CHECK(inv_permuted_idx.sizes() == token_expert_indices.sizes(),
+              "token_expert_indices shape must be same as inv_permuted_idx");
+  auto device = input.device();
+  auto n_token = input.sizes()[0];
+  auto n_hidden = input.sizes()[1];
+  auto expanded_rows = n_token * topk;
+  auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+  auto sorter_size = moe_permute_sort_workspace_size(expanded_rows, n_expert);
+  auto sort_workspace =
+      maybe_allocate_tensor(maybe_sort_workspace, {sorter_size}, torch::kInt8,
+                            device, "sort_workspace");
+  auto permuted_experts_id =
+      maybe_allocate_tensor(maybe_permuted_experts_id, topk_ids.sizes(),
+                            at::ScalarType::Int, device, "permuted_experts_id");
+  auto sorted_row_idx =
+      maybe_allocate_tensor(maybe_sorted_row_idx, inv_permuted_idx.sizes(),
+                            at::ScalarType::Int, device, "sorted_row_idx");
+
+  CubKeyValueSorter sorter{};
+  int64_t* valid_num_ptr = nullptr;
+  torch::Tensor topk_ids_for_sort = topk_ids;
+
+  if (expert_map.has_value()) {
+    const int* expert_map_ptr = get_ptr<int>(expert_map.value());
+    valid_num_ptr =
+        get_ptr<int64_t>(expert_first_token_offset) + n_local_expert;
+    topk_ids_for_sort =
+        maybe_allocate_tensor(maybe_topk_ids_for_sort, topk_ids.sizes(),
+                              at::ScalarType::Int, device, "topk_ids_for_sort");
+    topk_ids_for_sort.copy_(topk_ids);
+    preprocessTopkIdLauncher(get_ptr<int>(topk_ids_for_sort), n_token * topk,
+                             expert_map_ptr, n_expert, stream);
+  }
+
+  sortAndScanExpert(
+      get_ptr<const int>(topk_ids_for_sort), get_ptr<int>(token_expert_indices),
+      get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
+      get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
+      n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
+
+  MOE_DISPATCH(input.scalar_type(), [&] {
+    expandInputRowsKernelLauncher<scalar_t>(
+        get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
+        get_ptr<int>(sorted_row_idx), get_ptr<int>(inv_permuted_idx),
+        get_ptr<int>(permuted_idx), get_ptr<int64_t>(expert_first_token_offset),
+        n_token, valid_num_ptr, n_hidden, topk, n_local_expert, stream);
+  });
+}
+
 void moe_permute(
     const torch::Tensor& input,                      // [n_token, hidden]
     const torch::Tensor& topk_ids,                   // [n_token, topk]
@@ -18,65 +120,26 @@ void moe_permute(
     torch::Tensor& expert_first_token_offset,  // [n_local_expert + 1]
     torch::Tensor& inv_permuted_idx,           // [n_token, topk]
     torch::Tensor& permuted_idx) {             // [permute_size]
-  TORCH_CHECK(expert_first_token_offset.scalar_type() == at::ScalarType::Long,
-              "expert_first_token_offset must be int64");
-  TORCH_CHECK(topk_ids.scalar_type() == at::ScalarType::Int,
-              "topk_ids must be int32");
-  TORCH_CHECK(token_expert_indices.scalar_type() == at::ScalarType::Int,
-              "token_expert_indices must be int32");
-  TORCH_CHECK(inv_permuted_idx.scalar_type() == at::ScalarType::Int,
-              "inv_permuted_idx must be int32");
-  TORCH_CHECK(expert_first_token_offset.size(0) == n_local_expert + 1,
-              "expert_first_token_offset shape != n_local_expert+1")
-  TORCH_CHECK(inv_permuted_idx.sizes() == token_expert_indices.sizes(),
-              "token_expert_indices shape must be same as inv_permuted_idx");
-  auto n_token = input.sizes()[0];
-  auto n_hidden = input.sizes()[1];
-  auto stream = at::cuda::getCurrentCUDAStream().stream();
-  const long sorter_size =
-      CubKeyValueSorter::getWorkspaceSize(n_token * topk, n_expert);
-  auto sort_workspace = torch::empty(
-      {sorter_size},
-      torch::dtype(torch::kInt8).device(torch::kCUDA).requires_grad(false));
-  torch::Tensor topk_ids_for_sort = topk_ids;
-  auto permuted_experts_id = torch::empty_like(topk_ids);
-  auto sorted_row_idx = torch::empty_like(inv_permuted_idx);
+  moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert,
+                   n_local_expert, topk, permuted_input,
+                   expert_first_token_offset, inv_permuted_idx, permuted_idx,
+                   std::nullopt, std::nullopt, std::nullopt, std::nullopt);
+}
 
-  CubKeyValueSorter sorter{};
-  int64_t* valid_num_ptr = nullptr;
-  // pre-process kernel for expert-parallelism:
-  // no local expert id plus "n_expert" offset for priority to local expert
-  // map local expert id [n, .., n+n_local_expert-1] to [0, n_local_expert -1]
-  // For example, 4 expert with ep_size=2. ep_rank=1 owns global expert id
-  // [2,3] with expert_map[-1, -1, 0, 1], preprocess_topk_id  process topk_ids
-  // and map global expert id [2, 3] to local_expert id [0, 1] and map global
-  // expert id [0, 1] ( not in ep rank=1)  to [4, 5] by plus n_expert. This map
-  // operation is to make local expert high priority in following sort topk_ids
-  // and scan local expert_first_token_offset for each ep rank for next group
-  // gemm.
-  if (expert_map.has_value()) {
-    const int* expert_map_ptr = get_ptr<int>(expert_map.value());
-    valid_num_ptr =
-        get_ptr<int64_t>(expert_first_token_offset) + n_local_expert;
-    topk_ids_for_sort = topk_ids.clone();
-    preprocessTopkIdLauncher(get_ptr<int>(topk_ids_for_sort), n_token * topk,
-                             expert_map_ptr, n_expert, stream);
-  }
-  // expert sort topk expert id and scan expert id get expert_first_token_offset
-  sortAndScanExpert(
-      get_ptr<const int>(topk_ids_for_sort), get_ptr<int>(token_expert_indices),
-      get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
-      get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
-      n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
-
-  // dispatch expandInputRowsKernelLauncher
-  MOE_DISPATCH(input.scalar_type(), [&] {
-    expandInputRowsKernelLauncher<scalar_t>(
-        get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
-        get_ptr<int>(sorted_row_idx), get_ptr<int>(inv_permuted_idx),
-        get_ptr<int>(permuted_idx), get_ptr<int64_t>(expert_first_token_offset),
-        n_token, valid_num_ptr, n_hidden, topk, n_local_expert, stream);
-  });
+void moe_permute_with_scratch(
+    const torch::Tensor& input, const torch::Tensor& topk_ids,
+    const torch::Tensor& token_expert_indices,
+    const std::optional<torch::Tensor>& expert_map, int64_t n_expert,
+    int64_t n_local_expert, int64_t topk, torch::Tensor& permuted_input,
+    torch::Tensor& expert_first_token_offset, torch::Tensor& inv_permuted_idx,
+    torch::Tensor& permuted_idx, torch::Tensor& sort_workspace,
+    torch::Tensor& permuted_experts_id, torch::Tensor& sorted_row_idx,
+    torch::Tensor& topk_ids_for_sort) {
+  moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert,
+                   n_local_expert, topk, permuted_input,
+                   expert_first_token_offset, inv_permuted_idx, permuted_idx,
+                   sort_workspace, permuted_experts_id, sorted_row_idx,
+                   topk_ids_for_sort);
 }
 
 void moe_unpermute(
@@ -169,6 +232,12 @@ void shuffle_rows(const torch::Tensor& input_tensor,
 
 #else
 
+int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
+                                        int64_t n_expert) {
+  TORCH_CHECK(
+      false, "moe_permute_sort_workspace_size is not supported on CUDA < 12.0");
+}
+
 void moe_permute(const torch::Tensor& input, const torch::Tensor& topk_ids,
                  const torch::Tensor& token_expert_indices,
                  const std::optional<torch::Tensor>& expert_map,
@@ -179,6 +248,19 @@ void moe_permute(const torch::Tensor& input, const torch::Tensor& topk_ids,
   TORCH_CHECK(false, "moe_permute is not supported on CUDA < 12.0");
 }
 
+void moe_permute_with_scratch(
+    const torch::Tensor& input, const torch::Tensor& topk_ids,
+    const torch::Tensor& token_expert_indices,
+    const std::optional<torch::Tensor>& expert_map, int64_t n_expert,
+    int64_t n_local_expert, int64_t topk, torch::Tensor& permuted_input,
+    torch::Tensor& expert_first_token_offset, torch::Tensor& inv_permuted_idx,
+    torch::Tensor& permuted_idx, torch::Tensor& sort_workspace,
+    torch::Tensor& permuted_experts_id, torch::Tensor& sorted_row_idx,
+    torch::Tensor& topk_ids_for_sort) {
+  TORCH_CHECK(false,
+              "moe_permute_with_scratch is not supported on CUDA < 12.0");
+}
+
 void moe_unpermute(
     const torch::Tensor& permuted_hidden_states,
     const torch::Tensor& topk_weights, const torch::Tensor& inv_permuted_idx,
@@ -199,5 +281,6 @@ bool moe_permute_unpermute_supported() {
 
 TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
   m.impl("moe_permute", &moe_permute);
+  m.impl("moe_permute_with_scratch", &moe_permute_with_scratch);
   m.impl("moe_unpermute", &moe_unpermute);
 }

csrc/moe/torch_bindings.cpp

@@ -100,13 +100,26 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
       "expert_first_token_offset, Tensor! inv_permuted_idx, Tensor! "
       "permuted_idx)->()");
 
+  m.def(
+      "moe_permute_with_scratch(Tensor input, Tensor topk_ids,"
+      "Tensor token_expert_indices, Tensor? expert_map, int n_expert,"
+      "int n_local_expert,"
+      "int topk, Tensor! permuted_input, Tensor! "
+      "expert_first_token_offset, Tensor! inv_permuted_idx, Tensor! "
+      "permuted_idx, Tensor! sort_workspace, Tensor! permuted_experts_id, "
+      "Tensor! sorted_row_idx, Tensor! topk_ids_for_sort)->()");
+
   m.def(
       "moe_unpermute(Tensor permuted_hidden_states, Tensor topk_weights,"
       "Tensor inv_permuted_idx, Tensor? expert_first_token_offset, "
       "int topk, Tensor! hidden_states)->()");
 
   m.def("moe_permute_unpermute_supported() -> bool");
+  m.def(
+      "moe_permute_sort_workspace_size(int num_expanded_rows, int n_expert) -> "
+      "int");
   m.impl("moe_permute_unpermute_supported", &moe_permute_unpermute_supported);
+  m.impl("moe_permute_sort_workspace_size", &moe_permute_sort_workspace_size);
 
   // Row shuffle for MoE
   m.def(

csrc/ops.h

@@ -31,36 +31,6 @@ torch::Tensor weak_ref_tensor(torch::Tensor& tensor) {
   return new_tensor;
 }
 
-void paged_attention_v1(
-    torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache,
-    torch::Tensor& value_cache, int64_t num_kv_heads, double scale,
-    torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size,
-    int64_t max_seq_len, const std::optional<torch::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::Tensor& k_scale,
-    torch::Tensor& v_scale, const int64_t tp_rank,
-    const int64_t blocksparse_local_blocks,
-    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
-    const int64_t blocksparse_head_sliding_step);
-
-void paged_attention_v2(
-    torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
-    torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
-    torch::Tensor& value_cache, int64_t num_kv_heads, double scale,
-    torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size,
-    int64_t max_seq_len, const std::optional<torch::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::Tensor& k_scale,
-    torch::Tensor& v_scale, const int64_t tp_rank,
-    const int64_t blocksparse_local_blocks,
-    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
-    const int64_t blocksparse_head_sliding_step);
-
-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,
-    const std::optional<int64_t> prefill_tokens_with_context,
-    const std::optional<torch::Tensor>& output_scale = std::nullopt);
-
 // rms_norm and fused_add_rms_norm declarations also exist in
 // csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
 // because the CPU build still uses these torch::Tensor declarations.
@@ -70,30 +40,11 @@ void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
 void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
                         torch::Tensor& weight, double epsilon);
 
-void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::Tensor& q, torch::Tensor const& kv, torch::Tensor& k_cache,
+torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+    torch::Tensor const& q_in, torch::Tensor const& kv, torch::Tensor& k_cache,
     torch::Tensor const& slot_mapping, torch::Tensor const& position_ids,
-    torch::Tensor const& cos_sin_cache, double eps, int64_t cache_block_size);
-
-void apply_repetition_penalties_(torch::Tensor& logits,
-                                 const torch::Tensor& prompt_mask,
-                                 const torch::Tensor& output_mask,
-                                 const torch::Tensor& repetition_penalties);
-
-void top_k_per_row_prefill(const torch::Tensor& logits,
-                           const torch::Tensor& rowStarts,
-                           const torch::Tensor& rowEnds, torch::Tensor& indices,
-                           int64_t numRows, int64_t stride0, int64_t stride1,
-                           int64_t topK);
-
-void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,
-                          const torch::Tensor& seqLens, torch::Tensor& indices,
-                          int64_t numRows, int64_t stride0, int64_t stride1,
-                          int64_t topK);
-
-void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
-                     torch::Tensor& output, torch::Tensor& workspace, int64_t k,
-                     int64_t max_seq_len);
+    torch::Tensor const& cos_sin_cache, int64_t q_head_padded, double eps,
+    int64_t cache_block_size);
 
 void silu_and_mul_per_block_quant(torch::Tensor& out,
                                   torch::Tensor const& input,
@@ -111,7 +62,8 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
 
 void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
 
-void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
+void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit,
+                        double alpha = 1.0, double beta = 0.0);
 
 void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
                         torch::Tensor& scale);
@@ -149,22 +101,6 @@ void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
                                torch::Tensor& scales,
                                std::optional<torch::Tensor> const& azp);
 
-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_, 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, 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);
-
 torch::Tensor dynamic_4bit_int_moe_cpu(
     torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,
     torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I,
@@ -211,4 +147,12 @@ std::tuple<torch::Tensor, torch::Tensor> minimax_allreduce_rms_qk(
     torch::Tensor const& norm_weight_k, torch::Tensor workspace,
     int64_t const q_size, int64_t const kv_size, int64_t const rank,
     int64_t const nranks, double const eps);
+
+void run_minimax_m3_build_k2q_csr_with_schedule(
+    torch::Tensor q2k, torch::Tensor cu_q, torch::Tensor cu_k,
+    torch::Tensor row_ptr, torch::Tensor q_idx,
+    torch::Tensor scheduler_metadata, torch::Tensor work_count,
+    torch::Tensor qsplit_idx, torch::Tensor split_counts, int64_t topk,
+    int64_t blk_kv, int64_t total_rows, int64_t max_kv_blocks,
+    int64_t target_q_per_cta, int64_t work_capacity, int64_t max_seqlen_q);
 #endif

csrc/persistent_topk.cuh

@@ -126,10 +126,10 @@ struct RadixRowState {
 // ============================================================================
 
 struct PersistentTopKParams {
-  const float* __restrict__ input;  // [num_rows, stride]
-  int32_t* __restrict__ output;     // [num_rows, top_k]
-  int32_t* __restrict__ lengths;    // [num_rows]
-  RadixRowState* row_states;        // large path: per-group state
+  const float* __restrict__ input;      // [num_rows, stride]
+  int32_t* __restrict__ output;         // [num_rows, top_k]
+  const int32_t* __restrict__ lengths;  // [num_rows]
+  RadixRowState* row_states;            // large path: per-group state
   uint32_t num_rows;
   uint32_t stride;
   uint32_t top_k;           // actual k value for output stride
@@ -1269,9 +1269,11 @@ constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) {
 }
 
 template <typename DType, typename IdType, uint32_t MAX_K = 2048>
-cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
-                                        IdType* lengths, uint32_t num_rows,
-                                        uint32_t top_k_val, uint32_t max_len,
+cudaError_t FilteredTopKRaggedTransform(const DType* input,
+                                        IdType* output_indices,
+                                        const IdType* lengths,
+                                        uint32_t num_rows, uint32_t top_k_val,
+                                        uint32_t max_len,
                                         cudaStream_t stream = 0) {
   constexpr size_t smem_size = FILTERED_TOPK_SMEM_DYNAMIC;
   constexpr int MAX_VEC = 16 / sizeof(DType);

csrc/quantization/marlin/marlin.cu

@@ -400,10 +400,12 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
                 "Turing only support FP16 or INT8 activation.");
   }
   if (a_type == vllm::kFE4M3fn) {
+    TORCH_CHECK(major_capability * 10 + minor_capability >= 89,
+                "FP8 only support Ada Lovelace or newer GPUs.");
     TORCH_CHECK(
         major_capability * 10 + minor_capability == 89 ||
-            major_capability * 10 + minor_capability == 120,
-        "Marlin W4A8-FP8 only support SM89 or SM120 device (It is slower than "
+            major_capability == 12,
+        "Marlin W4A8-FP8 only support SM89 or SM12x device (It is slower than "
         "Marlin W4A16 on other devices).");
   }
 

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

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

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

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

csrc/rocm/ops.h

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

csrc/rocm/q_gemm_rdna3.cu

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

csrc/rocm/qdq_4_rdna3.cuh

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

csrc/rocm/skinny_gemms.cu

@@ -1277,6 +1277,12 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
         else
           WVSPLIT_TILE_CFG(64, 16, sYT, 4)
         break;
+      case 5:
+        if (use_wave32)
+          WVSPLIT_TILE_CFG(32, 16, sYT, 5)
+        else
+          WVSPLIT_TILE_CFG(64, 16, sYT, 5)
+        break;
       default:
         throw std::runtime_error(
             "Unsupported N value: " + std::to_string(M_in) + "," +

csrc/rocm/torch_bindings.cpp

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

csrc/torch_bindings.cpp

@@ -1,4 +1,7 @@
-#include "cache.h"
+// Provides torch::Tensor for ops.h (previously included transitively via
+// cache.h, which is no longer included here after cache ops moved to
+// _C_stable_libtorch).
+#include <torch/all.h>
 #include "cuda_utils.h"
 #include "ops.h"
 #include "core/registration.h"
@@ -33,50 +36,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
            &get_cuda_view_from_cpu_tensor);
 
-  // Attention ops
-  // Compute the attention between an input query and the cached
-  // keys/values using PagedAttention.
-  ops.def(
-      "paged_attention_v1("
-      "    Tensor! out, Tensor query, Tensor key_cache,"
-      "    Tensor value_cache, int num_kv_heads, float scale,"
-      "    Tensor block_tables, Tensor seq_lens, int block_size,"
-      "    int max_seq_len, Tensor? alibi_slopes,"
-      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
-      "    int tp_rank, int blocksparse_local_blocks,"
-      "    int blocksparse_vert_stride, int blocksparse_block_size,"
-      "    int blocksparse_head_sliding_step) -> ()");
-  ops.impl("paged_attention_v1", torch::kCUDA, &paged_attention_v1);
-
-  // PagedAttention V2.
-  ops.def(
-      "paged_attention_v2("
-      "    Tensor! out, Tensor! exp_sums, Tensor! max_logits,"
-      "    Tensor! tmp_out, Tensor query, Tensor key_cache,"
-      "    Tensor value_cache, int num_kv_heads, float scale,"
-      "    Tensor block_tables, Tensor seq_lens, int block_size,"
-      "    int max_seq_len, Tensor? alibi_slopes,"
-      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
-      "    int tp_rank, int blocksparse_local_blocks,"
-      "    int blocksparse_vert_stride, int blocksparse_block_size,"
-      "    int blocksparse_head_sliding_step) -> ()");
-  ops.impl("paged_attention_v2", torch::kCUDA, &paged_attention_v2);
-
-  // Merge attn states
-  // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
-  // can be used to combine partial attention results (in the split-KV case)
-  ops.def(
-      "merge_attn_states("
-      "    Tensor! output,"
-      "    Tensor!? output_lse,"
-      "    Tensor prefix_output,"
-      "    Tensor prefix_lse,"
-      "    Tensor suffix_output,"
-      "    Tensor suffix_lse,"
-      "    int!? prefill_tokens_with_context,"
-      "    Tensor? output_scale=None) -> ()");
-  ops.impl("merge_attn_states", torch::kCUDA, &merge_attn_states);
-
   // Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
   ops.def(
       "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
@@ -99,37 +58,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   // kernel launch.
   ops.def(
       "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
-      "Tensor! q, Tensor kv, Tensor! k_cache, "
+      "Tensor q_in, Tensor kv, Tensor! k_cache, "
       "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
-      "float eps, int cache_block_size) -> ()");
+      "int q_head_padded, float eps, int cache_block_size) -> Tensor");
   ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
            &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
 
-  // Apply repetition penalties to logits in-place
-  ops.def(
-      "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
-      "Tensor output_mask, Tensor repetition_penalties) -> ()");
-  ops.impl("apply_repetition_penalties_", torch::kCUDA,
-           &apply_repetition_penalties_);
-
-  // Optimized top-k per row operation
-  ops.def(
-      "top_k_per_row_prefill(Tensor logits, Tensor rowStarts, Tensor rowEnds, "
-      "Tensor! indices, int numRows, int stride0, "
-      "int stride1, int topK) -> ()");
-  ops.impl("top_k_per_row_prefill", torch::kCUDA, &top_k_per_row_prefill);
-
-  ops.def(
-      "top_k_per_row_decode(Tensor logits, int next_n, "
-      "Tensor seq_lens, Tensor! indices, "
-      "int numRows, int stride0, int stride1, int topK) -> ()");
-  ops.impl("top_k_per_row_decode", torch::kCUDA, &top_k_per_row_decode);
-
-  ops.def(
-      "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, "
-      "Tensor workspace, int k, int max_seq_len) -> ()");
-  ops.impl("persistent_topk", torch::kCUDA, &persistent_topk);
-
   // Quantization ops
 #ifndef USE_ROCM
 
@@ -230,25 +164,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 
 #endif
 
-  // Mamba selective scan kernel
-  ops.def(
-      "selective_scan_fwd(Tensor! u, Tensor! delta,"
-      "Tensor! A, Tensor! B, Tensor! C,"
-      "Tensor? D_, Tensor!? z_, Tensor? delta_bias_,"
-      "bool delta_softplus,"
-      "Tensor? query_start_loc,"
-      "Tensor? cache_indices,"
-      "Tensor? has_initial_state,"
-      "Tensor! ssm_states,"
-      "int null_block_id,"
-      "int block_size,"
-      "Tensor? block_idx_first_scheduled_token,"
-      "Tensor? block_idx_last_scheduled_token,"
-      "Tensor? initial_state_idx,"
-      "Tensor? cu_chunk_seqlen,"
-      "Tensor? last_chunk_indices) -> ()");
-  ops.impl("selective_scan_fwd", torch::kCUDA, &selective_scan_fwd);
-
 #ifndef USE_ROCM
   ops.def(
       "minimax_allreduce_rms("
@@ -272,118 +187,31 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "float eps) -> (Tensor, Tensor)");
   ops.impl("minimax_allreduce_rms_qk", torch::kCUDA, &minimax_allreduce_rms_qk);
 
+  ops.def(
+      "minimax_m3_build_k2q_csr_with_schedule("
+      "Tensor q2k,"
+      "Tensor cu_q,"
+      "Tensor cu_k,"
+      "Tensor! row_ptr,"
+      "Tensor! q_idx,"
+      "Tensor! scheduler_metadata,"
+      "Tensor! work_count,"
+      "Tensor! qsplit_idx,"
+      "Tensor! split_counts,"
+      "int topk,"
+      "int blk_kv,"
+      "int total_rows,"
+      "int max_kv_blocks,"
+      "int target_q_per_cta,"
+      "int work_capacity,"
+      "int max_seqlen_q) -> ()");
+  ops.impl("minimax_m3_build_k2q_csr_with_schedule", torch::kCUDA,
+           &run_minimax_m3_build_k2q_csr_with_schedule);
+
   //  conditionally compiled so impl in source file
 #endif
 }
 
-TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
-  // Cache ops
-  // Swap in (out) the cache blocks from src to dst.
-  cache_ops.def(
-      "swap_blocks(Tensor src, Tensor! dst,"
-      "            int block_size_in_bytes, Tensor block_mapping) -> ()");
-  cache_ops.impl("swap_blocks", torch::kCUDA, &swap_blocks);
-
-  // Batch swap: submit all block copies in a single driver call.
-  cache_ops.def(
-      "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs,"
-      "                  Tensor sizes,"
-      "                  bool is_src_access_order_any=False) -> ()");
-  cache_ops.impl("swap_blocks_batch", torch::kCPU, &swap_blocks_batch);
-
-  // Reshape the key and value tensors and cache them.
-  cache_ops.def(
-      "reshape_and_cache(Tensor key, Tensor value,"
-      "                  Tensor! key_cache, Tensor! value_cache,"
-      "                  Tensor slot_mapping,"
-      "                  str kv_cache_dtype,"
-      "                  Tensor k_scale, Tensor v_scale) -> ()");
-  cache_ops.impl("reshape_and_cache", torch::kCUDA, &reshape_and_cache);
-
-  // Reshape the key and value tensors and cache them.
-  cache_ops.def(
-      "reshape_and_cache_flash(Tensor key, Tensor value,"
-      "                        Tensor! key_cache,"
-      "                        Tensor! value_cache,"
-      "                        Tensor slot_mapping,"
-      "                        str kv_cache_dtype,"
-      "                        Tensor k_scale, Tensor v_scale) -> ()");
-  cache_ops.impl("reshape_and_cache_flash", torch::kCUDA,
-                 &reshape_and_cache_flash);
-
-  // Concat kv_c and k_pe and cache them.
-  cache_ops.def(
-      "concat_and_cache_mla(Tensor kv_c, Tensor k_pe,"
-      "                     Tensor! kv_cache,"
-      "                     Tensor slot_mapping,"
-      "                     str kv_cache_dtype,"
-      "                     Tensor scale) -> ()");
-  cache_ops.impl("concat_and_cache_mla", torch::kCUDA, &concat_and_cache_mla);
-
-  // Rotate Q and K, then write to kv cache for MLA
-  cache_ops.def(
-      "concat_and_cache_mla_rope_fused("
-      "                     Tensor positions,"
-      "                     Tensor! q_pe,"
-      "                     Tensor! k_pe,"
-      "                     Tensor kv_c,"
-      "                     Tensor cos_sin_cache,"
-      "                     bool is_neox,"
-      "                     Tensor slot_mapping,"
-      "                     Tensor! kv_cache,"
-      "                     str kv_cache_dtype,"
-      "                     Tensor kv_cache_scale) -> ()");
-  cache_ops.impl("concat_and_cache_mla_rope_fused", torch::kCUDA,
-                 &concat_and_cache_mla_rope_fused);
-
-  // Convert the key and value cache to fp8 data type.
-  cache_ops.def(
-      "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, "
-      "str kv_cache_dtype) -> ()");
-  cache_ops.impl("convert_fp8", torch::kCUDA, &convert_fp8);
-
-  // Gather cache blocks from src_cache to dst, dequantizing from
-  // src_cache's dtype to dst's dtype if necessary.
-  cache_ops.def(
-      "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, "
-      "                               Tensor block_table, Tensor cu_seq_lens, "
-      "                               Tensor token_to_seq, "
-      "                               int num_tokens, "
-      "                               str kv_cache_dtype, "
-      "                               Tensor scale, Tensor? seq_starts) -> ()");
-  cache_ops.impl("gather_and_maybe_dequant_cache", torch::kCUDA,
-                 &gather_and_maybe_dequant_cache);
-
-  cache_ops.def(
-      "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, "
-      "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()");
-  cache_ops.impl("cp_gather_cache", torch::kCUDA, &cp_gather_cache);
-
-  cache_ops.def(
-      "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, "
-      "Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int "
-      "batch_size) -> ()");
-  cache_ops.impl("cp_gather_and_upconvert_fp8_kv_cache", torch::kCUDA,
-                 &cp_gather_and_upconvert_fp8_kv_cache);
-
-  cache_ops.def(
-      "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor "
-      "slot_mapping, "
-      "int quant_block_size, str kv_cache_dtype) -> ()");
-  cache_ops.impl("indexer_k_quant_and_cache", torch::kCUDA,
-                 &indexer_k_quant_and_cache);
-
-  cache_ops.def(
-      "concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()");
-  cache_ops.impl("concat_mla_q", torch::kCUDA, &concat_mla_q);
-
-  cache_ops.def(
-      "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! "
-      "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()");
-  cache_ops.impl("cp_gather_indexer_k_quant_cache", torch::kCUDA,
-                 &cp_gather_indexer_k_quant_cache);
-}
-
 TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
   // Cuda utils
 

docker/Dockerfile

@@ -220,7 +220,8 @@ COPY use_existing_torch.py use_existing_torch.py
 COPY pyproject.toml pyproject.toml
 RUN --mount=type=cache,target=/opt/uv/cache \
     if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
-        sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' requirements/cuda.txt; \
+        sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' requirements/cuda.txt; \
+        sed -i 's/^humming-kernels\[cu13\]/humming-kernels[cu12]/' requirements/cuda.txt; \
     fi \
     && if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
         echo "Installing torch nightly..." \
@@ -746,7 +747,8 @@ COPY requirements/common.txt /tmp/common.txt
 COPY requirements/cuda.txt /tmp/requirements-cuda.txt
 RUN --mount=type=cache,target=/opt/uv/cache \
     if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
-        sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' /tmp/requirements-cuda.txt; \
+        sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' /tmp/requirements-cuda.txt; \
+        sed -i 's/^humming-kernels\[cu13\]/humming-kernels[cu12]/' /tmp/requirements-cuda.txt; \
     fi && \
     uv pip install --system -r /tmp/requirements-cuda.txt \
         --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') && \
@@ -755,7 +757,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
 # https://docs.flashinfer.ai/installation.html
 # From versions.json: .flashinfer.version
-ARG FLASHINFER_VERSION=0.6.11.post2
+ARG FLASHINFER_VERSION=0.6.12
 RUN --mount=type=cache,target=/opt/uv/cache \
     uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
         --extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
@@ -866,6 +868,7 @@ FROM vllm-base AS test
 ADD . /vllm-workspace/
 
 ARG PYTHON_VERSION
+ARG TARGETPLATFORM
 
 ARG PIP_INDEX_URL UV_INDEX_URL
 ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
@@ -905,6 +908,11 @@ RUN --mount=type=cache,target=/opt/uv/cache \
             --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
         else \
             echo "Installing dev requirements..." \
+            && if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
+                echo "Recompiling test requirements for arm64..." \
+                && uv pip compile requirements/test/cuda.in -o requirements/test/cuda.txt --index-strategy unsafe-best-match \
+                --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+            fi \
             && uv pip install --system -r requirements/dev.txt \
             --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
         fi \

docker/Dockerfile.cpu

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

docker/Dockerfile.nightly_torch

@@ -256,13 +256,13 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.
 
 
 # build flashinfer for torch nightly from source around 10 mins
-# release version: v0.6.11.post2
+# release version: v0.6.12
 # todo(elainewy): cache flashinfer build result for faster build
 ENV CCACHE_DIR=/root/.cache/ccache
 RUN --mount=type=cache,target=/root/.cache/ccache \
     --mount=type=cache,target=/root/.cache/uv \
     echo "git clone flashinfer..." \
-    && git clone --depth 1 --branch v0.6.11.post2 --recursive https://github.com/flashinfer-ai/flashinfer.git \
+    && git clone --depth 1 --branch v0.6.12 --recursive https://github.com/flashinfer-ai/flashinfer.git \
     && cd flashinfer \
     && git submodule update --init --recursive \
     && echo "finish git clone flashinfer..." \

docker/Dockerfile.rocm

@@ -135,13 +135,18 @@ ENV PATH="/root/.cargo/bin:${PATH}"
 # Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit
 # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
 ENV CARGO_BUILD_JOBS=4
+ENV CARGO_NET_RETRY=10
+ENV RUSTUP_MAX_RETRIES=10
 
-# Build the release binary. Cache cargo registry/git, and copy the binary out
-# so it persists into the image layer for later COPY --from=rust-build.
-RUN --mount=type=cache,target=/root/.cargo/registry \
-    --mount=type=cache,target=/root/.cargo/git \
+# Build the release binary. Cargo's registry/git caches can be written by
+# concurrent BuildKit jobs on shared workers, so lock those cache mounts while
+# keeping the cache benefit. Copy the binary out so it persists into the image
+# layer for later COPY --from=rust-build.
+RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \
+    --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \
     cd ${COMMON_WORKDIR}/vllm \
-    && VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh
+    && VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh \
+    && test -x /tmp/vllm-rs
 
 # -----------------------
 # vLLM build stages

docker/Dockerfile.rocm_base

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.2-complete
+ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete
 ARG TRITON_BRANCH="ba5c1517"
 ARG TRITON_REPO="https://github.com/ROCm/triton.git"
 ARG PYTORCH_BRANCH="8514f051" # release/2.10 as of 3/17
@@ -104,28 +104,6 @@ ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
 ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
 ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
 
-# torch profiler hotfix for 7.2.2: rebuild CLR with https://github.com/ROCm/rocm-systems/pull/5062
-# will be removed once we move to ROCm 7.2.3
-RUN apt-get update && apt-get install -y rocm-llvm-dev
-RUN pip install CppHeaderParser
-RUN git clone --no-checkout --filter=blob:none https://github.com/ROCm/rocm-systems /tmp/rocm-systems \
-    && cd /tmp/rocm-systems \
-    && git sparse-checkout init --cone \
-    && git sparse-checkout set projects/hip projects/clr \
-    && git checkout 35e8c7bf8911862e5389509800e65fdf125412b3 \
-    && export CLR_DIR=/tmp/rocm-systems/projects/clr \
-    && export HIP_DIR=/tmp/rocm-systems/projects/hip \
-    && mkdir -p $CLR_DIR/build && cd $CLR_DIR/build \
-    && cmake \
-        -DHIP_COMMON_DIR=$HIP_DIR \
-        -DCMAKE_PREFIX_PATH="/opt/rocm/" \
-        -DCLR_BUILD_HIP=ON \
-        -DCLR_BUILD_OCL=OFF \
-        -DHIP_PLATFORM=amd \
-        .. \
-    && make -j$(nproc) \
-    && make install \
-    && rm -rf /tmp/rocm-systems
 
 ###
 ### Triton Build
@@ -259,9 +237,8 @@ ARG AITER_REPO
 ARG USE_SCCACHE
 RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
     pip install /install/*.whl
-RUN git clone --recursive ${AITER_REPO}
+RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO}
 RUN cd aiter \
-    && git checkout ${AITER_BRANCH} \
     && git submodule update --init --recursive \
     && pip install -r requirements.txt
 RUN pip install pyyaml && cd aiter \

docker/versions.json

@@ -68,7 +68,7 @@
       "default": "true"
     },
     "FLASHINFER_VERSION": {
-      "default": "0.6.11.post2"
+      "default": "0.6.12"
     },
     "GDRCOPY_CUDA_VERSION": {
       "default": "12.8"

docs/benchmarking/cli.md

@@ -231,13 +231,21 @@ vllm bench serve \
 
 #### Custom Image Dataset
 
-If the image dataset you want to benchmark is not supported yet in vLLM, then you can benchmark on it using `CustomImageDataset`. At inference time, use the option `--dataset-name custom_image`. Your data needs to be in the `.jsonl` format and needs to have "prompt" and "image_files" fields per entry, e.g., `image_data.jsonl`:
+If the image dataset you want to benchmark is not supported yet in vLLM, then you can benchmark on it using `CustomImageDataset`. At inference time, use the option `--dataset-name custom_image`. Your data needs to be in the `.jsonl` format and can use "prompt" and "image_files" fields per entry, e.g., `image_data.jsonl`:
 
 ```json
 {"prompt": "How many animals are present in the given image?", "image_files": ["/path/to/image/folder/horsepony.jpg"]}
 {"prompt": "What colour is the bird shown in the image?", "image_files": ["/path/to/image/folder/flycatcher.jpeg"]}
 ```
 
+Every image listed in "image_files" is added to the request in the listed order after the prompt text. To preserve an interleaved order of text and images, use a "content" field with OpenAI-compatible content parts:
+
+```json
+{"content": [{"type": "text", "text": "Compare "}, {"type": "image", "image": "/path/to/image/folder/chart_a.png"}, {"type": "text", "text": " with "}, {"type": "image_url", "image_url": {"url": "/path/to/image/folder/chart_b.png"}}]}
+```
+
+The "image" shorthand accepts the same values as "image_files". The "image_url" field accepts either an OpenAI-style object with a "url" field or a URL string.
+
 ```bash
 # need a model with vision capability here
 vllm serve Qwen/Qwen2-VL-7B-Instruct
@@ -910,6 +918,41 @@ vllm bench serve \
 
 </details>
 
+### Replay Timed Traces
+
+<details class="admonition abstract" markdown="1">
+<summary>Show more</summary>
+
+Example of how to run traces which have timing information
+with them.
+
+#### Running MoonshotAI traces
+
+Start the server:
+
+```bash
+vllm serve Qwen/Qwen3.5-2B \
+--host 127.0.0.1 --port 8000
+```
+
+Run the benchmark:
+
+```bash
+# Download an example trace 
+# curl -L -o conversation_trace.jsonl \
+#https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/conversation_trace.jsonl 
+
+vllm bench serve --model Qwen/Qwen3.5-2B \  
+--dataset-name=timed_trace --num-prompts 100 --host 127.0.0.1 \
+--port 8000 --dataset-path ./conversation_trace.jsonl \
+--ignore-eos  --self-timed --timed-trace-chunk-hash-size 512 \
+--timed-trace-sec-multiplier 0.001 
+```
+
+This will replay the first 100 lines from the trace file `conversation.jsonl`.  
+
+</details>
+
 ### 🧪 Hashing Benchmarks
 
 <details class="admonition abstract" markdown="1">

docs/configuration/optimization.md

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

docs/contributing/profiling.md

@@ -201,43 +201,45 @@ The profiling traces generated by the continuous profiling workflow are publicly
 
 The Python standard library includes
 [cProfile](https://docs.python.org/3/library/profile.html) for profiling Python
-code. vLLM includes a couple of helpers that make it easy to apply it to a section of vLLM.
-Both the `vllm.utils.profiling.cprofile` and `vllm.utils.profiling.cprofile_context` functions can be
-used to profile a section of code.
+code.
 
-!!! note
-    The `vllm.utils.profiling` helpers are deprecated and will be removed in
-    `v0.21`. Please use Python's `cProfile` module directly instead.
+### Example usage - function call
 
-### Example usage - decorator
-
-The first helper is a Python decorator that can be used to profile a function.
-If a filename is specified, the profile will be saved to that file. If no filename is
-specified, profile data will be printed to stdout.
+If a filename is specified, the profile will be saved to that file. If no
+filename is specified, profile data can be printed to stdout.
 
 ```python
-from vllm.utils.profiling import cprofile
+import cProfile
+
 
-@cprofile("expensive_function.prof")
 def expensive_function():
     # some expensive code
     pass
+
+
+profiler = cProfile.Profile()
+profiler.runcall(expensive_function)
+profiler.dump_stats("expensive_function.prof")
 ```
 
-### Example Usage - context manager
-
-The second helper is a context manager that can be used to profile a block of
-code. Similar to the decorator, the filename is optional.
+### Example usage - context manager style
 
 ```python
-from vllm.utils.profiling import cprofile_context
+import cProfile
+
 
 def another_function():
     # more expensive code
     pass
 
-with cprofile_context("another_function.prof"):
+
+profiler = cProfile.Profile()
+profiler.enable()
+try:
     another_function()
+finally:
+    profiler.disable()
+    profiler.dump_stats("another_function.prof")
 ```
 
 ### Analyzing Profile Results

docs/design/attention_backends.md

@@ -170,14 +170,14 @@ Priority is **1 = highest** (tried first).
 | Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
 | `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
-| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
-| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
+| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
+| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
 | `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
 | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
 | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
 | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
-| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
+| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
 | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
 | `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | Any |
@@ -187,6 +187,18 @@ Priority is **1 = highest** (tried first).
 >
 > **\*** Specify the FlashAttention version via `--attention-config.flash_attn_version=2`, `3`, or `4`. Default is FA4 on SM100+ (Blackwell), FA3 on SM90 (Hopper), FA2 otherwise.
 
+## MiniMax M3 Sparse Attention Backends
+
+Block-sparse GQA backend used by MiniMax M3 sparse ("lightning indexer")
+layers. It is wired in directly by the model and is not part of the
+automatic priority lists above. A lightning indexer scores KV blocks, the
+top-k blocks (plus fixed init/local blocks) are selected, and attention
+attends only to those blocks; index keys live in a separate side cache.
+
+| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
+| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
+| `MINIMAX_M3_SPARSE` | bf16, fp16 | `bfloat16` | 128 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
+
 ## MLA (Multi-head Latent Attention) Backends
 
 MLA uses separate backends for prefill and decode phases.
@@ -205,8 +217,9 @@ hardware and configuration.
 | `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | DeepSeek R1 dims only |
 | `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | DeepSeek R1 dims only |
 
-> **‡** TRT-LLM Ragged is the default on Blackwell (SM100).
-> On other GPUs, FlashAttention is used as the default.
+> **‡** Automatic selection tries FlashAttention first. On Blackwell
+> (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then
+> TokenSpeed MLA. On other GPUs, only FlashAttention is considered.
 
 ### Decode Backends
 

docs/design/fusions.md

@@ -21,7 +21,6 @@ or just on the low or high end.
 | Fusion                                                                         | `PassConfig` flag            | Fused operations                               | Default at                     | E2E Speedup        | Fullgraph | `num_tokens` |
 | ------------------------------------------------------------------------------ | ---------------------------- | ---------------------------------------------- | ------------------------------ | ------------------ | --------- | ------------ |
 | [AllReduce + RMSNorm](#allreduce--rmsnorm-fuse_allreduce_rms)                  | `fuse_allreduce_rms`         | All-reduce → RMSNorm (+residual_add) (→ quant) | O2 (Hopper/Blackwell + TP > 1) | 5-20%              | No        | Low          |
-| [MiniMax QK Norm](#minimax-qk-norm-fuse_minimax_qk_norm)                       | `fuse_minimax_qk_norm`       | Q/K variance all-reduce → Q/K RMSNorm          | Off by default                 | 2-3%               | No        | Low          |
 | [Attention + Quant](#attention--quantization-fuse_attn_quant)                  | `fuse_attn_quant`            | Attention output → FP8/NVFP4 quant             | Off by default                 | 3-7%               | Yes       | Always       |
 | [MLA Attention + Quant](#attention--quantization-fuse_attn_quant)              | `fuse_attn_quant`            | MLA Attention output → FP8/NVFP4 quant         | Off by default                 | TBD                | Yes       | Always       |
 | [RoPE + KV-Cache Update](#rope--kv-cache-update-fuse_rope_kvcache)             | `fuse_rope_kvcache`          | Rotary embedding → KV cache write              | O2 (ROCm/AITER only)           | 2-4%               | No        | Low          |
@@ -42,7 +41,6 @@ The table below lists the quantization schemes supported by each fusion on each
 | Fusion                       | SM100 (Blackwell)                        | SM90 (Hopper)                            | SM89 (Ada)                               | SM80 (Ampere) | ROCm                                     |
 | ---------------------------- | ---------------------------------------- | ---------------------------------------- | ---------------------------------------- | ------------- | ---------------------------------------- |
 | `fuse_allreduce_rms`         | FP16/BF16, FP8 static, NVFP4             | FP16/BF16, FP8 static                    | —                                        | —             | —                                        |
-| `fuse_minimax_qk_norm`\*     | FP16/BF16                                | FP16/BF16                                | FP16/BF16                                | FP16/BF16     | —                                        |
 | `fuse_attn_quant`\*          | FP8 static\*, NVFP4\*                    | FP8 static\*                             | FP8 static\*                             | —             | FP8 static\*                             |
 | `fuse_attn_quant` (MLA)\*    | FP8 static\*, FP8 per-group\*, NVFP4\*   | FP8 static\*, FP8 per-group\*            | FP8 static\*, FP8 per-group\*            | —             | FP8 static\* (untested)                  |
 | `fuse_rope_kvcache`          | —                                        | —                                        | —                                        | —             | FP16/BF16                                |
@@ -58,9 +56,6 @@ The table below lists the quantization schemes supported by each fusion on each
 fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant)
 for per-backend details.
 
-\* `fuse_minimax_qk_norm` is a model-specific pass for `MiniMaxM2ForCausalLM`. It also requires
-tensor parallelism (`tp_size > 1`) and the CUDA custom op `minimax_allreduce_rms_qk`.
-
 † `enable_sp` and `fuse_gemm_comms` are only autoconfigured for SM90 today;
 other architectures support requires setting `PassConfig.sp_min_token_num` explicitly.
 SM100 support also requires setting `VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel`.
@@ -191,35 +186,6 @@ If these conditions are set, the fusion is enabled automatically for optimizatio
 
 - Pass: [`vllm/compilation/passes/fusion/rope_kvcache_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rope_kvcache_fusion.py)
 
-### MiniMax QK Norm (`fuse_minimax_qk_norm`)
-
-!!! info
-    This is a MiniMax-specific compile pass. It is currently only enabled when all of the following hold:
-    the model architecture is `MiniMaxM2ForCausalLM`, tensor parallelism is enabled (`tp_size > 1`),
-    and the CUDA custom op `minimax_allreduce_rms_qk` is available. It is not enabled by default at any
-    optimization level.
-
-**What it fuses.** Fuses the MiniMax M2 Q/K normalization path that performs an all-reduce over the
-per-token Q/K variances before applying RMS normalization to Q and K.
-
-This pass is distinct from [`enable_qk_norm_rope_fusion`](#qk-norm--rope-enable_qk_norm_rope_fusion):
-`fuse_minimax_qk_norm` targets MiniMax M2's tensor-parallel all-reduce + RMSNorm sequence, while
-`enable_qk_norm_rope_fusion` targets the later Q/K RMSNorm + RoPE sequence used by several other models.
-
-Example:
-
-```bash
-vllm serve MiniMaxAI/MiniMax-M2.5 \
-  --tensor-parallel-size 4 \
-  --compilation-config '{"mode": 3, "pass_config": {"fuse_minimax_qk_norm": true}}'
-```
-
-**Code locations.**
-
-- Pass: [`vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py)
-- CUDA op: [`csrc/minimax_reduce_rms_kernel.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/minimax_reduce_rms_kernel.cu) (`minimax_allreduce_rms_qk`)
-- Workspace helper: [`vllm/model_executor/layers/mamba/lamport_workspace.py`](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/mamba/lamport_workspace.py)
-
 ### Sequence Parallelism (`enable_sp`)
 
 **What it fuses.** Replaces all-reduce collectives with reduce-scatter + local RMSNorm + all-gather,

docs/features/disagg_prefill.md

@@ -19,25 +19,25 @@ Two main reasons:
 
 Please refer to [examples/disaggregated/disaggregated_prefill.sh](../../examples/disaggregated/disaggregated_prefill.sh) for the example usage of disaggregated prefilling.
 
-Now supports 6 types of connectors:
+Now supports 9 types of connectors:
 
 - **ExampleConnector**: refer to [examples/disaggregated/example_connector/run.sh](../../examples/disaggregated/example_connector/run.sh) for the example usage of ExampleConnector disaggregated prefilling.
 - **LMCacheConnectorV1**: refer to [examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission.
-- **NixlConnector**: refer to [tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh](../../tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh) for the example usage of NixlConnector disaggregated prefilling which support fully async send/recv. For detailed usage guide, see [NixlConnector Usage Guide](nixl_connector_usage.md). For feature compatibility details, see [NixlConnector Compatibility Matrix](nixl_connector_compatibility.md).
+- **NixlConnector**: refer to [tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh](../../tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh) for the example usage of NixlConnector disaggregated prefilling which support fully async send/recv. For detailed usage guide, see [NixlConnector Usage Guide](nixl_connector_usage.md). For feature compatibility details, see [NixlConnector Compatibility Matrix](nixl_connector_compatibility.md). You may specify one or multiple NIXL transfer backends, such as:
+
+  ```bash
+  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both", "kv_buffer_device":"cuda", "kv_connector_extra_config":{"backends":["UCX", "GDS"]}}'
+  ```
+
 - **P2pNcclConnector**: refer to [examples/disaggregated/p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh](../../examples/disaggregated/p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh) for the example usage of P2pNcclConnector disaggregated prefilling.
 - **MooncakeConnector**: refer to [examples/disaggregated/mooncake_connector/run_mooncake_connector.sh](../../examples/disaggregated/mooncake_connector/run_mooncake_connector.sh) for the example usage of MooncakeConnector disaggregated prefilling. For detailed usage guide, see [MooncakeConnector Usage Guide](mooncake_connector_usage.md).
+- **MoRIIOConnector** (ROCm only): see [MoRI-IO Usage Guide](moriio_connector_usage.md) for example usage and detailed documentation.
 - **MultiConnector**: take advantage of the kv_connector_extra_config: dict[str, Any] already present in KVTransferConfig to stash all the connectors we want in an ordered list of kwargs.such as:
 
   ```bash
   --kv-transfer-config '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both"},{"kv_connector":"ExampleConnector","kv_role":"kv_both","kv_connector_extra_config":{"shared_storage_path":"local_storage"}}]}}'
   ```
 
-For NixlConnector, you may also specify one or multiple NIXL_Backend. Such as:
-
-  ```bash
-  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both", "kv_buffer_device":"cuda", "kv_connector_extra_config":{"backends":["UCX", "GDS"]}}'
-  ```
-
 - **OffloadingConnector**: enable offloading of KV data to CPU memory, customizing the CPU block size (in tokens) and total CPU memory bytes to allocate:
 
   ```bash

docs/features/mooncake_connector_usage.md

@@ -31,7 +31,7 @@ vllm serve Qwen/Qwen2.5-7B-Instruct --port 8020 --kv-transfer-config '{"kv_conne
 ### Proxy
 
 ```bash
-python examples/disaggregated/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py --prefill http://192.168.0.2:8010 --decode http://192.168.0.3:8020
+python examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py --prefill http://192.168.0.2:8010 --decode http://192.168.0.3:8020
 ```
 
 Now you can send requests to the proxy server through port 8000.

docs/features/moriio_connector_usage.md

@@ -0,0 +1,264 @@
+# MoRIIOConnector Usage Guide
+
+`MoRIIOConnector` is a high-performance KV connector used for KV cache transfer in PD disaggregated deployments, built on ROCm's [MoRI-IO](https://github.com/rocm/mori) communication library for point-to-point communication with ultra-low overhead.
+
+## Prerequisites
+
+### Installation
+
+**Docker:** MoRI is shipped with the official ROCm vLLM image: `vllm/vllm-openai-rocm:nightly`.
+
+**Manual installation:** MoRI wheel can be installed with
+
+```bash
+pip install amd_mori
+```
+
+Refer to the [Dockerfile.rocm_base](../../docker/Dockerfile.rocm_base) for more information, or [official MoRI repository](https://github.com/rocm/mori) for instructions on how to build MoRI from source.
+
+For instructions on installing appropriate NIC userspace libraries, see [Installing NIC userspace libraries](#appendix-installing-nic-userspace-libraries).
+
+## Basic usage (single host)
+
+Start the proxy first; the producer and consumer instances will retry registration until the proxy is reachable.
+
+### Producer (prefiller) configuration
+
+Start a prefiller instance that produces KV caches
+
+```bash
+# Prefill instance (GPU 0-3) 
+export VLLM_ROCM_USE_AITER=1
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+export HIP_VISIBLE_DEVICES=0,1,2,3
+ 
+vllm serve Qwen/Qwen3-235B-A22B-FP8 \
+  -tp 4 \
+  --port 20005 \
+  --gpu-memory-utilization 0.9 \
+  --kv-transfer-config '{
+    "kv_connector": "MoRIIOConnector",
+    "kv_role": "kv_producer",
+    "kv_connector_extra_config": {
+      "proxy_ip": "127.0.0.1",
+      "proxy_ping_port": "36367",
+      "http_port": "20005",
+      "handshake_port": "6301",
+      "notify_port": "6105"
+    }
+  }'
+```
+
+### Consumer (decoder) configuration
+
+Start a decoder instance that consumes KV caches:
+
+```bash
+# Decode instance (GPU 4-7)
+export VLLM_ROCM_USE_AITER=1
+export CUDA_VISIBLE_DEVICES=4,5,6,7
+export HIP_VISIBLE_DEVICES=4,5,6,7
+
+vllm serve Qwen/Qwen3-235B-A22B-FP8 \
+  -tp 4 \
+  --port 40005 \
+  --gpu-memory-utilization 0.9 \
+  --kv-transfer-config '{
+    "kv_connector": "MoRIIOConnector",
+    "kv_role": "kv_consumer",
+    "kv_connector_extra_config": {
+      "proxy_ip": "127.0.0.1",
+      "http_port": "40005",
+      "proxy_ping_port": "36367",
+      "handshake_port": "7301",
+      "notify_port": "7501"
+    }
+  }'
+```
+
+### Proxy server
+
+The proxy fronts the producer and consumer instances and routes incoming requests to them. `vllm-router` is the recommended proxy; it can be installed manually or run as a Docker container. Note that the port `36367` below is the `proxy_ping_port` configured on each vLLM instance.
+
+**Docker:**
+
+```bash
+docker run \
+  --network host \
+  vllm/vllm-router:nightly \
+  vllm-router \
+  --vllm-pd-disaggregation \
+  --kv-connector moriio \
+  --vllm-discovery-address "0.0.0.0:36367"
+```
+
+**Manual install:**
+
+```bash
+pip install vllm-router
+vllm-router \
+  --vllm-pd-disaggregation \
+  --kv-connector moriio \
+  --vllm-discovery-address "0.0.0.0:36367"
+```
+
+Alternatively, you can use the reference implementation proxy shipped with vLLM:
+
+```bash
+cd <path_to>/vllm
+pip install quart aiohttp msgpack
+python examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
+```
+
+## Configuration
+
+The connector is configured at two levels: the application level and the transport level.
+
+### Application-level configuration
+
+**Modes:** MoRI has two modes of operation: WRITE and READ mode.
+
+- In WRITE mode, the producer actively pushes computed KV blocks after every layer into the consumer's memory.
+- In READ mode, the consumer pulls the KV blocks from the producer all at once, as soon as it has been notified those blocks are ready.
+
+WRITE mode is used by default. READ mode can be configured by setting `--kv-transfer-config.kv_connector_extra_config.read_mode true`.
+
+**Control-plane configuration:** MoRI moves KV bytes over RDMA/xGMI, but producers and consumers also need out-of-band TCP channels for handshake, block id exchange, liveness, and completion signaling. These keys live under `kv_connector_extra_config`:
+
+- `proxy_ip`: IP address of the disaggregation proxy/router that fronts the prefiller and decoder. Each vLLM instance uses it to register itself and to send heartbeats so the proxy knows where to route incoming requests.
+- `proxy_ping_port`: TCP port on `proxy_ip` where the proxy listens for instance heartbeats and registration messages. Used to detect dead vLLM instances and keep routing tables fresh.
+- `http_port`: HTTP port that this vLLM instance exposes its OpenAI-compatible API on. The proxy registers this port, and forwards user requests to this port once it has picked an instance.
+- `handshake_port`: TCP port used for the one-time MoRI engine handshake between a prefiller and a decoder. The two sides exchange RDMA engine descriptors here before any KV transfer can happen.
+- `notify_port`: TCP port used for control and synchronization messages between prefiller and decoder. Used differently in the two modes:
+    - WRITE mode: **Block allocation:** the decoder notifies the prefiller about its block ids, so the prefiller can push its computed KV blocks into the correct place on the decoder instance. **Completion:** once all blocks have been transferred, the prefiller notifies the decoder that it's safe to use its blocks.
+    - READ mode: **Completion:** once the decoder has read all blocks from the prefiller, it notifies the prefiller so it can free its KV cache blocks.
+
+!!! note
+    `notify_port` is used as a *base* port: each (DP rank, TP rank) pair within an instance uses `notify_port + offset` where the offset is based on the rank. Make sure the range starting at `notify_port` is free on the host.
+
+### Transport configuration
+
+MoRI has two transport backends: RDMA and xGMI. You can select backend using `--kv-transfer-config.kv_connector_extra_config.backend $BACKEND`, with `$BACKEND` being `rdma` or `xgmi`. RDMA is the default backend and should be used in multi-node deployments.
+
+The configuration options for each backend are as follows.
+
+#### RDMA backend
+
+- `qp_per_transfer`: number of RDMA Queue Pairs (QPs) used per transfer. More QPs let a single transfer be striped over multiple QPs to increase NIC concurrency, at the cost of more RDMA resources.
+- `post_batch_size`: how many RDMA Work Requests (WR) are batched into one `ibv_post_send` doorbell. Defaults to -1, meaning the backend default. Larger batches reduce the posting overhead per WR.
+- `num_workers`: number of worker threads MoRI uses to post and poll transfer completions.
+
+Advanced users can also configure MoRI itself using environment variables such as `MORI_IO_QP_MAX_SEND_WR`, `MORI_IO_QP_MAX_CQE`, etc. These are MoRI library variables and are separate from vLLM's own `VLLM_MORIIO_*` settings. Refer to the [MoRI repository](https://github.com/rocm/mori) for more information.
+
+#### xGMI backend
+
+Use xGMI when the prefiller and decoder run on the same physical host so transfers go over the AMD GPU fabric and skip the NIC entirely. Currently only configured using MoRI-specific environment variables; see the [MoRI repository](https://github.com/rocm/mori).
+
+## Multi-node deployment
+
+The example below shows how to run a 1P1D deployment on two nodes. We run the proxy on the same node as the prefill instance.
+
+### On both nodes
+
+```bash
+# Set on both nodes before running any command
+export PREFILL_IP=<node1-ip>
+export DECODE_IP=<node2-ip>
+```
+
+### On node 1
+
+Start the proxy first as described in [Proxy server](#proxy-server), then start the prefill instance:
+
+```bash
+docker run \
+  --name moriio-prefill \
+  --init --network host --ipc host --privileged \
+  --security-opt seccomp=unconfined \
+  --ulimit memlock=-1 --ulimit stack=67108864 --shm-size 256G \
+  --group-add video --group-add render \
+  --device /dev/kfd --device /dev/dri --device /dev/infiniband \
+  -e VLLM_ROCM_USE_AITER=1 \
+  vllm/vllm-openai-rocm:nightly \
+  deepseek-ai/DeepSeek-R1-0528 \
+    --port 8100 \
+    --tensor-parallel-size 8 \
+    --enable-expert-parallel \
+    --gpu-memory-utilization 0.8 \
+    --trust-remote-code \
+    --kv-transfer-config '{
+      "kv_connector": "MoRIIOConnector",
+      "kv_role": "kv_producer",
+      "kv_connector_extra_config": {
+        "proxy_ip": "'"${PREFILL_IP}"'",
+        "proxy_ping_port": "36367",
+        "http_port": "8100",
+        "handshake_port": "6301",
+        "notify_port": "61005"
+      }
+    }'
+```
+
+### On node 2
+
+Decode instance:
+
+```bash
+docker run \
+  --name moriio-decode \
+  --init --network host --ipc host --privileged \
+  --security-opt seccomp=unconfined \
+  --ulimit memlock=-1 --ulimit stack=67108864 --shm-size 256G \
+  --group-add video --group-add render \
+  --device /dev/kfd --device /dev/dri --device /dev/infiniband \
+  -e VLLM_ROCM_USE_AITER=1 \
+  vllm/vllm-openai-rocm:nightly \
+  deepseek-ai/DeepSeek-R1-0528 \
+    --port 8200 \
+    --tensor-parallel-size 8 \
+    --gpu-memory-utilization 0.8 \
+    --trust-remote-code \
+    --enable-expert-parallel \
+    --kv-transfer-config '{
+      "kv_connector": "MoRIIOConnector",
+      "kv_role": "kv_consumer",
+      "kv_connector_extra_config": {
+        "proxy_ip": "'"${PREFILL_IP}"'",
+        "proxy_ping_port": "36367",
+        "http_port": "8200",
+        "handshake_port": "6301",
+        "notify_port": "61005"
+      }
+    }'
+```
+
+## Troubleshooting
+
+### `availDevices.size() > 0` assertion failure
+
+**Problem:** vLLM fails to launch with the following log:
+
+```bash
+libibverbs: Warning: Driver bnxt_re does not support the kernel ABI of 6 (supports 1 to 1) for device /sys/class/infiniband/rdma4
+...
+ker: /app/mori/src/io/rdma/backend_impl.cpp: mori::io::RdmaManager::RdmaManager(const RdmaBackendConfig, application::RdmaContext *): Assertion `availDevices.size() > 0' failed.
+```
+
+**Fix:** The installed RDMA userspace libraries do not match the driver and firmware version installed on the host. You must install NIC userspace libraries corresponding to your RDMA kernel module and firmware version. See [Installing NIC userspace
+libraries](#appendix-installing-nic-userspace-libraries) for more information.
+
+## Appendix: installing NIC userspace libraries
+
+To run MoRI with RDMA, your environment must have the necessary RDMA userspace libraries installed that match the associated kernel module and firmware version.
+
+The official image `vllm/vllm-openai-rocm:nightly` comes pre-installed with userspace libraries for the following NICs and kernel module versions:
+
+- AINIC (AMD Pensando Pollara): version `1.117.3-hydra`, tested with `ioinic-dkms=25.11.1.001`
+- Thor2 (Broadcom): version `235.2.86.0`, tested with `bnxt-en-dkms=1.10.3.235.2.86.0`, `bnxt-re-dkms=235.2.86.0`
+
+Refer to [Dockerfile.rocm](../../docker/Dockerfile.rocm) for more details. For users with NICs, kernel modules, and/or FW other than those stated above we refer to
+the vendors' own installation instructions.
+
+## Further reading
+
+- [Next-Level Inference: Why Your Single-Node vLLM Setup Needs Prefill-Decode Disaggregation](https://vllm.ai/blog/2026-04-07-moriio-kv-connector).

docs/features/reasoning_outputs.md

@@ -15,6 +15,7 @@ vLLM currently supports the following reasoning models:
 | ------------ | ----------- | ---------------- | ----------- |
 | [Cohere Command A Reasoning](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025) | `cohere_command3` | `json`, `regex` | ✅ |
 | [DeepSeek R1 series](https://huggingface.co/collections/deepseek-ai/deepseek-r1-678e1e131c0169c0bc89728d) | `deepseek_r1` | `json`, `regex` | ❌ |
+| [Gemma 4 series](https://huggingface.co/google/gemma-4-26B-A4B-it) | `gemma4` | `json`, `regex` | ✅ |
 | [DeepSeek-V3.1](https://huggingface.co/collections/deepseek-ai/deepseek-v31-68a491bed32bd77e7fca048f) | `deepseek_v3` | `json`, `regex` | ❌ |
 | [ERNIE-4.5-VL series](https://huggingface.co/baidu/ERNIE-4.5-VL-28B-A3B-PT) | `ernie45` | `json`, `regex` | ❌ |
 | [ERNIE-4.5-21B-A3B-Thinking](https://huggingface.co/baidu/ERNIE-4.5-21B-A3B-Thinking) | `ernie45` | `json`, `regex` | ✅ |
@@ -29,6 +30,7 @@ vLLM currently supports the following reasoning models:
 !!! note
     IBM Granite 3.2 and DeepSeek-V3.1 reasoning is disabled by default; to enable it, you must also pass `thinking=True` in your `chat_template_kwargs`.
     The reasoning feature for the Qwen3 series is enabled by default. To disable it, you must pass `enable_thinking=False` in your `chat_template_kwargs`.
+    Gemma 4 reasoning is disabled by default; to enable it, pass `enable_thinking=True` in your `chat_template_kwargs` or set `reasoning_effort` (which enables it automatically).
     DeepSeek-V3.1 tool calling is supported in non-thinking mode.
     Holo2 reasoning is enabled by default. To disable it, you must also pass `thinking=False` in your `chat_template_kwargs`.
 
@@ -314,9 +316,44 @@ for output in outputs:
     print("text:", output.outputs[0].text)
 ```
 
+## Automatic `enable_thinking` Activation
+
+Some models (such as Gemma 4, DeepSeek-V4-Pro and IBM Granite 3.2) require `enable_thinking: true` in their chat template kwargs to activate thinking mode — without it, reasoning tokens are never generated regardless of other settings.
+
+When you set `reasoning_effort` in a Chat Completions request (or `reasoning.effort` in a Responses API request), vLLM automatically injects `enable_thinking` into the chat template kwargs:
+
+- `reasoning_effort` = `"low"`, `"medium"`, or `"high"` → `enable_thinking = true`
+- `reasoning_effort` = `"none"` → `enable_thinking = false`
+- `reasoning_effort` not set → `enable_thinking` is not injected (preserves existing behavior)
+
+This means you no longer need to manually pass `chat_template_kwargs: {"enable_thinking": true}` when using `reasoning_effort` — it is handled automatically.
+
+!!! note
+    If you explicitly set `enable_thinking` in `chat_template_kwargs`, your value takes priority over the automatic injection. This allows you to override the behavior if needed.
+
+    For models whose templates don't declare `enable_thinking` (e.g., DeepSeek R1), the injected kwarg is harmlessly filtered out by `resolve_chat_template_kwargs`.
+
+### Example
+
+```python
+from openai import OpenAI
+
+client = OpenAI(base_url="http://localhost:8000/v1", api_key="dummy")
+
+# reasoning_effort automatically enables thinking for models that need it
+response = client.chat.completions.create(
+    model="google/gemma-4-26B-A4B-it",
+    messages=[{"role": "user", "content": "What is 15 * 37?"}],
+    reasoning_effort="high",  # Automatically sets enable_thinking=true
+)
+
+print(response.choices[0].message.reasoning)
+print(response.choices[0].message.content)
+```
+
 ## Limitations
 
-- The reasoning content is only available for online serving's chat completion endpoint (`/v1/chat/completions`).
+- The reasoning content is only available for online serving's chat completion endpoint (`/v1/chat/completions`), Anthropic Messages API (`/v1/messages`) and the Responses API (`/v1/responses`).
 
 ## How to support a new reasoning model
 

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

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

docs/getting_started/quickstart.md

@@ -76,6 +76,15 @@ This guide will help you quickly get started with vLLM to perform:
     !!! note
         For more detailed instructions, including Docker, installing from source, and troubleshooting, please refer to the [vLLM on TPU documentation](https://docs.vllm.ai/projects/tpu/en/latest/).
 
+=== "Ascend NPU"
+
+    If you are using Ascend NPUs, you can run vLLM through [vLLM Ascend](https://github.com/vllm-project/vllm-ascend), a community-maintained hardware plugin.
+
+    Follow the installation instructions in the [vLLM Ascend quick start](https://docs.vllm.ai/projects/ascend/en/latest/quick_start.html).
+
+    !!! note
+        Ascend setup depends on your NPU hardware and CANN version. For supported versions, Docker images, and troubleshooting, please refer to the [vLLM Ascend documentation](https://docs.vllm.ai/projects/ascend/en/latest/).
+
 === "Apple Silicon (Mac)"
 
     If you are using Apple Silicon Macs, you can use vLLM-Metal for GPU-accelerated inference via Apple's Metal framework.

docs/governance/committers.md

@@ -17,6 +17,7 @@ Sorted alphabetically by GitHub handle:
 - [@bbrowning](https://github.com/bbrowning): Tool use and reasoning parser
 - [@benchislett](https://github.com/benchislett): Engine core and spec decode
 - [@bigPYJ1151](https://github.com/bigPYJ1151): Intel CPU/XPU integration
+- [@BugenZhao](https://github.com/BugenZhao): Rust frontend
 - [@chaunceyjiang](https://github.com/chaunceyjiang): Tool use and reasoning parser
 - [@DarkLight1337](https://github.com/DarkLight1337): Multimodality, API server
 - [@esmeetu](https://github.com/esmeetu): developer marketing, community
@@ -130,6 +131,8 @@ If you have PRs touching the area, please feel free to ping the area owner for r
     - @DarkLight1337
 - API Server: The OpenAI-compatible API server
     - @DarkLight1337, @njhill, @aarnphm, @simon-mo, @heheda12345 (Responses API)
+- Rust Frontend: The experimental API server in Rust
+    - @BugenZhao, @njhill
 - Batch Runner: The OpenAI-compatible batch runner
     - @simon-mo
 

docs/mkdocs/hooks/url_schemes.py

@@ -106,8 +106,8 @@ class UrlSchemesPreprocessor(Preprocessor):
             return f"[{gh_icon} {title}]({url})"
 
         markdown = "\n".join(lines)
-        markdown = relative_link.sub(replace_relative_link, markdown)
         markdown = github_link.sub(replace_github_link, markdown)
+        markdown = relative_link.sub(replace_relative_link, markdown)
         return markdown.split("\n")
 
 

docs/models/pooling_models/classify.md

@@ -299,7 +299,3 @@ Example configuration:
 ### Remove softmax from PoolingParams
 
 We have already removed `softmax` and `activation` from PoolingParams. Instead, use `use_activation`, since we allow `classify` and `token_classify` to use any activation function.
-
-### Remove `logit_bias` and `logit_scale`
-
-`logit_bias` and `logit_scale` are deprecated aliases for `logit_mean` and `logit_sigma` respectively. When using `logit_scale`, it is automatically converted to `logit_sigma = 1/logit_scale`. These deprecated parameters will be removed in v0.21.

docs/models/supported_models.md

@@ -428,7 +428,6 @@ th {
 | `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ |
 | `IQuestCoderForCausalLM` | IQuestCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Instruct`, etc. | | |
 | `IQuestLoopCoderForCausalLM` | IQuestLoopCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Loop-Instruct`, etc. | | |
-| `JAISLMHeadModel` | Jais | `inceptionai/jais-13b`, `inceptionai/jais-13b-chat`, `inceptionai/jais-30b-v3`, `inceptionai/jais-30b-chat-v3`, etc. | | ✅︎ |
 | `Jais2ForCausalLM` | Jais2 | `inceptionai/Jais-2-8B-Chat`, `inceptionai/Jais-2-70B-Chat`, etc. | | ✅︎ |
 | `JambaForCausalLM` | Jamba | `ai21labs/AI21-Jamba-1.5-Large`, `ai21labs/AI21-Jamba-1.5-Mini`, `ai21labs/Jamba-v0.1`, etc. | ✅︎ | ✅︎ |
 | `KimiLinearForCausalLM` | Kimi-Linear-48B-A3B-Base, Kimi-Linear-48B-A3B-Instruct | `moonshotai/Kimi-Linear-48B-A3B-Base`, `moonshotai/Kimi-Linear-48B-A3B-Instruct` | | ✅︎ |
@@ -551,6 +550,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `ChameleonForConditionalGeneration` | Chameleon | T + I | `facebook/chameleon-7b`, etc. | | ✅︎ |
 | `CheersForConditionalGeneration` | Cheers | T + I | `ai9stars/Cheers` | | ✅︎ |
 | `Cohere2VisionForConditionalGeneration` | Command A Vision | T + I<sup>+</sup> | `CohereLabs/command-a-vision-07-2025`, etc. | | ✅︎ |
+| `Cosmos3ForConditionalGeneration` | Cosmos3 (understanding tower) | T + I<sup>E+</sup> + V<sup>E+</sup> | `nvidia/Cosmos3-Nano` | | ✅︎ |
 | `DeepseekVLV2ForCausalLM` | DeepSeek-VL2 | T + I<sup>+</sup> | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ |
 | `DeepseekOCRForCausalLM` | DeepSeek-OCR | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR`, etc. | ✅︎ | ✅︎ |
 | `DeepseekOCR2ForCausalLM` | DeepSeek-OCR-2 | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR-2`, etc. | ✅︎ | ✅︎ |
@@ -633,6 +633,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `SmolVLMForConditionalGeneration` | SmolVLM2 | T + I | `SmolVLM2-2.2B-Instruct` | ✅︎ | |
 | `Step3VLForConditionalGeneration` | Step3-VL | T + I<sup>+</sup> | `stepfun-ai/step3` | | ✅︎ |
 | `StepVLForConditionalGeneration` | Step3-VL-10B | T + I<sup>+</sup> | `stepfun-ai/Step3-VL-10B` | | ✅︎ |
+| `Step3p7ForConditionalGeneration` | Step-3.7-Flash | T + I<sup>+</sup> | `stepfun-ai/Step-3.7-Flash` | | ✅︎ |
 | `TarsierForConditionalGeneration` | Tarsier | T + I<sup>E+</sup> | `omni-search/Tarsier-7b`, `omni-search/Tarsier-34b` | | ✅︎ |
 | `Tarsier2ForConditionalGeneration`<sup>^</sup> | Tarsier2 | T + I<sup>E+</sup> + V<sup>E+</sup> | `omni-research/Tarsier2-Recap-7b`, `omni-research/Tarsier2-7b-0115` | | ✅︎ |
 | `UltravoxModel` | Ultravox | T + A<sup>E+</sup> | `fixie-ai/ultravox-v0_5-llama-3_2-1b` | ✅︎ | ✅︎ |

docs/pre_run_check.sh

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

docs/serving/expert_parallel_deployment.md

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