Apply pre-commit formatting and use torch.accelerator API

- Apply ruff, clang-format formatting fixes
- Replace torch.cuda.set_device/device_count/synchronize with
  torch.accelerator equivalents per project convention

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Alexander Matveev <alexm-redhat@dgx-b200-02.mgmt.accl-001.lab.rdu2.dc.redhat.com>

Signed-off-by:  <>

.buildkite/hardware_tests/cpu.yaml

@@ -91,7 +91,7 @@ steps:
   - tests/quantization/test_cpu_wna16.py
   commands:
     - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
       pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs
       pytest -x -v -s tests/quantization/test_cpu_wna16.py"
       

.buildkite/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8.yaml

@@ -6,9 +6,7 @@ tasks:
     value: 0.7142
   - name: "exact_match,flexible-extract"
     value: 0.4579
-env_vars:
-  VLLM_USE_FLASHINFER_MOE_FP8: "1"
-  VLLM_FLASHINFER_MOE_BACKEND: "throughput"
+moe_backend: "flashinfer_cutlass"
 limit: 1319
 num_fewshot: 5
 max_model_len: 262144

.buildkite/lm-eval-harness/test_lm_eval_correctness.py

@@ -68,6 +68,10 @@ def launch_lm_eval(eval_config, tp_size):
     if current_platform.is_rocm() and "Nemotron-3" in eval_config["model_name"]:
         model_args += "attention_backend=TRITON_ATTN"
 
+    moe_backend = eval_config.get("moe_backend", None)
+    if moe_backend is not None:
+        model_args += f"moe_backend={moe_backend},"
+
     env_vars = eval_config.get("env_vars", None)
     with scoped_env_vars(env_vars):
         results = lm_eval.simple_evaluate(

.buildkite/release-pipeline.yaml

@@ -1,12 +1,25 @@
 # CUDA architecture lists — following PyTorch RELEASE.md
 # (https://github.com/pytorch/pytorch/blob/main/RELEASE.md)
 # SM86 included for broader Ampere coverage; SM89 for marlin fp8 support
+# These requested arches are filtered by CMake's CUDA_SUPPORTED_ARCHS before
+# per-kernel arch selection. Do not add +PTX here: top-level +PTX is stripped
+# during that filtering, so kernels that need PTX must request it locally.
 env:
-  CUDA_ARCH_X86: "7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX"
-  # aarch64 only architectures: 8.7 for Orin, 11.0 for Thor (since CUDA 13)
-  CUDA_ARCH_AARCH64: "8.0 8.7 8.9 9.0 10.0 11.0 12.0+PTX"
+  # for CUDA >=13, sm_100+ targets have family specifiers (see CMakeLists.txt)
+  # so targets like 10.3 and 12.1 are automatically supported with this list
+  CUDA_ARCH_X86: "7.5 8.0 8.6 8.9 9.0 10.0 12.0"
+  # aarch64-only targets: Orin (8.7), Thor (11.0, CUDA 13+)
+  CUDA_ARCH_AARCH64: "8.0 8.7 8.9 9.0 10.0 11.0 12.0"
+
+  # for CUDA <13, we need to specify all needed targets
+  # some targets (10.3, 12.1) are skipped to limit the wheel size (< 500MB)
+  # please use CUDA 13 wheels or compile yourself on these new devices
   CUDA_ARCH_X86_CU129: "7.5 8.0 8.6 8.9 9.0 10.0 12.0"
   CUDA_ARCH_AARCH64_CU129: "8.0 8.7 8.9 9.0 10.0 12.0"
+  
+  # pre-built mooncake wheels
+  # the manylinux_2_35 wheel has compatibility issue on Ubuntu 24.04
+  # so we use different wheels for the time being
   MOONCAKE_WHEEL_AARCH64_2_35: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_aarch64.whl"
   MOONCAKE_WHEEL_AARCH64_2_39: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_39_aarch64.whl"
   MOONCAKE_WHEEL_X86_64: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_x86_64.whl"
@@ -846,7 +859,6 @@ steps:
         allow_failure: true
       - step: build-cpu-release-image-arm64
         allow_failure: true
-    if: build.env("NIGHTLY") != "1"
 
   - label: "Publish release images to DockerHub"
     depends_on:

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

@@ -13,5 +13,8 @@ INPUT_FILE="$1"
 # Strip timestamps
 sed -i 's/^\[[0-9]\{4\}-[0-9]\{2\}-[0-9]\{2\}T[0-9]\{2\}:[0-9]\{2\}:[0-9]\{2\}Z\] //' "$INPUT_FILE"
 
+# Strip Buildkite inline timestamp markers (ESC _bk;t=<ms> BEL)
+sed -i 's/\x1B_bk;t=[0-9]*\x07//g' "$INPUT_FILE"
+
 # Strip colorization
 sed -i -r 's/\x1B\[[0-9;]*[mK]//g' "$INPUT_FILE"

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

@@ -1,74 +1,178 @@
 #!/bin/bash
-# Usage: ./ci-fetch-log.sh <buildkite_job_url> [output_file]
-#        ./ci-fetch-log.sh <build_number> <job_uuid> [output_file]
+# Fetch vLLM Buildkite CI logs (public; no login required).
 #
-# Downloads the raw log for a Buildkite job from the public, unauthenticated
-# /organizations/<org>/pipelines/<pipeline>/builds/<n>/jobs/<uuid>/download
-# endpoint, then strips ANSI/timestamps via ci-clean-log.sh.
+# Usage:
+#   ci-fetch-log.sh [--soft|--all] --pr [<PR>]  failed jobs in the PR's latest
+#                                               build (current branch if omitted)
+#   ci-fetch-log.sh [--soft|--all] <build_url>  failed jobs in that build
+#   ci-fetch-log.sh <job_url> [output]          one job; both #<job_uuid> and
+#                                               ?sid=<id> URL forms work
+#   ci-fetch-log.sh <build> <job_uuid> [output]
 #
-# Find <build_number> and <job_uuid> via:
-#   gh pr checks <PR> --repo vllm-project/vllm
-# Each failing row's URL is .../builds/<build_number>#<job_uuid>.
-#
-# Default output path: ci-<build>-<uuid_first_13_chars>.log (e.g.
-# ci-68478-019e6b07-daae.log). Jobs in the same build share the UUID's
-# first 8 chars, so the second segment is needed for uniqueness when
-# fetching multiple jobs in parallel. The script refuses to overwrite an
-# existing output file; pass an explicit path or set CI_FETCH_LOG_FORCE=1
-# to override.
+# --soft also fetches soft-failed jobs; --all fetches every finished job.
+# Saves each log as ci-<build>-<job-name>.log (ANSI/timestamps stripped) and
+# prints "<file>\t<job name>" per job. [output] is single-job only; "-"
+# streams to stdout. Existing files are kept; CI_FETCH_LOG_FORCE=1 refetches.
 
 set -euo pipefail
 
 ORG="vllm"
 PIPELINE="ci"
+UA="vllm-ci-fetch-log"
+UUID_RE='[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}'
 
 usage() {
-    echo "Usage: $0 <buildkite_job_url> [output_file]"
-    echo "       $0 <build_number> <job_uuid> [output_file]"
+    sed -n '2,15p' "$0" | sed 's/^# \{0,1\}//'
     exit 1
 }
 
-if [ $# -lt 1 ]; then usage; fi
+die() {
+    echo "$1" >&2
+    exit 1
+}
 
-if [[ "$1" == https://* ]]; then
+BUILD="" JOB="" SID="" OUT=""
+SCOPE="failed"
+
+while :; do
+    case "${1:-}" in
+    --soft) SCOPE="soft" ;;
+    --all) SCOPE="all" ;;
+    *) break ;;
+    esac
+    shift
+done
+
+case "${1:-}" in
+--pr)
+    PR="${2:-}"
+    # gh pr checks exits non-zero when checks are failing; that is the
+    # expected case here.
+    URL=$(gh pr checks ${PR:+"$PR"} --repo vllm-project/vllm 2>/dev/null |
+        grep -oE "https://buildkite.com/${ORG}/${PIPELINE}/builds/[0-9]+" |
+        sort -t/ -k7 -n | tail -1 || true)
+    [ -n "$URL" ] || die "No Buildkite build found via: gh pr checks ${PR:-<current branch>}"
+    BUILD="${URL##*/}"
+    ;;
+https://*)
     BUILD=$(echo "$1" | sed -nE 's#.*/builds/([0-9]+).*#\1#p')
-    JOB=$(echo "$1" | grep -oE '[0-9a-f]{8}-[0-9a-f-]+' | head -n 1)
+    JOB=$(echo "$1" | grep -oE "#${UUID_RE}" | head -n 1 | cut -c2- || true)
+    SID=$(echo "$1" | grep -oE "[?&]sid=${UUID_RE}" | head -n 1 | sed 's/.*sid=//' || true)
     OUT="${2:-}"
-else
-    if [ $# -lt 2 ]; then usage; fi
+    [ -n "$BUILD" ] || die "Could not parse build number from: $1"
+    ;;
+[0-9]*)
+    [ $# -ge 2 ] || usage
     BUILD="$1"
     JOB="$2"
     OUT="${3:-}"
-fi
-
-if [ -z "$BUILD" ] || [ -z "$JOB" ]; then
-    echo "Could not parse build number or job UUID from: $1" >&2
+    ;;
+*)
     usage
-fi
-
-# Jobs in the same build share the UUID's first segment, so include the
-# second segment (chars 9-13, e.g. "019e6b07-daae") to keep default filenames
-# unique when fetching multiple jobs from one build in parallel.
-if [ -z "$OUT" ]; then
-    OUT="ci-${BUILD}-${JOB:0:13}.log"
-fi
-
-if [ -e "$OUT" ] && [ -z "${CI_FETCH_LOG_FORCE:-}" ]; then
-    echo "Refusing to overwrite existing $OUT (set CI_FETCH_LOG_FORCE=1 or pass an explicit output path)." >&2
-    exit 1
-fi
+    ;;
+esac
 
 COOKIES=$(mktemp)
-trap 'rm -f "$COOKIES"' EXIT
+JOBS_TSV=$(mktemp)
+trap 'rm -f "$COOKIES" "$JOBS_TSV"' EXIT
 
-# Buildkite issues a session cookie on first hit; subsequent /download needs it.
-curl -fsSL -c "$COOKIES" -A "vllm-ci-fetch-log" \
+# Buildkite issues a session cookie on first hit; later requests need it.
+curl -fsSL -c "$COOKIES" -A "$UA" \
     "https://buildkite.com/${ORG}/${PIPELINE}/builds/${BUILD}" -o /dev/null
 
-curl -fsSL -b "$COOKIES" -A "vllm-ci-fetch-log" \
-    "https://buildkite.com/organizations/${ORG}/pipelines/${PIPELINE}/builds/${BUILD}/jobs/${JOB}/download" \
-    -o "$OUT"
+# The build's job list (id, step uuid, state, name) is served as JSON from
+# the user-facing /data/jobs endpoint. Flatten it to TSV for easy filtering:
+#   job_id  step_uuid  failed  soft_failed  finished  slug  name
+curl -fsSL -b "$COOKIES" -A "$UA" \
+    "https://buildkite.com/${ORG}/${PIPELINE}/builds/${BUILD}/data/jobs" |
+    python3 -c '
+import json, re, sys
 
-bash "$(dirname "$0")/ci-clean-log.sh" "$OUT"
+data = json.load(sys.stdin)
+if data.get("has_next_page"):
+    print("warning: job list is paginated; some jobs not shown", file=sys.stderr)
+for r in data["records"]:
+    if r.get("type") != "script":
+        continue
+    name = (r.get("name") or "").replace("\t", " ").replace("\n", " ")
+    slug = re.sub(r"[^a-z0-9]+", "-", name.lower()).strip("-")[:60]
+    print("\t".join([
+        r["id"],
+        r.get("step_uuid") or "",
+        str(r.get("passed") is False),
+        str(bool(r.get("soft_failed"))),
+        str(bool(r.get("finished_at"))),
+        slug,
+        name,
+    ]))
+' >"$JOBS_TSV" || die "Could not list jobs for build ${BUILD}"
 
-echo "$OUT"
+if [ -n "$SID" ] && [ -z "$JOB" ]; then
+    # The ?sid= in builds/<N>/list URLs is the *step* uuid, not the job uuid.
+    JOB=$(awk -F'\t' -v s="$SID" '$1 == s || $2 == s {print $1; exit}' "$JOBS_TSV")
+    [ -n "$JOB" ] || die "No job matching sid=${SID} in build ${BUILD}"
+fi
+
+fetch_job() { # <job_uuid> <output_file>
+    curl -fsSL -b "$COOKIES" -A "$UA" \
+        "https://buildkite.com/organizations/${ORG}/pipelines/${PIPELINE}/builds/${BUILD}/jobs/$1/download" \
+        -o "$2"
+    bash "$(dirname "$0")/ci-clean-log.sh" "$2"
+}
+
+if [ -n "$JOB" ]; then
+    # Single-job mode.
+    NAME=$(awk -F'\t' -v j="$JOB" '$1 == j {print $7; exit}' "$JOBS_TSV")
+    SLUG=$(awk -F'\t' -v j="$JOB" '$1 == j {print $6; exit}' "$JOBS_TSV")
+    [ -n "$OUT" ] || OUT="ci-${BUILD}-${SLUG:-${JOB:0:13}}.log"
+    if [ "$OUT" = "-" ]; then
+        TMP=$(mktemp)
+        fetch_job "$JOB" "$TMP"
+        cat "$TMP"
+        rm -f "$TMP"
+        exit 0
+    fi
+    if [ -e "$OUT" ] && [ -z "${CI_FETCH_LOG_FORCE:-}" ]; then
+        die "Refusing to overwrite existing ${OUT} (set CI_FETCH_LOG_FORCE=1 or pass an output path)."
+    fi
+    fetch_job "$JOB" "$OUT"
+    printf '%s\t%s\n' "$OUT" "${NAME:-$JOB}"
+    exit 0
+fi
+
+# Build-wide mode: fetch finished jobs matching $SCOPE.
+[ -z "$OUT" ] || die "[output_file] is only valid when fetching a single job."
+
+case "$SCOPE" in
+failed) FILTER='$3 == "True" && $4 == "False" && $5 == "True"' ;;
+soft) FILTER='$3 == "True" && $5 == "True"' ;;
+all) FILTER='$5 == "True"' ;;
+esac
+
+if [ "$SCOPE" = "failed" ]; then
+    SOFT=$(awk -F'\t' '$3 == "True" && $4 == "True"' "$JOBS_TSV" | wc -l)
+    [ "$SOFT" -eq 0 ] || echo "Skipping ${SOFT} soft-failed job(s); use --soft to include them." >&2
+fi
+
+FOUND=0
+EMITTED=" "
+while IFS=$'\t' read -r job_id _ _ _ _ slug name; do
+    FOUND=$((FOUND + 1))
+    out="ci-${BUILD}-${slug:-${job_id:0:13}}.log"
+    # Retries share a name with the original job; disambiguate by uuid.
+    case "$EMITTED" in
+    *" $out "*) out="ci-${BUILD}-${slug:-job}-${job_id:0:13}.log" ;;
+    esac
+    EMITTED="${EMITTED}${out} "
+    if [ -e "$out" ] && [ -z "${CI_FETCH_LOG_FORCE:-}" ]; then
+        echo "Keeping existing ${out} (set CI_FETCH_LOG_FORCE=1 to refetch)." >&2
+    elif ! fetch_job "$job_id" "$out"; then
+        echo "Failed to download log for job ${job_id} (${name})." >&2
+        continue
+    fi
+    printf '%s\t%s\n' "$out" "$name"
+done < <(awk -F'\t' "$FILTER" "$JOBS_TSV")
+
+if [ "$FOUND" -eq 0 ]; then
+    echo "No matching jobs in build ${BUILD} (scope: ${SCOPE})." >&2
+fi

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

@@ -110,6 +110,36 @@ install_uv() {
     | env UV_INSTALL_DIR="$CARGO_HOME/bin" sh
 }
 
+setup_pyo3_python() {
+  local python_version="${PYO3_PYTHON_VERSION:-3.12}"
+
+  log_section "Installing Python ${python_version} for PyO3 tests"
+  uv python install "$python_version"
+  PYO3_PYTHON="$(uv python find \
+    --managed-python \
+    --no-project \
+    --resolve-links \
+    "$python_version")"
+  export PYO3_PYTHON
+
+  local python_libdir
+  python_libdir="$("$PYO3_PYTHON" - <<'PY'
+import pathlib
+import sysconfig
+
+libdir = pathlib.Path(sysconfig.get_config_var("LIBDIR"))
+ldlibrary = sysconfig.get_config_var("LDLIBRARY")
+assert sysconfig.get_config_var("Py_ENABLE_SHARED") == 1
+assert ldlibrary
+assert (libdir / ldlibrary).exists(), libdir / ldlibrary
+print(libdir)
+PY
+)"
+
+  export LD_LIBRARY_PATH="${python_libdir}:${LD_LIBRARY_PATH:-}"
+  export LIBRARY_PATH="${python_libdir}:${LIBRARY_PATH:-}"
+}
+
 run_style_clippy() {
   install_cargo_sort
 
@@ -132,6 +162,7 @@ run_style_clippy() {
 
 run_tests() {
   install_uv
+  setup_pyo3_python
   install_cargo_nextest
 
   log_section "Running cargo nextest"

.buildkite/test-amd.yaml

@@ -398,7 +398,7 @@ steps:
   - tests/kernels/helion/
   - vllm/platforms/rocm.py
   commands:
-  - pip install helion==1.0.0
+  - pip install helion==1.1.0
   - pytest -v -s kernels/helion/
 
 - label: Kernels Mamba Test # TBD
@@ -2946,7 +2946,6 @@ steps:
   - vllm/model_executor/models/qwen3_5_mtp.py
   - vllm/transformers_utils/configs/qwen3_5.py
   - vllm/transformers_utils/configs/qwen3_5_moe.py
-  - vllm/model_executor/models/qwen.py
   - vllm/model_executor/models/qwen2.py
   - vllm/model_executor/models/qwen3.py
   - vllm/model_executor/models/qwen3_next.py
@@ -3184,7 +3183,6 @@ steps:
   - vllm/model_executor/models/qwen3_5_mtp.py
   - vllm/transformers_utils/configs/qwen3_5.py
   - vllm/transformers_utils/configs/qwen3_5_moe.py
-  - vllm/model_executor/models/qwen.py
   - vllm/model_executor/models/qwen2.py
   - vllm/model_executor/models/qwen3.py
   - vllm/model_executor/models/qwen3_next.py

.buildkite/test_areas/attention.yaml

@@ -15,7 +15,7 @@ steps:
     - pytest -v -s v1/attention
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 70
       depends_on:
       - image-build-amd

.buildkite/test_areas/disaggregated.yaml

@@ -61,6 +61,20 @@ steps:
     - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
+- label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs)
+  key: hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus
+  timeout_in_minutes: 25
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 2
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+    - vllm/v1/core/sched/
+    - vllm/v1/core/kv_cache_coordinator.py
+    - tests/v1/kv_connector/nixl_integration/
+  commands:
+    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+    - bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh
+
 - label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
   key: multiconnector-nixl-offloading-pd-accuracy-2-gpus
   timeout_in_minutes: 30

.buildkite/test_areas/engine.yaml

@@ -28,7 +28,7 @@ steps:
   - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 60
       depends_on:
       - image-build-amd
@@ -44,7 +44,7 @@ steps:
     - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 40
       depends_on:
       - image-build-amd

.buildkite/test_areas/entrypoints.yaml

@@ -46,7 +46,7 @@ steps:
   - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
 
@@ -63,7 +63,7 @@ steps:
   - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 80
       depends_on:
       - image-build-amd
@@ -82,7 +82,7 @@ steps:
   - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 80
       depends_on:
       - image-build-amd
@@ -104,7 +104,7 @@ steps:
   - pytest -v -s entrypoints/anthropic
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 60
       depends_on:
       - image-build-amd
@@ -165,7 +165,7 @@ steps:
   - pytest -s entrypoints/openai/correctness/
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
       source_file_dependencies:

.buildkite/test_areas/kernels.yaml

@@ -75,6 +75,19 @@ steps:
     - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
   parallelism: 2
 
+- label: Kernels Attention DiffKV Test (H100)
+  key: kernels-attention-diffkv-test-h100
+  timeout_in_minutes: 20
+  device: h100
+  num_devices: 1
+  source_file_dependencies:
+  - vllm/v1/attention/ops/triton_unified_attention_diffkv.py
+  - vllm/v1/attention/backends/triton_attn_diffkv.py
+  - vllm/v1/attention/backends/flash_attn_diffkv.py
+  - tests/kernels/attention/test_triton_unified_attention_diffkv.py
+  commands:
+    - pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py
+
 - label: Kernels Quantization Test %N
   key: kernels-quantization-test
   timeout_in_minutes: 90
@@ -87,7 +100,7 @@ steps:
   parallelism: 2
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       source_file_dependencies:
       - csrc/quantization/
       - vllm/model_executor/layers/quantization
@@ -224,7 +237,7 @@ steps:
   - vllm/utils/import_utils.py
   - tests/kernels/helion/
   commands:
-    - pip install helion==1.0.0
+    - pip install helion==1.1.0
     - pytest -v -s kernels/helion/
 
  

.buildkite/test_areas/lm_eval.yaml

@@ -14,7 +14,7 @@ steps:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 55
       depends_on:
       - image-build-amd

.buildkite/test_areas/misc.yaml

@@ -138,11 +138,26 @@ steps:
     - vllm/v1/spec_decode/extract_hidden_states.py
     - vllm/model_executor/models/extract_hidden_states.py
     - vllm/transformers_utils/configs/extract_hidden_states.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
     - tests/v1/kv_connector/extract_hidden_states_integration
   commands:
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pytest -v -s v1/kv_connector/extract_hidden_states_integration
 
+- label: Extract Hidden States Integration (2 GPUs)
+  key: extract-hidden-states-integration-2-gpus
+  timeout_in_minutes: 20
+  num_devices: 2
+  source_file_dependencies:
+    - vllm/v1/spec_decode/extract_hidden_states.py
+    - vllm/model_executor/models/extract_hidden_states.py
+    - vllm/transformers_utils/configs/extract_hidden_states.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
+    - tests/v1/kv_connector/extract_hidden_states_integration
+  commands:
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -v -s -m 'distributed' v1/kv_connector/extract_hidden_states_integration
+
 - label: Regression
   key: regression
   timeout_in_minutes: 20
@@ -293,6 +308,7 @@ steps:
   - vllm/transformers_utils/
   - vllm/utils/
   - vllm/v1/
+  - tests/test_envs.py
   - tests/test_inputs.py
   - tests/test_outputs.py
   - tests/test_pooling_params.py
@@ -300,24 +316,25 @@ steps:
   - tests/multimodal
   - tests/renderers
   - tests/standalone_tests/lazy_imports.py
-  - tests/tokenizers_
   - tests/reasoning
   - tests/tool_parsers
+  - tests/tokenizers_
   - tests/parser
   - tests/transformers_utils
   - tests/config
   device: cpu-small
   commands:
   - python3 standalone_tests/lazy_imports.py
+  - pytest -v -s test_envs.py
   - pytest -v -s test_inputs.py
   - pytest -v -s test_outputs.py
   - pytest -v -s test_pooling_params.py
   - pytest -v -s test_ray_env.py
   - pytest -v -s -m 'cpu_test' multimodal
   - pytest -v -s renderers
-  - pytest -v -s tokenizers_
   - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
   - pytest -v -s tool_parsers
+  - pytest -v -s tokenizers_
   - pytest -v -s parser
   - pytest -v -s transformers_utils
   - pytest -v -s config

.buildkite/test_areas/models_multimodal.yaml

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

.buildkite/test_areas/plugins.yaml

@@ -40,3 +40,17 @@ steps:
   - pytest -v -s plugins_tests/test_oot_registration_online.py # it needs a clean process
   - pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process
   - pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins
+
+
+- label: GGUF Plugin
+  key: gguf-plugin
+  device: h200_18gb
+  timeout_in_minutes: 30
+  soft_fail: true
+  optional: true
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization
+  - tests/plugins_tests/test_gguf_plugin.py
+  commands:
+    - pip install "vllm-gguf-plugin >= 0.0.2"
+    - pytest -v -s plugins_tests/gguf

.buildkite/test_areas/quantization.yaml

@@ -21,6 +21,18 @@ steps:
   - uv pip install --system conch-triton-kernels
   - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
 
+- label: Quantized Fusions
+  key: quantized-fusions
+  timeout_in_minutes: 30
+  source_file_dependencies:
+  - tests/fusion
+  - vllm/model_executor/layers/fusion
+  - vllm/model_executor/kernels/linear
+  - vllm/model_executor/layers/quantization/compressed_tensors
+  - vllm/model_executor/layers/quantization/modelopt.py
+  commands:
+    - pytest -v -s fusion/
+
 - label: Quantized MoE Test (B200)
   key: quantized-moe-test-b200
   timeout_in_minutes: 60

.buildkite/test_areas/spec_decode.yaml

@@ -39,7 +39,7 @@ steps:
     - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 65
       depends_on:
       - image-build-amd
@@ -78,7 +78,7 @@ steps:
     - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 65
       depends_on:
       - image-build-amd
@@ -103,7 +103,7 @@ steps:
     - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 50
       depends_on:
       - image-build-amd

.github/CODEOWNERS

@@ -80,7 +80,7 @@
 /vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery @NickLucche
 
 # Model runner V2
-/vllm/v1/worker/gpu @WoosukKwon @njhill
+/vllm/v1/worker/gpu @WoosukKwon @njhill @yewentao256
 /vllm/v1/worker/gpu/kv_connector.py @orozery
 
 # CI & building

.github/dependabot.yml

@@ -21,7 +21,6 @@ updates:
       - dependency-name: "torchvision"
       - dependency-name: "xformers"
       - dependency-name: "lm-format-enforcer"
-      - dependency-name: "gguf"
       - dependency-name: "compressed-tensors"
       - dependency-name: "ray[cgraph]" # Ray Compiled Graph
       - dependency-name: "lm-eval"

.github/mergify.yml

@@ -144,12 +144,12 @@ pull_request_rules:
     - label != stale
     - or:
       - files~=^examples/.*mistral.*\.py
-      - files~=^tests/.*mistral.*\.py
-      - files~=^vllm/model_executor/models/.*mistral.*\.py
+      - files~=^tests/.*(?:mistral|voxtral|mixtral|pixtral).*\.py
+      - files~=^vllm/model_executor/models/.*(?:mistral|voxtral|mixtral|pixtral).*\.py
       - files~=^vllm/reasoning/.*mistral.*\.py
       - files~=^vllm/tool_parsers/.*mistral.*\.py
-      - files~=^vllm/transformers_utils/.*mistral.*\.py
-      - title~=(?i)Mistral
+      - files~=^vllm/transformers_utils/.*(?:mistral|voxtral|pixtral).*\.py
+      - title~=(?i)(?:mistral|ministral|voxtral|mixtral|pixtral)
   actions:
     label:
       add:

.github/workflows/scripts/build.sh

@@ -9,7 +9,7 @@ PATH=${cuda_home}/bin:$PATH
 LD_LIBRARY_PATH=${cuda_home}/lib64:$LD_LIBRARY_PATH
 
 # Install requirements
-if [ "$(echo $2 | cut -d. -f1)" = "12" ]; then
+if [ "$(echo "$2" | cut -d. -f1)" = "12" ]; then
     sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' requirements/cuda.txt
 fi
 $python_executable -m pip install -r requirements/build/cuda.txt -r requirements/cuda.txt
@@ -17,7 +17,10 @@ $python_executable -m pip install -r requirements/build/cuda.txt -r requirements
 # Limit the number of parallel jobs to avoid OOM
 export MAX_JOBS=1
 # Make sure release wheels are built for the following architectures
-export TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX"
+# Do not add +PTX here: vLLM filters torch's top-level PTX flag when it
+# converts global gencode flags into per-kernel arch lists. If a specific
+# kernel needs PTX, add +PTX to that kernel's CMake arch list instead.
+export TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9 9.0 10.0 12.0"
 
 bash tools/check_repo.sh
 

.gitignore

@@ -15,6 +15,9 @@ vllm/third_party/flashmla/flash_mla_interface.py
 # DeepGEMM vendored package built from source
 vllm/third_party/deep_gemm/
 
+# fmha_sm100 vendored package built from source
+vllm/third_party/fmha_sm100/
+
 # triton jit
 .triton
 
@@ -233,7 +236,7 @@ actionlint
 shellcheck*/
 
 # Ignore moe/marlin_moe gen code
-csrc/moe/marlin_moe_wna16/kernel_*
+csrc/libtorch_stable/moe/marlin_moe_wna16/kernel_*
 
 # Ignore ep_kernels_workspace folder
 ep_kernels_workspace/

.pre-commit-config.yaml

@@ -21,7 +21,7 @@ repos:
   rev: v21.1.2
   hooks:
   - id: clang-format
-    exclude: 'csrc/(moe/topk_softmax_kernels.cu|libtorch_stable/quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*'
+    exclude: 'csrc/libtorch_stable/moe/topk_softmax_kernels.cu|vllm/third_party/.*'
     types_or: [c++, cuda]
     args: [--style=file, --verbose]
 - repo: https://github.com/DavidAnson/markdownlint-cli2

AGENTS.md

@@ -105,6 +105,26 @@ The line length limit for Python code is 88 characters. If you are not sure, use
 
 Use [Google-style docstrings](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings) (`Args:`/`Returns:`/`Raises:` sections), not reStructuredText/Sphinx fields (`:param:`, `:return:`, `:rtype:`).
 
+### Coding style guidelines
+
+Follow these rules for all code changes in this repository:
+
+- Try to match existing code style.
+- Code should be self-documenting and self-explanatory.
+- Keep comments and docstrings minimal and concise.
+- Assume the reader is familiar with vLLM.
+
+### Diagnosing CI failures
+
+Buildkite logs are public; no login needed. Details: [docs/contributing/ci/failures.md](docs/contributing/ci/failures.md).
+
+```bash
+# All failed-job logs for a PR's latest build (current branch's PR if omitted):
+.buildkite/scripts/ci-fetch-log.sh --pr <PR>
+# Any Buildkite build or job URL also works:
+.buildkite/scripts/ci-fetch-log.sh "<buildkite_url>"
+```
+
 ### Commit messages
 
 Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:

SECURITY.md

@@ -34,6 +34,15 @@ Vulnerabilities that cause denial of service or partial disruption, but do not a
 
 Minor issues such as informational disclosures, logging errors, non-exploitable flaws, or weaknesses that require local or high-privilege access and offer negligible impact. Examples include side channel attacks or hash collisions. These issues often have CVSS scores less than 4.0
 
+## Fix disclosure policy
+
+When a security report is accepted, the fix process depends on the severity:
+
+* **CRITICAL and HIGH severity**: Fixes are developed in a private security fork and coordinated with the prenotification group before public disclosure.
+* **MODERATE and LOW severity**: Fixes are developed and submitted as public pull requests. These issues do not require embargo since they do not enable arbitrary code execution or significant data breach, and public visibility accelerates community review and adoption of the fix.
+
+The vulnerability management team reserves the right to adjust the disclosure approach on a case-by-case basis, taking into account factors such as active exploitation, unusual attack surface, or coordination requirements with downstream vendors.
+
 ## Prenotification policy
 
 For certain security issues of CRITICAL, HIGH, or MODERATE severity level, we may prenotify certain organizations or vendors that ship vLLM. The purpose of this prenotification is to allow for a coordinated release of fixes for severe issues.

benchmarks/attention_benchmarks/README.md

@@ -108,7 +108,6 @@ python benchmark.py \
     --backends flash triton flashinfer \
     --batch-specs "q2k" "8q1s1k" "2q2k_32q1s1k" \
     --num-layers 10 \
-    --repeats 5 \
     --output-csv results.csv
 ```
 
@@ -164,14 +163,17 @@ python benchmark.py \
 # Model configuration
 --num-layers N                      # Number of layers
 --head-dim N                        # Head dimension
+--v-head-dim N                      # Value head dimension (defaults to --head-dim)
 --num-q-heads N                     # Query heads
 --num-kv-heads N                    # KV heads
 --block-size N                      # Block size
+--kv-lora-rank N                    # MLA KV LoRA rank
+--qk-nope-head-dim N                # MLA non-RoPE QK head dim
+--qk-rope-head-dim N                # MLA RoPE QK head dim
 
 # Benchmark settings
 --device DEVICE                     # Device (default: cuda:0)
---repeats N                         # Repetitions
---warmup-iters N                    # Warmup iterations
+--warmup-ms N                       # Warmup window in ms for triton do_bench
 --profile-memory                    # Profile memory usage
 
 # Parameter sweeps
@@ -211,8 +213,6 @@ config = BenchmarkConfig(
     num_kv_heads=1,
     block_size=128,
     device="cuda:0",
-    repeats=5,
-    warmup_iters=3,
 )
 
 # CUTLASS MLA with specific num_kv_splits
@@ -253,14 +253,10 @@ formatter.save_json(results, "output.json")
 
 ## Tips
 
-**1. Warmup matters** - Use `--warmup-iters 10` for stable results
+**1. Save results** - Always use `--output-csv` or `--output-json`
 
-**2. Multiple repeats** - Use `--repeats 20` for low variance
+**2. Test incrementally** - Start with `--num-layers 1`
 
-**3. Save results** - Always use `--output-csv` or `--output-json`
+**3. Extended grammar** - Leverage spec decode, chunked prefill patterns
 
-**4. Test incrementally** - Start with `--num-layers 1 --repeats 1`
-
-**5. Extended grammar** - Leverage spec decode, chunked prefill patterns
-
-**6. Parameter sweeps** - Use `--sweep-param` and `--sweep-values` to find optimal values
+**4. Parameter sweeps** - Use `--sweep-param` and `--sweep-values` to find optimal values

benchmarks/attention_benchmarks/benchmark.py

@@ -26,6 +26,9 @@ Examples:
 """
 
 import argparse
+import os
+import shutil
+import subprocess
 import sys
 from dataclasses import replace
 from pathlib import Path
@@ -83,13 +86,15 @@ def run_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:
         else:
             return run_standard_attention_benchmark(config)
     except Exception as e:
+        error_msg = str(e) or repr(e)
         return BenchmarkResult(
             config=config,
             mean_time=float("inf"),
+            median_time=float("inf"),
             std_time=0,
             min_time=float("inf"),
             max_time=float("inf"),
-            error=str(e),
+            error=error_msg,
         )
 
 
@@ -115,9 +120,12 @@ def run_model_parameter_sweep(
     """
     all_results = []
 
-    console.print(
-        f"[yellow]Model sweep mode: testing {sweep.param_name} = {sweep.values}[/]"
+    sweep_desc = (
+        f"{sweep.param_name} = {sweep.values}"
+        if sweep.param_name
+        else f"{len(sweep.values)} configurations"
     )
+    console.print(f"[yellow]Model sweep mode: testing {sweep_desc}[/]")
 
     total = len(backends) * len(batch_specs) * len(sweep.values)
 
@@ -125,9 +133,9 @@ def run_model_parameter_sweep(
         for backend in backends:
             for spec in batch_specs:
                 for value in sweep.values:
-                    # Create config with modified model parameter
+                    # Create config with modified model parameter(s)
                     config_args = base_config_args.copy()
-                    config_args[sweep.param_name] = value
+                    sweep.apply(config_args, value)
 
                     # Create config with original backend for running
                     clean_config = BenchmarkConfig(
@@ -144,13 +152,21 @@ def run_model_parameter_sweep(
                     all_results.append(result)
 
                     if not result.success:
+                        err_label = (
+                            f"{sweep.param_name}={value}"
+                            if sweep.param_name
+                            else f"{value}"
+                        )
                         console.print(
-                            f"[red]Error {backend} {spec} {sweep.param_name}="
-                            f"{value}: {result.error}[/]"
+                            f"[red]Error {backend} {spec} {err_label}"
+                            f": {result.error}[/]"
                         )
 
                     pbar.update(1)
 
+    if base_config_args.get("ncu_profile"):
+        return all_results
+
     # Display sweep results - create separate table for each parameter value
     console.print("\n[bold green]Model Parameter Sweep Results:[/]")
     formatter = ResultsFormatter(console)
@@ -184,7 +200,10 @@ def run_model_parameter_sweep(
     )
 
     for param_value in sorted_param_values:
-        console.print(f"\n[bold cyan]{sweep.param_name} = {param_value}[/]")
+        label = (
+            f"{sweep.param_name} = {param_value}" if sweep.param_name else param_value
+        )
+        console.print(f"\n[bold cyan]{label}[/]")
         param_results = by_param_value[param_value]
 
         # Create modified results with original backend names
@@ -200,8 +219,9 @@ def run_model_parameter_sweep(
         formatter.print_table(modified_results, backends, compare_to_fastest=True)
 
     # Show optimal backend for each (param_value, batch_spec) combination
+    sweep_name = sweep.param_name or "config"
     console.print(
-        f"\n[bold cyan]Optimal backend for each ({sweep.param_name}, batch_spec):[/]"
+        f"\n[bold cyan]Optimal backend for each ({sweep_name}, batch_spec):[/]"
     )
 
     # Group by (param_value, batch_spec)
@@ -236,7 +256,10 @@ def run_model_parameter_sweep(
     for param_value, spec in sorted_keys:
         # Print header when param value changes
         if param_value != current_param_value:
-            console.print(f"\n  [bold]{sweep.param_name}={param_value}:[/]")
+            header = (
+                f"{sweep.param_name}={param_value}" if sweep.param_name else param_value
+            )
+            console.print(f"\n  [bold]{header}:[/]")
             current_param_value = param_value
 
         results = by_param_and_spec[(param_value, spec)]
@@ -322,6 +345,9 @@ def run_parameter_sweep(
 
                     pbar.update(1)
 
+    if base_config_args.get("ncu_profile"):
+        return all_results
+
     # Display sweep results
     console.print("\n[bold green]Sweep Results:[/]")
     backend_labels = [sweep.get_label(b, v) for b in backends for v in sweep_values]
@@ -474,11 +500,35 @@ def main():
     parser.add_argument("--num-q-heads", type=int, default=32, help="Query heads")
     parser.add_argument("--num-kv-heads", type=int, default=8, help="KV heads")
     parser.add_argument("--block-size", type=int, default=16, help="Block size")
+    parser.add_argument(
+        "--v-head-dim",
+        type=int,
+        default=None,
+        help="Value head dimension (defaults to --head-dim if unset)",
+    )
+
+    # MLA-specific model dimensions
+    parser.add_argument(
+        "--kv-lora-rank", type=int, default=None, help="MLA KV LoRA rank"
+    )
+    parser.add_argument(
+        "--qk-nope-head-dim", type=int, default=None, help="MLA non-RoPE QK head dim"
+    )
+    parser.add_argument(
+        "--qk-rope-head-dim", type=int, default=None, help="MLA RoPE QK head dim"
+    )
 
     # Benchmark settings
     parser.add_argument("--device", default="cuda:0", help="Device")
-    parser.add_argument("--repeats", type=int, default=1, help="Repetitions")
-    parser.add_argument("--warmup-iters", type=int, default=3, help="Warmup iterations")
+    parser.add_argument(
+        "--warmup-ms",
+        type=int,
+        default=None,
+        help=(
+            "Warmup window in ms for triton's do_bench (default: triton's own). "
+            "Has no effect with CUDA graphs; pass --no-cuda-graphs to use it."
+        ),
+    )
     parser.add_argument("--profile-memory", action="store_true", help="Profile memory")
     parser.add_argument(
         "--kv-cache-dtype",
@@ -491,10 +541,33 @@ def main():
         action=argparse.BooleanOptionalAction,
         default=True,
         help=(
-            "Launch kernels with CUDA graphs to eliminate CPU overhead"
-            "in measurements (default: True)"
+            "Use triton do_bench_cudagraph (True) or do_bench (False) "
+            "for timing. CUDA graphs eliminate CPU launch overhead "
+            "(default: True)"
         ),
     )
+    parser.add_argument(
+        "--num-splits",
+        type=int,
+        default=None,
+        help="FlashAttention split-K factor (0=auto heuristic, 1=disabled, >1=force N)",
+    )
+    parser.add_argument(
+        "--ncu-profile",
+        action="store_true",
+        default=False,
+        help=(
+            "Enable Nsight Compute profiling mode. Automatically wraps the "
+            "script with ncu, capturing a profile with source correlation. "
+            "Use --ncu-output to set the output file name."
+        ),
+    )
+    parser.add_argument(
+        "--ncu-output",
+        type=str,
+        default="profile",
+        help="Output file name for ncu profile (default: 'profile').",
+    )
 
     # Parameter sweep (use YAML config for advanced sweeps)
     parser.add_argument(
@@ -576,23 +649,28 @@ def main():
             model = yaml_config["model"]
             args.num_layers = model.get("num_layers", args.num_layers)
             args.head_dim = model.get("head_dim", args.head_dim)
+            args.v_head_dim = model.get("v_head_dim", args.v_head_dim)
             args.num_q_heads = model.get("num_q_heads", args.num_q_heads)
             args.num_kv_heads = model.get("num_kv_heads", args.num_kv_heads)
             args.block_size = model.get("block_size", args.block_size)
+            # MLA-specific dimensions
+            args.kv_lora_rank = model.get("kv_lora_rank", args.kv_lora_rank)
+            args.qk_nope_head_dim = model.get("qk_nope_head_dim", args.qk_nope_head_dim)
+            args.qk_rope_head_dim = model.get("qk_rope_head_dim", args.qk_rope_head_dim)
 
         # Benchmark settings (top-level keys)
         if "device" in yaml_config:
             args.device = yaml_config["device"]
-        if "repeats" in yaml_config:
-            args.repeats = yaml_config["repeats"]
-        if "warmup_iters" in yaml_config:
-            args.warmup_iters = yaml_config["warmup_iters"]
+        if "warmup_ms" in yaml_config:
+            args.warmup_ms = yaml_config["warmup_ms"]
         if "profile_memory" in yaml_config:
             args.profile_memory = yaml_config["profile_memory"]
         if "kv_cache_dtype" in yaml_config:
             args.kv_cache_dtype = yaml_config["kv_cache_dtype"]
         if "cuda_graphs" in yaml_config:
             args.cuda_graphs = yaml_config["cuda_graphs"]
+        if "ncu_profile" in yaml_config:
+            args.ncu_profile = yaml_config["ncu_profile"]
 
         # Parameter sweep configuration
         if "parameter_sweep" in yaml_config:
@@ -612,7 +690,7 @@ def main():
         if "model_parameter_sweep" in yaml_config:
             sweep_config = yaml_config["model_parameter_sweep"]
             args.model_parameter_sweep = ModelParameterSweep(
-                param_name=sweep_config["param_name"],
+                param_name=sweep_config.get("param_name"),
                 values=sweep_config["values"],
                 label_format=sweep_config.get(
                     "label_format", "{backend}_{param_name}_{value}"
@@ -631,6 +709,32 @@ def main():
 
         console.print()
 
+    # Re-exec under ncu if --ncu-profile and not already inside ncu. This runs
+    # after YAML processing so ncu_profile set via config file is honored.
+    if args.ncu_profile and "_NCU_INNER" not in os.environ:
+        ncu = shutil.which("ncu")
+        if ncu is None:
+            print("Error: 'ncu' not found in PATH", file=sys.stderr)
+            sys.exit(1)
+        cmd = [
+            ncu,
+            "--profile-from-start",
+            "off",
+            "--set",
+            "full",
+            "--import-source",
+            "yes",
+            "-o",
+            args.ncu_output,
+            sys.executable,
+            *sys.argv,
+        ]
+        env = os.environ.copy()
+        env["CUTE_DSL_LINEINFO"] = "1"
+        env["_NCU_INNER"] = "1"
+        print(f"Launching: {' '.join(cmd)}")
+        sys.exit(subprocess.call(cmd, env=env))
+
     # Handle CLI-based parameter sweep (if not from YAML)
     if (
         (not hasattr(args, "parameter_sweep") or args.parameter_sweep is None)
@@ -655,6 +759,18 @@ def main():
     console.print(f"Batch specs: {', '.join(args.batch_specs)}")
     console.print(f"KV cache dtype: {args.kv_cache_dtype}")
     console.print(f"CUDA graphs: {args.cuda_graphs}")
+    if args.warmup_ms is not None and args.cuda_graphs:
+        console.print(
+            "[yellow]Warning: --warmup-ms is ignored with CUDA graphs "
+            "(do_bench_cudagraph warms up internally). Pass --no-cuda-graphs "
+            "to use it.[/]"
+        )
+    if args.num_splits == 0 and args.cuda_graphs:
+        console.print(
+            "[yellow]Warning: --num-splits 0 (FA3 heuristic) is not CUDA-graph "
+            "compatible and may fail or fall back. Pass --no-cuda-graphs or use "
+            "--num-splits >=1.[/]"
+        )
     console.print()
 
     init_workspace_manager(args.device)
@@ -662,6 +778,15 @@ def main():
     # Run benchmarks
     all_results = []
 
+    # Under ncu profiling the kernels run only to be captured by the profiler;
+    # timings are placeholder zeros, so the result tables and saved metrics are
+    # skipped. The Nsight Compute report (--ncu-output) holds the real data.
+    if args.ncu_profile:
+        console.print(
+            "[dim]ncu profiling enabled: result tables and saved metrics are "
+            "skipped (timings are placeholder zeros).[/]"
+        )
+
     # Handle special mode: decode_vs_prefill comparison
     if hasattr(args, "mode") and args.mode == "decode_vs_prefill":
         console.print("[yellow]Mode: Decode vs Prefill pipeline comparison[/]")
@@ -708,11 +833,11 @@ def main():
                         num_kv_heads=args.num_kv_heads,
                         block_size=args.block_size,
                         device=args.device,
-                        repeats=args.repeats,
-                        warmup_iters=args.warmup_iters,
                         profile_memory=args.profile_memory,
                         kv_cache_dtype=args.kv_cache_dtype,
                         use_cuda_graphs=args.cuda_graphs,
+                        ncu_profile=args.ncu_profile,
+                        warmup_ms=args.warmup_ms,
                     )
 
                     # Add decode pipeline config
@@ -749,6 +874,7 @@ def main():
                         result = BenchmarkResult(
                             config=config,
                             mean_time=timing["mean"],
+                            median_time=timing.get("median", timing["mean"]),
                             std_time=timing["std"],
                             min_time=timing["min"],
                             max_time=timing["max"],
@@ -770,6 +896,7 @@ def main():
                         result = BenchmarkResult(
                             config=config,
                             mean_time=float("inf"),
+                            median_time=float("inf"),
                             std_time=0,
                             min_time=float("inf"),
                             max_time=float("inf"),
@@ -779,6 +906,9 @@ def main():
 
                 pbar.update(1)
 
+        if args.ncu_profile:
+            return
+
         # Display decode vs prefill results
         console.print("\n[bold green]Decode vs Prefill Results:[/]")
 
@@ -858,15 +988,20 @@ def main():
         base_config_args = {
             "num_layers": args.num_layers,
             "head_dim": args.head_dim,
+            "v_head_dim": args.v_head_dim,
             "num_q_heads": args.num_q_heads,
             "num_kv_heads": args.num_kv_heads,
             "block_size": args.block_size,
             "device": args.device,
-            "repeats": args.repeats,
-            "warmup_iters": args.warmup_iters,
             "profile_memory": args.profile_memory,
             "kv_cache_dtype": args.kv_cache_dtype,
             "use_cuda_graphs": args.cuda_graphs,
+            "ncu_profile": args.ncu_profile,
+            "warmup_ms": args.warmup_ms,
+            "num_splits": args.num_splits,
+            "kv_lora_rank": args.kv_lora_rank,
+            "qk_nope_head_dim": args.qk_nope_head_dim,
+            "qk_rope_head_dim": args.qk_rope_head_dim,
         }
         all_results = run_model_parameter_sweep(
             backends,
@@ -882,15 +1017,17 @@ def main():
         base_config_args = {
             "num_layers": args.num_layers,
             "head_dim": args.head_dim,
+            "v_head_dim": args.v_head_dim,
             "num_q_heads": args.num_q_heads,
             "num_kv_heads": args.num_kv_heads,
             "block_size": args.block_size,
             "device": args.device,
-            "repeats": args.repeats,
-            "warmup_iters": args.warmup_iters,
             "profile_memory": args.profile_memory,
             "kv_cache_dtype": args.kv_cache_dtype,
             "use_cuda_graphs": args.cuda_graphs,
+            "ncu_profile": args.ncu_profile,
+            "warmup_ms": args.warmup_ms,
+            "num_splits": args.num_splits,
         }
         all_results = run_parameter_sweep(
             backends, args.batch_specs, base_config_args, args.parameter_sweep, console
@@ -914,15 +1051,17 @@ def main():
                             batch_spec=spec,
                             num_layers=args.num_layers,
                             head_dim=args.head_dim,
+                            v_head_dim=getattr(args, "v_head_dim", None),
                             num_q_heads=args.num_q_heads,
                             num_kv_heads=args.num_kv_heads,
                             block_size=args.block_size,
                             device=args.device,
-                            repeats=args.repeats,
-                            warmup_iters=args.warmup_iters,
                             profile_memory=args.profile_memory,
                             kv_cache_dtype=args.kv_cache_dtype,
                             use_cuda_graphs=args.cuda_graphs,
+                            ncu_profile=args.ncu_profile,
+                            warmup_ms=args.warmup_ms,
+                            num_splits=args.num_splits,
                         )
 
                         result = run_benchmark(config)
@@ -935,9 +1074,10 @@ def main():
 
                         pbar.update(1)
 
-            console.print("\n[bold green]Results:[/]")
-            formatter = ResultsFormatter(console)
-            formatter.print_table(decode_results, backends)
+            if not args.ncu_profile:
+                console.print("\n[bold green]Results:[/]")
+                formatter = ResultsFormatter(console)
+                formatter.print_table(decode_results, backends)
 
         # Run prefill backend comparison
         if prefill_backends:
@@ -962,9 +1102,8 @@ def main():
                             num_kv_heads=args.num_kv_heads,
                             block_size=args.block_size,
                             device=args.device,
-                            repeats=args.repeats,
-                            warmup_iters=args.warmup_iters,
                             profile_memory=args.profile_memory,
+                            warmup_ms=args.warmup_ms,
                             prefill_backend=pb,
                         )
 
@@ -980,16 +1119,17 @@ def main():
 
                         pbar.update(1)
 
-            console.print("\n[bold green]Prefill Backend Results:[/]")
-            formatter = ResultsFormatter(console)
-            formatter.print_table(
-                prefill_results, prefill_backends, compare_to_fastest=True
-            )
+            if not args.ncu_profile:
+                console.print("\n[bold green]Prefill Backend Results:[/]")
+                formatter = ResultsFormatter(console)
+                formatter.print_table(
+                    prefill_results, prefill_backends, compare_to_fastest=True
+                )
 
         all_results = decode_results + prefill_results
 
-    # Save results
-    if all_results:
+    # Save results (skip ncu profiling runs: timings are placeholder zeros)
+    if all_results and not args.ncu_profile:
         formatter = ResultsFormatter(console)
         if args.output_csv:
             formatter.save_csv(all_results, args.output_csv)

benchmarks/attention_benchmarks/common.py

@@ -15,6 +15,8 @@ from batch_spec import get_batch_type, parse_batch_spec
 from rich.console import Console
 from rich.table import Table
 
+from vllm.triton_utils import triton
+
 
 def batch_spec_sort_key(spec: str) -> tuple[int, int, int]:
     """
@@ -34,6 +36,30 @@ def batch_spec_sort_key(spec: str) -> tuple[int, int, int]:
         return (0, 0, 0)
 
 
+def run_do_bench(
+    benchmark_fn,
+    use_cuda_graphs: bool,
+    warmup_ms: int | None = None,
+) -> list[float]:
+    kwargs: dict[str, Any] = {"return_mode": "all"}
+    if use_cuda_graphs:
+        result = triton.testing.do_bench_cudagraph(benchmark_fn, **kwargs)
+    else:
+        if warmup_ms is not None:
+            kwargs["warmup"] = warmup_ms
+        result = triton.testing.do_bench(benchmark_fn, **kwargs)
+    return result
+
+
+def run_ncu_profile(benchmark_fn) -> None:
+    benchmark_fn()
+    torch.accelerator.synchronize()
+    torch.cuda.cudart().cudaProfilerStart()
+    benchmark_fn()
+    torch.accelerator.synchronize()
+    torch.cuda.cudart().cudaProfilerStop()
+
+
 # Mock classes for vLLM attention infrastructure
 
 
@@ -182,18 +208,37 @@ class ParameterSweep:
 
 @dataclass
 class ModelParameterSweep:
-    """Configuration for sweeping a model configuration parameter."""
+    """Configuration for sweeping model configuration parameter(s).
 
-    param_name: str  # Name of the model config parameter to sweep (e.g., "num_q_heads")
-    values: list[Any]  # List of values to test
-    label_format: str = "{backend}_{param_name}_{value}"  # Result label template
+    Supports two modes:
+    - Single param: param_name="head_dim", values=[128, 256, 512]
+    - Multi param: values=[{head_dim: 192, v_head_dim: 128}, {head_dim: 256}]
+      When values are dicts, each dict's keys are applied as config overrides.
+    """
+
+    param_name: str | None = None
+    values: list[Any] | None = None
+    label_format: str = "{backend}_{param_name}_{value}"
 
     def get_label(self, backend: str, value: Any) -> str:
         """Generate a label for a specific parameter value."""
+        if isinstance(value, dict):
+            return self.label_format.format(
+                backend=backend, param_name=self.param_name, value=value, **value
+            )
         return self.label_format.format(
             backend=backend, param_name=self.param_name, value=value
         )
 
+    def apply(self, config_args: dict, value: Any) -> None:
+        """Apply a sweep value to config args."""
+        if isinstance(value, dict):
+            config_args.update(value)
+        elif self.param_name is not None:
+            config_args[self.param_name] = value
+        else:
+            raise ValueError("param_name must be set if sweep values are not dicts")
+
 
 @dataclass
 class BenchmarkConfig:
@@ -208,10 +253,10 @@ class BenchmarkConfig:
     block_size: int
     device: str
     dtype: torch.dtype = torch.float16
-    repeats: int = 1
-    warmup_iters: int = 3
     profile_memory: bool = False
     use_cuda_graphs: bool = False
+    ncu_profile: bool = False
+    warmup_ms: int | None = None
 
     # "auto" or "fp8"
     kv_cache_dtype: str = "auto"
@@ -226,6 +271,7 @@ class BenchmarkConfig:
     # Backend-specific tuning
     num_kv_splits: int | None = None  # CUTLASS MLA
     reorder_batch_threshold: int | None = None  # FlashAttn MLA, FlashMLA
+    num_splits: int | None = None  # FlashAttention split-K (0=auto, 1=disabled)
 
 
 @dataclass
@@ -234,6 +280,7 @@ class BenchmarkResult:
 
     config: BenchmarkConfig
     mean_time: float  # seconds
+    median_time: float  # seconds
     std_time: float  # seconds
     min_time: float  # seconds
     max_time: float  # seconds
@@ -252,6 +299,7 @@ class BenchmarkResult:
         return {
             "config": asdict(self.config),
             "mean_time": self.mean_time,
+            "median_time": self.median_time,
             "std_time": self.std_time,
             "min_time": self.min_time,
             "max_time": self.max_time,

benchmarks/attention_benchmarks/configs/mla_decode.yaml

@@ -56,8 +56,6 @@ backends:
   - TOKENSPEED_MLA  # Blackwell + R1 dims + FP8 KV (use --kv-cache-dtype fp8)
 
 device: "cuda:0"
-repeats: 100
-warmup_iters: 10
 profile_memory: true
 
 # Backend-specific tuning

benchmarks/attention_benchmarks/configs/mla_mixed_batch.yaml

@@ -51,8 +51,6 @@ backends:
   - FLASHMLA         # Hopper only
 
 device: "cuda:0"
-repeats: 5
-warmup_iters: 3
 profile_memory: true
 
 # Analyze chunked prefill workspace size impact

benchmarks/attention_benchmarks/configs/mla_prefill.yaml

@@ -124,5 +124,3 @@ prefill_backends:
   - tokenspeed
 
 device: "cuda:0"
-repeats: 20
-warmup_iters: 5

benchmarks/attention_benchmarks/configs/mla_sparse_decode.yaml

@@ -53,6 +53,4 @@ backends:
   - FLASHINFER_MLA_SPARSE
 
 device: "cuda:0"
-repeats: 100
-warmup_iters: 10
 profile_memory: true

benchmarks/attention_benchmarks/configs/mla_sparse_prefill.yaml

@@ -57,6 +57,4 @@ backends:
   - FLASHINFER_MLA_SPARSE
 
 device: "cuda:0"
-repeats: 10
-warmup_iters: 3
 profile_memory: true

benchmarks/attention_benchmarks/configs/reorder_threshold.yaml

@@ -63,8 +63,6 @@ model:
 
 # Benchmark settings
 device: "cuda:0"
-repeats: 15          # More repeats for spec decode variance
-warmup_iters: 5
 profile_memory: false
 
 # Output

benchmarks/attention_benchmarks/configs/speculative_decode.yaml

@@ -49,8 +49,6 @@ backends:
 
 # Benchmark settings
 device: "cuda:0"
-repeats: 10  # More repeats for statistical significance
-warmup_iters: 5
 profile_memory: false
 
 # Test these threshold values for optimization

benchmarks/attention_benchmarks/configs/standard_attention.yaml

@@ -43,6 +43,4 @@ backends:
   - FLASHINFER
 
 device: "cuda:0"
-repeats: 5
-warmup_iters: 3
 profile_memory: false

benchmarks/attention_benchmarks/configs/standard_decode.yaml

@@ -0,0 +1,142 @@
+# Standard attention decode benchmark configuration
+# Sweeps num_q_heads and num_kv_heads to isolate effects of:
+#   1. GQA ratio (fixed num_q_heads=32, vary num_kv_heads)
+#   2. Absolute head count (fixed 4:1 ratio, vary scale)
+
+model:
+  num_layers: 32
+  num_q_heads: 32 # Base value, overridden by sweep
+  num_kv_heads: 8 # Base value, overridden by sweep
+  head_dim: 128
+  block_size: 16
+
+# Head count sweep: each entry overrides num_q_heads, num_kv_heads, and
+# head_dim where it differs from the base (128). Head counts are per-GPU
+# (i.e. after TP sharding).
+#
+# Group A — vary GQA ratio (fixed q=32, head_dim=128):
+#   32:32 (MHA), 32:8 (GQA 4:1), 32:4 (GQA 8:1), 32:1 (MQA)
+#
+# Groups B-E — real model configs at various TP degrees:
+#   Model              head_dim  Full    TP2      TP4      TP8
+#   Llama 3 8B           128    32:8   16:4      8:2      4:1
+#   Llama 3 70B          128    64:8   32:4     16:2      8:1
+#   GPT-OSS 120B          64    64:8   32:4     16:2      8:1
+#   Llama 3 405B         128   128:8   64:4     32:2     16:1
+model_parameter_sweep:
+  values:
+    # --- head_dim=128 (Llama 3 family) ---
+    - { num_q_heads: 32, num_kv_heads: 32, head_dim: 128 } # MHA 1:1
+    - { num_q_heads: 32, num_kv_heads: 1, head_dim: 128 } # MQA 32:1
+    - { num_q_heads: 4, num_kv_heads: 1, head_dim: 128 } # Llama 3 8B TP8
+    - { num_q_heads: 8, num_kv_heads: 2, head_dim: 128 } # Llama 3 8B TP4
+    - { num_q_heads: 16, num_kv_heads: 4, head_dim: 128 } # Llama 3 8B TP2
+    - { num_q_heads: 32, num_kv_heads: 8, head_dim: 128 } # Llama 3 8B TP1 / GQA 4:1
+    - { num_q_heads: 8, num_kv_heads: 1, head_dim: 128 } # Llama 3 70B TP8
+    - { num_q_heads: 16, num_kv_heads: 2, head_dim: 128 } # Llama 3 70B TP4
+    - { num_q_heads: 32, num_kv_heads: 4, head_dim: 128 } # Llama 3 70B TP2 / GQA 8:1
+    - { num_q_heads: 64, num_kv_heads: 8, head_dim: 128 } # Llama 3 70B TP1
+    - { num_q_heads: 16, num_kv_heads: 1, head_dim: 128 } # Llama 3 405B TP8
+    - { num_q_heads: 32, num_kv_heads: 2, head_dim: 128 } # Llama 3 405B TP4
+    - { num_q_heads: 64, num_kv_heads: 4, head_dim: 128 } # Llama 3 405B TP2
+    - { num_q_heads: 128, num_kv_heads: 8, head_dim: 128 } # Llama 3 405B TP1
+    # --- head_dim=64 (GPT-OSS 120B) ---
+    - { num_q_heads: 8, num_kv_heads: 1, head_dim: 64 } # GPT-OSS 120B TP8
+    - { num_q_heads: 16, num_kv_heads: 2, head_dim: 64 } # GPT-OSS 120B TP4
+    - { num_q_heads: 32, num_kv_heads: 4, head_dim: 64 } # GPT-OSS 120B TP2
+    - { num_q_heads: 64, num_kv_heads: 8, head_dim: 64 } # GPT-OSS 120B TP1
+  label_format: "{backend}_q{num_q_heads}kv{num_kv_heads}d{head_dim}"
+
+batch_specs:
+  # ---- batch_size x seq_len grid (decode: q_len=1) ----
+  # Small grid for quick iteration. Uncomment for full sweep.
+
+  # Batch size 1
+  - "q1s1k"
+  - "q1s512"
+  - "q1s2k"
+  - "q1s4k"
+  - "q1s8k"
+  - "q1s16k"
+  - "q1s32k"
+
+  # Batch size 2
+  - "2q1s512"
+  - "2q1s1k"
+  - "2q1s2k"
+  - "2q1s4k"
+  - "2q1s8k"
+  - "2q1s16k"
+  - "2q1s32k"
+
+  # Batch size 4
+  - "4q1s512"
+  - "4q1s1k"
+  - "4q1s2k"
+  - "4q1s4k"
+  - "4q1s8k"
+  - "4q1s16k"
+  - "4q1s32k"
+
+  # Batch size 8
+  - "8q1s1k"
+  - "8q1s512"
+  - "8q1s2k"
+  - "8q1s4k"
+  - "8q1s8k"
+  - "8q1s16k"
+  - "8q1s32k"
+
+  # Batch size 16
+  - "16q1s512"
+  - "16q1s1k"
+  - "16q1s2k"
+  - "16q1s4k"
+  - "16q1s8k"
+  - "16q1s16k"
+  - "16q1s32k"
+
+  # Batch size 32
+  - "32q1s512"
+  - "32q1s1k"
+  - "32q1s2k"
+  - "32q1s4k"
+  - "32q1s8k"
+  - "32q1s16k"
+  - "32q1s32k"
+
+  # Batch size 64
+  - "64q1s1k"
+  - "64q1s512"
+  - "64q1s2k"
+  - "64q1s4k"
+  - "64q1s8k"
+  - "64q1s16k"
+  - "64q1s32k"
+
+  # Batch size 128
+  - "128q1s512"
+  - "128q1s1k"
+  - "128q1s2k"
+  - "128q1s4k"
+  - "128q1s8k"
+  - "128q1s16k"
+  - "128q1s32k"
+
+  # Batch size 256
+  - "256q1s1k"
+  - "256q1s512"
+  - "256q1s2k"
+  - "256q1s4k"
+  - "256q1s8k"
+  - "256q1s16k"
+  - "256q1s32k"
+
+# Available backends: FLASH_ATTN, TRITON_ATTN, FLASHINFER
+backends:
+  - FLASH_ATTN
+  - TRITON_ATTN
+  - FLASHINFER
+
+device: "cuda:0"
+profile_memory: false

benchmarks/attention_benchmarks/configs/standard_prefill.yaml

@@ -0,0 +1,108 @@
+# Standard attention prefill benchmark configuration
+# Sweeps num_q_heads and num_kv_heads to isolate effects of:
+#   1. GQA ratio (fixed num_q_heads=32, vary num_kv_heads)
+#   2. Absolute head count (fixed 4:1 ratio, vary scale)
+
+model:
+  num_layers: 32
+  num_q_heads: 32 # Base value, overridden by sweep
+  num_kv_heads: 8 # Base value, overridden by sweep
+  head_dim: 128
+  block_size: 16
+
+# Head count sweep: each entry overrides num_q_heads, num_kv_heads, and
+# head_dim where it differs from the base (128). Head counts are per-GPU
+# (i.e. after TP sharding).
+#
+# Group A — vary GQA ratio (fixed q=32, head_dim=128):
+#   32:32 (MHA), 32:8 (GQA 4:1), 32:4 (GQA 8:1), 32:1 (MQA)
+#
+# Groups B-E — real model configs at various TP degrees:
+#   Model              head_dim  Full    TP2      TP4      TP8
+#   Llama 3 8B           128    32:8   16:4      8:2      4:1
+#   Llama 3 70B          128    64:8   32:4     16:2      8:1
+#   GPT-OSS 120B          64    64:8   32:4     16:2      8:1
+#   Llama 3 405B         128   128:8   64:4     32:2     16:1
+model_parameter_sweep:
+  values:
+    # --- head_dim=128 (Llama 3 family) ---
+    - { num_q_heads: 32, num_kv_heads: 32, head_dim: 128 } # MHA 1:1
+    - { num_q_heads: 32, num_kv_heads: 1, head_dim: 128 } # MQA 32:1
+    - { num_q_heads: 4, num_kv_heads: 1, head_dim: 128 } # Llama 3 8B TP8
+    - { num_q_heads: 8, num_kv_heads: 2, head_dim: 128 } # Llama 3 8B TP4
+    - { num_q_heads: 16, num_kv_heads: 4, head_dim: 128 } # Llama 3 8B TP2
+    - { num_q_heads: 32, num_kv_heads: 8, head_dim: 128 } # Llama 3 8B TP1 / GQA 4:1
+    - { num_q_heads: 8, num_kv_heads: 1, head_dim: 128 } # Llama 3 70B TP8
+    - { num_q_heads: 16, num_kv_heads: 2, head_dim: 128 } # Llama 3 70B TP4
+    - { num_q_heads: 32, num_kv_heads: 4, head_dim: 128 } # Llama 3 70B TP2 / GQA 8:1
+    - { num_q_heads: 64, num_kv_heads: 8, head_dim: 128 } # Llama 3 70B TP1
+    - { num_q_heads: 16, num_kv_heads: 1, head_dim: 128 } # Llama 3 405B TP8
+    - { num_q_heads: 32, num_kv_heads: 2, head_dim: 128 } # Llama 3 405B TP4
+    - { num_q_heads: 64, num_kv_heads: 4, head_dim: 128 } # Llama 3 405B TP2
+    - { num_q_heads: 128, num_kv_heads: 8, head_dim: 128 } # Llama 3 405B TP1
+    # --- head_dim=64 (GPT-OSS 120B) ---
+    - { num_q_heads: 8, num_kv_heads: 1, head_dim: 64 } # GPT-OSS 120B TP8
+    - { num_q_heads: 16, num_kv_heads: 2, head_dim: 64 } # GPT-OSS 120B TP4
+    - { num_q_heads: 32, num_kv_heads: 4, head_dim: 64 } # GPT-OSS 120B TP2
+    - { num_q_heads: 64, num_kv_heads: 8, head_dim: 64 } # GPT-OSS 120B TP1
+  label_format: "{backend}_q{num_q_heads}kv{num_kv_heads}d{head_dim}"
+
+batch_specs:
+  # ---- batch_size x prefill_len grid (prefill: q_len == seq_len) ----
+  # Total tokens = batch_size * prefill_len, and prefill compute scales with
+  # prefill_len^2, so the largest cells are expensive. Trim batch sizes or
+  # lengths for quick iteration.
+
+  # Batch size 1
+  - "q512"
+  - "q1k"
+  - "q2k"
+  - "q4k"
+  - "q8k"
+  - "q16k"
+  - "q32k"
+
+  # Batch size 2
+  - "2q512"
+  - "2q1k"
+  - "2q2k"
+  - "2q4k"
+  - "2q8k"
+  - "2q16k"
+  - "2q32k"
+
+  # Batch size 4
+  - "4q512"
+  - "4q1k"
+  - "4q2k"
+  - "4q4k"
+  - "4q8k"
+  - "4q16k"
+  - "4q32k"
+
+  # Batch size 8
+  - "8q512"
+  - "8q1k"
+  - "8q2k"
+  - "8q4k"
+  - "8q8k"
+  - "8q16k"
+  - "8q32k"
+
+  # Batch size 16
+  - "16q512"
+  - "16q1k"
+  - "16q2k"
+  - "16q4k"
+  - "16q8k"
+  - "16q16k"
+  - "16q32k"
+
+# Available backends: FLASH_ATTN, TRITON_ATTN, FLASHINFER
+backends:
+  - FLASH_ATTN
+  - TRITON_ATTN
+  - FLASHINFER
+
+device: "cuda:0"
+profile_memory: false

benchmarks/attention_benchmarks/mla_runner.py

@@ -8,6 +8,8 @@ This module provides helpers for running MLA backends without
 needing full VllmConfig integration.
 """
 
+import statistics
+
 import numpy as np
 import torch
 from batch_spec import parse_batch_spec
@@ -17,6 +19,8 @@ from common import (
     MockIndexer,
     MockKVBProj,
     MockLayer,
+    run_do_bench,
+    run_ncu_profile,
     setup_mla_dims,
 )
 
@@ -820,7 +824,7 @@ def _run_single_benchmark(
             num_prefill, mla_dims, query_fmt, device, torch.bfloat16
         )
 
-    # Build forward function
+    # Build forward function (runs a single decode/prefill pass)
     def forward_fn():
         results = []
         if has_decode:
@@ -839,44 +843,35 @@ def _run_single_benchmark(
             )
         return results[0] if len(results) == 1 else tuple(results)
 
-    # Warmup
-    for _ in range(config.warmup_iters):
-        forward_fn()
-    torch.accelerator.synchronize()
-
-    # Optionally capture a CUDA graph after warmup.
-    # Graph replay eliminates CPU launch overhead so timings reflect pure
-    # kernel time.
-    if config.use_cuda_graphs:
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            forward_fn()
-        benchmark_fn = graph.replay
-    else:
-        benchmark_fn = forward_fn
-
-    # Benchmark
-    times = []
-    for _ in range(config.repeats):
-        start = torch.cuda.Event(enable_timing=True)
-        end = torch.cuda.Event(enable_timing=True)
-
-        start.record()
+    def benchmark_fn():
         for _ in range(config.num_layers):
-            benchmark_fn()
-        end.record()
+            forward_fn()
 
-        torch.accelerator.synchronize()
-        elapsed_ms = start.elapsed_time(end)
-        times.append(elapsed_ms / 1000.0 / config.num_layers)
+    if config.ncu_profile:
+        run_ncu_profile(benchmark_fn)
+        return BenchmarkResult(
+            config=config,
+            mean_time=0.0,
+            median_time=0.0,
+            std_time=0.0,
+            min_time=0.0,
+            max_time=0.0,
+            throughput_tokens_per_sec=0.0,
+        )
+
+    all_ms = run_do_bench(benchmark_fn, config.use_cuda_graphs, config.warmup_ms)
+
+    # Convert ms to seconds per layer
+    times = [t / 1000.0 / config.num_layers for t in all_ms]
+    mean_time = statistics.mean(times)
 
-    mean_time = float(np.mean(times))
     return BenchmarkResult(
         config=config,
         mean_time=mean_time,
-        std_time=float(np.std(times)),
-        min_time=float(np.min(times)),
-        max_time=float(np.max(times)),
+        median_time=statistics.median(times),
+        std_time=statistics.stdev(times) if len(times) > 1 else 0.0,
+        min_time=min(times),
+        max_time=max(times),
         throughput_tokens_per_sec=total_q / mean_time if mean_time > 0 else 0,
     )
 

benchmarks/attention_benchmarks/runner.py

@@ -9,13 +9,20 @@ This module provides helpers for running standard attention backends
 """
 
 import logging
+import statistics
 import types
 from contextlib import contextmanager
 
-import numpy as np
 import torch
 from batch_spec import parse_batch_spec, reorder_for_flashinfer
-from common import BenchmarkConfig, BenchmarkResult, MockLayer, get_attention_scale
+from common import (
+    BenchmarkConfig,
+    BenchmarkResult,
+    MockLayer,
+    get_attention_scale,
+    run_do_bench,
+    run_ncu_profile,
+)
 
 from vllm.config import (
     CacheConfig,
@@ -208,6 +215,13 @@ def _create_backend_impl(
 
     scale = get_attention_scale(config.head_dim)
 
+    # Set v_head_dim for diff-headdim backends. Always reset (defaulting to
+    # head_dim) so a prior run's value doesn't leak into this one via the
+    # backend's class-level state.
+    if hasattr(backend_class, "set_head_size_v"):
+        v_dim = config.v_head_dim if config.v_head_dim is not None else config.head_dim
+        backend_class.set_head_size_v(v_dim)
+
     impl = backend_class.get_impl_cls()(
         num_heads=config.num_q_heads,
         head_size=config.head_dim,
@@ -300,6 +314,7 @@ def _create_input_tensors(
         from vllm.platforms import current_platform
 
         q_dtype = current_platform.fp8_dtype()
+    v_dim = config.v_head_dim if config.v_head_dim is not None else config.head_dim
     q_list = [
         torch.randn(
             total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype
@@ -313,9 +328,7 @@ def _create_input_tensors(
         for _ in range(config.num_layers)
     ]
     v_list = [
-        torch.randn(
-            total_q, config.num_kv_heads, config.head_dim, device=device, dtype=dtype
-        )
+        torch.randn(total_q, config.num_kv_heads, v_dim, device=device, dtype=dtype)
         for _ in range(config.num_layers)
     ]
     return q_list, k_list, v_list
@@ -389,14 +402,17 @@ def _run_single_benchmark(
     device: torch.device,
     dtype: torch.dtype,
 ) -> tuple:
-    """Run single benchmark iteration with warmup and timing loop."""
-    total_q = q_list[0].shape[0]
-    out = torch.empty(
-        total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype
-    )
+    """Run single benchmark using triton's do_bench_cudagraph/do_bench.
 
-    # Warmup
-    for _ in range(config.warmup_iters):
+    Returns:
+        (timing_stats, mem_stats) where timing_stats is a dict with
+        mean/std/min/max in seconds per layer.
+    """
+    total_q = q_list[0].shape[0]
+    v_dim = config.v_head_dim if config.v_head_dim is not None else config.head_dim
+    out = torch.empty(total_q, config.num_q_heads, v_dim, device=device, dtype=dtype)
+
+    def benchmark_fn():
         for i in range(config.num_layers):
             impl.forward(
                 layer,
@@ -407,52 +423,22 @@ def _run_single_benchmark(
                 attn_metadata,
                 output=out,
             )
-    torch.accelerator.synchronize()
 
-    # Optionally capture a CUDA graph after warmup.
-    # Graph replay eliminates CPU launch overhead so timings reflect pure
-    # kernel time.
-    if config.use_cuda_graphs:
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            for i in range(config.num_layers):
-                impl.forward(
-                    layer,
-                    q_list[i],
-                    k_list[i],
-                    v_list[i],
-                    cache_list[i],
-                    attn_metadata,
-                    output=out,
-                )
-        benchmark_fn = graph.replay
+    if config.ncu_profile:
+        run_ncu_profile(benchmark_fn)
+        timing_stats = dict.fromkeys(("mean", "median", "std", "min", "max"), 0.0)
     else:
+        all_ms = run_do_bench(benchmark_fn, config.use_cuda_graphs, config.warmup_ms)
 
-        def benchmark_fn():
-            for i in range(config.num_layers):
-                impl.forward(
-                    layer,
-                    q_list[i],
-                    k_list[i],
-                    v_list[i],
-                    cache_list[i],
-                    attn_metadata,
-                    output=out,
-                )
-
-    # Benchmark
-    times = []
-    for _ in range(config.repeats):
-        start = torch.cuda.Event(enable_timing=True)
-        end = torch.cuda.Event(enable_timing=True)
-
-        start.record()
-        benchmark_fn()
-        end.record()
-
-        torch.accelerator.synchronize()
-        elapsed_ms = start.elapsed_time(end)
-        times.append(elapsed_ms / 1000.0 / config.num_layers)  # seconds per layer
+        # Convert ms to seconds per layer
+        times = [t / 1000.0 / config.num_layers for t in all_ms]
+        timing_stats = {
+            "mean": statistics.mean(times),
+            "std": statistics.stdev(times) if len(times) > 1 else 0.0,
+            "min": min(times),
+            "max": max(times),
+            "median": statistics.median(times),
+        }
 
     mem_stats = {}
     if config.profile_memory:
@@ -461,7 +447,7 @@ def _run_single_benchmark(
             "reserved_mb": torch.accelerator.memory_reserved(device) / 1024**2,
         }
 
-    return times, mem_stats
+    return timing_stats, mem_stats
 
 
 # ============================================================================
@@ -541,6 +527,12 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
                 common_attn_metadata=common_metadata,
             )
 
+            # Override num_splits for split-K testing (FlashAttention only)
+            if config.num_splits is not None and hasattr(
+                attn_metadata, "max_num_splits"
+            ):
+                attn_metadata.max_num_splits = config.num_splits
+
             # Only quantize queries when the impl supports it
             quantize_query = config.kv_cache_dtype.startswith("fp8") and getattr(
                 impl, "supports_quant_query_input", False
@@ -553,7 +545,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
                 config, max_num_blocks, backend_class, device, dtype
             )
 
-            times, mem_stats = _run_single_benchmark(
+            timing_stats, mem_stats = _run_single_benchmark(
                 config,
                 impl,
                 layer,
@@ -566,15 +558,16 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
                 dtype,
             )
 
-    mean_time = np.mean(times)
+    mean_time = timing_stats["mean"]
     throughput = total_q / mean_time if mean_time > 0 else 0
 
     return BenchmarkResult(
         config=config,
         mean_time=mean_time,
-        std_time=np.std(times),
-        min_time=np.min(times),
-        max_time=np.max(times),
+        median_time=timing_stats["median"],
+        std_time=timing_stats["std"],
+        min_time=timing_stats["min"],
+        max_time=timing_stats["max"],
         throughput_tokens_per_sec=throughput,
         memory_allocated_mb=mem_stats.get("allocated_mb"),
         memory_reserved_mb=mem_stats.get("reserved_mb"),

benchmarks/benchmark_hidden_state_extraction.py

@@ -92,7 +92,6 @@ def run_baseline(
     llm = LLM(
         model=model,
         enable_prefix_caching=False,
-        enable_chunked_prefill=False,
         **extra_args,
     )
     sampling_params = SamplingParams(max_tokens=1)
@@ -194,7 +193,6 @@ async def _run_extraction_async(
     engine_args = AsyncEngineArgs(
         model=model,
         enable_prefix_caching=False,
-        enable_chunked_prefill=False,
         max_num_batched_tokens=40960,
         max_model_len=40960,
         speculative_config={

benchmarks/kernels/benchmark_device_communicators.py

@@ -33,6 +33,7 @@ from vllm.distributed.device_communicators.custom_all_reduce import CustomAllred
 from vllm.distributed.device_communicators.flashinfer_all_reduce import (
     FlashInferAllReduce,
 )
+from vllm.distributed.device_communicators.push_all_reduce import PushAllReduce
 from vllm.distributed.device_communicators.pynccl import (
     PyNcclCommunicator,
     register_nccl_symmetric_ops,
@@ -80,6 +81,7 @@ class CommunicatorBenchmark:
 
         # Initialize communicators
         self.custom_allreduce = None
+        self.push_ar_comm = None
         self.pynccl_comm = None
         self.symm_mem_comm = None
         self.symm_mem_comm_multimem = None
@@ -106,6 +108,23 @@ class CommunicatorBenchmark:
             )
             self.custom_allreduce = None
 
+        try:
+            self.push_ar_comm = PushAllReduce(
+                group=self.cpu_group,
+                device=self.device,
+                max_size=self.max_size_override,
+            )
+            if not self.push_ar_comm.disabled:
+                logger.info("Rank %s: PushAllReduce initialized", self.rank)
+            else:
+                logger.info("Rank %s: PushAllReduce disabled", self.rank)
+                self.push_ar_comm = None
+        except Exception as e:
+            logger.warning(
+                "Rank %s: Failed to initialize PushAllReduce: %s", self.rank, e
+            )
+            self.push_ar_comm = None
+
         try:
             self.pynccl_comm = PyNcclCommunicator(
                 group=self.cpu_group, device=self.device
@@ -216,6 +235,19 @@ class CommunicatorBenchmark:
                 )
             )
 
+        if self.push_ar_comm is not None:
+            comm = self.push_ar_comm
+            communicators.append(
+                (
+                    "push_ar",
+                    lambda t, c=comm: c.all_reduce(t),
+                    lambda t, c=comm: c.should_use(t),
+                    comm.capture(),
+                    {},
+                    None,
+                )
+            )
+
         if self.pynccl_comm is not None:
             comm = self.pynccl_comm
             communicators.append(

benchmarks/kernels/benchmark_flydsl_moe_w4a16.py

@@ -0,0 +1,277 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Copyright (c) 2025 FlyDSL Project Contributors
+
+import json
+import os
+
+import torch
+from aiter.test_common import run_perftest
+
+from vllm.model_executor.layers.fused_moe import fused_experts
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import (
+    int4_w4a16_moe_quant_config,
+)
+from vllm.model_executor.layers.fused_moe.fused_flydsl_moe import fused_flydsl_moe
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import (  # noqa: E501
+    compressed_tensors_moe_w4a16_flydsl,
+)
+from vllm.platforms import current_platform
+
+RoutingBuffers = tuple[
+    torch.Tensor,  # sorted_token_ids
+    torch.Tensor,  # sorted_weights
+    torch.Tensor,  # sorted_expert_ids
+    torch.Tensor,  # num_valid_ids (shape [1], i32)
+    int,  # sorted_size
+    int,  # blocks
+]
+
+MODEL_PARAMS_TO_TUNE = [
+    # (num_experts, inter_dim, hidden_size, topk)
+    (384, 256, 7168, 8),  # Kimi K2.5 TP=8
+    (384, 512, 7168, 8),  # Kimi K2.5 TP=4
+]
+
+NUM_TOKENS_TO_TUNE = [
+    1,
+    2,
+    4,
+    8,
+    16,
+    24,
+    32,
+    48,
+    64,
+    128,
+    256,
+    512,
+    1024,
+    2048,
+    4096,
+    8192,
+]
+
+TILE_M_SEARCH_SPACE = [16, 32, 64, 128, 256]
+TILE_N_SEARCH_SPACE = [16, 32, 64, 128, 256]
+TILE_K_SEARCH_SPACE = [16, 32, 64, 128, 256, 512]
+TILE_N2_SEARCH_SPACE = [16, 32, 64, 128, 256]
+TILE_K2_SEARCH_SPACE = [16, 32, 64, 128, 256, 512]
+
+TILE_CONFIGS = []
+for tile_m in TILE_M_SEARCH_SPACE:
+    for tile_n in TILE_N_SEARCH_SPACE:
+        for tile_k in TILE_K_SEARCH_SPACE:
+            for tile_n2 in TILE_N2_SEARCH_SPACE:
+                for tile_k2 in TILE_K2_SEARCH_SPACE:
+                    TILE_CONFIGS.append(
+                        {
+                            "tile_m": tile_m,
+                            "tile_n": tile_n,
+                            "tile_k": tile_k,
+                            "tile_n2": tile_n2,
+                            "tile_k2": tile_k2,
+                        }
+                    )
+
+
+def tune_flydsl_moe_w4a16(
+    device: str = "cuda", num_iters: int = 100, num_warmup: int = 10
+):
+    packed_factor = 8
+    w13_num_shards = 2
+    params_dtype = torch.bfloat16
+    group_size = 32
+    scale_factor = 0.01
+
+    for model_params in MODEL_PARAMS_TO_TUNE:
+        num_experts = model_params[0]
+        inter_dim = model_params[1]
+        hidden_size = model_params[2]
+        topk = model_params[3]
+        print(
+            f"\nTuning: num_experts={num_experts}, inter_dim={inter_dim}, "
+            f"hidden_size={hidden_size}, topk={topk}...\n"
+        )
+
+        w2_scales_size = inter_dim
+        num_groups_w2 = w2_scales_size // group_size
+        num_groups_w13 = hidden_size // group_size
+
+        w13_weight = torch.randint(
+            0,
+            255,
+            (num_experts, hidden_size // packed_factor, w13_num_shards * inter_dim),
+            dtype=torch.int32,
+            device=device,
+        )
+
+        w2_weight = torch.randint(
+            0,
+            255,
+            (num_experts, inter_dim // packed_factor, hidden_size),
+            dtype=torch.int32,
+            device=device,
+        )
+        w13_scale = scale_factor * torch.randn(
+            num_experts,
+            num_groups_w13,
+            w13_num_shards * inter_dim,
+            dtype=params_dtype,
+            device=device,
+        )
+        w2_scale = scale_factor * torch.randn(
+            num_experts, num_groups_w2, hidden_size, dtype=params_dtype, device=device
+        )
+
+        w13 = w13_weight
+        w13 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w13)
+        w13 = w13.view(-1).contiguous()
+
+        w2 = w2_weight
+        w2 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w2)
+        w2 = w2.view(-1).contiguous()
+
+        w13_scale_flydsl = w13_scale
+        w2_scale_flydsl = w2_scale
+
+        if group_size > 0 and w13_scale.dim() == 3 and w13_scale.shape[1] > 1:
+            E, G, N = w13_scale.shape
+            w13_scale_flydsl = (
+                w13_scale_flydsl.view(E, G // 2, 2, N)
+                .permute(0, 1, 3, 2)
+                .contiguous()
+                .view(-1)
+                .contiguous()
+            )
+        elif w13_scale.dim() == 3 and w13_scale.shape[1] == 1:
+            w13_scale_flydsl = w13_scale_flydsl.squeeze(1)
+
+        if group_size > 0 and w2_scale.dim() == 3 and w2_scale.shape[1] > 1:
+            E, G, N = w2_scale.shape
+            w2_scale_flydsl = (
+                w2_scale_flydsl.view(E, G // 2, 2, N)
+                .permute(0, 1, 3, 2)
+                .contiguous()
+                .view(-1)
+                .contiguous()
+            )
+        elif w2_scale.dim() == 3 and w2_scale.shape[1] == 1:
+            w2_scale_flydsl = w2_scale_flydsl.squeeze(1)
+
+        w13_scale_flydsl = w13_scale_flydsl.contiguous()
+        w2_scale_flydsl = w2_scale_flydsl.contiguous()
+
+        w13.is_shuffled = True
+        w2.is_shuffled = True
+
+        w13_weight_scale = w13_scale.transpose(1, 2).contiguous()
+        w2_weight_scale = w2_scale.transpose(1, 2).contiguous()
+        w13_weight_packed = w13_weight.transpose(1, 2).contiguous().view(torch.uint8)
+        w2_weight_packed = w2_weight.transpose(1, 2).contiguous().view(torch.uint8)
+
+        moe_quant_config = int4_w4a16_moe_quant_config(
+            w1_scale=w13_weight_scale,
+            w2_scale=w2_weight_scale,
+            w1_zp=None,
+            w2_zp=None,
+            block_shape=[0, group_size],
+        )
+
+        tuned_config = {}
+
+        for num_tokens in NUM_TOKENS_TO_TUNE:
+            score = torch.rand(
+                (num_tokens, num_experts), device=device, dtype=torch.float32
+            )
+            topk_vals, topk_ids = torch.topk(score, k=topk, dim=1)
+            topk_weights = torch.softmax(topk_vals, dim=1).to(torch.float32)
+            x = torch.randn(
+                (num_tokens, hidden_size), dtype=torch.bfloat16, device=device
+            )
+            us_best = float("inf")
+            for tile_config in TILE_CONFIGS:
+                try:
+                    tile_m = tile_config["tile_m"]
+                    tile_n = tile_config["tile_n"]
+                    tile_k = tile_config["tile_k"]
+                    tile_n2 = tile_config["tile_n2"]
+                    tile_k2 = tile_config["tile_k2"]
+
+                    model_dim = x.shape[1]
+                    assert model_dim % 64 == 0
+                    assert model_dim % tile_k == 0
+                    assert inter_dim % tile_n == 0
+                    assert model_dim % tile_n2 == 0
+                    assert inter_dim % tile_k2 == 0
+                    assert ((tile_m * tile_k2) % 256) == 0
+                    bytes_per_thread_x = (tile_m * tile_k2) // 256
+                    assert (bytes_per_thread_x % 4) == 0
+
+                    out, _us = run_perftest(
+                        fused_flydsl_moe,
+                        x,
+                        w13,
+                        w2,
+                        num_experts,
+                        inter_dim,
+                        topk_weights,
+                        topk_ids,
+                        num_iters=num_iters,
+                        num_warmup=num_warmup,
+                        w1_scale=w13_scale_flydsl,
+                        w2_scale=w2_scale_flydsl,
+                        topk=topk_weights.shape[-1],
+                        group_size=group_size,
+                        doweight_stage1=False,
+                        scale_is_bf16=True,
+                        config=tile_config,
+                    )
+                    torch.accelerator.synchronize()
+                except Exception:
+                    torch.accelerator.synchronize()
+                    continue
+                else:
+                    us = _us.item()
+                    if us < us_best:
+                        out_ref = fused_experts(
+                            x,
+                            w13_weight_packed,
+                            w2_weight_packed,
+                            topk_weights=topk_weights,
+                            topk_ids=topk_ids,
+                            activation=MoEActivation.SILU,
+                            apply_router_weight_on_input=False,
+                            global_num_experts=num_experts,
+                            expert_map=None,
+                            quant_config=moe_quant_config,
+                        )
+                        try:
+                            assert torch.allclose(out, out_ref, atol=0.5, rtol=0.1)
+                        except Exception:
+                            continue
+                        else:
+                            print(
+                                f"For [num_tokens={num_tokens}, num_experts={num_experts}, "  # noqa: E501
+                                f"inter_dim={inter_dim}] found new best "  # noqa: E501
+                                f"config={tile_config}, us={us:0.3f}"
+                            )
+                            us_best = us
+                            tuned_config[str(num_tokens)] = tile_config
+            device_name = current_platform.get_device_name().replace(" ", "_")
+            tuned_config_file_name = (
+                f"E={num_experts},N={inter_dim},device_name={device_name},"
+                f"dtype=int4_w4a16,backend=flydsl.json"
+            )
+            tuner_dir_path = os.path.dirname(os.path.realpath(__file__))
+            store_path = os.path.join(tuner_dir_path, tuned_config_file_name)
+            with open(store_path, "w") as f:
+                json.dump(tuned_config, f, indent=4)
+            print(
+                f"\nTuned config for num_tokens={num_tokens} was stored at {store_path}\n"  # noqa: E501
+            )
+
+
+if __name__ == "__main__":
+    tune_flydsl_moe_w4a16(device="cuda")

benchmarks/kernels/benchmark_moe.py

@@ -792,6 +792,12 @@ def get_model_params(config):
         topk = text_config.num_experts_per_tok
         intermediate_size = text_config.moe_intermediate_size
         hidden_size = text_config.hidden_size
+    elif architecture == "DiffusionGemmaForBlockDiffusion":
+        text_config = config.get_text_config()
+        E = text_config.num_experts
+        topk = text_config.top_k_experts
+        intermediate_size = text_config.moe_intermediate_size
+        hidden_size = text_config.hidden_size
     elif architecture == "HunYuanMoEV1ForCausalLM":
         E = config.num_experts
         topk = config.moe_topk[0]

benchmarks/kv_cache_watermark.sh

@@ -0,0 +1,248 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Reproducible demonstration of the KV cache watermark (`--watermark`) for
+# reducing preemption thrashing.
+#
+# The watermark is the fraction of total KV cache blocks the scheduler keeps
+# free when admitting a waiting/preempted request into the running queue.
+#
+# Why this workload triggers thrashing:
+#   Requests are admitted based on the KV cache they need *at admission time*.
+#   With `--scheduler-reserve-full-isl` (default) the input length is reserved up
+#   front, but the *output* length is unknown and unreserved. A decode-heavy
+#   workload (output >> input) at high concurrency therefore over-admits while
+#   requests are short, then runs out of KV cache as they all grow during decode
+#   -> the scheduler preempts (recompute) recently-admitted requests, re-prefills
+#   them later, and repeats. The watermark keeps a block of KV cache free so
+#   running requests can grow into it instead of triggering this churn.
+#
+# This script launches `vllm serve` under a deliberately KV-constrained config
+# and a decode-heavy workload, sweeping the watermark across several values, and
+# reports the preemption count (scraped from /metrics), throughput, and latency
+# percentiles for each. It then plots the results.
+#
+# Default workload: concurrency 200, input ~300 tokens, output ~4000 tokens
+# (+/- 20% variance), sized to run each config for ~5 minutes.
+#
+# Usage:
+#   benchmarks/kv_cache_watermark.sh
+#   MODEL=Qwen/Qwen2.5-14B-Instruct TP=2 benchmarks/kv_cache_watermark.sh
+#
+# Run inside the vLLM virtualenv (so `vllm` and `python` resolve to it).
+set -euo pipefail
+
+# ---- Config (override via environment) -------------------------------------
+MODEL=${MODEL:-Qwen/Qwen2.5-7B-Instruct}
+TP=${TP:-1}
+PORT=${PORT:-8000}
+URL="http://127.0.0.1:${PORT}"
+# Constrain the KV cache to a *near-critical* size: large enough that the engine
+# can run stably, but small enough that greedy over-admission tips it into
+# preemption thrashing. (Independent of GPU size, so the demo is reproducible.)
+# At the default workload this fits ~1.5x the mean concurrent KV demand.
+KV_CACHE_MEMORY_GB=${KV_CACHE_MEMORY_GB:-16}
+MAX_MODEL_LEN=${MAX_MODEL_LEN:-8192}
+MAX_NUM_SEQS=${MAX_NUM_SEQS:-256}
+# Optional weight loader (e.g. fastsafetensors on the GCP cluster).
+LOAD_FORMAT=${LOAD_FORMAT:-auto}
+# Decode-heavy workload: moderate input, long output, with length variance. The
+# long output means preempted requests have generated a lot before eviction, so
+# resuming them re-prefills a long sequence (high recomputation cost).
+INPUT_LEN=${INPUT_LEN:-1000}
+OUTPUT_LEN=${OUTPUT_LEN:-5000}
+RANGE_RATIO=${RANGE_RATIO:-0.2}
+CONCURRENCY=${CONCURRENCY:-128}
+# Enough prompts to keep each config saturated for ~5+ minutes.
+NUM_PROMPTS=${NUM_PROMPTS:-450}
+OUTDIR=${OUTDIR:-./watermark_bench_results}
+# Watermark fractions compared. "label value" per line; value=0 disables it.
+CONFIGS=${CONFIGS:-"off    0
+w0.02  0.02
+w0.05  0.05
+w0.10  0.10
+w0.15  0.15"}
+
+KV_CACHE_MEMORY_BYTES=$((KV_CACHE_MEMORY_GB * 1024 * 1024 * 1024))
+mkdir -p "$OUTDIR"
+
+SERVER_PID=""
+cleanup() { [[ -n "$SERVER_PID" ]] && kill "$SERVER_PID" 2>/dev/null || true; }
+trap cleanup EXIT
+
+scrape_preemptions() {
+  # Sum the vllm:num_preemptions_total counter across engines.
+  python - "${URL}/metrics" <<'PY'
+import sys, urllib.request
+total = 0.0
+try:
+    body = urllib.request.urlopen(sys.argv[1], timeout=10).read().decode("utf-8", "replace")
+    for line in body.splitlines():
+        if line.startswith("vllm:num_preemptions_total"):
+            total += float(line.rsplit(" ", 1)[-1])
+except Exception as e:  # noqa: BLE001
+    print(f"scrape error: {e}", file=sys.stderr)
+print(int(total))
+PY
+}
+
+wait_for_server() {
+  for _ in $(seq 1 300); do
+    if curl -s "${URL}/health" >/dev/null 2>&1; then return 0; fi
+    if ! kill -0 "$SERVER_PID" 2>/dev/null; then
+      echo "ERROR: server process exited during startup" >&2; return 1
+    fi
+    sleep 5
+  done
+  echo "ERROR: server did not become ready" >&2; return 1
+}
+
+run_one() {
+  local label=$1 watermark=$2
+  echo
+  echo "==================== watermark: ${label} (${watermark}) ===================="
+  vllm serve "$MODEL" \
+    --tensor-parallel-size "$TP" \
+    --load-format "$LOAD_FORMAT" \
+    --kv-cache-memory-bytes "$KV_CACHE_MEMORY_BYTES" \
+    --max-model-len "$MAX_MODEL_LEN" \
+    --max-num-seqs "$MAX_NUM_SEQS" \
+    --no-enable-prefix-caching \
+    --watermark "$watermark" \
+    --port "$PORT" >"${OUTDIR}/serve_${label}.log" 2>&1 &
+  SERVER_PID=$!
+  wait_for_server
+  sleep 5
+
+  local pre post
+  pre=$(scrape_preemptions)
+  vllm bench serve \
+    --backend vllm \
+    --base-url "$URL" \
+    --model "$MODEL" \
+    --dataset-name random \
+    --random-input-len "$INPUT_LEN" \
+    --random-output-len "$OUTPUT_LEN" \
+    --random-range-ratio "$RANGE_RATIO" \
+    --ignore-eos \
+    --num-prompts "$NUM_PROMPTS" \
+    --max-concurrency "$CONCURRENCY" \
+    --percentile-metrics "ttft,tpot,itl,e2el" \
+    --metric-percentiles "50,90,99" \
+    --save-result \
+    --result-dir "$OUTDIR" \
+    --result-filename "bench_${label}.json"
+  post=$(scrape_preemptions)
+  echo "${label} ${watermark} $((post - pre))" >>"${OUTDIR}/preemptions.txt"
+
+  kill "$SERVER_PID" 2>/dev/null || true
+  for _ in $(seq 1 60); do curl -s "${URL}/health" >/dev/null 2>&1 || break; sleep 2; done
+  SERVER_PID=""
+  sleep 10
+}
+
+: >"${OUTDIR}/preemptions.txt"
+while read -r label watermark; do
+  [[ -z "${label:-}" ]] && continue
+  run_one "$label" "$watermark"
+done <<<"$CONFIGS"
+
+echo
+echo "==================== summary ===================="
+python - "$OUTDIR" <<'PY'
+import json, os, sys
+outdir = sys.argv[1]
+pre = {}
+order = []
+for line in open(os.path.join(outdir, "preemptions.txt")):
+    label, watermark, n = line.split()
+    pre[label] = (float(watermark), int(n))
+    order.append(label)
+
+def g(d, *names):
+    for n in names:
+        if d.get(n) is not None:
+            return d[n]
+    return float("nan")
+
+cols = ["watermark", "frac", "preempt", "out_tok/s", "req/s",
+        "TTFT_p50", "TTFT_p99", "ITL_p99", "E2EL_p50"]
+print("  ".join(f"{c:>10}" for c in cols))
+rows = []
+for label in order:
+    watermark, n = pre[label]
+    d = json.load(open(os.path.join(outdir, f"bench_{label}.json")))
+    rows.append(dict(
+        label=label, watermark=watermark, preempt=n,
+        out_tok_s=g(d, "output_throughput"),
+        req_s=g(d, "request_throughput"),
+        ttft_p50=g(d, "p50_ttft_ms", "median_ttft_ms"),
+        ttft_p99=g(d, "p99_ttft_ms"),
+        itl_p99=g(d, "p99_itl_ms"),
+        e2el_p50=g(d, "p50_e2el_ms", "median_e2el_ms"),
+    ))
+    print("  ".join(f"{str(v):>10}" for v in [
+        label, watermark, n,
+        f"{rows[-1]['out_tok_s']:.0f}",
+        f"{rows[-1]['req_s']:.3f}",
+        f"{rows[-1]['ttft_p50']/1000:.2f}",
+        f"{rows[-1]['ttft_p99']/1000:.2f}",
+        f"{rows[-1]['itl_p99']:.2f}",
+        f"{rows[-1]['e2el_p50']/1000:.1f}",
+    ]))
+print("\n(TTFT/E2EL in seconds; ITL in ms. Lower preempt is better.)")
+
+# ---- Plot -------------------------------------------------------------------
+try:
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+except Exception as e:  # noqa: BLE001
+    print(f"\n(skip plot: matplotlib unavailable: {e})")
+    sys.exit(0)
+
+x = [r["watermark"] for r in rows]
+xt = [f"{r['watermark']:g}\n({r['label']})" for r in rows]
+idx = list(range(len(rows)))
+
+fig, axes = plt.subplots(2, 2, figsize=(12, 8))
+fig.suptitle(
+    f"KV cache watermark sweep — {os.path.basename(os.path.abspath(outdir))}",
+    fontsize=12,
+)
+
+ax = axes[0][0]
+ax.bar(idx, [r["preempt"] for r in rows], color="tab:red")
+ax.set_title("Preemptions (lower is better)")
+ax.set_ylabel("preemptions")
+ax.set_xticks(idx); ax.set_xticklabels(xt)
+
+ax = axes[0][1]
+ax.plot(idx, [r["out_tok_s"] for r in rows], "o-", color="tab:green")
+ax.set_title("Output throughput (higher is better)")
+ax.set_ylabel("tokens/s")
+ax.set_xticks(idx); ax.set_xticklabels(xt)
+
+ax = axes[1][0]
+ax.plot(idx, [r["itl_p99"] for r in rows], "o-", color="tab:blue")
+ax.set_title("Inter-token latency p99 (lower is better)")
+ax.set_ylabel("ITL p99 (ms)")
+ax.set_xlabel("watermark fraction")
+ax.set_xticks(idx); ax.set_xticklabels(xt)
+
+ax = axes[1][1]
+ax.plot(idx, [r["ttft_p50"] / 1000 for r in rows], "o-", label="TTFT p50")
+ax.plot(idx, [r["ttft_p99"] / 1000 for r in rows], "o-", label="TTFT p99")
+ax.plot(idx, [r["e2el_p50"] / 1000 for r in rows], "o-", label="E2EL p50")
+ax.set_title("Latency (lower is better)")
+ax.set_ylabel("seconds")
+ax.set_xlabel("watermark fraction")
+ax.set_xticks(idx); ax.set_xticklabels(xt)
+ax.legend()
+
+fig.tight_layout(rect=(0, 0, 1, 0.95))
+out_png = os.path.join(outdir, "watermark_results.png")
+fig.savefig(out_png, dpi=120)
+print(f"\nWrote plot: {out_png}")
+PY

build_rust.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# Build the vllm-rs Rust frontend binary.
+# Build vLLM Rust artifacts and install them into the vllm package.
 # Usage: ./build_rust.sh [--debug]
 #
 # By default builds in release mode. Pass --debug for faster compile times

cmake/cpu_extension.cmake

@@ -166,6 +166,10 @@ elseif (S390_FOUND)
         "-mtune=native")
 elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
     message(STATUS "RISC-V detected")
+    if(DEFINED VLLM_RVV_VLEN AND NOT VLLM_RVV_VLEN GREATER 0)
+        message(FATAL_ERROR
+            "VLLM_RVV_VLEN must be a positive integer; got '${VLLM_RVV_VLEN}'")
+    endif()
     # VLLM_RVV_VLEN selects the target VLEN. Auto-detected from /proc/cpuinfo
     # by default; override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256.
     if(NOT DEFINED VLLM_RVV_VLEN)
@@ -189,8 +193,7 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
                 "RISC-V RVV is available but VLEN could not be auto-detected. "
                 "Please specify VLEN explicitly:\n"
                 "  -DVLLM_RVV_VLEN=128   (for VLEN=128 hardware)\n"
-                "  -DVLLM_RVV_VLEN=256   (for VLEN=256 hardware, e.g. Spacemit X100)\n"
-                "  -DVLLM_RVV_VLEN=0     (force scalar, no RVV)")
+                "  -DVLLM_RVV_VLEN=256   (for VLEN=256 hardware, e.g. Spacemit X100)")
         endif()
     endif()
     if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0)
@@ -219,7 +222,7 @@ endif()
 
 
 # Build oneDNN for GEMM kernels
-if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
+if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND OR RVV_FP16_FOUND OR RVV_BF16_FOUND)
     # Fetch and build Arm Compute Library (ACL) as oneDNN's backend for AArch64
     # TODO [fadara01]: remove this once ACL can be fetched and built automatically as a dependency of oneDNN
     set(ONEDNN_AARCH64_USE_ACL OFF CACHE BOOL "")
@@ -435,6 +438,12 @@ if(USE_ONEDNN)
         ${VLLM_EXT_SRC})
 endif()
 
+if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
+    set(VLLM_EXT_SRC
+        "csrc/cpu/sgl-kernels/gemm_int4.cpp"
+        ${VLLM_EXT_SRC})
+endif()
+
 if (ENABLE_X86_ISA)
     set(VLLM_EXT_SRC_SGL
         "csrc/cpu/sgl-kernels/conv.cpp"

cmake/external_projects/fmha_sm100.cmake

@@ -0,0 +1,50 @@
+include(FetchContent)
+
+# If FMHA_SM100_SRC_DIR is set, fmha_sm100 is installed from that directory
+# instead of downloading. This is useful for local MSA development.
+if(DEFINED ENV{FMHA_SM100_SRC_DIR})
+  set(FMHA_SM100_SRC_DIR $ENV{FMHA_SM100_SRC_DIR})
+endif()
+
+if(FMHA_SM100_SRC_DIR)
+  FetchContent_Declare(
+    fmha_sm100
+    SOURCE_DIR ${FMHA_SM100_SRC_DIR}
+    CONFIGURE_COMMAND ""
+    BUILD_COMMAND ""
+  )
+else()
+  FetchContent_Declare(
+    fmha_sm100
+    GIT_REPOSITORY https://github.com/vllm-project/MSA.git
+    GIT_TAG 544eee5e09ae2dfa774d5b06739013f9b7402c57
+    GIT_PROGRESS TRUE
+    CONFIGURE_COMMAND ""
+    BUILD_COMMAND ""
+  )
+endif()
+
+FetchContent_GetProperties(fmha_sm100)
+if(NOT fmha_sm100_POPULATED)
+  FetchContent_Populate(fmha_sm100)
+endif()
+message(STATUS "fmha_sm100 is available at ${fmha_sm100_SOURCE_DIR}")
+
+add_custom_target(fmha_sm100)
+
+install(FILES
+  "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100/__init__.py"
+  "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100/sparse.py"
+  DESTINATION vllm/third_party/fmha_sm100
+  COMPONENT fmha_sm100)
+
+install(DIRECTORY "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100/cute/"
+  DESTINATION vllm/third_party/fmha_sm100/cute
+  COMPONENT fmha_sm100
+  FILES_MATCHING
+    REGEX "/__pycache__(/.*)?$" EXCLUDE
+    REGEX ".*\\.pyc$" EXCLUDE
+    PATTERN "example.py" EXCLUDE
+    PATTERN "test_*.py" EXCLUDE
+    PATTERN "*.py"
+    PATTERN "build_k2q_csr.cu")

cmake/external_projects/qutlass.cmake

@@ -32,21 +32,33 @@ endif()
 message(STATUS "[QUTLASS] QuTLASS is available at ${qutlass_SOURCE_DIR}")
 
 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-  cuda_archs_loose_intersection(QUTLASS_ARCHS "10.0f;12.0f" "${CUDA_ARCHS}")
+  cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0f" "${CUDA_ARCHS}")
+  cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}")
 else()
-  cuda_archs_loose_intersection(QUTLASS_ARCHS "12.0a;12.1a;10.0a;10.3a" "${CUDA_ARCHS}")
+  cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
+  cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0a;10.3a" "${CUDA_ARCHS}")
+endif()
+
+# QUTLASS uses TARGET_CUDA_ARCH as a single preprocessor selector for all its
+# sources. Do not compile a mixed SM100/SM120 arch list with one selector; prefer
+# SM100 when both families are requested because that is the primary deployed
+# target for this extension today.
+if(QUTLASS_SM100_ARCHS)
+  set(QUTLASS_ARCHS "${QUTLASS_SM100_ARCHS}")
+  set(QUTLASS_TARGET_CC 100)
+  if(QUTLASS_SM120_ARCHS)
+    message(WARNING
+      "[QUTLASS] Both SM100 and SM120 archs were requested; selecting SM100 "
+      "because TARGET_CUDA_ARCH is a single compile-time selector.")
+  endif()
+elseif(QUTLASS_SM120_ARCHS)
+  set(QUTLASS_ARCHS "${QUTLASS_SM120_ARCHS}")
+  set(QUTLASS_TARGET_CC 120)
+else()
+  set(QUTLASS_ARCHS)
 endif()
 
 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
-
-  if(QUTLASS_ARCHS MATCHES "10\\.(0a|3a|0f)")
-    set(QUTLASS_TARGET_CC 100)
-  elseif(QUTLASS_ARCHS MATCHES "12\\.[01][af]?")
-    set(QUTLASS_TARGET_CC 120)
-  else()
-    message(FATAL_ERROR "[QUTLASS] internal error parsing CUDA_ARCHS='${QUTLASS_ARCHS}'.")
-  endif()
-
   set(QUTLASS_SOURCES
     ${qutlass_SOURCE_DIR}/qutlass/csrc/bindings.cpp
     ${qutlass_SOURCE_DIR}/qutlass/csrc/gemm.cu

cmake/external_projects/vllm_flash_attn.cmake

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

cmake/utils.cmake

@@ -487,9 +487,9 @@ endfunction()
 
 function(cuda_archs_sm90plus OUT_CUDA_ARCHS TGT_CUDA_ARCHS)
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f" "${TGT_CUDA_ARCHS}")
+    cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f;12.0f" "${TGT_CUDA_ARCHS}")
   else()
-    cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a" "${TGT_CUDA_ARCHS}")
+    cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${TGT_CUDA_ARCHS}")
   endif()
   set(${OUT_CUDA_ARCHS} ${_archs} PARENT_SCOPE)
 endfunction()

csrc/cpu/cpu_attn_impl.hpp

@@ -822,8 +822,8 @@ struct AttentionInput {
       logits_buffer_t *__restrict__ logits_buffer,                          \
       float *__restrict__ partial_q_buffer, float *__restrict__ max_buffer, \
       float *__restrict__ sum_buffer, int32_t *__restrict__ block_table,    \
-      const int32_t kv_tile_start_pos, const int32_t kv_tile_end_pos,       \
-      const int32_t kv_tile_token_num,                                      \
+      const int32_t kv_end_pos, const int32_t kv_tile_start_pos,            \
+      const int32_t kv_tile_end_pos, const int32_t kv_tile_token_num,       \
       const int64_t kv_cache_num_blocks_stride, const int32_t q_head_num,   \
       const int32_t q_token_num, const int32_t q_tile_start_pos,            \
       const int32_t q_heads_per_kv, const int32_t block_size,               \
@@ -834,7 +834,7 @@ struct AttentionInput {
 
 #define CPU_ATTENTION_PARAMS                                                  \
   q_heads_buffer, k_head_cache_ptr, v_head_cache_ptr, logits_buffer,          \
-      partial_q_buffer, max_buffer, sum_buffer, block_table,                  \
+      partial_q_buffer, max_buffer, sum_buffer, block_table, kv_end_pos,      \
       kv_tile_start_pos, kv_tile_end_pos, kv_tile_token_num,                  \
       kv_cache_num_blocks_stride, q_head_num, q_token_num, q_tile_start_pos,  \
       q_heads_per_kv, block_size, left_window_size, right_window_size, scale, \
@@ -917,6 +917,7 @@ class AttentionMainLoop {
     //  - max_buffer: [MaxQHeadNumPerIteration, 1], store max logits
     //  - sum_buffer: [MaxQHeadNumPerIteration, 1], store sum of exp
     //  - block_table
+    //  - kv_end_pos: un-aligned end position of KV cache
     //  - kv_tile_start_pos: start position of KV cache, aligned to
     //  BlockSizeAlignment
     //  - kv_tile_end_pos: end position of KV cache, aligned to
@@ -1043,7 +1044,7 @@ class AttentionMainLoop {
         }
 
         apply_mask(logits_buffer, kv_tile_token_num, q_tile_start_pos,
-                   kv_tile_start_pos, kv_tile_end_pos, q_token_num,
+                   kv_end_pos, kv_tile_start_pos, kv_tile_end_pos, q_token_num,
                    q_heads_per_kv, left_window_size, right_window_size);
 
         // if (debug_info){
@@ -1126,7 +1127,7 @@ class AttentionMainLoop {
 
     void apply_mask(logits_buffer_t* __restrict__ logits_buffer,
                     const int64_t logits_buffer_stride,
-                    const int32_t q_tile_start_pos,
+                    const int32_t q_tile_start_pos, const int32_t kv_end_pos,
                     const int32_t kv_tile_start_pos,
                     const int32_t kv_tile_end_pos, const int32_t q_token_num,
                     const int32_t q_heads_per_kv,
@@ -1154,7 +1155,7 @@ class AttentionMainLoop {
                            std::max(kv_tile_start_pos,
                                     curr_token_pos + sliding_window_right + 1));
           }
-          return pos;
+          return std::min(pos, kv_end_pos);
         }();
 
         int32_t left_invalid_token_num = left_kv_pos - kv_tile_start_pos;
@@ -1789,7 +1790,7 @@ class AttentionMainLoop {
                   attn_impl.template execute_attention<Attention>(
                       curr_q_heads_buffer, curr_k_cache, curr_v_cache,
                       logits_buffer, curr_partial_q_buffer, curr_max_buffer,
-                      curr_sum_buffer, curr_block_table,
+                      curr_sum_buffer, curr_block_table, kv_end_pos,
                       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,

csrc/cpu/cpu_types_riscv_defs.hpp

@@ -57,6 +57,10 @@ typedef RVVTYPE(vfloat32, LMUL_512, _t) fixed_fp32x16_t
 typedef RVVTYPE(vfloat32, LMUL_1024, _t) fixed_fp32x32_t
     __attribute__((riscv_rvv_vector_bits(1024)));
 
+// int8
+typedef RVVTYPE(vint8, LMUL_128, _t) fixed_i8x16_t
+    __attribute__((riscv_rvv_vector_bits(128)));
+
 // int32
 typedef RVVTYPE(vint32, LMUL_256, _t) fixed_i32x8_t
     __attribute__((riscv_rvv_vector_bits(256)));

csrc/cpu/cpu_types_riscv_impl.hpp

@@ -9,10 +9,14 @@
 
 #include <algorithm>
 #include <cmath>
+#include <cstdint>
 #include <cstring>
 #include <iostream>
 #include <limits>
 #include <torch/all.h>
+
+#include "float_convert.hpp"
+
 namespace vec_op {
 
 // FP8 KV cache is not supported on RISC-V. These tag types and the
@@ -245,8 +249,7 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
     const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
     float tmp[8];
     for (int i = 0; i < 8; ++i) {
-      uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
-      std::memcpy(&tmp[i], &v, 4);
+      tmp[i] = bf16_to_float(u16[i]);
     }
     reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_256)(tmp, 8);
   }
@@ -256,9 +259,7 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
     RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
     uint16_t* u16 = static_cast<uint16_t*>(ptr);
     for (int i = 0; i < 8; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
+      u16[i] = float_to_bf16(tmp[i]);
     }
   }
   void save(void* ptr, int elem_num) const {
@@ -266,9 +267,7 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
     RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
     uint16_t* u16 = static_cast<uint16_t*>(ptr);
     for (int i = 0; i < elem_num; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
+      u16[i] = float_to_bf16(tmp[i]);
     }
   }
   void save_strided(void* ptr, ptrdiff_t stride) const {
@@ -277,10 +276,8 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
     uint8_t* u8 = static_cast<uint8_t*>(ptr);
     ptrdiff_t byte_stride = stride * sizeof(uint16_t);
     for (int i = 0; i < 8; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      uint16_t val = static_cast<uint16_t>(v >> 16);
-      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
+      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) =
+          float_to_bf16(tmp[i]);
     }
   }
 };
@@ -292,8 +289,7 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
     const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
     float tmp[16];
     for (int i = 0; i < 16; ++i) {
-      uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
-      std::memcpy(&tmp[i], &v, 4);
+      tmp[i] = bf16_to_float(u16[i]);
     }
     reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16);
   }
@@ -306,9 +302,7 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
     RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
     uint16_t* u16 = static_cast<uint16_t*>(ptr);
     for (int i = 0; i < 16; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
+      u16[i] = float_to_bf16(tmp[i]);
     }
   }
   void save(void* ptr, int elem_num) const {
@@ -316,9 +310,7 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
     RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
     uint16_t* u16 = static_cast<uint16_t*>(ptr);
     for (int i = 0; i < elem_num; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
+      u16[i] = float_to_bf16(tmp[i]);
     }
   }
   void save_strided(void* ptr, ptrdiff_t stride) const {
@@ -327,10 +319,8 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
     uint8_t* u8 = static_cast<uint8_t*>(ptr);
     ptrdiff_t byte_stride = stride * sizeof(uint16_t);
     for (int i = 0; i < 16; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      uint16_t val = static_cast<uint16_t>(v >> 16);
-      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
+      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) =
+          float_to_bf16(tmp[i]);
     }
   }
 };
@@ -343,8 +333,7 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
     const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
     float tmp[32];
     for (int i = 0; i < 32; ++i) {
-      uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
-      std::memcpy(&tmp[i], &v, 4);
+      tmp[i] = bf16_to_float(u16[i]);
     }
     reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp, 32);
   }
@@ -371,9 +360,7 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
     RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
     uint16_t* u16 = static_cast<uint16_t*>(ptr);
     for (int i = 0; i < 32; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
+      u16[i] = float_to_bf16(tmp[i]);
     }
   }
 
@@ -382,9 +369,7 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
     RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
     uint16_t* u16 = static_cast<uint16_t*>(ptr);
     for (int i = 0; i < elem_num; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
+      u16[i] = float_to_bf16(tmp[i]);
     }
   }
 
@@ -394,10 +379,8 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
     uint8_t* u8 = static_cast<uint8_t*>(ptr);
     ptrdiff_t byte_stride = stride * sizeof(uint16_t);
     for (int i = 0; i < 32; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      uint16_t val = static_cast<uint16_t>(v >> 16);
-      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
+      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) =
+          float_to_bf16(tmp[i]);
     }
   }
 };
@@ -734,10 +717,18 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
     return FP32Vec16(
         RVVI(__riscv_vfmax_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
   }
+  FP32Vec16 max(const FP32Vec16& b, const int elem_num) const {
+    return FP32Vec16(
+        RVVI(__riscv_vfmax_vv_f32, LMUL_512)(reg, b.reg, elem_num));
+  }
   FP32Vec16 min(const FP32Vec16& b) const {
     return FP32Vec16(
         RVVI(__riscv_vfmin_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
   }
+  FP32Vec16 min(const FP32Vec16& b, const int elem_num) const {
+    return FP32Vec16(
+        RVVI(__riscv_vfmin_vv_f32, LMUL_512)(reg, b.reg, elem_num));
+  }
   FP32Vec16 abs() const {
     return FP32Vec16(RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM));
   }
@@ -867,6 +858,27 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
   }
 };
 
+struct INT8Vec16 : public Vec<INT8Vec16> {
+  constexpr static int VEC_ELEM_NUM = 16;
+  fixed_i8x16_t reg;
+
+  explicit INT8Vec16(const FP32Vec16& vec) {
+    auto i32_vec =
+        RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_512)(vec.reg, VEC_ELEM_NUM);
+    auto i16_vec = RVVI(__riscv_vnclip_wx_i16, LMUL_256)(
+        i32_vec, 0, __RISCV_VXRM_RNU, VEC_ELEM_NUM);
+    reg = RVVI(__riscv_vnclip_wx_i8, LMUL_128)(i16_vec, 0, __RISCV_VXRM_RNU,
+                                               VEC_ELEM_NUM);
+  }
+
+  void save(int8_t* ptr) const {
+    RVVI(__riscv_vse8_v_i8, LMUL_128)(ptr, reg, VEC_ELEM_NUM);
+  }
+  void save(int8_t* ptr, int elem_num) const {
+    RVVI(__riscv_vse8_v_i8, LMUL_128)(ptr, reg, elem_num);
+  }
+};
+
 // ============================================================================
 // Type Traits & Global Helpers
 // ============================================================================
@@ -956,9 +968,7 @@ inline BF16Vec16::BF16Vec16(const FP32Vec16& v)
 #else
 template <>
 inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
-  uint32_t val;
-  std::memcpy(&val, &v, 4);
-  *reinterpret_cast<uint16_t*>(ptr) = static_cast<uint16_t>(val >> 16);
+  *reinterpret_cast<uint16_t*>(ptr) = float_to_bf16(v);
 }
 inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {}
 inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {}

csrc/cpu/cpu_types_vxe.hpp

@@ -3,7 +3,9 @@
 #define CPU_TYPES_VXE_HPP
 
 #include <vecintrin.h>
+#include <bit>
 #include <cmath>
+#include <cstdint>
 #include <limits>
 #include <torch/all.h>
 namespace vec_op {
@@ -817,8 +819,7 @@ inline void storeFP32<::c10::Half>(float v, ::c10::Half* ptr) {
   // intrinsics for FP32 to FP16 conversion does not use IEEE rounding and can
   // produce incorrect results for some inputs. Process each of the 4 vectors
   // separately.
-  uint32_t in;
-  std::memcpy(&in, &v, sizeof(in));
+  uint32_t in = std::bit_cast<uint32_t>(v);
 
   uint32_t s = (in & 0x80000000) >> 16;  // Sign
   uint32_t e = (in & 0x7F800000) >> 23;  // Exponent

csrc/cpu/float_convert.hpp

@@ -1,14 +1,15 @@
+#pragma once
 
-static float bf16_to_float(uint16_t bf16) {
+#include <bit>
+#include <cstdint>
+
+inline float bf16_to_float(uint16_t bf16) {
   uint32_t bits = static_cast<uint32_t>(bf16) << 16;
-  float fp32;
-  std::memcpy(&fp32, &bits, sizeof(fp32));
-  return fp32;
+  return std::bit_cast<float>(bits);
 }
 
-static uint16_t float_to_bf16(float fp32) {
-  uint32_t bits;
-  std::memcpy(&bits, &fp32, sizeof(fp32));
+inline uint16_t float_to_bf16(float fp32) {
+  uint32_t bits = std::bit_cast<uint32_t>(fp32);
   return static_cast<uint16_t>(bits >> 16);
 }
 
@@ -18,14 +19,13 @@ static uint16_t float_to_bf16(float fp32) {
  * Codes below copied from
  * https://github.com/PrincetonVision/marvin/tree/master/tools/tensorIO_matlab
  *************************************************/
-static uint16_t float_to_fp16(float fp32) {
+inline uint16_t float_to_fp16(float fp32) {
   uint16_t fp16;
 
-  unsigned x;
   unsigned u, remainder, shift, lsb, lsb_s1, lsb_m1;
   unsigned sign, exponent, mantissa;
 
-  std::memcpy(&x, &fp32, sizeof(fp32));
+  uint32_t x = std::bit_cast<uint32_t>(fp32);
   u = (x & 0x7fffffff);
 
   // Get rid of +NaN/-NaN case first.
@@ -77,12 +77,11 @@ static uint16_t float_to_fp16(float fp32) {
   return fp16;
 }
 
-static float fp16_to_float(uint16_t fp16) {
+inline float fp16_to_float(uint16_t fp16) {
   unsigned sign = ((fp16 >> 15) & 1);
   unsigned exponent = ((fp16 >> 10) & 0x1f);
   unsigned mantissa = ((fp16 & 0x3ff) << 13);
-  int temp;
-  float fp32;
+  uint32_t temp;
   if (exponent == 0x1f) { /* NaN or Inf */
     mantissa = (mantissa ? (sign = 0, 0x7fffff) : 0);
     exponent = 0xff;
@@ -101,6 +100,5 @@ static float fp16_to_float(uint16_t fp16) {
     exponent += 0x70;
   }
   temp = ((sign << 31) | (exponent << 23) | mantissa);
-  std::memcpy(&fp32, &temp, sizeof(temp));
-  return fp32;
+  return std::bit_cast<float>(temp);
 }

csrc/cpu/generate_cpu_attn_dispatch.py

@@ -11,7 +11,7 @@ import os
 HEAD_DIMS_32 = [32, 64, 96, 128, 160, 192, 224, 256, 512]
 
 # Head dimensions divisible by 16 but not 32 (VEC16 only)
-HEAD_DIMS_16 = [80, 112]
+HEAD_DIMS_16 = [48, 80, 112]
 
 # ISA types
 ISA_TYPES = {

csrc/cpu/sgl-kernels/gemm_int4.cpp

@@ -268,6 +268,23 @@ void _dequant_gemm_accum_small_M(
   _dequant_gemm_accum_small_M<M, N, ldb, sym_quant_act>(C, A, scales_a, qzeros_a, B, scales_b, qzeros_b, K, lda, ldc);
 #endif
 
+template <int64_t N, int64_t ldb>
+inline int32_t load_uint4_vnni(const uint8_t* __restrict__ B, int64_t k, int64_t n) {
+  // B is packed as [_block_k / 4, N / 2, 4] for VNNI4. Each byte stores two
+  // columns from adjacent 8-column groups for one K lane.
+  constexpr int64_t n_group_size = 8;
+  constexpr int64_t vnni_size = 4;
+  static_assert(N % (2 * n_group_size) == 0);
+
+  int64_t n_group = n / n_group_size;
+  int64_t ni = n % n_group_size;
+  int64_t ki = k % vnni_size;
+  int64_t k_base = k - ki;
+  int64_t packed_n = (n_group / 2) * n_group_size + ni;
+  uint8_t packed = B[k_base * ldb + packed_n * vnni_size + ki];
+  return (n_group % 2 == 0) ? (packed & 0x0f) : ((packed >> 4) & 0x0f);
+}
+
 template <int64_t N, int64_t ldb, bool sym_quant_act>
 void _dequant_gemm_accum(
     float* C,
@@ -321,7 +338,24 @@ void _dequant_gemm_accum(
   } else
 #endif
   {
-    TORCH_CHECK(false, "tinygemm_kernel: scalar path not implemented!");
+    for (int64_t m = 0; m < M; ++m) {
+      for (int64_t n = 0; n < N; ++n) {
+        int32_t acc = 0;
+        for (int64_t k = 0; k < K; ++k) {
+          int32_t b = load_uint4_vnni<N, ldb>(B, k, n) - qzeros_b[n];
+          if constexpr (sym_quant_act) {
+            const int8_t* A_s8 = reinterpret_cast<const int8_t*>(A);
+            acc += static_cast<int32_t>(A_s8[m * lda + k]) * b;
+          } else {
+            acc += static_cast<int32_t>(A[m * lda + k]) * b;
+          }
+        }
+        if constexpr (!sym_quant_act) {
+          acc -= qzeros_a[m] * compensation[n];
+        }
+        C[m * ldc + n] += static_cast<float>(acc) * scales_a[m] * scales_b[n];
+      }
+    }
   }
 }
 
@@ -496,9 +530,11 @@ void _da8w4_linear_impl(
         store_out<out_dtype, BLOCK_N>(C_tmp, output + mci * block_m * N + nc * BLOCK_N, m_size, N /*lda*/);
       }
     }
+#if defined(CPU_CAPABILITY_AVX512)
     if (use_brgemm) {
       at::native::cpublas::brgemm_release();
     }
+#endif
   });
 }
 

csrc/cpu/sgl-kernels/vec.h

@@ -245,7 +245,7 @@ quantize_row_int8(uint8_t* __restrict__ Aq, float& As, const scalar_t* __restric
 
   for (int64_t k = 0; k < K; ++k) {
     const float val = static_cast<float>(A[k]) * inv_scale;
-    Aq[k] = (uint8_t)(std::round(val)) + 128;
+    Aq[k] = static_cast<uint8_t>(static_cast<int32_t>(std::round(val)) + 128);
   }
   As = scale;
 }

csrc/cpu/torch_bindings.cpp

@@ -329,8 +329,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);
 
   // Quantization
-#if defined(__AVX512F__) || defined(__AVX2__) || \
-    (defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__)
+#if defined(__AVX512F__) || defined(__AVX2__) ||                               \
+    (defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__) || \
+    defined(__riscv_v)
   // Helper function to release oneDNN handlers
   ops.def("release_dnnl_matmul_handler(int handler) -> ()",
           &release_dnnl_matmul_handler);
@@ -428,19 +429,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.impl("int8_scaled_mm_with_quant", torch::kCPU,
            &int8_scaled_mm_with_quant);
 
-  // Adapted from sglang: INT4 W4A8 kernels
-  ops.def(
-      "convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor "
-      "scales, int quant_method_4bit) -> (Tensor, "
-      "Tensor, Tensor)");
-  ops.impl("convert_weight_packed_scale_zp", torch::kCPU,
-           &convert_weight_packed_scale_zp);
-
-  ops.def(
-      "int4_scaled_mm_cpu(Tensor(a0!) x, Tensor(a1!) w, Tensor(a2!) w_zeros, "
-      "Tensor(a3!) w_scales, Tensor? bias) -> Tensor");
-  ops.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu);
-
   // Adapted from sglang: FP8 W8A16 kernel
   ops.def(
       "fp8_scaled_mm_cpu(Tensor(a0!) mat1, Tensor(a1!) mat2, Tensor(a2!) "
@@ -467,6 +455,23 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
 #endif
 
+#if (defined(__AVX512BF16__) && defined(__AVX512F__) && \
+     defined(__AVX512VNNI__)) ||                        \
+    defined(__riscv)
+  // Adapted from sglang: INT4 W4A8 kernels
+  ops.def(
+      "convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor "
+      "scales, int quant_method_4bit) -> (Tensor, "
+      "Tensor, Tensor)");
+  ops.impl("convert_weight_packed_scale_zp", torch::kCPU,
+           &convert_weight_packed_scale_zp);
+
+  ops.def(
+      "int4_scaled_mm_cpu(Tensor(a0!) x, Tensor(a1!) w, Tensor(a2!) w_zeros, "
+      "Tensor(a3!) w_scales, Tensor? bias) -> Tensor");
+  ops.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu);
+#endif
+
   // Adapted from sglang: GDN kernels
   ops.def(
       "chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, "

csrc/cutlass_extensions/vllm_cutlass_library_extension.py

@@ -57,13 +57,13 @@ VLLMDataTypeVLLMScalarTypeTag: dict[VLLMDataType | DataType, str] = {
 }
 
 VLLMDataTypeTorchDataTypeTag: dict[VLLMDataType | DataType, str] = {
-    DataType.u8: "at::ScalarType::Byte",
-    DataType.s8: "at::ScalarType::Char",
-    DataType.e4m3: "at::ScalarType::Float8_e4m3fn",
-    DataType.s32: "at::ScalarType::Int",
-    DataType.f16: "at::ScalarType::Half",
-    DataType.bf16: "at::ScalarType::BFloat16",
-    DataType.f32: "at::ScalarType::Float",
+    DataType.u8: "torch::headeronly::ScalarType::Byte",
+    DataType.s8: "torch::headeronly::ScalarType::Char",
+    DataType.e4m3: "torch::headeronly::ScalarType::Float8_e4m3fn",
+    DataType.s32: "torch::headeronly::ScalarType::Int",
+    DataType.f16: "torch::headeronly::ScalarType::Half",
+    DataType.bf16: "torch::headeronly::ScalarType::BFloat16",
+    DataType.f32: "torch::headeronly::ScalarType::Float",
 }
 
 VLLMKernelScheduleTag: dict[MixedInputKernelScheduleType | KernelScheduleType, str] = {

csrc/libtorch_stable/activation_kernels.cu

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

csrc/libtorch_stable/cuda_vec_utils.cuh

@@ -21,7 +21,7 @@
 // together enable 256-bit (v8.u32) PTX load/store instructions.
 // Use for PTX instruction selection with architecture fallback paths.
 #if !defined(USE_ROCM) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \
-    defined(CUDA_VERSION) && CUDA_VERSION >= 12090
+    defined(CUDART_VERSION) && CUDART_VERSION >= 12090
   #define VLLM_256B_PTX_ENABLED 1
 #else
   #define VLLM_256B_PTX_ENABLED 0

csrc/libtorch_stable/dispatch_utils.h

@@ -30,6 +30,28 @@
   THO_DISPATCH_SWITCH(TYPE, NAME,                            \
                       VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
 
+#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(...)                  \
+  THO_DISPATCH_CASE(torch::headeronly::ScalarType::Byte, __VA_ARGS__)  \
+  THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__)  \
+  THO_DISPATCH_CASE(torch::headeronly::ScalarType::Short, __VA_ARGS__) \
+  THO_DISPATCH_CASE(torch::headeronly::ScalarType::Int, __VA_ARGS__)   \
+  THO_DISPATCH_CASE(torch::headeronly::ScalarType::Long, __VA_ARGS__)
+
+#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(...)      \
+  VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__)                 \
+  THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt16, __VA_ARGS__) \
+  THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt32, __VA_ARGS__) \
+  THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt64, __VA_ARGS__)
+
+#define VLLM_STABLE_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \
+  THO_DISPATCH_SWITCH(TYPE, NAME,                            \
+                      VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__))
+
+#define VLLM_STABLE_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(TYPE, NAME, ...) \
+  THO_DISPATCH_SWITCH(                                                    \
+      TYPE, NAME,                                                         \
+      VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(__VA_ARGS__))
+
 // FP8 type dispatch - ROCm uses FNUZ format, CUDA uses OCP format
 #ifdef USE_ROCM
   #define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...)                 \

csrc/libtorch_stable/fp32_router_gemm.cu

@@ -175,49 +175,52 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a,
 }
 
 // ---------------------------------------------------------------------------
-// Explicit instantiations: M=1..32, E=256, H=3072, for both input types
+// Explicit instantiations: M=1..32, for both input types, for the supported
+// (E, H) pairs:  (256, 3072) [MiniMax-M2/M2.5]  and  (128, 6144) [MiniMax-M3].
 // ---------------------------------------------------------------------------
 
-#define INSTANTIATE(T, M)                              \
-  template void invokeFp32RouterGemm<T, M, 256, 3072>( \
-      float*, T const*, float const*, cudaStream_t);
+#define INSTANTIATE(T, M, E, H)                                    \
+  template void invokeFp32RouterGemm<T, M, E, H>(float*, T const*, \
+                                                 float const*, cudaStream_t);
 
-#define INSTANTIATE_ALL(T) \
-  INSTANTIATE(T, 1)        \
-  INSTANTIATE(T, 2)        \
-  INSTANTIATE(T, 3)        \
-  INSTANTIATE(T, 4)        \
-  INSTANTIATE(T, 5)        \
-  INSTANTIATE(T, 6)        \
-  INSTANTIATE(T, 7)        \
-  INSTANTIATE(T, 8)        \
-  INSTANTIATE(T, 9)        \
-  INSTANTIATE(T, 10)       \
-  INSTANTIATE(T, 11)       \
-  INSTANTIATE(T, 12)       \
-  INSTANTIATE(T, 13)       \
-  INSTANTIATE(T, 14)       \
-  INSTANTIATE(T, 15)       \
-  INSTANTIATE(T, 16)       \
-  INSTANTIATE(T, 17)       \
-  INSTANTIATE(T, 18)       \
-  INSTANTIATE(T, 19)       \
-  INSTANTIATE(T, 20)       \
-  INSTANTIATE(T, 21)       \
-  INSTANTIATE(T, 22)       \
-  INSTANTIATE(T, 23)       \
-  INSTANTIATE(T, 24)       \
-  INSTANTIATE(T, 25)       \
-  INSTANTIATE(T, 26)       \
-  INSTANTIATE(T, 27)       \
-  INSTANTIATE(T, 28)       \
-  INSTANTIATE(T, 29)       \
-  INSTANTIATE(T, 30)       \
-  INSTANTIATE(T, 31)       \
-  INSTANTIATE(T, 32)
+#define INSTANTIATE_ALL(T, E, H) \
+  INSTANTIATE(T, 1, E, H)        \
+  INSTANTIATE(T, 2, E, H)        \
+  INSTANTIATE(T, 3, E, H)        \
+  INSTANTIATE(T, 4, E, H)        \
+  INSTANTIATE(T, 5, E, H)        \
+  INSTANTIATE(T, 6, E, H)        \
+  INSTANTIATE(T, 7, E, H)        \
+  INSTANTIATE(T, 8, E, H)        \
+  INSTANTIATE(T, 9, E, H)        \
+  INSTANTIATE(T, 10, E, H)       \
+  INSTANTIATE(T, 11, E, H)       \
+  INSTANTIATE(T, 12, E, H)       \
+  INSTANTIATE(T, 13, E, H)       \
+  INSTANTIATE(T, 14, E, H)       \
+  INSTANTIATE(T, 15, E, H)       \
+  INSTANTIATE(T, 16, E, H)       \
+  INSTANTIATE(T, 17, E, H)       \
+  INSTANTIATE(T, 18, E, H)       \
+  INSTANTIATE(T, 19, E, H)       \
+  INSTANTIATE(T, 20, E, H)       \
+  INSTANTIATE(T, 21, E, H)       \
+  INSTANTIATE(T, 22, E, H)       \
+  INSTANTIATE(T, 23, E, H)       \
+  INSTANTIATE(T, 24, E, H)       \
+  INSTANTIATE(T, 25, E, H)       \
+  INSTANTIATE(T, 26, E, H)       \
+  INSTANTIATE(T, 27, E, H)       \
+  INSTANTIATE(T, 28, E, H)       \
+  INSTANTIATE(T, 29, E, H)       \
+  INSTANTIATE(T, 30, E, H)       \
+  INSTANTIATE(T, 31, E, H)       \
+  INSTANTIATE(T, 32, E, H)
 
-INSTANTIATE_ALL(float)
-INSTANTIATE_ALL(__nv_bfloat16)
+INSTANTIATE_ALL(float, 256, 3072)
+INSTANTIATE_ALL(__nv_bfloat16, 256, 3072)
+INSTANTIATE_ALL(float, 128, 6144)
+INSTANTIATE_ALL(__nv_bfloat16, 128, 6144)
 
 #undef INSTANTIATE_ALL
 #undef INSTANTIATE

csrc/libtorch_stable/fp32_router_gemm_entry.cu

@@ -22,36 +22,42 @@ inline int getSMVersion() {
 
 }  // namespace
 
-static constexpr int FP32_NUM_EXPERTS = 256;
-static constexpr int FP32_HIDDEN_DIM = 3072;
 static constexpr int FP32_MAX_TOKENS = 32;
 
+// Supported (hidden_dim, num_experts) pairs (must match the instantiations in
+// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3.
+static inline bool fp32_router_gemm_supported(int hidden_dim, int num_experts) {
+  return (hidden_dim == 3072 && num_experts == 256) ||
+         (hidden_dim == 6144 && num_experts == 128);
+}
+
 // Forward declarations — 4 template params must match fp32_router_gemm.cu
 template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
 void invokeFp32RouterGemm(float* output, InputT const* mat_a,
                           float const* mat_b, cudaStream_t stream);
 
-// LoopUnroller templated on InputT
-template <typename InputT, int kBegin, int kEnd>
+// LoopUnroller templated on InputT, kNumExperts and kHiddenDim
+template <typename InputT, int kNumExperts, int kHiddenDim, int kBegin,
+          int kEnd>
 struct Fp32LoopUnroller {
   static void unroll(int num_tokens, float* output, InputT const* mat_a,
                      float const* mat_b, cudaStream_t stream) {
     if (num_tokens == kBegin) {
-      invokeFp32RouterGemm<InputT, kBegin, FP32_NUM_EXPERTS, FP32_HIDDEN_DIM>(
+      invokeFp32RouterGemm<InputT, kBegin, kNumExperts, kHiddenDim>(
           output, mat_a, mat_b, stream);
     } else {
-      Fp32LoopUnroller<InputT, kBegin + 1, kEnd>::unroll(num_tokens, output,
-                                                         mat_a, mat_b, stream);
+      Fp32LoopUnroller<InputT, kNumExperts, kHiddenDim, kBegin + 1,
+                       kEnd>::unroll(num_tokens, output, mat_a, mat_b, stream);
     }
   }
 };
 
-template <typename InputT, int kEnd>
-struct Fp32LoopUnroller<InputT, kEnd, kEnd> {
+template <typename InputT, int kNumExperts, int kHiddenDim, int kEnd>
+struct Fp32LoopUnroller<InputT, kNumExperts, kHiddenDim, kEnd, kEnd> {
   static void unroll(int num_tokens, float* output, InputT const* mat_a,
                      float const* mat_b, cudaStream_t stream) {
     if (num_tokens == kEnd) {
-      invokeFp32RouterGemm<InputT, kEnd, FP32_NUM_EXPERTS, FP32_HIDDEN_DIM>(
+      invokeFp32RouterGemm<InputT, kEnd, kNumExperts, kHiddenDim>(
           output, mat_a, mat_b, stream);
     } else {
       throw std::invalid_argument(
@@ -60,6 +66,23 @@ struct Fp32LoopUnroller<InputT, kEnd, kEnd> {
   }
 };
 
+// Dispatch over the supported (num_experts, hidden_dim) pairs.
+template <typename InputT>
+void dispatchFp32RouterGemm(int num_experts, int hidden_dim, int num_tokens,
+                            float* output, InputT const* mat_a,
+                            float const* mat_b, cudaStream_t stream) {
+  if (num_experts == 256 && hidden_dim == 3072) {
+    Fp32LoopUnroller<InputT, 256, 3072, 1, FP32_MAX_TOKENS>::unroll(
+        num_tokens, output, mat_a, mat_b, stream);
+  } else if (num_experts == 128 && hidden_dim == 6144) {
+    Fp32LoopUnroller<InputT, 128, 6144, 1, FP32_MAX_TOKENS>::unroll(
+        num_tokens, output, mat_a, mat_b, stream);
+  } else {
+    throw std::invalid_argument(
+        "fp32_router_gemm: unsupported (hidden_dim, num_experts) pair");
+  }
+}
+
 void fp32_router_gemm(
     torch::stable::Tensor& output,       // [num_tokens, num_experts]
     torch::stable::Tensor const& mat_a,  // [num_tokens, hidden_dim]
@@ -85,10 +108,10 @@ void fp32_router_gemm(
   STD_TORCH_CHECK(
       mat_a.size(1) == mat_b.size(1),
       "fp32_router_gemm: mat_a and mat_b must have the same hidden_dim");
-  STD_TORCH_CHECK(hidden_dim == FP32_HIDDEN_DIM,
-                  "fp32_router_gemm: expected hidden_dim=3072");
-  STD_TORCH_CHECK(num_experts == FP32_NUM_EXPERTS,
-                  "fp32_router_gemm: expected num_experts=256");
+  STD_TORCH_CHECK(
+      fp32_router_gemm_supported(hidden_dim, num_experts),
+      "fp32_router_gemm: supported (hidden_dim, num_experts) pairs are "
+      "(3072, 256) and (6144, 128)");
   STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS,
                   "fp32_router_gemm: num_tokens must be in [0, 32]");
   STD_TORCH_CHECK(
@@ -113,12 +136,13 @@ void fp32_router_gemm(
   if (mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
     auto const* mat_a_ptr =
         reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
-    Fp32LoopUnroller<__nv_bfloat16, 1, FP32_MAX_TOKENS>::unroll(
-        num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream);
+    dispatchFp32RouterGemm<__nv_bfloat16>(num_experts, hidden_dim, num_tokens,
+                                          out_ptr, mat_a_ptr, mat_b_ptr,
+                                          stream);
   } else {
     auto const* mat_a_ptr = reinterpret_cast<float const*>(mat_a.data_ptr());
-    Fp32LoopUnroller<float, 1, FP32_MAX_TOKENS>::unroll(
-        num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream);
+    dispatchFp32RouterGemm<float>(num_experts, hidden_dim, num_tokens, out_ptr,
+                                  mat_a_ptr, mat_b_ptr, stream);
   }
 }
 

csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu

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

csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu

@@ -18,14 +18,11 @@
  * limitations under the License.
  */
 
-#include <ATen/ATen.h>
-#include <ATen/cuda/CUDAContext.h>
+#include <cstdint>
 
 #include <cuda_bf16.h>
 #include <cuda_runtime.h>
 
-#include "dsv3_router_gemm_utils.h"
-
 // Custom FMA implementation using PTX assembly instructions
 __device__ __forceinline__ void fma(float2& d, float2 const& a, float2 const& b,
                                     float2 const& c) {

csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu

@@ -18,15 +18,25 @@
  * limitations under the License.
  */
 
-#include <ATen/ATen.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <torch/all.h>
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+
+#include "libtorch_stable/torch_utils.h"
 
 #include <cuda_bf16.h>
 #include <cuda_runtime.h>
 
-#include "core/registration.h"
-#include "dsv3_router_gemm_utils.h"
+#include <stdexcept>
+
+namespace {
+
+inline int getSMVersion() {
+  auto* props = get_device_prop();
+  return props->major * 10 + props->minor;
+}
+
+}  // namespace
 
 static constexpr int DEFAULT_NUM_EXPERTS = 256;
 static constexpr int KIMI_K2_NUM_EXPERTS = 384;
@@ -98,40 +108,47 @@ struct LoopUnroller<kEnd, kEnd, kNumExperts, kHiddenDim> {
   }
 };
 
-void dsv3_router_gemm(at::Tensor& output,       // [num_tokens, num_experts]
-                      const at::Tensor& mat_a,  // [num_tokens, hidden_dim]
-                      const at::Tensor& mat_b   // [num_experts, hidden_dim]
+void dsv3_router_gemm(
+    torch::stable::Tensor& output,       // [num_tokens, num_experts]
+    torch::stable::Tensor const& mat_a,  // [num_tokens, hidden_dim]
+    torch::stable::Tensor const& mat_b   // [num_experts, hidden_dim]
 ) {
-  TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2);
+  STD_TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2);
 
   const int num_tokens = mat_a.size(0);
   const int num_experts = mat_b.size(0);
   const int hidden_dim = mat_a.size(1);
 
-  TORCH_CHECK(mat_a.size(1) == mat_b.size(1),
-              "mat_a and mat_b must have the same hidden_dim");
-  TORCH_CHECK(hidden_dim == DEFAULT_HIDDEN_DIM,
-              "Expected hidden_dim=", DEFAULT_HIDDEN_DIM,
-              ", but got hidden_dim=", hidden_dim);
-  TORCH_CHECK(
+  STD_TORCH_CHECK(mat_a.size(1) == mat_b.size(1),
+                  "mat_a and mat_b must have the same hidden_dim");
+  STD_TORCH_CHECK(hidden_dim == DEFAULT_HIDDEN_DIM,
+                  "Expected hidden_dim=", DEFAULT_HIDDEN_DIM,
+                  ", but got hidden_dim=", hidden_dim);
+  STD_TORCH_CHECK(
       num_experts == DEFAULT_NUM_EXPERTS || num_experts == KIMI_K2_NUM_EXPERTS,
       "Expected num_experts=", DEFAULT_NUM_EXPERTS,
       " or num_experts=", KIMI_K2_NUM_EXPERTS,
       ", but got num_experts=", num_experts);
-  TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
-              "currently num_tokens must be less than or equal to 16 for "
-              "router_gemm");
-  TORCH_CHECK(mat_a.dtype() == at::kBFloat16, "mat_a must be bf16");
-  TORCH_CHECK(mat_b.dtype() == at::kBFloat16, "mat_b must be bf16");
-  TORCH_CHECK(output.dtype() == at::kFloat || output.dtype() == at::kBFloat16,
-              "output must be float32 or bf16");
+  STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
+                  "currently num_tokens must be less than or equal to 16 for "
+                  "router_gemm");
+  STD_TORCH_CHECK(
+      mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+      "mat_a must be bf16");
+  STD_TORCH_CHECK(
+      mat_b.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+      "mat_b must be bf16");
+  STD_TORCH_CHECK(
+      output.scalar_type() == torch::headeronly::ScalarType::Float ||
+          output.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+      "output must be float32 or bf16");
 
-  auto const sm = getSMVersion();
-  TORCH_CHECK(sm >= 90 && sm <= 103, "required SM_103 >= CUDA ARCH >= SM_90");
+  const int sm = getSMVersion();
+  STD_TORCH_CHECK(sm >= 90, "required CUDA ARCH >= SM_90");
 
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream(mat_a.get_device_index());
 
-  if (output.dtype() == at::kFloat) {
+  if (output.scalar_type() == torch::headeronly::ScalarType::Float) {
     if (num_experts == DEFAULT_NUM_EXPERTS) {
       LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
           unroll_float_output(
@@ -145,7 +162,7 @@ void dsv3_router_gemm(at::Tensor& output,       // [num_tokens, num_experts]
               reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
               reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
     }
-  } else if (output.dtype() == at::kBFloat16) {
+  } else if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
     if (num_experts == DEFAULT_NUM_EXPERTS) {
       LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
           unroll_bf16_output(
@@ -164,6 +181,6 @@ void dsv3_router_gemm(at::Tensor& output,       // [num_tokens, num_experts]
   }
 }
 
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-  m.impl("dsv3_router_gemm", &dsv3_router_gemm);
+STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) {
+  m.impl("dsv3_router_gemm", TORCH_BOX(&dsv3_router_gemm));
 }

csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu

@@ -18,14 +18,11 @@
  * limitations under the License.
  */
 
-#include <ATen/ATen.h>
-#include <ATen/cuda/CUDAContext.h>
+#include <cstdint>
 
 #include <cuda_bf16.h>
 #include <cuda_runtime.h>
 
-#include "dsv3_router_gemm_utils.h"
-
 // Custom FMA implementation using PTX assembly instructions
 __device__ __forceinline__ void fma(float2& d, float2 const& a, float2 const& b,
                                     float2 const& c) {

csrc/libtorch_stable/moe/grouped_topk_kernels.cu

@@ -18,9 +18,14 @@
  * limitations under the License.
  */
 #include "moeTopKFuncs.cuh"
-#include <c10/cuda/CUDAStream.h>
-#include <torch/all.h>
+
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+
+#include "libtorch_stable/torch_utils.h"
+
 #include <cmath>
+#include <tuple>
 #include <cuda_fp16.h>
 #include <cuda_bf16.h>
 #include <cuda/std/limits>
@@ -1001,38 +1006,40 @@ INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t, SCORING_NONE);
 }  // end namespace moe
 }  // namespace vllm
 
-std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
-    torch::Tensor const& scores, int64_t n_group, int64_t topk_group,
+std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
+    torch::stable::Tensor const& scores, int64_t n_group, int64_t topk_group,
     int64_t topk, bool renormalize, double routed_scaling_factor,
-    torch::Tensor const& bias, int64_t scoring_func = 0) {
-  auto data_type = scores.scalar_type();
-  auto bias_type = bias.scalar_type();
-  auto input_size = scores.sizes();
-  int64_t num_tokens = input_size[0];
-  int64_t num_experts = input_size[1];
-  TORCH_CHECK(input_size.size() == 2, "scores must be a 2D Tensor");
-  TORCH_CHECK(n_group > 0, "n_group must be positive");
-  TORCH_CHECK(topk > 0, "topk must be positive");
-  TORCH_CHECK(topk_group > 0, "topk_group must be positive");
-  TORCH_CHECK(topk_group <= n_group, "topk_group must be <= n_group");
-  TORCH_CHECK(num_experts % n_group == 0,
-              "num_experts should be divisible by n_group");
-  TORCH_CHECK(n_group <= 32,
-              "n_group should be smaller than or equal to 32 for now");
-  TORCH_CHECK(topk <= 32, "topk should be smaller than or equal to 32 for now");
-  TORCH_CHECK(topk <= topk_group * (num_experts / n_group),
-              "topk must be <= topk_group * (num_experts / n_group)");
-  TORCH_CHECK(scoring_func == vllm::moe::SCORING_NONE ||
-                  scoring_func == vllm::moe::SCORING_SIGMOID,
-              "scoring_func must be SCORING_NONE (0) or SCORING_SIGMOID (1)");
+    torch::stable::Tensor const& bias, int64_t scoring_func = 0) {
+  const auto data_type = scores.scalar_type();
+  const auto bias_type = bias.scalar_type();
+  STD_TORCH_CHECK(scores.dim() == 2, "scores must be a 2D Tensor");
+  const int64_t num_tokens = scores.size(0);
+  const int64_t num_experts = scores.size(1);
+  STD_TORCH_CHECK(n_group > 0, "n_group must be positive");
+  STD_TORCH_CHECK(topk > 0, "topk must be positive");
+  STD_TORCH_CHECK(topk_group > 0, "topk_group must be positive");
+  STD_TORCH_CHECK(topk_group <= n_group, "topk_group must be <= n_group");
+  STD_TORCH_CHECK(num_experts % n_group == 0,
+                  "num_experts should be divisible by n_group");
+  STD_TORCH_CHECK(n_group <= 32,
+                  "n_group should be smaller than or equal to 32 for now");
+  STD_TORCH_CHECK(topk <= 32,
+                  "topk should be smaller than or equal to 32 for now");
+  STD_TORCH_CHECK(topk <= topk_group * (num_experts / n_group),
+                  "topk must be <= topk_group * (num_experts / n_group)");
+  STD_TORCH_CHECK(
+      scoring_func == vllm::moe::SCORING_NONE ||
+          scoring_func == vllm::moe::SCORING_SIGMOID,
+      "scoring_func must be SCORING_NONE (0) or SCORING_SIGMOID (1)");
 
   // Always output float32 for topk_values (eliminates Python-side conversion)
-  torch::Tensor topk_values = torch::empty(
-      {num_tokens, topk}, torch::dtype(torch::kFloat32).device(torch::kCUDA));
-  torch::Tensor topk_indices = torch::empty(
-      {num_tokens, topk}, torch::dtype(torch::kInt32).device(torch::kCUDA));
+  auto topk_values = torch::stable::new_empty(
+      scores, {num_tokens, topk}, torch::headeronly::ScalarType::Float);
+  auto topk_indices = torch::stable::new_empty(
+      scores, {num_tokens, topk}, torch::headeronly::ScalarType::Int);
 
-  auto stream = c10::cuda::getCurrentCUDAStream(scores.get_device());
+  const cudaStream_t stream =
+      get_current_cuda_stream(scores.get_device_index());
   auto const sf = static_cast<vllm::moe::ScoringFunc>(scoring_func);
 
 #define LAUNCH_KERNEL_SF(T, BiasT, IdxT)                                      \
@@ -1057,7 +1064,7 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
             routed_scaling_factor, false, stream);                            \
         break;                                                                \
       default:                                                                \
-        throw std::invalid_argument("Unsupported scoring_func");              \
+        STD_TORCH_CHECK(false, "Unsupported scoring_func");                   \
         break;                                                                \
     }                                                                         \
   } while (0)
@@ -1065,17 +1072,18 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
 #define LAUNCH_KERNEL(T, IdxT)                                         \
   do {                                                                 \
     switch (bias_type) {                                               \
-      case torch::kFloat16:                                            \
+      case torch::headeronly::ScalarType::Half:                        \
         LAUNCH_KERNEL_SF(T, half, IdxT);                               \
         break;                                                         \
-      case torch::kFloat32:                                            \
+      case torch::headeronly::ScalarType::Float:                       \
         LAUNCH_KERNEL_SF(T, float, IdxT);                              \
         break;                                                         \
-      case torch::kBFloat16:                                           \
+      case torch::headeronly::ScalarType::BFloat16:                    \
         LAUNCH_KERNEL_SF(T, __nv_bfloat16, IdxT);                      \
         break;                                                         \
       default:                                                         \
-        throw std::invalid_argument(                                   \
+        STD_TORCH_CHECK(                                               \
+            false,                                                     \
             "Invalid bias dtype, only supports float16, float32, and " \
             "bfloat16");                                               \
         break;                                                         \
@@ -1083,22 +1091,22 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
   } while (0)
 
   switch (data_type) {
-    case torch::kFloat16:
+    case torch::headeronly::ScalarType::Half:
       // Handle Float16
       LAUNCH_KERNEL(half, int32_t);
       break;
-    case torch::kFloat32:
+    case torch::headeronly::ScalarType::Float:
       // Handle Float32
       LAUNCH_KERNEL(float, int32_t);
       break;
-    case torch::kBFloat16:
+    case torch::headeronly::ScalarType::BFloat16:
       // Handle BFloat16
       LAUNCH_KERNEL(__nv_bfloat16, int32_t);
       break;
     default:
       // Handle other data types
-      throw std::invalid_argument(
-          "Invalid dtype, only supports float16, float32, and bfloat16");
+      STD_TORCH_CHECK(
+          false, "Invalid dtype, only supports float16, float32, and bfloat16");
       break;
   }
 #undef LAUNCH_KERNEL

csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py

@@ -302,7 +302,7 @@ def generate_new_kernels():
     if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT:
         kernel_selector_str += (
             "else if (a_type == vllm::kFE4M3fn)\n"
-            "  TORCH_CHECK(false, "
+            "  STD_TORCH_CHECK(false, "
             '"marlin kernel with fp8 activation is not built.");'
         )
 

csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h

@@ -3,8 +3,8 @@
   #define MARLIN_NAMESPACE_NAME marlin_moe_wna16
 #endif
 
-#include "quantization/marlin/marlin.cuh"
-#include "quantization/marlin/marlin_dtypes.cuh"
+#include "libtorch_stable/quantization/marlin/marlin.cuh"
+#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh"
 #include "core/scalar_type.hpp"
 
 #define MARLIN_KERNEL_PARAMS                                          \

csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h

@@ -23,10 +23,10 @@
   #define MARLIN_NAMESPACE_NAME marlin_moe_wna16
 #endif
 
-#include "quantization/marlin/marlin.cuh"
-#include "quantization/marlin/marlin_dtypes.cuh"
-#include "quantization/marlin/dequant.h"
-#include "quantization/marlin/marlin_mma.h"
+#include "libtorch_stable/quantization/marlin/marlin.cuh"
+#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh"
+#include "libtorch_stable/quantization/marlin/dequant.h"
+#include "libtorch_stable/quantization/marlin/marlin_mma.h"
 #include "core/scalar_type.hpp"
 
 #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t)               \

csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu

@@ -24,7 +24,15 @@
 #endif
 
 #include "kernel.h"
-#include "core/registration.h"
+
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/ops.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/util/Exception.h>
+
+#include "libtorch_stable/torch_utils.h"
 
 #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t)               \
   static_assert(std::is_same<scalar_t, half>::value ||          \
@@ -350,18 +358,18 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
   bool m_block_size_8 = moe_block_size == 8;
   bool is_a_8bit = a_type.size_bits() == 8;
 
-  TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m,
-              ", ", prob_n, ", ", prob_k, "]");
+  STD_TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [",
+                  prob_m, ", ", prob_n, ", ", prob_k, "]");
 
   int group_blocks = 0;
   if (has_act_order) {
     if (is_k_full) {
-      TORCH_CHECK(group_size != -1);
+      STD_TORCH_CHECK(group_size != -1);
       group_blocks = group_size / 16;
-      TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
-                  " is not divisible by group_blocks = ", group_blocks);
+      STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
+                      " is not divisible by group_blocks = ", group_blocks);
     } else {
-      TORCH_CHECK(group_size == 0);
+      STD_TORCH_CHECK(group_size == 0);
       group_blocks = 0;
     }
   } else {
@@ -369,8 +377,8 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
       group_blocks = -1;
     } else {
       group_blocks = group_size / 16;
-      TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
-                  " is not divisible by group_blocks = ", group_blocks);
+      STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
+                      " is not divisible by group_blocks = ", group_blocks);
     }
   }
 
@@ -407,7 +415,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
     else if (moe_block_size == 64)
       kernel = permute_cols_kernel<64>;
     else
-      TORCH_CHECK(false, "unsupported moe_block_size ", moe_block_size);
+      STD_TORCH_CHECK(false, "unsupported moe_block_size ", moe_block_size);
 
     // avoid ">>>" being formatted to "> > >"
     // clang-format off
@@ -428,25 +436,25 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
   int max_shared_mem = 0;
   cudaDeviceGetAttribute(&max_shared_mem,
                          cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
-  TORCH_CHECK(max_shared_mem > 0);
+  STD_TORCH_CHECK(max_shared_mem > 0);
 
   int major_capability, minor_capability;
   cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor,
                          dev);
   cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor,
                          dev);
-  TORCH_CHECK(major_capability * 10 + minor_capability >= 75,
-              "marlin kernel only support Turing or newer GPUs.");
+  STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 75,
+                  "marlin kernel only support Turing or newer GPUs.");
   int stages = 4;
   if (major_capability == 7 && minor_capability == 5) {
     stages = 2;
-    TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8,
-                "Turing only support FP16 or INT8 activation.");
+    STD_TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8,
+                    "Turing only support FP16 or INT8 activation.");
   }
   if (a_type == vllm::kFE4M3fn) {
-    TORCH_CHECK(major_capability * 10 + minor_capability >= 89,
-                "FP8 only support Ada Lovelace or newer GPUs.");
-    TORCH_CHECK(
+    STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 89,
+                    "FP8 only support Ada Lovelace or newer GPUs.");
+    STD_TORCH_CHECK(
         major_capability * 10 + minor_capability == 89 ||
             major_capability == 12,
         "Marlin W4A8-FP8 only support SM89 or SM12x device (It is slower than "
@@ -460,10 +468,10 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
     thread_tfg = thread_config_t{thread_k, thread_n, thread_k * thread_n / 64};
     if (blocks_per_sm == -1) blocks_per_sm = 1;
     exec_cfg = exec_config_t{blocks_per_sm, thread_tfg};
-    TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n,
-                " is not divisible by thread_n = ", thread_n);
-    TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k,
-                " is not divisible by thread_k = ", thread_k);
+    STD_TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n,
+                    " is not divisible by thread_n = ", thread_n);
+    STD_TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k,
+                    " is not divisible by thread_k = ", thread_k);
   } else {
     // Auto config
     exec_cfg = determine_exec_config(
@@ -484,19 +492,19 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
   int thread_k_blocks = thread_k / 16;
   int thread_n_blocks = thread_n / 16;
 
-  TORCH_CHECK(is_valid_config(thread_tfg, m_block_size_8, thread_m_blocks,
-                              prob_m, prob_n, prob_k, num_bits, group_size,
-                              has_act_order, is_k_full, has_zp, is_zp_float,
-                              is_a_8bit, stages, max_shared_mem),
-              "Invalid thread config: thread_m_blocks = ", thread_m_blocks,
-              ", thread_k = ", thread_tfg.thread_k,
-              ", thread_n = ", thread_tfg.thread_n,
-              ", num_threads = ", thread_tfg.num_threads, " for MKN = [",
-              prob_m, ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits,
-              ", group_size = ", group_size,
-              ", has_act_order = ", has_act_order, ", is_k_full = ", is_k_full,
-              ", has_zp = ", has_zp, ", is_zp_float = ", is_zp_float,
-              ", max_shared_mem = ", max_shared_mem);
+  STD_TORCH_CHECK(
+      is_valid_config(thread_tfg, m_block_size_8, thread_m_blocks, prob_m,
+                      prob_n, prob_k, num_bits, group_size, has_act_order,
+                      is_k_full, has_zp, is_zp_float, is_a_8bit, stages,
+                      max_shared_mem),
+      "Invalid thread config: thread_m_blocks = ", thread_m_blocks,
+      ", thread_k = ", thread_tfg.thread_k,
+      ", thread_n = ", thread_tfg.thread_n,
+      ", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m, ", ",
+      prob_k, ", ", prob_n, "] and num_bits = ", num_bits,
+      ", group_size = ", group_size, ", has_act_order = ", has_act_order,
+      ", is_k_full = ", is_k_full, ", has_zp = ", has_zp,
+      ", is_zp_float = ", is_zp_float, ", max_shared_mem = ", max_shared_mem);
 
   int sh_cache_size =
       get_kernel_cache_size(thread_tfg, m_block_size_8, thread_m_blocks, prob_m,
@@ -509,13 +517,13 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
       num_threads, is_zp_float, stages);
 
   if (kernel == MarlinDefault) {
-    TORCH_CHECK(false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n,
-                ", ", prob_k, "]", ", has_act_order = ", has_act_order,
-                ", num_groups = ", num_groups, ", group_size = ", group_size,
-                ", thread_m_blocks = ", thread_m_blocks,
-                ", thread_n_blocks = ", thread_n_blocks,
-                ", thread_k_blocks = ", thread_k_blocks,
-                ", num_bits = ", num_bits);
+    STD_TORCH_CHECK(
+        false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, ", ",
+        prob_k, "]", ", has_act_order = ", has_act_order,
+        ", num_groups = ", num_groups, ", group_size = ", group_size,
+        ", thread_m_blocks = ", thread_m_blocks,
+        ", thread_n_blocks = ", thread_n_blocks,
+        ", thread_k_blocks = ", thread_k_blocks, ", num_bits = ", num_bits);
   }
 
   cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
@@ -532,75 +540,81 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
 
 }  // namespace MARLIN_NAMESPACE_NAME
 
-torch::Tensor moe_wna16_marlin_gemm(
-    torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
-    torch::Tensor& b_q_weight,
-    std::optional<torch::Tensor> const& b_bias_or_none, torch::Tensor& b_scales,
-    std::optional<torch::Tensor> const& a_scales_or_none,
-    std::optional<torch::Tensor> const& global_scale_or_none,
-    std::optional<torch::Tensor> const& b_zeros_or_none,
-    std::optional<torch::Tensor> const& g_idx_or_none,
-    std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
-    torch::Tensor& sorted_token_ids, torch::Tensor& expert_ids,
-    torch::Tensor& num_tokens_past_padded, torch::Tensor& topk_weights,
-    int64_t moe_block_size, int64_t top_k, bool mul_topk_weights,
-    vllm::ScalarTypeId const& b_type_id, int64_t size_m, int64_t size_n,
-    int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
-    bool is_zp_float, int64_t thread_k, int64_t thread_n,
+torch::stable::Tensor moe_wna16_marlin_gemm(
+    torch::stable::Tensor& a, std::optional<torch::stable::Tensor> c_or_none,
+    torch::stable::Tensor& b_q_weight,
+    std::optional<torch::stable::Tensor> const& b_bias_or_none,
+    torch::stable::Tensor& b_scales,
+    std::optional<torch::stable::Tensor> const& a_scales_or_none,
+    std::optional<torch::stable::Tensor> const& global_scale_or_none,
+    std::optional<torch::stable::Tensor> const& b_zeros_or_none,
+    std::optional<torch::stable::Tensor> const& g_idx_or_none,
+    std::optional<torch::stable::Tensor> const& perm_or_none,
+    torch::stable::Tensor& workspace, torch::stable::Tensor& sorted_token_ids,
+    torch::stable::Tensor& expert_ids,
+    torch::stable::Tensor& num_tokens_past_padded,
+    torch::stable::Tensor& topk_weights, int64_t moe_block_size, int64_t top_k,
+    bool mul_topk_weights, vllm::ScalarTypeId const& b_type_id, int64_t size_m,
+    int64_t size_n, int64_t size_k, bool is_k_full, bool use_atomic_add,
+    bool use_fp32_reduce, bool is_zp_float, int64_t thread_k, int64_t thread_n,
     int64_t blocks_per_sm) {
   vllm::ScalarTypeId a_type_id, c_type_id, s_type_id;
 
-  auto c_dtype = a.dtype();
-  if (a.scalar_type() == at::ScalarType::Half) {
+  auto c_dtype = a.scalar_type();
+  if (a.scalar_type() == torch::headeronly::ScalarType::Half) {
     a_type_id = vllm::kFloat16.id();
     c_type_id = vllm::kFloat16.id();
-  } else if (a.scalar_type() == at::ScalarType::BFloat16) {
+  } else if (a.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
     a_type_id = vllm::kBFloat16.id();
     c_type_id = vllm::kBFloat16.id();
   } else {
-    c_dtype = b_scales.dtype();
-    if (b_scales.scalar_type() == at::ScalarType::Half) {
+    c_dtype = b_scales.scalar_type();
+    if (b_scales.scalar_type() == torch::headeronly::ScalarType::Half) {
       c_type_id = vllm::kFloat16.id();
-    } else if (b_scales.scalar_type() == at::ScalarType::BFloat16) {
+    } else if (b_scales.scalar_type() ==
+               torch::headeronly::ScalarType::BFloat16) {
       c_type_id = vllm::kBFloat16.id();
     } else {
       c_type_id = vllm::kBFloat16.id();
 
-      TORCH_CHECK(c_or_none.has_value(), "c must be passed for W4A8-FP4");
-      torch::Tensor c = c_or_none.value();
-      c_dtype = c.dtype();
+      STD_TORCH_CHECK(c_or_none.has_value(), "c must be passed for W4A8-FP4");
+      torch::stable::Tensor c = c_or_none.value();
+      c_dtype = c.scalar_type();
 
-      if (c.scalar_type() == at::ScalarType::Half) {
+      if (c.scalar_type() == torch::headeronly::ScalarType::Half) {
         c_type_id = vllm::kFloat16.id();
-      } else if (c.scalar_type() == at::ScalarType::BFloat16) {
+      } else if (c.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
         c_type_id = vllm::kBFloat16.id();
       } else {
-        TORCH_CHECK(false, "unsupported c dtype");
+        STD_TORCH_CHECK(false, "unsupported c dtype");
       }
     }
 
-    if (a.scalar_type() == at::ScalarType::Float8_e4m3fn) {
+    if (a.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn) {
       a_type_id = vllm::kFE4M3fn.id();
-    } else if (a.scalar_type() == at::ScalarType::Char) {
+    } else if (a.scalar_type() == torch::headeronly::ScalarType::Char) {
       a_type_id = vllm::kS8.id();
     } else {
-      TORCH_CHECK(false, "unsupported `a` scalar_type");
+      STD_TORCH_CHECK(false, "unsupported `a` scalar_type");
     }
   }
 
   s_type_id = c_type_id;
   if (b_type_id == vllm::kFE2M1f.id()) {
-    if (b_scales.scalar_type() == at::ScalarType::Float8_e4m3fn) {
+    if (b_scales.scalar_type() ==
+        torch::headeronly::ScalarType::Float8_e4m3fn) {
       s_type_id = vllm::kFE4M3fn.id();
-    } else if (b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) {
+    } else if (b_scales.scalar_type() ==
+               torch::headeronly::ScalarType::Float8_e8m0fnu) {
       s_type_id = vllm::kFE8M0fnu.id();
     } else {
-      TORCH_CHECK(false,
-                  "When b_type = float4_e2m1f, b_scale scalar type must be",
-                  "float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4).");
+      STD_TORCH_CHECK(
+          false, "When b_type = float4_e2m1f, b_scale scalar type must be",
+          "float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4).");
     }
   } else if (b_type_id == vllm::kFE4M3fn.id() &&
-             b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) {
+             b_scales.scalar_type() ==
+                 torch::headeronly::ScalarType::Float8_e8m0fnu) {
     s_type_id = vllm::kFE8M0fnu.id();
   }
 
@@ -613,58 +627,60 @@ torch::Tensor moe_wna16_marlin_gemm(
   int num_experts = b_q_weight.size(0);
 
   if (moe_block_size != 8) {
-    TORCH_CHECK(moe_block_size % 16 == 0,
-                "unsupported moe_block_size=", moe_block_size);
-    TORCH_CHECK(moe_block_size >= 16 && moe_block_size <= 64,
-                "unsupported moe_block_size=", moe_block_size);
+    STD_TORCH_CHECK(moe_block_size % 16 == 0,
+                    "unsupported moe_block_size=", moe_block_size);
+    STD_TORCH_CHECK(moe_block_size >= 16 && moe_block_size <= 64,
+                    "unsupported moe_block_size=", moe_block_size);
   }
 
   // Verify A
-  TORCH_CHECK(a.size(0) == size_m, "Shape mismatch: a.size(0) = ", a.size(0),
-              ", size_m = ", size_m);
-  TORCH_CHECK(a.size(1) == size_k, "Shape mismatch: a.size(1) = ", a.size(1),
-              ", size_k = ", size_k);
+  STD_TORCH_CHECK(a.size(0) == size_m,
+                  "Shape mismatch: a.size(0) = ", a.size(0),
+                  ", size_m = ", size_m);
+  STD_TORCH_CHECK(a.size(1) == size_k,
+                  "Shape mismatch: a.size(1) = ", a.size(1),
+                  ", size_k = ", size_k);
 
   // Verify B
-  TORCH_CHECK(
+  STD_TORCH_CHECK(
       size_k % MARLIN_NAMESPACE_NAME::tile_size == 0, "size_k = ", size_k,
       " is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
-  TORCH_CHECK((size_k / MARLIN_NAMESPACE_NAME::tile_size) == b_q_weight.size(1),
-              "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1),
-              ", size_k = ", size_k,
-              ", tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
-  TORCH_CHECK(
+  STD_TORCH_CHECK(
+      (size_k / MARLIN_NAMESPACE_NAME::tile_size) == b_q_weight.size(1),
+      "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1),
+      ", size_k = ", size_k,
+      ", tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
+  STD_TORCH_CHECK(
       b_q_weight.size(2) % MARLIN_NAMESPACE_NAME::tile_size == 0,
       "b_q_weight.size(2) = ", b_q_weight.size(2),
       " is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
   int actual_size_n =
       (b_q_weight.size(2) / MARLIN_NAMESPACE_NAME::tile_size) * pack_factor;
-  TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n,
-              ", actual_size_n = ", actual_size_n);
+  STD_TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n,
+                  ", actual_size_n = ", actual_size_n);
 
   // Verify device and strides
-  TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
-  TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
+  STD_TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
+  STD_TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
 
-  TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU");
-  TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
+  STD_TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU");
+  STD_TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
 
-  TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
-  TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
+  STD_TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
+  STD_TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
 
-  torch::Tensor a_scales;
-  auto options = torch::TensorOptions().dtype(c_dtype).device(a.device());
-  auto options_fp32 =
-      torch::TensorOptions().dtype(at::kFloat).device(a.device());
+  torch::stable::Tensor a_scales;
+  constexpr auto kFloat = torch::headeronly::ScalarType::Float;
 
   if (a_scales_or_none.has_value()) {
     a_scales = a_scales_or_none.value();
-    TORCH_CHECK(a_type.size_bits() == 8,
-                "a_scales can only be used for 8bit activation.");
+    STD_TORCH_CHECK(a_type.size_bits() == 8,
+                    "a_scales can only be used for 8bit activation.");
   } else {
-    a_scales = torch::empty({0}, options_fp32);
-    TORCH_CHECK(a_type.size_bits() != 8,
-                "the a_scales parameter must be passed for 8bit activation.");
+    a_scales = torch::stable::new_empty(a, {0}, kFloat);
+    STD_TORCH_CHECK(
+        a_type.size_bits() != 8,
+        "the a_scales parameter must be passed for 8bit activation.");
   }
 
   // sms: number of SMs to use for the kernel
@@ -672,82 +688,84 @@ torch::Tensor moe_wna16_marlin_gemm(
   cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, a.get_device());
 
   // Alloc buffers
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
-  torch::Tensor c;
+  torch::stable::accelerator::DeviceGuard device_guard(a.get_device_index());
+  torch::stable::Tensor c;
   if (c_or_none.has_value()) {
     c = c_or_none.value();
-    TORCH_CHECK(c.device().is_cuda(), "c is not on GPU");
-    TORCH_CHECK(c.is_contiguous(), "c is not contiguous");
-    TORCH_CHECK(c.size(0) == size_m * top_k,
-                "Shape mismatch: c.size(0) = ", c.size(0),
-                ", size_m * topk = ", size_m * top_k);
-    TORCH_CHECK(c.size(1) == size_n, "Shape mismatch: c.size(1) = ", c.size(1),
-                ", size_n = ", size_n);
+    STD_TORCH_CHECK(c.device().is_cuda(), "c is not on GPU");
+    STD_TORCH_CHECK(c.is_contiguous(), "c is not contiguous");
+    STD_TORCH_CHECK(c.size(0) == size_m * top_k,
+                    "Shape mismatch: c.size(0) = ", c.size(0),
+                    ", size_m * topk = ", size_m * top_k);
+    STD_TORCH_CHECK(c.size(1) == size_n,
+                    "Shape mismatch: c.size(1) = ", c.size(1),
+                    ", size_n = ", size_n);
   } else {
-    c = torch::empty({size_m * top_k, size_n}, options);
+    c = torch::stable::new_empty(a, {size_m * top_k, size_n}, c_dtype);
   }
 
   // Alloc C tmp buffer that is going to be used for the global reduce
-  torch::Tensor c_tmp;
+  torch::stable::Tensor c_tmp;
   if (use_fp32_reduce && !use_atomic_add) {
     // max num of threadblocks is sms * 4
     long max_c_tmp_size = min(
         (long)size_n * sorted_token_ids.size(0),
         (long)sms * 4 * moe_block_size * MARLIN_NAMESPACE_NAME::max_thread_n);
     if (moe_block_size == 8) max_c_tmp_size *= 2;
-    c_tmp = torch::empty({max_c_tmp_size}, options_fp32);
+    c_tmp = torch::stable::new_empty(a, {max_c_tmp_size}, kFloat);
   } else {
-    c_tmp = torch::empty({0}, options_fp32);
+    c_tmp = torch::stable::new_empty(a, {0}, kFloat);
   }
 
   // Detect groupsize and act_order
   int num_groups = -1;
   int group_size = -1;
 
-  int rank = b_scales.sizes().size();
-  TORCH_CHECK(rank == 3, "b_scales rank = ", rank, " is not 3");
-  TORCH_CHECK(b_scales.size(2) == size_n, "b_scales dim 2 = ", b_scales.size(2),
-              " is not size_n = ", size_n);
+  int rank = b_scales.dim();
+  STD_TORCH_CHECK(rank == 3, "b_scales rank = ", rank, " is not 3");
+  STD_TORCH_CHECK(b_scales.size(2) == size_n,
+                  "b_scales dim 2 = ", b_scales.size(2),
+                  " is not size_n = ", size_n);
   num_groups = b_scales.size(1);
 
-  torch::Tensor g_idx, perm, a_tmp;
+  torch::stable::Tensor g_idx, perm, a_tmp;
   if (g_idx_or_none.has_value() && perm_or_none.has_value()) {
     g_idx = g_idx_or_none.value();
     perm = perm_or_none.value();
 
-    TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU");
-    TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous");
-    TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU");
-    TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
+    STD_TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU");
+    STD_TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous");
+    STD_TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU");
+    STD_TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
 
     // Verify g_idx and perm
-    TORCH_CHECK((g_idx.size(-1) == 0 && perm.size(-1) == 0) ||
-                    (g_idx.size(-1) == size_k && perm.size(-1) == size_k),
-                "Unexpected g_idx.size(-1) = ", g_idx.size(-1),
-                " and perm.size(-1) = ", perm.size(-1),
-                ", where size_k = ", size_k);
+    STD_TORCH_CHECK((g_idx.size(-1) == 0 && perm.size(-1) == 0) ||
+                        (g_idx.size(-1) == size_k && perm.size(-1) == size_k),
+                    "Unexpected g_idx.size(-1) = ", g_idx.size(-1),
+                    " and perm.size(-1) = ", perm.size(-1),
+                    ", where size_k = ", size_k);
   } else {
-    g_idx = torch::empty({0}, options);
-    perm = torch::empty({0}, options);
-    a_tmp = torch::empty({0}, options);
+    g_idx = torch::stable::new_empty(a, {0}, c_dtype);
+    perm = torch::stable::new_empty(a, {0}, c_dtype);
+    a_tmp = torch::stable::new_empty(a, {0}, c_dtype);
   }
   bool has_act_order = g_idx.size(-1) > 0 && perm.size(-1) > 0;
 
   if (has_act_order) {
-    a_tmp = torch::empty({size_m * top_k, size_k}, options);
+    a_tmp = torch::stable::new_empty(a, {size_m * top_k, size_k}, c_dtype);
     if (is_k_full) {
-      TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1");
-      TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k,
-                  ", is not divisible by num_groups = ", num_groups);
+      STD_TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1");
+      STD_TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k,
+                      ", is not divisible by num_groups = ", num_groups);
       group_size = size_k / num_groups;
     } else {
       group_size = 0;
     }
 
   } else {
-    a_tmp = torch::empty({0}, options);
+    a_tmp = torch::stable::new_empty(a, {0}, c_dtype);
     if (num_groups > 1) {
-      TORCH_CHECK(
+      STD_TORCH_CHECK(
           size_k % num_groups == 0, "size_k = ", size_k,
           ", is not divisible by b_scales.size(1) = ", b_scales.size(1));
       group_size = size_k / num_groups;
@@ -756,119 +774,125 @@ torch::Tensor moe_wna16_marlin_gemm(
     }
   }
 
-  torch::Tensor global_scale;
+  torch::stable::Tensor global_scale;
   if (global_scale_or_none.has_value()) {
     global_scale = global_scale_or_none.value();
-    TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn,
-                "global_scale can only be used for nvfp4 format.");
+    STD_TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn,
+                    "global_scale can only be used for nvfp4 format.");
   } else {
-    global_scale = torch::empty({0}, options_fp32);
-    TORCH_CHECK(!(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn),
-                "the global_scale parameter must be passed for nvfp4 format.");
+    global_scale = torch::stable::new_empty(a, {0}, kFloat);
+    STD_TORCH_CHECK(
+        !(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn),
+        "the global_scale parameter must be passed for nvfp4 format.");
   }
 
   bool has_bias = b_bias_or_none.has_value();
-  torch::Tensor b_bias;
+  torch::stable::Tensor b_bias;
   if (has_bias) {
     b_bias = b_bias_or_none.value();
-    TORCH_CHECK(b_bias.device().is_cuda(), "b_bias is not on GPU");
-    TORCH_CHECK(b_bias.is_contiguous(), "b_bias is not contiguous");
-    TORCH_CHECK(b_bias.size(1) == size_n, "b_bias.size(1) != size_n");
-    TORCH_CHECK(b_bias.stride(1) == 1, "b_bias.stride(1) != 1");
+    STD_TORCH_CHECK(b_bias.device().is_cuda(), "b_bias is not on GPU");
+    STD_TORCH_CHECK(b_bias.is_contiguous(), "b_bias is not contiguous");
+    STD_TORCH_CHECK(b_bias.size(1) == size_n, "b_bias.size(1) != size_n");
+    STD_TORCH_CHECK(b_bias.stride(1) == 1, "b_bias.stride(1) != 1");
   } else {
-    b_bias = torch::empty({0}, options);
+    b_bias = torch::stable::new_empty(a, {0}, c_dtype);
   }
 
-  torch::Tensor b_zeros;
+  torch::stable::Tensor b_zeros;
   if (b_zeros_or_none.has_value()) {
     b_zeros = b_zeros_or_none.value();
-    TORCH_CHECK(b_zeros.device().is_cuda(), "b_zeros is not on GPU");
-    TORCH_CHECK(b_zeros.is_contiguous(), "b_zeros is not contiguous");
+    STD_TORCH_CHECK(b_zeros.device().is_cuda(), "b_zeros is not on GPU");
+    STD_TORCH_CHECK(b_zeros.is_contiguous(), "b_zeros is not contiguous");
   } else {
-    b_zeros = torch::empty({0}, options);
+    b_zeros = torch::stable::new_empty(a, {0}, c_dtype);
   }
   bool has_zp = b_zeros.size(-1) > 0;
   if (has_zp) {
-    TORCH_CHECK(
+    STD_TORCH_CHECK(
         b_type == vllm::kU4 || b_type == vllm::kU8,
         "b_type must be u4 or u8 when has_zp = True. Got = ", b_type.str());
   } else {
-    TORCH_CHECK(b_type == vllm::kU4B8 || b_type == vllm::kU8B128 ||
-                    b_type == vllm::kS4 || b_type == vllm::kS8 ||
-                    b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f,
-                "b_type must be uint4b8, uint8b128, int4, int8, "
-                "float8_e4m3fn or float4_e2m1f when has_zp = False. Got = ",
-                b_type.str());
+    STD_TORCH_CHECK(b_type == vllm::kU4B8 || b_type == vllm::kU8B128 ||
+                        b_type == vllm::kS4 || b_type == vllm::kS8 ||
+                        b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f,
+                    "b_type must be uint4b8, uint8b128, int4, int8, "
+                    "float8_e4m3fn or float4_e2m1f when has_zp = False. Got = ",
+                    b_type.str());
   }
 
   if (has_zp && is_zp_float) {
-    TORCH_CHECK(a.scalar_type() == at::ScalarType::Half,
-                "Computation type must be float16 (half) when using float zero "
-                "points.");
+    STD_TORCH_CHECK(
+        a.scalar_type() == torch::headeronly::ScalarType::Half,
+        "Computation type must be float16 (half) when using float zero "
+        "points.");
   }
 
   // Verify b_zeros
   if (has_zp) {
-    int rank = b_zeros.sizes().size();
-    TORCH_CHECK(rank == 3, "b_zeros rank = ", rank, " is not 3");
+    int rank = b_zeros.dim();
+    STD_TORCH_CHECK(rank == 3, "b_zeros rank = ", rank, " is not 3");
     if (is_zp_float) {
-      TORCH_CHECK(b_zeros.size(2) == size_n,
-                  "b_zeros dim 2 = ", b_zeros.size(2),
-                  " is not size_n = ", size_n);
-      TORCH_CHECK(num_groups == b_zeros.size(1),
-                  "b_zeros dim 1 = ", b_zeros.size(1),
-                  " is not num_groups = ", num_groups);
-      TORCH_CHECK(num_groups != -1, "num_groups must be != -1");
+      STD_TORCH_CHECK(b_zeros.size(2) == size_n,
+                      "b_zeros dim 2 = ", b_zeros.size(2),
+                      " is not size_n = ", size_n);
+      STD_TORCH_CHECK(num_groups == b_zeros.size(1),
+                      "b_zeros dim 1 = ", b_zeros.size(1),
+                      " is not num_groups = ", num_groups);
+      STD_TORCH_CHECK(num_groups != -1, "num_groups must be != -1");
     } else {
-      TORCH_CHECK(b_zeros.size(1) == num_groups,
-                  "b_zeros dim 1 = ", b_zeros.size(1),
-                  " is not num_groups = ", num_groups);
-      TORCH_CHECK(b_zeros.size(2) == size_n / pack_factor,
-                  "b_zeros dim 2 = ", b_zeros.size(2),
-                  " is not size_n / pack_factor = ", size_n / pack_factor);
+      STD_TORCH_CHECK(b_zeros.size(1) == num_groups,
+                      "b_zeros dim 1 = ", b_zeros.size(1),
+                      " is not num_groups = ", num_groups);
+      STD_TORCH_CHECK(b_zeros.size(2) == size_n / pack_factor,
+                      "b_zeros dim 2 = ", b_zeros.size(2),
+                      " is not size_n / pack_factor = ", size_n / pack_factor);
     }
   }
 
   // Verify workspace size
-  TORCH_CHECK(size_n % MARLIN_NAMESPACE_NAME::min_thread_n == 0,
-              "size_n = ", size_n, ", is not divisible by min_thread_n = ",
-              MARLIN_NAMESPACE_NAME::min_thread_n);
+  STD_TORCH_CHECK(size_n % MARLIN_NAMESPACE_NAME::min_thread_n == 0,
+                  "size_n = ", size_n, ", is not divisible by min_thread_n = ",
+                  MARLIN_NAMESPACE_NAME::min_thread_n);
 
   int max_n_tiles = size_n / MARLIN_NAMESPACE_NAME::min_thread_n;
   int min_workspace_size = min(
       max_n_tiles * (int)(sorted_token_ids.size(0) / moe_block_size), sms * 4);
-  TORCH_CHECK(workspace.numel() >= min_workspace_size,
-              "workspace.numel = ", workspace.numel(),
-              " is below min_workspace_size = ", min_workspace_size);
+  STD_TORCH_CHECK(workspace.numel() >= min_workspace_size,
+                  "workspace.numel = ", workspace.numel(),
+                  " is below min_workspace_size = ", min_workspace_size);
 
   int dev = a.get_device();
 
-  TORCH_CHECK(a_scales.scalar_type() == at::ScalarType::Float,
-              "scalar type of a_scales must be float");
-  TORCH_CHECK(global_scale.scalar_type() == at::ScalarType::Float,
-              "scalar type of global_scale must be float");
+  STD_TORCH_CHECK(
+      a_scales.scalar_type() == torch::headeronly::ScalarType::Float,
+      "scalar type of a_scales must be float");
+  STD_TORCH_CHECK(
+      global_scale.scalar_type() == torch::headeronly::ScalarType::Float,
+      "scalar type of global_scale must be float");
   if (a_type.size_bits() == 16) {
-    TORCH_CHECK(
+    STD_TORCH_CHECK(
         a.scalar_type() == c.scalar_type(),
         "scalar type of a must be the same with c for 16 bit activation");
   }
 
   MARLIN_NAMESPACE_NAME::marlin_mm(
-      a.data_ptr(), b_q_weight.data_ptr(), c.data_ptr(), c_tmp.data_ptr(),
-      b_bias.data_ptr(), a_scales.data_ptr(), b_scales.data_ptr(),
-      global_scale.data_ptr(), b_zeros.data_ptr(), g_idx.data_ptr(),
-      perm.data_ptr(), a_tmp.data_ptr(), sorted_token_ids.data_ptr(),
-      expert_ids.data_ptr(), num_tokens_past_padded.data_ptr(),
-      topk_weights.data_ptr(), moe_block_size, num_experts, top_k,
-      mul_topk_weights, size_m, size_n, size_k, workspace.data_ptr(), a_type,
-      b_type, c_type, s_type, has_bias, has_act_order, is_k_full, has_zp,
-      num_groups, group_size, dev, at::cuda::getCurrentCUDAStream(dev),
+      a.const_data_ptr(), b_q_weight.const_data_ptr(), c.mutable_data_ptr(),
+      c_tmp.mutable_data_ptr(), b_bias.mutable_data_ptr(),
+      a_scales.mutable_data_ptr(), b_scales.mutable_data_ptr(),
+      global_scale.mutable_data_ptr(), b_zeros.mutable_data_ptr(),
+      g_idx.mutable_data_ptr(), perm.mutable_data_ptr(),
+      a_tmp.mutable_data_ptr(), sorted_token_ids.mutable_data_ptr(),
+      expert_ids.mutable_data_ptr(), num_tokens_past_padded.mutable_data_ptr(),
+      topk_weights.mutable_data_ptr(), moe_block_size, num_experts, top_k,
+      mul_topk_weights, size_m, size_n, size_k, workspace.mutable_data_ptr(),
+      a_type, b_type, c_type, s_type, has_bias, has_act_order, is_k_full,
+      has_zp, num_groups, group_size, dev, get_current_cuda_stream(dev),
       thread_k, thread_n, sms, blocks_per_sm, use_atomic_add, use_fp32_reduce,
       is_zp_float);
 
   return c;
 }
 
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-  m.impl("moe_wna16_marlin_gemm", &moe_wna16_marlin_gemm);
+STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) {
+  m.impl("moe_wna16_marlin_gemm", TORCH_BOX(&moe_wna16_marlin_gemm));
 }

csrc/libtorch_stable/moe/moe_align_sum_kernels.cu

@@ -1,14 +1,17 @@
-#include <torch/all.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
+#include <array>
 #include <cub/cub.cuh>
 
-#include <ATen/ATen.h>
-#include <ATen/cuda/Atomic.cuh>
+#include <cuda_runtime.h>
+#include <torch/csrc/stable/macros.h>
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/ops.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
 
-#include "../cuda_compat.h"
-#include "../dispatch_utils.h"
+#include "../../cuda_compat.h"
 #include "core/math.hpp"
+#include "libtorch_stable/dispatch_utils.h"
+#include "libtorch_stable/torch_utils.h"
 
 #define CEILDIV(x, y) (((x) + (y) - 1) / (y))
 
@@ -492,12 +495,13 @@ __global__ void moe_lora_align_block_size_small_batch_expert_kernel(
 
 // taken from
 // https://github.com/sgl-project/sglang/blob/8b5f83ed3b7d2a49ad5c5cd5aa61c5d502f47dbc
-void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
-                          int64_t block_size, torch::Tensor sorted_token_ids,
-                          torch::Tensor experts_ids,
-                          torch::Tensor num_tokens_post_pad,
-                          std::optional<torch::Tensor> maybe_expert_map) {
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+void moe_align_block_size(
+    torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size,
+    torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids,
+    torch::stable::Tensor num_tokens_post_pad,
+    std::optional<torch::stable::Tensor> maybe_expert_map) {
+  const cudaStream_t stream =
+      get_current_cuda_stream(topk_ids.get_device_index());
 
   int64_t padded_num_experts =
       ((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
@@ -506,19 +510,18 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
   threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
 
   // BlockScan uses 1024 threads and assigns one thread per expert.
-  TORCH_CHECK(padded_num_experts < 1024,
-              "padded_num_experts must be less than 1024");
-  auto options_int =
-      torch::TensorOptions().dtype(torch::kInt).device(topk_ids.device());
+  STD_TORCH_CHECK(padded_num_experts < 1024,
+                  "padded_num_experts must be less than 1024");
   bool has_expert_map = maybe_expert_map.has_value();
-  torch::Tensor expert_map;
+  torch::stable::Tensor expert_map;
   if (has_expert_map) {
     expert_map = maybe_expert_map.value();
   } else {
-    expert_map = torch::empty({0}, options_int);
+    expert_map = torch::stable::new_empty(topk_ids, {0},
+                                          torch::headeronly::ScalarType::Int);
   }
 
-  VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
+  VLLM_STABLE_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
       topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] {
         // calc needed amount of shared mem for `cumsum` tensors
         bool small_batch_expert_mode =
@@ -538,16 +541,17 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
                   scalar_t, fill_threads>;
           small_batch_expert_kernel<<<1, fill_threads + threads,
                                       shared_mem_size, stream>>>(
-              topk_ids.data_ptr<scalar_t>(),
-              sorted_token_ids.data_ptr<int32_t>(),
-              experts_ids.data_ptr<int32_t>(),
-              num_tokens_post_pad.data_ptr<int32_t>(),
-              expert_map.data_ptr<int32_t>(), num_experts, block_size,
-              topk_ids.numel(), sorted_token_ids.size(0), topk_ids.size(1),
-              has_expert_map);
+              reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()),
+              reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(experts_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(
+                  num_tokens_post_pad.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
+              num_experts, block_size, topk_ids.numel(),
+              sorted_token_ids.size(0), topk_ids.size(1), has_expert_map);
         } else {
-          torch::Tensor cumsum_buffer =
-              torch::empty({num_experts + 1}, options_int);
+          torch::stable::Tensor cumsum_buffer = torch::stable::new_empty(
+              topk_ids, {num_experts + 1}, torch::headeronly::ScalarType::Int);
           auto align_kernel = vllm::moe::moe_align_block_size_kernel<scalar_t>;
 
           size_t num_warps = CEILDIV(padded_num_experts, experts_per_warp);
@@ -558,14 +562,16 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
           // blockIdx.x == 0: counting experts and aligning
           // blockIdx.x == 1: filling sorted_token_ids
           align_kernel<<<2, threads, shared_mem_size, stream>>>(
-              topk_ids.data_ptr<scalar_t>(),
-              sorted_token_ids.data_ptr<int32_t>(),
-              experts_ids.data_ptr<int32_t>(),
-              num_tokens_post_pad.data_ptr<int32_t>(),
-              expert_map.data_ptr<int32_t>(), num_experts, padded_num_experts,
-              experts_per_warp, block_size, topk_ids.numel(),
-              cumsum_buffer.data_ptr<int32_t>(), sorted_token_ids.size(0),
-              topk_ids.size(1), has_expert_map);
+              reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()),
+              reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(experts_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(
+                  num_tokens_post_pad.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
+              num_experts, padded_num_experts, experts_per_warp, block_size,
+              topk_ids.numel(),
+              reinterpret_cast<int32_t*>(cumsum_buffer.mutable_data_ptr()),
+              sorted_token_ids.size(0), topk_ids.size(1), has_expert_map);
 
           const int block_threads = std::min(256, (int)threads);
           const int num_blocks =
@@ -577,9 +583,10 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
           auto sort_kernel =
               vllm::moe::count_and_sort_expert_tokens_kernel<scalar_t>;
           sort_kernel<<<gridDims, block_threads, 0, stream>>>(
-              topk_ids.data_ptr<scalar_t>(),
-              sorted_token_ids.data_ptr<int32_t>(),
-              cumsum_buffer.data_ptr<int32_t>(), expert_map.data_ptr<int32_t>(),
+              reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()),
+              reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(cumsum_buffer.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
               topk_ids.numel(), num_experts, sorted_token_ids.size(0),
               topk_ids.size(1), has_expert_map);
         }
@@ -588,33 +595,36 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
 
 void batched_moe_align_block_size(int64_t max_tokens_per_batch,
                                   int64_t block_size,
-                                  torch::Tensor const& batch_num_tokens,
-                                  torch::Tensor sorted_ids,
-                                  torch::Tensor batch_ids,
-                                  torch::Tensor num_tokens_post_pad) {
+                                  const torch::stable::Tensor& batch_num_tokens,
+                                  torch::stable::Tensor sorted_ids,
+                                  torch::stable::Tensor batch_ids,
+                                  torch::stable::Tensor num_tokens_post_pad) {
   namespace batched_kernel = vllm::moe::batched_moe_align_block_size;
 
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream =
+      get_current_cuda_stream(batch_num_tokens.get_device_index());
   int32_t const B = batch_num_tokens.size(0);
   int32_t const num_blocks_per_batch =
       round_to_next_multiple_of(max_tokens_per_batch, block_size) / block_size;
   int32_t const num_blocks = num_blocks_per_batch * B;
   int64_t const sorted_ids_size = num_blocks * block_size;
 
-  TORCH_CHECK(sorted_ids.size(0) == sorted_ids_size);
-  TORCH_CHECK(batch_ids.size(0) == sorted_ids_size / block_size);
-  TORCH_CHECK(num_tokens_post_pad.size(0) == 1);
-  TORCH_CHECK(B <= batched_kernel::num_threads);
+  STD_TORCH_CHECK(sorted_ids.size(0) == sorted_ids_size);
+  STD_TORCH_CHECK(batch_ids.size(0) == sorted_ids_size / block_size);
+  STD_TORCH_CHECK(num_tokens_post_pad.size(0) == 1);
+  STD_TORCH_CHECK(B <= batched_kernel::num_threads);
 
   batched_kernel::batched_moe_align_block_size_kernel<<<
       batched_kernel::num_blocks, batched_kernel::num_threads, 0, stream>>>(
-      B, max_tokens_per_batch, block_size, batch_num_tokens.data_ptr<int32_t>(),
-      sorted_ids.data_ptr<int32_t>(), batch_ids.data_ptr<int32_t>(),
-      num_tokens_post_pad.data_ptr<int32_t>());
+      B, max_tokens_per_batch, block_size,
+      reinterpret_cast<const int32_t*>(batch_num_tokens.const_data_ptr()),
+      reinterpret_cast<int32_t*>(sorted_ids.mutable_data_ptr()),
+      reinterpret_cast<int32_t*>(batch_ids.mutable_data_ptr()),
+      reinterpret_cast<int32_t*>(num_tokens_post_pad.mutable_data_ptr()));
 }
 
-void moe_sum(torch::Tensor& input,   // [num_tokens, topk, hidden_size]
-             torch::Tensor& output)  // [num_tokens, hidden_size]
+void moe_sum(torch::stable::Tensor& input,   // [num_tokens, topk, hidden_size]
+             torch::stable::Tensor& output)  // [num_tokens, hidden_size]
 {
   const int hidden_size = input.size(-1);
   const auto num_tokens = output.numel() / hidden_size;
@@ -622,77 +632,86 @@ void moe_sum(torch::Tensor& input,   // [num_tokens, topk, hidden_size]
 
   dim3 grid(num_tokens);
   dim3 block(std::min(hidden_size, 1024));
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(output));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      output.get_device_index());
+  const cudaStream_t stream =
+      get_current_cuda_stream(output.get_device_index());
 
   switch (topk) {
     case 2:
-      VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_sum_kernel", [&] {
-        vllm::moe::moe_sum_kernel<scalar_t, 2><<<grid, block, 0, stream>>>(
-            output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
-            hidden_size);
-      });
+      VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+          input.scalar_type(), "moe_sum_kernel", [&] {
+            vllm::moe::moe_sum_kernel<scalar_t, 2><<<grid, block, 0, stream>>>(
+                reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
+                reinterpret_cast<const scalar_t*>(input.const_data_ptr()),
+                hidden_size);
+          });
       break;
 
     case 3:
-      VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_sum_kernel", [&] {
-        vllm::moe::moe_sum_kernel<scalar_t, 3><<<grid, block, 0, stream>>>(
-            output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
-            hidden_size);
-      });
+      VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+          input.scalar_type(), "moe_sum_kernel", [&] {
+            vllm::moe::moe_sum_kernel<scalar_t, 3><<<grid, block, 0, stream>>>(
+                reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
+                reinterpret_cast<const scalar_t*>(input.const_data_ptr()),
+                hidden_size);
+          });
       break;
 
     case 4:
-      VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_sum_kernel", [&] {
-        vllm::moe::moe_sum_kernel<scalar_t, 4><<<grid, block, 0, stream>>>(
-            output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
-            hidden_size);
-      });
+      VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+          input.scalar_type(), "moe_sum_kernel", [&] {
+            vllm::moe::moe_sum_kernel<scalar_t, 4><<<grid, block, 0, stream>>>(
+                reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
+                reinterpret_cast<const scalar_t*>(input.const_data_ptr()),
+                hidden_size);
+          });
       break;
 
     default:
-      at::sum_out(output, input, 1);
+      torch::stable::sum_out(output, input, std::array<int64_t, 1>{1});
       break;
   }
 }
 
 void moe_lora_align_block_size(
-    torch::Tensor topk_ids, torch::Tensor token_lora_mapping,
+    torch::stable::Tensor topk_ids, torch::stable::Tensor token_lora_mapping,
     int64_t num_experts, int64_t block_size, int64_t max_loras,
     int64_t max_num_tokens_padded, int64_t max_num_m_blocks,
-    torch::Tensor sorted_token_ids, torch::Tensor expert_ids,
-    torch::Tensor num_tokens_post_pad, torch::Tensor adapter_enabled,
-    torch::Tensor lora_ids, std::optional<torch::Tensor> maybe_expert_map) {
+    torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids,
+    torch::stable::Tensor num_tokens_post_pad,
+    torch::stable::Tensor adapter_enabled, torch::stable::Tensor lora_ids,
+    std::optional<torch::stable::Tensor> maybe_expert_map) {
   const int topk_num = topk_ids.size(1);
 
-  TORCH_CHECK(block_size > 0, "block_size should be greater than 0. ");
+  STD_TORCH_CHECK(block_size > 0, "block_size should be greater than 0. ");
 
   int device_max_shared_mem;
-  auto dev = topk_ids.get_device();
+  int dev = topk_ids.get_device_index();
   cudaDeviceGetAttribute(&device_max_shared_mem,
                          cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream(dev);
 
   int64_t padded_num_experts =
       ((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
 
   // BlockScan uses 1024 threads and assigns one thread per expert.
-  TORCH_CHECK(padded_num_experts < 1024,
-              "padded_num_experts must be less than 1024");
+  STD_TORCH_CHECK(padded_num_experts < 1024,
+                  "padded_num_experts must be less than 1024");
 
-  auto options_int =
-      torch::TensorOptions().dtype(torch::kInt).device(topk_ids.device());
-  torch::Tensor token_mask =
-      torch::empty({max_loras * topk_ids.size(0)}, options_int);
+  torch::stable::Tensor token_mask =
+      torch::stable::new_empty(topk_ids, {max_loras * topk_ids.size(0)},
+                               torch::headeronly::ScalarType::Int);
   bool has_expert_map = maybe_expert_map.has_value();
-  torch::Tensor expert_map;
+  torch::stable::Tensor expert_map;
   if (has_expert_map) {
     expert_map = maybe_expert_map.value();
   } else {
-    expert_map = torch::empty({0}, options_int);
+    expert_map = torch::stable::new_empty(topk_ids, {0},
+                                          torch::headeronly::ScalarType::Int);
   }
 
-  VLLM_DISPATCH_INTEGRAL_TYPES(
+  VLLM_STABLE_DISPATCH_INTEGRAL_TYPES(
       topk_ids.scalar_type(), "moe_lora_align_sum_kernel", [&] {
         bool small_batch_expert_mode =
             (topk_ids.numel() < 1024) && (num_experts <= 64);
@@ -703,7 +722,7 @@ void moe_lora_align_block_size(
               (num_thread + 1) * num_experts * sizeof(int32_t) +
               (num_experts + 1) * sizeof(int32_t);
           if (shared_mem > device_max_shared_mem) {
-            TORCH_CHECK(false, "Shared memory usage exceeds device limit.");
+            STD_TORCH_CHECK(false, "Shared memory usage exceeds device limit.");
           }
 
           // threadIdx.x >= fill_threads: counting experts and aligning
@@ -714,7 +733,7 @@ void moe_lora_align_block_size(
           auto kernel =
               vllm::moe::moe_lora_align_block_size_small_batch_expert_kernel<
                   scalar_t, fill_threads>;
-          AT_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(
+          STD_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(
               (void*)kernel, shared_mem));
           // Grid size is (max_loras + 1) because active_lora_ids has length
           // max_loras + 1: sorted-unique values of token_lora_mapping, which
@@ -725,15 +744,21 @@ void moe_lora_align_block_size(
           // MoE-LoRA kernels. This mirrors the fix made for the Triton
           // _fused_moe_lora_kernel grid in vllm-project/vllm#32277.
           kernel<<<max_loras + 1, blockDim, shared_mem, stream>>>(
-              topk_ids.data_ptr<scalar_t>(),
-              token_lora_mapping.data_ptr<int32_t>(), block_size,
-              expert_map.data_ptr<int32_t>(), num_experts, max_loras,
-              topk_ids.numel(), max_num_tokens_padded, max_num_m_blocks,
-              sorted_token_ids.data_ptr<int32_t>(),
-              expert_ids.data_ptr<int32_t>(), topk_num,
-              num_tokens_post_pad.data_ptr<int32_t>(),
-              adapter_enabled.data_ptr<int32_t>(), lora_ids.data_ptr<int32_t>(),
-              token_mask.data_ptr<int32_t>(), has_expert_map);
+              reinterpret_cast<scalar_t*>(topk_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(token_lora_mapping.mutable_data_ptr()),
+              block_size,
+              reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
+              num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded,
+              max_num_m_blocks,
+              reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(expert_ids.mutable_data_ptr()),
+              topk_num,
+              reinterpret_cast<int32_t*>(
+                  num_tokens_post_pad.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(adapter_enabled.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(lora_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(token_mask.mutable_data_ptr()),
+              has_expert_map);
         } else {
           int num_thread = 1024;
           dim3 blockDim(num_thread);
@@ -742,8 +767,9 @@ void moe_lora_align_block_size(
           size_t shared_mem_size = num_warps * WARP_SIZE * sizeof(int32_t);
 
           // cumsum buffer
-          torch::Tensor cumsum =
-              torch::zeros({max_loras * (num_experts + 1)}, options_int);
+          torch::stable::Tensor cumsum = torch::stable::new_zeros(
+              topk_ids, {max_loras * (num_experts + 1)},
+              torch::headeronly::ScalarType::Int);
 
           auto align_kernel =
               vllm::moe::moe_lora_align_block_size_kernel<scalar_t>;
@@ -759,16 +785,23 @@ void moe_lora_align_block_size(
           // blockIdx.x % 2 == 1: filling sorted_token_ids
           align_kernel<<<(max_loras + 1) * 2, blockDim, shared_mem_size,
                          stream>>>(
-              topk_ids.data_ptr<scalar_t>(),
-              token_lora_mapping.data_ptr<int32_t>(), block_size,
-              expert_map.data_ptr<int32_t>(), num_experts, max_loras,
-              topk_ids.numel(), max_num_tokens_padded, max_num_m_blocks,
-              sorted_token_ids.data_ptr<int32_t>(),
-              expert_ids.data_ptr<int32_t>(), topk_num,
-              num_tokens_post_pad.data_ptr<int32_t>(),
-              adapter_enabled.data_ptr<int32_t>(), cumsum.data_ptr<int32_t>(),
-              WARP_SIZE, padded_num_experts, lora_ids.data_ptr<int32_t>(),
-              token_mask.data_ptr<int32_t>(), has_expert_map);
+              reinterpret_cast<scalar_t*>(topk_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(token_lora_mapping.mutable_data_ptr()),
+              block_size,
+              reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
+              num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded,
+              max_num_m_blocks,
+              reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(expert_ids.mutable_data_ptr()),
+              topk_num,
+              reinterpret_cast<int32_t*>(
+                  num_tokens_post_pad.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(adapter_enabled.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(cumsum.mutable_data_ptr()), WARP_SIZE,
+              padded_num_experts,
+              reinterpret_cast<int32_t*>(lora_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(token_mask.mutable_data_ptr()),
+              has_expert_map);
 
           const int block_threads = std::min(256, (int)num_thread);
           const int num_blocks =
@@ -785,12 +818,16 @@ void moe_lora_align_block_size(
               vllm::moe::lora_count_and_sort_expert_tokens_kernel<scalar_t>;
 
           sort_kernel<<<gridDims, block_threads, 0, stream>>>(
-              topk_ids.data_ptr<scalar_t>(),
-              sorted_token_ids.data_ptr<int32_t>(), cumsum.data_ptr<int32_t>(),
-              expert_map.data_ptr<int32_t>(), topk_ids.numel(), num_experts,
-              max_num_tokens_padded, topk_num, token_mask.data_ptr<int32_t>(),
-              max_loras, lora_ids.data_ptr<int32_t>(),
-              adapter_enabled.data_ptr<int32_t>(), has_expert_map);
+              reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()),
+              reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(cumsum.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
+              topk_ids.numel(), num_experts, max_num_tokens_padded, topk_num,
+              reinterpret_cast<int32_t*>(token_mask.mutable_data_ptr()),
+              max_loras,
+              reinterpret_cast<int32_t*>(lora_ids.mutable_data_ptr()),
+              reinterpret_cast<int32_t*>(adapter_enabled.mutable_data_ptr()),
+              has_expert_map);
         }
       });
 }

csrc/libtorch_stable/moe/moe_ops.h

@@ -0,0 +1,87 @@
+#pragma once
+
+#include <torch/csrc/stable/tensor.h>
+
+#include <optional>
+#include <tuple>
+
+void topk_softmax(torch::stable::Tensor& topk_weights,
+                  torch::stable::Tensor& topk_indices,
+                  torch::stable::Tensor& token_expert_indices,
+                  torch::stable::Tensor& gating_output, bool renormalize,
+                  std::optional<torch::stable::Tensor> bias);
+
+void topk_sigmoid(torch::stable::Tensor& topk_weights,
+                  torch::stable::Tensor& topk_indices,
+                  torch::stable::Tensor& token_expert_indices,
+                  torch::stable::Tensor& gating_output, bool renormalize,
+                  std::optional<torch::stable::Tensor> bias);
+
+void topk_softplus_sqrt(
+    torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices,
+    torch::stable::Tensor& token_expert_indices,
+    torch::stable::Tensor& gating_output, bool renormalize,
+    double routed_scaling_factor,
+    const std::optional<torch::stable::Tensor>& correction_bias,
+    const std::optional<torch::stable::Tensor>& input_ids,
+    const std::optional<torch::stable::Tensor>& tid2eid);
+
+void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output);
+
+void moe_align_block_size(
+    torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size,
+    torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids,
+    torch::stable::Tensor num_tokens_post_pad,
+    std::optional<torch::stable::Tensor> maybe_expert_map);
+
+void batched_moe_align_block_size(
+    int64_t max_tokens_per_batch, int64_t block_size,
+    const torch::stable::Tensor& expert_num_tokens,
+    torch::stable::Tensor sorted_ids, torch::stable::Tensor expert_ids,
+    torch::stable::Tensor num_tokens_post_pad);
+
+void moe_lora_align_block_size(
+    torch::stable::Tensor topk_ids, torch::stable::Tensor token_lora_mapping,
+    int64_t num_experts, int64_t block_size, int64_t max_loras,
+    int64_t max_num_tokens_padded, int64_t max_num_m_blocks,
+    torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids,
+    torch::stable::Tensor num_tokens_post_pad,
+    torch::stable::Tensor adapter_enabled, torch::stable::Tensor lora_ids,
+    std::optional<torch::stable::Tensor> maybe_expert_map);
+#ifndef USE_ROCM
+torch::stable::Tensor moe_wna16_gemm(
+    torch::stable::Tensor input, torch::stable::Tensor output,
+    torch::stable::Tensor b_qweight, torch::stable::Tensor b_scales,
+    std::optional<torch::stable::Tensor> b_qzeros,
+    std::optional<torch::stable::Tensor> topk_weights,
+    torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids,
+    torch::stable::Tensor num_tokens_post_pad, int64_t top_k,
+    int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N, int64_t BLOCK_SIZE_K,
+    int64_t bit);
+
+std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
+    const torch::stable::Tensor& scores, int64_t n_group, int64_t topk_group,
+    int64_t topk, bool renormalize, double routed_scaling_factor,
+    const torch::stable::Tensor& bias, int64_t scoring_func);
+#endif
+
+bool moe_permute_unpermute_supported();
+
+int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
+                                        int64_t num_expert);
+
+void shuffle_rows(const torch::stable::Tensor& input_tensor,
+                  const torch::stable::Tensor& dst2src_map,
+                  torch::stable::Tensor& output_tensor);
+
+#ifndef USE_ROCM
+// DeepSeek V3 optimized router GEMM kernel for SM90+
+// Computes output = mat_a @ mat_b.T where:
+//   mat_a: [num_tokens, hidden_dim] in bf16
+//   mat_b: [num_experts, hidden_dim] in bf16
+//   output: [num_tokens, num_experts] in bf16 or fp32
+// Supports num_tokens in [1, 16], num_experts in {256, 384}, hidden_dim = 7168
+void dsv3_router_gemm(torch::stable::Tensor& output,
+                      const torch::stable::Tensor& mat_a,
+                      const torch::stable::Tensor& mat_b);
+#endif

csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu

@@ -0,0 +1,319 @@
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/ops.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/util/Exception.h>
+
+#include "core/registration.h"
+#include "libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h"
+#include "libtorch_stable/torch_utils.h"
+
+#include <torch/csrc/stable/library.h>
+
+// moe_permute kernels require at least CUDA 12.0
+#if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)
+
+namespace {
+
+int64_t product_integers(torch::headeronly::IntHeaderOnlyArrayRef sizes) {
+  int64_t numel = 1;
+  for (int64_t s : sizes) {
+    numel *= s;
+  }
+  return numel;
+}
+
+torch::stable::Tensor maybe_allocate_tensor(
+    const std::optional<torch::stable::Tensor>& maybe_tensor,
+    torch::headeronly::IntHeaderOnlyArrayRef expected_sizes,
+    torch::headeronly::ScalarType dtype, torch::stable::Device device,
+    char const* name) {
+  auto expected_numel = product_integers(expected_sizes);
+  if (maybe_tensor.has_value()) {
+    auto tensor = maybe_tensor.value();
+    STD_TORCH_CHECK(tensor.device() == device, name,
+                    " must be on the same device");
+    STD_TORCH_CHECK(tensor.scalar_type() == dtype, name,
+                    " has incorrect dtype");
+    STD_TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+    STD_TORCH_CHECK(tensor.numel() >= expected_numel, name,
+                    " is too small for the requested shape");
+    auto flat_tensor = torch::stable::view(tensor, {tensor.numel()});
+    return torch::stable::view(
+        torch::stable::narrow(flat_tensor, 0, 0, expected_numel),
+        expected_sizes);
+  }
+  return torch::stable::empty(expected_sizes, dtype, std::nullopt, device);
+}
+
+}  // namespace
+
+int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
+                                        int64_t n_expert) {
+  return static_cast<int64_t>(
+      CubKeyValueSorter::getWorkspaceSize(num_expanded_rows, n_expert));
+}
+
+void moe_permute_impl(
+    const torch::stable::Tensor& input,                 // [n_token, hidden]
+    const torch::stable::Tensor& topk_ids,              // [n_token, topk]
+    const torch::stable::Tensor& token_expert_indices,  // [n_token, topk]
+    const std::optional<torch::stable::Tensor>& expert_map,  // [n_expert]
+    int64_t n_expert, int64_t n_local_expert, int64_t topk,
+    torch::stable::Tensor& permuted_input,  // [permuted_size, hidden]
+    torch::stable::Tensor& expert_first_token_offset,  // [n_local_expert + 1]
+    torch::stable::Tensor& inv_permuted_idx,           // [n_token, topk]
+    torch::stable::Tensor& permuted_idx,               // [permute_size]
+    const std::optional<torch::stable::Tensor>& maybe_sort_workspace,
+    const std::optional<torch::stable::Tensor>& maybe_permuted_experts_id,
+    const std::optional<torch::stable::Tensor>& maybe_sorted_row_idx,
+    const std::optional<torch::stable::Tensor>& maybe_topk_ids_for_sort) {
+  STD_TORCH_CHECK(expert_first_token_offset.scalar_type() ==
+                      torch::headeronly::ScalarType::Long,
+                  "expert_first_token_offset must be int64");
+  STD_TORCH_CHECK(topk_ids.scalar_type() == torch::headeronly::ScalarType::Int,
+                  "topk_ids must be int32");
+  STD_TORCH_CHECK(
+      token_expert_indices.scalar_type() == torch::headeronly::ScalarType::Int,
+      "token_expert_indices must be int32");
+  STD_TORCH_CHECK(
+      inv_permuted_idx.scalar_type() == torch::headeronly::ScalarType::Int,
+      "inv_permuted_idx must be int32");
+  STD_TORCH_CHECK(expert_first_token_offset.size(0) == n_local_expert + 1,
+                  "expert_first_token_offset shape != n_local_expert+1");
+  STD_TORCH_CHECK(
+      inv_permuted_idx.sizes().equals(token_expert_indices.sizes()),
+      "token_expert_indices shape must be same as inv_permuted_idx");
+
+  auto device = input.device();
+  auto n_token = input.sizes()[0];
+  auto n_hidden = input.sizes()[1];
+  auto expanded_rows = n_token * topk;
+  auto stream = get_current_cuda_stream(input.get_device_index());
+
+  auto sorter_size = moe_permute_sort_workspace_size(expanded_rows, n_expert);
+  auto sort_workspace = maybe_allocate_tensor(
+      maybe_sort_workspace, {sorter_size}, torch::headeronly::ScalarType::Char,
+      device, "sort_workspace");
+  auto permuted_experts_id = maybe_allocate_tensor(
+      maybe_permuted_experts_id, topk_ids.sizes(),
+      torch::headeronly::ScalarType::Int, device, "permuted_experts_id");
+  auto sorted_row_idx = maybe_allocate_tensor(
+      maybe_sorted_row_idx, inv_permuted_idx.sizes(),
+      torch::headeronly::ScalarType::Int, device, "sorted_row_idx");
+
+  CubKeyValueSorter sorter{};
+  int64_t* valid_num_ptr = nullptr;
+  torch::stable::Tensor topk_ids_for_sort = topk_ids;
+
+  if (expert_map.has_value()) {
+    const int* expert_map_ptr = get_ptr<int>(expert_map.value());
+    valid_num_ptr =
+        get_ptr<int64_t>(expert_first_token_offset) + n_local_expert;
+    topk_ids_for_sort = maybe_allocate_tensor(
+        maybe_topk_ids_for_sort, topk_ids.sizes(),
+        torch::headeronly::ScalarType::Int, device, "topk_ids_for_sort");
+    torch::stable::copy_(topk_ids_for_sort, topk_ids);
+    preprocessTopkIdLauncher(get_ptr<int>(topk_ids_for_sort), n_token * topk,
+                             expert_map_ptr, n_expert, stream);
+  }
+
+  sortAndScanExpert(
+      get_ptr<const int>(topk_ids_for_sort), get_ptr<int>(token_expert_indices),
+      get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
+      get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
+      n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
+
+  MOE_DISPATCH(input.scalar_type(), [&] {
+    expandInputRowsKernelLauncher<scalar_t>(
+        get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
+        get_ptr<int>(sorted_row_idx), get_ptr<int>(inv_permuted_idx),
+        get_ptr<int>(permuted_idx), get_ptr<int64_t>(expert_first_token_offset),
+        n_token, valid_num_ptr, n_hidden, topk, n_local_expert, stream);
+  });
+}
+
+void moe_permute(
+    const torch::stable::Tensor& input,                 // [n_token, hidden]
+    const torch::stable::Tensor& topk_ids,              // [n_token, topk]
+    const torch::stable::Tensor& token_expert_indices,  // [n_token, topk]
+    const std::optional<torch::stable::Tensor>& expert_map,  // [n_expert]
+    int64_t n_expert, int64_t n_local_expert, int64_t topk,
+    torch::stable::Tensor& permuted_input,  // [permuted_size, hidden]
+    torch::stable::Tensor& expert_first_token_offset,  // [n_local_expert + 1]
+    torch::stable::Tensor& inv_permuted_idx,           // [n_token, topk]
+    torch::stable::Tensor& permuted_idx) {             // [permute_size]
+  moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert,
+                   n_local_expert, topk, permuted_input,
+                   expert_first_token_offset, inv_permuted_idx, permuted_idx,
+                   std::nullopt, std::nullopt, std::nullopt, std::nullopt);
+}
+
+void moe_permute_with_scratch(
+    const torch::stable::Tensor& input, const torch::stable::Tensor& topk_ids,
+    const torch::stable::Tensor& token_expert_indices,
+    const std::optional<torch::stable::Tensor>& expert_map, int64_t n_expert,
+    int64_t n_local_expert, int64_t topk, torch::stable::Tensor& permuted_input,
+    torch::stable::Tensor& expert_first_token_offset,
+    torch::stable::Tensor& inv_permuted_idx,
+    torch::stable::Tensor& permuted_idx, torch::stable::Tensor& sort_workspace,
+    torch::stable::Tensor& permuted_experts_id,
+    torch::stable::Tensor& sorted_row_idx,
+    torch::stable::Tensor& topk_ids_for_sort) {
+  moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert,
+                   n_local_expert, topk, permuted_input,
+                   expert_first_token_offset, inv_permuted_idx, permuted_idx,
+                   sort_workspace, permuted_experts_id, sorted_row_idx,
+                   topk_ids_for_sort);
+}
+
+void moe_unpermute(
+    const torch::stable::Tensor&
+        permuted_hidden_states,                     // [n_token * topk, hidden]
+    const torch::stable::Tensor& topk_weights,      // [n_token, topk]
+    const torch::stable::Tensor& inv_permuted_idx,  // [n_token, topk]
+    const std::optional<torch::stable::Tensor>&
+        expert_first_token_offset,  // [n_local_expert+1]
+    int64_t topk,
+    torch::stable::Tensor& hidden_states) {  // [n_token, hidden]
+  STD_TORCH_CHECK(
+      permuted_hidden_states.scalar_type() == hidden_states.scalar_type(),
+      "permuted_hidden_states dtype must be same as hidden_states");
+
+  auto n_token = hidden_states.size(0);
+  auto n_hidden = hidden_states.size(1);
+  auto stream = get_current_cuda_stream(hidden_states.get_device_index());
+
+  int64_t const* valid_ptr = nullptr;
+  if (expert_first_token_offset.has_value()) {
+    int n_local_expert = expert_first_token_offset.value().size(0) - 1;
+    valid_ptr =
+        get_ptr<int64_t>(expert_first_token_offset.value()) + n_local_expert;
+  }
+
+  MOE_DISPATCH(hidden_states.scalar_type(), [&] {
+    finalizeMoeRoutingKernelLauncher<scalar_t, scalar_t>(
+        get_ptr<scalar_t>(permuted_hidden_states),
+        get_ptr<scalar_t>(hidden_states), get_ptr<float>(topk_weights),
+        get_ptr<int>(inv_permuted_idx), n_token, n_hidden, topk, valid_ptr,
+        stream);
+  });
+}
+
+template <typename T>
+__global__ void shuffleInputRowsKernel(const T* input,
+                                       const int32_t* dst2src_map, T* output,
+                                       int64_t num_src_rows,
+                                       int64_t num_dst_rows, int64_t num_cols) {
+  int64_t dest_row_idx = blockIdx.x;
+  int64_t const source_row_idx = dst2src_map[dest_row_idx];
+
+  if (blockIdx.x < num_dst_rows) {
+    // Load 128-bits per thread
+    constexpr int64_t ELEM_PER_THREAD = 128 / sizeof(T) / 8;
+    using DataElem = cutlass::Array<T, ELEM_PER_THREAD>;
+
+    // Duplicate and permute rows
+    auto const* source_row_ptr =
+        reinterpret_cast<DataElem const*>(input + source_row_idx * num_cols);
+    auto* dest_row_ptr =
+        reinterpret_cast<DataElem*>(output + dest_row_idx * num_cols);
+
+    int64_t const start_offset = threadIdx.x;
+    int64_t const stride = blockDim.x;
+    int64_t const num_elems_in_col = num_cols / ELEM_PER_THREAD;
+
+    for (int elem_index = start_offset; elem_index < num_elems_in_col;
+         elem_index += stride) {
+      dest_row_ptr[elem_index] = source_row_ptr[elem_index];
+    }
+  }
+}
+
+void shuffle_rows(const torch::stable::Tensor& input_tensor,
+                  const torch::stable::Tensor& dst2src_map,
+                  torch::stable::Tensor& output_tensor) {
+  STD_TORCH_CHECK(input_tensor.scalar_type() == output_tensor.scalar_type(),
+                  "Input and output tensors must have the same data type");
+
+  auto stream = get_current_cuda_stream(output_tensor.get_device_index());
+  const int64_t blocks = output_tensor.size(0);
+  const int64_t threads = 256;
+  const int64_t num_dest_rows = output_tensor.size(0);
+  const int64_t num_src_rows = input_tensor.size(0);
+  const int64_t num_cols = input_tensor.size(1);
+
+  STD_TORCH_CHECK(!(num_cols % (128 / input_tensor.element_size() / 8)),
+                  "num_cols must be divisible by 128 / "
+                  "input_tensor.element_size() / 8");
+
+  MOE_DISPATCH(input_tensor.scalar_type(), [&] {
+    shuffleInputRowsKernel<scalar_t><<<blocks, threads, 0, stream>>>(
+        reinterpret_cast<const scalar_t*>(input_tensor.const_data_ptr()),
+        reinterpret_cast<const int32_t*>(dst2src_map.const_data_ptr()),
+        reinterpret_cast<scalar_t*>(output_tensor.mutable_data_ptr()),
+        num_src_rows, num_dest_rows, num_cols);
+  });
+}
+
+#else
+
+int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
+                                        int64_t n_expert) {
+  STD_TORCH_CHECK(
+      false, "moe_permute_sort_workspace_size is not supported on CUDA < 12.0");
+}
+
+void moe_permute(const torch::stable::Tensor& input,
+                 const torch::stable::Tensor& topk_ids,
+                 const torch::stable::Tensor& token_expert_indices,
+                 const std::optional<torch::stable::Tensor>& expert_map,
+                 int64_t n_expert, int64_t n_local_expert, int64_t topk,
+                 torch::stable::Tensor& permuted_input,
+                 torch::stable::Tensor& expert_first_token_offset,
+                 torch::stable::Tensor& inv_permuted_idx,
+                 torch::stable::Tensor& permuted_idx) {
+  STD_TORCH_CHECK(false, "moe_permute is not supported on CUDA < 12.0");
+}
+
+void moe_permute_with_scratch(
+    const torch::stable::Tensor& input, const torch::stable::Tensor& topk_ids,
+    const torch::stable::Tensor& token_expert_indices,
+    const std::optional<torch::stable::Tensor>& expert_map, int64_t n_expert,
+    int64_t n_local_expert, int64_t topk, torch::stable::Tensor& permuted_input,
+    torch::stable::Tensor& expert_first_token_offset,
+    torch::stable::Tensor& inv_permuted_idx,
+    torch::stable::Tensor& permuted_idx, torch::stable::Tensor& sort_workspace,
+    torch::stable::Tensor& permuted_experts_id,
+    torch::stable::Tensor& sorted_row_idx,
+    torch::stable::Tensor& topk_ids_for_sort) {
+  STD_TORCH_CHECK(false,
+                  "moe_permute_with_scratch is not supported on CUDA < 12.0");
+}
+
+void moe_unpermute(
+    const torch::stable::Tensor& permuted_hidden_states,
+    const torch::stable::Tensor& topk_weights,
+    const torch::stable::Tensor& inv_permuted_idx,
+    const std::optional<torch::stable::Tensor>& expert_first_token_offset,
+    int64_t topk, torch::stable::Tensor& hidden_states) {
+  STD_TORCH_CHECK(false, "moe_unpermute is not supported on CUDA < 12.0");
+}
+
+#endif
+
+bool moe_permute_unpermute_supported() {
+#if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)
+  return true;
+#else
+  return false;
+#endif
+}
+
+STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) {
+  m.impl("moe_permute", TORCH_BOX(&moe_permute));
+  m.impl("moe_permute_with_scratch", TORCH_BOX(&moe_permute_with_scratch));
+  m.impl("moe_unpermute", TORCH_BOX(&moe_unpermute));
+}

csrc/libtorch_stable/moe/moe_wna16.cu

@@ -1,11 +1,14 @@
+#include <algorithm>
 
-#include <torch/all.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <ATen/cuda/CUDAContext.h>
 #include <cuda_runtime.h>
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/ops.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
 
 #include <cuda_fp16.h>
 #include <cuda_bf16.h>
+#include "libtorch_stable/torch_utils.h"
 #include "moe_wna16_utils.h"
 
 #define DIVIDE(x, size) (((x) + (size) - 1) / (size))
@@ -263,7 +266,7 @@ void run_moe_wna16_gemm(const scalar_t* input, scalar_t* output,
   }
 
   const int shared_mem_size = BLOCK_SIZE_M * BLOCK_SIZE_K * 2;
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream();
   kernel<<<gridDim, blockDim, shared_mem_size, stream>>>(
       input, output, b_qweight, b_scales, b_qzeros, topk_weights,
       sorted_token_ids, expert_ids, num_tokens_post_pad, num_experts,
@@ -271,17 +274,18 @@ void run_moe_wna16_gemm(const scalar_t* input, scalar_t* output,
       BLOCK_SIZE_K, has_zp, mul_topk_weight);
 }
 
-torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
-                             torch::Tensor b_qweight, torch::Tensor b_scales,
-                             std::optional<torch::Tensor> b_qzeros,
-                             std::optional<torch::Tensor> topk_weights,
-                             torch::Tensor sorted_token_ids,
-                             torch::Tensor expert_ids,
-                             torch::Tensor num_tokens_post_pad, int64_t top_k,
-                             int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
-                             int64_t BLOCK_SIZE_K, int64_t bit) {
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  output.zero_();
+torch::stable::Tensor moe_wna16_gemm(
+    torch::stable::Tensor input, torch::stable::Tensor output,
+    torch::stable::Tensor b_qweight, torch::stable::Tensor b_scales,
+    std::optional<torch::stable::Tensor> b_qzeros,
+    std::optional<torch::stable::Tensor> topk_weights,
+    torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids,
+    torch::stable::Tensor num_tokens_post_pad, int64_t top_k,
+    int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N, int64_t BLOCK_SIZE_K,
+    int64_t bit) {
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  torch::stable::zero_(output);
 
   const int num_experts = b_qweight.size(0);
   const int size_m = input.size(0);
@@ -291,52 +295,56 @@ torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
 
   int64_t EM = sorted_token_ids.size(0);
   if (size_m <= BLOCK_SIZE_M) {
-    EM = min(EM, size_m * BLOCK_SIZE_M * top_k);
+    EM = std::min(EM, size_m * BLOCK_SIZE_M * top_k);
   }
   const int num_token_blocks = (EM + BLOCK_SIZE_M - 1) / BLOCK_SIZE_M;
 
   const uint32_t* b_qzeros_ptr;
   if (b_qzeros.has_value())
-    b_qzeros_ptr = (const uint32_t*)b_qzeros.value().data_ptr<uint8_t>();
+    b_qzeros_ptr = (const uint32_t*)b_qzeros.value().const_data_ptr<uint8_t>();
   const float* topk_weights_ptr = nullptr;
   if (topk_weights.has_value())
-    topk_weights_ptr = (const float*)topk_weights.value().data_ptr<float>();
+    topk_weights_ptr =
+        (const float*)topk_weights.value().const_data_ptr<float>();
 
   int groups_per_block_row = BLOCK_SIZE_K / group_size;
-  TORCH_CHECK(bit == 4 || bit == 8, "bit must be 4 or 8");
-  TORCH_CHECK(size_k % BLOCK_SIZE_K == 0,
-              "size_k must divisible by BLOCK_SIZE_K");
-  TORCH_CHECK(BLOCK_SIZE_K % group_size == 0,
-              "BLOCK_SIZE_K must divisible by group_size");
-  TORCH_CHECK(BLOCK_SIZE_M <= 64, "BLOCK_SIZE_M must less or equal to 64");
-  TORCH_CHECK(groups_per_block_row == 1 || groups_per_block_row == 2 ||
-                  groups_per_block_row == 4 || groups_per_block_row == 8,
-              "BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8]");
+  STD_TORCH_CHECK(bit == 4 || bit == 8, "bit must be 4 or 8");
+  STD_TORCH_CHECK(size_k % BLOCK_SIZE_K == 0,
+                  "size_k must divisible by BLOCK_SIZE_K");
+  STD_TORCH_CHECK(BLOCK_SIZE_K % group_size == 0,
+                  "BLOCK_SIZE_K must divisible by group_size");
+  STD_TORCH_CHECK(BLOCK_SIZE_M <= 64, "BLOCK_SIZE_M must less or equal to 64");
+  STD_TORCH_CHECK(groups_per_block_row == 1 || groups_per_block_row == 2 ||
+                      groups_per_block_row == 4 || groups_per_block_row == 8,
+                  "BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8]");
 
-  if (input.scalar_type() == at::ScalarType::Half) {
+  if (input.scalar_type() == torch::headeronly::ScalarType::Half) {
     run_moe_wna16_gemm<half>(
-        (const half*)input.data_ptr<at::Half>(),
-        (half*)output.data_ptr<at::Half>(),
-        (const uint32_t*)b_qweight.data_ptr<uint8_t>(),
-        (const half*)b_scales.data_ptr<at::Half>(), b_qzeros_ptr,
-        topk_weights_ptr, sorted_token_ids.data_ptr<int32_t>(),
-        expert_ids.data_ptr<int32_t>(), num_tokens_post_pad.data_ptr<int32_t>(),
-        num_experts, group_size, num_token_blocks, top_k, size_m, size_n,
-        size_k, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, bit,
-        b_qzeros.has_value(), topk_weights.has_value());
-  } else if (input.scalar_type() == at::ScalarType::BFloat16) {
+        reinterpret_cast<const half*>(input.const_data_ptr()),
+        reinterpret_cast<half*>(output.mutable_data_ptr()),
+        (const uint32_t*)b_qweight.const_data_ptr<uint8_t>(),
+        reinterpret_cast<const half*>(b_scales.const_data_ptr()), b_qzeros_ptr,
+        topk_weights_ptr, sorted_token_ids.const_data_ptr<int32_t>(),
+        expert_ids.const_data_ptr<int32_t>(),
+        num_tokens_post_pad.const_data_ptr<int32_t>(), num_experts, group_size,
+        num_token_blocks, top_k, size_m, size_n, size_k, BLOCK_SIZE_M,
+        BLOCK_SIZE_N, BLOCK_SIZE_K, bit, b_qzeros.has_value(),
+        topk_weights.has_value());
+  } else if (input.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
     run_moe_wna16_gemm<nv_bfloat16>(
-        (const nv_bfloat16*)input.data_ptr<at::BFloat16>(),
-        (nv_bfloat16*)output.data_ptr<at::BFloat16>(),
-        (const uint32_t*)b_qweight.data_ptr<uint8_t>(),
-        (const nv_bfloat16*)b_scales.data_ptr<at::BFloat16>(), b_qzeros_ptr,
-        topk_weights_ptr, sorted_token_ids.data_ptr<int32_t>(),
-        expert_ids.data_ptr<int32_t>(), num_tokens_post_pad.data_ptr<int32_t>(),
-        num_experts, group_size, num_token_blocks, top_k, size_m, size_n,
-        size_k, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, bit,
-        b_qzeros.has_value(), topk_weights.has_value());
+        reinterpret_cast<const nv_bfloat16*>(input.const_data_ptr()),
+        reinterpret_cast<nv_bfloat16*>(output.mutable_data_ptr()),
+        (const uint32_t*)b_qweight.const_data_ptr<uint8_t>(),
+        reinterpret_cast<const nv_bfloat16*>(b_scales.const_data_ptr()),
+        b_qzeros_ptr, topk_weights_ptr,
+        sorted_token_ids.const_data_ptr<int32_t>(),
+        expert_ids.const_data_ptr<int32_t>(),
+        num_tokens_post_pad.const_data_ptr<int32_t>(), num_experts, group_size,
+        num_token_blocks, top_k, size_m, size_n, size_k, BLOCK_SIZE_M,
+        BLOCK_SIZE_N, BLOCK_SIZE_K, bit, b_qzeros.has_value(),
+        topk_weights.has_value());
   } else {
-    TORCH_CHECK(false, "moe_wna16_gemm only supports bfloat16 and float16");
+    STD_TORCH_CHECK(false, "moe_wna16_gemm only supports bfloat16 and float16");
   }
   return output;
 }

csrc/libtorch_stable/moe/permute_unpermute_kernels/dispatch.h

@@ -0,0 +1,60 @@
+#pragma once
+
+#include <cuda_fp8.h>
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/util/Exception.h>
+
+#define MOE_SWITCH(TYPE, ...)                                         \
+  const auto _st = (TYPE);                                            \
+  switch (_st) {                                                      \
+    __VA_ARGS__                                                       \
+    default:                                                          \
+      STD_TORCH_CHECK(false, "[moe permute]data type dispatch fail!") \
+  }
+
+#define MOE_DISPATCH_CASE(enum_type, ...)                  \
+  case enum_type: {                                        \
+    using scalar_t = ScalarType2CudaType<enum_type>::type; \
+    __VA_ARGS__();                                         \
+    break;                                                 \
+  }
+
+#define MOE_DISPATCH_FLOAT_CASE(...)                                           \
+  MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Float, __VA_ARGS__)         \
+  MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__)          \
+  MOE_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__)      \
+  MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e5m2, __VA_ARGS__)   \
+  MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
+  MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Byte, __VA_ARGS__)
+
+#define MOE_DISPATCH(TYPE, ...) \
+  MOE_SWITCH(TYPE, MOE_DISPATCH_FLOAT_CASE(__VA_ARGS__))
+
+template <torch::headeronly::ScalarType type>
+struct ScalarType2CudaType;
+
+template <>
+struct ScalarType2CudaType<torch::headeronly::ScalarType::Float> {
+  using type = float;
+};
+template <>
+struct ScalarType2CudaType<torch::headeronly::ScalarType::Half> {
+  using type = half;
+};
+template <>
+struct ScalarType2CudaType<torch::headeronly::ScalarType::BFloat16> {
+  using type = __nv_bfloat16;
+};
+// uint8 for packed fp4
+template <>
+struct ScalarType2CudaType<torch::headeronly::ScalarType::Byte> {
+  using type = uint8_t;
+};
+template <>
+struct ScalarType2CudaType<torch::headeronly::ScalarType::Float8_e5m2> {
+  using type = __nv_fp8_e5m2;
+};
+template <>
+struct ScalarType2CudaType<torch::headeronly::ScalarType::Float8_e4m3fn> {
+  using type = __nv_fp8_e4m3;
+};

csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu

@@ -1,5 +1,7 @@
+#include <cuda.h>
+#include <torch/headeronly/util/Exception.h>
 
-#include "moe_permute_unpermute_kernel.h"
+#include "libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h"
 
 // moe_permute kernels require at least CUDA 12.0
 #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)
@@ -48,9 +50,10 @@ void CubKeyValueSorter::run(void* workspace, size_t const workspace_size,
   size_t expected_ws_size = getWorkspaceSize(num_key_value_pairs, num_experts_);
   size_t actual_ws_size = workspace_size;
 
-  TORCH_CHECK(expected_ws_size <= workspace_size,
-              "[CubKeyValueSorter::run] The allocated workspace is too small "
-              "to run this problem.");
+  STD_TORCH_CHECK(
+      expected_ws_size <= workspace_size,
+      "[CubKeyValueSorter::run] The allocated workspace is too small "
+      "to run this problem.");
   cub::DeviceRadixSort::SortPairs(workspace, actual_ws_size, keys_in, keys_out,
                                   values_in, values_out, num_key_value_pairs, 0,
                                   num_bits_, stream);

csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h

@@ -2,23 +2,24 @@
 // reference from tensorrt_llm moe kernel implementation archive in
 // https://github.com/BBuf/tensorrt-llm-moe/tree/master
 
-#include <c10/core/ScalarType.h>
-#include <torch/all.h>
-#include "dispatch.h"
+#include <torch/csrc/stable/tensor.h>
+
 #include <cub/cub.cuh>
 #include <cub/device/device_radix_sort.cuh>
 #include <cub/util_type.cuh>
-#include "cutlass/numeric_size.h"
+
 #include "cutlass/array.h"
+#include "cutlass/numeric_size.h"
+#include "libtorch_stable/moe/permute_unpermute_kernels/dispatch.h"
 
 template <typename T>
-inline T* get_ptr(torch::Tensor& t) {
-  return reinterpret_cast<T*>(t.data_ptr());
+inline T* get_ptr(torch::stable::Tensor& t) {
+  return reinterpret_cast<T*>(t.mutable_data_ptr());
 }
 
 template <typename T>
-inline const T* get_ptr(const torch::Tensor& t) {
-  return reinterpret_cast<const T*>(t.data_ptr());
+inline const T* get_ptr(const torch::stable::Tensor& t) {
+  return reinterpret_cast<const T*>(t.const_data_ptr());
 }
 
 class CubKeyValueSorter {

csrc/libtorch_stable/moe/topk_softmax_kernels.cu

@@ -17,11 +17,16 @@
  * limitations under the License.
  */
 #include <type_traits>
-#include <torch/all.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-#include "../cuda_compat.h"
-#include "../cub_helpers.h"
+
+#include <cuda_runtime.h>
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/util/Exception.h>
+
+#include "../../cuda_compat.h"
+#include "../../cub_helpers.h"
+#include "libtorch_stable/torch_utils.h"
 
 #ifndef USE_ROCM
     #include <cuda_bf16.h>
@@ -713,7 +718,7 @@ void topkGatingKernelLauncher(
             break;
 #endif
         default: {
-            TORCH_CHECK(workspace != nullptr,
+            STD_TORCH_CHECK(workspace != nullptr,
                 "workspace must be provided for num_experts that are not a power of 2 or multiple of 64.");
             static constexpr int TPB = 256;
             if constexpr (SF == SCORING_SOFTMAX) {
@@ -723,7 +728,7 @@ void topkGatingKernelLauncher(
               moeSigmoid<TPB, InputType><<<num_tokens, TPB, 0, stream>>>(
                 gating_output, nullptr, workspace, num_experts);
             } else {
-                TORCH_CHECK(false, "Unsupported scoring func");
+                STD_TORCH_CHECK(false, "Unsupported scoring func");
             }
             moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
                 workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
@@ -738,63 +743,65 @@ void topkGatingKernelLauncher(
 
 template<typename ComputeType, vllm::moe::ScoringFunc SF>
 void dispatch_topk_launch(
-    torch::Tensor& gating_output,
-    torch::Tensor& topk_weights,
-    torch::Tensor& topk_indices,
-    torch::Tensor& token_expert_indices,
-    torch::Tensor& softmax_workspace,
+    torch::stable::Tensor& gating_output,
+    torch::stable::Tensor& topk_weights,
+    torch::stable::Tensor& topk_indices,
+    torch::stable::Tensor& token_expert_indices,
+    torch::stable::Tensor& softmax_workspace,
     int num_tokens, int num_experts, int topk, bool renormalize,
-    std::optional<torch::Tensor> bias,
+    std::optional<torch::stable::Tensor> bias,
     cudaStream_t stream)
  {
     const float* bias_ptr = nullptr;
     if (bias.has_value()) {
-      const torch::Tensor& bias_tensor = bias.value();
-      TORCH_CHECK(bias_tensor.scalar_type() == at::ScalarType::Float, "bias tensor must be float32");
-      TORCH_CHECK(bias_tensor.dim() == 1, "bias tensor must be 1D");
-      TORCH_CHECK(bias_tensor.size(0) == num_experts, "bias size mismatch, expected: ", num_experts);
-      TORCH_CHECK(bias_tensor.is_contiguous(), "bias tensor must be contiguous");
-      bias_ptr = bias_tensor.data_ptr<float>();
+      const torch::stable::Tensor& bias_tensor = bias.value();
+      STD_TORCH_CHECK(bias_tensor.scalar_type() == torch::headeronly::ScalarType::Float,
+                      "bias tensor must be float32");
+      STD_TORCH_CHECK(bias_tensor.dim() == 1, "bias tensor must be 1D");
+      STD_TORCH_CHECK(bias_tensor.size(0) == num_experts,
+                      "bias size mismatch, expected: ", num_experts);
+      STD_TORCH_CHECK(bias_tensor.is_contiguous(), "bias tensor must be contiguous");
+      bias_ptr = bias_tensor.const_data_ptr<float>();
     }
 
-    if (topk_indices.scalar_type() == at::ScalarType::Int) {
+    if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
         vllm::moe::topkGatingKernelLauncher<int, ComputeType, SF>(
-            reinterpret_cast<const ComputeType*>(gating_output.data_ptr()),
-            topk_weights.data_ptr<float>(),
-            topk_indices.data_ptr<int>(),
-            token_expert_indices.data_ptr<int>(),
-            softmax_workspace.data_ptr<float>(),
+            reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
+            topk_weights.mutable_data_ptr<float>(),
+            topk_indices.mutable_data_ptr<int>(),
+            token_expert_indices.mutable_data_ptr<int>(),
+            softmax_workspace.mutable_data_ptr<float>(),
             num_tokens, num_experts, topk, renormalize,
             bias_ptr, stream);
-    } else if (topk_indices.scalar_type() == at::ScalarType::UInt32) {
+    } else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) {
         vllm::moe::topkGatingKernelLauncher<uint32_t, ComputeType, SF>(
-            reinterpret_cast<const ComputeType*>(gating_output.data_ptr()),
-            topk_weights.data_ptr<float>(),
-            topk_indices.data_ptr<uint32_t>(),
-            token_expert_indices.data_ptr<int>(),
-            softmax_workspace.data_ptr<float>(),
+            reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
+            topk_weights.mutable_data_ptr<float>(),
+            topk_indices.mutable_data_ptr<uint32_t>(),
+            token_expert_indices.mutable_data_ptr<int>(),
+            softmax_workspace.mutable_data_ptr<float>(),
             num_tokens, num_experts, topk, renormalize,
             bias_ptr, stream);
     } else {
-        TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long);
+        STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long);
         vllm::moe::topkGatingKernelLauncher<int64_t, ComputeType, SF>(
-            reinterpret_cast<const ComputeType*>(gating_output.data_ptr()),
-            topk_weights.data_ptr<float>(),
-            topk_indices.data_ptr<int64_t>(),
-            token_expert_indices.data_ptr<int>(),
-            softmax_workspace.data_ptr<float>(),
+            reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
+            topk_weights.mutable_data_ptr<float>(),
+            topk_indices.mutable_data_ptr<int64_t>(),
+            token_expert_indices.mutable_data_ptr<int>(),
+            softmax_workspace.mutable_data_ptr<float>(),
             num_tokens, num_experts, topk, renormalize,
             bias_ptr, stream);
     }
 }
 
 void topk_softmax(
-    torch::Tensor& topk_weights,                // [num_tokens, topk]
-    torch::Tensor& topk_indices,                // [num_tokens, topk]
-    torch::Tensor& token_expert_indices,        // [num_tokens, topk]
-    torch::Tensor& gating_output,               // [num_tokens, num_experts]
+    torch::stable::Tensor& topk_weights,                // [num_tokens, topk]
+    torch::stable::Tensor& topk_indices,                // [num_tokens, topk]
+    torch::stable::Tensor& token_expert_indices,        // [num_tokens, topk]
+    torch::stable::Tensor& gating_output,               // [num_tokens, num_experts]
     bool renormalize,
-    std::optional<torch::Tensor> bias)
+    std::optional<torch::stable::Tensor> bias)
 {
     const int num_experts = gating_output.size(-1);
     const auto num_tokens = gating_output.numel() / num_experts;
@@ -804,35 +811,36 @@ void topk_softmax(
     const bool needs_workspace = !is_pow_2 || num_experts > 256;
     const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0;
 
-    const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
-    const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-    const auto workspace_options = gating_output.options().dtype(at::ScalarType::Float);
-    torch::Tensor softmax_workspace = torch::empty({workspace_size}, workspace_options);
+    torch::stable::accelerator::DeviceGuard guard(gating_output.get_device_index());
+    const cudaStream_t stream =
+        get_current_cuda_stream(gating_output.get_device_index());
+    auto softmax_workspace = torch::stable::new_empty(
+        gating_output, {workspace_size}, torch::headeronly::ScalarType::Float);
 
-    if (gating_output.scalar_type() == at::ScalarType::Float) {
+    if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
         dispatch_topk_launch<float, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
             token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
             bias, stream);
-    } else if (gating_output.scalar_type() == at::ScalarType::Half) {
+    } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
         dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
             token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
             bias, stream);
-    } else if (gating_output.scalar_type() == at::ScalarType::BFloat16) {
+    } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
         dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
             token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
             bias, stream);
     } else {
-        TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
+        STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
     }
 }
 
 void topk_sigmoid(
-    torch::Tensor& topk_weights,                // [num_tokens, topk]
-    torch::Tensor& topk_indices,                // [num_tokens, topk]
-    torch::Tensor& token_expert_indices,        // [num_tokens, topk]
-    torch::Tensor& gating_output,               // [num_tokens, num_experts]
+    torch::stable::Tensor& topk_weights,                // [num_tokens, topk]
+    torch::stable::Tensor& topk_indices,                // [num_tokens, topk]
+    torch::stable::Tensor& token_expert_indices,        // [num_tokens, topk]
+    torch::stable::Tensor& gating_output,               // [num_tokens, num_experts]
     bool renormalize,
-    std::optional<torch::Tensor> bias)
+    std::optional<torch::stable::Tensor> bias)
 {
     const int num_experts = gating_output.size(-1);
     const auto num_tokens = gating_output.numel() / num_experts;
@@ -842,24 +850,25 @@ void topk_sigmoid(
     const bool needs_workspace = !is_pow_2 || num_experts > 256;
     const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0;
 
-    const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
-    const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-    const auto workspace_options = gating_output.options().dtype(at::ScalarType::Float);
-    torch::Tensor workspace = torch::empty({workspace_size}, workspace_options);
+    torch::stable::accelerator::DeviceGuard guard(gating_output.get_device_index());
+    const cudaStream_t stream =
+        get_current_cuda_stream(gating_output.get_device_index());
+    auto workspace = torch::stable::new_empty(
+        gating_output, {workspace_size}, torch::headeronly::ScalarType::Float);
 
-    if (gating_output.scalar_type() == at::ScalarType::Float) {
+    if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
         dispatch_topk_launch<float, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
             token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
             bias, stream);
-    } else if (gating_output.scalar_type() == at::ScalarType::Half) {
+    } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
         dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
             token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
             bias, stream);
-    } else if (gating_output.scalar_type() == at::ScalarType::BFloat16) {
+    } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
         dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
             token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
             bias, stream);
     } else {
-        TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
+        STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
     }
 }

csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu

@@ -18,11 +18,16 @@
  * limitations under the License.
  */
 #include <type_traits>
-#include <torch/all.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-#include "../cuda_compat.h"
-#include "../cub_helpers.h"
+
+#include <cuda_runtime.h>
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/util/Exception.h>
+
+#include "../../cuda_compat.h"
+#include "../../cub_helpers.h"
+#include "libtorch_stable/torch_utils.h"
 #ifndef USE_ROCM
   #include <cuda_bf16.h>
   #include <cuda_fp16.h>
@@ -618,7 +623,7 @@ void topkGatingSoftplusSqrtKernelLauncher(
       LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
       break;
     default: {
-      TORCH_CHECK(false, "Unsupported expert number: ", num_experts);
+      STD_TORCH_CHECK(false, "Unsupported expert number: ", num_experts);
     }
   }
 }
@@ -628,100 +633,109 @@ void topkGatingSoftplusSqrtKernelLauncher(
 
 template <typename ComputeType>
 void dispatch_topk_softplus_sqrt_launch(
-    const ComputeType* gating_output, torch::Tensor& topk_weights,
-    torch::Tensor& topk_indices, torch::Tensor& token_expert_indices,
-    int num_tokens, int num_experts, int topk, bool renormalize,
-    double routed_scaling_factor,
-    const c10::optional<torch::Tensor>& correction_bias,
-    const c10::optional<torch::Tensor>& input_ids,
-    const c10::optional<torch::Tensor>& tid2eid, cudaStream_t stream) {
+    const ComputeType* gating_output, torch::stable::Tensor& topk_weights,
+    torch::stable::Tensor& topk_indices,
+    torch::stable::Tensor& token_expert_indices, int num_tokens,
+    int num_experts, int topk, bool renormalize, double routed_scaling_factor,
+    const std::optional<torch::stable::Tensor>& correction_bias,
+    const std::optional<torch::stable::Tensor>& input_ids,
+    const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream) {
   const float* bias_ptr = nullptr;
   if (correction_bias.has_value()) {
-    bias_ptr = correction_bias.value().data_ptr<float>();
+    bias_ptr = correction_bias.value().const_data_ptr<float>();
   }
   bool use_hash = false;
   if (tid2eid.has_value()) {
-    TORCH_CHECK(input_ids.has_value(), "input_ids is required for hash MoE");
+    STD_TORCH_CHECK(input_ids.has_value(),
+                    "input_ids is required for hash MoE");
     use_hash = true;
   }
-  if (topk_indices.scalar_type() == at::ScalarType::Int) {
+  if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
     const int* input_ids_ptr = nullptr;
     const int* tid2eid_ptr = nullptr;
     if (tid2eid.has_value()) {
-      input_ids_ptr = input_ids.value().data_ptr<int>();
-      tid2eid_ptr = tid2eid.value().data_ptr<int>();
+      input_ids_ptr = input_ids.value().const_data_ptr<int>();
+      tid2eid_ptr = tid2eid.value().const_data_ptr<int>();
     }
 
     vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int, ComputeType>(
-        gating_output, topk_weights.data_ptr<float>(),
-        topk_indices.data_ptr<int>(), token_expert_indices.data_ptr<int>(),
-        num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
-        bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
-  } else if (topk_indices.scalar_type() == at::ScalarType::UInt32) {
+        gating_output, topk_weights.mutable_data_ptr<float>(),
+        topk_indices.mutable_data_ptr<int>(),
+        token_expert_indices.mutable_data_ptr<int>(), num_tokens, num_experts,
+        topk, renormalize, routed_scaling_factor, bias_ptr, use_hash,
+        input_ids_ptr, tid2eid_ptr, stream);
+  } else if (topk_indices.scalar_type() ==
+             torch::headeronly::ScalarType::UInt32) {
     const uint32_t* input_ids_ptr = nullptr;
     const uint32_t* tid2eid_ptr = nullptr;
     if (tid2eid.has_value()) {
-      input_ids_ptr = input_ids.value().data_ptr<uint32_t>();
-      tid2eid_ptr = tid2eid.value().data_ptr<uint32_t>();
+      input_ids_ptr = input_ids.value().const_data_ptr<uint32_t>();
+      tid2eid_ptr = tid2eid.value().const_data_ptr<uint32_t>();
     }
     vllm::moe::topkGatingSoftplusSqrtKernelLauncher<uint32_t, ComputeType>(
-        gating_output, topk_weights.data_ptr<float>(),
-        topk_indices.data_ptr<uint32_t>(), token_expert_indices.data_ptr<int>(),
-        num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
-        bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
+        gating_output, topk_weights.mutable_data_ptr<float>(),
+        topk_indices.mutable_data_ptr<uint32_t>(),
+        token_expert_indices.mutable_data_ptr<int>(), num_tokens, num_experts,
+        topk, renormalize, routed_scaling_factor, bias_ptr, use_hash,
+        input_ids_ptr, tid2eid_ptr, stream);
   } else {
-    TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long);
+    STD_TORCH_CHECK(topk_indices.scalar_type() ==
+                    torch::headeronly::ScalarType::Long);
 
     const int64_t* input_ids_ptr = nullptr;
     const int64_t* tid2eid_ptr = nullptr;
     if (tid2eid.has_value()) {
-      input_ids_ptr = input_ids.value().data_ptr<int64_t>();
-      tid2eid_ptr = tid2eid.value().data_ptr<int64_t>();
+      input_ids_ptr = input_ids.value().const_data_ptr<int64_t>();
+      tid2eid_ptr = tid2eid.value().const_data_ptr<int64_t>();
     }
 
     vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int64_t, ComputeType>(
-        gating_output, topk_weights.data_ptr<float>(),
-        topk_indices.data_ptr<int64_t>(), token_expert_indices.data_ptr<int>(),
-        num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
-        bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
+        gating_output, topk_weights.mutable_data_ptr<float>(),
+        topk_indices.mutable_data_ptr<int64_t>(),
+        token_expert_indices.mutable_data_ptr<int>(), num_tokens, num_experts,
+        topk, renormalize, routed_scaling_factor, bias_ptr, use_hash,
+        input_ids_ptr, tid2eid_ptr, stream);
   }
 }
 
 void topk_softplus_sqrt(
-    torch::Tensor& topk_weights,          // [num_tokens, topk]
-    torch::Tensor& topk_indices,          // [num_tokens, topk]
-    torch::Tensor& token_expert_indices,  // [num_tokens, topk]
-    torch::Tensor& gating_output,         // [num_tokens, num_experts]
+    torch::stable::Tensor& topk_weights,          // [num_tokens, topk]
+    torch::stable::Tensor& topk_indices,          // [num_tokens, topk]
+    torch::stable::Tensor& token_expert_indices,  // [num_tokens, topk]
+    torch::stable::Tensor& gating_output,         // [num_tokens, num_experts]
     bool renormalize, double routed_scaling_factor,
-    const c10::optional<torch::Tensor>& correction_bias,
-    const c10::optional<torch::Tensor>& input_ids,
-    const c10::optional<torch::Tensor>& tid2eid) {
+    const std::optional<torch::stable::Tensor>& correction_bias,
+    const std::optional<torch::stable::Tensor>& input_ids,
+    const std::optional<torch::stable::Tensor>& tid2eid) {
   const int num_experts = gating_output.size(-1);
   const auto num_tokens = gating_output.numel() / num_experts;
   const int topk = topk_weights.size(-1);
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const torch::stable::accelerator::DeviceGuard guard(
+      gating_output.get_device_index());
+  const cudaStream_t stream =
+      get_current_cuda_stream(gating_output.get_device_index());
 
-  if (gating_output.scalar_type() == at::ScalarType::Float) {
+  if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
     dispatch_topk_softplus_sqrt_launch<float>(
-        gating_output.data_ptr<float>(), topk_weights, topk_indices,
+        gating_output.const_data_ptr<float>(), topk_weights, topk_indices,
         token_expert_indices, num_tokens, num_experts, topk, renormalize,
         routed_scaling_factor, correction_bias, input_ids, tid2eid, stream);
-  } else if (gating_output.scalar_type() == at::ScalarType::Half) {
+  } else if (gating_output.scalar_type() ==
+             torch::headeronly::ScalarType::Half) {
     dispatch_topk_softplus_sqrt_launch<__half>(
-        reinterpret_cast<const __half*>(gating_output.data_ptr<at::Half>()),
+        reinterpret_cast<const __half*>(gating_output.const_data_ptr()),
         topk_weights, topk_indices, token_expert_indices, num_tokens,
         num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
         input_ids, tid2eid, stream);
-  } else if (gating_output.scalar_type() == at::ScalarType::BFloat16) {
+  } else if (gating_output.scalar_type() ==
+             torch::headeronly::ScalarType::BFloat16) {
     dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>(
-        reinterpret_cast<const __nv_bfloat16*>(
-            gating_output.data_ptr<at::BFloat16>()),
+        reinterpret_cast<const __nv_bfloat16*>(gating_output.const_data_ptr()),
         topk_weights, topk_indices, token_expert_indices, num_tokens,
         num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
         input_ids, tid2eid, stream);
   } else {
-    TORCH_CHECK(false, "Unsupported gating_output data type: ",
-                gating_output.scalar_type());
+    STD_TORCH_CHECK(false, "Unsupported gating_output data type: ",
+                    gating_output.scalar_type());
   }
 }

csrc/libtorch_stable/moe/torch_bindings.cpp

@@ -1,32 +1,30 @@
 #include "core/registration.h"
 #include "moe_ops.h"
 
-TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
+#include <torch/csrc/stable/library.h>
+
+STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
   // Apply topk softmax to the gating outputs.
   m.def(
       "topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
       "token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "
       "bias) -> ()");
-  m.impl("topk_softmax", torch::kCUDA, &topk_softmax);
 
   // Apply topk sigmoid to the gating outputs.
   m.def(
       "topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
       "token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "
       "bias) -> ()");
-  m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);
 
   m.def(
       "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
       "token_expert_indices, Tensor gating_output, bool renormalize, float "
       "routed_scaling_factor, Tensor? "
       "bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
-  m.impl("topk_softplus_sqrt", torch::kCUDA, &topk_softplus_sqrt);
 
   // Calculate the result of moe by summing up the partial results
   // from all selected experts.
   m.def("moe_sum(Tensor input, Tensor! output) -> ()");
-  m.impl("moe_sum", torch::kCUDA, &moe_sum);
 
   // Aligning the number of tokens to be processed by each expert such
   // that it is divisible by the block size.
@@ -36,7 +34,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
       "                     Tensor! experts_ids,"
       "                     Tensor! num_tokens_post_pad,"
       "                     Tensor? maybe_expert_map) -> ()");
-  m.impl("moe_align_block_size", torch::kCUDA, &moe_align_block_size);
 
   // Aligning the number of tokens to be processed by each expert such
   // that it is divisible by the block size, but for the batched case.
@@ -46,8 +43,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
       "                     Tensor! sorted_token_ids,"
       "                     Tensor! experts_ids,"
       "                     Tensor! num_tokens_post_pad) -> ()");
-  m.impl("batched_moe_align_block_size", torch::kCUDA,
-         &batched_moe_align_block_size);
 
   // Aligning the number of tokens to be processed by each expert such
   // that it is divisible by the block size.
@@ -64,8 +59,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
       "                     Tensor !adapter_enabled,"
       "                     Tensor !lora_ids,"
       "                     Tensor? maybe_expert_map) -> () ");
-  m.impl("moe_lora_align_block_size", torch::kCUDA, &moe_lora_align_block_size);
-
 #ifndef USE_ROCM
   m.def(
       "moe_wna16_gemm(Tensor input, Tensor! output, Tensor b_qweight, "
@@ -75,8 +68,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
       "int top_k, int BLOCK_SIZE_M, int BLOCK_SIZE_N, int BLOCK_SIZE_K, "
       "int bit) -> Tensor");
 
-  m.impl("moe_wna16_gemm", torch::kCUDA, &moe_wna16_gemm);
-
   m.def(
       "moe_wna16_marlin_gemm(Tensor! a, Tensor? c_or_none,"
       "Tensor! b_q_weight, Tensor? b_bias_or_none,"
@@ -118,14 +109,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
   m.def(
       "moe_permute_sort_workspace_size(int num_expanded_rows, int n_expert) -> "
       "int");
-  m.impl("moe_permute_unpermute_supported", &moe_permute_unpermute_supported);
-  m.impl("moe_permute_sort_workspace_size", &moe_permute_sort_workspace_size);
 
   // Row shuffle for MoE
   m.def(
       "shuffle_rows(Tensor input_tensor, Tensor dst2src_map, Tensor! "
       "output_tensor) -> ()");
-  m.impl("shuffle_rows", torch::kCUDA, &shuffle_rows);
 
   // Apply grouped topk routing to select experts.
   m.def(
@@ -133,7 +121,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
       "topk_group, int topk, bool renormalize, float "
       "routed_scaling_factor, Tensor bias, int scoring_func) -> (Tensor, "
       "Tensor)");
-  m.impl("grouped_topk", torch::kCUDA, &grouped_topk);
 
   // DeepSeek V3 optimized router GEMM for SM90+
   m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
@@ -141,4 +128,30 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
 #endif
 }
 
-REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
+STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) {
+  m.impl("topk_softmax", TORCH_BOX(&topk_softmax));
+  m.impl("topk_sigmoid", TORCH_BOX(&topk_sigmoid));
+  m.impl("topk_softplus_sqrt", TORCH_BOX(&topk_softplus_sqrt));
+  m.impl("moe_sum", TORCH_BOX(&moe_sum));
+  m.impl("moe_align_block_size", TORCH_BOX(&moe_align_block_size));
+  m.impl("batched_moe_align_block_size",
+         TORCH_BOX(&batched_moe_align_block_size));
+  m.impl("moe_lora_align_block_size", TORCH_BOX(&moe_lora_align_block_size));
+#ifndef USE_ROCM
+  m.impl("moe_wna16_gemm", TORCH_BOX(&moe_wna16_gemm));
+  m.impl("shuffle_rows", TORCH_BOX(&shuffle_rows));
+  m.impl("grouped_topk", TORCH_BOX(&grouped_topk));
+#endif
+}
+
+#ifndef USE_ROCM
+// Primitive-only ops have no tensor to dispatch on.
+STABLE_TORCH_LIBRARY_IMPL(_moe_C, CompositeExplicitAutograd, m) {
+  m.impl("moe_permute_unpermute_supported",
+         TORCH_BOX(&moe_permute_unpermute_supported));
+  m.impl("moe_permute_sort_workspace_size",
+         TORCH_BOX(&moe_permute_sort_workspace_size));
+}
+#endif
+
+REGISTER_EXTENSION(_moe_C_stable_libtorch)

csrc/libtorch_stable/ops.h

@@ -281,6 +281,24 @@ minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
                          int64_t const nranks, double const eps);
 #endif
 
+// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV /
+// index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs
+// the index branch and scatters k/v/index_k into their paged caches.
+void fused_minimax_m3_qknorm_rope_kv_insert(
+    torch::stable::Tensor& qkv, torch::stable::Tensor const& q_norm_weight,
+    torch::stable::Tensor const& k_norm_weight,
+    torch::stable::Tensor const& cos_sin_cache,
+    torch::stable::Tensor const& positions, int64_t num_heads,
+    int64_t num_kv_heads, int64_t rotary_dim, double eps,
+    std::optional<torch::stable::Tensor> index_q_norm_weight,
+    std::optional<torch::stable::Tensor> index_k_norm_weight,
+    int64_t num_index_heads, std::optional<torch::stable::Tensor> slot_mapping,
+    std::optional<torch::stable::Tensor> index_slot_mapping,
+    std::optional<torch::stable::Tensor> kv_cache,
+    std::optional<torch::stable::Tensor> index_cache, int64_t block_size,
+    std::optional<torch::stable::Tensor> q_out,
+    std::optional<torch::stable::Tensor> index_q_out);
+
 // Sampler kernels (shared CUDA/ROCm)
 void apply_repetition_penalties_(
     torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
@@ -346,7 +364,8 @@ void free_shared_buffer(int64_t buffer);
 // Activation kernels (shared CUDA/ROCm)
 void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void silu_and_mul_clamp(torch::stable::Tensor& out,
-                        torch::stable::Tensor& input, double limit);
+                        torch::stable::Tensor& input, double limit,
+                        double alpha = 1.0, double beta = 0.0);
 void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_tanh_and_mul(torch::stable::Tensor& out,
@@ -397,35 +416,6 @@ torch::stable::Tensor gptq_gemm(torch::stable::Tensor a,
 void gptq_shuffle(torch::stable::Tensor q_weight, torch::stable::Tensor q_perm,
                   int64_t bit);
 
-// GGML kernels (shared CUDA/ROCm)
-torch::stable::Tensor ggml_dequantize(
-    torch::stable::Tensor W, int64_t type, int64_t m, int64_t n,
-    std::optional<torch::headeronly::ScalarType> const& dtype);
-
-torch::stable::Tensor ggml_mul_mat_vec_a8(torch::stable::Tensor W,
-                                          torch::stable::Tensor X, int64_t type,
-                                          int64_t row);
-
-torch::stable::Tensor ggml_mul_mat_a8(torch::stable::Tensor W,
-                                      torch::stable::Tensor X, int64_t type,
-                                      int64_t row);
-
-torch::stable::Tensor ggml_moe_a8(torch::stable::Tensor X,
-                                  torch::stable::Tensor W,
-                                  torch::stable::Tensor sorted_token_ids,
-                                  torch::stable::Tensor expert_ids,
-                                  torch::stable::Tensor num_tokens_post_padded,
-                                  int64_t type, int64_t row, int64_t top_k,
-                                  int64_t tokens);
-
-torch::stable::Tensor ggml_moe_a8_vec(torch::stable::Tensor X,
-                                      torch::stable::Tensor W,
-                                      torch::stable::Tensor topk_ids,
-                                      int64_t top_k, int64_t type, int64_t row,
-                                      int64_t tokens);
-
-int64_t ggml_moe_get_block_size(int64_t type);
-
 void paged_attention_v1(
     torch::stable::Tensor& out, torch::stable::Tensor& query,
     torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,

csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu

@@ -27,15 +27,24 @@
 #include <torch/csrc/stable/tensor.h>
 #include "libtorch_stable/torch_utils.h"
 #include "libtorch_stable/dispatch_utils.h"
+#include "libtorch_stable/cutlass_extensions/common.hpp"
 #include "../../cuda_vec_utils.cuh"
 #include "cuda_utils.h"
 
 #include "nvfp4_utils.cuh"
+
+#if defined(CUDART_VERSION) && CUDART_VERSION >= 12090
+  #define VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED 1
 static_assert(CVT_FP4_ELTS_PER_THREAD == 16,
               "MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)");
+#else
+  #define VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED 0
+#endif
 
 #include "libtorch_stable/launch_bounds_utils.h"
 
+#if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED
+
 namespace vllm {
 
 // MXFP4 block size constants
@@ -104,7 +113,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
                 &input_offset_by_experts[chunk_start + 12]));
         local_offsets[16] = __ldca(&input_offset_by_experts[chunk_start + 16]);
 
-#pragma unroll
+  #pragma unroll
         for (int i = 0; i < 16; i++) {
           if (rowIdx >= local_offsets[i] && rowIdx < local_offsets[i + 1]) {
             rowIdx_in_expert = rowIdx - local_offsets[i];
@@ -309,14 +318,14 @@ void mxfp4_quant_impl(void* output, void* output_scale, void* input,
 
 }  // namespace vllm
 
-/*Quantization entry for mxfp4 experts quantization*/
-#define CHECK_TH_CUDA(x, m) \
-  STD_TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
-#define CHECK_CONTIGUOUS(x, m) \
-  STD_TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
-#define CHECK_INPUT(x, m) \
-  CHECK_TH_CUDA(x, m);    \
-  CHECK_CONTIGUOUS(x, m);
+  /*Quantization entry for mxfp4 experts quantization*/
+  #define CHECK_TH_CUDA(x, m) \
+    STD_TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
+  #define CHECK_CONTIGUOUS(x, m) \
+    STD_TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
+  #define CHECK_INPUT(x, m) \
+    CHECK_TH_CUDA(x, m);    \
+    CHECK_CONTIGUOUS(x, m);
 
 constexpr auto HALF = torch::headeronly::ScalarType::Half;
 constexpr auto BF16 = torch::headeronly::ScalarType::BFloat16;
@@ -364,12 +373,28 @@ static void validate_mxfp4_experts_quant_inputs(
   STD_TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
 }
 
+#endif  // VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED
+
+static bool mxfp4_experts_quant_sm_supported(int64_t cuda_device_capability) {
+#if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED
+  return cuda_device_capability >= 100 && cuda_device_capability < 120;
+#else
+  return false;
+#endif
+}
+
 void mxfp4_experts_quant(
     torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
     torch::stable::Tensor const& input,
     torch::stable::Tensor const& input_offset_by_experts,
     torch::stable::Tensor const& output_scale_offset_by_experts,
     int64_t n_experts) {
+#if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED
+  int32_t sm = get_sm_version_num();
+  STD_TORCH_CHECK(mxfp4_experts_quant_sm_supported(sm),
+                  "No compiled MXFP4 experts quant kernel for SM ", sm,
+                  ". Recompile with SM10x/11x FP4 support and CUDA >= 12.9.");
+
   auto m_topk = input.size(0);
   auto k = input.size(1);
 
@@ -390,6 +415,10 @@ void mxfp4_experts_quant(
             output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
             stream);
       });
+#else
+  STD_TORCH_CHECK_NOT_IMPLEMENTED(false,
+                                  "MXFP4 experts quant requires CUDA >= 12.9.");
+#endif
 }
 
 void silu_and_mul_mxfp4_experts_quant(
@@ -398,6 +427,12 @@ void silu_and_mul_mxfp4_experts_quant(
     torch::stable::Tensor const& input_offset_by_experts,
     torch::stable::Tensor const& output_scale_offset_by_experts,
     int64_t n_experts) {
+#if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED
+  int32_t sm = get_sm_version_num();
+  STD_TORCH_CHECK(mxfp4_experts_quant_sm_supported(sm),
+                  "No compiled SiLU+Mul MXFP4 experts quant kernel for SM ", sm,
+                  ". Recompile with SM10x/11x FP4 support and CUDA >= 12.9.");
+
   auto m_topk = input.size(0);
   auto k_times_2 = input.size(1);
   STD_TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)");
@@ -420,13 +455,29 @@ void silu_and_mul_mxfp4_experts_quant(
             output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
             stream);
       });
+#else
+  STD_TORCH_CHECK_NOT_IMPLEMENTED(
+      false, "SiLU+Mul MXFP4 experts quant requires CUDA >= 12.9.");
+#endif
 }
 
-// Registered here (not torch_bindings.cpp) because VLLM_GPU_FLAGS is applied
-// only under COMPILE_LANGUAGE:CUDA, so ENABLE_NVFP4_SM100 is invisible to
-// .cpp files and cannot gate the registration from there.
+bool mxfp4_experts_quant_supported(int64_t cuda_device_capability) {
+  return mxfp4_experts_quant_sm_supported(cuda_device_capability);
+}
+
+STABLE_TORCH_LIBRARY_FRAGMENT(_C, m) {
+  m.def("mxfp4_experts_quant_supported(int cuda_device_capability) -> bool");
+}
+
+// Registered here so the CUDA 12.8 stub and CUDA 12.9+ implementation stay
+// tied to the same translation unit.
 STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
   m.impl("mxfp4_experts_quant", TORCH_BOX(&mxfp4_experts_quant));
   m.impl("silu_and_mul_mxfp4_experts_quant",
          TORCH_BOX(&silu_and_mul_mxfp4_experts_quant));
 }
+
+STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, m) {
+  m.impl("mxfp4_experts_quant_supported",
+         TORCH_BOX(&mxfp4_experts_quant_supported));
+}

csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh

@@ -22,15 +22,15 @@
 
 #include "../../cuda_vec_utils.cuh"
 
-#if defined(NVFP4_ENABLE_ELTS16) && defined(CUDA_VERSION) && \
-    CUDA_VERSION >= 12090
+#if defined(NVFP4_ENABLE_ELTS16) && defined(CUDART_VERSION) && \
+    CUDART_VERSION >= 12090
   #define ELTS_PER_THREAD 16
+  #define CVT_FP4_PACK16 1
 constexpr int CVT_FP4_ELTS_PER_THREAD = 16;
-constexpr bool CVT_FP4_PACK16 = true;
 #else
   #define ELTS_PER_THREAD 8
+  #define CVT_FP4_PACK16 0
 constexpr int CVT_FP4_ELTS_PER_THREAD = 8;
-constexpr bool CVT_FP4_PACK16 = false;
 #endif
 
 constexpr int CVT_FP4_SF_VEC_SIZE = 16;
@@ -237,21 +237,30 @@ __device__ __forceinline__ fp4_packed_t cvt_warp_fp16_to_fp4(
   // Get the final absolute maximum values.
   float vecMax = float(__hmax(localMax.x, localMax.y));
 
-  // Get the SF (max value of the vector / max value of e2m1).
-  // maximum value of e2m1 = 6.0.
-  // TODO: use half as compute data type.
-  float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f));
   // 8 bits representation of the SF.
+  float SFValue;
   uint8_t fp8SFVal;
-  // Write the SF to global memory (STG.8).
+
   if constexpr (UE8M0_SF) {
-    // Extract the 8 exponent bits from float32.
-    // float 32bits = 1 sign bit + 8 exponent bits + 23 mantissa bits.
-    uint32_t tmp = reinterpret_cast<uint32_t&>(SFValue) >> 23;
-    fp8SFVal = tmp & 0xff;
-    // Convert back to fp32.
-    reinterpret_cast<uint32_t&>(SFValue) = tmp << 23;
+    // OCP MX spec E8M0 scale computation (MXFP4 path):
+    // scale_exp = biased_exponent(round_up(vecMax)) - 2
+    //   -2 because max E2M1 value is 6.0 ≈ 2^2.58; we use 2^2=4 as the
+    //   safe divisor so that max_val / scale <= 6.0 for values near 2^n.
+    uint32_t max_bits = __float_as_uint(vecMax);
+    // Add rounding bias at mantissa bit 21 (equivalent to bf16 val_to_add=32
+    // at bit 5). Threshold: values with mantissa >= 0.75 (i.e. >= 1.75*2^n)
+    // round up to the next power of 2.
+    uint32_t rounded_bits = (max_bits + (1u << 21)) & 0xFF800000u;
+    uint32_t biased_exp = (rounded_bits >> 23) & 0xFFu;
+    uint32_t scale_exp = (biased_exp > 2u) ? (biased_exp - 2u) : 0u;
+    scale_exp = min(scale_exp, 254u);
+    fp8SFVal = static_cast<uint8_t>(scale_exp);
+    // Reconstruct scale as float32: scale = 2^(scale_exp - 127)
+    uint32_t sf_bits = scale_exp << 23;
+    SFValue = __uint_as_float(sf_bits);
   } else {
+    // NVFP4 path: scale = max / 6.0, stored as E4M3.
+    SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f));
     // Here SFValue is always positive, so E4M3 is the same as UE4M3.
     __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue);
     reinterpret_cast<__nv_fp8_e4m3&>(fp8SFVal) = tmp;
@@ -262,13 +271,21 @@ __device__ __forceinline__ fp4_packed_t cvt_warp_fp16_to_fp4(
   // Write the SF to global memory (STG.8).
   if (SFout) *SFout = fp8SFVal;
 
-  // Get the output scale.
-  // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) *
-  //                       reciprocal(SFScaleVal))
-  float outputScale =
-      SFValue != 0.0f ? reciprocal_approximate_ftz(
+  // Get the output scale (= 1 / SFValue for the MXFP4/UE8M0 path where
+  // SFScaleVal=1).  Use exact division for UE8M0 to ensure bit-exact scaling
+  // that matches the reference QDQ implementation (dividing by a power-of-2
+  // scale is exact in IEEE 754).
+  float outputScale;
+  if constexpr (UE8M0_SF) {
+    // SFValue is always a power of 2 for UE8M0, so 1/SFValue is exact.
+    outputScale = SFValue != 0.0f ? (1.0f / SFValue) : 0.0f;
+  } else {
+    // NVFP4 path: use fast approximate reciprocal (original behavior).
+    outputScale = SFValue != 0.0f
+                      ? reciprocal_approximate_ftz(
                             SFValue * reciprocal_approximate_ftz(SFScaleVal))
                       : 0.0f;
+  }
 
   // Convert the input to float.
   float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];

csrc/libtorch_stable/quantization/gguf/dequantize.cuh

@@ -1,571 +0,0 @@
-// copied and adapted from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/convert.cu
-// Dequant functions
-static __device__ __forceinline__ void dequantize_q4_0(const void * vx, const int ib, const int iqs, dfloat2 & v){
-    const block_q4_0 * x = (const block_q4_0 *) vx;
-
-    const dfloat d = x[ib].d;
-
-    const int vui = x[ib].qs[iqs];
-
-    v.x = __int2half_rn(vui & 0xF);
-    v.y = __int2half_rn(vui >> 4);
-
-    v = __hsub2(v, __floats2half2_rn(8.0f, 8.0f));
-    v = __hmul2(v, {d, d});
-}
-
-static __device__ __forceinline__ void dequantize_q4_1(const void * vx, const int ib, const int iqs, dfloat2 & v){
-    const block_q4_1 * x = (const block_q4_1 *) vx;
-
-    const dfloat d = __low2half(x[ib].dm);
-    const dfloat m = __high2half(x[ib].dm);
-
-    const int vui = x[ib].qs[iqs];
-
-    v.x = __int2half_rn(vui & 0xF);
-    v.y = __int2half_rn(vui >> 4);
-
-    v = __hmul2(v, {d, d});
-    v = __hadd2(v, {m, m});
-}
-
-static __device__ __forceinline__ void dequantize_q5_0(const void * vx, const int ib, const int iqs, dfloat2 & v){
-    const block_q5_0 * x = (const block_q5_0 *) vx;
-
-    const dfloat d = x[ib].d;
-
-    uint32_t qh;
-    memcpy(&qh, x[ib].qh, sizeof(qh));
-
-    const int xh_0 = ((qh >> (iqs +  0)) << 4) & 0x10;
-    const int xh_1 = ((qh >> (iqs + 12))     ) & 0x10;
-
-    v.x = __int2half_rn((x[ib].qs[iqs] & 0xf) | xh_0);
-    v.y = __int2half_rn((x[ib].qs[iqs] >>  4) | xh_1);
-
-    v = __hsub2(v, __floats2half2_rn(16.0f, 16.0f));
-    v = __hmul2(v, {d, d});
-}
-
-static __device__ __forceinline__ void dequantize_q5_1(const void * vx, const int ib, const int iqs, dfloat2 & v){
-    const block_q5_1 * x = (const block_q5_1 *) vx;
-
-    const dfloat d = __low2half(x[ib].dm);
-    const dfloat m = __high2half(x[ib].dm);
-
-    uint32_t qh;
-    memcpy(&qh, x[ib].qh, sizeof(qh));
-
-    const int xh_0 = ((qh >> (iqs +  0)) << 4) & 0x10;
-    const int xh_1 = ((qh >> (iqs + 12))     ) & 0x10;
-
-    v.x = __int2half_rn((x[ib].qs[iqs] & 0xf) | xh_0);
-    v.y = __int2half_rn((x[ib].qs[iqs] >>  4) | xh_1);
-
-    v = __hmul2(v, {d, d});
-    v = __hadd2(v, {m, m});
-}
-
-static __device__ __forceinline__ void dequantize_q8_0(const void * vx, const int ib, const int iqs, dfloat2 & v){
-    const block_q8_0 * x = (const block_q8_0 *) vx;
-
-    const dfloat d = x[ib].d;
-
-    v.x = __int2half_rn(x[ib].qs[iqs + 0]);
-    v.y = __int2half_rn(x[ib].qs[iqs + 1]);
-
-    v = __hmul2(v, {d, d});
-}
-
-template <int qk, int qr, dequantize_kernel_t dequantize_kernel, typename dst_t>
-static __global__ void dequantize_block(const void * __restrict__ vx, dst_t * __restrict__ y, const int k) {
-    const int i = 2*(blockDim.x*blockIdx.x + threadIdx.x);
-
-    if (i >= k) {
-        return;
-    }
-
-    const int ib = i/qk; // block index
-    const int iqs = (i%qk)/qr; // quant index
-    const int iybs = i - i%qk; // y block start index
-    const int y_offset = qr == 1 ? 1 : qk/2;
-
-    // dequantize
-    dfloat2 v;
-    dequantize_kernel(vx, ib, iqs, v);
-
-    y[iybs + iqs + 0]        = convert_from_half<dst_t>(v.x);
-    y[iybs + iqs + y_offset] = convert_from_half<dst_t>(v.y);
-}
-
-template<typename dst_t>
-static __global__ void dequantize_block_q2_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
-    const auto i   = blockIdx.x;
-    const block_q2_K * x = (const block_q2_K *) vx;
-
-    const auto tid = threadIdx.x;
-    const int n   = tid/32;
-    const int l   = tid - 32*n;
-    const int is  = 8*n + l/16;
-
-    const uint8_t q = x[i].qs[32*n + l];
-    dst_t * y = yy + i*QK_K + 128*n;
-
-    half dall = __low2half(x[i].dm);
-    half dmin = __high2half(x[i].dm);
-    y[l+ 0] = convert_from_half<dst_t>(__hsub(__hmul(dall, __int2half_rn((x[i].scales[is+0] & 0xF) * ((q >> 0) & 3))), __hmul(dmin,  __int2half_rn(x[i].scales[is+0] >> 4))));
-    y[l+32] = convert_from_half<dst_t>(__hsub(__hmul(dall, __int2half_rn((x[i].scales[is+2] & 0xF) * ((q >> 2) & 3))), __hmul(dmin,  __int2half_rn(x[i].scales[is+2] >> 4))));
-    y[l+64] = convert_from_half<dst_t>(__hsub(__hmul(dall, __int2half_rn((x[i].scales[is+4] & 0xF) * ((q >> 4) & 3))), __hmul(dmin,  __int2half_rn(x[i].scales[is+4] >> 4))));
-    y[l+96] = convert_from_half<dst_t>(__hsub(__hmul(dall, __int2half_rn((x[i].scales[is+6] & 0xF) * ((q >> 6) & 3))), __hmul(dmin,  __int2half_rn(x[i].scales[is+6] >> 4))));
-}
-
-template<typename dst_t>
-static __global__ void dequantize_block_q3_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
-    const auto i = blockIdx.x;
-    const block_q3_K * x = (const block_q3_K *) vx;
-
-    const auto r = threadIdx.x/4;
-    const int tid = r/2;
-    const int is0 = r%2;
-    const int l0 = 16*is0 + 4*(threadIdx.x%4);
-    const int n = tid / 4;
-    const int j = tid - 4*n;
-
-    uint8_t m = 1 << (4*n + j);
-    int is = 8*n + 2*j + is0;
-    int shift = 2*j;
-
-    int8_t us = is <  4 ? (x[i].scales[is-0] & 0xF) | (((x[i].scales[is+8] >> 0) & 3) << 4) :
-                is <  8 ? (x[i].scales[is-0] & 0xF) | (((x[i].scales[is+4] >> 2) & 3) << 4) :
-                is < 12 ? (x[i].scales[is-8] >>  4) | (((x[i].scales[is+0] >> 4) & 3) << 4) :
-                          (x[i].scales[is-8] >>  4) | (((x[i].scales[is-4] >> 6) & 3) << 4);
-    half d_all = x[i].d;
-    half dl = __hmul(d_all,  __int2half_rn(us - 32));
-
-    dst_t * y = yy + i*QK_K + 128*n + 32*j;
-    const uint8_t * q = x[i].qs + 32*n;
-    const uint8_t * hm = x[i].hmask;
-
-    for (int l = l0; l < l0+4; ++l) {
-        y[l] = convert_from_half<dst_t>(__hmul(dl,  __int2half_rn((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4))));
-    }
-}
-
-static inline __device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8_t & m) {
-    if (j < 4) {
-        d = q[j] & 63; m = q[j + 4] & 63;
-    } else {
-        d = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4);
-        m = (q[j+4] >>  4) | ((q[j-0] >> 6) << 4);
-    }
-}
-
-template<typename dst_t>
-static __global__ void dequantize_block_q4_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-    const block_q4_K * x = (const block_q4_K *) vx;
-
-    const auto i = blockIdx.x;
-
-    // assume 32 threads
-    const auto tid = threadIdx.x;
-    const int il  = tid/8;
-    const int ir  = tid%8;
-    const int is  = 2*il;
-    const int n   = 4;
-
-    dst_t * y = yy + i*QK_K + 64*il + n*ir;
-
-    const half dall = __low2half(x[i].dm);
-    const half dmin = __high2half(x[i].dm);
-
-    const uint8_t * q = x[i].qs + 32*il + n*ir;
-
-    uint8_t sc, m;
-    get_scale_min_k4(is + 0, x[i].scales, sc, m);
-    const half d1 = __hmul(dall, __int2half_rn(sc));
-    const half m1 = __hmul(dmin,  __int2half_rn(m));
-    get_scale_min_k4(is + 1, x[i].scales, sc, m);
-    const half d2 = __hmul(dall, __int2half_rn(sc));
-    const half m2 = __hmul(dmin, __int2half_rn(m));
-    for (int l = 0; l < n; ++l) {
-        y[l + 0] = convert_from_half<dst_t>(__hsub(__hmul(d1, __int2half_rn(q[l] & 0xF)), m1));
-        y[l +32] = convert_from_half<dst_t>(__hsub(__hmul(d2,  __int2half_rn(q[l] >> 4)), m2));
-    }
-}
-
-template<typename dst_t>
-static __global__ void dequantize_block_q5_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-    const block_q5_K * x = (const block_q5_K *) vx;
-
-    const auto i = blockIdx.x;
-
-    // assume 64 threads - this is very slightly better than the one below
-    const auto tid = threadIdx.x;
-    const int il  = tid/16;   // il is in 0...3
-    const int ir  = tid%16;   // ir is in 0...15
-    const int is  = 2*il;     // is is in 0...6
-
-    dst_t * y = yy + i*QK_K + 64*il + 2*ir;
-
-    const half dall = __low2half(x[i].dm);
-    const half dmin = __high2half(x[i].dm);
-
-    const uint8_t * ql = x[i].qs + 32*il + 2*ir;
-    const uint8_t * qh = x[i].qh + 2*ir;
-
-    uint8_t sc, m;
-    get_scale_min_k4(is + 0, x[i].scales, sc, m);
-    const half d1 = __hmul(dall, __int2half_rn(sc)); const half m1 = __hmul(dmin, __int2half_rn(m));
-    get_scale_min_k4(is + 1, x[i].scales, sc, m);
-    const half d2 = __hmul(dall, __int2half_rn(sc)); const half m2 = __hmul(dmin, __int2half_rn(m));
-
-    uint8_t   hm  = 1 << (2*il);
-    y[ 0] = convert_from_half<dst_t>(__hsub(__hmul(d1, __int2half_rn((ql[0] & 0xF) + (qh[0] & hm ? 16 : 0))), m1));
-    y[ 1] = convert_from_half<dst_t>(__hsub(__hmul(d1, __int2half_rn((ql[1] & 0xF) + (qh[1] & hm ? 16 : 0))), m1));
-    hm <<= 1;
-    y[32] = convert_from_half<dst_t>(__hsub(__hmul(d2, __int2half_rn((ql[0] >>  4) + (qh[0] & hm ? 16 : 0))), m2));
-    y[33] = convert_from_half<dst_t>(__hsub(__hmul(d2, __int2half_rn((ql[1] >>  4) + (qh[1] & hm ? 16 : 0))), m2));
-}
-
-template<typename dst_t>
-static __global__ void dequantize_block_q6_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-    const block_q6_K * x = (const block_q6_K *) vx;
-
-    const auto i = blockIdx.x;
-
-    // assume 64 threads - this is very slightly better than the one below
-    const auto tid = threadIdx.x;
-    const int ip  = tid/32;   // ip is 0 or 1
-    const int il  = tid - 32*ip; // 0...32
-    const int is  = 8*ip + il/16;
-
-    dst_t * y = yy + i*QK_K + 128*ip + il;
-
-    const half d = x[i].d;
-
-    const uint8_t * ql = x[i].ql + 64*ip + il;
-    const uint8_t   qh = x[i].qh[32*ip + il];
-    const int8_t  * sc = x[i].scales + is;
-
-    y[ 0] = convert_from_half<dst_t>(__hmul(d, __int2half_rn(sc[0] * ((int8_t)((ql[ 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32))));
-    y[32] = convert_from_half<dst_t>(__hmul(d, __int2half_rn(sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32))));
-    y[64] = convert_from_half<dst_t>(__hmul(d, __int2half_rn(sc[4] * ((int8_t)((ql[ 0]  >> 4) | (((qh >> 4) & 3) << 4)) - 32))));
-    y[96] = convert_from_half<dst_t>(__hmul(d, __int2half_rn(sc[6] * ((int8_t)((ql[32]  >> 4) | (((qh >> 6) & 3) << 4)) - 32))));
-}
-
-template<typename dst_t>
-static __global__ void dequantize_block_iq2_xxs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
-    const auto i   = blockIdx.x;
-    const block_iq2_xxs * x = (const block_iq2_xxs  *) vx;
-
-    const auto tid = threadIdx.x;
-    const int il = tid/8; // 0...3
-    const int ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const uint16_t * q2 = x[i].qs + 4*ib;
-    const uint8_t  * aux8 = (const uint8_t *)q2;
-    const uint8_t  * grid = (const uint8_t *)(iq2xxs_grid + aux8[il]);
-    const uint32_t aux32 = q2[2] | (q2[3] << 16);
-    const float d = __half2float(x[i].d) * (0.5f + (aux32 >> 28)) * 0.25f;
-    const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
-    for (int j = 0; j < 8; ++j) y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
-}
-
-template<typename dst_t>
-static __global__ void dequantize_block_iq2_xs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
-    const auto i   = blockIdx.x;
-    const block_iq2_xs * x = (const block_iq2_xs *) vx;
-
-    const auto tid = threadIdx.x;
-    const int il = tid/8; // 0...3
-    const int ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const uint16_t * q2 = x[i].qs + 4*ib;
-    const uint8_t  * grid = (const uint8_t *)(iq2xs_grid + (q2[il] & 511));
-    const float d = __half2float(x[i].d) * (0.5f + ((x[i].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
-    const uint8_t signs = ksigns_iq2xs[q2[il] >> 9];
-    for (int j = 0; j < 8; ++j) y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
-
-}
-
-template<typename dst_t>
-static __global__ void dequantize_block_iq2_s(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
-    const auto i   = blockIdx.x;
-    const block_iq2_s * x = (const block_iq2_s *) vx;
-
-    const auto tid = threadIdx.x;
-    const int il = tid/8; // 0...3
-    const int ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const uint8_t * grid = (const uint8_t *)(iq2s_grid + (x[i].qs[4*ib+il] | ((x[i].qh[ib] << (8-2*il)) & 0x300)));
-    const float d = __half2float(x[i].d) * (0.5f + ((x[i].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
-    const uint8_t signs = x[i].qs[QK_K/8+4*ib+il];
-    for (int j = 0; j < 8; ++j) y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
-}
-
-template<typename dst_t>
-static __global__ void dequantize_block_iq3_xxs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
-    const auto i   = blockIdx.x;
-    const block_iq3_xxs * x = (const block_iq3_xxs  *) vx;
-
-    const auto tid = threadIdx.x;
-    const int il = tid/8; // 0...3
-    const int ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const uint8_t  * q3 = x[i].qs + 8*ib;
-    const uint16_t * gas = (const uint16_t *)(x[i].qs + QK_K/4) + 2*ib;
-    const uint8_t  * grid1 = (const uint8_t *)(iq3xxs_grid + q3[2*il+0]);
-    const uint8_t  * grid2 = (const uint8_t *)(iq3xxs_grid + q3[2*il+1]);
-    const uint32_t aux32 = gas[0] | (gas[1] << 16);
-    const float d = __half2float(x[i].d) * (0.5f + (aux32 >> 28)) * 0.5f;
-    const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
-    for (int j = 0; j < 4; ++j) {
-        y[j+0] = d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f);
-        y[j+4] = d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f);
-    }
-}
-
-template<typename dst_t>
-static __global__ void dequantize_block_iq3_s(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
-    const auto i   = blockIdx.x;
-    const block_iq3_s * x = (const block_iq3_s *) vx;
-
-    const auto tid = threadIdx.x;
-    const int il = tid/8; // 0...3
-    const int ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const uint8_t * qs = x[i].qs + 8*ib;
-    const uint8_t * grid1 = (const uint8_t *)(iq3xs_grid + (qs[2*il+0] | ((x[i].qh[ib] << (8-2*il)) & 256)));
-    const uint8_t * grid2 = (const uint8_t *)(iq3xs_grid + (qs[2*il+1] | ((x[i].qh[ib] << (7-2*il)) & 256)));
-    const float d = __half2float(x[i].d) * (0.5f + ((x[i].scales[ib/2] >> 4*(ib%2)) & 0xf)) * 0.5f;
-    const uint8_t signs = x[i].signs[4*ib + il];
-    for (int j = 0; j < 4; ++j) {
-        y[j+0] = d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f);
-        y[j+4] = d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f);
-    }
-}
-
-template<typename dst_t>
-static __global__ void dequantize_block_iq1_s(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
-    const int64_t i   = blockIdx.x;
-    const block_iq1_s * x = (const block_iq1_s  *) vx;
-
-    const int64_t tid = threadIdx.x;
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const float delta = x[i].qh[ib] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA;
-    const float d = __half2float(x[i].d) * (2*((x[i].qh[ib] >> 12) & 7) + 1);
-    uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32;
-    grid32[0] = iq1s_grid_gpu[x[i].qs[4*ib+il] | (((x[i].qh[ib] >> 3*il) & 7) << 8)];
-    grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
-    grid32[0] &= 0x0f0f0f0f;
-    for (int j = 0; j < 8; ++j) {
-        y[j] = d * (q[j] + delta);
-    }
-}
-
-template<typename dst_t>
-static __global__ void dequantize_block_iq1_m(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
-    const int64_t i   = blockIdx.x;
-    const block_iq1_m * x = (const block_iq1_m  *) vx;
-
-    const int64_t tid = threadIdx.x;
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const uint16_t * sc = (const uint16_t *)x[i].scales;
-    iq1m_scale_t scale;
-    scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
-    const int64_t ib16 = 2*ib + il/2; // sc[ib16/4] >> 3*(ib16%4) -> sc[ib/2] >> 3*((2*ib+il/2)%4);
-    const float d = __half2float(scale.f16) * (2*((sc[ib16/4] >> 3*(ib16%4)) & 0x7) + 1);
-    const float delta = x[i].qh[2*ib+il/2] & (0x08 << 4*(il%2)) ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA;
-    uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32;
-    grid32[0] = iq1s_grid_gpu[x[i].qs[4*ib+il] | (((x[i].qh[2*ib+il/2] >> 4*(il%2)) & 7) << 8)];
-    grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
-    grid32[0] &= 0x0f0f0f0f;
-    for (int j = 0; j < 8; ++j) {
-        y[j] = d * (q[j] + delta);
-    }
-}
-
-template<typename dst_t>
-static __global__ void dequantize_block_iq4_nl(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
-    const auto i   = blockIdx.x;
-    const block_iq4_nl * x = (const block_iq4_nl *) vx + i*(QK_K/QK4_NL);
-
-    const auto tid = threadIdx.x;
-    const int il = tid/8; // 0...3
-    const int ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 4*il;
-    const uint8_t  * q4 = x[ib].qs + 4*il;
-    const float d = __half2float(x[ib].d);
-    for (int j = 0; j < 4; ++j) {
-        y[j+ 0] = d * kvalues_iq4nl[q4[j] & 0xf];
-        y[j+16] = d * kvalues_iq4nl[q4[j] >>  4];
-    }
-
-}
-
-template<typename dst_t>
-static __global__ void dequantize_block_iq4_xs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-    const auto i   = blockIdx.x;
-    const block_iq4_xs * x = (const block_iq4_xs *)vx;
-
-    const auto tid = threadIdx.x;
-    const int il = tid/8; // 0...3
-    const int ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 4*il;
-    const uint8_t  * q4 = x[i].qs + 16*ib + 4*il;
-    const float d = __half2float(x[i].d) * ((((x[i].scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((x[i].scales_h >> 2*ib) & 3) << 4)) - 32);
-    for (int j = 0; j < 4; ++j) {
-        y[j+ 0] = d * kvalues_iq4nl[q4[j] & 0xf];
-        y[j+16] = d * kvalues_iq4nl[q4[j] >>  4];
-    }
-}
-
-template <int qk, int qr, dequantize_kernel_t dequantize_kernel, typename dst_t>
-static void dequantize_block_cuda(const void * __restrict__ vx, dst_t * __restrict__ y, const int k, cudaStream_t stream) {
-    const int num_blocks = (k + 2*CUDA_DEQUANTIZE_BLOCK_SIZE - 1) / (2*CUDA_DEQUANTIZE_BLOCK_SIZE);
-    dequantize_block<qk, qr, dequantize_kernel><<<num_blocks, CUDA_DEQUANTIZE_BLOCK_SIZE, 0, stream>>>(vx, y, k);
-}
-
-template<typename dst_t>
-static void dequantize_row_q2_K_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
-    const int nb = k / QK_K;
-    dequantize_block_q2_K<<<nb, 64, 0, stream>>>(vx, y);
-}
-
-template<typename dst_t>
-static void dequantize_row_q3_K_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
-    const int nb = k / QK_K;
-    dequantize_block_q3_K<<<nb, 64, 0, stream>>>(vx, y);
-}
-
-template<typename dst_t>
-static void dequantize_row_q4_K_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
-    const int nb = k / QK_K;
-    dequantize_block_q4_K<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template<typename dst_t>
-static void dequantize_row_q5_K_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
-    const int nb = k / QK_K;
-    dequantize_block_q5_K<<<nb, 64, 0, stream>>>(vx, y);
-}
-
-template<typename dst_t>
-static void dequantize_row_q6_K_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
-    const int nb = k / QK_K;
-    dequantize_block_q6_K<<<nb, 64, 0, stream>>>(vx, y);
-}
-
-template<typename dst_t>
-static void dequantize_row_iq2_xxs_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
-    const int nb = k / QK_K;
-    dequantize_block_iq2_xxs<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template<typename dst_t>
-static void dequantize_row_iq2_xs_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
-    const int nb = k / QK_K;
-    dequantize_block_iq2_xs<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template<typename dst_t>
-static void dequantize_row_iq2_s_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
-    const int nb = k / QK_K;
-    dequantize_block_iq2_s<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template<typename dst_t>
-static void dequantize_row_iq3_xxs_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
-    const int nb = k / QK_K;
-    dequantize_block_iq3_xxs<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template<typename dst_t>
-static void dequantize_row_iq3_s_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
-    const int nb = k / QK_K;
-    dequantize_block_iq3_s<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template<typename dst_t>
-static void dequantize_row_iq1_s_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
-    const int nb = k / QK_K;
-    dequantize_block_iq1_s<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template<typename dst_t>
-static void dequantize_row_iq1_m_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
-    const int nb = k / QK_K;
-    dequantize_block_iq1_m<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template<typename dst_t>
-static void dequantize_row_iq4_nl_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
-    const int nb = (k + QK_K - 1) / QK_K;
-    dequantize_block_iq4_nl<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template<typename dst_t>
-static void dequantize_row_iq4_xs_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) {
-    const int nb = (k + QK_K - 1) / QK_K;
-    dequantize_block_iq4_xs<<<nb, 32, 0, stream>>>(vx, y);
-}
-
-template<typename dst_t>
-static to_cuda_ggml_t<dst_t> ggml_get_to_cuda(int64_t type) {
-    switch (type) {
-        case 2:
-            return dequantize_block_cuda<QK4_0, QR4_0, dequantize_q4_0>;
-        case 3:
-            return dequantize_block_cuda<QK4_1, QR4_1, dequantize_q4_1>;
-        case 6:
-            return dequantize_block_cuda<QK5_0, QR5_0, dequantize_q5_0>;
-        case 7:
-            return dequantize_block_cuda<QK5_1, QR5_1, dequantize_q5_1>;
-        case 8:
-            return dequantize_block_cuda<QK8_0, QR8_0, dequantize_q8_0>;
-        case 10:
-            return dequantize_row_q2_K_cuda;
-        case 11:
-            return dequantize_row_q3_K_cuda;
-        case 12:
-            return dequantize_row_q4_K_cuda;
-        case 13:
-            return dequantize_row_q5_K_cuda;
-        case 14:
-            return dequantize_row_q6_K_cuda;
-        case 16:
-            return dequantize_row_iq2_xxs_cuda;
-        case 17:
-            return dequantize_row_iq2_xs_cuda;
-        case 18:
-            return dequantize_row_iq3_xxs_cuda;
-        case 19:
-            return dequantize_row_iq1_s_cuda;
-        case 20:
-            return dequantize_row_iq4_nl_cuda;
-        case 21:
-            return dequantize_row_iq3_s_cuda;
-        case 22:
-            return dequantize_row_iq2_s_cuda;
-        case 23:
-            return dequantize_row_iq4_xs_cuda;
-        case 29:
-            return dequantize_row_iq1_m_cuda;
-        default:
-            return nullptr;
-    }
-}

csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu

@@ -1,557 +0,0 @@
-#include <cuda_fp16.h>
-#include <cuda_runtime.h>
-
-#include "../../../cuda_compat.h"
-#include "../../dispatch_utils.h"
-#include "../../torch_utils.h"
-
-#include <torch/csrc/stable/ops.h>
-
-#include "ggml-common.h"
-#include "vecdotq.cuh"
-#include "dequantize.cuh"
-#include "mmvq.cuh"
-#include "mmq.cuh"
-#include "moe.cuh"
-#include "moe_vec.cuh"
-
-// Q8 gemv
-template <typename scalar_t>
-static __global__ void quantize_q8_1(const scalar_t* __restrict__ x,
-                                     void* __restrict__ vy, const int kx,
-                                     const int kx_padded) {
-  const auto ix = blockDim.x * blockIdx.x + threadIdx.x;
-  if (ix >= kx_padded) {
-    return;
-  }
-  const auto iy = blockDim.y * blockIdx.y + threadIdx.y;
-  const int i_padded = iy * kx_padded + ix;
-
-  block_q8_1* y = (block_q8_1*)vy;
-
-  const int ib = i_padded / QK8_1;   // block index
-  const int iqs = i_padded % QK8_1;  // quant index
-
-  const float xi = ix < kx ? static_cast<float>(x[iy * kx + ix]) : 0.0f;
-  float amax = fabsf(xi);
-  float sum = xi;
-
-#pragma unroll
-  for (int mask = 16; mask > 0; mask >>= 1) {
-    amax = fmaxf(amax, VLLM_SHFL_XOR_SYNC_WIDTH(amax, mask, 32));
-    sum += VLLM_SHFL_XOR_SYNC_WIDTH(sum, mask, 32);
-  }
-
-  const float d = amax / 127;
-  const int8_t q = amax == 0.0f ? 0 : roundf(xi / d);
-
-  y[ib].qs[iqs] = q;
-
-  if (iqs > 0) {
-    return;
-  }
-
-  y[ib].ds.x = __float2half(d);
-  y[ib].ds.y = __float2half(sum);
-}
-
-template <typename scalar_t>
-static void quantize_row_q8_1_cuda(const scalar_t* x, void* vy, const int kx,
-                                   const int ky, cudaStream_t stream) {
-  const int64_t kx_padded = (kx + 512 - 1) / 512 * 512;
-  const int block_num_x =
-      (kx_padded + CUDA_QUANTIZE_BLOCK_SIZE - 1) / CUDA_QUANTIZE_BLOCK_SIZE;
-  constexpr int MAX_BLOCK_SIZE = 65535;
-  for (int off = 0; off < ky; off += MAX_BLOCK_SIZE) {
-    const int num_blocks_y = std::min(ky, off + MAX_BLOCK_SIZE) - off;
-    const dim3 num_blocks(block_num_x, num_blocks_y, 1);
-    const dim3 block_size(CUDA_DEQUANTIZE_BLOCK_SIZE, 1, 1);
-    quantize_q8_1<<<num_blocks, block_size, 0, stream>>>(
-        &x[off * kx], (int32_t*)vy + off * (kx_padded / 32 * 9), kx, kx_padded);
-  }
-}
-
-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_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);
-  });
-
-  return DW;
-}
-
-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 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::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 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);
-
-    switch (type) {
-      case 2:
-        ggml_mul_mat_q4_0_q8_1_cuda(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream);
-        break;
-      case 3:
-        ggml_mul_mat_q4_1_q8_1_cuda(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream);
-        break;
-      case 6:
-        ggml_mul_mat_q5_0_q8_1_cuda(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream);
-        break;
-      case 7:
-        ggml_mul_mat_q5_1_q8_1_cuda(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream);
-        break;
-      case 8:
-        ggml_mul_mat_q8_0_q8_1_cuda(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream);
-        break;
-      case 10:
-        ggml_mul_mat_q2_K_q8_1_cuda(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream);
-        break;
-      case 11:
-        ggml_mul_mat_q3_K_q8_1_cuda(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream);
-        break;
-      case 12:
-        ggml_mul_mat_q4_K_q8_1_cuda(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream);
-        break;
-      case 13:
-        ggml_mul_mat_q5_K_q8_1_cuda(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream);
-        break;
-      case 14:
-        ggml_mul_mat_q6_K_q8_1_cuda(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream);
-        break;
-    }
-  });
-  return Y;
-}
-
-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 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) {
-      case 2:
-        ggml_moe_q4_0_q8_1_cuda(
-            (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
-            tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
-        break;
-      case 3:
-        ggml_moe_q4_1_q8_1_cuda(
-            (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
-            tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
-        break;
-      case 6:
-        ggml_moe_q5_0_q8_1_cuda(
-            (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
-            tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
-        break;
-      case 7:
-        ggml_moe_q5_1_q8_1_cuda(
-            (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
-            tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
-        break;
-      case 8:
-        ggml_moe_q8_0_q8_1_cuda(
-            (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
-            tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
-        break;
-      case 10:
-        ggml_moe_q2_K_q8_1_cuda(
-            (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
-            tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
-        break;
-      case 11:
-        ggml_moe_q3_K_q8_1_cuda(
-            (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
-            tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
-        break;
-      case 12:
-        ggml_moe_q4_K_q8_1_cuda(
-            (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
-            tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
-        break;
-      case 13:
-        ggml_moe_q5_K_q8_1_cuda(
-            (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
-            tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
-        break;
-      case 14:
-        ggml_moe_q6_K_q8_1_cuda(
-            (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
-            (int*)expert_ids.data_ptr(),
-            (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
-            tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
-        break;
-    }
-  });
-  return Y;
-}
-
-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 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);
-    switch (type) {
-      case 2:
-        moe_vec_q4_0_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
-            col, row, quant_X.stride(0), stream);
-        break;
-      case 3:
-        moe_vec_q4_1_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
-            col, row, quant_X.stride(0), stream);
-        break;
-      case 6:
-        moe_vec_q5_0_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
-            col, row, quant_X.stride(0), stream);
-        break;
-      case 7:
-        moe_vec_q5_1_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
-            col, row, quant_X.stride(0), stream);
-        break;
-      case 8:
-        moe_vec_q8_0_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
-            col, row, quant_X.stride(0), stream);
-        break;
-      case 10:
-        moe_vec_q2_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
-            col, row, quant_X.stride(0), stream);
-        break;
-      case 11:
-        moe_vec_q3_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
-            col, row, quant_X.stride(0), stream);
-        break;
-      case 12:
-        moe_vec_q4_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
-            col, row, quant_X.stride(0), stream);
-        break;
-      case 13:
-        moe_vec_q5_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
-            col, row, quant_X.stride(0), stream);
-        break;
-      case 14:
-        moe_vec_q6_K_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
-            col, row, quant_X.stride(0), stream);
-        break;
-      case 16:
-        moe_vec_iq2_xxs_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
-            col, row, quant_X.stride(0), stream);
-        break;
-      case 17:
-        moe_vec_iq2_xs_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
-            col, row, quant_X.stride(0), stream);
-        break;
-      case 18:
-        moe_vec_iq3_xxs_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
-            col, row, quant_X.stride(0), stream);
-        break;
-      case 19:
-        moe_vec_iq1_s_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
-            col, row, quant_X.stride(0), stream);
-        break;
-      case 20:
-        moe_vec_iq4_nl_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
-            col, row, quant_X.stride(0), stream);
-        break;
-      case 21:
-        moe_vec_iq3_s_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
-            col, row, quant_X.stride(0), stream);
-        break;
-      case 22:
-        moe_vec_iq2_s_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
-            col, row, quant_X.stride(0), stream);
-        break;
-      case 23:
-        moe_vec_iq4_xs_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
-            col, row, quant_X.stride(0), stream);
-        break;
-      case 29:
-        moe_vec_iq1_m_q8_1_cuda<scalar_t>(
-            (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
-            (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
-            col, row, quant_X.stride(0), stream);
-        break;
-    }
-  });
-  return Y;
-}
-
-int64_t ggml_moe_get_block_size(int64_t type) {
-  switch (type) {
-    case 2:
-      return MOE_X_Q4_0;
-    case 3:
-      return MOE_X_Q4_1;
-    case 6:
-      return MOE_X_Q5_0;
-    case 7:
-      return MOE_X_Q5_1;
-    case 8:
-      return MOE_X_Q8_0;
-    case 10:
-      return MOE_X_Q2_K;
-    case 11:
-      return MOE_X_Q3_K;
-    case 12:
-      return MOE_X_Q4_K;
-    case 13:
-      return MOE_X_Q5_K;
-    case 14:
-      return MOE_X_Q6_K;
-  }
-  return 0;
-}

csrc/libtorch_stable/quantization/gguf/mmq.cuh

@@ -1,610 +0,0 @@
-// copied from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmq.cu
-template <typename scalar_t, int qk, int qr, int qi, bool need_sum, typename block_q_t, int mmq_x, int mmq_y, int nwarps,
-              allocate_tiles_cuda_t allocate_tiles, load_tiles_cuda_t load_tiles, int vdr, vec_dot_q_mul_mat_cuda_t vec_dot>
-static __device__ __forceinline__ void mul_mat_q(
-    const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst,
-    const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) {
-
-    const block_q_t  * x = (const block_q_t  *) vx;
-    const block_q8_1 * y = (const block_q8_1 *) vy;
-
-    const int blocks_per_row_x = ncols_x / qk;
-    const int blocks_per_col_y = nrows_y / QK8_1;
-    const int blocks_per_warp = WARP_SIZE_GGUF / qi;
-
-    const int & ncols_dst = ncols_y;
-
-    const auto row_dst_0 = blockIdx.x*mmq_y;
-    const int & row_x_0 = row_dst_0;
-
-    const auto col_dst_0 = blockIdx.y*mmq_x;
-    const int & col_y_0 = col_dst_0;
-
-    int   * tile_x_ql = nullptr;
-    half2 * tile_x_dm = nullptr;
-    int   * tile_x_qh = nullptr;
-    int   * tile_x_sc = nullptr;
-
-    allocate_tiles(&tile_x_ql, &tile_x_dm, &tile_x_qh, &tile_x_sc);
-
-    __shared__ int    tile_y_qs[mmq_x * WARP_SIZE_GGUF];
-    __shared__ half2  tile_y_ds[mmq_x * WARP_SIZE_GGUF/QI8_1];
-
-    float sum[mmq_y/WARP_SIZE_GGUF][mmq_x/nwarps] = {{0.0f}};
-
-    for (int ib0 = 0; ib0 < blocks_per_row_x; ib0 += blocks_per_warp) {
-
-        load_tiles(x + row_x_0*blocks_per_row_x + ib0, tile_x_ql, tile_x_dm, tile_x_qh, tile_x_sc,
-                   threadIdx.y, nrows_x-row_x_0-1, threadIdx.x, blocks_per_row_x);
-
-#pragma unroll
-        for (int ir = 0; ir < qr && ib0 + ir * blocks_per_warp/qr < blocks_per_row_x; ++ir) {
-            const auto kqs = ir*WARP_SIZE_GGUF + threadIdx.x;
-            const int kbxd = kqs / QI8_1;
-
-#pragma unroll
-            for (int i = 0; i < mmq_x; i += nwarps) {
-                const int col_y_eff = min(col_y_0 + threadIdx.y + i, ncols_y-1); // to prevent out-of-bounds memory accesses
-                const block_q8_1 * by0 = &y[col_y_eff*blocks_per_col_y + ib0 * (qk/QK8_1) + kbxd];
-                const int index_y = (threadIdx.y + i) * WARP_SIZE_GGUF + kqs % WARP_SIZE_GGUF;
-                tile_y_qs[index_y] = get_int_from_int8_aligned(by0->qs, threadIdx.x % QI8_1);
-            }
-
-#pragma unroll
-            for (int ids0 = 0; ids0 < mmq_x; ids0 += nwarps * QI8_1) {
-                const int ids = (ids0 + threadIdx.y * QI8_1 + threadIdx.x / (WARP_SIZE_GGUF/QI8_1)) % mmq_x;
-                const auto kby = threadIdx.x % (WARP_SIZE_GGUF/QI8_1);
-                const int col_y_eff = min(col_y_0 + ids, ncols_y-1);
-
-                // if the sum is not needed it's faster to transform the scale to f32 ahead of time
-                const half2 * dsi_src = &y[col_y_eff*blocks_per_col_y + ib0 * (qk/QK8_1) + ir*(WARP_SIZE_GGUF/QI8_1) + kby].ds;
-                half2       * dsi_dst = &tile_y_ds[ids * (WARP_SIZE_GGUF/QI8_1) + kby];
-                if (need_sum) {
-                    *dsi_dst = *dsi_src;
-                } else {
-                    float * dfi_dst = (float *) dsi_dst;
-                    *dfi_dst = __low2float(*dsi_src);
-                }
-            }
-
-            __syncthreads();
-
-// #pragma unroll // unrolling this loop causes too much register pressure
-            for (int k = ir*WARP_SIZE_GGUF/qr; k < (ir+1)*WARP_SIZE_GGUF/qr; k += vdr) {
-#pragma unroll
-                for (int j = 0; j < mmq_x; j += nwarps) {
-#pragma unroll
-                    for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) {
-                        sum[i/WARP_SIZE_GGUF][j/nwarps] += vec_dot(
-                            tile_x_ql, tile_x_dm, tile_x_qh, tile_x_sc, tile_y_qs, tile_y_ds,
-                            threadIdx.x + i, threadIdx.y + j, k);
-                    }
-                }
-            }
-            __syncthreads();
-        }
-    }
-
-#pragma unroll
-    for (int j = 0; j < mmq_x; j += nwarps) {
-        const auto col_dst = col_dst_0 + j + threadIdx.y;
-        if (col_dst >= ncols_dst) {
-            return;
-        }
-
-#pragma unroll
-        for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) {
-            const auto row_dst = row_dst_0 + threadIdx.x + i;
-            if (row_dst >= nrows_dst) {
-                continue;
-            }
-            dst[col_dst*nrows_dst + row_dst] = sum[i/WARP_SIZE_GGUF][j/nwarps];
-        }
-    }
-}
-
-#if defined(USE_ROCM)
-#define  MMQ_X_Q4_0  64
-#define  MMQ_Y_Q4_0  128
-#define NWARPS_Q4_0  8
-#else
-#define  MMQ_X_Q4_0 4
-#define  MMQ_Y_Q4_0 32
-#define NWARPS_Q4_0 4
-#endif
-
-template<typename scalar_t, bool need_check> static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q4_0, 2)
-#endif
-mul_mat_q4_0(
-    const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst,
-    const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) {
-    const int mmq_x  =  MMQ_X_Q4_0;
-    const int mmq_y  =  MMQ_Y_Q4_0;
-    const int nwarps = NWARPS_Q4_0;
-
-    mul_mat_q<scalar_t, QK4_0, QR4_0, QI4_0, true, block_q4_0, mmq_x, mmq_y, nwarps, allocate_tiles_q4_0<mmq_y>,
-        load_tiles_q4_0<mmq_y, nwarps, need_check>, VDR_Q4_0_Q8_1_MMQ, vec_dot_q4_0_q8_1_mul_mat>
-        (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template<typename scalar_t>
-static void ggml_mul_mat_q4_0_q8_1_cuda(
-    const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x,
-    const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) {
-
-    int mmq_x  =  MMQ_X_Q4_0;
-    int mmq_y  =  MMQ_Y_Q4_0;
-    int nwarps = NWARPS_Q4_0;
-
-    const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-    const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-    const dim3 block_nums(block_num_x, block_num_y, 1);
-    const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-    if (nrows_x % mmq_y == 0) {
-        const bool need_check = false;
-        mul_mat_q4_0<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    } else {
-        const bool need_check = true;
-        mul_mat_q4_0<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    }
-}
-
-#if defined(USE_ROCM)
-#define  MMQ_X_Q4_1 64
-#define  MMQ_Y_Q4_1 128
-#define NWARPS_Q4_1 8
-#else
-#define  MMQ_X_Q4_1 4
-#define  MMQ_Y_Q4_1 32
-#define NWARPS_Q4_1 4
-#endif
-
-template<typename scalar_t, bool need_check> static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q4_1, 2)
-#endif
-mul_mat_q4_1(
-    const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst,
-    const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) {
-    const int mmq_x  =  MMQ_X_Q4_1;
-    const int mmq_y  =  MMQ_Y_Q4_1;
-    const int nwarps = NWARPS_Q4_1;
-
-    mul_mat_q<scalar_t, QK4_1, QR4_1, QI4_1, true, block_q4_1, mmq_x, mmq_y, nwarps, allocate_tiles_q4_1<mmq_y>,
-        load_tiles_q4_1<mmq_y, nwarps, need_check>, VDR_Q4_1_Q8_1_MMQ, vec_dot_q4_1_q8_1_mul_mat>
-        (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template<typename scalar_t>
-static void ggml_mul_mat_q4_1_q8_1_cuda(
-    const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x,
-    const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) {
-
-    int mmq_x  =  MMQ_X_Q4_1;
-    int mmq_y  =  MMQ_Y_Q4_1;
-    int nwarps = NWARPS_Q4_1;
-
-    const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-    const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-    const dim3 block_nums(block_num_x, block_num_y, 1);
-    const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-    if (nrows_x % mmq_y == 0) {
-        const bool need_check = false;
-        mul_mat_q4_1<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    } else {
-        const bool need_check = true;
-        mul_mat_q4_1<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    }
-}
-
-#if defined(USE_ROCM)
-#define  MMQ_X_Q5_0 64
-#define  MMQ_Y_Q5_0 128
-#define NWARPS_Q5_0 8
-#else
-#define  MMQ_X_Q5_0 4
-#define  MMQ_Y_Q5_0 32
-#define NWARPS_Q5_0 4
-#endif
-
-template<typename scalar_t, bool need_check> static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q5_0, 2)
-#endif
-mul_mat_q5_0(
-    const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst,
-    const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) {
-    const int mmq_x  =  MMQ_X_Q5_0;
-    const int mmq_y  =  MMQ_Y_Q5_0;
-    const int nwarps = NWARPS_Q5_0;
-
-    mul_mat_q<scalar_t, QK5_0, QR5_0, QI5_0, false, block_q5_0, mmq_x, mmq_y, nwarps, allocate_tiles_q5_0<mmq_y>,
-        load_tiles_q5_0<mmq_y, nwarps, need_check>, VDR_Q5_0_Q8_1_MMQ, vec_dot_q5_0_q8_1_mul_mat>
-        (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template<typename scalar_t>
-static void ggml_mul_mat_q5_0_q8_1_cuda(
-    const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x,
-    const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) {
-
-    const int mmq_x  =  MMQ_X_Q5_0;
-    const int mmq_y  =  MMQ_Y_Q5_0;
-    const int nwarps = NWARPS_Q5_0;
-
-    const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-    const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-    const dim3 block_nums(block_num_x, block_num_y, 1);
-    const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-    if (nrows_x % mmq_y == 0) {
-        const bool need_check = false;
-        mul_mat_q5_0<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    } else {
-        const bool need_check = true;
-        mul_mat_q5_0<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    }
-}
-
-#if defined(USE_ROCM)
-#define  MMQ_X_Q5_1 64
-#define  MMQ_Y_Q5_1 128
-#define NWARPS_Q5_1 8
-#else
-#define  MMQ_X_Q5_1 4
-#define  MMQ_Y_Q5_1 32
-#define NWARPS_Q5_1 4
-#endif
-
-template<typename scalar_t, bool need_check> static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q5_1, 2)
-#endif
-mul_mat_q5_1(
-    const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst,
-    const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) {
-    const int mmq_x  =  MMQ_X_Q5_1;
-    const int mmq_y  =  MMQ_Y_Q5_1;
-    const int nwarps = NWARPS_Q5_1;
-
-    mul_mat_q<scalar_t, QK5_1, QR5_1, QI5_1, true, block_q5_1, mmq_x, mmq_y, nwarps, allocate_tiles_q5_1<mmq_y>,
-        load_tiles_q5_1<mmq_y, nwarps, need_check>, VDR_Q5_1_Q8_1_MMQ, vec_dot_q5_1_q8_1_mul_mat>
-        (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template<typename scalar_t>
-static void ggml_mul_mat_q5_1_q8_1_cuda(
-    const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x,
-    const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) {
-    const int mmq_x  =  MMQ_X_Q5_1;
-    const int mmq_y  =  MMQ_Y_Q5_1;
-    const int nwarps = NWARPS_Q5_1;
-
-    const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-    const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-    const dim3 block_nums(block_num_x, block_num_y, 1);
-    const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-    if (nrows_x % mmq_y == 0) {
-        const bool need_check = false;
-        mul_mat_q5_1<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    } else {
-        const bool need_check = true;
-        mul_mat_q5_1<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    }
-}
-
-#if defined(USE_ROCM)
-#define  MMQ_X_Q8_0 64
-#define  MMQ_Y_Q8_0 128
-#define NWARPS_Q8_0 8
-#else
-#define  MMQ_X_Q8_0 4
-#define  MMQ_Y_Q8_0 32
-#define NWARPS_Q8_0 4
-#endif
-
-template<typename scalar_t, bool need_check> static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q8_0, 2)
-#endif
-mul_mat_q8_0(
-    const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst,
-    const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) {
-    const int mmq_x  =  MMQ_X_Q8_0;
-    const int mmq_y  =  MMQ_Y_Q8_0;
-    const int nwarps = NWARPS_Q8_0;
-
-    mul_mat_q<scalar_t, QK8_0, QR8_0, QI8_0, false, block_q8_0, mmq_x, mmq_y, nwarps, allocate_tiles_q8_0<mmq_y>,
-        load_tiles_q8_0<mmq_y, nwarps, need_check>, VDR_Q8_0_Q8_1_MMQ, vec_dot_q8_0_q8_1_mul_mat>
-        (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template<typename scalar_t>
-static void ggml_mul_mat_q8_0_q8_1_cuda(
-    const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x,
-    const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) {
-    const int mmq_x  =  MMQ_X_Q8_0;
-    const int mmq_y  =  MMQ_Y_Q8_0;
-    const int nwarps = NWARPS_Q8_0;
-
-    const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-    const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-    const dim3 block_nums(block_num_x, block_num_y, 1);
-    const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-    if (nrows_x % mmq_y == 0) {
-        const bool need_check = false;
-        mul_mat_q8_0<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    } else {
-        const bool need_check = true;
-        mul_mat_q8_0<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    }
-}
-
-#if defined(USE_ROCM)
-#define  MMQ_X_Q2_K 64
-#define  MMQ_Y_Q2_K 128
-#define NWARPS_Q2_K 8
-#else
-#define  MMQ_X_Q2_K 4
-#define  MMQ_Y_Q2_K 32
-#define NWARPS_Q2_K 4
-#endif
-
-template<typename scalar_t, bool need_check> static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q2_K, 2)
-#endif
-mul_mat_q2_K(
-    const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst,
-    const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) {
-    const int mmq_x  =  MMQ_X_Q2_K;
-    const int mmq_y  =  MMQ_Y_Q2_K;
-    const int nwarps = NWARPS_Q2_K;
-
-    mul_mat_q<scalar_t, QK_K, QR2_K, QI2_K, false, block_q2_K, mmq_x, mmq_y, nwarps, allocate_tiles_q2_K<mmq_y>,
-        load_tiles_q2_K<mmq_y, nwarps, need_check>, VDR_Q2_K_Q8_1_MMQ, vec_dot_q2_K_q8_1_mul_mat>
-        (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template<typename scalar_t>
-static void ggml_mul_mat_q2_K_q8_1_cuda(
-    const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x,
-    const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) {
-    const int mmq_x  =  MMQ_X_Q2_K;
-    const int mmq_y  =  MMQ_Y_Q2_K;
-    const int nwarps = NWARPS_Q2_K;
-
-    const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-    const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-    const dim3 block_nums(block_num_x, block_num_y, 1);
-    const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-    if (nrows_x % mmq_y == 0) {
-        const bool need_check = false;
-        mul_mat_q2_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    } else {
-        const bool need_check = true;
-        mul_mat_q2_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    }
-}
-
-#if defined(USE_ROCM)
-#define  MMQ_X_Q3_K 64
-#define  MMQ_Y_Q3_K 128
-#define NWARPS_Q3_K 8
-#else
-#define  MMQ_X_Q3_K 4
-#define  MMQ_Y_Q3_K 32
-#define NWARPS_Q3_K 4
-#endif
-
-template<typename scalar_t, bool need_check> static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q3_K, 2)
-#endif
-mul_mat_q3_K(
-    const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst,
-    const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) {
-
-    const int mmq_x  =  MMQ_X_Q3_K;
-    const int mmq_y  =  MMQ_Y_Q3_K;
-    const int nwarps = NWARPS_Q3_K;
-
-    mul_mat_q<scalar_t, QK_K, QR3_K, QI3_K, false, block_q3_K, mmq_x, mmq_y, nwarps, allocate_tiles_q3_K<mmq_y>,
-        load_tiles_q3_K<mmq_y, nwarps, need_check>, VDR_Q3_K_Q8_1_MMQ, vec_dot_q3_K_q8_1_mul_mat>
-        (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template<typename scalar_t>
-static void ggml_mul_mat_q3_K_q8_1_cuda(
-    const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x,
-    const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) {
-
-    const int mmq_x  =  MMQ_X_Q3_K;
-    const int mmq_y  =  MMQ_Y_Q3_K;
-    const int nwarps = NWARPS_Q3_K;
-
-    const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-    const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-    const dim3 block_nums(block_num_x, block_num_y, 1);
-    const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-    if (nrows_x % mmq_y == 0) {
-        const bool need_check = false;
-        mul_mat_q3_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    } else {
-        const bool need_check = true;
-        mul_mat_q3_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    }
-}
-
-#if defined(USE_ROCM)
-#define  MMQ_X_Q4_K 64
-#define  MMQ_Y_Q4_K 128
-#define NWARPS_Q4_K 8
-#else
-#define  MMQ_X_Q4_K 4
-#define  MMQ_Y_Q4_K 32
-#define NWARPS_Q4_K 4
-#endif
-
-template<typename scalar_t, bool need_check> static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q4_K, 2)
-#endif
-mul_mat_q4_K(
-    const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst,
-    const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) {
-    const int mmq_x  =  MMQ_X_Q4_K;
-    const int mmq_y  =  MMQ_Y_Q4_K;
-    const int nwarps = NWARPS_Q4_K;
-
-    mul_mat_q<scalar_t, QK_K, QR4_K, QI4_K, true, block_q4_K, mmq_x, mmq_y, nwarps, allocate_tiles_q4_K<mmq_y>,
-        load_tiles_q4_K<mmq_y, nwarps, need_check>, VDR_Q4_K_Q8_1_MMQ, vec_dot_q4_K_q8_1_mul_mat>
-        (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template<typename scalar_t>
-static void ggml_mul_mat_q4_K_q8_1_cuda(
-    const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x,
-    const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) {
-    const int mmq_x  =  MMQ_X_Q4_K;
-    const int mmq_y  =  MMQ_Y_Q4_K;
-    const int nwarps = NWARPS_Q4_K;
-
-    const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-    const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-    const dim3 block_nums(block_num_x, block_num_y, 1);
-    const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-    if (nrows_x % mmq_y == 0) {
-        const bool need_check = false;
-        mul_mat_q4_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    } else {
-        const bool need_check = true;
-        mul_mat_q4_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    }
-}
-
-#if defined(USE_ROCM)
-#define  MMQ_X_Q5_K 64
-#define  MMQ_Y_Q5_K 128
-#define NWARPS_Q5_K 8
-#else
-#define  MMQ_X_Q5_K 4
-#define  MMQ_Y_Q5_K 32
-#define NWARPS_Q5_K 4
-#endif
-
-template<typename scalar_t, bool need_check> static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q5_K, 2)
-#endif
-mul_mat_q5_K(
-    const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst,
-    const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) {
-    const int mmq_x  =  MMQ_X_Q5_K;
-    const int mmq_y  =  MMQ_Y_Q5_K;
-    const int nwarps = NWARPS_Q5_K;
-
-    mul_mat_q<scalar_t, QK_K, QR5_K, QI5_K, true, block_q5_K, mmq_x, mmq_y, nwarps, allocate_tiles_q5_K<mmq_y>,
-        load_tiles_q5_K<mmq_y, nwarps, need_check>, VDR_Q5_K_Q8_1_MMQ, vec_dot_q5_K_q8_1_mul_mat>
-        (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template<typename scalar_t>
-static void ggml_mul_mat_q5_K_q8_1_cuda(
-    const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x,
-    const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) {
-
-    const int mmq_x  =  MMQ_X_Q5_K;
-    const int mmq_y  =  MMQ_Y_Q5_K;
-    const int nwarps = NWARPS_Q5_K;
-
-    const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-    const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-    const dim3 block_nums(block_num_x, block_num_y, 1);
-    const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-    if (nrows_x % mmq_y == 0) {
-        const bool need_check = false;
-        mul_mat_q5_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    } else {
-        const bool need_check = true;
-        mul_mat_q5_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    }
-}
-
-#if defined(USE_ROCM)
-#define  MMQ_X_Q6_K 64
-#define  MMQ_Y_Q6_K 128
-#define NWARPS_Q6_K 8
-#else
-#define  MMQ_X_Q6_K 4
-#define  MMQ_Y_Q6_K 32
-#define NWARPS_Q6_K 4
-#endif
-
-template<typename scalar_t, bool need_check> static __global__ void
-#if defined(USE_ROCM)
-__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q6_K, 2)
-#endif
-mul_mat_q6_K(
-    const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst,
-    const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) {
-    const int mmq_x  =  MMQ_X_Q6_K;
-    const int mmq_y  =  MMQ_Y_Q6_K;
-    const int nwarps = NWARPS_Q6_K;
-
-    mul_mat_q<scalar_t, QK_K, QR6_K, QI6_K, false, block_q6_K, mmq_x, mmq_y, nwarps, allocate_tiles_q6_K<mmq_y>,
-        load_tiles_q6_K<mmq_y, nwarps, need_check>, VDR_Q6_K_Q8_1_MMQ, vec_dot_q6_K_q8_1_mul_mat>
-        (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-}
-
-template<typename scalar_t>
-static void ggml_mul_mat_q6_K_q8_1_cuda(
-    const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x,
-    const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) {
-    const int mmq_x  =  MMQ_X_Q6_K;
-    const int mmq_y  =  MMQ_Y_Q6_K;
-    const int nwarps = NWARPS_Q6_K;
-
-    const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
-    const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x;
-    const dim3 block_nums(block_num_x, block_num_y, 1);
-    const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
-
-    if (nrows_x % mmq_y == 0) {
-        const bool need_check = false;
-        mul_mat_q6_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    } else {
-        const bool need_check = true;
-        mul_mat_q6_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>
-            (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst);
-    }
-}

csrc/libtorch_stable/quantization/gguf/mmvq.cuh

@@ -1,212 +0,0 @@
-// copied and adapted from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmvq.cu
-template <typename scalar_t, int qk, int qi, typename block_q_t, int vdr, vec_dot_q_cuda_t vec_dot_q_cuda>
-static __global__ void mul_mat_vec_q(const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst, const int ncols, const int nrows, const int nvecs) {
-    const auto row = blockIdx.x*blockDim.y + threadIdx.y;
-    const auto vec = blockIdx.y;
-
-    if (row >= nrows || vec >= nvecs) {
-        return;
-    }
-
-    const int blocks_per_row = ncols / qk;
-    const int blocks_per_warp = vdr * WARP_SIZE / qi;
-    const int nrows_y = (ncols + 512 - 1) / 512 * 512;
-
-
-    // partial sum for each thread
-    float tmp = 0.0f;
-
-    const block_q_t  * x = (const block_q_t  *) vx;
-    const block_q8_1 * y = (const block_q8_1 *) vy;
-
-    for (auto i = threadIdx.x / (qi/vdr); i < blocks_per_row; i += blocks_per_warp) {
-        const int ibx = row*blocks_per_row + i; // x block index
-
-        const int iby = vec*(nrows_y/QK8_1) + i * (qk/QK8_1); // y block index that aligns with ibx
-
-        const int iqs  = vdr * (threadIdx.x % (qi/vdr)); // x block quant index when casting the quants to int
-
-        tmp += vec_dot_q_cuda(&x[ibx], &y[iby], iqs);
-    }
-
-    // sum up partial sums and write back result
-#pragma unroll
-    for (int mask = WARP_SIZE/2; mask > 0; mask >>= 1) {
-        tmp += VLLM_SHFL_XOR_SYNC(tmp, mask);
-    }
-
-    if (threadIdx.x == 0) {
-        dst[vec*nrows + row] = tmp;
-    }
-}
-
-template<typename scalar_t>
-static void mul_mat_vec_q4_0_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
-    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-    const dim3 block_nums(block_num_y, nvecs, 1);
-    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-    mul_mat_vec_q<scalar_t, QK4_0, QI4_0, block_q4_0, VDR_Q4_0_Q8_1_MMVQ, vec_dot_q4_0_q8_1>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template<typename scalar_t>
-static void mul_mat_vec_q4_1_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
-    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-    const dim3 block_nums(block_num_y, nvecs, 1);
-    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-    mul_mat_vec_q<scalar_t, QK4_0, QI4_1, block_q4_1, VDR_Q4_1_Q8_1_MMVQ, vec_dot_q4_1_q8_1>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template<typename scalar_t>
-static void mul_mat_vec_q5_0_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
-    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-    const dim3 block_nums(block_num_y, nvecs, 1);
-    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-    mul_mat_vec_q<scalar_t, QK5_0, QI5_0, block_q5_0, VDR_Q5_0_Q8_1_MMVQ, vec_dot_q5_0_q8_1>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template<typename scalar_t>
-static void mul_mat_vec_q5_1_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
-    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-    const dim3 block_nums(block_num_y, nvecs, 1);
-    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-    mul_mat_vec_q<scalar_t, QK5_1, QI5_1, block_q5_1, VDR_Q5_1_Q8_1_MMVQ, vec_dot_q5_1_q8_1>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template<typename scalar_t>
-static void mul_mat_vec_q8_0_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
-    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-    const dim3 block_nums(block_num_y, nvecs, 1);
-    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-    mul_mat_vec_q<scalar_t, QK8_0, QI8_0, block_q8_0, VDR_Q8_0_Q8_1_MMVQ, vec_dot_q8_0_q8_1>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template<typename scalar_t>
-static void mul_mat_vec_q2_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
-    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-    const dim3 block_nums(block_num_y, nvecs, 1);
-    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-    mul_mat_vec_q<scalar_t, QK_K, QI2_K, block_q2_K, VDR_Q2_K_Q8_1_MMVQ, vec_dot_q2_K_q8_1>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template<typename scalar_t>
-static void mul_mat_vec_q3_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
-    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-    const dim3 block_nums(block_num_y, nvecs, 1);
-    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-    mul_mat_vec_q<scalar_t, QK_K, QI3_K, block_q3_K, VDR_Q3_K_Q8_1_MMVQ, vec_dot_q3_K_q8_1>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template<typename scalar_t>
-static void mul_mat_vec_q4_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
-    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-    const dim3 block_nums(block_num_y, nvecs, 1);
-    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-    mul_mat_vec_q<scalar_t, QK_K, QI4_K, block_q4_K, VDR_Q4_K_Q8_1_MMVQ, vec_dot_q4_K_q8_1>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template<typename scalar_t>
-static void mul_mat_vec_q5_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
-    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-    const dim3 block_nums(block_num_y, nvecs, 1);
-    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-    mul_mat_vec_q<scalar_t, QK_K, QI5_K, block_q5_K, VDR_Q5_K_Q8_1_MMVQ, vec_dot_q5_K_q8_1>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template<typename scalar_t>
-static void mul_mat_vec_q6_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
-    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-    const dim3 block_nums(block_num_y, nvecs, 1);
-    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-    mul_mat_vec_q<scalar_t, QK_K, QI6_K, block_q6_K, VDR_Q6_K_Q8_1_MMVQ, vec_dot_q6_K_q8_1>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template<typename scalar_t>
-static void mul_mat_vec_iq2_xxs_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
-    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-    const dim3 block_nums(block_num_y, nvecs, 1);
-    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-    mul_mat_vec_q<scalar_t, QK_K, QI2_XXS, block_iq2_xxs, 1, vec_dot_iq2_xxs_q8_1>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template<typename scalar_t>
-static void mul_mat_vec_iq2_xs_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
-    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-    const dim3 block_nums(block_num_y, nvecs, 1);
-    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-    mul_mat_vec_q<scalar_t, QK_K, QI2_XS, block_iq2_xs, 1, vec_dot_iq2_xs_q8_1>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template<typename scalar_t>
-static void mul_mat_vec_iq2_s_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
-    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-    const dim3 block_nums(block_num_y, nvecs, 1);
-    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-    mul_mat_vec_q<scalar_t, QK_K, QI2_S, block_iq2_s, 1, vec_dot_iq2_s_q8_1>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template<typename scalar_t>
-static void mul_mat_vec_iq3_xxs_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
-    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-    const dim3 block_nums(block_num_y, nvecs, 1);
-    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-    mul_mat_vec_q<scalar_t, QK_K, QI3_XXS, block_iq3_xxs, 1, vec_dot_iq3_xxs_q8_1>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template<typename scalar_t>
-static void mul_mat_vec_iq1_s_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
-    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-    const dim3 block_nums(block_num_y, nvecs, 1);
-    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-    mul_mat_vec_q<scalar_t, QK_K, QI1_S, block_iq1_s, 1, vec_dot_iq1_s_q8_1>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template<typename scalar_t>
-static void mul_mat_vec_iq1_m_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
-    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-    const dim3 block_nums(block_num_y, nvecs, 1);
-    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-    mul_mat_vec_q<scalar_t, QK_K, QI1_M, block_iq1_m, 1, vec_dot_iq1_m_q8_1>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template<typename scalar_t>
-static void mul_mat_vec_iq4_nl_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
-    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-    const dim3 block_nums(block_num_y, nvecs, 1);
-    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-    mul_mat_vec_q<scalar_t, QK4_NL, QI4_NL, block_iq4_nl, VDR_Q4_0_Q8_1_MMVQ, vec_dot_iq4_nl_q8_1>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template<typename scalar_t>
-static void mul_mat_vec_iq4_xs_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
-    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-    const dim3 block_nums(block_num_y, nvecs, 1);
-    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-    mul_mat_vec_q<scalar_t, QK_K, QI4_XS, block_iq4_xs, 1, vec_dot_iq4_xs_q8_1>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}
-
-template<typename scalar_t>
-static void mul_mat_vec_iq3_s_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
-    const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
-    const dim3 block_nums(block_num_y, nvecs, 1);
-    const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
-    mul_mat_vec_q<scalar_t, QK_K, QI3_XS, block_iq3_s, 1, vec_dot_iq3_s_q8_1>
-        <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
-}