[CI] Filter import-only files using function-level coverage

Skip files where only module-level code ran (imports, class defs)
but no named functions were actually called. Uses the
functions_called field from stripped coverage JSON.

Reduces false-positive mappings by ~78% — e.g. ompmultiprocessing.py
drops from 73 steps to 0 (only used on CPU but imported everywhere).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

.buildkite/hardware_tests/cpu.yaml

@@ -54,6 +54,20 @@ steps:
       pytest -x -v -s tests/models/language/generation -m cpu_model
       pytest -x -v -s tests/models/language/pooling -m cpu_model"
 
+- label: CPU-ModelRunnerV2 Tests
+  depends_on: []
+  device: intel_cpu
+  no_plugin: true
+  soft_fail: true
+  source_file_dependencies:
+  - vllm/v1/worker/cpu/
+  - vllm/v1/worker/gpu/
+  commands:
+    - |
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
+      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"
+
 - label: CPU-Quantization Model Tests
   depends_on: []
   device: intel_cpu

.buildkite/image_build/image_build.yaml

@@ -6,6 +6,48 @@ 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
+    # 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
+    # into the image. Runs as UID 2000:0 via a shell so we can verify
+    # filesystem perms + passwd/group file state + wrapper presence without
+    # triggering vLLM's GPU-requiring config-init path. The opt-in
+    # `vllm-openai-nonroot` target adds only `USER vllm`, `WORKDIR
+    # /home/vllm`, and an `ENTRYPOINT` override on top of these invariants;
+    # its build correctness is reviewed at the Dockerfile level. Wrapper
+    # logic is covered separately by the pre-commit hook
+    # `test-nonroot-entrypoint` (see .pre-commit-config.yaml).
+    - |
+      docker run --rm --user 2000:0 --entrypoint /bin/sh "$IMAGE_TAG" -ec '
+        if ! getent passwd 2000 | grep -q ^vllm:; then
+          echo FAIL: UID 2000 != vllm
+          exit 1
+        fi
+        if ! id -gn 2>/dev/null | grep -qx root; then
+          echo FAIL: GID 0 not root group
+          exit 1
+        fi
+        touch /home/vllm/.smoke && rm /home/vllm/.smoke
+        touch /opt/uv/cache/.smoke && rm /opt/uv/cache/.smoke
+        if ! test -x /usr/local/bin/vllm-nonroot-entrypoint.sh; then
+          echo FAIL: wrapper missing
+          exit 1
+        fi
+        if ! test -w /etc/passwd; then
+          echo FAIL: /etc/passwd not group-writable
+          exit 1
+        fi
+        if ! test -w /etc/group; then
+          echo FAIL: /etc/group not group-writable
+          exit 1
+        fi
+        echo non-root invariants OK
+      '
     retry:
       automatic:
         - exit_status: -1  # Agent was lost

.buildkite/image_build/image_build_hpu.sh

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

.buildkite/image_build/image_build_xpu.sh

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

.buildkite/intel_jobs/test-intel.yaml

@@ -66,3 +66,22 @@ steps:
         pytest -v -s v1/test_serial_utils.py &&
         pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py &&
         pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py'
+  - label: "XPU server test"
+    depends_on:
+      - image-build-xpu
+    timeout_in_minutes: 30
+    device: intel_gpu
+    no_plugin: true
+    env:
+      REGISTRY: "public.ecr.aws/q9t5s3a7"
+      REPO: "vllm-ci-test-repo"
+    source_file_dependencies:
+      - vllm/
+      - .buildkite/intel_jobs/test-intel.yaml 
+    commands:
+      - >-
+        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+        'pip install av &&
+        cd tests &&
+        pytest -v -s entrypoints/openai/chat_completion/test_audio_in_video.py &&
+        pytest -v -s benchmarks/test_serve_cli.py'

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

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

.buildkite/scripts/coverage/aggregate-coverage.py

@@ -0,0 +1,208 @@
+#!/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

@@ -0,0 +1,42 @@
+#!/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-intel-test.sh

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

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

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

.buildkite/test-amd.yaml

@@ -139,19 +139,6 @@ steps:
   - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py
   - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)'
 
-#--------------------------------------------------------  mi250 · benchmarks  ---------------------------------------------------------#
-
-- label: Benchmarks # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  working_dir: "/vllm-workspace/.buildkite"
-  source_file_dependencies:
-  - benchmarks/
-  - vllm/platforms/rocm.py
-  commands:
-  - bash scripts/run-benchmarks.sh
-
 #----------------------------------------------------------  mi250 · compile  ----------------------------------------------------------#
 
 - label: PyTorch Compilation Unit Tests # TBD
@@ -485,7 +472,7 @@ steps:
   - pytest -v -s model_executor -m '(not slow_test)'
   - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
 
-#----------------------------------------------------------  mi250 · models  -----------------------------------------------------------#
+#------------------------------------------------------  mi250 · models / basic  -------------------------------------------------------#
 
 - label: Basic Models Test (Other CPU) # TBD
   timeout_in_minutes: 180
@@ -546,6 +533,8 @@ steps:
   commands:
   - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
 
+#-----------------------------------------------------  mi250 · models / language  -----------------------------------------------------#
+
 - label: Language Models Test (MTEB) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -590,6 +579,8 @@ steps:
   - pip freeze | grep -E 'torch'
   - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
 
+#----------------------------------------------------  mi250 · models / multimodal  ----------------------------------------------------#
+
 - label: Multi-Modal Models (Extended Generation 2) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -976,18 +967,6 @@ steps:
 
 #--------------------------------------------------------  mi300 · benchmarks  ---------------------------------------------------------#
 
-- label: Benchmarks # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  optional: true
-  working_dir: "/vllm-workspace/.buildkite"
-  source_file_dependencies:
-  - benchmarks/
-  - vllm/platforms/rocm.py
-  commands:
-  - bash scripts/run-benchmarks.sh
-
 - label: Benchmarks CLI Test # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1759,7 +1738,7 @@ steps:
   - pytest -v -s -x lora/test_gptoss_tp.py
   - pytest -v -s -x lora/test_qwen35_densemodel_lora.py
 
-#----------------------------------------------------------  mi300 · models  -----------------------------------------------------------#
+#-----------------------------------------------------  mi300 · models / language  -----------------------------------------------------#
 
 - label: Language Models Test (Extended Pooling)  # TBD
   timeout_in_minutes: 180
@@ -1787,6 +1766,8 @@ steps:
   - pip freeze | grep -E 'torch'
   - pytest -v -s models/language -m 'core_model and (not slow_test)'
 
+#----------------------------------------------------  mi300 · models / multimodal  ----------------------------------------------------#
+
 - label: Multi-Modal Models (Extended Generation 1) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1892,10 +1873,11 @@ steps:
   - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
   - pytest -v -s models/multimodal/processing/test_tensor_schema.py
 
-- label: Multi-Modal Processor (CPU) # TBD
+- label: Multi-Modal Processor (CPU) %N # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_1
+  parallelism: 4
   no_gpu: true
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -1905,7 +1887,9 @@ steps:
   - tests/models/registry.py
   commands:
   - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-  - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py
+  - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+
+#-----------------------------------------------------  mi300 · models / quantized  -----------------------------------------------------#
 
 - label: Quantized Models Test # TBD
   timeout_in_minutes: 180
@@ -1921,7 +1905,31 @@ steps:
   commands:
   - pytest -v -s models/quantization
 
-- label: Transformers Nightly Models # TBD
+#--------------------------------------------------  mi300 · models / transformers  ---------------------------------------------------#
+
+- label: Transformers Nightly Models (Shardable) %N # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  parallelism: 4
+  optional: true
+  working_dir: "/vllm-workspace/"
+  source_file_dependencies:
+  - vllm/model_executor/models/
+  - vllm/model_executor/model_loader/
+  - vllm/multimodal/
+  - vllm/model_executor/layers/
+  - vllm/v1/attention/backends/
+  - vllm/v1/attention/selector.py
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  - tests/models/
+  commands:
+  - pip install --upgrade git+https://github.com/huggingface/transformers
+  - pytest -v -s tests/models/test_initialization.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+  - pytest -v -s tests/models/multimodal/processing/ --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+
+- label: Transformers Nightly Models (Single) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_1
@@ -1940,9 +1948,7 @@ steps:
   - examples/
   commands:
   - pip install --upgrade git+https://github.com/huggingface/transformers
-  - pytest -v -s tests/models/test_initialization.py
   - pytest -v -s tests/models/test_transformers.py
-  - pytest -v -s tests/models/multimodal/processing/
   - pytest -v -s tests/models/multimodal/test_mapping.py
   - python3 examples/basic/offline_inference/chat.py
   - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
@@ -2391,7 +2397,7 @@ steps:
   - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
   - DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
-- label: Hyrbid SSM NixlConnector PD accuracy tests (4 GPUs) # TBD
+- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_4
@@ -2593,7 +2599,7 @@ steps:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8
 
-#----------------------------------------------------------  mi325 · models  -----------------------------------------------------------#
+#-----------------------------------------------------  mi325 · models / language  -----------------------------------------------------#
 
 - label: Language Models Test (Extended Generation) # TBD
   timeout_in_minutes: 180
@@ -2624,6 +2630,8 @@ steps:
   - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
   - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
 
+#----------------------------------------------------  mi325 · models / multimodal  ----------------------------------------------------#
+
 - label: Multi-Modal Models (Extended Pooling) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
@@ -3043,7 +3051,7 @@ steps:
   commands:
     - pytest -v -s kernels/moe/test_deepep_moe.py
 
-#----------------------------------------------------------  mi355 · models  -----------------------------------------------------------#
+#-----------------------------------------------------  mi355 · models / language  -----------------------------------------------------#
 
 - label: Language Models Test (Extended Generation) # TBD
   timeout_in_minutes: 180
@@ -3111,6 +3119,8 @@ steps:
   - pip freeze | grep -E 'torch'
   - pytest -v -s models/language -m 'core_model and (not slow_test)'
 
+#----------------------------------------------------  mi355 · models / multimodal  ----------------------------------------------------#
+
 - label: Multi-Modal Models (Extended Generation 1) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
@@ -3182,6 +3192,8 @@ steps:
   - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model
   - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model
 
+#-----------------------------------------------------  mi355 · models / quantized  -----------------------------------------------------#
+
 - label: Quantized Models Test # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]

.buildkite/test_areas/kernels.yaml

@@ -150,9 +150,12 @@ steps:
   - csrc/quantization/fp4/
   - csrc/attention/mla/
   - csrc/quantization/cutlass_w8a8/moe/
-  - vllm/model_executor/layers/fused_moe/cutlass_moe.py
-  - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
-  - vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py
+  - vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
+  - vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py
+  - vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+  - vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+  - vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py
+  - vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_two_sided.py
   - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
   - vllm/v1/attention/backends/flashinfer.py
   - vllm/v1/attention/backends/mla/cutlass_mla.py
@@ -185,6 +188,7 @@ steps:
     - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
     - pytest -v -s tests/kernels/moe/test_flashinfer.py
     - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
+    - pytest -v -s tests/kernels/moe/test_trtllm_nvfp4_moe.py
     - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
     # e2e
     - pytest -v -s tests/models/quantization/test_nvfp4.py

.buildkite/test_areas/misc.yaml

@@ -367,6 +367,7 @@ steps:
     - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
     - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
     - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
+    - pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py
   
 - label: Acceptance Length Test (Large Models) # optional
   device: h200_35gb

.buildkite/test_areas/model_executor.yaml

@@ -3,7 +3,6 @@ depends_on:
   - image-build
 steps:
 - label: Model Executor
-  device: h200_35gb
   key: model-executor
   timeout_in_minutes: 35
   source_file_dependencies:

.buildkite/test_areas/rust_frontend.yaml

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

.buildkite/test_areas/rust_frontend_cargo.yaml

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

.dockerignore

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

.github/CODEOWNERS

@@ -170,13 +170,13 @@ mkdocs.yaml @hmellor
 /vllm/model_executor/layers/fla @ZJY0516 @vadiklyutiy
 
 # ROCm related: specify owner with write access to notify AMD folks for careful code review
-/vllm/**/*rocm* @tjtanaa
-/docker/Dockerfile.rocm* @gshtras @tjtanaa
-/vllm/v1/attention/backends/rocm*.py @gshtras @tjtanaa
-/vllm/v1/attention/backends/mla/rocm*.py @gshtras @tjtanaa
-/vllm/v1/attention/ops/rocm*.py @gshtras @tjtanaa
-/vllm/model_executor/layers/fused_moe/rocm*.py @gshtras @tjtanaa
-/csrc/rocm @gshtras @tjtanaa
+/vllm/**/*rocm* @tjtanaa @dllehr-amd
+/docker/Dockerfile.rocm* @tjtanaa @dllehr-amd
+/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
 /docs/**/*rocm* @tjtanaa

.gitignore

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

.pre-commit-config.yaml

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

CMakeLists.txt

@@ -311,20 +311,10 @@ set(VLLM_EXT_SRC
   "csrc/attention/paged_attention_v1.cu"
   "csrc/attention/paged_attention_v2.cu"
   "csrc/attention/merge_attn_states.cu"
-  "csrc/pos_encoding_kernels.cu"
-  "csrc/activation_kernels.cu"
-  "csrc/layernorm_kernels.cu"
-  "csrc/fused_qknorm_rope_kernel.cu"
-  "csrc/layernorm_quant_kernels.cu"
   "csrc/sampler.cu"
   "csrc/topk.cu"
   "csrc/cuda_view.cu"
-  "csrc/quantization/gptq/q_gemm.cu"
-  "csrc/quantization/w8a8/int8/scaled_quant.cu"
-  "csrc/quantization/w8a8/fp8/common.cu"
-  "csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
   "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
-  "csrc/quantization/gguf/gguf_kernel.cu"
   "csrc/quantization/activation_kernels.cu"
   "csrc/cuda_utils_kernels.cu"
   "csrc/custom_all_reduce.cu"
@@ -628,33 +618,38 @@ define_extension_target(
 # Setting this variable sidesteps the issue by calling the driver directly.
 target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
 
-# add OR VLLM_GPU_LANG STREQUAL "HIP" here once
-# https://github.com/vllm-project/vllm/issues/35163 is resolved
-if(VLLM_GPU_LANG STREQUAL "CUDA")
+if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
   #
   # _C_stable_libtorch extension (ops registered via STABLE_TORCH_LIBRARY)
   #
   set(VLLM_STABLE_EXT_SRC
     "csrc/libtorch_stable/torch_bindings.cpp"
-    "csrc/cutlass_extensions/common.cpp"
-    "csrc/cuda_utils_kernels.cu"
-    "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
-    "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
-    "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu")
+    "csrc/libtorch_stable/activation_kernels.cu"
+    "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
+    "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
+    "csrc/libtorch_stable/quantization/gptq/q_gemm.cu"
+    "csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu"
+    "csrc/libtorch_stable/pos_encoding_kernels.cu"
+    "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu"
+    "csrc/libtorch_stable/layernorm_kernels.cu"
+    "csrc/libtorch_stable/layernorm_quant_kernels.cu"
+    "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu")
 
   if(VLLM_GPU_LANG STREQUAL "CUDA")
     list(APPEND VLLM_STABLE_EXT_SRC
+      "csrc/cuda_utils_kernels.cu"
+      "csrc/cutlass_extensions/common.cpp"
+      "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
+      "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
+      "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu"
       "csrc/libtorch_stable/permute_cols.cu"
       "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
       "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
       "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu")
-  endif()
 
-  if(VLLM_GPU_LANG STREQUAL "CUDA")
     set_gencode_flags_for_srcs(
       SRCS "${VLLM_STABLE_EXT_SRC}"
       CUDA_ARCHS "${CUDA_ARCHS}")
-  endif()
 
   # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
@@ -1034,6 +1029,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     message(STATUS "Building hadacore")
   endif()
 
+  # if CUDA endif
+  endif()
+
   message(STATUS "Enabling C_stable extension.")
   define_extension_target(
     _C_stable_libtorch
@@ -1053,13 +1051,34 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   target_compile_definitions(_C_stable_libtorch PRIVATE
     TORCH_TARGET_VERSION=0x020A000000000000ULL)
 
-  # Needed to use cuda APIs from C-shim
-  target_compile_definitions(_C_stable_libtorch PRIVATE
-    USE_CUDA)
+  # Needed to use cuda/hip APIs from C-shim
+  if(VLLM_GPU_LANG STREQUAL "CUDA")
+    target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
+    # Needed by CUTLASS kernels
+    target_compile_definitions(_C_stable_libtorch PRIVATE
+      CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
+  elseif(VLLM_GPU_LANG STREQUAL "HIP")
+    target_compile_definitions(_C_stable_libtorch PRIVATE USE_ROCM)
+  endif()
 
-  # Needed by CUTLASS kernels
-  target_compile_definitions(_C_stable_libtorch PRIVATE
-    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
+  # On ROCm, _C_stable_libtorch calls raw HIP APIs (e.g. hipGetDevice in
+  # get_device_prop()) which must resolve to the same libamdhip64.so that
+  # PyTorch uses.  When PyTorch bundles its own copy (pip/conda wheels),
+  # the raw HIP calls would otherwise resolve to the system ROCm copy,
+  # initializing a second HIP runtime that corrupts device state (wrong
+  # device on DeviceGuard, core dumps on multi-GPU tests).
+  #
+  # If PyTorch doesn't bundle libamdhip64 (built from source against system
+  # ROCm), there is only one copy in the process and no action is needed —
+  # the HIP compiler already links the system libamdhip64 automatically.
+  if(VLLM_GPU_LANG STREQUAL "HIP")
+    find_library(_STABLE_TORCH_AMDHIP64 amdhip64
+      PATHS "${TORCH_INSTALL_PREFIX}/lib" NO_DEFAULT_PATH)
+    if(_STABLE_TORCH_AMDHIP64)
+      message(STATUS "Found PyTorch-bundled libamdhip64 at ${_STABLE_TORCH_AMDHIP64}")
+      target_link_libraries(_C_stable_libtorch PRIVATE ${_STABLE_TORCH_AMDHIP64})
+    endif()
+  endif()
 endif()
 
 #
@@ -1221,16 +1240,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       CUDA_ARCHS "${DSV3_ROUTER_GEMM_ARCHS}")
     list(APPEND VLLM_MOE_EXT_SRC "${DSV3_ROUTER_GEMM_SRC}")
     message(STATUS "Building DSV3 router GEMM kernel for archs: ${DSV3_ROUTER_GEMM_ARCHS}")
-
-    # DeepSeek V4 fused RMSNorm + router GEMV - same arch gating as DSV3.
-    set(DSV4_NORM_ROUTER_GEMM_SRC
-      "csrc/moe/dsv4_norm_router_gemm_entry.cu"
-      "csrc/moe/dsv4_norm_router_gemm_kernel.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${DSV4_NORM_ROUTER_GEMM_SRC}"
-      CUDA_ARCHS "${DSV3_ROUTER_GEMM_ARCHS}")
-    list(APPEND VLLM_MOE_EXT_SRC "${DSV4_NORM_ROUTER_GEMM_SRC}")
-    message(STATUS "Building DSV4 norm+router GEMV kernel for archs: ${DSV3_ROUTER_GEMM_ARCHS}")
   else()
     message(STATUS "Not building DSV3 router GEMM kernel as no compatible archs found"
                    " (requires SM90+ and CUDA >= 12.0)")
@@ -1270,6 +1279,12 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
     WITH_SOABI)
 endif()
 
+# Must run after the last HIP `define_extension_target` so every extension
+# has registered its sources.
+if (VLLM_GPU_LANG STREQUAL "HIP")
+  vllm_finalize_hipify_target()
+endif()
+
 # For CUDA and HIP builds also build the triton_kernels external package.
 if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     include(cmake/external_projects/triton_kernels.cmake)

benchmarks/benchmark_hidden_state_extraction.py

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

benchmarks/kernels/benchmark_fused_moe_lora_one_shot.py

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

benchmarks/kernels/benchmark_norm_router_gemm.py

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

benchmarks/kernels/benchmark_selective_state_update.py

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

build_rust.sh

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

cmake/external_projects/triton_kernels.cmake

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

cmake/utils.cmake

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

csrc/attention/dtype_fp8.cuh

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

csrc/cpu/cpu_attn_impl.hpp

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

csrc/cpu/cpu_attn_rvv.hpp

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

csrc/cpu/cpu_types_riscv_defs.hpp

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

csrc/cpu/generate_cpu_attn_dispatch.py

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

csrc/cuda_vec_utils.cuh

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

csrc/cutlass_extensions/torch_utils.hpp

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

csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu

@@ -87,6 +87,12 @@ constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1;  // 8 (7 real + 1 pad)
 constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2;  // 448 + 128 = 576
 constexpr float kFp8Max = 448.0f;
 
+#ifndef USE_ROCM
+// When num_tokens is less than this threshold,
+// run the reduced grid variant on cuda
+constexpr float NUM_TOKEN_CUTOFF = 1024;
+#endif
+
 // Per-warp layout:  32 lanes × 16 elems/lane = 512 elems = HEAD_DIM.
 constexpr int kNumLanes = 32;
 constexpr int kElemsPerLane = kHeadDim / kNumLanes;  // 16
@@ -112,6 +118,197 @@ __device__ __forceinline__ float warpSum(float val) {
   return val;
 }
 
+// ────────────────────────────────────────────────────────────────────────────
+// Per-slot inner pipeline
+// ────────────────────────────────────────────────────────────────────────────
+// Shared by both kernel variants: 1 CTA per (token, head) pair vs. 1 CTA per
+// token
+template <typename scalar_t_in>
+__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,
+    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);
+
+  // ── Decode the bf16 → 16 fp32 registers ─────────────────────────────
+  float elements[kElemsPerLane];
+  {
+    typename Converter::packed_hip_type const* p0 =
+        reinterpret_cast<typename Converter::packed_hip_type const*>(&v0);
+    typename Converter::packed_hip_type const* p1 =
+        reinterpret_cast<typename Converter::packed_hip_type const*>(&v1);
+#pragma unroll
+    for (int i = 0; i < 4; i++) {
+      float2 f2 = Converter::convert(p0[i]);
+      elements[2 * i] = f2.x;
+      elements[2 * i + 1] = f2.y;
+    }
+#pragma unroll
+    for (int i = 0; i < 4; i++) {
+      float2 f2 = Converter::convert(p1[i]);
+      elements[8 + 2 * i] = f2.x;
+      elements[8 + 2 * i + 1] = f2.y;
+    }
+  }
+
+  // ── Q branch: RMSNorm (no weight) ───────────────────────────────────
+  if (!isKV) {
+    float sumOfSquares = 0.0f;
+#pragma unroll
+    for (int i = 0; i < kElemsPerLane; i++) {
+      sumOfSquares += elements[i] * elements[i];
+    }
+    sumOfSquares = warpSum<float>(sumOfSquares);
+    float const rms_rcp =
+        rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
+#pragma unroll
+    for (int i = 0; i < kElemsPerLane; i++) {
+      elements[i] = elements[i] * rms_rcp;
+    }
+  }
+
+  // ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ─────────────────────────────
+  // All math in fp32.  cos_sin_cache is loaded as fp32 (its native storage).
+  bool const is_rope_lane = dim_base >= kNopeDim;
+  if (is_rope_lane) {
+    int64_t const pos = position_ids[tokenIdx];
+    constexpr int kHalfRope = kRopeDim / 2;
+    float const* cos_ptr = cos_sin_cache + pos * kRopeDim;
+    float const* sin_ptr = cos_ptr + kHalfRope;
+
+    int const rope_local_base = dim_base - kNopeDim;
+    int const half_base = rope_local_base >> 1;
+
+    // Load phase: 4 vectorized LDGs issue back-to-back.
+    float4 const c0 = *reinterpret_cast<float4 const*>(cos_ptr + half_base);
+    float4 const c1 = *reinterpret_cast<float4 const*>(cos_ptr + half_base + 4);
+    float4 const s0 = *reinterpret_cast<float4 const*>(sin_ptr + half_base);
+    float4 const s1 = *reinterpret_cast<float4 const*>(sin_ptr + half_base + 4);
+    float const cos_arr[8] = {c0.x, c0.y, c0.z, c0.w, c1.x, c1.y, c1.z, c1.w};
+    float const sin_arr[8] = {s0.x, s0.y, s0.z, s0.w, s1.x, s1.y, s1.z, s1.w};
+
+#pragma unroll
+    for (int p = 0; p < kElemsPerLane / 2; p++) {
+      float const x_even = elements[2 * p];
+      float const x_odd = elements[2 * p + 1];
+      elements[2 * p] = x_even * cos_arr[p] - x_odd * sin_arr[p];
+      elements[2 * p + 1] = x_even * sin_arr[p] + x_odd * cos_arr[p];
+    }
+  }
+
+  // ═══════════════════════════════════════════════════════════════════
+  // Q / KV branch dispatch. Restructured as if/else (no early `return`)
+  // so every code path lands at the same exit point — callers own PDL
+  // triggering and per-iteration buffer rotation.
+  // ═══════════════════════════════════════════════════════════════════
+  if (!isKV) {
+    // ── Q: cast back to bf16 and store. ────────────────────────────
+    uint4 out0, out1;
+    typename Converter::packed_hip_type* po0 =
+        reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
+    typename Converter::packed_hip_type* po1 =
+        reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
+#pragma unroll
+    for (int i = 0; i < 4; i++) {
+      po0[i] =
+          Converter::convert(make_float2(elements[2 * i], elements[2 * i + 1]));
+    }
+#pragma unroll
+    for (int i = 0; i < 4; i++) {
+      po1[i] = Converter::convert(
+          make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
+    }
+    scalar_t_in* dst =
+        q_inout +
+        (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
+        dim_base;
+    *reinterpret_cast<uint4*>(dst) = out0;
+    *reinterpret_cast<uint4*>(dst + 8) = out1;
+  } else {
+    // ── KV: FP8 quant on NoPE + bf16 store on RoPE + cache insert.
+    int64_t const slot_id = slot_mapping[tokenIdx];
+    if (slot_id >= 0) {
+      int64_t const block_idx = slot_id / cache_block_size;
+      int64_t const pos_in_block = slot_id % cache_block_size;
+      uint8_t* block_base =
+          k_cache + block_idx * static_cast<int64_t>(kv_block_stride);
+      uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes;
+      uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim;
+      uint8_t* token_scale_ptr =
+          block_base +
+          static_cast<int64_t>(cache_block_size) * kTokenDataBytes +
+          pos_in_block * kScaleBytesPerToken;
+
+#pragma unroll
+      for (int i = 0; i < kElemsPerLane; i++) {
+        elements[i] = Converter::convert(Converter::convert(elements[i]));
+      }
+
+      float local_absmax = 0.0f;
+#pragma unroll
+      for (int i = 0; i < kElemsPerLane; i++) {
+        local_absmax = fmaxf(local_absmax, fabsf(elements[i]));
+      }
+      float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f);
+      float const exponent = ceilf(log2f(absmax / kFp8Max));
+      float const inv_scale = exp2f(-exponent);
+
+      if (!is_rope_lane) {
+        uint8_t out_bytes[kElemsPerLane];
+#pragma unroll
+        for (int i = 0; i < kElemsPerLane; i++) {
+          float scaled = elements[i] * inv_scale;
+          scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
+#ifndef USE_ROCM
+          __nv_fp8_storage_t s =
+              __nv_cvt_float_to_fp8(scaled, __NV_SATFINITE, __NV_E4M3);
+          out_bytes[i] = static_cast<uint8_t>(s);
+#else
+          out_bytes[i] = rocm_cvt_float_to_fp8_e4m3(scaled);
+#endif
+        }
+        *reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
+            *reinterpret_cast<uint4 const*>(out_bytes);
+
+        if ((laneId & 3) == 0) {
+          int const q_block_idx = laneId >> 2;
+          float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f);
+          token_scale_ptr[q_block_idx] = static_cast<uint8_t>(encoded);
+        }
+        if (laneId == 0) {
+          token_scale_ptr[kNumQuantBlocks] = 0;
+        }
+      } else {
+        uint4 out0, out1;
+        typename Converter::packed_hip_type* po0 =
+            reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
+        typename Converter::packed_hip_type* po1 =
+            reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
+#pragma unroll
+        for (int i = 0; i < 4; i++) {
+          po0[i] = Converter::convert(
+              make_float2(elements[2 * i], elements[2 * i + 1]));
+        }
+#pragma unroll
+        for (int i = 0; i < 4; i++) {
+          po1[i] = Converter::convert(
+              make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
+        }
+        int const rope_local_base = dim_base - kNopeDim;
+        scalar_t_in* bf16_dst =
+            reinterpret_cast<scalar_t_in*>(token_bf16_ptr) + rope_local_base;
+        *reinterpret_cast<uint4*>(bf16_dst) = out0;
+        *reinterpret_cast<uint4*>(bf16_dst + 8) = out1;
+      }
+    }
+  }
+}
+
 // ────────────────────────────────────────────────────────────────────────────
 // Kernel
 // ────────────────────────────────────────────────────────────────────────────
@@ -149,8 +346,6 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
     return;
   } else {
 #endif
-    using Converter = vllm::_typeConvert<scalar_t_in>;
-
     int const warpsPerBlock = blockDim.x / 32;
     int const warpId = threadIdx.x / 32;
     int const laneId = threadIdx.x % 32;
@@ -176,10 +371,8 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
     // Dim range this lane owns within the 512-wide head.
     int const dim_base = laneId * kElemsPerLane;  // in [0, 512) step 16
 
-    // ── Load 16 bf16 → 16 fp32 registers (one 16-byte + one 16-byte LDG) ────
-    float elements[kElemsPerLane];
-    float sumOfSquares = 0.0f;
-
+    // 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;
@@ -189,196 +382,103 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
           dim_base;
       src_ptr = q_inout + q_row_offset;
     }
+    uint4 const v0 = *reinterpret_cast<uint4 const*>(src_ptr);
+    uint4 const v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
 
-    // Two 16-byte loads per thread (8 bf16 each).  Use uint4 as the vector
-    // type and bitcast to scalar_t_in packed pairs for conversion.
-    uint4 v0 = *reinterpret_cast<uint4 const*>(src_ptr);
-    uint4 v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
+    processDeepseekV4Slot<scalar_t_in>(
+        v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_inout,
+        k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size,
+        kv_block_stride);
 
-    {
-      typename Converter::packed_hip_type const* p0 =
-          reinterpret_cast<typename Converter::packed_hip_type const*>(&v0);
-      typename Converter::packed_hip_type const* p1 =
-          reinterpret_cast<typename Converter::packed_hip_type const*>(&v1);
-// Each packed_hip_type holds 2 bf16 → 4 packed = 8 elems per uint4.
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        float2 f2 = Converter::convert(p0[i]);
-        elements[2 * i] = f2.x;
-        elements[2 * i + 1] = f2.y;
-      }
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        float2 f2 = Converter::convert(p1[i]);
-        elements[8 + 2 * i] = f2.x;
-        elements[8 + 2 * i + 1] = f2.y;
-      }
-    }
-
-    // ── Q branch: RMSNorm with no weight (has_weight=False) ─────────────────
-    // Variance + rsqrt + multiply all in fp32, no intermediate bf16 round.
-    // The downstream bf16 round only happens at the final store.
-    if (!isKV) {
-#pragma unroll
-      for (int i = 0; i < kElemsPerLane; i++) {
-        sumOfSquares += elements[i] * elements[i];
-      }
-      sumOfSquares = warpSum<float>(sumOfSquares);
-      float const rms_rcp =
-          rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
-#pragma unroll
-      for (int i = 0; i < kElemsPerLane; i++) {
-        elements[i] = elements[i] * rms_rcp;
-      }
-    }
-
-    // ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ─────────────────────────────
-    // All math in fp32.  cos_sin_cache is loaded as fp32 (its native storage).
-    bool const is_rope_lane = dim_base >= kNopeDim;
-    if (is_rope_lane) {
-      int64_t const pos = position_ids[tokenIdx];
-      constexpr int kHalfRope = kRopeDim / 2;  // 32
-      float const* cos_ptr = cos_sin_cache + pos * kRopeDim;
-      float const* sin_ptr = cos_ptr + kHalfRope;
-
-      int const rope_local_base = dim_base - kNopeDim;  // in [0, 64) step 16
-#pragma unroll
-      for (int p = 0; p < kElemsPerLane / 2; p++) {
-        int const pair_dim = rope_local_base + 2 * p;
-        int const half_idx = pair_dim / 2;
-        float const cos_v = VLLM_LDG(cos_ptr + half_idx);
-        float const sin_v = VLLM_LDG(sin_ptr + half_idx);
-        float const x_even = elements[2 * p];
-        float const x_odd = elements[2 * p + 1];
-        elements[2 * p] = x_even * cos_v - x_odd * sin_v;
-        elements[2 * p + 1] = x_even * sin_v + x_odd * cos_v;
-      }
-    }
-
-    // ═══════════════════════════════════════════════════════════════════════
-    // Q branch: cast to bf16 and store back in place.
-    // ═══════════════════════════════════════════════════════════════════════
-    if (!isKV) {
-      uint4 out0, out1;
-      typename Converter::packed_hip_type* po0 =
-          reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
-      typename Converter::packed_hip_type* po1 =
-          reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        po0[i] = Converter::convert(
-            make_float2(elements[2 * i], elements[2 * i + 1]));
-      }
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        po1[i] = Converter::convert(
-            make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
-      }
-      scalar_t_in* dst =
-          q_inout +
-          (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
-          dim_base;
-      *reinterpret_cast<uint4*>(dst) = out0;
-      *reinterpret_cast<uint4*>(dst + 8) = out1;
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-      cudaTriggerProgrammaticLaunchCompletion();
+    cudaTriggerProgrammaticLaunchCompletion();
 #endif
-      return;
-    }
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
+  }
+#endif
+}
+
+// ────────────────────────────────────────────────────────────────────────────
+// Kernel
+// ────────────────────────────────────────────────────────────────────────────
+//
+// Grid: 1D, gridDim.x = num_tokens_full
+// Block: blockDim.x = 256 threads (8 warps per block) Each
+// warp handles one token, iterating over each head.
+// Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q
+// KV branch (RoPE + UE8M0 quant + insert)
+//
+template <typename scalar_t_in>
+__global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
+    scalar_t_in* __restrict__ q_inout,  // [N, H, 512] bf16, in place
+    scalar_t_in const* __restrict__ kv_in, uint8_t* __restrict__ k_cache,
+    int64_t const* __restrict__ slot_mapping,
+    int64_t const* __restrict__ position_ids,
+    float const* __restrict__ cos_sin_cache, float const eps,
+    int const num_tokens_full, int const num_tokens_insert,
+    int const num_heads_q, int const cache_block_size,
+    int const kv_block_stride) {
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
+  if constexpr (std::is_same_v<scalar_t_in, c10::BFloat16>) {
+    return;
+  } else {
+#endif
+    int const warpsPerBlock = blockDim.x / 32;
+    int const warpId = threadIdx.x / 32;
+    int const laneId = threadIdx.x % 32;
+
+    int const tokenIdx = blockIdx.x;
+    if (tokenIdx >= num_tokens_full) return;
 
-    // ═══════════════════════════════════════════════════════════════════════
-    // KV branch.
-    // ═══════════════════════════════════════════════════════════════════════
-    int64_t const slot_id = slot_mapping[tokenIdx];
-    if (slot_id < 0) {
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-      cudaTriggerProgrammaticLaunchCompletion();
+    cudaGridDependencySynchronize();
 #endif
-      return;
-    }
 
-    int64_t const block_idx = slot_id / cache_block_size;
-    int64_t const pos_in_block = slot_id % cache_block_size;
-    uint8_t* block_base =
-        k_cache + block_idx * static_cast<int64_t>(kv_block_stride);
-    uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes;
-    uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim;
-    uint8_t* token_scale_ptr =
-        block_base + static_cast<int64_t>(cache_block_size) * kTokenDataBytes +
-        pos_in_block * kScaleBytesPerToken;
+    int const dim_base = laneId * kElemsPerLane;  // in [0, 512) step 16
+    int const slot_end =
+        (tokenIdx >= num_tokens_insert) ? num_heads_q : (num_heads_q + 1);
 
-    // Round K to bf16 first, matching the unfused reference path where K is
-    // materialized as bf16 before K quantization.  absmax, clamp, and FP8
-    // quant below all run on these bf16-rounded values.
-#pragma unroll
-    for (int i = 0; i < kElemsPerLane; i++) {
-      elements[i] = Converter::convert(Converter::convert(elements[i]));
-    }
+    auto 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;
+      }
+      return q_inout +
+             (static_cast<int64_t>(tokenIdx) * num_heads_q +
+              static_cast<int64_t>(s)) *
+                 kHeadDim +
+             dim_base;
+    };
 
-    // Per-quant-block absmax must be computed by ALL 32 lanes (warp-collective
-    // shuffle requires full participation).  RoPE lanes contribute garbage,
-    // but their values are gated out below via `!is_rope_lane`.
-    float local_absmax = 0.0f;
-#pragma unroll
-    for (int i = 0; i < kElemsPerLane; i++) {
-      local_absmax = fmaxf(local_absmax, fabsf(elements[i]));
-    }
-    float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f);
-    float const exponent = ceilf(log2f(absmax / kFp8Max));
-    float const inv_scale = exp2f(-exponent);
+    if (warpId < slot_end) {
+      int curr_slot = warpId;
+      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);
 
-    if (!is_rope_lane) {
-      // ── NoPE lane: UE8M0 FP8 quant ───────────────────────────────────────
-      uint8_t out_bytes[kElemsPerLane];
-#pragma unroll
-      for (int i = 0; i < kElemsPerLane; i++) {
-        float scaled = elements[i] * inv_scale;
-        scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
-#ifndef USE_ROCM
-        __nv_fp8_storage_t s =
-            __nv_cvt_float_to_fp8(scaled, __NV_SATFINITE, __NV_E4M3);
-        out_bytes[i] = static_cast<uint8_t>(s);
-#else
-      out_bytes[i] = rocm_cvt_float_to_fp8_e4m3(scaled);
-#endif
-      }
-      // One 16-byte STG per lane.
-      *reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
-          *reinterpret_cast<uint4 const*>(out_bytes);
+      while (curr_slot < slot_end) {
+        int const next_slot = curr_slot + warpsPerBlock;
+        bool const has_next = (next_slot < slot_end);
+
+        // Prefetch src for the next slot
+        uint4 v0_next, v1_next;
+        if (has_next) {
+          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);
+        }
+
+        processDeepseekV4Slot<scalar_t_in>(
+            v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId,
+            num_heads_q, eps, q_inout, k_cache, slot_mapping, position_ids,
+            cos_sin_cache, cache_block_size, kv_block_stride);
+
+        // ── Buffer rotation: hand the prefetched LDGs to the next iter.
+        v0_curr = v0_next;
+        v1_curr = v1_next;
+        curr_slot = next_slot;
+      }  // while
+    }  // if (warpId < slot_end)
 
-      // Lane (4k) of each 4-lane group writes the scale byte for block k<7.
-      if ((laneId & 3) == 0) {
-        int const q_block_idx = laneId >> 2;  // 0..6 for NoPE lanes
-        float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f);
-        token_scale_ptr[q_block_idx] = static_cast<uint8_t>(encoded);
-      }
-      // Lane 0 also writes the padding byte at index 7.
-      if (laneId == 0) {
-        token_scale_ptr[kNumQuantBlocks] = 0;  // pad
-      }
-    } else {
-      // ── RoPE lane: cast back to bf16 and store to cache bf16 tail ────────
-      uint4 out0, out1;
-      typename Converter::packed_hip_type* po0 =
-          reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
-      typename Converter::packed_hip_type* po1 =
-          reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        po0[i] = Converter::convert(
-            make_float2(elements[2 * i], elements[2 * i + 1]));
-      }
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        po1[i] = Converter::convert(
-            make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
-      }
-      int const rope_local_base = dim_base - kNopeDim;  // in [0, 64)
-      scalar_t_in* bf16_dst =
-          reinterpret_cast<scalar_t_in*>(token_bf16_ptr) + rope_local_base;
-      *reinterpret_cast<uint4*>(bf16_dst) = out0;
-      *reinterpret_cast<uint4*>(bf16_dst + 8) = out1;
-    }
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
     cudaTriggerProgrammaticLaunchCompletion();
 #endif
@@ -430,11 +530,22 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
   config.attrs = attrs;
   config.numAttrs = (sm_version >= 90) ? 1 : 0;
 
-  cudaLaunchKernelEx(
-      &config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>,
-      q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
-      num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
-      kv_block_stride);
+  if (num_tokens_full < NUM_TOKEN_CUTOFF) {
+    cudaLaunchKernelEx(
+        &config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>,
+        q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
+        num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
+        kv_block_stride);
+  } 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,
+        kv_block_stride);
+  }
+
 #else
   // ROCm: use standard kernel launch syntax (no PDL/stream serialization)
   // clang-format off
@@ -508,4 +619,4 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
                 num_tokens_full, num_tokens_insert, num_heads_q,
                 cache_block_size_i, kv_block_stride, stream);
       });
-}
+}

csrc/libtorch_stable/activation_kernels.cu

@@ -1,12 +1,12 @@
-#include <ATen/cuda/CUDAContext.h>
-#include <torch/all.h>
-#include <c10/cuda/CUDAGuard.h>
+#include <cuda.h>
+#include <torch/csrc/stable/tensor.h>
 
 #include <cmath>
 
-#include "cuda_compat.h"
-#include "cuda_vec_utils.cuh"
+#include "../cuda_compat.h"
+#include "../cuda_vec_utils.cuh"
 #include "dispatch_utils.h"
+#include "torch_utils.h"
 
 namespace vllm {
 
@@ -210,64 +210,68 @@ packed_gelu_tanh_kernel(const packed_t& val) {
     return;                                                                    \
   }                                                                            \
   dim3 grid(num_tokens);                                                       \
-  int cc_major = at::cuda::getCurrentDeviceProperties()->major;                \
+  int cc_major = get_device_prop()->major;                                     \
   int support_vec =                                                            \
       (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128)            \
           ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE                           \
           : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE;                         \
-  int vec_size = support_vec / at::elementSize(dtype);                         \
+  int vec_size = support_vec / input.element_size();                           \
   const bool use_vec = (d % vec_size == 0);                                    \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
+  const torch::stable::accelerator::DeviceGuard device_guard(                  \
+      input.get_device_index());                                               \
+  const cudaStream_t stream = get_current_cuda_stream();                       \
   if (use_vec) {                                                               \
     dim3 block(std::min(d / vec_size, 1024));                                  \
     if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) {         \
-      VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {          \
+      VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {   \
         vllm::act_and_mul_kernel<                                              \
             scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
             KERNEL<scalar_t>,                                                  \
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
             ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>(     \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
+            out.mutable_data_ptr<scalar_t>(),                                  \
+            input.const_data_ptr<scalar_t>(), d, LIMIT);                       \
       });                                                                      \
     } else {                                                                   \
-      VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {          \
+      VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {   \
         vllm::act_and_mul_kernel<                                              \
             scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
             KERNEL<scalar_t>,                                                  \
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
             ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>(    \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
+            out.mutable_data_ptr<scalar_t>(),                                  \
+            input.const_data_ptr<scalar_t>(), d, LIMIT);                       \
       });                                                                      \
     }                                                                          \
   } else {                                                                     \
     dim3 block(std::min(d, 1024));                                             \
-    VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {            \
+    VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {     \
       vllm::act_and_mul_kernel<                                                \
           scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,        \
           KERNEL<scalar_t>,                                                    \
           PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
           ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>(            \
-          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);     \
+          out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),  \
+          d, LIMIT);                                                           \
     });                                                                        \
   }
 
-void silu_and_mul(torch::Tensor& out,    // [..., d]
-                  torch::Tensor& input)  // [..., 2 * d]
+void silu_and_mul(torch::stable::Tensor& out,    // [..., d]
+                  torch::stable::Tensor& input)  // [..., 2 * d]
 {
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
                                 true, false, 0.0f);
 }
 
-void silu_and_mul_clamp(torch::Tensor& out,    // [..., d]
-                        torch::Tensor& input,  // [..., 2 * d]
+void silu_and_mul_clamp(torch::stable::Tensor& out,    // [..., d]
+                        torch::stable::Tensor& input,  // [..., 2 * d]
                         double limit) {
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
                                 true, true, (float)limit);
 }
 
-void mul_and_silu(torch::Tensor& out,    // [..., d]
-                  torch::Tensor& input)  // [..., 2 * d]
+void mul_and_silu(torch::stable::Tensor& out,    // [..., d]
+                  torch::stable::Tensor& input)  // [..., 2 * d]
 {
   // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
   // applies the silu to the latter half of the input.
@@ -275,15 +279,15 @@ void mul_and_silu(torch::Tensor& out,    // [..., d]
                                 false, false, 0.0f);
 }
 
-void gelu_and_mul(torch::Tensor& out,    // [..., d]
-                  torch::Tensor& input)  // [..., 2 * d]
+void gelu_and_mul(torch::stable::Tensor& out,    // [..., d]
+                  torch::stable::Tensor& input)  // [..., 2 * d]
 {
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
                                 true, false, 0.0f);
 }
 
-void gelu_tanh_and_mul(torch::Tensor& out,    // [..., d]
-                       torch::Tensor& input)  // [..., 2 * d]
+void gelu_tanh_and_mul(torch::stable::Tensor& out,    // [..., d]
+                       torch::stable::Tensor& input)  // [..., 2 * d]
 {
   LAUNCH_ACTIVATION_GATE_KERNEL(
       vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
@@ -434,19 +438,20 @@ __global__ void swigluoai_and_mul_kernel(
     return;                                                                    \
   }                                                                            \
   dim3 grid(num_tokens);                                                       \
-  int cc_major = at::cuda::getCurrentDeviceProperties()->major;                \
+  int cc_major = get_device_prop()->major;                                     \
   int support_vec =                                                            \
       (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128)            \
           ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE                           \
           : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE;                         \
-  int vec_size = support_vec / at::elementSize(dtype);                         \
+  int vec_size = support_vec / input.element_size();                           \
   const bool use_vec = (d % vec_size == 0);                                    \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
+  const torch::stable::accelerator::DeviceGuard device_guard(                  \
+      input.get_device_index());                                               \
+  const cudaStream_t stream = get_current_cuda_stream();                       \
   if (use_vec) {                                                               \
     dim3 block(std::min(d / vec_size, 1024));                                  \
     if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) {         \
-      VLLM_DISPATCH_FLOATING_TYPES(                                            \
+      VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                     \
           dtype, "act_and_mul_kernel_with_param", [&] {                        \
             vllm::act_and_mul_kernel_with_param<                               \
                 scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,  \
@@ -454,11 +459,11 @@ __global__ void swigluoai_and_mul_kernel(
                 PACKED_KERNEL<                                                 \
                     typename vllm::PackedTypeConverter<scalar_t>::Type>,       \
                 true, true><<<grid, block, 0, stream>>>(                       \
-                out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d,       \
-                PARAM);                                                        \
+                out.mutable_data_ptr<scalar_t>(),                              \
+                input.const_data_ptr<scalar_t>(), d, PARAM);                   \
           });                                                                  \
     } else {                                                                   \
-      VLLM_DISPATCH_FLOATING_TYPES(                                            \
+      VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                     \
           dtype, "act_and_mul_kernel_with_param", [&] {                        \
             vllm::act_and_mul_kernel_with_param<                               \
                 scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,  \
@@ -466,45 +471,49 @@ __global__ void swigluoai_and_mul_kernel(
                 PACKED_KERNEL<                                                 \
                     typename vllm::PackedTypeConverter<scalar_t>::Type>,       \
                 true, false><<<grid, block, 0, stream>>>(                      \
-                out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d,       \
-                PARAM);                                                        \
+                out.mutable_data_ptr<scalar_t>(),                              \
+                input.const_data_ptr<scalar_t>(), d, PARAM);                   \
           });                                                                  \
     }                                                                          \
   } else {                                                                     \
     dim3 block(std::min(d, 1024));                                             \
-    VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel_with_param", [&] { \
-      vllm::act_and_mul_kernel_with_param<                                     \
-          scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,        \
-          KERNEL<scalar_t>,                                                    \
-          PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
-          false><<<grid, block, 0, stream>>>(                                  \
-          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, PARAM);     \
-    });                                                                        \
+    VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                       \
+        dtype, "act_and_mul_kernel_with_param", [&] {                          \
+          vllm::act_and_mul_kernel_with_param<                                 \
+              scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,    \
+              KERNEL<scalar_t>,                                                \
+              PACKED_KERNEL<                                                   \
+                  typename vllm::PackedTypeConverter<scalar_t>::Type>,         \
+              false><<<grid, block, 0, stream>>>(                              \
+              out.mutable_data_ptr<scalar_t>(),                                \
+              input.const_data_ptr<scalar_t>(), d, PARAM);                     \
+        });                                                                    \
   }
 
-#define LAUNCH_SIGLUOAI_AND_MUL(KERNEL, ALPHA, LIMIT)                          \
-  int d = input.size(-1) / 2;                                                  \
-  int64_t num_tokens = input.numel() / input.size(-1);                         \
-  dim3 grid(num_tokens);                                                       \
-  dim3 block(std::min(d, 1024));                                               \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
-  VLLM_DISPATCH_FLOATING_TYPES(                                                \
-      input.scalar_type(), "clamp_swiglu_kernel_with_params", [&] {            \
-        vllm::swigluoai_and_mul_kernel<scalar_t, KERNEL<scalar_t>>             \
-            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),             \
-                                         input.data_ptr<scalar_t>(), d, ALPHA, \
-                                         LIMIT);                               \
+#define LAUNCH_SIGLUOAI_AND_MUL(KERNEL, ALPHA, LIMIT)                         \
+  int d = input.size(-1) / 2;                                                 \
+  int64_t num_tokens = input.numel() / input.size(-1);                        \
+  dim3 grid(num_tokens);                                                      \
+  dim3 block(std::min(d, 1024));                                              \
+  const torch::stable::accelerator::DeviceGuard device_guard(                 \
+      input.get_device_index());                                              \
+  const cudaStream_t stream = get_current_cuda_stream();                      \
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                        \
+      input.scalar_type(), "clamp_swiglu_kernel_with_params", [&] {           \
+        vllm::swigluoai_and_mul_kernel<scalar_t, KERNEL<scalar_t>>            \
+            <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(),    \
+                                         input.const_data_ptr<scalar_t>(), d, \
+                                         ALPHA, LIMIT);                       \
       });
 
-void fatrelu_and_mul(torch::Tensor& out,    // [..., d],
-                     torch::Tensor& input,  // [..., 2 * d]
+void fatrelu_and_mul(torch::stable::Tensor& out,    // [..., d],
+                     torch::stable::Tensor& input,  // [..., 2 * d]
                      double threshold) {
   LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(
       vllm::fatrelu_kernel, vllm::packed_fatrelu_kernel, threshold);
 }
-void swigluoai_and_mul(torch::Tensor& out,    // [..., d]
-                       torch::Tensor& input,  // [..., 2 * d]
+void swigluoai_and_mul(torch::stable::Tensor& out,    // [..., d]
+                       torch::stable::Tensor& input,  // [..., 2 * d]
                        double alpha, double limit) {
   LAUNCH_SIGLUOAI_AND_MUL(vllm::swigluoai_and_mul, alpha, limit);
 }
@@ -559,45 +568,46 @@ __global__ void activation_kernel(
 }  // namespace vllm
 
 // Launch element-wise activation kernel.
-#define LAUNCH_ACTIVATION_KERNEL(KERNEL)                                 \
-  auto dtype = input.scalar_type();                                      \
-  int d = input.size(-1);                                                \
-  int64_t num_tokens = input.numel() / input.size(-1);                   \
-  if (num_tokens == 0) {                                                 \
-    return;                                                              \
-  }                                                                      \
-  dim3 grid(num_tokens);                                                 \
-  int cc_major = at::cuda::getCurrentDeviceProperties()->major;          \
-  int support_vec =                                                      \
-      (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128)      \
-          ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE                     \
-          : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE;                   \
-  int vec_size = support_vec / at::elementSize(dtype);                   \
-  const bool use_vec = (d % vec_size == 0);                              \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));      \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();          \
-  if (use_vec) {                                                         \
-    dim3 block(std::min(d / vec_size, 1024));                            \
-    if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) {   \
-      VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {     \
-        vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, true>  \
-            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \
-                                         input.data_ptr<scalar_t>(), d); \
-      });                                                                \
-    } else {                                                             \
-      VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {     \
-        vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, false> \
-            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \
-                                         input.data_ptr<scalar_t>(), d); \
-      });                                                                \
-    }                                                                    \
-  } else {                                                               \
-    dim3 block(std::min(d, 1024));                                       \
-    VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {       \
-      vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, false>         \
-          <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),         \
-                                       input.data_ptr<scalar_t>(), d);   \
-    });                                                                  \
+#define LAUNCH_ACTIVATION_KERNEL(KERNEL)                                       \
+  auto dtype = input.scalar_type();                                            \
+  int d = input.size(-1);                                                      \
+  int64_t num_tokens = input.numel() / input.size(-1);                         \
+  if (num_tokens == 0) {                                                       \
+    return;                                                                    \
+  }                                                                            \
+  dim3 grid(num_tokens);                                                       \
+  int cc_major = get_device_prop()->major;                                     \
+  int support_vec =                                                            \
+      (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128)            \
+          ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE                           \
+          : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE;                         \
+  int vec_size = support_vec / input.element_size();                           \
+  const bool use_vec = (d % vec_size == 0);                                    \
+  const torch::stable::accelerator::DeviceGuard device_guard(                  \
+      input.get_device_index());                                               \
+  const cudaStream_t stream = get_current_cuda_stream();                       \
+  if (use_vec) {                                                               \
+    dim3 block(std::min(d / vec_size, 1024));                                  \
+    if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) {         \
+      VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {    \
+        vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, true>        \
+            <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(),     \
+                                         input.const_data_ptr<scalar_t>(), d); \
+      });                                                                      \
+    } else {                                                                   \
+      VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {    \
+        vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, false>       \
+            <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(),     \
+                                         input.const_data_ptr<scalar_t>(), d); \
+      });                                                                      \
+    }                                                                          \
+  } else {                                                                     \
+    dim3 block(std::min(d, 1024));                                             \
+    VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {      \
+      vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, false>               \
+          <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(),       \
+                                       input.const_data_ptr<scalar_t>(), d);   \
+    });                                                                        \
   }
 
 namespace vllm {
@@ -625,20 +635,20 @@ __device__ __forceinline__ T gelu_quick_kernel(const T& x) {
 
 }  // namespace vllm
 
-void gelu_new(torch::Tensor& out,    // [..., d]
-              torch::Tensor& input)  // [..., d]
+void gelu_new(torch::stable::Tensor& out,    // [..., d]
+              torch::stable::Tensor& input)  // [..., d]
 {
   LAUNCH_ACTIVATION_KERNEL(vllm::gelu_new_kernel);
 }
 
-void gelu_fast(torch::Tensor& out,    // [..., d]
-               torch::Tensor& input)  // [..., d]
+void gelu_fast(torch::stable::Tensor& out,    // [..., d]
+               torch::stable::Tensor& input)  // [..., d]
 {
   LAUNCH_ACTIVATION_KERNEL(vllm::gelu_fast_kernel);
 }
 
-void gelu_quick(torch::Tensor& out,    // [..., d]
-                torch::Tensor& input)  // [..., d]
+void gelu_quick(torch::stable::Tensor& out,    // [..., d]
+                torch::stable::Tensor& input)  // [..., d]
 {
   LAUNCH_ACTIVATION_KERNEL(vllm::gelu_quick_kernel);
 }

csrc/libtorch_stable/dispatch_utils.h

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

csrc/libtorch_stable/fused_qknorm_rope_kernel.cu

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

csrc/libtorch_stable/layernorm_kernels.cu

@@ -1,11 +1,12 @@
-#include "type_convert.cuh"
-#include "dispatch_utils.h"
-#include "cub_helpers.h"
-#include "core/batch_invariant.hpp"
-#include "libtorch_stable/quantization/vectorization_utils.cuh"
+#include <numeric>
 
-#include <torch/cuda.h>
-#include <c10/cuda/CUDAGuard.h>
+#include "torch_utils.h"
+
+#include "../cub_helpers.h"
+#include "../core/batch_invariant.hpp"
+#include "../type_convert.cuh"
+#include "dispatch_utils.h"
+#include "quantization/vectorization_utils.cuh"
 
 namespace vllm {
 
@@ -189,16 +190,16 @@ fused_add_rms_norm_kernel(
 
 }  // namespace vllm
 
-void rms_norm(torch::Tensor& out,     // [..., hidden_size]
-              torch::Tensor& input,   // [..., hidden_size]
-              torch::Tensor& weight,  // [hidden_size]
+void rms_norm(torch::stable::Tensor& out,     // [..., hidden_size]
+              torch::stable::Tensor& input,   // [..., hidden_size]
+              torch::stable::Tensor& weight,  // [hidden_size]
               double epsilon) {
-  TORCH_CHECK(out.is_contiguous());
+  STD_TORCH_CHECK(out.is_contiguous());
   if (input.stride(-1) != 1) {
-    input = input.contiguous();
+    input = torch::stable::contiguous(input);
   }
-  TORCH_CHECK(input.stride(-1) == 1);
-  TORCH_CHECK(weight.is_contiguous());
+  STD_TORCH_CHECK(input.stride(-1) == 1);
+  STD_TORCH_CHECK(weight.is_contiguous());
 
   int hidden_size = input.size(-1);
 
@@ -213,45 +214,49 @@ void rms_norm(torch::Tensor& out,     // [..., hidden_size]
   // For large num_tokens, use smaller blocks to increase SM concurrency.
   const int max_block_size = (num_tokens < 256) ? 1024 : 256;
   dim3 grid(num_tokens);
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  VLLM_DISPATCH_RANK234(num_dims, [&] {
-    VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_kernel", [&] {
-      const int calculated_vec_size =
-          std::gcd(16 / sizeof(scalar_t), hidden_size);
-      const int block_size =
-          std::min(hidden_size / calculated_vec_size, max_block_size);
-      dim3 block(block_size);
-      VLLM_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
-        vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank>
-            <<<grid, block, 0, stream>>>(
-                out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
-                input_stride_d2, input_stride_d3, input_stride_d4,
-                input_shape_d2, input_shape_d3, weight.data_ptr<scalar_t>(),
-                epsilon, num_tokens, hidden_size);
-      });
-    });
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+  VLLM_STABLE_DISPATCH_RANK234(num_dims, [&] {
+    VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+        input.scalar_type(), "rms_norm_kernel", [&] {
+          const int calculated_vec_size =
+              std::gcd(16 / sizeof(scalar_t), hidden_size);
+          const int block_size =
+              std::min(hidden_size / calculated_vec_size, max_block_size);
+          dim3 block(block_size);
+          VLLM_STABLE_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
+            vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank>
+                <<<grid, block, 0, stream>>>(
+                    out.mutable_data_ptr<scalar_t>(),
+                    input.const_data_ptr<scalar_t>(), input_stride_d2,
+                    input_stride_d3, input_stride_d4, input_shape_d2,
+                    input_shape_d3, weight.const_data_ptr<scalar_t>(), epsilon,
+                    num_tokens, hidden_size);
+          });
+        });
   });
 }
 
-#define LAUNCH_FUSED_ADD_RMS_NORM(width)                                    \
-  VLLM_DISPATCH_FLOATING_TYPES(                                             \
-      input.scalar_type(), "fused_add_rms_norm_kernel", [&] {               \
-        vllm::fused_add_rms_norm_kernel<scalar_t, width>                    \
-            <<<grid, block, 0, stream>>>(                                   \
-                input.data_ptr<scalar_t>(), input_stride,                   \
-                residual.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(), \
-                epsilon, num_tokens, hidden_size);                          \
+#define LAUNCH_FUSED_ADD_RMS_NORM(width)                                \
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                  \
+      input.scalar_type(), "fused_add_rms_norm_kernel", [&] {           \
+        vllm::fused_add_rms_norm_kernel<scalar_t, width>                \
+            <<<grid, block, 0, stream>>>(                               \
+                input.mutable_data_ptr<scalar_t>(), input_stride,       \
+                residual.mutable_data_ptr<scalar_t>(),                  \
+                weight.const_data_ptr<scalar_t>(), epsilon, num_tokens, \
+                hidden_size);                                           \
       });
 
-void fused_add_rms_norm(torch::Tensor& input,     // [..., hidden_size]
-                        torch::Tensor& residual,  // [..., hidden_size]
-                        torch::Tensor& weight,    // [hidden_size]
+void fused_add_rms_norm(torch::stable::Tensor& input,     // [..., hidden_size]
+                        torch::stable::Tensor& residual,  // [..., hidden_size]
+                        torch::stable::Tensor& weight,    // [hidden_size]
                         double epsilon) {
-  TORCH_CHECK(weight.scalar_type() == input.scalar_type());
-  TORCH_CHECK(input.scalar_type() == residual.scalar_type());
-  TORCH_CHECK(residual.is_contiguous());
-  TORCH_CHECK(weight.is_contiguous());
+  STD_TORCH_CHECK(weight.scalar_type() == input.scalar_type());
+  STD_TORCH_CHECK(input.scalar_type() == residual.scalar_type());
+  STD_TORCH_CHECK(residual.is_contiguous());
+  STD_TORCH_CHECK(weight.is_contiguous());
   int hidden_size = input.size(-1);
   int64_t input_stride = input.stride(-2);
   int num_tokens = input.numel() / hidden_size;
@@ -263,8 +268,9 @@ void fused_add_rms_norm(torch::Tensor& input,     // [..., hidden_size]
      hiding on global mem ops. */
   const int max_block_size = (num_tokens < 256) ? 1024 : 256;
   dim3 block(std::min(hidden_size, max_block_size));
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
   /*If the tensor types are FP16/BF16, try to use the optimized kernel
     with packed + vectorized ops.
     Max optimization is achieved with a width-8 vector of FP16/BF16s

csrc/libtorch_stable/layernorm_quant_kernels.cu

@@ -5,15 +5,16 @@
  * Currently, only static fp8 quantization is supported.
  */
 
-#include "type_convert.cuh"
-#include "quantization/w8a8/fp8/common.cuh"
-#include "dispatch_utils.h"
-#include "cub_helpers.h"
-#include "core/batch_invariant.hpp"
-#include "libtorch_stable/quantization/vectorization_utils.cuh"
+#include <numeric>
 
-#include <torch/cuda.h>
-#include <c10/cuda/CUDAGuard.h>
+#include "torch_utils.h"
+
+#include "../cub_helpers.h"
+#include "../core/batch_invariant.hpp"
+#include "../quantization/w8a8/fp8/common.cuh"
+#include "../type_convert.cuh"
+#include "dispatch_utils.h"
+#include "quantization/vectorization_utils.cuh"
 
 namespace vllm {
 
@@ -202,12 +203,13 @@ fused_add_rms_norm_static_fp8_quant_kernel(
 
 }  // namespace vllm
 
-void rms_norm_static_fp8_quant(torch::Tensor& out,     // [..., hidden_size]
-                               torch::Tensor& input,   // [..., hidden_size]
-                               torch::Tensor& weight,  // [hidden_size]
-                               torch::Tensor& scale,   // [1]
-                               double epsilon) {
-  TORCH_CHECK(out.is_contiguous());
+void rms_norm_static_fp8_quant(
+    torch::stable::Tensor& out,     // [..., hidden_size]
+    torch::stable::Tensor& input,   // [..., hidden_size]
+    torch::stable::Tensor& weight,  // [hidden_size]
+    torch::stable::Tensor& scale,   // [1]
+    double epsilon) {
+  STD_TORCH_CHECK(out.is_contiguous());
   int hidden_size = input.size(-1);
   int input_stride = input.stride(-2);
   int num_tokens = input.numel() / hidden_size;
@@ -215,24 +217,26 @@ void rms_norm_static_fp8_quant(torch::Tensor& out,     // [..., hidden_size]
   // For large num_tokens, use smaller blocks to increase SM concurrency.
   const int max_block_size = (num_tokens < 256) ? 1024 : 256;
   dim3 grid(num_tokens);
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  VLLM_DISPATCH_FLOATING_TYPES(
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "rms_norm_kernel_scalar_type", [&] {
-        VLLM_DISPATCH_FP8_TYPES(
+        VLLM_STABLE_DISPATCH_FP8_TYPES(
             out.scalar_type(), "rms_norm_kernel_fp8_type", [&] {
               const int calculated_vec_size =
                   std::gcd(16 / sizeof(scalar_t), hidden_size);
               const int block_size =
                   std::min(hidden_size / calculated_vec_size, max_block_size);
               dim3 block(block_size);
-              VLLM_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
+              VLLM_STABLE_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
                 vllm::rms_norm_static_fp8_quant_kernel<scalar_t, fp8_t,
                                                        vec_size>
                     <<<grid, block, 0, stream>>>(
-                        out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
-                        input_stride, weight.data_ptr<scalar_t>(),
-                        scale.data_ptr<float>(), epsilon, num_tokens,
+                        out.mutable_data_ptr<fp8_t>(),
+                        input.const_data_ptr<scalar_t>(), input_stride,
+                        weight.const_data_ptr<scalar_t>(),
+                        scale.const_data_ptr<float>(), epsilon, num_tokens,
                         hidden_size);
               });
             });
@@ -240,30 +244,32 @@ void rms_norm_static_fp8_quant(torch::Tensor& out,     // [..., hidden_size]
 }
 
 #define LAUNCH_FUSED_ADD_RMS_NORM(width)                                     \
-  VLLM_DISPATCH_FLOATING_TYPES(                                              \
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                       \
       input.scalar_type(), "fused_add_rms_norm_kernel_scalar_type", [&] {    \
-        VLLM_DISPATCH_FP8_TYPES(                                             \
+        VLLM_STABLE_DISPATCH_FP8_TYPES(                                      \
             out.scalar_type(), "fused_add_rms_norm_kernel_fp8_type", [&] {   \
               vllm::fused_add_rms_norm_static_fp8_quant_kernel<scalar_t,     \
                                                                width, fp8_t> \
                   <<<grid, block, 0, stream>>>(                              \
-                      out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),     \
-                      input_stride, residual.data_ptr<scalar_t>(),           \
-                      weight.data_ptr<scalar_t>(), scale.data_ptr<float>(),  \
-                      epsilon, num_tokens, hidden_size);                     \
+                      out.mutable_data_ptr<fp8_t>(),                         \
+                      input.mutable_data_ptr<scalar_t>(), input_stride,      \
+                      residual.mutable_data_ptr<scalar_t>(),                 \
+                      weight.const_data_ptr<scalar_t>(),                     \
+                      scale.const_data_ptr<float>(), epsilon, num_tokens,    \
+                      hidden_size);                                          \
             });                                                              \
       });
 void fused_add_rms_norm_static_fp8_quant(
-    torch::Tensor& out,       // [..., hidden_size],
-    torch::Tensor& input,     // [..., hidden_size]
-    torch::Tensor& residual,  // [..., hidden_size]
-    torch::Tensor& weight,    // [hidden_size]
-    torch::Tensor& scale,     // [1]
+    torch::stable::Tensor& out,       // [..., hidden_size],
+    torch::stable::Tensor& input,     // [..., hidden_size]
+    torch::stable::Tensor& residual,  // [..., hidden_size]
+    torch::stable::Tensor& weight,    // [hidden_size]
+    torch::stable::Tensor& scale,     // [1]
     double epsilon) {
-  TORCH_CHECK(out.is_contiguous());
-  TORCH_CHECK(residual.is_contiguous());
-  TORCH_CHECK(residual.scalar_type() == input.scalar_type());
-  TORCH_CHECK(weight.scalar_type() == input.scalar_type());
+  STD_TORCH_CHECK(out.is_contiguous());
+  STD_TORCH_CHECK(residual.is_contiguous());
+  STD_TORCH_CHECK(residual.scalar_type() == input.scalar_type());
+  STD_TORCH_CHECK(weight.scalar_type() == input.scalar_type());
   int hidden_size = input.size(-1);
   int input_stride = input.stride(-2);
   int num_tokens = input.numel() / hidden_size;
@@ -275,8 +281,9 @@ void fused_add_rms_norm_static_fp8_quant(
      hiding on global mem ops. */
   const int max_block_size = (num_tokens < 256) ? 1024 : 256;
   dim3 block(std::min(hidden_size, max_block_size));
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
   /*If the tensor types are FP16/BF16, try to use the optimized kernel
     with packed + vectorized ops.
     Max optimization is achieved with a width-8 vector of FP16/BF16s

csrc/libtorch_stable/ops.h

@@ -166,3 +166,139 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
 
 torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x,
                                          bool inplace);
+
+// Layernorm kernels (shared CUDA/ROCm)
+void rms_norm(torch::stable::Tensor& out, torch::stable::Tensor& input,
+              torch::stable::Tensor& weight, double epsilon);
+
+void fused_add_rms_norm(torch::stable::Tensor& input,
+                        torch::stable::Tensor& residual,
+                        torch::stable::Tensor& weight, double epsilon);
+
+// Layernorm-quant kernels (shared CUDA/ROCm)
+void rms_norm_static_fp8_quant(torch::stable::Tensor& out,
+                               torch::stable::Tensor& input,
+                               torch::stable::Tensor& weight,
+                               torch::stable::Tensor& scale, double epsilon);
+
+void fused_add_rms_norm_static_fp8_quant(torch::stable::Tensor& out,
+                                         torch::stable::Tensor& input,
+                                         torch::stable::Tensor& residual,
+                                         torch::stable::Tensor& weight,
+                                         torch::stable::Tensor& scale,
+                                         double epsilon);
+
+// Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm)
+void rms_norm_dynamic_per_token_quant(
+    torch::stable::Tensor& out, torch::stable::Tensor const& input,
+    torch::stable::Tensor const& weight, torch::stable::Tensor& scales,
+    double const var_epsilon, std::optional<torch::stable::Tensor> scale_ub,
+    std::optional<torch::stable::Tensor> residual);
+
+void rms_norm_per_block_quant(torch::stable::Tensor& out,
+                              torch::stable::Tensor const& input,
+                              torch::stable::Tensor const& weight,
+                              torch::stable::Tensor& scales,
+                              double const var_epsilon,
+                              std::optional<torch::stable::Tensor> scale_ub,
+                              std::optional<torch::stable::Tensor> residual,
+                              int64_t group_size, bool is_scale_transposed);
+
+// Positional encoding kernels (shared CUDA/ROCm)
+void rotary_embedding(torch::stable::Tensor& positions,
+                      torch::stable::Tensor& query,
+                      std::optional<torch::stable::Tensor> key,
+                      int64_t head_size, torch::stable::Tensor& cos_sin_cache,
+                      bool is_neox, int64_t rope_dim_offset, bool inverse);
+
+void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q,
+                        int64_t num_heads_k, int64_t num_heads_v,
+                        int64_t head_dim, double eps,
+                        torch::stable::Tensor& q_weight,
+                        torch::stable::Tensor& k_weight,
+                        torch::stable::Tensor& cos_sin_cache, bool is_neox,
+                        torch::stable::Tensor& position_ids,
+                        int64_t forced_token_heads_per_warp);
+
+// 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);
+void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input);
+void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
+void gelu_tanh_and_mul(torch::stable::Tensor& out,
+                       torch::stable::Tensor& input);
+void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
+                     double threshold);
+void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
+                       double alpha = 1.702, double limit = 7.0);
+void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input);
+void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input);
+void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input);
+
+// INT8 quantization kernels (shared CUDA/ROCm)
+void static_scaled_int8_quant(torch::stable::Tensor& out,
+                              torch::stable::Tensor const& input,
+                              torch::stable::Tensor const& scale,
+                              std::optional<torch::stable::Tensor> const& azp);
+
+void dynamic_scaled_int8_quant(torch::stable::Tensor& out,
+                               torch::stable::Tensor const& input,
+                               torch::stable::Tensor& scales,
+                               std::optional<torch::stable::Tensor> const& azp);
+
+// FP8 quantization kernels (shared CUDA/ROCm)
+void static_scaled_fp8_quant(
+    torch::stable::Tensor& out, torch::stable::Tensor const& input,
+    torch::stable::Tensor const& scale,
+    std::optional<torch::headeronly::IntHeaderOnlyArrayRef> group_shape =
+        std::nullopt);
+
+void dynamic_scaled_fp8_quant(torch::stable::Tensor& out,
+                              torch::stable::Tensor const& input,
+                              torch::stable::Tensor& scale);
+
+void dynamic_per_token_scaled_fp8_quant(
+    torch::stable::Tensor& out, torch::stable::Tensor const& input,
+    torch::stable::Tensor& scale,
+    std::optional<torch::stable::Tensor> const& scale_ub);
+
+// GPTQ kernels (shared CUDA/ROCm)
+torch::stable::Tensor gptq_gemm(torch::stable::Tensor a,
+                                torch::stable::Tensor b_q_weight,
+                                torch::stable::Tensor b_gptq_qzeros,
+                                torch::stable::Tensor b_gptq_scales,
+                                torch::stable::Tensor b_g_idx, bool use_exllama,
+                                bool use_v2_format, int64_t bit);
+
+void gptq_shuffle(torch::stable::Tensor q_weight, torch::stable::Tensor q_perm,
+                  int64_t bit);
+
+// GGML kernels (shared CUDA/ROCm)
+torch::stable::Tensor ggml_dequantize(
+    torch::stable::Tensor W, int64_t type, int64_t m, int64_t n,
+    std::optional<torch::headeronly::ScalarType> const& dtype);
+
+torch::stable::Tensor ggml_mul_mat_vec_a8(torch::stable::Tensor W,
+                                          torch::stable::Tensor X, int64_t type,
+                                          int64_t row);
+
+torch::stable::Tensor ggml_mul_mat_a8(torch::stable::Tensor W,
+                                      torch::stable::Tensor X, int64_t type,
+                                      int64_t row);
+
+torch::stable::Tensor ggml_moe_a8(torch::stable::Tensor X,
+                                  torch::stable::Tensor W,
+                                  torch::stable::Tensor sorted_token_ids,
+                                  torch::stable::Tensor expert_ids,
+                                  torch::stable::Tensor num_tokens_post_padded,
+                                  int64_t type, int64_t row, int64_t top_k,
+                                  int64_t tokens);
+
+torch::stable::Tensor ggml_moe_a8_vec(torch::stable::Tensor X,
+                                      torch::stable::Tensor W,
+                                      torch::stable::Tensor topk_ids,
+                                      int64_t top_k, int64_t type, int64_t row,
+                                      int64_t tokens);
+
+int64_t ggml_moe_get_block_size(int64_t type);

csrc/libtorch_stable/pos_encoding_kernels.cu

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

csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu

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

csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu

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

csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu

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

csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh

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

csrc/libtorch_stable/quantization/fused_kernels/quant_conversions.cuh

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

csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu

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

csrc/libtorch_stable/quantization/gptq/q_gemm.cu

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

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

@@ -1,11 +1,9 @@
-#include "common.cuh"
-#include "dispatch_utils.h"
-#include "cub_helpers.h"
-#include "libtorch_stable/quantization/vectorization_utils.cuh"
-#include <c10/cuda/CUDAGuard.h>
-#include <ATen/cuda/Exceptions.h>
-#include <tuple>
-
+#include "../../../../quantization/w8a8/fp8/common.cuh"
+#include "../../../dispatch_utils.h"
+#include "../../../../cub_helpers.h"
+#include "../../vectorization_utils.cuh"
+#include "../../../torch_utils.h"
+#include <torch/csrc/stable/macros.h>
 namespace vllm {
 
 // STRIDE_I_ZERO: true if scale_stride_i == 0 (per-tensor or per-channel)
@@ -183,16 +181,16 @@ __global__ void dynamic_per_token_scaled_fp8_quant_kernel_strided(
 }  // namespace vllm
 
 void static_scaled_fp8_quant(
-    torch::Tensor& out,          // [..., d]
-    torch::Tensor const& input,  // [..., d]
-    torch::Tensor const& scale,  // various shapes
-    std::optional<std::tuple<int64_t, int64_t>>
-        opt_group_shape)  // optional explicit (group_m, group_n)
+    torch::stable::Tensor& out,          // [..., d]
+    torch::stable::Tensor const& input,  // [..., d]
+    torch::stable::Tensor const& scale,  // various shapes
+    std::optional<torch::headeronly::IntHeaderOnlyArrayRef>
+        opt_group_shape)  // optional explicit [group_m, group_n]
 {
-  TORCH_CHECK(input.stride(-1) == 1,
-              "last dimension of input must be contiguous");
-  TORCH_CHECK(out.stride(-1) == 1,
-              "last dimension of output must be contiguous");
+  STD_TORCH_CHECK(input.stride(-1) == 1,
+                  "last dimension of input must be contiguous");
+  STD_TORCH_CHECK(out.stride(-1) == 1,
+                  "last dimension of output must be contiguous");
 
   const int hidden_size = input.size(-1);              // N (columns)
   const int num_tokens = input.numel() / hidden_size;  // M (rows)
@@ -212,13 +210,18 @@ void static_scaled_fp8_quant(
   } else if (scale.dim() == 1) {
     // 1D scale: require explicit group_shape to disambiguate per-channel vs
     // per-token (avoids edge case where num_tokens == hidden_size)
-    TORCH_CHECK(opt_group_shape.has_value(),
-                "1D scale requires explicit group_shape to disambiguate "
-                "per-channel vs per-token quantization. "
-                "Use group_shape=(-1, 1) for per-channel or group_shape=(1, "
-                "-1) for per-token.");
+    STD_TORCH_CHECK(
+        opt_group_shape.has_value(),
+        "1D scale requires explicit group_shape to disambiguate "
+        "per-channel vs per-token quantization. "
+        "Use group_shape=(-1, 1) for per-channel or group_shape=(1, "
+        "-1) for per-token.");
+    STD_TORCH_CHECK(opt_group_shape->size() == 2,
+                    "group_shape must have exactly 2 elements, got ",
+                    opt_group_shape->size());
 
-    const auto& [opt_group_m, opt_group_n] = opt_group_shape.value();
+    const auto opt_group_m = (*opt_group_shape)[0];
+    const auto opt_group_n = (*opt_group_shape)[1];
     group_m = opt_group_m == -1 ? num_tokens : static_cast<int>(opt_group_m);
     group_n = opt_group_n == -1 ? hidden_size : static_cast<int>(opt_group_n);
 
@@ -228,11 +231,11 @@ void static_scaled_fp8_quant(
     const int64_t expected_scale_n = hidden_size / group_n;
     const int64_t expected_scale_numel = expected_scale_m * expected_scale_n;
 
-    TORCH_CHECK(scale_len == expected_scale_numel, "1D scale length (",
-                scale_len, ") does not match expected size (",
-                expected_scale_numel, ") for group_shape (", opt_group_m, ", ",
-                opt_group_n, ") with input shape (", num_tokens, ", ",
-                hidden_size, ")");
+    STD_TORCH_CHECK(scale_len == expected_scale_numel, "1D scale length (",
+                    scale_len, ") does not match expected size (",
+                    expected_scale_numel, ") for group_shape (", opt_group_m,
+                    ", ", opt_group_n, ") with input shape (", num_tokens, ", ",
+                    hidden_size, ")");
 
     // For 1D scale, determine strides based on which dim is trivial
     // Scale indexing: scale[gi * scale_stride_i + gj * scale_stride_j]
@@ -248,7 +251,7 @@ void static_scaled_fp8_quant(
       scale_stride_i = scale.stride(0);
       scale_stride_j = 0;
     } else {
-      TORCH_CHECK(
+      STD_TORCH_CHECK(
           false,
           "1D scale can only be used when one of the scale dimensions is 1. "
           "For 2D group scaling, use a 2D scale tensor.");
@@ -259,10 +262,12 @@ void static_scaled_fp8_quant(
     const int64_t scale_size_0 = scale.size(0);
     const int64_t scale_size_1 = scale.size(1);
 
-    TORCH_CHECK(num_tokens % scale_size_0 == 0, "num_tokens (", num_tokens,
-                ") must be divisible by scale.size(0) (", scale_size_0, ")");
-    TORCH_CHECK(hidden_size % scale_size_1 == 0, "hidden_size (", hidden_size,
-                ") must be divisible by scale.size(1) (", scale_size_1, ")");
+    STD_TORCH_CHECK(num_tokens % scale_size_0 == 0, "num_tokens (", num_tokens,
+                    ") must be divisible by scale.size(0) (", scale_size_0,
+                    ")");
+    STD_TORCH_CHECK(hidden_size % scale_size_1 == 0, "hidden_size (",
+                    hidden_size, ") must be divisible by scale.size(1) (",
+                    scale_size_1, ")");
 
     // Infer from 2D scale shape
     int inferred_group_m = num_tokens / scale_size_0;
@@ -270,16 +275,21 @@ void static_scaled_fp8_quant(
 
     // Use explicit if provided, otherwise use inferred
     if (opt_group_shape.has_value()) {
-      const auto& [opt_group_m, opt_group_n] = opt_group_shape.value();
+      STD_TORCH_CHECK(opt_group_shape->size() == 2,
+                      "group_shape must have exactly 2 elements, got ",
+                      opt_group_shape->size());
+      const auto opt_group_m = (*opt_group_shape)[0];
+      const auto opt_group_n = (*opt_group_shape)[1];
       group_m = opt_group_m == -1 ? num_tokens : static_cast<int>(opt_group_m);
       group_n = opt_group_n == -1 ? hidden_size : static_cast<int>(opt_group_n);
 
       // Validate explicit matches inferred
-      TORCH_CHECK(group_m == inferred_group_m && group_n == inferred_group_n,
-                  "Explicit group_shape (", opt_group_m, ", ", opt_group_n,
-                  ") does not match inferred group shape (", inferred_group_m,
-                  ", ", inferred_group_n, ") from 2D scale tensor shape (",
-                  scale_size_0, ", ", scale_size_1, ")");
+      STD_TORCH_CHECK(
+          group_m == inferred_group_m && group_n == inferred_group_n,
+          "Explicit group_shape (", opt_group_m, ", ", opt_group_n,
+          ") does not match inferred group shape (", inferred_group_m, ", ",
+          inferred_group_n, ") from 2D scale tensor shape (", scale_size_0,
+          ", ", scale_size_1, ")");
     } else {
       group_m = inferred_group_m;
       group_n = inferred_group_n;
@@ -288,8 +298,8 @@ void static_scaled_fp8_quant(
     scale_stride_i = scale.stride(0);
     scale_stride_j = scale.stride(1);
   } else {
-    TORCH_CHECK(false, "scale must be 0D, 1D, or 2D tensor, but got ",
-                scale.dim(), "D");
+    STD_TORCH_CHECK(false, "scale must be 0D, 1D, or 2D tensor, but got ",
+                    scale.dim(), "D");
   }
 
   const int block_size = 256;
@@ -299,37 +309,39 @@ void static_scaled_fp8_quant(
   const int64_t in_row_stride = input.stride(-2);
   const int64_t out_row_stride = out.stride(-2);
 
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
 
   // Dispatch to template-specialized kernel based on stride pattern
-  VLLM_DISPATCH_FLOATING_TYPES(
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "scaled_fp8_quant_kernel_scalar_type", [&] {
-        VLLM_DISPATCH_FP8_TYPES(
+        VLLM_STABLE_DISPATCH_FP8_TYPES(
             out.scalar_type(), "scaled_fp8_quant_kernel_fp8_type", [&] {
-              VLLM_DISPATCH_BOOL(scale_stride_i == 0, S0_ZERO, [&] {
-                VLLM_DISPATCH_BOOL(scale_stride_j == 0, S1_ZERO, [&] {
+              VLLM_STABLE_DISPATCH_BOOL(scale_stride_i == 0, S0_ZERO, [&] {
+                VLLM_STABLE_DISPATCH_BOOL(scale_stride_j == 0, S1_ZERO, [&] {
                   vllm::scaled_fp8_quant_kernel_strided_group_shape<
                       scalar_t, fp8_t, S0_ZERO, S1_ZERO>
                       <<<grid, block, 0, stream>>>(
-                          out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
-                          scale.data_ptr<float>(), hidden_size, in_row_stride,
-                          out_row_stride, group_m, group_n, scale_stride_i,
-                          scale_stride_j);
+                          out.mutable_data_ptr<fp8_t>(),
+                          input.const_data_ptr<scalar_t>(),
+                          scale.const_data_ptr<float>(), hidden_size,
+                          in_row_stride, out_row_stride, group_m, group_n,
+                          scale_stride_i, scale_stride_j);
                 });
               });
             });
       });
 }
 
-void dynamic_scaled_fp8_quant(torch::Tensor& out,          // [..., d]
-                              torch::Tensor const& input,  // [..., d]
-                              torch::Tensor& scale)        // [1]
+void dynamic_scaled_fp8_quant(torch::stable::Tensor& out,          // [..., d]
+                              torch::stable::Tensor const& input,  // [..., d]
+                              torch::stable::Tensor& scale)        // [1]
 {
-  TORCH_CHECK(input.stride(-1) == 1,
-              "last dimension of input must be contiguous");
-  TORCH_CHECK(out.stride(-1) == 1,
-              "last dimension of output must be contiguous");
+  STD_TORCH_CHECK(input.stride(-1) == 1,
+                  "last dimension of input must be contiguous");
+  STD_TORCH_CHECK(out.stride(-1) == 1,
+                  "last dimension of output must be contiguous");
 
   const int hidden_size = input.size(-1);
   const int num_tokens = input.numel() / hidden_size;
@@ -340,40 +352,43 @@ void dynamic_scaled_fp8_quant(torch::Tensor& out,          // [..., d]
   const int64_t in_row_stride = input.stride(-2);
   const int64_t out_row_stride = out.stride(-2);
 
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
 
   // scale tensor should be initialised to <=0 before reduction
-  AT_CUDA_CHECK(
-      cudaMemsetAsync(scale.data_ptr<float>(), 0, sizeof(float), stream));
+  STD_CUDA_CHECK(cudaMemsetAsync(scale.mutable_data_ptr<float>(), 0,
+                                 sizeof(float), stream));
 
-  VLLM_DISPATCH_FLOATING_TYPES(
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "scaled_fp8_quant_kernel_scalar_type", [&] {
-        VLLM_DISPATCH_FP8_TYPES(
+        VLLM_STABLE_DISPATCH_FP8_TYPES(
             out.scalar_type(), "scaled_fp8_quant_kernel_fp8_type", [&] {
               vllm::segmented_max_reduction_strided<scalar_t, fp8_t>
                   <<<grid, block, 0, stream>>>(
-                      scale.data_ptr<float>(), input.data_ptr<scalar_t>(),
-                      hidden_size, in_row_stride,
-                      static_cast<int64_t>(num_tokens));
+                      scale.mutable_data_ptr<float>(),
+                      input.const_data_ptr<scalar_t>(), hidden_size,
+                      in_row_stride, static_cast<int64_t>(num_tokens));
 
               vllm::scaled_fp8_quant_kernel_strided_dynamic<scalar_t, fp8_t>
-                  <<<grid, block, 0, stream>>>(
-                      out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
-                      scale.data_ptr<float>(), hidden_size, in_row_stride,
-                      out_row_stride);
+                  <<<grid, block, 0, stream>>>(out.mutable_data_ptr<fp8_t>(),
+                                               input.const_data_ptr<scalar_t>(),
+                                               scale.const_data_ptr<float>(),
+                                               hidden_size, in_row_stride,
+                                               out_row_stride);
             });
       });
 }
 
 void dynamic_per_token_scaled_fp8_quant(
-    torch::Tensor& out,          // [..., d]
-    torch::Tensor const& input,  // [..., d]
-    torch::Tensor& scales, std::optional<at::Tensor> const& scale_ub) {
-  TORCH_CHECK(input.stride(-1) == 1,
-              "last dimension of input must be contiguous");
-  TORCH_CHECK(out.stride(-1) == 1,
-              "last dimension of output must be contiguous");
+    torch::stable::Tensor& out,          // [..., d]
+    torch::stable::Tensor const& input,  // [..., d]
+    torch::stable::Tensor& scales,
+    std::optional<torch::stable::Tensor> const& scale_ub) {
+  STD_TORCH_CHECK(input.stride(-1) == 1,
+                  "last dimension of input must be contiguous");
+  STD_TORCH_CHECK(out.stride(-1) == 1,
+                  "last dimension of output must be contiguous");
 
   const int hidden_size = input.size(-1);
   const int num_tokens = input.numel() / hidden_size;
@@ -384,20 +399,24 @@ void dynamic_per_token_scaled_fp8_quant(
   const int64_t in_row_stride = input.stride(-2);
   const int64_t out_row_stride = out.stride(-2);
 
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  VLLM_DISPATCH_FLOATING_TYPES(
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
       input.scalar_type(),
       "dynamic_per_token_scaled_fp8_quant_kernel_scalar_type", [&] {
-        VLLM_DISPATCH_FP8_TYPES(
+        VLLM_STABLE_DISPATCH_FP8_TYPES(
             out.scalar_type(),
             "dynamic_per_token_scaled_fp8_quant_kernel_fp8_type", [&] {
-              vllm::dynamic_per_token_scaled_fp8_quant_kernel_strided<
-                  scalar_t, fp8_t><<<grid, block, 0, stream>>>(
-                  out.data_ptr<fp8_t>(), scales.data_ptr<float>(),
-                  input.data_ptr<scalar_t>(),
-                  scale_ub.has_value() ? scale_ub->data_ptr<float>() : nullptr,
-                  hidden_size, in_row_stride, out_row_stride);
+              vllm::dynamic_per_token_scaled_fp8_quant_kernel_strided<scalar_t,
+                                                                      fp8_t>
+                  <<<grid, block, 0, stream>>>(
+                      out.mutable_data_ptr<fp8_t>(),
+                      scales.mutable_data_ptr<float>(),
+                      input.const_data_ptr<scalar_t>(),
+                      scale_ub.has_value() ? scale_ub->const_data_ptr<float>()
+                                           : nullptr,
+                      hidden_size, in_row_stride, out_row_stride);
             });
       });
 }

csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu

@@ -1,12 +1,11 @@
-#include <ATen/cuda/CUDAContext.h>
-#include <torch/all.h>
-#include <c10/cuda/CUDAGuard.h>
+#include <torch/csrc/stable/tensor.h>
 
 #include <cmath>
 
-#include "dispatch_utils.h"
-#include "libtorch_stable/quantization/vectorization_utils.cuh"
-#include "cub_helpers.h"
+#include "../../../dispatch_utils.h"
+#include "../../../torch_utils.h"
+#include "../../vectorization_utils.cuh"
+#include "../../../../cub_helpers.h"
 
 static inline __device__ int8_t float_to_int8_rn(float x) {
 #ifdef USE_ROCM
@@ -263,66 +262,73 @@ __global__ void dynamic_scaled_int8_azp_quant_kernel(
 
 }  // namespace vllm
 
-void static_scaled_int8_quant(torch::Tensor& out,          // [..., hidden_size]
-                              torch::Tensor const& input,  // [..., hidden_size]
-                              torch::Tensor const& scale,
-                              std::optional<torch::Tensor> const& azp) {
-  TORCH_CHECK(input.is_contiguous());
-  TORCH_CHECK(out.is_contiguous());
-  TORCH_CHECK(scale.numel() == 1);
-  TORCH_CHECK(!azp || azp->numel() == 1);
+void static_scaled_int8_quant(
+    torch::stable::Tensor& out,          // [..., hidden_size]
+    torch::stable::Tensor const& input,  // [..., hidden_size]
+    torch::stable::Tensor const& scale,
+    std::optional<torch::stable::Tensor> const& azp) {
+  STD_TORCH_CHECK(input.is_contiguous());
+  STD_TORCH_CHECK(out.is_contiguous());
+  STD_TORCH_CHECK(scale.numel() == 1);
+  STD_TORCH_CHECK(!azp || azp->numel() == 1);
 
   int const hidden_size = input.size(-1);
   int const num_tokens = input.numel() / hidden_size;
   dim3 const grid(num_tokens);
   dim3 const block(std::min(hidden_size, 256));
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  VLLM_DISPATCH_FLOATING_TYPES(
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "static_scaled_int8_quant_kernel", [&] {
         if (!azp) {
           vllm::static_scaled_int8_quant_kernel<scalar_t, float>
-              <<<grid, block, 0, stream>>>(
-                  input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
-                  scale.data_ptr<float>(), hidden_size);
+              <<<grid, block, 0, stream>>>(input.const_data_ptr<scalar_t>(),
+                                           out.mutable_data_ptr<int8_t>(),
+                                           scale.const_data_ptr<float>(),
+                                           hidden_size);
         } else {
           vllm::static_scaled_int8_azp_quant_kernel<scalar_t, float, int32_t>
               <<<grid, block, 0, stream>>>(
-                  input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
-                  scale.data_ptr<float>(), azp->data_ptr<int32_t>(),
-                  hidden_size);
+                  input.const_data_ptr<scalar_t>(),
+                  out.mutable_data_ptr<int8_t>(), scale.const_data_ptr<float>(),
+                  azp->const_data_ptr<int32_t>(), hidden_size);
         }
       });
 }
 
 void dynamic_scaled_int8_quant(
-    torch::Tensor& out,          // [..., hidden_size]
-    torch::Tensor const& input,  // [..., hidden_size]
-    torch::Tensor& scales, std::optional<torch::Tensor> const& azp) {
-  TORCH_CHECK(input.is_contiguous());
-  TORCH_CHECK(out.is_contiguous());
-  TORCH_CHECK(scales.is_contiguous());
-  TORCH_CHECK(!azp || azp->is_contiguous());
+    torch::stable::Tensor& out,          // [..., hidden_size]
+    torch::stable::Tensor const& input,  // [..., hidden_size]
+    torch::stable::Tensor& scales,
+    std::optional<torch::stable::Tensor> const& azp) {
+  STD_TORCH_CHECK(input.is_contiguous());
+  STD_TORCH_CHECK(out.is_contiguous());
+  STD_TORCH_CHECK(scales.is_contiguous());
+  STD_TORCH_CHECK(!azp || azp->is_contiguous());
 
   int const hidden_size = input.size(-1);
   int const num_tokens = input.numel() / hidden_size;
   dim3 const grid(num_tokens);
   dim3 const block(std::min(hidden_size, 256));
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  VLLM_DISPATCH_FLOATING_TYPES(
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "dynamic_scaled_int8_quant_kernel", [&] {
         if (!azp) {
           vllm::dynamic_scaled_int8_quant_kernel<scalar_t, float>
-              <<<grid, block, 0, stream>>>(
-                  input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
-                  scales.data_ptr<float>(), hidden_size);
+              <<<grid, block, 0, stream>>>(input.const_data_ptr<scalar_t>(),
+                                           out.mutable_data_ptr<int8_t>(),
+                                           scales.mutable_data_ptr<float>(),
+                                           hidden_size);
         } else {
           vllm::dynamic_scaled_int8_azp_quant_kernel<scalar_t, float, int32_t>
-              <<<grid, block, 0, stream>>>(
-                  input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
-                  scales.data_ptr<float>(), azp->data_ptr<int32_t>(),
-                  hidden_size);
+              <<<grid, block, 0, stream>>>(input.const_data_ptr<scalar_t>(),
+                                           out.mutable_data_ptr<int8_t>(),
+                                           scales.mutable_data_ptr<float>(),
+                                           azp->mutable_data_ptr<int32_t>(),
+                                           hidden_size);
         }
       });
 }

csrc/libtorch_stable/torch_bindings.cpp

@@ -266,6 +266,162 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
   // Hadamard transforms
   // conditionally compiled so impl registration is in source file
   ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor");
+
+  // Apply Root Mean Square (RMS) Normalization to the input tensor.
+  ops.def(
+      "rms_norm(Tensor! result, Tensor input, Tensor weight, float epsilon) -> "
+      "()");
+
+  // In-place fused Add and RMS Normalization.
+  ops.def(
+      "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor weight, "
+      "float epsilon) -> ()");
+
+  // Layernorm-quant
+  // Apply Root Mean Square (RMS) Normalization to the input tensor.
+  ops.def(
+      "rms_norm_static_fp8_quant(Tensor! result, Tensor input, Tensor weight, "
+      "Tensor scale, float epsilon) -> "
+      "()");
+
+  // In-place fused Add and RMS Normalization.
+  ops.def(
+      "fused_add_rms_norm_static_fp8_quant(Tensor! result, Tensor input, "
+      "Tensor! residual, Tensor weight, "
+      "Tensor scale, float epsilon) -> ()");
+
+  // Fused Layernorm + Quant kernels
+  ops.def(
+      "rms_norm_dynamic_per_token_quant(Tensor! result, Tensor input, "
+      "Tensor weight, Tensor! scale, float epsilon, "
+      "Tensor? scale_ub, Tensor!? residual) -> ()");
+
+  // Fused Layernorm + Block quant kernels
+  ops.def(
+      "rms_norm_per_block_quant(Tensor! result, Tensor input, "
+      "Tensor weight, Tensor! scale, float epsilon, "
+      "Tensor? scale_ub, Tensor!? residual, int group_size, "
+      "bool is_scale_transposed) -> ()");
+
+  // Rotary embedding
+  // Apply GPT-NeoX or GPT-J style rotary embedding to query and key.
+  ops.def(
+      "rotary_embedding(Tensor positions, Tensor! query,"
+      "                 Tensor!? key, int head_size,"
+      "                 Tensor cos_sin_cache, bool is_neox, int "
+      "rope_dim_offset=0, bool inverse=False) -> ()");
+
+  // Function for fused QK Norm and RoPE
+  ops.def(
+      "fused_qk_norm_rope(Tensor! qkv, int num_heads_q, "
+      "int num_heads_k, int num_heads_v, int head_dim, float eps, "
+      "Tensor q_weight, Tensor k_weight, Tensor cos_sin_cache, "
+      "bool is_neox, Tensor position_ids, "
+      "int forced_token_heads_per_warp=-1) -> ()");
+
+  // Activation ops
+  // Activation function used in SwiGLU.
+  ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
+
+  ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
+
+  // SwiGLU activation with input clamping.
+  ops.def(
+      "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
+      "-> ()");
+
+  // Activation function used in GeGLU with `none` approximation.
+  ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
+
+  // Activation function used in GeGLU with `tanh` approximation.
+  ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()");
+
+  // FATReLU implementation.
+  ops.def("fatrelu_and_mul(Tensor! out, Tensor input, float threshold) -> ()");
+
+  ops.def(
+      "swigluoai_and_mul(Tensor! out, Tensor input, float alpha=1.702, float "
+      "limit=7.0) "
+      "-> ()");
+
+  // GELU implementation used in GPT-2.
+  ops.def("gelu_new(Tensor! out, Tensor input) -> ()");
+
+  // Approximate GELU implementation.
+  ops.def("gelu_fast(Tensor! out, Tensor input) -> ()");
+
+  // Quick GELU implementation.
+  ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
+
+  // Compute int8 quantized tensor for given scaling factor.
+  ops.def(
+      "static_scaled_int8_quant(Tensor! result, Tensor input, Tensor scale,"
+      "Tensor? azp) -> ()");
+
+  // Compute int8 quantized tensor and scaling factor
+  ops.def(
+      "dynamic_scaled_int8_quant(Tensor! result, Tensor input, Tensor! scale, "
+      "Tensor!? azp) -> ()");
+
+  // Compute FP8 quantized tensor for given scaling factor.
+  // Supports per-tensor, per-channel, per-token, and arbitrary 2D group
+  // scaling. Optional group_m/group_n specify the group shape explicitly;
+  // required for 1D scales to disambiguate per-channel vs per-token.
+  ops.def(
+      "static_scaled_fp8_quant(Tensor! result, Tensor input, Tensor scale, "
+      "int[]? group_shape=None) -> ()");
+
+  // Compute dynamic-per-tensor FP8 quantized tensor and scaling factor.
+  ops.def(
+      "dynamic_scaled_fp8_quant(Tensor! result, Tensor input, Tensor! scale) "
+      "-> "
+      "()");
+
+  // Compute dynamic-per-token FP8 quantized tensor and scaling factor.
+  ops.def(
+      "dynamic_per_token_scaled_fp8_quant(Tensor! result, Tensor input, "
+      "Tensor! scale, Tensor? scale_ub) -> "
+      "()");
+
+  // Quantized GEMM for GPTQ.
+  // Note: even though the C++ inferred schema is correct for this op, it seems
+  // to prevent the meta function registry.
+  ops.def(
+      "gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, "
+      "Tensor b_gptq_scales, Tensor b_g_idx, bool use_exllama, bool "
+      "use_v2_format, int bit) "
+      "-> Tensor");
+
+  // Post processing for GPTQ.
+  ops.def("gptq_shuffle(Tensor! q_weight, Tensor q_perm, int bit) -> ()");
+
+  // Dequantization for GGML.
+  ops.def(
+      "ggml_dequantize(Tensor W, int type, SymInt m, SymInt n, ScalarType? "
+      "dtype) -> Tensor");
+
+  // mmvq kernel for GGML.
+  ops.def(
+      "ggml_mul_mat_vec_a8(Tensor W, Tensor X, int type, SymInt row) "
+      "-> Tensor");
+
+  // mmq kernel for GGML.
+  ops.def(
+      "ggml_mul_mat_a8(Tensor W, Tensor X, int type, SymInt row) -> Tensor");
+
+  // moe kernel for GGML.
+  ops.def(
+      "ggml_moe_a8(Tensor X, Tensor W, "
+      "Tensor sorted_token_ids, Tensor expert_ids, Tensor "
+      "num_tokens_post_padded, "
+      "int type, SymInt row, SymInt top_k, SymInt tokens) -> Tensor");
+
+  ops.def(
+      "ggml_moe_a8_vec(Tensor X, Tensor W, "
+      "Tensor topk_ids, int top_k, "
+      "int type, SymInt row, SymInt tokens) -> Tensor");
+
+  ops.def("ggml_moe_get_block_size(int type) -> int");
 }
 
 STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
@@ -312,6 +468,57 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   // AllSpark ops: conditionally compiled so impl registrations are in source
   // files (allspark_repack.cu and allspark_qgemm_w8a16.cu)
 #endif
+
+  // 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));
+
+  // Layernorm-quant kernels (shared CUDA/ROCm)
+  ops.impl("rms_norm_static_fp8_quant", TORCH_BOX(&rms_norm_static_fp8_quant));
+  ops.impl("fused_add_rms_norm_static_fp8_quant",
+           TORCH_BOX(&fused_add_rms_norm_static_fp8_quant));
+
+  // Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm)
+  ops.impl("rms_norm_dynamic_per_token_quant",
+           TORCH_BOX(&rms_norm_dynamic_per_token_quant));
+  ops.impl("rms_norm_per_block_quant", TORCH_BOX(&rms_norm_per_block_quant));
+
+  // 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));
+
+  // Activation kernels (shared CUDA/ROCm)
+  ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul));
+  ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu));
+  ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul));
+  ops.impl("gelu_tanh_and_mul", TORCH_BOX(&gelu_tanh_and_mul));
+  ops.impl("fatrelu_and_mul", TORCH_BOX(&fatrelu_and_mul));
+  ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul));
+  ops.impl("gelu_new", TORCH_BOX(&gelu_new));
+  ops.impl("gelu_fast", TORCH_BOX(&gelu_fast));
+  ops.impl("gelu_quick", TORCH_BOX(&gelu_quick));
+  ops.impl("silu_and_mul_with_clamp", TORCH_BOX(&silu_and_mul_clamp));
+
+  // INT8 quantization kernels
+  ops.impl("static_scaled_int8_quant", TORCH_BOX(&static_scaled_int8_quant));
+  ops.impl("dynamic_scaled_int8_quant", TORCH_BOX(&dynamic_scaled_int8_quant));
+
+  // FP8 quantization kernels
+  ops.impl("static_scaled_fp8_quant", TORCH_BOX(&static_scaled_fp8_quant));
+  ops.impl("dynamic_scaled_fp8_quant", TORCH_BOX(&dynamic_scaled_fp8_quant));
+  ops.impl("dynamic_per_token_scaled_fp8_quant",
+           TORCH_BOX(&dynamic_per_token_scaled_fp8_quant));
+
+  // GPTQ kernels
+  ops.impl("gptq_gemm", TORCH_BOX(&gptq_gemm));
+  ops.impl("gptq_shuffle", TORCH_BOX(&gptq_shuffle));
+
+  // GGML kernels
+  ops.impl("ggml_dequantize", TORCH_BOX(&ggml_dequantize));
+  ops.impl("ggml_mul_mat_vec_a8", TORCH_BOX(&ggml_mul_mat_vec_a8));
+  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));
 }
 
 // These capability-check functions take only primitive args (no tensors), so
@@ -329,6 +536,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) {
   ops.impl("cutlass_scaled_mm_supports_fp4",
            TORCH_BOX(&cutlass_scaled_mm_supports_fp4));
 #endif
+
+  // GGML block size lookup (no tensor args)
+  ops.impl("ggml_moe_get_block_size", TORCH_BOX(&ggml_moe_get_block_size));
 }
 
 REGISTER_EXTENSION(_C_stable_libtorch)

csrc/libtorch_stable/torch_utils.h

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

csrc/moe/dsv3_router_gemm_bf16_out.cu

@@ -182,7 +182,7 @@ void invokeRouterGemmBf16Output(__nv_bfloat16* output, T const* mat_a,
   config.stream = stream;
   cudaLaunchAttribute attrs[1];
   attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-  attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
+  attrs[0].val.programmaticStreamSerializationAllowed = 1;
   config.numAttrs = 1;
   config.attrs = attrs;
   cudaLaunchKernelEx(

csrc/moe/dsv3_router_gemm_float_out.cu

@@ -182,7 +182,7 @@ void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b,
   config.stream = stream;
   cudaLaunchAttribute attrs[1];
   attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-  attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
+  attrs[0].val.programmaticStreamSerializationAllowed = 1;
   config.numAttrs = 1;
   config.attrs = attrs;
   cudaLaunchKernelEx(

csrc/moe/dsv3_router_gemm_utils.h

@@ -29,15 +29,3 @@ inline int getSMVersion() {
   auto* props = at::cuda::getCurrentDeviceProperties();
   return props->major * 10 + props->minor;
 }
-
-inline bool getEnvEnablePDL() {
-  static std::once_flag flag;
-  static bool enablePDL = false;
-  std::call_once(flag, [&]() {
-    if (getSMVersion() >= 90) {
-      const char* env = std::getenv("TRTLLM_ENABLE_PDL");
-      enablePDL = env && env[0] == '1' && env[1] == '\0';
-    }
-  });
-  return enablePDL;
-}

csrc/moe/dsv4_norm_router_gemm.h

@@ -1,30 +0,0 @@
-/*
- * Fused RMSNorm + router GEMV for DeepSeek V4.
- *
- * Computes in a single kernel:
- *   normed_x[m,k]   = x[m,k] * rsqrt(mean(x[m]^2) + eps) * norm_weight[k]
- *   router_logits[m,n] = sum_k(normed_x[m,k] * gate_weight[n,k])
- *
- * The GEMV body mirrors the algorithm in csrc/moe/dsv3_router_gemm_*.cu
- * (warp butterfly + smem cross-warp reduction, fp32 accumulation, PDL on
- * SM90+).  Blocks 0..kNumTokens-1 each materialize one token's normed_x
- * row to global memory using the algebraic identity
- *      logits[m,n] = rsqrt[m] * sum_k(x[m,k] * nw[k] * gw[n,k])
- * which lets every block produce its column of logits before normed_x
- * exists in gmem.
- *
- * Logits output is fp32 only — DeepSeek V4 router gate is hard-coded to
- * fp32 (vllm/model_executor/models/deepseek_v4.py:749).
- */
-
-#pragma once
-
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-
-#include "dsv3_router_gemm_utils.h"
-
-template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
-void invokeNormRouterGemm(float* logits, __nv_bfloat16* normed_x, T const* x,
-                          T const* norm_weight, T const* gate_weight, float eps,
-                          cudaStream_t stream);

csrc/moe/dsv4_norm_router_gemm_entry.cu

@@ -1,130 +0,0 @@
-/*
- * TORCH op entry for the fused RMSNorm + router GEMV kernel
- * (DeepSeek V4 Pro).  This op is DSV4-Pro-specific: the kernel is
- * instantiated only for ``num_experts == 384`` and ``hidden_dim ==
- * 7168``.  Other configurations (e.g. DSV4-Flash with H=4096) must
- * fall back to the unfused ``rms_norm`` + ``dsv3_router_gemm`` path.
- */
-
-#include <ATen/ATen.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <torch/all.h>
-
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-
-#include "core/registration.h"
-#include "dsv4_norm_router_gemm.h"
-
-namespace {
-
-// DSV4-Pro hard-coded shape constants.  Renamed from the earlier
-// ``kKimiK2NumExperts`` to avoid the misleading impression that this
-// kernel targets Kimi K2 — 384 happens to match Kimi K2's gate but the
-// intent here is DSV4-Pro.
-constexpr int kDsv4NumExperts = 384;
-constexpr int kDsv4HiddenDim = 7168;
-
-template <int kBegin, int kEnd>
-struct LoopUnroller {
-  static void unroll(int num_tokens, float* logits, __nv_bfloat16* normed_x,
-                     __nv_bfloat16 const* x, __nv_bfloat16 const* norm_weight,
-                     __nv_bfloat16 const* gate_weight, float eps,
-                     cudaStream_t stream) {
-    if (num_tokens == kBegin) {
-      invokeNormRouterGemm<__nv_bfloat16, kBegin, kDsv4NumExperts,
-                           kDsv4HiddenDim>(logits, normed_x, x, norm_weight,
-                                           gate_weight, eps, stream);
-    } else {
-      LoopUnroller<kBegin + 1, kEnd>::unroll(num_tokens, logits, normed_x, x,
-                                             norm_weight, gate_weight, eps,
-                                             stream);
-    }
-  }
-};
-
-template <int kEnd>
-struct LoopUnroller<kEnd, kEnd> {
-  static void unroll(int num_tokens, float* logits, __nv_bfloat16* normed_x,
-                     __nv_bfloat16 const* x, __nv_bfloat16 const* norm_weight,
-                     __nv_bfloat16 const* gate_weight, float eps,
-                     cudaStream_t stream) {
-    if (num_tokens == kEnd) {
-      invokeNormRouterGemm<__nv_bfloat16, kEnd, kDsv4NumExperts,
-                           kDsv4HiddenDim>(logits, normed_x, x, norm_weight,
-                                           gate_weight, eps, stream);
-    } else {
-      throw std::invalid_argument(
-          "Invalid num_tokens, only supports 1 to 16 for "
-          "dsv4_norm_router_gemm");
-    }
-  }
-};
-
-}  // namespace
-
-void dsv4_norm_router_gemm(at::Tensor& logits,    // [num_tokens, E] fp32
-                           at::Tensor& normed_x,  // [num_tokens, H] bf16
-                           at::Tensor const& x,   // [num_tokens, H] bf16
-                           at::Tensor const& norm_weight,  // [H] bf16
-                           at::Tensor const& gate_weight,  // [E, H] bf16
-                           double eps) {
-  TORCH_CHECK(x.dim() == 2 && norm_weight.dim() == 1 && gate_weight.dim() == 2,
-              "x must be 2D, norm_weight 1D, gate_weight 2D");
-  TORCH_CHECK(logits.dim() == 2 && normed_x.dim() == 2,
-              "logits and normed_x must be 2D");
-
-  int const num_tokens = x.size(0);
-  int const hidden_dim = x.size(1);
-  int const num_experts = gate_weight.size(0);
-
-  TORCH_CHECK(hidden_dim == kDsv4HiddenDim,
-              "Expected hidden_dim=", kDsv4HiddenDim,
-              " (DSV4-Pro), but got hidden_dim=", hidden_dim);
-  TORCH_CHECK(gate_weight.size(1) == hidden_dim,
-              "gate_weight.shape[1] must equal x.shape[1]");
-  TORCH_CHECK(norm_weight.size(0) == hidden_dim,
-              "norm_weight.shape[0] must equal x.shape[1]");
-  TORCH_CHECK(num_experts == kDsv4NumExperts,
-              "Expected num_experts=", kDsv4NumExperts,
-              " (DSV4-Pro), but got num_experts=", num_experts);
-  TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
-              "num_tokens must be in [1, 16] for dsv4_norm_router_gemm");
-
-  TORCH_CHECK(x.dtype() == at::kBFloat16, "x must be bf16");
-  TORCH_CHECK(norm_weight.dtype() == at::kBFloat16, "norm_weight must be bf16");
-  TORCH_CHECK(gate_weight.dtype() == at::kBFloat16, "gate_weight must be bf16");
-  TORCH_CHECK(normed_x.dtype() == at::kBFloat16, "normed_x must be bf16");
-  TORCH_CHECK(logits.dtype() == at::kFloat,
-              "logits must be float32 (DSV4 router output is hard-coded fp32)");
-
-  TORCH_CHECK(normed_x.size(0) == num_tokens && normed_x.size(1) == hidden_dim,
-              "normed_x must be [num_tokens, hidden_dim]");
-  TORCH_CHECK(logits.size(0) == num_tokens && logits.size(1) == num_experts,
-              "logits must be [num_tokens, num_experts]");
-
-  TORCH_CHECK(x.is_contiguous() && norm_weight.is_contiguous() &&
-                  gate_weight.is_contiguous() && normed_x.is_contiguous() &&
-                  logits.is_contiguous(),
-              "all tensors must be contiguous");
-
-  auto const sm = getSMVersion();
-  TORCH_CHECK(sm >= 90 && sm <= 103,
-              "dsv4_norm_router_gemm requires SM_90 <= CUDA ARCH <= SM_103");
-
-  cudaStream_t const stream = at::cuda::getCurrentCUDAStream();
-
-  auto* logits_ptr = reinterpret_cast<float*>(logits.mutable_data_ptr());
-  auto* nx_ptr = reinterpret_cast<__nv_bfloat16*>(normed_x.mutable_data_ptr());
-  auto* x_ptr = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr());
-  auto* nw_ptr = reinterpret_cast<__nv_bfloat16 const*>(norm_weight.data_ptr());
-  auto* gw_ptr = reinterpret_cast<__nv_bfloat16 const*>(gate_weight.data_ptr());
-  float const eps_f = static_cast<float>(eps);
-
-  LoopUnroller<1, 16>::unroll(num_tokens, logits_ptr, nx_ptr, x_ptr, nw_ptr,
-                              gw_ptr, eps_f, stream);
-}
-
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-  m.impl("dsv4_norm_router_gemm", &dsv4_norm_router_gemm);
-}

csrc/moe/dsv4_norm_router_gemm_kernel.cu

@@ -1,249 +0,0 @@
-/*
- * Fused RMSNorm + router GEMV for DeepSeek V4 (logits are fp32; bf16
- * output is unsupported because DSV4 hard-codes fp32 logits).  See
- * dsv4_norm_router_gemm.h for the math.
- *
- * The GEMV body mirrors csrc/moe/dsv3_router_gemm_float_out.cu (warp
- * butterfly reduction + smem cross-warp reduction, fp32 accumulation,
- * 128-thread block, PDL on SM90+).  RMSNorm is folded into the same
- * pass via the identity
- *   logits[m,n] = rsqrt[m] * sum_k(x[m,k] * nw[k] * gw[n,k])
- * so x is read exactly once per block during the GEMV phase.  Blocks
- * 0..kNumTokens-1 each materialize one row of normed_x for downstream
- * experts / shared_experts to consume.
- */
-
-#include <ATen/ATen.h>
-#include <ATen/cuda/CUDAContext.h>
-
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-
-#include "dsv4_norm_router_gemm.h"
-
-namespace {
-
-// Convert 8 bf16 values packed in uint4 into 8 floats. Mirrors the helper
-// in dsv3_router_gemm_float_out.cu (kept local so the dsv3 file stays
-// untouched).
-template <int VPT>
-__device__ __forceinline__ void bf16_uint4_to_float8(uint4 const& vec,
-                                                     float* dst) {
-  __nv_bfloat16* bf16_ptr =
-      reinterpret_cast<__nv_bfloat16*>(const_cast<uint4*>(&vec));
-#pragma unroll
-  for (int i = 0; i < VPT; i++) {
-    dst[i] = __bfloat162float(bf16_ptr[i]);
-  }
-}
-
-template <typename T, int kBlockSize, int VPT, int kNumTokens, int kNumExperts,
-          int kHiddenDim>
-__global__ __launch_bounds__(128, 1) void norm_router_gemm_kernel(
-    float* __restrict__ logits, __nv_bfloat16* __restrict__ normed_x,
-    T const* __restrict__ x, T const* __restrict__ norm_weight,
-    T const* __restrict__ gate_weight, float eps) {
-  static_assert(kBlockSize == 128, "kernel assumes blockDim.x == 128");
-  static_assert(kHiddenDim % (VPT * kBlockSize) == 0,
-                "kHiddenDim must be a multiple of VPT * kBlockSize");
-
-  int const n_idx = blockIdx.x;
-  int const tid = threadIdx.x;
-  constexpr int kWarpSize = 32;
-  constexpr int kNumWarps = kBlockSize / kWarpSize;
-  constexpr int k_elems_per_iter = VPT * kBlockSize;
-  constexpr int k_iterations = kHiddenDim / k_elems_per_iter;
-
-  T const* gw_col = gate_weight + n_idx * kHiddenDim;
-
-  // Per-thread accumulators — fp32 throughout, matching dsv3 / layernorm.
-  float partial[kNumTokens] = {};
-  float ss[kNumTokens] = {};
-
-  // Cross-warp reduction scratch.
-  __shared__ float sm_partial[kNumTokens][kNumWarps];
-  __shared__ float sm_ss[kNumTokens][kNumWarps];
-  __shared__ float s_rsqrt[kNumTokens];
-
-  int k_bases[k_iterations];
-#pragma unroll
-  for (int ki = 0; ki < k_iterations; ki++) {
-    k_bases[ki] = ki * k_elems_per_iter + tid * VPT;
-  }
-
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.wait;");
-#endif
-
-  // ---- Phase 1: single pass over x, accumulate partial GEMV and ss. ----
-#pragma unroll
-  for (int ki = 0; ki < k_iterations; ki++) {
-    int const k_base = k_bases[ki];
-
-    uint4 nw_vec = *reinterpret_cast<uint4 const*>(norm_weight + k_base);
-    float nw_f[VPT];
-    bf16_uint4_to_float8<VPT>(nw_vec, nw_f);
-
-    uint4 b_vec = *reinterpret_cast<uint4 const*>(gw_col + k_base);
-    float b_f[VPT];
-    bf16_uint4_to_float8<VPT>(b_vec, b_f);
-
-#pragma unroll
-    for (int m = 0; m < kNumTokens; m++) {
-      uint4 a_vec =
-          *reinterpret_cast<uint4 const*>(x + m * kHiddenDim + k_base);
-      float a_f[VPT];
-      bf16_uint4_to_float8<VPT>(a_vec, a_f);
-
-#pragma unroll
-      for (int k = 0; k < VPT; k++) {
-        float a = a_f[k];
-        ss[m] += a * a;
-        partial[m] += a * nw_f[k] * b_f[k];
-      }
-    }
-  }
-
-  // ---- Phase 2: warp butterfly reduction for both ss[] and partial[]. ----
-  int const warpId = tid / kWarpSize;
-  int const laneId = tid % kWarpSize;
-
-#pragma unroll
-  for (int m = 0; m < kNumTokens; m++) {
-    float p = partial[m];
-    float s = ss[m];
-
-    p += __shfl_xor_sync(0xffffffff, p, 16);
-    s += __shfl_xor_sync(0xffffffff, s, 16);
-    p += __shfl_xor_sync(0xffffffff, p, 8);
-    s += __shfl_xor_sync(0xffffffff, s, 8);
-    p += __shfl_xor_sync(0xffffffff, p, 4);
-    s += __shfl_xor_sync(0xffffffff, s, 4);
-    p += __shfl_xor_sync(0xffffffff, p, 2);
-    s += __shfl_xor_sync(0xffffffff, s, 2);
-    p += __shfl_xor_sync(0xffffffff, p, 1);
-    s += __shfl_xor_sync(0xffffffff, s, 1);
-
-    if (laneId == 0) {
-      sm_partial[m][warpId] = p;
-      sm_ss[m][warpId] = s;
-    }
-  }
-
-  __syncthreads();
-
-  // ---- Phase 3: tid 0 finalises the reduction, writes logits, stashes
-  //               rsqrt[m] in smem for phase 4. ----
-  if (tid == 0) {
-#pragma unroll
-    for (int m = 0; m < kNumTokens; m++) {
-      float p_sum = 0.0f;
-      float s_sum = 0.0f;
-#pragma unroll
-      for (int w = 0; w < kNumWarps; w++) {
-        p_sum += sm_partial[m][w];
-        s_sum += sm_ss[m][w];
-      }
-      // Order matches layernorm_kernels.cu: rsqrtf(variance / H + eps).
-      // Use division (not multiply-by-reciprocal) to avoid an extra ULP
-      // mismatch with the reference RMSNorm.
-      float rs = rsqrtf(s_sum / static_cast<float>(kHiddenDim) + eps);
-      s_rsqrt[m] = rs;
-      logits[m * kNumExperts + n_idx] = p_sum * rs;
-    }
-  }
-
-  __syncthreads();
-
-  // ---- Phase 4: spread normed_x writes across blocks 0..kNumTokens-1.
-  //              Each writer block handles exactly one token row,
-  //              avoiding the long tail of block 0 doing all M rows.
-  //              Every block has every token's rsqrt[] in s_rsqrt
-  //              already (computed independently in phase 3), so no
-  //              cross-block synchronization is required. ----
-  if (n_idx < kNumTokens) {
-    int const m_writer = n_idx;
-    float const rs = s_rsqrt[m_writer];
-    __nv_bfloat16 const* x_row = x + m_writer * kHiddenDim;
-    __nv_bfloat16* normed_row = normed_x + m_writer * kHiddenDim;
-
-#pragma unroll
-    for (int ki = 0; ki < k_iterations; ki++) {
-      int const k_base = k_bases[ki];
-
-      uint4 nw_vec = *reinterpret_cast<uint4 const*>(norm_weight + k_base);
-      float nw_f[VPT];
-      bf16_uint4_to_float8<VPT>(nw_vec, nw_f);
-
-      uint4 a_vec = *reinterpret_cast<uint4 const*>(x_row + k_base);
-      float a_f[VPT];
-      bf16_uint4_to_float8<VPT>(a_vec, a_f);
-
-      uint4 normed_vec;
-      __nv_bfloat16* np = reinterpret_cast<__nv_bfloat16*>(&normed_vec);
-#pragma unroll
-      for (int k = 0; k < VPT; k++) {
-        np[k] = __float2bfloat16(a_f[k] * rs * nw_f[k]);
-      }
-      *reinterpret_cast<uint4*>(normed_row + k_base) = normed_vec;
-    }
-  }
-
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.launch_dependents;");
-#endif
-}
-
-}  // namespace
-
-template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
-void invokeNormRouterGemm(float* logits, __nv_bfloat16* normed_x, T const* x,
-                          T const* norm_weight, T const* gate_weight, float eps,
-                          cudaStream_t stream) {
-  constexpr int VPT = 16 / sizeof(T);
-  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,
-                     norm_router_gemm_kernel<T, kBlockSize, VPT, kNumTokens,
-                                             kNumExperts, kHiddenDim>,
-                     logits, normed_x, x, norm_weight, gate_weight, eps);
-}
-
-// Template instantiations — DSV4-Pro is the only supported configuration:
-// num_experts=384, hidden_dim=7168.  Other shapes (e.g. DSV4-Flash with
-// hidden_dim=4096) fall back to the unfused path on the Python side.
-#define INSTANTIATE(M)                                                    \
-  template void invokeNormRouterGemm<__nv_bfloat16, M, 384, 7168>(        \
-      float*, __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, \
-      __nv_bfloat16 const*, float, cudaStream_t);
-
-INSTANTIATE(1)
-INSTANTIATE(2)
-INSTANTIATE(3)
-INSTANTIATE(4)
-INSTANTIATE(5)
-INSTANTIATE(6)
-INSTANTIATE(7)
-INSTANTIATE(8)
-INSTANTIATE(9)
-INSTANTIATE(10)
-INSTANTIATE(11)
-INSTANTIATE(12)
-INSTANTIATE(13)
-INSTANTIATE(14)
-INSTANTIATE(15)
-INSTANTIATE(16)
-
-#undef INSTANTIATE

csrc/moe/moe_ops.h

@@ -75,12 +75,4 @@ void shuffle_rows(const torch::Tensor& input_tensor,
 // Supports num_tokens in [1, 16], num_experts in {256, 384}, hidden_dim = 7168
 void dsv3_router_gemm(torch::Tensor& output, const torch::Tensor& mat_a,
                       const torch::Tensor& mat_b);
-
-// Fused RMSNorm + router GEMV for DeepSeek V4. Produces both:
-//   normed_x[m,k]      = x[m,k] * rsqrt(mean(x[m]^2) + eps) * norm_weight[k]
-//   logits[m,n]        = sum_k(normed_x[m,k] * gate_weight[n,k])
-// in a single kernel launch. Same dim/dtype constraints as dsv3_router_gemm.
-void dsv4_norm_router_gemm(at::Tensor& logits, at::Tensor& normed_x,
-                           at::Tensor const& x, at::Tensor const& norm_weight,
-                           at::Tensor const& gate_weight, double eps);
 #endif

csrc/moe/torch_bindings.cpp

@@ -125,12 +125,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
   // DeepSeek V3 optimized router GEMM for SM90+
   m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
   // conditionally compiled so impl registration is in source file
-
-  // DeepSeek V4 fused RMSNorm + router GEMV for SM90+
-  m.def(
-      "dsv4_norm_router_gemm(Tensor! logits, Tensor! normed_x, Tensor x, "
-      "Tensor norm_weight, Tensor gate_weight, float eps) -> ()");
-  // conditionally compiled so impl registration is in source file
 #endif
 }
 

csrc/ops.h

@@ -61,19 +61,15 @@ void merge_attn_states(
     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.
 void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
               double epsilon);
 
 void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
                         torch::Tensor& weight, double epsilon);
 
-void fused_qk_norm_rope(torch::Tensor& qkv, int64_t num_heads_q,
-                        int64_t num_heads_k, int64_t num_heads_v,
-                        int64_t head_dim, double eps, torch::Tensor& q_weight,
-                        torch::Tensor& k_weight, torch::Tensor& cos_sin_cache,
-                        bool is_neox, torch::Tensor& position_ids,
-                        int64_t forced_token_heads_per_warp);
-
 void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
     torch::Tensor& q, torch::Tensor const& kv, torch::Tensor& k_cache,
     torch::Tensor const& slot_mapping, torch::Tensor const& position_ids,
@@ -99,37 +95,15 @@ 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);
 
-void rms_norm_static_fp8_quant(torch::Tensor& out, torch::Tensor& input,
-                               torch::Tensor& weight, torch::Tensor& scale,
-                               double epsilon);
-
-void fused_add_rms_norm_static_fp8_quant(torch::Tensor& out,
-                                         torch::Tensor& input,
-                                         torch::Tensor& residual,
-                                         torch::Tensor& weight,
-                                         torch::Tensor& scale, double epsilon);
-
-void rms_norm_dynamic_per_token_quant(torch::Tensor& out,
-                                      torch::Tensor const& input,
-                                      torch::Tensor const& weight,
-                                      torch::Tensor& scales,
-                                      double const epsilon,
-                                      std::optional<torch::Tensor> scale_ub,
-                                      std::optional<torch::Tensor> residual);
-
-void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
-                              torch::Tensor const& weight,
-                              torch::Tensor& scales, double const epsilon,
-                              std::optional<torch::Tensor> scale_ub,
-                              std::optional<torch::Tensor> residual,
-                              int64_t group_size, bool is_scale_transposed);
-
 void silu_and_mul_per_block_quant(torch::Tensor& out,
                                   torch::Tensor const& input,
                                   torch::Tensor& scales, int64_t group_size,
                                   std::optional<torch::Tensor> scale_ub,
                                   bool is_scale_transposed);
 
+// rotary_embedding also exist in csrc/libtorch_stable/ops.h (torch::stable
+// ABI for CUDA). It remains here because the CPU build still uses these
+// torch::Tensor declarations.
 void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
                       std::optional<torch::Tensor> key, int64_t head_size,
                       torch::Tensor& cos_sin_cache, bool is_neox,
@@ -149,17 +123,10 @@ void persistent_masked_m_silu_mul_quant(
     at::Tensor& y_s,           // (E, T, H//group_size) [OUT]
     bool use_ue8m0);
 
-void mul_and_silu(torch::Tensor& out, torch::Tensor& input);
-
 void gelu_and_mul(torch::Tensor& out, torch::Tensor& input);
 
 void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input);
 
-void fatrelu_and_mul(torch::Tensor& out, torch::Tensor& input,
-                     double threshold);
-void swigluoai_and_mul(torch::Tensor& out, torch::Tensor& input,
-                       double alpha = 1.702, double limit = 7.0);
-
 void gelu_new(torch::Tensor& out, torch::Tensor& input);
 
 void gelu_fast(torch::Tensor& out, torch::Tensor& input);
@@ -174,28 +141,6 @@ void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,
 
 torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
 
-torch::Tensor ggml_dequantize(torch::Tensor W, int64_t type, int64_t m,
-                              int64_t n,
-                              std::optional<at::ScalarType> const& dtype);
-
-torch::Tensor ggml_mul_mat_vec_a8(torch::Tensor W, torch::Tensor X,
-                                  int64_t type, int64_t row);
-
-torch::Tensor ggml_mul_mat_a8(torch::Tensor W, torch::Tensor X, int64_t type,
-                              int64_t row);
-
-torch::Tensor ggml_moe_a8(torch::Tensor X, torch::Tensor W,
-                          torch::Tensor sorted_token_ids,
-                          torch::Tensor expert_ids,
-                          torch::Tensor num_tokens_post_padded, int64_t type,
-                          int64_t row, int64_t top_k, int64_t tokens);
-
-torch::Tensor ggml_moe_a8_vec(torch::Tensor X, torch::Tensor W,
-                              torch::Tensor topk_ids, int64_t top_k,
-                              int64_t type, int64_t row, int64_t tokens);
-
-int64_t ggml_moe_get_block_size(int64_t type);
-
 void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
                               torch::Tensor const& scale,
                               std::optional<torch::Tensor> const& azp);
@@ -204,24 +149,6 @@ void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
                                torch::Tensor& scales,
                                std::optional<torch::Tensor> const& azp);
 
-torch::Tensor gptq_gemm(torch::Tensor a, torch::Tensor b_q_weight,
-                        torch::Tensor b_gptq_qzeros,
-                        torch::Tensor b_gptq_scales, torch::Tensor b_g_idx,
-                        bool use_exllama, bool use_v2_format, int64_t bit);
-
-void gptq_shuffle(torch::Tensor q_weight, torch::Tensor q_perm, int64_t bit);
-
-void static_scaled_fp8_quant(
-    torch::Tensor& out, torch::Tensor const& input, torch::Tensor const& scale,
-    std::optional<std::tuple<int64_t, int64_t>> group_shape = std::nullopt);
-
-void dynamic_scaled_fp8_quant(torch::Tensor& out, torch::Tensor const& input,
-                              torch::Tensor& scale);
-
-void dynamic_per_token_scaled_fp8_quant(
-    torch::Tensor& out, torch::Tensor const& input, torch::Tensor& scale,
-    std::optional<torch::Tensor> const& scale_ub);
-
 void selective_scan_fwd(
     const torch::Tensor& u, const torch::Tensor& delta, const torch::Tensor& A,
     const torch::Tensor& B, const torch::Tensor& C,

csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu

@@ -5,8 +5,7 @@
 #include <c10/cuda/CUDAGuard.h>
 
 #include "../../dispatch_utils.h"
-#include "quant_conversions.cuh"
-#include "../w8a8/fp8/common.cuh"
+#include "libtorch_stable/quantization/fused_kernels/quant_conversions.cuh"
 
 namespace vllm {
 

csrc/quantization/utils.cuh

@@ -7,23 +7,23 @@
  */
 
 #include <cmath>
-#include <torch/types.h>
+#include <torch/headeronly/macros/Macros.h>
 
 #ifndef USE_ROCM
-  #include <c10/util/Float8_e4m3fn.h>
+  #include <torch/headeronly/util/Float8_e4m3fn.h>
   #define MAYBE_HOST_DEVICE C10_HOST_DEVICE
 #else
-  #include <ATen/hip/HIPContext.h>
-  #include <c10/util/Float8_e4m3fn.h>
-  #include <c10/util/Float8_e4m3fnuz.h>
+  #include <torch/headeronly/util/Float8_e4m3fn.h>
+  #include <torch/headeronly/util/Float8_e4m3fnuz.h>
   // ROCm doesn't seem to need C10_HOST_DEVICE for static constexpr
   #define MAYBE_HOST_DEVICE
 #endif
 
 template <typename T,
-          typename = std::enable_if_t<std::is_same_v<T, c10::Float8_e4m3fn> ||
-                                      std::is_same_v<T, c10::Float8_e4m3fnuz> ||
-                                      std::is_same_v<T, int8_t>>>
+          typename = std::enable_if_t<
+              std::is_same_v<T, torch::headeronly::Float8_e4m3fn> ||
+              std::is_same_v<T, torch::headeronly::Float8_e4m3fnuz> ||
+              std::is_same_v<T, int8_t>>>
 struct quant_type_max {
   static constexpr T val() { return std::numeric_limits<T>::max(); }
 };
@@ -31,9 +31,10 @@ struct quant_type_max {
 // Using the default max value from pytorch (240.0 0x7F) will cause accuracy
 // issues when running dynamic quantization. Here use 224.0 0x7E for rocm.
 template <>
-struct quant_type_max<c10::Float8_e4m3fnuz> {
-  static constexpr c10::Float8_e4m3fnuz val() {
-    return c10::Float8_e4m3fnuz(0x7E, c10::Float8_e4m3fnuz::from_bits());
+struct quant_type_max<torch::headeronly::Float8_e4m3fnuz> {
+  static constexpr torch::headeronly::Float8_e4m3fnuz val() {
+    return torch::headeronly::Float8_e4m3fnuz(
+        0x7E, torch::headeronly::Float8_e4m3fnuz::from_bits());
   }
 };
 
@@ -42,9 +43,10 @@ MAYBE_HOST_DEVICE static constexpr T quant_type_max_v =
     quant_type_max<T>::val();
 
 template <typename T,
-          typename = std::enable_if_t<std::is_same_v<T, c10::Float8_e4m3fn> ||
-                                      std::is_same_v<T, c10::Float8_e4m3fnuz> ||
-                                      std::is_same_v<T, int8_t>>>
+          typename = std::enable_if_t<
+              std::is_same_v<T, torch::headeronly::Float8_e4m3fn> ||
+              std::is_same_v<T, torch::headeronly::Float8_e4m3fnuz> ||
+              std::is_same_v<T, int8_t>>>
 struct min_scaling_factor {
   C10_DEVICE C10_ALWAYS_INLINE static float val() {
     return 1.0f / (quant_type_max_v<T> * 512.0f);

csrc/quantization/w8a8/fp8/common.cuh

@@ -5,6 +5,19 @@
 
 #include <cmath>
 
+// This header is shared between _C and _C_stable_libtorch targets.
+// torch_utils.h provides get_device_prop(). We need to pass USE_CUDA
+// to the .so to expose some of the shims used by torch_utils.h. For now
+// this is only done for _C_stable_libtorch and not for _C, so we use the
+// non stable at::cuda::getCurrentDeviceProperties for _C for now.
+#ifdef TORCH_TARGET_VERSION
+  #include "../../../libtorch_stable/torch_utils.h"
+#else
+  #ifdef USE_ROCM
+    #include <ATen/hip/HIPContext.h>
+  #endif
+#endif
+
 #ifndef USE_ROCM
   #include "nvidia/quant_utils.cuh"
 #else
@@ -18,7 +31,11 @@ static bool is_fp8_ocp() {
 #ifndef USE_ROCM
   return true;
 #else
-  auto dprops = at::cuda::getCurrentDeviceProperties();
+  #ifdef TORCH_TARGET_VERSION
+  auto* dprops = get_device_prop();
+  #else
+  auto* dprops = at::cuda::getCurrentDeviceProperties();
+  #endif
   std::string device_arch = dprops->gcnArchName;
   size_t substring = device_arch.find("gfx94");
   return substring == std::string::npos;

csrc/torch_bindings.cpp

@@ -77,17 +77,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "    Tensor? output_scale=None) -> ()");
   ops.impl("merge_attn_states", torch::kCUDA, &merge_attn_states);
 
-  // Activation ops
-  // Activation function used in SwiGLU.
-  ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
-  ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);
-
-  // SwiGLU activation with input clamping.
-  ops.def(
-      "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
-      "-> ()");
-  ops.impl("silu_and_mul_with_clamp", torch::kCUDA, &silu_and_mul_clamp);
-
+  // Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
   ops.def(
       "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
   ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
@@ -104,61 +94,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.impl("silu_and_mul_per_block_quant", torch::kCUDA,
            &silu_and_mul_per_block_quant);
 
-  ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
-  ops.impl("mul_and_silu", torch::kCUDA, &mul_and_silu);
-
-  // Activation function used in GeGLU with `none` approximation.
-  ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
-  ops.impl("gelu_and_mul", torch::kCUDA, &gelu_and_mul);
-
-  // Activation function used in GeGLU with `tanh` approximation.
-  ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()");
-  ops.impl("gelu_tanh_and_mul", torch::kCUDA, &gelu_tanh_and_mul);
-
-  // FATReLU implementation.
-  ops.def("fatrelu_and_mul(Tensor! out, Tensor input, float threshold) -> ()");
-  ops.impl("fatrelu_and_mul", torch::kCUDA, &fatrelu_and_mul);
-
-  ops.def(
-      "swigluoai_and_mul(Tensor! out, Tensor input, float alpha=1.702, float "
-      "limit=7.0) "
-      "-> ()");
-  ops.impl("swigluoai_and_mul", torch::kCUDA, &swigluoai_and_mul);
-
-  // GELU implementation used in GPT-2.
-  ops.def("gelu_new(Tensor! out, Tensor input) -> ()");
-  ops.impl("gelu_new", torch::kCUDA, &gelu_new);
-
-  // Approximate GELU implementation.
-  ops.def("gelu_fast(Tensor! out, Tensor input) -> ()");
-  ops.impl("gelu_fast", torch::kCUDA, &gelu_fast);
-
-  // Quick GELU implementation.
-  ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
-  ops.impl("gelu_quick", torch::kCUDA, &gelu_quick);
-
-  // Layernorm
-  // Apply Root Mean Square (RMS) Normalization to the input tensor.
-  ops.def(
-      "rms_norm(Tensor! result, Tensor input, Tensor weight, float epsilon) -> "
-      "()");
-  ops.impl("rms_norm", torch::kCUDA, &rms_norm);
-
-  // In-place fused Add and RMS Normalization.
-  ops.def(
-      "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor weight, "
-      "float epsilon) -> ()");
-  ops.impl("fused_add_rms_norm", torch::kCUDA, &fused_add_rms_norm);
-
-  // Function for fused QK Norm and RoPE
-  ops.def(
-      "fused_qk_norm_rope(Tensor! qkv, int num_heads_q, "
-      "int num_heads_k, int num_heads_v, int head_dim, float eps, "
-      "Tensor q_weight, Tensor k_weight, Tensor cos_sin_cache, "
-      "bool is_neox, Tensor position_ids, "
-      "int forced_token_heads_per_warp=-1) -> ()");
-  ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);
-
   // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
   // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
   // kernel launch.
@@ -195,48 +130,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "Tensor workspace, int k, int max_seq_len) -> ()");
   ops.impl("persistent_topk", torch::kCUDA, &persistent_topk);
 
-  // Layernorm-quant
-  // Apply Root Mean Square (RMS) Normalization to the input tensor.
-  ops.def(
-      "rms_norm_static_fp8_quant(Tensor! result, Tensor input, Tensor weight, "
-      "Tensor scale, float epsilon) -> "
-      "()");
-  ops.impl("rms_norm_static_fp8_quant", torch::kCUDA,
-           &rms_norm_static_fp8_quant);
-
-  // In-place fused Add and RMS Normalization.
-  ops.def(
-      "fused_add_rms_norm_static_fp8_quant(Tensor! result, Tensor input, "
-      "Tensor! residual, Tensor weight, "
-      "Tensor scale, float epsilon) -> ()");
-  ops.impl("fused_add_rms_norm_static_fp8_quant", torch::kCUDA,
-           &fused_add_rms_norm_static_fp8_quant);
-
-  // Fused Layernorm + Quant kernels
-  ops.def(
-      "rms_norm_dynamic_per_token_quant(Tensor! result, Tensor input, "
-      "Tensor weight, Tensor! scale, float epsilon, "
-      "Tensor? scale_ub, Tensor!? residual) -> ()");
-  ops.impl("rms_norm_dynamic_per_token_quant", torch::kCUDA,
-           &rms_norm_dynamic_per_token_quant);
-
-  // Fused Layernorm + Block quant kernels
-  ops.def(
-      "rms_norm_per_block_quant(Tensor! result, Tensor input, "
-      "Tensor weight, Tensor! scale, float epsilon, "
-      "Tensor? scale_ub, Tensor!? residual, int group_size, "
-      "bool is_scale_transposed) -> ()");
-  ops.impl("rms_norm_per_block_quant", torch::kCUDA, &rms_norm_per_block_quant);
-
-  // Rotary embedding
-  // Apply GPT-NeoX or GPT-J style rotary embedding to query and key.
-  ops.def(
-      "rotary_embedding(Tensor positions, Tensor! query,"
-      "                 Tensor!? key, int head_size,"
-      "                 Tensor cos_sin_cache, bool is_neox, int "
-      "rope_dim_offset=0, bool inverse=False) -> ()");
-  ops.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);
-
   // Quantization ops
 #ifndef USE_ROCM
 
@@ -318,39 +211,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 
 #endif
 
-  // Dequantization for GGML.
-  ops.def(
-      "ggml_dequantize(Tensor W, int type, SymInt m, SymInt n, ScalarType? "
-      "dtype) -> Tensor");
-  ops.impl("ggml_dequantize", torch::kCUDA, &ggml_dequantize);
-
-  // mmvq kernel for GGML.
-  ops.def(
-      "ggml_mul_mat_vec_a8(Tensor W, Tensor X, int type, SymInt row) "
-      "-> Tensor");
-  ops.impl("ggml_mul_mat_vec_a8", torch::kCUDA, &ggml_mul_mat_vec_a8);
-
-  // mmq kernel for GGML.
-  ops.def(
-      "ggml_mul_mat_a8(Tensor W, Tensor X, int type, SymInt row) -> Tensor");
-  ops.impl("ggml_mul_mat_a8", torch::kCUDA, &ggml_mul_mat_a8);
-
-  // moe kernel for GGML.
-  ops.def(
-      "ggml_moe_a8(Tensor X, Tensor W, "
-      "Tensor sorted_token_ids, Tensor expert_ids, Tensor "
-      "num_tokens_post_padded, "
-      "int type, SymInt row, SymInt top_k, SymInt tokens) -> Tensor");
-  ops.impl("ggml_moe_a8", torch::kCUDA, &ggml_moe_a8);
-
-  ops.def(
-      "ggml_moe_a8_vec(Tensor X, Tensor W, "
-      "Tensor topk_ids, int top_k, "
-      "int type, SymInt row, SymInt tokens) -> Tensor");
-  ops.impl("ggml_moe_a8_vec", torch::kCUDA, &ggml_moe_a8_vec);
-
-  ops.def("ggml_moe_get_block_size", &ggml_moe_get_block_size);
-
 #ifndef USE_ROCM
   // Expert-specialization mxfp8 blockscaled grouped quantization (SM100+).
   ops.def(
@@ -370,57 +230,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 
 #endif
 
-  // Quantized GEMM for GPTQ.
-  // Note: even though the C++ inferred schema is correct for this op, it seems
-  // to prevent the meta function registry.
-  ops.def(
-      "gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, "
-      "Tensor b_gptq_scales, Tensor b_g_idx, bool use_exllama, bool "
-      "use_v2_format, int bit) "
-      "-> Tensor");
-  ops.impl("gptq_gemm", torch::kCUDA, &gptq_gemm);
-
-  // Post processing for GPTQ.
-  ops.def("gptq_shuffle(Tensor! q_weight, Tensor q_perm, int bit) -> ()");
-  ops.impl("gptq_shuffle", torch::kCUDA, &gptq_shuffle);
-
-  // Compute FP8 quantized tensor for given scaling factor.
-  // Supports per-tensor, per-channel, per-token, and arbitrary 2D group
-  // scaling. Optional group_m/group_n specify the group shape explicitly;
-  // required for 1D scales to disambiguate per-channel vs per-token.
-  ops.def(
-      "static_scaled_fp8_quant(Tensor! result, Tensor input, Tensor scale, "
-      "(int, int)? group_shape=None) -> ()");
-  ops.impl("static_scaled_fp8_quant", torch::kCUDA, &static_scaled_fp8_quant);
-
-  // Compute dynamic-per-tensor FP8 quantized tensor and scaling factor.
-  ops.def(
-      "dynamic_scaled_fp8_quant(Tensor! result, Tensor input, Tensor! scale) "
-      "-> "
-      "()");
-  ops.impl("dynamic_scaled_fp8_quant", torch::kCUDA, &dynamic_scaled_fp8_quant);
-
-  // Compute dynamic-per-token FP8 quantized tensor and scaling factor.
-  ops.def(
-      "dynamic_per_token_scaled_fp8_quant(Tensor! result, Tensor input, "
-      "Tensor! scale, Tensor? scale_ub) -> "
-      "()");
-  ops.impl("dynamic_per_token_scaled_fp8_quant", torch::kCUDA,
-           &dynamic_per_token_scaled_fp8_quant);
-
-  // Compute int8 quantized tensor for given scaling factor.
-  ops.def(
-      "static_scaled_int8_quant(Tensor! result, Tensor input, Tensor scale,"
-      "Tensor? azp) -> ()");
-  ops.impl("static_scaled_int8_quant", torch::kCUDA, &static_scaled_int8_quant);
-
-  // Compute int8 quantized tensor and scaling factor
-  ops.def(
-      "dynamic_scaled_int8_quant(Tensor! result, Tensor input, Tensor! scale, "
-      "Tensor!? azp) -> ()");
-  ops.impl("dynamic_scaled_int8_quant", torch::kCUDA,
-           &dynamic_scaled_int8_quant);
-
   // Mamba selective scan kernel
   ops.def(
       "selective_scan_fwd(Tensor! u, Tensor! delta,"

csrc/torch_utils.h

@@ -0,0 +1,17 @@
+#pragma once
+
+// Shared TORCH_UTILS_CHECK across both libtorch stable and unstable source
+// files. Keep this header free of CUTLASS/CUTE so attention/quant headers can
+// use it.
+//
+// If TORCH_TARGET_VERSION is defined, we are building _C_stable_libtorch.so so
+// use STD_TORCH_CHECK via header-only.
+// Otherwise, use TORCH_CHECK via torch/all.h.
+
+#ifdef TORCH_TARGET_VERSION
+  #include <torch/headeronly/util/Exception.h>
+  #define TORCH_UTILS_CHECK STD_TORCH_CHECK
+#else
+  #include <torch/all.h>
+  #define TORCH_UTILS_CHECK TORCH_CHECK
+#endif

csrc/type_convert.cuh

@@ -1,8 +1,10 @@
 #pragma once
 
-#include <torch/all.h>
+#include <torch/headeronly/util/BFloat16.h>
+#include <torch/headeronly/util/Half.h>
 
 #ifndef USE_ROCM
+  #include <cuda.h>
   #include <cuda_bf16.h>
   #include <cuda_fp16.h>
 #else
@@ -191,4 +193,4 @@ struct alignas(16) _f16Vec {
     return result;
   }
 };
-}  // namespace vllm
+}  // namespace vllm

docker/Dockerfile

@@ -105,6 +105,23 @@ ARG BUILD_OS
 
 ENV DEBIAN_FRONTEND=noninteractive
 
+# Environment for uv
+# Declared BEFORE the installer + `uv venv` invocations below so the uv
+# binary, managed Python, download cache, and /opt/venv all land under
+# /opt/uv instead of /root/.local/. Without this, the venv created at
+# build time hardlinks back to /root/.local/share/uv/python and
+# descendants of this stage (`build`, `dev`, `csrc-build`,
+# `extensions-build`) inherit a root-owned, non-root-unreadable layout.
+# See #15174, #15359, #31959. Child stages inherit these via Dockerfile
+# `ENV` unless they override them explicitly.
+ENV UV_HTTP_TIMEOUT=500
+ENV UV_INDEX_STRATEGY="unsafe-best-match"
+ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
+ENV UV_CACHE_DIR=/opt/uv/cache
+ENV UV_INSTALL_DIR=/opt/uv/bin
+ENV PATH="/opt/venv/bin:/opt/uv/bin:$PATH"
+ENV VIRTUAL_ENV="/opt/venv"
+
 # Install system dependencies including build tools.
 # The Ubuntu path uses apt + deadsnakes-via-uv for Python; the manylinux path
 # (AlmaLinux 8, e.g. pytorch/manylinux2_28-builder) uses dnf and the Python
@@ -145,15 +162,21 @@ RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
 
 # Install uv and bootstrap /opt/venv. Both paths converge on /opt/venv so all
 # downstream stages stay distro-agnostic.
-RUN curl -LsSf https://astral.sh/uv/install.sh | sh \
+RUN mkdir -p "${UV_PYTHON_INSTALL_DIR}" "${UV_CACHE_DIR}" "${UV_INSTALL_DIR}" \
+    && chmod -R a+rX /opt/uv \
+    && curl -LsSf https://astral.sh/uv/install.sh | sh \
+    # `--seed` installs pip/setuptools/wheel into the venv so `python3 -m
+    # pip` works regardless of how uv happens to link the venv back to the
+    # managed Python install (which, at a non-default UV_PYTHON_INSTALL_DIR,
+    # doesn't always expose ensurepip via the default venv layout).
     && if [ "${BUILD_OS}" = "manylinux" ]; then \
            # manylinux images ship Python at /opt/python/cpXY-cpXY/; point uv
            # at the matching interpreter rather than letting it fetch one.
            PYV_NODOT=$(echo ${PYTHON_VERSION} | tr -d '.') \
            && MANYLINUX_PY=/opt/python/cp${PYV_NODOT}-cp${PYV_NODOT}/bin/python${PYTHON_VERSION} \
-           && $HOME/.local/bin/uv venv /opt/venv --python "$MANYLINUX_PY"; \
+           && uv venv --seed /opt/venv --python "$MANYLINUX_PY"; \
        else \
-           $HOME/.local/bin/uv venv /opt/venv --python ${PYTHON_VERSION}; \
+           uv venv --seed /opt/venv --python ${PYTHON_VERSION}; \
        fi \
     && rm -f /usr/bin/python3 /usr/bin/python3-config /usr/bin/pip \
     && ln -sf /opt/venv/bin/python3 /usr/bin/python3 \
@@ -161,13 +184,10 @@ RUN curl -LsSf https://astral.sh/uv/install.sh | sh \
     && ln -sf /opt/venv/bin/pip /usr/bin/pip \
     && python3 --version && python3 -m pip --version
 
-# Activate virtual environment and add uv to PATH
-ENV PATH="/opt/venv/bin:/root/.local/bin:$PATH"
-ENV VIRTUAL_ENV="/opt/venv"
-
-# Environment for uv
-ENV UV_HTTP_TIMEOUT=500
-ENV UV_INDEX_STRATEGY="unsafe-best-match"
+# UV_LINK_MODE=copy applies to subsequent `uv pip install` RUNs (avoids
+# hardlink failures with BuildKit cache mounts); it must not be set during
+# `uv venv` above, which relies on hardlinking /opt/venv back to the
+# managed Python source so ensurepip / `python3 -m pip` still resolve.
 ENV UV_LINK_MODE=copy
 
 # Verify GCC version
@@ -198,7 +218,7 @@ COPY requirements/common.txt requirements/common.txt
 COPY requirements/cuda.txt requirements/cuda.txt
 COPY use_existing_torch.py use_existing_torch.py
 COPY pyproject.toml pyproject.toml
-RUN --mount=type=cache,target=/root/.cache/uv \
+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; \
     fi \
@@ -218,7 +238,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # Track PyTorch lib versions used during build and match in downstream instances.
 # We do this for both nightly and release so we can strip dependencies/*.txt as needed.
 # Otherwise library dependencies can upgrade/downgrade torch incorrectly.
-RUN --mount=type=cache,target=/root/.cache/uv \
+RUN --mount=type=cache,target=/opt/uv/cache \
     uv pip freeze | grep -i "^torch=\|^torchvision=\|^torchaudio=" > torch_lib_versions.txt \
     && TORCH_LIB_VERSIONS=$(cat torch_lib_versions.txt | xargs) \
     && echo "Installed torch libs: ${TORCH_LIB_VERSIONS}"
@@ -231,6 +251,56 @@ ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX'
 ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
 #################### BUILD BASE IMAGE ####################
 
+#################### RUST BUILD IMAGE ####################
+# Build the Rust frontend (`vllm-rs`) in a dedicated stage so the main wheel
+# build stage doesn't need the rust toolchain, protoc, or the rust source.
+# This stage runs in parallel with csrc-build/extensions-build.
+FROM ${BUILD_BASE_IMAGE} AS rust-build
+ARG BUILD_OS
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Install a basic C toolchain (some rust crates compile C in their build.rs
+# scripts) and unzip (used to extract the pinned protoc release below).
+RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
+        dnf install -y --setopt=install_weak_deps=False \
+            ca-certificates curl git gcc gcc-c++ make unzip \
+        && dnf clean all && rm -rf /var/cache/dnf; \
+    else \
+        apt-get update -y \
+        && apt-get install -y --no-install-recommends \
+            ca-certificates curl git build-essential unzip \
+        && rm -rf /var/lib/apt/lists/*; \
+    fi
+
+COPY tools/install_protoc.sh /tmp/install_protoc.sh
+RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
+
+# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
+RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
+    sh -s -- -y --profile minimal --default-toolchain none
+ENV PATH="/root/.cargo/bin:${PATH}"
+
+WORKDIR /workspace
+
+# Copy only the rust workspace — the binary is the sole artifact we need.
+COPY rust rust
+COPY rust-toolchain.toml rust-toolchain.toml
+COPY build_rust.sh build_rust.sh
+
+# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
+# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
+ENV CARGO_BUILD_JOBS=4
+
+# Build the release binary. Cache cargo registry/git and target/, but copy the
+# binary out of the target/ cache mount 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 \
+    --mount=type=cache,target=/workspace/rust/target \
+    VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
+#################### RUST BUILD IMAGE ####################
+
 #################### CSRC BUILD IMAGE ####################
 FROM base AS csrc-build
 ARG TARGETPLATFORM
@@ -254,7 +324,7 @@ ENV UV_INDEX_STRATEGY="unsafe-best-match"
 # Use copy mode to avoid hardlink failures with Docker cache mounts
 ENV UV_LINK_MODE=copy
 
-RUN --mount=type=cache,target=/root/.cache/uv \
+RUN --mount=type=cache,target=/opt/uv/cache \
     if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
         echo "Installing build requirements without torch..." \
         && python3 use_existing_torch.py --prefix \
@@ -299,7 +369,7 @@ ARG VLLM_MAIN_CUDA_VERSION=""
 ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+csrc.build"
 
 # Use existing torch for nightly builds
-RUN --mount=type=cache,target=/root/.cache/uv \
+RUN --mount=type=cache,target=/opt/uv/cache \
     if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
         python3 use_existing_torch.py --prefix; \
     fi
@@ -315,7 +385,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # Build the vLLM wheel
 # if USE_SCCACHE is set, use sccache to speed up compilation
 # AWS credentials mounted at ~/.aws/credentials for sccache S3 auth (optional)
-RUN --mount=type=cache,target=/root/.cache/uv \
+RUN --mount=type=cache,target=/opt/uv/cache \
     --mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \
     if [ "$USE_SCCACHE" = "1" ]; then \
         echo "Installing sccache..." \
@@ -349,7 +419,7 @@ ARG vllm_target_device="cuda"
 ENV VLLM_TARGET_DEVICE=${vllm_target_device}
 ENV CCACHE_DIR=/root/.cache/ccache
 RUN --mount=type=cache,target=/root/.cache/ccache \
-    --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=cache,target=/opt/uv/cache \
     if [ "$USE_SCCACHE" != "1" ]; then \
         # Clean any existing CMake artifacts
         rm -rf .deps && \
@@ -381,7 +451,7 @@ COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.
 # Defaults moved here from tools/ep_kernels/install_python_libraries.sh for centralized version management
 ARG DEEPEP_COMMIT_HASH=73b6ea4
 ARG NVSHMEM_VER
-RUN --mount=type=cache,target=/root/.cache/uv \
+RUN --mount=type=cache,target=/opt/uv/cache \
     mkdir -p /tmp/ep_kernels_workspace/dist && \
     export TORCH_CUDA_ARCH_LIST='9.0a 10.0a' && \
     /tmp/install_python_libraries.sh \
@@ -415,7 +485,7 @@ ENV UV_INDEX_STRATEGY="unsafe-best-match"
 # Use copy mode to avoid hardlink failures with Docker cache mounts
 ENV UV_LINK_MODE=copy
 
-RUN --mount=type=cache,target=/root/.cache/uv \
+RUN --mount=type=cache,target=/opt/uv/cache \
     if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
         echo "Installing build requirements without torch..." \
         && python3 use_existing_torch.py --prefix \
@@ -435,6 +505,10 @@ WORKDIR /workspace
 COPY --from=csrc-build /workspace/dist /precompiled-wheels
 COPY . .
 
+# Drop the pre-built rust frontend binary into the source tree. setup.py
+# detects it and ships it as-is, skipping the local cargo build.
+COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
+
 ARG GIT_REPO_CHECK=0
 RUN --mount=type=bind,source=.git,target=.git \
     if [ "$GIT_REPO_CHECK" != "0" ]; then bash tools/check_repo.sh ; fi
@@ -446,15 +520,16 @@ ENV VLLM_TARGET_DEVICE=${vllm_target_device}
 ENV VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX=1
 
 # Use existing torch for nightly builds
-RUN --mount=type=cache,target=/root/.cache/uv \
+RUN --mount=type=cache,target=/opt/uv/cache \
     if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
         python3 use_existing_torch.py --prefix; \
     fi
 
 # Build the vLLM wheel
-RUN --mount=type=cache,target=/root/.cache/uv \
+RUN --mount=type=cache,target=/opt/uv/cache \
     --mount=type=bind,source=.git,target=.git \
     if [ "${vllm_target_device}" = "cuda" ]; then \
+        export VLLM_USE_PRECOMPILED=1; \
         export VLLM_PRECOMPILED_WHEEL_LOCATION=$(ls /precompiled-wheels/*.whl); \
     fi && \
     python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38
@@ -509,7 +584,7 @@ COPY requirements/test/cuda.txt requirements/test/cuda.txt
 COPY requirements/dev.txt requirements/dev.txt
 COPY use_existing_torch.py use_existing_torch.py
 COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt
-RUN --mount=type=cache,target=/root/.cache/uv \
+RUN --mount=type=cache,target=/opt/uv/cache \
     if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
         echo "Installing dev requirements plus torch nightly..." \
         && python3 use_existing_torch.py --prefix \
@@ -609,9 +684,50 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
 RUN python3 -m pip install uv
 
 # Environment for uv
+# Redirect uv's managed Python and download cache out of /root/ so downstream
+# images (`FROM vllm/vllm-openai` + `USER <uid>`) and direct non-root runs
+# (`docker run --user <uid>:<gid>`) can read and execute them. See #15174,
+# #15359, #31959.
 ENV UV_HTTP_TIMEOUT=500
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE=copy
+ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
+ENV UV_CACHE_DIR=/opt/uv/cache
+RUN mkdir -p "${UV_PYTHON_INSTALL_DIR}" "${UV_CACHE_DIR}" \
+    && chgrp -R 0 /opt/uv \
+    && chmod -R g+rwX,a+rX /opt/uv
+
+# ----------------------------------------------------------------------
+# Non-root support (opt-in)
+# ----------------------------------------------------------------------
+# Create a conventional `vllm` user (UID 2000, GID 0) so the image can be
+# run under `--user 2000:0` or the opt-in `vllm-openai-nonroot` target.
+#
+# Design notes:
+#   * GID 0 + group-writable cache dirs follow the OpenShift arbitrary-UID
+#     pattern, so any UID that is a member of group 0 at runtime can write
+#     to /home/vllm and /opt/uv without additional chown work.
+#   * The default `vllm-openai` image keeps `USER root`, so every existing
+#     `docker run vllm/vllm-openai ...` / K8s manifest / `FROM vllm/vllm-openai`
+#     + `RUN uv pip install --system ...` flow is unchanged.
+#   * The entrypoint wrapper below is only used by `vllm-openai-nonroot`; it
+#     handles the OpenShift arbitrary-UID case (UID not in /etc/passwd).
+# See #31959 and docs/deployment/docker.md.
+RUN useradd --uid 2000 --gid 0 --create-home --home-dir /home/vllm \
+        --shell /bin/bash vllm \
+    && mkdir -p /home/vllm/.cache /home/vllm/.config \
+    && chown -R 2000:0 /home/vllm \
+    && chmod -R g+rwX /home/vllm \
+    # Allow the entrypoint wrapper to append a /etc/passwd entry for an
+    # arbitrary runtime UID that shares GID 0. Without this, `whoami`, bash's
+    # `\u` prompt, `id -un`, and anything else that calls `getpwuid()`
+    # directly return "I have no name!" for OpenShift-style arbitrary UIDs.
+    # This matches the convention used by Red Hat UBI base images.
+    && chgrp 0 /etc/passwd /etc/group \
+    && chmod g=u /etc/passwd /etc/group
+COPY docker/entrypoints/vllm-nonroot-entrypoint.sh \
+    /usr/local/bin/vllm-nonroot-entrypoint.sh
+RUN chmod 0755 /usr/local/bin/vllm-nonroot-entrypoint.sh
 
 # Enable CUDA forward compatibility by setting '-e VLLM_ENABLE_CUDA_COMPATIBILITY=1'
 # Only needed for datacenter/professional GPUs with older drivers.
@@ -628,7 +744,7 @@ ENV VLLM_ENABLE_CUDA_COMPATIBILITY=0
 ARG PYTORCH_CUDA_INDEX_BASE_URL
 COPY requirements/common.txt /tmp/common.txt
 COPY requirements/cuda.txt /tmp/requirements-cuda.txt
-RUN --mount=type=cache,target=/root/.cache/uv \
+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; \
     fi && \
@@ -640,7 +756,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # https://docs.flashinfer.ai/installation.html
 # From versions.json: .flashinfer.version
 ARG FLASHINFER_VERSION=0.6.11.post2
-RUN --mount=type=cache,target=/root/.cache/uv \
+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 '.')
 
@@ -672,7 +788,7 @@ ARG BITSANDBYTES_VERSION_X86=0.46.1
 ARG BITSANDBYTES_VERSION_ARM64=0.42.0
 ARG TIMM_VERSION=">=1.0.17"
 ARG RUNAI_MODEL_STREAMER_VERSION=">=0.15.7"
-RUN --mount=type=cache,target=/root/.cache/uv \
+RUN --mount=type=cache,target=/opt/uv/cache \
     if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
         BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_ARM64}"; \
     else \
@@ -697,7 +813,7 @@ ARG PYTORCH_NIGHTLY
 # Check whether to install torch nightly instead of release for this build.
 COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt
 RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist \
-    --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=cache,target=/opt/uv/cache \
     if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
         echo "Installing torch nightly..." \
         && uv pip install --system $(cat torch_lib_versions.txt | xargs) --pre \
@@ -711,7 +827,7 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist
         --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
     fi
 
-RUN --mount=type=cache,target=/root/.cache/uv \
+RUN --mount=type=cache,target=/opt/uv/cache \
 . /etc/environment && \
 uv pip list
 
@@ -720,7 +836,7 @@ ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
 
 # Install EP kernels wheels (DeepEP) that have been built in the `build` stage
 RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm-workspace/ep_kernels/dist \
-    --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=cache,target=/opt/uv/cache \
     uv pip install --system ep_kernels/dist/*.whl --verbose \
         --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
 
@@ -775,7 +891,7 @@ COPY requirements/test/cuda.txt requirements/test/cuda.txt
 COPY requirements/dev.txt requirements/dev.txt
 COPY use_existing_torch.py use_existing_torch.py
 COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt
-RUN --mount=type=cache,target=/root/.cache/uv \
+RUN --mount=type=cache,target=/opt/uv/cache \
     CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
     if [ "$CUDA_MAJOR" -ge 12 ]; then \
         if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
@@ -795,7 +911,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     fi
 
 # install development dependencies (for testing)
-RUN --mount=type=cache,target=/root/.cache/uv \
+RUN --mount=type=cache,target=/opt/uv/cache \
     uv pip install --system -e tests/vllm_test_utils
 
 # enable fast downloads from hf (for testing)
@@ -835,7 +951,7 @@ ENV UV_HTTP_TIMEOUT=500
 # install kv_connectors if requested
 ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX'
 ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
-RUN --mount=type=cache,target=/root/.cache/uv \
+RUN --mount=type=cache,target=/opt/uv/cache \
     --mount=type=bind,source=requirements/kv_connectors.txt,target=/tmp/kv_connectors.txt,ro \
     CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
     CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-'); \
@@ -903,5 +1019,32 @@ ENTRYPOINT ["./sagemaker-entrypoint.sh"]
 
 FROM vllm-openai-base AS vllm-openai
 
+# To run the image as non-root, either build the `vllm-openai-nonroot` target
+# below, or in a derived Dockerfile uncomment the following line and ensure
+# any additional layers chgrp-0 / chmod-g+rwX paths they write to. The `vllm`
+# user (UID 2000, GID 0) is already created in the `vllm-base` stage.
+# See docs/deployment/docker.md.
+# USER vllm
 ENTRYPOINT ["vllm", "serve"]
 #################### OPENAI API SERVER ####################
+
+#################### OPENAI API SERVER (NON-ROOT, OPT-IN) ####################
+# Non-root-ready variant of `vllm-openai`. Built via:
+#   docker build --target vllm-openai-nonroot -t vllm:openai-nonroot \
+#       -f docker/Dockerfile .
+#
+# Runtime behavior:
+#   * Default USER is `vllm` (UID 2000, GID 0) created in `vllm-base`.
+#   * HOME is /home/vllm, pre-created group-0-writable so arbitrary UIDs in
+#     group 0 (OpenShift / `--user <uid>:0`) can also use the image.
+#   * Entrypoint wrapper handles the "UID not in /etc/passwd" case for truly
+#     arbitrary UIDs by falling back HOME/USER to sane writable defaults.
+#   * All cache/config envs (HF_HOME, VLLM_CACHE_ROOT, TRITON_CACHE_DIR, ...)
+#     remain unset so their library defaults resolve to $HOME/.cache/... ,
+#     which is writable.
+FROM vllm-openai AS vllm-openai-nonroot
+
+USER vllm
+WORKDIR /home/vllm
+ENTRYPOINT ["/usr/local/bin/vllm-nonroot-entrypoint.sh"]
+#################### OPENAI API SERVER (NON-ROOT, OPT-IN) ####################

docker/Dockerfile.cpu

@@ -27,15 +27,19 @@ WORKDIR /workspace
 ARG PYTHON_VERSION=3.12
 ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"
 
+ARG max_jobs=32
+ENV MAX_JOBS=${max_jobs}
+
 # Install minimal dependencies and uv
 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     --mount=type=cache,target=/var/lib/apt,sharing=locked \
     apt-get update -y \
-    && apt-get install -y --no-install-recommends sudo ccache git curl wget ca-certificates \
+    && apt-get install -y --no-install-recommends sudo ccache git curl wget ca-certificates zlib1g-dev \
     gcc-12 g++-12 libtcmalloc-minimal4 libnuma-dev ffmpeg libsm6 libxext6 libgl1 jq lsof make xz-utils \
     && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 10 --slave /usr/bin/g++ g++ /usr/bin/g++-12 \
     && curl -LsSf https://astral.sh/uv/install.sh | sh
 
+# Compiler and linker environment
 ENV CC=/usr/bin/gcc-12 CXX=/usr/bin/g++-12
 ENV CCACHE_DIR=/root/.cache/ccache
 ENV CMAKE_CXX_COMPILER_LAUNCHER=ccache
@@ -80,12 +84,48 @@ FROM base-${TARGETARCH} AS base
 
 RUN echo 'ulimit -c 0' >> ~/.bashrc
 
+######################### RUST BUILD IMAGE #########################
+# Build the Rust frontend (`vllm-rs`) in a dedicated stage so the wheel build
+# stage doesn't need the rust toolchain or protoc. This stage runs in parallel
+# with the main vllm-build stage.
+FROM ubuntu:22.04 AS rust-build
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update -y \
+    && apt-get install -y --no-install-recommends \
+        ca-certificates curl git build-essential unzip \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY tools/install_protoc.sh /tmp/install_protoc.sh
+RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
+
+# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
+RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
+    sh -s -- -y --profile minimal --default-toolchain none
+ENV PATH="/root/.cargo/bin:${PATH}"
+
+WORKDIR /workspace
+
+# Copy only the rust workspace — the binary is the sole artifact we need.
+COPY rust rust
+COPY rust-toolchain.toml rust-toolchain.toml
+COPY build_rust.sh build_rust.sh
+
+# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
+# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
+ENV CARGO_BUILD_JOBS=4
+
+# Build the release binary. Cache cargo registry/git and target/, but copy the
+# binary out of the target/ cache mount 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 \
+    --mount=type=cache,target=/workspace/rust/target \
+    VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
+
 ######################### BUILD IMAGE #########################
 FROM base AS vllm-build
 
-ARG max_jobs=32
-ENV MAX_JOBS=${max_jobs}
-
 ARG GIT_REPO_CHECK=0
 # Support for cross-compilation with x86 ISA including AVX2 and AVX512: docker build --build-arg VLLM_CPU_X86="true" ...
 ARG VLLM_CPU_X86=0
@@ -114,6 +154,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 
 COPY . .
 
+# Drop the pre-built rust frontend binary into the source tree. setup.py
+# detects it and ships it as-is, skipping the local cargo build.
+COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
+
 RUN if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh ; fi
 
 RUN --mount=type=cache,target=/root/.cache/uv \
@@ -213,8 +257,7 @@ 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 dist/*.whl && \
-    uv pip install "vllm[audio]"
+    uv pip install "$(realpath dist/*.whl)[audio,triton-cpu]"
 
 # Add labels to document build configuration
 LABEL org.opencontainers.image.title="vLLM CPU"

docker/Dockerfile.nightly_torch

@@ -94,6 +94,41 @@ RUN cat torch_build_versions.txt
 
 #################### BASE BUILD IMAGE ####################
 
+#################### RUST BUILD IMAGE ####################
+# Build the Rust frontend (`vllm-rs`) in a dedicated stage so the wheel build
+# stage doesn't need the rust toolchain or protoc.
+FROM ubuntu:22.04 AS rust-build
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update -y \
+    && apt-get install -y --no-install-recommends \
+        ca-certificates curl git build-essential unzip \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY tools/install_protoc.sh /tmp/install_protoc.sh
+RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
+
+# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
+RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
+    sh -s -- -y --profile minimal --default-toolchain none
+ENV PATH="/root/.cargo/bin:${PATH}"
+
+WORKDIR /workspace
+
+COPY rust rust
+COPY rust-toolchain.toml rust-toolchain.toml
+COPY build_rust.sh build_rust.sh
+
+# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
+# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
+ENV CARGO_BUILD_JOBS=4
+
+RUN --mount=type=cache,target=/root/.cargo/registry \
+    --mount=type=cache,target=/root/.cargo/git \
+    --mount=type=cache,target=/workspace/rust/target \
+    VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
+#################### RUST BUILD IMAGE ####################
+
 #################### WHEEL BUILD IMAGE ####################
 FROM base AS build
 ARG TARGETPLATFORM
@@ -104,6 +139,10 @@ ENV UV_HTTP_TIMEOUT=500
 
 COPY . .
 
+# Drop the pre-built rust frontend binary into the source tree. setup.py
+# detects it and ships it as-is, skipping the local cargo build.
+COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
+
 RUN python3 use_existing_torch.py
 
 RUN --mount=type=cache,target=/root/.cache/uv \

docker/Dockerfile.rocm

@@ -111,9 +111,47 @@ ONBUILD RUN git clone ${VLLM_REPO} \
                && git fetch upstream ; fi
 FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm
 
+# -----------------------
+# Rust build stage
+# Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages
+# don't need the rust toolchain or protoc. Runs in parallel with the main wheel
+# build for faster end-to-end builds.
+FROM fetch_vllm AS rust-build
+ARG COMMON_WORKDIR
+
+# protoc is used by tonic-build/prost-build.
+RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
+        ca-certificates curl unzip \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY tools/install_protoc.sh /tmp/install_protoc.sh
+RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
+
+# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
+RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
+    sh -s -- -y --profile minimal --default-toolchain none
+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
+
+# 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 \
+    cd ${COMMON_WORKDIR}/vllm \
+    && VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh
+
 # -----------------------
 # vLLM build stages
 FROM fetch_vllm AS build_vllm
+ARG COMMON_WORKDIR
+
+# Drop the pre-built rust frontend binary into the source tree. setup.py
+# detects it and ships it as-is, skipping the local cargo build.
+COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
+
 # Build vLLM (setup.py auto-detects sccache in PATH)
 RUN cd vllm \
     && python3 -m pip install -r requirements/rocm.txt \
@@ -293,6 +331,10 @@ FROM fetch_vllm AS build_vllm_wheel_release
 
 ARG COMMON_WORKDIR
 
+# Drop the pre-built rust frontend binary into the source tree. setup.py
+# detects it and ships it as-is, skipping the local cargo build.
+COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
+
 # Create /install directory for custom wheels
 RUN mkdir -p /install
 

docker/Dockerfile.xpu

@@ -1,3 +1,37 @@
+######################### RUST BUILD IMAGE #########################
+# Build the Rust frontend (`vllm-rs`) in a dedicated stage so the main image
+# doesn't need the rust toolchain or protoc. Runs in parallel with vllm-base.
+FROM ubuntu:22.04 AS rust-build
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update -y \
+    && apt-get install -y --no-install-recommends \
+        ca-certificates curl git build-essential unzip \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY tools/install_protoc.sh /tmp/install_protoc.sh
+RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
+
+# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
+RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
+    sh -s -- -y --profile minimal --default-toolchain none
+ENV PATH="/root/.cargo/bin:${PATH}"
+
+WORKDIR /workspace
+
+COPY rust rust
+COPY rust-toolchain.toml rust-toolchain.toml
+COPY build_rust.sh build_rust.sh
+
+# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
+# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
+ENV CARGO_BUILD_JOBS=4
+
+RUN --mount=type=cache,target=/root/.cargo/registry \
+    --mount=type=cache,target=/root/.cargo/git \
+    --mount=type=cache,target=/workspace/rust/target \
+    VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
+
 FROM intel/deep-learning-essentials:2025.3.2-0-devel-ubuntu24.04 AS vllm-base
 
 WORKDIR /workspace/
@@ -99,6 +133,11 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
 
 COPY . .
+
+# Drop the pre-built rust frontend binary into the source tree. setup.py
+# detects it and ships it as-is, skipping the local cargo build.
+COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
+
 ARG GIT_REPO_CHECK=0
 RUN --mount=type=bind,source=.git,target=.git \
     if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh; fi

docker/entrypoints/test_vllm_nonroot_entrypoint.sh

@@ -0,0 +1,266 @@
+#!/bin/sh
+# Shell-level unit test for vllm-nonroot-entrypoint.sh.
+#
+# Runs on the host (no Docker, no GPU) by stubbing `vllm` with a shim that
+# dumps its env + argv instead of actually serving. Exercises the wrapper's
+# HOME/USER fallback behavior that can't be easily tested from buildkite
+# (which would need a GPU to run `vllm serve --help`).
+#
+# Usage:
+#   bash docker/entrypoints/test_vllm_nonroot_entrypoint.sh
+# Exits non-zero on the first failed assertion.
+
+set -eu
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+WRAPPER="${SCRIPT_DIR}/vllm-nonroot-entrypoint.sh"
+
+if [ ! -x "$WRAPPER" ]; then
+    echo "FAIL: wrapper not found or not executable: $WRAPPER" >&2
+    exit 1
+fi
+
+WORKDIR="$(mktemp -d)"
+trap 'rm -rf "$WORKDIR"' EXIT
+
+# Stub `vllm` on PATH. It dumps env + argv + cwd to stdout so we can assert.
+mkdir -p "$WORKDIR/bin"
+cat > "$WORKDIR/bin/vllm" <<'EOF'
+#!/bin/sh
+echo "ARGV=$*"
+echo "HOME=${HOME-__unset__}"
+echo "USER=${USER-__unset__}"
+echo "LOGNAME=${LOGNAME-__unset__}"
+echo "PWD=$(pwd)"
+EOF
+chmod +x "$WORKDIR/bin/vllm"
+
+run_wrapper() {
+    # Usage: run_wrapper <output_file> <env_kv>... -- <wrapper_arg>...
+    _out="$1"; shift
+    _env=""
+    while [ "${1:-}" != "--" ]; do
+        _env="$_env $1"; shift
+    done
+    shift
+    env -i PATH="$WORKDIR/bin:/usr/bin:/bin" $_env "$WRAPPER" "$@" > "$_out"
+}
+
+fail() { echo "FAIL: $*" >&2; echo "--- stdout ---" >&2; cat "$1" >&2; exit 1; }
+
+expect_default_home() {
+    _out="$1"
+    _case="$2"
+    if [ -w /home/vllm ]; then
+        expected_home="/home/vllm"
+        grep -q "^HOME=$expected_home\$" "$_out" \
+            || fail "$_out" "$_case: HOME not set to $expected_home"
+    else
+        expected_home="/tmp/vllm-home.XXXXXX"
+        grep -Eq '^HOME=/tmp/vllm-home\.[^/]+$' "$_out" \
+            || fail "$_out" "$_case: HOME not set to $expected_home"
+    fi
+}
+
+# -----------------------------------------------------------------------------
+# Case 1: writable HOME and USER both set -> wrapper must leave them alone.
+# -----------------------------------------------------------------------------
+case1_home="$WORKDIR/case1-home"
+mkdir -p "$case1_home"
+out="$WORKDIR/case1.out"
+run_wrapper "$out" "HOME=$case1_home" "USER=alice" "LOGNAME=alice" -- --model foo
+grep -q "^HOME=$case1_home\$" "$out" || fail "$out" "case1: HOME not preserved"
+grep -q "^USER=alice\$" "$out" || fail "$out" "case1: USER not preserved"
+grep -q "^LOGNAME=alice\$" "$out" || fail "$out" "case1: LOGNAME not preserved"
+grep -q "^ARGV=serve --model foo\$" "$out" || fail "$out" "case1: ARGV wrong"
+echo "PASS: case1 (writable HOME + USER preserved)"
+
+# -----------------------------------------------------------------------------
+# Case 2: HOME unset -> falls back to /home/vllm if writable, else
+# /tmp/vllm-home.XXXXXX.
+# -----------------------------------------------------------------------------
+# The wrapper checks whether the real /home/vllm exists and is writable. On
+# dev machines /home/vllm typically does NOT exist, so the
+# wrapper should fall to /tmp/vllm-home.XXXXXX.
+out="$WORKDIR/case2.out"
+run_wrapper "$out" -- --model bar
+expect_default_home "$out" "case2"
+grep -q "^USER=vllm\$" "$out" || fail "$out" "case2: USER not defaulted to vllm"
+grep -q "^LOGNAME=vllm\$" "$out" || fail "$out" "case2: LOGNAME not defaulted to vllm"
+grep -q "^ARGV=serve --model bar\$" "$out" || fail "$out" "case2: ARGV wrong"
+echo "PASS: case2 (unset HOME falls back to $expected_home, USER defaulted)"
+
+# -----------------------------------------------------------------------------
+# Case 3: HOME set but unwritable -> must also fall back.
+# -----------------------------------------------------------------------------
+ro_home="$WORKDIR/ro-home"
+mkdir -p "$ro_home"
+chmod 0500 "$ro_home"
+out="$WORKDIR/case3.out"
+run_wrapper "$out" "HOME=$ro_home" -- --model baz
+expect_default_home "$out" "case3"
+grep -q "^USER=vllm\$" "$out" || fail "$out" "case3: USER not defaulted"
+chmod 0700 "$ro_home"
+echo "PASS: case3 (unwritable HOME overridden)"
+
+# -----------------------------------------------------------------------------
+# Case 4: USER set but LOGNAME unset -> LOGNAME mirrors USER.
+# -----------------------------------------------------------------------------
+case4_home="$WORKDIR/case4-home"
+mkdir -p "$case4_home"
+out="$WORKDIR/case4.out"
+run_wrapper "$out" "HOME=$case4_home" "USER=carol" -- --model qux
+grep -q "^USER=carol\$" "$out" || fail "$out" "case4: USER not preserved"
+grep -q "^LOGNAME=carol\$" "$out" || fail "$out" "case4: LOGNAME not mirrored from USER"
+echo "PASS: case4 (LOGNAME mirrors USER when unset)"
+
+# -----------------------------------------------------------------------------
+# Case 5: /etc/passwd is writable AND the current UID is not in it -> wrapper
+# appends a synthetic entry. Uses the VLLM_PASSWD_FILE test hook so we don't
+# touch the real /etc/passwd.
+# -----------------------------------------------------------------------------
+fake_passwd="$WORKDIR/fake-passwd"
+: > "$fake_passwd"  # empty file, current UID definitely not present
+case5_home="$WORKDIR/case5-home"
+mkdir -p "$case5_home"
+out="$WORKDIR/case5.out"
+run_wrapper "$out" "HOME=$case5_home" "VLLM_PASSWD_FILE=$fake_passwd" -- --model foo
+current_uid="$(id -u)"
+current_gid="$(id -g)"
+expected_line="vllm:x:${current_uid}:${current_gid}:vllm:${case5_home}:/bin/bash"
+grep -Fx "$expected_line" "$fake_passwd" > /dev/null \
+    || { echo "FAIL: case5: expected line not found in fake passwd:"; echo "  expected: $expected_line"; echo "  file contents:"; cat "$fake_passwd"; exit 1; }
+echo "PASS: case5 (passwd entry appended for arbitrary UID)"
+
+# -----------------------------------------------------------------------------
+# Case 6: /etc/passwd is writable but current UID already has an entry ->
+# wrapper must NOT duplicate the entry.
+# -----------------------------------------------------------------------------
+fake_passwd="$WORKDIR/fake-passwd-prepopulated"
+printf 'vllm:x:%s:%s:vllm:/home/vllm:/bin/bash\n' "$current_uid" "$current_gid" > "$fake_passwd"
+out="$WORKDIR/case6.out"
+run_wrapper "$out" "HOME=$case5_home" "VLLM_PASSWD_FILE=$fake_passwd" -- --model foo
+line_count="$(wc -l < "$fake_passwd")"
+# NOTE: wc may count 0 or 1 depending on trailing newline; accept 1.
+# More robust: count lines matching our UID.
+uid_lines="$(grep -c ":${current_uid}:" "$fake_passwd" || true)"
+[ "$uid_lines" = "1" ] \
+    || { echo "FAIL: case6: expected exactly one entry for UID $current_uid, got $uid_lines"; cat "$fake_passwd"; exit 1; }
+echo "PASS: case6 (existing passwd entry not duplicated)"
+
+# -----------------------------------------------------------------------------
+# Case 7: /etc/passwd is NOT writable -> wrapper must NOT crash, just skip.
+# Skipped when running as root, because root's DAC override means [ -w ... ]
+# is always true regardless of mode bits -- the case can't be simulated.
+# In the real deployment (non-root UID inside the container) this IS the
+# relevant behavior and is what `_passwd_file is not writable` encodes.
+# -----------------------------------------------------------------------------
+if [ "$(id -u)" = "0" ]; then
+    echo "SKIP: case7 (running as root; DAC override makes unwritable check meaningless)"
+else
+    fake_passwd="$WORKDIR/ro-passwd"
+    : > "$fake_passwd"
+    chmod 0444 "$fake_passwd"
+    out="$WORKDIR/case7.out"
+    run_wrapper "$out" "HOME=$case5_home" "VLLM_PASSWD_FILE=$fake_passwd" -- --model foo
+    # File must remain empty (no write happened) and the wrapper exec'd
+    # `vllm serve` successfully (stdout contains ARGV line).
+    [ ! -s "$fake_passwd" ] \
+        || { echo "FAIL: case7: RO passwd file was modified"; cat "$fake_passwd"; exit 1; }
+    grep -q "^ARGV=serve --model foo\$" "$out" || fail "$out" "case7: wrapper didn't exec vllm"
+    chmod 0600 "$fake_passwd"
+    echo "PASS: case7 (unwritable passwd file tolerated)"
+fi
+
+# -----------------------------------------------------------------------------
+# Case 8: caller's writable CWD is preserved — wrapper must NOT chdir to HOME
+# when cwd is usable. Protects relative-path workflows like
+# `docker run -w /models ... --model ./llama.gguf`.
+# -----------------------------------------------------------------------------
+case8_home="$WORKDIR/case8-home"
+mkdir -p "$case8_home"
+case8_cwd="$WORKDIR/case8-cwd"
+mkdir -p "$case8_cwd"
+out="$WORKDIR/case8.out"
+(cd "$case8_cwd" && run_wrapper "$out" "HOME=$case8_home" "USER=alice" "LOGNAME=alice" -- --model ./relpath)
+grep -q "^PWD=$case8_cwd\$" "$out" \
+    || fail "$out" "case8: writable cwd not preserved (got $(grep '^PWD=' "$out"))"
+grep -q "^ARGV=serve --model \\./relpath\$" "$out" \
+    || fail "$out" "case8: relative argv not preserved"
+echo "PASS: case8 (writable cwd preserved; relative argv still resolves from caller's cwd)"
+
+# -----------------------------------------------------------------------------
+# Case 9: read-only cwd is ALSO preserved. A caller who mounts a read-only
+# model directory at the container's cwd (e.g. `docker run -w /models` with
+# /models bind-mounted ro) expects relative argv like `--model ./foo.gguf`
+# to resolve against /models. An earlier version of this wrapper rewrote
+# read-only cwd to $HOME and broke that workflow; this case guards against
+# the regression returning.
+# -----------------------------------------------------------------------------
+case9_home="$WORKDIR/case9-home"
+mkdir -p "$case9_home"
+case9_ro="$WORKDIR/case9-ro"
+mkdir -p "$case9_ro"
+chmod 0555 "$case9_ro"
+out="$WORKDIR/case9.out"
+(cd "$case9_ro" && run_wrapper "$out" "HOME=$case9_home" "USER=alice" "LOGNAME=alice" -- --model ./foo)
+grep -q "^PWD=$case9_ro\$" "$out" \
+    || fail "$out" "case9: read-only cwd was rewritten (got $(grep '^PWD=' "$out"))"
+grep -q "^ARGV=serve --model \\./foo\$" "$out" \
+    || fail "$out" "case9: relative argv not preserved"
+chmod 0700 "$case9_ro"
+echo "PASS: case9 (read-only cwd preserved; relative argv still resolves from caller's cwd)"
+
+# -----------------------------------------------------------------------------
+# Case 10: truly inaccessible cwd (no search bit) DOES fall back to $HOME.
+# Skipped as root because DAC override lets root cd into 0000 directories.
+# -----------------------------------------------------------------------------
+if [ "$(id -u)" = "0" ]; then
+    echo "SKIP: case10 (running as root; DAC override makes inaccessible cwd untestable)"
+else
+    case10_home="$WORKDIR/case10-home"
+    mkdir -p "$case10_home"
+    case10_cwd="$WORKDIR/case10-cwd"
+    mkdir -p "$case10_cwd"
+    out="$WORKDIR/case10.out"
+    # Make cwd genuinely inaccessible (mode 0000 = no search bit -> cd .
+    # fails with EACCES). Use absolute paths for chmod so our own test
+    # cleanup still works without needing search perm on the dir.
+    (
+        cd "$case10_cwd"
+        chmod 0000 "$case10_cwd"
+        run_wrapper "$out" "HOME=$case10_home" "USER=alice" "LOGNAME=alice" -- --model foo
+    )
+    chmod 0700 "$case10_cwd"
+    grep -q "^PWD=$case10_home\$" "$out" \
+        || fail "$out" "case10: inaccessible cwd not overridden to HOME (got $(grep '^PWD=' "$out"))"
+    echo "PASS: case10 (inaccessible cwd falls back to \$HOME)"
+fi
+
+# -----------------------------------------------------------------------------
+# Case 11: if /tmp cannot create a private fallback dir, wrapper uses /tmp as
+# the last-resort HOME instead of leaving HOME empty under set -eu.
+# -----------------------------------------------------------------------------
+if [ -w /home/vllm ]; then
+    echo "SKIP: case11 (/home/vllm is writable; mktemp fallback path is not used)"
+else
+    cat > "$WORKDIR/bin/mktemp" <<'EOF'
+#!/bin/sh
+exit 1
+EOF
+    chmod +x "$WORKDIR/bin/mktemp"
+
+    out="$WORKDIR/case11.out"
+    run_wrapper "$out" -- --model no-mktemp
+    rm -f "$WORKDIR/bin/mktemp"
+
+    grep -q "^HOME=/tmp\$" "$out" \
+        || fail "$out" "case11: mktemp failure did not fall back to /tmp"
+    grep -q "^USER=vllm\$" "$out" || fail "$out" "case11: USER not defaulted"
+    grep -q "^LOGNAME=vllm\$" "$out" || fail "$out" "case11: LOGNAME not defaulted"
+    grep -q "^ARGV=serve --model no-mktemp\$" "$out" || fail "$out" "case11: ARGV wrong"
+    echo "PASS: case11 (mktemp failure falls back to /tmp)"
+fi
+
+echo ""
+echo "ALL CASES PASSED."

docker/entrypoints/vllm-nonroot-entrypoint.sh

@@ -0,0 +1,87 @@
+#!/bin/sh
+# Entrypoint wrapper for the opt-in `vllm-openai-nonroot` image.
+#
+# The image also ships a `vllm` user (UID 2000, GID 0) with HOME /home/vllm
+# and a group-0-writable home directory. When the container is launched with
+# `--user 2000:0` (or any other UID in group 0) the passwd entry is enough on
+# its own: Docker picks up HOME=/home/vllm, getpass.getuser() resolves to
+# "vllm", and every cache dir (HF, Triton, Inductor, vLLM, Numba, Outlines)
+# that defaults to `$HOME/.cache/...` lands in a writable location.
+#
+# This wrapper exists for the *arbitrary-UID* case (e.g. OpenShift's
+# `runAsUser: 1000540000` Restricted Pod Security Standard) where the caller
+# UID is not in /etc/passwd at all. In that case:
+#   * $HOME may be unset or resolve to "/" (unwritable).
+#   * getpass.getuser() falls back to pwd.getpwuid() -> KeyError.
+#
+# The wrapper re-points $HOME to /home/vllm when writable, /tmp/vllm-home.XXXXXX
+# otherwise, and defaults $USER to "vllm" so the pwd-lookup path is never
+# taken. Everything else is forwarded to `vllm serve`.
+#
+# Non-empty caller-set env vars (HOME, USER, LOGNAME) are preserved, so
+# existing K8s manifests and `docker run -e ...` keep working unchanged.
+# Unset or empty values fall through to the wrapper's defaults, matching
+# what shell code typically expects from "unset".
+
+set -eu
+
+if [ -z "${HOME:-}" ] || [ ! -w "${HOME}" ]; then
+    if [ -w /home/vllm ]; then
+        export HOME=/home/vllm
+    else
+        if _h="$(mktemp -d /tmp/vllm-home.XXXXXX 2>/dev/null)"; then
+            export HOME="$_h"
+            chmod 0700 "$HOME" 2>/dev/null || true
+        else
+            export HOME=/tmp
+        fi
+        unset _h
+    fi
+fi
+
+# Preserve the caller's cwd whenever it's still usable. A read-only mount
+# (e.g. `docker run -w /models ... --model ./llama.gguf` where /models is
+# the user's model share) is a legitimate, usable cwd — vllm only needs to
+# *read* relative paths from there. We only fall back to $HOME when the
+# cwd itself is truly inaccessible (no search bit, deleted inode, mount
+# gone, etc.), which is when `cd .` actually fails.
+#
+# This is the accessibility check, not a writability check; the latter
+# would silently rewrite cwd for any read-only workflow and break relative
+# argv like `--model ./llama.gguf`, `--chat-template ./t.jinja`, relative
+# TLS cert paths, etc.
+if ! cd . 2>/dev/null; then
+    cd "$HOME"
+fi
+
+# getpass.getuser() prefers $USER/$LOGNAME/etc. before hitting getpwuid();
+# setting it here makes the "UID not in passwd" path a no-op for everything
+# in the process tree.
+if [ -z "${USER:-}" ]; then
+    export USER=vllm
+fi
+if [ -z "${LOGNAME:-}" ]; then
+    export LOGNAME="$USER"
+fi
+
+# Shell-level tooling (`whoami`, bash's `\u` prompt, `id -un`, `sudo`) does
+# NOT consult $USER; it calls getpwuid(geteuid()) directly. For arbitrary
+# runtime UIDs in OpenShift-style deploys this returns "I have no name!".
+# If /etc/passwd is group-0 writable (set at build time) and doesn't yet
+# have an entry for this UID, append a synthetic one so every downstream
+# consumer sees a consistent "vllm" identity.
+#
+# We parse the passwd file directly instead of calling `getent` because
+# the container's NSS is typically just files anyway, and this lets us
+# unit-test via the VLLM_PASSWD_FILE hook (undocumented; production uses
+# /etc/passwd).
+_passwd_file="${VLLM_PASSWD_FILE:-/etc/passwd}"
+_uid="$(id -u)"
+if [ -w "$_passwd_file" ] \
+    && ! awk -F: -v u="$_uid" '$3==u {found=1; exit} END {exit !found}' "$_passwd_file" 2>/dev/null; then
+    printf 'vllm:x:%s:%s:vllm:%s:/bin/bash\n' \
+        "$_uid" "$(id -g)" "$HOME" >> "$_passwd_file"
+fi
+unset _uid _passwd_file
+
+exec vllm serve "$@"

docs/cli/.nav.yml

@@ -6,3 +6,5 @@ nav:
   - run-batch.md
   - vllm bench:
     - bench/**/*.md
+  - vllm launch:
+    - launch/**/*.md

docs/cli/README.md

@@ -9,7 +9,7 @@ vllm --help
 Available Commands:
 
 ```bash
-vllm {chat,complete,serve,bench,collect-env,run-batch}
+vllm {chat,complete,serve,launch,bench,collect-env,run-batch}
 ```
 
 ## serve
@@ -37,24 +37,36 @@ vllm serve meta-llama/Llama-2-7b-hf --uds /tmp/vllm.sock
 Check with --help for more options:
 
 ```bash
-# To list all groups
-vllm serve --help=listgroup
+# To list all flags
+vllm serve --help=all
 
-# To view a argument group
+# To view an argument group
 vllm serve --help=ModelConfig
 
 # To view a single argument
 vllm serve --help=max-num-seqs
 
-# To search by keyword
+# To search by keyword or flag name
 vllm serve --help=max
-
-# To view full help with pager (less/more)
-vllm serve --help=page
 ```
 
 See [vllm serve](./serve.md) for the full reference of all available arguments.
 
+## launch
+
+Launch individual vLLM components.
+
+```bash
+# Launch the rendering server component
+vllm launch render meta-llama/Llama-3.2-1B-Instruct
+
+# Inspect all available flags for the render component
+vllm launch render --help=all
+```
+
+See [vllm launch render](./launch/render.md) for the current launch
+component reference.
+
 ## chat
 
 Generate chat completions via the running API server.

docs/cli/launch/render.md

@@ -0,0 +1,22 @@
+# vllm launch render
+
+## Overview
+
+`vllm launch render` starts a GPU-less rendering server for preprocessing and
+postprocessing only.
+
+```bash
+vllm launch render meta-llama/Llama-3.2-1B-Instruct --port 8100
+```
+
+This command reuses the standard serving parser, so model, frontend,
+networking, and related CLI options follow the same conventions as
+[`vllm serve`](../serve.md).
+
+## JSON CLI Arguments
+
+--8<-- "docs/cli/json_tip.inc.md"
+
+## Arguments
+
+--8<-- "docs/generated/argparse/launch_render.inc.md"

docs/configuration/optimization.md

@@ -270,31 +270,34 @@ Known supported models (with corresponding benchmarks):
 
 ## Input Processing
 
-### fastokens Tokenizer Mode
+### fastokens Backend
 
 By default vLLM uses the standard Hugging Face `tokenizers` library to power
-the fast tokenizer (`--tokenizer-mode hf`). For BPE tokenizers (Qwen, Llama,
-DeepSeek, GPT-OSS, etc.) you can switch to the
-[fastokens](https://github.com/crusoecloud/fastokens) Rust backend, a drop-in
-replacement that's substantially faster on encode/decode and on streaming
-detokenization:
+the fast tokenizer. For BPE tokenizers (Qwen, Llama, DeepSeek, GPT-OSS, etc.)
+you can switch to the [fastokens](https://github.com/crusoecloud/fastokens)
+Rust backend, a drop-in replacement that's substantially faster on
+encode/decode and on streaming detokenization. Enable it by setting
+`VLLM_USE_FASTOKENS=1`:
 
 ```console
-vllm serve Qwen/Qwen3-8B --tokenizer-mode fastokens
+VLLM_USE_FASTOKENS=1 vllm serve Qwen/Qwen3-8B
 ```
 
 Equivalent in the offline API:
 
 ```python
+import os
+os.environ["VLLM_USE_FASTOKENS"] = "1"
+
 from vllm import LLM
-llm = LLM(model="Qwen/Qwen3-8B", tokenizer_mode="fastokens")
+llm = LLM(model="Qwen/Qwen3-8B")
 ```
 
-The `fastokens` Python package must be installed; if it isn't, vLLM raises
-a clear `ImportError` at tokenizer load. `fastokens` loads a Hugging Face
-fast tokenizer with its inner Rust tokenizer replaced by the fastokens shim,
-so it is mutually exclusive with non-HF modes such as `mistral` or
-`deepseek_v32`.
+The `fastokens` Python package (>= 0.2.0) must be installed; if it isn't,
+vLLM raises a clear `ImportError` at tokenizer load. The override applies to
+any `--tokenizer-mode` that ends up loading an HF fast tokenizer (`hf`,
+`deepseek_v32`, `deepseek_v4`, `qwen_vl`, …). Modes that don't use the HF
+fast tokenizer (`mistral`, `grok2`, `kimi_audio`) ignore the flag.
 
 Tokenizer-bound workloads — long shared prefixes, bursty short prompts,
 batch detokenization — see the largest wins. If your bottleneck is GPU

docs/contributing/README.md

@@ -43,6 +43,13 @@ If you are only developing vLLM's Python code, install vLLM using:
 VLLM_USE_PRECOMPILED=1 uv pip install -e .
 ```
 
+To rebuild only the Rust frontend binary:
+
+```bash
+./build_rust.sh          # release build
+./build_rust.sh --debug  # faster build for development
+```
+
 If you are developing vLLM's Python and CUDA/C++ code, install Pytorch first:
 
 ```bash
@@ -297,6 +304,18 @@ review process:
   resources. The reviewer will add `ready` label to the PR when the PR is
   ready to merge or a full CI run is needed.
 
+### Escalating Stalled Contributions
+
+If you have an important contribution that has not yet received maintainer attention, please email us at:
+
+<pr-review-request@vllm.ai>
+
+Using a verifiable company or university email, include:
+
+- your production or research use case
+- the problem you encountered
+- how your contribution addresses it
+
 ## Thank You
 
 Finally, thank you for taking the time to read these guidelines and for your interest in contributing to vLLM.

docs/contributing/ci/nightly_builds.md

@@ -136,10 +136,10 @@ When installing vLLM with `VLLM_USE_PRECOMPILED=1`, the `setup.py` script:
 3. **Selects compatible wheel** based on:
     - Package name (`vllm`)
     - Platform tag (architecture match)
-4. **Downloads and extracts** precompiled binaries from the wheel:
-    - C++ extension modules (`.so` files)
-    - Flash Attention Python modules
-    - Triton kernel Python files
+4. **Downloads and extracts** precompiled artifacts from the wheel:
+    - Native extension modules (`.so` files)
+    - The `vllm-rs` Rust frontend binary
+    - Flash Attention Python modules and Triton/FlashMLA Python files
 5. **Patches package_data** to include extracted files in the installation
 
 !!! note "What is the base commit?"

docs/deployment/docker.md

@@ -8,6 +8,64 @@ toc_depth: 2
 
 --8<-- "docs/getting_started/installation/gpu.md:pre-built-images"
 
+## Run as a non-root user
+
+The CUDA `vllm/vllm-openai` image runs as root by default for backward
+compatibility. It is also prepared to run as the built-in `vllm` user
+(UID 2000, GID 0):
+
+```bash
+docker run --rm --gpus all \
+    --user 2000:0 \
+    -p 8000:8000 \
+    vllm/vllm-openai:latest \
+    meta-llama/Llama-3.1-8B-Instruct
+```
+
+When mounting model or cache volumes for a non-root container, mount writable
+paths under `/home/vllm` instead of `/root`. For example, mount the Hugging
+Face cache at `/home/vllm/.cache/huggingface` and make the mounted directory
+writable by group 0.
+
+```bash
+docker run --rm --gpus all \
+    --user 2000:0 \
+    -v ~/.cache/huggingface:/home/vllm/.cache/huggingface \
+    -p 8000:8000 \
+    vllm/vllm-openai:latest \
+    meta-llama/Llama-3.1-8B-Instruct
+```
+
+To build an image that defaults to the non-root `vllm` user, use the opt-in
+`vllm-openai-nonroot` target:
+
+```bash
+docker build --target vllm-openai-nonroot \
+    -t vllm-openai-nonroot:local \
+    -f docker/Dockerfile .
+
+docker run --rm --gpus all \
+    -p 8000:8000 \
+    vllm-openai-nonroot:local \
+    meta-llama/Llama-3.1-8B-Instruct
+```
+
+The `vllm-openai-nonroot` target also supports OpenShift-style arbitrary UIDs
+when the runtime UID is a member of group 0. In Kubernetes manifests, set the
+container security context accordingly and keep mounted cache/model paths
+writable by group 0:
+
+```yaml
+securityContext:
+  runAsNonRoot: true
+  runAsUser: 1000540000
+  runAsGroup: 0
+  fsGroup: 0
+```
+
+Runtime UIDs outside group 0 are not part of the documented support matrix
+because they may be unable to write to `/home/vllm` or `/opt/uv/cache`.
+
 ## Build image from source
 
 --8<-- "docs/getting_started/installation/gpu.md:build-image-from-source"

docs/design/attention_backends.md

@@ -170,8 +170,8 @@ 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 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
-| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
+| `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 |
 | `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 |
@@ -219,7 +219,7 @@ MLA decode backends are selected using the standard
 | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
 | `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
 | `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
-| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
+| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
 | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
 | `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |

docs/design/huggingface_integration.md

@@ -21,7 +21,7 @@ Let's say we want to serve the popular Qwen model by running `vllm serve Qwen/Qw
 
 Beyond that, there are two more things vLLM depends on Hugging Face for.
 
-1. **Tokenizer**: vLLM uses the tokenizer from Hugging Face to tokenize the input text. The tokenizer is loaded using [AutoTokenizer.from_pretrained](https://huggingface.co/docs/transformers/en/model_doc/auto#transformers.AutoTokenizer.from_pretrained) with the `model` argument as the model name and the `--revision` argument as the revision. It is also possible to use a tokenizer from another model by specifying the `--tokenizer` argument in the `vllm serve` command. Other relevant arguments are `--tokenizer-revision` and `--tokenizer-mode`. Setting `--tokenizer-mode fastokens` swaps in a drop-in Rust BPE backend for the HF fast tokenizer (see [fastokens Tokenizer Mode](../configuration/optimization.md#fastokens-tokenizer-mode)). Please check Hugging Face's documentation for the meaning of these arguments. This part of the logic can be found in the [get_tokenizer](https://github.com/vllm-project/vllm/blob/127c07480ecea15e4c2990820c457807ff78a057/vllm/transformers_utils/tokenizer.py#L87) function. After obtaining the tokenizer, notably, vLLM will cache some expensive attributes of the tokenizer in [vllm.tokenizers.hf.get_cached_tokenizer][].
+1. **Tokenizer**: vLLM uses the tokenizer from Hugging Face to tokenize the input text. The tokenizer is loaded using [AutoTokenizer.from_pretrained](https://huggingface.co/docs/transformers/en/model_doc/auto#transformers.AutoTokenizer.from_pretrained) with the `model` argument as the model name and the `--revision` argument as the revision. It is also possible to use a tokenizer from another model by specifying the `--tokenizer` argument in the `vllm serve` command. Other relevant arguments are `--tokenizer-revision` and `--tokenizer-mode`. Setting `VLLM_USE_FASTOKENS=1` swaps in a drop-in Rust BPE backend for any HF fast tokenizer loaded by vLLM (see [fastokens Backend](../configuration/optimization.md#fastokens-backend)). Please check Hugging Face's documentation for the meaning of these arguments. This part of the logic can be found in the [get_tokenizer](https://github.com/vllm-project/vllm/blob/127c07480ecea15e4c2990820c457807ff78a057/vllm/transformers_utils/tokenizer.py#L87) function. After obtaining the tokenizer, notably, vLLM will cache some expensive attributes of the tokenizer in [vllm.tokenizers.hf.get_cached_tokenizer][].
 
 2. **Model weight**: vLLM downloads the model weight from the Hugging Face model hub using the `model` argument as the model name and the `--revision` argument as the revision. vLLM provides the argument `--load-format` to control what files to download from the model hub. By default, it will try to load the weights in the safetensors format and fall back to the PyTorch bin format if the safetensors format is not available. We can also pass `--load-format dummy` to skip downloading the weights.
     - It is recommended to use the safetensors format, as it is efficient for loading in distributed inference and also safe from arbitrary code execution. See the [documentation](https://huggingface.co/docs/safetensors/en/index) for more information on the safetensors format. This part of the logic can be found [here](https://github.com/vllm-project/vllm/blob/10b67d865d92e376956345becafc249d4c3c0ab7/vllm/model_executor/model_loader/loader.py#L385). Please note that:

docs/design/nixl_kv_cache_lease.md

@@ -0,0 +1,136 @@
+# NIXL KV Cache Lease Renewal
+
+In disaggregated prefill/decode deployments, the Prefill instance (P) must hold KV cache blocks in GPU memory after completing a prefill, waiting for the Decode instance (D) to read them via RDMA. A mechanism is needed to determine when those blocks can safely be freed when D isn't able to retrieve them. This mechanism was introduced in [PR #41383](https://github.com/vllm-project/vllm/pull/41383).
+
+## Motivation
+
+### The single-timeout problem
+
+The original design used a single, large timeout (`VLLM_NIXL_ABORT_REQUEST_TIMEOUT`, default 480s) to control how long P retains KV blocks. When D crashed or disconnected, P would hold onto potentially several GBs of "dead" blocks for up to 8 minutes before reclaiming them. During this window, subsequent requests hitting P would find reduced cache capacity and experience degraded performance.
+
+### The overloading problem
+
+Simply lowering the timeout introduces a different failure mode. Under traffic surges, requests can sit in D's waiting queue for a long time before being scheduled. If the fixed timeout on P is too short, blocks get freed before D ever has a chance to read them --- causing unnecessary recomputation and wasted prefill work.
+
+### Solution: lease renewal via heartbeats
+
+The lease renewal mechanism addresses both problems simultaneously. P grants a **short initial lease** (default 30s) when prefill completes. While a request is **queued or in-flight** on D, D **periodically sends heartbeats** to P extending the lease. If D crashes and stops heartbeating, P reclaims blocks within seconds of the last heartbeat rather than waiting minutes. If D is merely overloaded, the heartbeats keep the blocks alive for as long as needed.
+
+## How It Works
+
+### Lease lifecycle
+
+When P finishes a prefill, it pins the KV blocks with an initial lease duration (`kv_lease_duration`, default 30s). From that point, the blocks are held until either:
+
+1. **D completes the KV transfer** --- P receives a read-completion notification and frees the blocks immediately.
+2. **D keeps heartbeating** --- each heartbeat extends the lease by `lease_duration * 2/3` (~20s), keeping blocks alive indefinitely while D is healthy.
+3. **No heartbeat arrives** --- the lease expires and P reclaims the blocks.
+
+### Piggybacking on NIXL notifications
+
+Rather than introducing a new transport channel, heartbeats reuse NIXL's existing notification system (`send_notif` / `get_new_notifs`). The notification medium is backend-specific, with automatic fallback from IB/RoCE to TCP already handled by NIXL. Each single heartbeat message sent from D to a particular P renews all requests pinned in P on behalf of that D --- in other words, a single batched message per iteration renews the lease of multiple requests.
+
+### Scheduler-side tracking (D)
+
+A critical insight is that heartbeating must start **as soon as a request enters D's scheduler** --- not when it gets scheduled for execution. Under heavy load, a request may sit in the waiting queue for much longer than the initial lease duration, and the gap between arrival and scheduling is unbounded.
+
+To achieve this, D's connector (`NixlConnectorScheduler`) hooks into the scheduler via `on_new_request()`. When a request with `do_remote_prefill=True` arrives, the connector immediately starts tracking it for heartbeats. Requests are grouped by `remote_engine_id` for efficient batching. On each scheduler step, heartbeat metadata is packaged into `NixlConnectorMetadata` and sent to the worker, throttled by a heartbeat interval of `lease_duration // 6` (~5s).
+
+Tracking stops when either the KV transfer completes (via `update_connector_output`) or the request finishes/aborts (via `request_finished`).
+
+### Timing and simplicity
+
+Heartbeat sending and processing happen **in the forward loop**, not in a background thread. This means timing is not millisecond-precise --- a long model forward pass will delay heartbeats. However, the lease durations are configured with sufficient margin: with default settings, the heartbeat interval (~5s) and lease extension (~20s) are at least an order of magnitude larger than a typical forward pass. This avoids lock complexity between threads while keeping the design simple and extensible.
+
+## Happy Path
+
+```mermaid
+sequenceDiagram
+    participant R as Routing Proxy
+    participant P as Prefill Instance
+    participant D as Decode Instance
+
+    R->>P: Request (do_remote_decode=True)
+    P->>P: Run prefill
+    P->>P: Grant lease (30s)
+    P->>R: Response (with kv_transfer_params)
+
+    R->>D: Request (do_remote_prefill=True)
+    note over D: Request enters waiting queue
+    D->>D: on_new_request() starts tracking
+
+    loop Every ~5s (heartbeat interval)
+        D->>P: Heartbeat (extend lease)
+        P->>P: Lease extended by ~20s
+    end
+
+    note over D: Request scheduled for execution
+    D->>P: KV transfer (RDMA read)
+    P-->D: Transfer complete
+    D->>D: Stop heartbeating
+    P->>P: Free KV blocks
+```
+
+## Decode Instance Crash
+
+```mermaid
+sequenceDiagram
+    participant R as Routing Proxy
+    participant P as Prefill Instance
+    participant D as Decode Instance
+
+    R->>P: Request (do_remote_decode=True)
+    P->>P: Run prefill (holds onto KVs with lease)
+    P->>R: Response
+
+    R->>D: Request (do_remote_prefill=True)
+    D->>P: Heartbeat (extend lease)
+    D->>P: Heartbeat (extend lease)
+    note over D: D crashes
+    note over P: No heartbeat received
+    P->>P: Lease expires (~20s, not 480s)
+    P->>P: Free KV blocks
+```
+
+### Worker-side sending and receiving
+
+**On D (sending):** During `start_load_kv()` (called every forward pass), the worker reads `metadata.heartbeat_by_engine` and sends batched heartbeat notifications to each remote P engine. If D hasn't yet handshaked with P for a given engine (common for requests still in the waiting queue), it triggers a **proactive handshake** in a background thread.
+The heartbeat is deferred to the next step once the handshake completes --- the early handshake also **speeds up the eventual KV transfer.**
+
+**On P (receiving):** In `_get_new_notifs()`, P's worker checks incoming NIXL notifications. Messages starting with `"HB:"` are routed to `_handle_heartbeat()`, which extends the lease expiry for each referenced request using `max(old_expiry, now + lease_extension)`. This ensures leases are never accidentally shortened.
+
+## Bidirectional KV Transfer
+
+For multi-turn conversations, [bidirectional KV transfer](../features/disagg_prefill.md) allows D to cache KV blocks that P can pull from on subsequent turns. Since the timing of the next conversational turn is **client-dependent** (not controlled by the system), the heartbeat-based lease mechanism does not apply here. Instead, a separate `decoder_kv_blocks_ttl` (default 480s) provides a simple fixed timeout for blocks cached on D. If the client takes too long to continue the conversation, the blocks expire and P recomputes. Future work may extend a symmetric heartbeat mechanism to this case.
+
+## Key Design Decisions
+
+- **Per-request leasing, not per-instance.** P has no notion of which D its KV blocks belong to --- block ownership is only resolved after prefill completes and the router selects a D. Leasing at the request level avoids coupling P/D selection in the load balancer. In practice, D batches lease extensions toward the same P by grouping requests with the same `remote_engine_id`.
+
+- **NIXL notifications as transport.** Heartbeats reuse the existing `send_notif`/`get_new_notifs` system rather than adding ZMQ connections or API changes. The notification medium is backend-specific with IB/RoCE-to-TCP fallback already handled, making heartbeats work across any NIXL-supported transport.
+
+- **No background thread.** Heartbeat sending and processing happen in the forward loop (`start_load_kv` / `get_finished`). This avoids lock complexity between threads. Lease durations provide sufficient margin over forward-pass latency (seconds vs. milliseconds).
+
+- **Proactive handshake.** When D needs to heartbeat a P engine it hasn't connected to yet (common for requests still in the waiting queue), it triggers an early handshake in a background thread. This also speeds up the eventual KV transfer.
+
+- **Heterogeneous TP support.** When P TP > D TP (e.g., P TP=4, D TP=2), a single D worker pulls from multiple P workers. Heartbeats must be sent to all P workers for a given engine. Conversely, when D TP > P TP, a single P receives notifications from multiple Ds, which simply refreshes the TTL multiple times with no downside.
+
+## Configuration
+
+The lease mechanism is controlled through `kv_connector_extra_config` in `--kv-transfer-config`:
+
+| Parameter               | Default | Description                                                                                                                                                   |
+|-------------------------|---------|---------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| `kv_lease_duration`     | 30s     | Initial lease duration on P. Heartbeat interval and extension amount are derived automatically (`interval = duration // 6`, `extension = duration * 2 // 3`). |
+| `decoder_kv_blocks_ttl` | 480s    | TTL for KV blocks cached on D in bidirectional transfer mode. Simple fixed timeout, not renewed via heartbeats.                                               |
+
+```bash
+vllm serve <MODEL> \
+  --kv-transfer-config '{
+    "kv_connector": "NixlConnector",
+    "kv_role": "kv_both",
+    "kv_connector_extra_config": {"kv_lease_duration": 60}
+  }'
+```
+
+For full NixlConnector configuration details, see the [NixlConnector Usage Guide](../features/nixl_connector_usage.md).

docs/features/mooncake_store_connector_usage.md

@@ -204,7 +204,6 @@ the vLLM JSON config.
 
 - `load_async` (bool): Enable asynchronous loading for better compute-I/O overlap. Default: `true`.
 - `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`.
-- `discard_partial_chunks` (bool): Discard partial block chunks during store. Default: `true`.
 - `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`.
 
 ## Notes

docs/features/nixl_connector_usage.md

@@ -136,6 +136,170 @@ python tests/v1/kv_connector/nixl_integration/toy_proxy_server.py \
     - In bidirectional mode, the decoder caches KV blocks for multi-turn conversations. This TTL controls how long those blocks are held before being released. Unlike the prefiller lease, this TTL is not renewed via heartbeats.
     - Example: `--kv-transfer-config '{"kv_connector_extra_config": {"decoder_kv_blocks_ttl": 600}}'`
 
+## Bidirectional KV Transfer (Multi-turn)
+
+In standard disaggregated prefilling, KV cache flows in one direction: Prefill (P) computes the KV cache and Decode (D) reads from P. For multi-turn conversations this is wasteful — D already holds the KV cache corresponding to the generated tokens from prior turns, yet P must recompute it from scratch on every new turn. Bidirectional KV transfer lets P **pull** existing KV blocks from D via RDMA before computing only the new tokens, significantly reducing Time-To-First-Token (TTFT) for long-prefill such as **multi-turn heavy scenarios**.
+
+### How it works
+
+The feature relies on a **stateful proxy** that sits between the client and the P/D instances. The proxy tracks `kv_transfer_params` returned by D at the end of each turn, and attaches them to the next turn's request so P knows which blocks to pull from D.
+
+```mermaid
+sequenceDiagram
+    participant Client
+    participant Proxy
+    participant P as Prefill (P)
+    participant D as Decode (D)
+
+    rect rgb(240, 240, 250)
+    note right of Client: Turn 1 — Cache Miss
+    Client->>Proxy: chat request + conversation_id
+    Proxy->>P: request (no remote blocks)
+    activate P
+    note over P: full prefill
+    P-->>Proxy: kv_transfer_params (P's blocks)
+    deactivate P
+    Proxy->>D: request + P's kv_transfer_params
+    activate D
+    D-->P: RDMA read (D pulls KV from P)
+    note over D: decode
+    D-->>Proxy: stream response + kv_transfer_params
+    deactivate D
+    note over Proxy: cache D's kv_transfer_params
+    Proxy-->>Client: response
+    end
+
+    rect rgb(255, 245, 235)
+    note right of Client: Turn 2+ — Cache Hit (Bidirectional)
+    Client->>Proxy: chat request + conversation_id
+    note over Proxy: lookup cached D blocks
+    Proxy->>P: request + D's remote_block_ids
+    activate P
+    P-->D: RDMA read (P pulls KV from D)
+    note over P: prefill new tokens only
+    P-->>Proxy: kv_transfer_params (P's blocks)
+    deactivate P
+    Proxy->>D: request + P's kv_transfer_params
+    activate D
+    D-->P: RDMA read (D pulls new KV from P)
+    note over D: decode
+    D-->>Proxy: stream response + kv_transfer_params
+    deactivate D
+    note over Proxy: update cached kv_transfer_params
+    Proxy-->>Client: response
+    end
+```
+
+**Turn 1 (cache miss):**
+
+1. Client sends a chat request with a `conversation_id` to the proxy.
+2. Proxy forwards the request to P with no remote block info — P computes the full KV cache.
+3. Proxy forwards the request to D along with P's `kv_transfer_params` (block IDs, engine ID, host/port).
+4. D reads KV blocks from P via RDMA (peer-to-peer pull), then generates the response.
+5. D streams the response back through the proxy. The final chunk includes D's own `kv_transfer_params`.
+6. Proxy caches D's `kv_transfer_params` keyed by `conversation_id`, then returns the response to the client.
+
+**Turn 2+ (cache hit — bidirectional):**
+
+1. Client sends the next turn with the same `conversation_id`.
+2. Proxy looks up cached `kv_transfer_params` from the previous turn and attaches D's `remote_block_ids` to the request sent to P.
+3. P reads the existing KV cache from D via RDMA (D→P pull), then computes KV only for the new tokens.
+4. Proxy forwards the request to D with P's updated `kv_transfer_params`.
+5. D reads the new KV blocks from P, generates the response, and returns updated `kv_transfer_params` which the proxy caches for the next turn.
+
+### Configuration
+
+Enable bidirectional KV transfer by setting `bidirectional_kv_xfer` in `kv_connector_extra_config` on **both** P and D instances:
+
+```bash
+vllm serve <MODEL> \
+  --kv-transfer-config '{
+    "kv_connector": "NixlConnector",
+    "kv_role": "kv_both",
+    "kv_connector_extra_config": {
+      "bidirectional_kv_xfer": true
+    }
+  }'
+```
+
+Additional configuration options in `kv_connector_extra_config`:
+
+| Parameter | Default | Description |
+| --------- | ------- | ----------- |
+| `bidirectional_kv_xfer` | `false` | Enable bidirectional D→P KV transfer. |
+| `kv_recompute_threshold` | `64` | Minimum number of remote tokens required to trigger a D→P pull. Below this threshold, P recomputes locally instead of pulling (to amortize transfer latency). |
+| `decoder_kv_blocks_ttl` | `480` | TTL (seconds) for KV blocks cached on D for bidirectional reuse. Blocks are released after this duration. Not renewed via heartbeats. |
+
+### Multi-turn proxy setup
+
+Use the provided multi-turn proxy to manage `kv_transfer_params` caching across conversation turns:
+
+```bash
+python examples/disaggregated/disaggregated_serving/disagg_proxy_multiturn.py \
+  --host 0.0.0.0 --port 8000 \
+  --prefiller-host <P_IP> --prefiller-port 8100 \
+  --decoder-host <D_IP> --decoder-port 8200
+```
+
+The proxy supports multiple P and D instances via round-robin:
+
+```bash
+python examples/disaggregated/disaggregated_serving/disagg_proxy_multiturn.py \
+  --host 0.0.0.0 --port 8000 \
+  --prefiller-hosts <P_IP1> <P_IP2> --prefiller-ports 8100 8100 \
+  --decoder-hosts <D_IP1> <D_IP2> --decoder-ports 8200 8200
+```
+
+### Client usage
+
+Include a `conversation_id` field in the request body to enable cross-turn KV reuse. Without it, the proxy cannot link turns and falls back to full recomputation.
+
+```bash
+# Turn 1
+curl http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Qwen/Qwen3-0.6B",
+    "conversation_id": "session-42",
+    "messages": [
+      {"role": "user", "content": "What is vLLM?"}
+    ]
+  }'
+
+# Turn 2 — same conversation_id triggers bidirectional KV pull
+curl http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Qwen/Qwen3-0.6B",
+    "conversation_id": "session-42",
+    "messages": [
+      {"role": "user", "content": "What is vLLM?"},
+      {"role": "assistant", "content": "vLLM is a high-throughput LLM serving engine..."},
+      {"role": "user", "content": "How does disaggregated prefilling work?"}
+    ]
+  }'
+```
+
+!!! note
+    The `conversation_id` field is a non-standard extension to the OpenAI API. It is consumed by the proxy and not forwarded to the vLLM engine.
+
+### Limitations
+
+- Requires a stateful proxy (or equivalent router) to track and forward `kv_transfer_params` between turns.
+- Currently supported on CUDA with device-buffer KV cache. Host-buffer support (e.g., for Intel XPU) is planned for future work.
+
+!!! warning "Reasoning models with stripped thinking traces"
+    When using reasoning models (e.g. DeepSeek-R1) that produce thinking traces
+    (`<think>...</think>`), D's KV blocks cover the full token sequence including
+    thinking tokens. If the client strips thinking traces from the conversation
+    history before sending the next turn, the prompt P receives will be missing
+    tokens from the middle of what D generated. The block-alignment logic assumes
+    P's prompt is a prefix of D's sequence, so pulling KV blocks from D in this
+    case transfers cache computed for the wrong token positions, producing
+    incorrect results.
+
+    We currently assume the router is able to detect such mismatch across turns. See [#43094](https://github.com/vllm-project/vllm/issues/43094). 
+
 ## Multi-Instance Setup
 
 ### Multiple Prefiller Instances on Different Machines
@@ -211,6 +375,10 @@ The `kv_load_failure_policy` setting controls how the system handles failures wh
 !!! warning
     Using `kv_load_failure_policy="recompute"` can lead to performance degradation in production deployments. When KV loads fail, the decode instance will execute prefill work with decode-optimized configurations, which is inefficient and defeats the purpose of disaggregated prefilling. This also increases tail latency for other ongoing decode requests.
 
+### For NVIDIA GB-series GPUs
+
+GB-series GPUs support multi-node NVLink. NIXL supports this capability, but KVCache must be registered as VMM during KVCache registration. To enable this feature, you need to set `--enable-cumem-allocator` or `--enable-sleep-mode` flags, and set `UCX_CUDA_IPC_ENABLE_MNNVL: 'y'` env. Otherwise, NIXL can only use RDMA/TCP for cross-node KVCache transfers.
+
 ## Experimental Feature
 
 ### Heterogeneous KV Layout support

docs/features/speculative_decoding/README.md

@@ -15,6 +15,7 @@ vLLM supports a variety of methods of speculative decoding. Model-based methods
 - [Multi-Layer Perceptron](mlp.md)
 - [N-Gram](n_gram.md)
 - [Suffix Decoding](suffix.md)
+- [Hidden State Extraction](extract_hidden_states.md)
 - [Custom Proposer Backend (Experimental)](#custom-proposer-backend-experimental)
 
 ## Method Selection at a Glance

docs/features/speculative_decoding/extract_hidden_states.md

@@ -0,0 +1,86 @@
+# Hidden State Extraction
+
+The Hidden State Extraction feature allows vLLM to save intermediate layer activations from a target model during inference. This is useful for training [EAGLE](eagle.md)-style draft models, knowledge distillation, or offline analysis of model internals.
+
+!!! note
+    It is possible to save the last-layer's output hidden states by passing `num_hidden_layers` as a layer id. Note that these are _not_ normalized using the output norm.
+
+## Offline Example
+
+```python
+import tempfile
+
+from vllm import LLM, SamplingParams
+from vllm.config.kv_transfer import KVTransferConfig
+from vllm.distributed.kv_transfer.kv_connector.v1 import (
+    example_hidden_states_connector,
+)
+
+with tempfile.TemporaryDirectory() as tmpdir:
+    llm = LLM(
+        model="Qwen/Qwen3-8B",
+        enable_chunked_prefill=False,
+        speculative_config={
+            "method": "extract_hidden_states",
+            "num_speculative_tokens": 1,
+            "draft_model_config": {
+                "hf_config": {
+                    "eagle_aux_hidden_state_layer_ids": [1, 2, 3, 4],
+                },
+            },
+        },
+        kv_transfer_config=KVTransferConfig(
+            kv_connector="ExampleHiddenStatesConnector",
+            kv_role="kv_producer",
+            kv_connector_extra_config={
+                "shared_storage_path": tmpdir,
+            },
+        ),
+    )
+
+    outputs = llm.generate(
+        ["The future of AI is"],
+        SamplingParams(max_tokens=1),
+    )
+
+    for output in outputs:
+        path = output.kv_transfer_params["hidden_states_path"]
+        obj = example_hidden_states_connector.load_hidden_states(path)
+        print(f"token_ids: {obj['token_ids'].shape}")
+        print(f"hidden_states: {obj['hidden_states'].shape}")
+```
+
+A complete example is available at [`examples/features/speculative_decoding/extract_hidden_states_offline.py`](../../../examples/features/speculative_decoding/extract_hidden_states_offline.py).
+
+## Online Example
+
+For improved performance, it is recommended to use a RAM-mounted file system such as `/dev/shm/` for online usage in which the client cleans up the files soon after they are generated.
+
+```bash
+vllm serve Qwen/Qwen3-8B \
+    --speculative_config '{"method": "extract_hidden_states", "num_speculative_tokens": 1, "draft_model_config": {"hf_config": {"eagle_aux_hidden_state_layer_ids": [1, 2, 3, 4]}}}' \
+    --kv_transfer_config '{"kv_connector": "ExampleHiddenStatesConnector", "kv_role": "kv_producer", "kv_connector_extra_config": {"shared_storage_path": "/dev/shm/hidden_states"}}' \
+    --no-enable-chunked-prefill
+```
+
+## Configuration
+
+The `kv_connector_extra_config` dict accepts these options:
+
+| Parameter | Default | Description |
+| --- | --- | --- |
+| `shared_storage_path` | `/tmp` | Directory where hidden state files are saved |
+| `num_writer_threads` | `8` | Thread pool size for async disk writes |
+| `use_synchronization_lock` | `True` | Use file locks so concurrent readers block until writes complete. Can be disabled for batch generation where synchronization is not needed. |
+
+## Output Format
+
+Each request produces a `.safetensors` file containing:
+
+- **`hidden_states`** — shape `[num_tokens, num_extracted_layers, hidden_size]`
+- **`token_ids`** — shape `[num_tokens]`
+
+The file path is returned in `output.kv_transfer_params["hidden_states_path"]`. Use `load_hidden_states()` from the connector module to read the file with proper synchronization.
+
+!!! note
+    Chunked prefill is not compatible with this feature and must be disabled.