Merge branch 'main' into wentao-optimize-per-token-group-quant

2026-06-12 16:00:35 +00:00
parent c4bf7dadb9 a30addc754
commit 349af588f0
283 changed files with 22825 additions and 8616 deletions
@@ -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"
      
@@ -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"
@@ -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
@@ -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
@@ -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
@@ -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/

 
@@ -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
@@ -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:
@@ -358,145 +358,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    SRCS "${VLLM_EXT_SRC}"
    CUDA_ARCHS "${CUDA_ARCHS}")

-  # Only build Marlin kernels if we are building for at least some compatible archs.
-  # Keep building Marlin for 9.0 as there are some group sizes and shapes that
-  # are not supported by Machete yet.
-
-  # marlin arches for fp16 output
-  # Family-conditional 12.0f (one cubin for SM12x family) requires CUDA >= 13.0;
-  # fall back to architecture-specific 12.0a;12.1a on CUDA < 13.0 (e.g. 12.8).
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
-  endif()
-  # marlin has limited support for turing
-  cuda_archs_loose_intersection(MARLIN_SM75_ARCHS "7.5" "${CUDA_ARCHS}")
-  # marlin arches for bf16 output (we need 9.0 for bf16 atomicAdd PTX)
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
-  endif()
-  # marlin arches for fp8 input
-  # - sm80 doesn't support fp8 computation
-  # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
-  # so we only enable fp8 computation for SM89 (e.g. RTX 40x0)  and 12.0 (e.g. RTX 50x0)
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0a;12.1a" "${CUDA_ARCHS}")
-  endif()
-  # marlin arches for other files
-  cuda_archs_loose_intersection(MARLIN_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}")
-
-  if (MARLIN_OTHER_ARCHS)
-
-    #
-    # For the Marlin kernels we automatically generate sources for various
-    # preselected input type pairs and schedules.
-    # Generate sources:
-    set(MARLIN_GEN_SCRIPT
-      ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/marlin/generate_kernels.py)
-    file(MD5 ${MARLIN_GEN_SCRIPT} MARLIN_GEN_SCRIPT_HASH)
-    list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR)
-    set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})")
-
-    message(STATUS "Marlin generation script hash: ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
-    message(STATUS "Last run Marlin generate script hash: $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
-
-    if (NOT DEFINED CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
-        OR NOT $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH} STREQUAL ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH})
-      execute_process(
-        COMMAND ${CMAKE_COMMAND} -E env
-        PYTHONPATH=$ENV{PYTHONPATH}
-          ${Python_EXECUTABLE} ${MARLIN_GEN_SCRIPT} ${CUDA_ARCHS_STR}
-        RESULT_VARIABLE marlin_generation_result
-        OUTPUT_VARIABLE marlin_generation_result
-        OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log
-        ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log
-      )
-
-      if (NOT marlin_generation_result EQUAL 0)
-        message(FATAL_ERROR "Marlin generation failed."
-                            " Result: \"${marlin_generation_result}\""
-                            "\nCheck the log for details: "
-                            "${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log")
-      else()
-        set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
-            CACHE STRING "Last run Marlin generate script hash and arch" FORCE)
-        message(STATUS "Marlin generation completed successfully.")
-      endif()
-    else()
-      message(STATUS "Marlin generation script has not changed, skipping generation.")
-    endif()
-
-    if (MARLIN_ARCHS)
-      file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/quantization/marlin/sm80_kernel_*_float16.cu")
-      set_gencode_flags_for_srcs(
-        SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}"
-        CUDA_ARCHS "${MARLIN_ARCHS}")
-      if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
-        set_source_files_properties(${MARLIN_TEMPLATE_KERNEL_SRC}
-          PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
-      endif()
-      list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC})
-
-      file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/quantization/marlin/sm80_kernel_*_bfloat16.cu")
-      set_gencode_flags_for_srcs(
-        SRCS "${MARLIN_TEMPLATE_BF16_KERNEL_SRC}"
-        CUDA_ARCHS "${MARLIN_BF16_ARCHS}")
-      if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
-        set_source_files_properties(${MARLIN_TEMPLATE_BF16_KERNEL_SRC}
-          PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
-      endif()
-      list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_BF16_KERNEL_SRC})
-    endif()
-
-    if (MARLIN_SM75_ARCHS)
-      file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/quantization/marlin/sm75_kernel_*.cu")
-      set_gencode_flags_for_srcs(
-        SRCS "${MARLIN_TEMPLATE_SM75_KERNEL_SRC}"
-        CUDA_ARCHS "${MARLIN_SM75_ARCHS}")
-      if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
-        set_source_files_properties(${MARLIN_TEMPLATE_SM75_KERNEL_SRC}
-          PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
-      endif()
-      list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_SM75_KERNEL_SRC})
-    endif()
-
-    if (MARLIN_FP8_ARCHS)
-      file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/quantization/marlin/sm89_kernel_*.cu")
-      set_gencode_flags_for_srcs(
-        SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}"
-        CUDA_ARCHS "${MARLIN_FP8_ARCHS}")
-      if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
-        set_source_files_properties(${MARLIN_TEMPLATE_FP8_KERNEL_SRC}
-          PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
-      endif()
-      list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_FP8_KERNEL_SRC})
-    endif()
-
-    set(MARLIN_SRCS
-       "csrc/quantization/marlin/marlin.cu"
-       "csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu"
-       "csrc/quantization/marlin/gptq_marlin_repack.cu"
-       "csrc/quantization/marlin/awq_marlin_repack.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${MARLIN_SRCS}"
-      CUDA_ARCHS "${MARLIN_OTHER_ARCHS}")
-    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
-      set_source_files_properties(${MARLIN_SRCS}
-        PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
-    endif()
-    list(APPEND VLLM_EXT_SRC "${MARLIN_SRCS}")
-
-    message(STATUS "Building Marlin kernels for archs: ${MARLIN_OTHER_ARCHS}")
-  else()
-    message(STATUS "Not building Marlin kernels as no compatible archs found"
-                   " in CUDA target architectures")
-  endif()
-
  # Expert-specialization MXFP8 blockscaled grouped kernels (SM100+).
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
    cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
@@ -524,76 +385,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    endif()
  endif()

-  #
-  # Machete kernels
-
-  # The machete kernels only work on hopper and require CUDA 12.0 or later.
-  # Only build Machete kernels if we are building for something compatible with sm90a
-  cuda_archs_loose_intersection(MACHETE_ARCHS "9.0a" "${CUDA_ARCHS}")
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND MACHETE_ARCHS)
-    #
-    # For the Machete kernels we automatically generate sources for various
-    # preselected input type pairs and schedules.
-    # Generate sources:
-    set(MACHETE_GEN_SCRIPT
-      ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/machete/generate.py)
-    file(MD5 ${MACHETE_GEN_SCRIPT} MACHETE_GEN_SCRIPT_HASH)
-
-    message(STATUS "Machete generation script hash: ${MACHETE_GEN_SCRIPT_HASH}")
-    message(STATUS "Last run machete generate script hash: $CACHE{MACHETE_GEN_SCRIPT_HASH}")
-
-    if (NOT DEFINED CACHE{MACHETE_GEN_SCRIPT_HASH}
-        OR NOT $CACHE{MACHETE_GEN_SCRIPT_HASH} STREQUAL ${MACHETE_GEN_SCRIPT_HASH})
-      execute_process(
-        COMMAND ${CMAKE_COMMAND} -E env
-        PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_extensions/:${CUTLASS_DIR}/python/:${VLLM_PYTHON_PATH}:$ENV{PYTHONPATH}
-          ${Python_EXECUTABLE} ${MACHETE_GEN_SCRIPT}
-        RESULT_VARIABLE machete_generation_result
-        OUTPUT_VARIABLE machete_generation_output
-        OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log
-        ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log
-      )
-
-      if (NOT machete_generation_result EQUAL 0)
-        message(FATAL_ERROR "Machete generation failed."
-                            " Result: \"${machete_generation_result}\""
-                            "\nCheck the log for details: "
-                            "${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log")
-      else()
-        set(MACHETE_GEN_SCRIPT_HASH ${MACHETE_GEN_SCRIPT_HASH}
-            CACHE STRING "Last run machete generate script hash" FORCE)
-        message(STATUS "Machete generation completed successfully.")
-      endif()
-    else()
-      message(STATUS "Machete generation script has not changed, skipping generation.")
-    endif()
-
-    # Add machete generated sources
-    file(GLOB MACHETE_GEN_SOURCES "csrc/quantization/machete/generated/*.cu")
-    list(APPEND VLLM_EXT_SRC ${MACHETE_GEN_SOURCES})
-
-    # forward compatible
-    set_gencode_flags_for_srcs(
-      SRCS "${MACHETE_GEN_SOURCES}"
-      CUDA_ARCHS "${MACHETE_ARCHS}")
-
-    list(APPEND VLLM_EXT_SRC
-      csrc/quantization/machete/machete_pytorch.cu)
-
-    message(STATUS "Building Machete kernels for archs: ${MACHETE_ARCHS}")
-  else()
-    if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0
-        AND MACHETE_ARCHS)
-      message(STATUS "Not building Machete kernels as CUDA Compiler version is "
-                     "not >= 12.0, we recommend upgrading to CUDA 12.0 or "
-                     "later if you intend on running w4a16 quantized models on "
-                     "Hopper.")
-    else()
-      message(STATUS "Not building Machete kernels as no compatible archs "
-                     "found in CUDA target architectures")
-    endif()
-  endif()
-


 # if CUDA endif
@@ -672,10 +463,218 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
      "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu"
      "csrc/libtorch_stable/minimax_reduce_rms_kernel.cu")

+    #
+    # Machete kernels
+    #
+    # The machete kernels only work on hopper and require CUDA 12.0 or later.
+    # Only build Machete kernels if we are building for something compatible with sm90a
+    cuda_archs_loose_intersection(MACHETE_ARCHS "9.0a" "${CUDA_ARCHS}")
+    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND MACHETE_ARCHS)
+      #
+      # For the Machete kernels we automatically generate sources for various
+      # preselected input type pairs and schedules.
+      # Generate sources:
+      set(MACHETE_GEN_SCRIPT
+        ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/quantization/machete/generate.py)
+      file(MD5 ${MACHETE_GEN_SCRIPT} MACHETE_GEN_SCRIPT_HASH)
+
+      message(STATUS "Machete generation script hash: ${MACHETE_GEN_SCRIPT_HASH}")
+      message(STATUS "Last run machete generate script hash: $CACHE{MACHETE_GEN_SCRIPT_HASH}")
+
+      if (NOT DEFINED CACHE{MACHETE_GEN_SCRIPT_HASH}
+          OR NOT $CACHE{MACHETE_GEN_SCRIPT_HASH} STREQUAL ${MACHETE_GEN_SCRIPT_HASH})
+        execute_process(
+          COMMAND ${CMAKE_COMMAND} -E env
+          PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_extensions/:${CUTLASS_DIR}/python/:${VLLM_PYTHON_PATH}:$ENV{PYTHONPATH}
+            ${Python_EXECUTABLE} ${MACHETE_GEN_SCRIPT}
+          RESULT_VARIABLE machete_generation_result
+          OUTPUT_VARIABLE machete_generation_output
+          OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log
+          ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log
+        )
+
+        if (NOT machete_generation_result EQUAL 0)
+          message(FATAL_ERROR "Machete generation failed."
+                              " Result: \"${machete_generation_result}\""
+                              "\nCheck the log for details: "
+                              "${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log")
+        else()
+          set(MACHETE_GEN_SCRIPT_HASH ${MACHETE_GEN_SCRIPT_HASH}
+              CACHE STRING "Last run machete generate script hash" FORCE)
+          message(STATUS "Machete generation completed successfully.")
+        endif()
+      else()
+        message(STATUS "Machete generation script has not changed, skipping generation.")
+      endif()
+
+      # Add machete generated sources
+      file(GLOB MACHETE_GEN_SOURCES "csrc/libtorch_stable/quantization/machete/generated/*.cu")
+      list(APPEND VLLM_STABLE_EXT_SRC ${MACHETE_GEN_SOURCES})
+
+      # forward compatible
+      set_gencode_flags_for_srcs(
+        SRCS "${MACHETE_GEN_SOURCES}"
+        CUDA_ARCHS "${MACHETE_ARCHS}")
+
+      list(APPEND VLLM_STABLE_EXT_SRC
+        csrc/libtorch_stable/quantization/machete/machete_pytorch.cu)
+      message(STATUS "Building Machete kernels for archs: ${MACHETE_ARCHS}")
+    else()
+      if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0
+          AND MACHETE_ARCHS)
+        message(STATUS "Not building Machete kernels as CUDA Compiler version is "
+                       "not >= 12.0, we recommend upgrading to CUDA 12.0 or "
+                       "later if you intend on running w4a16 quantized models on "
+                       "Hopper.")
+      else()
+        message(STATUS "Not building Machete kernels as no compatible archs "
+                       "found in CUDA target architectures")
+      endif()
+    endif()
+
    set_gencode_flags_for_srcs(
      SRCS "${VLLM_STABLE_EXT_SRC}"
      CUDA_ARCHS "${CUDA_ARCHS}")

+    # Only build Marlin kernels if we are building for at least some compatible archs.
+    # Keep building Marlin for 9.0 as there are some group sizes and shapes that
+    # are not supported by Machete yet.
+
+    # marlin arches for fp16 output
+    # Family-conditional 12.0f (one cubin for SM12x family) requires CUDA >= 13.0;
+    # fall back to architecture-specific 12.0a;12.1a on CUDA < 13.0 (e.g. 12.8).
+    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+      cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0f" "${CUDA_ARCHS}")
+    else()
+      cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
+    endif()
+    # marlin has limited support for turing
+    cuda_archs_loose_intersection(MARLIN_SM75_ARCHS "7.5" "${CUDA_ARCHS}")
+    # marlin arches for bf16 output (we need 9.0 for bf16 atomicAdd PTX)
+    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+      cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0f" "${CUDA_ARCHS}")
+    else()
+      cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
+    endif()
+    # marlin arches for fp8 input
+    # - sm80 doesn't support fp8 computation
+    # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
+    # so we only enable fp8 computation for SM89 (e.g. RTX 40x0)  and 12.0 (e.g. RTX 50x0)
+    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+      cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0f" "${CUDA_ARCHS}")
+    else()
+      cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0a;12.1a" "${CUDA_ARCHS}")
+    endif()
+    # marlin arches for other files
+    cuda_archs_loose_intersection(MARLIN_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}")
+
+    if (MARLIN_OTHER_ARCHS)
+
+      #
+      # For the Marlin kernels we automatically generate sources for various
+      # preselected input type pairs and schedules.
+      # Generate sources:
+      set(MARLIN_GEN_SCRIPT
+        ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/quantization/marlin/generate_kernels.py)
+      file(MD5 ${MARLIN_GEN_SCRIPT} MARLIN_GEN_SCRIPT_HASH)
+      list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR)
+      set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})")
+
+      message(STATUS "Marlin generation script hash: ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
+      message(STATUS "Last run Marlin generate script hash: $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
+
+      if (NOT DEFINED CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
+          OR NOT $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH} STREQUAL ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH})
+        execute_process(
+          COMMAND ${CMAKE_COMMAND} -E env
+          PYTHONPATH=$ENV{PYTHONPATH}
+            ${Python_EXECUTABLE} ${MARLIN_GEN_SCRIPT} ${CUDA_ARCHS_STR}
+          RESULT_VARIABLE marlin_generation_result
+          OUTPUT_VARIABLE marlin_generation_result
+          OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log
+          ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log
+        )
+
+        if (NOT marlin_generation_result EQUAL 0)
+          message(FATAL_ERROR "Marlin generation failed."
+                              " Result: \"${marlin_generation_result}\""
+                              "\nCheck the log for details: "
+                              "${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log")
+        else()
+          set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
+              CACHE STRING "Last run Marlin generate script hash and arch" FORCE)
+          message(STATUS "Marlin generation completed successfully.")
+        endif()
+      else()
+        message(STATUS "Marlin generation script has not changed, skipping generation.")
+      endif()
+
+      if (MARLIN_ARCHS)
+        file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm80_kernel_*_float16.cu")
+        set_gencode_flags_for_srcs(
+          SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}"
+          CUDA_ARCHS "${MARLIN_ARCHS}")
+        if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
+          set_source_files_properties(${MARLIN_TEMPLATE_KERNEL_SRC}
+            PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
+        endif()
+        list(APPEND VLLM_STABLE_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC})
+
+        file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm80_kernel_*_bfloat16.cu")
+        set_gencode_flags_for_srcs(
+          SRCS "${MARLIN_TEMPLATE_BF16_KERNEL_SRC}"
+          CUDA_ARCHS "${MARLIN_BF16_ARCHS}")
+        if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
+          set_source_files_properties(${MARLIN_TEMPLATE_BF16_KERNEL_SRC}
+            PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
+        endif()
+        list(APPEND VLLM_STABLE_EXT_SRC ${MARLIN_TEMPLATE_BF16_KERNEL_SRC})
+      endif()
+
+      if (MARLIN_SM75_ARCHS)
+        file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm75_kernel_*.cu")
+        set_gencode_flags_for_srcs(
+          SRCS "${MARLIN_TEMPLATE_SM75_KERNEL_SRC}"
+          CUDA_ARCHS "${MARLIN_SM75_ARCHS}")
+        if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
+          set_source_files_properties(${MARLIN_TEMPLATE_SM75_KERNEL_SRC}
+            PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
+        endif()
+        list(APPEND VLLM_STABLE_EXT_SRC ${MARLIN_TEMPLATE_SM75_KERNEL_SRC})
+      endif()
+
+      if (MARLIN_FP8_ARCHS)
+        file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm89_kernel_*.cu")
+        set_gencode_flags_for_srcs(
+          SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}"
+          CUDA_ARCHS "${MARLIN_FP8_ARCHS}")
+        if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
+          set_source_files_properties(${MARLIN_TEMPLATE_FP8_KERNEL_SRC}
+            PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
+        endif()
+        list(APPEND VLLM_STABLE_EXT_SRC ${MARLIN_TEMPLATE_FP8_KERNEL_SRC})
+      endif()
+
+      set(MARLIN_SRCS
+         "csrc/libtorch_stable/quantization/marlin/marlin.cu"
+         "csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu"
+         "csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu"
+         "csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu")
+      set_gencode_flags_for_srcs(
+        SRCS "${MARLIN_SRCS}"
+        CUDA_ARCHS "${MARLIN_OTHER_ARCHS}")
+      if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
+        set_source_files_properties(${MARLIN_SRCS}
+          PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
+      endif()
+      list(APPEND VLLM_STABLE_EXT_SRC "${MARLIN_SRCS}")
+
+      message(STATUS "Building Marlin kernels for archs: ${MARLIN_OTHER_ARCHS}")
+    else()
+      message(STATUS "Not building Marlin kernels as no compatible archs found"
+                     " in CUDA target architectures")
+    endif()
+
  # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
@@ -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]
@@ -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
@@ -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
@@ -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] = {
@@ -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                                          \
@@ -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)               \
@@ -6,7 +6,7 @@

 #include <iostream>

-#include "quantization/marlin/marlin_dtypes.cuh"
+#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh"
 using marlin::MarlinScalarType2;

 namespace allspark {
@@ -39,10 +39,10 @@ namespace machete {
 {% for impl_config in impl_configs %}
 {% set type_sig = gen_type_sig(impl_config.types) -%}
 {% for s in impl_config.schedules %}
-extern torch::Tensor impl_{{type_sig}}_sch_{{gen_sch_sig(s)}}(MMArgs);
+extern torch::stable::Tensor impl_{{type_sig}}_sch_{{gen_sch_sig(s)}}(MMArgs);
 {%- endfor %}

-torch::Tensor mm_dispatch_{{type_sig}}(MMArgs args) {
+torch::stable::Tensor mm_dispatch_{{type_sig}}(MMArgs args) {
  [[maybe_unused]] auto M = args.A.size(0);
  [[maybe_unused]] auto N = args.B.size(1);
  [[maybe_unused]] auto K = args.A.size(1);
@@ -59,14 +59,14 @@ torch::Tensor mm_dispatch_{{type_sig}}(MMArgs args) {
  if (*args.maybe_schedule == "{{ gen_sch_sig(s) }}")
    return impl_{{type_sig}}_sch_{{ gen_sch_sig(s) }}(args);
  {%- endfor %}
-  TORCH_CHECK_NOT_IMPLEMENTED(false, "machete_gemm(..) is not implemented for "
+  STD_TORCH_CHECK_NOT_IMPLEMENTED(false, "machete_gemm(..) is not implemented for "
                                     "schedule = ", *args.maybe_schedule);
 }
 {%- endfor %}


-static inline std::optional<at::ScalarType> maybe_scalartype(
-    std::optional<at::Tensor> const& t) {
+static inline std::optional<torch::headeronly::ScalarType> maybe_scalartype(
+    std::optional<torch::stable::Tensor> const& t) {
    if (!t) {
      return std::nullopt;
    } else {
@@ -74,7 +74,7 @@ static inline std::optional<at::ScalarType> maybe_scalartype(
    };
 }

-torch::Tensor mm_dispatch(MMArgs args) {
+torch::stable::Tensor mm_dispatch(MMArgs args) {
  auto out_type = args.maybe_out_type.value_or(args.A.scalar_type());
  auto a_type = args.A.scalar_type();
  auto maybe_g_scales_type = maybe_scalartype(args.maybe_group_scales);
@@ -105,19 +105,19 @@ torch::Tensor mm_dispatch(MMArgs args) {
  }
  {%- endfor %}
  
-  TORCH_CHECK_NOT_IMPLEMENTED(
+  STD_TORCH_CHECK_NOT_IMPLEMENTED(
    false, "machete_mm(..) is not implemented for "
-    "a_type=", args.A.scalar_type(),
+    "a_type=", torch::headeronly::toString(args.A.scalar_type()),
    ", b_type=", args.b_type.str(),
-    ", out_type=", out_type,
+    ", out_type=", torch::headeronly::toString(out_type),
    ", with_group_scale_type=", maybe_g_scales_type
-        ? toString(*maybe_g_scales_type) : "None",
+        ? torch::headeronly::toString(*maybe_g_scales_type) : "None",
    ", with_group_zeropoint_type=", maybe_g_zeros_type
-        ? toString(*maybe_g_zeros_type) : "None",
+        ? torch::headeronly::toString(*maybe_g_zeros_type) : "None",
    ", with_channel_scale_type=", maybe_ch_scales_type
-        ? toString(*maybe_ch_scales_type) : "None",
+        ? torch::headeronly::toString(*maybe_ch_scales_type) : "None",
    ", with_token_scale_type=", maybe_tok_scales_type
-        ? toString(*maybe_tok_scales_type) : "None",
+        ? torch::headeronly::toString(*maybe_tok_scales_type) : "None",
    "; implemented types are: \\n",
    {%- for impl_config in impl_configs %}
    {% set t = impl_config.types -%}
@@ -197,7 +197,7 @@ using Kernel_{{type_sig}} = MacheteKernelTemplate<

 {% for sch in schs %}
 {% set sch_sig = gen_sch_sig(sch) -%}
-torch::Tensor 
+torch::stable::Tensor 
 impl_{{type_sig}}_sch_{{sch_sig}}(MMArgs args) {
  return run_impl<Kernel_{{type_sig}}<sch_{{sch_sig}}>>(args);
 }
@@ -212,7 +212,7 @@ PREPACK_TEMPLATE = """

 namespace machete {

-torch::Tensor prepack_B_dispatch(PrepackBArgs args) {
+torch::stable::Tensor prepack_B_dispatch(PrepackBArgs args) {
  auto convert_type = args.maybe_group_scales_type.value_or(args.a_type);
  {%- for t in types %}
  {% set b_type = unsigned_type_with_bitwidth(t.b_num_bits) %}
@@ -231,12 +231,12 @@ torch::Tensor prepack_B_dispatch(PrepackBArgs args) {
  }
  {%- endfor %}
  
-  TORCH_CHECK_NOT_IMPLEMENTED(false, 
+  STD_TORCH_CHECK_NOT_IMPLEMENTED(false, 
    "prepack_B_dispatch(..) is not implemented for "
-    "atype = ", args.a_type,
+    "atype = ", torch::headeronly::toString(args.a_type),
    ", b_type = ", args.b_type.str(),
    ", with_group_scales_type= ", args.maybe_group_scales_type ? 
-        toString(*args.maybe_group_scales_type) : "None");
+        torch::headeronly::toString(*args.maybe_group_scales_type) : "None");
 }

 }; // namespace machete
@@ -1,8 +1,6 @@
 #pragma once

-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <torch/all.h>
+#include <torch/headeronly/util/Exception.h>

 // clang-format off
 // The cutlass include order matters (annoyingly)
@@ -175,19 +173,23 @@ struct MacheteKernelTemplate {

  static Arguments create_arguments(
      cudaStream_t stream,
-      torch::Tensor const& A,  // MxK matrix
-      torch::Tensor const& B,  // KxN prepacked matrix
-      torch::Tensor& D,        // MxN matrix
-      std::optional<torch::Tensor> const& maybe_g_scales,  // scale_KxN matrix
-      std::optional<torch::Tensor> const& maybe_g_zeros,   // scale_KxN matrix
+      torch::stable::Tensor const& A,  // MxK matrix
+      torch::stable::Tensor const& B,  // KxN prepacked matrix
+      torch::stable::Tensor& D,        // MxN matrix
+      std::optional<torch::stable::Tensor> const&
+          maybe_g_scales,  // scale_KxN matrix
+      std::optional<torch::stable::Tensor> const&
+          maybe_g_zeros,  // scale_KxN matrix
      std::optional<int64_t> maybe_group_size,
-      std::optional<torch::Tensor> const& maybe_ch_scales,   // len N vector
-      std::optional<torch::Tensor> const& maybe_tok_scales)  // len M vector
+      std::optional<torch::stable::Tensor> const&
+          maybe_ch_scales,  // len N vector
+      std::optional<torch::stable::Tensor> const&
+          maybe_tok_scales)  // len M vector
  {
    static_assert(!with_group_zeropoints || with_group_scales);

    int M = A.size(0), N = B.size(1), K = A.size(1);
-    TORCH_CHECK(D.size(0) == M && D.size(1) == N);
+    STD_TORCH_CHECK(D.size(0) == M && D.size(1) == N);

    auto layout_A = make_cute_layout<StrideA>(A, "A");
    auto layout_D = make_cute_layout<StrideD>(D, "D");
@@ -216,29 +218,29 @@ struct MacheteKernelTemplate {
        maybe_group_size == -1 ? K : maybe_group_size.value_or(K);
    int const scale_k = (K + group_size - 1) / group_size;

-    TORCH_CHECK(size<0>(layout_A) == M && size<1>(layout_A) == K);
-    TORCH_CHECK(size<0>(layout_D) == M && size<1>(layout_D) == N);
+    STD_TORCH_CHECK(size<0>(layout_A) == M && size<1>(layout_A) == K);
+    STD_TORCH_CHECK(size<0>(layout_D) == M && size<1>(layout_D) == N);

    if constexpr (with_group_scales) {
-      TORCH_CHECK(S_group_ptr && layout_S_group);
-      TORCH_CHECK((size<0>(*layout_S_group) == scale_k &&
-                   size<1>(*layout_S_group) == N));
+      STD_TORCH_CHECK(S_group_ptr && layout_S_group);
+      STD_TORCH_CHECK((size<0>(*layout_S_group) == scale_k &&
+                       size<1>(*layout_S_group) == N));
    } else {
-      TORCH_CHECK(!S_group_ptr, "Scales not supported");
+      STD_TORCH_CHECK(!S_group_ptr, "Scales not supported");
    }

    if constexpr (with_group_zeropoints) {
-      TORCH_CHECK(Z_group_ptr && layout_Z_group);
-      TORCH_CHECK((size<0>(*layout_Z_group) == scale_k &&
-                   size<1>(*layout_Z_group) == N));
-      TORCH_CHECK(layout_S_group && *layout_Z_group == *layout_S_group,
-                  "Scales and zeros must have the same layout");
+      STD_TORCH_CHECK(Z_group_ptr && layout_Z_group);
+      STD_TORCH_CHECK((size<0>(*layout_Z_group) == scale_k &&
+                       size<1>(*layout_Z_group) == N));
+      STD_TORCH_CHECK(layout_S_group && *layout_Z_group == *layout_S_group,
+                      "Scales and zeros must have the same layout");
    } else {
-      TORCH_CHECK(!Z_group_ptr, "Zeropoints not supported");
+      STD_TORCH_CHECK(!Z_group_ptr, "Zeropoints not supported");
    }

    if constexpr (with_channel_scales || with_token_scales) {
-      TORCH_CHECK(
+      STD_TORCH_CHECK(
          (maybe_ch_scales->numel() == N || maybe_ch_scales->numel() == 1) &&
          (maybe_tok_scales->numel() == M || maybe_tok_scales->numel() == 1));
    }
@@ -298,11 +300,12 @@ struct MacheteKernelTemplate {
    Gemm gemm_op;

    cutlass::Status status = gemm_op.initialize(args, workspace, stream);
-    TORCH_CHECK(status == cutlass::Status::kSuccess,
-                "Machete kernel failed to initialize workspace");
+    STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
+                    "Machete kernel failed to initialize workspace");

    status = gemm_op.run(stream);
-    TORCH_CHECK(status == cutlass::Status::kSuccess, "Machete kernel failed");
+    STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
+                    "Machete kernel failed");
  }
 };

@@ -0,0 +1,80 @@
+#pragma once
+
+#include "machete_mm_kernel.cuh"
+#include "cutlass_extensions/torch_utils.hpp"
+#include "core/scalar_type.hpp"
+#include "libtorch_stable/torch_utils.h"
+
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+
+#include <optional>
+#include <string>
+#include <vector>
+
+namespace machete {
+
+struct MMArgs {
+  torch::stable::Tensor const& A;
+  torch::stable::Tensor const& B;
+  vllm::ScalarType const& b_type;
+  std::optional<torch::headeronly::ScalarType> const& maybe_out_type;
+  std::optional<torch::stable::Tensor> const& maybe_group_scales;
+  std::optional<torch::stable::Tensor> const& maybe_group_zeros;
+  std::optional<int64_t> maybe_group_size;
+  std::optional<torch::stable::Tensor> const& maybe_channel_scales;
+  std::optional<torch::stable::Tensor> const& maybe_token_scales;
+  std::optional<std::string> maybe_schedule;
+};
+
+struct SupportedSchedulesArgs {
+  torch::headeronly::ScalarType a_type;
+  vllm::ScalarType b_type;
+  std::optional<torch::headeronly::ScalarType> maybe_group_scales_type;
+  std::optional<torch::headeronly::ScalarType> maybe_group_zeros_type;
+  std::optional<torch::headeronly::ScalarType> maybe_channel_scales_type;
+  std::optional<torch::headeronly::ScalarType> maybe_token_scales_type;
+  std::optional<torch::headeronly::ScalarType> maybe_out_type;
+};
+
+torch::stable::Tensor mm_dispatch(MMArgs args);
+
+std::vector<std::string> supported_schedules_dispatch(
+    SupportedSchedulesArgs args);
+
+template <typename MacheteKernel>
+torch::stable::Tensor run_impl(MMArgs args) {
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      args.A.get_device_index());
+
+  auto device = args.A.device();
+  auto stream = get_current_cuda_stream(device.index());
+
+  int M = args.A.size(0);
+  int N = args.B.size(1);
+  int K = args.A.size(1);
+
+  // Allocate output
+  torch::stable::Tensor D = torch::stable::empty(
+      {M, N}, equivalent_scalar_type_v<typename MacheteKernel::ElementD>,
+      std::nullopt, device);
+
+  auto arguments = MacheteKernel::create_arguments(
+      stream,  //
+      args.A, args.B, D, args.maybe_group_scales, args.maybe_group_zeros,
+      args.maybe_group_size, args.maybe_channel_scales,
+      args.maybe_token_scales);
+  STD_TORCH_CHECK(MacheteKernel::can_implement(arguments),
+                  "Machete kernel cannot be run with these arguments");
+
+  size_t workspace_size = MacheteKernel::get_workspace_size(arguments);
+  torch::stable::Tensor workspace =
+      torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
+                           std::nullopt, device);
+
+  MacheteKernel::run(arguments, workspace.mutable_data_ptr(), stream);
+
+  return D;
+};
+
+};  // namespace machete
@@ -3,6 +3,7 @@
 #include "machete_mm_kernel.cuh"
 #include "cutlass_extensions/cute_utils.cuh"
 #include "cutlass_extensions/torch_utils.hpp"
+#include <torch/headeronly/util/Exception.h>

 namespace machete {

@@ -60,8 +61,8 @@ static void prepack_B_template(
  auto ilvd_NKbNbKL_to_offset =
      PrepackedLayoutB::ilvd_NKbNbKL_to_offset(shape(B_layout));

-  TORCH_CHECK(size<0>(B_layout) % size<0>(TileShapeNKL{}) == 0);
-  TORCH_CHECK(size<1>(B_layout) % size<1>(TileShapeNKL{}) == 0);
+  STD_TORCH_CHECK(size<0>(B_layout) % size<0>(TileShapeNKL{}) == 0);
+  STD_TORCH_CHECK(size<1>(B_layout) % size<1>(TileShapeNKL{}) == 0);

  auto N_tiles = size<0>(B_layout) / size<0>(TileShapeNKL{});
  auto K_tiles = size<1>(B_layout) / size<1>(TileShapeNKL{});
@@ -3,39 +3,47 @@
 #include "machete_prepack_kernel.cuh"
 #include "cutlass_extensions/torch_utils.hpp"
 #include "core/scalar_type.hpp"
+#include "libtorch_stable/torch_utils.h"
+
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+
+#include <optional>

 namespace machete {

 struct PrepackBArgs {
-  torch::Tensor const& B;
-  at::ScalarType a_type;
+  torch::stable::Tensor const& B;
+  torch::headeronly::ScalarType a_type;
  vllm::ScalarType b_type;
-  std::optional<at::ScalarType> maybe_group_scales_type;
+  std::optional<torch::headeronly::ScalarType> maybe_group_scales_type;
 };

 template <typename PrepackedLayoutB>
-torch::Tensor prepack_impl(torch::Tensor const B) {
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(B));
+torch::stable::Tensor prepack_impl(torch::stable::Tensor const& B) {
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      B.get_device_index());
  using ElementB = typename PrepackedLayoutB::ElementB;
  using PPBlockShape_NK = typename PrepackedLayoutB::PPBlockShape_NK;

  auto device = B.device();
-  auto stream = at::cuda::getCurrentCUDAStream(device.index());
+  auto stream = get_current_cuda_stream(device.index());
  auto B_ptr = static_cast<ElementB const*>(B.const_data_ptr());
  // elements per storage item for B
  auto eles_per_storage =
-      (B.dtype().itemsize() * 8) / cute::sizeof_bits_v<ElementB>;
+      (B.element_size() * 8) / cute::sizeof_bits_v<ElementB>;

  // torch B passed in is/should be (packed_K,N), the kernel expects (N,K,L) (to
  // match cutlass using (N,K,L) for B), so we transpose B to (N,packed_K,L)
-  auto Bt_packed = B.t();
+  auto Bt_packed = torch::stable::transpose(B, 0, 1);

-  TORCH_CHECK(
+  STD_TORCH_CHECK(
      (B.size(0) * eles_per_storage) % size<1>(PPBlockShape_NK{}) == 0,
      "B.shape[0] (in terms of unpacked elements) must be a multiple of ",
      size<1>(PPBlockShape_NK{}));
-  TORCH_CHECK(B.size(1) % size<0>(PPBlockShape_NK{}) == 0,
-              "B.shape[1] must be a multiple of ", size<0>(PPBlockShape_NK{}));
+  STD_TORCH_CHECK(B.size(1) % size<0>(PPBlockShape_NK{}) == 0,
+                  "B.shape[1] must be a multiple of ",
+                  size<0>(PPBlockShape_NK{}));

  using StrideB = cutlass::detail::TagToStrideB_t<cutlass::layout::ColumnMajor>;
  auto const l_Bt_packed = make_cute_layout<StrideB>(Bt_packed, "B");
@@ -49,7 +57,7 @@ torch::Tensor prepack_impl(torch::Tensor const B) {
  //   new_shape = (N, packed_K, L) * (1, eles_per_storage, 1) -> (N, K, L)
  //   new_stride = (s0, s1, s2) * (eles_per_storage, 1, eles_per_storage)
  //                 when s1 == 1
-  TORCH_CHECK(stride<1>(l_Bt_packed) == 1);
+  STD_TORCH_CHECK(stride<1>(l_Bt_packed) == 1);
  // clang-format off
  auto const layout_Bt = make_layout(
      transform_with_idx(l_Bt_packed.shape(), [&](auto ele, auto idx) {
@@ -61,7 +69,9 @@ torch::Tensor prepack_impl(torch::Tensor const B) {
  // clang-format on

  // Allocate output
-  torch::Tensor D = torch::empty_like(B, {}, at::MemoryFormat::Contiguous);
+  torch::stable::Tensor D = torch::stable::empty(
+      B.sizes(), B.scalar_type(), std::nullopt, B.device(), std::nullopt,
+      torch::headeronly::MemoryFormat::Contiguous);

  prepack_B_template<PrepackedLayoutB>(
      stream, B_ptr, layout_Bt, static_cast<ElementB*>(D.mutable_data_ptr()));
@@ -69,6 +79,6 @@ torch::Tensor prepack_impl(torch::Tensor const B) {
  return D;
 };

-torch::Tensor prepack_B_dispatch(PrepackBArgs args);
+torch::stable::Tensor prepack_B_dispatch(PrepackBArgs args);

 };  // namespace machete
@@ -1,9 +1,5 @@
 #pragma once

-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <torch/all.h>
-
 // clang-format off
 // The cutlass include order matters (annoyingly)

@@ -0,0 +1,77 @@
+#include "machete_mm_launcher.cuh"
+#include "machete_prepack_launcher.cuh"
+#include "core/scalar_type.hpp"
+
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+
+namespace machete {
+
+using namespace vllm;
+
+std::vector<std::string> supported_schedules(
+    torch::headeronly::ScalarType a_type, int64_t b_type_id,
+    std::optional<torch::headeronly::ScalarType> maybe_group_scales_type,
+    std::optional<torch::headeronly::ScalarType> maybe_group_zeros_type,
+    std::optional<torch::headeronly::ScalarType> maybe_channel_scales_type,
+    std::optional<torch::headeronly::ScalarType> maybe_token_scales_type,
+    std::optional<torch::headeronly::ScalarType> maybe_out_type) {
+  ScalarType const b_type = ScalarType::from_id(b_type_id);
+  return supported_schedules_dispatch({
+      .a_type = a_type,
+      .b_type = b_type,
+      .maybe_group_scales_type = maybe_group_scales_type,
+      .maybe_group_zeros_type = maybe_group_zeros_type,
+      .maybe_channel_scales_type = maybe_channel_scales_type,
+      .maybe_token_scales_type = maybe_token_scales_type,
+      .maybe_out_type = maybe_out_type,
+  });
+}
+
+torch::stable::Tensor mm(
+    torch::stable::Tensor const& A, torch::stable::Tensor const& B,
+    int64_t b_type_id,
+    std::optional<torch::headeronly::ScalarType> const& maybe_out_type,
+    std::optional<torch::stable::Tensor> const& maybe_group_scales,
+    std::optional<torch::stable::Tensor> const& maybe_group_zeros,
+    std::optional<int64_t> maybe_group_size,
+    std::optional<torch::stable::Tensor> const& maybe_channel_scales,
+    std::optional<torch::stable::Tensor> const& maybe_token_scales,
+    std::optional<std::string> maybe_schedule) {
+  ScalarType const b_type = ScalarType::from_id(b_type_id);
+  return mm_dispatch({.A = A,
+                      .B = B,
+                      .b_type = b_type,
+                      .maybe_out_type = maybe_out_type,
+                      .maybe_group_scales = maybe_group_scales,
+                      .maybe_group_zeros = maybe_group_zeros,
+                      .maybe_group_size = maybe_group_size,
+                      .maybe_channel_scales = maybe_channel_scales,
+                      .maybe_token_scales = maybe_token_scales,
+                      .maybe_schedule = maybe_schedule});
+}
+
+torch::stable::Tensor prepack_B(
+    torch::stable::Tensor const& B, torch::headeronly::ScalarType const& a_type,
+    int64_t b_type_id,
+    std::optional<torch::headeronly::ScalarType> const&
+        maybe_group_scales_type) {
+  ScalarType const b_type = ScalarType::from_id(b_type_id);
+  return prepack_B_dispatch(
+      {.B = B,
+       .a_type = a_type,
+       .b_type = b_type,
+       .maybe_group_scales_type = maybe_group_scales_type});
+}
+
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
+  m.impl("machete_prepack_B", TORCH_BOX(&prepack_B));
+  m.impl("machete_mm", TORCH_BOX(&mm));
+}
+
+STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, m) {
+  m.impl("machete_supported_schedules", TORCH_BOX(&supported_schedules));
+}
+
+};  // namespace machete
@@ -1,6 +1,13 @@
 #include "marlin.cuh"

-#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"

 namespace marlin {

@@ -218,56 +225,55 @@ __global__ void awq_marlin_repack_kernel(
            b_q_weight_ptr, out_ptr, size_k, size_n);                      \
  }

-torch::Tensor awq_marlin_repack(torch::Tensor& b_q_weight, int64_t size_k,
-                                int64_t size_n, int64_t num_bits,
-                                bool is_a_8bit) {
+torch::stable::Tensor awq_marlin_repack(torch::stable::Tensor& b_q_weight,
+                                        int64_t size_k, int64_t size_n,
+                                        int64_t num_bits, bool is_a_8bit) {
  // Verify compatibility with marlin tile of 16x64
-  TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k,
-              " is not divisible by tile_k_size = ", marlin::tile_k_size);
-  TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n,
-              " is not divisible by tile_n_size = ", marlin::tile_n_size);
+  STD_TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k,
+                  " is not divisible by tile_k_size = ", marlin::tile_k_size);
+  STD_TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n,
+                  " is not divisible by tile_n_size = ", marlin::tile_n_size);

-  TORCH_CHECK(num_bits == 4 || num_bits == 8,
-              "num_bits must be 4 or 8. Got = ", num_bits);
+  STD_TORCH_CHECK(num_bits == 4 || num_bits == 8,
+                  "num_bits must be 4 or 8. Got = ", num_bits);
  int const pack_factor = 32 / num_bits;

  // Verify B
-  TORCH_CHECK(b_q_weight.size(0) == size_k,
-              "b_q_weight.size(0) = ", b_q_weight.size(0),
-              " is not size_k = ", size_k);
-  TORCH_CHECK((size_n / pack_factor) == b_q_weight.size(1),
-              "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1),
-              ", size_n = ", size_n, ", pack_factor = ", pack_factor);
+  STD_TORCH_CHECK(b_q_weight.size(0) == size_k,
+                  "b_q_weight.size(0) = ", b_q_weight.size(0),
+                  " is not size_k = ", size_k);
+  STD_TORCH_CHECK((size_n / pack_factor) == b_q_weight.size(1),
+                  "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1),
+                  ", size_n = ", size_n, ", pack_factor = ", pack_factor);

  // Verify device and strides
-  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");
-  TORCH_CHECK(b_q_weight.dtype() == at::kInt, "b_q_weight type is not kInt");
+  STD_TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight is not on GPU");
+  STD_TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
+  STD_TORCH_CHECK(
+      b_q_weight.scalar_type() == torch::headeronly::ScalarType::Int,
+      "b_q_weight type is not kInt");
+
+  const int32_t device_index = b_q_weight.get_device_index();
+  torch::stable::accelerator::DeviceGuard device_guard(device_index);
+  const cudaStream_t stream = get_current_cuda_stream(device_index);

  // Alloc buffers
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(b_q_weight));
-  auto options = torch::TensorOptions()
-                     .dtype(b_q_weight.dtype())
-                     .device(b_q_weight.device());
-  torch::Tensor out = torch::empty(
+  torch::stable::Tensor out = torch::stable::empty(
      {size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor},
-      options);
+      b_q_weight.scalar_type(), std::nullopt, b_q_weight.device());

  // Get ptrs
  uint32_t const* b_q_weight_ptr =
-      reinterpret_cast<uint32_t const*>(b_q_weight.data_ptr());
-  uint32_t* out_ptr = reinterpret_cast<uint32_t*>(out.data_ptr());
+      reinterpret_cast<uint32_t const*>(b_q_weight.const_data_ptr());
+  uint32_t* out_ptr = reinterpret_cast<uint32_t*>(out.mutable_data_ptr());

-  // Get dev info
-  int dev = b_q_weight.get_device();
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream(dev);
  int blocks;
-  cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev);
+  cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device_index);

  int max_shared_mem = 0;
  cudaDeviceGetAttribute(&max_shared_mem,
-                         cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
-  TORCH_CHECK(max_shared_mem > 0);
+                         cudaDevAttrMaxSharedMemoryPerBlockOptin, device_index);
+  STD_TORCH_CHECK(max_shared_mem > 0);

  if (false) {
  }
@@ -276,13 +282,13 @@ torch::Tensor awq_marlin_repack(torch::Tensor& b_q_weight, int64_t size_k,
  CALL_IF(4, true)
  CALL_IF(8, true)
  else {
-    TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits,
-                ", is_a_8bit = ", is_a_8bit);
+    STD_TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits,
+                    ", is_a_8bit = ", is_a_8bit);
  }

  return out;
 }

-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-  m.impl("awq_marlin_repack", &awq_marlin_repack);
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
+  m.impl("awq_marlin_repack", TORCH_BOX(&awq_marlin_repack));
 }
@@ -303,7 +303,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.");'
        )

@@ -1,6 +1,13 @@
 #include "marlin.cuh"

-#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"

 namespace marlin {

@@ -275,64 +282,66 @@ __global__ void gptq_marlin_repack_kernel(
            b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n);             \
  }

-torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm,
-                                 int64_t size_k, int64_t size_n,
-                                 int64_t num_bits, bool is_a_8bit) {
+torch::stable::Tensor gptq_marlin_repack(torch::stable::Tensor& b_q_weight,
+                                         torch::stable::Tensor& perm,
+                                         int64_t size_k, int64_t size_n,
+                                         int64_t num_bits, bool is_a_8bit) {
  // Verify compatibility with marlin tile of 16x64
-  TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k,
-              " is not divisible by tile_k_size = ", marlin::tile_k_size);
-  TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n,
-              " is not divisible by tile_n_size = ", marlin::tile_n_size);
+  STD_TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k,
+                  " is not divisible by tile_k_size = ", marlin::tile_k_size);
+  STD_TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n,
+                  " is not divisible by tile_n_size = ", marlin::tile_n_size);

-  TORCH_CHECK(num_bits == 4 || num_bits == 8,
-              "num_bits must be 4 or 8. Got = ", num_bits);
+  STD_TORCH_CHECK(num_bits == 4 || num_bits == 8,
+                  "num_bits must be 4 or 8. Got = ", num_bits);
  int const pack_factor = 32 / num_bits;

  // Verify B
-  TORCH_CHECK((size_k / pack_factor) == b_q_weight.size(0),
-              "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0),
-              ", size_k = ", size_k, ", pack_factor = ", pack_factor);
-  TORCH_CHECK(b_q_weight.size(1) == size_n,
-              "b_q_weight.size(1) = ", b_q_weight.size(1),
-              " is not size_n = ", size_n);
+  STD_TORCH_CHECK((size_k / pack_factor) == b_q_weight.size(0),
+                  "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0),
+                  ", size_k = ", size_k, ", pack_factor = ", pack_factor);
+  STD_TORCH_CHECK(b_q_weight.size(1) == size_n,
+                  "b_q_weight.size(1) = ", b_q_weight.size(1),
+                  " is not size_n = ", size_n);

  // Verify device and strides
-  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");
-  TORCH_CHECK(b_q_weight.dtype() == at::kInt, "b_q_weight type is not kInt");
+  STD_TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight is not on GPU");
+  STD_TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
+  STD_TORCH_CHECK(
+      b_q_weight.scalar_type() == torch::headeronly::ScalarType::Int,
+      "b_q_weight type is not kInt");

-  TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU");
-  TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
-  TORCH_CHECK(perm.dtype() == at::kInt, "perm type is not at::kInt");
+  STD_TORCH_CHECK(perm.is_cuda(), "perm is not on GPU");
+  STD_TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
+  STD_TORCH_CHECK(perm.scalar_type() == torch::headeronly::ScalarType::Int,
+                  "perm type is not at::kInt");
+
+  const int32_t device_index = b_q_weight.get_device_index();
+  torch::stable::accelerator::DeviceGuard device_guard(device_index);
+  const cudaStream_t stream = get_current_cuda_stream(device_index);

  // Alloc buffers
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(b_q_weight));
-  auto options = torch::TensorOptions()
-                     .dtype(b_q_weight.dtype())
-                     .device(b_q_weight.device());
-  torch::Tensor out = torch::empty(
+  torch::stable::Tensor out = torch::stable::empty(
      {size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor},
-      options);
+      b_q_weight.scalar_type(), std::nullopt, b_q_weight.device());

  // Detect if there is act_order
  bool has_perm = perm.size(0) != 0;

  // Get ptrs
  uint32_t const* b_q_weight_ptr =
-      reinterpret_cast<uint32_t const*>(b_q_weight.data_ptr());
-  uint32_t const* perm_ptr = reinterpret_cast<uint32_t const*>(perm.data_ptr());
-  uint32_t* out_ptr = reinterpret_cast<uint32_t*>(out.data_ptr());
+      reinterpret_cast<uint32_t const*>(b_q_weight.const_data_ptr());
+  uint32_t const* perm_ptr =
+      reinterpret_cast<uint32_t const*>(perm.const_data_ptr());
+  uint32_t* out_ptr = reinterpret_cast<uint32_t*>(out.mutable_data_ptr());

-  // Get dev info
-  int dev = b_q_weight.get_device();
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream(dev);
  int blocks;
-  cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev);
+  cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device_index);

  int max_shared_mem = 0;
  cudaDeviceGetAttribute(&max_shared_mem,
-                         cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
-  TORCH_CHECK(max_shared_mem > 0);
+                         cudaDevAttrMaxSharedMemoryPerBlockOptin, device_index);
+  STD_TORCH_CHECK(max_shared_mem > 0);

  if (false) {
  }
@@ -345,13 +354,13 @@ torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm,
  CALL_IF(8, false, true)

  else {
-    TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits,
-                ", has_perm = ", has_perm, ", is_a_8bit = ", is_a_8bit);
+    STD_TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits,
+                    ", has_perm = ", has_perm, ", is_a_8bit = ", is_a_8bit);
  }

  return out;
 }

-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-  m.impl("gptq_marlin_repack", &gptq_marlin_repack);
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
+  m.impl("gptq_marlin_repack", TORCH_BOX(&gptq_marlin_repack));
 }
@@ -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 ||          \
@@ -46,19 +54,22 @@ __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,

 }  // namespace marlin

-torch::Tensor 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& b_zeros_or_none,
-    std::optional<torch::Tensor> const& g_idx_or_none,
-    std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
-    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) {
-  TORCH_CHECK_NOT_IMPLEMENTED(false,
-                              "marlin_gemm(..) requires CUDA_ARCH >= 7.5");
-  return torch::empty({1, 1});
+torch::stable::Tensor 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, 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) {
+  STD_TORCH_CHECK_NOT_IMPLEMENTED(false,
+                                  "marlin_gemm(..) requires CUDA_ARCH >= 7.5");
+  return torch::stable::empty({1, 1});
 }

 #else
@@ -323,18 +334,18 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
               int thread_n_init, int sms, bool use_atomic_add,
               bool use_fp32_reduce, bool is_zp_float) {
  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 {
@@ -342,8 +353,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);
    }
  }

@@ -384,25 +395,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 "
@@ -432,10 +443,10 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
    if (thread_k != -1 && thread_n != -1) {
      thread_tfg = thread_config_t{thread_k, thread_n, default_threads};
      exec_cfg = exec_config_t{1, 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(
@@ -474,7 +485,7 @@ 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(
+    STD_TORCH_CHECK(
        is_valid_config(thread_tfg, thread_m_blocks, prob_m_split, prob_n,
                        prob_k, num_bits, group_size, has_act_order, is_k_full,
                        has_zp, is_zp_float, is_a_8bit, stages,
@@ -495,14 +506,15 @@ 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,
-                  ", prob_m_split = ", prob_m_split,
-                  ", thread_m_blocks = ", thread_m_blocks,
-                  ", thread_n_blocks = ", thread_n_blocks,
-                  ", thread_k_blocks = ", thread_k_blocks,
-                  ", num_threads = ", num_threads, ", 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,
+          ", prob_m_split = ", prob_m_split,
+          ", thread_m_blocks = ", thread_m_blocks,
+          ", thread_n_blocks = ", thread_n_blocks,
+          ", thread_k_blocks = ", thread_k_blocks,
+          ", num_threads = ", num_threads, ", num_bits = ", num_bits);
    }

    cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
@@ -530,71 +542,76 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,

 }  // namespace marlin

-torch::Tensor 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,
-    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) {
+torch::stable::Tensor 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, 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) {
  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_scalar_type = 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_scalar_type = 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_scalar_type = 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();
  }

@@ -606,54 +623,58 @@ torch::Tensor marlin_gemm(
  int pack_factor = 32 / b_type.size_bits();

  // 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(0),
-              "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0),
-              ", 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(0),
+      "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0),
+      ", size_k = ", size_k,
+      ", tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
+  STD_TORCH_CHECK(
      b_q_weight.size(1) % MARLIN_NAMESPACE_NAME::tile_size == 0,
      "b_q_weight.size(1) = ", b_q_weight.size(1),
      " is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
  int actual_size_n =
      (b_q_weight.size(1) / 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.stride(1) == 1, "A.stride(1) is not 1");
+  STD_TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
+  STD_TORCH_CHECK(a.stride(1) == 1, "A.stride(1) is not 1");
  // We use int4 (16 bytes) to load A, so A must aligned to 16 bytes
-  TORCH_CHECK(a.stride(0) % 8 == 0, "A.stride(0) must divisible by 8");
-  TORCH_CHECK(((uint64_t)a.data_ptr()) % 16 == 0, "A must aligned to 16 bytes");
+  STD_TORCH_CHECK(a.stride(0) % 8 == 0, "A.stride(0) must divisible by 8");
+  STD_TORCH_CHECK(((uint64_t)a.const_data_ptr()) % 16 == 0,
+                  "A must aligned to 16 bytes");

-  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;
+  const auto device = a.device();

  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::empty({0}, torch::headeronly::ScalarType::Float,
+                                    std::nullopt, device);
+    STD_TORCH_CHECK(
+        a_type.size_bits() != 8,
+        "the a_scales parameter must be passed for 8bit activation.");
  }

  // thread_k: `k` size of a thread_tile in `weights` (can usually be left as
@@ -664,84 +685,93 @@ torch::Tensor marlin_gemm(
  int thread_n = -1;
  // sms: number of SMs to use for the kernel
  int sms = -1;
-  cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, a.get_device());
+  const int32_t device_index = a.get_device_index();
+  cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, device_index);

  // Alloc buffers
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
-  torch::Tensor c;
+  torch::stable::accelerator::DeviceGuard device_guard(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, "Shape mismatch: c.size(0) = ", c.size(0),
-                ", size_m = ", size_m);
-    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,
+                    "Shape mismatch: c.size(0) = ", c.size(0),
+                    ", size_m = ", size_m);
+    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, size_n}, options);
+    c = torch::stable::empty({size_m, size_n}, c_scalar_type, std::nullopt,
+                             device);
  }
  if (size_m == 0) return c;

  // 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) {
    int max_m_block_size = (size_m + 16 - 1) / 16 * 16;
    max_m_block_size = min(max_m_block_size, 64);
    int max_c_tmp_size =
        sms * max_m_block_size * MARLIN_NAMESPACE_NAME::max_thread_n;
-    c_tmp = torch::empty({max_c_tmp_size}, options_fp32);
+    c_tmp = torch::stable::empty({max_c_tmp_size},
+                                 torch::headeronly::ScalarType::Float,
+                                 std::nullopt, device);
  } else {
-    c_tmp = torch::empty({0}, options_fp32);
+    c_tmp = torch::stable::empty({0}, torch::headeronly::ScalarType::Float,
+                                 std::nullopt, device);
  }

  // Detect groupsize and act_order
  int num_groups = -1;
  int group_size = -1;

-  int rank = b_scales.sizes().size();
-  TORCH_CHECK(rank == 2, "b_scales rank = ", rank, " is not 2");
-  TORCH_CHECK(b_scales.size(1) == size_n, "b_scales dim 1 = ", b_scales.size(1),
-              " is not size_n = ", size_n);
+  int rank = b_scales.dim();
+  STD_TORCH_CHECK(rank == 2, "b_scales rank = ", rank, " is not 2");
+  STD_TORCH_CHECK(b_scales.size(1) == size_n,
+                  "b_scales dim 1 = ", b_scales.size(1),
+                  " is not size_n = ", size_n);
  num_groups = b_scales.size(0);

-  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::empty({0}, c_scalar_type, std::nullopt, device);
+    perm = torch::stable::empty({0}, c_scalar_type, std::nullopt, device);
+    a_tmp = torch::stable::empty({0}, c_scalar_type, std::nullopt, device);
  }
  bool has_act_order = g_idx.size(-1) > 0 && perm.size(-1) > 0;

  if (has_act_order) {
-    a_tmp = torch::empty({size_m, size_k}, options);
+    a_tmp = torch::stable::empty({size_m, size_k}, c_scalar_type, std::nullopt,
+                                 device);
    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::empty({0}, c_scalar_type, std::nullopt, device);
    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(0) = ", b_scales.size(0));
      group_size = size_k / num_groups;
@@ -750,109 +780,114 @@ torch::Tensor 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::empty(
+        {0}, torch::headeronly::ScalarType::Float, std::nullopt, device);
+    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(0) == size_n, "b_bias.size(0) != size_n");
-    TORCH_CHECK(b_bias.stride(0) == 1, "b_bias.stride(0) != 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(0) == size_n, "b_bias.size(0) != size_n");
+    STD_TORCH_CHECK(b_bias.stride(0) == 1, "b_bias.stride(0) != 1");
  } else {
-    b_bias = torch::empty({0}, options);
+    b_bias = torch::stable::empty({0}, c_scalar_type, std::nullopt, device);
  }

-  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::empty({0}, c_scalar_type, std::nullopt, device);
  }
  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 == 2, "b_zeros rank = ", rank, " is not 2");
+    int rank = b_zeros.dim();
+    STD_TORCH_CHECK(rank == 2, "b_zeros rank = ", rank, " is not 2");
    if (is_zp_float) {
-      TORCH_CHECK(b_zeros.size(1) == size_n,
-                  "b_zeros dim 1 = ", b_zeros.size(1),
-                  " is not size_n = ", size_n);
-      TORCH_CHECK(num_groups == b_zeros.size(0),
-                  "b_zeros dim 0 = ", b_zeros.size(0),
-                  " is not num_groups = ", num_groups);
-      TORCH_CHECK(num_groups != -1, "num_groups must be != -1");
+      STD_TORCH_CHECK(b_zeros.size(1) == size_n,
+                      "b_zeros dim 1 = ", b_zeros.size(1),
+                      " is not size_n = ", size_n);
+      STD_TORCH_CHECK(num_groups == b_zeros.size(0),
+                      "b_zeros dim 0 = ", b_zeros.size(0),
+                      " is not num_groups = ", num_groups);
+      STD_TORCH_CHECK(num_groups != -1, "num_groups must be != -1");
    } else {
-      TORCH_CHECK(b_zeros.size(0) == num_groups,
-                  "b_zeros dim 0 = ", b_zeros.size(0),
-                  " is not num_groups = ", num_groups);
-      TORCH_CHECK(b_zeros.size(1) == size_n / pack_factor,
-                  "b_zeros dim 1 = ", b_zeros.size(1),
-                  " is not size_n / pack_factor = ", size_n / pack_factor);
+      STD_TORCH_CHECK(b_zeros.size(0) == num_groups,
+                      "b_zeros dim 0 = ", b_zeros.size(0),
+                      " is not num_groups = ", num_groups);
+      STD_TORCH_CHECK(b_zeros.size(1) == size_n / pack_factor,
+                      "b_zeros dim 1 = ", b_zeros.size(1),
+                      " 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 min_workspace_size = sms;
-  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::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(), size_m, size_n, size_k, a.stride(0),
-      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), thread_k, thread_n, sms,
+      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(), size_m, size_n, size_k, a.stride(0),
+      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, device_index,
+      get_current_cuda_stream(device_index), thread_k, thread_n, sms,
      use_atomic_add, use_fp32_reduce, is_zp_float);

  return c;
@@ -860,6 +895,6 @@ torch::Tensor marlin_gemm(

 #endif

-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-  m.impl("marlin_gemm", &marlin_gemm);
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
+  m.impl("marlin_gemm", TORCH_BOX(&marlin_gemm));
 }
@@ -2,14 +2,6 @@

 #ifndef _marlin_cuh
  #define _marlin_cuh
-  // These torch headers are only needed by non-stable callers (e.g. ops.cu).
-  // Guard them so that stable ABI targets can still include marlin.cuh
-  // for Vec, constants, and cp_async helpers without pulling in torch/all.h.
-  #ifndef TORCH_TARGET_VERSION
-    #include <torch/all.h>
-    #include <ATen/cuda/CUDAContext.h>
-    #include <c10/cuda/CUDAGuard.h>
-  #endif
  #include <cuda.h>
  #include <cuda_fp16.h>
  #include <cuda_runtime.h>
@@ -0,0 +1,118 @@
+
+#include "marlin.cuh"
+
+#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"
+
+// for only non-zp format (like gptq)
+__global__ void marlin_int4_fp8_preprocess_kernel_without_zp(
+    // qweight: (size_k * size_n // 8,)
+    const int32_t* __restrict__ qweight,
+    // output: same shape with qweight
+    int32_t* __restrict__ output) {
+  int32_t val = qweight[blockIdx.x * 32 + threadIdx.x];
+  int32_t new_val = 0;
+
+#pragma unroll
+  for (int32_t i = 0; i < 8; i++) {
+    int32_t single_val = val & 0xF;
+    single_val = single_val >= 8 ? single_val - 8 : 15 - single_val;
+    new_val |= single_val << (i * 4);
+    val >>= 4;
+  }
+
+  output[blockIdx.x * 32 + threadIdx.x] = new_val;
+}
+
+// for awq format only (with zp and with awq weight layout)
+__global__ void marlin_int4_fp8_preprocess_kernel_awq(
+    // AWQ qweight: (size_k, size_n // 8)
+    const int32_t* __restrict__ qweight,
+    // output: same shape with qweight
+    int32_t* __restrict__ output,
+    // AWQ zeros: (size_k // group_size, size_n // 8)
+    const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k,
+    int32_t group_size) {
+  int32_t val =
+      qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y];
+  int32_t zero =
+      qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 +
+             blockIdx.y];
+  int32_t new_val = 0;
+
+#pragma unroll
+  for (int32_t i = 0; i < 8; i++) {
+    int32_t single_val = val & 0xF;
+    int32_t single_zero = zero & 0xF;
+
+    single_val =
+        single_val >= single_zero ? single_val - single_zero : 15 - single_val;
+    new_val |= single_val << (i * 4);
+    val >>= 4;
+    zero >>= 4;
+  }
+
+  output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val;
+}
+
+torch::stable::Tensor marlin_int4_fp8_preprocess(
+    torch::stable::Tensor& qweight,
+    std::optional<torch::stable::Tensor> qzeros_or_none, bool inplace) {
+  STD_TORCH_CHECK(qweight.is_cuda(), "qweight is not on GPU");
+  STD_TORCH_CHECK(qweight.scalar_type() == torch::headeronly::ScalarType::Int,
+                  "qweight.dtype != torch.int32");
+
+  const int32_t device_index = qweight.get_device_index();
+  torch::stable::accelerator::DeviceGuard device_guard(device_index);
+  const cudaStream_t stream = get_current_cuda_stream(device_index);
+
+  torch::stable::Tensor output =
+      inplace ? qweight : torch::stable::empty_like(qweight);
+
+  if (!qzeros_or_none.has_value()) {
+    STD_TORCH_CHECK(qweight.numel() * 8 % 256 == 0,
+                    "qweight.numel() * 8 % 256 != 0");
+
+    int blocks = qweight.numel() * 8 / 256;
+    marlin_int4_fp8_preprocess_kernel_without_zp<<<blocks, 32, 0, stream>>>(
+        reinterpret_cast<const int32_t*>(qweight.const_data_ptr()),
+        reinterpret_cast<int32_t*>(output.mutable_data_ptr()));
+  } else {
+    int32_t size_k = qweight.size(0);
+    int32_t size_n = qweight.size(1) * 8;
+    torch::stable::Tensor qzeros = qzeros_or_none.value();
+
+    STD_TORCH_CHECK(size_k % 32 == 0, "size_k % 32 != 0");
+    STD_TORCH_CHECK(qzeros.is_cuda(), "qzeros is not on GPU");
+    STD_TORCH_CHECK(qzeros.scalar_type() == torch::headeronly::ScalarType::Int,
+                    "qweight.dtype != torch.int32");
+    STD_TORCH_CHECK(qzeros.get_device_index() == device_index,
+                    "qzeros is not on the same device with qweight");
+
+    int32_t group_size = qweight.size(0) / qzeros.size(0);
+    STD_TORCH_CHECK(qweight.size(1) == qzeros.size(1),
+                    "qweight.size(1) != qzeros.size(1)");
+    STD_TORCH_CHECK(qweight.size(0) % qzeros.size(0) == 0,
+                    "qweight.size(0) % qzeros.size(0) != 0");
+    STD_TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0");
+
+    dim3 blocks(size_k / 32, size_n / 8);
+    marlin_int4_fp8_preprocess_kernel_awq<<<blocks, 32, 0, stream>>>(
+        reinterpret_cast<const int32_t*>(qweight.const_data_ptr()),
+        reinterpret_cast<int32_t*>(output.mutable_data_ptr()),
+        reinterpret_cast<const int32_t*>(qzeros.const_data_ptr()), size_n,
+        size_k, group_size);
+  }
+
+  return output;
+}
+
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
+  m.impl("marlin_int4_fp8_preprocess", TORCH_BOX(&marlin_int4_fp8_preprocess));
+}
@@ -33,6 +33,68 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {

  // TODO: Remove this once ROCm upgrade to torch 2.11.
  ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
+
+  // Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
+  ops.def(
+      "machete_supported_schedules("
+      "   ScalarType a_type,"
+      "   int b_type,"
+      "   ScalarType? maybe_group_scales_type,"
+      "   ScalarType? maybe_group_zeros_type,"
+      "   ScalarType? maybe_channel_scales_type,"
+      "   ScalarType? maybe_token_scales_type,"
+      "   ScalarType? maybe_out_type"
+      ") -> str[]");
+  ops.def(
+      "machete_mm("
+      "   Tensor A,"
+      "   Tensor B,"
+      "   int b_type,"
+      "   ScalarType? out_type,"
+      "   Tensor? group_scales,"
+      "   Tensor? group_zeros,"
+      "   int?    group_size,"
+      "   Tensor? channel_scales,"
+      "   Tensor? token_scales,"
+      "   str?    schedule"
+      ") -> Tensor");
+  ops.def(
+      "machete_prepack_B("
+      "   Tensor B,"
+      "   ScalarType a_type,"
+      "   int b_type,"
+      "   ScalarType? group_scales_type"
+      ") -> Tensor");
+  // conditionally compiled so impl registration is in source file
+
+  // Marlin GEMM
+  ops.def(
+      "marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
+      "Tensor? b_bias_or_none,Tensor b_scales, "
+      "Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, "
+      "Tensor? "
+      "g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_type_id, "
+      "SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, "
+      "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor");
+  // conditionally compiled so impl registrations are in source file
+
+  // gptq_marlin repack from GPTQ.
+  ops.def(
+      "gptq_marlin_repack(Tensor b_q_weight, Tensor perm, "
+      "SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor");
+  // conditionally compiled so impl registrations are in source file
+
+  // awq_marlin repack from AWQ.
+  ops.def(
+      "awq_marlin_repack(Tensor b_q_weight, SymInt size_k, "
+      "SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor");
+  // conditionally compiled so impl registrations are in source file
+
+  // preprocess W-int4A-fp8 weight for marlin kernel
+  ops.def(
+      "marlin_int4_fp8_preprocess(Tensor qweight, "
+      "Tensor? qzeros_or_none, bool inplace) -> Tensor");
+  // conditionally compiled so impl registrations are in source file
 #endif

 #ifndef USE_ROCM
@@ -1,75 +0,0 @@
-#pragma once
-
-#include <torch/all.h>
-#include <Python.h>
-
-#include "machete_mm_kernel.cuh"
-#include "cutlass_extensions/torch_utils.hpp"
-#include "core/scalar_type.hpp"
-
-namespace machete {
-
-struct MMArgs {
-  torch::Tensor const& A;
-  torch::Tensor const& B;
-  vllm::ScalarType const& b_type;
-  std::optional<at::ScalarType> const& maybe_out_type;
-  std::optional<torch::Tensor> const& maybe_group_scales;
-  std::optional<torch::Tensor> const& maybe_group_zeros;
-  std::optional<int64_t> maybe_group_size;
-  std::optional<torch::Tensor> const& maybe_channel_scales;
-  std::optional<torch::Tensor> const& maybe_token_scales;
-  std::optional<std::string> maybe_schedule;
-};
-
-struct SupportedSchedulesArgs {
-  at::ScalarType a_type;
-  vllm::ScalarType b_type;
-  std::optional<at::ScalarType> maybe_group_scales_type;
-  std::optional<at::ScalarType> maybe_group_zeros_type;
-  std::optional<at::ScalarType> maybe_channel_scales_type;
-  std::optional<at::ScalarType> maybe_token_scales_type;
-  std::optional<at::ScalarType> maybe_out_type;
-};
-
-torch::Tensor mm_dispatch(MMArgs args);
-
-std::vector<std::string> supported_schedules_dispatch(
-    SupportedSchedulesArgs args);
-
-template <typename MacheteKernel>
-torch::Tensor run_impl(MMArgs args) {
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(args.A));
-
-  auto device = args.A.device();
-  auto stream = at::cuda::getCurrentCUDAStream(device.index());
-
-  int M = args.A.size(0);
-  int N = args.B.size(1);
-  int K = args.A.size(1);
-
-  // Allocate output
-  torch::Tensor D = torch::empty(
-      {M, N},
-      torch::TensorOptions()
-          .dtype(equivalent_scalar_type_v<typename MacheteKernel::ElementD>)
-          .device(device));
-
-  auto arguments = MacheteKernel::create_arguments(
-      stream,  //
-      args.A, args.B, D, args.maybe_group_scales, args.maybe_group_zeros,
-      args.maybe_group_size, args.maybe_channel_scales,
-      args.maybe_token_scales);
-  TORCH_CHECK(MacheteKernel::can_implement(arguments),
-              "Machete kernel cannot be run with these arguments");
-
-  size_t workspace_size = MacheteKernel::get_workspace_size(arguments);
-  torch::Tensor workspace = torch::empty(
-      workspace_size, torch::TensorOptions().dtype(torch::kU8).device(device));
-
-  MacheteKernel::run(arguments, workspace.mutable_data_ptr(), stream);
-
-  return D;
-};
-
-};  // namespace machete
@@ -1,73 +0,0 @@
-#include "machete_mm_launcher.cuh"
-#include "machete_prepack_launcher.cuh"
-#include "core/scalar_type.hpp"
-
-#include "core/registration.h"
-
-namespace machete {
-
-using namespace vllm;
-
-std::vector<std::string> supported_schedules(
-    at::ScalarType a_type, int64_t b_type_id,
-    std::optional<at::ScalarType> maybe_group_scales_type,
-    std::optional<at::ScalarType> maybe_group_zeros_type,
-    std::optional<at::ScalarType> maybe_channel_scales_type,
-    std::optional<at::ScalarType> maybe_token_scales_type,
-    std::optional<at::ScalarType> maybe_out_type) {
-  ScalarType const b_type = ScalarType::from_id(b_type_id);
-  return supported_schedules_dispatch({
-      .a_type = a_type,
-      .b_type = b_type,
-      .maybe_group_scales_type = maybe_group_scales_type,
-      .maybe_group_zeros_type = maybe_group_zeros_type,
-      .maybe_channel_scales_type = maybe_channel_scales_type,
-      .maybe_token_scales_type = maybe_token_scales_type,
-      .maybe_out_type = maybe_out_type,
-  });
-}
-
-torch::Tensor mm(torch::Tensor const& A, torch::Tensor const& B,
-                 int64_t b_type_id,
-                 std::optional<at::ScalarType> const& maybe_out_type,
-                 std::optional<torch::Tensor> const& maybe_group_scales,
-                 std::optional<torch::Tensor> const& maybe_group_zeros,
-                 std::optional<int64_t> maybe_group_size,
-                 std::optional<torch::Tensor> const& maybe_channel_scales,
-                 std::optional<torch::Tensor> const& maybe_token_scales,
-                 std::optional<std::string> maybe_schedule) {
-  ScalarType const b_type = ScalarType::from_id(b_type_id);
-  return mm_dispatch({.A = A,
-                      .B = B,
-                      .b_type = b_type,
-                      .maybe_out_type = maybe_out_type,
-                      .maybe_group_scales = maybe_group_scales,
-                      .maybe_group_zeros = maybe_group_zeros,
-                      .maybe_group_size = maybe_group_size,
-                      .maybe_channel_scales = maybe_channel_scales,
-                      .maybe_token_scales = maybe_token_scales,
-                      .maybe_schedule = maybe_schedule});
-}
-
-torch::Tensor prepack_B(
-    torch::Tensor const& B, at::ScalarType const& a_type, int64_t b_type_id,
-    std::optional<at::ScalarType> const& maybe_group_scales_type) {
-  ScalarType const b_type = ScalarType::from_id(b_type_id);
-  return prepack_B_dispatch(
-      {.B = B,
-       .a_type = a_type,
-       .b_type = b_type,
-       .maybe_group_scales_type = maybe_group_scales_type});
-}
-
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-  m.impl("machete_prepack_B", &prepack_B);
-  m.impl("machete_mm", &mm);
-}
-
-// use CatchAll since supported_schedules has no tensor arguments
-TORCH_LIBRARY_IMPL(TORCH_EXTENSION_NAME, CatchAll, m) {
-  m.impl("machete_supported_schedules", &supported_schedules);
-}
-
-};  // namespace machete
@@ -1,106 +0,0 @@
-
-
-#include "marlin.cuh"
-
-#include "core/registration.h"
-
-// for only non-zp format (like gptq)
-__global__ void marlin_int4_fp8_preprocess_kernel_without_zp(
-    // qweight: (size_k * size_n // 8,)
-    const int32_t* __restrict__ qweight,
-    // output: same shape with qweight
-    int32_t* __restrict__ output) {
-  int32_t val = qweight[blockIdx.x * 32 + threadIdx.x];
-  int32_t new_val = 0;
-
-#pragma unroll
-  for (int32_t i = 0; i < 8; i++) {
-    int32_t single_val = val & 0xF;
-    single_val = single_val >= 8 ? single_val - 8 : 15 - single_val;
-    new_val |= single_val << (i * 4);
-    val >>= 4;
-  }
-
-  output[blockIdx.x * 32 + threadIdx.x] = new_val;
-}
-
-// for awq format only (with zp and with awq weight layout)
-__global__ void marlin_int4_fp8_preprocess_kernel_awq(
-    // AWQ qweight: (size_k, size_n // 8)
-    const int32_t* __restrict__ qweight,
-    // output: same shape with qweight
-    int32_t* __restrict__ output,
-    // AWQ zeros: (size_k // group_size, size_n // 8)
-    const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k,
-    int32_t group_size) {
-  int32_t val =
-      qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y];
-  int32_t zero =
-      qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 +
-             blockIdx.y];
-  int32_t new_val = 0;
-
-#pragma unroll
-  for (int32_t i = 0; i < 8; i++) {
-    int32_t single_val = val & 0xF;
-    int32_t single_zero = zero & 0xF;
-
-    single_val =
-        single_val >= single_zero ? single_val - single_zero : 15 - single_val;
-    new_val |= single_val << (i * 4);
-    val >>= 4;
-    zero >>= 4;
-  }
-
-  output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val;
-}
-
-torch::Tensor marlin_int4_fp8_preprocess(
-    torch::Tensor& qweight, std::optional<torch::Tensor> qzeros_or_none,
-    bool inplace) {
-  TORCH_CHECK(qweight.device().is_cuda(), "qweight is not on GPU");
-  TORCH_CHECK(qweight.scalar_type() == at::ScalarType::Int,
-              "qweight.dtype != torch.int32");
-
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(qweight));
-
-  torch::Tensor output = inplace ? qweight : torch::empty_like(qweight);
-
-  if (!qzeros_or_none.has_value()) {
-    TORCH_CHECK(qweight.numel() * 8 % 256 == 0,
-                "qweight.numel() * 8 % 256 != 0");
-
-    int blocks = qweight.numel() * 8 / 256;
-    marlin_int4_fp8_preprocess_kernel_without_zp<<<blocks, 32>>>(
-        (const int32_t*)qweight.data_ptr(), (int32_t*)output.data_ptr());
-  } else {
-    int32_t size_k = qweight.size(0);
-    int32_t size_n = qweight.size(1) * 8;
-    torch::Tensor qzeros = qzeros_or_none.value();
-
-    TORCH_CHECK(size_k % 32 == 0, "size_k % 32 != 0");
-    TORCH_CHECK(qzeros.device().is_cuda(), "qzeros is not on GPU");
-    TORCH_CHECK(qzeros.scalar_type() == at::ScalarType::Int,
-                "qweight.dtype != torch.int32");
-    TORCH_CHECK(device_of(qweight) == device_of(qzeros),
-                "qzeros is not on the same device with qweight");
-
-    int32_t group_size = qweight.size(0) / qzeros.size(0);
-    TORCH_CHECK(qweight.size(1) == qzeros.size(1),
-                "qweight.size(1) != qzeros.size(1)");
-    TORCH_CHECK(qweight.size(0) % qzeros.size(0) == 0,
-                "qweight.size(0) % qzeros.size(0) != 0");
-    TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0");
-
-    dim3 blocks(size_k / 32, size_n / 8);
-    marlin_int4_fp8_preprocess_kernel_awq<<<blocks, 32>>>(
-        (const int32_t*)qweight.data_ptr(), (int32_t*)output.data_ptr(),
-        (const int32_t*)qzeros.data_ptr(), size_n, size_k, group_size);
-  }
-
-  return output;
-}
-
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-  m.impl("marlin_int4_fp8_preprocess", &marlin_int4_fp8_preprocess);
-}
@@ -68,68 +68,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  // custom types:
  // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA

-  // Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
-  ops.def(
-      "machete_supported_schedules("
-      "   ScalarType a_type,"
-      "   int b_type,"
-      "   ScalarType? maybe_group_scales_type,"
-      "   ScalarType? maybe_group_zeros_type,"
-      "   ScalarType? maybe_channel_scales_type,"
-      "   ScalarType? maybe_token_scales_type,"
-      "   ScalarType? maybe_out_type"
-      ") -> str[]");
-  ops.def(
-      "machete_mm("
-      "   Tensor A,"
-      "   Tensor B,"
-      "   int b_type,"
-      "   ScalarType? out_type,"
-      "   Tensor? group_scales,"
-      "   Tensor? group_zeros,"
-      "   int?    group_size,"
-      "   Tensor? channel_scales,"
-      "   Tensor? token_scales,"
-      "   str?    schedule"
-      ") -> Tensor");
-  ops.def(
-      "machete_prepack_B("
-      "   Tensor B,"
-      "   ScalarType a_type,"
-      "   int b_type,"
-      "   ScalarType? group_scales_type"
-      ") -> Tensor");
-  // conditionally compiled so impl registration is in source file
-
-  // Marlin Optimized Quantized GEMM (supports GPTQ, AWQ, FP8, NVFP4, MXFP4).
-  ops.def(
-      "marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
-      "Tensor? b_bias_or_none,Tensor b_scales, "
-      "Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, "
-      "Tensor? "
-      "g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_type_id, "
-      "SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, "
-      "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor");
-  // conditionally compiled so impl registration is in source file
-
-  // gptq_marlin repack from GPTQ.
-  ops.def(
-      "gptq_marlin_repack(Tensor b_q_weight, Tensor perm, "
-      "SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor");
-  // conditionally compiled so impl registrations are in source file
-
-  // awq_marlin repack from AWQ.
-  ops.def(
-      "awq_marlin_repack(Tensor b_q_weight, SymInt size_k, "
-      "SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor");
-  // conditionally compiled so impl registrations are in source file
-
-  // preprocess W-int4A-fp8 weight for marlin kernel
-  ops.def(
-      "marlin_int4_fp8_preprocess(Tensor qweight, "
-      "Tensor? qzeros_or_none, bool inplace) -> Tensor");
-  // conditionally compiled so impl registrations are in source file
-
 #endif
 }

@@ -37,7 +37,7 @@ th {
 | HuggingFace-HumanEval | ✅ | ✅ | `openai/openai_humaneval` |
 | HuggingFace-GSM8K | ✅ | ✅ | `openai/gsm8k` |
 | HuggingFace-Blazedit | ✅ | ✅ | `vdaita/edit_5k_char`, `vdaita/edit_10k_char` |
-| HuggingFace-ASR | ✅ | ✅ | `openslr/librispeech_asr`, `facebook/voxpopuli`,  `LIUM/tedlium`, `edinburghcstr/ami`,        `speechcolab/gigaspeech`,        `kensho/spgispeech` |
+| HuggingFace-ASR | ✅ | ✅ | `openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech`, `ArtificialAnalysis/Earnings22-Cleaned-AA`, `D4nt3/esb-datasets-earnings22-validation-tiny-filtered` |
 | Spec Bench | ✅ | ✅ | `wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl` |
 | SPEED-Bench | ✅ | ✅ | `curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -` |
 | Custom | ✅ | ✅ | Local file: `data.jsonl` |
@@ -532,7 +532,7 @@ vllm bench serve \
    --blazedit-max-distance 0.99
 ```

-`openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech`
+`openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech`, `ArtificialAnalysis/Earnings22-Cleaned-AA`, `D4nt3/esb-datasets-earnings22-validation-tiny-filtered`

 ```bash
 vllm bench serve \
@@ -60,15 +60,21 @@ the failure?

 ## Logs Wrangling

-Download a job's log (no Buildkite login required):
-
+Logs are public; no Buildkite login needed.
 [.buildkite/scripts/ci-fetch-log.sh](../../../.buildkite/scripts/ci-fetch-log.sh)
+saves each log as `ci-<build>-<job-name>.log`, stripped of timestamps and
+ANSI codes:

 ```bash
-# Find the failing job. Each row's URL is .../builds/<N>#<job_uuid>:
-gh pr checks <PR> --repo vllm-project/vllm
+# All failed jobs in a PR's latest build (current branch's PR if omitted):
+.buildkite/scripts/ci-fetch-log.sh --pr <PR>

-# Download + strip timestamps/ANSI in one step:
+# All failed jobs in a build (--soft also includes soft-failed jobs;
+# --all fetches every finished job):
+.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/<N>"
+
+# One job — `gh pr checks` URLs (#<job_uuid>) and web UI URLs (?sid=) both
+# work; pass "-" as a second argument to stream to stdout:
 .buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/<N>#<job_uuid>"
 ```

@@ -180,7 +180,8 @@ Priority is **1 = highest** (tried first).
 | `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
 | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
-| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | Any |
+| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ✅ | ❌ | All | Any |
+| `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
 | `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |

 > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
@@ -82,6 +82,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra

 | Architecture | Models | CG for Image | CG for Video |
 | ------------ | ------ | ------------ | ------------ |
+| `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | - |
 | `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ |
 | `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ |
 | `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ |
@@ -114,6 +115,14 @@ vllm serve Qwen/Qwen3-VL-32B \
  --compilation-config '{"cudagraph_mm_encoder": true}'
 ```

+For `Llama 4` (image only):
+
+```bash
+vllm serve meta-llama/Llama-4-Scout-17B-16E-Instruct \
+  --limit-mm-per-prompt '{"image": 1}' \
+  --compilation-config '{"cudagraph_mm_encoder": true}'
+```
+
 With explicit budgets:

 ```bash
@@ -0,0 +1,256 @@
+# NIXL push-mode KV transfer
+
+The default NIXL connector is **pull-based**: the decode (D) instance
+reads KV blocks from the prefill (P) instance via `NIXL READ` after
+prefill completes. `NixlPushConnector` adds a **push-based** alternative
+in which P writes the KV blocks directly into D's pre-allocated memory
+via `NIXL WRITE`.
+
+This document describes the threading, queues, and scheduling
+interactions specific to the push design. The pull-mode design is
+unchanged; the push connector reuses the same handshake, NIXL agent
+setup, and metadata path wherever possible.
+
+## High-level flow
+
+```mermaid
+sequenceDiagram
+    autonumber
+    participant Client
+    participant Proxy
+    participant DSched as D Scheduler
+    participant DWorker as D Worker (main)
+    participant DWriter as D Writer
+    participant PWriter as P Writer
+    participant PWorker as P Worker (main)
+    participant PSched as P Scheduler
+
+    Client->>Proxy: POST /v1/completions
+    Proxy->>PSched: prefill leg (do_remote_decode=True, max_tokens=1)
+    Proxy->>DSched: decode leg (do_remote_prefill=True, P coordinates)
+
+    note over DSched,DWriter: D side - register blocks with P
+    DSched->>DSched: update_state_after_alloc, stash registration, arm watchdog
+    DSched->>DWorker: build_connector_meta -> meta.push_registrations
+    DWorker->>DWriter: enqueue (req_id, reg_data) on _reg_send_inbox
+    DWriter->>PWriter: NIXL send_notif PUSH_REG msgpack
+
+    note over PSched,PWriter: P side - prefill, stage finished blocks
+    PSched->>PSched: request_finished, stash blocks
+    PSched->>PWorker: build_connector_meta -> meta.push_finished_blocks
+    PWorker->>PWriter: enqueue (req_id, blocks) on _finished_blocks_inbox
+
+    note over PWriter: P writer matches and WRITEs
+    PWriter->>PWriter: get_new_notifs returns PUSH_REG, route via _handle_push_reg_notif
+    alt PUSH_REG and finished blocks both present
+        PWriter->>PWriter: pop matching pair, fire WRITE
+    else only one side present
+        PWriter->>PWriter: stash and wait, self-poll only when blocks unmatched
+    end
+    PWriter->>PWriter: ensure D handshake (one-time)
+    PWriter->>DWriter: NIXL WRITE direct to D GPU + completion notif
+
+    note over DWorker,DWriter: D side - completion accounting
+    DWriter-->>DWorker: forward HB and completion notifs via _pending_completion_notifs
+    DWorker->>DWorker: _get_new_notifs drains, HB extends lease, completion marks recv done
+    DWorker->>DSched: update_connector_output(finished_recving)
+    DSched->>DSched: clear watchdog deadline
+
+    note over PWorker,PWriter: P side - reclaim
+    PWorker->>PWorker: get_finished, drain _sending_transfers, queue eviction
+    PWriter->>PWriter: drain _evict_finished_inbox, drop stale state
+    PWorker->>PSched: update_connector_output(finished_sending)
+    PSched->>PSched: free lease
+
+    DWorker-->>Proxy: stream decode tokens
+    Proxy-->>Client: response
+```
+
+## Threads
+
+``NixlPushConnectorWorker`` introduces a single dedicated background
+thread per worker (i.e. per TP rank), named ``nixl-push-writer``.
+Each owns the new push-specific NIXL operations on its rank:
+
+* ``nixl_wrapper.get_new_notifs()`` — receive notifications.
+* ``nixl_wrapper.send_notif(...)`` for the ``PUSH_REG:<msgpack>`` (D
+  side) and for the per-WRITE completion notif (P side).
+* ``nixl_wrapper.make_prepped_xfer(...) / transfer(...)`` — submit the
+  WRITE itself.
+
+Heartbeats continue to go out from the engine main thread via the
+existing base-worker ``_send_heartbeats`` plumbing inside
+``start_load_kv``.
+
+### Wake model
+
+The writer thread blocks on ``_push_writer_wake`` (a
+``threading.Event``) when it has no work. Three callers set the
+event:
+
+1. **``start_load_kv``** (worker main thread, called once per engine
+   step with the scheduler's metadata) — sets the wake only when the
+   step actually hands the writer new work, i.e. when
+   ``meta.push_registrations`` or ``meta.push_finished_blocks`` is
+   non-empty. This is the wake for new transfers.
+2. **``get_finished``** (worker main thread, called once per engine
+   step to report completions) — always sets the wake. The writer is
+   the sole consumer of ``nixl_wrapper.get_new_notifs()`` for push,
+   so this gives it a chance to drain inbound notifs (heartbeats from
+   D, completion notifs after a WRITE, late-arriving ``PUSH_REG``)
+   even when there is no new metadata to act on.
+3. **Handshake-completion callback** (background handshake executor
+   thread) — when a deferred D→P handshake finishes successfully, the
+   future's done-callback re-enqueues the registration onto
+   ``_reg_send_inbox`` and sets the wake so the corresponding
+   ``send_notif`` runs on the writer (we never call ``send_notif`` from
+   the executor thread). On this second pass ``_ensure_handshake``
+   returns ``None`` (the agent is now connected), so the writer sends
+   the ``PUSH_REG`` directly. If the handshake *failed*, the callback
+   fails the request instead of re-enqueuing, so there is no retry
+   loop.
+
+In addition to event-driven wakes, the writer self-polls at
+``_PUSH_WRITER_POLL_INTERVAL_MS = 1.0`` ms while there are P-side
+finished blocks waiting for an unmatched ``PUSH_REG``.
+
+When a request completes on P (lease expires or the WRITE finishes),
+``get_finished`` enqueues the request id onto ``_evict_finished_inbox``,
+which the writer drains to drop stale ``_push_finished_blocks`` /
+``_pending_d_registrations`` and stop self-polling.
+
+## Writer-local matching tables
+
+| Table                          | Owner            | Holds                                                                  |
+|--------------------------------|------------------|------------------------------------------------------------------------|
+| `_pending_d_registrations`     | writer           | D registrations received from a remote D, waiting for P's blocks       |
+| `_push_finished_blocks`        | writer           | P blocks staged by the scheduler, waiting for a remote D registration  |
+
+Either side can arrive first. The writer matches in both directions:
+when a ``PUSH_REG`` arrives we look up ``_push_finished_blocks``, and
+when finished blocks arrive we look up ``_pending_d_registrations``.
+Both lookups try an exact ``request_id`` match first, then fall back
+to comparing the ids after stripping the trailing per-engine random
+suffix (via ``get_base_request_id``). The fallback exists because the
+proxy hands the same ``X-Request-Id`` to both legs, so P and D wrap it
+into the same ``cmpl-<uuid>-<index>`` form and differ only by the
+8-hex randomization suffix that ``input_processor.assign_request_id``
+appends per engine. Stripping just that suffix normalizes both sides
+to the same id while preserving the completion index (so multi-prompt
+sub-requests stay distinct). It also works whether or not
+``VLLM_DISABLE_REQUEST_ID_RANDOMIZATION`` is set, which matters since
+that env var is slated for removal upstream.
+
+## Wire format
+
+A push registration is sent as a NIXL notification:
+
+```text
+PUSH_REG:<msgpack-encoded dict>
+```
+
+Fields in the dict:
+
+| Field                | Set by | Meaning                                                                |
+|----------------------|--------|------------------------------------------------------------------------|
+| ``request_id``       | D      | D's own vLLM request id; P's match key, echoed in the completion notif |
+| ``decode_engine_id`` | D      | D's engine id (P uses this for the reverse handshake)                  |
+| ``decode_host``      | D      | D's NIXL side-channel host                                             |
+| ``decode_port``      | D      | D's NIXL side-channel port                                             |
+| ``decode_tp_size``   | D      | D's tensor-parallel size                                               |
+| ``local_block_ids``  | D      | per-group lists of D's *logical* block ids (preallocated)              |
+| ``remote_engine_id`` | D      | P's engine id (for the existing P-side handshake)                      |
+| ``remote_host``      | D      | P's NIXL side-channel host                                             |
+| ``remote_port``      | D      | P's NIXL side-channel port                                             |
+| ``remote_tp_size``   | D      | P's tensor-parallel size                                               |
+
+D ships **logical** block ids; P expands them to physical block ids at
+WRITE-submission time using the ratio learned during the NIXL
+handshake (`remote_physical_blocks_per_logical`). This matches the
+pull-mode contract — schedulers ship logical ids, workers expand to
+physical at submission.
+
+The completion notif sent from P to D after a WRITE is the existing
+`<request_id>:<tp_size>` format used in pull mode (here ``request_id``
+is D's own request id, taken from the registration), so the D-side
+accounting code is unchanged.
+
+## Scheduler-side responsibilities
+
+`NixlPushConnectorScheduler` extends the base scheduler with:
+
+* **D side** — `update_state_after_alloc` stashes registration data in
+  `_push_pending_registrations` and arms a soft watchdog
+  (`_push_registration_deadlines`). `build_connector_meta` drains the
+  stash into `meta.push_registrations` and any expired entries are
+  dropped with a warning.
+* **P side** — `request_finished` stashes block IDs in
+  `_finished_request_blocks` (for the lease and for
+  `has_pending_push_work`) and `_newly_finished_push_blocks` (for the
+  next worker step via `meta.push_finished_blocks`).
+* **Both sides** — `has_pending_push_work` keeps the engine main loop
+  stepping while there is in-flight push state, so the writer always
+  gets at least one wake per step.
+
+`update_connector_output`:
+
+* `finished_sending` (P side) clears the lease entry.
+* `finished_recving` (D side) clears the watchdog deadline.
+
+## Timeouts and watchdogs
+
+Two per-request timers are armed on the scheduler:
+
+* **D-side registration watchdog** — ``_push_registration_deadlines``.
+  If a registered request does not see a push completion within
+  ``push_registration_timeout`` seconds (defaults to
+  ``decoder_kv_blocks_ttl``), ``build_connector_meta`` drops the stale
+  registration and the pending entry, logs a warning, and stops trying
+  to resend the registration. The corresponding request remains tracked
+  in ``_reqs_need_recv``; it is the engine's request-level abort path
+  (or the user / proxy timing out the HTTP call) that ultimately fails
+  the request.
+* **P-side block lease** — same ``_kv_lease_duration`` used by pull
+  mode. ``request_finished`` sets the expiration in ``_reqs_need_send``
+  and ``update_connector_output(finished_sending=...)`` clears it on
+  successful WRITE. Stale leases are reaped by ``get_finished`` in the
+  base worker, which then enqueues the eviction onto
+  ``_evict_finished_inbox`` so the writer also stops self-polling.
+
+## Failure handling
+
+* **D-side handshake failure (P→D handshake before sending PUSH_REG)** —
+  the future's done-callback calls ``_handle_failed_transfer(rid, None)``,
+  which marks D's pre-allocated blocks invalid and enqueues onto
+  ``_failed_recv_reqs`` so the next ``get_finished`` reports the
+  request as a failed recv. Same recv-side accounting as pull mode.
+* **D-side ``send_notif`` failure when shipping the PUSH_REG to P** —
+  identical handling: ``_handle_failed_transfer`` marks the recv as
+  failed.
+* **P-side WRITE submission failure** — the WRITE handle (if any) is
+  released and ``xfer_stats.record_failed_transfer()`` bumps the
+  failure counter. We deliberately do not call
+  ``_handle_failed_transfer`` here: ``req_id`` on the P side has no
+  entry in ``_recving_metadata`` (P is not the receiver), so the
+  helper would put a P-local request id into ``_failed_recv_reqs``
+  and trip the assertion in the base worker's ``get_finished``. The
+  outbound WRITE is dropped on the floor; D's lease watchdog handles
+  the missing completion.
+
+## Summary
+
+The push design is a small, well-contained extension on top of the
+existing NIXL connector:
+
+* one new connector class, one new scheduler class, one new worker
+  class — all subclasses of the existing base classes;
+* one dedicated background thread per worker;
+* a few cross-thread queues, each with a single consumer (the writer);
+  most have one producer, except ``_reg_send_inbox``, which is fed both
+  by the engine main thread (new registrations) and by the
+  handshake-completion callback (registrations replayed after their
+  D→P handshake finishes);
+* one new notification type (`PUSH_REG:<msgpack>`).
+
+Behavior on the engine main thread is otherwise unchanged. The writer
+thread is event-driven and idle when there is no push work.
@@ -423,6 +423,54 @@ To enable this feature:
 --kv-transfer-config '{..., "kv_connector_extra_config": {"enable_cross_layers_blocks": "True"}}'
 ```

+## Metrics Reference
+
+vLLM periodically logs a `KV Transfer metrics` line summarising NIXL transfer
+activity for the last reporting interval. Example output:
+
+```text
+KV Transfer metrics: Num successful transfers=4, Avg xfer time (ms)=1.381,
+P90 xfer time (ms)=2.601, Avg post time (ms)=0.672, P90 post time (ms)=0.801,
+Avg MB per transfer=2.25, Throughput (MB/s)=1629.549, Avg number of descriptors=72.0
+```
+
+The table below describes each field. All timing values cover only the
+successful transfers recorded in the current interval; failed transfers are
+counted separately via Prometheus (see
+[Prometheus metrics](#prometheus-metrics) below).
+
+| Metric | Unit | Description |
+| -------- | ------ | ------------- |
+| `Num successful transfers` | count | Number of NIXL KV-block transfers that completed without error during the interval. A transfer corresponds to one prefill request's worth of KV cache being moved from the prefiller to the decoder (or vice versa in bidirectional mode). |
+| `Avg xfer time (ms)` | ms | Mean end-to-end transfer duration (`xferDuration` in NIXL telemetry, converted from µs). Measured from when the request is posted to when the backend reports completion, so it includes both the posting step and the actual data movement. |
+| `P90 xfer time (ms)` | ms | 90th-percentile transfer duration. Use this to identify tail latency: a large gap between average and P90 suggests occasional stragglers (e.g., network congestion or large KV blocks). |
+| `Avg post time (ms)` | ms | Mean time to submit the transfer request to the RDMA backend (`postDuration` in NIXL telemetry). This is the synchronous cost of posting work to the NIC queue (descriptor setup, etc.) before the async data movement begins. |
+| `P90 post time (ms)` | ms | 90th-percentile request-posting duration. Elevated P90 here (with low xfer P90) points to overhead in submitting requests rather than in the data transfer itself. |
+| `Avg MB per transfer` | MB | Mean payload size per transfer, computed as `total bytes transferred / number of transfers`. Reflects the average KV cache footprint of a single request (sequence length × layers × head dimension × dtype bytes). |
+| `Throughput (MB/s)` | MB/s | Effective bandwidth over the interval: `total MB transferred / total xfer time (s)` across all successful transfers. This is aggregate throughput, not per-request bandwidth. |
+| `Avg number of descriptors` | count | Mean number of NIXL memory descriptors (scatter-gather segments) submitted per transfer. More descriptors indicate more fragmented or larger KV cache allocations; very high counts can increase descriptor-registration overhead. |
+
+### Prometheus metrics
+
+In addition to the periodic log line, the following Prometheus metrics are
+exported when NixlConnector is active:
+
+| Metric name | Type | Description |
+| ------------- | ------ | ------------- |
+| `vllm:nixl_xfer_time_seconds` | Histogram | Per-transfer RDMA copy duration (seconds). |
+| `vllm:nixl_post_time_seconds` | Histogram | Time to submit the transfer request to the RDMA backend (seconds). |
+| `vllm:nixl_bytes_transferred` | Histogram | Bytes moved per transfer. |
+| `vllm:nixl_num_descriptors` | Histogram | Descriptor count per transfer. |
+| `vllm:nixl_num_failed_transfers` | Counter | Cumulative count of failed NIXL KV-block transfers. |
+| `vllm:nixl_num_failed_notifications` | Counter | Cumulative count of failed completion notifications (`send_notif`). |
+| `vllm:nixl_num_kv_expired_reqs` | Counter | Requests whose KV blocks expired on the prefiller before the decoder read them (tracked on the P instance). |
+
+!!! tip
+    High `vllm:nixl_num_kv_expired_reqs` indicates that the prefiller's lease
+    duration (`kv_lease_duration`) is too short for your network or workload.
+    Increase it via `--kv-transfer-config '{"kv_connector_extra_config":
+    {"kv_lease_duration": <seconds>}}'`.
+
 ## Example Scripts/Code

 Refer to these example scripts in the vLLM repository:
@@ -115,18 +115,28 @@ Whether vLLM enforces the tool parameter schema during generation depends on the
 | --- | --- | --- |
 | Named function | Yes (via structured outputs backend) | Arguments are guaranteed to be valid JSON conforming to the function's parameter schema. |
 | `"required"` | Yes (via structured outputs backend) | Same as named function. The model must produce at least one tool call. |
-| `"auto"` | No | The model generates freely. A tool-call parser extracts tool calls from the raw text. Arguments may be malformed or not match the schema. |
+| `"auto"` | Depends on the parser | Model-specific structural-tag parsers can constrain tool-call arguments with structured outputs. Other parsers generate freely and extract tool calls from raw text. |
 | `"none"` | N/A | No tool calls are produced. |

-When schema conformance matters, prefer `tool_choice="required"` or named function calling over `"auto"`.
+### Strict Mode

-### Strict Mode (`strict` parameter)
+Strict tool calling makes function-call arguments adhere to the function schema instead of relying only on best-effort parsing. vLLM implements strict tool calling for structural-tag based tool parsers by using the structured outputs backend under the hood.

-The [OpenAI API](https://platform.openai.com/docs/guides/function-calling#strict-mode) supports a `strict` field on function definitions. When set to `true`, OpenAI uses constrained decoding to guarantee that tool-call arguments match the function schema, even in `tool_choice="auto"` mode.
+For best compatibility with strict schema enforcement, define tool parameter schemas in the OpenAI strict-schema style:

-vLLM **does not implement** `strict` mode today. The `strict` field is accepted in requests (to avoid breaking clients that set it), but it has no effect on decoding behavior. In auto mode, argument validity depends entirely on the model's output quality and the parser's extraction logic.
+* Set `additionalProperties` to `false` for each object in `parameters`.
+* Mark all fields in `properties` as required.
+* Represent optional fields by allowing `null`, for example `{"type": ["string", "null"]}`.

-Tracking issues: [#15526](https://github.com/vllm-project/vllm/issues/15526), [#16313](https://github.com/vllm-project/vllm/issues/16313).
+vLLM controls structural-tag strict tool calling with the `VLLM_ENFORCE_STRICT_TOOL_CALLING` environment variable. It defaults to `true`.
+
+```bash
+VLLM_ENFORCE_STRICT_TOOL_CALLING=false vllm serve ...
+```
+
+When this variable is `true`, structural-tag based tool parsers attach a structural tag to the request, so the structured outputs backend can constrain the model-specific tool-call format and function-call arguments. When it is `false`, vLLM does not attach structural tags for tool calling. In that case, `tool_choice="auto"` falls back to best-effort parser extraction from the raw model output, and no structural-tag constraint is applied.
+
+This environment variable only affects structural-tag based tool calling. It does not change schema-derived structured outputs used by named function calling or `tool_choice="required"`.

 ## Automatic Function Calling

@@ -146,7 +156,7 @@ from HuggingFace; and you can find an example of this in a `tokenizer_config.jso
 If your favorite tool-calling model is not supported, please feel free to contribute a parser & tool use chat template!

 !!! note
-    With `tool_choice="auto"`, tool-call arguments are extracted from the model's raw text output by the selected parser. No schema-level constraint is applied during decoding, so arguments may occasionally be malformed or violate the function's parameter schema. See [Constrained Decoding Behavior](#constrained-decoding-behavior) for details.
+    With `tool_choice="auto"`, schema-level constraint depends on the selected parser and `VLLM_ENFORCE_STRICT_TOOL_CALLING`. Structural-tag parsers can enforce tool-call constraints when it is `true`; when it is `false`, or when the selected parser has no structural-tag support, vLLM extracts tool calls from raw text, so arguments may occasionally be malformed or violate the function's parameter schema.

 ### Hermes Models (`hermes`)

@@ -423,7 +423,6 @@ th {
 | `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
 | `HYV3ForCausalLM` | HY3 | `tencent/Hy3-preview-Base`, `tencent/Hy3-preview` | ✅︎ | ✅︎ |
 | `HyperCLOVAXForCausalLM` | HyperCLOVAX-SEED-Think-14B | `naver-hyperclovax/HyperCLOVAX-SEED-Think-14B` | ✅︎ | ✅︎ |
-| `InternLMForCausalLM` | InternLM | `internlm/internlm-7b`, `internlm/internlm-chat-7b`, etc. | ✅︎ | ✅︎ |
 | `InternLM2ForCausalLM` | InternLM2 | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc. | ✅︎ | ✅︎ |
 | `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ |
 | `IQuestCoderForCausalLM` | IQuestCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Instruct`, etc. | | |
@@ -578,7 +577,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `InternS1ForConditionalGeneration` | Intern-S1 | T + I<sup>E+</sup> + V<sup>E+</sup> | `internlm/Intern-S1`, `internlm/Intern-S1-mini`, etc. | ✅︎ | ✅︎ |
 | `InternS1ProForConditionalGeneration` | Intern-S1-Pro | T + I<sup>E+</sup> + V<sup>E+</sup> | `internlm/Intern-S1-Pro`, etc. | ✅︎ | ✅︎ |
 | `InternS2PreviewForConditionalGeneration` | Intern-S2-Preview | T + I<sup>E+</sup> + V<sup>E+</sup> | `internlm/Intern-S2-Preview`, etc. | ✅︎ | ✅︎ |
-| `InternVLChatModel` | InternVL 3.5, InternVL 3.0, InternVideo 2.5, InternVL 2.5, Mono-InternVL, InternVL 2.0 | T + I<sup>E+</sup> + (V<sup>E+</sup>) | `OpenGVLab/InternVL3_5-14B`, `OpenGVLab/InternVL3-9B`, `OpenGVLab/InternVideo2_5_Chat_8B`, `OpenGVLab/InternVL2_5-4B`, `OpenGVLab/Mono-InternVL-2B`, `OpenGVLab/InternVL2-4B`, etc. | ✅︎ | ✅︎ |
+| `InternVLChatModel` | InternVL 3.5, InternVL 3.0, InternVideo 2.5, InternVL 2.5, InternVL 2.0 | T + I<sup>E+</sup> + (V<sup>E+</sup>) | `OpenGVLab/InternVL3_5-14B`, `OpenGVLab/InternVL3-9B`, `OpenGVLab/InternVideo2_5_Chat_8B`, `OpenGVLab/InternVL2_5-4B`, `OpenGVLab/InternVL2-4B`, etc. | ✅︎ | ✅︎ |
 | `InternVLForConditionalGeneration` | InternVL 3.0 (HF format) | T + I<sup>E+</sup> + V<sup>E+</sup> | `OpenGVLab/InternVL3-1B-hf`, etc. | ✅︎ | ✅︎ |
 | `KananaVForConditionalGeneration` | Kanana-V | T + I<sup>+</sup> | `kakaocorp/kanana-1.5-v-3b-instruct`, etc. | | ✅︎ |
 | `KeyeForConditionalGeneration` | Keye-VL-8B-Preview | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-8B-Preview` | ✅︎ | ✅︎ |
@@ -0,0 +1,429 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Push-mode disaggregated prefilling proxy demo.
+
+Companion to ``disagg_proxy_demo.py`` (pull mode). The client-facing API is
+the same; the difference is in how P and D coordinate the KV transfer:
+
+* Pull mode: proxy forwards P's ``kv_transfer_params`` (including
+  ``remote_block_ids``) to D, and D pulls KV from P via NIXL READ.
+* Push mode: proxy hands D **only** P's coordinates
+  (``remote_engine_id``, ``remote_host``, ``remote_port``, ``tp_size``)
+  and the shared ``remote_request_id``. D registers its locally allocated
+  blocks with P over a NIXL notification; P then pushes the KV to D via
+  NIXL WRITE.
+
+Launch multiple vLLM instances configured with ``NixlPushConnector`` and
+matching ``engine_id`` / ``side_channel_port``, then start this proxy:
+
+  python3 examples/disaggregated/disaggregated_serving/\
+disagg_proxy_pushconnector_demo.py \
+       --model $model_name \
+       --prefill localhost:8100 \
+       --decode  localhost:8200 \
+       --prefill-engine-id prefill-engine-001 \
+       --prefill-kv-host  10.0.0.1 \
+       --prefill-side-channel-port 5600 \
+       --prefill-tp-size 1 \
+       --port 8000
+"""
+
+import argparse
+import contextlib
+import ipaddress
+import itertools
+import json
+import logging
+import os
+import sys
+import uuid
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+
+import aiohttp
+import uvicorn
+from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+
+AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60)
+logger = logging.getLogger()
+logging.basicConfig(level=logging.INFO)
+
+
+class SchedulingPolicy(ABC):
+    @abstractmethod
+    def schedule(self, cycler: itertools.cycle):
+        raise NotImplementedError("Scheduling Proxy is not set.")
+
+
+class RoundRobinSchedulingPolicy(SchedulingPolicy):
+    def schedule(self, cycler: itertools.cycle) -> str:
+        return next(cycler)
+
+
+class PushProxy:
+    """Push-mode proxy.
+
+    The structure mirrors the pull-mode ``Proxy`` in
+    ``disagg_proxy_demo.py``: an APIRouter with ``/v1/completions``,
+    ``/v1/chat/completions``, ``/status`` and ``/instances/add``, plus
+    round-robin scheduling across multiple P / D instances.
+
+    Push-specific differences are confined to the request-handling
+    methods (``create_completion`` / ``create_chat_completion``):
+
+    * D's ``kv_transfer_params`` is built from CLI-provided P
+      coordinates instead of being derived from P's response.
+    * P and D requests are issued concurrently — D registers blocks and
+      waits while P prefills and pushes.
+    """
+
+    def __init__(
+        self,
+        prefill_instances: list[str],
+        decode_instances: list[str],
+        model: str,
+        scheduling_policy: SchedulingPolicy,
+        prefill_engine_id: str,
+        prefill_kv_host: str,
+        prefill_side_channel_port: int,
+        prefill_tp_size: int,
+        custom_create_completion: Callable[[Request], StreamingResponse] | None = None,
+        custom_create_chat_completion: Callable[[Request], StreamingResponse]
+        | None = None,
+    ):
+        self.prefill_instances = prefill_instances
+        self.decode_instances = decode_instances
+        self.prefill_cycler = itertools.cycle(prefill_instances)
+        self.decode_cycler = itertools.cycle(decode_instances)
+        self.model = model
+        self.scheduling_policy = scheduling_policy
+
+        # Push-mode metadata: D needs P's coordinates up-front. Pull mode
+        # learns these from P's response; push mode uses CLI args because
+        # D issues its registration before P responds.
+        self.push_metadata = {
+            "do_remote_decode": False,
+            "do_remote_prefill": True,
+            "remote_engine_id": prefill_engine_id,
+            "remote_host": prefill_kv_host,
+            "remote_port": prefill_side_channel_port,
+            "tp_size": prefill_tp_size,
+        }
+
+        self.custom_create_completion = custom_create_completion
+        self.custom_create_chat_completion = custom_create_chat_completion
+        self.router = APIRouter()
+        self.setup_routes()
+
+    # ── routes ──────────────────────────────────────────────────────── #
+
+    def setup_routes(self):
+        self.router.post(
+            "/v1/completions", dependencies=[Depends(self.validate_json_request)]
+        )(
+            self.custom_create_completion
+            if self.custom_create_completion
+            else self.create_completion
+        )
+        self.router.post(
+            "/v1/chat/completions", dependencies=[Depends(self.validate_json_request)]
+        )(
+            self.custom_create_chat_completion
+            if self.custom_create_chat_completion
+            else self.create_chat_completion
+        )
+        self.router.get("/status", response_class=JSONResponse)(self.get_status)
+
+    async def validate_json_request(self, raw_request: Request):
+        content_type = raw_request.headers.get("content-type", "").lower()
+        if content_type != "application/json":
+            raise HTTPException(
+                status_code=415,
+                detail="Unsupported Media Type: Only 'application/json' is allowed",
+            )
+
+    # ── HTTP forwarding ─────────────────────────────────────────────── #
+
+    async def forward_request(self, url, data, headers, use_chunked=True):
+        async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
+            try:
+                async with session.post(
+                    url=url, json=data, headers=headers
+                ) as response:
+                    if 200 <= response.status < 300 or 400 <= response.status < 500:
+                        if use_chunked:
+                            async for chunk_bytes in response.content.iter_chunked(
+                                1024
+                            ):
+                                yield chunk_bytes
+                        else:
+                            yield await response.read()
+                    else:
+                        error_content = await response.text()
+                        with contextlib.suppress(json.JSONDecodeError):
+                            error_content = json.loads(error_content)
+                        logger.error(
+                            "Request failed with status %s: %s",
+                            response.status,
+                            error_content,
+                        )
+                        raise HTTPException(
+                            status_code=response.status,
+                            detail=f"Request failed with status {response.status}: "
+                            f"{error_content}",
+                        )
+            except aiohttp.ClientError as e:
+                logger.error("ClientError occurred: %s", str(e))
+                raise HTTPException(
+                    status_code=502,
+                    detail="Bad Gateway: Error communicating with upstream server.",
+                ) from e
+            except Exception as e:
+                logger.error("Unexpected error: %s", str(e))
+                raise HTTPException(status_code=500, detail=str(e)) from e
+
+    def schedule(self, cycler: itertools.cycle) -> str:
+        return self.scheduling_policy.schedule(cycler)
+
+    async def get_status(self):
+        return {
+            "mode": "push",
+            "prefill_node_count": len(self.prefill_instances),
+            "decode_node_count": len(self.decode_instances),
+            "prefill_nodes": self.prefill_instances,
+            "decode_nodes": self.decode_instances,
+            "prefill_engine_id": self.push_metadata["remote_engine_id"],
+            "prefill_kv_host": self.push_metadata["remote_host"],
+            "prefill_side_channel_port": self.push_metadata["remote_port"],
+            "prefill_tp_size": self.push_metadata["tp_size"],
+        }
+
+    # ── push-mode request handling ──────────────────────────────────── #
+
+    def _build_decode_kv_params(self, request_id: str) -> dict:
+        """Push-mode kv_transfer_params for D.
+
+        ``remote_block_ids`` is intentionally omitted: D allocates its
+        own blocks and registers them with P; P determines the
+        prefill-side block IDs and ships them via the WRITE.
+        """
+        params = self.push_metadata.copy()
+        params["remote_request_id"] = request_id
+        return params
+
+    def _common_headers(self, request_id: str) -> dict:
+        h = {"X-Request-Id": request_id}
+        api_key = os.environ.get("OPENAI_API_KEY")
+        if api_key:
+            h["Authorization"] = f"Bearer {api_key}"
+        return h
+
+    async def _push_completion(self, raw_request: Request, path: str):
+        """Shared body for /v1/completions and /v1/chat/completions.
+
+        Push mode fires P and D concurrently:
+          * P runs a normal prefill (max_tokens=1, do_remote_decode=True).
+          * D runs the decode (do_remote_prefill=True, no remote_block_ids).
+
+        D blocks waiting for P's WRITE; the response streamed back to the
+        client is the decode output from D.
+        """
+        request = await raw_request.json()
+        request_id = str(uuid.uuid4())
+
+        # Prefill leg (max_tokens=1, signals P to keep KV around for D).
+        prefill_request = request.copy()
+        prefill_request["max_tokens"] = 1
+        if "max_completion_tokens" in prefill_request:
+            prefill_request["max_completion_tokens"] = 1
+        prefill_request["kv_transfer_params"] = {
+            "do_remote_decode": True,
+            "do_remote_prefill": False,
+            "remote_engine_id": None,
+            "remote_block_ids": None,
+            "remote_host": None,
+            "remote_port": None,
+        }
+
+        # Decode leg (push mode: no remote_block_ids).
+        decode_request = request.copy()
+        decode_request["kv_transfer_params"] = self._build_decode_kv_params(request_id)
+
+        prefill_instance = self.schedule(self.prefill_cycler)
+        decode_instance = self.schedule(self.decode_cycler)
+        headers = self._common_headers(request_id)
+
+        # Fire prefill; we don't read its body but must drain the
+        # connection so the upstream server can free its slot.
+        async for _ in self.forward_request(
+            f"http://{prefill_instance}{path}", prefill_request, headers
+        ):
+            continue
+
+        generator = self.forward_request(
+            f"http://{decode_instance}{path}", decode_request, headers
+        )
+        return StreamingResponse(generator)
+
+    async def create_completion(self, raw_request: Request):
+        try:
+            return await self._push_completion(raw_request, "/v1/completions")
+        except HTTPException:
+            raise
+        except Exception:
+            exc_info = sys.exc_info()
+            print("Error occurred in disagg push proxy server")
+            print(exc_info)
+            raise
+
+    async def create_chat_completion(self, raw_request: Request):
+        try:
+            return await self._push_completion(raw_request, "/v1/chat/completions")
+        except HTTPException:
+            raise
+        except Exception:
+            exc_info = sys.exc_info()
+            error_messages = [str(e) for e in exc_info if e]
+            print("Error occurred in disagg push proxy server")
+            print(error_messages)
+            return StreamingResponse(
+                content=iter(error_messages), media_type="text/event-stream"
+            )
+
+
+class PushProxyServer:
+    def __init__(
+        self,
+        args: argparse.Namespace,
+        scheduling_policy: SchedulingPolicy | None = None,
+        create_completion: Callable[[Request], StreamingResponse] | None = None,
+        create_chat_completion: Callable[[Request], StreamingResponse] | None = None,
+    ):
+        self.validate_parsed_serve_args(args)
+        self.port = args.port
+        self.proxy_instance = PushProxy(
+            prefill_instances=[] if args.prefill is None else args.prefill,
+            decode_instances=[] if args.decode is None else args.decode,
+            model=args.model,
+            scheduling_policy=(
+                scheduling_policy
+                if scheduling_policy is not None
+                else RoundRobinSchedulingPolicy()
+            ),
+            prefill_engine_id=args.prefill_engine_id,
+            prefill_kv_host=args.prefill_kv_host,
+            prefill_side_channel_port=args.prefill_side_channel_port,
+            prefill_tp_size=args.prefill_tp_size,
+            custom_create_completion=create_completion,
+            custom_create_chat_completion=create_chat_completion,
+        )
+
+    def validate_parsed_serve_args(self, args: argparse.Namespace):
+        if not args.prefill:
+            raise ValueError("Please specify at least one prefill node.")
+        if not args.decode:
+            raise ValueError("Please specify at least one decode node.")
+        if not args.prefill_engine_id:
+            raise ValueError(
+                "--prefill-engine-id is required in push mode (it must match "
+                "the engine_id passed to the prefill vLLM instance via "
+                "--kv-transfer-config)."
+            )
+        if not args.prefill_kv_host:
+            raise ValueError(
+                "--prefill-kv-host is required in push mode (the IP / host "
+                "that the prefill vLLM advertises on its NIXL side channel)."
+            )
+        self.validate_instances(args.prefill)
+        self.validate_instances(args.decode)
+
+    def validate_instances(self, instances: list):
+        for instance in instances:
+            if len(instance.split(":")) != 2:
+                raise ValueError(f"Invalid instance format: {instance}")
+            host, port = instance.split(":")
+            try:
+                if host != "localhost":
+                    ipaddress.ip_address(host)
+                port = int(port)
+                if not (0 < port < 65536):
+                    raise ValueError(f"Invalid port number in instance: {instance}")
+            except Exception as e:
+                raise ValueError(f"Invalid instance {instance}: {str(e)}") from e
+
+    def run_server(self):
+        app = FastAPI()
+        app.include_router(self.proxy_instance.router)
+        config = uvicorn.Config(app, port=self.port, loop="uvloop")
+        server = uvicorn.Server(config)
+        server.run()
+
+
+def parse_args():
+    parser = argparse.ArgumentParser("vLLM disaggregated push-mode proxy server.")
+    parser.add_argument("--model", "-m", type=str, required=True, help="Model name")
+
+    parser.add_argument(
+        "--prefill",
+        "-p",
+        type=str,
+        nargs="+",
+        help="List of prefill node URLs (host:port)",
+    )
+
+    parser.add_argument(
+        "--decode",
+        "-d",
+        type=str,
+        nargs="+",
+        help="List of decode node URLs (host:port)",
+    )
+
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=8000,
+        help="Server port number",
+    )
+
+    # Push-mode specific: P's coordinates that D needs in advance.
+    parser.add_argument(
+        "--prefill-engine-id",
+        type=str,
+        required=True,
+        help=(
+            "engine_id of the prefill vLLM instance (must match "
+            "--kv-transfer-config engine_id on the prefill server)"
+        ),
+    )
+    parser.add_argument(
+        "--prefill-kv-host",
+        type=str,
+        required=True,
+        help=(
+            "IP / host the prefill vLLM advertises on its NIXL side "
+            "channel (VLLM_NIXL_SIDE_CHANNEL_HOST)"
+        ),
+    )
+    parser.add_argument(
+        "--prefill-side-channel-port",
+        type=int,
+        default=5600,
+        help="NIXL side channel port on the prefill node "
+        "(VLLM_NIXL_SIDE_CHANNEL_PORT, default 5600)",
+    )
+    parser.add_argument(
+        "--prefill-tp-size",
+        type=int,
+        default=1,
+        help="Tensor parallel size of the prefill vLLM instance",
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    proxy_server = PushProxyServer(args=args)
+    proxy_server.run_server()
@@ -2532,6 +2532,7 @@ MODELS_NEED_VIDEO_METADATA = [


 MODELS_SUPPORT_VIT_CUDA_GRAPH = [
+    "llama4",
    "internvl_chat",
    "qwen2_5_vl",
    "qwen3_vl",
@@ -25,7 +25,7 @@ outlines_core == 0.2.14
 # required for outlines backend disk cache
 diskcache == 5.6.3
 lark == 1.2.2
-xgrammar >= 0.2.0, < 1.0.0; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le"
+xgrammar >= 0.2.1, < 1.0.0; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le"
 typing_extensions >= 4.10
 filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317
 partial-json-parser # used for parsing partial JSON outputs
@@ -1367,7 +1367,7 @@ word2number==1.1
    # via lm-eval
 wrapt==2.1.2
    # via smart-open
-xgrammar==0.2.0
+xgrammar==0.2.1
    # via
    #   -c requirements/common.txt
    #   -r requirements/test/../common.txt
@@ -468,6 +468,8 @@ impl ServeArgs {
            self.runtime.model.clone(),
            self.runtime.max_model_len,
            self.runtime.language_model_only,
+            self.runtime.disable_log_stats,
+            self.runtime.shutdown_timeout,
            handshake_port,
        )
    }
@@ -100,6 +100,40 @@ fn serve_args_auto_forward_enable_lora_to_python() {
    assert_eq!(args.managed_engine.python_args, vec!["--enable-lora"]);
 }

+#[test]
+fn serve_args_forward_shutdown_timeout_to_managed_engine() {
+    let cli = Cli::try_parse_from([
+        "vllm-rs",
+        "serve",
+        "Qwen/Qwen3-0.6B",
+        "--shutdown-timeout",
+        "60",
+    ])
+    .unwrap();
+
+    let Command::Serve(args) = cli.command else {
+        panic!("expected serve args");
+    };
+    assert_eq!(args.runtime.shutdown_timeout, 60);
+
+    let config = args.to_managed_engine_config(5555);
+    assert_eq!(config.python_args, vec!["--shutdown-timeout", "60"]);
+}
+
+#[test]
+fn serve_args_forward_disable_log_stats_to_managed_engine() {
+    let cli = Cli::try_parse_from(["vllm-rs", "serve", "Qwen/Qwen3-0.6B", "--disable-log-stats"])
+        .unwrap();
+
+    let Command::Serve(args) = cli.command else {
+        panic!("expected serve args");
+    };
+    assert!(args.runtime.disable_log_stats);
+
+    let config = args.to_managed_engine_config(5555);
+    assert_eq!(config.python_args, vec!["--disable-log-stats"]);
+}
+
 #[test]
 fn serve_args_auto_forward_python_multi_char_alias_without_separator() {
    let cli = Cli::try_parse_from(["vllm-rs", "serve", "Qwen/Qwen3-0.6B", "-tp", "2"]).unwrap();
@@ -72,6 +72,8 @@ impl ManagedEngineArgs {
        model: String,
        max_model_len: Option<u32>,
        language_model_only: bool,
+        disable_log_stats: bool,
+        shutdown_timeout: u64,
        handshake_port: u16,
    ) -> ManagedEngineConfig {
        let mut python_args = self.python_args;
@@ -83,6 +85,15 @@ impl ManagedEngineArgs {
        if language_model_only {
            python_args.push("--language-model-only".to_string());
        }
+        if disable_log_stats {
+            python_args.push("--disable-log-stats".to_string());
+        }
+        // we must pass through shutdown_timeout to the engine,
+        // otherwise inflight requests get aborted on shutdown
+        if shutdown_timeout > 0 {
+            python_args.push("--shutdown-timeout".to_string());
+            python_args.push(shutdown_timeout.to_string());
+        }
        if let Some(data_parallel_size_local) = self.data_parallel_size_local {
            python_args.push("--data-parallel-size-local".to_string());
            python_args.push(data_parallel_size_local.to_string());
@@ -1,6 +1,7 @@
 use axum::Json;
 use axum::http::StatusCode;
 use axum::response::{IntoResponse, Response};
+use thiserror_ext::AsReport as _;
 use thiserror_ext::{Construct, Macro};

 use crate::routes::openai::utils::types::{ErrorDetail, ErrorResponse};
@@ -72,3 +73,84 @@ impl IntoResponse for ApiError {
        (self.status_code(), Json(self.to_error_response())).into_response()
    }
 }
+
+/// Classify a text-pipeline submit failure: tokenized-prompt validation
+/// failures (the prompt is too long for the model, or empty after
+/// tokenization) are the client's fault and map to HTTP 400, mirroring the
+/// Python frontend. Everything else stays an internal 500.
+pub fn text_submit_error(context: &'static str, error: vllm_text::Error) -> ApiError {
+    if is_prompt_validation_error(&error) {
+        return invalid_request!("{error}");
+    }
+    server_error!("{}: {}", context, error.to_report_string())
+}
+
+/// Like [`text_submit_error`], for the chat pipeline (which both wraps the
+/// text errors and raises its own prompt-length variant).
+pub fn chat_submit_error(context: &'static str, error: vllm_chat::Error) -> ApiError {
+    match &error {
+        vllm_chat::Error::PromptTooLong { .. } => invalid_request!("{error}"),
+        vllm_chat::Error::Text(text_error) if is_prompt_validation_error(text_error) => {
+            invalid_request!("{error}")
+        }
+        _ => server_error!("{}: {}", context, error.to_report_string()),
+    }
+}
+
+fn is_prompt_validation_error(error: &vllm_text::Error) -> bool {
+    matches!(
+        error,
+        vllm_text::Error::PromptTooLong { .. }
+            | vllm_text::Error::EmptyPromptTokenIds { .. }
+            // An empty tokenized prompt detected later, at request prepare
+            // time, surfaces through the transparent Llm wrapper.
+            | vllm_text::Error::Llm(vllm_llm::Error::EmptyPromptTokenIds { .. })
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn prompt_too_long_maps_to_invalid_request() {
+        let error = vllm_text::Error::PromptTooLong {
+            max_model_len: 8192,
+            prompt_len: 9000,
+        };
+        let api_error = text_submit_error("failed to submit completion request", error);
+        assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);
+        let response = api_error.to_error_response();
+        assert_eq!(response.error.error_type, "invalid_request_error");
+        assert!(response.error.message.contains("8192"));
+        assert!(response.error.message.contains("9000"));
+    }
+
+    #[test]
+    fn chat_wrapped_prompt_too_long_maps_to_invalid_request() {
+        let error = vllm_chat::Error::Text(vllm_text::Error::PromptTooLong {
+            max_model_len: 8192,
+            prompt_len: 9000,
+        });
+        let api_error = chat_submit_error("failed to submit chat request", error);
+        assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);
+    }
+
+    #[test]
+    fn llm_wrapped_empty_prompt_maps_to_invalid_request() {
+        let error = vllm_text::Error::Llm(vllm_llm::Error::EmptyPromptTokenIds {
+            request_id: "req-1".to_string(),
+        });
+        let api_error = text_submit_error("failed to submit completion request", error);
+        assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);
+    }
+
+    #[test]
+    fn other_submit_errors_stay_internal() {
+        let error = vllm_text::Error::Tokenizer("backend exploded".to_string());
+        let api_error = text_submit_error("failed to submit completion request", error);
+        assert_eq!(api_error.status_code(), StatusCode::INTERNAL_SERVER_ERROR);
+        let response = api_error.to_error_response();
+        assert!(response.error.message.starts_with("failed to submit completion request:"));
+    }
+}
@@ -28,7 +28,7 @@ use self::types::{
    GenerateResponseStreamChoice, GenerateStreamResponse,
 };
 use crate::config::ApiServerOptions;
-use crate::error::{ApiError, bail_server_error, server_error};
+use crate::error::{ApiError, bail_server_error, server_error, text_submit_error};
 use crate::routes::openai::utils::logprobs::clamp_logprob;
 use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb, Usage};
 use crate::routes::openai::utils::validated_json::ValidatedJson;
@@ -65,11 +65,8 @@ pub async fn generate(
    {
        Ok(stream) => stream,
        Err(error) => {
-            return server_error!(
-                "failed to submit raw generate request: {}",
-                error.to_report_string()
-            )
-            .into_response();
+            return text_submit_error("failed to submit raw generate request", error)
+                .into_response();
        }
    };

@@ -25,7 +25,7 @@ use vllm_engine_core_client::protocol::StopReason;

 use self::convert::{ResponseOptions, prepare_chat_request};
 use crate::config::ApiServerOptions;
-use crate::error::{ApiError, bail_server_error, server_error};
+use crate::error::{ApiError, bail_server_error, chat_submit_error, server_error};
 use crate::routes::openai::chat_completions::types::{
    AssistantRole, ChatCompletionChoice, ChatCompletionMessage, ChatCompletionRequest,
    ChatCompletionResponse, ChatCompletionStreamChoice, ChatCompletionStreamResponse,
@@ -37,6 +37,7 @@ use crate::routes::openai::utils::logprobs::{
 use crate::routes::openai::utils::types::{
    ChatLogProbs, FunctionCallDelta, FunctionCallResponse, ToolCall, ToolCallDelta, Usage,
 };
+use crate::routes::openai::utils::usage::ContinuousUsage;
 use crate::routes::openai::utils::validated_json::ValidatedJson;
 use crate::state::AppState;
 use crate::utils::{resolve_request_context, unix_timestamp};
@@ -76,11 +77,7 @@ pub async fn chat_completions(
        match state.chat.chat(prepared.chat_request).instrument(request_span.clone()).await {
            Ok(stream) => stream,
            Err(error) => {
-                return server_error!(
-                    "failed to submit chat request: {}",
-                    error.to_report_string()
-                )
-                .into_response();
+                return chat_submit_error("failed to submit chat request", error).into_response();
            }
        };

@@ -129,6 +126,8 @@ async fn collect_chat_completion(
    ResponseOptions {
        // Ignored: non-streaming responses always include usage.
        include_usage: _,
+        // Ignored: non-streaming responses are collected before usage is attached.
+        include_continuous_usage: _,
        requested_logprobs,
        include_prompt_logprobs,
        include_reasoning,
@@ -249,6 +248,7 @@ async fn chat_completion_chunk_stream(
    }: ApiServerOptions,
    ResponseOptions {
        include_usage,
+        include_continuous_usage,
        requested_logprobs,
        // Ignored: chat streaming prompt logprobs are rejected for Python parity.
        include_prompt_logprobs: _,
@@ -265,33 +265,47 @@ async fn chat_completion_chunk_stream(
    // starts or ends, omit its token metadata as well as its visible delta.
    let mut inside_hidden_reasoning = false;
    let mut suppress_current_update_metadata = false;
+    let mut continuous_usage = ContinuousUsage::default();
+
+    /// Yield a chunk with optional continuous usage attached.
+    macro_rules! yield_chunk {
+        ($chunk:expr) => {{
+            let mut chunk = $chunk;
+            if include_continuous_usage {
+                chunk.usage = Some(continuous_usage.to_usage());
+            }
+            y.yield_ok(chunk).await;
+        }};
+    }

    // If the client requested logprobs or token_ids, we need to buffer chunks until
    // we receive the separate `LogprobsDelta` event, so that we can emit one
    // combined chunk with both the semantic delta and its per-update metadata.
-    let mut pending_chunk =
-        (requested_logprobs || return_token_ids).then(PendingChatChunk::default);
+    // Continuous usage also buffers so the token count from `LogprobsDelta` can
+    // be attached to the matching semantic chunk.
+    let mut pending_chunk = (requested_logprobs || return_token_ids || include_continuous_usage)
+        .then(PendingChatChunk::default);

    while let Some(next) = stream.next().await {
        match next {
            Ok(ChatEvent::Start {
                prompt_token_ids, ..
            }) => {
+                continuous_usage.set_prompt_tokens(prompt_token_ids.len());
                let mut chunk = start_chunk(&request_id, &response_model, created);
                if return_token_ids {
                    chunk.prompt_token_ids = Some(prompt_token_ids.to_vec());
                }
-                y.yield_ok(chunk).await;
+                yield_chunk!(chunk);
                // When echo=true, emit the last assistant message content as a delta chunk.
                if let Some(echo_text) = &echo {
-                    y.yield_ok(block_delta_chunk(
+                    yield_chunk!(block_delta_chunk(
                        &request_id,
                        &response_model,
                        created,
                        AssistantBlockKind::Text,
                        echo_text.clone(),
-                    ))
-                    .await;
+                    ));
                }
            }
            Ok(ChatEvent::BlockDelta { kind, delta, .. }) => {
@@ -301,14 +315,13 @@ async fn chat_completion_chunk_stream(
                    if let Some(pending_chunk) = pending_chunk.as_mut() {
                        pending_chunk.push_block_delta(kind, delta);
                    } else {
-                        y.yield_ok(block_delta_chunk(
+                        yield_chunk!(block_delta_chunk(
                            &request_id,
                            &response_model,
                            created,
                            kind,
                            delta,
-                        ))
-                        .await;
+                        ));
                    }
                } else {
                    suppress_current_update_metadata = true;
@@ -318,6 +331,8 @@ async fn chat_completion_chunk_stream(
                logprobs,
                token_ids,
            }) => {
+                let delta_token_count = token_ids.len();
+                continuous_usage.add_output_tokens(delta_token_count);
                let include_metadata =
                    !suppress_current_update_metadata && !inside_hidden_reasoning;
                suppress_current_update_metadata = false;
@@ -339,16 +354,15 @@ async fn chat_completion_chunk_stream(
                    if let Some(chunk) =
                        pending_chunk.take_chunk(&request_id, &response_model, created)
                    {
-                        y.yield_ok(chunk).await;
+                        yield_chunk!(chunk);
                    }
                } else if let Some(logprobs) = openai_logprobs {
-                    y.yield_ok(logprobs_only_chunk(
+                    yield_chunk!(logprobs_only_chunk(
                        &request_id,
                        &response_model,
                        created,
                        logprobs,
-                    ))
-                    .await;
+                    ));
                }
            }
            Ok(ChatEvent::BlockStart { kind, .. }) => {
@@ -376,15 +390,14 @@ async fn chat_completion_chunk_stream(
                if let Some(pending_chunk) = pending_chunk.as_mut() {
                    pending_chunk.push_tool_call_start(tool_index, id, name);
                } else {
-                    y.yield_ok(tool_call_start_chunk(
+                    yield_chunk!(tool_call_start_chunk(
                        &request_id,
                        &response_model,
                        created,
                        tool_index,
                        id,
                        name,
-                    ))
-                    .await;
+                    ));
                }
            }
            Ok(ChatEvent::ToolCallArgumentsDelta { index, delta }) => {
@@ -392,21 +405,20 @@ async fn chat_completion_chunk_stream(
                if let Some(pending_chunk) = pending_chunk.as_mut() {
                    pending_chunk.push_tool_call_arguments(tool_index, delta);
                } else {
-                    y.yield_ok(tool_call_arguments_chunk(
+                    yield_chunk!(tool_call_arguments_chunk(
                        &request_id,
                        &response_model,
                        created,
                        tool_index,
                        delta,
-                    ))
-                    .await;
+                    ));
                }
            }
            Ok(ChatEvent::ToolCallEnd { .. }) => {
                debug!("ending current tool call");
            }
            Ok(ChatEvent::Done {
-                usage,
+                usage: final_usage,
                finish_reason,
                ..
            }) => {
@@ -414,18 +426,23 @@ async fn chat_completion_chunk_stream(
                    info!(
                        stream = true,
                        model = %response_model,
-                        prompt_tokens = usage.prompt_token_count,
-                        output_tokens = usage.output_token_count,
+                        prompt_tokens = final_usage.prompt_token_count,
+                        output_tokens = final_usage.output_token_count,
                        finish_reason = finish_reason.as_str(),
                        "chat completion finished"
                    );
                }

+                continuous_usage.set_final_counts(
+                    final_usage.prompt_token_count,
+                    final_usage.output_token_count,
+                );
+
                if let Some(pending_chunk) = pending_chunk.as_mut()
                    && let Some(chunk) =
                        pending_chunk.take_chunk(&request_id, &response_model, created)
                {
-                    y.yield_ok(chunk).await;
+                    yield_chunk!(chunk);
                }

                match final_chunk(
@@ -435,7 +452,7 @@ async fn chat_completion_chunk_stream(
                    finish_reason,
                    saw_tool_calls,
                ) {
-                    Ok(chunk) => y.yield_ok(chunk).await,
+                    Ok(chunk) => yield_chunk!(chunk),
                    Err(error) => {
                        error!(
                            error = %error.to_error_response().error.message,
@@ -450,7 +467,7 @@ async fn chat_completion_chunk_stream(
                        &request_id,
                        &response_model,
                        created,
-                        Usage::from_token_usage(usage, enable_prompt_tokens_details),
+                        Usage::from_token_usage(final_usage, enable_prompt_tokens_details),
                    ))
                    .await;
                }
@@ -33,6 +33,8 @@ pub(super) struct PreparedRequest {
 pub(super) struct ResponseOptions {
    /// Whether the caller asked for the final streamed usage chunk.
    pub include_usage: bool,
+    /// Whether every streamed chunk should carry cumulative usage.
+    pub include_continuous_usage: bool,
    /// Whether the caller requested output logprobs on chat choices.
    pub requested_logprobs: bool,
    /// Whether the caller requested top-level prompt logprobs.
@@ -82,6 +84,12 @@ pub(super) fn prepare_chat_request(
    let include_usage = (request.stream_options.as_ref())
        .and_then(|options| options.include_usage)
        .unwrap_or(false);
+    let include_continuous_usage = include_usage
+        && request
+            .stream_options
+            .as_ref()
+            .and_then(|options| options.continuous_usage_stats)
+            .unwrap_or(false);
    let requested_logprobs = request.logprobs;

    // Auto-enable prompt logprobs for non-streaming echo, matching Python vLLM's
@@ -154,6 +162,7 @@ pub(super) fn prepare_chat_request(
        response_model,
        options: ResponseOptions {
            include_usage,
+            include_continuous_usage,
            requested_logprobs,
            include_prompt_logprobs,
            include_reasoning,
@@ -375,8 +384,8 @@ mod tests {
        AssistantRole, ChatCompletionMessage, ChatCompletionRequest,
    };
    use crate::routes::openai::utils::types::{
-        ChatMessage, ContentPart, Function, FunctionCallResponse, ImageUrl, MessageContent, Tool,
-        ToolCall, ToolChoice, ToolChoiceValue, VideoUrl,
+        ChatMessage, ContentPart, Function, FunctionCallResponse, ImageUrl, MessageContent,
+        StreamOptions, Tool, ToolCall, ToolChoice, ToolChoiceValue, VideoUrl,
    };
    use crate::utils::{ResolvedRequestContext, resolve_request_context};

@@ -456,6 +465,46 @@ mod tests {
        assert_eq!(prepared.chat_request.tool_choice, ChatToolChoice::Auto);
    }

+    #[test]
+    fn prepare_chat_request_maps_stream_usage_and_token_format_options() {
+        let mut request = base_request();
+        request.return_tokens_as_token_ids = Some(true);
+        request.stream_options = Some(StreamOptions {
+            include_usage: Some(true),
+            continuous_usage_stats: Some(true),
+        });
+
+        let prepared = prepare_chat_request(
+            request,
+            &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
+            ResolvedRequestContext::default(),
+        )
+        .expect("request is valid");
+
+        assert!(prepared.options.include_usage);
+        assert!(prepared.options.include_continuous_usage);
+        assert!(prepared.options.return_tokens_as_token_ids);
+    }
+
+    #[test]
+    fn prepare_chat_request_gates_continuous_usage_on_include_usage() {
+        let mut request = base_request();
+        request.stream_options = Some(StreamOptions {
+            include_usage: None,
+            continuous_usage_stats: Some(true),
+        });
+
+        let prepared = prepare_chat_request(
+            request,
+            &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
+            ResolvedRequestContext::default(),
+        )
+        .expect("request is valid");
+
+        assert!(!prepared.options.include_usage);
+        assert!(!prepared.options.include_continuous_usage);
+    }
+
    #[test]
    fn prepare_chat_request_keeps_optional_sampling_fields_unset() {
        let prepared = prepare_chat_request(
@@ -137,15 +137,6 @@ pub(super) fn validate_request_compat(
        "repetition_detection is not supported.",
    )?;

-    if let Some(options) = &request.stream_options
-        && options.continuous_usage_stats.is_some()
-    {
-        bail_invalid_request!(
-            param = "stream_options",
-            "continuous_usage_stats is not supported."
-        );
-    }
-
    Ok(())
 }

@@ -25,12 +25,13 @@ use super::utils::logprobs::{
 };
 use super::utils::types::Usage;
 use crate::config::ApiServerOptions;
-use crate::error::{ApiError, bail_server_error, server_error};
+use crate::error::{ApiError, bail_server_error, server_error, text_submit_error};
 use crate::routes::openai::completions::types::{
    CompletionChoice, CompletionRequest, CompletionResponse, CompletionSseChunk,
    CompletionStreamChoice, CompletionStreamResponse,
 };
 use crate::routes::openai::utils::types::LogProbs;
+use crate::routes::openai::utils::usage::ContinuousUsage;
 use crate::routes::openai::utils::validated_json::ValidatedJson;
 use crate::state::AppState;
 use crate::utils::{resolve_request_context, unix_timestamp};
@@ -74,11 +75,7 @@ pub async fn completions(
    {
        Ok(stream) => stream,
        Err(error) => {
-            return server_error!(
-                "failed to submit completion request: {}",
-                error.to_report_string()
-            )
-            .into_response();
+            return text_submit_error("failed to submit completion request", error).into_response();
        }
    };

@@ -127,6 +124,8 @@ async fn collect_completion(
    ResponseOptions {
        // Ignored: non-streaming responses always include usage.
        include_usage: _,
+        // Ignored: non-streaming responses are collected before usage is attached.
+        include_continuous_usage: _,
        echo,
        requested_logprobs,
        include_prompt_logprobs,
@@ -218,6 +217,7 @@ async fn completion_chunk_stream(
    }: ApiServerOptions,
    ResponseOptions {
        include_usage,
+        include_continuous_usage,
        echo,
        requested_logprobs,
        // Ignored: streaming prompt logprobs are rejected for Python parity.
@@ -230,6 +230,18 @@ async fn completion_chunk_stream(
    pin_mut!(stream);
    let mut visible_text_len = 0_u32;
    let mut first_chunk = true;
+    let mut continuous_usage = ContinuousUsage::default();
+
+    /// Yield a chunk with optional continuous usage attached.
+    macro_rules! yield_chunk {
+        ($chunk:expr) => {{
+            let mut chunk = $chunk;
+            if include_continuous_usage {
+                chunk.usage = Some(continuous_usage.to_usage());
+            }
+            y.yield_ok(CompletionSseChunk::Chunk(chunk)).await;
+        }};
+    }

    while let Some(next) = stream.next().await {
        match next {
@@ -237,6 +249,7 @@ async fn completion_chunk_stream(
                prompt_token_ids, ..
            }) => {
                debug!("completion stream started");
+                continuous_usage.set_prompt_tokens(prompt_token_ids.len());
                if let Some(prompt) = echo.as_ref() {
                    visible_text_len = text_len(prompt);
                    let mut chunk =
@@ -247,7 +260,7 @@ async fn completion_chunk_stream(
                        }
                        first_chunk = false;
                    }
-                    y.yield_ok(CompletionSseChunk::Chunk(chunk)).await;
+                    yield_chunk!(chunk);
                } else if return_token_ids {
                    // Emit a chunk with prompt_token_ids in the first streaming response
                    let mut chunk =
@@ -256,7 +269,7 @@ async fn completion_chunk_stream(
                        choice.prompt_token_ids = Some(prompt_token_ids.to_vec());
                    }
                    first_chunk = false;
-                    y.yield_ok(CompletionSseChunk::Chunk(chunk)).await;
+                    yield_chunk!(chunk);
                }
            }
            Ok(DecodedTextEvent::TextDelta {
@@ -281,10 +294,12 @@ async fn completion_chunk_stream(
                    None
                };
                let mut chunk = delta_chunk(&request_id, &response_model, created, delta, logprobs);
+                let delta_token_count = token_ids.len();
+                continuous_usage.add_output_tokens(delta_token_count);
                if return_token_ids && let Some(choice) = chunk.choices.first_mut() {
                    choice.token_ids = Some(token_ids);
                }
-                y.yield_ok(CompletionSseChunk::Chunk(chunk)).await;
+                yield_chunk!(chunk);
                visible_text_len = visible_text_len.saturating_add(delta_text_len);

                if let Some(finished) = finished {
@@ -298,13 +313,17 @@ async fn completion_chunk_stream(
                            "completion finished"
                        );
                    }
-                    y.yield_ok(CompletionSseChunk::Chunk(final_chunk(
+                    continuous_usage.set_final_counts(
+                        finished.usage.prompt_token_count,
+                        finished.usage.output_token_count,
+                    );
+                    let final_chunk = final_chunk(
                        &request_id,
                        &response_model,
                        created,
                        finished.finish_reason,
-                    )?))
-                    .await;
+                    )?;
+                    yield_chunk!(final_chunk);

                    if include_usage {
                        y.yield_ok(CompletionSseChunk::Usage(usage_chunk(
@@ -25,6 +25,8 @@ pub(super) struct PreparedRequest {
 pub(super) struct ResponseOptions {
    /// Whether the caller asked for the final streamed usage chunk.
    pub include_usage: bool,
+    /// Whether every streamed chunk should carry cumulative usage.
+    pub include_continuous_usage: bool,
    /// Original text prompt that should be echoed back northbound when
    /// `echo=true`.
    pub echo: Option<String>,
@@ -74,6 +76,12 @@ pub(super) fn prepare_completion_request(
    let include_usage = (request.stream_options.as_ref())
        .and_then(|options| options.include_usage)
        .unwrap_or(false);
+    let include_continuous_usage = include_usage
+        && request
+            .stream_options
+            .as_ref()
+            .and_then(|options| options.continuous_usage_stats)
+            .unwrap_or(false);
    let include_prompt_logprobs = prompt_logprobs.is_some();
    let echo = request.echo.then(|| request.prompt.as_text().cloned()).flatten();

@@ -129,6 +137,7 @@ pub(super) fn prepare_completion_request(
        response_model,
        options: ResponseOptions {
            include_usage,
+            include_continuous_usage,
            echo,
            requested_logprobs: request.logprobs,
            include_prompt_logprobs,
@@ -247,6 +256,55 @@ mod tests {
        assert!(!prepared.text_request.decode_options.skip_special_tokens);
    }

+    #[test]
+    fn prepare_completion_request_maps_stream_usage_and_token_format_options() {
+        let request: CompletionRequest = serde_json::from_value(json!({
+            "model": "Qwen/Qwen1.5-0.5B-Chat",
+            "prompt": "hello",
+            "stream": true,
+            "stream_options": {
+                "include_usage": true,
+                "continuous_usage_stats": true
+            },
+            "return_tokens_as_token_ids": true
+        }))
+        .expect("parse request");
+
+        let prepared = prepare_completion_request(
+            request,
+            &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
+            ResolvedRequestContext::default(),
+        )
+        .expect("prepare");
+
+        assert!(prepared.options.include_usage);
+        assert!(prepared.options.include_continuous_usage);
+        assert!(prepared.options.return_tokens_as_token_ids);
+    }
+
+    #[test]
+    fn prepare_completion_request_gates_continuous_usage_on_include_usage() {
+        let request: CompletionRequest = serde_json::from_value(json!({
+            "model": "Qwen/Qwen1.5-0.5B-Chat",
+            "prompt": "hello",
+            "stream": true,
+            "stream_options": {
+                "continuous_usage_stats": true
+            }
+        }))
+        .expect("parse request");
+
+        let prepared = prepare_completion_request(
+            request,
+            &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
+            ResolvedRequestContext::default(),
+        )
+        .expect("prepare");
+
+        assert!(!prepared.options.include_usage);
+        assert!(!prepared.options.include_continuous_usage);
+    }
+
    #[test]
    fn prepare_completion_request_accepts_text_echo() {
        let request: CompletionRequest = serde_json::from_value(json!({
@@ -95,15 +95,6 @@ pub(super) fn validate_request_compat(
        );
    }

-    if let Some(options) = &request.stream_options
-        && options.continuous_usage_stats.is_some()
-    {
-        bail_invalid_request!(
-            param = "stream_options",
-            "continuous_usage_stats is not supported."
-        );
-    }
-
    Ok(())
 }

@@ -2,4 +2,5 @@ pub mod logprobs;
 pub mod structured_outputs;
 pub mod token_ids;
 pub mod types;
+pub mod usage;
 pub mod validated_json;
@@ -0,0 +1,35 @@
+use super::types::Usage;
+
+/// Tracks cumulative token counts for OpenAI streaming chunks.
+///
+/// This helper is intentionally only a counter. Callers decide whether to
+/// attach `counts()` to each streamed data chunk, while final usage-only chunks
+/// should still be built from the authoritative terminal `TokenUsage`.
+#[derive(Debug, Clone, Default)]
+pub(crate) struct ContinuousUsage {
+    prompt_tokens: usize,
+    output_tokens: usize,
+}
+
+impl ContinuousUsage {
+    /// Record the prompt-token count reported when a stream starts.
+    pub(crate) fn set_prompt_tokens(&mut self, prompt_tokens: usize) {
+        self.prompt_tokens = prompt_tokens;
+    }
+
+    /// Add newly decoded output tokens to the running completion count.
+    pub(crate) fn add_output_tokens(&mut self, output_tokens: usize) {
+        self.output_tokens = self.output_tokens.saturating_add(output_tokens);
+    }
+
+    /// Replace the running counts with the final counts reported by generation.
+    pub(crate) fn set_final_counts(&mut self, prompt_tokens: usize, output_tokens: usize) {
+        self.prompt_tokens = prompt_tokens;
+        self.output_tokens = output_tokens;
+    }
+
+    /// Build a streaming usage snapshot without prompt cache details.
+    pub(crate) fn to_usage(&self) -> Usage {
+        Usage::from_counts(self.prompt_tokens, self.output_tokens, None)
+    }
+}
@@ -151,6 +151,14 @@ fn sse_data_payloads(text: &str) -> Vec<&str> {
    text.lines().filter_map(|line| line.strip_prefix("data: ")).collect()
 }

+fn sse_json_payloads(text: &str) -> Vec<serde_json::Value> {
+    sse_data_payloads(text)
+        .into_iter()
+        .filter(|payload| *payload != "[DONE]")
+        .map(|payload| serde_json::from_str(payload).expect("sse json payload"))
+        .collect()
+}
+
 type TestFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;

 fn boxed_test_future<'a>(future: impl Future<Output = ()> + Send + 'a) -> TestFuture<'a> {
@@ -2341,6 +2349,60 @@ async fn include_usage_adds_final_usage_chunk_before_done() {
    assert_eq!(usage_chunk["usage"]["total_tokens"], 25);
 }

+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn stream_continuous_usage_stats_adds_usage_to_chat_chunks() {
+    let (app, engine_task) = test_app_with_stream_output_specs(default_stream_output_specs()).await;
+    let response = app
+        .clone()
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/v1/chat/completions")
+                .header("content-type", "application/json")
+                .body(Body::from(
+                    json!({
+                        "model": "Qwen/Qwen1.5-0.5B-Chat",
+                        "stream": true,
+                        "stream_options": {
+                            "include_usage": true,
+                            "continuous_usage_stats": true
+                        },
+                        "messages": [{"role": "user", "content": "hello"}]
+                    })
+                    .to_string(),
+                ))
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+
+    assert_eq!(response.status(), StatusCode::OK);
+
+    let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
+    engine_task.await.expect("mock engine task");
+    let text = String::from_utf8(body.to_vec()).expect("utf8 body");
+    let payloads = sse_json_payloads(&text);
+
+    assert!(
+        payloads.iter().all(|payload| payload.get("usage").is_some()),
+        "{text}"
+    );
+    assert!(
+        payloads.iter().any(|payload| {
+            payload["choices"].as_array().is_some_and(|choices| !choices.is_empty())
+                && payload["usage"]["completion_tokens"] == json!(1)
+        }),
+        "{text}"
+    );
+    let usage_chunk = payloads
+        .iter()
+        .find(|payload| payload["choices"] == json!([]))
+        .expect("final usage chunk");
+    assert_eq!(usage_chunk["usage"]["prompt_tokens"], 22);
+    assert_eq!(usage_chunk["usage"]["completion_tokens"], 3);
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn stream_without_include_usage_keeps_existing_shape() {
@@ -3434,6 +3496,59 @@ async fn completions_happy_path_returns_sse_stream() {
    assert_eq!(usage_chunk["usage"]["completion_tokens"], 3);
 }

+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn completions_stream_continuous_usage_stats_adds_usage_to_chunks() {
+    let (app, engine_task) = test_app_with_stream_output_specs(default_stream_output_specs()).await;
+    let response = app
+        .clone()
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/v1/completions")
+                .header("content-type", "application/json")
+                .body(Body::from(
+                    json!({
+                        "model": "Qwen/Qwen1.5-0.5B-Chat",
+                        "prompt": "hello",
+                        "stream": true,
+                        "stream_options": {
+                            "include_usage": true,
+                            "continuous_usage_stats": true
+                        }
+                    })
+                    .to_string(),
+                ))
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+
+    assert_eq!(response.status(), StatusCode::OK);
+
+    let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
+    engine_task.await.expect("mock engine task");
+    let text = String::from_utf8(body.to_vec()).expect("utf8 body");
+    let payloads = sse_json_payloads(&text);
+
+    assert!(
+        payloads.iter().all(|payload| payload.get("usage").is_some()),
+        "{text}"
+    );
+    assert!(
+        payloads.iter().any(|payload| {
+            payload["choices"].as_array().is_some_and(|choices| !choices.is_empty())
+                && payload["usage"]["completion_tokens"] == json!(1)
+        }),
+        "{text}"
+    );
+    let usage_chunk = payloads
+        .iter()
+        .find(|payload| payload["choices"] == json!([]))
+        .expect("final usage chunk");
+    assert_eq!(usage_chunk["usage"]["completion_tokens"], 3);
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn completions_echo_stream_emits_separate_prompt_chunk() {
@@ -1229,7 +1229,7 @@ setup(
        # NOTE: When updating helion version, also update CI files:
        #   - .buildkite/test_areas/kernels.yaml
        #   - .buildkite/test-amd.yaml
-        "helion": ["helion==1.0.0"],
+        "helion": ["helion==1.1.0"],
        # Optional deps for gRPC server (vllm serve --grpc)
        "grpc": ["smg-grpc-servicer[vllm] >= 0.5.2"],
        # Optional deps for OpenTelemetry tracing
@@ -0,0 +1,200 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+from pathlib import Path
+from typing import Protocol, cast
+
+import numpy as np
+import pytest
+import soundfile as sf
+
+import vllm.benchmarks.datasets.datasets as datasets_module
+import vllm.benchmarks.lib.endpoint_request_func as request_func_module
+from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+
+class _ReadableBinary(Protocol):
+    def read(self, size: int = -1) -> bytes: ...
+
+
+class _TokenizedPrompt:
+    def __init__(self, prompt: str) -> None:
+        self.input_ids = prompt.split()
+
+
+class _Tokenizer:
+    def __init__(self, name_or_path: str = "openai/whisper-large-v3") -> None:
+        self.name_or_path = name_or_path
+
+    def __call__(self, prompt: str) -> _TokenizedPrompt:
+        return _TokenizedPrompt(prompt)
+
+
+def _write_wav(path: Path, duration_s: float = 0.1, sample_rate: int = 16_000) -> None:
+    num_samples = int(duration_s * sample_rate)
+    sf.write(path, np.zeros(num_samples, dtype=np.float32), sample_rate)
+
+
+class _FakeFormData:
+    def __init__(self) -> None:
+        self.fields: list[tuple[str, object, dict[str, str]]] = []
+
+    def add_field(self, name: str, value: object, **kwargs: str) -> None:
+        self.fields.append((name, value, kwargs))
+
+
+class _FakeContent:
+    async def iter_any(self):
+        yield b'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n'
+        yield b'data: {"usage":{"completion_tokens":1}}\n\n'
+        yield b"data: [DONE]\n\n"
+
+
+class _FakeResponse:
+    def __init__(self) -> None:
+        self.status = 200
+        self.reason = "OK"
+        self.content = _FakeContent()
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        return False
+
+
+class _FakeSession:
+    def __init__(self) -> None:
+        self.uploaded_bytes: bytes | None = None
+        self.upload_filename: str | None = None
+        self.fields: list[tuple[str, object, dict[str, str]]] | None = None
+
+    def post(self, *, url: str, data: _FakeFormData, headers: dict[str, str]):
+        del url, headers
+        self.fields = list(data.fields)
+        _, file_obj, file_kwargs = self.fields[0]
+        file_obj = cast(_ReadableBinary, file_obj)
+        self.uploaded_bytes = file_obj.read()
+        self.upload_filename = file_kwargs.get("filename")
+        return _FakeResponse()
+
+
+def test_asr_dataset_sample_handles_local_audio_paths(tmp_path: Path) -> None:
+    audio_path = tmp_path / "earnings.wav"
+    _write_wav(audio_path, duration_s=0.1)
+
+    dataset = object.__new__(datasets_module.ASRDataset)
+    dataset.data = [
+        {
+            "audio": {
+                "path": str(audio_path),
+                "bytes": None,
+            },
+            "text": "quarterly earnings call",
+        }
+    ]
+
+    samples = dataset.sample(
+        tokenizer=_Tokenizer(),
+        num_requests=1,
+        output_len=32,
+        asr_min_audio_len_sec=0.0,
+        asr_max_audio_len_sec=1.0,
+    )
+
+    assert len(samples) == 1
+    assert samples[0].multi_modal_data == {"audio_path": str(audio_path)}
+    assert (
+        samples[0].prompt == "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>"
+    )
+
+
+def test_asr_dataset_sample_handles_embedded_audio_bytes(tmp_path: Path) -> None:
+    audio_path = tmp_path / "earnings.wav"
+    _write_wav(audio_path, duration_s=0.1)
+
+    dataset = object.__new__(datasets_module.ASRDataset)
+    dataset.data = [
+        {
+            "audio": {
+                "path": None,
+                "bytes": audio_path.read_bytes(),
+            },
+            "text": "quarterly earnings call",
+        }
+    ]
+
+    samples = dataset.sample(
+        tokenizer=_Tokenizer(),
+        num_requests=1,
+        output_len=32,
+        asr_min_audio_len_sec=0.0,
+        asr_max_audio_len_sec=1.0,
+    )
+
+    assert len(samples) == 1
+    assert isinstance(samples[0].multi_modal_data, dict)
+    audio, sample_rate = samples[0].multi_modal_data["audio"]
+    assert sample_rate == 16_000
+    assert isinstance(audio, np.ndarray)
+    assert audio.size > 0
+
+
+def test_async_request_openai_audio_handles_local_audio_paths(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    audio_path = tmp_path / "earnings.wav"
+    _write_wav(audio_path, duration_s=0.25)
+
+    monkeypatch.setattr(request_func_module.aiohttp, "FormData", _FakeFormData)
+    session = _FakeSession()
+    request_input = RequestFuncInput(
+        prompt="",
+        api_url="http://localhost:8000/v1/audio/transcriptions",
+        prompt_len=1,
+        output_len=32,
+        model="openai/whisper-large-v3",
+        multi_modal_content={"audio_path": str(audio_path)},
+    )
+
+    output = asyncio.run(
+        request_func_module.async_request_openai_audio(request_input, session)
+    )
+
+    assert session.upload_filename == audio_path.name
+    assert session.uploaded_bytes == audio_path.read_bytes()
+    assert output.success is True
+    assert output.generated_text == "hello"
+    assert output.output_tokens == 1
+    assert output.input_audio_duration == pytest.approx(0.25, abs=1e-2)
+
+
+def test_async_request_openai_audio_handles_decoded_audio_arrays(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(request_func_module.aiohttp, "FormData", _FakeFormData)
+    session = _FakeSession()
+    request_input = RequestFuncInput(
+        prompt="",
+        api_url="http://localhost:8000/v1/audio/transcriptions",
+        prompt_len=1,
+        output_len=32,
+        model="openai/whisper-large-v3",
+        multi_modal_content={
+            "audio": (np.zeros(1_600, dtype=np.float32), 16_000),
+        },
+    )
+
+    output = asyncio.run(
+        request_func_module.async_request_openai_audio(request_input, session)
+    )
+
+    assert session.upload_filename == "audio.wav"
+    assert session.uploaded_bytes is not None
+    assert output.success is True
+    assert output.generated_text == "hello"
@@ -124,8 +124,6 @@ TEXT_GENERATION_MODELS = {
    "EleutherAI/pythia-1.4b": PPTestSettings.fast(),
    "ibm/PowerLM-3b": PPTestSettings.fast(),
    "ibm/PowerMoE-3b": PPTestSettings.fast(),
-    # Uses Llama
-    # "internlm/internlm-chat-7b": PPTestSettings.fast(),
    "internlm/internlm2-chat-7b": PPTestSettings.fast(),
    "ai21labs/Jamba-tiny-dev": PPTestSettings.fast(),
    "pfnet/plamo-2-1b": PPTestSettings.fast(),
@@ -6,13 +6,29 @@ Tests the image source handling and tool_result content parsing in
 AnthropicServingMessages._convert_anthropic_to_openai_request().

 Also covers extended-thinking edge cases such as ``redacted_thinking``
-blocks echoed back by Anthropic clients.
+blocks echoed back by Anthropic clients, and streaming conversion in
+``message_stream_converter``.
 """

+import json
+from unittest.mock import MagicMock
+
+import pytest
+
 from vllm.entrypoints.anthropic.protocol import (
    AnthropicMessagesRequest,
 )
 from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    ChatCompletionResponseStreamChoice,
+    ChatCompletionStreamResponse,
+)
+from vllm.entrypoints.openai.engine.protocol import (
+    DeltaFunctionCall,
+    DeltaMessage,
+    DeltaToolCall,
+    UsageInfo,
+)

 _convert = AnthropicServingMessages._convert_anthropic_to_openai_request
 _img_url = AnthropicServingMessages._convert_image_source_to_url
@@ -775,3 +791,208 @@ class TestInlineSystemMessageInMessagesArray:
        assert result.messages[0]["role"] == "system"
        assert result.messages[0]["content"] == "Top-level prompt.Inline hint."
        assert result.messages[1]["role"] == "user"
+
+
+# ======================================================================
+# Streaming conversion: message_stream_converter
+# ======================================================================
+
+
+def _make_stream_converter():
+    obj = MagicMock(spec=AnthropicServingMessages)
+    obj.stop_reason_map = {
+        "stop": "end_turn",
+        "length": "max_tokens",
+        "tool_calls": "tool_use",
+    }
+    obj.message_stream_converter = (
+        AnthropicServingMessages.message_stream_converter.__get__(obj)
+    )
+    return obj
+
+
+def _parse_sse_events(raw_events: list[str]) -> list[tuple[str, dict]]:
+    results = []
+    for raw in raw_events:
+        headers = dict(
+            line.split(": ", 1) for line in raw.strip().split("\n") if ": " in line
+        )
+        if "event" in headers and "data" in headers:
+            results.append((headers["event"], json.loads(headers["data"])))
+    return results
+
+
+def _make_stream_chunk(
+    *,
+    delta: DeltaMessage | None = None,
+    finish_reason: str | None = None,
+    choices: list[ChatCompletionResponseStreamChoice] | None = None,
+    usage: UsageInfo | None = None,
+) -> str:
+    if choices is None:
+        choices = [
+            ChatCompletionResponseStreamChoice(
+                index=0,
+                delta=delta or DeltaMessage(),
+                finish_reason=finish_reason,
+            )
+        ]
+    chunk = ChatCompletionStreamResponse(
+        id="chatcmpl-test",
+        created=0,
+        model="test-model",
+        choices=choices,
+        usage=usage,
+    )
+    return f"data: {chunk.model_dump_json()}"
+
+
+def _tc(*, args, id=None, name=None):
+    return DeltaToolCall(
+        index=0,
+        id=id,
+        function=DeltaFunctionCall(name=name, arguments=args),
+    )
+
+
+class TestMessageStreamConverterToolUseContentBuffering:
+    """Regression test for tool_use arguments being silently dropped.
+
+    With speculative decoding or multi-token prediction, a single delta
+    can carry both the final tool_call argument fragment and trailing
+    content.
+    """
+
+    @pytest.mark.asyncio
+    async def test_tool_use_args_not_dropped_when_content_in_same_chunk(
+        self,
+    ):
+        async def sse_input():
+            yield _make_stream_chunk(
+                delta=DeltaMessage(role="assistant"),
+                usage=UsageInfo(prompt_tokens=10, total_tokens=10),
+            )
+            yield _make_stream_chunk(
+                delta=DeltaMessage(
+                    tool_calls=[
+                        _tc(id="call_abc123", name="read_file", args=""),
+                    ]
+                )
+            )
+            yield _make_stream_chunk(
+                delta=DeltaMessage(
+                    tool_calls=[
+                        _tc(args='{"path":"/tmp/f"'),
+                    ]
+                )
+            )
+            # BUG TRIGGER: final tool_call args and trailing content in
+            # one delta, as happens with spec decoding / multi-token
+            # prediction where multiple tokens land in a single chunk.
+            yield _make_stream_chunk(
+                delta=DeltaMessage(
+                    content="\nOkay",
+                    tool_calls=[_tc(args="}")],
+                )
+            )
+            yield _make_stream_chunk(finish_reason="tool_calls")
+            yield _make_stream_chunk(
+                choices=[],
+                usage=UsageInfo(
+                    prompt_tokens=10,
+                    total_tokens=30,
+                    completion_tokens=20,
+                ),
+            )
+            yield "data: [DONE]"
+
+        converter = _make_stream_converter()
+        output = []
+        async for event in converter.message_stream_converter(sse_input()):
+            output.append(event)
+
+        events = _parse_sse_events(output)
+
+        assert events[0][0] == "message_start"
+
+        arg_fragments = [
+            data["delta"]["partial_json"]
+            for _, data in events
+            if data.get("delta", {}).get("type") == "input_json_delta"
+        ]
+        full_args = "".join(arg_fragments)
+        assert full_args == '{"path":"/tmp/f"}'
+
+        text_deltas = [
+            data["delta"]["text"]
+            for _, data in events
+            if data.get("delta", {}).get("type") == "text_delta"
+        ]
+        assert text_deltas == ["\nOkay"]
+
+        block_starts = [
+            (data["content_block"]["type"], data.get("index"))
+            for ev_type, data in events
+            if ev_type == "content_block_start"
+        ]
+        assert block_starts[0] == ("tool_use", 0)
+        assert block_starts[1] == ("text", 1)
+
+        msg_deltas = [data for ev_type, data in events if ev_type == "message_delta"]
+        assert msg_deltas[0]["delta"]["stop_reason"] == "tool_use"
+
+        assert events[-1][0] == "message_stop"
+
+    @pytest.mark.asyncio
+    async def test_buffered_content_flushed_on_done_without_usage_chunk(self):
+        """Content buffered during tool_use must be emitted even if the
+        stream jumps straight from finish_reason to [DONE], skipping the
+        empty-choices usage chunk."""
+
+        async def sse_input():
+            yield _make_stream_chunk(
+                delta=DeltaMessage(role="assistant"),
+                usage=UsageInfo(prompt_tokens=10, total_tokens=10),
+            )
+            yield _make_stream_chunk(
+                delta=DeltaMessage(
+                    tool_calls=[
+                        _tc(id="call_xyz", name="get_weather", args=""),
+                    ]
+                )
+            )
+            yield _make_stream_chunk(
+                delta=DeltaMessage(
+                    tool_calls=[_tc(args='{"city":"NYC"}')],
+                )
+            )
+            yield _make_stream_chunk(
+                delta=DeltaMessage(content="\nDone"),
+                finish_reason="tool_calls",
+            )
+            # No empty-choices usage chunk — go straight to [DONE].
+            yield "data: [DONE]"
+
+        converter = _make_stream_converter()
+        output = []
+        async for event in converter.message_stream_converter(sse_input()):
+            output.append(event)
+
+        events = _parse_sse_events(output)
+
+        text_deltas = [
+            data["delta"]["text"]
+            for _, data in events
+            if data.get("delta", {}).get("type") == "text_delta"
+        ]
+        assert text_deltas == ["\nDone"]
+
+        block_starts = [
+            data["content_block"]["type"]
+            for ev_type, data in events
+            if ev_type == "content_block_start"
+        ]
+        assert "tool_use" in block_starts
+        assert "text" in block_starts
+
+        assert events[-1][0] == "message_stop"
@@ -250,6 +250,7 @@ async def k2_client(k2_server):


@pytest.mark.asyncio
+@pytest.mark.skip(reason="Skipping Kimi K2 tool ID test")
@pytest.mark.parametrize("model_name", [MODEL_NAME])
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.parametrize("tool_choice", ["required"])
@@ -442,7 +443,7 @@ async def test_named_tool_use(
        if delta.role:
            assert delta.role == "assistant"
        assert delta.content is None or len(delta.content) == 0
-        if delta.tool_calls:
+        if delta.tool_calls and delta.tool_calls[0].function.arguments:
            output.append(delta.tool_calls[0].function.arguments)
        if chunk.choices[0].finish_reason is not None:
            finish_reason_count += 1
@@ -38,12 +38,12 @@ from vllm.entrypoints.serve.render.serving import OpenAIServingRender
 from vllm.exceptions import VLLMValidationError
 from vllm.inputs import TokensPrompt
 from vllm.outputs import CompletionOutput, RequestOutput
+from vllm.parser import HarmonyParser
 from vllm.renderers.hf import HfRenderer
 from vllm.renderers.mistral import MistralRenderer
 from vllm.tokenizers import get_tokenizer
 from vllm.tokenizers.mistral import MistralTokenizer
 from vllm.tokenizers.registry import cached_tokenizer_from_config
-from vllm.tool_parsers import ToolParserManager
 from vllm.v1.engine.async_llm import AsyncLLM

 GPT_OSS_MODEL_NAME = "openai/gpt-oss-20b"
@@ -575,7 +575,13 @@ def _build_serving_render(
    )


-def _build_serving_chat(engine: AsyncLLM) -> OpenAIServingChat:
+def _build_serving_chat(
+    engine: AsyncLLM,
+    *,
+    reasoning_parser: str = "",
+    tool_parser: str | None = None,
+    enable_auto_tools: bool = False,
+) -> OpenAIServingChat:
    models = OpenAIServingModels(
        engine_client=engine,
        base_model_paths=BASE_MODEL_PATHS,
@@ -590,6 +596,9 @@ def _build_serving_chat(engine: AsyncLLM) -> OpenAIServingChat:
        chat_template=CHAT_TEMPLATE,
        chat_template_content_format="auto",
        request_logger=None,
+        reasoning_parser=reasoning_parser,
+        tool_parser=tool_parser,
+        enable_auto_tools=enable_auto_tools,
    )

    return serving_chat
@@ -637,7 +646,7 @@ async def test_serving_chat_returns_correct_model_name():
    serving_chat = _build_serving_chat(mock_engine)
    messages = [{"role": "user", "content": "what is 1+1?"}]

-    async def return_model_name(*args):
+    async def return_model_name(*args, **kwargs):
        return args[3]

    serving_chat.chat_completion_full_generator = return_model_name
@@ -1210,15 +1219,21 @@ class TestServingChatWithHarmony:
        mock_engine = MagicMock(spec=AsyncLLM)
        mock_engine.errored = False
        mock_engine.model_config = MockModelConfig()
+        mock_engine.model_config.hf_config = MockHFConfig(model_type="gpt_oss")
+        mock_engine.model_config.hf_text_config = MockHFConfig(model_type="gpt_oss")
        mock_engine.input_processor = MagicMock()
        mock_engine.renderer = _build_renderer(mock_engine.model_config)
        return mock_engine

    @pytest.fixture()
    def serving_chat(self, mock_engine) -> OpenAIServingChat:
-        chat = _build_serving_chat(mock_engine)
-        chat.use_harmony = True
-        chat.tool_parser = ToolParserManager.get_tool_parser("openai")
+        chat = _build_serving_chat(
+            mock_engine,
+            reasoning_parser="openai_gptoss",
+            tool_parser="openai",
+            enable_auto_tools=True,
+        )
+        assert chat.parser_cls is HarmonyParser
        return chat

    def mock_request_output_from_req_and_token_ids(
@@ -1277,6 +1292,7 @@ class TestServingChatWithHarmony:
        stream: bool = False,
    ) -> ChatCompletionResponse:
        harmony_token_ids = get_encoding().encode(harmony_str, allowed_special="all")
+        tokenizer = get_tokenizer(GPT_OSS_MODEL_NAME)

        async def result_generator():
            if stream:
@@ -1304,11 +1320,12 @@ class TestServingChatWithHarmony:
            request_id=req.request_id,
            model_name=req.model,
            conversation=[],
-            tokenizer=get_tokenizer(req.model),
+            tokenizer=tokenizer,
            request_metadata=RequestResponseMetadata(
                request_id=req.request_id,
                model_name=req.model,
            ),
+            chat_template_kwargs=serving_chat._effective_chat_template_kwargs(req),
        )

        if stream:
@@ -1316,11 +1333,18 @@ class TestServingChatWithHarmony:
        return await result

    @pytest.mark.asyncio
-    async def test_simple_chat(self, serving_chat, stream):
+    @pytest.mark.parametrize(
+        "include_reasoning", [True, False], ids=["with_reasoning", "no_reasoning"]
+    )
+    async def test_simple_chat(self, serving_chat, stream, include_reasoning):
        messages = [{"role": "user", "content": "what is 1+1?"}]

        # Test the Harmony messages for the first turn's input
-        req = ChatCompletionRequest(model=MODEL_NAME, messages=messages)
+        req = ChatCompletionRequest(
+            model=MODEL_NAME,
+            messages=messages,
+            include_reasoning=include_reasoning,
+        )
        input_messages, _ = (
            serving_chat.openai_serving_render._make_request_with_harmony(req)
        )
@@ -1342,7 +1366,11 @@ class TestServingChatWithHarmony:
        response = await self.generate_response_from_harmony_str(
            serving_chat, req, response_str, stream=stream
        )
-        verify_chat_response(response, content=final_str, reasoning=reasoning_str)
+        verify_chat_response(
+            response,
+            content=final_str,
+            reasoning=reasoning_str if include_reasoning else None,
+        )

        # Add the output messages from the first turn as input to the second turn
        for choice in response.choices:
@@ -1,471 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for harmony streaming delta extraction.
-"""
-
-from dataclasses import dataclass, field
-from unittest.mock import patch
-
-import pytest
-
-from vllm.entrypoints.openai.chat_completion.stream_harmony import (
-    TokenState,
-    extract_harmony_streaming_delta,
-)
-
-
-@dataclass
-class MockMessage:
-    """Mock message object for testing."""
-
-    channel: str | None = None
-    recipient: str | None = None
-
-
-@dataclass
-class MockStreamableParser:
-    """Mock StreamableParser for testing without openai_harmony dependency."""
-
-    messages: list[MockMessage] = field(default_factory=list)
-
-
-class TestExtractHarmonyStreamingDelta:
-    """Tests for extract_harmony_streaming_delta function."""
-
-    @pytest.mark.parametrize(
-        "delta_text,expected_content",
-        [
-            ("Hello, world!", "Hello, world!"),
-            ("", ""),
-        ],
-    )
-    def test_final_channel_returns_content_delta(self, delta_text, expected_content):
-        """Test that final channel returns a DeltaMessage with content."""
-        parser = MockStreamableParser()
-
-        # Updated to use TokenState list
-        token_states = [TokenState(channel="final", recipient=None, text=delta_text)]
-
-        delta_message, tools_streamed = extract_harmony_streaming_delta(
-            harmony_parser=parser,
-            token_states=token_states,
-            prev_recipient=None,
-            include_reasoning=False,
-        )
-
-        assert delta_message is not None
-        assert delta_message.content == expected_content
-        assert tools_streamed is False
-
-    @pytest.mark.parametrize(
-        "include_reasoning,expected_has_message",
-        [
-            (True, True),
-            (False, False),
-        ],
-    )
-    def test_analysis_channel_reasoning(self, include_reasoning, expected_has_message):
-        """Test analysis channel respects include_reasoning flag."""
-        parser = MockStreamableParser()
-        text = "Let me think..."
-        token_states = [TokenState(channel="analysis", recipient=None, text=text)]
-
-        delta_message, tools_streamed = extract_harmony_streaming_delta(
-            harmony_parser=parser,
-            token_states=token_states,
-            prev_recipient=None,
-            include_reasoning=include_reasoning,
-        )
-
-        if expected_has_message:
-            assert delta_message is not None
-            assert delta_message.reasoning == text
-        else:
-            assert delta_message is None
-        assert tools_streamed is False
-
-    @pytest.mark.parametrize("channel", ["commentary", "analysis"])
-    @patch("vllm.entrypoints.openai.chat_completion.stream_harmony.make_tool_call_id")
-    def test_new_tool_call(self, mock_make_tool_call_id, channel):
-        """Test new tool call creation when recipient changes."""
-        mock_make_tool_call_id.return_value = "call_test123"
-        parser = MockStreamableParser()
-
-        token_states = [
-            TokenState(channel=channel, recipient="functions.get_weather", text="")
-        ]
-
-        delta_message, tools_streamed = extract_harmony_streaming_delta(
-            harmony_parser=parser,
-            token_states=token_states,
-            prev_recipient=None,
-            include_reasoning=False,
-        )
-
-        assert delta_message is not None
-        assert len(delta_message.tool_calls) == 1
-        tool_call = delta_message.tool_calls[0]
-        assert tool_call.id == "call_test123"
-        assert tool_call.type == "function"
-        assert tool_call.function.name == "get_weather"
-        assert tool_call.function.arguments == ""
-        assert tool_call.index == 0
-        assert tools_streamed is True
-
-    @pytest.mark.parametrize("channel", ["commentary", "analysis"])
-    def test_tool_call_argument_streaming(self, channel):
-        """Test streaming tool call arguments (same recipient)."""
-        parser = MockStreamableParser()
-        args_text = '{"location": "Paris"}'
-
-        token_states = [
-            TokenState(
-                channel=channel,
-                recipient="functions.get_weather",
-                text=args_text,
-            )
-        ]
-
-        delta_message, tools_streamed = extract_harmony_streaming_delta(
-            harmony_parser=parser,
-            token_states=token_states,
-            prev_recipient="functions.get_weather",
-            include_reasoning=False,
-        )
-
-        assert delta_message is not None
-        tool_call = delta_message.tool_calls[0]
-        assert tool_call.id is None
-        assert tool_call.function.arguments == args_text
-        assert tool_call.index == 0
-        assert tools_streamed is True
-
-    @pytest.mark.parametrize("channel", ["commentary", "analysis"])
-    def test_tool_call_empty_arguments_returns_none(self, channel):
-        """Test empty delta_text with same recipient returns None."""
-        parser = MockStreamableParser()
-
-        token_states = [
-            TokenState(channel=channel, recipient="functions.get_weather", text="")
-        ]
-
-        delta_message, tools_streamed = extract_harmony_streaming_delta(
-            harmony_parser=parser,
-            token_states=token_states,
-            prev_recipient="functions.get_weather",
-            include_reasoning=False,
-        )
-
-        assert delta_message is None
-        assert tools_streamed is False
-
-    def test_tool_call_index_from_previous_messages(self):
-        """Test tool call index accounts for previous function messages."""
-        messages = [
-            MockMessage(channel="analysis", recipient=None),  # Not counted
-            MockMessage(channel="commentary", recipient="functions.tool1"),  # Counted
-            MockMessage(channel="final", recipient=None),  # Not counted
-        ]
-        parser = MockStreamableParser(messages=messages)
-
-        token_states = [
-            TokenState(
-                channel="commentary",
-                recipient="functions.tool2",
-                text="args",
-            )
-        ]
-
-        delta_message, _ = extract_harmony_streaming_delta(
-            harmony_parser=parser,
-            token_states=token_states,
-            prev_recipient="functions.tool2",
-            include_reasoning=False,
-        )
-
-        assert delta_message.tool_calls[0].index == 1
-
-    def test_returns_preambles_as_content(self):
-        """Test that commentary with no recipient (preamble) is user content."""
-        parser = MockStreamableParser()
-        delta_text = "some text"
-
-        token_states = [
-            TokenState(channel="commentary", recipient=None, text=delta_text)
-        ]
-
-        delta_message, tools_streamed = extract_harmony_streaming_delta(
-            harmony_parser=parser,
-            token_states=token_states,
-            prev_recipient=None,
-            include_reasoning=True,
-        )
-
-        assert delta_message.content == delta_text
-        assert tools_streamed is False
-
-    @pytest.mark.parametrize("channel", ["commentary", "analysis"])
-    @patch("vllm.entrypoints.openai.chat_completion.stream_harmony.make_tool_call_id")
-    def test_new_tool_call_dotted_function_name(self, mock_make_tool_call_id, channel):
-        mock_make_tool_call_id.return_value = "call_dotted123"
-        parser = MockStreamableParser()
-
-        token_states = [TokenState(channel=channel, recipient="math.sum", text="")]
-
-        delta_message, tools_streamed = extract_harmony_streaming_delta(
-            harmony_parser=parser,
-            token_states=token_states,
-            prev_recipient=None,
-            include_reasoning=False,
-        )
-
-        assert delta_message is not None
-        assert len(delta_message.tool_calls) == 1
-        tool_call = delta_message.tool_calls[0]
-        assert tool_call.id == "call_dotted123"
-        assert tool_call.type == "function"
-        assert tool_call.function.name == "math.sum"
-        assert tool_call.function.arguments == ""
-        assert tool_call.index == 0
-        assert tools_streamed is True
-
-    @pytest.mark.parametrize(
-        "channel,recipient",
-        [
-            (None, None),
-            ("unknown_channel", None),
-            ("commentary", "browser.search"),
-            ("commentary", "assistant"),
-        ],
-    )
-    def test_returns_none_for_invalid_inputs(self, channel, recipient):
-        """Test that invalid channel/recipient combinations return None."""
-        parser = MockStreamableParser()
-
-        token_states = [
-            TokenState(channel=channel, recipient=recipient, text="some text")
-        ]
-
-        delta_message, tools_streamed = extract_harmony_streaming_delta(
-            harmony_parser=parser,
-            token_states=token_states,
-            prev_recipient=None,
-            include_reasoning=True,
-        )
-
-        assert delta_message is None
-        assert tools_streamed is False
-
-    def test_consecutive_token_grouping(self):
-        """
-        Test that consecutive tokens with the same channel/recipient
-        are merged into a single processing group.
-        """
-        parser = MockStreamableParser()
-        token_states = [
-            TokenState("final", None, "H"),
-            TokenState("final", None, "el"),
-            TokenState("final", None, "lo"),
-            TokenState("final", None, ","),
-            TokenState("final", None, " World"),
-        ]
-
-        delta_message, _ = extract_harmony_streaming_delta(
-            harmony_parser=parser,
-            token_states=token_states,
-            prev_recipient=None,
-            include_reasoning=False,
-        )
-
-        assert delta_message is not None
-        assert delta_message.content == "Hello, World"
-
-    @patch("vllm.entrypoints.openai.chat_completion.stream_harmony.make_tool_call_id")
-    def test_complex_batch_permutation(self, mock_make_id):
-        """
-        Test a complex permutation: Reasoning -> Tool Call -> Content.
-        This verifies that multiple distinct actions in one batch
-        are all captured in the single DeltaMessage.
-        """
-        mock_make_id.return_value = "call_batch_test"
-        parser = MockStreamableParser()
-
-        token_states = [
-            # 1. Reasoning
-            TokenState("analysis", None, "Reasoning about query..."),
-            # 2. Tool Calling
-            TokenState("commentary", "functions.search", '{"query":'),
-            TokenState("commentary", "functions.search", ' "vllm"}'),
-            # 3. Final Content
-            TokenState("final", None, "."),
-        ]
-
-        delta_message, tools_streamed = extract_harmony_streaming_delta(
-            harmony_parser=parser,
-            token_states=token_states,
-            prev_recipient=None,
-            include_reasoning=True,
-        )
-
-        assert delta_message is not None
-
-        assert delta_message.reasoning == "Reasoning about query..."
-
-        # We expect 2 objects for 1 logical tool call:
-        # 1. The definition (id, name, type)
-        # 2. The arguments payload
-        assert len(delta_message.tool_calls) == 2
-
-        header = delta_message.tool_calls[0]
-        payload = delta_message.tool_calls[1]
-
-        assert header.function.name == "search"
-        assert header.id == "call_batch_test"
-        assert header.index == 0
-
-        assert payload.index == 0
-        assert payload.function.arguments == '{"query": "vllm"}'
-
-        assert delta_message.content == "."
-        assert tools_streamed is True
-
-    @patch("vllm.entrypoints.openai.chat_completion.stream_harmony.make_tool_call_id")
-    def test_tool_call_index_consistency_with_ongoing_call(self, mock_make_id):
-        """
-        Test that an ongoing tool call continuation and subsequent new calls
-        maintain correct indexing when interleaved with content.
-        """
-        mock_make_id.side_effect = ["id_b", "id_c"]
-
-        messages = [
-            MockMessage(channel="commentary", recipient="functions.previous_tool")
-        ]
-        parser = MockStreamableParser(messages=messages)
-
-        token_states = [
-            TokenState("commentary", "functions.tool_a", '{"key_a": "val_a"}'),
-            TokenState("final", None, "Thinking..."),
-            TokenState("commentary", "functions.tool_b", '{"key_b": "val_b"}'),
-            TokenState("final", None, " Thinking again..."),
-            TokenState("commentary", "functions.tool_c", '{"key_c": "val_c"}'),
-        ]
-
-        delta_message, _ = extract_harmony_streaming_delta(
-            harmony_parser=parser,
-            token_states=token_states,
-            prev_recipient="functions.tool_a",
-            include_reasoning=False,
-        )
-
-        assert delta_message is not None
-
-        tool_a_deltas = [t for t in delta_message.tool_calls if t.index == 1]
-        assert len(tool_a_deltas) > 0
-        assert tool_a_deltas[0].id is None
-        assert tool_a_deltas[0].function.arguments == '{"key_a": "val_a"}'
-
-        tool_b_header = next(t for t in delta_message.tool_calls if t.id == "id_b")
-        assert tool_b_header.index == 2
-        tool_b_args = next(
-            t for t in delta_message.tool_calls if t.index == 2 and t.id is None
-        )
-        assert tool_b_args.function.arguments == '{"key_b": "val_b"}'
-
-        tool_c_start = next(t for t in delta_message.tool_calls if t.id == "id_c")
-        assert tool_c_start.index == 3
-        tool_c_args = next(
-            t for t in delta_message.tool_calls if t.index == 3 and t.id is None
-        )
-        assert tool_c_args.function.arguments == '{"key_c": "val_c"}'
-
-        assert delta_message.content == "Thinking... Thinking again..."
-
-
-class TestToolCallsOnNonStandardChannels:
-    """Tool calls are detected by recipient, not channel.
-
-    Models sometimes emit tool calls on unexpected channels (e.g. ``comment``
-    instead of ``commentary``).  These tests verify that the streaming delta
-    extraction is channel-agnostic for tool call detection.
-    """
-
-    @patch("vllm.entrypoints.openai.chat_completion.stream_harmony.make_tool_call_id")
-    def test_prefixed_tool_call_on_comment_channel(self, mock_make_tool_call_id):
-        mock_make_tool_call_id.return_value = "call_comment_chan"
-        parser = MockStreamableParser()
-
-        token_states = [
-            TokenState(channel="comment", recipient="functions.get_weather", text="")
-        ]
-
-        delta_message, tools_streamed = extract_harmony_streaming_delta(
-            harmony_parser=parser,
-            token_states=token_states,
-            prev_recipient=None,
-            include_reasoning=False,
-        )
-
-        assert delta_message is not None
-        assert len(delta_message.tool_calls) == 1
-        assert delta_message.tool_calls[0].function.name == "get_weather"
-        assert tools_streamed is True
-
-    @patch("vllm.entrypoints.openai.chat_completion.stream_harmony.make_tool_call_id")
-    def test_bare_tool_call_on_comment_channel(self, mock_make_tool_call_id):
-        mock_make_tool_call_id.return_value = "call_bare_comment"
-        parser = MockStreamableParser()
-
-        token_states = [TokenState(channel="comment", recipient="get_weather", text="")]
-
-        delta_message, tools_streamed = extract_harmony_streaming_delta(
-            harmony_parser=parser,
-            token_states=token_states,
-            prev_recipient=None,
-            include_reasoning=False,
-        )
-
-        assert delta_message is not None
-        assert len(delta_message.tool_calls) == 1
-        assert delta_message.tool_calls[0].function.name == "get_weather"
-        assert tools_streamed is True
-
-    def test_tool_call_arguments_on_comment_channel(self):
-        parser = MockStreamableParser()
-        args_text = '{"location": "Paris"}'
-
-        token_states = [
-            TokenState(
-                channel="comment", recipient="functions.get_weather", text=args_text
-            )
-        ]
-
-        delta_message, tools_streamed = extract_harmony_streaming_delta(
-            harmony_parser=parser,
-            token_states=token_states,
-            prev_recipient="functions.get_weather",
-            include_reasoning=False,
-        )
-
-        assert delta_message is not None
-        assert delta_message.tool_calls[0].function.arguments == args_text
-        assert tools_streamed is True
-
-    def test_base_index_counts_tool_calls_on_comment_channel(self):
-        messages = [
-            MockMessage(channel="comment", recipient="functions.tool1"),
-        ]
-        parser = MockStreamableParser(messages=messages)
-
-        token_states = [
-            TokenState(channel="commentary", recipient="functions.tool2", text="args")
-        ]
-
-        delta_message, _ = extract_harmony_streaming_delta(
-            harmony_parser=parser,
-            token_states=token_states,
-            prev_recipient="functions.tool2",
-            include_reasoning=False,
-        )
-
-        assert delta_message.tool_calls[0].index == 1
@@ -58,9 +58,12 @@ async def test_single_completion(client: openai.AsyncOpenAI, model_name: str) ->
    choice = completion.choices[0]
    assert len(choice.text) >= 5
    assert choice.finish_reason == "length"
-    assert completion.usage == openai.types.CompletionUsage(
-        completion_tokens=5, prompt_tokens=6, total_tokens=11
-    )
+    assert completion.usage is not None
+    assert completion.usage.completion_tokens == 5
+    assert completion.usage.prompt_tokens == 6
+    assert completion.usage.total_tokens == 11
+    assert completion.usage.prompt_tokens_details is not None
+    assert completion.usage.prompt_tokens_details.cached_tokens == 0

    # test using token IDs
    completion = await client.completions.create(
@@ -11,12 +11,10 @@ from vllm.entrypoints.openai.parser.harmony_utils import (
    auto_drop_analysis_messages,
    create_tool_definition,
    extract_function_from_recipient,
-    get_encoding,
    get_system_message,
    has_custom_tools,
    is_function_recipient,
    parse_chat_input_to_harmony_message,
-    parse_chat_output,
 )
 from vllm.entrypoints.openai.responses.harmony import (
    response_input_to_harmony,
@@ -941,110 +939,6 @@ class TestAutoDropAnalysisMessages:
        assert cleaned_messages == messages[1:]


-class TestParseChatOutput:
-    def test_parse_chat_output_interrupted_first_message(self) -> None:
-        harmony_str = "<|channel|>final<|message|>I'm in the middle of answering"
-        token_ids = get_encoding().encode(harmony_str, allowed_special="all")
-        reasoning, final_content, _ = parse_chat_output(token_ids)
-        assert reasoning is None
-        assert final_content == "I'm in the middle of answering"
-
-    def test_parse_chat_output_interrupted_reasoning_first_message(self) -> None:
-        harmony_str = "<|channel|>analysis<|message|>I'm in the middle of thinking"
-        token_ids = get_encoding().encode(harmony_str, allowed_special="all")
-        reasoning, final_content, _ = parse_chat_output(token_ids)
-        assert reasoning == "I'm in the middle of thinking"
-        assert final_content is None
-
-    def test_parse_chat_output_complete_reasoning_interrupted_content(self) -> None:
-        harmony_str = (
-            "<|channel|>analysis<|message|>I'm thinking.<|end|>"
-            "<|start|>assistant<|channel|>final"
-            "<|message|>I'm in the middle of answering"
-        )
-        token_ids = get_encoding().encode(harmony_str, allowed_special="all")
-        reasoning, final_content, _ = parse_chat_output(token_ids)
-        assert reasoning == "I'm thinking."
-        assert final_content == "I'm in the middle of answering"
-
-    def test_parse_chat_output_complete_content(self) -> None:
-        harmony_str = "<|channel|>final<|message|>The answer is 4.<|end|>"
-        token_ids = get_encoding().encode(harmony_str, allowed_special="all")
-        reasoning, final_content, _ = parse_chat_output(token_ids)
-        assert reasoning is None
-        assert final_content == "The answer is 4."
-
-    def test_parse_chat_output_complete_commentary(self) -> None:
-        harmony_str = (
-            "<|channel|>commentary<|message|>I need to call some tools.<|end|>"
-        )
-        token_ids = get_encoding().encode(harmony_str, allowed_special="all")
-        reasoning, final_content, _ = parse_chat_output(token_ids)
-        assert reasoning is None
-        assert final_content == "I need to call some tools."
-
-    def test_parse_chat_output_complete_reasoning(self) -> None:
-        harmony_str = (
-            "<|channel|>analysis<|message|>I've thought hard about this.<|end|>"
-        )
-        token_ids = get_encoding().encode(harmony_str, allowed_special="all")
-        reasoning, final_content, _ = parse_chat_output(token_ids)
-        assert reasoning == "I've thought hard about this."
-        assert final_content is None
-
-    def test_parse_chat_output_complete_reasoning_and_content(self) -> None:
-        harmony_str = (
-            "<|channel|>analysis<|message|>I've thought hard about this.<|end|>"
-            "<|start|>assistant<|channel|>final<|message|>The answer is 4.<|end|>"
-        )
-        token_ids = get_encoding().encode(harmony_str, allowed_special="all")
-        reasoning, final_content, _ = parse_chat_output(token_ids)
-        assert reasoning == "I've thought hard about this."
-        assert final_content == "The answer is 4."
-
-    def test_parse_chat_output_commentary_with_recipient_excluded(self) -> None:
-        """Commentary with a recipient (tool call) should not appear in
-        final_content — those are handled separately by the tool parser.
-
-        The first message is a preamble (visible), the second is a tool
-        call (excluded). Only the preamble should appear in final_content.
-        """
-        harmony_str = (
-            "<|channel|>commentary"
-            "<|message|>Let me check the weather.<|end|>"
-            "<|start|>assistant to=functions.get_weather"
-            "<|channel|>commentary"
-            '<|message|>{"location": "SF"}<|end|>'
-        )
-        token_ids = get_encoding().encode(harmony_str, allowed_special="all")
-        reasoning, final_content, _ = parse_chat_output(token_ids)
-        assert reasoning is None
-        assert final_content == "Let me check the weather."
-
-    def test_parse_chat_output_interrupted_preamble(self) -> None:
-        """Partial/interrupted preamble (commentary without recipient) should
-        appear in final_content, not reasoning."""
-        harmony_str = "<|channel|>commentary<|message|>I'll search for that"
-        token_ids = get_encoding().encode(harmony_str, allowed_special="all")
-        reasoning, final_content, _ = parse_chat_output(token_ids)
-        assert reasoning is None
-        assert final_content == "I'll search for that"
-
-    def test_parse_chat_output_preamble_then_final(self) -> None:
-        """Preamble followed by a final message should both appear in
-        final_content, joined by newline."""
-        harmony_str = (
-            "<|channel|>commentary"
-            "<|message|>Let me look that up.<|end|>"
-            "<|start|>assistant<|channel|>final"
-            "<|message|>The answer is 42.<|end|>"
-        )
-        token_ids = get_encoding().encode(harmony_str, allowed_special="all")
-        reasoning, final_content, _ = parse_chat_output(token_ids)
-        assert reasoning is None
-        assert final_content == "Let me look that up.\nThe answer is 42."
-
-
 def test_has_custom_tools() -> None:
    assert not has_custom_tools(set())
    assert not has_custom_tools({"web_search_preview", "code_interpreter", "container"})
@@ -512,3 +512,46 @@ async def test_stream_prompt_tokens_details():
    usage_chunk = parsed[-2]
    assert usage_chunk["choices"] == []
    assert usage_chunk["usage"]["prompt_tokens_details"]["cached_tokens"] == 2
+
+
+@pytest.mark.asyncio
+async def test_stream_prompt_tokens_details_zero_cached():
+    """enable_prompt_tokens_details includes cached_tokens=0 in final usage.
+
+    Regression test for https://github.com/vllm-project/vllm/issues/44377:
+    zero cached tokens must not be treated as falsy and omitted.
+    """
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output(
+            "req-1",
+            token_ids=[10],
+            finish_reason="stop",
+            finished=True,
+            num_cached_tokens=0,
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_serving_tokens(engine, enable_prompt_tokens_details=True)
+
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=10),
+        model=MODEL_NAME,
+        stream=True,
+        stream_options=StreamOptions(include_usage=True),
+    )
+
+    response = await serving.serve_tokens(request)
+    chunks = []
+    async for chunk in response:
+        chunks.append(chunk)
+
+    parsed = _parse_sse_chunks(chunks)
+    # Usage-only chunk (before [DONE])
+    usage_chunk = parsed[-2]
+    assert usage_chunk["choices"] == []
+    # Zero cached tokens must be present, not omitted
+    assert usage_chunk["usage"]["prompt_tokens_details"] is not None
+    assert usage_chunk["usage"]["prompt_tokens_details"]["cached_tokens"] == 0
@@ -289,6 +289,17 @@ async def test_metrics_exist(
            continue
        assert metric in response.text

+    cache_config_samples = [
+        sample
+        for family in text_string_to_metric_families(response.text)
+        if family.name == "vllm:cache_config_info"
+        for sample in family.samples
+    ]
+    assert cache_config_samples
+    for sample in cache_config_samples:
+        assert sample.labels.get("kv_cache_size_tokens") not in (None, "None", "")
+        assert sample.labels.get("kv_cache_max_concurrency") not in (None, "None", "")
+

@pytest.mark.asyncio
 async def test_abort_metrics_reset(
@@ -16,10 +16,11 @@ from statistics import mean, median
 import pytest
 import soundfile
 import torch
-from datasets import load_dataset
+from datasets import Audio, load_dataset
 from evaluate import load
 from transformers.models.whisper.english_normalizer import EnglishTextNormalizer

+from vllm.benchmarks.datasets.datasets import ASRDataset
 from vllm.multimodal.audio import get_audio_duration
 from vllm.tokenizers import get_tokenizer

@@ -38,6 +39,20 @@ def to_bytes(y, sr):
    return buffer


+def load_audio_sample(audio):
+    # Avoid torchcodec in CI by decoding dataset audio with soundfile.
+    if "array" in audio and "sampling_rate" in audio:
+        return audio["array"], audio["sampling_rate"]
+
+    if audio.get("path"):
+        return soundfile.read(audio["path"], dtype="float32")
+
+    if audio.get("bytes") is not None:
+        return soundfile.read(io.BytesIO(audio["bytes"]), dtype="float32")
+
+    raise ValueError("Audio sample did not contain array, path, or bytes data")
+
+
 # not all models have a normalizer so use the one from whisper as a standard option
 normalizer_model_info = HF_EXAMPLE_MODELS.find_hf_info("openai/whisper-large-v3")
 normalizer_tokenizer = get_tokenizer(
@@ -48,7 +63,7 @@ normalizer_tokenizer = get_tokenizer(
 normalizer = EnglishTextNormalizer(normalizer_tokenizer.english_spelling_normalizer)


-async def transcribe_audio(client, tokenizer, y, sr):
+async def transcribe_audio(client, tokenizer, y, sr, extra_body=None):
    # Send loaded audio directly instead of loading from disk,
    # don't account for that time though
    with to_bytes(y, sr) as f:
@@ -58,6 +73,7 @@ async def transcribe_audio(client, tokenizer, y, sr):
            model=tokenizer.name_or_path,
            language="en",
            temperature=0.0,
+            extra_body=extra_body,
        )
        end_time = time.perf_counter()
        # NOTE there's no streaming in transcriptions, can't measure ttft
@@ -68,17 +84,21 @@ async def transcribe_audio(client, tokenizer, y, sr):
    return latency, num_output_tokens, transcription.text


-async def bound_transcribe(sem, client, tokenizer, audio, reference):
+async def bound_transcribe(
+    sem, client, tokenizer, audio, sr, reference, extra_body=None
+):
    # Use semaphore to limit concurrent requests.
    async with sem:
-        result = await transcribe_audio(client, tokenizer, *audio)
+        result = await transcribe_audio(
+            client, tokenizer, audio, sr, extra_body=extra_body
+        )
        # Normalize *english* output/reference for evaluation.
        out = normalizer(result[2])
        ref = normalizer(reference)
        return result[:2] + (out, ref)


-async def process_dataset(model, client, data, concurrent_request):
+async def process_dataset(model, client, data, concurrent_request, extra_body=None):
    sem = asyncio.Semaphore(concurrent_request)

    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
@@ -89,14 +109,16 @@ async def process_dataset(model, client, data, concurrent_request):
    )

    # Warmup call as the first `load_audio` server-side is quite slow.
-    audio, sr = data[0]["audio"]["array"], data[0]["audio"]["sampling_rate"]
-    _ = await bound_transcribe(sem, client, tokenizer, (audio, sr), "")
+    audio, sr = load_audio_sample(data[0]["audio"])
+    _ = await bound_transcribe(sem, client, tokenizer, audio, sr, "", extra_body)

    tasks: list[asyncio.Task] = []
    for sample in data:
-        audio, sr = sample["audio"]["array"], sample["audio"]["sampling_rate"]
+        audio, sr = load_audio_sample(sample["audio"])
        task = asyncio.create_task(
-            bound_transcribe(sem, client, tokenizer, (audio, sr), sample["text"])
+            bound_transcribe(
+                sem, client, tokenizer, audio, sr, sample["text"], extra_body
+            )
        )
        tasks.append(task)
    return await asyncio.gather(*tasks)
@@ -121,19 +143,36 @@ def print_performance_metrics(results, total_time):


 def add_duration(sample):
-    y, sr = sample["audio"]["array"], sample["audio"]["sampling_rate"]
+    y, sr = load_audio_sample(sample["audio"])
    sample["duration_ms"] = get_audio_duration(y=y, sr=sr) * 1000
    return sample


-def load_hf_dataset(dataset_repo: str, split="validation", **hf_kwargs):
-    ## Load and filter the dataset
-    dataset = load_dataset(dataset_repo, split=split, **hf_kwargs)
-    if "duration_ms" not in dataset[0]:
-        # compute duration to filter
+def load_asr_dataset_rows(dataset_repo: str, split="validation", **hf_kwargs):
+    if dataset_repo in ASRDataset.SUPPORTED_DATASET_PATHS:
+        asr_dataset_kwargs = {
+            "dataset_path": dataset_repo,
+            "dataset_split": split,
+            "disable_shuffle": True,
+            "no_stream": True,
+        }
+        for key in ("dataset_subset", "hf_name", "trust_remote_code"):
+            if key in hf_kwargs:
+                asr_dataset_kwargs[key] = hf_kwargs[key]
+        return ASRDataset(**asr_dataset_kwargs).data
+
+    return load_dataset(dataset_repo, split=split, **hf_kwargs)
+
+
+def load_shortform_eval_dataset(dataset_repo: str, split="validation", **hf_kwargs):
+    ## Load and filter the dataset.
+    dataset = load_asr_dataset_rows(dataset_repo, split=split, **hf_kwargs)
+    dataset = dataset.cast_column("audio", Audio(decode=False))
+    if "duration_ms" not in dataset.column_names:
+        # Compute duration to filter.
        dataset = dataset.map(add_duration)

-    # Whisper max supported duration
+    # Whisper max supported duration.
    dataset = dataset.filter(lambda example: example["duration_ms"] < 30000)
    return dataset

@@ -145,11 +184,16 @@ def run_evaluation(
    max_concurrent_reqs: int,
    n_examples: int = -1,
    print_metrics: bool = True,
+    extra_body=None,
 ):
    if n_examples > 0:
        dataset = dataset.select(range(n_examples))
    start = time.perf_counter()
-    results = asyncio.run(process_dataset(model, client, dataset, max_concurrent_reqs))
+    results = asyncio.run(
+        process_dataset(
+            model, client, dataset, max_concurrent_reqs, extra_body=extra_body
+        )
+    )
    end = time.perf_counter()
    total_time = end - start
    print(f"Total Test Time: {total_time:.4f} seconds")
@@ -164,6 +208,106 @@ def run_evaluation(
    return wer_score


+LONGFORM_DATASET_REPO = ASRDataset.EARNINGS22_CLEANED_DATASET
+LONGFORM_DATASET_SPLIT = "test"
+LONGFORM_NUM_SAMPLES = 6
+
+
+def load_longform_dataset():
+    dataset = load_asr_dataset_rows(
+        LONGFORM_DATASET_REPO,
+        split=LONGFORM_DATASET_SPLIT,
+    )
+    assert len(dataset) >= LONGFORM_NUM_SAMPLES
+    return dataset.select(range(LONGFORM_NUM_SAMPLES))
+
+
+async def transcribe_audio_path(client, tokenizer, audio_path: str, extra_body=None):
+    with open(audio_path, "rb") as f:
+        start_time = time.perf_counter()
+        transcription = await client.audio.transcriptions.create(
+            file=f,
+            model=tokenizer.name_or_path,
+            language="en",
+            temperature=0.0,
+            extra_body=extra_body,
+        )
+        end_time = time.perf_counter()
+
+    latency = end_time - start_time
+    num_output_tokens = len(
+        tokenizer(transcription.text, add_special_tokens=False).input_ids
+    )
+    return latency, num_output_tokens, transcription.text
+
+
+async def bound_transcribe_path(
+    sem, client, tokenizer, audio_path, reference, extra_body=None
+):
+    async with sem:
+        result = await transcribe_audio_path(
+            client, tokenizer, audio_path, extra_body=extra_body
+        )
+        out = normalizer(result[2])
+        ref = normalizer(reference)
+        return result[:2] + (out, ref)
+
+
+async def process_longform_dataset(
+    model, client, data, concurrent_request, extra_body=None
+):
+    sem = asyncio.Semaphore(concurrent_request)
+
+    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
+    tokenizer = get_tokenizer(
+        model,
+        tokenizer_mode=model_info.tokenizer_mode,
+        trust_remote_code=model_info.trust_remote_code,
+    )
+
+    warmup_path = data[0]["audio"]["path"]
+    _ = await bound_transcribe_path(sem, client, tokenizer, warmup_path, "", extra_body)
+
+    tasks: list[asyncio.Task] = []
+    for sample in data:
+        audio_path = sample["audio"]["path"]
+        task = asyncio.create_task(
+            bound_transcribe_path(
+                sem, client, tokenizer, audio_path, sample["text"], extra_body
+            )
+        )
+        tasks.append(task)
+    return await asyncio.gather(*tasks)
+
+
+def run_longform_evaluation(
+    model: str,
+    client,
+    dataset,
+    max_concurrent_reqs: int,
+    print_metrics: bool = True,
+    extra_body=None,
+):
+    start = time.perf_counter()
+    results = asyncio.run(
+        process_longform_dataset(
+            model, client, dataset, max_concurrent_reqs, extra_body=extra_body
+        )
+    )
+    end = time.perf_counter()
+    total_time = end - start
+    print(f"Total Test Time: {total_time:.4f} seconds")
+    if print_metrics:
+        print_performance_metrics(results, total_time)
+
+    predictions = [res[2] for res in results]
+    references = [res[3] for res in results]
+    wer = load("wer")
+    wer_score = 100 * wer.compute(references=references, predictions=predictions)
+    print("WER:", wer_score)
+    return wer_score
+
+
 # alternatives "openai/whisper-large-v2", "openai/whisper-large-v3-turbo"..
 # NOTE: Expected WER measured with equivalent hf.transformers args:
 # whisper-large-v3 + esb-datasets-earnings22-validation-tiny-filtered.
@@ -184,7 +328,6 @@ def test_wer_correctness(
 ):
    model_name, expected_wer = model_config
    model_info = HF_EXAMPLE_MODELS.find_hf_info(model_name)
-    # TODO refactor to use `ASRDataset`
    server_args = [
        "--enforce-eager",
        f"--tokenizer_mode={model_info.tokenizer_mode}",
@@ -197,7 +340,7 @@ def test_wer_correctness(
        model_name,
        server_args,
    ) as remote_server:
-        dataset = load_hf_dataset(dataset_repo)
+        dataset = load_shortform_eval_dataset(dataset_repo)

        if not max_concurrent_request:
            # No max concurrency
@@ -216,3 +359,42 @@ def test_wer_correctness(

        if expected_wer:
            torch.testing.assert_close(wer, expected_wer, atol=1e-1, rtol=1e-2)
+
+
+# 14-22mins of 6 audio samples of total ~115 mins and just 37MB.
+# checks for long audio transcription correctness and RMS split.
+@pytest.mark.parametrize(
+    "model_config",
+    [("openai/whisper-large-v3", 9.5)],
+)
+def test_long_audio_wer_correctness(model_config):
+    model_name, expected_wer = model_config
+    model_info = HF_EXAMPLE_MODELS.find_hf_info(model_name)
+    server_args = [
+        f"--tokenizer_mode={model_info.tokenizer_mode}",
+    ]
+
+    if model_info.trust_remote_code:
+        server_args.append("--trust-remote-code")
+
+    # 1800 seconds is 30 minutes
+    env_dict = {
+        "VLLM_MAX_AUDIO_DECODE_DURATION_S": "1800",
+    }
+
+    with RemoteOpenAIServer(
+        model_name,
+        server_args,
+        env_dict=env_dict,
+    ) as remote_server:
+        dataset = load_longform_dataset()
+        client = remote_server.get_async_client()
+        wer = run_longform_evaluation(
+            model=model_name,
+            client=client,
+            dataset=dataset,
+            max_concurrent_reqs=LONGFORM_NUM_SAMPLES,
+        )
+
+    print(f"Expected WER: {expected_wer}, Actual WER: {wer}")
+    torch.testing.assert_close(wer, expected_wer, atol=1e-1, rtol=1e-2)
@@ -29,8 +29,10 @@ def test_sparse_flashmla_metadata_smoke():
        topk=topk,
        is_fp8_kvcache=True,
    )
-    assert tile_md.dtype == torch.int32
-    assert num_splits.dtype == torch.int32
+    assert isinstance(tile_md, fm.FlashMLASchedMeta)
+    assert tile_md.tile_scheduler_metadata is None
+    assert tile_md.num_splits is None
+    assert num_splits is None


 def test_sparse_flashmla_decode_smoke():
@@ -116,7 +118,7 @@ def test_sparse_flashmla_prefill_smoke():
    kv = torch.zeros((s_kv, h_kv, d_qk), dtype=torch.bfloat16, device=device)
    indices = torch.zeros((s_q, h_kv, topk), dtype=torch.int32, device=device)

-    out, max_logits, lse = fm.flash_mla_sparse_prefill(q, kv, indices, 1.0, d_v)
+    out, max_logits, lse = fm.flash_mla_sparse_fwd(q, kv, indices, 1.0, d_v)
    assert out.shape == (s_q, h_q, d_v)
    assert max_logits.shape == (s_q, h_q)
    assert lse.shape == (s_q, h_q)
@@ -0,0 +1,318 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for per-request causal/non-causal attention (mixed batches).
+
+Validates that both triton and flash-attention backends correctly handle
+batches where some sequences use causal masking and others use non-causal
+(bidirectional) masking — needed by DiffusionGemma.
+"""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import set_random_seed
+
+# Mixed causal/non-causal attention is only validated on a subset of GPUs:
+# the Triton path on Hopper (SM90) and B200 (SM100); the FA4 path on Hopper
+# (SM90) only.
+_device_capability = current_platform.get_device_capability()
+_major = _device_capability.major if _device_capability is not None else None
+
+NUM_HEADS = [(4, 4), (8, 2)]
+HEAD_SIZES = [128]
+BLOCK_SIZES = [16]
+DTYPES = [torch.bfloat16]
+
+
+def ref_paged_attn(
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    query_lens: list[int],
+    kv_lens: list[int],
+    block_tables: torch.Tensor,
+    scale: float,
+    per_seq_causal: list[bool],
+    sliding_window: int | None = None,
+) -> torch.Tensor:
+    num_seqs = len(query_lens)
+    block_tables_np = block_tables.cpu().numpy()
+    _, block_size, num_kv_heads, head_size = key_cache.shape
+
+    outputs: list[torch.Tensor] = []
+    start_idx = 0
+    for i in range(num_seqs):
+        query_len = query_lens[i]
+        kv_len = kv_lens[i]
+        q = query[start_idx : start_idx + query_len]
+        q = q * scale
+
+        num_kv_blocks = (kv_len + block_size - 1) // block_size
+        block_indices = block_tables_np[i, :num_kv_blocks]
+        k = key_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len]
+        v = value_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len]
+
+        if q.shape[1] != k.shape[1]:
+            k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1)
+            v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
+
+        attn = torch.einsum("qhd,khd->hqk", q, k).float()
+
+        if per_seq_causal[i]:
+            mask = torch.triu(
+                torch.ones(query_len, kv_len, device=attn.device),
+                diagonal=kv_len - query_len + 1,
+            ).bool()
+        else:
+            mask = torch.zeros(query_len, kv_len, device=attn.device).bool()
+
+        if sliding_window is not None:
+            sw_mask = (
+                torch.triu(
+                    torch.ones(query_len, kv_len, device=attn.device),
+                    diagonal=kv_len - (query_len + sliding_window) + 1,
+                )
+                .bool()
+                .logical_not()
+            )
+            mask |= sw_mask
+
+        attn.masked_fill_(mask, float("-inf"))
+        attn = torch.softmax(attn, dim=-1).to(v.dtype)
+        out = torch.einsum("hqk,khd->qhd", attn, v)
+        outputs.append(out)
+        start_idx += query_len
+
+    return torch.cat(outputs, dim=0)
+
+
+# ---- Triton backend test ----
+
+
+@pytest.mark.skipif(
+    _major not in (9, 10),
+    reason="Triton mixed causal attention requires Hopper (SM90) or B200 (SM100).",
+)
+@pytest.mark.parametrize(
+    "seq_lens",
+    [[(1, 128), (5, 64), (1, 256)]],
+)
+@pytest.mark.parametrize(
+    "per_seq_causal",
+    [[True, False, True], [False, True, False], [True, True, False]],
+)
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
+@pytest.mark.parametrize("head_size", HEAD_SIZES)
+@pytest.mark.parametrize("block_size", BLOCK_SIZES)
+@pytest.mark.parametrize("dtype", DTYPES)
+@torch.inference_mode()
+def test_triton_mixed_causal(
+    seq_lens: list[tuple[int, int]],
+    per_seq_causal: list[bool],
+    num_heads: tuple[int, int],
+    head_size: int,
+    dtype: torch.dtype,
+    block_size: int,
+):
+    if not current_platform.is_cuda():
+        pytest.skip("Triton attention requires CUDA")
+
+    from vllm.v1.attention.ops.triton_unified_attention import unified_attention
+
+    set_random_seed(42)
+    device = "cuda"
+
+    num_query_heads, num_kv_heads = num_heads
+    assert len(seq_lens) == len(per_seq_causal)
+
+    query_lens = [s[0] for s in seq_lens]
+    kv_lens = [s[1] for s in seq_lens]
+    num_seqs = len(seq_lens)
+
+    num_query_tokens = sum(query_lens)
+    max_kv_len = max(kv_lens)
+    max_num_blocks = (max_kv_len + block_size - 1) // block_size
+    num_blocks = max_num_blocks * num_seqs + 10
+
+    scale = head_size**-0.5
+    query = torch.randn(
+        num_query_tokens, num_query_heads, head_size, dtype=dtype, device=device
+    )
+    key_cache = torch.randn(
+        num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device
+    )
+    value_cache = torch.randn(
+        num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device
+    )
+
+    block_tables_list = []
+    for i in range(num_seqs):
+        n_blocks = (kv_lens[i] + block_size - 1) // block_size
+        blocks = list(range(i * max_num_blocks, i * max_num_blocks + n_blocks))
+        blocks += [0] * (max_num_blocks - n_blocks)
+        block_tables_list.append(blocks)
+    block_tables = torch.tensor(block_tables_list, dtype=torch.int32, device=device)
+
+    cu_seqlens_q = torch.zeros(num_seqs + 1, dtype=torch.int32, device=device)
+    for i, ql in enumerate(query_lens):
+        cu_seqlens_q[i + 1] = cu_seqlens_q[i] + ql
+
+    seqused_k = torch.tensor(kv_lens, dtype=torch.int32, device=device)
+    max_seqlen_q = max(query_lens)
+    max_seqlen_k = max(kv_lens)
+
+    causal_tensor = torch.tensor(per_seq_causal, dtype=torch.bool, device=device)
+
+    output = torch.empty_like(query)
+    unified_attention(
+        q=query,
+        k=key_cache,
+        v=value_cache,
+        out=output,
+        cu_seqlens_q=cu_seqlens_q,
+        max_seqlen_q=max_seqlen_q,
+        seqused_k=seqused_k,
+        max_seqlen_k=max_seqlen_k,
+        softmax_scale=scale,
+        causal=causal_tensor,
+        window_size=(-1, -1),
+        block_table=block_tables,
+        softcap=0.0,
+        q_descale=None,
+        k_descale=1.0,
+        v_descale=1.0,
+    )
+
+    ref_output = ref_paged_attn(
+        query,
+        key_cache,
+        value_cache,
+        query_lens,
+        kv_lens,
+        block_tables,
+        scale,
+        per_seq_causal,
+    )
+
+    torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)
+
+
+# ---- Flash Attention 4 backend test (native per_seq_causal) ----
+
+
+@pytest.mark.skipif(
+    _major != 9,
+    reason="FA4 mixed causal attention requires Hopper (SM90).",
+)
+@pytest.mark.parametrize(
+    "seq_lens",
+    [[(1, 128), (5, 64), (1, 256)]],
+)
+@pytest.mark.parametrize(
+    "per_seq_causal",
+    [[True, False, True], [False, True, False]],
+)
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
+@pytest.mark.parametrize("head_size", HEAD_SIZES)
+@pytest.mark.parametrize("block_size", BLOCK_SIZES)
+@pytest.mark.parametrize("dtype", DTYPES)
+@torch.inference_mode()
+def test_flash_attn4_mixed_causal(
+    seq_lens: list[tuple[int, int]],
+    per_seq_causal: list[bool],
+    num_heads: tuple[int, int],
+    head_size: int,
+    dtype: torch.dtype,
+    block_size: int,
+):
+    if not current_platform.is_cuda():
+        pytest.skip("Flash attention requires CUDA")
+
+    try:
+        from vllm.vllm_flash_attn import (
+            fa_version_unsupported_reason,
+            flash_attn_varlen_func,
+            is_fa_version_supported,
+        )
+    except ImportError:
+        pytest.skip("vllm_flash_attn not available")
+
+    if not is_fa_version_supported(4):
+        reason = fa_version_unsupported_reason(4)
+        pytest.skip(f"FA4 not supported: {reason}")
+
+    set_random_seed(42)
+    device = "cuda"
+
+    num_query_heads, num_kv_heads = num_heads
+    assert len(seq_lens) == len(per_seq_causal)
+
+    query_lens = [s[0] for s in seq_lens]
+    kv_lens = [s[1] for s in seq_lens]
+    num_seqs = len(seq_lens)
+
+    num_query_tokens = sum(query_lens)
+    max_kv_len = max(kv_lens)
+    max_num_blocks = (max_kv_len + block_size - 1) // block_size
+    num_blocks = max_num_blocks * num_seqs + 10
+
+    scale = head_size**-0.5
+    query = torch.randn(
+        num_query_tokens, num_query_heads, head_size, dtype=dtype, device=device
+    )
+    key_cache = torch.randn(
+        num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device
+    )
+    value_cache = torch.randn(
+        num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device
+    )
+
+    block_tables_list = []
+    for i in range(num_seqs):
+        n_blocks = (kv_lens[i] + block_size - 1) // block_size
+        blocks = list(range(i * max_num_blocks, i * max_num_blocks + n_blocks))
+        blocks += [0] * (max_num_blocks - n_blocks)
+        block_tables_list.append(blocks)
+    block_tables = torch.tensor(block_tables_list, dtype=torch.int32, device=device)
+
+    cu_seqlens_q = torch.zeros(num_seqs + 1, dtype=torch.int32, device=device)
+    for i, ql in enumerate(query_lens):
+        cu_seqlens_q[i + 1] = cu_seqlens_q[i] + ql
+
+    seqused_k = torch.tensor(kv_lens, dtype=torch.int32, device=device)
+    per_seq_causal_tensor = torch.tensor(
+        per_seq_causal, dtype=torch.int32, device=device
+    )
+
+    ref_output = ref_paged_attn(
+        query,
+        key_cache,
+        value_cache,
+        query_lens,
+        kv_lens,
+        block_tables,
+        scale,
+        per_seq_causal,
+    )
+
+    output = torch.empty_like(query)
+    flash_attn_varlen_func(
+        q=query,
+        k=key_cache,
+        v=value_cache,
+        out=output,
+        cu_seqlens_q=cu_seqlens_q,
+        max_seqlen_q=max(query_lens),
+        seqused_k=seqused_k,
+        max_seqlen_k=max(kv_lens),
+        softmax_scale=scale,
+        # The kernel must be compiled causal for `dynamic_causal` to take effect.
+        causal=True,
+        block_table=block_tables,
+        softcap=0.0,
+        dynamic_causal=per_seq_causal_tensor,
+        fa_version=4,
+    )
+
+    torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)
@@ -10,6 +10,25 @@ pytestmark = pytest.mark.skipif(
    not current_platform.is_rocm(), reason="Only used by ROCm"
 )

+
+def _on_gfx950() -> bool:
+    if not current_platform.is_rocm():
+        return False
+    try:
+        from vllm.platforms.rocm import _ON_GFX950
+
+        return bool(_ON_GFX950)
+    except Exception:
+        return False
+
+
+# The flash-decode split-K decode path is only tuned for AMD gfx950; other
+# architectures take the fallback decode kernel, so its tests are skipped there.
+requires_gfx950 = pytest.mark.skipif(
+    not _on_gfx950(),
+    reason="split-K decode kernel is only tuned for AMD gfx950",
+)
+
 NOPE_HEAD_DIM = 448
 ROPE_HEAD_DIM = 64
 HEAD_DIM = NOPE_HEAD_DIM + ROPE_HEAD_DIM
@@ -156,6 +175,20 @@ def _ref_sparse_decode_ragged(
    return out.to(torch.bfloat16)


+def _ragged_from_rows(
+    rows: list[list[int]], device: torch.device
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Flatten per-query slot lists into ragged (indices, indptr) tensors."""
+    flat = [slot for row in rows for slot in row]
+    indptr = [0]
+    for row in rows:
+        indptr.append(indptr[-1] + len(row))
+    return (
+        torch.tensor(flat, dtype=torch.int32, device=device),
+        torch.tensor(indptr, dtype=torch.int32, device=device),
+    )
+
+
 def _ref_combine_topk_swa_ragged(
    device: torch.device,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
@@ -375,3 +408,110 @@ def test_combine_topk_swa_indices_ragged() -> None:
    )
    torch.testing.assert_close(actual_indptr, expected_indptr)
    torch.testing.assert_close(actual_lens, expected_lens)
+
+
+@requires_gfx950
+@torch.inference_mode()
+def test_decode_num_splits_heuristic(monkeypatch) -> None:
+    """Split-count heuristic added with the flash-decode split-K decode path."""
+    from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
+
+    # Pin the CU count so the heuristic is deterministic off-device.
+    monkeypatch.setattr(mod, "_decode_cu_count", lambda: 256)
+
+    # A batch that already fills the device should not be split.
+    assert mod._decode_num_splits(256, 1, avg_main_len=128.0, avg_extra_len=0.0) == 1
+    # A tiny batch on a large device should split to add parallelism.
+    assert mod._decode_num_splits(2, 1, avg_main_len=256.0, avg_extra_len=0.0) > 1
+
+    # The chosen count always stays within the searched [1, 16] range, and a
+    # zero-length workload never splits (no work to parallelize).
+    for num_queries in (1, 4, 24, 224, 1024):
+        splits = mod._decode_num_splits(
+            num_queries, 1, avg_main_len=512.0, avg_extra_len=128.0
+        )
+        assert 1 <= splits <= 16
+    assert mod._decode_num_splits(2, 1, avg_main_len=0.0, avg_extra_len=0.0) >= 1
+
+
+@requires_gfx950
+@pytest.mark.parametrize("num_splits", [1, 2, 3, 4, 8])
+@pytest.mark.parametrize("with_extra", [True, False])
+@pytest.mark.parametrize("with_sink", [True, False])
+@torch.inference_mode()
+def test_sparse_attn_decode_split_k_kernel(
+    monkeypatch, num_splits: int, with_extra: bool, with_sink: bool
+) -> None:
+    """Flash-decode split-K decode path (partial + reduce kernels).
+
+    This path is the gfx950 production path (``_ON_GFX950``), so the test only
+    runs on gfx950. The split count is pinned so the partial/reduce kernels are
+    exercised across split counts. ``num_splits=8`` drives splits past the
+    shortest segment length, covering the empty-split edge case handled by the
+    reduce kernel.
+    """
+    from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
+
+    device = torch.device("cuda")
+    torch.manual_seed(7)
+    block_size = 4
+    num_heads = 3
+
+    main_rows = [[0, 2, 4, 6, 1, 3, 7, 5], [4, 1, 6, 0, 2]]
+    num_queries = len(main_rows)
+    q = (
+        torch.randn(
+            num_queries, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device
+        )
+        * 0.125
+    )
+    main_kv = torch.randn(8, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
+    main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size)
+    main_indices, main_indptr = _ragged_from_rows(main_rows, device)
+
+    extra_rows: list[list[int]] | None = None
+    extra_cache: torch.Tensor | None = None
+    extra_indices: torch.Tensor | None = None
+    extra_indptr: torch.Tensor | None = None
+    if with_extra:
+        rows = [[1, 3, 0, 5, 2, 4], [3, 0, 6]]
+        extra_kv = torch.randn(7, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
+        extra_rows = rows
+        extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size)
+        extra_indices, extra_indptr = _ragged_from_rows(rows, device)
+
+    attn_sink = (
+        torch.tensor([-0.1, 0.0, 0.1], dtype=torch.float32, device=device)
+        if with_sink
+        else None
+    )
+    scale = HEAD_DIM**-0.5
+
+    # Pin the split count so each parametrized value is exercised deterministically.
+    monkeypatch.setattr(mod, "_decode_num_splits", lambda *args, **kwargs: num_splits)
+
+    actual = mod._rocm_sparse_attn_decode_ragged_triton(
+        q=q,
+        main_cache=main_cache,
+        main_indices=main_indices,
+        main_indptr=main_indptr,
+        scale=scale,
+        attn_sink=attn_sink,
+        nope_head_dim=NOPE_HEAD_DIM,
+        rope_head_dim=ROPE_HEAD_DIM,
+        extra_cache=extra_cache,
+        extra_indices=extra_indices,
+        extra_indptr=extra_indptr,
+    )
+    expected = _ref_sparse_decode_ragged(
+        q=q,
+        main_cache=main_cache,
+        main_rows=main_rows,
+        scale=scale,
+        attn_sink=attn_sink,
+        block_size=block_size,
+        extra_cache=extra_cache,
+        extra_rows=extra_rows,
+    )
+
+    torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
@@ -0,0 +1,189 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Unit tests for the Triton DiffKV unified-attention kernel.
+"""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+from vllm.utils.math_utils import next_power_of_2
+from vllm.utils.torch_utils import (
+    canonicalize_singleton_dim_strides,
+    set_random_seed,
+)
+from vllm.v1.attention.backends.fa_utils import (
+    get_flash_attn_version,
+    is_flash_attn_varlen_func_available,
+)
+from vllm.v1.attention.ops.triton_unified_attention_diffkv import (
+    unified_attention_diffkv,
+)
+
+DEVICE_TYPE = current_platform.device_type
+
+# (num_query_heads, num_kv_heads): MHA, GQA, and the num_kv_heads==1
+# (degenerate-stride) case.
+NUM_HEADS = [(4, 4), (8, 2), (5, 1)]
+# (head_size_qk, head_size_v).  (192, 128) is the canonical asymmetric
+# DiffKV shape; FA4 on Blackwell only supports head_size>128 when it is
+# 192, and FA3 on Hopper supports it too -- so this pair is runnable on
+# both.  (128, 128) keeps the equal-dim path covered through the DiffKV
+# kernel.
+HEAD_SIZES = [(128, 128), (192, 128)]
+BLOCK_SIZES = [16]
+DTYPES = [torch.bfloat16]
+
+NUM_BLOCKS = 2048
+
+# 0: 2D decode kernel; 8: 3D (split-KV) decode kernel.
+SEQ_THRESHOLD_3D_VALUES = [0, 8]
+
+NUM_PAR_SOFTMAX_SEGMENTS = 16
+
+
+def _alloc_segm_buffers(seq_threshold_3D: int, num_query_heads: int, head_size_v: int):
+    """Allocate the split-KV softmax scratch (last dim == head_size_v)."""
+    head_size_v_padded = next_power_of_2(head_size_v)
+    segm_output = torch.empty(
+        (
+            seq_threshold_3D,
+            num_query_heads,
+            NUM_PAR_SOFTMAX_SEGMENTS,
+            head_size_v_padded,
+        ),
+        dtype=torch.float32,
+    )
+    segm_max = torch.empty(
+        (seq_threshold_3D, num_query_heads, NUM_PAR_SOFTMAX_SEGMENTS),
+        dtype=torch.float32,
+    )
+    segm_expsum = torch.empty(
+        (seq_threshold_3D, num_query_heads, NUM_PAR_SOFTMAX_SEGMENTS),
+        dtype=torch.float32,
+    )
+    return segm_output, segm_max, segm_expsum
+
+
+@pytest.mark.parametrize(
+    "seq_lens",
+    [
+        [(1, 1328), (5, 18), (129, 463)],  # mixed prefill + decode
+        [(1, 523), (1, 37), (1, 2011)],  # decode-only (exercises 3D path)
+    ],
+)
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
+@pytest.mark.parametrize("head_sizes", HEAD_SIZES)
+@pytest.mark.parametrize("block_size", BLOCK_SIZES)
+@pytest.mark.parametrize("sliding_window", [None, 128])
+@pytest.mark.parametrize("soft_cap", [None, 50.0])
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
+@torch.inference_mode()
+def test_triton_unified_attn_diffkv_vs_fa(
+    seq_lens: list[tuple[int, int]],
+    num_heads: tuple[int, int],
+    head_sizes: tuple[int, int],
+    sliding_window: int | None,
+    soft_cap: float | None,
+    dtype: torch.dtype,
+    block_size: int,
+    seq_threshold_3D: int,
+) -> None:
+    head_size_qk, head_size_v = head_sizes
+
+    # DiffKV requires FA3 (Hopper) / FA4 (Blackwell) as the reference.
+    fa_version = get_flash_attn_version(head_size=head_size_qk, head_size_v=head_size_v)
+    if not is_flash_attn_varlen_func_available() or fa_version not in (3, 4):
+        pytest.skip(f"FA DiffKV needs FA3/FA4 (got version {fa_version}).")
+
+    from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func
+
+    torch.set_default_device(DEVICE_TYPE)
+    set_random_seed(0)
+
+    num_seqs = len(seq_lens)
+    query_lens = [x[0] for x in seq_lens]
+    kv_lens = [x[1] for x in seq_lens]
+    num_query_heads, num_kv_heads = num_heads
+    assert num_query_heads % num_kv_heads == 0
+    max_query_len = max(query_lens)
+    max_kv_len = max(kv_lens)
+    window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
+    scale = head_size_qk**-0.5
+
+    query = torch.randn(sum(query_lens), num_query_heads, head_size_qk, dtype=dtype)
+    # Packed KV cache: [num_blocks, block_size, num_kv_heads, hqk + hv].
+    kv_cache = torch.randn(
+        NUM_BLOCKS,
+        block_size,
+        num_kv_heads,
+        head_size_qk + head_size_v,
+        dtype=dtype,
+    )
+    key_cache = kv_cache[..., :head_size_qk]
+    value_cache = kv_cache[..., head_size_qk:]
+
+    cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
+        dim=0, dtype=torch.int32
+    )
+    kv_lens_t = torch.tensor(kv_lens, dtype=torch.int32)
+
+    max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
+    block_tables = torch.randint(
+        0, NUM_BLOCKS, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
+    )
+
+    # ---- FlashAttention DiffKV (ground truth) ---------------------------
+    # Mirror the backend: fix degenerate strides on size-1 dims so FA's
+    # TMA path sees ≥16-byte-aligned strides (matters for num_kv_heads==1).
+    fa_k = canonicalize_singleton_dim_strides(key_cache)
+    fa_v = canonicalize_singleton_dim_strides(value_cache)
+    fa_out = torch.empty(sum(query_lens), num_query_heads, head_size_v, dtype=dtype)
+    flash_attn_varlen_func(
+        q=query,
+        k=fa_k,
+        v=fa_v,
+        out=fa_out,
+        cu_seqlens_q=cu_query_lens,
+        max_seqlen_q=max_query_len,
+        seqused_k=kv_lens_t,
+        max_seqlen_k=max_kv_len,
+        softmax_scale=scale,
+        causal=True,
+        window_size=list(window_size),
+        block_table=block_tables,
+        softcap=soft_cap if soft_cap is not None else 0,
+        fa_version=fa_version,
+    )
+
+    # ---- Triton DiffKV --------------------------------------------------
+    segm_output, segm_max, segm_expsum = _alloc_segm_buffers(
+        seq_threshold_3D, num_query_heads, head_size_v
+    )
+    triton_out = torch.empty(sum(query_lens), num_query_heads, head_size_v, dtype=dtype)
+    unified_attention_diffkv(
+        q=query,
+        k=key_cache,
+        v=value_cache,
+        out=triton_out,
+        cu_seqlens_q=cu_query_lens,
+        seqused_k=kv_lens_t,
+        softmax_scale=scale,
+        causal=True,
+        window_size=window_size,
+        block_table=block_tables,
+        softcap=soft_cap if soft_cap is not None else 0,
+        max_seqlen_q=max_query_len,
+        seq_threshold_3D=seq_threshold_3D,
+        num_par_softmax_segments=NUM_PAR_SOFTMAX_SEGMENTS,
+        softmax_segm_output=segm_output,
+        softmax_segm_max=segm_max,
+        softmax_segm_expsum=segm_expsum,
+    )
+
+    (
+        torch.testing.assert_close(triton_out, fa_out, atol=2e-2, rtol=2e-2),
+        f"triton vs FA max abs diff: {torch.max(torch.abs(triton_out - fa_out))}",
+    )
@@ -0,0 +1,243 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the per_token_group_fp8_quant helion kernel
+
+Run `pytest tests/kernels/helion/test_per_token_group_fp8_quant.py`.
+"""
+
+from typing import Any
+
+import pytest
+import torch
+from torch._subclasses.fake_tensor import FakeTensorMode
+
+from tests.kernels.helion.utils import skip_if_platform_unsupported
+from tests.kernels.quant_utils import FP8_DTYPE
+from vllm.kernels.helion.case_key import CaseKey
+from vllm.kernels.helion.config_manager import ConfigManager
+from vllm.kernels.helion.ops.per_token_group_fp8_quant import (
+    _pick_cache,
+    baseline,
+    per_token_group_fp8_quant,
+    pick_config,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    get_fp8_min_max,
+)
+from vllm.utils.import_utils import has_helion
+
+if not has_helion():
+    pytest.skip(
+        "Helion is not installed. Install with: pip install vllm[helion]",
+        allow_module_level=True,
+    )
+
+
+def _generate_fake_input(
+    num_tokens: int, hidden_size: int, group_size: int
+) -> tuple[Any, ...]:
+    with FakeTensorMode():
+        input = torch.randn(
+            (num_tokens, hidden_size), device="cuda", dtype=torch.bfloat16
+        )
+        output_q = torch.empty(input.shape, device=input.device, dtype=FP8_DTYPE)
+        output_s = torch.empty(
+            (num_tokens, hidden_size // group_size),
+            device=input.device,
+            dtype=torch.float32,
+        )
+        use_ue8m0 = False
+        column_major = False
+        fp8_min, fp8_max = get_fp8_min_max()
+        eps = 1e-10
+        args = (
+            input,
+            output_q,
+            output_s,
+            group_size,
+            eps,
+            fp8_min,
+            fp8_max,
+            use_ue8m0,
+            column_major,
+        )
+        return args
+
+
+@pytest.fixture(autouse=True)
+def reset_config_manager_singleton():
+    ConfigManager.reset_instance()
+    ConfigManager()
+    yield
+    ConfigManager.reset_instance()
+
+
+class TestPerTokenGroupFp8QuantConfigPicker:
+    def setup_method(self):
+        _pick_cache.clear()
+
+    def test_config_picker_exact_match(self):
+        config_keys = [
+            CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}),
+            CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}),
+        ]
+
+        args = _generate_fake_input(16, 4096, 128)
+        selected_key = pick_config(args, config_keys)
+        assert selected_key == CaseKey(
+            {"hidden_size": 4096, "group_size": 128, "num_tokens": 16}
+        )
+
+    def test_config_picker_closest_match(self):
+        config_keys = [
+            CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}),
+            CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 32}),
+            CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 16}),
+            CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 32}),
+            CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 16}),
+            CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 32}),
+            CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}),
+            CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 32}),
+        ]
+
+        args = _generate_fake_input(20, 3000, 70)
+        selected_key = pick_config(args, config_keys)
+        assert selected_key == CaseKey(
+            {"hidden_size": 2048, "group_size": 64, "num_tokens": 32}
+        )
+
+    def test_config_picker_no_configs(self):
+        config_keys: list[dict] = []
+
+        args = _generate_fake_input(16, 4096, 128)
+        selected_key = pick_config(args, config_keys)
+        assert selected_key is None
+
+    def test_config_picker_fallback_to_largest(self):
+        config_keys = [
+            CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}),
+            CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 32}),
+            CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 16}),
+            CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 32}),
+            CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 16}),
+            CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 32}),
+            CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}),
+            CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 32}),
+        ]
+
+        args = _generate_fake_input(64, 8192, 256)
+        selected_key = pick_config(args, config_keys)
+        assert selected_key == CaseKey(
+            {"hidden_size": 4096, "group_size": 128, "num_tokens": 32}
+        )
+
+
+class TestPerTokenGroupFp8QuantCorrectness:
+    @pytest.mark.parametrize(
+        "shape", [(31, 128), (32, 128), (63, 256), (64, 256), (16, 512), (2048, 5120)]
+    )
+    @pytest.mark.parametrize("column_major", [False, True])
+    @pytest.mark.parametrize("tma_aligned", [False, True])
+    @pytest.mark.parametrize("scale_ue8m0", [False, True])
+    @pytest.mark.parametrize("group_size", [64, 128])
+    def test_per_token_group_fp8_quant(
+        self,
+        shape,
+        column_major: bool,
+        tma_aligned: bool,
+        scale_ue8m0: bool,
+        group_size: int,
+    ):
+        skip_if_platform_unsupported("per_token_group_fp8_quant")
+
+        torch.manual_seed(42)
+        num_tokens, hidden_size = shape
+        fp8_min, fp8_max = get_fp8_min_max()
+        eps = 1e-10
+        input = (
+            torch.randn((num_tokens, hidden_size), device="cuda", dtype=torch.bfloat16)
+            * 8
+        )
+        ref_q = torch.empty(input.shape, device=input.device, dtype=FP8_DTYPE)
+        ops_q = ref_q.clone()
+
+        groups_per_row = hidden_size // group_size
+        if column_major:
+            if tma_aligned:
+                tma_alignment = 4
+                tma_aligned_m = (
+                    (num_tokens + tma_alignment - 1) // tma_alignment * tma_alignment
+                )
+                shape = (num_tokens, groups_per_row)
+                stride = (1, tma_aligned_m)
+                ref_s = torch.empty_strided(
+                    shape, stride, device=input.device, dtype=torch.float32
+                )
+            else:
+                ref_s = torch.empty(
+                    (groups_per_row, num_tokens),
+                    device=input.device,
+                    dtype=torch.float32,
+                ).transpose(0, 1)
+        else:
+            ref_s = torch.empty(
+                (num_tokens, groups_per_row), device=input.device, dtype=torch.float32
+            )
+
+        ops_s = ref_s.clone()
+
+        baseline(
+            input,
+            ref_q,
+            ref_s,
+            group_size,
+            eps,
+            fp8_min,
+            fp8_max,
+            scale_ue8m0,
+            column_major,
+            tma_aligned,
+        )
+        per_token_group_fp8_quant(
+            input,
+            ops_q,
+            ops_s,
+            group_size,
+            eps,
+            fp8_min,
+            fp8_max,
+            scale_ue8m0,
+            column_major,
+            tma_aligned,
+        )
+
+        assert torch.allclose(ref_s, ops_s)
+        # allow 1 ULP difference
+        assert (
+            ref_q.view(torch.uint8).to(torch.int16)
+            - ops_q.view(torch.uint8).to(torch.int16)
+        ).abs().max() <= 1
+
+
+class TestPerTokenGroupFp8QuantIntegration:
+    def test_kernel_registration_integration(self):
+        from vllm.kernels.helion.register import get_registered_kernels
+
+        registered_kernels = get_registered_kernels()
+        assert "per_token_group_fp8_quant" in registered_kernels
+
+        kernel_wrapper = registered_kernels["per_token_group_fp8_quant"]
+        assert kernel_wrapper.op_name == "per_token_group_fp8_quant"
+        assert kernel_wrapper._config_picker is not None
+        assert kernel_wrapper._mutates_args == ["output_q", "output_s"]
+
+    def test_fake_impl_functionality(self):
+        skip_if_platform_unsupported("per_token_group_fp8_quant")
+        from vllm.kernels.helion.register import get_registered_kernels
+
+        registered_kernels = get_registered_kernels()
+        kernel_wrapper = registered_kernels["per_token_group_fp8_quant"]
+        fake_impl = kernel_wrapper._fake_impl
+
+        args = _generate_fake_input(16, 4096, 128)
+        assert fake_impl(*args) is None
@@ -713,6 +713,7 @@ class TestHelionKernelWrapper:

        new_op = Mock()
        registered_ops: dict[str, Mock] = {}
+        mutates_args = ["y"]

        class MockNamespace:
            def __getattr__(self, name):
@@ -748,6 +749,7 @@ class TestHelionKernelWrapper:
                raw_kernel_func=sample_kernel,
                op_name="test_kernel",
                fake_impl=fake_impl,
+                mutates_args=mutates_args,
                config_picker=default_picker,
            )
            result = wrapper._get_or_register_custom_op()
@@ -755,6 +757,7 @@ class TestHelionKernelWrapper:
            mock_register.assert_called_once()
            assert result is new_op
            assert mock_register.call_args[1]["op_func"] is mock_decorated
+            assert mock_register.call_args[1]["mutates_args"] is mutates_args


 class TestKernelRegistry:
@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Helion Kernel test utils"""
+
+import pytest
+import torch
+
+from vllm.kernels.helion.config_manager import ConfigManager
+
+
+def skip_if_platform_unsupported(op_name: str):
+    try:
+        from vllm.kernels.helion.utils import get_canonical_gpu_name
+
+        if not torch.cuda.is_available():
+            pytest.skip("CUDA not available")
+
+        platform = get_canonical_gpu_name()
+
+        try:
+            config_manager = ConfigManager.get_instance()
+        except RuntimeError:
+            config_manager = ConfigManager()
+
+        configs = config_manager.get_platform_configs(op_name, platform)
+        if len(configs) == 0:
+            pytest.skip(f"Current GPU platform not supported for {op_name} kernel")
+
+    except (ImportError, RuntimeError, KeyError):
+        pytest.skip(f"Error detecting platform support for {op_name} kernel")
@@ -496,5 +496,258 @@ def test_mxfp4_cpu_fused_moe_bias_swiglu(M, N, K, E, topk, seed):
    torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)


+# INT4 W4A16 group-quantized MoE
+
+
+def _pack_int4_gptq(w_int4: torch.Tensor) -> torch.Tensor:
+    """Pack INT4 values [N, K] → [N, K//8] int32 along K dim (GPTQ format)."""
+    N, K = w_int4.shape
+    assert K % 8 == 0
+    w = w_int4.to(torch.int32)
+    w_packed = torch.zeros(N, K // 8, dtype=torch.int32)
+    for j in range(8):
+        w_packed |= (w[:, j::8] & 0xF) << (j * 4)
+    return w_packed
+
+
+def _pack_int4_awq(w_int4: torch.Tensor) -> torch.Tensor:
+    """Pack INT4 values [..., N] → [..., N//8] int32 along last dim (AWQ format)."""
+    # AWQ packing bitshifts: indices {0,4,1,5,2,6,3,7} * 4 bits each
+    _AWQ_BITSHIFTS = [0, 16, 4, 20, 8, 24, 12, 28]
+
+    N = w_int4.shape[-1]
+    assert N % 8 == 0
+    w = w_int4.to(torch.int32)
+    w_packed = torch.zeros(*w.shape[:-1], N // 8, dtype=torch.int32)
+    for j, shift in enumerate(_AWQ_BITSHIFTS):
+        w_packed |= (w[..., j::8] & 0xF) << shift
+    return w_packed
+
+
+def _ref_int4_moe(
+    a: torch.Tensor,
+    w1_int4: torch.Tensor,
+    w2_int4: torch.Tensor,
+    w1_zeros: torch.Tensor | None,
+    w2_zeros: torch.Tensor | None,
+    w1_s: torch.Tensor,
+    w2_s: torch.Tensor,
+    topk_weight: torch.Tensor,
+    topk_ids: torch.Tensor,
+    group_size: int,
+) -> torch.Tensor:
+    """Reference INT4 W4A16 group-quantized fused MoE in pure torch."""
+    B = a.shape[0]
+    topk = topk_ids.size(1)
+    K_out = a.shape[1]
+
+    out = torch.zeros(B, topk, K_out, dtype=torch.float32)
+    for b in range(B):
+        for t in range(topk):
+            eid = topk_ids[b, t].item()
+            x = a[b : b + 1].float()
+
+            # Dequantize w1: [K, 2*N], groups along K (input dim)
+            K_dim = w1_int4.shape[1]
+            w1_dq = torch.zeros(K_dim, w1_int4.shape[2], dtype=torch.float32)
+            for g in range(w1_s.shape[1]):
+                k_start = g * group_size
+                k_end = min((g + 1) * group_size, K_dim)
+                zp = w1_zeros[eid, g, :].float() if w1_zeros is not None else 8.0
+                w1_dq[k_start:k_end, :] = (
+                    w1_int4[eid, k_start:k_end, :].float() - zp
+                ) * w1_s[eid, g, :].float()
+
+            ic = torch.matmul(x, w1_dq)  # [1, K] @ [K, 2*N] → [1, 2*N]
+            ic = _silu_and_mul(ic)  # [1, N]
+
+            # Dequantize w2: [N, K], groups along N (input dim)
+            N_dim = w2_int4.shape[1]
+            w2_dq = torch.zeros(N_dim, w2_int4.shape[2], dtype=torch.float32)
+            for g in range(w2_s.shape[1]):
+                n_start = g * group_size
+                n_end = min((g + 1) * group_size, N_dim)
+                zp = w2_zeros[eid, g, :].float() if w2_zeros is not None else 8.0
+                w2_dq[n_start:n_end, :] = (
+                    w2_int4[eid, n_start:n_end, :].float() - zp
+                ) * w2_s[eid, g, :].float()
+
+            oc = torch.matmul(ic, w2_dq)  # [1, N] @ [N, K] → [1, K]
+            out[b, t] = oc.squeeze(0)
+
+    return (out * topk_weight.unsqueeze(-1)).sum(dim=1).to(a.dtype)
+
+
+def _make_int4_moe_weights(E, N, K, group_size, quant_algo):
+    """Create INT4 MoE weights in GPTQ or AWQ packed format.
+
+    Canonical layout (input × output):
+      w1_int4: [E, K, 2*N]  w2_int4: [E, N, K]
+
+    GPTQ packed (pack transposed weight along input/K dim):
+      w1_packed: [E, K//8, 2*N]  w2_packed: [E, N//8, K]
+      zeros: actual int4 zero points, same packing as weights
+
+    AWQ packed (pack along output/N dim):
+      w1_packed: [E, K, 2*N//8]  w2_packed: [E, N, K//8]
+      zeros: actual int4 zero points, same packing as weights
+
+    Returns:
+        w1_int4, w2_int4,
+        w1_packed, w2_packed,
+        w1_zeros, w2_zeros,
+        w1_zeros_packed, w2_zeros_packed,
+        w1_s, w2_s
+    """
+    w1_int4 = torch.randint(0, 16, (E, K, 2 * N), dtype=torch.int32)
+    w2_int4 = torch.randint(0, 16, (E, N, K), dtype=torch.int32)
+
+    num_groups_w1 = K // group_size
+    num_groups_w2 = N // group_size
+    w1_s = (
+        torch.randn(E, num_groups_w1, 2 * N, dtype=torch.bfloat16) * 0.01
+    ).abs() + 0.001
+    w2_s = (torch.randn(E, num_groups_w2, K, dtype=torch.bfloat16) * 0.01).abs() + 0.001
+
+    if quant_algo == ops.CPUQuantAlgo.GPTQ:
+        # Pack: canonical [E, K, 2*N] → transpose [E, 2*N, K] → GPTQ pack
+        # [E, 2*N, K//8] → transpose [E, K//8, 2*N]
+        w1_t = w1_int4.transpose(1, 2).contiguous()  # [E, 2*N, K]
+        w1_packed = (
+            torch.stack([_pack_int4_gptq(w1_t[e]) for e in range(E)])
+            .transpose(1, 2)
+            .contiguous()
+        )  # [E, K//8, 2*N]
+        w2_t = w2_int4.transpose(1, 2).contiguous()  # [E, K, N]
+        w2_packed = (
+            torch.stack([_pack_int4_gptq(w2_t[e]) for e in range(E)])
+            .transpose(1, 2)
+            .contiguous()
+        )  # [E, N//8, K]
+        w1_zeros = w2_zeros = None
+        w1_zeros_packed = torch.full(
+            (E, num_groups_w1, 2 * N // 8), 0x77777777, dtype=torch.int32
+        )
+        w2_zeros_packed = torch.full(
+            (E, num_groups_w2, K // 8), 0x77777777, dtype=torch.int32
+        )
+    else:  # AWQ
+        # Asymmetric: actual zero points, packed along output dim.
+        w1_zeros = torch.randint(1, 15, (E, num_groups_w1, 2 * N), dtype=torch.int32)
+        w2_zeros = torch.randint(1, 15, (E, num_groups_w2, K), dtype=torch.int32)
+        w1_packed = torch.stack(
+            [_pack_int4_awq(w1_int4[e]) for e in range(E)]
+        )  # [E, K, 2*N//8]
+        w2_packed = torch.stack(
+            [_pack_int4_awq(w2_int4[e]) for e in range(E)]
+        )  # [E, N, K//8]
+        w1_zeros_packed = torch.stack(
+            [_pack_int4_awq(w1_zeros[e]) for e in range(E)]
+        )  # [E, K//gs, 2*N//8]
+        w2_zeros_packed = torch.stack(
+            [_pack_int4_awq(w2_zeros[e]) for e in range(E)]
+        )  # [E, N//gs, K//8]
+
+    return (
+        w1_int4,
+        w2_int4,
+        w1_packed,
+        w2_packed,
+        w1_zeros,
+        w2_zeros,
+        w1_zeros_packed,
+        w2_zeros_packed,
+        w1_s,
+        w2_s,
+    )
+
+
+INT4_MOE_CONFIGS = [
+    # (N, K, E, topk, group_size)
+    (256, 512, 8, 2, 128),
+    (512, 256, 8, 2, 128),
+    (512, 512, 8, 4, 128),
+    (768, 2048, 8, 2, 128),
+]
+
+
+@pytest.mark.parametrize("M", [1, 2, 64, 121])
+@pytest.mark.parametrize("N,K,E,topk,group_size", INT4_MOE_CONFIGS)
+@pytest.mark.parametrize("quant_algo", [ops.CPUQuantAlgo.GPTQ, ops.CPUQuantAlgo.AWQ])
+@pytest.mark.parametrize("seed", [0])
+def test_int4_w4a16_cpu_fused_moe(M, N, K, E, topk, group_size, quant_algo, seed):
+    """Test fused_experts_cpu INT4 W4A16 for both GPTQ and AWQ quant formats."""
+    set_random_seed(seed)
+
+    a = torch.randn(M, K, dtype=torch.bfloat16) / (0.5 * K**0.5)
+    (
+        w1_int4,
+        w2_int4,
+        w1_packed,
+        w2_packed,
+        w1_zeros,
+        w2_zeros,
+        w1_zeros_packed,
+        w2_zeros_packed,
+        w1_s,
+        w2_s,
+    ) = _make_int4_moe_weights(E, N, K, group_size, quant_algo)
+
+    score = torch.randn(M, E, dtype=torch.bfloat16)
+    score = torch.softmax(score, dim=-1, dtype=torch.float32)
+    topk_weight, topk_ids = torch.topk(score, topk)
+    topk_ids = topk_ids.to(torch.int32)
+
+    ref_out = _ref_int4_moe(
+        a,
+        w1_int4,
+        w2_int4,
+        w1_zeros,
+        w2_zeros,
+        w1_s,
+        w2_s,
+        topk_weight,
+        topk_ids,
+        group_size,
+    )
+
+    from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
+        prepare_int4_moe_layer_for_cpu,
+    )
+
+    (blocked_w1, blocked_w2, blocked_s1, blocked_s2, blocked_z1, blocked_z2) = (
+        prepare_int4_moe_layer_for_cpu(
+            w1_packed,
+            w2_packed,
+            w1_s,
+            w2_s,
+            quant_algo=quant_algo,
+            w13_zeros=w1_zeros_packed,
+            w2_zeros=w2_zeros_packed,
+        )
+    )
+
+    out = ops.fused_experts_cpu(
+        a.clone(),
+        blocked_w1,
+        blocked_w2,
+        topk_weight,
+        topk_ids,
+        False,  # inplace
+        ops.CPUQuantMethod.INT4_W4A8,
+        blocked_s1,
+        blocked_s2,
+        blocked_z1,
+        blocked_z2,
+        None,  # block_size
+        None,  # w1_bias
+        None,  # w2_bias
+        None,  # alpha
+        None,  # limit
+        True,  # is_vnni
+    )
+    torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
+
+
 if __name__ == "__main__":
    sys.exit(pytest.main([__file__, "-v"]))
@@ -27,6 +27,7 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
    per_token_group_quant_fp8,
 )
+from vllm.platforms import current_platform
 from vllm.utils.import_utils import has_deep_ep
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.worker.workspace import init_workspace_manager
@@ -64,7 +65,7 @@ def make_weights(
        return w1, w2, None, None

    # per-out-channel weight quantization
-    assert dtype == torch.float8_e4m3fn
+    assert dtype == current_platform.fp8_dtype()
    w1 = torch.empty((e, 2 * n, k), device="cuda", dtype=torch.float16)
    w2 = torch.empty((e, k, n), device="cuda", dtype=torch.float16)

@@ -105,9 +106,11 @@ class TestTensors:
    @staticmethod
    def make(config: TestConfig, low_latency_mode: bool) -> "TestTensors":
        # TODO (varun) - check that float16 works ?
-        assert config.dtype in [torch.bfloat16, torch.float8_e4m3fn]
+        assert config.dtype in [torch.bfloat16, current_platform.fp8_dtype()]
        token_dtype = (
-            torch.bfloat16 if config.dtype == torch.float8_e4m3fn else config.dtype
+            torch.bfloat16
+            if config.dtype == current_platform.fp8_dtype()
+            else config.dtype
        )
        rank_tokens = (
            torch.randn((config.m, config.k), device="cuda", dtype=token_dtype) / 10
@@ -216,10 +219,10 @@ def deep_ep_moe_impl(
        return expert_map.to(device=device, dtype=torch.int32)

    hidden_size = test_tensors.rank_tokens.size(1)
-    is_quantized = w1.dtype == torch.float8_e4m3fn
+    is_quantized = w1.dtype == current_platform.fp8_dtype()
    q_dtype = None
    if is_quantized:
-        q_dtype = torch.float8_e4m3fn
+        q_dtype = current_platform.fp8_dtype()

    out_hidden_states = torch.empty_like(test_tensors.rank_tokens)
    total_num_tokens = test_tensors.rank_tokens.size(0)
@@ -318,7 +321,7 @@ def torch_moe_impl(
            .to(a.dtype)
        )

-    is_quantized = w1.dtype == torch.float8_e4m3fn
+    is_quantized = w1.dtype == current_platform.fp8_dtype()
    a_dtype = a.dtype
    if is_quantized:
        w1 = w1.to(dtype=torch.float32) * w1_scale
@@ -367,7 +370,7 @@ def _deep_ep_moe(
            "FP8 dispatch interface is available only in low-latency mode"
        )

-    is_quantized = w1.dtype == torch.float8_e4m3fn
+    is_quantized = w1.dtype == current_platform.fp8_dtype()
    device_idx = torch.accelerator.current_device_index()
    w1 = w1.to(device=device_idx)
    w2 = w2.to(device=device_idx)
@@ -441,7 +444,7 @@ MNKs = [
    (222, 1024, 2048),
 ]

-DTYPES = [torch.bfloat16, torch.float8_e4m3fn]
+DTYPES = [torch.bfloat16, current_platform.fp8_dtype()]


@pytest.mark.parametrize("dtype", DTYPES)
@@ -496,7 +499,7 @@ MNKs = [
    (64, 1024, 2560),
    (222, 1024, 2560),
 ]
-DTYPES = [torch.float8_e4m3fn, torch.bfloat16]
+DTYPES = [current_platform.fp8_dtype(), torch.bfloat16]
 USE_FP8_DISPATCH = [True, False]


@@ -1,11 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+import types
+from unittest.mock import patch
+
 import pytest

 from vllm import SamplingParams
 from vllm.config.load import LoadConfig
 from vllm.model_executor.model_loader import get_model_loader
+from vllm.model_executor.model_loader import runai_streamer_loader as rsl

 load_format = "runai_streamer"
 test_model = "openai-community/gpt2"
@@ -53,3 +57,24 @@ def test_runai_model_loader_download_files_gcs(
    with vllm_runner(test_gcs_model, load_format=load_format) as llm:
        deserialized_outputs = llm.generate(prompts, sampling_params)
        assert deserialized_outputs
+
+
+def test_runai_passes_revision_by_name():
+    # revision must reach download_safetensors_index_file_from_hf as the
+    # ``revision`` keyword, not the positional ``subfolder`` slot.
+    fake_self = types.SimpleNamespace(
+        load_config=types.SimpleNamespace(download_dir="/cache", ignore_patterns=[])
+    )
+    with (
+        patch.object(rsl, "is_runai_obj_uri", return_value=False),
+        patch.object(rsl, "download_weights_from_hf", return_value="/folder"),
+        patch.object(
+            rsl, "list_safetensors", return_value=["/folder/model.safetensors"]
+        ),
+        patch.object(rsl, "download_safetensors_index_file_from_hf") as mock_idx,
+    ):
+        rsl.RunaiModelStreamerLoader._prepare_weights(fake_self, "org/model", "myrev")
+
+    mock_idx.assert_called_once()
+    assert mock_idx.call_args.kwargs.get("revision") == "myrev"
+    assert "myrev" not in mock_idx.call_args.args
@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm.config.compilation import CompilationMode
+from vllm.model_executor.models import deepseek_v2 as deepseek_mod
+from vllm.model_executor.models import mistral_large_3_eagle as eagle_mod
+
+
+class DummyPPGroup:
+    world_size = 1
+    is_first_rank = True
+    is_last_rank = True
+
+
+class DummyEmbedding(nn.Module):
+    def __init__(self, vocab_size, hidden_size, *args, **kwargs):
+        super().__init__()
+        self.hidden_size = hidden_size
+
+    def forward(self, input_ids):
+        return torch.zeros(
+            (*input_ids.shape, self.hidden_size),
+            dtype=torch.float32,
+            device=input_ids.device,
+        )
+
+
+class DummyLinear(nn.Module):
+    def __init__(self, in_features, out_features, *args, **kwargs):
+        super().__init__()
+        self.out_features = out_features
+
+    def forward(self, x):
+        return torch.zeros(
+            (*x.shape[:-1], self.out_features),
+            dtype=x.dtype,
+            device=x.device,
+        )
+
+
+class DummyNorm(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+    def forward(self, hidden_states, residual=None):
+        return hidden_states, residual
+
+
+class DummyDecoderLayer(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+    def forward(self, positions, hidden_states, residual, llama_4_scaling=None):
+        return hidden_states, residual
+
+
+def make_vllm_config(
+    *, model_type="mistral3", qk_nope_head_dim=128, qk_rope_head_dim=64
+):
+    hf_config = SimpleNamespace(
+        model_type=model_type,
+        first_k_dense_replace=0,
+        vocab_size=32000,
+        hidden_size=16,
+        num_hidden_layers=1,
+        rms_norm_eps=1e-5,
+        qk_nope_head_dim=qk_nope_head_dim,
+        qk_rope_head_dim=qk_rope_head_dim,
+    )
+
+    return SimpleNamespace(
+        model_config=SimpleNamespace(hf_config=hf_config),
+        quant_config=None,
+        parallel_config=SimpleNamespace(
+            eplb_config=SimpleNamespace(num_redundant_experts=0),
+        ),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
+        cache_config=None,
+        compilation_config=SimpleNamespace(mode=CompilationMode.NONE),
+    )
+
+
+@pytest.fixture(autouse=True)
+def patch_heavy_modules(monkeypatch):
+    monkeypatch.setattr(eagle_mod, "get_pp_group", lambda: DummyPPGroup())
+    monkeypatch.setattr(deepseek_mod, "get_pp_group", lambda: DummyPPGroup())
+
+    monkeypatch.setattr(eagle_mod, "VocabParallelEmbedding", DummyEmbedding)
+    monkeypatch.setattr(eagle_mod, "RowParallelLinear", DummyLinear)
+    monkeypatch.setattr(eagle_mod, "RMSNorm", DummyNorm)
+    monkeypatch.setattr(eagle_mod, "DeepseekV2DecoderLayer", DummyDecoderLayer)
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize(
+    ("model_type", "qk_nope_head_dim", "qk_rope_head_dim", "expected_use_mha"),
+    [
+        # MLA-style config: should not use MHA.
+        ("mistral3", 128, 64, False),
+        # No MLA dims: should use MHA, matching DeepseekV2Model.__init__ logic.
+        ("mistral3", 0, 0, True),
+        # DeepSeek model type always uses MHA by the parent logic.
+        ("deepseek", 128, 64, True),
+    ],
+)
+def test_eagle_mistral_large3_initializes_deepseek_runtime_attrs(
+    model_type,
+    qk_nope_head_dim,
+    qk_rope_head_dim,
+    expected_use_mha,
+):
+    vllm_config = make_vllm_config(
+        model_type=model_type,
+        qk_nope_head_dim=qk_nope_head_dim,
+        qk_rope_head_dim=qk_rope_head_dim,
+    )
+
+    model = eagle_mod.EagleMistralLarge3Model(vllm_config=vllm_config)
+
+    assert model.aux_hidden_state_layers == ()
+    assert model.use_mha is expected_use_mha
+
+    # Add this if your fix also copies num_redundant_experts from
+    # DeepseekV2Model.__init__.
+    assert model.num_redundant_experts == 0
+
+
+@pytest.mark.cpu_test
+def test_eagle_mistral_large3_forward_reuses_deepseek_parent_forward():
+    vllm_config = make_vllm_config()
+    model = eagle_mod.EagleMistralLarge3Model(vllm_config=vllm_config)
+
+    input_ids = torch.tensor([[1, 2, 3]])
+    positions = torch.tensor([[0, 1, 2]])
+    hidden_states = torch.zeros((1, 3, 16))
+
+    output = model(input_ids, positions, hidden_states)
+
+    assert isinstance(output, torch.Tensor)
+    assert output.shape == hidden_states.shape
@@ -604,8 +604,6 @@ VLM_TEST_SETTINGS = {
        models=[
            "OpenGVLab/InternVL2-1B",
            "OpenGVLab/InternVL2-2B",
-            # FIXME: Config cannot be loaded in transformers 4.52
-            # "OpenGVLab/Mono-InternVL-2B",
        ],
        test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
        prompt_formatter=lambda img_prompt: f"<|im_start|>User\n{img_prompt}<|im_end|>\n<|im_start|>Assistant\n",  # noqa: E501
@@ -55,6 +55,26 @@ def step3_vl_chat_template(content: str) -> str:


 MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
+    "llama4": VitCudagraphTestConfig(
+        model="meta-llama/Llama-4-Scout-17B-16E-Instruct",
+        modalities=["image"],
+        image_prompt=(
+            "<|begin_of_text|><|header_start|>user<|header_end|>\n\n"
+            "<|image|>What is in this image?<|eot|>"
+            "<|header_start|>assistant<|header_end|>\n\n"
+        ),
+        max_model_len=4096,
+        max_tokens=32,
+        max_num_seqs=2,
+        vllm_runner_kwargs={
+            "load_format": "dummy",
+            "hf_overrides": partial(
+                dummy_hf_overrides,
+                model_arch="Llama4ForConditionalGeneration",
+            ),
+        },
+        marks=[pytest.mark.core_model],
+    ),
    "internvl": VitCudagraphTestConfig(
        model="OpenGVLab/InternVL3-1B",
        num_video_frames=8,
@@ -5,12 +5,16 @@
 Run `pytest tests/quantization/test_bitsandbytes.py`.
 """

+import types
+from unittest.mock import MagicMock, patch
+
 import pytest
 from packaging.version import Version
 from transformers import BitsAndBytesConfig
 from transformers import __version__ as TRANSFORMERS_VERSION

 from tests.quantization.utils import is_quant_method_supported
+from vllm.model_executor.model_loader import bitsandbytes_loader as bnb
 from vllm.platforms import current_platform

 from ...utils import compare_two_settings, multi_gpu_test
@@ -300,3 +304,27 @@ def validate_generated_texts(
            f"HF Output: '{hf_str}'\n"
            f"vLLM Output: '{vllm_str}'"
        )
+
+
+def test_bitsandbytes_passes_revision_by_name():
+    # revision must reach download_safetensors_index_file_from_hf as the
+    # ``revision`` keyword, not a positional slot.
+    fake_self = types.SimpleNamespace(
+        load_config=types.SimpleNamespace(download_dir="/cache"),
+        _get_weight_files=MagicMock(
+            return_value=("/folder", ["/folder/model.safetensors"], "*.safetensors")
+        ),
+    )
+    with (
+        patch.object(bnb, "download_safetensors_index_file_from_hf") as mock_idx,
+        patch.object(
+            bnb,
+            "filter_duplicate_safetensors_files",
+            return_value=["/folder/model.safetensors"],
+        ),
+    ):
+        bnb.BitsAndBytesModelLoader._prepare_weights(fake_self, "org/model", "myrev")
+
+    mock_idx.assert_called_once()
+    assert mock_idx.call_args.kwargs.get("revision") == "myrev"
+    assert "myrev" not in mock_idx.call_args.args
@@ -338,23 +338,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
        "naver-hyperclovax/HyperCLOVAX-SEED-Think-14B",
        min_transformers_version="5.9.0",
    ),
-    "InternLMForCausalLM": _HfExamplesInfo(
-        "internlm/internlm-chat-7b", trust_remote_code=True
-    ),
    "InternLM2ForCausalLM": _HfExamplesInfo(
        "internlm/internlm2-chat-7b", trust_remote_code=True
    ),
-    "InternLM2VEForCausalLM": _HfExamplesInfo(
-        "OpenGVLab/Mono-InternVL-2B",
-        trust_remote_code=True,
-        max_transformers_version="4.57",
-        transformers_version_reason={
-            "vllm": (
-                "Custom config cannot be loaded with Transformers "
-                "v5 because `vision_config` is not always set"
-            )
-        },
-    ),
    "InternLM3ForCausalLM": _HfExamplesInfo(
        "internlm/internlm3-8b-instruct", trust_remote_code=True
    ),
@@ -901,6 +887,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {
    ),
    "FuyuForCausalLM": _HfExamplesInfo("adept/fuyu-8b"),
    "Gemma3ForConditionalGeneration": _HfExamplesInfo("google/gemma-3-4b-it"),
+    "DiffusionGemmaForBlockDiffusion": _HfExamplesInfo(
+        "google/diffusiongemma-26B-A4B-it",
+        trust_remote_code=True,
+    ),
    "Gemma4ForConditionalGeneration": _HfExamplesInfo(
        "google/gemma-4-E2B-it",
        min_transformers_version="5.5.0",
@@ -486,6 +486,7 @@ def dummy_hf_overrides(
                "Gemma3nForConditionalGeneration",
                "Gemma4ForCausalLM",
                "Gemma4ForConditionalGeneration",
+                "DiffusionGemmaForBlockDiffusion",
            )
            else 1
        )
@@ -506,12 +507,13 @@ def dummy_hf_overrides(
    # Only set MoE related config when the model has MoE layers.
    # Otherwise all models detected as MoE by _get_transformers_backend_cls.
    if model_arch_config.num_experts > 0:
+        num_experts_per_tok = 1 if model_arch == "Llama4ForConditionalGeneration" else 2
        update_dict.update(
            {
                "num_experts": num_experts,
-                "num_experts_per_tok": 2,
+                "num_experts_per_tok": num_experts_per_tok,
                # Kimi uses `num_experts_per_token`.
-                "num_experts_per_token": 2,
+                "num_experts_per_token": num_experts_per_tok,
                "num_local_experts": num_experts,
                # Otherwise there will not be any expert layers
                "first_k_dense_replace": 0,
@@ -558,7 +560,8 @@ def dummy_hf_overrides(
        )

    # e.g.: Qwen/Qwen2-Audio-7B-Instruct
-    if hasattr(hf_config, "audio_config"):
+    # audio_config may exist but be None (e.g. audio-less Gemma4 variants).
+    if getattr(hf_config, "audio_config", None) is not None:
        hf_config.audio_config.update(
            {
                "num_layers": 1,
@@ -0,0 +1,736 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+from collections.abc import Sequence
+
+import pytest
+from openai_harmony import (
+    Conversation,
+    Message,
+    RenderConversationConfig,
+    Role,
+)
+from transformers import AutoTokenizer
+
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.entrypoints.openai.engine.protocol import FunctionCall
+from vllm.entrypoints.openai.parser.harmony_utils import (
+    get_encoding,
+)
+from vllm.parser.harmony import HarmonyParser
+from vllm.parser.parser_manager import ParserManager
+
+REASONING_MODEL_NAME = "openai/gpt-oss-20b"
+
+
+@pytest.fixture(scope="module")
+def gpt_oss_tokenizer():
+    return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME)
+
+
+@pytest.fixture
+def harmony_parser(gpt_oss_tokenizer):
+    parser_cls = ParserManager.get_parser(
+        tool_parser_name="openai",
+        reasoning_parser_name="openai_gptoss",
+        enable_auto_tools=True,
+        model_name=REASONING_MODEL_NAME,
+        is_harmony=True,
+    )
+    assert parser_cls is HarmonyParser
+    return parser_cls(gpt_oss_tokenizer)
+
+
+@pytest.fixture
+def chat_request():
+    return ChatCompletionRequest(
+        model="openai/gpt-oss-20b",
+        messages=[{"role": "user", "content": "Hello"}],
+    )
+
+
+def encode_output(harmony_str: str) -> list[int]:
+    return get_encoding().encode(harmony_str, allowed_special="all")
+
+
+def assistant(content: str, channel: str) -> Message:
+    return Message.from_role_and_content(Role.ASSISTANT, content).with_channel(channel)
+
+
+def tool_call(
+    recipient: str,
+    content: str,
+    channel: str = "commentary",
+    content_type: str | None = "json",
+) -> Message:
+    message = assistant(content, channel).with_recipient(recipient)
+    return message if content_type is None else message.with_content_type(content_type)
+
+
+def get_model_output_tokens(
+    prompt_messages: Sequence[Message],
+    response_messages: Sequence[Message],
+) -> list[int]:
+    enc = get_encoding()
+    # Keep analysis messages when synthesizing model-output-only token sequences
+    # for parser tests; the default render path drops them after a later final turn.
+    config = RenderConversationConfig(auto_drop_analysis=False)
+    prompt_ids = enc.render_conversation_for_completion(
+        Conversation.from_messages(list(prompt_messages)),
+        Role.ASSISTANT,
+        config=config,
+    )
+    full_ids = enc.render_conversation_for_completion(
+        Conversation.from_messages([*prompt_messages, *response_messages]),
+        Role.ASSISTANT,
+        config=config,
+    )
+    assert full_ids[: len(prompt_ids)] == prompt_ids
+    return full_ids[len(prompt_ids) :]
+
+
+def get_text(msg: Message) -> str:
+    return msg.content[0].text if msg.content else ""
+
+
+def tool_call_tuples(tool_calls: list[FunctionCall] | None) -> list[tuple[str, str]]:
+    return [] if tool_calls is None else [(tc.name, tc.arguments) for tc in tool_calls]
+
+
+def tool_call_headers(delta_message) -> list:
+    if delta_message is None or not delta_message.tool_calls:
+        return []
+    return [
+        tool_call
+        for tool_call in delta_message.tool_calls
+        if tool_call.function and tool_call.function.name
+    ]
+
+
+def tool_call_payloads(delta_message) -> list:
+    if delta_message is None or not delta_message.tool_calls:
+        return []
+    return [
+        tool_call
+        for tool_call in delta_message.tool_calls
+        if tool_call.function and tool_call.function.arguments
+    ]
+
+
+def combined_tool_arguments(delta_message) -> dict[int, str]:
+    combined: dict[int, str] = {}
+    for tool_call in tool_call_payloads(delta_message):
+        combined.setdefault(tool_call.index, "")
+        combined[tool_call.index] += tool_call.function.arguments
+    return combined
+
+
+class TestParse:
+    # Rendered conversation outputs.
+
+    def test_reasoning_only(self, harmony_parser, chat_request):
+        prompt = [Message.from_role_and_content(Role.USER, "Why?")]
+        response = [assistant("This is reasoning", "analysis")]
+
+        reasoning, content, tool_calls = harmony_parser.parse(
+            "",
+            chat_request,
+            model_output_token_ids=get_model_output_tokens(prompt, response),
+        )
+
+        assert reasoning == "This is reasoning"
+        assert content is None
+        assert tool_calls is None
+
+    def test_content_only(self, harmony_parser, chat_request):
+        prompt = [Message.from_role_and_content(Role.USER, "Hello")]
+        response = [assistant("This is a test", "final")]
+
+        reasoning, content, tool_calls = harmony_parser.parse(
+            "",
+            chat_request,
+            model_output_token_ids=get_model_output_tokens(prompt, response),
+        )
+
+        assert reasoning is None
+        assert content == "This is a test"
+        assert tool_calls is None
+
+    def test_reasoning_and_content(self, harmony_parser, chat_request):
+        prompt = [Message.from_role_and_content(Role.USER, "What is 2+2?")]
+        response = [
+            assistant("I should think first.", "analysis"),
+            assistant("The answer is 4.", "final"),
+        ]
+
+        reasoning, content, tool_calls = harmony_parser.parse(
+            "",
+            chat_request,
+            model_output_token_ids=get_model_output_tokens(prompt, response),
+        )
+
+        assert reasoning == "I should think first."
+        assert content == "The answer is 4."
+        assert tool_calls is None
+
+    @pytest.mark.parametrize(
+        "tool_args",
+        [
+            '{"location": "Tokyo"}',
+            '{\n"location": "Tokyo"\n}',
+        ],
+    )
+    @pytest.mark.parametrize("tool_channel", ["commentary", "analysis"])
+    def test_single_tool_call(
+        self, harmony_parser, chat_request, tool_args, tool_channel
+    ):
+        prompt = [
+            Message.from_role_and_content(Role.USER, "What is the weather in Tokyo?")
+        ]
+        response = [tool_call("functions.get_current_weather", tool_args, tool_channel)]
+
+        reasoning, content, tool_calls = harmony_parser.parse(
+            "",
+            chat_request,
+            model_output_token_ids=get_model_output_tokens(prompt, response),
+        )
+
+        assert reasoning is None
+        assert content is None
+        assert tool_call_tuples(tool_calls) == [
+            ("get_current_weather", json.dumps({"location": "Tokyo"}))
+        ]
+
+    def test_multiple_tool_calls_varied_formats(self, harmony_parser, chat_request):
+        prompt = [
+            Message.from_role_and_content(
+                Role.USER, "What is the weather in Tokyo based on where I'm at?"
+            )
+        ]
+        response = [
+            tool_call("functions.get_current_weather", '{"location": "Tokyo"}'),
+            tool_call("functions.get_user_location", '{"location": "Tokyo"}'),
+            tool_call(
+                "functions.no_content_type",
+                '{"location": "Tokyo"}',
+                content_type=None,
+            ),
+            tool_call("functions.not_json_no_content_type", "foo", content_type=None),
+            tool_call("functions.empty_args", "{}"),
+            tool_call("functions.no_args", ""),
+        ]
+
+        _, content, tool_calls = harmony_parser.parse(
+            "",
+            chat_request,
+            model_output_token_ids=get_model_output_tokens(prompt, response),
+        )
+
+        assert content is None
+        assert tool_call_tuples(tool_calls) == [
+            ("get_current_weather", json.dumps({"location": "Tokyo"})),
+            ("get_user_location", json.dumps({"location": "Tokyo"})),
+            ("no_content_type", json.dumps({"location": "Tokyo"})),
+            ("not_json_no_content_type", "foo"),
+            ("empty_args", json.dumps({})),
+            ("no_args", ""),
+        ]
+
+    def test_tool_call_bare_recipient(self, harmony_parser, chat_request):
+        prompt = [Message.from_role_and_content(Role.USER, "Weather?")]
+        response = [tool_call("get_current_weather", '{"location": "Tokyo"}')]
+
+        _, _, tool_calls = harmony_parser.parse(
+            "",
+            chat_request,
+            model_output_token_ids=get_model_output_tokens(prompt, response),
+        )
+
+        assert tool_call_tuples(tool_calls) == [
+            ("get_current_weather", json.dumps({"location": "Tokyo"}))
+        ]
+
+    def test_multiple_tool_calls_bare_recipients(self, harmony_parser, chat_request):
+        prompt = [Message.from_role_and_content(Role.USER, "Use both tools.")]
+        response = [
+            tool_call("get_current_weather", '{"location": "Tokyo"}'),
+            tool_call("get_user_location", "{}"),
+        ]
+
+        _, _, tool_calls = harmony_parser.parse(
+            "",
+            chat_request,
+            model_output_token_ids=get_model_output_tokens(prompt, response),
+        )
+
+        assert tool_call_tuples(tool_calls) == [
+            ("get_current_weather", json.dumps({"location": "Tokyo"})),
+            ("get_user_location", json.dumps({})),
+        ]
+
+    def test_assistant_recipient_not_tool(self, harmony_parser, chat_request):
+        prompt = [Message.from_role_and_content(Role.USER, "Hello")]
+        response = [
+            tool_call("assistant", "Some tool response", content_type=None),
+            assistant("Here is the answer", "final"),
+        ]
+
+        reasoning, content, tool_calls = harmony_parser.parse(
+            "",
+            chat_request,
+            model_output_token_ids=get_model_output_tokens(prompt, response),
+        )
+
+        assert reasoning is None
+        assert content == "Here is the answer"
+        assert tool_calls is None
+
+    def test_tool_call_dotted_name(self, harmony_parser, chat_request):
+        prompt = [Message.from_role_and_content(Role.USER, "Compute 2+3")]
+        response = [tool_call("math.sum", '{"a": 2, "b": 3}')]
+
+        _, _, tool_calls = harmony_parser.parse(
+            "",
+            chat_request,
+            model_output_token_ids=get_model_output_tokens(prompt, response),
+        )
+
+        assert tool_call_tuples(tool_calls) == [
+            ("math.sum", json.dumps({"a": 2, "b": 3}))
+        ]
+
+    def test_tool_calls_with_final_content(self, harmony_parser, chat_request):
+        prompt = [Message.from_role_and_content(Role.USER, "What is the weather?")]
+        response = [
+            assistant("User asked about the weather.", "analysis"),
+            tool_call("functions.get_current_weather", '{"location": "Tokyo"}'),
+            assistant("This tool call will get the weather.", "final"),
+        ]
+
+        reasoning, content, tool_calls = harmony_parser.parse(
+            "",
+            chat_request,
+            model_output_token_ids=get_model_output_tokens(prompt, response),
+        )
+
+        assert reasoning == "User asked about the weather."
+        assert content == "This tool call will get the weather."
+        assert tool_call_tuples(tool_calls) == [
+            ("get_current_weather", json.dumps({"location": "Tokyo"}))
+        ]
+
+    # Raw/truncated Harmony output streams.
+
+    def test_interrupted_first_message(self, harmony_parser, chat_request):
+        reasoning, content, tool_calls = harmony_parser.parse(
+            "",
+            chat_request,
+            model_output_token_ids=encode_output(
+                "<|channel|>final<|message|>I'm in the middle of answering"
+            ),
+        )
+
+        assert reasoning is None
+        assert content == "I'm in the middle of answering"
+        assert tool_calls is None
+
+    def test_interrupted_reasoning_first_message(self, harmony_parser, chat_request):
+        reasoning, content, tool_calls = harmony_parser.parse(
+            "",
+            chat_request,
+            model_output_token_ids=encode_output(
+                "<|channel|>analysis<|message|>I'm in the middle of thinking"
+            ),
+        )
+
+        assert reasoning == "I'm in the middle of thinking"
+        assert content is None
+        assert tool_calls is None
+
+    def test_truncated_output(self, harmony_parser, chat_request):
+        reasoning, content, tool_calls = harmony_parser.parse(
+            "",
+            chat_request,
+            model_output_token_ids=encode_output(
+                "<|channel|>analysis<|message|>I'm thinking.<|end|>"
+                "<|start|>assistant<|channel|>final<|message|>"
+                "I'm in the middle of answering"
+            ),
+        )
+
+        assert reasoning == "I'm thinking."
+        assert content == "I'm in the middle of answering"
+        assert tool_calls is None
+
+    @pytest.mark.parametrize(
+        ("harmony_str", "expected_content"),
+        [
+            (
+                "<|channel|>commentary<|message|>I'll search for that",
+                "I'll search for that",
+            ),
+            (
+                "<|channel|>commentary<|message|>Let me look that up.<|end|>"
+                "<|start|>assistant<|channel|>final<|message|>The answer is 42.<|end|>",
+                "Let me look that up.\nThe answer is 42.",
+            ),
+        ],
+    )
+    def test_commentary_preambles(
+        self,
+        harmony_parser,
+        chat_request,
+        harmony_str,
+        expected_content,
+    ):
+        reasoning, content, tool_calls = harmony_parser.parse(
+            "",
+            chat_request,
+            model_output_token_ids=encode_output(harmony_str),
+        )
+
+        assert reasoning is None
+        assert content == expected_content
+        assert tool_calls is None
+
+    def test_commentary_with_recipient_excluded(self, harmony_parser, chat_request):
+        reasoning, content, tool_calls = harmony_parser.parse(
+            "",
+            chat_request,
+            model_output_token_ids=encode_output(
+                "<|channel|>commentary"
+                "<|message|>Let me check the weather.<|end|>"
+                "<|start|>assistant to=functions.get_weather"
+                "<|channel|>commentary"
+                '<|message|>{"location": "SF"}<|end|>'
+            ),
+        )
+
+        assert reasoning is None
+        assert content == "Let me check the weather."
+        assert tool_call_tuples(tool_calls) == [
+            ("get_weather", json.dumps({"location": "SF"}))
+        ]
+
+
+class TestParseDelta:
+    def test_basic(self, gpt_oss_tokenizer, chat_request):
+        parser = HarmonyParser(gpt_oss_tokenizer)
+
+        first_delta = parser.parse_delta(
+            delta_text="",
+            delta_token_ids=encode_output("<|channel|>analysis<|message|>Thinking"),
+            request=chat_request,
+            finished=False,
+        )
+        second_delta = parser.parse_delta(
+            delta_text="",
+            delta_token_ids=encode_output(
+                "<|end|><|start|>assistant<|channel|>final<|message|>Answer"
+            ),
+            request=chat_request,
+            finished=False,
+        )
+
+        assert first_delta is not None
+        assert first_delta.reasoning == "Thinking"
+        assert first_delta.content is None
+        assert second_delta is not None
+        assert second_delta.content == "Answer"
+        assert second_delta.reasoning is None
+
+    def test_multi_token(self, gpt_oss_tokenizer, chat_request):
+        parser = HarmonyParser(gpt_oss_tokenizer)
+
+        delta = parser.parse_delta(
+            delta_text="",
+            delta_token_ids=encode_output("<|channel|>final<|message|>Hello, world!"),
+            request=chat_request,
+            finished=False,
+        )
+
+        assert delta is not None
+        assert delta.content == "Hello, world!"
+        assert delta.reasoning is None
+        assert not delta.tool_calls
+
+    @pytest.mark.parametrize("tool_channel", ["commentary", "analysis"])
+    def test_tool_call_split_across_deltas(
+        self, gpt_oss_tokenizer, chat_request, tool_channel
+    ):
+        parser = HarmonyParser(gpt_oss_tokenizer)
+
+        first_delta = parser.parse_delta(
+            delta_text="",
+            delta_token_ids=encode_output(
+                "<|channel|>analysis<|message|>Thinking<|end|>"
+                f"<|start|>assistant to=functions.get_weather<|channel|>{tool_channel}"
+                '<|constrain|>json<|message|>{"location": '
+            ),
+            request=chat_request,
+            finished=False,
+        )
+        second_delta = parser.parse_delta(
+            delta_text="",
+            delta_token_ids=encode_output('"Paris"}<|call|>'),
+            request=chat_request,
+            finished=False,
+        )
+
+        assert first_delta is not None
+        assert first_delta.reasoning == "Thinking"
+        assert first_delta.content is None
+        assert [tool.function.name for tool in tool_call_headers(first_delta)] == [
+            "get_weather"
+        ]
+        assert combined_tool_arguments(first_delta) == {0: '{"location": '}
+        assert {tool.index for tool in first_delta.tool_calls} == {0}
+
+        assert second_delta is not None
+        assert second_delta.reasoning is None
+        assert second_delta.content is None
+        assert not tool_call_headers(second_delta)
+        assert combined_tool_arguments(second_delta) == {0: '"Paris"}'}
+        assert {tool.index for tool in second_delta.tool_calls} == {0}
+
+    def test_commentary_preamble_streaming(self, gpt_oss_tokenizer, chat_request):
+        parser = HarmonyParser(gpt_oss_tokenizer)
+
+        delta = parser.parse_delta(
+            delta_text="",
+            delta_token_ids=encode_output(
+                "<|channel|>commentary<|message|>I'll search for that"
+            ),
+            request=chat_request,
+            finished=False,
+        )
+
+        assert delta is not None
+        assert delta.content == "I'll search for that"
+        assert delta.reasoning is None
+        assert not delta.tool_calls
+
+    def test_multiple_choices(self, gpt_oss_tokenizer, chat_request):
+        parser_a = HarmonyParser(gpt_oss_tokenizer)
+        parser_b = HarmonyParser(gpt_oss_tokenizer)
+
+        delta_a = parser_a.parse_delta(
+            delta_text="",
+            delta_token_ids=encode_output(
+                "<|channel|>analysis<|message|>Check weather<|end|>"
+                "<|start|>assistant to=functions.get_weather<|channel|>commentary"
+                '<|constrain|>json<|message|>{"location": "Paris"}'
+            ),
+            request=chat_request,
+            finished=False,
+        )
+        delta_b = parser_b.parse_delta(
+            delta_text="",
+            delta_token_ids=encode_output(
+                "<|channel|>analysis<|message|>Check time<|end|>"
+                "<|start|>assistant to=functions.get_time<|channel|>commentary"
+                '<|constrain|>json<|message|>{"timezone": "UTC"}'
+            ),
+            request=chat_request,
+            finished=False,
+        )
+
+        assert [tool.function.name for tool in tool_call_headers(delta_a)] == [
+            "get_weather"
+        ]
+        assert [tool.function.name for tool in tool_call_headers(delta_b)] == [
+            "get_time"
+        ]
+        assert {tool.index for tool in delta_a.tool_calls} == {0}
+        assert {tool.index for tool in delta_b.tool_calls} == {0}
+
+    def test_dotted_function_name(self, gpt_oss_tokenizer, chat_request):
+        parser = HarmonyParser(gpt_oss_tokenizer)
+
+        delta = parser.parse_delta(
+            delta_text="",
+            delta_token_ids=encode_output(
+                "<|channel|>analysis<|message|>Compute this<|end|>"
+                "<|start|>assistant to=math.sum<|channel|>commentary"
+                '<|constrain|>json<|message|>{"a": 2, "b": 3}'
+            ),
+            request=chat_request,
+            finished=False,
+        )
+
+        assert delta is not None
+        assert [tool.function.name for tool in tool_call_headers(delta)] == ["math.sum"]
+        assert {tool.index for tool in delta.tool_calls} == {0}
+
+    @pytest.mark.parametrize("recipient", ["assistant", "browser"])
+    def test_builtin_recipient_skipped(
+        self,
+        gpt_oss_tokenizer,
+        chat_request,
+        recipient,
+    ):
+        parser = HarmonyParser(gpt_oss_tokenizer)
+        prompt = [Message.from_role_and_content(Role.USER, "Hello")]
+        response = [tool_call(recipient, "Ignore this", content_type=None)]
+
+        delta = parser.parse_delta(
+            delta_text="",
+            delta_token_ids=get_model_output_tokens(prompt, response),
+            request=chat_request,
+            finished=False,
+        )
+
+        assert delta is None
+
+    def test_cross_channel_with_tool(self, gpt_oss_tokenizer, chat_request):
+        parser = HarmonyParser(gpt_oss_tokenizer)
+
+        delta = parser.parse_delta(
+            delta_text="",
+            delta_token_ids=encode_output(
+                "<|channel|>analysis<|message|>Reasoning about query...<|end|>"
+                "<|start|>assistant to=functions.search<|channel|>commentary"
+                '<|constrain|>json<|message|>{"query": "vllm"}<|call|>'
+                "<|start|>assistant<|channel|>final<|message|>Done"
+            ),
+            request=chat_request,
+            finished=False,
+        )
+
+        assert delta is not None
+        assert delta.reasoning == "Reasoning about query..."
+        assert delta.content == "Done"
+        assert [tool.function.name for tool in tool_call_headers(delta)] == ["search"]
+        assert combined_tool_arguments(delta) == {0: '{"query": "vllm"}'}
+
+    def test_tool_index_across_calls(self, gpt_oss_tokenizer, chat_request):
+        parser = HarmonyParser(gpt_oss_tokenizer)
+
+        first_delta = parser.parse_delta(
+            delta_text="",
+            delta_token_ids=encode_output(
+                "<|channel|>analysis<|message|>Thinking<|end|>"
+                "<|start|>assistant to=functions.get_weather<|channel|>commentary"
+                '<|constrain|>json<|message|>{"location": "Paris"}<|call|>'
+            ),
+            request=chat_request,
+            finished=False,
+        )
+        second_delta = parser.parse_delta(
+            delta_text="",
+            delta_token_ids=encode_output(
+                "<|start|>assistant to=functions.get_time<|channel|>commentary"
+                '<|constrain|>json<|message|>{"timezone": "UTC"}<|call|>'
+            ),
+            request=chat_request,
+            finished=False,
+        )
+
+        assert [tool.index for tool in tool_call_headers(first_delta)] == [0]
+        assert [tool.index for tool in tool_call_headers(second_delta)] == [1]
+        assert [tool.function.name for tool in tool_call_headers(second_delta)] == [
+            "get_time"
+        ]
+
+    def test_multi_tool_interleaved(self, gpt_oss_tokenizer, chat_request):
+        parser = HarmonyParser(gpt_oss_tokenizer)
+
+        first_delta = parser.parse_delta(
+            delta_text="",
+            delta_token_ids=encode_output(
+                "<|channel|>analysis<|message|>Plan<|end|>"
+                "<|start|>assistant to=functions.tool_a<|channel|>commentary"
+                '<|constrain|>json<|message|>{"a": 1}<|call|>'
+                "<|start|>assistant to=functions.tool_b<|channel|>commentary"
+                '<|constrain|>json<|message|>{"b": '
+            ),
+            request=chat_request,
+            finished=False,
+        )
+        second_delta = parser.parse_delta(
+            delta_text="",
+            delta_token_ids=encode_output("2"),
+            request=chat_request,
+            finished=False,
+        )
+        third_delta = parser.parse_delta(
+            delta_text="",
+            delta_token_ids=encode_output(
+                "}<|call|><|start|>assistant<|channel|>final<|message|>Done<|end|>"
+                "<|start|>assistant to=functions.tool_c<|channel|>commentary"
+                '<|constrain|>json<|message|>{"c": 3}'
+            ),
+            request=chat_request,
+            finished=False,
+        )
+
+        assert [tool.index for tool in tool_call_headers(first_delta)] == [0, 1]
+        assert combined_tool_arguments(first_delta) == {
+            0: '{"a": 1}',
+            1: '{"b": ',
+        }
+
+        assert second_delta is not None
+        assert [tool.index for tool in tool_call_payloads(second_delta)] == [1]
+        assert combined_tool_arguments(second_delta) == {1: "2"}
+
+        assert third_delta is not None
+        assert third_delta.content == "Done"
+        assert combined_tool_arguments(third_delta) == {
+            1: "}",
+            2: '{"c": 3}',
+        }
+        assert [tool.index for tool in tool_call_headers(third_delta)] == [2]
+
+
+class TestProcessChunk:
+    def test_empty(self, harmony_parser):
+        result = harmony_parser.process_chunk([])
+        assert result.segments == []
+        assert result.reasoning_token_count == 0
+
+    def test_single_channel(self, harmony_parser):
+        result = harmony_parser.process_chunk(
+            encode_output("<|channel|>final<|message|>Hello")
+        )
+
+        assert [
+            (s.channel, s.recipient, s.delta) for s in result.segments if s.delta
+        ] == [("final", None, "Hello")]
+
+    def test_cross_channel(self, harmony_parser):
+        result = harmony_parser.process_chunk(
+            encode_output(
+                "<|channel|>analysis<|message|>Think<|end|>"
+                "<|start|>assistant<|channel|>final<|message|>Answer"
+            )
+        )
+
+        assert [
+            (s.channel, s.recipient, s.delta) for s in result.segments if s.delta
+        ] == [
+            ("analysis", None, "Think"),
+            ("final", None, "Answer"),
+        ]
+
+    def test_multi_boundary(self, harmony_parser):
+        result = harmony_parser.process_chunk(
+            encode_output(
+                "<|channel|>analysis<|message|>One<|end|>"
+                "<|start|>assistant<|channel|>final<|message|>Two<|end|>"
+            )
+        )
+
+        boundary_segments = [
+            segment
+            for segment in result.segments
+            if segment.completed_message is not None
+        ]
+        assert [
+            (segment.completed_message.channel, get_text(segment.completed_message))
+            for segment in boundary_segments
+        ] == [
+            ("analysis", "One"),
+            ("final", "Two"),
+        ]
@@ -2,13 +2,31 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 import json
+import os

 import pytest

-from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
-from vllm.parser.abstract_parser import DelegatingParser
-from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser
-from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser
+_STRICT_TOOL_CALLING_ENV = "VLLM_ENFORCE_STRICT_TOOL_CALLING"
+_STRICT_TOOL_CALLING_ENV_VALUE = os.environ.get(_STRICT_TOOL_CALLING_ENV)
+os.environ[_STRICT_TOOL_CALLING_ENV] = "0"
+
+from vllm.entrypoints.openai.chat_completion.protocol import (  # noqa: E402
+    ChatCompletionRequest,
+)
+from vllm.parser.abstract_parser import DelegatingParser  # noqa: E402
+from vllm.reasoning.basic_parsers import (  # noqa: E402
+    BaseThinkingReasoningParser,
+)
+from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser  # noqa: E402
+
+
+@pytest.fixture(scope="module", autouse=True)
+def restore_strict_tool_calling_env():
+    yield
+    if _STRICT_TOOL_CALLING_ENV_VALUE is None:
+        os.environ.pop(_STRICT_TOOL_CALLING_ENV, None)
+    else:
+        os.environ[_STRICT_TOOL_CALLING_ENV] = _STRICT_TOOL_CALLING_ENV_VALUE


 class ThinkReasoningParser(BaseThinkingReasoningParser):
--- a/Show More
+++ b/Show More